v4.5 tag update (#3202)

* Python DSL examples reorganization.

* v4.5 tag update.

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/README.md

@@ -0,0 +1,25 @@
+# CUTLASS Tutorial Examples for Blackwell GEMM
+
+This folder contains tutorial examples demonstrating how to write performant GEMM (General Matrix Multiplication) kernels using Tensor Cores on NVIDIA Blackwell GPUs.
+
+## Overview
+
+The examples showcase different scenarios and optimization techniques for implementing GEMM operations:
+
+- Basic FP16 GEMM implementation
+- Software Pipeline optimizations
+- Tensor Core utilization
+- Thread/warp/block level parallelism
+
+## Examples
+
+### tutorial_fp16_gemm_0.py
+
+A basic example showing:
+- FP16 GEMM implementation using Tensor Cores
+- TMA (Tensor Memory Access) for efficient data loading
+- SMEM (Shared Memory) layouts and access patterns
+- Usage of ``cutlass.range(..., prefetch_stages=...)`` to replace boilerplate code for multi-stage software pipeline
+
+With some minor optimization tricks
+- Tiling Epilogue to avoid bursty write out and reduce register pressure

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/fp16_gemm_0.py

@@ -0,0 +1,441 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 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.
+
+import argparse
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.runtime import from_dlpack
+
+"""
+The first tutorial GEMM demonstrating a simple kernel implementation in CuTeDSL
+
+This dense GEMM kernel is implemented in just over 200 lines of code.
+With large tile sizes, it can achieve very high performance on 8k×8k×8k problem sizes.
+It can serve as a starting point to help users quickly experiment
+with optimizations for challenges that may arise with other problem sizes.
+
+To run this example:
+.. code-block:: bash
+
+    python examples/blackwell/tutorial_gemm/fp16_gemm_0.py  \
+      --mnk 8192,8192,8192
+
+Constraints for this example:
+* The problem size of m and n must be divisible by the tile size m & n (128, 256)
+"""
+
+io_dtype = cutlass.Float16
+acc_dtype = cutlass.Float32
+mma_inst_shape_mnk = (128, 256, 16)
+mma_tiler_mnk = (128, 256, 64)
+threads_per_cta = 128
+
+# Pipeline stage configuration
+ab_stages = 4
+acc_stage = 1
+
+
+@cute.struct
+class SharedStorage:
+    ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
+    acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stage * 2]
+    tmem_holding_buf: cutlass.Int32
+
+
+@cute.kernel
+def kernel(
+    tiled_mma: cute.TiledMma,
+    tma_atom_a: cute.CopyAtom,
+    mA_mkl: cute.Tensor,
+    tma_atom_b: cute.CopyAtom,
+    mB_nkl: cute.Tensor,
+    mC_mnl: cute.Tensor,
+    a_smem_layout: cute.ComposedLayout,
+    b_smem_layout: cute.ComposedLayout,
+):
+    # Current thread/warp/block coordinates
+    tidx, _, _ = cute.arch.thread_idx()
+    warp_idx = cute.arch.warp_idx()
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+    bidx, bidy, _ = cute.arch.block_idx()
+    mma_coord_mnk = (bidx, bidy, None)
+
+    #
+    # 1. Prepare args
+    #
+
+    # Allocate SMEM
+    smem = cutlass.utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+    sA = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=a_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=a_smem_layout.inner,
+    )
+    sB = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=b_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=b_smem_layout.inner,
+    )
+
+    # Allocate all TMEM columns
+    tmem_alloc_barrier = pipeline.NamedBarrier(
+        barrier_id=1,
+        num_threads=threads_per_cta,
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buf.ptr,
+        barrier_for_retrieve=tmem_alloc_barrier,
+    )
+    num_tmem_cols = 512
+    tmem.allocate(num_tmem_cols)
+
+    # Prefetch tma descriptor
+    if warp_idx == 0:
+        cpasync.prefetch_descriptor(tma_atom_a)
+        cpasync.prefetch_descriptor(tma_atom_b)
+
+    # Pipeline configuration
+    num_tma_copy_bytes = cute.size_in_bytes(
+        io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2])
+    ) + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
+    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+        num_stages=ab_stages,
+        producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+        consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+        tx_count=num_tma_copy_bytes,
+        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+    ).make_participants()
+    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+        num_stages=acc_stage,
+        producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+        consumer_group=pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            threads_per_cta,
+        ),
+        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+    ).make_participants()
+
+    # Partition tensors for MMA and make fragments
+    # (bM, bK, RestK)
+    gA = cute.local_tile(mA_mkl, mma_tiler_mnk, mma_coord_mnk, proj=(1, None, 1))
+    # (bN, bK, RestK)
+    gB = cute.local_tile(mB_nkl, mma_tiler_mnk, mma_coord_mnk, proj=(None, 1, 1))
+    # (bM, bN)
+    gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
+    thr_mma = tiled_mma.get_slice(0)
+    # (MMA, MMA_M, MMA_K)
+    tCgA = thr_mma.partition_A(gA)
+    # (MMA, MMA_N, MMA_K)
+    tCgB = thr_mma.partition_B(gB)
+    # (MMA, MMA_M, MMA_N)
+    tCgC = thr_mma.partition_C(gC)
+    # (MMA, MMA_M, MMA_K)
+    tCrA = tiled_mma.make_fragment_A(sA)
+    # (MMA, MMA_N, MMA_K)
+    tCrB = tiled_mma.make_fragment_B(sB)
+    # (MMA, MMA_M, MMA_N)
+    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+    # (MMA, MMA_M, MMA_N)
+    tCtAcc = tiled_mma.make_fragment_C(acc_shape)
+    # Partition tensors for TMA; This requires the tensors partitioned for MMA
+    tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_a,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sA, 0, 3),
+        cute.group_modes(tCgA, 0, 3),
+    )
+    tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_b,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sB, 0, 3),
+        cute.group_modes(tCgB, 0, 3),
+    )
+
+    # CTA-wide sync before retrieving the pointer to the start of the allocated TMEM
+    # Only warp 0 does the allocation so we need to sync before retrieving the TMEM start address
+    tmem.wait_for_alloc()
+    tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+    # Swap the pointer in tCtAcc
+    tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc.layout)
+
+    subtile_cnt = 4
+    # (EpiTile)
+    epi_tiler = (
+        (cute.size(tCtAcc, mode=[0, 0]), cute.size(tCtAcc, mode=[0, 1]) // subtile_cnt),
+    )
+    # (EpiTile, NumTiles)
+    tCtAcc_epi = cute.zipped_divide(tCtAcc, epi_tiler)
+    # (EpiTile, NumTiles)
+    gC_epi = cute.zipped_divide(tCgC, epi_tiler)
+
+    # Every thread loads 64 x fp32
+    tmem_atom = cute.make_copy_atom(
+        tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
+        cutlass.Float32,
+    )
+    tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_atom, tCtAcc_epi[None, 0])
+    tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
+
+    # (TmemCpy,NumTmemCpy,NumTiles)
+    tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
+    # (TmemCpy,NumTmemCpy,NumTiles)
+    tDgC = tmem_thr_copy.partition_D(gC_epi)
+
+    # (TmemCpy,NumTmemCpy)
+    tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
+    # (TmemCpy,NumTmemCpy)
+    tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
+
+    #
+    # 2. Main loop
+    #
+    num_k_tiles = cute.size(gA, mode=[2])
+    if warp_idx == 0:
+        # Wait for a empty accumulator buffer
+        acc_empty = acc_producer.acquire_and_advance()
+        for k_tile_idx in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
+            # Issue TMA loads
+            ab_empty = ab_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_a,
+                tAgA[(None, ab_empty.count)],
+                tAsA[(None, ab_empty.index)],
+                tma_bar_ptr=ab_empty.barrier,
+            )
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, ab_empty.count)],
+                tBsB[(None, ab_empty.index)],
+                tma_bar_ptr=ab_empty.barrier,
+            )
+
+            # Execute one K-block worth of MMA instructions
+            ab_full = ab_consumer.wait_and_advance()
+            num_k_blocks = cute.size(tCrA, mode=[2])
+            for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                k_block_coord = (None, None, k_block_idx, ab_full.index)
+                cute.gemm(
+                    tiled_mma,
+                    tCtAcc,
+                    tCrA[k_block_coord],
+                    tCrB[k_block_coord],
+                    tCtAcc,
+                )
+                tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+            # Signal that the A/B buffers have been consumed and are ready for the next load
+            ab_full.release()
+
+        # Signal that the accumulator is fully computed
+        acc_empty.commit()
+
+    #
+    # 3. Epilogue
+    #
+
+    # Release TMEM allocation lock
+    tmem.relinquish_alloc_permit()
+
+    # Wait for the accumulator buffer to be full
+    acc_full = acc_consumer.wait_and_advance()
+
+    # TMEM -> RMEM -> GEMM
+    # Sub-tiling for better instruction-level parallelism
+    for i in cutlass.range(cute.size(tDtC, mode=[2])):
+        cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
+        tCrC.store(tCrAcc.load().to(io_dtype))
+        cute.autovec_copy(tCrC, tDgC[None, None, i])
+    acc_full.release()
+
+    # Deallocate TMEM
+    pipeline.sync(barrier_id=1)
+    tmem.free(tmem_ptr)
+
+
+@cute.jit
+def host_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
+    # Construct tiled MMA
+    op = tcgen05.MmaF16BF16Op(
+        io_dtype,
+        acc_dtype,
+        mma_inst_shape_mnk,
+        tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        tcgen05.OperandMajorMode.K,
+        tcgen05.OperandMajorMode.K,
+    )
+    tiled_mma = cute.make_tiled_mma(op)
+
+    # Construct SMEM layouts for A and B
+    a_smem_layout = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a.element_type,
+        ab_stages,
+    )
+    b_smem_layout = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b.element_type,
+        ab_stages,
+    )
+    a_smem_layout_one_stage = cute.select(a_smem_layout, mode=[0, 1, 2])
+    b_smem_layout_one_stage = cute.select(b_smem_layout, mode=[0, 1, 2])
+
+    # Construct TMA load atoms
+    op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
+    a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        op,
+        a,
+        a_smem_layout_one_stage,
+        mma_tiler_mnk,
+        tiled_mma,
+    )
+    b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        op,
+        b,
+        b_smem_layout_one_stage,
+        mma_tiler_mnk,
+        tiled_mma,
+    )
+
+    # Pretty prints kernel attributes useful for debugging
+    # print(f"a            = {cute.pretty_str(a)}")
+    # print(f"b            = {cute.pretty_str(b)}")
+    # print(f"c            = {cute.pretty_str(c)}")
+    # print(f"tiled_mma    = {cute.pretty_str(tiled_mma)}")
+    # print(f"a_tma_atom   = {cute.pretty_str(a_tma_atom)}")
+    # print(f"b_tma_atom   = {cute.pretty_str(b_tma_atom)}")
+    # print(f"a_tma_tensor = {cute.pretty_str(a_tma_tensor)}")
+    # print(f"b_tma_tensor = {cute.pretty_str(b_tma_tensor)}")
+
+    # Launch the kernel
+    grid_shape = cute.ceil_div((*c.layout.shape, 1), mma_tiler_mnk[:2])
+    kernel(
+        tiled_mma,
+        a_tma_atom,
+        a_tma_tensor,
+        b_tma_atom,
+        b_tma_tensor,
+        c,
+        a_smem_layout,
+        b_smem_layout,
+    ).launch(
+        grid=grid_shape,
+        block=(threads_per_cta, 1, 1),
+    )
+
+
+def run_dense_gemm(
+    mnk: Tuple[int, int, int],
+    tolerance: float,
+):
+    global torch, cutlass_torch
+    import torch
+    import cutlass.torch as cutlass_torch
+
+    print("===================================================================")
+    print("Running Blackwell fp16 GEMM example 0 with:")
+    print(f"  mnk:       {mnk}")
+    print(f"  tolerance: {tolerance}")
+    print("===================================================================")
+    print()
+
+    m, n, k = mnk
+    torch.manual_seed(1111)
+
+    # Make K-major tensors (torch tensors are row-major)
+    def make_tensors(mn, k, dtype):
+        shape = (mn, k)
+        return (
+            torch.empty(*shape, dtype=torch.int32)
+            .random_(-2, 2)
+            .to(dtype=dtype, device="cuda")
+        )
+
+    a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
+    b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
+    c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
+    a_tensor = (
+        from_dlpack(a, assumed_align=32)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=k)
+    )
+    b_tensor = (
+        from_dlpack(b, assumed_align=32)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=k)
+    )
+    c_tensor = (
+        from_dlpack(c, assumed_align=32)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=n)
+    )
+
+    # Entry point to the host JIT function
+    host_function(a_tensor, b_tensor, c_tensor, no_cache=True)
+
+    # Compute reference result and verify
+    ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu()
+
+    torch.testing.assert_close(
+        c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
+    )
+
+
+if __name__ == "__main__":
+
+    def parse_comma_separated_ints(s: str):
+        try:
+            return [int(x.strip()) for x in s.split(",")]
+        except ValueError:
+            raise argparse.ArgumentTypeError(
+                "Invalid format. Expected comma-separated integers."
+            )
+
+    from cuda.bindings import driver as cu_driver
+
+    cu_driver.cuInit(0)
+    err, device_count = cu_driver.cuDeviceGetCount()
+    if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
+        raise RuntimeError("A GPU is required to run this example")
+
+    parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 0")
+    parser.add_argument(
+        "--mnk",
+        type=parse_comma_separated_ints,
+        default=[8192, 8192, 8192],
+        help="MNK dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    args = parser.parse_args()
+    if len(args.mnk) != 3:
+        parser.error("--mnk must contain exactly 3 values")
+    if args.mnk[0] % mma_tiler_mnk[0] != 0 or args.mnk[1] % mma_tiler_mnk[1] != 0:
+        parser.error("m n must be divisible by mma_tiler_mn")
+
+    run_dense_gemm(
+        args.mnk,
+        args.tolerance,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/fp16_gemm_1.py

@@ -0,0 +1,533 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 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.
+
+# This is the second tutorial GEMM. It builds on the first tutorial by adding 2CTA MMA
+# instructions with a 2x1 cluster.
+
+
+import argparse
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.runtime import from_dlpack
+
+"""
+The second tutorial GEMM demonstrating a simple kernel implementation in CuTeDSL
+
+With large tile sizes, it can as well achieve very high performance on 8kx8kx8k problem sizes.
+Compared with fp16_gemm_0.py, this example adds 2CTA MMA & TMA multicast supports.
+For fp16_gemm_0.py running at relative high SM frequency, the dram latency will be a potential performance issue.
+This example can achieve better performance than fp16_gemm_0.py due to:
+
+1. The 2CTA MMA can reduce B tensor smem size which allows larger ab stages to hide dram latency.
+
+For both 1CTA & 2CTA, one stage of A tensor smem size is 128x64xsizeof(float16)=16KB
+Situation for B is different.
+For 1CTA, one stage of B tensor smem size is 256x64xsizeof(float16)=32KB,
+while for 2CTA, we can only take half size, i.e. 16KB.
+So, the maxmimum AB stage for 1CTA is 227 // (16 + 32) = 4, while for 2CTA is 227 // (16 + 16) = 7.
+The latency hiding capability is 512 * (4 - 1) = 1.5K cycles for 1CTA, while 512 * (7 - 1) = 3K cycles for 2CTA.
+
+2. The L2 traffic is reduced due to the TMA multicast.
+
+For a (m, n) cluster shape, the L2 traffic for one tile is 16KB / n + 32KB / m.
+
+16KB / 1 + 32KB / 2 = 24KB for 2x1 cluster shape
+16KB / 4 + 32KB / 4 = 12KB for 4x4 cluster shape
+
+If no TMA multicast enabled, the L2 traffic for one tile is less than 16KB + 32KB = 48KB, which depends on hardware optimization.
+
+The first one can provide large latency hiding capability while the second one can reduce the data ready time.
+These two factors should be considered for latency/memory throughput bound cases.
+
+
+To run this example:
+.. code-block:: bash
+
+    python examples/blackwell/tutorial_gemm/fp16_gemm_1.py  \
+      --mnk 8192,8192,8192
+
+Constraints for this example:
+* The problem size of m and n must be divisible by the tile size m & n (256, 256)
+"""
+
+io_dtype = cutlass.Float16
+acc_dtype = cutlass.Float32
+cluster_shape_mnk = (2, 1, 1)
+mma_inst_shape_mnk = (256, 256, 16)
+mma_tiler_mnk = (256, 256, 64)
+threads_per_cta = 128
+
+# Pipeline stage configuration
+ab_stages = 7
+acc_stage = 1
+
+
+@cute.struct
+class SharedStorage:
+    ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
+    acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stage * 2]
+    tmem_dealloc_mbar: cutlass.Int64
+    tmem_holding_buf: cutlass.Int32
+
+
+@cute.kernel()
+def kernel(
+    tiled_mma: cute.TiledMma,
+    tma_atom_a: cute.CopyAtom,
+    mA_mkl: cute.Tensor,
+    tma_atom_b: cute.CopyAtom,
+    mB_nkl: cute.Tensor,
+    mC_mnl: cute.Tensor,
+    a_smem_layout: cute.ComposedLayout,
+    b_smem_layout: cute.ComposedLayout,
+    cta_layout_vmnk: cute.Layout,
+):
+    # Current thread/warp/block coordinates
+    tidx, _, _ = cute.arch.thread_idx()
+    warp_idx = cute.arch.warp_idx()
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+    bidx, bidy, _ = cute.arch.block_idx()
+    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+    cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+    mma_coord_vmnk = (
+        bidx % cute.size(cta_layout_vmnk, mode=[0]),
+        bidx // cute.size(cta_layout_vmnk, mode=[0]),
+        bidy,
+        None,
+    )
+    mma_coord_mnk = mma_coord_vmnk[1:]
+
+    #
+    # 1. Prepare args
+    #
+
+    # Allocate SMEM
+    smem = cutlass.utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+    sA = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=a_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=a_smem_layout.inner,
+    )
+    sB = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=b_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=b_smem_layout.inner,
+    )
+
+    # Prefetch tma descriptor
+    if warp_idx == 0:
+        cpasync.prefetch_descriptor(tma_atom_a)
+        cpasync.prefetch_descriptor(tma_atom_b)
+
+    # Pipeline configuration
+    num_tma_copy_bytes = (
+        cute.size_in_bytes(io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2]))
+        + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
+    ) * cute.size(cta_layout_vmnk, mode=[0])
+    num_mcast_ctas_a = cute.size(cta_layout_vmnk.shape[2])
+    num_mcast_ctas_b = cute.size(cta_layout_vmnk.shape[1])
+    num_tma_producer = num_mcast_ctas_a + num_mcast_ctas_b - 1
+    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+        num_stages=ab_stages,
+        producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+        consumer_group=pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_tma_producer
+        ),
+        tx_count=num_tma_copy_bytes,
+        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+        num_stages=acc_stage,
+        producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+        consumer_group=pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            cute.size(cta_layout_vmnk, mode=[0]) * threads_per_cta,
+        ),
+        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Partition tensors for MMA and make fragments
+    # (bM, bK, RestK)
+    gA = cute.local_tile(mA_mkl, mma_tiler_mnk, mma_coord_mnk, proj=(1, None, 1))
+    # (bN, bK, RestK)
+    gB = cute.local_tile(mB_nkl, mma_tiler_mnk, mma_coord_mnk, proj=(None, 1, 1))
+    # (bM, bN)
+    gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
+    thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
+    # (MMA, MMA_M, MMA_K)
+    tCgA = thr_mma.partition_A(gA)
+    # (MMA, MMA_N, MMA_K)
+    tCgB = thr_mma.partition_B(gB)
+    # (MMA, MMA_M, MMA_N)
+    tCgC = thr_mma.partition_C(gC)
+    # (MMA, MMA_M, MMA_K)
+    tCrA = tiled_mma.make_fragment_A(sA)
+    # (MMA, MMA_N, MMA_K)
+    tCrB = tiled_mma.make_fragment_B(sB)
+    # (MMA, MMA_M, MMA_N)
+    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+    # (MMA, MMA_M, MMA_N)
+    tCtAcc = tiled_mma.make_fragment_C(acc_shape)
+
+    # Partition tensors for TMA; This requires the tensors partitioned for MMA
+    tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_a,
+        cta_in_cluster_coord_vmnk[2],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+        cute.group_modes(sA, 0, 3),
+        cute.group_modes(tCgA, 0, 3),
+    )
+    tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_b,
+        cta_in_cluster_coord_vmnk[1],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+        cute.group_modes(sB, 0, 3),
+        cute.group_modes(tCgB, 0, 3),
+    )
+    tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+    )
+    tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
+    )
+
+    # Allocate TMEM and swap the pointer in tCtAcc
+    tmem_alloc_barrier = pipeline.NamedBarrier(
+        barrier_id=1,
+        num_threads=threads_per_cta,
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buf.ptr,
+        barrier_for_retrieve=tmem_alloc_barrier,
+        is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
+        two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+    )
+    num_tmem_cols = 512
+    tmem.allocate(num_tmem_cols)
+
+    # CTA-wide sync before retrieving the pointer to the start of the allocated TMEM
+    # Only warp 0 does the allocation so we need to sync before retrieving the TMEM start address
+    tmem.wait_for_alloc()
+    tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+    # Swap the pointer in tCtAcc
+    tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc.layout)
+
+    subtile_cnt = 4
+    # (EpiTile)
+    epi_tiler = (
+        (cute.size(tCtAcc, mode=[0, 0]), cute.size(tCtAcc, mode=[0, 1]) // subtile_cnt),
+    )
+    # (EpiTile, NumTiles)
+    tCtAcc_epi = cute.zipped_divide(tCtAcc, epi_tiler)
+    # (EpiTile, NumTiles)
+    gC_epi = cute.zipped_divide(tCgC, epi_tiler)
+
+    # Every thread loads 64 x fp32
+    tmem_atom = cute.make_copy_atom(
+        tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
+        cutlass.Float32,
+    )
+    tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_atom, tCtAcc_epi[None, 0])
+    tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
+
+    # (TmemCpy,NumTmemCpy,NumTiles)
+    tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
+    # (TmemCpy,NumTmemCpy,NumTiles)
+    tDgC = tmem_thr_copy.partition_D(gC_epi)
+
+    # (TmemCpy,NumTmemCpy)
+    tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
+    # (TmemCpy,NumTmemCpy)
+    tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
+
+    #
+    # 2. Main loop
+    #
+    is_leader_cta = mma_coord_vmnk[0] == 0
+    num_k_tiles = cute.size(gA, mode=[2])
+    if warp_idx == 0:
+        # Wait for a empty accumulator buffer
+        if is_leader_cta:
+            acc_producer.acquire_and_advance()
+        for _ in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
+            # Issue TMA loads
+            ab_empty = ab_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_a,
+                tAgA[(None, ab_empty.count)],
+                tAsA[(None, ab_empty.index)],
+                tma_bar_ptr=ab_empty.barrier,
+                mcast_mask=tma_mcast_mask_a,
+            )
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, ab_empty.count)],
+                tBsB[(None, ab_empty.index)],
+                tma_bar_ptr=ab_empty.barrier,
+                mcast_mask=tma_mcast_mask_b,
+            )
+
+            # Execute one K-block worth of MMA instructions
+            if is_leader_cta:
+                ab_full = ab_consumer.wait_and_advance()
+                # Execute one K-block worth of MMA instructions
+                num_k_blocks = cute.size(tCrA, mode=[2])
+                for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                    k_block_coord = (None, None, k_block_idx, ab_full.index)
+                    cute.gemm(
+                        tiled_mma,
+                        tCtAcc,
+                        tCrA[k_block_coord],
+                        tCrB[k_block_coord],
+                        tCtAcc,
+                    )
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                ab_full.release()
+
+        # Signal that the accumulator is fully computed
+        if is_leader_cta:
+            acc_producer.commit()
+
+    #
+    # 3. Epilogue
+    #
+
+    #  Release TMEM allocation lock
+    tmem.relinquish_alloc_permit()
+
+    # Wait for the accumulator buffer to be full
+    acc_full = acc_consumer.wait_and_advance()
+    # TMEM -> RMEM -> GEMM
+    # Sub-tiling for better instruction-level parallelism
+    for i in cutlass.range(cute.size(tDtC, mode=[2])):
+        cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
+        tCrC.store(tCrAcc.load().to(io_dtype))
+        cute.autovec_copy(tCrC, tDgC[None, None, i])
+    acc_full.release()
+
+    # Ensure used buffers are properly synchronized before producer exit.
+    # This could avoid the invalid dsmem access due to early leading CTA exit.
+    if warp_idx == 0:
+        ab_producer.tail()
+        if is_leader_cta:
+            acc_producer.tail()
+
+    # Deallocate TMEM
+    pipeline.sync(barrier_id=1)
+    tmem.free(tmem_ptr)
+
+
+@cute.jit
+def host_function(
+    a: cute.Tensor,
+    b: cute.Tensor,
+    c: cute.Tensor,
+):
+    # Construct tiled MMA
+    op = tcgen05.MmaF16BF16Op(
+        io_dtype,
+        acc_dtype,
+        mma_inst_shape_mnk,
+        tcgen05.CtaGroup.TWO,
+        tcgen05.OperandSource.SMEM,
+        tcgen05.OperandMajorMode.K,
+        tcgen05.OperandMajorMode.K,
+    )
+    tiled_mma = cute.make_tiled_mma(op)
+
+    # Construct SMEM layouts for A and B
+    a_smem_layout = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a.element_type,
+        ab_stages,
+    )
+    b_smem_layout = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b.element_type,
+        ab_stages,
+    )
+    a_smem_layout_one_stage = cute.select(a_smem_layout, mode=[0, 1, 2])
+    b_smem_layout_one_stage = cute.select(b_smem_layout, mode=[0, 1, 2])
+
+    # Construct the VMNK layout
+    cta_layout_mnk = cute.make_layout(cluster_shape_mnk)
+    cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
+
+    # Construct TMA load atoms
+    op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
+    a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        op,
+        a,
+        a_smem_layout_one_stage,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,  # take the layout and extract the shape internally
+    )
+    b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        op,
+        b,
+        b_smem_layout_one_stage,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+
+    grid_shape = cute.round_up(
+        cute.ceil_div(
+            (*c.layout.shape, 1), (mma_tiler_mnk[0] // 2, *mma_tiler_mnk[1:])
+        ),
+        cluster_shape_mnk,
+    )
+
+    # Pretty prints kernel attributes useful for debugging
+    # print(f"a               = {cute.pretty_str(a)}")
+    # print(f"b               = {cute.pretty_str(b)}")
+    # print(f"c               = {cute.pretty_str(c)}")
+    # print(f"tiled_mma       = {cute.pretty_str(tiled_mma)}")
+    # print(f"a_smem_layout   = {cute.pretty_str(a_smem_layout)}")
+    # print(f"b_smem_layout   = {cute.pretty_str(b_smem_layout)}")
+    # print(f"cta_layout_mnk  = {cute.pretty_str(cta_layout_mnk)}")
+    # print(f"cta_layout_vmnk = {cute.pretty_str(cta_layout_vmnk)}")
+    # print(f"a_tma_atom   = {cute.pretty_str(a_tma_atom)}")
+    # print(f"b_tma_atom   = {cute.pretty_str(b_tma_atom)}")
+    # print(f"a_tma_tensor    = {cute.pretty_str(a_tma_tensor)}")
+    # print(f"b_tma_tensor    = {cute.pretty_str(b_tma_tensor)}")
+    # cute.printf("grid_shape = {}", grid_shape)
+
+    # Launch the kernel
+    kernel(
+        tiled_mma,
+        a_tma_atom,
+        a_tma_tensor,
+        b_tma_atom,
+        b_tma_tensor,
+        c,
+        a_smem_layout,
+        b_smem_layout,
+        cta_layout_vmnk,
+    ).launch(
+        grid=grid_shape,
+        block=[threads_per_cta, 1, 1],
+        cluster=cluster_shape_mnk,
+    )
+
+
+def run_dense_gemm(
+    mnk: Tuple[int, int, int],
+    tolerance: float,
+):
+    global torch, cutlass_torch
+    import torch
+    import cutlass.torch as cutlass_torch
+
+    print("===================================================================")
+    print("Running Blackwell fp16 GEMM example 1 with:")
+    print(f"  mnk:       {mnk}")
+    print(f"  tolerance: {tolerance}")
+    print("===================================================================")
+    print()
+
+    m, n, k = mnk
+    torch.manual_seed(1111)
+
+    # Make K-major tensors (torch tensors are row-major)
+    def make_tensors(mn, k, dtype):
+        shape = (mn, k)
+        return (
+            torch.empty(*shape, dtype=torch.int32)
+            .random_(-2, 2)
+            .to(device="cuda", dtype=dtype)
+        )
+
+    a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
+    b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
+    c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
+    a_tensor = (
+        from_dlpack(a, assumed_align=32)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=k)
+    )
+    b_tensor = (
+        from_dlpack(b, assumed_align=32)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=k)
+    )
+    c_tensor = (
+        from_dlpack(c, assumed_align=32)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=n)
+    )
+
+    # Entry point to the host JIT function
+    host_function(
+        a_tensor,
+        b_tensor,
+        c_tensor,
+        no_cache=True,
+    )
+
+    # Compute reference result and verify
+    ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu()
+    torch.testing.assert_close(
+        c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
+    )
+
+
+if __name__ == "__main__":
+
+    def parse_comma_separated_ints(s: str):
+        try:
+            return [int(x.strip()) for x in s.split(",")]
+        except ValueError:
+            raise argparse.ArgumentTypeError(
+                "Invalid format. Expected comma-separated integers."
+            )
+
+    from cuda.bindings import driver as cu_driver
+
+    cu_driver.cuInit(0)
+    err, device_count = cu_driver.cuDeviceGetCount()
+    if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
+        raise RuntimeError("A GPU is required to run this example")
+
+    parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 1")
+    parser.add_argument(
+        "--mnk",
+        type=parse_comma_separated_ints,
+        default=[8192, 8192, 8192],
+        help="MNK dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    args = parser.parse_args()
+    if len(args.mnk) != 3:
+        parser.error("--mnk must contain exactly 3 values")
+    if args.mnk[0] % mma_tiler_mnk[0] != 0 or args.mnk[1] % mma_tiler_mnk[1] != 0:
+        parser.error("m n must be divisible by mma_tiler_mn")
+
+    run_dense_gemm(
+        args.mnk,
+        args.tolerance,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/fp16_gemm_2.py

@@ -0,0 +1,687 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 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.
+
+# This is the third tutorial GEMM. It further enhances the second tutorial by adding warp
+# specialization for TMA, MMA, and epilogue warps.
+
+
+import argparse
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.runtime import from_dlpack
+
+"""
+The third tutorial GEMM demonstrates a simple kernel implementation in CuTeDSL.
+
+It further enhances fp16_gemm_1.py by adding warp specialization for TMA, MMA, and epilogue warps. 
+In the epilogue warp, we use TMA store instead of regular copy to store the result from registers to global memory.
+
+This example can achieve better performance than fp16_gemm_1.py due to:
+1. We use warp specialization(WS) to overlap the memory loads and MMA computations.
+
+Core concept of WS is to specialize warps with different tasks (e.g., DMA, MMA, epilogue),
+therefore, different warps in a CTA must communicate with each other.
+
+Warp specialization's benefit comes from task parallelism between warps in the CTA.
+For example, the DMA warps proceed to start loading A/B tensors for the next K-block as soon as they finish loading the current K-block.
+While the MMA warps are computing the result of the current K-block. So the dram latency is hidden.
+The dram latency can also be hidden by prefetch in non-WS version,
+but WS version has better instruction level parallelism as different types of instructions are issued in different warps.
+For example, in non-WS version, tmem allocation and TMA loads are both issued in the same warp, TMA loads only issue after tmem allocation is finished.
+But in WS version, tmem allocation and TMA loads are issued in different warps, tmem allocation can be overlapped with TMA loads.
+
+2. We use TMA store instead of regular copy to store the results from registers to global memory.
+
+To store the results from registers to global memory using TMA actually requires two steps:
+1). Write the tile from registers to shared memory
+2). Write the tile from shared memory to global memory
+Here we continue to use epiolgue subtiles, one reason is that it reduces the shared memory usage in the epilogue,
+and another reason is that it can hide the st.shared latency, that is the st.shared of the next subtile can be overlapped with the TMA store of the current subtile.
+
+For large mma tile size, the mainloop performance between Non-WS and WS version could be similar if there are enough ab_stages to hide the dram latency.
+The performance gain of WS version mainly comes from the prologue and epilogue in this case.
+That means, if k-dimension is small, then the performance of WS version will be obviously better than non-WS version.
+For small mma tile size, we may also see better mainloop performance for WS version.
+This is because there are ALU instructions (preparation work for MMA) for each MMA instruction, and ALU proportion is higher for small mma tile size.
+In Non-WS version, warp 0 will issue the ALU operations for both TMA and MMA instruction, while in WS version, they are issued in different warps,
+so less ALU instructions are issued in MMA warp, and mma instructions can be issued more efficiently.
+
+To run this example:
+.. code-block:: bash
+    python examples/blackwell/tutorial_gemm/fp16_gemm_2.py  \
+      --mnk 8192,8192,8192
+
+Constraints for this example:
+* The problem size of m and n must be divisible by the tile size m & n (256, 256)
+"""
+
+io_dtype = cutlass.Float16
+acc_dtype = cutlass.Float32
+use_2cta_instrs = True
+cluster_shape_mnk = (2, 1, 1) if use_2cta_instrs else (1, 1, 1)
+mma_inst_shape_mnk = (256, 256, 16)
+mma_tiler_mnk = (256, 256, 64)
+threads_in_epilogue = 128  # epilogue threads per cta
+
+# Pipeline stage configuration
+ab_stages = 6
+epi_stages = 2
+acc_stages = 1
+
+
+@cute.struct
+class SharedStorage:
+    ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
+    acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stages * 2]
+    tmem_dealloc_mbar: cutlass.Int64
+    tmem_holding_buffer: cutlass.Int32
+
+
+@cute.kernel()
+def kernel(
+    tiled_mma: cute.TiledMma,
+    tma_atom_a: cute.CopyAtom,
+    mA_mkl: cute.Tensor,
+    tma_atom_b: cute.CopyAtom,
+    mB_nkl: cute.Tensor,
+    tma_atom_c: cute.CopyAtom,
+    mC_mnl: cute.Tensor,
+    a_smem_layout: cute.ComposedLayout,
+    b_smem_layout: cute.ComposedLayout,
+    c_smem_layout_kind: cutlass.Constexpr,
+    epi_smem_layout_staged: cute.ComposedLayout,
+    epi_tile: cute.Tile,
+    cta_layout_vmnk: cute.Layout,
+):
+    warp_idx = cute.arch.warp_idx()
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, bidy, _ = cute.arch.block_idx()
+
+    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+    cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+
+    mma_coord_vmnk = (
+        bidx % cute.size(cta_layout_vmnk, mode=[0]),
+        bidx // cute.size(cta_layout_vmnk, mode=[0]),
+        bidy,
+        None,
+    )
+    mma_coord_mnk = mma_coord_vmnk[1:]
+    is_leader_cta = mma_coord_vmnk[0] == 0
+
+    epilogue_warp_ids = (
+        0,
+        1,
+        2,
+        3,
+    )
+    mma_warp_id = 4
+    tma_warp_id = 5
+
+    #
+    # 1. Prepare args
+    #
+
+    # Allocate SMEM
+    smem = cutlass.utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+    sA = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=a_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=a_smem_layout.inner,
+    )
+    sB = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=b_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=b_smem_layout.inner,
+    )
+    sC = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=epi_smem_layout_staged.outer,
+        byte_alignment=128,
+        swizzle=epi_smem_layout_staged.inner,
+    )
+
+    # Prefetch tma descriptor
+    if warp_idx == tma_warp_id:
+        cpasync.prefetch_descriptor(tma_atom_a)
+        cpasync.prefetch_descriptor(tma_atom_b)
+        cpasync.prefetch_descriptor(tma_atom_c)
+
+    # As many participants as the number of threads issuing the MMA in the same row and column
+    # Substract one to not count twice the same thread
+    num_mcast_participants = (
+        cute.size(cta_layout_vmnk, mode=[1]) + cute.size(cta_layout_vmnk, mode=[2]) - 1
+    )
+
+    # Mcast mask initialization
+    tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+    )
+    tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
+    )
+
+    # Partition tensors for MMA and make fragments
+    # (bM, bK, RestK)
+    gA = cute.local_tile(mA_mkl, mma_tiler_mnk, mma_coord_mnk, proj=(1, None, 1))
+    # (bN, bK, RestK)
+    gB = cute.local_tile(mB_nkl, mma_tiler_mnk, mma_coord_mnk, proj=(None, 1, 1))
+    # (bM, bN)
+    gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
+
+    thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
+    # (MMA, MMA_M, MMA_K, RestK)
+    tCgA = thr_mma.partition_A(gA)
+    # (MMA, MMA_N, MMA_K, RestK)
+    tCgB = thr_mma.partition_B(gB)
+    # (MMA, MMA_M, MMA_N)
+    tCgC = thr_mma.partition_C(gC)
+
+    # (MMA, MMA_M, MMA_K, STAGE)
+    tCrA = tiled_mma.make_fragment_A(sA)
+    # (MMA, MMA_N, MMA_K, STAGE)
+    tCrB = tiled_mma.make_fragment_B(sB)
+
+    # (MMA, MMA_M, MMA_N)
+    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+    # (MMA, MMA_M, MMA_N)
+    tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+
+    # Barrier 1 for epilogue synchronization
+    epilogue_sync_barrier = pipeline.NamedBarrier(
+        barrier_id=1,
+        num_threads=threads_in_epilogue,
+    )
+
+    # Only MMA warp and epilogue warps participate in TMEM allocation synchronization
+    # TMA warp does NOT participate
+    tmem_alloc_barrier = pipeline.NamedBarrier(
+        barrier_id=2,
+        num_threads=32
+        * len((mma_warp_id, *epilogue_warp_ids)),  # 5 warps = 160 threads
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buffer.ptr,
+        barrier_for_retrieve=tmem_alloc_barrier,
+        allocator_warp_id=epilogue_warp_ids[0],
+        is_two_cta=True if use_2cta_instrs else False,
+        two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+    )
+
+    # Partition tensors for TMA; This requires the tensors partitioned for MMA
+    tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_a,
+        cta_in_cluster_coord_vmnk[2],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+        cute.group_modes(sA, 0, 3),
+        cute.group_modes(tCgA, 0, 3),
+    )
+
+    tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_b,
+        cta_in_cluster_coord_vmnk[1],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+        cute.group_modes(sB, 0, 3),
+        cute.group_modes(tCgB, 0, 3),
+    )
+
+    # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N)
+    tCgC_epi = cute.flat_divide(tCgC[((None, None), 0, 0)], epi_tile)
+
+    tCsC, tCgC_tma = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_c,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sC, 0, 2),
+        cute.group_modes(tCgC_epi, 0, 2),
+    )
+
+    num_tma_copy_bytes = (
+        cute.size_in_bytes(io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2]))
+        + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
+    ) * cute.size(cta_layout_vmnk, mode=[0])
+
+    # Threads/warps participating in the mainloop pipeline
+    mainloop_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    mainloop_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread, size=num_mcast_participants
+    )
+
+    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+        num_stages=ab_stages,
+        producer_group=mainloop_pipeline_producer_group,
+        consumer_group=mainloop_pipeline_consumer_group,
+        tx_count=num_tma_copy_bytes,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Threads/warps participating in the accumulator pipeline
+    acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread,
+        size=cute.size(cta_layout_vmnk, mode=[0]) * len(epilogue_warp_ids),
+    )
+
+    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+        num_stages=acc_stages,
+        producer_group=acc_pipeline_producer_group,
+        consumer_group=acc_pipeline_consumer_group,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    #
+    # Main loop
+    #
+
+    num_k_tiles = cute.size(gA, mode=[2])
+
+    # TMA warp
+    if warp_idx == tma_warp_id:
+        for k_tile_idx in range(num_k_tiles):
+            # Wait for A/B buffers to be empty before loading into them
+            handle = ab_producer.acquire_and_advance()
+
+            # Issue TMA loads
+            cute.copy(
+                tma_atom_a,
+                tAgA[(None, k_tile_idx)],
+                tAsA[(None, handle.index)],
+                tma_bar_ptr=handle.barrier,
+                mcast_mask=tma_mcast_mask_a,
+            )
+            cute.copy(
+                tma_atom_b,
+                tBgB[(None, k_tile_idx)],
+                tBsB[(None, handle.index)],
+                tma_bar_ptr=handle.barrier,
+                mcast_mask=tma_mcast_mask_b,
+            )
+
+        # This mbarrier_wait is preventing threadblocks within a set of dependent threadblocks within the cluster
+        # (dependent in the context of the TMA/MMA synchronization pattern) to exit early making
+        # a late tcgen05 commit_arrive illegal
+        ab_producer.tail()
+
+    # MMA warp
+    elif warp_idx == mma_warp_id:
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        # Wait for an empty accumulator buffer
+        if is_leader_cta:
+            acc_empty = acc_producer.acquire_and_advance()
+
+            for k_tile_idx in range(num_k_tiles):
+                # Wait for TMA copies to complete
+                handle = ab_consumer.wait_and_advance()
+
+                # Execute one K-block worth of MMA instructions
+                num_k_blocks = cute.size(tCrA, mode=[2])
+                for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                    k_block_coord = (None, None, k_block_idx, handle.index)
+                    cute.gemm(
+                        tiled_mma,
+                        tCtAcc,
+                        tCrA[k_block_coord],
+                        tCrB[k_block_coord],
+                        tCtAcc,
+                    )
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                # Signal that the A/B buffers have been consumed and are ready for the next load
+                handle.release()
+
+            # Signal that the accumulator is fully computed
+            acc_empty.commit()
+
+    # Epilogue warps
+    elif warp_idx < mma_warp_id:
+        # Allocate TMEM (only epilogue warp 0 actually allocates)
+        num_tmem_cols = 512
+        tmem.allocate(num_tmem_cols)
+
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        # Initialize TMA store pipeline for epilogue
+        epilogue_pipeline_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            size=128,
+        )
+        epilogue_pipeline = pipeline.PipelineTmaStore.create(
+            num_stages=epi_stages,
+            producer_group=epilogue_pipeline_producer_group,
+        )
+
+        # Wait for the accumulator buffer to be full
+        acc_consumer.wait_and_advance()
+
+        copy_atom_t2r = cute.make_copy_atom(
+            tcgen05.Ld16x256bOp(tcgen05.Repetition.x8)
+            if mma_tiler_mnk[0] == 64
+            else tcgen05.Ld32x32bOp(tcgen05.Repetition.x32),
+            cutlass.Float32,
+        )
+
+        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N)
+        tCtAcc_epi = cute.flat_divide(
+            tCtAcc[((None, None), 0, 0)],
+            epi_tile,
+        )
+
+        # Tiled copy for TMEM -> RMEM load
+        tiled_copy_t2r = tcgen05.make_tmem_copy(
+            copy_atom_t2r, tCtAcc_epi[(None, None, 0, 0)]
+        )
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+        tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc_epi)
+        # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+        tTR_gC = thr_copy_t2r.partition_D(tCgC_epi)
+        # (T2R, T2R_M, T2R_N)
+        tTR_rAcc = cute.make_rmem_tensor(
+            tTR_gC[(None, None, None, 0, 0)].shape, cutlass.Float32
+        )
+        tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+
+        # Copy atom and tiled copy for RMEM -> SMEM load
+        copy_atom_r2s = cutlass.utils.blackwell_helpers.get_smem_store_op(
+            c_smem_layout_kind, cutlass.Float32, cutlass.Float32, tiled_copy_t2r
+        )
+        tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+        # (R2S, R2S_M, R2S_N, PIPE_D)
+        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+        tRS_sC = thr_copy_r2s.partition_D(sC)
+
+        tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
+        tRS_rC = cute.make_rmem_tensor(tRS_rAcc.shape, io_dtype)
+
+        tCgC_grouped = cute.group_modes(tCgC_tma, 1, cute.rank(tCgC_tma))
+
+        subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+
+        # Epilogue tiling loop
+        for subtile_idx in cutlass.range(subtile_cnt):
+            # TMEM -> RMEM
+            tTR_tAcc_slice = tTR_tAcc[(None, None, None, subtile_idx)]
+            cute.copy(tiled_copy_t2r, tTR_tAcc_slice, tTR_rAcc)
+
+            # RMEM -> SMEM
+            c_buffer = subtile_idx % epi_stages
+            tRS_sC_slice = tRS_sC[(None, None, None, c_buffer)]
+
+            # type conversion
+            tRS_rC.store(tRS_rAcc.load().to(io_dtype))
+
+            cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC_slice)
+
+            # Memory fence and barrier to ensure shared memory stores are visible to TMA stores
+            cute.arch.fence_view_async_shared()
+            epilogue_sync_barrier.arrive_and_wait()
+            # SMEM -> GMEM
+            if warp_idx == epilogue_warp_ids[0]:
+                cute.copy(
+                    tma_atom_c,
+                    tCsC[(None, c_buffer)],
+                    tCgC_grouped[(None, subtile_idx)],
+                )
+
+                epilogue_pipeline.producer_commit()
+                epilogue_pipeline.producer_acquire()
+            epilogue_sync_barrier.arrive_and_wait()
+
+        epilogue_pipeline.producer_tail()
+
+        # Dealloc the tensor memory buffer
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+
+@cute.jit
+def host_function(
+    a: cute.Tensor,
+    b: cute.Tensor,
+    c: cute.Tensor,
+):
+    #
+    # Construct tiled MMA
+    #
+
+    op = tcgen05.MmaF16BF16Op(
+        io_dtype,
+        acc_dtype,
+        mma_inst_shape_mnk,
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        tcgen05.OperandMajorMode.K,
+        tcgen05.OperandMajorMode.K,
+    )
+    tiled_mma = cute.make_tiled_mma(op)
+
+    #
+    # Construct SMEM layouts for A and B
+    #
+
+    a_smem_layout = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a.element_type,
+        ab_stages,
+    )
+    b_smem_layout = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b.element_type,
+        ab_stages,
+    )
+
+    # c_smem_layout_kind is an enum for row/column major, not a CuTe layout
+    c_smem_layout_kind = utils.LayoutEnum.from_tensor(c)
+
+    #
+    # Construct the VMNK layout
+    #
+
+    cta_layout_mnk = cute.make_layout(cluster_shape_mnk)
+    cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
+
+    #
+    # Construct TMA load atoms
+    #
+
+    op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
+    )
+    a_smem_layout_slice = cute.slice_(a_smem_layout, (None, None, None, 0))
+    a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        op,
+        a,
+        a_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,  # take the layout and extract the shape internally
+
+    )
+    b_smem_layout_slice = cute.slice_(b_smem_layout, (None, None, None, 0))
+    b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        op,
+        b,
+        b_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+
+    cta_tile_shape_mnk = (
+        mma_tiler_mnk[0] // cute.size(tiled_mma.thr_id),
+        mma_tiler_mnk[1],
+        mma_tiler_mnk[2],
+    )
+
+    epi_tile = utils.compute_epilogue_tile_shape(
+        cta_tile_shape_mnk,
+        use_2cta_instrs,
+        c_smem_layout_kind,
+        io_dtype,
+    )
+
+    epi_smem_layout_staged = cutlass.utils.blackwell_helpers.make_smem_layout_epi(
+        io_dtype,
+        c_smem_layout_kind,
+        epi_tile,
+        epi_stages,
+    )
+    epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
+
+    c_tma_atom, c_tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom(
+        cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(),
+        c,
+        epi_smem_layout,
+        epi_tile,
+    )
+
+    #
+    # Launch the kernel
+    #
+
+    grid_shape = cute.round_up(
+        (
+            cute.ceil_div(
+                c.layout.shape[0], mma_tiler_mnk[0] // (2 if use_2cta_instrs else 1)
+            ),
+            cute.ceil_div(c.layout.shape[1], mma_tiler_mnk[1]),
+            1,
+        ),
+        cluster_shape_mnk,
+    )
+
+    kernel(
+        tiled_mma,
+        a_tma_atom,
+        a_tma_tensor,
+        b_tma_atom,
+        b_tma_tensor,
+        c_tma_atom,
+        c_tma_tensor,
+        a_smem_layout,
+        b_smem_layout,
+        c_smem_layout_kind,
+        epi_smem_layout_staged,
+        epi_tile,
+        cta_layout_vmnk,
+    ).launch(
+        grid=grid_shape,
+        block=[192, 1, 1],
+        cluster=cluster_shape_mnk,
+    )
+
+
+def run_dense_gemm(
+    mnk: Tuple[int, int, int],
+    tolerance: float,
+):
+    global torch, cutlass_torch
+    import torch
+    import cutlass.torch as cutlass_torch
+
+    print("===================================================================")
+    print("Running Blackwell fp16 GEMM example 2 with:")
+    print(f"  mnk:       {mnk}")
+    print(f"  tolerance: {tolerance}")
+    print("===================================================================")
+    print()
+
+    m, n, k = mnk
+    torch.manual_seed(1111)
+
+    # Make K-major tensors (torch tensors are row-major)
+    def make_tensors(mn, k, dtype):
+        shape = (mn, k)
+        return (
+            torch.empty(*shape, dtype=torch.int32)
+            .random_(-2, 2)
+            .to(device="cuda", dtype=dtype)
+        )
+
+    a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
+    b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
+    c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
+    a_memref = from_dlpack(a).mark_layout_dynamic()
+    b_memref = from_dlpack(b).mark_layout_dynamic()
+    c_memref = from_dlpack(c).mark_layout_dynamic()
+
+    # Entry point to the host JIT function
+    host_function(
+        a_memref,
+        b_memref,
+        c_memref,
+        no_cache=True,
+    )
+
+    # Compute reference result and verify
+    ref = (torch.einsum("mk,nk->mn", a, b)).cpu()
+    torch.testing.assert_close(
+        c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
+    )
+
+
+if __name__ == "__main__":
+
+    def parse_comma_separated_ints(s: str) -> list[int]:
+        try:
+            return [int(x.strip()) for x in s.split(",")]
+        except ValueError:
+            raise argparse.ArgumentTypeError(
+                "Invalid format. Expected comma-separated integers."
+            )
+
+    from cuda.bindings import driver as cu_driver
+
+    cu_driver.cuInit(0)
+    err, device_count = cu_driver.cuDeviceGetCount()
+    if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
+        raise RuntimeError("A GPU is required to run this example")
+
+    parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 2")
+    parser.add_argument(
+        "--mnk",
+        type=parse_comma_separated_ints,
+        default=(8192, 8192, 8192),
+        help="MNK dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    args = parser.parse_args()
+    if len(args.mnk) != 3:
+        parser.error("--mnk must contain exactly 3 values")
+
+    run_dense_gemm(
+        args.mnk,
+        args.tolerance,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/fp16_gemm_3.py

@@ -0,0 +1,779 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 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.
+
+# This is the third tutorial GEMM. It further enhances the second tutorial by adding warp
+# specialization for TMA, MMA, and epilogue warps.
+
+
+import argparse
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.runtime import from_dlpack
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+"""
+The third tutorial GEMM demonstrates a simple kernel implementation in CuTeDSL.
+
+Compared to fp16_gemm_2.py, this kernel uses a static persistent tile scheduler (StaticPersistentTileScheduler).
+The static scheduler simplifies work distribution by assigning tiles to CTAs in a fixed, deterministic order,
+suitable for well-partitioned workloads. With static scheduling,
+the persistent clusters can stay on the GPU throughout kernel execution and process multiple tiles, hiding prologue and epilogue costs.
+Notes that the static scheduler is susceptible to workload imbalance if the resources of some SMs are unavailable,
+which is why we add a dynamic scheduler in the next example (fp16_gemm_3_1.py).
+
+Therefore, for larger problem sizes the performance will be more advantageous, since a larger problem size leads to more tiles,
+and the prologue/epilogue can be hidden among different tiles.
+This is especially true when the main loop is relatively short and the prologue/epilogue accounts for a large proportion of the work,
+in which case the performance gains become even more significant.
+
+To run this example:
+.. code-block:: bash
+    python examples/blackwell/tutorial_gemm/fp16_gemm_3.py  \
+      --mnk 8192,8192,8192
+
+Constraints for this example:
+* The problem size of m and n must be divisible by the tile size m & n (256, 256)
+"""
+
+io_dtype = cutlass.Float16
+acc_dtype = cutlass.Float32
+use_2cta_instrs = True
+cluster_shape_mnk = (2, 1, 1) if use_2cta_instrs else (1, 1, 1)
+mma_inst_shape_mnk = (256, 256, 16)
+mma_tiler_mnk = (256, 256, 64)
+threads_in_epilogue = 128  # epilogue threads per cta
+
+# Pipeline stage configuration
+ab_stages = 6
+epi_stages = 2
+acc_stages = 2
+
+# Scheduler
+scheduler_type = utils.StaticPersistentTileScheduler
+
+
+@cute.struct
+class SharedStorage:
+    ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
+    acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stages * 2]
+    tmem_dealloc_mbar: cutlass.Int64
+    tmem_holding_buffer: cutlass.Int32
+
+
+@cute.kernel()
+def kernel(
+    tiled_mma: cute.TiledMma,
+    tma_atom_a: cute.CopyAtom,
+    mA_mkl: cute.Tensor,
+    tma_atom_b: cute.CopyAtom,
+    mB_nkl: cute.Tensor,
+    tma_atom_c: cute.CopyAtom,
+    mC_mnl: cute.Tensor,
+    a_smem_layout: cute.ComposedLayout,
+    b_smem_layout: cute.ComposedLayout,
+    c_smem_layout_kind: cutlass.Constexpr,
+    epi_smem_layout_staged: cute.ComposedLayout,
+    epi_tile: cute.Tile,
+    cta_layout_vmnk: cute.Layout,
+    tile_sched_params: utils.PersistentTileSchedulerParams,
+):
+    warp_idx = cute.arch.warp_idx()
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+    cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+
+    mma_tile_coord_v = bidx % cute.size(cta_layout_vmnk, mode=[0])
+    is_leader_cta = mma_tile_coord_v == 0
+
+    epilogue_warp_ids = (
+        0,
+        1,
+        2,
+        3,
+    )
+    mma_warp_id = 4
+    tma_warp_id = 5
+
+    epilog_sync_bar_id = 1
+    tmem_alloc_sync_bar_id = 2
+
+    # Prefetch tma descriptor
+    if warp_idx == tma_warp_id:
+        cpasync.prefetch_descriptor(tma_atom_a)
+        cpasync.prefetch_descriptor(tma_atom_b)
+        cpasync.prefetch_descriptor(tma_atom_c)
+
+    # As many participants as the number of threads issuing the MMA in the same row and column
+    # Substract one to not count twice the same thread
+    num_mcast_participants = (
+        cute.size(cta_layout_vmnk, mode=[1]) + cute.size(cta_layout_vmnk, mode=[2]) - 1
+    )
+
+    # Mcast mask initialization
+    tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+    )
+    tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
+    )
+
+    # Allocate SMEM
+    smem = cutlass.utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+
+    # Barrier 1 for epilogue synchronization
+    epilogue_sync_barrier = pipeline.NamedBarrier(
+        barrier_id=epilog_sync_bar_id,
+        num_threads=threads_in_epilogue,
+    )
+
+    # Only MMA warp and epilogue warps participate in TMEM allocation synchronization
+    # TMA warp does NOT participate
+    tmem_alloc_barrier = pipeline.NamedBarrier(
+        barrier_id=tmem_alloc_sync_bar_id,
+        num_threads=32
+        * len((mma_warp_id, *epilogue_warp_ids)),  # 5 warps = 160 threads
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buffer.ptr,
+        barrier_for_retrieve=tmem_alloc_barrier,
+        allocator_warp_id=epilogue_warp_ids[0],
+        is_two_cta=True,
+        two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+    )
+
+    num_tma_copy_bytes = (
+        cute.size_in_bytes(io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2]))
+        + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
+    ) * cute.size(cta_layout_vmnk, mode=[0])
+
+    # Threads/warps participating in the mainloop pipeline
+    mainloop_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    mainloop_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread, size=num_mcast_participants
+    )
+
+    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+        num_stages=ab_stages,
+        producer_group=mainloop_pipeline_producer_group,
+        consumer_group=mainloop_pipeline_consumer_group,
+        tx_count=num_tma_copy_bytes,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Threads/warps participating in the accumulator pipeline
+    acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread,
+        size=cute.size(cta_layout_vmnk, mode=[0]) * len(epilogue_warp_ids),
+    )
+
+    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+        num_stages=acc_stages,
+        producer_group=acc_pipeline_producer_group,
+        consumer_group=acc_pipeline_consumer_group,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Cluster arrive after barrier init
+    pipeline_init_arrive(cluster_shape_mn=cluster_shape_mnk, is_relaxed=True)
+
+    # Allocate SMEM
+    sA = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=a_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=a_smem_layout.inner,
+    )
+    sB = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=b_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=b_smem_layout.inner,
+    )
+    sC = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=epi_smem_layout_staged.outer,
+        byte_alignment=128,
+        swizzle=epi_smem_layout_staged.inner,
+    )
+
+    # Partition tensors for MMA and make fragments
+    # (bM, bK, RestM, RestK)
+    gA = cute.local_tile(
+        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None)
+    )
+    # (bN, bK, RestN, RestK)
+    gB = cute.local_tile(
+        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None)
+    )
+    # (bM, bN, RestM, RestN)
+    gC = cute.local_tile(
+        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None)
+    )
+
+    thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+    # (MMA, MMA_M, MMA_K, RestM, RestK)
+    tCgA = thr_mma.partition_A(gA)
+    # (MMA, MMA_N, MMA_K, RestN, RestK)
+    tCgB = thr_mma.partition_B(gB)
+    # (MMA, MMA_M, MMA_N, RestM, RestN)
+    tCgC = thr_mma.partition_C(gC)
+
+    # (MMA, MMA_M, MMA_K, STAGE)
+    tCrA = tiled_mma.make_fragment_A(sA)
+    # (MMA, MMA_N, MMA_K, STAGE)
+    tCrB = tiled_mma.make_fragment_B(sB)
+
+    # (MMA, MMA_M, MMA_N)
+    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+    # (MMA, MMA_M, MMA_N, STAGE)
+    tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages))
+
+    # Partition tensors for TMA; This requires the tensors partitioned for MMA
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), RestM, RestK)
+    tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_a,
+        cta_in_cluster_coord_vmnk[2],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+        cute.group_modes(sA, 0, 3),
+        cute.group_modes(tCgA, 0, 3),
+    )
+
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), RestN, RestK)
+    tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_b,
+        cta_in_cluster_coord_vmnk[1],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+        cute.group_modes(sB, 0, 3),
+        cute.group_modes(tCgB, 0, 3),
+    )
+
+    gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile)
+
+    tCsC, tCgC_tma = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_c,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sC, 0, 2),
+        cute.group_modes(gC_epi, 0, 2),
+    )
+
+    # Cluster wait before starting work
+    pipeline_init_wait(cluster_shape_mn=cluster_shape_mnk)
+
+    tile_sched = scheduler_type.create(
+        tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+    )
+    work_tile = tile_sched.initial_work_tile_info()
+
+    #
+    # Main loop
+    #
+
+    num_k_tiles = cute.size(gA, mode=[3])
+
+    # TMA warp
+    if warp_idx == tma_warp_id:
+        #
+        # Persistent tile scheduling loop
+        #
+        while work_tile.is_valid_tile:
+            # Get tile coord from tile scheduler
+            cur_tile_coord = work_tile.tile_idx
+            mma_tile_coord_mnl = (
+                cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                cur_tile_coord[1],
+                cur_tile_coord[2],
+            )
+
+            # Slice to per mma tile index
+            # ((atom_v, rest_v), RestK)
+            tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None)]
+
+            # ((atom_v, rest_v), RestK)
+            tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None)]
+
+            # Tma load loop
+            for k_tile_idx in range(num_k_tiles):
+                # Wait for A/B buffers to be empty before loading into them
+                handle = ab_producer.acquire_and_advance()
+
+                # Issue TMA loads
+                cute.copy(
+                    tma_atom_a,
+                    tAgA_slice[(None, k_tile_idx)],
+                    tAsA[(None, handle.index)],
+                    tma_bar_ptr=handle.barrier,
+                    mcast_mask=tma_mcast_mask_a,
+                )
+                cute.copy(
+                    tma_atom_b,
+                    tBgB_slice[(None, k_tile_idx)],
+                    tBsB[(None, handle.index)],
+                    tma_bar_ptr=handle.barrier,
+                    mcast_mask=tma_mcast_mask_b,
+                )
+
+            # Advance to next k_tile
+            tile_sched.advance_to_next_work()
+            work_tile = tile_sched.get_current_work()
+
+        # This mbarrier_wait is preventing threadblocks within a set of dependent threadblocks within the cluster
+        # (dependent in the context of the TMA/MMA synchronization pattern) to exit early making
+        # a late tcgen05 commit_arrive illegal
+        ab_producer.tail()
+
+    # MMA warp
+    elif warp_idx == mma_warp_id:
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        while work_tile.is_valid_tile:
+            if is_leader_cta:
+                # Wait for accumulator buffer empty
+                acc_empty = acc_producer.acquire_and_advance()
+
+                # Set tensor memory buffer for current tile
+                # (MMA, MMA_M, MMA_N)
+                tCtAcc = tCtAcc_base[(None, None, None, acc_empty.index)]
+
+                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                for k_tile_idx in range(num_k_tiles):
+                    # Wait for TMA copies to complete
+                    handle = ab_consumer.wait_and_advance()
+
+                    # Execute one K-block worth of MMA instructions
+                    num_k_blocks = cute.size(tCrA, mode=[2])
+                    for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                        k_block_coord = (None, None, k_block_idx, handle.index)
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc,
+                            tCrA[k_block_coord],
+                            tCrB[k_block_coord],
+                            tCtAcc,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                    # Signal that the A/B buffers have been consumed and are ready for the next load
+                    handle.release()
+
+                # Signal that the accumulator is fully computed
+                acc_empty.commit()
+
+            # Advance to next tile
+            tile_sched.advance_to_next_work()
+            work_tile = tile_sched.get_current_work()
+
+        # Wait for accumulator buffer empty
+        acc_producer.tail()
+
+    # Epilogue warps
+    elif warp_idx < mma_warp_id:
+        # Allocate TMEM (only epilogue warp 0 actually allocates)
+        num_tmem_cols = 512
+        tmem.allocate(num_tmem_cols)
+
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        # Initialize TMA store pipeline for epilogue
+        epilogue_pipeline_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            size=128,
+        )
+        epilogue_pipeline = pipeline.PipelineTmaStore.create(
+            num_stages=epi_stages,
+            producer_group=epilogue_pipeline_producer_group,
+        )
+
+        copy_atom_t2r = cute.make_copy_atom(
+            tcgen05.Ld32x32bOp(tcgen05.Repetition.x32, tcgen05.Pack.NONE),
+            cutlass.Float32,
+        )
+
+        while work_tile.is_valid_tile:
+            # Get tile coord from tile scheduler
+            cur_tile_coord = work_tile.tile_idx
+            mma_tile_coord_mnl = (
+                cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                cur_tile_coord[1],
+                cur_tile_coord[2],
+            )
+
+            # Wait for accumulator buffer full
+            acc_full = acc_consumer.wait_and_advance()
+
+            # Set tensor memory buffer for current tile
+            # (MMA, MMA_M, MMA_N)
+            tCtAcc = tCtAcc_base[(None, None, None, acc_full.index)]
+
+            # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N)
+            tCtAcc_epi = cute.flat_divide(
+                tCtAcc[((None, None), 0, 0)],  # why 0,0 ?
+                epi_tile,
+            )
+
+            mma_tile_coord_mn = cute.slice_(mma_tile_coord_mnl, (None, None, 0))
+            # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN)
+            tCgC_epi = cute.flat_divide(
+                tCgC[((None, None), 0, 0, *mma_tile_coord_mn)], epi_tile
+            )
+
+            tCgC_tma_cur_tile = tCgC_tma[(None, None, None, *mma_tile_coord_mn)]
+
+            # Tiled copy for TMEM -> RMEM load
+            tiled_copy_t2r = tcgen05.make_tmem_copy(
+                copy_atom_t2r, tCtAcc_epi[(None, None, 0, 0)]
+            )
+            thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+            # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+            tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc_epi)
+            # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+            tTR_gC = thr_copy_t2r.partition_D(tCgC_epi)
+            # (T2R, T2R_M, T2R_N)
+            tTR_rAcc = cute.make_rmem_tensor(
+                tTR_gC[(None, None, None, 0, 0)].shape, cutlass.Float32
+            )
+            tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+
+            # Copy atom and tiled copy for RMEM -> SMEM load
+            copy_atom_r2s = cutlass.utils.blackwell_helpers.get_smem_store_op(
+                c_smem_layout_kind, cutlass.Float32, cutlass.Float32, tiled_copy_t2r
+            )
+            tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+            # (R2S, R2S_M, R2S_N, PIPE_D)
+            thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+            tRS_sC = thr_copy_r2s.partition_D(sC)
+
+            tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
+            tRS_rC = cute.make_rmem_tensor(tRS_rAcc.shape, io_dtype)
+            tCgC_grouped = cute.group_modes(
+                tCgC_tma_cur_tile, 1, cute.rank(tCgC_tma_cur_tile)
+            )
+
+            subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+
+            # Epilogue tiling loop
+            for subtile_idx in cutlass.range(subtile_cnt):
+                # TMEM -> RMEM
+                tTR_tAcc_slice = tTR_tAcc[(None, None, None, subtile_idx)]
+                cute.copy(tiled_copy_t2r, tTR_tAcc_slice, tTR_rAcc)
+
+                # RMEM -> SMEM
+                c_buffer = subtile_idx % epi_stages
+                tRS_sC_slice = tRS_sC[(None, None, None, c_buffer)]
+
+                # type conversion
+                tRS_rC.store(tRS_rAcc.load().to(io_dtype))
+
+                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC_slice)
+
+                # Memory fence and barrier to ensure shared memory stores are visible to TMA stores
+                cute.arch.fence_view_async_shared()
+                epilogue_sync_barrier.arrive_and_wait()
+                # SMEM -> GMEM
+                if warp_idx == epilogue_warp_ids[0]:
+                    cute.copy(
+                        tma_atom_c,
+                        tCsC[(None, c_buffer)],
+                        tCgC_grouped[(None, subtile_idx)],
+                    )
+
+                    epilogue_pipeline.producer_commit()
+                    epilogue_pipeline.producer_acquire()
+                epilogue_sync_barrier.arrive_and_wait()
+
+            # Async arrive accumulator buffer empty
+            with cute.arch.elect_one():
+                acc_full.release()
+
+            # Advance to next tile
+            tile_sched.advance_to_next_work()
+            work_tile = tile_sched.get_current_work()
+
+        # Wait for C store complete
+        epilogue_pipeline.producer_tail()
+
+        # Dealloc the tensor memory buffer
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+
+def compute_grid(
+    c: cute.Tensor,
+    mma_tiler_mnk: Tuple[int, int, int],
+    cluster_shape_mnk: Tuple[int, int, int],
+    max_active_clusters: cutlass.Constexpr,
+) -> Tuple[
+    utils.PersistentTileSchedulerParams,
+    Tuple[int, int, int],
+]:
+    c_shape = cute.slice_(mma_tiler_mnk, (None, None, 0))
+    gc = cute.zipped_divide(c, tiler=c_shape)
+    num_ctas_mn = gc[(0, (None, None))].shape
+
+    tile_sched_params = utils.PersistentTileSchedulerParams(
+        (*num_ctas_mn, 1), cluster_shape_mnk
+    )
+    grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+        tile_sched_params, max_active_clusters
+    )
+    return tile_sched_params, grid
+
+
+@cute.jit
+def host_function(
+    a: cute.Tensor,
+    b: cute.Tensor,
+    c: cute.Tensor,
+    max_active_clusters: cutlass.Constexpr,
+):
+    #
+    # Construct tiled MMA
+    #
+
+    op = tcgen05.MmaF16BF16Op(
+        io_dtype,
+        acc_dtype,
+        mma_inst_shape_mnk,
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        tcgen05.OperandMajorMode.K,
+        tcgen05.OperandMajorMode.K,
+    )
+    tiled_mma = cute.make_tiled_mma(op)
+
+    #
+    # Construct SMEM layouts for A and B
+    #
+
+    a_smem_layout = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a.element_type,
+        ab_stages,
+    )
+    b_smem_layout = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b.element_type,
+        ab_stages,
+    )
+
+    # c_smem_layout_kind is an enum for row/column major, not a CuTe layout
+    c_smem_layout_kind = utils.LayoutEnum.from_tensor(c)
+
+    #
+    # Construct the VMNK layout
+    #
+
+    cta_layout_mnk = cute.make_layout(cluster_shape_mnk)
+    cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
+
+    #
+    # Construct TMA load atoms
+    #
+
+    op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
+    )
+    a_smem_layout_slice = cute.slice_(a_smem_layout, (None, None, None, 0))
+    a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        op,
+        a,
+        a_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+    b_smem_layout_slice = cute.slice_(b_smem_layout, (None, None, None, 0))
+    b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        op,
+        b,
+        b_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+
+    cta_tile_shape_mnk = (
+        mma_tiler_mnk[0] // cute.size(tiled_mma.thr_id),
+        mma_tiler_mnk[1],
+        mma_tiler_mnk[2],
+    )
+
+    epi_tile = utils.compute_epilogue_tile_shape(
+        cta_tile_shape_mnk,
+        use_2cta_instrs,
+        c_smem_layout_kind,
+        io_dtype,
+    )
+
+    epi_smem_layout_staged = cutlass.utils.blackwell_helpers.make_smem_layout_epi(
+        io_dtype,
+        c_smem_layout_kind,
+        epi_tile,
+        epi_stages,
+    )
+
+    epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
+    c_tma_atom, c_tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom(
+        cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(),
+        c,
+        epi_smem_layout,
+        epi_tile,
+    )
+
+    #
+    # Launch the kernel
+    #
+
+    tile_sched_params, grid_shape = compute_grid(
+        c,
+        cta_tile_shape_mnk,
+        cluster_shape_mnk,
+        max_active_clusters,
+    )
+
+    kernel(
+        tiled_mma,
+        a_tma_atom,
+        a_tma_tensor,
+        b_tma_atom,
+        b_tma_tensor,
+        c_tma_atom,
+        c_tma_tensor,
+        a_smem_layout,
+        b_smem_layout,
+        c_smem_layout_kind,
+        epi_smem_layout_staged,
+        epi_tile,
+        cta_layout_vmnk,
+        tile_sched_params,
+    ).launch(
+        grid=grid_shape,
+        block=[192, 1, 1],
+        cluster=cluster_shape_mnk,
+    )
+
+
+def run_dense_gemm(
+    mnk: Tuple[int, int, int],
+    tolerance: float,
+):
+    global torch, cutlass_torch
+    import torch
+    import cutlass.torch as cutlass_torch
+
+    print("===================================================================")
+    print("Running Blackwell fp16 GEMM example 3 with:")
+    print(f"  mnk:       {mnk}")
+    print(f"  tolerance: {tolerance}")
+    print("===================================================================")
+    print()
+
+    m, n, k = mnk
+    torch.manual_seed(1111)
+
+    # Make K-major tensors (torch tensors are row-major)
+    def make_tensors(mn, k, dtype):
+        shape = (mn, k)
+        return (
+            # torch.empty(*shape, dtype=torch.int32)
+            # .random_(-2, 2)
+            torch.ones(*shape, dtype=torch.int32)
+            .to(device="cuda", dtype=dtype)
+        )
+
+    a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
+    b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
+    c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
+    a_memref = from_dlpack(a).mark_layout_dynamic()
+    b_memref = from_dlpack(b).mark_layout_dynamic()
+    c_memref = from_dlpack(c).mark_layout_dynamic()
+
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
+        cluster_shape_mnk[0] * cluster_shape_mnk[1]
+    )
+
+    # Entry point to the host JIT function
+    host_function(
+        a_memref,
+        b_memref,
+        c_memref,
+        max_active_clusters,
+        no_cache=True,
+    )
+
+    # Compute reference result and verify
+    ref = (torch.einsum("mk,nk->mn", a, b)).cpu()
+    torch.testing.assert_close(
+        c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
+    )
+
+
+if __name__ == "__main__":
+
+    def parse_comma_separated_ints(s: str):
+        try:
+            return [int(x.strip()) for x in s.split(",")]
+        except ValueError:
+            raise argparse.ArgumentTypeError(
+                "Invalid format. Expected comma-separated integers."
+            )
+
+    from cuda.bindings import driver as cu_driver
+
+    cu_driver.cuInit(0)
+    err, device_count = cu_driver.cuDeviceGetCount()
+    if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
+        raise RuntimeError("A GPU is required to run this example")
+
+    parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 3")
+    parser.add_argument(
+        "--mnk",
+        type=parse_comma_separated_ints,
+        default=(8192, 8192, 8192),
+        help="MNK dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    args = parser.parse_args()
+    if len(args.mnk) != 3:
+        parser.error("--mnk must contain exactly 3 values")
+
+    run_dense_gemm(
+        args.mnk,
+        args.tolerance,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/fp16_gemm_3_1.py

@@ -0,0 +1,891 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 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.
+
+# This is the third tutorial GEMM. It further enhances the second tutorial by adding warp
+# specialization for TMA, MMA, and epilogue warps.
+
+
+import argparse
+from typing import Tuple, Union
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.runtime import from_dlpack
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+"""
+The third tutorial GEMM demonstrates a simple kernel implementation in CuTeDSL.
+
+Compared to fp16_gemm_3.py, this kernel uses a dynamic persistent tile scheduler (ClcDynamicPersistentTileScheduler).
+The dynamic scheduler is more flexible than the static scheduler, as it can handle workload imbalance better.
+
+To run this example:
+.. code-block:: bash
+    python examples/blackwell/tutorial_gemm/fp16_gemm_3_1.py  \
+      --mnk 8192,8192,8192
+
+Constraints for this example:
+* The problem size of m and n must be divisible by the tile size m & n (256, 256)
+"""
+
+io_dtype = cutlass.Float16
+acc_dtype = cutlass.Float32
+use_2cta_instrs = True
+cluster_shape_mnk = (2, 1, 1) if use_2cta_instrs else (1, 1, 1)
+mma_inst_shape_mnk = (256, 256, 16)
+mma_tiler_mnk = (256, 256, 64)
+threads_in_epilogue = 128  # epilogue threads per cta
+
+# Pipeline stage configuration
+ab_stages = 6
+epi_stages = 2
+acc_stages = 2
+num_clc_stage = 1
+
+# Scheduler
+use_clc_dynamic_scheduler = True
+scheduler_type = (
+    utils.ClcDynamicPersistentTileScheduler
+    if use_clc_dynamic_scheduler
+    else utils.StaticPersistentTileScheduler
+)
+# Response size is 4B * 4 elements
+num_clc_response_bytes = 16
+
+
+@cute.struct
+class SharedStorage:
+    ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
+    acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stages * 2]
+    tmem_dealloc_mbar: cutlass.Int64
+    tmem_holding_buffer: cutlass.Int32
+    # Only for CLC Dynamic Scheduler
+    clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+    clc_response: cute.struct.MemRange[cutlass.Int32, 4]
+
+
+@cute.kernel()
+def kernel(
+    tiled_mma: cute.TiledMma,
+    tma_atom_a: cute.CopyAtom,
+    mA_mkl: cute.Tensor,
+    tma_atom_b: cute.CopyAtom,
+    mB_nkl: cute.Tensor,
+    tma_atom_c: cute.CopyAtom,
+    mC_mnl: cute.Tensor,
+    a_smem_layout: cute.ComposedLayout,
+    b_smem_layout: cute.ComposedLayout,
+    c_smem_layout_kind: cutlass.Constexpr,
+    epi_smem_layout_staged: cute.ComposedLayout,
+    epi_tile: cute.Tile,
+    cta_layout_vmnk: cute.Layout,
+    tile_sched_params: Union[
+        utils.ClcDynamicPersistentTileSchedulerParams,
+        utils.PersistentTileSchedulerParams,
+    ],
+):  
+    warp_idx = cute.arch.warp_idx()
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+    cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+
+    mma_tile_coord_v = bidx % cute.size(cta_layout_vmnk, mode=[0])
+    is_leader_cta = mma_tile_coord_v == 0
+
+    epilogue_warp_ids = (
+        0,
+        1,
+        2,
+        3,
+    )
+    mma_warp_id = 4
+    tma_warp_id = 5
+    # sched_warp_id only for dynamic scheduler
+    sched_warp_id = 6
+
+    epilog_sync_bar_id = 1
+    tmem_alloc_sync_bar_id = 2
+
+    # Prefetch tma descriptor
+    if warp_idx == tma_warp_id:
+        cpasync.prefetch_descriptor(tma_atom_a)
+        cpasync.prefetch_descriptor(tma_atom_b)
+        cpasync.prefetch_descriptor(tma_atom_c)
+
+    # As many participants as the number of threads issuing the MMA in the same row and column
+    # Substract one to not count twice the same thread
+    num_mcast_participants = (
+        cute.size(cta_layout_vmnk, mode=[1]) + cute.size(cta_layout_vmnk, mode=[2]) - 1
+    )
+
+    # Mcast mask initialization
+    tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+    )
+    tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
+    )
+
+    # Allocate SMEM
+    smem = cutlass.utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+
+    # Barrier 1 for epilogue synchronization
+    epilogue_sync_barrier = pipeline.NamedBarrier(
+        barrier_id=epilog_sync_bar_id,
+        num_threads=threads_in_epilogue,
+    )
+
+    # Only MMA warp and epilogue warps participate in TMEM allocation synchronization
+    # TMA warp does NOT participate
+    tmem_alloc_barrier = pipeline.NamedBarrier(
+        barrier_id=tmem_alloc_sync_bar_id,
+        num_threads=32
+        * len((mma_warp_id, *epilogue_warp_ids)),  # 5 warps = 160 threads
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buffer.ptr,
+        barrier_for_retrieve=tmem_alloc_barrier,
+        allocator_warp_id=epilogue_warp_ids[0],
+        is_two_cta=True,
+        two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+    )
+
+    num_tma_copy_bytes = (
+        cute.size_in_bytes(io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2]))
+        + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
+    ) * cute.size(cta_layout_vmnk, mode=[0])
+
+    # Threads/warps participating in the mainloop pipeline
+    mainloop_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    mainloop_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread, size=num_mcast_participants
+    )
+
+    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+        num_stages=ab_stages,
+        producer_group=mainloop_pipeline_producer_group,
+        consumer_group=mainloop_pipeline_consumer_group,
+        tx_count=num_tma_copy_bytes,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Threads/warps participating in the accumulator pipeline
+    acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread,
+        size=cute.size(cta_layout_vmnk, mode=[0]) * len(epilogue_warp_ids),
+    )
+
+    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+        num_stages=acc_stages,
+        producer_group=acc_pipeline_producer_group,
+        consumer_group=acc_pipeline_consumer_group,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Initialize clc_pipeline (barrier) and states
+    # ONLY for CLC Dynamic Scheduler
+    if cutlass.const_expr(use_clc_dynamic_scheduler):
+        clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        cluster_size = cute.size(cluster_shape_mnk)
+        num_clc_consumer_threads = 32 * len(
+            (
+                sched_warp_id,
+                *(
+                    cluster_size
+                    * (
+                        mma_warp_id,
+                        tma_warp_id,
+                        *epilogue_warp_ids,
+                    )
+                ),
+            )
+        )
+        clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_clc_consumer_threads
+        )
+        clc_pipeline = pipeline.PipelineClcFetchAsync.create(
+            barrier_storage=storage.clc_mbar_ptr.data_ptr(),
+            num_stages=num_clc_stage,
+            producer_group=clc_pipeline_producer_group,
+            consumer_group=clc_pipeline_consumer_group,
+            tx_count=num_clc_response_bytes,
+            cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
+        )
+        # Initial clc response pointer
+        clc_response_ptr = storage.clc_response.data_ptr()
+
+        clc_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, num_clc_stage
+        )
+    else:
+        clc_pipeline = None
+        clc_response_ptr = None
+        clc_consumer_state = None
+
+    # Cluster arrive after barrier init
+    pipeline_init_arrive(cluster_shape_mn=cluster_shape_mnk, is_relaxed=True)
+
+    # Allocate SMEM
+    sA = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=a_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=a_smem_layout.inner,
+    )
+    sB = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=b_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=b_smem_layout.inner,
+    )
+    sC = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=epi_smem_layout_staged.outer,
+        byte_alignment=128,
+        swizzle=epi_smem_layout_staged.inner,
+    )
+
+    # Partition tensors for MMA and make fragments
+    # (bM, bK, RestM, RestK)
+    gA = cute.local_tile(
+        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None)
+    )
+    # (bN, bK, RestN, RestK)
+    gB = cute.local_tile(
+        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None)
+    )
+    # (bM, bN, RestM, RestN)
+    gC = cute.local_tile(
+        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None)
+    )
+
+    thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+    # (MMA, MMA_M, MMA_K, RestM, RestK)
+    tCgA = thr_mma.partition_A(gA)
+    # (MMA, MMA_N, MMA_K, RestN, RestK)
+    tCgB = thr_mma.partition_B(gB)
+    # (MMA, MMA_M, MMA_N, RestM, RestN)
+    tCgC = thr_mma.partition_C(gC)
+
+    # (MMA, MMA_M, MMA_K, STAGE)
+    tCrA = tiled_mma.make_fragment_A(sA)
+    # (MMA, MMA_N, MMA_K, STAGE)
+    tCrB = tiled_mma.make_fragment_B(sB)
+
+    # (MMA, MMA_M, MMA_N)
+    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+    # (MMA, MMA_M, MMA_N, STAGE)
+    tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages))
+
+    # Partition tensors for TMA; This requires the tensors partitioned for MMA
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), RestM, RestK)
+    tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_a,
+        cta_in_cluster_coord_vmnk[2],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+        cute.group_modes(sA, 0, 3),
+        cute.group_modes(tCgA, 0, 3),
+    )
+
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), RestN, RestK)
+    tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_b,
+        cta_in_cluster_coord_vmnk[1],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+        cute.group_modes(sB, 0, 3),
+        cute.group_modes(tCgB, 0, 3),
+    )
+
+    gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile)
+
+    tCsC, tCgC_tma = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_c,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sC, 0, 2),
+        cute.group_modes(gC_epi, 0, 2),
+    )
+
+    # Cluster wait before starting work
+    pipeline_init_wait(cluster_shape_mn=cluster_shape_mnk)
+
+    # Construct the scheduler
+    if cutlass.const_expr(use_clc_dynamic_scheduler):
+        tile_sched = scheduler_type.create(
+            tile_sched_params,
+            cute.arch.block_idx(),
+            cute.arch.grid_dim(),
+            clc_response_ptr,
+        )
+    else:
+        tile_sched = scheduler_type.create(
+            tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+        )
+    work_tile = tile_sched.initial_work_tile_info()
+
+    #
+    # Main loop
+    #
+
+    num_k_tiles = cute.size(gA, mode=[3])
+
+    # TMA warp
+    if warp_idx == tma_warp_id:
+        #
+        # Persistent tile scheduling loop
+        #
+        while work_tile.is_valid_tile:
+            # Get tile coord from tile scheduler
+            cur_tile_coord = work_tile.tile_idx
+            mma_tile_coord_mnl = (
+                cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                cur_tile_coord[1],
+                cur_tile_coord[2],
+            )
+
+            # Slice to per mma tile index
+            # ((atom_v, rest_v), RestK)
+            tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None)]
+
+            # ((atom_v, rest_v), RestK)
+            tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None)]
+
+            # Tma load loop
+            for k_tile_idx in range(num_k_tiles):
+                # Wait for A/B buffers to be empty before loading into them
+                handle = ab_producer.acquire_and_advance()
+
+                # Issue TMA loads
+                cute.copy(
+                    tma_atom_a,
+                    tAgA_slice[(None, k_tile_idx)],
+                    tAsA[(None, handle.index)],
+                    tma_bar_ptr=handle.barrier,
+                    mcast_mask=tma_mcast_mask_a,
+                )
+                cute.copy(
+                    tma_atom_b,
+                    tBgB_slice[(None, k_tile_idx)],
+                    tBsB[(None, handle.index)],
+                    tma_bar_ptr=handle.barrier,
+                    mcast_mask=tma_mcast_mask_b,
+                )
+
+            # Advance to next k_tile
+            if cutlass.const_expr(use_clc_dynamic_scheduler):
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            else:
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+        # This mbarrier_wait is preventing threadblocks within a set of dependent threadblocks within the cluster
+        # (dependent in the context of the TMA/MMA synchronization pattern) to exit early making
+        # a late tcgen05 commit_arrive illegal
+        ab_producer.tail()
+
+    # Sched warp (only for dynamic scheduler)
+    if cutlass.const_expr(use_clc_dynamic_scheduler):
+        is_first_cta_in_cluster = cta_rank_in_cluster == 0
+
+        if warp_idx == sched_warp_id and is_first_cta_in_cluster:
+            # Persistent tile scheduling loop
+            clc_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.ProducerConsumer, num_clc_stage
+            )
+            while work_tile.is_valid_tile:
+                # Advance to next tile
+                clc_pipeline.producer_acquire(clc_producer_state)
+                mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state)
+                tile_sched.advance_to_next_work(mbarrier_addr)
+                clc_producer_state.advance()
+
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            clc_pipeline.producer_tail(clc_producer_state)
+
+    # MMA warp
+    if warp_idx == mma_warp_id:
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        while work_tile.is_valid_tile:
+            if is_leader_cta:
+                # Wait for accumulator buffer empty
+                acc_empty = acc_producer.acquire_and_advance()
+
+                # Set tensor memory buffer for current tile
+                # (MMA, MMA_M, MMA_N)
+                tCtAcc = tCtAcc_base[(None, None, None, acc_empty.index)]
+
+                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+                for k_tile_idx in range(num_k_tiles):
+                    # Wait for TMA copies to complete
+                    handle = ab_consumer.wait_and_advance()
+
+                    # Execute one K-block worth of MMA instructions
+                    num_k_blocks = cute.size(tCrA, mode=[2])
+                    for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                        k_block_coord = (None, None, k_block_idx, handle.index)
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc,
+                            tCrA[k_block_coord],
+                            tCrB[k_block_coord],
+                            tCtAcc,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                    # Signal that the A/B buffers have been consumed and are ready for the next load
+                    handle.release()
+
+                # Signal that the accumulator is fully computed
+                acc_empty.commit()
+
+            # Advance to next tile
+            if cutlass.const_expr(use_clc_dynamic_scheduler):
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            else:
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+        # Wait for accumulator buffer empty
+        acc_producer.tail()
+
+    # Epilogue warps
+    if warp_idx < mma_warp_id:
+        # Allocate TMEM (only epilogue warp 0 actually allocates)
+        num_tmem_cols = 512
+        tmem.allocate(num_tmem_cols)
+
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        # Initialize TMA store pipeline for epilogue
+        epilogue_pipeline_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            size=128,
+        )
+        epilogue_pipeline = pipeline.PipelineTmaStore.create(
+            num_stages=epi_stages,
+            producer_group=epilogue_pipeline_producer_group,
+        )
+
+        copy_atom_t2r = cute.make_copy_atom(
+            tcgen05.Ld32x32bOp(tcgen05.Repetition.x32, tcgen05.Pack.NONE),
+            cutlass.Float32,
+        )
+
+        while work_tile.is_valid_tile:
+            # Get tile coord from tile scheduler
+            cur_tile_coord = work_tile.tile_idx
+            mma_tile_coord_mnl = (
+                cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                cur_tile_coord[1],
+                cur_tile_coord[2],
+            )
+
+            # Wait for accumulator buffer full
+            acc_full = acc_consumer.wait_and_advance()
+
+            # Set tensor memory buffer for current tile
+            # (MMA, MMA_M, MMA_N)
+            tCtAcc = tCtAcc_base[(None, None, None, acc_full.index)]
+
+            # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N)
+            tCtAcc_epi = cute.flat_divide(
+                tCtAcc[((None, None), 0, 0)],  # why 0,0 ?
+                epi_tile,
+            )
+
+            mma_tile_coord_mn = cute.slice_(mma_tile_coord_mnl, (None, None, 0))
+            # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN)
+            tCgC_epi = cute.flat_divide(
+                tCgC[((None, None), 0, 0, *mma_tile_coord_mn)], epi_tile
+            )
+
+            tCgC_tma_cur_tile = tCgC_tma[(None, None, None, *mma_tile_coord_mn)]
+
+            # Tiled copy for TMEM -> RMEM load
+            tiled_copy_t2r = tcgen05.make_tmem_copy(
+                copy_atom_t2r, tCtAcc_epi[(None, None, 0, 0)]
+            )
+            thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+            # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+            tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc_epi)
+            # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+            tTR_gC = thr_copy_t2r.partition_D(tCgC_epi)
+            # (T2R, T2R_M, T2R_N)
+            tTR_rAcc = cute.make_rmem_tensor(
+                tTR_gC[(None, None, None, 0, 0)].shape, cutlass.Float32
+            )
+            tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+
+            # Copy atom and tiled copy for RMEM -> SMEM load
+            copy_atom_r2s = cutlass.utils.blackwell_helpers.get_smem_store_op(
+                c_smem_layout_kind, cutlass.Float32, cutlass.Float32, tiled_copy_t2r
+            )
+            tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+            # (R2S, R2S_M, R2S_N, PIPE_D)
+            thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+            tRS_sC = thr_copy_r2s.partition_D(sC)
+
+            tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
+            tRS_rC = cute.make_rmem_tensor(tRS_rAcc.shape, io_dtype)
+            tCgC_grouped = cute.group_modes(
+                tCgC_tma_cur_tile, 1, cute.rank(tCgC_tma_cur_tile)
+            )
+
+            subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+
+            # Epilogue tiling loop
+            for subtile_idx in cutlass.range(subtile_cnt):
+                # TMEM -> RMEM
+                tTR_tAcc_slice = tTR_tAcc[(None, None, None, subtile_idx)]
+                cute.copy(tiled_copy_t2r, tTR_tAcc_slice, tTR_rAcc)
+
+                # RMEM -> SMEM
+                c_buffer = subtile_idx % epi_stages
+                tRS_sC_slice = tRS_sC[(None, None, None, c_buffer)]
+
+                # type conversion
+                tRS_rC.store(tRS_rAcc.load().to(io_dtype))
+
+                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC_slice)
+
+                # Memory fence and barrier to ensure shared memory stores are visible to TMA stores
+                cute.arch.fence_view_async_shared()
+                epilogue_sync_barrier.arrive_and_wait()
+                # SMEM -> GMEM
+                if warp_idx == epilogue_warp_ids[0]:
+                    cute.copy(
+                        tma_atom_c,
+                        tCsC[(None, c_buffer)],
+                        tCgC_grouped[(None, subtile_idx)],
+                    )
+
+                    epilogue_pipeline.producer_commit()
+                    epilogue_pipeline.producer_acquire()
+                epilogue_sync_barrier.arrive_and_wait()
+
+            # Async arrive accumulator buffer empty
+            with cute.arch.elect_one():
+                acc_full.release()
+
+            # Advance to next tile
+            if cutlass.const_expr(use_clc_dynamic_scheduler):
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            else:
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+        # Wait for C store complete
+        epilogue_pipeline.producer_tail()
+
+        # Dealloc the tensor memory buffer
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+
+def compute_grid(
+    c: cute.Tensor,
+    mma_tiler_mnk: Tuple[int, int, int],
+    cluster_shape_mnk: Tuple[int, int, int],
+    scheduler_type: Union[
+        utils.StaticPersistentTileScheduler, utils.ClcDynamicPersistentTileScheduler
+    ],
+    max_active_clusters: cutlass.Constexpr,
+) -> Tuple[
+    Union[
+        utils.ClcDynamicPersistentTileSchedulerParams,
+        utils.PersistentTileSchedulerParams,
+    ],
+    Tuple[int, int, int],
+]:
+    c_shape = cute.slice_(mma_tiler_mnk, (None, None, 0))
+    gc = cute.zipped_divide(c, tiler=c_shape)
+    num_ctas_mn = gc[(0, (None, None))].shape
+
+    if cutlass.const_expr(
+        issubclass(scheduler_type, utils.ClcDynamicPersistentTileScheduler)
+    ):
+        tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams(
+            (*num_ctas_mn, 1), cluster_shape_mnk
+        )
+        grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params)
+    else:
+        tile_sched_params = utils.PersistentTileSchedulerParams(
+            (*num_ctas_mn, 1), cluster_shape_mnk
+        )
+        grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+            tile_sched_params, max_active_clusters
+        )
+    return tile_sched_params, grid
+
+
+@cute.jit
+def host_function(
+    a: cute.Tensor,
+    b: cute.Tensor,
+    c: cute.Tensor,
+    max_active_clusters: cutlass.Constexpr,
+):
+    #
+    # Construct tiled MMA
+    #
+
+    op = tcgen05.MmaF16BF16Op(
+        io_dtype,
+        acc_dtype,
+        mma_inst_shape_mnk,
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        tcgen05.OperandMajorMode.K,
+        tcgen05.OperandMajorMode.K,
+    )
+    tiled_mma = cute.make_tiled_mma(op)
+
+    #
+    # Construct SMEM layouts for A and B
+    #
+
+    a_smem_layout = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a.element_type,
+        ab_stages,
+    )
+    b_smem_layout = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b.element_type,
+        ab_stages,
+    )
+
+    # c_smem_layout_kind is an enum for row/column major, not a CuTe layout
+    c_smem_layout_kind = utils.LayoutEnum.from_tensor(c)
+
+    #
+    # Construct the VMNK layout
+    #
+
+    cta_layout_mnk = cute.make_layout(cluster_shape_mnk)
+    cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
+
+    #
+    # Construct TMA load atoms
+    #
+
+    op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
+    )
+    a_smem_layout_slice = cute.slice_(a_smem_layout, (None, None, None, 0))
+    a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        op,
+        a,
+        a_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+
+    )
+    b_smem_layout_slice = cute.slice_(b_smem_layout, (None, None, None, 0))
+    b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        op,
+        b,
+        b_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+
+    cta_tile_shape_mnk = (
+        mma_tiler_mnk[0] // cute.size(tiled_mma.thr_id),
+        mma_tiler_mnk[1],
+        mma_tiler_mnk[2],
+    )
+
+    epi_tile = utils.compute_epilogue_tile_shape(
+        cta_tile_shape_mnk,
+        use_2cta_instrs,
+        c_smem_layout_kind,
+        io_dtype,
+    )
+
+    epi_smem_layout_staged = cutlass.utils.blackwell_helpers.make_smem_layout_epi(
+        io_dtype,
+        c_smem_layout_kind,
+        epi_tile,
+        epi_stages,
+    )
+
+    epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
+    c_tma_atom, c_tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom(
+        cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(),
+        c,
+        epi_smem_layout,
+        epi_tile,
+    )
+
+    #
+    # Launch the kernel
+    #
+
+    tile_sched_params, grid_shape = compute_grid(
+        c,
+        cta_tile_shape_mnk,
+        cluster_shape_mnk,
+        scheduler_type,
+        max_active_clusters,
+    )
+
+    kernel(
+        tiled_mma,
+        a_tma_atom,
+        a_tma_tensor,
+        b_tma_atom,
+        b_tma_tensor,
+        c_tma_atom,
+        c_tma_tensor,
+        a_smem_layout,
+        b_smem_layout,
+        c_smem_layout_kind,
+        epi_smem_layout_staged,
+        epi_tile,
+        cta_layout_vmnk,
+        tile_sched_params,
+    ).launch(
+        grid=grid_shape,
+        block=[224, 1, 1] if use_clc_dynamic_scheduler else [192, 1, 1],
+        cluster=cluster_shape_mnk,
+    )
+
+
+def run_dense_gemm(
+    mnk: Tuple[int, int, int],
+    tolerance: float,
+):
+    global torch, cutlass_torch
+    import torch
+    import cutlass.torch as cutlass_torch
+
+    print("===================================================================")
+    print("Running Blackwell fp16 GEMM example 3_1 with:")
+    print(f"  mnk:       {mnk}")
+    print(f"  tolerance: {tolerance}")
+    print("===================================================================")
+    print()
+
+    m, n, k = mnk
+    torch.manual_seed(1111)
+
+    # Make K-major tensors (torch tensors are row-major)
+    def make_tensors(mn, k, dtype):
+        shape = (mn, k)
+        return (
+            torch.empty(*shape, dtype=torch.int32)
+            .random_(-2, 2)
+            .to(device="cuda", dtype=dtype)
+        )
+
+    a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
+    b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
+    c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
+    a_memref = from_dlpack(a).mark_layout_dynamic()
+    b_memref = from_dlpack(b).mark_layout_dynamic()
+    c_memref = from_dlpack(c).mark_layout_dynamic()
+
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
+        cluster_shape_mnk[0] * cluster_shape_mnk[1]
+    )
+
+    # Entry point to the host JIT function
+    host_function(
+        a_memref,
+        b_memref,
+        c_memref,
+        max_active_clusters,
+        no_cache=True,
+    )
+
+    # Compute reference result and verify
+    ref = (torch.einsum("mk,nk->mn", a, b)).cpu()
+    torch.testing.assert_close(
+        c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
+    )
+
+
+if __name__ == "__main__":
+
+    def parse_comma_separated_ints(s: str):
+        try:
+            return [int(x.strip()) for x in s.split(",")]
+        except ValueError:
+            raise argparse.ArgumentTypeError(
+                "Invalid format. Expected comma-separated integers."
+            )
+
+    from cuda.bindings import driver as cu_driver
+
+    cu_driver.cuInit(0)
+    err, device_count = cu_driver.cuDeviceGetCount()
+    if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
+        raise RuntimeError("A GPU is required to run this example")
+
+    parser = argparse.ArgumentParser(description="Blackwell fp16 GEMM example 3_1")
+    parser.add_argument(
+        "--mnk",
+        type=parse_comma_separated_ints,
+        default=(8192, 8192, 8192),
+        help="MNK dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    args = parser.parse_args()
+    if len(args.mnk) != 3:
+        parser.error("--mnk must contain exactly 3 values")
+
+    run_dense_gemm(
+        args.mnk,
+        args.tolerance,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/fp16_gemm_5.py

@@ -0,0 +1,929 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 - 2026 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.
+
+# This is the fifth tutorial GEMM (5). It extends fp16_gemm_3_1.py by adding TMA prefetch.
+# TMA prefetch uses cute.prefetch() to bring data into L2 cache before TMA copy needs it,
+# helping to hide DRAM latency for memory-bound workloads.
+
+
+import argparse
+from typing import Tuple, Union
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.runtime import from_dlpack
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+
+"""
+The fifth tutorial GEMM (5) demonstrates TMA prefetch optimization in CuTeDSL.
+
+TMA Prefetch uses cute.prefetch() to bring data from DRAM into L2 cache ahead of time.
+This helps hide DRAM latency, which is particularly beneficial for memory-bound workloads.
+
+TMA Prefetch consists of two phases:
+1. Initial Phase: Before the TMA load loop starts, prefetch the first `prefetch_dist`
+   k-tiles into L2 cache. This primes the cache before any TMA copies begin.
+2. Rolling Phase: During each iteration of the TMA load loop, after issuing TMA copy
+   for the current k-tile, prefetch the k-tile that is `prefetch_dist` ahead.
+
+Key differences from fp16_gemm_3_1.py:
+1. Added cute.prefetch() calls to bring data into L2 cache before TMA copy
+2. Initial prefetch loop before the main TMA load loop
+3. Rolling prefetch inside the TMA load loop to keep L2 primed
+
+To run this example:
+.. code-block:: bash
+    python examples/blackwell/tutorial_gemm/fp16_gemm_5.py  \
+      --mnk 8192,8192,8192
+
+Constraints for this example:
+* The problem size of m and n must be divisible by the tile size m & n (256, 256)
+"""
+
+io_dtype = cutlass.Float16
+acc_dtype = cutlass.Float32
+use_2cta_instrs = True
+cluster_shape_mnk = (2, 2, 1) if use_2cta_instrs else (1, 1, 1)
+mma_inst_shape_mnk = (256, 64, 16)
+mma_tiler_mnk = (256, 64, 64)
+threads_in_epilogue = 128  # epilogue threads per cta
+
+# Pipeline stage configuration
+ab_stages = 10
+epi_stages = 2
+acc_stages = 2
+num_clc_stage = 1
+
+# Scheduler
+use_clc_dynamic_scheduler = True
+scheduler_type = (
+    utils.ClcDynamicPersistentTileScheduler
+    if use_clc_dynamic_scheduler
+    else utils.StaticPersistentTileScheduler
+)
+# Response size is 4B * 4 elements
+num_clc_response_bytes = 16
+
+
+@cute.struct
+class SharedStorage:
+    ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, ab_stages * 2]
+    acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, acc_stages * 2]
+    tmem_dealloc_mbar: cutlass.Int64
+    tmem_holding_buffer: cutlass.Int32
+    # Only for CLC Dynamic Scheduler
+    clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
+    clc_response: cute.struct.MemRange[cutlass.Int32, 4]
+
+
+@cute.kernel()
+def kernel(
+    tiled_mma: cute.TiledMma,
+    tma_atom_a: cute.CopyAtom,
+    mA_mkl: cute.Tensor,
+    tma_atom_b: cute.CopyAtom,
+    mB_nkl: cute.Tensor,
+    tma_atom_c: cute.CopyAtom,
+    mC_mnl: cute.Tensor,
+    a_smem_layout: cute.ComposedLayout,
+    b_smem_layout: cute.ComposedLayout,
+    c_smem_layout_kind: cutlass.Constexpr,
+    epi_smem_layout_staged: cute.ComposedLayout,
+    epi_tile: cute.Tile,
+    cta_layout_vmnk: cute.Layout,
+    tile_sched_params: Union[
+        utils.ClcDynamicPersistentTileSchedulerParams,
+        utils.PersistentTileSchedulerParams,
+    ],
+):
+    warp_idx = cute.arch.warp_idx()
+    warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+    cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+
+    mma_tile_coord_v = bidx % cute.size(cta_layout_vmnk, mode=[0])
+    is_leader_cta = mma_tile_coord_v == 0
+
+    epilogue_warp_ids = (
+        0,
+        1,
+        2,
+        3,
+    )
+    mma_warp_id = 4
+    tma_warp_id = 5
+    # sched_warp_id only for dynamic scheduler
+    sched_warp_id = 6
+
+    epilog_sync_bar_id = 1
+    tmem_alloc_sync_bar_id = 2
+
+    # Prefetch tma descriptor
+    if warp_idx == tma_warp_id:
+        cpasync.prefetch_descriptor(tma_atom_a)
+        cpasync.prefetch_descriptor(tma_atom_b)
+        cpasync.prefetch_descriptor(tma_atom_c)
+
+    # As many participants as the number of threads issuing the MMA in the same row and column
+    # Substract one to not count twice the same thread
+    num_mcast_participants = (
+        cute.size(cta_layout_vmnk, mode=[1]) + cute.size(cta_layout_vmnk, mode=[2]) - 1
+    )
+
+    # Mcast mask initialization
+    tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+    )
+    tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
+        cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
+    )
+
+    # Allocate SMEM
+    smem = cutlass.utils.SmemAllocator()
+    storage = smem.allocate(SharedStorage)
+
+    # Barrier 1 for epilogue synchronization
+    epilogue_sync_barrier = pipeline.NamedBarrier(
+        barrier_id=epilog_sync_bar_id,
+        num_threads=threads_in_epilogue,
+    )
+
+    # Only MMA warp and epilogue warps participate in TMEM allocation synchronization
+    # TMA warp does NOT participate
+    tmem_alloc_barrier = pipeline.NamedBarrier(
+        barrier_id=tmem_alloc_sync_bar_id,
+        num_threads=32
+        * len((mma_warp_id, *epilogue_warp_ids)),  # 5 warps = 160 threads
+    )
+    tmem = utils.TmemAllocator(
+        storage.tmem_holding_buffer.ptr,
+        barrier_for_retrieve=tmem_alloc_barrier,
+        allocator_warp_id=epilogue_warp_ids[0],
+        is_two_cta=True,
+        two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+    )
+
+    num_tma_copy_bytes = (
+        cute.size_in_bytes(io_dtype, cute.select(a_smem_layout, mode=[0, 1, 2]))
+        + cute.size_in_bytes(io_dtype, cute.select(b_smem_layout, mode=[0, 1, 2]))
+    ) * cute.size(cta_layout_vmnk, mode=[0])
+
+    # Threads/warps participating in the mainloop pipeline
+    mainloop_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    mainloop_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread, size=num_mcast_participants
+    )
+
+    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+        num_stages=ab_stages,
+        producer_group=mainloop_pipeline_producer_group,
+        consumer_group=mainloop_pipeline_consumer_group,
+        tx_count=num_tma_copy_bytes,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Threads/warps participating in the accumulator pipeline
+    acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+    acc_pipeline_consumer_group = pipeline.CooperativeGroup(
+        pipeline.Agent.Thread,
+        size=cute.size(cta_layout_vmnk, mode=[0]) * len(epilogue_warp_ids),
+    )
+
+    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+        num_stages=acc_stages,
+        producer_group=acc_pipeline_producer_group,
+        consumer_group=acc_pipeline_consumer_group,
+        cta_layout_vmnk=cta_layout_vmnk,
+    ).make_participants()
+
+    # Initialize clc_pipeline (barrier) and states
+    # ONLY for CLC Dynamic Scheduler
+    if cutlass.const_expr(use_clc_dynamic_scheduler):
+        clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        cluster_size = cute.size(cluster_shape_mnk)
+        num_clc_consumer_threads = 32 * len(
+            (
+                sched_warp_id,
+                *(
+                    cluster_size
+                    * (
+                        mma_warp_id,
+                        tma_warp_id,
+                        *epilogue_warp_ids,
+                    )
+                ),
+            )
+        )
+        clc_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_clc_consumer_threads
+        )
+        clc_pipeline = pipeline.PipelineClcFetchAsync.create(
+            barrier_storage=storage.clc_mbar_ptr.data_ptr(),
+            num_stages=num_clc_stage,
+            producer_group=clc_pipeline_producer_group,
+            consumer_group=clc_pipeline_consumer_group,
+            tx_count=num_clc_response_bytes,
+            cta_layout_vmnk=cta_layout_vmnk,
+            defer_sync=True,
+        )
+        # Initial clc response pointer
+        clc_response_ptr = storage.clc_response.data_ptr()
+
+        clc_consumer_state = pipeline.make_pipeline_state(
+            pipeline.PipelineUserType.Consumer, num_clc_stage
+        )
+    else:
+        clc_pipeline = None
+        clc_response_ptr = None
+        clc_consumer_state = None
+
+    # Cluster arrive after barrier init
+    pipeline_init_arrive(cluster_shape_mn=cluster_shape_mnk, is_relaxed=True)
+
+    # Allocate SMEM
+    sA = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=a_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=a_smem_layout.inner,
+    )
+    sB = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=b_smem_layout.outer,
+        byte_alignment=128,
+        swizzle=b_smem_layout.inner,
+    )
+    sC = smem.allocate_tensor(
+        element_type=io_dtype,
+        layout=epi_smem_layout_staged.outer,
+        byte_alignment=128,
+        swizzle=epi_smem_layout_staged.inner,
+    )
+
+    # Partition tensors for MMA and make fragments
+    # (bM, bK, RestM, RestK)
+    gA = cute.local_tile(
+        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None)
+    )
+    # (bN, bK, RestN, RestK)
+    gB = cute.local_tile(
+        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None)
+    )
+    # (bM, bN, RestM, RestN)
+    gC = cute.local_tile(
+        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None)
+    )
+
+    thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+    # (MMA, MMA_M, MMA_K, RestM, RestK)
+    tCgA = thr_mma.partition_A(gA)
+    # (MMA, MMA_N, MMA_K, RestN, RestK)
+    tCgB = thr_mma.partition_B(gB)
+    # (MMA, MMA_M, MMA_N, RestM, RestN)
+    tCgC = thr_mma.partition_C(gC)
+
+    # (MMA, MMA_M, MMA_K, STAGE)
+    tCrA = tiled_mma.make_fragment_A(sA)
+    # (MMA, MMA_N, MMA_K, STAGE)
+    tCrB = tiled_mma.make_fragment_B(sB)
+
+    # (MMA, MMA_M, MMA_N)
+    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+    # (MMA, MMA_M, MMA_N, STAGE)
+    tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages))
+
+    # Partition tensors for TMA; This requires the tensors partitioned for MMA
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), RestM, RestK)
+    tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_a,
+        cta_in_cluster_coord_vmnk[2],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+        cute.group_modes(sA, 0, 3),
+        cute.group_modes(tCgA, 0, 3),
+    )
+
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), RestN, RestK)
+    tBsB, tBgB = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_b,
+        cta_in_cluster_coord_vmnk[1],
+        cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+        cute.group_modes(sB, 0, 3),
+        cute.group_modes(tCgB, 0, 3),
+    )
+
+    gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile)
+
+    tCsC, tCgC_tma = cute.nvgpu.cpasync.tma_partition(
+        tma_atom_c,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(sC, 0, 2),
+        cute.group_modes(gC_epi, 0, 2),
+    )
+
+    # Cluster wait before starting work
+    pipeline_init_wait(cluster_shape_mn=cluster_shape_mnk)
+
+    # Construct the scheduler
+    if cutlass.const_expr(use_clc_dynamic_scheduler):
+        tile_sched = scheduler_type.create(
+            tile_sched_params,
+            cute.arch.block_idx(),
+            cute.arch.grid_dim(),
+            clc_response_ptr,
+        )
+    else:
+        tile_sched = scheduler_type.create(
+            tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+        )
+    work_tile = tile_sched.initial_work_tile_info()
+
+    #
+    # Main loop
+    #
+
+    num_k_tiles = cute.size(gA, mode=[3])
+
+    # Prefetch distance: how many k-tiles ahead to prefetch into L2 cache
+    # This helps hide DRAM latency by bringing data to L2 before TMA copy needs it
+    prefetch_dist = ab_stages
+
+    # TMA warp with prefetch
+    if warp_idx == tma_warp_id:
+        #
+        # Persistent tile scheduling loop
+        #
+        while work_tile.is_valid_tile:
+            # Get tile coord from tile scheduler
+            cur_tile_coord = work_tile.tile_idx
+            mma_tile_coord_mnl = (
+                cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                cur_tile_coord[1],
+                cur_tile_coord[2],
+            )
+
+            # Slice to per mma tile index
+            # ((atom_v, rest_v), RestK)
+            tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None)]
+
+            # ((atom_v, rest_v), RestK)
+            tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None)]
+
+            # =========================================================
+            # TMA Prefetch - Initial Phase
+            # =========================================================
+            # Prefetch the first `prefetch_dist` k-tiles into L2 cache
+            # This primes the cache before TMA copies start
+            for pf_k_tile in cutlass.range(
+                cutlass.min(prefetch_dist, num_k_tiles), unroll=1
+            ):
+                cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)])
+                cute.prefetch(tma_atom_b, tBgB_slice[(None, pf_k_tile)])
+
+            # =========================================================
+            # TMA Load Loop with Rolling Prefetch
+            # =========================================================
+            for k_tile_idx in range(num_k_tiles):
+                # Wait for A/B buffers to be empty before loading into them
+                handle = ab_producer.acquire_and_advance()
+
+                # Issue TMA loads (use k_tile_idx like fp16_gemm_3_1.py)
+                cute.copy(
+                    tma_atom_a,
+                    tAgA_slice[(None, k_tile_idx)],
+                    tAsA[(None, handle.index)],
+                    tma_bar_ptr=handle.barrier,
+                    mcast_mask=tma_mcast_mask_a,
+                )
+                cute.copy(
+                    tma_atom_b,
+                    tBgB_slice[(None, k_tile_idx)],
+                    tBsB[(None, handle.index)],
+                    tma_bar_ptr=handle.barrier,
+                    mcast_mask=tma_mcast_mask_b,
+                )
+
+                # Rolling prefetch: prefetch future k-tiles into L2 cache
+                # This keeps the L2 primed as we progress through the K dimension
+                if k_tile_idx + prefetch_dist < num_k_tiles:
+                    future_k_tile = k_tile_idx + prefetch_dist
+                    cute.prefetch(tma_atom_a, tAgA_slice[(None, future_k_tile)])
+                    cute.prefetch(tma_atom_b, tBgB_slice[(None, future_k_tile)])
+
+            # Advance to next tile
+            if cutlass.const_expr(use_clc_dynamic_scheduler):
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            else:
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+        # This mbarrier_wait is preventing threadblocks within a set of dependent threadblocks within the cluster
+        # (dependent in the context of the TMA/MMA synchronization pattern) to exit early making
+        # a late tcgen05 commit_arrive illegal
+        ab_producer.tail()
+
+    # Sched warp (only for dynamic scheduler)
+    if cutlass.const_expr(use_clc_dynamic_scheduler):
+        is_first_cta_in_cluster = cta_rank_in_cluster == 0
+
+        if warp_idx == sched_warp_id and is_first_cta_in_cluster:
+            # Persistent tile scheduling loop
+            clc_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.ProducerConsumer, num_clc_stage
+            )
+            while work_tile.is_valid_tile:
+                # Advance to next tile
+                clc_pipeline.producer_acquire(clc_producer_state)
+                mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state)
+                tile_sched.advance_to_next_work(mbarrier_addr)
+                clc_producer_state.advance()
+
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            clc_pipeline.producer_tail(clc_producer_state)
+
+    # MMA warp
+    if warp_idx == mma_warp_id:
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        while work_tile.is_valid_tile:
+            if is_leader_cta:
+                # Wait for accumulator buffer empty
+                acc_empty = acc_producer.acquire_and_advance()
+
+                # Set tensor memory buffer for current tile
+                # (MMA, MMA_M, MMA_N)
+                tCtAcc = tCtAcc_base[(None, None, None, acc_empty.index)]
+
+                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+
+                for k_tile_idx in range(num_k_tiles):
+                    # Wait for TMA copies to complete
+                    handle = ab_consumer.wait_and_advance()
+
+                    # Execute one K-block worth of MMA instructions
+                    num_k_blocks = cute.size(tCrA, mode=[2])
+                    for k_block_idx in cutlass.range_constexpr(num_k_blocks):
+                        k_block_coord = (None, None, k_block_idx, handle.index)
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc,
+                            tCrA[k_block_coord],
+                            tCrB[k_block_coord],
+                            tCtAcc,
+                        )
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                    # Signal that the A/B buffers have been consumed and are ready for the next load
+                    handle.release()
+
+                # Signal that the accumulator is fully computed
+                acc_empty.commit()
+
+            # Advance to next tile
+            if cutlass.const_expr(use_clc_dynamic_scheduler):
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            else:
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+        # Wait for accumulator buffer empty
+        acc_producer.tail()
+
+    # Epilogue warps
+    if warp_idx < mma_warp_id:
+        # Allocate TMEM (only epilogue warp 0 actually allocates)
+        num_tmem_cols = 512
+        tmem.allocate(num_tmem_cols)
+
+        # Wait for TMEM allocation and retrieve pointer
+        tmem.wait_for_alloc()
+        tmem_ptr = tmem.retrieve_ptr(acc_dtype)
+        # (MMA, MMA_M, MMA_N, STAGE)
+        tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
+
+        # Initialize TMA store pipeline for epilogue
+        epilogue_pipeline_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread,
+            size=128,
+        )
+        epilogue_pipeline = pipeline.PipelineTmaStore.create(
+            num_stages=epi_stages,
+            producer_group=epilogue_pipeline_producer_group,
+        )
+
+        copy_atom_t2r = cute.make_copy_atom(
+            tcgen05.Ld32x32bOp(tcgen05.Repetition.x32, tcgen05.Pack.NONE),
+            cutlass.Float32,
+        )
+
+        while work_tile.is_valid_tile:
+            # Get tile coord from tile scheduler
+            cur_tile_coord = work_tile.tile_idx
+            mma_tile_coord_mnl = (
+                cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                cur_tile_coord[1],
+                cur_tile_coord[2],
+            )
+
+            # Wait for accumulator buffer full
+            acc_full = acc_consumer.wait_and_advance()
+
+            # Set tensor memory buffer for current tile
+            # (MMA, MMA_M, MMA_N)
+            tCtAcc = tCtAcc_base[(None, None, None, acc_full.index)]
+
+            # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N)
+            tCtAcc_epi = cute.flat_divide(
+                tCtAcc[((None, None), 0, 0)],  # why 0,0 ?
+                epi_tile,
+            )
+
+            mma_tile_coord_mn = cute.slice_(mma_tile_coord_mnl, (None, None, 0))
+            # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN)
+            tCgC_epi = cute.flat_divide(
+                tCgC[((None, None), 0, 0, *mma_tile_coord_mn)], epi_tile
+            )
+
+            tCgC_tma_cur_tile = tCgC_tma[(None, None, None, *mma_tile_coord_mn)]
+
+            # Tiled copy for TMEM -> RMEM load
+            tiled_copy_t2r = tcgen05.make_tmem_copy(
+                copy_atom_t2r, tCtAcc_epi[(None, None, 0, 0)]
+            )
+            thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+            # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+            tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc_epi)
+            # (T2R, T2R_M, T2R_N, EPI_M, EPI_N)
+            tTR_gC = thr_copy_t2r.partition_D(tCgC_epi)
+            # (T2R, T2R_M, T2R_N)
+            tTR_rAcc = cute.make_rmem_tensor(
+                tTR_gC[(None, None, None, 0, 0)].shape, cutlass.Float32
+            )
+            tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
+
+            # Copy atom and tiled copy for RMEM -> SMEM load
+            copy_atom_r2s = cutlass.utils.blackwell_helpers.get_smem_store_op(
+                c_smem_layout_kind, cutlass.Float32, cutlass.Float32, tiled_copy_t2r
+            )
+            tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
+            # (R2S, R2S_M, R2S_N, PIPE_D)
+            thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
+            tRS_sC = thr_copy_r2s.partition_D(sC)
+
+            tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
+            tRS_rC = cute.make_rmem_tensor(tRS_rAcc.shape, io_dtype)
+            tCgC_grouped = cute.group_modes(
+                tCgC_tma_cur_tile, 1, cute.rank(tCgC_tma_cur_tile)
+            )
+
+            subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
+
+            # Epilogue tiling loop
+            for subtile_idx in cutlass.range(subtile_cnt):
+                # TMEM -> RMEM
+                tTR_tAcc_slice = tTR_tAcc[(None, None, None, subtile_idx)]
+                cute.copy(tiled_copy_t2r, tTR_tAcc_slice, tTR_rAcc)
+
+                # RMEM -> SMEM
+                c_buffer = subtile_idx % epi_stages
+                tRS_sC_slice = tRS_sC[(None, None, None, c_buffer)]
+
+                # type conversion
+                tRS_rC.store(tRS_rAcc.load().to(io_dtype))
+
+                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC_slice)
+
+                # Memory fence and barrier to ensure shared memory stores are visible to TMA stores
+                cute.arch.fence_view_async_shared()
+                epilogue_sync_barrier.arrive_and_wait()
+                # SMEM -> GMEM
+                if warp_idx == epilogue_warp_ids[0]:
+                    cute.copy(
+                        tma_atom_c,
+                        tCsC[(None, c_buffer)],
+                        tCgC_grouped[(None, subtile_idx)],
+                    )
+
+                    epilogue_pipeline.producer_commit()
+                    epilogue_pipeline.producer_acquire()
+                epilogue_sync_barrier.arrive_and_wait()
+
+            # Async arrive accumulator buffer empty
+            with cute.arch.elect_one():
+                acc_full.release()
+
+            # Advance to next tile
+            if cutlass.const_expr(use_clc_dynamic_scheduler):
+                clc_pipeline.consumer_wait(clc_consumer_state)
+                work_tile = tile_sched.get_current_work()
+                clc_pipeline.consumer_release(clc_consumer_state)
+                clc_consumer_state.advance()
+            else:
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+        # Wait for C store complete
+        epilogue_pipeline.producer_tail()
+
+        # Dealloc the tensor memory buffer
+        tmem.relinquish_alloc_permit()
+        tmem.free(tmem_ptr)
+
+
+def compute_grid(
+    c: cute.Tensor,
+    mma_tiler_mnk: Tuple[int, int, int],
+    cluster_shape_mnk: Tuple[int, int, int],
+    scheduler_type: Union[
+        utils.StaticPersistentTileScheduler, utils.ClcDynamicPersistentTileScheduler
+    ],
+    max_active_clusters: cutlass.Constexpr,
+) -> Tuple[
+    Union[
+        utils.ClcDynamicPersistentTileSchedulerParams,
+        utils.PersistentTileSchedulerParams,
+    ],
+    Tuple[int, int, int],
+]:
+    c_shape = cute.slice_(mma_tiler_mnk, (None, None, 0))
+    gc = cute.zipped_divide(c, tiler=c_shape)
+    num_ctas_mn = gc[(0, (None, None))].shape
+
+    if cutlass.const_expr(
+        issubclass(scheduler_type, utils.ClcDynamicPersistentTileScheduler)
+    ):
+        tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams(
+            (*num_ctas_mn, 1), cluster_shape_mnk
+        )
+        grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params)
+    else:
+        tile_sched_params = utils.PersistentTileSchedulerParams(
+            (*num_ctas_mn, 1), cluster_shape_mnk
+        )
+        grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+            tile_sched_params, max_active_clusters
+        )
+    return tile_sched_params, grid
+
+
+@cute.jit
+def host_function(
+    a: cute.Tensor,
+    b: cute.Tensor,
+    c: cute.Tensor,
+    max_active_clusters: cutlass.Constexpr,
+):
+    #
+    # Construct tiled MMA
+    #
+
+    op = tcgen05.MmaF16BF16Op(
+        io_dtype,
+        acc_dtype,
+        mma_inst_shape_mnk,
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE,
+        tcgen05.OperandSource.SMEM,
+        tcgen05.OperandMajorMode.K,
+        tcgen05.OperandMajorMode.K,
+    )
+    tiled_mma = cute.make_tiled_mma(op)
+
+    #
+    # Construct SMEM layouts for A and B
+    #
+
+    a_smem_layout = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a.element_type,
+        ab_stages,
+    )
+    b_smem_layout = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b.element_type,
+        ab_stages,
+    )
+
+    # c_smem_layout_kind is an enum for row/column major, not a CuTe layout
+    c_smem_layout_kind = utils.LayoutEnum.from_tensor(c)
+
+    #
+    # Construct the VMNK layout
+    #
+
+    cta_layout_mnk = cute.make_layout(cluster_shape_mnk)
+    cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
+
+    #
+    # Construct TMA load atoms
+    #
+
+    op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
+        tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
+    )
+    a_smem_layout_slice = cute.slice_(a_smem_layout, (None, None, None, 0))
+    tma_atom_a, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
+        op,
+        a,
+        a_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+    b_smem_layout_slice = cute.slice_(b_smem_layout, (None, None, None, 0))
+    tma_atom_b, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
+        op,
+        b,
+        b_smem_layout_slice,
+        mma_tiler_mnk,
+        tiled_mma,
+        cta_layout_vmnk.shape,
+    )
+
+    cta_tile_shape_mnk = (
+        mma_tiler_mnk[0] // cute.size(tiled_mma.thr_id),
+        mma_tiler_mnk[1],
+        mma_tiler_mnk[2],
+    )
+
+    epi_tile = utils.compute_epilogue_tile_shape(
+        cta_tile_shape_mnk,
+        use_2cta_instrs,
+        c_smem_layout_kind,
+        io_dtype,
+    )
+
+    epi_smem_layout_staged = cutlass.utils.blackwell_helpers.make_smem_layout_epi(
+        io_dtype,
+        c_smem_layout_kind,
+        epi_tile,
+        epi_stages,
+    )
+
+    epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
+    tma_atom_c, c_tma_tensor = cute.nvgpu.cpasync.make_tiled_tma_atom(
+        cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp(),
+        c,
+        epi_smem_layout,
+        epi_tile,
+    )
+
+    #
+    # Launch the kernel
+    #
+
+    tile_sched_params, grid_shape = compute_grid(
+        c,
+        cta_tile_shape_mnk,
+        cluster_shape_mnk,
+        scheduler_type,
+        max_active_clusters,
+    )
+
+    kernel(
+        tiled_mma,
+        tma_atom_a,
+        a_tma_tensor,
+        tma_atom_b,
+        b_tma_tensor,
+        tma_atom_c,
+        c_tma_tensor,
+        a_smem_layout,
+        b_smem_layout,
+        c_smem_layout_kind,
+        epi_smem_layout_staged,
+        epi_tile,
+        cta_layout_vmnk,
+        tile_sched_params,
+    ).launch(
+        grid=grid_shape,
+        block=[224, 1, 1] if use_clc_dynamic_scheduler else [192, 1, 1],
+        cluster=cluster_shape_mnk,
+    )
+
+
+def run_dense_gemm(
+    mnk: Tuple[int, int, int],
+    tolerance: float,
+):
+    global torch, cutlass_torch
+    import torch
+    import cutlass.torch as cutlass_torch
+
+    print("===================================================================")
+    print("Running Blackwell fp16 GEMM example 5 (with TMA prefetch):")
+    print(f"  mnk:       {mnk}")
+    print(f"  tolerance: {tolerance}")
+    print("===================================================================")
+    print()
+
+    m, n, k = mnk
+    torch.manual_seed(1111)
+
+    # Make K-major tensors (torch tensors are row-major)
+    def make_tensors(mn, k, dtype):
+        shape = (mn, k)
+        return (
+            torch.empty(*shape, dtype=torch.int32)
+            .random_(-2, 2)
+            .to(device="cuda", dtype=dtype)
+        )
+
+    a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
+    b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
+    c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
+    a_memref = from_dlpack(a).mark_layout_dynamic()
+    b_memref = from_dlpack(b).mark_layout_dynamic()
+    c_memref = from_dlpack(c).mark_layout_dynamic()
+
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
+        cluster_shape_mnk[0] * cluster_shape_mnk[1]
+    )
+
+    # Entry point to the host JIT function
+    host_function(
+        a_memref,
+        b_memref,
+        c_memref,
+        max_active_clusters,
+        no_cache=True,
+    )
+
+    # Compute reference result and verify
+    ref = (torch.einsum("mk,nk->mn", a, b)).cpu()
+    torch.testing.assert_close(
+        c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
+    )
+
+
+if __name__ == "__main__":
+
+    def parse_comma_separated_ints(s: str):
+        try:
+            return [int(x.strip()) for x in s.split(",")]
+        except ValueError:
+            raise argparse.ArgumentTypeError(
+                "Invalid format. Expected comma-separated integers."
+            )
+
+    from cuda.bindings import driver as cu_driver
+
+    cu_driver.cuInit(0)
+    err, device_count = cu_driver.cuDeviceGetCount()
+    if err != cu_driver.CUresult.CUDA_SUCCESS or device_count < 1:
+        raise RuntimeError("A GPU is required to run this example")
+
+    parser = argparse.ArgumentParser(
+        description="Blackwell fp16 GEMM example 5 (with TMA prefetch)"
+    )
+    parser.add_argument(
+        "--mnk",
+        type=parse_comma_separated_ints,
+        default=(8192, 8192, 8192),
+        help="MNK dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    args = parser.parse_args()
+    if len(args.mnk) != 3:
+        parser.error("--mnk must contain exactly 3 values")
+
+    run_dense_gemm(
+        args.mnk,
+        args.tolerance,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/nvfp4_gemm_0.py

@@ -0,0 +1,778 @@
+# Copyright (c) 2025 - 2026 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 os
+import sys
+from typing import Type, Tuple
+import cuda.bindings.driver as cuda
+
+import torch
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.torch as cutlass_torch
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass.cute.runtime import make_ptr
+
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    examples_dir = os.path.join(current_dir, "..", "..", "..", "..")
+    if examples_dir not in sys.path:
+        sys.path.insert(0, examples_dir)
+
+from cute.blackwell.tutorial.tutorial_gemm.utils import create_parser, run
+
+mma_tiler_mn = (128, 256)
+mma_inst_shape_k = 64
+ab_dtype = cutlass.Float4E2M1FN
+sf_dtype = cutlass.Float8E4M3FN
+c_dtype = cutlass.Float16
+sf_vec_size = 16
+
+"""
+The first tutorial NVFP4 block-scaled batched GEMM demonstrating a simple kernel implementation in CuTeDSL
+
+This example demonstrates the kernel implementation of block-scaled batched GEMM with NVFP4 data type.
+With large tile sizes (128x256x256), it can achieve very high performance on 8k×8k×8k problem sizes.
+It can serve as a starting point to help users quickly experiment with optimizations for 
+challenges that may arise with other problem sizes.
+
+To run this example:
+.. code-block:: bash
+
+    python examples/blackwell/tutorial_gemm/nvfp4_gemm_0.py  \
+      --mnkl 8192,8192,8192,1 --do_benchmark
+
+Constraints for this example:
+* The problem size of m, n and k must be divisible by the tile size m&n&k (128,256,256)
+* The scaling factor vector size is 16.
+* The A/B matrices have data contiguous on the k dimension.
+* The C matrix has data contiguous on the n dimension.
+* The A/B matrix data type is Float4E2M1FN.
+* The SFA/SFB matrix data type is Float8E4M3FN.
+"""
+
+
+class Sm100BlockScaledDenseGemmKernel:
+    def __init__(self):
+        self.threads_per_cta = 128
+        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
+        self.num_tmem_alloc_cols = 512
+
+        # set stages for ab_pipeline and acc_pipeline
+        self.num_acc_stage = 1
+        self.num_ab_stage = 4
+
+    @cute.jit
+    def __call__(
+        self,
+        a_ptr: cute.Pointer,
+        b_ptr: cute.Pointer,
+        sfa_ptr: cute.Pointer,
+        sfb_ptr: cute.Pointer,
+        c_ptr: cute.Pointer,
+        problem_size: tuple,
+        stream: cuda.CUstream,
+        epilogue_op: cutlass.Constexpr = lambda x: x,
+    ):
+        # setup static attributes before smem/grid/tma computation
+        self.c_layout = utils.LayoutEnum.ROW_MAJOR
+        m, n, k, l = problem_size
+
+        # Setup attributes that depend on gemm inputs
+        mma_inst_tile_k = 4
+        self.mma_tiler = (
+            mma_tiler_mn[0],
+            mma_tiler_mn[1],
+            mma_inst_shape_k * mma_inst_tile_k,
+        )
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0],
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+
+        a_tensor = cute.make_tensor(
+            a_ptr,
+            cute.make_layout(
+                (m, cute.assume(k, 32), l),
+                stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
+            ),
+        )
+        b_tensor = cute.make_tensor(
+            b_ptr,
+            cute.make_layout(
+                (n, cute.assume(k, 32), l),
+                stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
+            ),
+        )
+        # make address offset of c_tensor 256bit aligned,
+        # so that epilogue could use vectorized store with larger vector size.
+        c_tensor = cute.make_tensor(
+            c_ptr,
+            cute.make_layout(
+                (cute.assume(m, 32), cute.assume(n, 16), l),
+                stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)),
+            ),
+        )
+        # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
+        # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
+        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            a_tensor.shape, sf_vec_size
+        )
+        sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
+
+        # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
+        sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            b_tensor.shape, sf_vec_size
+        )
+        sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
+
+        mma_op = tcgen05.MmaMXF4NVF4Op(
+            sf_dtype,
+            (*mma_tiler_mn, mma_inst_shape_k),
+            tcgen05.CtaGroup.ONE,
+            tcgen05.OperandSource.SMEM,
+        )
+        tiled_mma = cute.make_tiled_mma(mma_op)
+
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((1, 1, 1)),
+            (tiled_mma.thr_id.shape,),
+        )
+
+        # Compute A/B/SFA/SFB/C shared memory layout
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            self.mma_tiler,
+            ab_dtype,
+            self.num_ab_stage,
+        )
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            self.mma_tiler,
+            ab_dtype,
+            self.num_ab_stage,
+        )
+        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            self.num_ab_stage,
+        )
+        self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            self.num_ab_stage,
+        )
+
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        # TMA load for A
+        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
+            cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+            a_tensor,
+            a_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for B
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+            b_tensor,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # TMA load for SFA
+        sfa_smem_layout = cute.slice_(
+            self.sfa_smem_layout_staged, (None, None, None, 0)
+        )
+        tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
+            cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+            sfa_tensor,
+            sfa_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        # TMA load for SFB
+        sfb_smem_layout = cute.slice_(
+            self.sfb_smem_layout_staged, (None, None, None, 0)
+        )
+        tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
+            cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+            sfb_tensor,
+            sfb_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        # Compute TMA load bytes
+        a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
+        b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
+        sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
+        sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
+        self.num_tma_load_bytes = (
+            a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
+        ) * atom_thr_size
+
+        # Compute grid size
+        grid = (
+            cute.ceil_div(c_tensor.shape[0], self.cta_tile_shape_mnk[0]),
+            cute.ceil_div(c_tensor.shape[1], self.cta_tile_shape_mnk[1]),
+            c_tensor.shape[2],
+        )
+
+        # Launch the kernel synchronously
+        self.kernel(
+            tiled_mma,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_sfa,
+            tma_tensor_sfa,
+            tma_atom_sfb,
+            tma_tensor_sfb,
+            c_tensor,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            epilogue_op,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=(1, 1, 1),
+            stream=stream,
+        )
+        return
+
+    # GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,
+        tma_atom_sfa: cute.CopyAtom,
+        mSFA_mkl: cute.Tensor,
+        tma_atom_sfb: cute.CopyAtom,
+        mSFB_nkl: cute.Tensor,
+        mC_mnl: cute.Tensor,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_staged: cute.Layout,
+        epilogue_op: cutlass.Constexpr,
+    ):
+        """
+        GPU device kernel performing the batched GEMM computation.
+        """
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        tidx, _, _ = cute.arch.thread_idx()
+
+        #
+        # Setup cta/thread coordinates
+        #
+        # Coords inside cluster
+        bidx, bidy, bidz = cute.arch.block_idx()
+
+        # Coords outside cluster
+        cta_coord = (bidx, bidy, bidz)
+        mma_tile_coord_mnl = (
+            cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
+            cta_coord[1],
+            cta_coord[2],
+        )
+
+        #
+        # Define shared storage for kernel
+        #
+        @cute.struct
+        class SharedStorage:
+            ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
+            acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
+            tmem_holding_buf: cutlass.Int32
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sA = smem.allocate_tensor(
+            element_type=ab_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        # (MMA, MMA_N, MMA_K, STAGE)
+        sB = smem.allocate_tensor(
+            element_type=ab_dtype,
+            layout=b_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=b_smem_layout_staged.inner,
+        )
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sSFA = smem.allocate_tensor(
+            element_type=sf_dtype,
+            layout=sfa_smem_layout_staged,
+            byte_alignment=128,
+        )
+        # (MMA, MMA_N, MMA_K, STAGE)
+        sSFB = smem.allocate_tensor(
+            element_type=sf_dtype,
+            layout=sfb_smem_layout_staged,
+            byte_alignment=128,
+        )
+
+        #
+        # Initialize mainloop ab_pipeline, acc_pipeline and their states
+        #
+        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+            num_stages=self.num_ab_stage,
+            producer_group=ab_pipeline_producer_group,
+            consumer_group=ab_pipeline_consumer_group,
+            tx_count=self.num_tma_load_bytes,
+        ).make_participants()
+        acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+            num_stages=self.num_acc_stage,
+            producer_group=ab_pipeline_producer_group,
+            consumer_group=pipeline.CooperativeGroup(
+                pipeline.Agent.Thread,
+                self.threads_per_cta,
+            ),
+        ).make_participants()
+
+        #
+        # Local_tile partition global tensors
+        #
+        # (bM, bK, RestM, RestK, RestL)
+        gA_mkl = cute.local_tile(
+            mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        # (bN, bK, RestN, RestK, RestL)
+        gB_nkl = cute.local_tile(
+            mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+        )
+        gSFA_mkl = cute.local_tile(
+            mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        gSFB_nkl = cute.local_tile(
+            mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+        )
+        # (bM, bN, RestM, RestN, RestL)
+        gC_mnl = cute.local_tile(
+            mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
+        )
+        k_tile_cnt = cute.size(gA_mkl, mode=[3])
+
+        #
+        # Partition global tensor for TiledMMA_A/B/SFA/SFB/C
+        #
+        # (MMA, MMA_M, MMA_K, RestK)
+        thr_mma = tiled_mma.get_slice(0)
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgA = thr_mma.partition_A(gA_mkl)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        tCgB = thr_mma.partition_B(gB_nkl)
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgSFA = thr_mma.partition_A(gSFA_mkl)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        tCgSFB = thr_mma.partition_B(gSFB_nkl)
+        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+        tCgC = thr_mma.partition_C(gC_mnl)
+
+        #
+        # Partition global/shared tensor for TMA load A/B/SFA/SFB
+        #
+        # TMA load A partition_S/D
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsA, tAgA = cpasync.tma_partition(
+            tma_atom_a,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(tCgA, 0, 3),
+        )
+        # TMA load B partition_S/D
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestN, RestK, RestL)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sB, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+
+        #  TMA load partition for SFA tensor
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsSFA, tAgSFA = cpasync.tma_partition(
+            tma_atom_sfa,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(tCgSFA, 0, 3),
+        )
+        tAsSFA = cute.filter_zeros(tAsSFA)
+        tAgSFA = cute.filter_zeros(tAgSFA)
+
+        # TMA load partition for SFB tensor
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestN, RestK, RestL)
+        tBsSFB, tBgSFB = cpasync.tma_partition(
+            tma_atom_sfb,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sSFB, 0, 3),
+            cute.group_modes(tCgSFB, 0, 3),
+        )
+        tBsSFB = cute.filter_zeros(tBsSFB)
+        tBgSFB = cute.filter_zeros(tBgSFB)
+
+        #
+        # Partition shared/tensor memory tensor for TiledMMA_A/B/C
+        #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        tCrA = tiled_mma.make_fragment_A(sA)
+        # (MMA, MMA_N, MMA_K, STAGE)
+        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)
+        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+
+        #
+        # Alloc tensor memory buffer
+        #
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=self.threads_per_cta,
+        )
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=tmem_alloc_barrier,
+        )
+        tmem.allocate(self.num_tmem_alloc_cols)
+        tmem.wait_for_alloc()
+        acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
+        tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+
+        #
+        # Make SFA/SFB tmem tensor
+        #
+        # Get SFA tmem ptr
+        sfa_tmem_ptr = cute.recast_ptr(
+            acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
+            dtype=sf_dtype,
+        )
+        # (MMA, MMA_M, MMA_K)
+        tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+        )
+        tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+        # Get SFB tmem ptr
+        sfb_tmem_ptr = cute.recast_ptr(
+            acc_tmem_ptr
+            + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+            + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
+            dtype=sf_dtype,
+        )
+        # (MMA, MMA_N, MMA_K)
+        tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
+        )
+        tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+
+        #
+        # Partition for S2T copy of SFA/SFB
+        #
+        # Make S2T CopyAtom
+        copy_atom_s2t = cute.make_copy_atom(
+            tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
+            sf_dtype,
+        )
+        # (MMA, MMA_MN, MMA_K, STAGE)
+        tCsSFA_compact = cute.filter_zeros(sSFA)
+        # (MMA, MMA_MN, MMA_K)
+        tCtSFA_compact = cute.filter_zeros(tCtSFA)
+        tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
+        thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+            tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
+        )
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
+
+        # (MMA, MMA_MN, MMA_K, STAGE)
+        tCsSFB_compact = cute.filter_zeros(sSFB)
+        # (MMA, MMA_MN, MMA_K)
+        tCtSFB_compact = cute.filter_zeros(tCtSFB)
+        tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
+        thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+            tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
+        )
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
+
+        #
+        # Slice to per mma tile index
+        #
+        # ((atom_v, rest_v), RestK)
+        tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+        # ((atom_v, rest_v), RestK)
+        tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+        # ((atom_v, rest_v), RestK)
+        tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+        # ((atom_v, rest_v), RestK)
+        tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+
+        #
+        # Execute Data copy and Math computation in the k_tile loop
+        #
+        if warp_idx == 0:
+            # Wait for accumulator buffer empty
+            acc_empty = acc_producer.acquire_and_advance()
+            # Set ACCUMULATE field to False for the first k_tile iteration
+            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+            # Execute k_tile loop
+            for k_tile in cutlass.range(
+                k_tile_cnt, prefetch_stages=self.num_ab_stage - 2
+            ):
+                # Wait for AB buffer empty
+                ab_empty = ab_producer.acquire_and_advance()
+
+                # TMA load for A/B/SFA/SFB
+                cute.copy(
+                    tma_atom_a,
+                    tAgA[(None, ab_empty.count)],
+                    tAsA[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                )
+                cute.copy(
+                    tma_atom_b,
+                    tBgB[(None, ab_empty.count)],
+                    tBsB[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                )
+                cute.copy(
+                    tma_atom_sfa,
+                    tAgSFA[(None, ab_empty.count)],
+                    tAsSFA[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                )
+                cute.copy(
+                    tma_atom_sfb,
+                    tBgSFB[(None, ab_empty.count)],
+                    tBsSFB[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                )
+
+                # Wait for AB buffer full
+                ab_full = ab_consumer.wait_and_advance()
+
+                #  Copy SFA/SFB to tmem
+                s2t_stage_coord = (None, None, None, None, ab_full.index)
+                tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
+                tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
+                cute.copy(
+                    tiled_copy_s2t_sfa,
+                    tCsSFA_compact_s2t_staged,
+                    tCtSFA_compact_s2t,
+                )
+                cute.copy(
+                    tiled_copy_s2t_sfb,
+                    tCsSFB_compact_s2t_staged,
+                    tCtSFB_compact_s2t,
+                )
+
+                # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
+                num_kblocks = cute.size(tCrA, mode=[2])
+                for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
+                    kblock_coord = (
+                        None,
+                        None,
+                        kblock_idx,
+                        ab_full.index,
+                    )
+
+                    # Set SFA/SFB tensor to tiled_mma
+                    sf_kblock_coord = (None, None, kblock_idx)
+                    tiled_mma.set(
+                        tcgen05.Field.SFA,
+                        tCtSFA[sf_kblock_coord].iterator,
+                    )
+                    tiled_mma.set(
+                        tcgen05.Field.SFB,
+                        tCtSFB[sf_kblock_coord].iterator,
+                    )
+
+                    cute.gemm(
+                        tiled_mma,
+                        tCtAcc,
+                        tCrA[kblock_coord],
+                        tCrB[kblock_coord],
+                        tCtAcc,
+                    )
+                    # Enable accumulate on tCtAcc after first kblock
+                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                # Async arrive AB buffer empty
+                ab_full.release()
+            acc_empty.commit()
+
+        #
+        # Epilogue
+        # Partition for epilogue
+        #
+        op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
+        copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
+        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc)
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        # (T2R_M, T2R_N, EPI_M, EPI_M)
+        tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc)
+        # (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        tTR_gC = thr_copy_t2r.partition_D(tCgC)
+        # (T2R_M, T2R_N, EPI_M, EPI_N）
+        tTR_rAcc = cute.make_rmem_tensor(
+            tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32
+        )
+        # (T2R_M, T2R_N, EPI_M, EPI_N）
+        tTR_rC = cute.make_rmem_tensor(
+            tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype
+        )
+        # STG Atom
+        simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype)
+        tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)]
+
+        # Release TMEM allocation lock
+        tmem.relinquish_alloc_permit()
+
+        # Wait for accumulator buffer full
+        acc_full = acc_consumer.wait_and_advance()
+
+        # Copy accumulator to register
+        cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc)
+        acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype))
+        tTR_rC.store(acc_vec)
+        # Store C to global memory
+        cute.copy(simt_atom, tTR_rC, tTR_gC)
+
+        acc_full.release()
+
+        # Deallocate TMEM
+        cute.arch.barrier()
+        tmem.free(acc_tmem_ptr)
+
+        return
+
+
+def run_nvfp4_gemm(
+    mnkl: Tuple[int, int, int, int],
+    tolerance: float,
+    do_benchmark: bool = False,
+    warmup_iterations: int = 10,
+    iterations: int = 100,
+    use_cold_l2: bool = True,
+):
+    run(
+        gemm_class=Sm100BlockScaledDenseGemmKernel,
+        ab_dtype=ab_dtype,
+        sf_dtype=sf_dtype,
+        c_dtype=c_dtype,
+        sf_vec_size=sf_vec_size,
+        mma_tiler_mn=mma_tiler_mn,
+        cluster_shape_mnk=(1, 1, 1),
+        mnkl=mnkl,
+        tolerance=tolerance,
+        do_benchmark=do_benchmark,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+        use_cold_l2=use_cold_l2,
+    )
+
+
+if __name__ == "__main__":
+    parser = create_parser()
+    args = parser.parse_args()
+
+    if len(args.mnkl) != 4:
+        parser.error("--mnkl must contain exactly 4 values")
+
+    m, n, k, _ = args.mnkl
+    if m % mma_tiler_mn[0] != 0:
+        parser.error("m must be multiples of mma_tiler_mn[0] (got m={})".format(m))
+    if n % mma_tiler_mn[1] != 0:
+        parser.error("n must be multiples of mma_tiler_mn[1] (got n={})".format(n))
+    if k % 256 != 0:
+        parser.error("k must be a multiple of 256 (got k={})".format(k))
+
+    run_nvfp4_gemm(
+        args.mnkl,
+        args.tolerance,
+        args.do_benchmark,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/nvfp4_gemm_1.py

@@ -0,0 +1,934 @@
+# Copyright (c) 2025 - 2026 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.
+
+# This is the second tutorial nvfp4 GEMM. It builds on the first tutorial by adding 2CTA MMA
+# instructions with a 2x1 cluster.
+
+import argparse
+import os
+import sys
+from typing import Type, Tuple
+import cuda.bindings.driver as cuda
+
+import torch
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync, tcgen05
+import cutlass.torch as cutlass_torch
+import cutlass.utils.blackwell_helpers as sm100_utils
+import cutlass.utils.blockscaled_layout as blockscaled_utils
+from cutlass.cute.runtime import from_dlpack, make_ptr
+
+if __name__ == "__main__":
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    examples_dir = os.path.join(current_dir, "..", "..", "..", "..")
+    if examples_dir not in sys.path:
+        sys.path.insert(0, examples_dir)
+
+from cute.blackwell.tutorial.tutorial_gemm.utils import create_parser, run
+
+mma_tiler_mn = (256, 256)
+mma_inst_shape_k = 64
+ab_dtype = cutlass.Float4E2M1FN
+sf_dtype = cutlass.Float8E4M3FN
+c_dtype = cutlass.Float16
+sf_vec_size = 16
+cluster_shape_mnk = (2, 1, 1)
+
+"""
+The second tutorial further improves the performance of NVFP4 block-scaled batched GEMM
+by adding 2CTA instructions and TMA multicast optimizations.
+
+(1) The 2 CTA instructions could reduce the smem size requirement for B tensor,
+increased num_ab_stage and improves the latency hiding capability.
+
+For both 1CTA and 2CTA, the shared memory (smem) size per stage for the A, sfA, and sfB tensors is the same:
+- For the A tensor, each stage requires 128 x 256 x sizeof(float4) = 16KB.
+- For the sfA tensor, each stage requires 128 x (256 / 16) x sizeof(float8) = 2KB.
+- For the sfB tensor, each stage requires 256 x (256 / 16) x sizeof(float8) = 4KB.
+
+The situation is different for the B tensor:
+- In the 1CTA case, each stage for the B tensor requires 256 x 256 x sizeof(float4) = 32KB.
+- In the 2CTA case, only half this size is needed, i.e., 128 x 256 x sizeof(float4) = 16KB.
+
+Therefore, the maximum number of AB stages is:
+- For 1CTA: 227 // (16 + 32 + 2 + 4) = 4
+- For 2CTA: 227 // (16 + 16 + 2 + 4) = 5
+
+The latency hiding capability is:
+- 1CTA: 512 * (4 - 1) = 1.5K cycles
+- 2CTA: 512 * (5 - 1) = 2K cycles
+
+(2) TMA multicast can help reduce L2 cache traffic.
+
+Without TMA multicast, the L2 traffic per tile is typically 16KB + 32KB = 48KB (possibly less in practice, depending on hardware optimizations).
+With TMA multicast in a cluster of shape (m, n), the L2 traffic per tile is reduced to 16KB / n + 32KB / m.
+For example:
+- In a 2x1 cluster: 16KB / 1 + 32KB / 2 = 24KB per tile
+- In a 4x4 cluster: 16KB / 4 + 32KB / 4 = 12KB per tile
+
+The first approach offers substantial capacity for hiding latency, whereas the second reduces the time required for data to become ready.
+Both could be tried when the workload is latency-bound or limited by memory throughput.
+
+To run this example:
+.. code-block:: bash
+
+    python examples/blackwell/tutorial_gemm/nvfp4_gemm_1.py  \
+      --mnkl 8192,8192,8192,1 --do_benchmark
+
+Constraints for this example:
+* The problem size of m, n and k must be divisible by the tile size m&n&k (256, 256, 256)
+* The scaling factor vector size is 16.
+* The A/B matrices have data contiguous on the k dimension.
+* The C matrix has data contiguous on the n dimension.
+* The A/B matrix data type is Float4E2M1FN.
+* The SFA/SFB matrix data type is Float8E4M3FN.
+"""
+
+
+class Sm100BlockScaledDenseGemmKernel:
+    def __init__(self):
+        self.threads_per_cta = 128
+        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
+        self.num_tmem_alloc_cols = 512
+
+        # set stages for ab_pipeline and acc_pipeline
+        self.num_acc_stage = 1
+        self.num_ab_stage = 5
+
+    @cute.jit
+    def __call__(
+        self,
+        a_ptr: cute.Pointer,
+        b_ptr: cute.Pointer,
+        sfa_ptr: cute.Pointer,
+        sfb_ptr: cute.Pointer,
+        c_ptr: cute.Pointer,
+        problem_size: tuple,
+        stream: cuda.CUstream,
+        epilogue_op: cutlass.Constexpr = lambda x: x,
+    ):
+        # setup static attributes before smem/grid/tma computation
+        self.c_layout = utils.LayoutEnum.ROW_MAJOR
+        m, n, k, l = problem_size
+
+        self.use_2cta_instrs = False if mma_tiler_mn[0] == 128 else True
+
+        # Setup attributes that depend on gemm inputs
+        mma_inst_tile_k = 4
+        self.mma_tiler = (
+            mma_tiler_mn[0],
+            mma_tiler_mn[1],
+            mma_inst_shape_k * mma_inst_tile_k,
+        )
+
+        self.mma_inst_shape_sfb = (
+            mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1),
+            mma_tiler_mn[1],
+            mma_inst_shape_k,
+        )
+        self.mma_tiler_sfb = (
+            self.mma_inst_shape_sfb[0],
+            self.mma_inst_shape_sfb[1],
+            mma_inst_shape_k * mma_inst_tile_k,
+        )
+
+        a_tensor = cute.make_tensor(
+            a_ptr,
+            cute.make_layout(
+                (m, cute.assume(k, 32), l),
+                stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
+            ),
+        )
+        b_tensor = cute.make_tensor(
+            b_ptr,
+            cute.make_layout(
+                (n, cute.assume(k, 32), l),
+                stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
+            ),
+        )
+        # 256bit aligned. row_major
+        c_tensor = cute.make_tensor(
+            c_ptr,
+            cute.make_layout(
+                (cute.assume(m, 32), cute.assume(n, 16), l),
+                stride=(cute.assume(n, 16), 1, cute.assume(m * n, 512)),
+            ),
+        )
+
+        # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
+        # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
+        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            a_tensor.shape, sf_vec_size
+        )
+        sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
+
+        # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
+        sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
+            b_tensor.shape, sf_vec_size
+        )
+        sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
+
+        mma_op = tcgen05.MmaMXF4NVF4Op(
+            sf_dtype,
+            (*mma_tiler_mn, mma_inst_shape_k),
+            tcgen05.CtaGroup.ONE if not self.use_2cta_instrs else tcgen05.CtaGroup.TWO,
+            tcgen05.OperandSource.SMEM,
+        )
+        tiled_mma = cute.make_tiled_mma(mma_op)
+
+        # (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
+        # Note sfB don't support share among 2ctas
+        sfb_mma_op = tcgen05.MmaMXF4NVF4Op(
+            sf_dtype,
+            self.mma_inst_shape_sfb,
+            tcgen05.CtaGroup.ONE,
+            tcgen05.OperandSource.SMEM,
+        )
+        tiled_mma_sfb = cute.make_tiled_mma(sfb_mma_op)
+
+        self.cta_tile_shape_mnk = (
+            self.mma_tiler[0] // (2 if self.use_2cta_instrs else 1),
+            self.mma_tiler[1],
+            self.mma_tiler[2],
+        )
+        self.cta_tile_shape_mnk_sfb = (
+            self.mma_tiler_sfb[0] // (2 if self.use_2cta_instrs else 1),
+            self.mma_tiler_sfb[1],
+            self.mma_tiler_sfb[2],
+        )
+
+        self.cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout(cluster_shape_mnk),
+            (tiled_mma.thr_id.shape,),
+        )
+        self.cluster_layout_sfb_vmnk = cute.tiled_divide(
+            cute.make_layout(cluster_shape_mnk),
+            (tiled_mma_sfb.thr_id.shape,),
+        )
+
+        # Compute number of multicast CTAs for A/B
+        self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+        self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
+        self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
+        self.is_a_mcast = self.num_mcast_ctas_a > 1
+        self.is_b_mcast = self.num_mcast_ctas_b > 1
+        self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
+
+        # Compute A/B/SFA/SFB/C shared memory layout
+        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+            tiled_mma,
+            self.mma_tiler,
+            ab_dtype,
+            self.num_ab_stage,
+        )
+        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+            tiled_mma,
+            self.mma_tiler,
+            ab_dtype,
+            self.num_ab_stage,
+        )
+        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            self.num_ab_stage,
+        )
+        self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            self.num_ab_stage,
+        )
+
+        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+
+        a_op = sm100_utils.cluster_shape_to_tma_atom_A(
+            cluster_shape_mnk[:2], tiled_mma.thr_id
+        )
+        # TMA load for A
+        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
+            a_op,
+            a_tensor,
+            a_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+        # TMA load for B
+        b_op = sm100_utils.cluster_shape_to_tma_atom_B(
+            cluster_shape_mnk[:2], tiled_mma.thr_id
+        )
+        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+            b_op,
+            b_tensor,
+            b_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+        )
+
+        # TMA load for SFA
+        sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
+            cluster_shape_mnk[:2], tiled_mma.thr_id
+        )
+        sfa_smem_layout = cute.slice_(
+            self.sfa_smem_layout_staged, (None, None, None, 0)
+        )
+        tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
+            sfa_op,
+            sfa_tensor,
+            sfa_smem_layout,
+            self.mma_tiler,
+            tiled_mma,
+            self.cluster_layout_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        # TMA load for SFB
+        sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
+            cluster_shape_mnk[:2], tiled_mma.thr_id
+        )
+        sfb_smem_layout = cute.slice_(
+            self.sfb_smem_layout_staged, (None, None, None, 0)
+        )
+        tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
+            sfb_op,
+            sfb_tensor,
+            sfb_smem_layout,
+            self.mma_tiler_sfb,
+            tiled_mma_sfb,
+            self.cluster_layout_sfb_vmnk.shape,
+            internal_type=cutlass.Int16,
+        )
+
+        # Compute TMA load bytes
+        a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
+        b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
+        sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
+        sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
+        self.num_tma_load_bytes = (
+            a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
+        ) * atom_thr_size
+
+        # Compute grid size
+        grid = cute.round_up(
+            cute.ceil_div(
+                (c_tensor.layout.shape),
+                (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], 1),
+            ),
+            cluster_shape_mnk,
+        )
+
+        # Launch the kernel
+        self.kernel(
+            tiled_mma,
+            tiled_mma_sfb,
+            tma_atom_a,
+            tma_tensor_a,
+            tma_atom_b,
+            tma_tensor_b,
+            tma_atom_sfa,
+            tma_tensor_sfa,
+            tma_atom_sfb,
+            tma_tensor_sfb,
+            c_tensor,
+            self.a_smem_layout_staged,
+            self.b_smem_layout_staged,
+            self.sfa_smem_layout_staged,
+            self.sfb_smem_layout_staged,
+            self.cluster_layout_vmnk,
+            self.cluster_layout_sfb_vmnk,
+            epilogue_op,
+        ).launch(
+            grid=grid,
+            block=[self.threads_per_cta, 1, 1],
+            cluster=cluster_shape_mnk,
+            stream=stream,
+        )
+        return
+
+    # GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        tiled_mma: cute.TiledMma,
+        tiled_mma_sfb: cute.TiledMma,
+        tma_atom_a: cute.CopyAtom,
+        mA_mkl: cute.Tensor,
+        tma_atom_b: cute.CopyAtom,
+        mB_nkl: cute.Tensor,
+        tma_atom_sfa: cute.CopyAtom,
+        mSFA_mkl: cute.Tensor,
+        tma_atom_sfb: cute.CopyAtom,
+        mSFB_nkl: cute.Tensor,
+        mC_mnl: cute.Tensor,
+        a_smem_layout_staged: cute.ComposedLayout,
+        b_smem_layout_staged: cute.ComposedLayout,
+        sfa_smem_layout_staged: cute.Layout,
+        sfb_smem_layout_staged: cute.Layout,
+        cta_layout_vmnk: cute.Layout,
+        cta_layout_sfb_vmnk: cute.Layout,
+        epilogue_op: cutlass.Constexpr,
+    ):
+        """
+        GPU device kernel performing the batched GEMM computation.
+        """
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        tidx, _, _ = cute.arch.thread_idx()
+
+        #
+        # Setup cta/thread coordinates
+        #
+        # Coords inside cluster
+        bidx, bidy, bidz = cute.arch.block_idx()
+        cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+        cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
+        cta_in_cluster_coord_sfb_vmnk = cta_layout_sfb_vmnk.get_flat_coord(
+            cta_rank_in_cluster
+        )
+
+        # Coords outside cluster
+        mma_tile_coord_vmnk = (
+            bidx % cute.size(cta_layout_vmnk, mode=[0]),
+            bidx // cute.size(cta_layout_vmnk, mode=[0]),
+            bidy,
+            bidz,
+        )
+        mma_tile_coord_mnl = mma_tile_coord_vmnk[1:]
+        is_leader_cta = mma_tile_coord_vmnk[0] == 0
+
+        #
+        # Define shared storage for kernel
+        #
+        @cute.struct
+        class SharedStorage:
+            ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
+            acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
+            tmem_dealloc_mbar: cutlass.Int64
+            tmem_holding_buf: cutlass.Int32
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(SharedStorage)
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sA = smem.allocate_tensor(
+            element_type=ab_dtype,
+            layout=a_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=a_smem_layout_staged.inner,
+        )
+        # (MMA, MMA_N, MMA_K, STAGE)
+        sB = smem.allocate_tensor(
+            element_type=ab_dtype,
+            layout=b_smem_layout_staged.outer,
+            byte_alignment=128,
+            swizzle=b_smem_layout_staged.inner,
+        )
+        # (MMA, MMA_M, MMA_K, STAGE)
+        sSFA = smem.allocate_tensor(
+            element_type=sf_dtype,
+            layout=sfa_smem_layout_staged,
+            byte_alignment=128,
+        )
+        # (MMA, MMA_N, MMA_K, STAGE)
+        sSFB = smem.allocate_tensor(
+            element_type=sf_dtype,
+            layout=sfb_smem_layout_staged,
+            byte_alignment=128,
+        )
+
+        #
+        # Compute multicast mask for A/B/SFA/SFB buffer full
+        #
+        a_full_mcast_mask = None
+        b_full_mcast_mask = None
+        sfa_full_mcast_mask = None
+        sfb_full_mcast_mask = None
+        if cutlass.const_expr(
+            self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs
+        ):
+            a_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+            )
+            b_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
+            )
+            sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
+            )
+            sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                cta_layout_sfb_vmnk, cta_in_cluster_coord_sfb_vmnk, mcast_mode=1
+            )
+
+        #
+        # Initialize mainloop ab_pipeline, acc_pipeline and their states
+        #
+        num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        ab_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, num_tma_producer
+        )
+        ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+            barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+            num_stages=self.num_ab_stage,
+            producer_group=ab_pipeline_producer_group,
+            consumer_group=ab_pipeline_consumer_group,
+            tx_count=self.num_tma_load_bytes,
+            cta_layout_vmnk=cta_layout_vmnk,
+        ).make_participants()
+        acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+            barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+            num_stages=self.num_acc_stage,
+            producer_group=ab_pipeline_producer_group,
+            consumer_group=pipeline.CooperativeGroup(
+                pipeline.Agent.Thread,
+                self.threads_per_cta * (2 if self.use_2cta_instrs else 1),
+            ),
+            cta_layout_vmnk=cta_layout_vmnk,
+        ).make_participants()
+
+        #
+        # Local_tile partition global tensors
+        #
+        # (bM, bK, RestM, RestK, RestL)
+        gA_mkl = cute.local_tile(
+            mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        # (bN, bK, RestN, RestK, RestL)
+        gB_nkl = cute.local_tile(
+            mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+        )
+        gSFA_mkl = cute.local_tile(
+            mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
+        )
+        gSFB_nkl = cute.local_tile(
+            mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
+        )
+        # (bM, bN, RestM, RestN, RestL)
+        gC_mnl = cute.local_tile(
+            mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
+        )
+        k_tile_cnt = cute.size(gA_mkl, mode=[3])
+
+        #
+        # Partition global tensor for TiledMMA_A/B/SFA/SFB/C
+        #
+        # (MMA, MMA_M, MMA_K, RestK)
+        thr_mma = tiled_mma.get_slice(mma_tile_coord_vmnk[0])
+        thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_vmnk[0])
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgA = thr_mma.partition_A(gA_mkl)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        tCgB = thr_mma.partition_B(gB_nkl)
+        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+        tCgSFA = thr_mma.partition_A(gSFA_mkl)
+        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+        # tCgSFB = thr_mma.partition_B(gSFB_nkl)
+        tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
+        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+        tCgC = thr_mma.partition_C(gC_mnl)
+
+        #
+        # Partition global/shared tensor for TMA load A/B/SFA/SFB
+        #
+        # TMA load A partition_S/D
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsA, tAgA = cpasync.tma_partition(
+            tma_atom_a,
+            # 0,
+            # cute.make_layout(1),
+            cta_in_cluster_coord_vmnk[2],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+            cute.group_modes(sA, 0, 3),
+            cute.group_modes(tCgA, 0, 3),
+        )
+        # TMA load B partition_S/D
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestN, RestK, RestL)
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_b,
+            # 0,
+            # cute.make_layout(1),
+            cta_in_cluster_coord_vmnk[1],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])),
+            cute.group_modes(sB, 0, 3),
+            cute.group_modes(tCgB, 0, 3),
+        )
+
+        #  TMA load SFA partition_S/D
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestM, RestK, RestL)
+        tAsSFA, tAgSFA = cpasync.tma_partition(
+            tma_atom_sfa,
+            # 0,
+            # cute.make_layout(1),
+            cta_in_cluster_coord_vmnk[2],
+            cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])),
+            cute.group_modes(sSFA, 0, 3),
+            cute.group_modes(tCgSFA, 0, 3),
+        )
+        tAsSFA = cute.filter_zeros(tAsSFA)
+        tAgSFA = cute.filter_zeros(tAgSFA)
+
+        # TMA load SFB partition_S/D
+        # ((atom_v, rest_v), STAGE)
+        # ((atom_v, rest_v), RestN, RestK, RestL)
+        sfb_cta_layout = cute.make_layout(
+            cute.slice_(cta_layout_sfb_vmnk, (0, None, 0, 0)).shape
+        )
+        tBsSFB, tBgSFB = cpasync.tma_partition(
+            tma_atom_sfb,
+            cta_in_cluster_coord_sfb_vmnk[1],
+            sfb_cta_layout,
+            cute.group_modes(sSFB, 0, 3),
+            cute.group_modes(tCgSFB, 0, 3),
+        )
+        tBsSFB = cute.filter_zeros(tBsSFB)
+        tBgSFB = cute.filter_zeros(tBgSFB)
+
+        #
+        # Partition shared/tensor memory tensor for TiledMMA_A/B/C
+        #
+        # (MMA, MMA_M, MMA_K, STAGE)
+        tCrA = tiled_mma.make_fragment_A(sA)
+        # (MMA, MMA_N, MMA_K, STAGE)
+        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)
+        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+
+        #
+        # Alloc tensor memory buffer
+        #
+        tmem_alloc_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=self.threads_per_cta,
+        )
+        tmem = utils.TmemAllocator(
+            storage.tmem_holding_buf.ptr,
+            barrier_for_retrieve=tmem_alloc_barrier,
+            is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
+            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
+        )
+        tmem.allocate(self.num_tmem_alloc_cols)
+        tmem.wait_for_alloc()
+        acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
+        tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+
+        #
+        # Make SFA/SFB tmem tensor
+        #
+        # Get SFA tmem ptr
+        sfa_tmem_ptr = cute.recast_ptr(
+            acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
+            dtype=sf_dtype,
+        )
+        # (MMA, MMA_M, MMA_K)
+        tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+        )
+        tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+        # Get SFB tmem ptr
+        sfb_tmem_ptr = cute.recast_ptr(
+            acc_tmem_ptr
+            + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+            + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
+            dtype=sf_dtype,
+        )
+        # (MMA, MMA_N, MMA_K)
+        tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+            tiled_mma,
+            self.mma_tiler,
+            sf_vec_size,
+            cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
+        )
+        tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+
+        #
+        # Partition for S2T copy of SFA/SFB
+        #
+        # Make S2T CopyAtom
+        copy_atom_s2t = cute.make_copy_atom(
+            tcgen05.Cp4x32x128bOp(
+                tcgen05.CtaGroup.ONE
+                if not self.use_2cta_instrs
+                else tcgen05.CtaGroup.TWO
+            ),
+            sf_dtype,
+        )
+        # (MMA, MMA_MN, MMA_K, STAGE)
+        tCsSFA_compact = cute.filter_zeros(sSFA)
+        # (MMA, MMA_MN, MMA_K)
+        tCtSFA_compact = cute.filter_zeros(tCtSFA)
+        tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
+        thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+            tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
+        )
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
+
+        # (MMA, MMA_MN, MMA_K, STAGE)
+        tCsSFB_compact = cute.filter_zeros(sSFB)
+        # (MMA, MMA_MN, MMA_K)
+        tCtSFB_compact = cute.filter_zeros(tCtSFB)
+        tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
+        thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+        tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+            tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
+        )
+        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
+
+        #
+        # Slice to per mma tile index
+        #
+        # ((atom_v, rest_v), RestK)
+        tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+        # ((atom_v, rest_v), RestK)
+        tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+        # ((atom_v, rest_v), RestK)
+        tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+        # ((atom_v, rest_v), RestK)
+        tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+
+        #
+        # Execute Data copy and Math computation in the k_tile loop
+        #
+        if warp_idx == 0:
+            # Wait for accumulator buffer empty
+            if is_leader_cta:
+                acc_producer.acquire_and_advance()
+
+            # Set ACCUMULATE field to False for the first k_tile iteration
+            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+            # Execute k_tile loop
+            for k_tile in cutlass.range(
+                k_tile_cnt, prefetch_stages=self.num_ab_stage - 2
+            ):
+                # Wait for AB buffer empty
+                ab_empty = ab_producer.acquire_and_advance()
+
+                #  TMA load A/B/SFA/SFB
+                cute.copy(
+                    tma_atom_a,
+                    tAgA[(None, ab_empty.count)],
+                    tAsA[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                    mcast_mask=a_full_mcast_mask,
+                )
+                cute.copy(
+                    tma_atom_b,
+                    tBgB[(None, ab_empty.count)],
+                    tBsB[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                    mcast_mask=b_full_mcast_mask,
+                )
+                cute.copy(
+                    tma_atom_sfa,
+                    tAgSFA[(None, ab_empty.count)],
+                    tAsSFA[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                    mcast_mask=sfa_full_mcast_mask,
+                )
+                cute.copy(
+                    tma_atom_sfb,
+                    tBgSFB[(None, ab_empty.count)],
+                    tBsSFB[(None, ab_empty.index)],
+                    tma_bar_ptr=ab_empty.barrier,
+                    mcast_mask=sfb_full_mcast_mask,
+                )
+
+                if is_leader_cta:
+                    # Wait for AB buffer full
+                    ab_full = ab_consumer.wait_and_advance()
+
+                    #  Copy SFA/SFB to tmem
+                    s2t_stage_coord = (None, None, None, None, ab_full.index)
+                    tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
+                    tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
+                    cute.copy(
+                        tiled_copy_s2t_sfa,
+                        tCsSFA_compact_s2t_staged,
+                        tCtSFA_compact_s2t,
+                    )
+                    cute.copy(
+                        tiled_copy_s2t_sfb,
+                        tCsSFB_compact_s2t_staged,
+                        tCtSFB_compact_s2t,
+                    )
+
+                    # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
+                    num_kblocks = cute.size(tCrA, mode=[2])
+                    for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
+                        kblock_coord = (
+                            None,
+                            None,
+                            kblock_idx,
+                            ab_full.index,
+                        )
+
+                        # Set SFA/SFB tensor to tiled_mma
+                        sf_kblock_coord = (None, None, kblock_idx)
+                        tiled_mma.set(
+                            tcgen05.Field.SFA,
+                            tCtSFA[sf_kblock_coord].iterator,
+                        )
+                        tiled_mma.set(
+                            tcgen05.Field.SFB,
+                            tCtSFB[sf_kblock_coord].iterator,
+                        )
+
+                        cute.gemm(
+                            tiled_mma,
+                            tCtAcc,
+                            tCrA[kblock_coord],
+                            tCrB[kblock_coord],
+                            tCtAcc,
+                        )
+                        # Enable accumulate on tCtAcc after first kblock
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+
+                    # Async arrive AB buffer empty
+                    ab_full.release()
+            if is_leader_cta:
+                acc_producer.commit()
+
+        #
+        # Epilogue
+        # Partition for epilogue
+        #
+        # x32 or x128 all is ok.
+        op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
+        copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
+        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc)
+        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+        # (T2R_M, T2R_N, EPI_M, EPI_M)
+        tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc)
+        # (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
+        tTR_gC = thr_copy_t2r.partition_D(tCgC)
+        # (T2R_M, T2R_N, EPI_M, EPI_N）
+        tTR_rAcc = cute.make_rmem_tensor(
+            tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32
+        )
+        # (T2R_M, T2R_N, EPI_M, EPI_N）
+        tTR_rC = cute.make_rmem_tensor(
+            tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype
+        )
+        # STG Atom
+        simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype)
+        tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)]
+
+        # Wait for accumulator buffer full
+        acc_full = acc_consumer.wait_and_advance()
+
+        # Copy accumulator to register
+        cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc)
+        acc_vec = epilogue_op(tTR_rAcc.load().to(c_dtype))
+        tTR_rC.store(acc_vec)
+        # Store C to global memory
+        cute.copy(simt_atom, tTR_rC, tTR_gC)
+
+        acc_full.release()
+
+        # Ensure used buffers are properly synchronized before producer exit.
+        # This could avoid the invalid dsmem access due to early leading CTA exit.
+        if warp_idx == 0:
+            ab_producer.tail()
+            if is_leader_cta:
+                acc_producer.tail()
+
+        # Deallocate TMEM
+        cute.arch.barrier()
+        tmem.free(acc_tmem_ptr)
+
+        return
+
+
+def run_nvfp4_gemm(
+    mnkl: Tuple[int, int, int, int],
+    tolerance: float,
+    warmup_iterations: int = 10,
+    iterations: int = 100,
+    use_cold_l2: bool = True,
+    do_benchmark: bool = False,
+):
+    run(
+        gemm_class=Sm100BlockScaledDenseGemmKernel,
+        ab_dtype=ab_dtype,
+        sf_dtype=sf_dtype,
+        c_dtype=c_dtype,
+        sf_vec_size=sf_vec_size,
+        mma_tiler_mn=mma_tiler_mn,
+        cluster_shape_mnk=cluster_shape_mnk,
+        mnkl=mnkl,
+        tolerance=tolerance,
+        do_benchmark=do_benchmark,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+        use_cold_l2=use_cold_l2,
+    )
+
+
+if __name__ == "__main__":
+    parser = create_parser()
+    args = parser.parse_args()
+    if len(args.mnkl) != 4:
+        parser.error("--mnkl must contain exactly 4 values")
+
+    m, n, k, _ = args.mnkl
+    if m % mma_tiler_mn[0] != 0:
+        parser.error("M must be multiples of mma_tiler_mn[0] (got m={})".format(m))
+    if n % mma_tiler_mn[1] != 0:
+        parser.error("N must be multiples of mma_tiler_mn[1] (got n={})".format(n))
+    if k % 256 != 0:
+        parser.error("k must be a multiple of 256 (got k={})".format(k))
+
+    run_nvfp4_gemm(
+        args.mnkl,
+        args.tolerance,
+        do_benchmark=args.do_benchmark,
+    )
+    print("PASS")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_gemm/utils.py

