[mxfp8 moe training] cuda blocked layout kernel handling for skinnier scale tensors

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] auto-select chunk_width in cuda blocked layout kernel #3658
• ->[mxfp8 moe training] cuda blocked layout kernel handling for skinnier scale tensors #3656

---

Summary

Fixes #3636

Context

Bug report shows the user has a scales tensor with 44 columns, which this kernel does not handle well for 2 reasons:

1. TMA 2d transfers require stride be divisible by 16, and for this row major scale tensor, stride would be 44, resulting in a non-informative CUDA error when creating the Tensormap.
   • "globalStrides array, which specifies tensor stride of each of the lower tensorRank - 1 dimensions in bytes, must be a multiple of 16 and less than 2^40.'" - source
2. Even for tensors with stride % 16 == 0, the minimum CHUNK_WIDTH the kernel handles is 64 (this corresponds to a model dim of 2048, as 2048/32==64). For smaller models like DSV3 16b with intermediate dim = 1408 in the MoE layer, this kernel will fail.

Fix

• Add validation with clear error messages enforcing scale_cols % 16 == 0, which will direct the user to use the more flexible, but slower, Triton kernel if they have some odd model dim that results in scale_cols not divisible by 16.
• Support chunk_size 16, 32 to handle smaller model dims.

Tests

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

[pytorch-bot]

🔗 Helpful Links

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

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

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

[drisspg]