[mxfp8 moe training] mxfp8 a2a working e2e in torchtitan llama4 training; improve tests + bench scripts

[danielvegamyhre]

Summary

• With these changes, the mxfp8 a2a working e2e in torchtitan Llama4 training (using this PR in torchtitan: [mxfp8 MoE training] Support mxfp8 all to all in expert parallel torchtitan#1765)
• Perf is currently worse than bf16 baseline to due d2h sync caused by aten::item call resulting from extracting the actual tokens from the overallocated symmetric memory grad_output buffer. This sym mem buff must be overallocated to account for the fact that the corresponding output from a2a fwd will be variable size.
• Therefore, my thinking is:
  1. Use this impl in experimental DSV3 model: This impl is more suitable for the experimental DSV3 no-sync model which natively supports this preallocation method by passing the full padded output/grad_input to downstream ops like grouped_mm, scatter_add etc unmodified. With a couple small changes to this mxfp8 impl (e.g., just returning full padded output and grad_input) we can use it there. The reason this method is experimental is because if while it avoids d2h sync, if there is enough skew in expert routing, the job will crash due to insufficient sym mem buffer space to write to during token exchange.
  2. Add new impl for non-experimental DSV3/Llama4 models: we can add a simpler mxfp8 a2a impl that just kicks off 2 async all_to_all_single_autograds on the e4m3 data and e8m0 scales.

Changes

• When integrating the mxfp8 a2a kernel with torchtitan I hit some CUDA IMA errors, which I've fixed with more robust bounds checking.
• Disable compile for forward()/backward() at method level, due to compile not playing nicely with class variables
• Compile to_mx and to_dtype in fwd and bwd
• Fix name of unit test
• Update bench script to measure real default_a2a and mxfp8_a2a impls used in torchitan (being added in [mxfp8 MoE training] Support mxfp8 all to all in expert parallel torchtitan#1765)
• Add option to profile run in bench script

Benchmarks

input_shape         num_splits    bf16_ms    mxfp8_ms
----------------  ------------  ---------  ----------
(16, 8192, 5120)             8     10.684     62.2852

Limitations

• Extracting actual tokens from overallocated device buffer at end of forward() and backward() causes d2h syncs, hurting perf. Need to think about ways to avoid this.

[pytorch-bot]

🔗 Helpful Links

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

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

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

[danielvegamyhre]

@kwen2501 @vkuzo I need try a different approach to get better perf (see PR description) but would like to land these incremental changes, which contain a mxfp8 a2a impl that is now at least e2e functional in torchtitan training.

[vkuzo]

we can add a simpler mxfp8 a2a impl that just kicks off 2 async all_to_all_single_autograds on the e4m3 data and e8m0 scales

for a general utility I'd start with this and then iterate, sounds simpler

[vkuzo]

stamp