[Gemm] Add gemm_sm120.py

Author: tridao

benchmarks/benchmark_gemm.py

@@ -14,7 +14,7 @@
 from cutlass.cute.runtime import from_dlpack, make_ptr
 from cutlass import Int32, Boolean
 
-from quack.gemm_default_epi import GemmDefaultSm90, GemmDefaultSm100
+from quack.gemm_default_epi import GemmDefaultSm90, GemmDefaultSm100, GemmDefaultSm120
 from quack.gemm_sm90 import TileSchedulerOptions
 
 from quack.cute_dsl_utils import get_device_capacity
@@ -141,6 +141,7 @@ def parse_arguments() -> argparse.Namespace:
     parser.add_argument("--gather_A", action="store_true", help="Gather A")
     parser.add_argument("--add_to_output", action="store_true", help="Add to output")
     parser.add_argument("--fp8_fast_accum", action="store_true", help="FP8 fast accum")
+    parser.add_argument("--sm120", action="store_true", help="Use SM120 warp-level MMA (on SM90 HW)")
     parser.add_argument("--skip_ref_check", action="store_true", help="Skip reference checking")
 
     args = parser.parse_args()
@@ -181,6 +182,7 @@ def run(
     gather_A: bool,
     add_to_output: bool,
     fp8_fast_accum: bool,
+    sm120: bool = False,
     **kwargs,
 ):
     """
@@ -235,7 +237,7 @@ def run(
     # Unpack parameters
     m, n, k, l = mnkl
     cluster_shape_mnk = (*cluster_shape_mn, 1)
-    GemmCls = GemmDefaultSm100 if is_sm100 else GemmDefaultSm90
+    GemmCls = GemmDefaultSm100 if is_sm100 else (GemmDefaultSm120 if sm120 else GemmDefaultSm90)
 
     # Skip unsupported types
     if not GemmCls.is_valid_dtypes(
@@ -377,6 +379,15 @@ def create_and_permute_tensor(l, mode0, mode1, is_mode0_major, dtype, is_dynamic
             gather_A=gather_A,
             use_clc_persistence=dynamic_persistent,
         )
+    elif sm120:
+        gemm = GemmCls(
+            acc_dtype,
+            a_dtype,
+            tile_shape_mn,
+            cluster_shape_mnk,
+            is_persistent=persistent,
+            gather_A=gather_A,
+        )
     else:
         gemm = GemmCls(
             acc_dtype,
@@ -600,5 +611,6 @@ def fn():
         args.gather_A,
         args.add_to_output,
         args.fp8_fast_accum,
+        args.sm120,
     )
     print("PASS")

quack/gemm_default_epi.py

@@ -10,6 +10,7 @@
 from quack.epi_ops import Scalar, RowVecLoad, ColVecLoad
 from quack.gemm_sm90 import GemmSm90
 from quack.gemm_sm100 import GemmSm100
+from quack.gemm_sm120 import GemmSm120
 from quack.rounding import RoundingMode
 import quack.utils as utils
 
@@ -101,3 +102,7 @@ class GemmDefaultSm90(GemmDefaultEpiMixin, GemmSm90):
 
 class GemmDefaultSm100(GemmDefaultEpiMixin, GemmSm100):
     pass
+
+
+class GemmDefaultSm120(GemmDefaultEpiMixin, GemmSm120):
+    pass

quack/gemm_sm100.py

@@ -1,3 +1,4 @@
+# Copyright (c) 2025-2026, Tri Dao.
 # Based on the cute-dsl example:
 # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
 
@@ -210,6 +211,10 @@ def __init__(
         self.epi_load_warp_id = self.ab_load_warp_id + self.num_ab_load_warps
         self.scheduler_warp_id = self.epi_load_warp_id + 1
         self.num_epi_warps = len(self.epilog_warp_id)
+        self.epilogue_barrier = pipeline.NamedBarrier(
+            barrier_id=int(NamedBarrierGemm.Epilogue),
+            num_threads=self.num_epi_warps * cute.arch.WARP_SIZE,
+        )
         # Register reallocation for gather_A (3 warp groups, 504 regs total, 168 per WG default).
         # Heavy epilogues (e.g. colvec_reduce in DGated) override these to avoid register spilling.
         # Without gather_A there are only 2 WGs (512 total, 256 per WG = max), no reallocation needed.
@@ -1393,11 +1398,6 @@ def kernel(
             # (MMA, MMA_M, MMA_N, STAGE)
             tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
 
-            epilogue_barrier = pipeline.NamedBarrier(
-                barrier_id=int(NamedBarrierGemm.Epilogue),
-                num_threads=self.num_epi_warps * cute.arch.WARP_SIZE,
-            )
-
             # Partition for epilogue
             epi_tidx = tidx
             tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition(
@@ -1479,7 +1479,7 @@ def kernel(
                     copy_C,
                     tile_coord_mnkl,
                     varlen_manager,
-                    epilogue_barrier,
+                    self.epilogue_barrier,
                     tile_scheduler,
                     epi_tidx,
                     is_tma_warp,