enable torch.compile for mxfp8_cublas recipe

[vkuzo]

Summary:

This PR enables MXLinear with mxfp8_cublas recipe to use
torch.compile.

The current approach is a short term workaround until
pytorch/pytorch#148461 is done. Since we can't
use e8m0 in torchinductor or triton yet, we create a custom op wrapper
around torch._scaled_mm which takes uint8 scales and does the cast to
e8m0 inside the wrapper, where torchinductor can't see it.

Test Plan:

// this now works (although performance is not ideal due to https://github.com/pytorch/ao/issues/1788)
python benchmarks/float8/profile_lowp_training.py ~/local/tmp/20250305_test --mx_recipe_name mxfp8_cublas

// we can also uncomment the hardware check and 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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit cb824b2 with merge base 4a5ab2d ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  RuntimeError: Command docker exec -t d82fad38907e1f018983a56d9edcb9c68ad5af5bbf13f047d1c920bc7a91dab7 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch==2.7.0.dev20250122 --index-... / linux-job (gh)
  RuntimeError: Command docker exec -t c380c7931a275914ec68fa225674119be2760cfb80f252afebc6923f94c35f42 /exec failed with exit code 1

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