[mxfp8 moe training] fallback cuda kernel for when input doesn't meet 2d TMA constraints

[danielvegamyhre]

[mxfp8 moe training] fallback cuda kernel for when input doesn't meet 2d TMA constraints

Context

• In mx_block_rearrange_2d_M_groups_cuda does not support scales_tensor with M-group size = 44 #3636 the user reported a bug training DeepSeekV3 16b. TL;DR is the CUDA blocked layout kernel for groups along M doesn't support input tensors which don't meet 2d TMA constraint of stride being a multiple of 16 bytes.
• Short term mitigation ([mxfp8 moe training] cuda blocked layout kernel handling for skinnier scale tensors #3656) was to assert this constraint is met, and throw an informative error prompting the user to use the slower but more flexible Triton kernel for this instead.
• This PR lands a proper fix, which is a simple non-pipelined CUDA kernel that is:
  • ~5.5x faster than Triton for DSV3 16b shapes, but handles arbitrary column width
  • Dispatching to use faster pipelined kernel if we can, otherwise fallback to this simpler kernel.

Tests

• pytest test/prototype/moe_training/test_kernels.py -v -s -k cuda_mx_block

Benchmarks

input_shape      chunks_per_tb    torch_time_us    triton_time_us    cuda_time_us  triton_speedup    cuda_speedup
-------------  ---------------  ---------------  ----------------  --------------  ----------------  --------------
(131072, 44)                 1          1105.65            123.94           19.46  8.92x             56.83x
(131072, 44)                 4          1101.41            109.57           19.46  10.05x            56.61x
(131072, 44)                 8          1042.5             207.97           19.46  5.01x             53.58x

E2E in Torchtitan on DSV3 16b I see an extra ~3% TPS speedup with single node dp2ep and the mxfp8_wgrad_with_hp + mxfp8 all2all/expert parallel enabled.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 7 Pending

As of commit 11d5e00 with merge base 28306f0 ():
💚 Looks good so far! There are no failures yet. 💚

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

[slayton58]

Some merge/rebase cleanup to sort, then LGTM

[slayton58]

Remove merge detritus :)