fix mix MatMul kernels fail with fp64 inputs due to accumulator type …

src/flag_gems/ops/addmm.py

@@ -47,6 +47,7 @@ def addmm_kernel(
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
+    IS_FP64: tl.constexpr = False,
 ):
     pid_m = tle.program_id(0)
     pid_n = tle.program_id(1)
@@ -57,7 +58,10 @@ def addmm_kernel(
     a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
     b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
 
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    if IS_FP64:
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float64)
+    else:
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = tl.load(
             a_ptrs,
@@ -69,6 +73,9 @@ def addmm_kernel(
             mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) & (offs_bn[None, :] < N),
             other=0.0,
         )
+        if IS_FP64:
+            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
@@ -131,6 +138,7 @@ def addmm(bias, mat1, mat2, *, beta=1, alpha=1):
             bias.stride(1),
             out.stride(0),
             out.stride(1),
+            IS_FP64=mat1.dtype == torch.float64,
         )
     return out
 
@@ -182,5 +190,6 @@ def addmm_out(bias, mat1, mat2, *, beta=1, alpha=1, out=None):
             bias.stride(1),
             out.stride(0),
             out.stride(1),
+            IS_FP64=mat1.dtype == torch.float64,
         )
     return out

src/flag_gems/ops/baddbmm.py

@@ -49,6 +49,7 @@ def baddbmm_kernel(
     bias_batch_stride: tl.constexpr,
     bias_M_stride: tl.constexpr,
     bias_N_stride: tl.constexpr,
+    IS_FP64: tl.constexpr = False,
 ):
     # batch offsets
     pid_b = tle.program_id(2)
@@ -89,7 +90,10 @@ def baddbmm_kernel(
     o_ptrs = O + offs_m[:, None] * N + offs_n[None, :]
 
     num_iters = tl.cdiv(K, TILE_K)
-    accumulator = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+    if IS_FP64:
+        accumulator = tl.zeros((TILE_M, TILE_N), dtype=tl.float64)
+    else:
+        accumulator = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
     for _ in range(num_iters):
         if DIVISIBLE_K:
             if DIVISIBLE_M:
@@ -173,6 +177,7 @@ def forward(ctx, bias, A, B, beta, alpha):
                 bias_batch_stride=bias_batch_stride,
                 bias_M_stride=bias_M_stride,
                 bias_N_stride=bias_N_stride,
+                IS_FP64=A.dtype == torch.float64,
             )
         return out
 

src/flag_gems/ops/bmm.py

@@ -54,6 +54,7 @@ def bmm_kernel(
     DIVISIBLE_M: tl.constexpr,
     DIVISIBLE_N: tl.constexpr,
     DIVISIBLE_K: tl.constexpr,
+    IS_FP64: tl.constexpr = False,
 ):
     # batch offsets
     pid_b = tle.program_id(2)
@@ -96,7 +97,10 @@ def bmm_kernel(
     o_ptrs = O + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
 
     num_iters = tl.cdiv(K, TILE_K)
-    o = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+    if IS_FP64:
+        o = tl.zeros((TILE_M, TILE_N), dtype=tl.float64)
+    else:
+        o = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
     for _ in range(num_iters):
         if DIVISIBLE_K:
             if DIVISIBLE_M:
@@ -168,6 +172,7 @@ def bmm(A, B):
             out.stride(0),
             out.stride(1),
             out.stride(2),
+            IS_FP64=A.dtype == torch.float64,
         )
     return out
 
@@ -201,5 +206,6 @@ def bmm_out(A, B, out):
             out.stride(0),
             out.stride(1),
             out.stride(2),
+            IS_FP64=A.dtype == torch.float64,
         )
     return out

src/flag_gems/ops/mm.py

@@ -49,6 +49,7 @@ def mm_kernel_general(
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
     GROUP_M: tl.constexpr,
+    IS_FP64: tl.constexpr = False,
 ):
     # matrix multiplication
     pid = tle.program_id(0)
@@ -69,15 +70,21 @@ def mm_kernel_general(
     rn = rn.to(tl.int64)
     prev_multiple = prev_multiple_of(K, BLOCK_K)
 
-    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    if IS_FP64:
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
+    else:
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for start_k in range(0, prev_multiple, BLOCK_K):
         rk = (start_k + tl.arange(0, BLOCK_K)).to(tl.int64)
         a = tl.load(A + (ram[:, None] * stride_am + rk[None, :] * stride_ak))
         b = tl.load(B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn))
         if a.dtype != b.dtype:
             a = a.to(C.dtype.element_ty)
             b = b.to(C.dtype.element_ty)
