[mxfp8 moe training] add benchmark for e2e mxfp8 EP pipeline

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] add torch.compile support #3606
• ->[mxfp8 moe training] add benchmark for e2e mxfp8 EP pipeline #3585
• [mxfp8 moe training] add e2e integration test of full mxfp8 EP pipeline #3584
• [mxfp8 moe training] handle pre-quantized inputs/grads in forward/backward of _MXFP8GroupedMM autograd func #3583
• [mxfp8 moe training] mxfp8 a2a_combine autograd function #3582
• [mxfp8 moe training] mxfp8 unpermute autograd function #3581
• [mxfp8 moe training] mxfp8 token permute autograd func + triton kernels #3580
• [mxfp8 moe training] mxfp8 a2a_dispatch autograd function #3579

---

[mxfp8 moe training] add benchmark for e2e mxfp8 EP pipeline

This stack creates a set of differentiable MXFP8 expert parallel primitives for MoE training:

• a2a_dispatch
• permute
• unpermute
• a2a_combine

The key idea is instead of waiting to quantize to MXFP8 until directly before the grouped GEMMs, we can quantize earlier, before the a2a collectives, to speed up these exposed comms. This requires staying in MXFP8 through the token permutation, then feeding those permuted fp8 data + scales into the MXFP8 grouped GEMM, which produces bf16 outputs.

We do this in both the forward pass and backward pass.

Design

Each color corresponds to a separate autograd function:

Benchmarks

This benchmark measures the entire forward + backward pass of the full EP pipeline:

Forward:

• a2a_dispatch
• permute
• MXFP8 grouped GEMM
• unpermute
• a2a_combine

Backward:

• a2a_combine backward()
• unpermute backward()
• MXFP8 grouped GEMM backward()
• permute backward()
• a2a_dispatch backward()

For DeepSeekV3 shapes we see:

• 1.4x speedup over vanilla bf16 pipeline (with torch._grouped_mm)
• 1.22x speedup over existing torchao pipeline (with bf16 all2alls, mxfp8 grouped GEMMs)

Note this is WITHOUT the mxfp8 grouped GEMM improvements landing in pytorch core soon, so we expect even higher speedups then! We have lower speedups than usual right now due to improvments landing to bf16 grouped gemm but none to mxfp8 grouped gemm yet.

Versus vanilla bf16 pipeline:

Expert Parallelism Pipeline Benchmark Results
World Size: 2
========================================================================================================================
+----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+
|   tokens |   dim |   hidden_dim |   num_experts |   fwd_bf16_ms |   fwd_mxfp8_ms | fwd_speedup   |   bwd_bf16_ms |   bwd_mxfp8_ms | bwd_speedup   | total_speedup   |
+==========+=======+==============+===============+===============+================+===============+===============+================+===============+=================+
|   131072 |  8192 |         5120 |             8 |        17.537 |         11.7   | 1.50x         |        29.303 |         21.563 | 1.36x         | 1.41x           |
+----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+
|   131072 |  7168 |         2048 |             8 |         9.285 |          6.387 | 1.45x         |        13.718 |         10.004 | 1.37x         | 1.40x           |
+----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+

Versus existing torchao strategy (bf16 all2alls, quantize to mxfp8 directly before grouped gemms)

========================================================================================================================
Expert Parallelism Pipeline Benchmark Results
World Size: 2
========================================================================================================================
+----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+
|   tokens |   dim |   hidden_dim |   num_experts |   fwd_bf16_ms |   fwd_mxfp8_ms | fwd_speedup   |   bwd_bf16_ms |   bwd_mxfp8_ms | bwd_speedup   | total_speedup   |
+==========+=======+==============+===============+===============+================+===============+===============+================+===============+=================+
|   131072 |  8192 |         5120 |             8 |        13.472 |         11.947 | 1.13x         |        25.453 |         21.891 | 1.16x         | 1.15x           |
+----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+
|   131072 |  7168 |         2048 |             8 |         7.903 |          6.388 | 1.24x         |        12.278 |         10.021 | 1.23x         | 1.23x           |
+----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+

[pytorch-bot]

🔗 Helpful Links

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

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

[vkuzo]

didn't read closely, stamp for prototype