[mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] bench and profile mxfp8 a2a fwd and bwd separately #3203
• [mxfp8 moe training] simplify e8m0 -> fp32 calc #3201
• ->[mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all #3197
• [mxfp8 moe training] add triton kernel for mxfp8 dequantization #3195

---

[mxfp8 moe training] integrate triton quant/dequant kernels into mxfp8 all to all

Test plan

• pytest test/prototype/moe_training/mxfp8/test_mxfp8_a2a.py -k ToMXFP8AllToAllVDequantTest -s

Benchmarks

Flat perf

input_shape         num_splits    fwd_bf16_ms    fwd_mxfp8_ms    bwd_bf16_ms    bwd_mxfp8_ms
----------------  ------------  -------------  --------------  -------------  --------------
(16, 8192, 5120)             4        2.38517         2.59848         7.5735          7.7693

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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