FP8 splitgemm user defined triton kernel

[msaroufim]

$ python test_fp8.py 
..
----------------------------------------------------------------------
Ran 2 tests in 8.207s

This is borrowing the kernel from https://github.com/pytorch-labs/applied-ai/blob/main/kernels/triton/inference/fp8/splitk_gemm_fp8.py

1. Added a test case for numerics
2. Added a test for user defined triton functions to guarantee this does not graph break
3. I had to reduce block sizes otherwise I was running into OOMs on H100
4. There's still a bunch we could do with an autotuner
5. Baseline should also be how much of this we can codgenerate via compile

We had an issue with testing this in our CI but it works fine locally and in pytorch CI pytorch/pytorch#126982

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/263

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 113ecde with merge base 5e28109 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[msaroufim]

This is passing locally on nightlies but failing CI, I'll check what's up - I suspect it's some caching issue with compile

[vkuzo]

this is interesting, wondering if we have benchmarks / accuracy data on how this compares to cuBLAS float8 gemm?

[msaroufim]

I don't believe the announcement baselined vs cuBLAS but perhaps @AdnanHoque can shed some more detail

[AdnanHoque]

Hey! Yes we did compare to cuBLAS and got gain in the small M <= 16 regime.

Here were the speedups (all NCU timings) with the original tile sizes, on H100-SXM: