[scaled grouped mm] integrate triton kernels into differentiable scaled grouped mm

[danielvegamyhre]

Prior PR in stack: #2064

Summary

• Integrate triton kernels from [scaled grouped mm] add triton kernels for float8 rowwise quantization with per-group/jagged scales #2064 into differentiable scaled grouped mm
• Add benchmarking script for e2e fwd+bwd of differentiable scaled grouped mm and share results
• Move old CPU loop based implementations of per-group scaling to test utils folder, since they're only needed for verifying numerical accuracy now.

Performance

TL;DR there is ~1.25x - 37x speedup using the triton kernels (most shapes speedups fell between 2x-6x).

CPU loop:

A_shape        B_shape           high_precision_dtype      time_us
-------------  ----------------  ----------------------  ---------
(256, 4096)    (4, 4096, 4096)   torch.bfloat16            4324.35
(256, 4096)    (8, 4096, 4096)   torch.bfloat16            8197.2
(256, 4096)    (16, 4096, 4096)  torch.bfloat16           15830.9
(4096, 4096)   (4, 4096, 4096)   torch.bfloat16            5211.11
(4096, 4096)   (8, 4096, 4096)   torch.bfloat16            9003.3
(4096, 4096)   (16, 4096, 4096)  torch.bfloat16           16720.7
(65536, 4096)  (4, 4096, 4096)   torch.bfloat16           31257
(65536, 4096)  (8, 4096, 4096)   torch.bfloat16           34253.2
(65536, 4096)  (16, 4096, 4096)  torch.bfloat16           40141.3

Triton kernels:

A_shape        B_shape           high_precision_dtype      time_us
-------------  ----------------  ----------------------  ---------
(256, 4096)    (4, 4096, 4096)   torch.bfloat16            835.657
(256, 4096)    (8, 4096, 4096)   torch.bfloat16            835.657
(256, 4096)    (16, 4096, 4096)  torch.bfloat16            832.382
(4096, 4096)   (4, 4096, 4096)   torch.bfloat16            830.429
(4096, 4096)   (8, 4096, 4096)   torch.bfloat16           4666.3
(4096, 4096)   (16, 4096, 4096)  torch.bfloat16           7502.13
(65536, 4096)  (4, 4096, 4096)   torch.bfloat16            840.335
(65536, 4096)  (8, 4096, 4096)   torch.bfloat16          26359.4
(65536, 4096)  (16, 4096, 4096)  torch.bfloat16          27910.4

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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! (6 Unrelated Failures)

As of commit 6821e44 with merge base 9af2a45 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (similar failure)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2

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

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

• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh) (trunk failure)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• 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) (trunk failure)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh) (trunk failure)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2

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

[danielvegamyhre]

Dr CI confirmed test failures are unrelated to this change (and I manually confirmed as well, they are QAT related)