[mxfp8 moe training] mxfp8 token permute autograd func + triton kernels

[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 token permute autograd func + triton kernels

Tests

• PYTHONPATH=/home/$USER/ao:$PYTHONPATH pytest test/prototype/moe_training/ep/test_permute.py
• 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/3580

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