enable torch.compile for mxfp8_cublas recipe

Summary:

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

The current approach is a short term workaround until
pytorch/pytorch#147873 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 #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:

ghstack-source-id: e5687e3
ghstack-comment-id: 2701679811
Pull Request resolved: #1841

Author: vkuzo

test/prototype/mx_formats/test_mx_linear.py

@@ -11,6 +11,7 @@
 import torch
 import torch.nn as nn
 
+from torchao.float8.float8_utils import is_row_major
 from torchao.prototype.mx_formats.config import (
     MXLinearConfig,
     MXLinearRecipeName,
@@ -22,6 +23,7 @@
     swap_linear_with_mx_inference_linear,
     swap_linear_with_mx_linear,
 )
+from torchao.prototype.mx_formats.mx_ops import _scaled_mm_with_uint8_scales
 from torchao.quantization.utils import compute_error
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_4,
@@ -169,11 +171,18 @@ def test_activation_checkpointing():
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 @pytest.mark.skipif(
-    is_sm_at_least_100(), reason="triton does not work yet on CUDA capability 10.0"
+    is_sm_at_least_100(),
+    reason="triton does not work yet on CUDA capability 10.0",
 )
 @pytest.mark.parametrize(
     "recipe_name",
-    ["mxfp8_emulated", "mxfp4_emulated", "mxfp8_cutlass", "mxfp4_cutlass"],
+    [
+        "mxfp8_emulated",
+        "mxfp4_emulated",
+        "mxfp8_cublas",
+        "mxfp8_cutlass",
+        "mxfp4_cutlass",
+    ],
 )
 @pytest.mark.parametrize("bias", [False, True])
 # TODO(future PR): figure out why torch.compile does not match eager when
@@ -190,9 +199,9 @@ def test_linear_compile(recipe_name, bias):
         if not is_sm_at_least_100():
             pytest.skip("CUDA capability >= 10.0 required for MX gemms")
 
-    if bias and recipe_name in ["mxfp8_cutlass", "mxfp4_cutlass"]:
+    if bias and recipe_name in ["mxfp8_cublas", "mxfp8_cutlass", "mxfp4_cutlass"]:
         # TODO(future PR): fix this, things are clearly broken with bias=True
-        pytest.skip("this test is broken for cutlass recipes with bias=True")
+        pytest.skip("this test is broken for non-emulated recipes with bias=True")
 
     M, K, N = 128, 256, 512
     input_shape = (M, K)
@@ -285,6 +294,60 @@ def test_inference_compile_simple(elem_dtype):
         assert sqnr >= 13.5
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    is_sm_at_least_100(),
+    reason="triton does not work yet on CUDA capability 10.0",
+)
+@pytest.mark.skipif(
+    not is_sm_at_least_100(),
+    reason="MX gemms require CUDA capability 10.0",
+)
+def test_scaled_mm_wrapper():
+    # today, e8m0 isn't supported in torchinductor or triton
+    # for now, work around this by creating a wrapper around torch._scaled_mm
+    # which takes uint8 scales, and reinterprets them as e8m0 inside the wrapper
+
+    M, K, N = 128, 256, 512
+    BLOCK_SIZE = 32
+    a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn)
+    b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn)
+
+    a_scale = torch.ones(M, K // BLOCK_SIZE, device="cuda", dtype=torch.float8_e8m0fnu)
+    b_scale = torch.ones(N, K // BLOCK_SIZE, device="cuda", dtype=torch.float8_e8m0fnu)
+
+    out = torch._scaled_mm(a, b.t(), a_scale, b_scale, out_dtype=torch.bfloat16)
+
+    def wrapped(a, b, a_scale, b_scale, out_dtype):
+        if is_row_major(b.stride()):
+            b = b.t().contiguous().t()
+        res = _scaled_mm_with_uint8_scales(a, b, a_scale, b_scale, out_dtype=out_dtype)
+        return res
+
+    wrapped = torch.compile(wrapped)
+
+    # correct memory format of `b`
+    out2 = wrapped(
+        a,
+        b.t(),
+        a_scale.view(torch.uint8),
+        b_scale.view(torch.uint8),
+        out_dtype=torch.bfloat16,
+    )
+    torch.testing.assert_close(out, out2, atol=0, rtol=0)
+
+    # incorrect memory format of `b`
+    b_col_major = b.t().contiguous().t()
+    out3 = wrapped(
+        a,
+        b_col_major.t(),
+        a_scale.view(torch.uint8),
+        b_scale.view(torch.uint8),
+        out_dtype=torch.bfloat16,
+    )
+    torch.testing.assert_close(out, out3, atol=0, rtol=0)
+
+
 def test_filter_fn():
     m1 = nn.Sequential(
         nn.Linear(32, 32),

torchao/prototype/mx_formats/mx_ops.py

@@ -36,6 +36,35 @@
 MX_OPS_TABLE: Dict[Any, Any] = {}
 
 
+@torch.library.custom_op("mylib::_scaled_mm_with_uint8_scales", mutates_args=())
+def _scaled_mm_with_uint8_scales(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    a_scale: torch.Tensor,
+    b_scale: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """
+    Until https://github.com/pytorch/pytorch/issues/147873 is done, we need to
+    work around the lack of support for `torch.float8_e8m0fnu` in
+    torchinductor. We do so by hiding the cast of scales to e8m0 inside a
+    custom op.
+    """
+    # cast back to e8m0 where torchinductor can't see it
+    a_scale = a_scale.view(torch.float8_e8m0fnu)
+    b_scale = b_scale.view(torch.float8_e8m0fnu)
+    res = torch._scaled_mm(a, b, a_scale, b_scale, out_dtype=out_dtype)
+    return res
+
+
+@_scaled_mm_with_uint8_scales.register_fake
+def _(a, b, a_scale, b_scale, out_dtype):
+    m, k = a.shape
+    k2, n = b.shape
+    res = torch.empty(m, n, dtype=out_dtype, device=a.device)
+    return res
+
+
 def implements(aten_ops):
     """Register aten ops to the mx op table"""
 
@@ -83,11 +112,11 @@ def mx_mm(aten_op, args, kwargs=None):
         if a._elem_dtype == torch.float8_e4m3fn:
             assert b._elem_dtype == torch.float8_e4m3fn
             if a._gemm_kernel_choice is MXGemmKernelChoice.CUBLAS:
-                res = torch._scaled_mm(
+                res = _scaled_mm_with_uint8_scales(
                     a._data,
                     b._data,
-                    a_scale_block.view(torch.float8_e8m0fnu),
-                    b_scale_block.view(torch.float8_e8m0fnu),
+                    a_scale_block,
+                    b_scale_block,
                     out_dtype=torch.bfloat16,
                 )
             else: