[scaled grouped mm] add triton kernels for float8 rowwise quantization with per-group/jagged scales

[danielvegamyhre]

Next PR in stack: #2077

Motivation

• In the first iteration of the differentiable scaled grouped mm prototype, to perform float8 quantization with per-group/"jagged" scales, we naively iterate through the offsets on the CPU in for loop, launching separate scaling kernels for each subtensor (see here). This approach requires host-device sync, to loop over the values via offs.tolist() on the CPU, which negatively impacts performance.
• Solution: do all of this in a single triton kernel to avoid host-device sync and improve performance.

Summary

This PR adds 2 Triton kernels:

• _triton_fp8_row_major_jagged_rowwise_scales: performs dynamic float8 quantization with row-wise scales for each group (as determined by the offsets), and writes the output in row-major memory layout. This kernel is to be used for quantizing left operand of the scaled grouped mm in the backward pass of the differentiable scaled grouped mm, to compute grad_weight (see here).
• _triton_fp8_col_major_jagged_colwise_scales: performs dynamic float8 quantization with row-wise scales computed over logical columns for each group (as determined by the offsets), and writes the output in col-major memory layout. This kernel is to be used for quantization of the right operand of the scaled grouped mm in the backward pass of the differentiable scaled grouped mm, to compute grad_weight (see here).

Test plan

• Added unit tests confirming triton kernels produce same outputs as corresponding pure pytorch implementations.

Next steps

• Measure e2e differentiable scaled grouped mm perf with vs without these changes
• Integrate the triton kernels into the autograd func
• Do any additional perf improvements as necessary
• Add JaggedFloat8Tensor subclass w/ override for torch._grouped_mm (bf16 op)
• Add user-facing prototype API for converting MoE layers to use float8 training

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [Infra] Jobs got frequently cancelled, sometimes mid-checkout

❌ 4 New Failures, 4 Pending

As of commit d818d51 with merge base f788897 ():

NEW FAILURES - The following jobs have failed:

• 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)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• 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)
  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)
  test/quantization/test_qat.py::TestQAT::test_qat_8da4w_prepare_vs_convert_2
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  test_identical_to_IntXQuantizationAwareTrainingConfig_{'weight_dtype': torch.int8, 'granularity': PerGroup(group_size=32), 'mapping_type': <MappingType.SYMMETRIC: 1>, 'scale_dtype': torch.bfloat16, 'model_dtype': torch.bfloat16}

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

[danielvegamyhre]

cc @drisspg @ngimel I'd appreciate your thoughts/reviews on this when you have time

[danielvegamyhre]

also cc @vkuzo for visibility upon return from recharge

[danielvegamyhre]

confirmed test failures are unrelated to this change