improve fp8 blockwise gemm perf

[danielvegamyhre]

Stacked PRs:

• integrate torch._scaled_mm into Float8BlockwiseLinear and add bench script #2785
• ->improve fp8 blockwise gemm perf #2784

---

improve fp8 blockwise gemm perf

Summary

• Add benchmark script comparing triton vs torch._scaled_mm for blockwise_1x128_128x128 fp8 gemm
• Minor autotuner config changes to triton gemm
• scaled_mm is slower for very small shapes (e.g., dim=256) but ~6x faster for llama4 shapes (dim=5120)! Next step is to migrate torchao blockwise autograd func to use scaled_mm.

fp8 blockwise 1x128_128x128 gemm

• This is the gemm used for:
  • output = input @ weight.t()
  • grad_input = grad_output @ weight
• Benching shapes representative of w1/w3 and w2 in llama4

    M     N     K  out_dtype         bf16_mm_us    fp8_triton_us    fp8_scaled_mm_us    bf16 tflops/sec    triton tflops/sec    scaled_mm tflops/sec
-----  ----  ----  --------------  ------------  ---------------  ------------------  -----------------  -------------------  ----------------------
16640  5120  8192  torch.bfloat16       2294.62          3494.42             1383.54            608.319              399.456                 1008.91
16640  8192  5120  torch.bfloat16       2322.43          3800.48             1392.75            601.036              367.286                 1002.23

fp8 blockwise 1x128_128x1 gemm

• This is the gemm used for:
  • grad_weight = grad_output_t @ input
• Benching shapes representative of w1/w3 and w2 in llama4

    M     N     K  out_dtype         bf16_mm_us    fp8_triton_us    fp8_scaled_mm_us    bf16 tflops/sec    triton tflops/sec    scaled_mm tflops/sec
-----  ----  ----  --------------  ------------  ---------------  ------------------  -----------------  -------------------  ----------------------
16640  5120  8192  torch.bfloat16        2333.1          2572.58             1726.98            598.286              542.594                 808.271
16640  8192  5120  torch.bfloat16        2290.9          2754.64             1729.66            609.309              506.732                 807.015

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Cancelled Job

As of commit c288547 with merge base 1526dfe ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  Process completed with exit code 1.

CANCELLED JOB - The following job was cancelled. Please retry:

• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  ##[error]The operation was canceled.

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

[danielvegamyhre]

@vkuzo I'm OOO Monday and Tuesday, let's cancel our meeting this week. (I'm not sure if meetings are autodeclined for PTO so I wanted to reach out to make sure)

[danielvegamyhre]

confirmed test failure is unrelated to change and talked with feature owner about fixing