Revert "Fix mixed-dtype matmul: correct output type conversion in kernels"

This reverts commit 635a742.

Author: wuhulalala

src/flag_gems/ops/addmm.py

@@ -10,21 +10,6 @@
 from flag_gems.utils import broadcastable_to, libentry, libtuner
 from flag_gems.utils import triton_lang_extension as tle
 
-_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
-
-
-def get_higher_dtype(a, b):
-    if a is b:
-        return a
-    assert a in _ordered_datatypes
-    assert b in _ordered_datatypes
-    for d in _ordered_datatypes:
-        if a is d:
-            return b
-        if b is d:
-            return a
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -91,9 +76,6 @@ def addmm_kernel(
         if IS_FP64:
             a = a.to(tl.float32)
             b = b.to(tl.float32)
-        if a.dtype != b.dtype:
-            a = a.to(tl.float32)
-            b = b.to(tl.float32)
         accumulator += tl.dot(a, b, allow_tf32=False)
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
@@ -106,7 +88,7 @@ def addmm_kernel(
     bias = tl.load(i_ptrs, mask=c_mask, other=0.0)
 
     accumulator = accumulator * alpha + bias * beta
-    c = accumulator.to(c_ptr.dtype.element_ty)
+    c = accumulator.to(bias.dtype)
     tl.store(c_ptrs, c, mask=c_mask)
 
 
@@ -130,8 +112,7 @@ def addmm(bias, mat1, mat2, *, beta=1, alpha=1):
     )
     mat1 = mat1.contiguous()
     # mat2 = mat2.contiguous()
-    c_dtype = get_higher_dtype(mat1.dtype, mat2.dtype)
-    out = torch.empty((M, N), device=mat1.device, dtype=c_dtype)
+    out = torch.empty((M, N), device=mat1.device, dtype=mat1.dtype)
     bias = bias.broadcast_to(out.shape)
 
     grid = lambda META: (
@@ -170,8 +151,7 @@ def addmm_out(bias, mat1, mat2, *, beta=1, alpha=1, out=None):
     M, K = mat1.shape
     _, N = mat2.shape
     if out is None:
-        c_dtype = get_higher_dtype(mat1.dtype, mat2.dtype)
-        out = torch.empty((M, N), device=mat1.device, dtype=c_dtype)
+        out = torch.empty((M, N), device=mat1.device, dtype=mat1.dtype)
     else:
         assert out.shape == (M, N), "Incompatible output shape"
     logger.debug(

src/flag_gems/ops/baddbmm.py

@@ -12,21 +12,6 @@
 from .bmm import bmm
 from .mul import mul
 
-_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
-
-
-def get_higher_dtype(a, b):
-    if a is b:
-        return a
-    assert a in _ordered_datatypes
-    assert b in _ordered_datatypes
-    for d in _ordered_datatypes:
-        if a is d:
-            return b
-        if b is d:
-            return a
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -131,9 +116,6 @@ def baddbmm_kernel(
                 mask_b = mask_k[:, None] & mask_n[None, :]
         a = tl.load(a_ptrs, mask=mask_a)
         b = tl.load(b_ptrs, mask=mask_b)
-        if a.dtype != b.dtype:
-            a = a.to(tl.float32)
-            b = b.to(tl.float32)
         accumulator += tl.dot(a, b, allow_tf32=False)
         offs_k += TILE_K
         a_ptrs += TILE_K
@@ -152,7 +134,7 @@ def baddbmm_kernel(
 
     bi = tl.load(bias_ptrs, mask=mask_c)
     out = accumulator * alpha + bi * beta
-    o = out.to(o_ptrs.dtype.element_ty)
+    o = out.to(bi.dtype)
     tl.store(o_ptrs, o, mask=mask_c)
 
 
@@ -169,8 +151,7 @@ def forward(ctx, bias, A, B, beta, alpha):
         _, _, N = B.shape
         A = A.contiguous()
         B = B.contiguous()
-        c_dtype = get_higher_dtype(A.dtype, B.dtype)
-        out = torch.empty((batch, M, N), dtype=c_dtype, device=A.device)
+        out = torch.empty((batch, M, N), dtype=A.dtype, device=A.device)
 
         bbias = torch.broadcast_to(bias, (batch, M, N)).contiguous()
         bias_batch_stride = bbias.stride(0)

src/flag_gems/ops/bmm.py

@@ -10,21 +10,6 @@
 from flag_gems.utils import libentry, libtuner
 from flag_gems.utils import triton_lang_extension as tle
 
-_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
-
-
-def get_higher_dtype(a, b):
-    if a is b:
-        return a
-    assert a in _ordered_datatypes
-    assert b in _ordered_datatypes
-    for d in _ordered_datatypes:
-        if a is d:
-            return b
-        if b is d:
-            return a
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -144,9 +129,6 @@ def bmm_kernel(
         a_ptrs += TILE_K * stride_ak
         b_ptrs += TILE_K * stride_bk
 
-        if a.dtype != b.dtype:
-            a = a.to(tl.float32)
-            b = b.to(tl.float32)
         o += tl.dot(a, b, allow_tf32=False)
 
     if DIVISIBLE_M and DIVISIBLE_N:
@@ -166,8 +148,7 @@ def bmm(A, B):
     assert A.shape[2] == B.shape[1], "K dim mismatch"
     batch, M, K = A.shape
     _, _, N = B.shape
-    c_dtype = get_higher_dtype(A.dtype, B.dtype)
-    out = torch.empty((batch, M, N), dtype=c_dtype, device=A.device)
+    out = torch.empty((batch, M, N), dtype=A.dtype, device=A.device)
 
     grid_fn = lambda meta: (
         triton.cdiv(meta["M"], meta["TILE_M"]),

src/flag_gems/runtime/backend/_nvidia/hopper/ops/mm.py

@@ -175,16 +175,13 @@ def mm_kernel_general(
         for k in range(0, tl.cdiv(K, BLOCK_K)):
             a = a_desc.load([offset_am.to(tl.int32), offset_k.to(tl.int32)])
             b = b_desc.load([offset_k.to(tl.int32), offset_bn.to(tl.int32)])
-            if a.dtype != b.dtype:
-                a = a.to(tl.float32)
-                b = b.to(tl.float32)
             if IS_FP64:
                 acc += tl.dot(a, b, allow_tf32=False)
             else:
                 acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
             offset_k += BLOCK_K
 
-        acc = acc.to(C.dtype.element_ty)
+        acc = acc.to(a_desc.dtype)
         c_desc.store([offset_am.to(tl.int32), offset_bn.to(tl.int32)], acc)
 
     else:

tests/test_addmm.py

@@ -17,7 +17,7 @@
         (15, 160, 1024),
         (495, 5333, 71),
     ]
-    FLOAT_DTYPES = utils.ALL_FLOAT_DTYPES
+    FLOAT_DTYPES = utils.FLOAT_DTYPES
 
 
 @pytest.mark.addmm
@@ -28,8 +28,6 @@
 def test_addmm(monkeypatch, M, N, K, scalar, dtype, b_column_major):
     if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
         pytest.skip("Skiping fp32 addmm test on tsingmicro platform")
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
 
     if flag_gems.vendor_name == "mthreads":
         monkeypatch.env("MUSA_ENABLE_SQMMA", "1")
@@ -69,8 +67,6 @@ def test_addmm(monkeypatch, M, N, K, scalar, dtype, b_column_major):
 def test_addmm_out(M, N, K, scalar, dtype):
     if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
         pytest.skip("Skiping fp32 addmm_out test on tsingmicro platform")
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
 
     mat1 = torch.randn((M, K), dtype=dtype, device=flag_gems.device)
     mat2 = torch.randn((K, N), dtype=dtype, device=flag_gems.device)
@@ -97,27 +93,3 @@ def test_addmm_out(M, N, K, scalar, dtype):
         torch.addmm(bias2, mat1, mat2, alpha=alpha, beta=beta, out=out)
 
     utils.gems_assert_close(out, ref_out, dtype, reduce_dim=K)
-
-
-MIXED_DTYPE_PAIRS = [
-    (torch.float16, torch.float32),
-    (torch.float32, torch.float16),
-]
-
-
-@pytest.mark.addmm
-@pytest.mark.parametrize("M, N, K", MNK_SHAPES)
-@pytest.mark.parametrize("dtype_a, dtype_b", MIXED_DTYPE_PAIRS)
-def test_addmm_mixed_dtype(M, N, K, dtype_a, dtype_b):
-    mat1 = torch.randn((M, K), dtype=dtype_a, device=flag_gems.device)
-    mat2 = torch.randn((K, N), dtype=dtype_b, device=flag_gems.device)
-    bias = torch.randn((N,), dtype=dtype_a, device=flag_gems.device)
-    ref_mat1 = utils.to_reference(mat1, True)
-    ref_mat2 = utils.to_reference(mat2, True)
-    ref_bias = utils.to_reference(bias, True)
-
-    ref_out = torch.addmm(ref_bias, ref_mat1, ref_mat2)
-    with flag_gems.use_gems():
-        res_out = torch.addmm(bias, mat1, mat2)
-
-    utils.gems_assert_close(res_out, ref_out, torch.float32, reduce_dim=K)

tests/test_baddbmm.py

@@ -3,7 +3,7 @@
 
 import flag_gems
 
-from .accuracy_utils import ALL_FLOAT_DTYPES as ORIG_ALL_FLOAT_DTYPES
+from .accuracy_utils import ALL_FLOAT_DTYPES as ORIG_FLOAT_DTYPES
 from .accuracy_utils import SCALARS, gems_assert_close, to_reference
 from .conftest import QUICK_MODE
 
@@ -18,7 +18,7 @@
         (15, 160, 1024),
         (495, 5333, 71),
     ]
-    FLOAT_DTYPES = ORIG_ALL_FLOAT_DTYPES
+    FLOAT_DTYPES = ORIG_FLOAT_DTYPES
 
 GNK_SHAPES = [(16, 512, 2048), (16, 2560, 2048), (64, 2048, 128)]
 
@@ -37,8 +37,6 @@
 @pytest.mark.parametrize("scalar", SCALARS)
 @pytest.mark.parametrize("dtype", FLOAT_DTYPES)
 def test_baddbmm(monkeypatch, M, N, K, scalar, dtype):
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
     if flag_gems.vendor_name == "mthreads" and dtype in [torch.float16, torch.bfloat16]:
         monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
     batch = 4
@@ -62,8 +60,6 @@ def test_baddbmm(monkeypatch, M, N, K, scalar, dtype):
 @pytest.mark.parametrize("scalar", SCALARS)
 @pytest.mark.parametrize("dtype", FLOAT_DTYPES)
 def test_baddbmm_backward(M, N, K, scalar, dtype):
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
     batch = 2
     mat1 = torch.randn(
         (batch, M, K), dtype=dtype, device=flag_gems.device, requires_grad=True
@@ -95,27 +91,3 @@ def test_baddbmm_backward(M, N, K, scalar, dtype):
     gems_assert_close(res_in_bias, ref_in_bias, dtype, reduce_dim=K)
     gems_assert_close(res_in_grad1, ref_in_grad1, dtype, reduce_dim=N)
     gems_assert_close(res_in_grad2, ref_in_grad2, dtype, reduce_dim=M)
-
-
-MIXED_DTYPE_PAIRS = [
-    (torch.float16, torch.float32),
-    (torch.float32, torch.float16),
-]
-
-
-@pytest.mark.baddbmm
-@pytest.mark.parametrize("M, N, K", MNK_SHAPES)
-@pytest.mark.parametrize("dtype_a, dtype_b", MIXED_DTYPE_PAIRS)
-def test_baddbmm_mixed_dtype(M, N, K, dtype_a, dtype_b):
-    batch = 4
-    mat1 = torch.randn((batch, M, K), dtype=dtype_a, device=flag_gems.device)
-    mat2 = torch.randn((batch, K, N), dtype=dtype_b, device=flag_gems.device)
-    bias = torch.randn((N,), dtype=dtype_a, device=flag_gems.device)
-    ref_mat1 = to_reference(mat1, True)
-    ref_mat2 = to_reference(mat2, True)
-    ref_bias = to_reference(bias, True)
-
-    ref_out = torch.baddbmm(ref_bias, ref_mat1, ref_mat2)
-    res_out = flag_gems.baddbmm(bias, mat1, mat2)
-
-    gems_assert_close(res_out, ref_out, torch.float32, reduce_dim=K)

tests/test_bmm.py

@@ -21,7 +21,7 @@
         (15, 160, 1024),
         (495, 5333, 71),
     ]
-    FLOAT_DTYPES = utils.ALL_FLOAT_DTYPES
+    FLOAT_DTYPES = utils.FLOAT_DTYPES
 
 
 @pytest.mark.bmm
@@ -30,8 +30,6 @@
 def test_bmm(M, N, K, dtype):
     if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
         pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
 
     if flag_gems.vendor_name == "mthreads":
         os.environ["MUSA_ENABLE_SQMMA"] = "1"
@@ -64,8 +62,6 @@ def test_bmm(M, N, K, dtype):
 def test_bmm_non_contiguous(M, N, K, dtype):
     if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
         pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
 
     if flag_gems.vendor_name == "kunlunxin":
         torch.manual_seed(0)
@@ -98,8 +94,6 @@ def test_bmm_non_contiguous(M, N, K, dtype):
 def test_bmm_out(M, N, K, dtype):
     if flag_gems.vendor_name == "tsingmicro" and dtype == torch.float32:
         pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
-    if dtype == torch.float64 and torch.cuda.get_device_capability()[0] < 9:
-        pytest.skip("tl.dot does not support fp64 on compute capability < 9.0")
 
     if flag_gems.vendor_name == "kunlunxin":
         torch.manual_seed(0)
@@ -119,30 +113,3 @@ def test_bmm_out(M, N, K, dtype):
         torch.bmm(mat1, mat2, out=out)
 
     utils.gems_assert_close(out, ref_out, dtype, reduce_dim=K)
-
-
-MIXED_DTYPE_PAIRS = [
-    (torch.float16, torch.float32),
-    (torch.float32, torch.float16),
-]
-
-
-@pytest.mark.bmm
-@pytest.mark.parametrize("M, N, K", MNK_SHAPES)
-@pytest.mark.parametrize("dtype_a, dtype_b", MIXED_DTYPE_PAIRS)
-def test_bmm_mixed_dtype(M, N, K, dtype_a, dtype_b):
-    if flag_gems.vendor_name == "tsingmicro" and (
-        dtype_a == torch.float32 or dtype_b == torch.float32
-    ):
-        pytest.skip("Skiping fp32 bmm test on tsingmicro platform")
-    batch = 4
-    mat1 = torch.randn((batch, M, K), dtype=dtype_a, device=flag_gems.device)
-    mat2 = torch.randn((batch, K, N), dtype=dtype_b, device=flag_gems.device)
-    ref_mat1 = utils.to_reference(mat1, True)
-    ref_mat2 = utils.to_reference(mat2, True)
-
-    ref_out = torch.bmm(ref_mat1, ref_mat2)
-    with flag_gems.use_gems():
-        res_out = torch.bmm(mat1, mat2)
-
-    utils.gems_assert_close(res_out, ref_out, torch.float32, reduce_dim=K)