v4.3 tag release update. (#2789)

examples/python/CuTeDSL/ampere/call_bypass_dlpack.py

@@ -74,9 +74,11 @@ with stride-1 which propagate alignment incorrectly.
 
 """
 
-# Add the current directory to sys.path
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-from tensorop_gemm import TensorOpGemm
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from ampere.tensorop_gemm import TensorOpGemm
 
 
 @cute.jit

examples/python/CuTeDSL/ampere/call_from_jit.py

@@ -67,10 +67,11 @@ import cutlass.cute as cute
 from cutlass.torch import dtype as torch_dtype
 from cutlass.cute.runtime import make_ptr
 
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
 
-# Add the current directory to sys.path
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-from tensorop_gemm import TensorOpGemm
+from ampere.tensorop_gemm import TensorOpGemm
 
 
 class BufferWithLayout:

examples/python/CuTeDSL/ampere/call_with_tvm_ffi.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import sys
+import os
+import torch
+import time
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+"""Demonstrates calling off-the-shelf kernels with TVM FFI without DLPack.
+
+This example shows how to compile CuTe JIT function with fake tensors then run it with TVM FFI.
+"""
+
+if __name__ == "__main__":
+    # Add the current directory to sys.path
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+from ampere.tensorop_gemm import TensorOpGemm
+
+
+def compile_op(use_tvm_ffi: bool = True):
+    from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor
+
+    a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
+    b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
+    c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
+    a = make_fake_compact_tensor(
+        cutlass.Float16, a_shape, stride_order=(1, 0, 2), assumed_align=16
+    )
+    b = make_fake_compact_tensor(
+        cutlass.Float16, b_shape, stride_order=(1, 0, 2), assumed_align=16
+    )
+    c = make_fake_compact_tensor(
+        cutlass.Float16, c_shape, stride_order=(1, 0, 2), assumed_align=16
+    )
+
+    tensor_op_gemm = TensorOpGemm(
+        cutlass.Float16, cutlass.Float16, cutlass.Float32, (2, 2, 1)
+    )
+    compiled_fn = cute.compile(
+        tensor_op_gemm, a, b, c, options="--enable-tvm-ffi" if use_tvm_ffi else ""
+    )
+    return compiled_fn
+
+
+def run_op(compiled_fn, mnkl, *, use_tvm_ffi: bool = True):
+    print("\nRunning TensorOpGemm test with:")
+    print(f"Tensor dimensions: {mnkl}")
+    torch.manual_seed(1112)
+    # (M,K,L)
+    a_torch = torch.randn(
+        mnkl[3], mnkl[0], mnkl[2], dtype=torch.float16, device="cuda"
+    ).permute(1, 2, 0)
+    # (N,K,L)
+    b_torch = torch.randn(
+        mnkl[3], mnkl[1], mnkl[2], dtype=torch.float16, device="cuda"
+    ).permute(1, 2, 0)
+    # (N,M,L)
+    c_torch = torch.randn(
+        mnkl[3], mnkl[0], mnkl[1], dtype=torch.float16, device="cuda"
+    ).permute(1, 2, 0)
+
+    print("Input tensor shapes:")
+    print(f"a: {a_torch.shape}, dtype: {a_torch.dtype}")
+    print(f"b: {b_torch.shape}, dtype: {b_torch.dtype}")
+    print(f"c: {c_torch.shape}, dtype: {c_torch.dtype}\n")
+    if not use_tvm_ffi:
+        a = from_dlpack(a_torch).mark_layout_dynamic(leading_dim=1)
+        b = from_dlpack(b_torch).mark_layout_dynamic(leading_dim=1)
+        c = from_dlpack(c_torch).mark_layout_dynamic(leading_dim=1)
+    else:
+        a = a_torch
+        b = b_torch
+        c = c_torch
+    # pass in torch tensor as input
+    compiled_fn(a, b, c)
+    torch.cuda.synchronize()
+
+    # measure the launch overhead of tvm ffi function
+    repeat = 100
+    start_time = time.time()
+    for i in range(repeat):
+        compiled_fn(a, b, c)
+    end_time = time.time()
+    print(
+        f"Launch overhead of tvm ffi function: {(end_time - start_time) / repeat} seconds"
+    )
+
+    ref = torch.einsum("mkl,nkl->mnl", a_torch, b_torch)
+    torch.testing.assert_close(c_torch, ref, atol=1e-05, rtol=1e-05)
+    print("\n[DSL INFO] Results verified successfully!")
+    print(f"First few elements of result: \n{c_torch[:3, :3, :3]}")
+
+
+if __name__ == "__main__":
+    compiled_fn = compile_op(use_tvm_ffi=False)
+    run_op(compiled_fn, [512, 512, 256, 1], use_tvm_ffi=False)
+    compiled_fn = compile_op(use_tvm_ffi=True)
+    run_op(compiled_fn, [512, 512, 256, 1], use_tvm_ffi=True)

examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py

@@ -39,6 +39,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.utils.blackwell_helpers as sm100_utils
 
 import math
@@ -220,22 +221,18 @@ class BlockwiseGemmKernel:
         self.num_regs_epilogue_warps = 216
         self.num_regs_acc_update_warps = 216
 
-        # Set barrier for cta sync, epilogue sync and tmem ptr sync
-        self.cta_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=1,
-            num_threads=self.threads_per_cta,
-        )
+        # Set barrier for epilogue sync and tmem ptr sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=2,
+            barrier_id=1,
             num_threads=32 * len(self.epilog_warp_id),
         )
         self.tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=3,
+            barrier_id=2,
             num_threads=32
             * len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)),
         )
         self.sched_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=4,
+            barrier_id=3,
             num_threads=self.threads_per_warp,
         )
         self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -703,6 +700,7 @@ class BlockwiseGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize mainloop scale_pipeline (barrier) and states
@@ -719,6 +717,7 @@ class BlockwiseGemmKernel:
             num_stages=self.num_scale_stage,
             producer_group=scale_pipeline_producer_group,
             consumer_group=scale_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Initialize acc_pipeline (barrier) and states
@@ -735,6 +734,7 @@ class BlockwiseGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize epilogue pipeline (barrier) and states
@@ -751,6 +751,7 @@ class BlockwiseGemmKernel:
             num_stages=1,
             producer_group=epi_pipeline_producer_group,
             consumer_group=epi_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Initialize tile info pipeline (barrier) and states
@@ -767,6 +768,7 @@ class BlockwiseGemmKernel:
             num_stages=self.num_tile_stage,
             producer_group=tile_info_pipeline_producer_group,
             consumer_group=tile_info_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Tensor memory dealloc barrier init
@@ -779,8 +781,7 @@ class BlockwiseGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C/Scale
@@ -968,10 +969,7 @@ class BlockwiseGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Specialized Schedule warp

examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py

@@ -39,6 +39,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.cute.runtime import from_dlpack
 
@@ -236,22 +237,18 @@ class BlockwiseContiguousGroupedGemmKernel:
         self.num_regs_epilogue_warps = 216
         self.num_regs_acc_update_warps = 216
 
-        # Set barrier for cta sync, epilogue sync and tmem ptr sync
-        self.cta_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=1,
-            num_threads=self.threads_per_cta,
-        )
+        # Set barrier for epilogue sync and tmem ptr sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=2,
+            barrier_id=1,
             num_threads=32 * len(self.epilog_warp_id),
         )
         self.tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=3,
+            barrier_id=2,
             num_threads=32
             * len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)),
         )
         self.sched_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=4,
+            barrier_id=3,
             num_threads=self.threads_per_warp,
         )
         self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -724,6 +721,7 @@ class BlockwiseContiguousGroupedGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize mainloop scale_pipeline (barrier) and states
@@ -740,6 +738,7 @@ class BlockwiseContiguousGroupedGemmKernel:
             num_stages=self.num_scale_stage,
             producer_group=scale_pipeline_producer_group,
             consumer_group=scale_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Initialize acc_pipeline (barrier) and states
@@ -756,6 +755,7 @@ class BlockwiseContiguousGroupedGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize epilogue pipeline (barrier) and states
@@ -772,6 +772,7 @@ class BlockwiseContiguousGroupedGemmKernel:
             num_stages=1,
             producer_group=epi_pipeline_producer_group,
             consumer_group=epi_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Initialize tile info pipeline (barrier) and states
@@ -788,6 +789,7 @@ class BlockwiseContiguousGroupedGemmKernel:
             num_stages=self.num_tile_stage,
             producer_group=tile_info_pipeline_producer_group,
             consumer_group=tile_info_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Tensor memory dealloc barrier init
@@ -800,8 +802,7 @@ class BlockwiseContiguousGroupedGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C/Scale
@@ -989,10 +990,7 @@ class BlockwiseContiguousGroupedGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Specialized Schedule warp

examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py

@@ -39,6 +39,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.cute.runtime import from_dlpack
 
@@ -235,22 +236,18 @@ class BlockwiseMaskedGroupedGemmKernel:
         self.num_regs_epilogue_warps = 216
         self.num_regs_acc_update_warps = 216
 
-        # Set barrier id for cta sync, epilogue sync and tmem ptr sync
-        self.cta_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=1,
-            num_threads=self.threads_per_cta,
-        )
+        # Set barrier id for epilogue sync and tmem ptr sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=2,
+            barrier_id=1,
             num_threads=32 * len(self.epilog_warp_id),
         )
         self.tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=3,
+            barrier_id=2,
             num_threads=32
             * len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)),
         )
         self.sched_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=4,
+            barrier_id=3,
             num_threads=self.threads_per_warp,
         )
         self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -723,6 +720,7 @@ class BlockwiseMaskedGroupedGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize mainloop scale_pipeline (barrier) and states
@@ -739,6 +737,7 @@ class BlockwiseMaskedGroupedGemmKernel:
             num_stages=self.num_scale_stage,
             producer_group=scale_pipeline_producer_group,
             consumer_group=scale_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Initialize acc_pipeline (barrier) and states
@@ -755,6 +754,7 @@ class BlockwiseMaskedGroupedGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize epilogue pipeline (barrier) and states
@@ -771,6 +771,7 @@ class BlockwiseMaskedGroupedGemmKernel:
             num_stages=1,
             producer_group=epi_pipeline_producer_group,
             consumer_group=epi_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Initialize tile info pipeline (barrier) and states
@@ -787,6 +788,7 @@ class BlockwiseMaskedGroupedGemmKernel:
             num_stages=self.num_tile_stage,
             producer_group=tile_info_pipeline_producer_group,
             consumer_group=tile_info_pipeline_consumer_group,
+            defer_sync=True,
         )
 
         # Tensor memory dealloc barrier init
@@ -799,8 +801,7 @@ class BlockwiseMaskedGroupedGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C/Scale
@@ -988,10 +989,7 @@ class BlockwiseMaskedGroupedGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Specialized Schedule warp

examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py

@@ -38,6 +38,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.utils.blockscaled_layout as blockscaled_utils
 from cutlass.cute.runtime import from_dlpack
@@ -208,17 +209,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
         self.threads_per_cta = 32 * len(
             (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
         )
-        # Set barrier id for cta sync, epilogue sync and tmem ptr sync
-        self.cta_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=1,
-            num_threads=self.threads_per_cta,
-        )
+        # Set barrier id for epilogue sync and tmem ptr sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=2,
+            barrier_id=1,
             num_threads=32 * len(self.epilog_warp_id),
         )
         self.tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=3,
+            barrier_id=2,
             num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
         )
         self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -288,6 +285,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
             self.mma_tiler[1],
             self.mma_tiler[2],
         )
+        self.cta_tile_shape_mnk_sfb = (
+            self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
+            self.mma_tiler_sfb[1],
+            self.mma_tiler_sfb[2],
+        )
 
         # Compute cluster layout
         self.cluster_layout_vmnk = cute.tiled_divide(
@@ -314,6 +316,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
             self.c_layout,
             self.c_dtype,
         )
+        self.epi_tile_n = cute.size(self.epi_tile[1])
 
         # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory
         self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages(
@@ -362,6 +365,19 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
             self.num_c_stage,
         )
 
+        # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
+        self.overlapping_accum = self.num_acc_stage == 1
+
+        # Compute number of TMEM columns for SFA/SFB/Accumulator
+        sf_atom_mn = 32
+        self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
+        self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
+        self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
+        self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
+
+        # Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue
+        self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n
+
     @cute.jit
     def __call__(
         self,
@@ -640,6 +656,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
             block=[self.threads_per_cta, 1, 1],
             cluster=(*self.cluster_shape_mn, 1),
             stream=stream,
+            min_blocks_per_mp=1,
         )
         return
 
@@ -726,6 +743,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize acc_pipeline (barrier) and states
@@ -742,6 +760,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Tensor memory dealloc barrier init
@@ -754,8 +773,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/SFA/SFB/C
@@ -910,18 +928,34 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
         tCrB = tiled_mma.make_fragment_B(sB)
         # (MMA, MMA_M, MMA_N)
         acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
-        # (MMA, MMA_M, MMA_N, STAGE)
-        tCtAcc_fake = tiled_mma.make_fragment_C(
-            cute.append(acc_shape, self.num_acc_stage)
-        )
+        if cutlass.const_expr(self.overlapping_accum):
+            num_acc_stage_overlapped = 2
+            tCtAcc_fake = tiled_mma.make_fragment_C(
+                cute.append(acc_shape, num_acc_stage_overlapped)
+            )
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtAcc_fake = cute.make_tensor(
+                tCtAcc_fake.iterator,
+                cute.make_layout(
+                    tCtAcc_fake.shape,
+                    stride = (
+                        tCtAcc_fake.stride[0],
+                        tCtAcc_fake.stride[1],
+                        tCtAcc_fake.stride[2],
+                        (256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1]
+                    ) 
+                )
+            )
+        else:
+            # (MMA, MMA_M, MMA_N, STAGE)
+            tCtAcc_fake = tiled_mma.make_fragment_C(
+                cute.append(acc_shape, self.num_acc_stage)
+            )
 
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Specialized TMA load warp
@@ -1057,7 +1091,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
 
             # Make SFA tmem tensor
             sfa_tmem_ptr = cute.recast_ptr(
-                acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base),
+                acc_tmem_ptr + self.num_accumulator_tmem_cols,
                 dtype=self.sf_dtype,
             )
             # (MMA, MMA_M, MMA_K)
@@ -1071,9 +1105,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
 
             # Make SFB tmem tensor
             sfb_tmem_ptr = cute.recast_ptr(
-                acc_tmem_ptr
-                + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
-                + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
+                acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
                 dtype=self.sf_dtype,
             )
             # (MMA, MMA_N, MMA_K)
@@ -1122,9 +1154,15 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                     cur_tile_coord[2],
                 )
 
+                # Get accumulator stage index
+                if cutlass.const_expr(self.overlapping_accum):
+                    acc_stage_index = acc_producer_state.phase ^ 1
+                else:
+                    acc_stage_index = acc_producer_state.index
+
                 # Set tensor memory buffer for current tile
                 # (MMA, MMA_M, MMA_N)
-                tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
+                tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
 
                 # Peek (try_wait) AB buffer full for k_tile = 0
                 ab_consumer_state.reset_count()
@@ -1146,8 +1184,8 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                     offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
                     shifted_ptr = cute.recast_ptr(
                         acc_tmem_ptr
-                        + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
-                        + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+                         + self.num_accumulator_tmem_cols
+                         + self.num_sfa_tmem_cols
                         + offset,
                         dtype=self.sf_dtype,
                     )
@@ -1156,9 +1194,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                     # Move in increments of 64 columns of SFB
                     offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
                     shifted_ptr = cute.recast_ptr(
-                        acc_tmem_ptr
-                        + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
-                        + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+                        acc_tmem_ptr 
+                        + self.num_accumulator_tmem_cols
+                        + self.num_sfa_tmem_cols
                         + offset,
                         dtype=self.sf_dtype,
                     )
@@ -1350,10 +1388,17 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                     )
                 ]
 
+                # Get accumulator stage index
+                if cutlass.const_expr(self.overlapping_accum):
+                    acc_stage_index = acc_consumer_state.phase
+                    reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
+                else:
+                    acc_stage_index = acc_consumer_state.index
+
                 # Set tensor memory buffer for current tile
                 # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
                 tTR_tAcc = tTR_tAcc_base[
-                    (None, None, None, None, None, acc_consumer_state.index)
+                    (None, None, None, None, None, acc_stage_index)
                 ]
 
                 #
@@ -1370,12 +1415,27 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                 subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
                 num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
                 for subtile_idx in cutlass.range(subtile_cnt):
+                    real_subtile_idx = subtile_idx
+                    if cutlass.const_expr(self.overlapping_accum):
+                        if reverse_subtile:
+                            real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx
                     #
                     # Load accumulator from tensor memory buffer to register
                     #
-                    tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
+                    tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
                     cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
 
+                    #
+                    # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled
+                    #
+                    if cutlass.const_expr(self.overlapping_accum):
+                        if subtile_idx == self.iter_acc_early_release_in_epilogue:
+                            # Fence for TMEM load
+                            cute.arch.fence_view_async_tmem_load()
+                            with cute.arch.elect_one():
+                                acc_pipeline.consumer_release(acc_consumer_state)
+                            acc_consumer_state.advance()
+
                     #
                     # Convert to C type
                     #
@@ -1386,7 +1446,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                     #
                     # Store C to shared memory
                     #
-                    c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage
+                    c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage
                     cute.copy(
                         tiled_copy_r2s,
                         tRS_rC,
@@ -1406,7 +1466,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                         cute.copy(
                             tma_atom_c,
                             bSG_sC[(None, c_buffer)],
-                            bSG_gC[(None, subtile_idx)],
+                            bSG_gC[(None, real_subtile_idx)],
                         )
                         # Fence and barrier to make sure shared memory store is visible to TMA store
                         c_pipeline.producer_commit()
@@ -1416,9 +1476,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
                 #
                 # Async arrive accumulator buffer empty
                 #
-                with cute.arch.elect_one():
-                    acc_pipeline.consumer_release(acc_consumer_state)
-                acc_consumer_state.advance()
+                if cutlass.const_expr(not self.overlapping_accum):
+                    with cute.arch.elect_one():
+                        acc_pipeline.consumer_release(acc_consumer_state)
+                    acc_consumer_state.advance()
 
                 #
                 # Advance to next tile
@@ -2286,6 +2347,7 @@ def run(
         c_tensor,
         max_active_clusters,
         current_stream,
+        options=f"--opt-level 2",
     )
 
     # Compute reference result

examples/python/CuTeDSL/blackwell/dense_gemm.py

@@ -36,6 +36,7 @@ import cutlass
 import cutlass.cute as cute
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils.blackwell_helpers as sm100_utils
@@ -536,6 +537,7 @@ class DenseGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         ).make_participants()
 
         # Initialize acc_pipeline (barrier) and states
@@ -549,6 +551,7 @@ class DenseGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
         acc_producer_state = pipeline.make_pipeline_state(
             pipeline.PipelineUserType.Producer, self.num_acc_stage
@@ -569,8 +572,7 @@ class DenseGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C
@@ -686,8 +688,7 @@ class DenseGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
+        pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
 
         # Alloc tensor memory buffer
         tmem.allocate(self.num_tmem_alloc_cols)

examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py

@@ -38,6 +38,7 @@ import cutlass.cute.testing as testing
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.cute.nvgpu import cpasync, tcgen05
 
@@ -232,9 +233,8 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
             )
         )
         # Set barrier id for cta sync, epilogue sync and tmem ptr sync
-        self.cta_sync_bar_id = 1
-        self.epilog_sync_bar_id = 2
-        self.tmem_alloc_sync_bar_id = 3
+        self.epilog_sync_bar_id = 1
+        self.tmem_alloc_sync_bar_id = 2
         self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
 
     def _setup_attributes(self):
@@ -635,6 +635,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Initialize acc_pipeline (barrier) and states
@@ -651,6 +652,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Load C pipeline
@@ -665,6 +667,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
             producer_group=c_producer_group,
             consumer_group=c_consumer_group,
             tx_count=self.tma_c_load_bytes,
+            defer_sync=True,
         )
 
         tmem_alloc_barrier = pipeline.NamedBarrier(
@@ -681,8 +684,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C/D
@@ -796,9 +798,6 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
 
         # Named barriers
         #
-        cta_sync_barrier = pipeline.NamedBarrier(
-            self.cta_sync_bar_id, self.threads_per_cta
-        )
         epilog_sync_barrier = pipeline.NamedBarrier(
             self.epilog_sync_bar_id, 32 * len(self.epilog_warp_ids)
         )
@@ -806,10 +805,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Specialized TMA load warp

examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py

@@ -29,16 +29,14 @@
 import argparse
 from typing import Optional, Tuple, Type, Union
 
-import torch
 import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
 import cutlass.cute.testing as testing
-import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
-import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 from cutlass.cute.nvgpu import cpasync, tcgen05
 
 """
@@ -152,10 +150,10 @@ def _compute_stages(
     num_c_stage = 2 if use_tma_store else 0
 
     # Calculate smem layout and size for one stage of A, B, and C with 1-stage
-    a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
+    a_smem_layout_stage_one = utils.sm100.make_smem_layout_a(
         tiled_mma, mma_tiler_mnk, a_dtype, 1
     )
-    b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
+    b_smem_layout_staged_one = utils.sm100.make_smem_layout_b(
         tiled_mma, mma_tiler_mnk, b_dtype, 1
     )
 
@@ -229,14 +227,14 @@ class PersistentDenseGemmKernel:
         - Cluster shape M must be multiple of 2 if use_2cta_instrs=True
         - Cluster shape M/N must be positive and power of 2, total cluster size <= 16
 
-    Example:
-        >>> gemm = PersistentDenseGemmKernel(
-        ...     acc_dtype=cutlass.Float32,
-        ...     use_2cta_instrs=True,
-        ...     mma_tiler_mn=(128, 128),
-        ...     cluster_shape_mn=(2, 2)
-        ... )
-        >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream)
+    **Example:**
+        gemm = PersistentDenseGemmKernel(
+            acc_dtype=cutlass.Float32,
+            use_2cta_instrs=True,
+            mma_tiler_mn=(128, 128),
+            cluster_shape_mn=(2, 2)
+        )
+        gemm(a, b, c, max_active_clusters, stream)
     """
 
     def __init__(
@@ -316,7 +314,7 @@ class PersistentDenseGemmKernel:
         - Computing tensor memory allocation columns
         """
         # Configure tiled mma
-        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+        tiled_mma = utils.sm100.make_trivial_tiled_mma(
             self.a_dtype,
             self.a_major_mode,
             self.b_major_mode,
@@ -353,7 +351,7 @@ class PersistentDenseGemmKernel:
 
         # Compute epilogue subtile
         if cutlass.const_expr(self.use_tma_store):
-            self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
+            self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
                 self.cta_tile_shape_mnk,
                 self.use_2cta_instrs,
                 self.c_layout,
@@ -364,7 +362,7 @@ class PersistentDenseGemmKernel:
 
         c_smem_layout = None
         if cutlass.const_expr(self.use_tma_store):
-            c_smem_layout = sm100_utils.make_smem_layout_epi(
+            c_smem_layout = utils.sm100.make_smem_layout_epi(
                 self.c_dtype, self.c_layout, self.epi_tile, 1
             )
 
@@ -382,16 +380,16 @@ class PersistentDenseGemmKernel:
         )
 
         # Compute A/B/C shared memory layout
-        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+        self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
             tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage
         )
-        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+        self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
             tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage
         )
 
         self.c_smem_layout_staged = None
         if self.use_tma_store:
