[mxfp8 moe training] mxfp8_on_device_all_to_all_v kernel

[danielvegamyhre]

Summary

• Trace analysis of Llama4 mxfp8 training on b200 (FSDP=4, EP=4, seq_len=8192, local_batch_size=8) shows exposed a2a comms constitute a huge % of total runtime (see trace).
• We can improve this by doing dynamic quant on inputs -> a2a on mxfp8 data and scales -> dequantize outputs, which should reduce bytes sent over the network by ~45%. There will be some overhead from dynamic quant/dequant, but this is on the order of microseconds, and the a2a here is 200+ms, so should be a good trade-off and solid net perf benefit.
• My main concern here is how I'm allocating symmetric memory buffers and necessary rendezvous overhead - the idea is for this to only happen once on the first forward pass then be re-used, but remains to be seen if this works as I expect in practice.

Changes

Using the bf16 all_to_all_v kernel from torchtitan as a starting point, this PR implements:

• _mxfp8_on_device_all_to_all_v kernel using Triton + Symmetric Memory for inter-device comms
• Autograd func wrapper to make it differentiable
• Unit tests verifying correctness of both forward and backward pass, by comparing against standard fp32 all_to_all_single_autograd and validating SQNR is high (>28).
• I also refactored to separate quantization kernels and comms kernels into separate files/modules.

Benchmarks

input_shape        num_splits    bf16_us    mxfp8_us
---------------  ------------  ---------  ----------
(8, 8192, 5120)             8     4032.3     1256.14

Next steps

• Refactor to use NVSHMEM apis directly for inter-node a2a
• torchtitan integration by extending this configurable EP a2a pattern
• Extend to all_to_all_v_2d to fuse in the local token shuffle so we can do it in low precision, moving fewer bytes around -> making it faster, then dequantize only after the token shuffle.
• Look into loop unrolling, which is used in the bf16 implementation

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [Maintenance] MacOS runners update

⏳ No Failures, 3 Pending

As of commit 0f381e3 with merge base d2fae7a ():
💚 Looks good so far! There are no failures yet. 💚

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

[kwen2501]

LGTM! Thanks for quantizing the kernel!

[kwen2501]

note that this works only in a case where number of blocks >= number of ranks.
Otherwise, we'd need to use the few blocks we have to loop over the remote_rank space.

[kwen2501]

Curious, how big is dim_scaling_groups in general?

[danielvegamyhre]

For an input shape (M, K), it is the number of columns in the scales tensor, which will be shape (M, K//block_size), where block_size usually is 32.

I may rename the variable to be something more clear. (At minimum I will clarify in the doc string)