[mxfp8 moe training] default to triton kernel for dim0 cast

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] _to_mxfp8_then_scaled_grouped_mm wrapper that accepts keyword args #3561
• ->[mxfp8 moe training] default to triton kernel for dim0 cast #3560

---

[mxfp8 moe training] default to triton kernel for dim0 cast

torch.compile is still slow for RCEIL (see pytorch/pytorch#170635). Since we are migrating to RCEIL as the default for MXFP8 training, we should default to the triton dim0 cast kernel which has decent performance (~6k TB on 1000W B200).

(torch) dev@gpu-dev-8f27b069:~/ao$ PYTHONPATH=/home/dev/ao:$PYTHONPATH python benchmarks/mx_formats/cast_bench.py  --mode dim0_mxfp8_triton_rceil
W0102 04:28:08.388000 67699 site-packages/torch/_library/triton.py:222] _dequant_mxfp8_kernel not in collector.assignments
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.11.0.dev20260101+cu128
triton version: 3.6.0
mode: dim0_mxfp8_triton_rceil
time_us 136.25599443912506
mem_bw_gbps 5971.810483270397

[pytorch-bot]

🔗 Helpful Links

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

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

✅ No Failures

As of commit 3ac311d with merge base 8bb433e ():
💚 Looks good so far! There are no failures yet. 💚

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