add wrapper with dtensor handling for mxfp8 dim1 cast kernel

[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

---

add wrapper with dtensor handling for mxfp8 dim1 cast kernel

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 4 Pending

As of commit fcda719 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.

[vkuzo]

add a test?

[danielvegamyhre]

add a test?

I couldn't really add a meaningful test until these pieces were integrated into MXLinear in the subsequent PRs - where I then expanded the existing tests in test_mx_linear.py and test_mx_dtensor.py to cover both the "triton, cuda, None" options for dim1 cast kernel choice.