[mxfp8 moe training] update benchmarks and tests; simplify per group blocked swizzle ref function

[danielvegamyhre]

Stacked PRs:

• ->[mxfp8 moe training] update benchmarks and tests; simplify per group blocked swizzle ref function #3286
• [mxfp8 moe training] compute prefix sum of group sizes inside kernel intead of precomputing #3285

---

[mxfp8 moe training] update benchmarks and tests; simplify per group blocked swizzle ref function

Changes

• Update benchmark_2d_3d_grouped_gemm.py to use torch._scaled_grouped_mm instead of the fbgemm custom op, now that it is integrated in core
• Simplify torch_to_blocked_2d_M_groups to not require K param as it is not needed, update tests and benchmarks accordingly
• Update bench_triton_mx_block_rearrange_2d_M_groups.py to also bench larger, more realistic total_M dim

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm failures during provisioning step due to network issues

✅ No Failures

As of commit 494df3f with merge base 01374eb ():
💚 Looks good so far! There are no failures yet. 💚

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