[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.

✅ No Failures

As of commit 527317f with merge base cd21d0e ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

should the comment say scale or max_abs? also row-wise should really say "across dim0"?

[danielvegamyhre]

the dim1 kernel said "column-wise" here, so i just changed it to say "row-wise" here for consistency. dim0/dim1 is also fine. The comment shouldn't "maximum absolute values" for sure though, I will remove that from both kernels.

[vkuzo]

should this either have a mask or a TODO to add masking?

[danielvegamyhre]

I added masking in #3131 where I refactored/simplified the kernel, but I'll go ahead and move it here

[drisspg]

remove please

[danielvegamyhre]

Woops, was debugging an import error, thanks

[drisspg]

is this storing in swizzled layout?

[danielvegamyhre]

No, we have separate kernels for that. Low hanging fruit could be to write directly to blocked swizzled format here, like Cursor mentioned they're doing in this blog https://cursor.com/blog/kernels

[drisspg]

Its a non zero perf win for activation casting in inference. We can leave as a follow up but should be pretty easy to just have aconstexpr path do this