@@ -0,0 +1,366 @@
+# Copyright (c) 2025 - 2026 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
+from typing import Tuple
+
+import torch
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.torch as cutlass_torch
+from cutlass.cute.runtime import make_ptr
+
+
+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 create_parser():
+    parser = argparse.ArgumentParser(
+        description="Example of Sm100 Dense BlockScaled GEMM."
+    )
+    parser.add_argument(
+        "--mnkl",
+        type=parse_comma_separated_ints,
+        default=(8192, 8192, 8192, 8),
+        help="mnkl dimensions (comma-separated)",
+    )
+    parser.add_argument(
+        "--tolerance", type=float, default=1e-01, help="Tolerance for validation"
+    )
+    parser.add_argument(
+        "--do_benchmark", action="store_true", default=False, help="Do benchmark test"
+    )
+    return parser
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+# Helper function to create scale factor tensor SFA/SFB
+# for 1x16 block scaled wise use case and follow the layout requirement
+# defined in https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout
+@cute.jit
+def cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
+    sf_ref_ptr: cute.Pointer,
+    sf_mma_ptr: cute.Pointer,
+    mn: int,
+    sf_k: int,
+    l: int,
+    mma_shape: tuple,
+):
+    mma_permute_order = (3, 4, 1, 5, 2, 0)
+    permuted_shape = tuple(mma_shape[i] for i in mma_permute_order)
+    cute_layout = cute.make_ordered_layout(permuted_shape, order=(2, 1, 4, 0, 3, 5))
+
+    sf_ref_tensor = cute.make_tensor(
+        sf_ref_ptr, cute.make_layout((mn, sf_k, l), stride=(sf_k, 1, mn * sf_k))
+    )
+    sf_mma_tensor = cute.make_tensor(sf_mma_ptr, cute_layout)
+
+    sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3)
+    sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3)
+    for i in cutlass.range(cute.size(sf_ref_tensor)):
+        mkl_coord = sf_ref_tensor.layout.get_hier_coord(i)
+        sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord]
+    pass
+
+
+def to_blocked(input_matrix):
+    rows, cols = input_matrix.shape
+
+    # Please ensure rows and cols are multiples of 128 and 4 respectively
+    n_row_blocks = ceil_div(rows, 128)
+    n_col_blocks = ceil_div(cols, 4)
+
+    padded = input_matrix
+    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
+    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
+
+    return rearranged.flatten()
+
+
+def run(
+    gemm_class,
+    ab_dtype,
+    sf_dtype,
+    c_dtype,
+    sf_vec_size,
+    mma_tiler_mn,
+    cluster_shape_mnk,
+    mnkl: Tuple[int, int, int, int],
+    tolerance: float,
+    warmup_iterations: int = 10,
+    iterations: int = 100,
+    use_cold_l2: bool = True,
+    do_benchmark: bool = False,
+):
+    """
+    Prepare A/B/SFA/SFB/C tensors, launch GPU kernel, and reference checking.
+    """
+    print("=" * 60)
+    print("Launching Blackwell Dense BlockScaled GEMM Test")
+    print("-" * 60)
+    print(f"Input dimensions (m, n, k, l): {mnkl}")
+    print(f"    m (rows): {mnkl[0]}")
+    print(f"    n (cols): {mnkl[1]}")
+    print(f"    k (inner): {mnkl[2]}")
+    print(f"    l (batch): {mnkl[3]}")
+    print(f"Data Types & Precision:")
+    print(f"    Input matrices (A, B): {ab_dtype}")
+    print(f"    Scale factors (SFA, SFB): {sf_dtype}")
+    print(f"    Output matrix (C): {c_dtype}")
+    print(f"    Scale factor vector size: {sf_vec_size}")
+    print("Tile and cluster configuration:")
+    print(f"    MMA tiler (M, N, K): {mma_tiler_mn}")
+    print(f"     Cluster shape (M, N, K): {cluster_shape_mnk}")
+    print(f"Validation tolerance: {tolerance}")
+    print(f"Do benchmark: {do_benchmark}")
+    print("=" * 60)
+
+    # Unpack parameters
+    m, n, k, l = mnkl
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("GPU is required to run this example!")
+
+    torch.manual_seed(1111)
+
+    # Create tensor A/B/C
+    a_ref = torch.randint(
+        0, 2, (l, m, k // 2), dtype=torch.uint8, device="cuda"
+    ).permute(1, 2, 0)
+    b_ref = torch.randint(
+        0, 2, (l, n, k // 2), dtype=torch.uint8, device="cuda"
+    ).permute(1, 2, 0)
+    # a_ref = torch.ones((l, m, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0)
+    # b_ref = torch.ones((l, n, k // 2), dtype=torch.uint8, device="cuda").permute(1, 2, 0)
+    a_ref_f4 = a_ref.view(torch.float4_e2m1fn_x2)
+    b_ref_f4 = b_ref.view(torch.float4_e2m1fn_x2)
+
+    c_tensor = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute(
+        1, 2, 0
+    )
+
+    # Create a torch tensor for scale factor tensor of A and B
+    def create_ref_scale_factor_tensor(l, mn, sf_k):
+        """
+        Create the reference scale factor tensor on CPU.
+        Returns the reshaped/pruned tensor ready for ref computation and its original permuted form.
+        """
+        ref_shape = (l, mn, sf_k)
+        ref_permute_order = (1, 2, 0)
+        ref_f8_random_int = torch.randint(1, 3, ref_shape, dtype=torch.int8)
+        ref_f8_torch_tensor_cpu = ref_f8_random_int.to(dtype=torch.float8_e4m3fn)
+        # permute to match ref_permute_order
+        ref_f8_torch_tensor_cpu_permuted = ref_f8_torch_tensor_cpu.permute(
+            *ref_permute_order
+        )
+        return ref_f8_torch_tensor_cpu_permuted
+
+    # Copy the reference scale factor tensor to the CUTE-format scale factor tensor
+    def create_cute_scale_factor_tensor(l, mn, sf_k, ref_f8_torch_tensor_cpu_permuted):
+        """
+        Create the CUTE-format scale factor tensor on CUDA based on the reference tensor.
+        """
+        atom_m = (32, 4)
+        atom_k = 4
+        mma_shape = (
+            l,  # batch size
+            ceil_div(mn, atom_m[0] * atom_m[1]),
+            ceil_div(sf_k, atom_k),
+            atom_m[0],
+            atom_m[1],
+            atom_k,
+        )
+        mma_permute_order = (3, 4, 1, 5, 2, 0)
+
+        # Generate a random int8 tensor, then convert to float8_e4m3fn
+        rand_int_tensor = torch.randint(0, 2, mma_shape, dtype=torch.int8)
+        cute_f8_torch_tensor_cpu = rand_int_tensor.to(dtype=torch.float8_e4m3fn)
+        # Permute according to mma_permute_order
+        cute_f8_torch_tensor_cpu = cute_f8_torch_tensor_cpu.permute(*mma_permute_order)
+
+        # Call the helper function to do layout conversion
+        cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
+            make_ptr(
+                cutlass.Float8E4M3FN,
+                ref_f8_torch_tensor_cpu_permuted.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=32,
+            ),
+            make_ptr(
+                cutlass.Float8E4M3FN,
+                cute_f8_torch_tensor_cpu.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=32,
+            ),
+            mn,
+            sf_k,
+            l,
+            mma_shape,
+        )
+        return cute_f8_torch_tensor_cpu.cuda()
+
+    sf_k = ceil_div(k, sf_vec_size)
+    sfa_ref = create_ref_scale_factor_tensor(l, m, sf_k)
+    sfb_ref = create_ref_scale_factor_tensor(l, n, sf_k)
+    # sfa_ref.fill_(1)
+    # sfb_ref.fill_(1)
+    sfa_tensor = create_cute_scale_factor_tensor(l, m, sf_k, sfa_ref)
+    sfb_tensor = create_cute_scale_factor_tensor(l, n, sf_k, sfb_ref)
+
+    # Configure gemm kernel
+    gemm = gemm_class()
+    # Initialize Stream
+    current_stream = cutlass_torch.default_stream()
+    a_ptr = make_ptr(
+        ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
+    )
+    b_ptr = make_ptr(
+        ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
+    )
+    c_ptr = make_ptr(
+        c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+    )
+    sfa_ptr = make_ptr(
+        sf_dtype, sfa_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+    )
+    sfb_ptr = make_ptr(
+        sf_dtype, sfb_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+    )
+    # Compile gemm kernel
+    compiled_gemm = cute.compile(
+        gemm,
+        a_ptr,
+        b_ptr,
+        sfa_ptr,
+        sfb_ptr,
+        c_ptr,
+        (m, n, k, l),
+        current_stream,
+    )
+    # Launch GPU kernel
+    compiled_gemm(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream)
+    # For batch l, do (m, k, l) @ (n, k, l).T along k for each batch.
+    # Result: (m, n, l)
+    # Allocate ref as (l, m, n) with n-contiguous layout, then permute to (m, n, l)
+    ref = torch.empty(
+        (l, m, n),
+        dtype=torch.float16,
+        device="cuda",
+    ).permute(1, 2, 0)
+    for l_idx in range(l):
+        # Convert the scale factor tensor to blocked format
+        scale_a = to_blocked(sfa_ref[:, :, l_idx])
+        scale_b = to_blocked(sfb_ref[:, :, l_idx])
+        # (m, k) @ (n, k).T -> (m, n)
+        res = torch._scaled_mm(
+            a_ref_f4[:, :, l_idx],
+            b_ref_f4[:, :, l_idx].transpose(0, 1),
+            scale_a.cuda(),
+            scale_b.cuda(),
+            bias=None,
+            out_dtype=torch.float16,
+        )
+        ref[:, :, l_idx] = res
+    torch.testing.assert_close(c_tensor, ref, atol=tolerance, rtol=1e-02)
+
+    if do_benchmark:
+
+        def generate_tensors():
+            a_ptr = make_ptr(
+                ab_dtype, a_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
+            )
+            b_ptr = make_ptr(
+                ab_dtype, b_ref_f4.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
+            )
+            c_ptr = make_ptr(
+                c_dtype, c_tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
+            )
+            sfa_ptr = make_ptr(
+                sf_dtype,
+                sfa_tensor.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=32,
+            )
+            sfb_ptr = make_ptr(
+                sf_dtype,
+                sfb_tensor.data_ptr(),
+                cute.AddressSpace.gmem,
+                assumed_align=32,
+            )
+            args = cute.testing.JitArguments(
+                a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l), current_stream
+            )
+            args.add_to_scope([a_ref_f4, b_ref_f4, sfa_tensor, sfb_tensor, c_tensor])
+            return args
+
+        workspace_count = 1
+        if use_cold_l2:
+            one_workspace_bytes = (
+                a_ref_f4.numel() * a_ref_f4.element_size()
+                + b_ref_f4.numel() * b_ref_f4.element_size()
+                + sfa_tensor.numel() * sfa_tensor.element_size()
+                + sfb_tensor.numel() * sfb_tensor.element_size()
+                + c_tensor.numel() * c_tensor.element_size()
+            )
+            workspace_count = cute.testing.get_workspace_count(
+                one_workspace_bytes, warmup_iterations, iterations
+            )
+
+        # Return execution time in microseconds
+        time = cute.testing.benchmark(
+            compiled_gemm,
+            workspace_generator=generate_tensors,
+            workspace_count=workspace_count,
+            stream=current_stream,
+            warmup_iterations=warmup_iterations,
+            iterations=iterations,
+        )
+
+        print(f"Execution time: {time} us")
+        peta_flops = (4 * m * n * k * l) / (time * 1e-6) / 1e9 / 1000000
+        print(f"FLOPS: {peta_flops} PFLOPS")
+        bytes_transfer = (
+            2 * m * k / 2 * l * a_ref_f4.element_size()
+            + 2 * n * k / 2 * l * b_ref_f4.element_size()
+            + 2 * m * n * l * c_tensor.element_size()
+            + 2 * m * sf_k * l * sfa_tensor.element_size()
+            + 2 * n * sf_k * l * sfb_tensor.element_size()
+        )
+        print(f"Bytes: {bytes_transfer} Bytes")
+        bandwidth = bytes_transfer / time * 1e-3
+        print(f"BW: {bandwidth} GB/s")

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_tma/README.md

@@ -0,0 +1,385 @@
+# CUTLASS Tutorial Examples for Blackwell TMA
+
+## TMA V0: Understanding tma_partition - The Foundation of TMA Operations
+
+This example demonstrates the fundamental building blocks of Tensor Memory Accelerator (TMA) operations on NVIDIA Blackwell (SM100) architecture using CuTe DSL. It focuses on understanding the `tma_partition` interface, which is essential for all TMA operations.
+
+### Key Concepts
+
+*   **TMA Load (Global → Shared)**: Asynchronous bulk data transfer from Global Memory to Shared Memory using TMA hardware
+*   **TMA Store (Shared → Global)**: Asynchronous bulk data transfer from Shared Memory back to Global Memory
+*   **tma_partition**: Core interface that prepares tensors for TMA operations by partitioning them according to TMA atom layout requirements
+*   **group_modes**: Tensor mode grouping to define the TMA atom shape - crucial for proper tma_partition usage
+*   **mbarrier Synchronization**: Hardware barriers (`mbarrier_init`, `mbarrier_arrive`, `mbarrier_wait`) for synchronizing asynchronous TMA operations
+
+### Kernel Architecture
+
+The `Sm100SimpleCopyKernel` performs a simple tile-based copy operation to illustrate TMA fundamentals:
+
+#### Configuration
+
+*   **tile_shape**: Fixed at (128, 128) - Tile dimensions (M, N)
+*   **cluster_shape_mn**: Fixed at (1, 1) - Single CTA execution (no cluster parallelism)
+*   **Shared Memory**: Single buffer sized to hold one tile (tile_m × tile_n elements)
+*   **Synchronization**: Single mbarrier for TMA load completion
+
+#### Key Components
+
+1.  **TMA Descriptor Creation**: Creates TMA atoms (`tma_atom_src`, `tma_atom_dst`) that encapsulate TMA hardware instructions
+2.  **Shared Memory Layout**: Row-major layout `(tile_m, tile_n):(tile_n, 1)` for simplicity
+3.  **Barrier Management**: Single barrier coordinates TMA load completion before processing
+
+### Execution Flow
+
+1.  **Initialization**:
+    *   Allocate shared memory buffer for one tile
+    *   Initialize mbarrier with `elect_one()` to ensure proper synchronization semantics
+    *   Set barrier to expect TMA transaction bytes (`tile_m × tile_n × element_size`)
+    *   Synchronize all threads after barrier initialization
+
+2.  **Tensor Preparation**:
+    ```
+    Tile global tensors into (tile_m, tile_n) blocks
+    Apply group_modes to combine tile dimensions into Mode 0 (TMA atom)
+    
+    Example: gSrc_tiled with shape (128, 128, 4, 2)
+    After group_modes(_, 0, 2): ((128, 128), 4, 2)
+                                  └───┬────┘  └─┬─┘
+                                   Mode 0    Rest modes
+    ```
+
+3.  **TMA Partition**:
+    ```python
+    tAsA, tAgA = tma_partition(
+        tma_atom_src,          # TMA operation atom
+        cta_id=0,              # CTA ID within cluster
+        cta_layout,            # Cluster layout
+        group_modes(smem_tensor, 0, 2),   # SMEM view (Mode 0 = atom)
+        group_modes(gSrc_tiled, 0, 2)     # Global view (Mode 0 = atom)
+    )
+    # Returns:
+    # tAsA: SMEM view with TMA internal layout
+    # tAgA: Global view with shape ((TMA_Layout), grid_m, grid_n)
+    ```
+
+4.  **Tile Selection**:
+    ```python
+    # Select specific tile for this CTA
+    tAgA_cta = tAgA[(None, bidx, bidy)]
+    # None: keep entire TMA atom (Mode 0)
+    # bidx, bidy: index into grid dimensions
+    ```
+
+5.  **TMA Load** (Global → Shared):
+    ```python
+    cute.copy(tma_atom_src, tAgA_cta, tAsA, tma_bar_ptr=barrier_ptr)
+    # Arrive on barrier (producer signals completion)
+    # Wait on barrier (all threads wait for TMA completion)
+    ```
+
+6.  **TMA Store** (Shared → Global):
+    ```python
+    cute.copy(tma_atom_dst, tAsA, tBgB_cta)
+    # Synchronous store completes before kernel exit
+    ```
+
+### Understanding tma_partition
+
+The `tma_partition` function is the key to TMA operations. The tutorial includes comprehensive inline diagrams showing:
+
+*   **Input Tensor Shapes**: How tensors are organized before partitioning
+*   **group_modes Effect**: How mode grouping creates the required atom structure
+*   **Partition Output**: The structure of partitioned tensors for TMA operations
+*   **Indexing Pattern**: How to select CTA-specific tiles from partitioned views
+
+See lines ~155-255 in `tma_v0.py` for detailed visual explanations.
+
+### Configuration Parameters
+
+*   `tile_shape`: Fixed at (128, 128) - Tile dimensions (M, N)
+*   `cluster_shape_mn`: Fixed at (1, 1) - Single CTA execution
+*   `threads_per_cta`: 32 - Single warp (all threads participate in barriers)
+*   `buffer_align_bytes`: 1024 - Shared memory alignment
+
+### Usage
+
+Run the copy kernel with custom matrix dimensions:
+
+```bash
+# Basic usage (default: 512×128 matrix)
+python tma_v0.py
+
+# Custom dimensions
+python tma_v0.py --M 1024 --N 2048
+
+# With custom benchmark iterations
+python tma_v0.py --M 4096 --N 4096 --num_warmup 10 --num_iters 50
+```
+
+#### Example Code
+
+```python
+from tma_v0 import run_tma_copy
+
+# Run copy on 1024×2048 matrix
+run_tma_copy(M=1024, N=2048)
+# Output: Performance metrics and verification result
+```
+
+### Performance Considerations
+
+*   **Single-stage operation**: No pipelining - simple sequential load → store pattern
+*   **Educational focus**: Designed for understanding TMA fundamentals, not peak performance
+*   **Barrier synchronization**: Demonstrates proper mbarrier usage patterns
+*   **Foundation for V1/V2**: Concepts learned here are essential for understanding multi-stage pipelines and warp specialization in V1 and V2
+
+## TMA V1: Matrix Transpose with Producer-Consumer Pattern
+
+This example demonstrates a TMA-based matrix transpose using producer-consumer synchronization with mbarriers on NVIDIA Blackwell (SM100) architecture.
+
+### Key Concepts
+
+*   **Producer-Consumer Pattern**: Different warps coordinate through mbarrier synchronization
+*   **TMA Operations**: Asynchronous bulk data transfer between Global and Shared Memory
+*   **Shared Memory Swizzle**: Optimized layouts to avoid bank conflicts during transpose
+*   **Warp Specialization**: Dedicated warps for loading, transposing, and storing
+*   **mbarrier Synchronization**: Hardware barriers coordinate asynchronous operations
+
+### Kernel Architecture
+
+The `Sm100MatrixTransposeKernel` performs a tiled matrix transpose (M×N → N×M) with the following design:
+
+#### Warp Roles
+
+1.  **TMA Load Warp** (Warp 4, Producer): Issues TMA load operations from Global Memory to Shared Memory buffer `sA`
+2.  **Transpose Warps** (Warps 0-3, 4 warps, Consumer/Producer): 
+    *   Wait for `sA` to be filled by TMA load (consumer of load_mbar)
+    *   Transpose data from `sA` → Registers → `sB`
+    *   Signal completion to TMA Store warp (producer for store_mbar)
+3.  **TMA Store Warp** (Warp 5, Consumer): 
+    *   Wait for `sB` to be ready (consumer of store_mbar)
+    *   Issue TMA store operations from `sB` to Global Memory
+
+#### Synchronization Barriers
+
+The kernel uses two mbarrier instances for producer-consumer coordination:
+
+1.  **load_mbar_ptr**: Synchronizes TMA Load → Transpose Warps
+    *   Producer: TMA Load Warp (arrives after TMA completes)
+    *   Consumer: Transpose Warps (wait before reading `sA`)
+    *   Expected arrivals: 1 (from TMA load warp)
+    *   Expected transactions: `tile_m × tile_n × element_size` bytes
+
+2.  **store_mbar_ptr**: Synchronizes Transpose Warps → TMA Store
+    *   Producer: Transpose Warps (each warp arrives after writing to `sB`)
+    *   Consumer: TMA Store Warp (waits before issuing TMA store)
+    *   Expected arrivals: 4 (one from each transpose warp)
+
+#### Execution Flow
+
+1.  **Initialization**:
+    *   All warps participate in barrier initialization (thread 0 initializes)
+    *   Allocate two shared memory buffers: `sA` (row-major) and `sB` (column-major)
+    *   Create TMA descriptors for source and transposed destination
+    *   Initialize `load_mbar` with expected count of 1 and transaction bytes
+    *   Initialize `store_mbar` with expected count of 4 (number of transpose warps)
+
+2.  **TMA Load Warp** (Producer for Load Pipeline):
+    ```
+    partition source tensor by tile shape
+    issue TMA load: Global[block_tile] → sA
+    arrive on load_mbar to signal completion
+    ```
+
+3.  **Transpose Warps** (Consumer for Load, Producer for Store):
+    ```
+    wait on load_mbar for TMA load to complete
+    
+    partition sA for reading (each thread handles subset)
+    copy data: sA → Registers
+    partition sB for writing
+    copy data: Registers → sB (transpose happens via layout)
+    
+    fence to ensure smem writes are visible
+    synchronize with trans_sync_barrier
+    [elect one thread] arrive on store_mbar to signal completion
+    ```
+
+4.  **TMA Store Warp** (Consumer for Store Pipeline):
+    ```
+    wait on store_mbar for transpose to complete
+    partition destination tensor by tile shape
+    issue TMA store: sB → Global[block_tile]
+    ```
+
+### Key Features
+
+*   **Simple Producer-Consumer Model**: Clear separation of concerns with dedicated warps
+*   **Efficient Synchronization**: Hardware mbarriers minimize synchronization overhead
+*   **Memory Layout Optimization**: Swizzled layouts prevent bank conflicts
+*   **Transposition via Layouts**: Transpose is achieved through different memory layouts for `sA` and `sB`
+
+### Configuration Parameters
+
+*   `tile_shape`: Fixed at (128, 128) - Tile dimensions (M, N)
+*   `cluster_shape_mn`: Fixed at (1, 1) - Single CTA execution
+*   Warp count: 6 warps (1 TMA Load + 4 Transpose + 1 TMA Store)
+
+### Usage
+
+Run the transpose kernel with custom matrix dimensions:
+
+```bash
+# Basic usage (128×128 matrix)
+python tma_v1.py
+
+# Custom dimensions
+python tma_v1.py --M 1024 --N 2048
+
+# With custom benchmark iterations
+python tma_v1.py --M 4096 --N 4096 --num_warmup 10 --num_iters 50
+```
+
+#### Example Code
+
+```python
+from tma_v1 import run_transpose
+
+# Run transpose on 1024×2048 matrix
+run_transpose(M=1024, N=2048)
+# Output: Performance metrics and verification result
+```
+
+### Performance Considerations
+
+*   **Single-stage pipeline**: Simpler than multi-stage but with potential for idle warps
+*   **Warp specialization**: Clear roles minimize synchronization complexity
+*   **Good for learning**: Demonstrates fundamental TMA and mbarrier concepts
+
+## TMA V2: Transpose with Multi-Stage Pipeline
+
+This example demonstrates a TMA implementation with multi-stage pipelining for efficient matrix transpose operations on NVIDIA Blackwell (SM100) architecture.
+
+### Key Concepts
+
+*   **Multi-Stage Pipeline**: Multiple buffers enable overlapping TMA loads, computation (transpose), and TMA stores to hide memory latency
+*   **Pipeline Abstractions**: `PipelineTmaAsync` and `PipelineTmaStore` provide producer-consumer synchronization
+*   **Persistent Tile Scheduler**: Efficient work distribution across CTAs for dynamic load balancing
+*   **Shared Memory Swizzle**: Optimized shared memory layouts to avoid bank conflicts
+*   **Warp Specialization**: Different warps handle loading and transposing; first transpose warp also handles storing
+
+### Kernel Architecture
+
+The `Sm100MatrixTransposeKernelV2` performs a tiled matrix transpose (M×N → N×M) with the following design:
+
+#### Warp Roles
+
+1.  **TMA Load Warp** (Producer): Loads tiles from Global Memory to Shared Memory buffer `sA` using TMA load operations
+2.  **Transpose Warps** (4 warps, Consumer/Producer): 
+    *   Wait for data in `sA` from load pipeline
+    *   Transpose data from `sA` → Registers → `sB`
+    *   Synchronize via named barrier to ensure all transpose warps complete
+    *   First transpose warp (`trans_warp_id[0]`) issues TMA store operations from `sB` to Global Memory
+
+#### Pipeline Stages
+
+The kernel uses two separate pipelines for maximum parallelism:
+
+1.  **Load Pipeline**: `TMA Load Warp` (producer) → `Transpose Warps` (consumer)
+    *   Multi-stage buffer `sA` (automatically computed based on available SMEM)
+    *   Enables prefetching multiple tiles while processing current tile
+2.  **Store Pipeline**: `Transpose Warps` (producer) → `First Transpose Warp` (consumer, issues TMA store)
+    *   Multi-stage buffer `sB` (automatically computed based on available SMEM)
+    *   First transpose warp handles TMA store after all transpose warps finish writing to `sB`
+
+#### Stage Calculation
+
+The kernel automatically computes the optimal number of pipeline stages using `_compute_stages()`:
+
+*   Calculates bytes needed per stage for `sA` (tile_m × tile_n) and `sB` (tile_m × tile_n)
+*   Reserves space for pipeline barriers and metadata (~1KB)
+*   Divides remaining shared memory by bytes per stage
+*   Clamps to 2-8 stages (2 for double buffering minimum, 8 for diminishing returns)
+
+#### Execution Flow
+
+1.  **Initialization**:
+    *   Allocate multi-stage shared memory buffers (`sA_staged`, `sB_staged`)
+    *   Create TMA descriptors for source and transposed destination
+    *   Initialize load and store pipelines with barrier synchronization
+    *   Use persistent tile scheduler to distribute work tiles across CTAs
+
+2.  **TMA Load Warp** (Producer for Load Pipeline):
+    ```
+    for each tile assigned by scheduler:
+        acquire next available stage in load pipeline
+        issue TMA load: Global[tile] → sA[stage]
+        advance to next stage
+    ```
+
+3.  **Transpose Warps** (Consumer for Load, Producer for Store):
+    ```
+    for each tile assigned by scheduler:
+        wait for load pipeline to fill current stage
+        copy data: sA[load_stage] → Registers
+        release load pipeline stage
+        
+        transpose and write: Registers → sB[store_stage]
+        fence to ensure smem writes are visible
+        synchronize all transpose warps with barrier
+        
+        [first transpose warp only] issue TMA store: sB[stage] → Global[tile]
+        [first transpose warp only] commit to store pipeline
+        [first transpose warp only] acquire next store pipeline stage
+        synchronize all transpose warps with barrier
+    ```
+
+4.  **Pipeline Teardown**:
+    *   Producer/consumer tail operations ensure all in-flight operations complete
+
+### Key Features
+
+*   **Automatic Stage Optimization**: Kernel calculates optimal number of stages based on tile size, data type, and available shared memory
+*   **Persistent Tile Scheduler**: Efficient work distribution across CTAs for dynamic load balancing
+*   **Memory Layout Optimization**: Uses appropriate swizzling for row-major input and column-major transposed output
+*   **Efficient Synchronization**: Named barriers for intra-CTA coordination, mbarriers for pipeline stages
+
+### Configuration Parameters
+
+*   `tile_shape`: Fixed at (128, 128) - Tile dimensions (M, N)
+*   `cluster_shape_mn`: Fixed at (1, 1) - Single CTA execution
+*   Number of pipeline stages: Automatically computed based on available SMEM
+
+### Usage
+
+Run the transpose kernel with custom matrix dimensions:
+
+```bash
+# Basic usage (128×128 matrix)
+python tma_v2.py
+
+# Custom dimensions
+python tma_v2.py --M 1024 --N 2048
+
+python tma_v2.py --M 1024 --N 2048 --num_warmup 10 --num_iters 50
+```
+
+#### Example Code
+
+```python
+from tma_v2 import run_transpose
+
+# Run transpose on 1024×2048 matrix
+run_transpose(M=1024, N=2048)
+# Output: "TransposeSuccess!" if verification passes
+```
+
+### Performance Considerations
+
+*   **Multi-stage pipelining** hides memory latency by overlapping loads, computation, and stores
+*   **Persistent scheduling** provides better load balancing for irregular matrix sizes
+*   **Warp specialization** maximizes throughput: TMA load warp handles all loads, transpose warps handle computation, and first transpose warp handles stores
+
+## TMA V3: TMA With MMA (Tensor Cores)
+
+TBD

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_tma/tma_v0.py

