feat(kernel): Add batch matmul API

tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gemm/mm_a16w16_gfx950.py

@@ -1041,7 +1041,7 @@ def _get_partial_scratch(
 def _check_supported_dense16_shape(M: int, N: int, K: int) -> None:
     if M <= 0 or N <= 0 or K <= 0:
         raise ValueError(
-            "dense16 Gluon GEMM requires positive M/N/K, got " f"M={M}, N={N}, K={K}"
+            f"dense16 Gluon GEMM requires positive M/N/K, got M={M}, N={N}, K={K}"
         )
     if M > MAX_M:
         raise ValueError(f"dense16 Gluon GEMM requires M <= {MAX_M}; got M={M}")
@@ -1057,6 +1057,29 @@ def _check_supported_dense16_shape(M: int, N: int, K: int) -> None:
         )
 
 
+def _resolve_output(
+    A: torch.Tensor,
+    M: int,
+    N: int,
+    out_dtype: torch.dtype,
+    out: torch.Tensor | None,
+    op_name: str,
+) -> torch.Tensor:
+    if out is None:
+        return torch.empty((M, N), device=A.device, dtype=out_dtype)
+    if out.shape != (M, N):
+        raise ValueError(
+            f"{op_name} output must have shape {(M, N)}, got {tuple(out.shape)}"
+        )
+    if out.dtype != out_dtype:
+        raise TypeError(f"{op_name} output dtype must be {out_dtype}, got {out.dtype}")
+    if out.device != A.device:
+        raise ValueError(f"{op_name} output must be on {A.device}, got {out.device}")
+    if out.stride(-1) != 1:
+        raise ValueError(f"{op_name} output must be contiguous in the last dimension")
+    return out
+
+
 def _use_warp_reduce_smallm(M: int, _N: int, K: int) -> bool:
     return (M in (1, 2) and K <= 2048) or (M == 4 and K <= 1024)
 
@@ -1142,6 +1165,7 @@ def gluon_mm_a16w16_warp_reduce_smallm_gfx950(
     out_dtype: torch.dtype,
     *,
     alpha: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Compute small-M dense ``A @ B.T`` with one warp per output.
 
@@ -1180,7 +1204,7 @@ def gluon_mm_a16w16_warp_reduce_smallm_gfx950(
         )
     _check_supported_dense16_shape(M, N, K)
 
-    C = torch.empty((M, N), device=A.device, dtype=out_dtype)
+    C = _resolve_output(A, M, N, out_dtype, out, "small-M dense16 warp-reduce GEMM")
     total_outputs = M * N
     grid = (triton.cdiv(total_outputs, WARP_REDUCE_OUTPUTS),)
     _warp_reduce_smallm_kernel[grid](
@@ -1202,7 +1226,7 @@ def gluon_mm_a16w16_warp_reduce_smallm_gfx950(
     )
 
     if alpha is not None:
-        C = C * alpha.to(dtype=C.dtype)
+        C.mul_(alpha.to(device=C.device, dtype=C.dtype))
     return C
 
 
@@ -1212,6 +1236,7 @@ def gluon_mm_a16w16_mfma_lds_smallm_gfx950(
     out_dtype: torch.dtype,
     *,
     alpha: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Compute high-K small-M dense ``A @ B.T`` with MFMA and partial-sum reduction.
 
@@ -1224,8 +1249,7 @@ def gluon_mm_a16w16_mfma_lds_smallm_gfx950(
     """
     if A.ndim != 2 or B.ndim != 2:
         raise ValueError(
-            "small-M dense16 MFMA LDS GEMM expects 2D inputs, got "
-            f"{A.ndim=} {B.ndim=}"
+            f"small-M dense16 MFMA LDS GEMM expects 2D inputs, got {A.ndim=} {B.ndim=}"
         )
     if A.dtype not in _SUPPORTED_DTYPES or B.dtype not in _SUPPORTED_DTYPES:
         raise TypeError(
@@ -1261,8 +1285,13 @@ def gluon_mm_a16w16_mfma_lds_smallm_gfx950(
             f"got M={M}, N={N}, K={K}"
         )
 
-    C_full = torch.empty((MFMA_LDS_REDUCE_M, N), device=A.device, dtype=out_dtype)
-    C = C_full[:M, :]
+    if out is not None:
+        _resolve_output(A, M, N, out_dtype, out, "small-M dense16 MFMA LDS GEMM")
+    if out is not None and M == MFMA_LDS_REDUCE_M:
+        C = out
+    else:
+        C_full = torch.empty((MFMA_LDS_REDUCE_M, N), device=A.device, dtype=out_dtype)
+        C = C_full[:M, :]
     split_k = _choose_mfma_lds_split_k(K)
     num_n_tiles = triton.cdiv(N, MFMA_LDS_BLOCK_N)
     total_work = num_n_tiles * split_k
@@ -1310,7 +1339,10 @@ def gluon_mm_a16w16_mfma_lds_smallm_gfx950(
     )
 
     if alpha is not None:
-        C = C * alpha.to(dtype=C.dtype)
+        C.mul_(alpha.to(device=C.device, dtype=C.dtype))
+    if out is not None and C is not out:
+        out.copy_(C)
+        return out
     return C
 
 
@@ -1320,6 +1352,7 @@ def gluon_mm_a16w16_mfma_lds_mediumm_gfx950(
     out_dtype: torch.dtype,
     *,
     alpha: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Compute tuned medium-M dense ``A @ B.T`` with MFMA/LDS tiling.
 
@@ -1364,7 +1397,7 @@ def gluon_mm_a16w16_mfma_lds_mediumm_gfx950(
         )
 
     block_m, block_n, block_k, warps_m, warps_n, num_buffers = config
-    C = torch.empty((M, N), device=A.device, dtype=out_dtype)
+    C = _resolve_output(A, M, N, out_dtype, out, "medium-M dense16 MFMA LDS GEMM")
     grid = (triton.cdiv(M, block_m) * triton.cdiv(N, block_n),)
     _mfma_lds_mediumm_kernel[grid](
         A,
@@ -1391,7 +1424,7 @@ def gluon_mm_a16w16_mfma_lds_mediumm_gfx950(
     )
 
     if alpha is not None:
-        C = C * alpha.to(dtype=C.dtype)
+        C.mul_(alpha.to(device=C.device, dtype=C.dtype))
     return C
 
 
@@ -1401,6 +1434,7 @@ def gluon_mm_a16w16_gfx950(
     out_dtype: torch.dtype,
     *,
     alpha: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor | None:
     """Compute dense16 GEMM ``A @ B.T`` on gfx950 when supported.
 
@@ -1444,6 +1478,7 @@ def gluon_mm_a16w16_gfx950(
             B,
             out_dtype,
             alpha=alpha,
+            out=out,
         )
     if M > MAX_M:
         return None
@@ -1454,19 +1489,22 @@ def gluon_mm_a16w16_gfx950(
             B,
             out_dtype,
             alpha=alpha,
+            out=out,
         )
     if _use_mfma_lds_smallm(M, N, K):
         return gluon_mm_a16w16_mfma_lds_smallm_gfx950(
             A,
             B,
             out_dtype,
             alpha=alpha,
+            out=out,
         )
     if _use_mfma_lds_mediumm(M, N, K):
         return gluon_mm_a16w16_mfma_lds_mediumm_gfx950(
             A,
             B,
             out_dtype,
             alpha=alpha,
+            out=out,
         )
     return None

tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gemm/mm_a16w16_largem_gfx950.py

@@ -394,12 +394,41 @@ def _supports_largem_shape(M: int, N: int, K: int) -> bool:
     )
 
 
+def _resolve_largem_output(
+    A: torch.Tensor,
+    M: int,
+    N: int,
+    out_dtype: torch.dtype,
+    out: torch.Tensor | None,
+) -> torch.Tensor:
+    if out is None:
+        return torch.empty((M, N), device=A.device, dtype=out_dtype)
+    if out.shape != (M, N):
+        raise ValueError(
+            f"large-M dense16 output must have shape {(M, N)}, got {tuple(out.shape)}"
+        )
+    if out.dtype != out_dtype:
+        raise TypeError(
+            f"large-M dense16 output dtype must be {out_dtype}, got {out.dtype}"
+        )
+    if out.device != A.device:
+        raise ValueError(
+            f"large-M dense16 output must be on {A.device}, got {out.device}"
+        )
+    if out.stride(-1) != 1:
+        raise ValueError(
+            "large-M dense16 output must be contiguous in the last dimension"
+        )
+    return out
+
+
 def gluon_mm_a16w16_largem_gfx950(
     A: torch.Tensor,
     B: torch.Tensor,
     out_dtype: torch.dtype,
     *,
     alpha: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor | None:
     """Compute large, aligned dense16 ``A @ B.T`` with an 8-wave MFMA/LDS tile."""
     if A.ndim != 2 or B.ndim != 2:
@@ -416,7 +445,7 @@ def gluon_mm_a16w16_largem_gfx950(
     if K_b != K or not _supports_largem_shape(M, N, K):
         return None
 
-    C = torch.empty((M, N), device=A.device, dtype=out_dtype)
+    C = _resolve_largem_output(A, M, N, out_dtype, out)
     grid_m = triton.cdiv(M, LARGEM_BLOCK_M)
     grid_n = triton.cdiv(N, LARGEM_BLOCK_N)
     grid_mn = grid_m * grid_n
@@ -446,5 +475,5 @@ def gluon_mm_a16w16_largem_gfx950(
     )
 
     if alpha is not None:
-        C = C * alpha.to(dtype=C.dtype)
+        C.mul_(alpha.to(device=C.device, dtype=C.dtype))
     return C

tokenspeed-kernel-amd/test/ops/test_mm_a16w16_gfx950.py

@@ -76,6 +76,39 @@ def test_dense16_kernel_variant_correctness(
     torch.testing.assert_close(out, torch.mm(a, b.T), atol=1e-2, rtol=1e-2)
 
 
+@pytest.mark.parametrize("kernel,shape", _CORRECTNESS_CASES)
+def test_dense16_kernel_variant_writes_strided_out(
+    kernel, shape: tuple[int, int, int]
+) -> None:
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    m, n, k = shape
+    a = torch.randn((m, k), device="cuda", dtype=dtype) * 0.25
+    b = torch.randn((n, k), device="cuda", dtype=dtype) * 0.25
+    backing = torch.empty((m, n + 17), device="cuda", dtype=dtype)
+    out = backing[:, :n]
+
+    actual = kernel(a, b, dtype, out=out)
+
+    assert actual is out
+    torch.testing.assert_close(out, torch.mm(a, b.T), atol=1e-2, rtol=1e-2)
+
+
+def test_splitk_smallm_out_handles_padded_reducer_rows() -> None:
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    m, n, k = 2, 256, 2048
+    a = torch.randn((m, k), device="cuda", dtype=dtype) * 0.25
+    b = torch.randn((n, k), device="cuda", dtype=dtype) * 0.25
+    backing = torch.empty((m, n + 17), device="cuda", dtype=dtype)
+    out = backing[:, :n]
+
+    actual = gluon_mm_a16w16_mfma_lds_smallm_gfx950(a, b, dtype, out=out)
+
+    assert actual is out
+    torch.testing.assert_close(out, torch.mm(a, b.T), atol=1e-2, rtol=1e-2)
+
+
 def test_use_warp_reduce_covers_small_k_decode_shapes() -> None:
     assert _use_warp_reduce_smallm(1, 1280, 1024)
     assert _use_warp_reduce_smallm(2, 2560, 2048)

tokenspeed-kernel/python/tokenspeed_kernel/__init__.py

@@ -43,7 +43,7 @@
     rel_mha_plan,
     rel_mha_prefill,
 )
-from tokenspeed_kernel.ops.gemm import mm
+from tokenspeed_kernel.ops.gemm import bmm, mm
 from tokenspeed_kernel.ops.moe import moe_apply, moe_plan, moe_process_weights
 from tokenspeed_kernel.ops.quantization import (
     quantize_fp8,
@@ -60,6 +60,7 @@
     # exceptions
     "NoKernelFoundError",
     # gemm
+    "bmm",
     "mm",
     # attention
     "mha_plan",

tokenspeed-kernel/python/tokenspeed_kernel/benchmark/throughput.py

@@ -54,6 +54,21 @@ def gemm_mm_bytes(M: int, N: int, K: int, dtype: torch.dtype) -> int:
         element_size = _dtype_nbytes(dtype)
         return (M * K + K * N + M * N) * element_size
 
+    @staticmethod
+    def gemm_bmm_flops(batch: int, M: int, N: int, K: int) -> int:
+        return batch * ThroughputCalculator.gemm_mm_flops(M, N, K)
+
+    @staticmethod
+    def gemm_bmm_bytes(
+        batch: int,
+        M: int,
+        N: int,
+        K: int,
+        dtype: torch.dtype,
+    ) -> int:
+        element_size = _dtype_nbytes(dtype)
+        return (batch * M * K + batch * N * K + batch * M * N) * element_size
+
     @staticmethod
     def attn_flops(
         batch: int,
@@ -105,6 +120,28 @@ def compute(
             bandwidth = bytes_moved / seconds / 1e9
             return tflops, bandwidth
 
+        if op_family == "gemm" and op_mode == "bmm":
+            batch = _shape_int(shape_params, "B", "batch")
+            M = _shape_int(shape_params, "M")
+            N = _shape_int(shape_params, "N")
+            K = _shape_int(shape_params, "K")
+            if any(value is None for value in (batch, M, N, K)):
+                return None, None
+
+            flops = ThroughputCalculator.gemm_bmm_flops(batch, M, N, K)
+            bytes_moved = ThroughputCalculator.gemm_bmm_bytes(
+                batch,
+                M,
+                N,
+                K,
+                dtype,
+            )
+            seconds = latency_us * 1e-6
+
+            tflops = flops / seconds / 1e12
+            bandwidth = bytes_moved / seconds / 1e9
+            return tflops, bandwidth
+
         if op_family in {"attention", "attn"} and op_mode == "decode":
             batch = _shape_int(
                 shape_params,