wmma: refactor to remove wmma_func and create TC funcs as needed (#3945)

* wmma: refactor to remove wmma_func and create TC funcs as needed

* test_linearizer: disable bf16 CUDA during emulation testing

* cstyle: clean up creation of CUDA vec dtypes

* extra/gemm: add option to accumulate to bfloat16

* cleanups

* benchmark: add CUDA bfloat16 matmul

* more cleanups
This commit is contained in:
Francis Lam
2024-03-27 13:43:09 -07:00
committed by GitHub
parent 88b24df40a
commit 7c5729a3bd
9 changed files with 101 additions and 92 deletions

View File

@@ -108,7 +108,9 @@ jobs:
- name: Test speed vs torch
run: CUDA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
- name: Run Tensor Core GEMM
run: CUDA=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
run: |
CUDA=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
CUDA=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
- name: Run LLaMA
run: |
CUDA=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -140,6 +142,7 @@ jobs:
onnx_inference_speed.csv
torch_speed.txt
matmul.txt
matmul_bfloat16.txt
llama_unjitted.txt
llama_jitted.txt
llama_beam.txt