[mxfp8 moe training] integrate mxfp8 grouped gemm and triton kernels for scale conversion to blocked format

[danielvegamyhre]

Summary

This PR integrates all the grouped GEMMs and triton kernels for per group scale conversions landed recently (details below) into mxfp8 MoE training:

• We landed 2d-2d support for mxfp8 grouped gemm in FBGEMM (Add 2d-2d support to MXFP8 Grouped GEMM FBGEMM#4816) then integrated into torch._scaled_grouped_mm (MXFP8 grouped GEMM support for torch._scaled_grouped_mm + submodule bump pytorch#162209)
• [mxfp8 moe training] add per group blocked scale kernels #2886 and [mxfp8 moe training] add triton kernel for blocked swizzled 3d weight scales #2894 added triton kernels needed to convert each group's mxfp8 scales to blocked swizzled format while keeping all data on the device / avoiding a d2h sync, for the 2d-3d grouped gemm
• [mxfp8 moe training] per group scale conversion to blocked format with groups along K dim (for 2d2d grouped gemm) #2956 is adding additional triton kernels needed for the per group scale conversion for the 2d-2d grouped gemm

Test plan

• pytest test/prototype/moe_training/test_scaled_grouped_mm.py -k test_mxfp8_grouped_gemm_with_dq_fwd_bwd -s

Next steps

• Make compatible with compile w/ custom ops
• Microbenchmarks
• torchtitan integration -> e2e benchmarks

[pytorch-bot]

🔗 Helpful Links

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

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

[drisspg]

nit row -> col

[drisspg]

nit: not needed comment