[moe fp8 training] test and bench new faster method for per group rowwise scaling

[danielvegamyhre]

Stacked PRs:

• [moe fp8 training] fused reduction kernel along dim1 for 3d expert weights in backward #2865
• [moe fp8 training] use transpose method when quantizing to avoid uncoalesced gmem accesses #2864
• ->[moe fp8 training] test and bench new faster method for per group rowwise scaling #2863

---

[moe fp8 training] test and bench new faster method for per group rowwise scaling

Summary

• Per group rowwise scaling kernel used in the backward pass for MoE fp8 rowwise training is slow, due to uncoalesced global accesses arising from the row major layout of the tensor conflicting with the memory access pattern in the kernel.
• To get around this, we can compute per group colwise scales on the transpose of the input, then transpose the result, to effectively compute per group rowwise scales.

Changes

• Add tests and benchmarks for this transpose based method
• Make separate kernel benchmark scripts for rowwise vs colwise per group scaling, rather than having 1 combined one.
• Update mem bw calcs

Benchmarks

• Benchmarking shows the new method is ~3x faster across all shapes tested for rowwise

Mg,N             n_groups    torch_loop_time_us    triton_time_us    triton_transpose_us    torch_mem_bw_gbps    triton_mem_bw_gbps    triton_transpose_mem_bw_gbps  triton_speedup    triton_transpose_speedup
-------------  ----------  --------------------  ----------------  ---------------------  -------------------  --------------------  ------------------------------  ----------------  --------------------------
(16640, 8192)           1               1705.38           2422.82                768.416              399.681               281.328                         887.029  0.70x             2.22x
(16640, 8192)          16               3078.21           2310.27                686.912              221.59                295.246                         992.993  1.33x             4.48x
(16640, 8192)          64              10639.2            2177.7                 784.288               64.26                313.943                         871.71   4.89x             13.57x

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit a1c0745 with merge base 253d65a ():
💚 Looks good so far! There are no failures yet. 💚

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