[not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm

[vkuzo]

Summary:

Requires https://github.com/pytorch/pytorch/pull/145562/files

None of this is for land - just testing for now as we work on a long
term support plan.

Status:

performance

• [done] torch._scaled_mm performance as expected: seeing about 2x speedup from bf16 to mxfp8, and 4x from bf16 to nvfp4
• [not in this PR] fast scaling kernels

numerics - mxfp8

• [done] hello world
• [done] check numerics with A and B set to eye-one and scales set to all-ones
• [done] numerics match fp32 baseline and mx emulation with a single tile (M == K == N == 128), with torch.randn inputs
• [done] numerics match fp32 baseline for multiple tiles, of various shapes
• [done] (not in this PR) train llama with real MX gemm and get decent numerics

numerics - nvfp4

• [done] hello world
• [not done] repeat all steps of mxfp8 numerics testing

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> [not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm #1625
• skip failing MX tests on cuda capability 10.0 #1624

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 9 New Failures

As of commit c8d768e with merge base 47f96f1 ():

NEW FAILURES - The following jobs have failed:

• Code Analysis with Ruff / build (3.9) (gh)
• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 7f882dc6bdea44aebef76ec502652f71705f531d9cd9e88ee21c4b172fd5b3d7 /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 75b4f324e1450773eca8fc1ecc615fcebe886d69a9d730d0bede86073ee66a64 /exec failed with exit code 2
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 54a8ce1343fbbd998a0901d5a48eb4f00c09f0448677eac433d791d8fc9d450c /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 1ff80f76f1b5372b57ba5ea5bca2db252473d7d6932ca80f43e03083b9749b7d /exec failed with exit code 2
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 3b9b5bc189e93b66b274daa35332ecdff8c1067f493c318efc5d94b97ea05752 /exec failed with exit code 2
• 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 a6fe078e551843efca80082a93dedba34635440c451208b06338a60a4b359ff2 /exec failed with exit code 2
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch==2.7.0.dev20250122 --index-url https://down... / linux-job (gh)
  RuntimeError: Command docker exec -t 37e9a2e0e4faaa4faff07949a20d9d057ad77a7b4e69a7126a60beb324815415 /exec failed with exit code 2
• 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 e3f78c4a3ff9cdcafbc114f0da32f5418b4de063ddd9a0ea495d4f195cbc4760 /exec failed with exit code 2

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

[vkuzo]

scale_A

[vkuzo]

scale_A