@@ -0,0 +1,409 @@
+# Copyright (c) 2024 - 2026 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
+from typing import Type, Union
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cute.runtime import from_dlpack
+import torch
+
+"""
+TMA V0: Understanding tma_partition - The Foundation of TMA Operations
+
+This tutorial demonstrates TMA (Tensor Memory Accelerator) operations through
+a simple copy kernel. It focuses on understanding the fundamental tma_partition
+interface, which is the key to all TMA operations.
+
+What This Tutorial Covers:
+1. TMA Load (Global Memory -> Shared Memory)
+2. TMA Store (Shared Memory -> Global Memory)  
+3. Barrier synchronization for TMA (elect_one, mbarrier_init, mbarrier_arrive, mbarrier_wait)
+4. Detailed explanation of tma_partition with visual diagrams
+
+Key Learning Points:
+- tma_partition: How it transforms tensors for TMA operations
+- group_modes: Why and how to group tensor modes to define TMA atom shape
+- Indexing: How to select specific tiles from partitioned tensors
+- Data flow: Complete visualization from input tensors to TMA copy
+
+Visual Diagrams:
+See line ~155 for comprehensive diagrams showing:
+- Input tensor shapes and transformations
+- group_modes effect on tensor layouts
+- tma_partition output structure
+- Complete data flow from global/shared memory to TMA copy
+- Indexing pattern for CTA-specific tiles
+
+Example Usage:
+```bash
+python cutlass_ir/compiler/python/examples/cute/blackwell/tutorial/tutorial_tma/tma_v0.py
+```
+"""
+
+
+class Sm100SimpleCopyKernel:
+    def __init__(self):
+        """
+        Initializes the configuration for a Blackwell TMA copy kernel.
+        """
+        self.tile_shape = (128, 128)
+        self.tile_m, self.tile_n = self.tile_shape
+        self.cluster_shape_mn = (1, 1)
+        self.threads_per_cta = 32
+        self.buffer_align_bytes = 1024
+
+    @cute.jit
+    def __call__(self, src: cute.Tensor, dst: cute.Tensor):
+        if cutlass.const_expr(src.element_type != dst.element_type):
+            raise TypeError("Source and destination element types must match")
+
+        self.dtype: Type[cutlass.Numeric] = src.element_type
+
+        # layout for each cta: (tile_m, tile_n):(tile_n, 1)
+        smem_layout = cute.make_layout(
+            (self.tile_m, self.tile_n), stride=(self.tile_n, 1)
+        )
+
+        @cute.struct
+        class SharedStorage:
+            barrier_storage: cute.struct.MemRange[cutlass.Int64, 1]
+            smem_data: cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(smem_layout)],
+                self.buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+        self.num_tma_load_bytes = cute.size_in_bytes(self.dtype, smem_layout)
+
+        # cta_tiler: the per-CTA tile extents (M, N) used by TMA.
+        # Note: smem_layout may include swizzle or composed layout,
+        # so we use product_each(...) to take the product along each logical dimension and get
+        # the final (tile_m, tile_n) extents expected by TMA.
+        # In this simple example, smem_layout.shape == (tile_m, tile_n), so product_each(...) is
+        # just (tile_m, tile_n).
+        cta_tiler = cute.product_each(smem_layout.shape)
+
+        tma_atom_src, tma_tensor_src = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(), src, smem_layout, cta_tiler
+        )
+
+        tma_atom_dst, tma_tensor_dst = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(), dst, smem_layout, cta_tiler
+        )
+
+        # Grid shape is now (M/TileM, N/TileN)
+        grid_shape = cute.ceil_div((*src.layout.shape, 1), self.tile_shape)
+
+        self.kernel(
+            tma_atom_src, tma_tensor_src, tma_atom_dst, tma_tensor_dst, smem_layout
+        ).launch(
+            grid=grid_shape,
+            block=(self.threads_per_cta, 1, 1),
+            cluster=(*self.cluster_shape_mn, 1),
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tma_atom_src: cute.CopyAtom,
+        tma_tensor_src: cute.Tensor,
+        tma_atom_dst: cute.CopyAtom,
+        tma_tensor_dst: cute.Tensor,
+        smem_layout: Union[cute.Layout, cute.ComposedLayout],
+    ):
+        bidx, bidy, _ = cute.arch.block_idx()
+
+        # Allocate Shared Memory
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        # Initialize barrier for TMA synchronization
+        barrier_ptr = storage.barrier_storage.data_ptr()
+
+        # Initialize the barrier: elect_one ensures only one thread executes this
+        # Note: We must use elect_one() instead of "if tid == 0" because:
+        # - elect_one() provides proper synchronization semantics
+        # - It ensures all threads are aware that exactly one thread is executing
+        # - It prevents race conditions and provides memory ordering guarantees
+        with cute.arch.elect_one():
+            cute.arch.mbarrier_init(barrier_ptr, 1)
+            cute.arch.mbarrier_expect_tx(barrier_ptr, self.num_tma_load_bytes)
+
+        # Fence ensures init/expect_tx are visible before proceeding
+        cute.arch.mbarrier_init_fence()
+        cute.arch.barrier()
+
+        # Tile the (M, N) tensor: ((TileM, TileN), M/TileM, N/TileN)
+        gSrc_tiled = cute.local_tile(
+            tma_tensor_src, (self.tile_m, self.tile_n), (None, None)
+        )
+        gDst_tiled = cute.local_tile(
+            tma_tensor_dst, (self.tile_m, self.tile_n), (None, None)
+        )
+
+        smem_tensor = storage.smem_data.get_tensor(smem_layout)
+
+        # ======================================================================
+        # TMA Partition: Tensor Preparation for TMA Operations
+        # ======================================================================
+        #
+        # tma_partition prepares tensors for TMA copy by partitioning them
+        # according to the TMA atom's internal layout requirements.
+        #
+        # Signature:
+        #   tma_partition(atom, cta_id, cta_layout, smem_tensor, gmem_tensor)
+        #     -> (smem_view, gmem_view)
+        #
+        # Key Requirement: Mode 0 of both tensors must represent the TMA atom
+        #
+        # Example: M=512, N=128, TileM=128, TileN=64
+        #
+        # Input Tensors:
+        #   gSrc_tiled:  (128, 64, 4, 2)      # 4 separate modes
+        #                 └──┬──┘ └──┬──┘
+        #                   Tile    Grid
+        #
+        #   smem_tensor: (128, 64)            # 2 separate modes
+        #                 └──┬──┘
+        #                   Tile
+        #
+        # Apply group_modes(tensor, 0, 2) to group first 2 modes:
+        #
+        #   group_modes(gSrc_tiled, 0, 2)   =>  ((128, 64), 4, 2)
+        #                                         └───┬───┘
+        #                                         Mode 0 = Atom
+        #
+        #   group_modes(smem_tensor, 0, 2)  =>  ((128, 64),)
+        #                                         └───┬───┘
+        #                                         Mode 0 = Atom
+        #
+        # After tma_partition:
+        #
+        #   tAsA: SMEM view with TMA internal layout
+        #     Shape: ((TMA_Layout),)
+        #     - TMA_Layout: Swizzled/banked layout for efficient SMEM access
+        #
+        #   tAgA: Global view preserving rest modes
+        #     Shape: ((TMA_Layout), 4, 2)
+        #            └─────┬─────┘ └──┬──┘
+        #            TMA atom      Rest modes (grid)
+        #
+        # Usage Pattern:
+        #   1. Group modes to define atom: group_modes(tensor, 0, 2)
+        #   2. Call tma_partition: tAsA, tAgA = tma_partition(...)
+        #   3. Select tile for CTA: tAgA_cta = tAgA[(None, bidx, bidy)]
+        #      - None: keep entire atom
+        #      - bidx, bidy: index into rest modes
+        #   4. Issue TMA copy: cute.copy(atom, tAgA_cta, tAsA)
+        #
+        # Visual Data Flow:
+        #
+        #   Global Memory (512x128)           Shared Memory (128x64)
+        #   ┌────────────────────┐            ┌──────────────┐
+        #   │  ┌───┬───┐         │            │              │
+        #   │  │0,0│0,1│         │            │  smem_tensor │
+        #   │  ├───┼───┤         │            │  (128, 64)   │
+        #   │  │1,0│1,1│ 4x2     │            │              │
+        #   │  ├───┼───┤ tiles   │            └──────────────┘
+        #   │  │2,0│2,1│         │                    │
+        #   │  ├───┼───┤         │                    │ group_modes
+        #   │  │3,0│3,1│         │                    ↓
+        #   │  └───┴───┘         │            ((128, 64),)
+        #   └────────────────────┘                    │
+        #            │                                 │
+        #            │ gSrc_tiled                      │
+        #            │ (128, 64, 4, 2)                │
+        #            ↓                                 │
+        #     group_modes(_, 0, 2)                    │
+        #            ↓                                 │
+        #     ((128, 64), 4, 2)                       │
+        #            │                                 │
+        #            └────────┬────────────────────────┘
+        #                     ↓
+        #              tma_partition
+        #                     ↓
+        #      ┌──────────────┴──────────────┐
+        #      │                             │
+        #   tAgA                          tAsA
+        #   ((TMA_Layout), 4, 2)         ((TMA_Layout),)
+        #      │
+        #      ↓ tAgA[(None, bidx, bidy)]
+        #   tAgA_cta
+        #   ((TMA_Layout),)
+        #
+        # ======================================================================
+
+        # TMA Load partition
+        # Here we only use 1x1 cluster, so cta_id is 0 and cta_layout is (1).
+        # More details about how to set cta_coord and cta_layout can be found in the tma_v4.py
+        # Note: Smem and gemm should have the same size (atom element size) in the first rank
+        tAsA, tAgA = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_src,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(smem_tensor, 0, 2),
+            cute.group_modes(gSrc_tiled, 0, 2),
+        )
+
+        # TMA Store partition
+        # Same process as TMA Load, but for destination tensor
+        # Partitions gDst_tiled and smem_tensor according to TMA Store atom
+        _, tBgB = cute.nvgpu.cpasync.tma_partition(
+            tma_atom_dst,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(smem_tensor, 0, 2),
+            cute.group_modes(gDst_tiled, 0, 2),
+        )
+
+        # Select specific tile for this CTA from partitioned global views
+        # Input:  tAgA with shape ((TMA_Layout), 4, 2)
+        # Output: tAgA_cta with shape ((TMA_Layout),)
+        # The (None, bidx, bidy) indexing:
+        #   - None: keeps the entire TMA atom layout (mode 0)
+        #   - bidx: selects from rest mode 1 (M dimension grid)
+        #   - bidy: selects from rest mode 2 (N dimension grid)
+        tAgA_cta = tAgA[(None, bidx, bidy)]
+        tBgB_cta = tBgB[(None, bidx, bidy)]
+
+        # ---------- TMA Load: Global -> Shared ----------
+        cute.copy(
+            tma_atom_src,
+            tAgA_cta,  # Source (TMA Tensor View)
+            tAsA,  # Dest (SMEM Tensor View)
+            tma_bar_ptr=barrier_ptr,
+        )
+
+        # Signal arrival on the barrier after TMA is issued
+        with cute.arch.elect_one():
+            cute.arch.mbarrier_arrive(barrier_ptr)
+
+        # Wait for TMA to complete
+        cute.arch.mbarrier_wait(barrier_ptr, 0)
+
+        # ---------- TMA Store: Shared -> Global ----------
+        cute.copy(
+            tma_atom_dst,
+            tAsA,  # Source (SMEM Tensor View)
+            tBgB_cta,  # Dest (Global Tensor View)
+        )
+
+
+def run_tma_copy(M, N, num_warmup=5, num_iters=20):
+    """
+    Run TMA copy kernel with performance measurement.
+
+    Args:
+        M: Matrix dimension M
+        N: Matrix dimension N
+        num_warmup: Number of warmup iterations
+        num_iters: Number of timing iterations
+    """
+    # Create tensors with shape (M, N)
+    a = torch.randn((M, N), dtype=torch.float16, device="cuda")
+    b = torch.zeros((M, N), dtype=torch.float16, device="cuda")
+
+    # Notice: We declare N-dimension as the leading dimension should be divisible by 16
+    a_cute = (
+        from_dlpack(a, assumed_align=16)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=16)
+    )
+    b_cute = (
+        from_dlpack(b, assumed_align=16)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=16)
+    )
+
+    copy_kernel = Sm100SimpleCopyKernel()
+    compiled_kernel = cute.compile(copy_kernel, a_cute, b_cute)
+
+    # Warmup runs
+    for _ in range(num_warmup):
+        compiled_kernel(a_cute, b_cute)
+    torch.cuda.synchronize()
+
+    # Timed runs
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    start_event.record()
+    for _ in range(num_iters):
+        compiled_kernel(a_cute, b_cute)
+    end_event.record()
+    torch.cuda.synchronize()
+
+    # Calculate performance metrics
+    elapsed_time_ms = start_event.elapsed_time(end_event)
+    avg_time_ms = elapsed_time_ms / num_iters
+
+    # Calculate throughput
+    # For copy: read M*N elements + write M*N elements
+    bytes_per_element = a.element_size()
+    total_bytes = 2 * M * N * bytes_per_element  # Read + Write
+    throughput_gb_s = (total_bytes / 1e9) / (avg_time_ms / 1000)
+
+    # Print performance metrics
+    print(f"Matrix size: {M}×{N}")
+    print(f"Tile shape: {copy_kernel.tile_shape}")
+    print(f"Average time: {avg_time_ms:.4f} ms")
+    print(f"Throughput: {throughput_gb_s:.2f} GB/s")
+
+    # Verify
+    if torch.allclose(a, b, atol=1e-3):
+        print("Verification: PASSED ✓")
+    else:
+        print("Verification: FAILED ✗")
+        diff = (a - b).abs()
+        print(f"Max diff: {diff.max()}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="TMA V0: Understanding tma_partition - The Foundation of TMA Operations"
+    )
+    parser.add_argument("--M", type=int, default=512, help="Matrix dimension M")
+    parser.add_argument("--N", type=int, default=128, help="Matrix dimension N")
+    parser.add_argument(
+        "--num_warmup", type=int, default=5, help="Number of warmup iterations"
+    )
+    parser.add_argument(
+        "--num_iters", type=int, default=20, help="Number of timing iterations"
+    )
+    args = parser.parse_args()
+
+    run_tma_copy(
+        M=args.M,
+        N=args.N,
+        num_warmup=args.num_warmup,
+        num_iters=args.num_iters,
+    )

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_tma/tma_v1.py

