[mxfp8 moe] add support for fbgemm 2d-3d mx8mx8bf16 grouped gemm

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] add per group blocked scale kernels #2886
• [mxfp8 moe training] add grouped gemm benchmark script #2882
• [mxfp8 moe training] refactor all var names with suffix _mx to _data for clarity #2879
• ->[mxfp8 moe] add support for fbgemm 2d-3d mx8mx8bf16 grouped gemm #2848

---

[mxfp8 moe] add support for fbgemm 2d-3d mx8mx8bf16 grouped gemm

Summary

• fbgemm recently added a 2d-3d mxfp8 grouped gemm in Enable MXFP8 grouped GEMM FBGEMM#4710
• This PR integrates the new gemm into the MoE training code base for the following 2d-3d grouped gemms:
  • output = input @ weight^T
  • grad_input = grad_output @ weight
• Add new mxfp8 utils to_blocked_per_group_2d (for input scales) and to_blocked_per_group_3d (for weight scales). These are pytorch reference implementations that are not performant. We can implement equivalent triton kernels for them later.
• Notes on fbgemm API and pytorch grouped mm API:
  • x must be shape (Mg, K) and row major / contiguous
  • x scales must have been preprocessed to have per-group blocked layout and be contiguous
  • weights must be shape (E, N, K) and row major / contiguous
  • weight scales must have been pre-processed to have per-group blocked layout and be contiguous
  • group sizes is a vector containing the size of each token group in the x tensor
  • starting_row_after_padding corresponds to x_scales tensor and must be size len(group_sizes) + 1 where the first starting row is always 0, and each value corresponds to the starting row of group[i] in the x_scales tensor AFTER padding
• Refactor _emulated_mxfp8_scaled_grouped_mm_2d_3d to have same function signature and input constraints as the fbgemm API

Test plan

• pytest test/prototype/moe_training/test_scaled_grouped_mm.py -k mx
• pytest test/prototype/moe_training/test_training.py -k mx

[pytorch-bot]

🔗 Helpful Links

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

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

• Multiple CI trunk failures after landing https://github.com/pytorch/pytorch/pull/161002

✅ You can merge normally! (1 Unrelated Failure)

As of commit f70cc90 with merge base 8722c0c ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Run Regression Tests on ROCm / test-nightly (ROCM Nightly, linux.rocm.gpu.gfx942.2, --pre torch --index-url https://download.pyt... / linux-job (gh) (matched linux rule in flaky-rules.json)
  Error response from daemon: Get "https://registry-1.docker.io/v2/": net/http: request canceled while waiting for connection (Client.Timeout exceeded while awaiting headers)

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

[pytorch-bot]

This PR needs to be approved by an authorized maintainer before merge.

[danielvegamyhre]

@drisspg @vkuzo this is working for all tested cases now and is ready for review/land

[vkuzo]

nit: _mx should be for a combination of raw data and scale, if B_mx is just the data then better to call it something else

[vkuzo]

is this the same function as the one we have in torchao?

ao/torchao/prototype/mx_formats/utils.py

Line 18 in 6f035e8

def to_blocked(input_matrix, use_triton_kernel: bool = False) -> Tensor:

[danielvegamyhre]

Will test if they're the same and replace if we can - I was having trouble getting the kernel working without CUDA errors so was trying to minimize differences between fbgemm unit test code and this torchao code path.

[drisspg]

should be the exact same

[vkuzo]

stamping since this is prototype