[mxfp8 moe training] add triton kernel for blocked swizzled 3d weight scales

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] use dim1 cast cuda kernel in bwd #2897
• ->[mxfp8 moe training] add triton kernel for blocked swizzled 3d weight scales #2894
• [mxfp8 moe training] add per group blocked scale kernels #2886

---

[mxfp8 moe training] add triton kernel for blocked swizzled 3d weight scales

Summary

• Add kernel to convert 3d mxfp8 weight scales to blocked swizzled format
• This kernel is simpler because we can make a modified version of this one used for 2d linear weights, where we just have (1) different launch params to parallelize across groups, and (2) different stride calculations in the kernel, updating the base pointer to account for the group we're in.
• I could have potentially expanded the existing dense model kernel to do this since there's no variable group sizes, but i didn't because (1) kernel needs different launch grid and additional params passed into it, so we'd have a weird if-else launching different kernels with different grids and different params. And (2) I would like to keep all the MoE training code in the prototype/moe_training to make the code base as self-contained as possible for now.
• Note that there are no data dependent group sizes so technically torch.compile should be competitive, maybe we don't need this kernel?

Test plan

• sanitize pytest test/prototype/moe_training/test_kernels.py -k scales_3d

Benchmarks

• The kernel is completely unoptimized and can be improved. NCU flags uncoalesced global accesses and partial waves, which should provide a nice speedup if we resolve.
  • As an aside, I've noticed with triton loading 2d blocks of data, uncoalesced loads from GMEM seem to always happen with row major data but not column major data. This surprises me, because with a 2d block of data in either row major or col major, we should always be able to get coalesced loads by assigning threads in each warp to load either a row or column of data, depending on which direction is contiguous. Not sure if the compiler isn't able to figure this out, or if it's me that is misunderstanding something.

input_shape        torch_time_us    triton_time_us    torch_mem_bw_gbps    triton_mem_bw_gbps  triton_speedup
---------------  ---------------  ----------------  -------------------  --------------------  ----------------
(1, 8192, 160)            19.424            10.112              134.959               259.241  1.92x
(2, 8192, 160)            16.384            10.432              320                   502.577  1.57x
(4, 8192, 160)            28.736            13.216              364.9                 793.414  2.17x
(8, 8192, 160)           133.184            20.512              157.463              1022.4    6.49x
(16, 8192, 160)          104.096            35.648              402.927              1176.59   2.92x
(1, 5120, 256)            26.56             10.08                98.699               260.063  2.63x
(2, 5120, 256)            25.248            10.688              207.655               490.539  2.36x
(4, 5120, 256)           104.512            14.464              100.331               724.956  7.23x
(8, 5120, 256)           544.8              22.56                38.494               929.589  24.15x
(16, 5120, 256)          134.144            38.912              312.672              1077.89   3.45x

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[drisspg]

For this one we can just call the other triton kernl impl though right?

[danielvegamyhre]

We could potentially modify the other impl but the result would be rather ugly IMO, see below (from PR description, sorry it was a bit long):

"I could have potentially expanded the existing dense model kernel to do this since there's no variable group sizes, but i didn't because (1) kernel needs different launch grid and additional params passed into it, so we'd have a weird if-else launching different kernels with different grids and different params. And (2) I would like to keep all the MoE training code in the prototype/moe_training to make the code base as self-contained as possible for now."

What do you think?

[drisspg]

maybe... you could also just factor out the main kernel into a series of triton jit functions and then the kernel builds them up

non blocking

[danielvegamyhre]

Hmm, that could work. Sounds good