[Example] Generalize TCGEN05 GEMMs for Thor (sm110a) and add Stream-K scheduling

README.md

@@ -200,7 +200,7 @@ print("GEMM + ReLU passed.")
 - **GEMM and quantization:** [GEMM](https://github.com/tile-ai/tilelang/tree/main/examples/gemm), [grouped GEMM](https://github.com/tile-ai/tilelang/tree/main/examples/grouped_gemm), [FP8 GEMM](https://github.com/tile-ai/tilelang/tree/main/examples/gemm_fp8), [dequantization GEMM](https://github.com/tile-ai/tilelang/tree/main/examples/dequantize_gemm), and [block-scaled GEMM](https://github.com/tile-ai/tilelang/tree/main/examples/blockscaled_gemm_sm100)
 - **Attention and sequence models:** [FlashAttention](https://github.com/tile-ai/tilelang/tree/main/examples/flash_attention), [Flash Decoding](https://github.com/tile-ai/tilelang/tree/main/examples/flash_decoding), [block-sparse attention](https://github.com/tile-ai/tilelang/tree/main/examples/blocksparse_attention), [linear attention](https://github.com/tile-ai/tilelang/tree/main/examples/linear_attention), and [GDN](https://github.com/tile-ai/tilelang/tree/main/examples/gdn)
 - **Model workloads:** [DeepSeek MLA](https://github.com/tile-ai/tilelang/tree/main/examples/deepseek_mla), [DeepSeek V3.2](https://github.com/tile-ai/tilelang/tree/main/examples/deepseek_v32), [DeepSeek V4](https://github.com/tile-ai/tilelang/tree/main/examples/deepseek_v4), and [DeepSeek mHC](https://github.com/tile-ai/tilelang/tree/main/examples/deepseek_mhc)
-- **Architecture-specific kernels:** [AMD](https://github.com/tile-ai/tilelang/tree/main/examples/amd), [SM100](https://github.com/tile-ai/tilelang/tree/main/examples/gemm_sm100), and [SM120](https://github.com/tile-ai/tilelang/tree/main/examples/gemm_sm120)
+- **Architecture-specific kernels:** [AMD](https://github.com/tile-ai/tilelang/tree/main/examples/amd), [TCGEN05](https://github.com/tile-ai/tilelang/tree/main/examples/gemm_tcgen05), and [SM120](https://github.com/tile-ai/tilelang/tree/main/examples/gemm_sm120)
 - **Compiler and debugging tools:** [analyzer](https://github.com/tile-ai/tilelang/tree/main/examples/analyze), [layout visualization](https://github.com/tile-ai/tilelang/tree/main/examples/plot_layout), [AutoDD](https://github.com/tile-ai/tilelang/tree/main/examples/autodd), and [IKET](https://github.com/tile-ai/tilelang/tree/main/examples/iket)
 
 Browse the [complete examples directory](https://github.com/tile-ai/tilelang/tree/main/examples) for additional operators, tests, and architecture-specific implementations.

examples/gemm_sm100/README.md → examples/gemm_tcgen05/README.md

@@ -1,6 +1,6 @@
-# TileLang SM100 Support (Preview)
+# TileLang TCGEN05 Support
 
-This directory contains examples for TileLang's experimental SM100 architecture support. **This is a preview version** with limited functionality.
+This directory contains examples for TileLang's experimental TCGEN05 support on compatible NVIDIA architectures. **This is a preview version** with limited functionality.
 
 ## Current Limitations (Manual Implementation Required)
 
@@ -21,7 +21,7 @@ T.tcgen05_gemm(A_shared, B_shared, C_tmem, trans_A, trans_B, mbar=mbar, clear_ac
 T.mbarrier_wait_parity(mbar, k%2)  # Manual phase calculation required
 ```
 
-TileLang now has a conservative `InjectTcgen05Fence` pass on SM100+ that can
+TileLang now has a conservative `InjectTcgen05Fence` pass on TCGEN05-capable targets that can
 insert `tcgen05_before_thread_sync()` / `tcgen05_after_thread_sync()` around:
 - `tvm_storage_sync("shared"|"shared.dyn")`
 - linear `mbarrier_wait_parity(...) -> tcgen05/TMEM use` regions
@@ -42,6 +42,15 @@ Demonstrates TCGEN5MMA operations with:
 ### Traditional MMA Example (`gemm_mma.py`)
 Shows standard MMA operations that work across architectures for comparison.
 
+### Persistent TCGEN05 GEMM (`gemm_tcgen5mma_ws_persistent.py`)
+Warp-specialized persistent kernels (single-CTA and 2-CTA) that use the static
+PersistentTileScheduler, with optional TMA stores.
+
+### Stream-K TCGEN05 GEMM (`gemm_tcgen5mma_ws_persistent_streamk.py`)
+Persistent 2-CTA kernel with Stream-K scheduling: the under-filled tail wave
+is split along K, peers publish partial accumulations to a workspace, and the
+final peer fixes up and writes each output tile.
+
 ## Code Example
 
 The following code is based on `gemm_tcgen5mma.py`, demonstrating TCGEN5MMA matrix multiplication:

examples/gemm_sm100/gemm_tcgen5mma_ws.py → examples/gemm_tcgen05/gemm_tcgen5mma_ws.py

@@ -142,7 +142,7 @@ def gemm_2cta(A, B, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype,
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Non-persistent warp-specialized SM100 GEMM")
+    parser = argparse.ArgumentParser(description="Non-persistent warp-specialized TCGEN05 GEMM")
     parser.add_argument("--m", type=int, default=8192)
     parser.add_argument("--n", type=int, default=8192)
     parser.add_argument("--k", type=int, default=8192)
@@ -182,8 +182,18 @@ def main():
     torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
     print("All checks passed. ✅")
 
-    tl_latency = do_bench(lambda: kernel(a, b, **kwargs), backend="cupti")
-    torch_latency = do_bench(lambda: a @ b, backend="cupti")
+    tl_latency = do_bench(
+        lambda: kernel(a, b, **kwargs),
+        _n_warmup=50,
+        _n_repeat=50,
+        backend="cupti",
+    )
+    torch_latency = do_bench(
+        lambda: a @ b,
+        _n_warmup=50,
+        _n_repeat=50,
+        backend="cupti",
+    )
     print(f"Tilelang latency: {tl_latency} ms")
     print(f"Flops: {2 * M * N * K / (tl_latency / 1e3) / 1e12} TFLOPS")
     print(f"Torch latency: {torch_latency} ms")

examples/gemm_sm100/gemm_tcgen5mma_ws_clc.py → examples/gemm_tcgen05/gemm_tcgen5mma_ws_clc.py

@@ -350,7 +350,10 @@ def ref_program(A, B):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="CLC-scheduled persistent warp-specialized SM100 GEMM")
+    parser = argparse.ArgumentParser(description="CLC-scheduled persistent warp-specialized TCGEN05 GEMM")
+    parser.add_argument("--m", type=int, default=8192)
+    parser.add_argument("--n", type=int, default=8192)
+    parser.add_argument("--k", type=int, default=8192)
     parser.add_argument("--block_M", type=int, default=128)
     parser.add_argument("--block_N", type=int, default=256)
     parser.add_argument("--store_block_N", type=int, default=64, help="block_N for C_shared")
@@ -360,8 +363,7 @@ def ref_program(A, B):
     parser.add_argument("--use_tma_store", action=argparse.BooleanOptionalAction, default=True)
     args = parser.parse_args()
 
-    group_size = args.group_size
-    use_tma_store = args.use_tma_store
+    M, N, K = args.m, args.n, args.k
     base_args = (
         args.block_M,
         args.block_N,
@@ -373,24 +375,40 @@ def ref_program(A, B):
         args.num_stages,
     )
 
-    for M, N, K in ((4096, 4096, 4096), (8192, 8192, 8192), (16384, 16384, 16384)):
-        a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
-        b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
-        ref = ref_program(a, b)
-        flops = 2 * M * N * K
-
-        c = gemm_clc_persistent_2cta(a, b, *base_args, group_size, use_tma_store)
-        torch.testing.assert_close(c, ref, rtol=1e-2, atol=1e-2)
-        ms_base = do_bench(lambda a=a, b=b: gemm_clc_persistent_2cta(a, b, *base_args, group_size, use_tma_store), backend="event")
-
-        ms_torch = do_bench(lambda a=a, b=b: a @ b, backend="event")
-        print(f"M=N=K={M:<6}  base: {flops / (ms_base / 1e3) / 1e12:6.1f} TFLOPS  torch: {flops / (ms_torch / 1e3) / 1e12:6.1f} TFLOPS")
-
-        for clc in (2, 3, 4):
-            c2 = gemm_clc_persistent_2cta_pipelined_clc(a, b, *base_args, clc, group_size, use_tma_store)
-            torch.testing.assert_close(c2, ref, rtol=1e-2, atol=1e-2)
-            ms_pipe = do_bench(
-                lambda a=a, b=b, clc=clc: gemm_clc_persistent_2cta_pipelined_clc(a, b, *base_args, clc, group_size, use_tma_store),
-                backend="event",
-            )
-            print(f"             pipe(clc={clc}): {flops / (ms_pipe / 1e3) / 1e12:6.1f} TFLOPS  ({ms_base / ms_pipe:.3f}x vs base)")
+    a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
+    b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
+    ref = ref_program(a, b)
+    flops = 2 * M * N * K
+
+    c = gemm_clc_persistent_2cta(a, b, *base_args, args.group_size, args.use_tma_store)
+    torch.testing.assert_close(c, ref, rtol=1e-2, atol=1e-2)
+    ms_base = do_bench(
+        lambda: gemm_clc_persistent_2cta(a, b, *base_args, args.group_size, args.use_tma_store),
+        backend="event",
+    )
+
+    ms_torch = do_bench(lambda: a @ b, backend="event")
+    print(f"M={M}, N={N}, K={K}  base: {flops / (ms_base / 1e3) / 1e12:6.1f} TFLOPS  torch: {flops / (ms_torch / 1e3) / 1e12:6.1f} TFLOPS")
+
+    for clc_stages in (2, 3, 4):
+        c2 = gemm_clc_persistent_2cta_pipelined_clc(
+            a,
+            b,
+            *base_args,
+            clc_stages,
+            args.group_size,
+            args.use_tma_store,
+        )
+        torch.testing.assert_close(c2, ref, rtol=1e-2, atol=1e-2)
+        ms_pipe = do_bench(
+            lambda clc_stages=clc_stages: gemm_clc_persistent_2cta_pipelined_clc(
+                a,
+                b,
+                *base_args,
+                clc_stages,
+                args.group_size,
+                args.use_tma_store,
+            ),
+            backend="event",
+        )
+        print(f"             pipe(clc={clc_stages}): {flops / (ms_pipe / 1e3) / 1e12:6.1f} TFLOPS  ({ms_base / ms_pipe:.3f}x vs base)")

examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py → examples/gemm_tcgen05/gemm_tcgen5mma_ws_persistent.py

@@ -1,4 +1,4 @@
-# Persistent, warp-specialized SM100 GEMM using PersistentTileScheduler.
+# Persistent, warp-specialized TCGEN05 GEMM using PersistentTileScheduler.
 #
 # Each warp role (TMA / tcgen5 MMA / epilogue) creates its own scheduler instance
 # and drives a single ``while sched.valid()`` loop. The scheduler owns the only
@@ -28,6 +28,7 @@ def gemm_persistent(
     out_dtype,
     accum_dtype,
     num_stages,
+    group_size,
     use_tma_store,
 ):
     M, N, K = T.const("M, N, K")
@@ -41,7 +42,6 @@ def gemm_persistent(
     n_blocks = T.ceildiv(N, block_N)
     assert K % (2 * block_K) == 0  # for simplicity
     k_blocks = T.ceildiv(K, block_K)
-    group_size = 8
     assert n_blocks % (2 * group_size) == 0  # Please adjust group_size if not satisfied
 
     with T.Kernel(sm_num, threads=256) as (block_id):
@@ -145,6 +145,7 @@ def gemm_persistent_2cta(
     out_dtype,
     accum_dtype,
     num_stages,
+    group_size,
     use_tma_store,
 ):
     M, N, K = T.const("M, N, K")
@@ -158,7 +159,6 @@ def gemm_persistent_2cta(
     n_blocks = T.ceildiv(N, block_N)
     assert K % (2 * block_K) == 0  # for simplicity
     k_blocks = T.ceildiv(K, block_K)
-    group_size = 8  # in cluster
     assert n_blocks % (2 * group_size) == 0  # Please adjust group_size if not satisfied
     cluster_size = 2
 
@@ -257,7 +257,7 @@ def gemm_persistent_2cta(
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Persistent warp-specialized SM100 GEMM")
+    parser = argparse.ArgumentParser(description="Persistent warp-specialized TCGEN05 GEMM")
     parser.add_argument("--m", type=int, default=8192)
     parser.add_argument("--n", type=int, default=8192)
     parser.add_argument("--k", type=int, default=8192)
@@ -266,6 +266,7 @@ def main():
     parser.add_argument("--block_K", type=int, default=64)
     parser.add_argument("--store_block_N", type=int, default=64, help="block_N for C_shared")
     parser.add_argument("--num_stages", type=int, default=None, help="pipeline stages (default: 6 with 2cta, else 4)")
+    parser.add_argument("--group_size", type=int, default=8)
     parser.add_argument("--enable_2cta_tcgen5mma", action=argparse.BooleanOptionalAction, default=True)
     parser.add_argument("--use_tma_store", action=argparse.BooleanOptionalAction, default=True)
     args = parser.parse_args()
@@ -287,25 +288,33 @@ def main():
         "out_dtype": out_dtype,
         "accum_dtype": accum_dtype,
         "num_stages": num_stages,
+        "group_size": args.group_size,
         "use_tma_store": args.use_tma_store,
     }
 
     a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
     b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
-    print(kernel.get_kernel_source(a, b, **kwargs))
-    c = kernel(a, b, **kwargs)
+    kernel_args = (a, b)
+
+    print(kernel.get_kernel_source(*kernel_args, **kwargs))
+    c = kernel(*kernel_args, **kwargs)
 
     ref_c = (a.to(torch.float) @ b.to(torch.float)).to(torch.bfloat16)
     torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
     print("All checks passed. ✅")
 
     tl_latency = do_bench(
-        lambda: kernel(a, b, **kwargs),
+        lambda: kernel(*kernel_args, **kwargs),
+        _n_warmup=50,
+        _n_repeat=50,
+        backend="cupti",
+    )
+    torch_latency = do_bench(
+        lambda: a @ b,
         _n_warmup=50,
         _n_repeat=50,
         backend="cupti",
     )
-    torch_latency = do_bench(lambda: a @ b, backend="cupti")
     print(f"Tilelang latency: {tl_latency} ms")
     print(f"Flops: {2 * M * N * K / (tl_latency / 1e3) / 1e12} TFLOPS")
     print(f"Torch latency: {torch_latency} ms")