[FEAT] Add bf16 gemm

[vadiklyutiy]

FI has fp8 and fp4 gemm implementations. But there is no bf16 one.

Original issue was found in vllm and described in vllm-project/vllm#27173.

In short, torch.nn.functional.linear is not optimal for small batch sizes. Torch team said that they just call cuBLAS.

It makes sense to support bp16 gemm and do tuning through cuBLAS, cutlass, cuDNN, internal FI implementation as it done for fp8 and fp4 cases.

Performance result and script for measurement are in vllm-project/vllm#27173.

[raayandhar]

I'm happy to try and tackle this as OS contributor, at least using the CUTLASS backend, as I already see for FP8 and FP4 it does not seem very difficult. I am GPU poor so if there is significant tuning/etc this may not be appropriate task for me.

Also, my understanding is that there is TGV GEMM (which issue poster uses as well), for fp16/bf16 inp/out, so it seems like fp16 gemm is already supported for SM100, unless I am misunderstanding? Also, in the script I see that it is using bf16, not fp16. So maybe issue-poster means bf16 instead?

[vadiklyutiy]

that there is TGV GEMM

yes, I used TGV Gemm for comparison. I don't know how good cutlass or cudnn for this case.

I see that it is using bf16, not fp16

Only bf16 used in LLM (default city syndrome :)).
I changed title and issue description.

[raayandhar]

Going to give bf16 cutlass GEMM a shot, see how it will do at small batch sizes

[vadiklyutiy]

Just wondering, is there a progress with this issue?

[raayandhar]

Just wondering, is there a progress with this issue?

Hi, yes, sorry was a bit busy. I have a PR out #2070 that is 90% done, I'm planning on wrapping it up this weekend and getting the numbers for you.

[raayandhar]

@vadiklyutiy now the PR is awaiting review. The numbers at batch=64 are better than all the others, but elsewhere we are slightly worse. I'm a total CUTLASS newbie, so it might be possible to get even better performance. The results and my thoughts are at my PR. Let me know what you think.