@@ -0,0 +1,462 @@
+# Copyright (c) 2025 - 2026 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
+from typing import Tuple, Type
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cute.runtime import from_dlpack
+import cutlass.pipeline as pipeline
+import torch
+
+"""
+TMA Matrix Transpose with Producer-Consumer Pattern: TMA load -> S2R --> R2S -> TMA store
+
+Warp Roles (Producer-Consumer Pattern):
+- Producer: TMA Load Warp (Warp 4) - Loads from Global to Shared A
+- Consumer: Transpose Warps (Warp 0-3) - Wait for load, then transpose sA -> sB
+- Consumer: TMA Store Warp (Warp 5) - Wait for transpose, then store sB to Global
+
+Synchronization:
+1. load_mbar_ptr: TMA Load (producer) -> Transpose Warps (consumer)
+2. store_mbar_ptr: Transpose Warps (producer) -> TMA Store (consumer)
+
+This demonstrates how different warps can use shared memory barriers to coordinate
+producer-consumer relationships.
+"""
+
+
+class Sm100MatrixTransposeKernelV1:
+    def __init__(self):
+        self.tile_shape = (128, 128)
+        self.tile_m, self.tile_n = self.tile_shape
+        self.cluster_shape_mn = (1, 1)
+        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)
+
+        # Set specialized warp ids based on tile_shape
+        self.num_trans_warps = 4  # Maximum number of transpose warps
+        self.trans_warp_id = tuple(range(self.num_trans_warps))
+        self.tma_load_warp_id = self.num_trans_warps
+        self.tma_store_warp_id = self.num_trans_warps + 1
+        self.threads_per_cta = 32 * len(
+            (self.tma_store_warp_id, self.tma_load_warp_id, *self.trans_warp_id)
+        )
+        self.num_trans_threads = 32 * len(self.trans_warp_id)
+
+        self.trans_tile = (self.tile_shape[0] // self.num_trans_warps, 8)
+
+        # Set barriers for producer-consumer sync
+        # Barrier 1: Trans warps sync (for internal coordination)
+        self.trans_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=32 * len(self.trans_warp_id),
+        )
+        # Barrier 2: TMA Store warp waits after Trans warps finish
+        self.store_barrier = pipeline.NamedBarrier(
+            barrier_id=2,
+            num_threads=32,  # Only TMA store warp
+        )
+        self.buffer_align_bytes = 1024
+
+    @cute.jit
+    def __call__(self, src: cute.Tensor, dst: cute.Tensor):
+        if cutlass.const_expr(src.element_type != dst.element_type):
+            raise TypeError("Source and destination element types must match")
+
+        self.dtype: Type[cutlass.Numeric] = src.element_type
+
+        # Create transposed view of dst for TMA descriptor
+        # dst is (N, M), we want to view it as (M, N) transposed
+        transed_dst = cute.make_tensor(
+            dst.iterator,
+            cute.make_layout(
+                (dst.shape[1], dst.shape[0]), stride=(dst.stride[1], dst.stride[0])
+            ),
+        )
+
+        # row-major smem layout for sA (tile_m, tile_n)
+        smem_layout_sA = sm100_utils.make_smem_layout(
+            utils.LayoutEnum.from_tensor(src).mma_major_mode(),
+            (self.tile_m, self.tile_n),
+            self.dtype,
+            1,
+        )
+
+        # col-major smem layout for sB (tile_n, tile_m)
+        # sB should match the transposed destination layout
+        smem_layout_sB = sm100_utils.make_smem_layout(
+            utils.LayoutEnum.from_tensor(transed_dst).mma_major_mode(),
+            (self.tile_m, self.tile_n),
+            self.dtype,
+            1,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            # Barrier for TMA Load: producer (TMA) -> consumer (Trans warps)
+            load_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1]
+            # Barrier for TMA Store: producer (Trans warps) -> consumer (TMA Store)
+            store_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1]
+            # Single shared memory buffer (sA and sB are different views of this)
+            sA: cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(smem_layout_sA)], 128
+            ]
+            sB: cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(smem_layout_sB)], 128
+            ]
+
+        self.shared_storage = SharedStorage
+        self.num_tma_load_bytes = cute.size_in_bytes(self.dtype, smem_layout_sA)
+
+        # TMA Atoms
+        # Use swizzled layout for TMA atom to handle swizzling during load
+        tma_atom_src, tma_tensor_src = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileG2SOp(),
+            src,
+            smem_layout_sA,
+            (self.tile_m, self.tile_n),
+        )
+
+        tma_atom_dst, tma_tensor_dst = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(),
+            transed_dst,
+            smem_layout_sB,
+            (self.tile_m, self.tile_n),
+        )
+
+        grid_shape = cute.ceil_div((*src.layout.shape, 1), self.tile_shape)
+        self.kernel(
+            tma_atom_src,
+            tma_tensor_src,
+            tma_atom_dst,
+            tma_tensor_dst,
+            smem_layout_sA,
+            smem_layout_sB,
+        ).launch(
+            grid=grid_shape,
+            block=(self.threads_per_cta, 1, 1),
+            cluster=self.cluster_shape_mnk,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tma_atom_load: cute.CopyAtom,
+        tma_tensor_src: cute.Tensor,
+        tma_atom_store: cute.CopyAtom,
+        tma_tensor_dst: cute.Tensor,
+        smem_layout_sA: cute.ComposedLayout,
+        smem_layout_sB: cute.ComposedLayout,
+    ):
+        bidx, bidy, _ = cute.arch.block_idx()
+        tidx, _, _ = cute.arch.thread_idx()
+
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+        # Allocate Shared Memory
+        # We need two buffers for transpose:
+        # sA: Source Tile (swizzled)
+        # sB: Destination Tile (swizzled) for TMA Store
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+        sA = storage.sA.get_tensor(smem_layout_sA.outer, swizzle=smem_layout_sA.inner)
+        # sA = cute.make_tensor(storage.sA.iterator, smem_layout_sA)
+        sB = storage.sB.get_tensor(smem_layout_sB.outer, swizzle=smem_layout_sB.inner)
+        self.num_tma_load_bytes = cute.size_in_bytes(self.dtype, smem_layout_sA)
+        load_mbar_ptr = storage.load_mbar_ptr.data_ptr()
+        store_mbar_ptr = storage.store_mbar_ptr.data_ptr()
+
+        # ------------------------------------------------------------------
+        # Initialize Barriers (all warps participate in initialization)
+        # ------------------------------------------------------------------
+        if tidx == 0:
+            # Barrier for TMA Load: expect 1 arrive (from TMA warp after TMA completes)
+            cute.arch.mbarrier_init(load_mbar_ptr, 1)
+            cute.arch.mbarrier_expect_tx(load_mbar_ptr, self.num_tma_load_bytes)
+
+            # Barrier for TMA Store: expect arrival from Trans warps
+            cute.arch.mbarrier_init(store_mbar_ptr, len(self.trans_warp_id))
+        cute.arch.mbarrier_init_fence()
+
+        # Sync all warps after barrier initialization
+        cute.arch.barrier()
+        # ------------------------------------------------------------------
+        # PRODUCER: TMA Load Warp (G -> sA)
+        # ------------------------------------------------------------------
+        if warp_idx == self.tma_load_warp_id:
+            # Issue TMA Load
+            # ((TileM, TileK), loopM, LoopK)
+            gA = cute.local_tile(tma_tensor_src, self.tile_shape, (None, None))
+            # ((TileM, TileK), loopM, LoopK)
+            tAsA, tAgA = cpasync.tma_partition(
+                tma_atom_load,
+                0,
+                cute.make_layout(1),
+                cute.group_modes(sA, 0, 2),
+                cute.group_modes(gA, 0, 2),
+            )
+
+            cute.copy(
+                tma_atom_load,
+                tAgA[(None, bidx, bidy)],
+                tAsA[(None, 0)],
+                tma_bar_ptr=load_mbar_ptr,
+            )
+
+            # Arrive on mbarrier to satisfy the init count of 1
+            with cute.arch.elect_one():
+                cute.arch.mbarrier_arrive(load_mbar_ptr)
+
+        # ------------------------------------------------------------------
+        # CONSUMER: Transpose Warps (sA -> Reg -> sB)
+        # ------------------------------------------------------------------
+        if warp_idx < self.tma_load_warp_id:
+            trans_tid = tidx % self.num_trans_threads
+
+            # Wait for TMA Load to complete (consumer wait on load_mbar)
+            cute.arch.mbarrier_wait(load_mbar_ptr, 0)
+
+            atom = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(),
+                self.dtype,
+                num_bits_per_copy=self.dtype.width,  # Copy one element at a time
+            )
+
+            copy_elems = 1
+
+            # Use SAME thread layout for both read and write
+            # Transpose happens through sB_transposed layout view
+
+            # TV layout notation: T = thread-id, V = value-lane within that thread.
+            # In this example `copy_elems = 1` and `thread_layout` has shape (T, V) = (num_trans_threads, 1),
+            # so V is always 0 (only one value-lane per thread).
+            #
+            # Linearization rule (row-major by strides):
+            #   idx(Ti, Vj) = i + j * num_trans_threads
+            # Therefore here:
+            #   idx(Ti, V0) = i
+            #
+            # Diagram (V is the column, T is the row; showing the first two threads):
+            #
+            #            V0
+            #          ┌──────┐
+            #   T0     │ T0V0 │  -> idx 0
+            #   T1     │ T1V0 │  -> idx 1
+            #   ...    │  ... │
+            #          └──────┘
+            #
+            thread_layout = cute.make_layout(
+                (self.num_trans_threads, 1),
+                stride=(1, self.num_trans_threads),
+            )
+            value_layout = cute.make_layout((1, copy_elems))
+            # Build a "tiled copy" operator that defines the per-thread copy mapping (T,V) for this warp-group:
+            # - It is used twice below via `thr_copy.partition_S(...)` and `thr_copy.partition_D(...)` to
+            #   create matching per-thread views of the source and destination tensors.
+            # - With the SAME (T,V) mapping, the actual transpose is achieved by changing the tensor view
+            #   (`sA` vs `sB` / `sB_transposed`), not by changing which threads perform the copies.
+            tiled_copy = cute.make_tiled_copy_tv(atom, thread_layout, value_layout)
+            thr_copy = tiled_copy.get_slice(trans_tid)
+
+            # Partition sA (source) for reading
+            tCsA = thr_copy.partition_S(sA)
+            # When to use `tiled_copy.retile(...)`:
+            # - Use it when you allocate/build a register tensor yourself (or slice/reshape it) and its
+            #   internal layout doesn't match the TV layout expected by `tiled_copy` for copy-in/out.
+            # Why we don't use it here:
+            # - `cute.make_fragment_like(tCsA)` creates an rmem fragment with the same per-thread shape/layout
+            #   as `tCsA`, so it already matches `tiled_copy`'s view and can be copied into directly.
+            tCrA = cute.make_fragment_like(tCsA)
+            cute.copy(tiled_copy, tCsA, tCrA)
+
+            # Partition sB for writing
+            tCsB = thr_copy.partition_D(sB)
+
+            # Write from register to sB
+            cute.copy(tiled_copy, tCrA, tCsB)
+
+            # Fence and barrier to make sure shared memory store is visible to TMA store
+            cute.arch.fence_proxy(
+                "async.shared",
+                space="cta",
+            )
+            self.trans_sync_barrier.arrive_and_wait()
+
+            # Trans warps signal TMA Store warp: "sB is ready!"
+            with cute.arch.elect_one():
+                cute.arch.mbarrier_arrive(store_mbar_ptr)
+
+        # ------------------------------------------------------------------
+        # CONSUMER: TMA Store Warp (sB -> G)
+        # ------------------------------------------------------------------
+        if warp_idx == self.tma_store_warp_id:
+            # Wait for Trans warp to complete (consumer wait on store_mbar)
+            cute.arch.mbarrier_wait(store_mbar_ptr, 0)
+
+            gDst_cta = cute.local_tile(
+                tma_tensor_dst, (self.tile_m, self.tile_n), (None, None)
+            )
+            tBsB, tBgB = cpasync.tma_partition(
+                tma_atom_store,
+                0,
+                cute.make_layout(1),
+                cute.group_modes(sB, 0, 2),
+                cute.group_modes(gDst_cta, 0, 2),
+            )
+            cute.copy(tma_atom_store, tBsB[(None, 0)], tBgB[(None, bidx, bidy)])
+
+
+def run_transpose(M, N, num_warmup=5, num_iters=20):
+    """
+    Run TMA transpose kernel with performance measurement.
+
+    Args:
+        M: Matrix dimension M
+        N: Matrix dimension N
+        num_warmup: Number of warmup iterations
+        num_iters: Number of timing iterations
+
+    Performance Metrics:
+        - Throughput: Actual achieved bandwidth in GB/s
+        - Theoretical BW: Peak memory bandwidth (2048 B/clk × 4000 MHz = 8.192 TB/s)
+        - Bandwidth Efficiency: Percentage of theoretical peak achieved
+    """
+    torch.manual_seed(1111)
+    # Input (M, N)
+    input_data = torch.randn((M, N), device="cuda", dtype=torch.float16)
+    # Output (N, M)
+    output_data = torch.zeros((N, M), device="cuda", dtype=torch.float16)
+
+    # CuTe Wrappers
+    tensor_src = (
+        from_dlpack(input_data, assumed_align=16)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=16)
+    )
+    tensor_dst = (
+        from_dlpack(output_data, assumed_align=16)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=16)
+    )
+
+    transpose_kernel = Sm100MatrixTransposeKernelV1()
+
+    print("Start kernel compilation...")
+
+    # Compile and Run
+    compiled_kernel = cute.compile(
+        transpose_kernel, tensor_src, tensor_dst, options="--generate-line-info"
+    )
+
+    print("Start kernel warmup...")
+    # Warmup runs
+    for _ in range(num_warmup):
+        compiled_kernel(tensor_src, tensor_dst)
+    torch.cuda.synchronize()
+    print("Kernel warmup completed.")
+
+    # Timed runs
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    start_event.record()
+    for _ in range(num_iters):
+        compiled_kernel(tensor_src, tensor_dst)
+    end_event.record()
+    torch.cuda.synchronize()
+
+    # Calculate performance metrics
+    elapsed_time_ms = start_event.elapsed_time(end_event)
+    avg_time_ms = elapsed_time_ms / num_iters
+
+    # Calculate throughput
+    # For transpose: read M*N elements + write M*N elements
+    bytes_per_element = input_data.element_size()
+    total_bytes = 2 * M * N * bytes_per_element  # Read + Write
+    throughput_gb_s = (total_bytes / 1e9) / (avg_time_ms / 1000)
+
+    # Theoretical bandwidth limit
+    # Blackwell: 2048 B/clk at 4000 MHz
+    bytes_per_clk = 2048
+    freq_mhz = 4000
+    theoretical_bw_gb_s = bytes_per_clk * freq_mhz * 1e6 / 1e9  # Convert to GB/s
+    theoretical_bw_tb_s = theoretical_bw_gb_s / 1000  # Convert to TB/s
+    bandwidth_efficiency = (throughput_gb_s / theoretical_bw_gb_s) * 100  # Percentage
+
+    # Print performance metrics
+    print(f"Matrix size: {M}×{N}")
+    print(f"Tile shape: {transpose_kernel.tile_shape}")
+    print(f"Average time: {avg_time_ms:.4f} ms")
+    print(f"Throughput: {throughput_gb_s:.2f} GB/s")
+    print(
+        f"Theoretical BW: {theoretical_bw_tb_s:.2f} TB/s ({theoretical_bw_gb_s:.2f} GB/s)"
+    )
+    print(f"Bandwidth Efficiency: {bandwidth_efficiency:.2f}%")
+
+    # Verification
+    expected = input_data.t()
+    if torch.allclose(output_data, expected, atol=1e-2):
+        print("Verification: PASSED ✓")
+    else:
+        print("Verification: FAILED ✗")
+        print(f"Max diff: {(output_data - expected).abs().max()}")
+
+
+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."
+            )
+
+    parser = argparse.ArgumentParser(
+        description="TMA Matrix Transpose with Producer-Consumer Pattern"
+    )
+    parser.add_argument("--M", type=int, default=128, help="Matrix dimension M")
+    parser.add_argument("--N", type=int, default=128, help="Matrix dimension N")
+    parser.add_argument(
+        "--num_warmup", type=int, default=5, help="Number of warmup iterations"
+    )
+    parser.add_argument(
+        "--num_iters", type=int, default=20, help="Number of timing iterations"
+    )
+    args = parser.parse_args()
+
+    run_transpose(
+        args.M,
+        args.N,
+        num_warmup=args.num_warmup,
+        num_iters=args.num_iters,
+    )

