[mxfp8 moe training] increase num_warps in mxfp8 a2a comms kernel

[danielvegamyhre]

Higher num warps to hide nvlink latency and increase occupancy. Runtime improves by ~13x in benchmark from ~2000ms to ~150ms. (num_warps was originally 16 but changed it to 1 for debugging something, and am now changing it back - just providing context).

Still, perf is much slower than a standalone all_to_all_single_autograd - 10ms vs 150ms for shape (16,8192,5120) with 8 splits. However, the benefit is we avoid the d2h sync required to compute input/output splits for all_to_all_single and get them on the CPU/host as the impl requires.

TODO: bench against bf16 a2a triton+symmetric memory impl to verify the dynamic quant saving network bandwidth is actually an improvement in practice.

[pytorch-bot]

🔗 Helpful Links

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

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

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