[moe training] fix uncoalesced global accesses in per group rowwise scaling kernel

[danielvegamyhre]

Stacked PRs:

• ->[moe training] fix uncoalesced global accesses in per group rowwise scaling kernel #2761
• [moe training] use llama4 shapes for kernel benchmarks #2756
• [moe training] fix scaling type bug; refactor distributed tests #2749
• [moe training] use custom ops instead of wrap_triton for fp8 rowwise kernels #2734
• [moe training] update tests for torchtitan moe refactor #2733

---

[moe training] fix uncoalesced global accesses in per group rowwise scaling kernel

Summary

• NCU analysis showed uncoalesced global accesses in per-group rowwise scaling kernel, that occur when writing the output data. This is because the output tensor was in row major memory layout, as required for the LHS operator of the scaled grouped gemm kernel.
• I did some testing and it turns out that when number of experts/groups is >= 64, it is actually faster to write to column major, then transform the output to row major afterwards. For smaller number of experts/groups, it's best to just write to row major and accept the uncoalesced writes.

Performance

Before:

input_shape      n_groups  high_precision_dtype      torch_time_us    triton_time_us  triton_speedup
-------------  ----------  ----------------------  ---------------  ----------------  ----------------
(16640, 5120)          16  torch.bfloat16                  5157.71           574.464  8.98x
(16640, 5120)          32  torch.bfloat16                 10792.7            757.536  14.25x
(16640, 5120)          64  torch.bfloat16                 20967.1           1356.64   15.46x
(16640, 5120)         128  torch.bfloat16                 42147.9           2467.57   17.08x

After:

input_shape      n_groups  high_precision_dtype      torch_time_us    triton_time_us  triton_speedup
-------------  ----------  ----------------------  ---------------  ----------------  ----------------
(16640, 5120)          16  torch.bfloat16                  5207.62           599.744  8.68x
(16640, 5120)          32  torch.bfloat16                 10647.6            670.416  15.88x
(16640, 5120)          64  torch.bfloat16                 21124             1355.79   15.58x
(16640, 5120)         128  torch.bfloat16                 42005.6           1408.99   29.81x

[pytorch-bot]

🔗 Helpful Links

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

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

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