[mxfp8 moe training] add per group blocked scale kernels

[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 per group blocked scale kernels for 2d input activations

Summary

• We currently use a pytorch loop based impl for converting mxfp8 scales to block swizzled format on a per-group basis
• This of course is not optimal for perf and results in d2h sync
• This PR implements a Triton kernel which does this conversion, without doing a d2h sync then looping on the host
• Note: to simplify the kernel, we pre-compute the start row of each group in the block padded output scales tensor (see compute_per_group_blocked_scale_offsets), but this is just a couple standard torch ops and shouldn't cause a d2h sync. There is probably still room for optimization here by doing this in the kernel somehow, but we'll take things one step at a time.

Test plan

• pytest test/prototype/moe_training/test_kernels.py -k blocked

Performance

• Low memory bandwidth utilization but 14x faster than existing torch implementation

input_shape      torch_time_us    triton_time_us    torch_mem_bw_gbps    triton_mem_bw_gbps  triton_speedup
-------------  ---------------  ----------------  -------------------  --------------------  ----------------
(16640, 160)           866.848            60.416                6.261                89.831  14.35x

[pytorch-bot]

🔗 Helpful Links

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

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

[danielvegamyhre]

Note for reviewer: I think we can probably do this without a loop, and just parallelize across row blocks as well (like in the original impl for dense models). Need to think about it some more.

[drisspg]

lets add as a follow up / todo