[mxfp8 moe training] integrate blocked scale kernels into training code

[danielvegamyhre]

Stacked PRs:

• ->[mxfp8 moe training] integrate blocked scale kernels into training code #2900
• [mxfp8 moe training] use dim1 cast cuda kernel in bwd #2897
• [mxfp8 moe training] add triton kernel for blocked swizzled 3d weight scales #2894
• [mxfp8 moe training] add per group blocked scale kernels #2886

---

[mxfp8 moe training] integrate blocked scale kernels into training code

Summary

• integrate triton kernels for converting scales to blocked format into training code
• wrap triton kernels in custom ops for compile composability
• uncomment fp8 rowwise test, which was previously left commented

Test plan

• sanitize pytest test/prototype/moe_training/test_training.py -s

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 1032499 with merge base f0cca99 ():
💚 Looks good so far! There are no failures yet. 💚

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