support for 2d-2d emulated mxfp8 grouped gemm

[danielvegamyhre]

Stacked PRs:

• Make scaling type configurable for MoE training #2642
• backward pass for differentiable mxfp8 grouped gemm with dynamic quant #2639
• ->support for 2d-2d emulated mxfp8 grouped gemm #2632

---

Summary

• Add torch implementations of per-token-group scaling for mxfp8 (needed for 2d-2d grouped gemm calculation of dW in backward)
• Add support for 2d-2d emulated mxfp8 grouped gemm
• Add unit tests verifying numerical accuracy

Test plan

• pytest test/prototype/moe_training/test_scaled_grouped_mm.py

[pytorch-bot]

🔗 Helpful Links

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

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

• Failing to create ec2 instances -> lots of queueing

❌ 1 New Failure, 1 Cancelled Job, 1 Pending

As of commit f6c2c3a with merge base 3c466f8 ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  Process completed with exit code 1.

CANCELLED JOB - The following job was cancelled. Please retry:

• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  ##[error]The operation was canceled.

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