integrate mxfp8 dim1 cast kernel choice enum into MXLinear

[danielvegamyhre]

Stacked PRs:

• ->integrate mxfp8 dim1 cast kernel choice enum into MXLinear #2554
• modify triton_scale_swizzle kernel to accept column major inputs #2553
• add wrapper with dtensor handling for mxfp8 dim1 cast kernel #2552
• register custom sharding for mxfp8 dim1 cast cuda kernel #2551

---

integrate mxfp8 dim1 cast kernel choice enum into MXLinear

Summary

• Add MXFP8Dim1CastKernelChoice enum and replace all uses of boolean flag use_fp8_dim1_cast_triton_kernel with it. (Default to Triton for now)
• Update tests accordingly and verify they are passing.

Test plan

• pytest test/prototype/mx_formats/test_mx_linear.py -k eager_vs_hp
• pytest test/prototype/mx_formats/test_mx_linear.py -k compile

Next steps

• Integrate into torchtitan for e2e fsdp training tests once this stack lands. Torchtitan PR: [mxpf8] Make mxfp8 dim1 cast kernel configurable torchtitan#1401
• Dtensor tests still having issues both with Triton and CUDA: ./test/prototype/mx_formats/test_mx_dtensor.sh
  • Triton error (known issue): RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()
  • Cuda error in dtensor op dispatch: assert res.ndim == 0, "output tensor should be scalar!"

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 5 Pending

As of commit b7c508c with merge base 95d13d5 ():
💚 Looks good so far! There are no failures yet. 💚

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

[danielvegamyhre]

cc @vkuzo for review of this stack. Changes are tested via the tests ran in "test plan" section, which have been updated to test all 3 dim1 cast kernel choices (none, triton, cuda).

[vkuzo]

nit: name it to support dim0_dim1 cast if that is added in the future? MXFP8CastKernelChoice?

[vkuzo]

how would someone use torch.compile generated kernel?

[danielvegamyhre]

By setting mxfp8_dim1_cast_kernel_choice=None. Perhaps we should make it more explicit, by adding a 3rd enum option for MXFP8CastKernelChoice.TORCH?