-            self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
+            self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi(
                 self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage
             )
 
@@ -447,7 +445,7 @@ class PersistentDenseGemmKernel:
         # Setup attributes that dependent on gemm inputs
         self._setup_attributes()
 
-        tiled_mma = sm100_utils.make_trivial_tiled_mma(
+        tiled_mma = utils.sm100.make_trivial_tiled_mma(
             self.a_dtype,
             self.a_major_mode,
             self.b_major_mode,
@@ -458,7 +456,7 @@ class PersistentDenseGemmKernel:
         atom_thr_size = cute.size(tiled_mma.thr_id.shape)
 
         # Setup TMA load for A
-        a_op = sm100_utils.cluster_shape_to_tma_atom_A(
+        a_op = utils.sm100.cluster_shape_to_tma_atom_A(
             self.cluster_shape_mn, tiled_mma.thr_id
         )
         a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
@@ -475,7 +473,7 @@ class PersistentDenseGemmKernel:
         )
 
         # Setup TMA load for B
-        b_op = sm100_utils.cluster_shape_to_tma_atom_B(
+        b_op = utils.sm100.cluster_shape_to_tma_atom_B(
             self.cluster_shape_mn, tiled_mma.thr_id
         )
         b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
@@ -614,6 +612,7 @@ class PersistentDenseGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         ).make_participants()
 
         # Initialize acc_pipeline (barrier) and states
@@ -630,6 +629,7 @@ class PersistentDenseGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         tmem_alloc_barrier = pipeline.NamedBarrier(
@@ -652,8 +652,7 @@ class PersistentDenseGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C
@@ -761,10 +760,7 @@ class PersistentDenseGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
 
         #
         # Specialized TMA load warp
@@ -1297,7 +1293,7 @@ class PersistentDenseGemmKernel:
         :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
         """
         # Make tiledCopy for tensor memory load
-        copy_atom_t2r = sm100_utils.get_tmem_load_op(
+        copy_atom_t2r = utils.sm100.get_tmem_load_op(
             self.cta_tile_shape_mnk,
             self.c_layout,
             self.c_dtype,
@@ -1354,7 +1350,7 @@ class PersistentDenseGemmKernel:
             - tRS_sC: The partitioned tensor C (smem destination)
         :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
         """
-        copy_atom_r2s = sm100_utils.get_smem_store_op(
+        copy_atom_r2s = utils.sm100.get_smem_store_op(
             self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r
         )
         tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
@@ -1617,97 +1613,187 @@ class PersistentDenseGemmKernel:
                 is_valid = False
         return is_valid
 
-    def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool:
-        """Check if the given tensors can be implemented by this kernel.
+    def can_implement(
+        self,
+        mnkl: Tuple[int, int, int, int],
+        ab_dtype: Type[cutlass.Numeric],
+        c_dtype: Type[cutlass.Numeric],
+        a_major: str,
+        b_major: str,
+        c_major: str,
+    ) -> bool:
+        """
+        Determine if the given tensor configuration can be implemented by this kernel.
 
-        :param a: Input tensor A
-        :type a: cute.Tensor
-        :param b: Input tensor B
-        :type b: cute.Tensor
-        :param c: Output tensor C
-        :type c: cute.Tensor
-
-        :return: True if the gemm supports the given config, False otherwise
+        :param mnkl: Problem size as a tuple (M, N, K, L).
+        :type mnkl: Tuple[int, int, int, int]
+        :param ab_dtype: Data type for input tensors A and B.
+        :type ab_dtype: Type[cutlass.Numeric]
+        :param c_dtype: Data type for output tensor C.
+        :type c_dtype: Type[cutlass.Numeric]
+        :param a_major: Major dimension of the A tensor layout ("m" or "k").
+        :type a_major: str
+        :param b_major: Major dimension of the B tensor layout ("n" or "k").
+        :type b_major: str
+        :param c_major: Major dimension of the C tensor layout ("m" or "n").
+        :type c_major: str
+        :return: True if the kernel supports the given configuration, False otherwise.
         :rtype: bool
         """
-        m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2]
 
-        # infer a_major, b_major, c_major
-        is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a()
-        is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b()
-        is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c()
-        a_major = "m" if is_m_major_a else "k"
-        b_major = "n" if is_n_major_b else "k"
-        c_major = "m" if is_m_major_c else "n"
-
-        can_implement = True
         # Skip unsupported types
-        if not self.is_valid_dtypes(a.element_type, c.element_type):
-            can_implement = False
+        if not self.is_valid_dtypes(ab_dtype, c_dtype):
+            return False
+
         # Skip invalid mma tile shape and cluster shape
         if not self.is_valid_mma_tiler_and_cluster_shape():
-            can_implement = False
+            return False
+
+        # Unpack mnkl for clarity in calling the epilog check
+        m, n, k, l = mnkl
         # Skip illegal problem shape for load/store alignment
         if not self.is_valid_tensor_alignment(
-            m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major
+            m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major
         ):
-            can_implement = False
+            return False
         # Skip invalid epilogue store option
         if not self.is_valid_epilog_store_option(m, n):
-            can_implement = False
+            return False
 
-        return can_implement
+        return True
 
 
-def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype):
-    torch.manual_seed(1111)
+@cute.jit
+def bmm(
+    gemm_op: cutlass.Constexpr,
+    a: cute.Tensor,  # (l, m, k)
+    b: cute.Tensor,  # (l, k, n)
+    c: cute.Tensor,  # (l, m, n)
+    max_active_clusters: cutlass.Constexpr,
+    stream: cuda.CUstream,
+    epilogue_op: cutlass.Constexpr = lambda x: x,
+):
+    """
+    Wrapper API for persistent GEMM kernel to follow the convention of PyTorch's batch matrix-multiply (bmm).
 
-    a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype)
-    b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype)
-    c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype)
+    Internally, the tensors are permuted to match CuTe's convention:
+      - a: (m, k, l)
+      - b: (n, k, l)
+      - c: (m, n, l)
 
-    a_tensor, _ = cutlass_torch.cute_tensor_like(
-        a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
+    :param gemm_op: Kernel operation, expects (a, b, c, max_active_clusters, stream, epilogue_op)
+    :type gemm_op: cutlass.Constexpr
+    :param a: Input tensor of shape (l, m, k)
+    :type a: cute.Tensor
+    :param b: Input tensor of shape (l, k, n)
+    :type b: cute.Tensor
+    :param c: Output tensor of shape (l, m, n)
+    :type c: cute.Tensor
+    :param max_active_clusters: Maximum number of hardware clusters to launch
+    :type max_active_clusters: cutlass.Constexpr
+    :param epilogue_op: Optional elementwise lambda function to apply per output element, defaults to identity
+    :type epilogue_op: cutlass.Constexpr, optional
+    """
+    # (l,m,k) -> (m,k,l)
+    a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0]))
+    # (l,k,n) -> (n,k,l)
+    b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0]))
+    # (l,m,n) -> (m,n,l)
+    c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0]))
+
+    gemm_op(a, b, c, max_active_clusters, stream, epilogue_op)
+
+
+def compile_bmm(
+    gemm_op: PersistentDenseGemmKernel,
+    a_dtype: Type[cutlass.Numeric],
+    b_dtype: Type[cutlass.Numeric],
+    c_dtype: Type[cutlass.Numeric],
+    a_major: str,
+    b_major: str,
+    c_major: str,
+    max_active_clusters: cutlass.Constexpr,
+    stream: cuda.CUstream,
+    epilogue_op: cutlass.Constexpr = lambda x: x,
+    options: str = "",
+):
+    from cutlass.cute.runtime import make_fake_compact_tensor
+
+    a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
+    b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
+    c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
+
+    if a_major == "k":
+        a_order = (2, 1, 0)  # k is leading dimension
+    elif a_major == "m":
+        a_order = (2, 0, 1)  # m is leading dimension
+
+    if b_major == "n":
+        b_order = (2, 1, 0)  # n is leading dimension
+    elif b_major == "k":
+        b_order = (2, 0, 1)  # k is leading dimension
+
+    if c_major == "n":
+        c_order = (2, 1, 0)  # n is leading dimension
+    elif c_major == "m":
+        c_order = (2, 0, 1)  # m is leading dimension
+
+    a = make_fake_compact_tensor(
+        a_dtype, a_shape, stride_order=a_order, assumed_align=16
     )
-    b_tensor, _ = cutlass_torch.cute_tensor_like(
-        b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
+    b = make_fake_compact_tensor(
+        b_dtype, b_shape, stride_order=b_order, assumed_align=16
     )
-    c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
-        c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16
+    c = make_fake_compact_tensor(
+        c_dtype, c_shape, stride_order=c_order, assumed_align=16
     )
 
+    return cute.compile(
+        bmm, gemm_op, a, b, c, max_active_clusters, stream, epilogue_op, options=options
+    )
+
+
+def prepare_tensors(
+    mnkl: Tuple[int, int, int, int],
+    ab_dtype: Type[cutlass.Numeric],
+    c_dtype: Type[cutlass.Numeric],
+    a_major: str,
+    b_major: str,
+    c_major: str,
+    init_random: bool = True,
+):
+    import torch
+    from cutlass.torch import dtype as torch_dtype
+
+    m, n, k, l = mnkl
+
+    if a_major == "k":
+        a = torch.empty((l, m, k), dtype=torch.float32, device="cuda")
+    elif a_major == "m":
+        a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
+
+    if b_major == "n":
+        b = torch.empty((l, k, n), dtype=torch.float32, device="cuda")
+    elif b_major == "k":
+        b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1)
+
+    if c_major == "n":
+        c = torch.empty((l, m, n), dtype=torch.float32, device="cuda")
+    elif c_major == "m":
+        c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
+
+    if init_random:
+        a.random_(-2, 3)
+        b.random_(-2, 3)
+        c.random_(-2, 3)
+
     return (
-        a_tensor,
-        b_tensor,
-        c_tensor,
-        a_torch_cpu,
-        b_torch_cpu,
-        c_torch_cpu,
-        c_torch_gpu,
+        a.to(dtype=torch_dtype(ab_dtype)),
+        b.to(dtype=torch_dtype(ab_dtype)),
+        c.to(dtype=torch_dtype(c_dtype)),
     )
 
 
-def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance):
-    # Copy gpu result back
-    kernel_result = c_torch_gpu.cpu()
-
-    # Compute reference result
-    ref = torch.einsum(
-        "mkl,nkl->mnl",
-        a_torch_cpu.to(dtype=torch.float32),
-        b_torch_cpu.to(dtype=torch.float32),
-    )
-
-    # Convert ref to c_dtype
-    _, ref_torch_gpu = cutlass_torch.cute_tensor_like(
-        ref, c_dtype, is_dynamic_layout=True, assumed_align=16
-    )
-    ref_result = ref_torch_gpu.cpu()
-
-    # Assert close results
-    torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05)
-
-
 def run(
     mnkl: Tuple[int, int, int, int],
     ab_dtype: Type[cutlass.Numeric],
@@ -1725,48 +1811,55 @@ def run(
     iterations: int = 1,
     skip_ref_check: bool = False,
     use_cold_l2: bool = False,
+    use_tvm_ffi: bool = False,
+    benchmark: bool = False,
     **kwargs,
 ):
-    """Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking.
+    """
+    Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking.
 
-    This function prepares input tensors, configures and launches the persistent GEMM kernel,
-    optionally performs reference validation, and benchmarks the execution performance.
+    Prepares input tensors, configures and launches the persistent GEMM kernel,
+    optionally performs reference validation, and benchmarks execution.
 
-    :param mnkl: Problem size (M, N, K, L)
+    :param mnkl: Problem size as a tuple (M, N, K, L).
     :type mnkl: Tuple[int, int, int, int]
-    :param ab_dtype: Data type for input tensors A and B
+    :param ab_dtype: Data type for input tensors A and B.
     :type ab_dtype: Type[cutlass.Numeric]
-    :param c_dtype: Data type for output tensor C
+    :param c_dtype: Data type for output tensor C.
     :type c_dtype: Type[cutlass.Numeric]
-    :param acc_dtype: Data type for accumulation during matrix multiplication
+    :param acc_dtype: Accumulator data type for the matrix multiplication.
     :type acc_dtype: Type[cutlass.Numeric]
-    :param a_major/b_major/c_major: Memory layout of tensor A/B/C
-    :type a_major/b_major/c_major: str
-    :param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the
-        default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters.
+    :param a_major: Memory layout of tensor A.
+    :type a_major: str
+    :param b_major: Memory layout of tensor B.
+    :type b_major: str
+    :param c_major: Memory layout of tensor C.
+    :type c_major: str
+    :param mma_tiler_mn: MMA tiling size (M, N), defaults to (256, 256).
     :type mma_tiler_mn: Tuple[int, int], optional
-    :param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the
-        default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters.
+    :param cluster_shape_mn: Cluster shape (M, N), defaults to (2, 1).
     :type cluster_shape_mn: Tuple[int, int], optional
-    :param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner
-        will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters.
+    :param use_2cta_instrs: Whether to use 2CTA MMA instructions, defaults to True.
     :type use_2cta_instrs: bool, optional
-    :param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use
-        the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters.
+    :param use_tma_store: Whether to use TMA store, defaults to True.
     :type use_tma_store: bool, optional
-    :param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01
+    :param tolerance: Tolerance for reference validation, defaults to 1e-01.
     :type tolerance: float, optional
-    :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0
+    :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0.
     :type warmup_iterations: int, optional
-    :param iterations: Number of benchmark iterations to run, defaults to 1
+    :param iterations: Number of benchmark iterations to run, defaults to 1.
     :type iterations: int, optional
-    :param skip_ref_check: Whether to skip reference result validation, defaults to False
+    :param skip_ref_check: Whether to skip reference result validation, defaults to False.
     :type skip_ref_check: bool, optional
-    :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False
+    :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False.
     :type use_cold_l2: bool, optional
-    :raises RuntimeError: If CUDA GPU is not available
-    :raises ValueError: If the configuration is invalid or unsupported by the kernel
-    :return: Execution time of the GEMM kernel
+    :param use_tvm_ffi: Whether to use TVM FFI for the kernel, defaults to False.
+    :type use_tvm_ffi: bool, optional
+    :param benchmark: Whether to only benchmark the kernel, defaults to False.
+    :type benchmark: bool, optional
+    :raises RuntimeError: If CUDA GPU is not available.
+    :raises ValueError: If the configuration is invalid or unsupported by the kernel.
+    :return: Execution time of the GEMM kernel.
     :rtype: float
     """
     print("Running Blackwell Persistent Dense GEMM test with:")
@@ -1781,9 +1874,24 @@ def run(
     print(f"Iterations: {iterations}")
     print(f"Skip reference checking: {skip_ref_check}")
     print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
+    print(f"Use TVM FFI: {'True' if use_tvm_ffi else 'False'}")
 
-    # Unpack parameters
-    m, n, k, l = mnkl
+    import torch
+    from cutlass.torch import dtype as torch_dtype
+
+    # Build GEMM object
+    gemm = PersistentDenseGemmKernel(
+        acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store
+    )
+    can_implement = gemm.can_implement(
+        mnkl, ab_dtype, c_dtype, a_major, b_major, c_major
+    )
+    if not can_implement:
+        raise testing.CantImplementError(
+            f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, "
+            f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, "
+            f"use_tma_store = {use_tma_store}"
+        )
 
     if not torch.cuda.is_available():
         raise RuntimeError("GPU is required to run this example!")
@@ -1793,59 +1901,75 @@ def run(
     # Get the raw stream pointer as a CUstream
     current_stream = cuda.CUstream(torch_stream.cuda_stream)
 
-    a_tensor, b_tensor, c_tensor, a_torch_cpu, b_torch_cpu, c_torch_cpu, c_torch_gpu = (
-        create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype)
-    )
-
-    # Build GEMM object
-    gemm = PersistentDenseGemmKernel(
-        acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store
-    )
-
     # Check if configuration can be implemented
-    can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor)
-    if not can_implement:
-        raise ValueError(
-            f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, "
-            f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, "
-            f"use_tma_store = {use_tma_store}"
-        )
     max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
         cluster_shape_mn[0] * cluster_shape_mn[1]
     )
-    compiled_gemm = cute.compile(
-        gemm, a_tensor, b_tensor, c_tensor, max_active_clusters, current_stream
+
+    options = []
+    if use_tvm_ffi:
+        options.append("--enable-tvm-ffi")
+
+    compiled_fn = compile_bmm(
+        gemm,
+        ab_dtype,
+        ab_dtype,
+        c_dtype,
+        a_major,
+        b_major,
+        c_major,
+        max_active_clusters,
+        current_stream,
+        options=",".join(options),
     )
 
+    # Run and verify BMM with torch
+    a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major)
+
     if not skip_ref_check:
-        compiled_gemm(a_tensor, b_tensor, c_tensor, current_stream)
-        compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance)
+        # Use small random number for deterministic result for reference check
+        compiled_fn(a, b, c, torch_stream)
+
+        # Manually quantize to be comparable
+        ref = (
+            torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32))
+            .to(dtype=torch_dtype(c_dtype))
+            .to(dtype=torch.float32)
+        )
+        torch.testing.assert_close(
+            c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03
+        )
+
+    if not benchmark:
+        return 0
 
     def generate_tensors():
