enable compile with mxfp8 and mxfp4 cutlass gemm

[vkuzo]

Summary:

Enables torch.compile e2e on MXLinear with the mxfp8_cutlass and mxfp4_cutlass recipes. To do this, we need to add a tag on the op definitions to make inductor honor the col-major layout.

Note that something else is broken with these ops and bias=True - leaving that for a future PR.

Test Plan:

python benchmarks/float8/profile_lowp_training.py ~/local/tmp/20250305_test --mx_recipe_name mxfp8_cutlass

# uncomment the hardware check, then can also run the unit test
pytest test/prototype/mx_formats -s -k test_linear_compile

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• enable mxfp8_cublas recipe in roofline script #1843
• enable torch.compile for mxfp8_cublas recipe #1841
• -> enable compile with mxfp8 and mxfp4 cutlass gemm #1838

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 1 Pending

As of commit 8f82ae7 with merge base 23c3162 ():
💚 Looks good so far! There are no failures yet. 💚

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

[drisspg]

Whoops I meant to add this before

Good catch