[mxfp8 training] use triton kernel for dim0 scaling with RCEIL rounding

[danielvegamyhre]

Stacked PRs:

• [mxfp8 training] fix bug in cast_bench.py for to_mx + RCEIL #3569
• ->[mxfp8 training] use triton kernel for dim0 scaling with RCEIL rounding #3568

---

Summary

• For MXLinear, use triton dim0 quantization kernel for RCEIL scaling mode until torch.compile is faster (Use cvt.rp.satfinite.ue8m0x2.f32 PTX instruction in Inductor codegen for mxfp8 quantization pytorch#170635)
• Update validation so tests pass for RCEIL

Tests

• pytest test/prototype/mx_formats/test_mx_linear.py

Benchmarks (e2e linear sigmoid fwd + bwd)

Command: python benchmarks/float8/float8_roofline.py --mx_recipe_name="mxfp8_cublas_rceil" --outfile "outfile.csv"

Before:

   fwd_M  fwd_K  fwd_N  r_bf16_gemm_s  r_fp8_gemm_s  r_fp8_ovhd_s  ...  b_fp8_gemm_s  b_bf16_e2e_s  b_fp8_e2e_s  b_fp8_e2e_spdp  rb_bf16_gemm_ratio  rb_fp8_gemm_ratio
0   1024   1024   1024       6.00e-06      6.00e-06      1.20e-05  ...      1.77e-05      1.49e-05     4.87e-05            0.31                0.41               0.34
1   2048   2048   2048       2.94e-05      1.47e-05      1.20e-05  ...      3.52e-05      4.56e-05     9.29e-05            0.49                0.65               0.42
2   4096   4096   4096       2.35e-04      1.17e-04      4.10e-05  ...      1.37e-04      2.80e-04     3.12e-04            0.90                0.82               0.85
3   8192   8192   8192       1.88e-03      9.40e-04      1.64e-04  ...      8.34e-04      1.92e-03     1.56e-03            1.23                0.97               1.13
4  16384  16384  16384       1.50e-02      7.52e-03      6.56e-04  ...      9.34e-03      1.65e-02     1.30e-02            1.27                0.96               0.81

After:

fwd_M  fwd_K  fwd_N  r_bf16_gemm_s  r_fp8_gemm_s  r_fp8_ovhd_s  ...  b_fp8_gemm_s  b_bf16_e2e_s  b_fp8_e2e_s  b_fp8_e2e_spdp  rb_bf16_gemm_ratio  rb_fp8_gemm_ratio
0   1024   1024   1024       6.00e-06      6.00e-06      1.20e-05  ...      1.75e-05      1.49e-05     4.49e-05            0.33                0.41               0.34
1   2048   2048   2048       2.94e-05      1.47e-05      1.20e-05  ...      3.53e-05      4.54e-05     7.55e-05            0.60                0.65               0.42
2   4096   4096   4096       2.35e-04      1.17e-04      4.10e-05  ...      1.37e-04      2.79e-04     2.44e-04            1.14                0.82               0.86
3   8192   8192   8192       1.88e-03      9.40e-04      1.64e-04  ...      8.35e-04      1.91e-03     1.29e-03            1.48                0.97               1.13
4  16384  16384  16384       1.50e-02      7.52e-03      6.56e-04  ...      9.22e-03      1.72e-02     1.19e-02            1.45                0.98               0.82

Kernel microbenchmarks

• to_mx + torch.compile = ~3.7 TB/s
• triton = ~5.8 TB/s
  See: https://github.com/pytorch/ao/tree/main/torchao/prototype/mx_formats#quantization-kernel-microbenchmarks

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• B200 runners are down due to network issues

❌ 4 New Failures

As of commit cf0acd8 with merge base 8d47813 ():

NEW FAILURES - The following jobs have failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  RuntimeError: Command docker exec -t b2cfad6725a72eeb0631781c3714920b14d2cca99235deeefbc1462de46840d8 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.8, linux.g5.12xlarge.nvidia.gpu, torch==2.8.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t eec334f34ee9ebfd130a66b9ce9adcb38f669c24307bdef8c318c27ee93ce0d3 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.9, linux.g5.12xlarge.nvidia.gpu, torch==2.9.1, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t ab2cf947195076d0e0a197bfc7d907a6593f95c7aba1dd0d1fc55e5c10900d85 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t 57f622c1e9e9025262d6dfd78129cf6a676cccb4bbda9b51f1d84f21fe6f7086 /exec failed with exit code 1

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

[vkuzo]

what do benchmarks look like

[danielvegamyhre]

what do benchmarks look like

Updated PR description with benchmarks, they are better

[vkuzo]

Updated PR description with benchmarks, they are better

that looks good for no-fusion, do we have metrics on what this looks like with layer fusion? The script has a flag to enable it, or you could use LLaMa 3 8B in torchtitan.

[vkuzo]

IMO we should ensure the codebase stays consistent with how the various config flags are used. Before this PR, mxfp8_cast_kernel_choice controls the dim1 cast. I think good options are:

1. make mxfp8_cast_kernel_choice control dim1 and dim0 casts (update names/docs as needed)
2. or, make a separate mxfp8_cast_kernel_choice thing for dim0, update names as needed
3. or, just delete all the configs and make the logic use custom kernels

The current PR is a bit confusing given the rest of the feature because it leaves the dim1 config as is but defaults the dim0 logic to a custom kernel.

I don't really like Option 3 because we should be able to quickly test compile, so Option 1 or Option 2 would sound good to me.

[vkuzo]

this isn't how the dim1 config works, let's make sure all of these flags operate in a consistent way. You can change the dim1 logic if needed (everything is prototype).