[mxfp8 moe training] compute prefix sum of group sizes inside kernel intead of precomputing

[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] compute prefix sum of group sizes inside kernel intead of precomputing

Context

Currently when converting scales to blocked swizzled format, we precompute the new start index of each padded group. However, this creates a dependency on torch.compile, which we rely on codgen fast triton kernels for these prefix sums, otherwise we inject slow eager mode ops into the hot path of the quantized grouped gemms, resulting in net slowdown in eager.

Given some users don't want to use torch.compile, and given sometimes there are bugs blocking the use of torch.compile, we should support eager mode execution with good perf.

Changes

• Do what we do in the mxfp8 scaled grouped gemm kernels in fbgemm, and compute these group offsets inside the kernel themselves. The work is duplicated, but it is such a tiny amount of work that perf impact is negligible (i.e., O(group_size) prefix sum where group_size is the local number of experts after EP is applied, so <= 8).

Testing

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 1cfb0b3 with merge base f856d36 ():
💚 Looks good so far! There are no failures yet. 💚

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