[mxfp8 moe training] mxfp8 a2a_combine autograd function

[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] mxfp8 a2a_combine autograd function

Tests

• There are no useful tests to add for forward() here, because the autograd function forward() uses the same code as the bf16 impl for forward. The difference where MXFP8 is used is in backward.
• Tests for backward require integration with other EP pipeline components, so this is tested in the integration test added in [mxfp8 moe training] add e2e integration test of full mxfp8 EP pipeline #3584

[pytorch-bot]

🔗 Helpful Links

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

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