[fp8 blockwise] load 2d chunks for groupwise quant to enable coalesced gmem accesses

[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

---

[fp8 blockwise] load 2d chunks for groupwise quant to enable coalesced gmem accesses

NCU flagged uncoalesced gmem accesses when loading input tensor data. I was loading in skinny tiles of size (1 x blocksize) or (blocksize x 1), which can cause uncoalesced loads and stores depending on the underlying memory layout.

By loading 2d chunks of (num_groups, blocksize) or (blocksize, num_groups) we can enable the uncoalesced loads to be grouped into several coalesced loads.

This results in a 20-33% speedup.

Kernel performance

• NCU reports all quantization kernels used for blockwise linear running at 80-90% peak memory bandwidth
• Runtime numbers TBD, we don't have quantization kernel specific benchmark scripts yet: [fp8 blockwise training] add benchmarking scripts comparing triton quantization kernels vs torch.compile #2828

Individual linear fwd + bwd

bf16: 7690 us

fp8 blockwise: 6302 us (22% speedup)

E2E training performance

Eager llama3 8b with FDSP=4 on H100s

• bf16: ~6080 TPS
• fp8 blockwise: ~6170 TPS

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 ba52b69 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.