[mxfp8 moe training] add triton kernel for mxfp8 quantization along dim0

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] use mxfp8 dim0 kernel in mxfp8 all2all; simplify mxfp8 dim0 kernel #3131
• [mxfp8 moe training] integrate mxfp8 dim0 triton kernel #3129
• ->[mxfp8 moe training] add triton kernel for mxfp8 quantization along dim0 #3128

---

[mxfp8 moe training] add triton kernel for mxfp8 quantization along dim0

Summary

• torch.compile codegen has had on and off perf regressions during the development process of mxfp8 moe training, it would be nice to have a simple triton kernel to have consistently good perf we can use

Test plan

• pytest test/prototype/mx_formats/test_kernels.py

Benchmarks

existing torch.compile/to_mx() benchmark:

(torch) [danvm@devgpu031.atn1 ~/ao/benchmarks (main)]$ CUDA_VISIBLE_DEVICES=7 python mx_formats/cast_bench.py --mode dim0_mxfp8_floor
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.10.0.dev20251008+cu128
triton version: 3.5.0
mode: dim0_mxfp8_floor
time_us 156.5759927034378
mem_bw_gbps 5196.805474139168

new triton dim0 mxfp8 kernel (~10% higher peak memory bandwidth utilization):

(torch) [danvm@devgpu031.atn1 ~/ao/benchmarks (danielvegamyhre/stack/76)]$ CUDA_VISIBLE_DEVICES=7 python mx_formats/cast_bench.py --mode dim0_mxfp8_triton_floor
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.10.0.dev20251008+cu128
triton version: 3.5.0
mode: dim0_mxfp8_triton_floor
time_us 140.35199582576752
mem_bw_gbps 5797.530496182742

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 5 New Failures

As of commit c62b0f0 with merge base cd21d0e ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU 2.6, linux.4xlarge, torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t c96ecadf0c43dc67d99a51737b381741183ade6062ad15b425df36c0fbb73602 /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t a6e7b0752094426b3272b2ae820ffe19a9a9f8eb3793fb4a2d84ed2a115c9778 /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 0e64c7f7ebe852385850f113b34dff0644da1be8c8a2dbcee703b4af6b5e9278 /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 5b31142f155f74692a0aebaada8bca1cb6490164fb2edf868efb5486a347341d /exec failed with exit code 2
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t f9ca3a89843fab03623d1bea76678d18fb02c0f8ff1a32a056a09a8ff66ac9e6 /exec failed with exit code 2

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