Fix mixed-dtype matmul: correct output type conversion in addmm/baddbmm kernels

Author: wuhulalala

src/flag_gems/ops/addmm.py

@@ -10,6 +10,21 @@
 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__)
 
 
@@ -76,6 +91,9 @@ 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
@@ -88,7 +106,7 @@ def addmm_kernel(
     bias = tl.load(i_ptrs, mask=c_mask, other=0.0)
 
     accumulator = accumulator * alpha + bias * beta
-    c = accumulator.to(bias.dtype)
+    c = accumulator.to(c_ptr.dtype.element_ty)
     tl.store(c_ptrs, c, mask=c_mask)
 
 
@@ -112,7 +130,8 @@ def addmm(bias, mat1, mat2, *, beta=1, alpha=1):
     )
     mat1 = mat1.contiguous()
     # mat2 = mat2.contiguous()
-    out = torch.empty((M, N), device=mat1.device, dtype=mat1.dtype)
+    c_dtype = get_higher_dtype(mat1.dtype, mat2.dtype)
+    out = torch.empty((M, N), device=mat1.device, dtype=c_dtype)
     bias = bias.broadcast_to(out.shape)
 
     grid = lambda META: (
@@ -151,7 +170,8 @@ def addmm_out(bias, mat1, mat2, *, beta=1, alpha=1, out=None):
     M, K = mat1.shape
     _, N = mat2.shape
     if out is None:
-        out = torch.empty((M, N), device=mat1.device, dtype=mat1.dtype)
+        c_dtype = get_higher_dtype(mat1.dtype, mat2.dtype)
+        out = torch.empty((M, N), device=mat1.device, dtype=c_dtype)
     else:
         assert out.shape == (M, N), "Incompatible output shape"
     logger.debug(

src/flag_gems/ops/baddbmm.py

@@ -12,6 +12,21 @@
 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__)
 
 
@@ -116,6 +131,9 @@ 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
@@ -134,7 +152,7 @@ def baddbmm_kernel(
 
     bi = tl.load(bias_ptrs, mask=mask_c)
     out = accumulator * alpha + bi * beta
-    o = out.to(bi.dtype)
+    o = out.to(o_ptrs.dtype.element_ty)
     tl.store(o_ptrs, o, mask=mask_c)
 
 
@@ -151,7 +169,8 @@ def forward(ctx, bias, A, B, beta, alpha):
         _, _, N = B.shape
         A = A.contiguous()
         B = B.contiguous()
-        out = torch.empty((batch, M, N), dtype=A.dtype, device=A.device)
+        c_dtype = get_higher_dtype(A.dtype, B.dtype)
+        out = torch.empty((batch, M, N), dtype=c_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,6 +10,21 @@
 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__)
 
 
@@ -129,6 +144,9 @@ 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:
@@ -148,7 +166,8 @@ def bmm(A, B):
     assert A.shape[2] == B.shape[1], "K dim mismatch"
     batch, M, K = A.shape
     _, _, N = B.shape
-    out = torch.empty((batch, M, N), dtype=A.dtype, device=A.device)
+    c_dtype = get_higher_dtype(A.dtype, B.dtype)
+    out = torch.empty((batch, M, N), dtype=c_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,13 +175,16 @@ 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(a_desc.dtype)
+        acc = acc.to(C.dtype.element_ty)
         c_desc.store([offset_am.to(tl.int32), offset_bn.to(tl.int32)], acc)
 
     else:

tests/test_addmm.py

@@ -97,3 +97,27 @@ 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

@@ -95,3 +95,27 @@ 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

@@ -119,3 +119,30 @@ 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)

tests/test_mm.py

@@ -52,3 +52,29 @@ def test_mm(M, N, K, dtype, b_column_major):
         res_out = torch.mm(mat1, mat2)
 
     utils.gems_assert_close(res_out, ref_out, dtype, reduce_dim=K)
+
+
+MIXED_DTYPE_PAIRS = [
+    (torch.float16, torch.float32),
+    (torch.float32, torch.float16),
+]
+
+
+@pytest.mark.mm
+@pytest.mark.parametrize("M, N, K", MNK_SHAPES)
+@pytest.mark.parametrize("dtype_a, dtype_b", MIXED_DTYPE_PAIRS)
+def test_mm_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 addmm_out test on tsingmicro platform")
+    mat1 = torch.randn((M, K), dtype=dtype_a, device=flag_gems.device)
+    mat2 = torch.randn((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.mm(ref_mat1, ref_mat2)
+    with flag_gems.use_gems():
+        res_out = torch.mm(mat1, mat2)
+
+    utils.gems_assert_close(res_out, ref_out, torch.float32, reduce_dim=K)