-        a_tensor, _ = cutlass_torch.cute_tensor_like(
-            a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
+        init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8]
+        a, b, c = prepare_tensors(
+            mnkl,
+            ab_dtype,
+            c_dtype,
+            a_major,
+            b_major,
+            c_major,
+            init_random=not init_normal,
         )
-        b_tensor, _ = cutlass_torch.cute_tensor_like(
-            b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
-        )
-        c_tensor, _ = cutlass_torch.cute_tensor_like(
-            c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16
-        )
-        return testing.JitArguments(a_tensor, b_tensor, c_tensor, current_stream)
+        return testing.JitArguments(a, b, c, torch_stream)
 
     workspace_count = 1
     if use_cold_l2:
         one_workspace_bytes = (
-            a_torch_cpu.numel() * a_torch_cpu.element_size()
-            + b_torch_cpu.numel() * b_torch_cpu.element_size()
-            + c_torch_cpu.numel() * c_torch_cpu.element_size()
+            a.numel() * a.element_size()
+            + b.numel() * b.element_size()
+            + c.numel() * c.element_size()
         )
         workspace_count = testing.get_workspace_count(
             one_workspace_bytes, warmup_iterations, iterations
         )
 
-    exec_time = testing.benchmark(
-        compiled_gemm,
+    # Return execution time in microseconds
+    return testing.benchmark(
+        compiled_fn,
         workspace_generator=generate_tensors,
         workspace_count=workspace_count,
         stream=current_stream,
@@ -1853,18 +1977,17 @@ def run(
         iterations=iterations,
     )
 
-    return exec_time  # Return execution time in microseconds
+
+def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
+    try:
+        return tuple(int(x.strip()) for x in s.split(","))
+    except ValueError:
+        raise argparse.ArgumentTypeError(
+            "Invalid format. Expected comma-separated integers."
+        )
 
 
-if __name__ == "__main__":
-
-    def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
-        try:
-            return tuple(int(x.strip()) for x in s.split(","))
-        except ValueError:
-            raise argparse.ArgumentTypeError(
-                "Invalid format. Expected comma-separated integers."
-            )
+def prepare_parser():
 
     parser = argparse.ArgumentParser(
         description="Example of Dense Persistent GEMM on Blackwell."
@@ -1872,19 +1995,13 @@ if __name__ == "__main__":
 
     parser.add_argument(
         "--mnkl",
-        type=parse_comma_separated_ints,
+        type=_parse_comma_separated_ints,
         default=(256, 256, 512, 1),
         help="mnkl dimensions (comma-separated)",
     )
-    parser.add_argument(
-        "--mma_tiler_mn",
-        type=parse_comma_separated_ints,
-        default=(128, 128),
-        help="Mma tile shape (comma-separated)",
-    )
     parser.add_argument(
         "--cluster_shape_mn",
-        type=parse_comma_separated_ints,
+        type=_parse_comma_separated_ints,
         default=(1, 1),
         help="Cluster shape (comma-separated)",
     )
@@ -1905,6 +2022,9 @@ if __name__ == "__main__":
     parser.add_argument(
         "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
     )
+    parser.add_argument(
+        "--benchmark", action="store_true", help="Only benchmark the kernel"
+    )
     parser.add_argument(
         "--warmup_iterations", type=int, default=0, help="Warmup iterations"
     )
@@ -1923,6 +2043,24 @@ if __name__ == "__main__":
         default=False,
         help="Use circular buffer tensor sets to ensure L2 cold cache",
     )
+    parser.add_argument(
+        "--use_tvm_ffi",
+        action="store_true",
+        default=False,
+        help="Enable TVM FFI for the kernel, defaults to False using CuTe DSL's native runtime",
+    )
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = prepare_parser()
+    parser.add_argument(
+        "--mma_tiler_mn",
+        type=_parse_comma_separated_ints,
+        default=(128, 128),
+        help="Mma tile shape (comma-separated)",
+    )
 
     args = parser.parse_args()
 
@@ -1952,5 +2090,7 @@ if __name__ == "__main__":
         args.iterations,
         args.skip_ref_check,
         args.use_cold_l2,
+        args.use_tvm_ffi,
+        args.benchmark,
     )
     print("PASS")

examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py

@@ -36,6 +36,7 @@ import cutlass
 import cutlass.cute as cute
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils.blackwell_helpers as sm100_utils
@@ -535,6 +536,7 @@ class DenseGemmKernel:
             consumer_group=ab_pipeline_consumer_group,
             tx_count=self.num_tma_load_bytes,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         ).make_participants()
 
         # Initialize acc_pipeline (barrier) and states
@@ -549,6 +551,7 @@ class DenseGemmKernel:
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
         acc_producer_state = pipeline.make_pipeline_state(
             pipeline.PipelineUserType.Producer, self.num_acc_stage
@@ -569,8 +572,7 @@ class DenseGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C
@@ -686,8 +688,7 @@ class DenseGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         # Alloc tensor memory buffer
         tmem.allocate(self.num_tmem_alloc_cols)

examples/python/CuTeDSL/blackwell/fmha.py

@@ -48,9 +48,11 @@ import cutlass.cute.testing as testing
 from cutlass.cute.runtime import from_dlpack
 from cutlass.cute.typing import Int32, Int64, Float32
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(current_dir, ".."))
-from utils import fmha_helpers as fmha_utils
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from helpers import fmha_helpers as fmha_utils
 
 """
 A fused multi-head attention (FMHA) example for the NVIDIA Blackwell SM100 architecture using CUTE DSL

examples/python/CuTeDSL/blackwell/fmha_bwd.py

@@ -49,9 +49,11 @@ import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.cute.runtime import from_dlpack
 from cutlass.cute.typing import Int32, Float32, Float8E4M3FN, Float16, BFloat16, Boolean
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(current_dir, ".."))
-from utils import fmha_helpers as fmha_utils
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from helpers import fmha_helpers as fmha_utils
 
 """
 A fused multi-head attention (FMHA) backward pass example for the NVIDIA Blackwell SM100 architecture using CUTE DSL

examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py

@@ -40,6 +40,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.utils.blockscaled_layout as blockscaled_utils
 from cutlass.cute.runtime import from_dlpack
@@ -177,22 +178,18 @@ class Sm100GroupedBlockScaledGemmKernel:
         self.threads_per_cta = 32 * len(
             (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
         )
-        # Set barrier for cta sync, epilogue sync and tmem ptr sync
-        self.cta_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=1,
-            num_threads=self.threads_per_cta,
-        )
+        # Set barrier for epilogue sync and tmem ptr sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=2,
+            barrier_id=1,
             num_threads=32 * len(self.epilog_warp_id),
         )
         self.tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=3,
+            barrier_id=2,
             num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
         )
         # Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion
         self.tensormap_ab_init_barrier = pipeline.NamedBarrier(
-            barrier_id=4,
+            barrier_id=3,
             num_threads=64,
         )
         self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -646,6 +643,7 @@ class Sm100GroupedBlockScaledGemmKernel:
             cluster=(*self.cluster_shape_mn, 1),
             smem=self.shared_storage.size_in_bytes(),
             stream=stream,
+            min_blocks_per_mp=1,
         )
         return
 
@@ -781,11 +779,9 @@ class Sm100GroupedBlockScaledGemmKernel:
                     cute.arch.mbarrier_init(
                         tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
                     )
-        cute.arch.mbarrier_init_fence()
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/SFA/SFB/C
@@ -944,10 +940,7 @@ class Sm100GroupedBlockScaledGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Get tensormap buffer address
@@ -2894,6 +2887,7 @@ def run(
         tensor_of_tensormap,
         max_active_clusters,
         current_stream,
+        options=f"--opt-level 2",
     )
 
     # reference check

examples/python/CuTeDSL/blackwell/grouped_gemm.py

@@ -39,6 +39,7 @@ import cutlass.cute as cute
 import cutlass.cute.testing as testing
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.utils.blackwell_helpers as sm100_utils
 import cutlass.torch as cutlass_torch
@@ -153,22 +154,18 @@ class GroupedGemmKernel:
         self.threads_per_cta = 32 * len(
             (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
         )
-        # Set barrier for cta sync, epilog sync, tmem ptr sync and tensormap update sync
-        self.cta_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=1,
-            num_threads=self.threads_per_cta,
-        )
+        # Set barrier for epilog sync, tmem ptr sync and tensormap update sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
-            barrier_id=2,
+            barrier_id=1,
             num_threads=32 * len(self.epilog_warp_id),
         )
         self.tmem_alloc_barrier = pipeline.NamedBarrier(
-            barrier_id=3,
+            barrier_id=2,
             num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
         )
         # Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion
         self.tensormap_ab_init_barrier = pipeline.NamedBarrier(
-            barrier_id=4,
+            barrier_id=3,
             num_threads=32 * (len(self.epilog_warp_id) + 1),
         )
         self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -586,11 +583,9 @@ class GroupedGemmKernel:
             is_two_cta=use_2cta_instrs,
             two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
         )
-        cute.arch.mbarrier_init_fence()
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         #
         # Setup smem tensor A/B/C
@@ -718,10 +713,7 @@ class GroupedGemmKernel:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         #
         # Get tensormap buffer address

examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py

@@ -27,6 +27,8 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
+import os
+import sys
 import argparse
 from typing import List, Type, Tuple, Optional
 import cuda.bindings.driver as cuda
@@ -39,21 +41,22 @@ import cutlass.cute as cute
 import cutlass.cute.testing as testing
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 from cutlass.cute.nvgpu import cpasync, tcgen05
 import cutlass.torch as cutlass_torch
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.cute.runtime import from_dlpack
 
-import sys
-from pathlib import Path
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, "../.."))
 
-sys.path.append(str(Path(__file__).resolve().parent))
-from mamba2_ssd_reference import (
+from blackwell.mamba2_ssd.mamba2_ssd_reference import (
     ssd_reference_fp32_all,
     ssd_reference_lowprecision_intermediates,
     analyze_relative_diffs,
 )
-from mamba2_ssd_tile_scheduler import (
+from blackwell.mamba2_ssd.mamba2_ssd_tile_scheduler import (
     Mamba2SSDTileSchedulerParams,
     Mamba2SSDTileScheduler,
 )
@@ -811,12 +814,10 @@ class SSDKernel:
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mnk) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
 
         # Cluster wait before tmem alloc
-        if cute.size(self.cluster_shape_mnk) > 1:
-            cute.arch.cluster_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
 
         tmem_alloc_barrier = pipeline.NamedBarrier(
             barrier_id=0,
@@ -2574,6 +2575,7 @@ class SSDKernel:
                 consumer_group=x_consumer_group,
                 tx_count=self.num_x_load_bytes,
                 barrier_storage=x_full_mbar_ptr,
+                defer_sync=True,
             )
         else:
             x_consumer_group_umma = pipeline.CooperativeGroup(
@@ -2590,6 +2592,7 @@ class SSDKernel:
                 consumer_group_async=x_consumer_group_async,
                 tx_count=self.num_x_load_bytes,
                 barrier_storage=x_full_mbar_ptr,
+                defer_sync=True,
             )
 
     def make_and_init_b_pipeline(self, b_full_mbar_ptr):
@@ -2609,6 +2612,7 @@ class SSDKernel:
             consumer_group_async=b_consumer_group_async,
             tx_count=self.num_b_load_bytes,
             barrier_storage=b_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_c_pipeline(self, c_full_mbar_ptr):
@@ -2624,6 +2628,7 @@ class SSDKernel:
             consumer_group=c_consumer_group,
             tx_count=self.num_c_load_bytes,
             barrier_storage=c_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_deltas_pipeline(self, deltas_full_mbar_ptr):
@@ -2643,6 +2648,7 @@ class SSDKernel:
             consumer_group=deltas_consumer_group,
             tx_count=self.num_delta_load_bytes + self.num_cumsum_delta_load_bytes,
             barrier_storage=deltas_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_d_pipeline(self, d_full_mbar_ptr):
@@ -2662,6 +2668,7 @@ class SSDKernel:
                 consumer_group=d_consumer_group,
                 tx_count=self.num_d_load_bytes,
                 barrier_storage=d_full_mbar_ptr,
+                defer_sync=True,
             )
 
     def make_and_init_intra1_acc_pipeline(self, intra1_acc_full_mbar_ptr):
@@ -2676,6 +2683,7 @@ class SSDKernel:
             producer_group=intra1_acc_producer_group,
             consumer_group=intra1_acc_consumer_group,
             barrier_storage=intra1_acc_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_intra2_q_pipeline(self, intra2_q_full_mbar_ptr):
@@ -2690,6 +2698,7 @@ class SSDKernel:
             producer_group=intra2_q_producer_group,
             consumer_group=intra2_q_consumer_group,
             barrier_storage=intra2_q_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_intra2_acc_pipeline(self, intra2_acc_full_mbar_ptr):
@@ -2704,6 +2713,7 @@ class SSDKernel:
             producer_group=intra2_acc_producer_group,
             consumer_group=intra2_acc_consumer_group,
             barrier_storage=intra2_acc_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_inter1_b_pipeline(self, inter1_b_full_mbar_ptr):
@@ -2718,6 +2728,7 @@ class SSDKernel:
             producer_group=inter1_b_producer_group,
             consumer_group=inter1_b_consumer_group,
             barrier_storage=inter1_b_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_inter1_acc_pipeline(self, inter1_acc_full_mbar_ptr):
@@ -2732,6 +2743,7 @@ class SSDKernel:
             producer_group=inter1_acc_producer_group,
             consumer_group=inter1_acc_consumer_group,
             barrier_storage=inter1_acc_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_inter2_p_pipeline(self, inter2_p_full_mbar_ptr):
@@ -2746,6 +2758,7 @@ class SSDKernel:
             producer_group=inter2_p_producer_group,
             consumer_group=inter2_p_consumer_group,
             barrier_storage=inter2_p_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def make_and_init_inter2_acc_pipeline(self, inter2_acc_full_mbar_ptr):
@@ -2760,6 +2773,7 @@ class SSDKernel:
             producer_group=inter2_acc_producer_group,
             consumer_group=inter2_acc_consumer_group,
             barrier_storage=inter2_acc_full_mbar_ptr,
+            defer_sync=True,
         )
 
     def tma_partition_for_mma_b_operand(

examples/python/CuTeDSL/blackwell/mixed_input_gemm.py

@@ -27,7 +27,6 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import argparse
-from enum import Enum, auto
 from math import log2, ceil
 from typing import Optional, Union
 
@@ -37,9 +36,12 @@ import cuda.bindings.driver as cuda
 import cutlass
 import cutlass.cute as cute
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.mixed_input_helpers as mixed_input_utils
+from cutlass.utils.mixed_input_helpers import TransformMode
 import cutlass.cute.testing as testing
 from cutlass.cute.nvgpu import cpasync, tcgen05
 from cutlass.cute.runtime import from_dlpack
@@ -130,15 +132,6 @@ Besides the requirements from the Blackwell dense GEMM example, there are some c
 """
 
 
-class TransformMode(Enum):
-    """
-    An enumeration for the possible transform modes of a mixed-input GEMM.
-    """
-
-    ConvertOnly = auto()
-    ConvertScale = auto()
-
-
 class MixedInputGemmKernel:
     """
     Mixed-input GEMM kernel for NVIDIA Blackwell SM100 architecture.
@@ -226,7 +219,7 @@ class MixedInputGemmKernel:
             + 1
         )
 
-        # Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync
+        # Set barrier id for epilogue sync, tmem ptr sync, and transform sync
         self.epilog_sync_barrier = pipeline.NamedBarrier(
             1, 32 * len(self.epilog_warp_id)
         )
@@ -234,7 +227,6 @@ class MixedInputGemmKernel:
         self.transform_sync_barrier = pipeline.NamedBarrier(
             3, 32 * len(self.transform_warp_id)
         )
-        self.cta_sync_barrier = pipeline.NamedBarrier(4, self.threads_per_cta)
 
         self.smem_buffer_align_bytes = 1024
 
@@ -255,7 +247,9 @@ class MixedInputGemmKernel:
         - Computing tensor memory allocation columns
         """
         # Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM)
-        self.transform_a_source = self._get_transform_a_source(self.a_major_mode)
+        self.transform_a_source = mixed_input_utils.get_transform_a_source(
+            self.a_major_mode
+        )
         tiled_mma = sm100_utils.make_trivial_tiled_mma(
             self.mma_dtype,
             self.a_major_mode,
@@ -346,7 +340,7 @@ class MixedInputGemmKernel:
             self.smem_layout_a,
             self.smem_layout_a_transform,
             self.smem_layout_b,
-        ) = self._compute_smem_layout(
+        ) = mixed_input_utils.compute_smem_layout(
             tiled_mma,
             self.mma_tiler,
             self.a_dtype,
@@ -358,11 +352,20 @@ class MixedInputGemmKernel:
         self.smem_layout_scale_per_stage = None
         self.smem_layout_scale = None
         if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale):
-            # Get smem layout for scale tensor
+            # Get scale tile shape and smem layout for scale tensor
             (
+                self.scale_tile_shape,
                 self.smem_layout_scale_per_stage,
                 self.smem_layout_scale,
-            ) = self.get_smem_layout_scale()
+            ) = mixed_input_utils.get_smem_layout_scale(
+                self.mma_tiler,
+                self.use_2cta_instrs,
+                self.scale_granularity_m,
+                self.scale_granularity_k,
+                self.scale_major_mode,
+                self.a_scale_dtype,
+                self.num_scale_load2trans_stage,
+            )
 
     def _validate_inputs(
         self,
@@ -448,7 +451,12 @@ class MixedInputGemmKernel:
         self.c_layout = utils.LayoutEnum.from_tensor(c)
         if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale):
             # Get gmem layout for scale tensor
-            self.gmem_layout_scale = self.get_gmem_layout_scale(a.shape)
+            self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale(
+                a.shape,
+                self.scale_granularity_m,
+                self.scale_granularity_k,
+                self.scale_major_mode,
+            )
 
         # Validate inputs
         self._validate_inputs(a, a_scale, b, c)
@@ -466,8 +474,12 @@ class MixedInputGemmKernel:
             self.transform_a_source,
         )
         # Set up gmem copy atoms for A, scale, and B
-        a_op = self._get_tma_atom_kind(self.is_a_mcast, self.use_2cta_instrs, False)
-        b_op = self._get_tma_atom_kind(self.is_b_mcast, self.use_2cta_instrs, True)
+        a_op = mixed_input_utils.get_tma_atom_kind(
+            self.is_a_mcast, self.use_2cta_instrs, False
+        )
+        b_op = mixed_input_utils.get_tma_atom_kind(
+            self.is_b_mcast, self.use_2cta_instrs, True
+        )
         a_scale_op = a_op
         # Deduce TMA copy atom and TMA tensor for A, scale, and B
         smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0))
@@ -650,7 +662,6 @@ class MixedInputGemmKernel:
             grid=grid,
             block=[self.threads_per_cta, 1, 1],
             cluster=(*self.cluster_shape_mn, 1),
-            smem=self.shared_storage.size_in_bytes(),
             stream=stream,
             min_blocks_per_mp=1,
         )
@@ -730,6 +741,7 @@ class MixedInputGemmKernel:
             cta_layout_vmnk=cluster_layout_vmnk,
             tidx=transform_thread_idx,
             mcast_mode_mn=(1, 0),  # multicast for A will only happen on the M-mode
+            defer_sync=True,
         )
         # Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading
         # of scale, and the transformation of A
@@ -753,6 +765,7 @@ class MixedInputGemmKernel:
                     1,
                     0,
                 ),  # multicast for scale_a will only happen on the M-mode
+                defer_sync=True,
             )
         # Initialize transform2mma pipeline, which tracks the dependencies between the transformation
         # of A and MMA's consumption of transformed A
@@ -766,6 +779,7 @@ class MixedInputGemmKernel:
             ),
             consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
         # Initialize pipeline for tensor B load to MMA
         # MMA warp informs TMA warp to proceed to load next tile of B tensor
@@ -779,6 +793,7 @@ class MixedInputGemmKernel:
             tx_count=self.num_tma_load_bytes_b,
             cta_layout_vmnk=cluster_layout_vmnk,
             mcast_mode_mn=(0, 1),  # multicast for B will only happen on the N-mode
+            defer_sync=True,
         )
         # Initialize accumulator pipeline, which tracks the dependencies between
         # MMA's computation of accumulators and epilogue warps' consumption of accumulators
@@ -790,6 +805,7 @@ class MixedInputGemmKernel:
                 pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id)
             ),
             cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
         )
 
         # Tensor memory dealloc barrier init
@@ -802,8 +818,7 @@ class MixedInputGemmKernel:
         )
 
         # Cluster arrive after barrier init
-        if cutlass.const_expr(cute.size(self.cluster_shape_mn) > 1):
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         # Setup smem tensor A/scale/B/C
         sC = (
@@ -897,7 +912,7 @@ class MixedInputGemmKernel:
             cute.dice(self.mma_tiler, (1, None, 1))
         )
         # Setup copy atom to store transformed A into tensor memory or shared memory
-        copy_atom_a_transform = self._get_copy_atom_a_transform(
+        copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform(
             self.mma_dtype,
             self.use_2cta_instrs,
             self.transform_a_source,
@@ -928,7 +943,7 @@ class MixedInputGemmKernel:
             tCsS = thr_mma.partition_A(sS_input)
             # ((atom_v, rest_v), STAGE)
             # ((atom_v, rest_v), loopM, loopK, loopL)
-            tSsS, tSgS = self.scale_tma_partition(
+            tSsS, tSgS = mixed_input_utils.scale_tma_partition(
                 tCsS,
                 tCgS,
                 tma_atom_s,
@@ -959,10 +974,7 @@ class MixedInputGemmKernel:
         )
 
         # Cluster wait before TMEM alloc and ensure pipelines are ready
-        if cutlass.const_expr(cute.size(self.cluster_shape_mn) > 1):
-            cute.arch.cluster_wait()
-        else:
-            self.cta_sync_barrier.arrive_and_wait()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         # TMEM allocation
         tmem.allocate(self.num_tmem_alloc_cols)
@@ -1145,7 +1157,7 @@ class MixedInputGemmKernel:
                 dst_copy_a,
                 tAsA_input,
                 tAsA_transform,
-            ) = self.transform_partition(
+            ) = mixed_input_utils.transform_partition(
                 self.transform_a_source,
                 self.scale_mode,
                 copy_atom_a_input,
@@ -1173,8 +1185,10 @@ class MixedInputGemmKernel:
             tSrS_copy = None
             tSrS = None
             if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale):
-                smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = self.scale_partition(
-                    src_copy_a, tCsS, transform_local_tidx, self.mma_dtype
+                smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = (
+                    mixed_input_utils.scale_partition(
+                        src_copy_a, tCsS, transform_local_tidx, self.mma_dtype
+                    )
                 )
                 assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), (
                     "tSrS and tArA have different leading dimension"
@@ -1582,328 +1596,6 @@ class MixedInputGemmKernel:
             if cutlass.const_expr(self.use_tma_store):
                 c_pipeline.producer_tail()
 
-    def scale_tma_partition(
-        self,
-        tCsS: cute.Tensor,
-        tCgS: cute.Tensor,
-        tma_atom_s: cute.CopyAtom,
-        block_in_cluster_coord_vmnk: cute.Coord,
-        scale_cta_layout: cute.Layout,
-    ) -> tuple[cute.Tensor, cute.Tensor]:
-        """
-        Perform TMA partition for scale tensor.
-
-        This method partitions the gobal memory and shared memory buffer for scale tensor for TMA load.
-
-        :param tCsS: Input scale shared memory tensor
-        :type tCsS: cute.Tensor
-        :param tCgS: Input scale global memory tensor
-        :type tCgS: cute.Tensor
-        :param tma_atom_s: TMA copy atom for scale tensor
-        :type tma_atom_s: cute.CopyAtom
-        :param block_in_cluster_coord_vmnk: CTA coord in the cluster
-        :type block_in_cluster_coord_vmnk: cute.Coord
-        :param scale_cta_layout: Layout of CTA from the view of the scale tensor
-        :type scale_cta_layout: cute.Layout
-
-        :return: A tuple containing (tSsS, tSgS) where:
-            - tSsS: Partitioned scale tensor in shared memory
-            - tSgS: Partitioned scale tensor in global memory
-        :rtype: tuple[cute.Tensor, cute.Tensor]
-        """
-        tSsS, tSgS = cpasync.tma_partition(
-            tma_atom_s,
-            block_in_cluster_coord_vmnk[2],
-            scale_cta_layout,
-            cute.group_modes(tCsS, 0, 3),
-            cute.group_modes(tCgS, 0, 3),
-        )
-        # Add rest_v mode
-        # ((atom_v, rest_v), STAGE)
-        # ((atom_v, rest_v), loopM, loopK, loopL)
-        tSsS = cute.make_tensor(
-            tSsS.iterator,
-            cute.make_layout(
-                ((tSsS.layout.shape[0], 1), *tSsS.layout.shape[1:]),
-                stride=(
-                    (tSsS.layout.stride[0], 0),
-                    *tSsS.layout.stride[1:],
-                ),
-            ),
-        )
-        tSgS = cute.make_tensor(
-            tSgS.iterator,
-            cute.make_layout(
-                ((tSgS.layout.shape[0], 1), *tSgS.layout.shape[1:]),
-                stride=(
-                    (tSgS.layout.stride[0], 0),
-                    *tSgS.layout.stride[1:],
-                ),
-            ),
-        )
-        return tSsS, tSgS
-
-    def transform_partition(
-        self,
-        transform_a_source: tcgen05.OperandSource,
-        scale_mode: TransformMode,
-        copy_atom_a_input: cute.CopyAtom,
-        copy_atom_a_transform: cute.CopyAtom,
-        sA_input: cute.Tensor,
-        A_transform: cute.Tensor,
-        transform_local_tidx: cutlass.Int32,
-    ) -> tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]:
-        """
-        Partition tensors for transform input and output.
-
-        This method sets up the copy atoms and partitions the shared/tensor memory
-        for the transformation of tensor A.
-
-        :param transform_a_source: Where the transformed tensor A is stored (TMEM or SMEM)
-        :type transform_a_source: tcgen05.OperandSource
-        :param scale_mode: The transform mode (ConvertOnly or ConvertScale)
-        :type scale_mode: TransformMode
-        :param copy_atom_a_input: Copy atom for loading A from shared memory
-        :type copy_atom_a_input: cute.CopyAtom
-        :param copy_atom_a_transform: Copy atom for storing transformed A
-        :type copy_atom_a_transform: cute.CopyAtom
-        :param sA_input: Input tensor A in shared memory
-        :type sA_input: cute.Tensor
-        :param A_transform: Transformed tensor A in tensor or shared memory
-        :type A_transform: cute.Tensor
-        :param transform_local_tidx: Local thread index for transformation warps
-        :type transform_local_tidx: cutlass.Int32
-
-        :return: A tuple containing (src_copy_a, dst_copy_a, tAsA_input, tA_transform) where:
-            - src_copy_a: Tiled copy for source tensor
-            - dst_copy_a: Tiled copy for destination tensor
-            - tAsA_input: Partitioned input tensor A
-            - tA_transform: Partitioned transformed tensor A
-        :rtype: tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]
-        """
-        if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM):
-            if cutlass.const_expr(
-                cute.size(A_transform, mode=[0, 0]) == 128
-                and cute.size(sA_input, mode=[0, 0]) == 64
-            ):
-                tensor_input = cute.make_tensor(
-                    sA_input.iterator,
-                    cute.logical_product(
-                        sA_input.layout,
-                        ((cute.make_layout(2, stride=0), None), None, None, None),
-                    ),
-                )
-            else:
-                tensor_input = sA_input
-            reg2tmem_tiled_copy = tcgen05.make_tmem_copy(
-                copy_atom_a_transform, A_transform[(None, None, None, 0)]
-            )
-            thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(
-                transform_local_tidx
-            )
-            partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input)
-            partitioned_tensor_transform = thr_reg2tmem_tiled_copy.partition_D(
-                A_transform
-            )
-            src_copy_a = (
-                cute.make_tiled_copy_S(copy_atom_a_input, reg2tmem_tiled_copy)
-                if scale_mode is TransformMode.ConvertScale
-                else None
-            )
-            dst_copy_a = reg2tmem_tiled_copy
-            tAsA_input = partitioned_tensor_input
-            tA_transform = partitioned_tensor_transform
-        elif cutlass.const_expr(transform_a_source == tcgen05.OperandSource.SMEM):
-            # Construct tiled_copy satisfying 8 contiguous elts per copy atom
-            reg2smem_tiled_copy = cute.make_cotiled_copy(
-                copy_atom_a_transform,
-                cute.make_layout((128, 8), stride=(8, 1)),
-                A_transform[(None, None, None, 0)].layout,
-            )
-            thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(
-                transform_local_tidx
-            )
-            partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(sA_input)
-            partitioned_tensor_transform = thr_reg2smem_tiled_copy.partition_D(
-                A_transform
-            )
-            src_copy_a = (
-                cute.make_tiled_copy_S(copy_atom_a_input, reg2smem_tiled_copy)
-                if scale_mode is TransformMode.ConvertScale
-                else None
-            )
-            # auto-vec copy is enough for copy from register to shared memory here
-            dst_copy_a = None
-            tAsA_input = partitioned_tensor_input
-            tA_transform = partitioned_tensor_transform
-        return src_copy_a, dst_copy_a, tAsA_input, tA_transform
-
-    def scale_partition(
-        self,
-        src_copy_a: cute.TiledCopy,
-        tCsS: cute.Tensor,
-        transform_local_tidx: cutlass.Int32,
-        mma_dtype: type[cutlass.Numeric],
-    ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]:
-        """
-        Partition the scale tensor for transformation.
-
-        This method prepares the copy atom and partitions the shared memory for the scale tensor.
-
-        :param src_copy_a: Tiled copy for the source tensor
-        :type src_copy_a: cute.TiledCopy
-        :param tCsS: Scale tensor in shared memory
-        :type tCsS: cute.Tensor
-        :param transform_local_tidx: Local thread index for transformation warps
-        :type transform_local_tidx: cutlass.Int32
-        :param mma_dtype: Data type for the MMA operation
-        :type mma_dtype: type[cutlass.Numeric]
-
-        :return: A tuple containing (smem_thr_copy_S, tSsS_trans, tSrS) where:
-            - smem_thr_copy_S: Tiled copy for the scale tensor
-            - tSsS_trans: Partitioned scale tensor for transformation
-            - tSrS_copy: Register fragment for the scale tensor
-            - tSrS: view of scale tensor used for transformation computation
-        :rtype: tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]
-        """
-        smem_thr_copy_S = None
-        tSsS_trans = None
-        tSrS = None
-        # Partition scale tensor
-        smem_thr_copy_S = src_copy_a.get_slice(transform_local_tidx)
-        tSsS_trans = smem_thr_copy_S.partition_S(tCsS)
-        # Construct register fragment for scale tensor
-        tSsS_layout_per_stage = tSsS_trans[(None, None, None, None, 0)].layout
-        # tSrS for copy
-        tSrS_copy = cute.make_rmem_tensor(
-            cute.filter_zeros(tSsS_layout_per_stage).shape, mma_dtype
-        )
-        # tSrS view for transformation computation
-        tSrS = cute.make_tensor(
-            tSrS_copy.iterator,
-            cute.make_layout(
-                tSsS_layout_per_stage.shape, stride=tSrS_copy.layout.stride
-            ),
-        )
-        return smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS
-
-    def get_gmem_layout_scale(
-        self, scale_shape_mkl: tuple[int, int, int]
-    ) -> cute.Layout:
-        """
-        Get the layout of the scale tensor in global memory.
-
-        :param scale_shape_mkl: The shape of the scale tensor (M, K, L).
-        :type scale_shape_mkl: tuple[int, int, int]
-
-        :return: The layout of the scale tensor in global memory.
-        :rtype: cute.Layout
-        """
-        m, k, l = scale_shape_mkl
-        shape_scale = (
-            (self.scale_granularity_m, cute.ceil_div(m, self.scale_granularity_m)),
-            (self.scale_granularity_k, cute.ceil_div(k, self.scale_granularity_k)),
-        )
-        if cutlass.const_expr(self.scale_major_mode == tcgen05.OperandMajorMode.MN):
-            layout_mk = cute.make_layout(
-                shape_scale,
-                stride=(
-                    (0, 1),
-                    (0, cute.size(shape_scale[0][1])),
-                ),
-            )
-        else:
-            layout_mk = cute.make_layout(
-                shape_scale,
-                stride=(
-                    (0, cute.size(shape_scale[1][1])),
-                    (0, 1),
-                ),
-            )
-        return cute.make_layout(
-            (*layout_mk.shape, l),
-            stride=(*layout_mk.stride, cute.cosize(layout_mk)),
-        )
-
-    def get_smem_layout_scale(self) -> tuple[cute.ComposedLayout, cute.ComposedLayout]:
-        """
-        Get the layout of the scale tensor in shared memory.
-
-        :return: A tuple containing:
-            - smem_layout_scale_per_stage: Shared memory layout for scale tensor per stage
-            - smem_layout_scale: Shared memory layout for scale tensor
-        :rtype: tuple[cute.ComposedLayout, cute.ComposedLayout]
-        """
-        self.scale_tile_shape = (
-            (
-                cute.size(self.mma_tiler[0]) // 2
-                if self.use_2cta_instrs
-                else cute.size(self.mma_tiler[0])
-            ),
-            cute.size(self.mma_tiler[2]),
-        )
-        size_mn = self.scale_tile_shape[0]
-        size_k = self.scale_tile_shape[1]
-        smem_size_mn = (
-            self.scale_granularity_m if self.scale_granularity_m < size_mn else size_mn
-        )
-        smem_size_k = (
-            self.scale_granularity_k if self.scale_granularity_k < size_k else size_k
-        )
-        div_mn = cute.ceil_div(size_mn, smem_size_mn)
-        div_k = cute.ceil_div(size_k, smem_size_k)
-        smem_atom_shape = (
-            (smem_size_mn, div_mn),
-            (smem_size_k, div_k),
-        )
-        if cutlass.const_expr(self.scale_major_mode == tcgen05.OperandMajorMode.MN):
-            outer_layout = cute.make_layout(
-                smem_atom_shape,
-                stride=(
-                    (0, 1),
-                    (0, div_mn),
-                ),
-            )
-        else:
-            outer_layout = cute.make_layout(
-                smem_atom_shape,
-                stride=(
-                    (0, div_k),
-                    (0, 1),
-                ),
-            )
-        # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom
-        smem_layout_scale_per_stage = cute.make_composed_layout(
-            cute.make_swizzle(0, 4, 3), 0, outer_layout
-        )
-        assert cute.rank(smem_layout_scale_per_stage) == 2, (
-            "Scale layout must be rank 2"
-        )
-
-        assert (
-            cute.size(self.mma_tiler[0])
-            % cute.size(smem_layout_scale_per_stage.outer[0])
-            == 0
-        ), "smem_layout_scale_per_stage must equal the tile shape."
-        assert (
-            cute.size(self.mma_tiler[2])
-            % cute.size(smem_layout_scale_per_stage.outer[1])
-            == 0
-        ), "smem_layout_scale_per_stage must evenly divide tile k shape."
-        # Shared memory buffer for scale must be at least 128B to satisfy TMA requirement
-        assert (
-            cute.size_in_bytes(self.a_scale_dtype, smem_layout_scale_per_stage) >= 128
-        ), "smem size for scale must be at least 128B"
-        # Scale layout in smem with multiple stages
-        smem_layout_scale = cute.append(
-            smem_layout_scale_per_stage,
-            cute.make_layout(
-                (self.num_scale_load2trans_stage),
-                stride=(cute.cosize(smem_layout_scale_per_stage.outer)),
-            ),
-        )
-        return smem_layout_scale_per_stage, smem_layout_scale
-
     def epilog_gmem_copy_and_partition(
         self,
         tidx: cutlass.Int32,
@@ -2285,126 +1977,6 @@ class MixedInputGemmKernel:
             num_tmem_a_cols,
         )
 
-    @staticmethod
-    def _compute_smem_layout(
-        tiled_mma: cute.TiledMma,
-        mma_tiler_mnk: tuple[int, int, int],
-        a_dtype: type[cutlass.Numeric],
-        b_dtype: type[cutlass.Numeric],
-        load2trans_stage_count: int,
-        trans2mma_stage_count: int,
-    ) -> tuple[
-        cute.ComposedLayout,
-        cute.ComposedLayout,
-        cute.ComposedLayout,
-    ]:
-        """
-        Compute shared memory layouts for tensor A, transformed A and tensor B.
-
-        :param tiled_mma: The tiled MMA object defining the core computation.
-        :type tiled_mma: cute.TiledMma
-        :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
-        :type mma_tiler_mnk: tuple[int, int, int]
-        :param a_dtype: Data type of operand A.
-        :type a_dtype: type[cutlass.Numeric]
-        :param b_dtype: Data type of operand B.
-        :type b_dtype: type[cutlass.Numeric]
-        :param load2trans_stage_count: Number of stages for load-to-transform pipeline.
-        :type load2trans_stage_count: int
-        :param trans2mma_stage_count: Number of stages for transform-to-MMA pipeline.
-        :type trans2mma_stage_count: int
-
-        :return: A tuple containing:
-            - smem_layout_a: Shared memory layout for tensor A
-            - smem_layout_a_transform: Shared memory layout for transformed tensor A
-            - smem_layout_b: Shared memory layout for tensor B
-        :rtype: tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]
-        """
-        smem_layout_a = sm100_utils.make_smem_layout_a(
-            tiled_mma,
-            mma_tiler_mnk,
-            a_dtype,
-            load2trans_stage_count,
-        )
-        smem_layout_a_transform = sm100_utils.make_smem_layout_a(
-            tiled_mma,
-            mma_tiler_mnk,
-            tiled_mma.op.a_dtype,
-            trans2mma_stage_count,
-        )
-        smem_layout_b = sm100_utils.make_smem_layout_b(
-            tiled_mma,
-            mma_tiler_mnk,
-            b_dtype,
-            load2trans_stage_count,
-        )
-        return (
-            smem_layout_a,
-            smem_layout_a_transform,
-            smem_layout_b,
-        )
-
-    @staticmethod
-    def _get_transform_a_source(
-        a_major_mode: tcgen05.OperandMajorMode,
-    ) -> tcgen05.OperandSource:
-        """
-        Determine the operand source for transformed A tensor based on the operand major mode.
-        """
-        if cutlass.const_expr(a_major_mode == tcgen05.OperandMajorMode.K):
-            return tcgen05.OperandSource.TMEM
-        else:
-            return tcgen05.OperandSource.SMEM
-
-    @staticmethod
-    def _get_tma_atom_kind(
-        mcast: cutlass.Boolean,
-        use_2cta_instrs: bool,
-        is_b: bool,
-    ) -> Union[
-        cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp
-    ]:
-        """
-        Get the TMA atom kind based on 1) whether it's a multicast operation,
-        2) whether 2CTA tcgen05.mma instruction is enabled, and
-        3) whether it's a B tensor
-        """
-        # Not using .2CTA instructions for tensor A as the consumer is threads on different CTAs
-        cta_group = (
-            tcgen05.CtaGroup.TWO if (use_2cta_instrs and is_b) else tcgen05.CtaGroup.ONE
-        )
-        if cutlass.const_expr(mcast):
-            return cpasync.CopyBulkTensorTileG2SMulticastOp(cta_group)
-        return cpasync.CopyBulkTensorTileG2SOp(cta_group)
-
-    @staticmethod
-    def _get_copy_atom_a_transform(
-        mma_dtype: type[cutlass.Numeric],
-        use_2cta_instrs: bool,
-        transform_a_source: tcgen05.OperandSource,
-        a_smem_shape: cute.Shape,
-        a_dtype: type[cutlass.Numeric],
-    ) -> cute.CopyAtom:
-        """
-        Determine the copy atom for transformed A tensor based on the operand source and tile size.
-        """
-        if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM):
-            if cutlass.const_expr(
-                cute.size(a_smem_shape[0][0]) == 64 and (not use_2cta_instrs)
-            ):
-                copy_op_r2t = tcgen05.St16x256bOp(
-                    tcgen05.Repetition(1), tcgen05.Unpack.NONE
-                )
-            else:
-                copy_op_r2t = tcgen05.St32x32bOp(
-                    tcgen05.Repetition(8), tcgen05.Unpack.NONE
-                )
-            return cute.make_copy_atom(copy_op_r2t, mma_dtype)
-        else:
-            return cute.make_copy_atom(
-                cute.nvgpu.CopyUniversalOp(), a_dtype, num_bits_per_copy=32
-            )
-
     @staticmethod
     def _compute_grid(
         c: cute.Tensor,
@@ -2429,29 +2001,6 @@ class MixedInputGemmKernel:
 
         return tile_sched_params, grid
 
-    def is_valid_scale_granularity(
-        scale_granularity_m: int,
-        scale_granularity_k: int,
-        a_dtype: type[cutlass.Numeric],
-        k: int,
-        mma_tiler_k: int,
-    ) -> bool:
-        """
-        Check if the scale granularity settings are valid for the given data type and problem size.
-        """
-        if a_dtype.width == 8:
-            # No scale tensor for 8bit data type A
-            if not (scale_granularity_m == 0 and scale_granularity_k == 0):
-                return False
-        elif a_dtype.width == 4:
-            if scale_granularity_m != 1 or (
-                scale_granularity_k == 0
-                or k % scale_granularity_k != 0
-                or scale_granularity_k % mma_tiler_k != 0
-            ):
-                return False
-        return True
-
     def is_valid_tensor_alignment(
         m: int,
         n: int,
@@ -2566,7 +2115,7 @@ class MixedInputGemmKernel:
             mma_tiler, cluster_shape_mn, use_2cta_instrs
         ):
             return False
-        if not MixedInputGemmKernel.is_valid_scale_granularity(
+        if not mixed_input_utils.is_valid_scale_granularity(
             scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2]
         ):
             return False
@@ -2634,7 +2183,11 @@ def create_i4_tensor_and_scale(
         m, num_scales, scale_granularity_k, l
     )
     # Get elements with maximum absolute value to compute scaling factors
-    a_max = torch.maximum(ref / up_4b, ref / lb_4b)
+    a_max = (
+        torch.maximum(ref / up_4b, ref / lb_4b)
+        if dtype == cutlass.Int4
+        else torch.maximum(ref / up_4b)
+    )
     a_scales, _ = torch.max(a_max, dim=2, keepdim=True)
     a_scale_inv = torch.where(a_scales == 0, 0, 1 / a_scales)
     a_quant = ref * a_scale_inv
@@ -2668,17 +2221,6 @@ def create_i4_tensor_and_scale(
     )
 
 
-def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int:
-    """
-    Calculate the largest power of 2 divisibility factor for memory alignment.
-    """
-    # Check the largest power of 2 factor of contiguous_dim_size
-    for i in range(int(log2(contiguous_dim_size)), 0, -1):
-        if contiguous_dim_size % (2**i) == 0:
-            return min(2**i, upper_bound)
-    return 1
-
-
 def create_tensor_a(
     l: int,
     m: int,
@@ -2688,7 +2230,7 @@ def create_tensor_a(
     scale_granularity_m: int = 0,
     scale_granularity_k: int = 0,
     transformed_dtype: Optional[type[cutlass.Numeric]] = None,
-) -> tuple[cute.Tensor, cute.Tensor, torch.Tensor, torch.Tensor]:
+) -> tuple[cute.Tensor, Optional[cute.Tensor], torch.Tensor, Optional[torch.Tensor]]:
     """
     Create tensor A and scale tensor.
     """
@@ -2710,7 +2252,7 @@ def create_tensor_a(
             a_dtype,
             scale_granularity_m,
             scale_granularity_k,
-            divisibility=get_divisibility(m if a_major == "m" else k),
+            divisibility=mixed_input_utils.get_divisibility(m if a_major == "m" else k),
             transformed_dtype=transformed_dtype,
         )
     else:
@@ -2725,7 +2267,9 @@ def create_tensor_a(
             a_torch_cpu,
             a_dtype,
             is_dynamic_layout=True,
-            assumed_align=get_divisibility(m if a_major == "m" else k),
+            assumed_align=mixed_input_utils.get_divisibility(
+                m if a_major == "m" else k
+            ),
         )
     return a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu
 
@@ -2774,18 +2318,18 @@ def create_tensors(
         b_torch_cpu,
         b_dtype,
         is_dynamic_layout=True,
-        assumed_align=get_divisibility(n if b_major == "n" else k),
+        assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k),
     )
     c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
         c_torch_cpu,
         c_dtype,
         is_dynamic_layout=True,
-        assumed_align=get_divisibility(m if c_major == "m" else n),
+        assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
     )
     c_tensor = c_tensor.mark_compact_shape_dynamic(
         mode=(0 if c_major == "m" else 1),
         stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1),
-        divisibility=get_divisibility(m if c_major == "m" else n),
+        divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
     )
 
     return (
@@ -2966,25 +2510,36 @@ def run(
 
     def generate_tensors():
         a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a(
-            l, m, k, a_major, a_dtype, scale_granularity_m, scale_granularity_k, b_dtype
+            l,
+            m,
+            k,
+            a_major,
+            a_dtype,
+            scale_granularity_m,
+            scale_granularity_k,
+            b_dtype,
         )
         b_tensor, _ = cutlass_torch.cute_tensor_like(
             b_torch_cpu,
             b_dtype,
             is_dynamic_layout=True,
-            assumed_align=get_divisibility(n if b_major == "n" else k),
+            assumed_align=mixed_input_utils.get_divisibility(
+                n if b_major == "n" else k
+            ),
         )
         c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype)
         c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
             c_torch_cpu,
             c_dtype,
             is_dynamic_layout=True,
-            assumed_align=get_divisibility(m if c_major == "m" else n),
+            assumed_align=mixed_input_utils.get_divisibility(
+                m if c_major == "m" else n
+            ),
         )
         c_tensor = c_tensor.mark_compact_shape_dynamic(
             mode=(0 if c_major == "m" else 1),
             stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1),
-            divisibility=get_divisibility(m if c_major == "m" else n),
+            divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
         )
         return testing.JitArguments(
             a_tensor, a_scale_tensor, b_tensor, c_tensor, current_stream

examples/python/CuTeDSL/blackwell/mla.py

@@ -43,6 +43,7 @@ import cutlass.cute.nvgpu.tcgen05 as tcgen05
 import cutlass.cute.nvgpu.cpasync as cpasync
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass.cute.runtime import from_dlpack
@@ -970,7 +971,6 @@ class BlackwellMultiHeadLatentAttentionForward:
             num_tmem_dealloc_threads = self.threads_per_warp * self.num_compute_warps
             with cute.arch.elect_one():
                 cute.arch.mbarrier_init(tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads)
-        cute.arch.mbarrier_init_fence()
 
         load_q_pipeline = self.make_and_init_load_qkv_pipeline(
             storage.load_q_mbar_ptr.data_ptr(),
@@ -1004,8 +1004,7 @@ class BlackwellMultiHeadLatentAttentionForward:
             )
 
         # Cluster arrive after barrier init
-        if cutlass.const_expr(cute.size(self.cluster_shape_mnk) > 1):
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
 
         # Generate smem tensor Q/KC/VC/exchange
         # (MMA, MMA_H, MMA_R, PIPE)
@@ -1035,10 +1034,7 @@ class BlackwellMultiHeadLatentAttentionForward:
         #
         # Cluster wait before tensor memory alloc
         #
-        if cutlass.const_expr(cute.size(self.cluster_shape_mnk) > 1):
-            cute.arch.cluster_wait()
-        else:
-            pipeline.sync(barrier_id=4)
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
 
         # ///////////////////////////////////////////////////////////////////////////////
         #  Load warps, including page table and data tensors
@@ -2046,8 +2042,9 @@ class BlackwellMultiHeadLatentAttentionForward:
 
         # wait cpasync arrive until the last stage
         load_kv_pipeline = common_params.load_kv_pipeline
-        if copy_in_flight_count == self.load_kv_stage:
-            cute.arch.cp_async_wait_group(self.load_kv_stage - 1)
+        release_distance = 2
+        if copy_in_flight_count == self.load_kv_stage - release_distance:
+            cute.arch.cp_async_wait_group(self.load_kv_stage - release_distance - 1)
             load_kv_pipeline.producer_commit(load_kv_commit_state)
             load_kv_commit_state.advance()
             copy_in_flight_count -= 1
@@ -3922,6 +3919,7 @@ class BlackwellMultiHeadLatentAttentionForward:
                 producer_group=load_qkv_producer_group,
                 consumer_group=load_qkv_consumer_group,
                 cta_layout_vmnk=cta_layout_vmnk,
+                defer_sync=True,
             )
         else:
             load_qkv_producer_group = pipeline.CooperativeGroup(
@@ -3937,6 +3935,7 @@ class BlackwellMultiHeadLatentAttentionForward:
                 consumer_group=load_qkv_consumer_group,
                 tx_count=tx_count,
                 cta_layout_vmnk=cta_layout_vmnk,
+                defer_sync=True,
             )
 
     def make_and_init_mma_s_pipeline(
@@ -3971,6 +3970,7 @@ class BlackwellMultiHeadLatentAttentionForward:
             producer_group=mma_s_producer_group,
             consumer_group=mma_s_consumer_group,
             cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
         )
 
     def make_and_init_p_mma_pipeline(
@@ -4005,6 +4005,7 @@ class BlackwellMultiHeadLatentAttentionForward:
             producer_group=p_mma_producer_group,
             consumer_group=p_mma_consumer_group,
             cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
         )
 
     def make_and_init_p_cor_pipeline(
@@ -4033,6 +4034,7 @@ class BlackwellMultiHeadLatentAttentionForward:
             num_stages=self.p_cor_stage,
             producer_group=p_cor_producer_group,
             consumer_group=p_cor_consumer_group,
+            defer_sync=True,
         )
 
     def make_and_init_mma_o_pipeline(
@@ -4067,6 +4069,7 @@ class BlackwellMultiHeadLatentAttentionForward:
             producer_group=mma_o_producer_group,
             consumer_group=mma_o_consumer_group,
             cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
         )
 
     def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr):
@@ -4091,6 +4094,7 @@ class BlackwellMultiHeadLatentAttentionForward:
             num_stages=self.load_pt_stage,
             producer_group=load_pt_producer_group,
             consumer_group=load_pt_consumer_group,
+            defer_sync=True,
         )
 
     @staticmethod

examples/python/CuTeDSL/cute/fake_tensor.py

@@ -0,0 +1,48 @@
+import cutlass
+import cutlass.cute as cute
+
+
+"""
+Example of using fake tensors in CuTe.
+
+This script demonstrates how to use fake tensors in CuTe to drive compilation without creating actual tensors
+from frameworks like PyTorch or TensorFlow.
+
+Run this file directly to see the output type information.
+"""
+
+
+@cute.jit
+def print_tensor_type(t: cute.Tensor):
+    print(t)
+
+
+def run():
+    from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_tensor
+
+    shape = (3, 4)
+    a = make_fake_compact_tensor(cutlass.Float16, (3, 4), stride_order=(1, 0))
+    cute.compile(print_tensor_type, a)
+
+    # 32-bit symbolic integer with divisibility 8
+    shape = (3, cute.sym_int32(divisibility=8))
+    a = make_fake_compact_tensor(cutlass.Float16, shape, stride_order=(1, 0))
+    cute.compile(print_tensor_type, a)
+
+    # with static stride
+    a = make_fake_tensor(cutlass.Float16, shape, stride=(4, 1))
+    cute.compile(print_tensor_type, a)
+
+    # with dynamic stride using 32bit integer
+    stride = (cute.sym_int32(divisibility=8), 1)
+    a = make_fake_tensor(cutlass.Float16, shape, stride=stride)
+    cute.compile(print_tensor_type, a)
+
+    # with dynamic stride using 64bit integer
+    stride = (cute.sym_int64(divisibility=8), 1)
+    a = make_fake_tensor(cutlass.Float16, shape, stride=stride)
+    cute.compile(print_tensor_type, a)
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/cute/ffi/CMakeLists.txt

@@ -35,20 +35,26 @@ find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
 
 # Get Python site-packages directory using Python
 execute_process(
-    COMMAND ${Python3_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])"
-    OUTPUT_VARIABLE Python_SITE_PACKAGES
+    COMMAND ${Python3_EXECUTABLE} -c "import sys, sysconfig; print(';'.join([sysconfig.get_paths()['purelib'],sysconfig.get_paths(vars={'base': sys.base_prefix, 'platbase': sys.base_prefix})['purelib']]))"
+    OUTPUT_VARIABLE Python_SITE_PACKAGES_PATHS
     OUTPUT_STRIP_TRAILING_WHITESPACE
 )
 
-message(STATUS "Python site-packages directory: ${Python_SITE_PACKAGES}")
+message(STATUS "Python site-packages directories: ${Python_SITE_PACKAGES_PATHS}")
 
 # Add nanobind path to CMAKE_PREFIX_PATH
-list(APPEND CMAKE_PREFIX_PATH ${Python_SITE_PACKAGES}/nanobind/cmake)
+foreach(path IN LISTS Python_SITE_PACKAGES_PATHS)
+    if(EXISTS "${path}/nanobind/cmake")
+        message(STATUS "Adding nanobind cmake path: ${path}/nanobind/cmake")
+        list(APPEND CMAKE_PREFIX_PATH "${path}/nanobind/cmake")
+        break()
+    endif()
+endforeach()
 
 # Find nanobind
 find_package(nanobind)
 if(NOT nanobind_FOUND)
-    message(FATAL_ERROR 
+    message(FATAL_ERROR
         "nanobind not found!\n"
         "Please install nanobind with: pip install nanobind\n"
     )

examples/python/CuTeDSL/cute/ffi/jit_argument.py

@@ -243,12 +243,7 @@ import tempfile
 import torch
 
 
-def run_test(tmpdir=None, cmake_args=""):
-    # Skip cleanup if user provides tmpdir
-    cleanup = tmpdir is None
-    # Initialize temporary build directory
-    tmpdir = tmpdir or tempfile.mkdtemp()
-
+def run_test(tmpdir=None, cmake_args="", cleanup=True):
     try:
         current_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -256,8 +251,6 @@ def run_test(tmpdir=None, cmake_args=""):
         subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True)
         subprocess.run(["cmake", "--build", tmpdir], check=True)
 
-        sys.path.append(tmpdir)
-
         from tensor import make_tensor, pycapsule_get_pointer
 
         # Mock test tensor and corresponding C structure for this example
@@ -314,4 +307,13 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    run_test(tmpdir=args.tmp_dir, cmake_args=args.cmake_args)
+    if args.tmp_dir:
+        tmp_dir = args.tmp_dir
+        cleanup = False
+    else:
+        tmp_dir = tempfile.mkdtemp()
+        cleanup = True
+
+    sys.path.append(tmp_dir)
+
+    run_test(tmp_dir, args.cmake_args, cleanup)

examples/python/CuTeDSL/cute/print_latex.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import argparse
+import cutlass
+import cutlass.cute as cute
+from cutlass.utils import print_latex, print_latex_tv
+from cutlass import for_generate, yield_out
+
+"""
+A Latex Printing Example using CuTe DSL.
+This example prints latex for a given layout or thread value layout.
+The primary goal for this example is to demonstrate how to dump latex, which can then be
+turned into an image in your favorite latex compiler.
+To run this example:
+.. code-block:: bash
+    python examples/python/CuteDSL/cute/print_latex.py
+    python examples/python/CuteDSL/cute/print_latex.py --tv_layout
+To compile, pipe the output to a file and use a tool like pdflatex:
+.. code-block:: bash
+    python examples/python/CuTeDSL/cute/print_latex.py > latex.tex
+    pdflatex latex.tex
+"""
+
+
+@cute.jit
+def main(print_tv_layout: cutlass.Constexpr[bool]):
+    # Note: only support compile time printing layouts
+    if cutlass.const_expr(print_tv_layout):
+        thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+        val_layout = cute.make_ordered_layout((4, 1), order=(1, 0))
+        tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+        print_latex_tv(tv_layout, tiler_mn)
+    else:
+        layout = cute.make_layout((10, 10))
+        print_latex(layout)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="example of print latex and print latex tv"
+    )
+    parser.add_argument("--tv_layout", action="store_true")
+
+    args = parser.parse_args()
+
+    main(args.tv_layout)

examples/python/CuTeDSL/cute/torch_fake_tensor.py

@@ -65,7 +65,7 @@ def print_tensor(t: cute.Tensor):
     cute.print_tensor(t)
 
 
-if __name__ == "__main__":
+def run():
     from torch._subclasses.fake_tensor import FakeTensorMode
 
     shape = (3, 4)
@@ -75,3 +75,7 @@ if __name__ == "__main__":
 
     real_tensor = torch.randn(shape, dtype=torch.float32)
     compiled_fn(from_dlpack(real_tensor))
+
+
+if __name__ == "__main__":
+    run()

examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Export the compiled function to a shared library
+3. Load the shared library and use the compiled function to work with torch.Tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/cute/tvm_ffi/aot_export.py
+    # run example to use in torch
+    python examples/cute/tvm_ffi/aot_use_in_torch.py
+    # run example to use in jax
+    python examples/cute/tvm_ffi/aot_use_in_jax.py
+    # run example to use in c++ bundle
+    bash examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+"""
+
+from pathlib import Path
+import torch
+import os
+import subprocess
+import tvm_ffi
+import torch
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    a_cute = from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    b_cute = from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+
+    object_file_path = "./build/add_one.o"
+    lib_path = "./build/add_one.so"
+    compiled_add_one.export_to_c(object_file_path, function_name="add_one")
+    shared_libs = cute.runtime.find_runtime_libraries(enable_tvm_ffi=True)
+    # compile the object file to a shared library
+    cmd = ["gcc", "-shared", "-o", lib_path, object_file_path, *shared_libs]
+    print(cmd)
+    subprocess.run(cmd, check=True)
+    print(f"Successfully created shared library: {lib_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp

@@ -0,0 +1,97 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+ *
+ * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+ * property and proprietary rights in and to this material, related
+ * documentation and any modifications thereto. Any use, reproduction,
+ * disclosure or distribution of this material and related documentation
+ * without an express license agreement from NVIDIA CORPORATION or
+ * its affiliates is strictly prohibited.
+ */
+
+// clang-format on
+// This example shows how to interface with an AOT compiled function in a C++
+// bundle. to build and run the example, run the following command in project
+// root bash
+// examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+
+#include <cuda_runtime.h>
+#include <iostream>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/module.h>
+#include <vector>
+
+namespace ffi = tvm::ffi;
+
+struct CUDANDAlloc {
+  void AllocData(DLTensor *tensor) {
+    size_t data_size = ffi::GetDataSize(*tensor);
+    void *ptr = nullptr;
+    cudaError_t err = cudaMalloc(&ptr, data_size);
+    TVM_FFI_ICHECK_EQ(err, cudaSuccess)
+        << "cudaMalloc failed: " << cudaGetErrorString(err);
+    tensor->data = ptr;
+  }
+
+  void FreeData(DLTensor *tensor) {
+    if (tensor->data != nullptr) {
+      cudaError_t err = cudaFree(tensor->data);
+      TVM_FFI_ICHECK_EQ(err, cudaSuccess)
+          << "cudaFree failed: " << cudaGetErrorString(err);
+      tensor->data = nullptr;
+    }
+  }
+};
+
+inline ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) {
+  return ffi::Tensor::FromNDAlloc(CUDANDAlloc(), shape, dtype, device);
+}
+
+// symbol from the shared library
+extern "C" int __tvm_ffi_add_one(void *, const TVMFFIAny *, int32_t,
+                                 TVMFFIAny *);
+
+// Redirects into the exported function in object
+void CallAddOne(ffi::TensorView x, ffi::TensorView y) {
+  tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y);
+}
+
+int main() {
+  DLDataType f32_dtype{kDLFloat, 32, 1};
+  DLDevice cuda_device{kDLCUDA, 0};
+
+  constexpr int ARRAY_SIZE = 10;
+
+  ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
+  ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
+
+  std::vector<float> host_x(ARRAY_SIZE);
+  for (int i = 0; i < ARRAY_SIZE; ++i) {
+    host_x[i] = static_cast<float>(i);
+  }
+
+  size_t nbytes = host_x.size() * sizeof(float);
+  cudaError_t err =
+      cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice);
+  TVM_FFI_ICHECK_EQ(err, cudaSuccess)
+      << "cudaMemcpy host to device failed: " << cudaGetErrorString(err);
+
+  // Call into the FFI function; tensors remain on device because they carry a
+  // kDLCUDA device tag.
+  CallAddOne(x, y);
+
+  std::vector<float> host_y(host_x.size());
+  err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost);
+  TVM_FFI_ICHECK_EQ(err, cudaSuccess)
+      << "cudaMemcpy device to host failed: " << cudaGetErrorString(err);
+
+  std::cout << "y after add_one_cuda(x, y)" << std::endl;
+  for (float value : host_y) {
+    std::cout << value << " ";
+  }
+  std::cout << std::endl;
+  return 0;
+}

examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#!/bin/bash
+CUDA_DIALECT_PATH="build/lib/"
+export LD_LIBRARY_PATH=${CUDA_DIALECT_PATH}:`tvm-ffi-config --libdir`
+
+CUDA_HOME=/usr/local/cuda
+SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp"
+
+echo "Compiling the executable..."
+g++ -o build/aot_use_in_cpp_bundle \
+  -I${CUDA_HOME}/include \
+  `tvm-ffi-config --cxxflags`  \
+  ${SOURCE_FILE} build/add_one.o \
+  -L${CUDA_DIALECT_PATH} \
+  -L${CUDA_HOME}/lib64 \
+  -lcuda_dialect_runtime -lcuda -lcudart \
+  `tvm-ffi-config --ldflags` \
+  `tvm-ffi-config --libs`
+
+echo "Running the executable..."
+./build/aot_use_in_cpp_bundle

examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import jax
+import jax.numpy as jnp
+import jax_tvm_ffi
+import cutlass.cute as cute
+
+
+# now load it back
+def main():
+    a_jax = jnp.arange(10, dtype=jnp.float32)
+    b_jax = jnp.zeros(10, dtype=jnp.float32)
+    lib_path = "./build/add_one.so"
+    aot_mod = cute.runtime.load_module(lib_path)
+    jax_tvm_ffi.register_ffi_target("add_one_cute", aot_mod.add_one, platform="gpu")
+    b_jax = jax.ffi.ffi_call(
+        "add_one_cute",
+        jax.ShapeDtypeStruct(a_jax.shape, a_jax.dtype),
+        vmap_method="broadcast_all",
+    )(a_jax)
+    print("result of b after aot_mod.add_one(a, b)")
+    print(b_jax)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import cutlass.cute as cute
+import torch
+
+# now load it back
+def main():
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    lib_path = "./build/add_one.so"
+    aot_mod = cute.runtime.load_module(lib_path)
+    aot_mod.add_one(a_torch, b_torch)
+    print("result of b after aot_mod.add_one(a, b)")
+    print(b_torch)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Directly use the compiled function to work with torch.Tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/cute/tvm_ffi/error_reporting.py
+"""
+import torch
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    a_cute = from_dlpack(a_torch, enable_tvm_ffi=True)
+    b_cute = from_dlpack(b_torch, enable_tvm_ffi=True)
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+    # should raise an error because of shape mismatch
+    compiled_add_one(torch.arange(5, dtype=torch.float32, device="cuda"), b_cute)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Directly use the compiled function to work with JAX
+
+To run this example:
+
+.. code-block:: bash
+
+    pip install jax-tvm-ffi
+    pip install jax[cuda13]
+
+    python examples/cute/tvm_ffi/jit_and_use_in_jax.py
+"""
+
+import jax
+from jax import numpy as jnp
+import jax_tvm_ffi
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_jax = jnp.arange(
+        10,
+        dtype=jnp.float32,
+    )
+    b_jax = jnp.zeros(
+        10,
+        dtype=jnp.float32,
+    )
+    a_cute = from_dlpack(a_jax, enable_tvm_ffi=True).mark_layout_dynamic()
+    b_cute = from_dlpack(b_jax, enable_tvm_ffi=True).mark_layout_dynamic()
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+
+    # register the compiled function to JAX as a FFI target
+    jax_tvm_ffi.register_ffi_target("add_one_cute", compiled_add_one, platform="gpu")
+    a_jax = jnp.arange(10, dtype=jnp.float32)
+    # call the compiled function using JAX FFI
+    b_jax = jax.ffi.ffi_call(
+        "add_one_cute",
+        jax.ShapeDtypeStruct(a_jax.shape, a_jax.dtype),
+        vmap_method="broadcast_all",
+    )(a_jax)
+    print("result of b_jax after add_one_cute")
+    print(b_jax)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""Example demonstrating how to use TVM-FFI ABI with CuTe.
+
+This example shows how to:
+1. Compile a CuTe function with "--enable-tvm-ffi" option
+2. Directly use the compiled function to work with torch.Tensor
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/cute/tvm_ffi/jit_and_use_in_torch.py
+"""
+import torch
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+
+@cute.kernel
+def device_add_one(a: cute.Tensor, b: cute.Tensor):
+    for i in range(a.shape[0]):
+        b[i] = a[i] + 1
+
+
+@cute.jit
+def add_one(a: cute.Tensor, b: cute.Tensor):
+    """b = a + 1"""
+    device_add_one(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1))
+
+
+def main():
+    # compile the kernel with "--enable-tvm-ffi" option
+    a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
+    b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
+    a_cute = from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    b_cute = from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic()
+    # compile the kernel with "--enable-tvm-ffi" option
+    compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
+    # run the compiled function by passing in cute.Tensor as input
+    # you need to set enable_tvm_ffi=True for now
+    compiled_add_one(a_cute, b_cute)
+    # print the result
+    print("result of b after compiled_add_one(a, b)")
+    print(b_torch)
+    a_torch = a_torch + 1
+    # We can directly pass in torch.Tensor as input
+    # the call overhead is optimized so it is very fast to pass in torch.Tensor as input
+    # takes about less than 0.5us per call likely in terms of API overhead
+    compiled_add_one(a_torch, b_torch)
+    # print the result
+    print("result of b after compiled_add_one(a, b)")
+    print(b_torch)
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/helpers/fmha_helpers.py

@@ -0,0 +1,975 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# Use of this software is governed by the terms and conditions of the
+# NVIDIA End User License Agreement (EULA), available at:
+# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
+#
+# Any use, reproduction, disclosure, or distribution of this software
+# and related documentation outside the scope permitted by the EULA
+# is strictly prohibited.
+
+import enum
+from typing import Tuple, Optional
+import cutlass
+from cutlass.cute.typing import Boolean
+
+from cutlass.cutlass_dsl import (
+    Int32,
+    Float32,
+    min,
+    extract_mlir_values,
+    new_from_mlir_values,
+)
+from cutlass.utils.hardware_info import HardwareInfo
+from cutlass.utils import WorkTileInfo
+import cutlass.cute as cute
+
+##############################################################################
+# Fmha static tile scheduler
+##############################################################################
+
+
+class FmhaStaticTileSchedulerParams:
+    """A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler.
+
+    This class holds the configuration parameters needed to initialize and configure
+    the tile scheduler for FMHA operations.
+
+    :ivar is_persistent: Whether to use persistent kernel mode.
+    :type is_persistent: bool
+    :ivar problem_shape_mbh: Problem shape in (M, B, H) format.
+    :type problem_shape_mbh: cute.Shape
+    """
+
+    def __init__(
+        self,
+        is_persistent: bool,
+        problem_shape_mbh: cute.Shape,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        Initializes the FmhaStaticTileSchedulerParams with the given parameters.
+
+        :param is_persistent: Whether to use persistent kernel mode.
+        :type is_persistent: bool
+        :param problem_shape_mbh: Problem shape in (M, B, H) format.
+        :type problem_shape_mbh: cute.Shape
+        """
+        self.is_persistent = is_persistent
+        self.problem_shape_mbh = problem_shape_mbh
+        self._loc = loc
+        self._ip = ip
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.problem_shape_mbh]:
+            obj_values = extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        for obj, n_items in zip([self.problem_shape_mbh], self._values_pos):
+            obj_list.append(new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return FmhaStaticTileSchedulerParams(
+            self.is_persistent, *(tuple(obj_list)), loc=self._loc
+        )
+
+
+class FmhaStaticTileScheduler:
+    """A static tile scheduler for FMHA (Fused Multi-Head Attention) operations.
+
+    This class manages the scheduling of work tiles for FMHA kernels, supporting
+    both persistent and non-persistent kernel modes. It tracks the current work
+    position and advances through the problem space efficiently.
+
+    :ivar _params: Scheduler parameters.
+    :type _params: FmhaStaticTileSchedulerParams
+    :ivar _blk_coord: Block coordinates.
+    :type _blk_coord: cute.Coord
+    :ivar _grid_shape: Grid shape for the kernel.
+    :type _grid_shape: cute.Shape
+    :ivar _is_persistent: Whether to use persistent kernel mode.
+    :type _is_persistent: bool
+    :ivar _current_work_linear_idx: Current linear work index.
+    :type _current_work_linear_idx: Int32
+    :ivar _problem_shape_mbh: Problem shape in (M, B, H) format.
+    :type _problem_shape_mbh: cute.Layout
+    :ivar _num_blocks: Number of blocks in the problem.
+    :type _num_blocks: Int32
+    :ivar _is_first_block: Whether this is the first block.
+    :type _is_first_block: bool
+    :ivar num_persistent_sm: Number of persistent SMs.
+    :type num_persistent_sm: Int32
+    """
+
+    def __init__(
+        self,
+        params: FmhaStaticTileSchedulerParams,
+        current_work_linear_idx: Int32,
+        blk_coord: cute.Coord,
+        grid_shape: cute.Shape,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        Initializes the FmhaStaticTileScheduler with the given parameters.
+
+        :param params: Scheduler parameters.
+        :type params: FmhaStaticTileSchedulerParams
+        :param current_work_linear_idx: Current linear work index.
+        :type current_work_linear_idx: Int32
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param grid_shape: Grid shape for the kernel.
+        :type grid_shape: cute.Shape
+        """
+        self._params = params
+        self._blk_coord = blk_coord
+        self._grid_shape = grid_shape
+        self._is_persistent = params.is_persistent
+        self._current_work_linear_idx = current_work_linear_idx
+        self._problem_shape_mbh = cute.make_layout(
+            params.problem_shape_mbh, loc=loc, ip=ip
+        )
+        self._num_blocks = cute.size(self._problem_shape_mbh, loc=loc, ip=ip)
+        self._is_first_block = True
+        self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip)
+        self._loc = loc
+        self._ip = ip
+
+    # called by host
+    @staticmethod
+    def get_grid_shape(
+        params: FmhaStaticTileSchedulerParams,
+        *,
+        loc=None,
+        ip=None,
+    ) -> cute.Shape:
+        """
+        Determine the grid shape for the FMHA kernel.
+
+        For persistent kernels, the grid shape is limited by the number of SMs
+        (Streaming Multiprocessors) available on the device. For non-persistent
+        kernels, the grid shape matches the problem shape.
+
+        :param params: Scheduler parameters.
+        :type params: FmhaStaticTileSchedulerParams
+
+        :return: Grid shape as (M, B, H) tuple.
+        :rtype: cute.Shape
+        """
+        if params.is_persistent:
+            hardware_info = HardwareInfo()
+            sm_count = hardware_info.get_device_multiprocessor_count()
+            return (
+                min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)),
+                1,
+                1,
+            )
+        else:
+            return params.problem_shape_mbh
+
+    @staticmethod
+    def check_valid_work_for_seqlen_q(
+        q_tiler: int,
+        current_idx: Int32,
+        seqlen_q: Int32,
+    ) -> Boolean:
+        """
+        Check if the current work index is valid for the given query sequence length.
+
+        This method verifies that the current work tile index multiplied by the
+        query tiler size is within the bounds of the query sequence length.
+
+        :param q_tiler: Query tiler size.
+        :type q_tiler: int
+        :param current_idx: Current work index.
+        :type current_idx: Int32
+        :param seqlen_q: Query sequence length.
+        :type seqlen_q: Int32
+
+        :return: True if the work is valid, False otherwise.
+        :rtype: Boolean
+        """
+        return current_idx * q_tiler < seqlen_q
+
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        """
+        Get information about the current work tile.
+
+        Determines if the current work is valid and computes the tile coordinates
+        based on whether the kernel is persistent or non-persistent.
+
+        :return: WorkTileInfo containing tile coordinates and validity flag.
+        :rtype: WorkTileInfo
+        """
+        is_valid = (
+            self._current_work_linear_idx < self._num_blocks
+            if self._is_persistent
+            else self._is_first_block
+        )
+
+        blk_coord = (0, 0, 0)
+        if self._is_persistent:
+            blk_coord = self._problem_shape_mbh.get_hier_coord(
+                self._current_work_linear_idx, loc=loc, ip=ip
+            )
+        else:
+            blk_coord = self._blk_coord
+
+        # cur_tile_coord is (mid, 0, (bid, hid))
+        cur_tile_coord = (
+            blk_coord[0],
+            0,
+            (blk_coord[1], blk_coord[2]),
+        )
+
+        return WorkTileInfo(cur_tile_coord, is_valid)
+
+    def initial_work_tile_info(self, *, loc=None, ip=None):
+        """
+        Get the initial work tile information.
+
+        :return: Initial WorkTileInfo.
+        :rtype: WorkTileInfo
+        """
+        return self.get_current_work(loc=loc, ip=ip)
+
+    def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None):
+        """
+        Advance to the next work tile.
+
+        For persistent kernels, advances by the number of persistent SMs.
+        For non-persistent kernels, marks that the first block has been processed.
+
+        :param advance_count: Number of steps to advance (default: 1).
+        :type advance_count: int
+        """
+        if self._is_persistent:
+            self._current_work_linear_idx += advance_count * self.num_persistent_sm
+        self._is_first_block = False
+
+    def __extract_mlir_values__(self):
+        values = extract_mlir_values(self._params)
+        values.extend(extract_mlir_values(self._current_work_linear_idx))
+        values.extend(extract_mlir_values(self._blk_coord))
+        values.extend(extract_mlir_values(self._grid_shape))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 10
+        new_params = new_from_mlir_values(self._params, values[0:3])
+        new_current_work_linear_idx = new_from_mlir_values(
+            self._current_work_linear_idx, [values[3]]
+        )
+        new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7])
+        new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:])
+        return FmhaStaticTileScheduler(
+            new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape
+        )
+
+
+def create_fmha_static_tile_scheduler(
+    params: FmhaStaticTileSchedulerParams,
+    blk_coord: cute.Coord,
+    grid_shape: cute.Shape,
+) -> FmhaStaticTileScheduler:
+    """
+    Create a new FMHA static tile scheduler.
+
+    :param params: Scheduler parameters.
+    :type params: FmhaStaticTileSchedulerParams
+    :param blk_coord: Block coordinates.
+    :type blk_coord: cute.Coord
+    :param grid_shape: Grid shape.
+    :type grid_shape: cute.Shape
+
+    :return: New FmhaStaticTileScheduler instance.
+    :rtype: FmhaStaticTileScheduler
+    """
+    return FmhaStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape)
+
+
+def create_fmha_static_tile_scheduler_params(
+    is_persistent: bool,
+    problem_shape_mbh: cute.Shape,
+) -> FmhaStaticTileSchedulerParams:
+    """
+    Create FMHA static tile scheduler parameters.
+
+    :param is_persistent: Whether to use persistent kernel mode.
+    :type is_persistent: bool
+    :param problem_shape_mbh: Problem shape in (M, B, H) format.
+    :type problem_shape_mbh: cute.Shape
+
+    :return: New FmhaStaticTileSchedulerParams instance.
+    :rtype: FmhaStaticTileSchedulerParams
+    """
+    return FmhaStaticTileSchedulerParams(is_persistent, problem_shape_mbh)
+
+
+def compute_grid(
+    o_shape: cute.Shape,
+    cta_tiler: Tuple[int, int, int],
+    is_persistent: bool,
+) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]:
+    """
+    Compute grid parameters for FMHA operation.
+
+    This function calculates the appropriate grid shape and scheduler parameters
+    based on the output tensor shape, CTA (Cooperative Thread Array) tiler,
+    and whether to use persistent kernel mode.
+
+    The output tensor o has shape (s, d, ((h_r, h_k), b)) where:
+    - s: sequence length
+    - d: head dimension
+    - h_r: number of heads for query
+    - h_k: number of heads for key
+    - b: batch size
+
+    :param o_shape: Output tensor shape for grid computation.
+    :type o_shape: cute.Shape
+    :param cta_tiler: CTA tiler dimensions (M, N, K).
+    :type cta_tiler: Tuple[int, int, int]
+    :param is_persistent: Whether to use persistent kernel mode.
+    :type is_persistent: bool
+
+    :return: Tuple of (scheduler_params, grid_shape).
+    :rtype: Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]
+    """
+    tile_sched_params = create_fmha_static_tile_scheduler_params(
+        is_persistent,
+        (
+            cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]),
+            cute.size(o_shape[2][0]),
+            cute.size(o_shape[2][1]),
+        ),
+    )
+    grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params)
+
+    return tile_sched_params, grid
+
+
+##############################################################################
+# Fused Mask
+##############################################################################
+
+
+class MaskEnum(enum.Enum):
+    """Enumeration of mask types for FMHA operations.
+
+    - RESIDUAL_MASK: Residual mask for handling variable sequence lengths
+    - WINDOW_MASK: Window mask for attention which also includes causal and no mask
+    - WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k
+    - WINDOW_MASK_BWD: Window mask for backward pass
+    - WINDOW_MASK_BWD_INFERENCE: Same as the window mask for backward pass, but has the limitation that the end of q is aligned with the end of k
+    """
+
+    RESIDUAL_MASK = enum.auto()
+    RESIDUAL_MASK_BWD = enum.auto()
+    WINDOW_MASK = enum.auto()
+    WINDOW_MASK_INFERENCE = enum.auto()
+    WINDOW_MASK_BWD = enum.auto()
+    WINDOW_MASK_BWD_INFERENCE = enum.auto()
+
+
+class FusedMask:
+    """A fused mask implementation for FMHA operations.
+
+    This class handles different types of attention masks including no mask,
+    residual mask for variable sequence lengths, and causal mask for
+    autoregressive attention patterns.
+
+    The class provides methods to:
+    - Calculate trip counts for different mask types
+    - Apply masks to attention scores
+    - Handle masked and unmasked trip calculations
+    """
+
+    def get_trip_count(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Calculate the number of trips needed for the current block.
+
+        The trip count depends on the mask type and the block coordinates.
+        For causal masks, it considers the autoregressive constraint.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+
+        :return: Number of trips needed.
+        :rtype: Int32
+        """
+        result = 0
+        offset = 0
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        if cutlass.const_expr(mask_type == MaskEnum.RESIDUAL_MASK):
+            result = cute.ceil_div(seqlen_k, tile_shape[1])
+        if cutlass.const_expr(mask_type is MaskEnum.RESIDUAL_MASK_BWD):
+            result = cute.ceil_div(seqlen_q, tile_shape[0])
+        if cutlass.const_expr(
+            mask_type == MaskEnum.WINDOW_MASK
+            or mask_type == MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is None):
+                result = cute.ceil_div(seqlen_k, tile_shape[1])
+            else:
+                max_idx_q = (blk_coord[0] + 1) * tile_shape[0]
+                idx_k = max_idx_q + offset + window_size_right
+                tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1])
+                max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1])
+                result = min(max_blocks_k, tmp_blocks_k)
+        if cutlass.const_expr(
+            mask_type == MaskEnum.WINDOW_MASK_BWD
+            or mask_type == MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_left is None):
+                result = cute.ceil_div(seqlen_q, tile_shape[0])
+            else:
+                max_idx_k = (blk_coord[1] + 1) * tile_shape[1]
+                idx_k = max_idx_k + offset + window_size_left
+                tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0])
+                max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0])
+                result = min(max_blocks_q, tmp_blocks_q)
+        start_block = FusedMask.get_trip_start(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        result = result - start_block
+        return result
+
+    @cute.jit
+    def get_trip_start(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Get the start of the trip for the current block.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        """
+        result = 0
+        offset = 0
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        if cutlass.const_expr(
+            mask_type is MaskEnum.WINDOW_MASK
+            or mask_type is MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_left is not None):
+                min_idx_q = blk_coord[0] * tile_shape[0]
+                idx_k = min_idx_q + offset - window_size_left
+                tmp_blocks_k = idx_k // tile_shape[1]
+                result = max(tmp_blocks_k, result)
+        if cutlass.const_expr(
+            mask_type is MaskEnum.WINDOW_MASK_BWD
+            or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is not None):
+                min_idx_k = blk_coord[1] * tile_shape[1]
+                idx_q = min_idx_k + offset - window_size_right
+                tmp_blocks_q = idx_q // tile_shape[0]
+                result = max(tmp_blocks_q, result)
+        return result
+
+    @cute.jit
+    def get_leading_mask_id(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Tuple[Int32, Int32]:
+        """
+        Get the begin and end tile idx for the leading mask.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+
+        :return: Tuple of (begin, end) tile idx for the leading mask.
+        :rtype: Tuple[Int32, Int32]
+        """
+        offset = 0
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        leading_mask_begin = FusedMask.get_trip_start(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        trip_count = FusedMask.get_trip_count(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+
+        leading_mask_end = leading_mask_begin
+        if cutlass.const_expr(
+            mask_type is MaskEnum.WINDOW_MASK
+            or mask_type is MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_left is not None):
+                min_idx_q = (
+                    (blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left
+                )
+                leading_mask_end = min(
+                    cute.ceil_div(min_idx_q, tile_shape[1]) - 1,
+                    trip_count + leading_mask_begin - 1,
+                )
+            else:
+                leading_mask_end = leading_mask_begin - 1
+        elif cutlass.const_expr(
+            mask_type is MaskEnum.WINDOW_MASK_BWD
+            or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is not None):
+                min_idx_k = (
+                    (blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right
+                )
+                leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1
+            else:
+                leading_mask_end = leading_mask_begin - 1
+        return leading_mask_begin, leading_mask_end
+
+    @cute.jit
+    def get_trailing_mask_id(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Tuple[Optional[Int32], Optional[Int32]]:
+        """
+        Get the begin and end tile idx for the trailing mask.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+
+        :return: Tuple of (begin, end) tile idx for the trailing mask.
+        :rtype: Tuple[Int32, Int32]
+        """
+        offset = 0
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_INFERENCE):
+            offset = seqlen_k - seqlen_q
+        if cutlass.const_expr(mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE):
+            offset = seqlen_q - seqlen_k
+        trip_start = FusedMask.get_trip_start(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+        trip_count = FusedMask.get_trip_count(
+            mask_type,
+            blk_coord,
+            tile_shape,
+            seqlen_q,
+            seqlen_k,
+            window_size_left,
+            window_size_right,
+        )
+
+        trailing_mask_begin, trailing_mask_end = None, None
+        if cutlass.const_expr(
+            mask_type is MaskEnum.WINDOW_MASK
+            or mask_type is MaskEnum.WINDOW_MASK_INFERENCE
+        ):
+            if cutlass.const_expr(window_size_right is not None):
+                min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right
+                trailing_mask_begin = min(
+                    min_idx_q // tile_shape[1], trip_count + trip_start - 1
+                )
+                trailing_mask_end = trip_count + trip_start - 1
+            else:
+                # last tile, we always apply mask on it regardless whether it's a residual tile
+                trailing_mask_begin = trip_count + trip_start - 1
+                trailing_mask_end = trip_count + trip_start - 1
+        else:
+            if cutlass.const_expr(window_size_left is not None):
+                min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1
+                max_idx_k = (
+                    (blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left
+                )
+                trailing_mask_begin = min(
+                    cute.ceil_div(min_idx_k, tile_shape[0]) - 1,
+                    trip_count + trip_start - 1,
+                )
+                trailing_mask_end = min(
+                    cute.ceil_div(max_idx_k, tile_shape[0]) - 1,
+                    trip_count + trip_start - 1,
+                )
+            else:
+                # last tile, we always apply mask on it regardless whether it's a residual tile
+                trailing_mask_begin = trip_count + trip_start - 1
+                trailing_mask_end = trip_count + trip_start - 1
+
+        return trailing_mask_begin, trailing_mask_end
+
+    @cute.jit
+    def get_masked_leading_count(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Calculate the number of masked trips for the leading mask.
+
+        This is used for blocks that need special handling due to masking.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+
+        :return: Number of masked trips.
+        :rtype: Int32
+        """
+        result = 0
+        if cutlass.const_expr(
+            mask_type is not MaskEnum.RESIDUAL_MASK
+            and mask_type is not MaskEnum.RESIDUAL_MASK_BWD
+        ):
+            if cutlass.const_expr(
+                window_size_left is not None or window_size_right is not None
+            ):
+                leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id(
+                    mask_type,
+                    blk_coord,
+                    tile_shape,
+                    seqlen_q,
+                    seqlen_k,
+                    window_size_left,
+                    window_size_right,
+                )
+                result = max(leading_mask_end - leading_mask_begin + 1, 0)
+
+        return result
+
+    @cute.jit
+    def get_masked_trailing_count(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+        rem_count: Optional[Int32] = 0,
+    ) -> Int32:
+        """
+        Calculate the number of masked trips for the trailing mask.
+
+        This is used for blocks that need special handling due to masking.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+        :param rem_count: Remaining count from previous calculations.
+        :type rem_count: Int32
+
+        :return: Number of masked trips.
+        :rtype: Int32
+        """
+        result = 0
+
+        if cutlass.const_expr(
+            mask_type is not MaskEnum.RESIDUAL_MASK
+            and mask_type is not MaskEnum.RESIDUAL_MASK_BWD
+        ):
+            if cutlass.const_expr(
+                window_size_left is not None or window_size_right is not None
+            ):
+                trailing_mask_begin, trailing_mask_end = FusedMask.get_trailing_mask_id(
+                    mask_type,
+                    blk_coord,
+                    tile_shape,
+                    seqlen_q,
+                    seqlen_k,
+                    window_size_left,
+                    window_size_right,
+                )
+                leading_mask_begin, leading_mask_end = FusedMask.get_leading_mask_id(
+                    mask_type,
+                    blk_coord,
+                    tile_shape,
+                    seqlen_q,
+                    seqlen_k,
+                    window_size_left,
+                    window_size_right,
+                )
+                if cutlass.const_expr(
+                    trailing_mask_begin is not None and trailing_mask_end is not None
+                ):
+                    if trailing_mask_begin <= leading_mask_end:
+                        result = max(trailing_mask_end - leading_mask_end, 0)
+                    else:
+                        result = max(trailing_mask_end - trailing_mask_begin + 1, 0)
+        else:
+            if seqlen_k % tile_shape[1] != 0:
+                result = 1
+            else:
+                result = 0
+
+        return result + rem_count
+
+    @cute.jit
+    def get_unmasked_trip_count(
+        mask_type: MaskEnum,
+        blk_coord: cute.Coord,
+        tile_shape: cute.Shape,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[Int32] = None,
+        window_size_right: Optional[Int32] = None,
+    ) -> Int32:
+        """
+        Calculate the number of unmasked trips for the current block.
+
+        This represents the number of trips that don't require special
+        masking treatment.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param blk_coord: Block coordinates.
+        :type blk_coord: cute.Coord
+        :param tile_shape: Shape of the tile.
+        :type tile_shape: cute.Shape
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Int32
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[Int32]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[Int32]
+
+        :return: Number of unmasked trips.
+        :rtype: Int32
+        """
+        result = (
+            FusedMask.get_trip_count(
+                mask_type,
+                blk_coord,
+                tile_shape,
+                seqlen_q,
+                seqlen_k,
+                window_size_left,
+                window_size_right,
+            )
+            - FusedMask.get_masked_leading_count(
+                mask_type,
+                blk_coord,
+                tile_shape,
+                seqlen_q,
+                seqlen_k,
+                window_size_left,
+                window_size_right,
+            )
+            - FusedMask.get_masked_trailing_count(
+                mask_type,
+                blk_coord,
+                tile_shape,
+                seqlen_q,
+                seqlen_k,
+                window_size_left,
+                window_size_right,
+                0,
+            )
+        )
+        return result
+
+    @cute.jit
+    def apply_mask(
+        mask_type: MaskEnum,
+        acc_qk: cute.Tensor,
+        index_qk: cute.Tensor,
+        seqlen_q: Int32,
+        seqlen_k: Int32,
+        window_size_left: Optional[int] = None,
+        window_size_right: Optional[int] = None,
+        index_transform: cutlass.Constexpr = lambda index_q, index_k: (
+            index_q,
+            index_k,
+        ),
+    ):
+        """
+        Apply the appropriate mask to the attention scores.
+
+        This method modifies the attention scores (acc_qk) based on the mask type
+        and the positions in the index tensor.
+
+        :param mask_type: Type of mask to use
+        :type mask_type: utils.MaskEnum
+        :param acc_qk: Accumulated QK attention scores tensor.
+        :type acc_qk: cute.Tensor
+        :param index_qk: Index tensor containing position information.
+        :type index_qk: cute.Tensor
+        :param seqlen_k: Key sequence length for attention computation.
+        :type seqlen_k: Int32
+        :param seqlen_q: Query sequence length for attention computation.
+        :type seqlen_q: Optional[int]
+        :param window_size_left: Left-side sliding window size for attention masking.
+        :type window_size_left: Optional[int]
+        :param window_size_right: Right-side sliding window size for attention masking.
+        :type window_size_right: Optional[int]
+        """
+
+        tidx, tidy, tidx = cute.arch.thread_idx()
+        offset = 0
+        offset = (
+            seqlen_k - seqlen_q
+            if cutlass.const_expr(
+                mask_type is MaskEnum.WINDOW_MASK_INFERENCE
+                or mask_type is MaskEnum.WINDOW_MASK_BWD_INFERENCE
+            )
+            else 0
+        )
+        for i in cutlass.range_constexpr(cute.size(acc_qk)):
+            index_q, index_k = index_transform(*index_qk[i])
+            if cutlass.const_expr(
+                window_size_left is not None or window_size_right is not None
+            ):
+                if cutlass.const_expr(window_size_left is None):
+                    if index_q + offset + window_size_right < index_k:
+                        acc_qk[i] = -Float32.inf
+                    if index_k >= seqlen_k or index_q >= seqlen_q:  # residual mask
+                        acc_qk[i] = -Float32.inf
+                elif cutlass.const_expr(window_size_right is None):
+                    if index_q + offset - window_size_left > index_k:
+                        acc_qk[i] = -Float32.inf
+                    if index_k >= seqlen_k or index_q >= seqlen_q:  # residual mask
+                        acc_qk[i] = -Float32.inf
+                else:
+                    max_K_index = min(index_q + offset + window_size_right, seqlen_k)
+                    min_K_index = max(0, index_q + offset - window_size_left)
+                    if index_k > max_K_index or index_k < min_K_index:
+                        acc_qk[i] = -Float32.inf
+                    if index_k >= seqlen_k or index_q >= seqlen_q:  # residual mask
+                        acc_qk[i] = -Float32.inf
+
+            if cutlass.const_expr(
+                mask_type == MaskEnum.RESIDUAL_MASK
+                or mask_type == MaskEnum.RESIDUAL_MASK_BWD
+            ):
+                if index_k >= seqlen_k or index_q >= seqlen_q:
+                    acc_qk[i] = -Float32.inf

examples/python/CuTeDSL/helpers/sparse_utils.py

@@ -0,0 +1,457 @@
+import numpy as np
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+import torch
+
+
+@cute.jit
+def print_tensor_dlpack(src: cute.Tensor):
+    print(src)
+    cute.print_tensor(src)
+
+
+# Sparse emulation
+class SparseEmulation:
+    def __init__(self, M: int, N: int, K: int, L: int):
+        self.M = M
+        self.N = N
+        self.K = K
+        self.L = L
+
+    @cute.jit
+    def __call__(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor):
+        """Sparse emulation"""
+        num_threads = 128
+        grid = (cute.ceil_div(self.M, num_threads), 1, 1)
+        block = (num_threads, 1, 1)
+        self.kernel(a, b, d, e).launch(grid=grid, block=block)
+        return
+
+    @cute.kernel
+    def kernel(self, a: cute.Tensor, b: cute.Tensor, d: cute.Tensor, e: cute.Tensor):
+        """CUDA kernel to emulate sparse tensor core"""
+        tidx, tidy, tidz = cute.arch.thread_idx()
+        bidx, bidy, bidz = cute.arch.block_idx()
+
+        row_idx = tidx + bidx * self.M
+        meta_idx = self.K // 4 // 8
+        if row_idx < self.M:
+            # each thread process 1 row
+            for col in range(self.N):
+                # each meta_idx stands for 32 elements
+                for e_idx in range(meta_idx):
+                    meta_val = e[(row_idx, e_idx)]
+                    for k in range(8):
+                        # each k stands for 4 elements
+                        meta_row = (meta_val >> (k * 4)) & 0xF
+                        idx0 = meta_row & 0x3
+                        idx1 = (meta_row >> 2) & 0x3
+                        # calculate the idx in b tensor which has value in A tensor
+                        km = e_idx * 16 + k * 2
+                        km_1 = km + 1
+                        kn = e_idx * 32 + k * 4 + idx0
+                        kn_1 = e_idx * 32 + k * 4 + idx1
+                        d[row_idx, col] += a[row_idx, km] * b[col, kn]
+                        d[row_idx, col] += a[row_idx, km_1] * b[col, kn_1]
+        return
+
+
+# Compressor
+# compress a sparse tensor to a dense tensor && generate metadata
+class Compressor:
+    def __init__(self, M: int, K: int, L: int):
+        self.M = M
+        self.K = K
+        self.L = L
+        self.pos_map = {
+            0x4: [0, 1],
+            0x8: [0, 2],
+            0xC: [0, 3],
+            0x9: [1, 2],
+            0xD: [1, 3],
+            0xE: [2, 3],
+        }
+
+    @cute.jit
+    def _init__(self, a: cute.Tensor):
+        self.__init__(a.shape[0], a.shape[1], a.shape[2])
+
+    def compress(self, a, a_compressed, meta, run_on_cpu: bool):
+        if run_on_cpu:
+            if a.device.type != "cpu":
+                raise ValueError("a must be on cpu")
+            return self.__compress_on_cpu(a, a_compressed, meta)
+        else:
+            if a.device.type != "cuda":
+                raise ValueError("a must be on cuda")
+            return self.__compress_on_cuda(a, a_compressed, meta)
+
+    def __compress_on_cpu(self, a, a_compressed, meta):
+        """
+        compress the tensor on cpu
+        # Convert to 4-bit metadata value
+        # The metadata value represents which 2 elements are non-zero
+        # 0x4: [1,1,0,0] - first two elements are non-zero
+        # 0x8: [1,0,1,0] - first and third elements are non-zero
+        # 0xC: [1,0,0,1] - first and fourth elements are non-zero
+        # 0x9: [0,1,1,0] - second and third elements are non-zero
+        # 0xD: [0,1,0,1] - second and fourth elements are non-zero
+        # 0xE: [0,0,1,1] - third and fourth elements are non-zero
+        # special case:
+        # [0,0,0,0]  == [0,0,1,1]
+        # [1,0,0,0]  == [1,0,0,1]
+        # [0,1,0,0]  == [0,1,0,1]
+        # [0,0,1,0]  == [0,0,1,1]
+        # [0,0,0,1]  == [0,0,1,1]
+        """
+        M, K = a.shape
+        assert a_compressed.shape == (
+            M,
+            K // 2,
+        ), f"Expected a_compressed shape {(M, K // 2)}, got {a_compressed.shape}"
+        assert meta.shape == (
+            M,
+            K // 4 // 8,
+        ), f"Expected meta shape {(M, K // 4 // 8)}, got {meta.shape}"
+        for m in range(M):
+            k_meta = 0
+            for k in range(0, K, 4):
+                chunk = a[m, k : k + 4]
+
+                non_zero_indices = torch.nonzero(chunk).squeeze()
+                meta_val = 0xE
+                if torch.equal(non_zero_indices, torch.tensor([0, 1])):
+                    meta_val = 0x4
+                elif torch.equal(non_zero_indices, torch.tensor([0, 2])):
+                    meta_val = 0x8
+                elif torch.equal(non_zero_indices, torch.tensor([0, 3])) or torch.equal(
+                    non_zero_indices, torch.tensor(0)
+                ):
+                    meta_val = 0xC
+                elif torch.equal(non_zero_indices, torch.tensor([1, 2])):
+                    meta_val = 0x9
+                elif torch.equal(non_zero_indices, torch.tensor([1, 3])) or torch.equal(
+                    non_zero_indices, torch.tensor(1)
+                ):
+                    meta_val = 0xD
+                elif torch.equal(non_zero_indices, torch.tensor([2, 3])) or torch.equal(
+                    non_zero_indices, torch.tensor(2)
+                ):
+                    meta_val = 0xE
+                elif torch.equal(non_zero_indices, torch.tensor([])) or torch.equal(
+                    non_zero_indices, torch.tensor(3)
+                ):
+                    meta_val = 0xE
+                else:
+                    raise ValueError(f"Invalid non-zero pattern: {non_zero_indices}")
+                meta_idx = k // 4 // 8
+                meta_bit_pos = (k // 4) % 8
+                if k_meta == meta_idx:
+                    k_meta = meta_idx + 1
+                    meta[m, meta_idx] = 0
+                meta[m, meta_idx] |= meta_val << (meta_bit_pos * 4)
+                compressed_idx = k // 2
+                index = self.pos_map[meta_val]
+                a_compressed[m, compressed_idx] = chunk[index[0]]
+                a_compressed[m, compressed_idx + 1] = chunk[index[1]]
+
+    def __compress_on_cuda(self, a, a_compressed, meta):
+        """
+        compress the tensor on cuda
+        """
+        a_tensor = from_dlpack(a)
+        a_compressed_tensor = from_dlpack(a_compressed)
+        meta_tensor = from_dlpack(meta)
+        self.compress_on_cuda_impl(a_tensor, a_compressed_tensor, meta_tensor)
+        return
+
+    @cute.jit
+    def compress_on_cuda_impl(
+        self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor
+    ):
+        """Compress the input tensor using the metadata"""
+        num_threads = 128
+        grid = (cute.ceil_div(self.M, num_threads), 1, 1)
+        block = (num_threads, 1, 1)
+        self.compressor_impl(a, a_compressed, meta).launch(grid=grid, block=block)
+
+    @cute.kernel
+    def compressor_impl(
+        self, a: cute.Tensor, a_compressed: cute.Tensor, meta: cute.Tensor
+    ):
+        """CUDA kernel to compress the tensor"""
+        tidx, tidy, tidz = cute.arch.thread_idx()
+        bidx, bidy, bidz = cute.arch.block_idx()
+        m = a.shape[0]
+        k = a.shape[1]
+
+        # each thread process 1 row
+        row_idx = tidx + bidx * self.M
+        meta_idx = self.K // 4 // 8
+        if row_idx < self.M:
+            # each meta_idx stands for 32 elements
+            for i in range(meta_idx):
+                meta[row_idx, i] = 0
+                # each k stands for 4 elements
+                for j in range(8):
+                    val = a[row_idx, i * 32 + j * 4]
+                    val_1 = a[row_idx, i * 32 + j * 4 + 1]
+                    val_2 = a[row_idx, i * 32 + j * 4 + 2]
+                    val_3 = a[row_idx, i * 32 + j * 4 + 3]
+                    value_idx = 0
+                    value_idx_1 = 0
+                    value_idx_2 = 0
+                    value_idx_3 = 0
+                    pos0 = 0
+                    pos1 = 0
+                    if val != 0:
+                        value_idx = 1
+                        pos0 = 0
+                    if val_1 != 0:
+                        value_idx_1 = 1
+                    if val_2 != 0:
+                        value_idx_2 = 1
+                    if val_3 != 0:
+                        value_idx_3 = 1
+                    pos = [value_idx, value_idx_1, value_idx_2, value_idx_3]
+                    tmp = 0
+                    if pos == [0, 0, 0, 0]:
+                        tmp = 0xE
+                        pos0 = 2
+                        pos1 = 3
+                    elif pos == [1, 0, 0, 0]:
+                        tmp = 0xC
+                        pos0 = 0
+                        pos1 = 3
+                    elif pos == [0, 1, 0, 0]:
+                        tmp = 0xD
+                        pos0 = 1
+                        pos1 = 3
+                    elif pos == [0, 0, 1, 0]:
+                        tmp = 0xE
+                        pos0 = 2
+                        pos1 = 3
+                    elif pos == [0, 0, 0, 1]:
+                        tmp = 0xE
+                        pos0 = 2
+                        pos1 = 3
+                    elif pos == [1, 1, 0, 0]:
+                        tmp = 0x4
+                        pos0 = 0
+                        pos1 = 1
+                    elif pos == [1, 0, 1, 0]:
+                        tmp = 0x8
+                        pos0 = 0
+                        pos1 = 2
+                    elif pos == [1, 0, 0, 1]:
+                        tmp = 0xC
+                        pos0 = 0
+                        pos1 = 3
+                    elif pos == [0, 1, 1, 0]:
+                        tmp = 0x9
+                        pos0 = 1
+                        pos1 = 2
+                    elif pos == [0, 1, 0, 1]:
+                        tmp = 0xD
+                        pos0 = 1
+                        pos1 = 3
+                    elif pos == [0, 0, 1, 1]:
+                        tmp = 0xE
+                        pos0 = 2
+                        pos1 = 3
+                    # cute.printf(row_idx, cutlass.Float32(val), cutlass.Float32(val_1), cutlass.Float32(val_2), cutlass.Float32(val_3), tmp)
+                    meta[row_idx, i] |= tmp << (j * 4)
+
+                    a_compressed[row_idx, i * 16 + j * 2] = a[
+                        row_idx, i * 32 + j * 4 + pos0
+                    ]
+                    a_compressed[row_idx, i * 16 + j * 2 + 1] = a[
+                        row_idx, i * 32 + j * 4 + pos1
+                    ]
+
+        return
+
+
+# SparseUtils is used to generate sparse tensor
+# format torch.Tensor
+class SparseUtils:
+    #!brief: SparseUtils is used to generate sparse tensor
+    #!param: M: int, K: int, L: int, dtype: cutlass.DataType
+    def __init__(self, M: int, K: int, L: int, dtype):
+        self.M = M
+        self.K = K
+        self.L = L
+        self.dtype = dtype
+        self.meta_data = self._generate_meta_data_4_2()
+        self._use_specific_meta_data = False
+
+    #!brief: cast cutlass.DataType to torch.Tensor
+    def _get_type(self):
+        if self.dtype == cutlass.Float16:
+            return torch.float16
+        elif self.dtype == cutlass.Float32:
+            return torch.float32
+        elif self.dtype == cutlass.Int8:
+            return torch.int8
+        else:
+            raise ValueError(f"Unsupported dtype: {self.dtype}")
+
+    def _generate_meta_data_4_2(self):
+        # metadata for 4:2 sparse will in range( 4,8,9,c,d,e)
+        # represents
+        # 0: [1,1,0,0] no zero pos 00,01 -> 0100 = 4
+        # 1: [1,0,1,0] no zero pos 00,10 -> 1000 = 8
+        # 2: [1,0,0,1] no zero pos 00,11 -> 1100 = c
+        # 3: [0,1,1,0] no zero pos 01,10 -> 1001 = 9
+        # 4: [0,1,0,1] no zero pos 01,11 -> 1101 = d
+        # 5: [0,0,1,1] no zero pos 10,11 -> 1011 = e
+        meta_value = [0x4, 0x8, 0x9, 0xC, 0xD, 0xE]
+        # 4:2 sparse, so each chunk is 4 elements, map to 4 bits
+        K_NumChunk = self.K // 4
+        meta_data = np.random.choice(
+            meta_value, size=(self.M, K_NumChunk), replace=True
+        )
+        meta_data = torch.from_numpy(
+            np.array(meta_data).astype(np.uint8).reshape(self.M, K_NumChunk)
+        )
+        return meta_data
+
+    #!brief: pack meta data
+    def _pack_meta_data(self):
+        tmp = []
+        K_NumChunk = self.K // 4
+        for i in range(self.M):
+            for j in range(K_NumChunk // 8):
+                v = 0
+                for k in range(8):
+                    vv = int(self.meta_data[i, j * 8 + k] & 0xF)
+                    tt = vv << (k * 4)
+                    v = v | tt
+                tmp.append(v)
+        # debug print
+        # print([hex(vt) for vt in tmp])
+        result = torch.from_numpy(
+            np.array(tmp).astype(np.uint32).reshape(self.M, K_NumChunk // 8)
+        )
+        return result
+
+    #!brief: use specific meta data
+    def use_specific_meta_data(self, meta_data: torch.Tensor = None):
+        if meta_data is not None:
+            self.meta_data = meta_data
+        self._use_specific_meta_data = True
+
+    #!brief: generate sparse tensor with tensor
+    #!param: a: torch.Tensor
+    #!param: run_on_cpu: bool
+    #!return: torch.Tensor
+    def generate_sparse_4_2_tensor_with_tensor(self, a, run_on_cpu):
+        if run_on_cpu:
+            if a.device.type != "cpu":
+                raise ValueError("a must be on cpu")
+            return self.__generate_sparse_tensor_cpu(a)
+        else:
+            if a.device.type != "cuda":
+                raise ValueError("a must be on cuda")
+            a_tensor = from_dlpack(a)
+            packed_meta_data = self._pack_meta_data()
+            meta_tensor = from_dlpack(packed_meta_data.cuda())
+            self.__generate_sparse_tensor_cuda(a_tensor, meta_tensor)
+            return a
+
+    #!brief: generate sparse tensor
+    #!param: run_on_cpu: bool
+    #!return: torch.Tensor
+    def generate_4_2_sparse_tensor(self, run_on_cpu):
+        dtype = self._get_type()
+        a = torch.empty(self.M, self.K).random_(-5, 5).to(dtype)
+        if run_on_cpu:
+            return self.generate_sparse_4_2_tensor_with_tensor(a, run_on_cpu)
+        else:
+            return self.generate_sparse_4_2_tensor_with_tensor(a.cuda(), run_on_cpu)
+
+    #!brief: generate sparse tensor on cpu
+    #!param: a: torch.Tensor
+    #!return: torch.Tensor
+    def __generate_sparse_tensor_cpu(self, a):
+        if not self._use_specific_meta_data:
+            for m in range(self.M):
+                for k in range(0, self.K, 4):
+                    # random choose 2 zero positions
+                    zero_indices = torch.randperm(4)[:2]
+                    a[m, k + zero_indices[0]] = 0
+                    a[m, k + zero_indices[1]] = 0
+            return a
+        else:
+            # use specific meta data
+            tensor_mask = []
+            for i in range(self.M):
+                for j in range(self.K // 4):
+                    meta_val = self.meta_data[i, j]
+                    tmp = []
+                    if meta_val == 0x4:
+                        tmp = [1, 1, 0, 0]
+                    elif meta_val == 0x8:
+                        tmp = [1, 0, 1, 0]
+                    elif meta_val == 0xC:
+                        tmp = [1, 0, 0, 1]
+                    elif meta_val == 0x9:
+                        tmp = [0, 1, 1, 0]
+                    elif meta_val == 0xD:
+                        tmp = [0, 1, 0, 1]
+                    elif meta_val == 0xE:
+                        tmp = [0, 0, 1, 1]
+                    tensor_mask.extend(tmp)
+            a = torch.reshape(a, (-1,))
+            mask = torch.tensor(tensor_mask)
+            a = a * mask
+            a = torch.reshape(a, (self.M, self.K))
+            return a
+
+    @cute.jit
+    def __generate_sparse_tensor_cuda(self, a: cute.Tensor, meta: cute.Tensor):
+        """Generate a sparse tensor from a dense tensor using metadata"""
+        assert a.shape[0] == self.M and a.shape[1] == self.K
+        assert meta.shape[0] == self.M and meta.shape[1] == self.K // 4 // 8
+        num_threads = 128
+        grid = (cute.ceil_div(self.M, num_threads), 1, 1)
+        block = (num_threads, 1, 1)
+        self.kernel(a, meta).launch(grid=grid, block=block)
+
+    @cute.kernel
+    def kernel(self, a: cute.Tensor, meta: cute.Tensor):
+        """Apply sparsity mask to input tensor using metadata"""
+        tidx, tidy, tidz = cute.arch.thread_idx()
+        bidx, bidy, bidz = cute.arch.block_idx()
+
+        # each thread process 1 ro
+        row_idx = tidx + bidx * self.M
+        meta_idx = self.K // 4 // 8
+        # each thread process 1 row
+        if row_idx < self.M:
+            # iterate over each chunk(32 elements)
+            for i in range(meta_idx):
+                meta_val = meta[(row_idx, i)]
+                # iterate over each sparse pattern(4 elements)
+                for j in range(8):
+                    meta_row = (meta_val >> (j * 4)) & 0xF
+                    idx0 = meta_row & 0x3
+                    idx1 = (meta_row >> 2) & 0x3
+                    r_id0 = 0
+                    r_id1 = 0
+                    # r_id is the idx that value is 0
+                    if idx0 >= 2 and idx1 >= 2:
+                        r_id0 = 0
+                        r_id1 = 1
+                    elif idx0 <= 1 and idx1 <= 1:
+                        r_id0 = 2
+                        r_id1 = 3
+                    else:
+                        r_id0 = idx0 ^ 0b1
+                        r_id1 = idx1 ^ 0b1
+                    row_id0 = r_id0 + i * 32 + j * 4
+                    row_id1 = r_id1 + i * 32 + j * 4
+                    a[row_idx, row_id0] = self.dtype(0.0)
+                    a[row_idx, row_id1] = self.dtype(0.0)
+        return

examples/python/CuTeDSL/helpers/test_sparse_utils.py

@@ -0,0 +1,104 @@
+import sparse_utils as su
+import cutlass
+import torch
+from cutlass.cute.runtime import from_dlpack
+import numpy as np
+import pytest
+
+
+@pytest.mark.L0
+def test_sparse_cpu():
+    M = 128
+    N = 32
+    K = 32
+    L = 1
+    debug = False
+    # generate sparse tensor
+    a = torch.empty(M, K).random_(-5, 5).to(torch.float16)
+    sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16)
+    if debug:
+        sparse_utils.use_specific_meta_data()
+    a_gen_from_cpu = sparse_utils.generate_sparse_4_2_tensor_with_tensor(a, True)
+    # print(a_gen_from_cpu)
+    # generate compressed tensor and meta data
+    a_compressed_cpu = torch.empty(M, K // 2).to(torch.float16)
+    meta_data_cpu = torch.empty(M, K // 4 // 8).to(torch.uint32)
+    compressor = su.Compressor(M, K, L)
+    compressor.compress(a_gen_from_cpu, a_compressed_cpu, meta_data_cpu, True)
+    # # test with gemm
+    b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda()
+    d = torch.empty(M, N).zero_().to(torch.float16).cuda()
+    b_tensor = from_dlpack(b)
+    d_tensor = from_dlpack(d)
+    a_compressed_cpu_tensor = from_dlpack(a_compressed_cpu.cuda())
+    meta_data_cpu_tensor = from_dlpack(meta_data_cpu.cuda())
+    sparse_emulation = su.SparseEmulation(M, N, K, 1)
+    sparse_emulation(a_compressed_cpu_tensor, b_tensor, d_tensor, meta_data_cpu_tensor)
+
+    ref = torch.einsum("mk,nk->mn", a_gen_from_cpu.cpu(), b.cpu())
+    if debug:
+        a_ori = a_gen_from_cpu.cpu().numpy()
+        np.savetxt("a.txt", a_ori, fmt="%f")
+        a_compressed_cpu_ori = a_compressed_cpu.cpu().numpy()
+        np.savetxt("a_compressed_cpu.txt", a_compressed_cpu_ori, fmt="%f")
+        meta_data_cpu_ori = meta_data_cpu.cpu().numpy()
+        np.savetxt("meta_data_cpu.txt", meta_data_cpu_ori, fmt="%f")
+        d_ori = d.cpu().numpy()
+        np.savetxt("d.txt", d_ori, fmt="%f")
+        ref_ori = ref.cpu().numpy()
+        np.savetxt("ref.txt", ref_ori, fmt="%f")
+    torch.testing.assert_close(d.cpu(), ref)
+    print("cpu d == ref")
+
+
+@pytest.mark.L0
+def test_sparse_cuda():
+    M = 128
+    N = 32
+    K = 32
+    L = 1
+    debug = False
+    sparse_utils = su.SparseUtils(M, K, L, cutlass.Float16)
+    if debug:
+        sparse_utils.use_specific_meta_data()
+    # generate sparse tensor
+    a = torch.empty(M, K).random_(-5, 5).to(torch.float16).cuda()
+    a_gen_from_cuda = sparse_utils.generate_4_2_sparse_tensor(False)
+    # print(a_gen_from_cuda)
+    # generate compressed tensor and meta data
+    a_compressed_cuda = torch.empty(M, K // 2).to(torch.float16).cuda()
+    meta_data_cuda = torch.empty(M, K // 4 // 8).to(torch.uint32).cuda()
+    compressor = su.Compressor(M, K, L)
+    compressor.compress(a_gen_from_cuda, a_compressed_cuda, meta_data_cuda, False)
+    # test with gemm
+    b = torch.empty(N, K).random_(-5, 5).to(torch.float16).cuda()
+    d = torch.empty(M, N).zero_().to(torch.float16).cuda()
+    b_tensor = from_dlpack(b)
+    d_tensor = from_dlpack(d)
+    a_compressed_cuda_tensor = from_dlpack(a_compressed_cuda)
+    meta_data_cuda_tensor = from_dlpack(meta_data_cuda)
+    sparse_emulation = su.SparseEmulation(M, N, K, 1)
+    sparse_emulation(
+        a_compressed_cuda_tensor, b_tensor, d_tensor, meta_data_cuda_tensor
+    )
+
+    ref = torch.einsum("mk,nk->mn", a_gen_from_cuda.cpu(), b.cpu())
+    if debug:
+        a_ori = a_gen_from_cuda.cpu().numpy()
+        np.savetxt("a.txt", a_ori, fmt="%f")
+        a_compressed_cuda_ori = a_compressed_cuda.cpu().numpy()
+        np.savetxt("a_compressed_cuda.txt", a_compressed_cuda_ori, fmt="%f")
+        meta_data_cuda_ori = meta_data_cuda.cpu().numpy()
+        np.savetxt("meta_data_cuda.txt", meta_data_cuda_ori, fmt="%f")
+        d_ori = d.cpu().numpy()
+        np.savetxt("d.txt", d_ori, fmt="%f")
+        ref_ori = ref.cpu().numpy()
+        np.savetxt("ref.txt", ref_ori, fmt="%f")
+    torch.testing.assert_close(d.cpu(), ref)
+    print("cuda d == ref")
+
+
+if __name__ == "__main__":
+    cutlass.cuda.initialize_cuda_context()
+    test_sparse_cpu()
+    test_sparse_cuda()

examples/python/CuTeDSL/hopper/dense_gemm.py

@@ -38,6 +38,7 @@ import cutlass.cute as cute
 import cutlass.cute.testing as testing
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 from cutlass.cute.runtime import from_dlpack
 import cutlass.utils.hopper_helpers as sm90_utils
@@ -624,11 +625,11 @@ class HopperWgmmaGemmKernel:
             consumer_group=mainloop_pipeline_consumer_group,
             tx_count=tma_copy_bytes,
             cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
         )
 
         #  Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         # ///////////////////////////////////////////////////////////////////////////////
         #  Generate smem tensor A/B
@@ -717,10 +718,7 @@ class HopperWgmmaGemmKernel:
         #  Cluster wait
         # ///////////////////////////////////////////////////////////////////////////////
         # cluster wait for barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
         # /////////////////////////////////////////////////////////////////////////////
         #  Prefetch
         # /////////////////////////////////////////////////////////////////////////////

examples/python/CuTeDSL/hopper/dense_gemm_persistent.py

@@ -37,6 +37,7 @@ import cutlass
 import cutlass.cute as cute
 import cutlass.cute.testing as testing
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 import cutlass.utils as utils
 import cutlass.utils.hopper_helpers as sm90_utils
@@ -634,11 +635,11 @@ class HopperWgmmaGemmPersistentKernel:
             consumer_group=mainloop_pipeline_consumer_group,
             tx_count=tma_copy_bytes,
             cta_layout_vmnk=cute.make_layout((1, *cta_layout_mnk.shape)),
+            defer_sync=True,
         )
 
         # Cluster arrive after barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_arrive_relaxed()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
 
         # Generate smem tensor A/B
         sA = storage.sA.get_tensor(
@@ -718,10 +719,7 @@ class HopperWgmmaGemmPersistentKernel:
         k_tile_cnt = cute.size(gA_mkl, mode=[3])
 
         # Cluster wait for barrier init
-        if cute.size(self.cluster_shape_mn) > 1:
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
 
         is_dma_warp_group = warp_group_idx < self.num_dma_warp_groups
         if is_dma_warp_group:

examples/python/CuTeDSL/hopper/fmha.py

@@ -95,15 +95,18 @@ import cutlass.cute.testing as testing
 import cutlass.cute.nvgpu.warpgroup as warpgroup
 import cutlass.utils as utils
 import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
 import cutlass.torch as cutlass_torch
 from cutlass._mlir.dialects import math as _math
 
 import cutlass.utils.hopper_helpers as sm90_utils
 from cutlass.cute.runtime import from_dlpack
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(current_dir, ".."))
-from utils import fmha_helpers as fmha_utils
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, os.path.join(current_dir, ".."))
+
+from helpers import fmha_helpers as fmha_utils
 
 
 class HopperFusedMultiHeadAttentionForward:
@@ -648,11 +651,8 @@ class HopperFusedMultiHeadAttentionForward:
         # We need this to guarantee that the Pipeline init is visible
         # To all producers and consumer blocks in the Cluster
         # and to finish smem init
-        if cute.size(self.cluster_shape_mnk) > 1:
-            cute.arch.cluster_arrive_relaxed()
-            cute.arch.cluster_wait()
-        else:
-            cute.arch.sync_threads()
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
 
         if warp_idx == 0:
             cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q)
@@ -1646,6 +1646,7 @@ class HopperFusedMultiHeadAttentionForward:
             producer_group=load_q_producer_group,
             consumer_group=load_q_consumer_group,
             tx_count=self.tma_copy_q_bytes,
+            defer_sync=True,
         ).make_participants()
 
     def make_and_init_load_kv_pipeline(self, load_kv_mbar_ptr):
@@ -1663,6 +1664,7 @@ class HopperFusedMultiHeadAttentionForward:
             producer_group=load_kv_producer_group,
             consumer_group=load_kv_consumer_group,
             tx_count=self.tma_copy_kv_bytes,
+            defer_sync=True,
         ).make_participants()
 
     def make_and_init_tma_store_pipeline(self):
@@ -1686,6 +1688,7 @@ class HopperFusedMultiHeadAttentionForward:
                 pipeline.Agent.Thread,
                 self.num_threads_per_warp_group,
             ),
+            defer_sync=True,
         )
 
     @staticmethod