examples/python/CuTeDSL/cute/blackwell/tutorial/tutorial_tma/tma_v2.py

@@ -0,0 +1,648 @@
+# Copyright (c) 2025 - 2026 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
+from typing import Tuple, Type, Union
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as utils
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cute.runtime import from_dlpack
+import cutlass.pipeline as pipeline
+from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+import torch
+
+"""
+TMA Matrix Transpose with Multi-Stage Pipeline(v2)
+
+This version extends tma_v1.py with:
+1. Multi-stage pipeline: Multiple buffers for pipelining TMA loads and stores
+2. Pipeline abstraction: Using PipelineTmaAsync for proper producer-consumer coordination
+3. Persistent tile scheduler: Efficient work distribution across CTAs
+
+Key Improvements over v1:
+- Multi-stage buffers enable overlapping TMA loads, computation, and TMA stores
+- Pipeline objects provide cleaner synchronization semantics
+
+Note: TMA multicast is NOT used because each CTA must process different input tiles.
+
+Warp Roles:
+- Producer: TMA Load Warp - Loads from Global to Shared memory (multi-stage, multi-tile)
+- Consumer: Transpose Warps - Wait for load, transpose sA -> sB (multi-tile)
+- Consumer: TMA Store Warp - Wait for transpose, store sB to Global (multi-stage, multi-tile)
+
+Pipeline Stages:
+1. Load Pipeline: TMA Load (producer) -> Transpose Warps (consumer)
+2. Store Pipeline: Transpose Warps (producer) -> TMA Store (consumer)
+"""
+
+
+class Sm100MatrixTransposeKernelV2:
+    def __init__(
+        self,
+    ):
+        """
+        Initialize the TMA transpose kernel with multi-stage pipeline support.
+
+        Args:
+            tile_shape: Tile dimensions (M, N)
+            cluster_shape_mn: Cluster shape for parallel CTA execution (M, N)
+
+        Note:
+            - Each CTA processes different tiles independently
+            - Stage counts are automatically computed based on available shared memory
+            - Persistent scheduler distributes work across CTAs in the cluster
+        """
+        self.tile_shape = (128, 128)
+        self.tile_m, self.tile_n = self.tile_shape
+        self.cluster_shape_mn = (1, 1)
+        self.cluster_shape_mnl = (*self.cluster_shape_mn, 1)
+
+        # Set specialized warp ids based on tile_shape
+        # For 128x128 tile, use 4 transpose warps (same as v1)
+        self.max_trans_warps = 4  # Maximum number of transpose warps
+        self.num_trans_warps = self.max_trans_warps  # Use all transpose warps
+        self.trans_warp_id = tuple(range(self.num_trans_warps))
+        self.tma_load_warp_id = self.num_trans_warps
+        self.threads_per_cta = 32 * len((self.tma_load_warp_id, *self.trans_warp_id))
+        self.num_trans_threads = 32 * len(self.trans_warp_id)
+
+        # Set barriers for producer-consumer sync
+        # Barrier 1: Trans warps sync (for internal coordination)
+        self.trans_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=32 * len(self.trans_warp_id),
+        )
+        self.buffer_align_bytes = 128
+
+        # Get shared memory capacity for stage computation
+        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
+
+    @staticmethod
+    def _compute_stages(
+        tile_m: int,
+        tile_n: int,
+        dtype: Type[cutlass.Numeric],
+        smem_capacity: int,
+    ) -> Tuple[int, int]:
+        """
+        Compute the number of load and store stages based on shared memory capacity.
+
+        Strategy:
+        1. Calculate bytes per stage for load (sA) and store (sB) buffers
+        2. Reserve space for barriers and alignment
+        3. Divide remaining smem by bytes per stage to get max stages
+        4. Clamp to reasonable min/max values
+
+        Args:
+            tile_m: Tile dimension M
+            tile_n: Tile dimension N
+            dtype: Data type of the tensors
+            smem_capacity: Total shared memory capacity in bytes
+
+        Returns:
+            Tuple of (num_load_stages, num_store_stages)
+        """
+        # Calculate bytes per tile (assuming row-major and col-major layouts)
+        bytes_per_element = dtype.width // 8
+
+        # For sA (load buffer): tile_m x tile_n elements
+        sA_bytes_per_stage = tile_m * tile_n * bytes_per_element
+
+        # For sB (store buffer): tile_m x tile_n elements (transposed)
+        sB_bytes_per_stage = tile_m * tile_n * bytes_per_element
+
+        # Reserve space for barriers and other metadata
+        # Each barrier: 8 bytes (Int64)
+        # Estimate: max 16 barriers (load + store stages * 2) + alignment
+        reserved_bytes = 1024  # Conservative estimate
+
+        # Available space for staging buffers
+        available_smem = smem_capacity - reserved_bytes
+
+        # Calculate max stages we can fit
+        # We need space for both load and store stages
+        total_bytes_per_stage_pair = sA_bytes_per_stage + sB_bytes_per_stage
+
+        # Max stages (same for load and store for simplicity)
+        max_stages = available_smem // total_bytes_per_stage_pair
+
+        # Clamp to reasonable values
+        # Min: 2 stages for basic double buffering
+        # Max: 8 stages (diminishing returns beyond this)
+        num_stages = max(2, min(max_stages, 8))
+
+        return num_stages, num_stages
+
+    @staticmethod
+    def _compute_grid(
+        c: cute.Tensor,
+        cta_tile_shape_mn: Tuple[int, int],
+        cluster_shape_mn: Tuple[int, int],
+        max_active_clusters: cutlass.Constexpr,
+    ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]:
+        """Use persistent tile scheduler to compute the grid size for the output tensor C.
+
+        :param c: The output tensor C
+        :type c: cute.Tensor
+        :param cta_tile_shape_mn: The shape (M, N) of the CTA tile.
+        :type cta_tile_shape_mn: tuple[int, int]
+        :param cluster_shape_mn: Shape of each cluster in M, N dimensions.
+        :type cluster_shape_mn: tuple[int, int]
+        :param max_active_clusters: Maximum number of active clusters.
+        :type max_active_clusters: cutlass.Constexpr
+
+        :return: A tuple containing:
+            - tile_sched_params: Parameters for the persistent tile scheduler.
+            - grid: Grid shape for kernel launch.
+        :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]
+        """
+        c_shape = cute.slice_(cta_tile_shape_mn, (None, None))
+        gc = cute.zipped_divide(c, tiler=c_shape)
+        num_ctas_mn = gc[(0, (None, None))].shape
+        cluster_shape_mnl = (*cluster_shape_mn, 1)
+        num_ctas_mnl = (*num_ctas_mn, 1)
+        tile_sched_params = utils.PersistentTileSchedulerParams(
+            num_ctas_mnl, cluster_shape_mnl
+        )
+        grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+            tile_sched_params, max_active_clusters
+        )
+
+        return tile_sched_params, grid
+
+    @cute.jit
+    def __call__(
+        self, src: cute.Tensor, dst: cute.Tensor, max_active_clusters: cutlass.Constexpr
+    ):
+        if cutlass.const_expr(src.element_type != dst.element_type):
+            raise TypeError("Source and destination element types must match")
+
+        self.dtype: Type[cutlass.Numeric] = src.element_type
+
+        # Compute optimal stage counts based on tile size and dtype
+        self.num_load_stages, self.num_store_stages = self._compute_stages(
+            self.tile_m,
+            self.tile_n,
+            self.dtype,
+            self.smem_capacity,
+        )
+
+        # Create transposed view of dst for TMA descriptor
+        # dst is (N, M), we want to view it as (M, N) transposed
+        transed_dst = cute.make_tensor(
+            dst.iterator,
+            cute.make_layout(
+                (dst.shape[1], dst.shape[0]), stride=(dst.stride[1], dst.stride[0])
+            ),
+        )
+
+        # Create multi-stage layouts for load and store buffers
+        # row-major smem layout for sA (tile_m, tile_n)
+        smem_layout_sA_staged = sm100_utils.make_smem_layout(
+            utils.LayoutEnum.from_tensor(src).mma_major_mode(),
+            (self.tile_m, self.tile_n),
+            self.dtype,
+            self.num_load_stages,
+        )
+
+        # col-major smem layout for sB (tile_n, tile_m)
+        # sB should match the transposed destination layout
+        smem_layout_sB_staged = sm100_utils.make_smem_layout(
+            utils.LayoutEnum.from_tensor(transed_dst).mma_major_mode(),
+            (self.tile_m, self.tile_n),
+            self.dtype,
+            self.num_store_stages,
+        )
+
+        @cute.struct
+        class SharedStorage:
+            # Pipeline barriers for multi-stage load
+            load_full_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.num_load_stages
+            ]
+            load_empty_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.num_load_stages
+            ]
+
+            # Pipeline barriers for multi-stage store
+            store_full_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.num_store_stages
+            ]
+            store_empty_mbar_ptr: cute.struct.MemRange[
+                cutlass.Int64, self.num_store_stages
+            ]
+
+            # Multi-stage shared memory buffers
+            sA: cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(smem_layout_sA_staged)],
+                self.buffer_align_bytes,
+            ]
+            sB: cute.struct.Align[
+                cute.struct.MemRange[self.dtype, cute.cosize(smem_layout_sB_staged)],
+                self.buffer_align_bytes,
+            ]
+
+        self.shared_storage = SharedStorage
+        a_smem_layout = cute.slice_(smem_layout_sA_staged, (None, None, 0))
+        self.num_tma_load_bytes = cute.size_in_bytes(self.dtype, a_smem_layout)
+
+        # TMA Atoms
+        # Each CTA loads its own tile independently
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+
+        cluster_layout_vmnk = cute.tiled_divide(
+            cute.make_layout((*self.cluster_shape_mn, 1)), (1,)
+        )
+
+        tma_atom_src, tma_tensor_src = cpasync.make_tiled_tma_atom(
+            tma_load_op,
+            src,
+            smem_layout_sA_staged,
+            (self.tile_m, self.tile_n),
+        )
+
+        tma_atom_dst, tma_tensor_dst = cpasync.make_tiled_tma_atom(
+            cpasync.CopyBulkTensorTileS2GOp(),
+            transed_dst,
+            smem_layout_sB_staged,
+            (self.tile_m, self.tile_n),
+        )
+        tile_sched_params, grid_shape = self._compute_grid(
+            transed_dst,
+            (self.tile_m, self.tile_n),
+            self.cluster_shape_mn,
+            max_active_clusters,
+        )
+
+        self.kernel(
+            tma_atom_src,
+            tma_tensor_src,
+            tma_atom_dst,
+            tma_tensor_dst,
+            smem_layout_sA_staged,
+            smem_layout_sB_staged,
+            cluster_layout_vmnk,
+            tile_sched_params,
+        ).launch(
+            grid=grid_shape,
+            block=(self.threads_per_cta, 1, 1),
+            cluster=self.cluster_shape_mnl,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        tma_atom_load: cute.CopyAtom,
+        tma_tensor_src: cute.Tensor,
+        tma_atom_store: cute.CopyAtom,
+        tma_tensor_dst: cute.Tensor,
+        smem_layout_sA_staged: Union[cute.Layout, cute.ComposedLayout],
+        smem_layout_sB_staged: Union[cute.Layout, cute.ComposedLayout],
+        cluster_layout_vmnk: cute.Layout,
+        tile_sched_params: utils.PersistentTileSchedulerParams,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+        # ---------------- Shared mem & staged buffers ----------------
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage)
+
+        sA_staged = storage.sA.get_tensor(
+            smem_layout_sA_staged.outer, swizzle=smem_layout_sA_staged.inner
+        )
+        sB_staged = storage.sB.get_tensor(
+            smem_layout_sB_staged.outer, swizzle=smem_layout_sB_staged.inner
+        )
+
+        load_mbar_ptr = storage.load_full_mbar_ptr.data_ptr()
+        _store_mbar_ptr = storage.store_full_mbar_ptr.data_ptr()
+
+        load_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        load_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_trans_warps
+        )
+
+        load_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=load_mbar_ptr,
+            num_stages=self.num_load_stages,
+            producer_group=load_producer_group,
+            consumer_group=load_consumer_group,
+            tx_count=self.num_tma_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+
+        store_producer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self.num_trans_threads
+        )
+        store_pipeline = pipeline.PipelineTmaStore.create(
+            num_stages=self.num_store_stages,
+            producer_group=store_producer_group,
+        )
+
+        # Critical: Initialize pipeline barriers across cluster
+        # This must happen after pipeline creation and before any producer/consumer work
+        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
+        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
+
+        gA = cute.local_tile(tma_tensor_src, self.tile_shape, (None, None))
+
+        _cta_layout = cute.make_layout(
+            cute.slice_(cluster_layout_vmnk, (0, None, None, 0)).shape
+        )
+        # ((TileM, TileK), loopM, LoopK)
+        tAsA, tAgA = cpasync.tma_partition(
+            tma_atom_load,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sA_staged, 0, 2),
+            cute.group_modes(gA, 0, 2),
+        )
+
+        gDst_cta = cute.local_tile(tma_tensor_dst, self.tile_shape, (None, None))
+        tBsB, tBgB = cpasync.tma_partition(
+            tma_atom_store,
+            0,
+            cute.make_layout(1),
+            cute.group_modes(sB_staged, 0, 2),
+            cute.group_modes(gDst_cta, 0, 2),
+        )
+
+        # ------------------------------------------------------------------
+        # PRODUCER: TMA Load Warp (G -> sA)
+        # ------------------------------------------------------------------
+        if warp_idx == self.tma_load_warp_id:
+            tile_sched = utils.StaticPersistentTileScheduler.create(
+                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+            )
+            work_tile = tile_sched.initial_work_tile_info()
+
+            load_producer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Producer, self.num_load_stages
+            )
+
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+                tAgA_slice = tAgA[(None, cur_tile_coord[0], cur_tile_coord[1])]
+
+                load_pipeline.producer_acquire(load_producer_state)
+
+                cute.copy(
+                    tma_atom_load,
+                    tAgA_slice,
+                    tAsA[(None, load_producer_state.index)],
+                    tma_bar_ptr=load_pipeline.producer_get_barrier(load_producer_state),
+                )
+
+                load_producer_state.advance()
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+            load_pipeline.producer_tail(load_producer_state)
+
+        # ------------------------------------------------------------------
+        # CONSUMER: Transpose Warps (sA -> Reg -> sB)
+        # ------------------------------------------------------------------
+        if warp_idx < self.tma_load_warp_id:
+            tile_sched = utils.StaticPersistentTileScheduler.create(
+                tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+            )
+            work_tile = tile_sched.initial_work_tile_info()
+
+            trans_tid = tidx % self.num_trans_threads
+
+            load_consumer_state = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_load_stages
+            )
+
+            while work_tile.is_valid_tile:
+                # Get tile coord from tile scheduler
+                cur_tile_coord = work_tile.tile_idx
+
+                # Wait for load pipeline to have data
+                load_pipeline.consumer_wait(load_consumer_state)
+
+                atom = cute.make_copy_atom(
+                    cute.nvgpu.CopyUniversalOp(),
+                    self.dtype,
+                    num_bits_per_copy=self.dtype.width,
+                )
+                copy_elems = 1
+
+                thread_layout = cute.make_layout(
+                    (self.num_trans_threads, 1),
+                    stride=(1, self.num_trans_threads),
+                )
+                value_layout = cute.make_layout((1, copy_elems))
+                tiled_copy = cute.make_tiled_copy_tv(atom, thread_layout, value_layout)
+                thr_copy = tiled_copy.get_slice(trans_tid)
+
+                # sA -> Reg
+                tCsA = thr_copy.partition_S(sA_staged)
+
+                tCrA = cute.make_rmem_tensor(
+                    tCsA[(None, None, None, 0)].shape, self.dtype
+                )
+                tCrA = tiled_copy.retile(tCrA)
+
+                cute.copy(
+                    tiled_copy,
+                    tCsA[(None, None, None, load_consumer_state.index)],
+                    tCrA,
+                )
+
+                # release load pipeline
+                load_pipeline.consumer_release(load_consumer_state)
+                load_consumer_state.advance()
+
+                index = tile_sched.num_tiles_executed % self.num_store_stages
+                # Reg -> sB
+                tCsB = thr_copy.partition_D(sB_staged)
+                cute.copy(tiled_copy, tCrA, tCsB[(None, None, None, index)])
+
+                # Fence to ensure smem writes are visible
+                cute.arch.fence_proxy(
+                    "async.shared",
+                    space="cta",
+                )
+                self.trans_sync_barrier.arrive_and_wait()
+
+                if warp_idx == self.trans_warp_id[0]:
+                    cute.copy(
+                        tma_atom_store,
+                        tBsB[(None, index)],
+                        tBgB[(None, cur_tile_coord[0], cur_tile_coord[1])],
+                    )
+                    store_pipeline.producer_commit()
+                    store_pipeline.producer_acquire()
+                self.trans_sync_barrier.arrive_and_wait()
+
+                tile_sched.advance_to_next_work()
+                work_tile = tile_sched.get_current_work()
+
+            self.trans_sync_barrier.arrive_and_wait()
+
+            store_pipeline.producer_tail()
+
+
+def run_transpose(M, N, max_active_clusters=0, num_warmup=5, num_iters=20):
+    """
+    Run TMA transpose kernel with automatic stage calculation and performance measurement.
+
+    Args:
+        M: Matrix dimension M
+        N: Matrix dimension N
+        max_active_clusters: Maximum number of active clusters (0 for auto)
+        num_warmup: Number of warmup iterations
+        num_iters: Number of timing iterations
+
+    Performance Metrics:
+        - Throughput: Actual achieved bandwidth in GB/s
+        - Theoretical BW: Peak memory bandwidth (2048 B/clk × 4000 MHz = 8.192 TB/s)
+        - Bandwidth Efficiency: Percentage of theoretical peak achieved
+    """
+    torch.manual_seed(1111)
+    # Input (M, N)
+    input_data = torch.randn((M, N), device="cuda", dtype=torch.float16)
+    # Output (N, M)
+    output_data = torch.zeros((N, M), device="cuda", dtype=torch.float16)
+
+    # CuTe Wrappers
+    tensor_src = (
+        from_dlpack(input_data, assumed_align=16)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=16)
+    )
+    tensor_dst = (
+        from_dlpack(output_data, assumed_align=16)
+        .mark_layout_dynamic(leading_dim=1)
+        .mark_compact_shape_dynamic(mode=1, divisibility=16)
+    )
+
+    transpose_kernel = Sm100MatrixTransposeKernelV2()
+
+    max_active_clusters = utils.HardwareInfo().get_max_active_clusters(1)
+    # Compile and Run
+    compiled_kernel = cute.compile(
+        transpose_kernel,
+        tensor_src,
+        tensor_dst,
+        max_active_clusters,
+        options="--generate-line-info",
+    )
+
+    # Warmup runs
+    for _ in range(num_warmup):
+        compiled_kernel(tensor_src, tensor_dst)
+    torch.cuda.synchronize()
+
+    # Timed runs
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    start_event.record()
+    for _ in range(num_iters):
+        compiled_kernel(tensor_src, tensor_dst)
+    end_event.record()
+    torch.cuda.synchronize()
+
+    # Calculate performance metrics
+    elapsed_time_ms = start_event.elapsed_time(end_event)
+    avg_time_ms = elapsed_time_ms / num_iters
+
+    # Calculate throughput
+    # For transpose: read M*N elements + write M*N elements
+    bytes_per_element = input_data.element_size()
+    total_bytes = 2 * M * N * bytes_per_element  # Read + Write
+    throughput_gb_s = (total_bytes / 1e9) / (avg_time_ms / 1000)
+
+    # Theoretical bandwidth limit
+    # Blackwell: 2048 B/clk at 4000 MHz
+    bytes_per_clk = 2048
+    freq_mhz = 4000
+    theoretical_bw_gb_s = bytes_per_clk * freq_mhz * 1e6 / 1e9  # Convert to GB/s
+    theoretical_bw_tb_s = theoretical_bw_gb_s / 1000  # Convert to TB/s
+    bandwidth_efficiency = (throughput_gb_s / theoretical_bw_gb_s) * 100  # Percentage
+
+    # Print computed stage counts after compilation
+    print(f"Matrix size: {M}×{N}")
+    print(f"Tile shape: {transpose_kernel.tile_shape}")
+    print(
+        f"Computed stages: Load={transpose_kernel.num_load_stages}, Store={transpose_kernel.num_store_stages}"
+    )
+    print(f"Average time: {avg_time_ms:.4f} ms")
+    print(f"Throughput: {throughput_gb_s:.2f} GB/s")
+    print(
+        f"Theoretical BW: {theoretical_bw_tb_s:.2f} TB/s ({theoretical_bw_gb_s:.2f} GB/s)"
+    )
+    print(f"Bandwidth Efficiency: {bandwidth_efficiency:.2f}%")
+
+    # Verification
+    expected = input_data.t()
+    if torch.allclose(output_data, expected, atol=1e-2):
+        print("Verification: PASSED ✓")
+    else:
+        print("Verification: FAILED ✗")
+        print(f"Max diff: {(output_data - expected).abs().max()}")
+
+
+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."
+            )
+
+    parser = argparse.ArgumentParser(
+        description="TMA Matrix Transpose with Multi-Stage Pipeline and Cluster Support (v2)"
+    )
+    parser.add_argument("--M", type=int, default=128, help="Matrix dimension M")
+    parser.add_argument("--N", type=int, default=128, help="Matrix dimension N")
+    parser.add_argument(
+        "--num_warmup", type=int, default=5, help="Number of warmup iterations"
+    )
+    parser.add_argument(
+        "--num_iters", type=int, default=20, help="Number of timing iterations"
+    )
+    args = parser.parse_args()
+
+    run_transpose(
+        args.M,
+        args.N,
+        num_warmup=args.num_warmup,
+        num_iters=args.num_iters,
+    )