-        acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
+        if IS_FP64:
+            acc += tl.dot(a, b, allow_tf32=False)
+        else:
+            acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
 
     # loop peeling
     rk = (prev_multiple + tl.arange(0, BLOCK_K)).to(tl.int64)
@@ -95,7 +102,10 @@ def mm_kernel_general(
     if a.dtype != b.dtype:
         a = a.to(C.dtype.element_ty)
         b = b.to(C.dtype.element_ty)
-    acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
+    if IS_FP64:
+        acc += tl.dot(a, b, allow_tf32=False)
+    else:
+        acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
 
     acc = acc.to(C.dtype.element_ty)
     # rematerialize rm and rn to save registers
@@ -107,7 +117,7 @@ def mm_kernel_general(
     tl.store(C, acc, mask=mask)
 
 
-_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32]
+_ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
 
 
 def get_higher_dtype(a, b):
@@ -152,6 +162,7 @@ def general_mm(a, b, c, M, N, K):
             c.stride(0),
             c.stride(1),
             GROUP_M=8,
+            IS_FP64=a.dtype == torch.float64,
         )
     return c
 

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

@@ -119,6 +119,7 @@ def mm_kernel_general(
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
     GROUP_M: tl.constexpr,
+    IS_FP64: tl.constexpr = False,
 ):
     # matrix multiplication
     pid = tle.program_id(0)
@@ -167,11 +168,17 @@ def mm_kernel_general(
             block_shape=[BLOCK_M, BLOCK_N],
         )
 
-        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        if IS_FP64:
+            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
+        else:
+            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
         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)])
-            acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
+            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)
@@ -187,15 +194,21 @@ def mm_kernel_general(
         rn = rn.to(tl.int64)
         prev_multiple = prev_multiple_of(K, BLOCK_K)
 
-        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        if IS_FP64:
+            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
+        else:
+            acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
         for start_k in range(0, prev_multiple, BLOCK_K):
             rk = (start_k + tl.arange(0, BLOCK_K)).to(tl.int64)
             a = tl.load(A + (ram[:, None] * stride_am + rk[None, :] * stride_ak))
             b = tl.load(B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn))
             if a.dtype != b.dtype:
                 a = a.to(C.dtype.element_ty)
                 b = b.to(C.dtype.element_ty)
-            acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
+            if IS_FP64:
+                acc += tl.dot(a, b, allow_tf32=False)
+            else:
+                acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
 
         # loop peeling
         rk = (prev_multiple + tl.arange(0, BLOCK_K)).to(tl.int64)
@@ -213,7 +226,10 @@ def mm_kernel_general(
         if a.dtype != b.dtype:
             a = a.to(C.dtype.element_ty)
             b = b.to(C.dtype.element_ty)
-        acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
+        if IS_FP64:
+            acc += tl.dot(a, b, allow_tf32=False)
+        else:
+            acc += tl.dot(a, b, out_dtype=tl.float32, allow_tf32=False)
 
         acc = acc.to(C.dtype.element_ty)
         # rematerialize rm and rn to save registers
@@ -343,7 +359,7 @@ def mm_kernel_general_host_tma(
 
 
 def get_higher_dtype(a, b):
-    _ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32]
+    _ordered_datatypes = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
 
     if a is b:
         return a
@@ -435,6 +451,7 @@ def alloc_fn(size: int, align: int, stream: Optional[int]):
                 c.stride(0),
                 c.stride(1),
                 GROUP_M=8,
+                IS_FP64=a.dtype == torch.float64,
             )
     return c
 
@@ -471,6 +488,7 @@ def gemv_kernel(
     stride_bk,
     BLOCK_M: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    IS_FP64: tl.constexpr = False,
 ):
     """Optimized kernel for matrix-vector multiplication (N=1 case)"""
     pid = tl.program_id(0)
@@ -481,7 +499,10 @@ def gemv_kernel(
     row_mask = row_offset < M
 
     # Accumulator for this block of rows
-    acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
+    if IS_FP64:
+        acc = tl.zeros((BLOCK_M,), dtype=tl.float64)
+    else:
+        acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
 
     # Iterate over K dimension
     for k_start in range(0, K, BLOCK_K):
@@ -525,6 +546,7 @@ def gemv_mm(a, b, c, M, K):
             a.stride(0),
             a.stride(1),
             b.stride(0),
+            IS_FP64=a.dtype == torch.float64,
         )
     return c
 

tests/test_mm.py

@@ -22,6 +22,7 @@
     ]
     FLOAT_DTYPES = utils.FLOAT_DTYPES
 
+
 MK_SHAPES = (
     [(1, 32)]
     if QUICK_MODE