integrate torch._scaled_mm into Float8BlockwiseLinear and add bench script

[danielvegamyhre]

Stacked PRs:

• [fp8 blockwise] wrap triton quantization kernels in custom ops for torch.compile compatibility #2829
• [fp8 blockwise] load 2d chunks for groupwise quant to enable coalesced gmem accesses #2827
• use shared bench + profile utils in blockwise fwd bwd bench script #2826
• ->integrate torch._scaled_mm into Float8BlockwiseLinear and add bench script #2785

---

integrate torch._scaled_mm into Float8BlockwiseLinear and add bench script

Summary

• Add use_triton flag to Float8BlockwiseLinear` that defaults to False. When True, use triton for gemms. When False, use torch._scaled_mm for gemms.
• Compute scales in column major format, for compatibility with torch._scaled_mm.
• Add some autotuner configs for IO bounds dynamic quantization kernels.
• Rename triton kernels to be prefixed with "triton_" for clarity
• Write out reciprocal scales directly from triton quant kernels, rather than writing regular scales then doing 1.0/scales to pass in reciprocals to GEMM for rescaling outputs. This is necessary because in some cases the division op (1.0/scales) changes the memory layout from col major to row major, which triggers assertions on mem layout

Benchmarks

Benchmarking an eager linear forward + backward pass with llama4 shapes, using torch._scaled_mm improves perf by ~2x but still slower than bf16. Next step is to look at the trace to see which quantization kernels are slowing things down.

    M     N     K  out_dtype         bf16_mm_linear_us    fp8_triton_linear_us    fp8_scaled_mm_linear_us
-----  ----  ----  --------------  -------------------  ----------------------  -------------------------
16640  5120  8192  torch.bfloat16              5349.79                 12643.5                    7746.98
16640  8192  5120  torch.bfloat16              6100.62                 12912.6                    7783.31

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (3 Unrelated Failures)

As of commit 3c78e2a with merge base 253d65a ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh) (trunk failure)
  test/integration/test_integration.py::TestAutoQuant::test_autoquant_hp_float
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh) (trunk failure)
  test/integration/test_integration.py::TestAutoQuant::test_autoquant_hp_float
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/integration/test_integration.py::TestAutoQuant::test_autoquant_hp_float

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

[drisspg]

Do we actually need these triton kernels? It seems like scaled_mm is pretty universally, is there a valid case for the small shapes you mentioned previously

I guess the other valid argument is that some people aren't using new enough triton to have this fp8 mm support

[danielvegamyhre]

Do we actually need these triton kernels? It seems like scaled_mm is pretty universally, is there a valid case for the small shapes you mentioned previously

I guess the other valid argument is that some people aren't using new enough triton to have this fp8 mm support

No, scaled_mm is definitely better, I just wanted to keep the triton gemms so as a side-project for learning I can try optimizing them to match scaled_mm using additional features like TMA, etc.

[drisspg]

!

[danielvegamyhre]

Yeah this shouldn't be commented out, this was mistakenly left in a previous PR. So fixing it here.

[drisspg]

we shouldnt be printing though.. maybe logging but forsure not printing

[danielvegamyhre]

Agreed, I had this as logging originally but changed to print when I noticed the logs didn't show up in torchtitan training runs when I was trying to get ScaledGroupedMMTensor to work e2e, and wanted to focus on the task at hand rather than debug whatever torchtitan is doing with the logging module.

I switched back to logging statements in this PR: #2835

[danielvegamyhre]

@drisspg I landed #2835 and rebased on top of main, so the print statements are logging now.