[mxfp8] support RCEIL in triton_to_mxfp8_dim0 kernel with inline PTX

[danielvegamyhre]

Summary

• Support RCEIL in triton_to_mxfp8_dim0 kernel with inline PTX, for consistency with other kernels. We are moving toward RCEIL as the sensible default for accuracy + perf.
• After this change, for MoE training code, all quantization ops have an option that supports RCEIL.

Reference:

ao/torchao/csrc/cuda/mx_kernels/mxfp8_quantize.cuh

Line 211 in 8555713

__device__ __forceinline__ e8m0_t float_to_e8m0(float val) {

Tests

• pytest test/prototype/mx_formats/test_kernels.py -k dim0
• pytest pytest test/prototype/mx_formats/test_mx_linear.py

Benchmarks

[danvm@devgpu031.atn1 ~/ao (rceil2)]$ cd /home/danvm/ao && PYTHONPATH=/home/danvm/ao:$PYTHONPATH conda run -n torch2 CUDA_VISIBLE_DEVICES=7 python benchmarks/mx_formats/cast_bench.py --mode dim0_mxfp8_triton_rceil
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.11.0.dev20251216+cu128
triton version: 3.6.0
mode: dim0_mxfp8_triton_rceil
time_us 139.45600390434265
mem_bw_gbps 5834.779093183679

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit dc219c2 with merge base b9e5780 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_fsdp2

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

[danielvegamyhre]

@drisspg CI is green now if you want to take another look (only failing test is unrelated).

I had to require sm100 for some tests, namely triton dim0 and dim1 kernels, for both rceil and floor, since they share the "calculate scale" triton jit func and just compiling it at all with the inline PTX requires sm100 now