[mxfp8 moe training] add grouped gemm benchmark script

[danielvegamyhre]

Stacked PRs:

• [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] add grouped gemm benchmark script #2882

---

[mxfp8 moe training] add grouped gemm benchmark script

Usage

• python benchmarks/prototype/moe_training/bench_2d-3d_grouped_gemm.py

Summary

• Delete old grouped gemm benchmark script, add new cleaner one.
• Result on llama4 shapes:
  • bf16 = ~1550us
  • mxfp8 = ~1370 us (optimization work to do here...)

[pytorch-bot]

🔗 Helpful Links

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

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