Add Float8BlockwiseLinear with Triton kernels for quantization and GEMMs

[danielvegamyhre]

Stacked PRs:

• ->Add Float8BlockwiseLinear with Triton kernels for quantization and GEMMs #2592

---

Summary

• Add "Float8BlockwiseLinear" and make it differentiable with autograd func to support training
• Add Triton kernels for (1) activation quant, (2) weight quant, and (3) GEMM based on these DeepGemm inference kernel and rename GEMM to explicitly include expected scaling granularity of operands in the function names (fp8_gemm => blockwise_fp8_gemm_1x128_128x128).
• Add new Triton kernel need for backward: blockwise_fp8_gemm_1x128_1x128_kernel for dW calculation where both left and right operands have activation scaling granularity (1 x block_size). This is a modified version of the kernel above, so it accepts 1x128 scaling for both operands.
• GEMM kernels do accumulation in fp32 and cast output to bfloat16.
• Modify all quantization kernels to use EPS to guard against NaNs upon division by 0.
• Add tests verifying numerics by enforcing reasonable SQNR
• Added benchmarking script for comparing Triton kernels vs FBGEMM vs DeepGEMM quantization kernels.

Why use Triton kernels instead of DeepGEMM cutlass kernels?

The GEMM APIs in @vkuzo's PoC here no longer exist in DeepGemm. I tried using the new GEMM APIs (fp8_gemm_nt etc), and:

• On B200, with both Vasiliy's PR and my PR, got device-side asserts on this line, that were not immediately clear how to resolve.
• On H100, I only tried Vasiliy's PR, but got undefined symbols error from CUDA, despite using CUDA toolkit 12.8+ as stated in the readme.

Since our only goal is a functional skeleton and not performance, rather than spend more time on this, I just used the existing Triton kernels we had and made a modified GEMM (1 line change) to support blockwise_fp8_gemm_1x128_1x128_kernel.

If we want to replace these Triton GEMMs with the Cutlass ones later to see if perf is better (it probably is), we can do that.

Note on numerics

Interestingly, the reference DeepGemm triton quantization kernels do NOT use EPS/clamping to prevent division by 0. This resulted in my unit tests passing (where inputs were from a normal distributed), but NaNs occuring in TorchTitan training runs, where actual activation values sometimes had amax of 0.

I updated the kernels to use the same EPS guards as torchao.float8, and this fixed the Nans.

Test plan

• pytest test/prototype/blockwise_fp8/test_blockwise_kernels.py
• pytest test/prototype/blockwise_fp8/test_blockwise_linear.py

Torchtitan PoC integration results

• Logs showing converted model, nice decreasing loss, but very low throughput (~5.1k vs ~9.8k for bfloat16).
• I think perf may be bad because of (1) all the eager .t().contiguous() type of transformations done, and (2) the DeepGemm triton GEMM requires the B tensor be in row-major, then does strided reads to get it in col-major in SMEM to do the fp8 wgmma? Same thing as with fp8 flex attention I think..

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 4 New Failures

As of commit d0631c0 with merge base 0e00df3 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  RuntimeError: Command docker exec -t e503a5df82188fe6380fff510f60214b6ce0f70cab538b858498530f6382800b /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 72789b784af213b122e26e814459fef0f9df88b2557f2346cf240ae5d1c1dc06 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 67bc12085d9e5af4d856f7a55f60be2656d8cd37092355003b732f4314ed6ebb /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t c3bf9f49cfae227ab80a717efcf9c0b5133a122283515633cff36da0777f2f2c /exec failed with exit code 1

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

[danielvegamyhre]

cc @vkuzo @drisspg for review

[drisspg]

Can you use sqnr everywhere match w/ existing numerics testing

[danielvegamyhre]

Updated to use SQNR

[drisspg]

since we have an optional runtime dependency on fbgemm can we just call their kernel directly?

[danielvegamyhre]

Yes that is the desired end state. For now I have tried and have had repeated problems getting it to work so far (fbgemm-gpu-genai), e.g. undefined symbols. Tried on both H100 and B200 and got different undefined symbol errors

[drisspg]

(32768, 32768)  (1,128)          5732.42      7960.21        7830.56
(32768, 32768)  (128,128)       13692.4        669.664       7831.14

this number is kinda weird to me, do you have memory bandwidth calcs? I dont immediately get why there is a 10x delta in group wise vs blockwise

[danielvegamyhre]

this number is kinda weird to me, do you have memory bandwidth calcs? I dont immediately get why there is a 10x delta in group wise vs blockwise

Yeah I agree it's odd, will try adding some mem bw calcs, was thinking about checking with Josh / fbgemm team as well if perhapst here is a different kernel they use for activation quant.

[danielvegamyhre]

@drisspg @vkuzo this is ready for another look, numerics look good and torchtitan loss curve looks good. perf is bad for llama3 8b right now.

(accidentally squashed my stack into a single pr with stack-pr, sorry for the large PR, I can break it up if necessary)