v4.5 tag update (#3202)

* Python DSL examples reorganization.

* v4.5 tag update.

examples/python/CuTeDSL/cute/blackwell/kernel/attention/mamba2_ssd/mamba2_ssd_reference.py

@@ -0,0 +1,397 @@
+# 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 torch
+import torch.nn.functional as F
+
+
+def ssd_reference_fp32_all(x, a, delta, B, C, Y_out, Fstate_out, D, has_d, d_has_hdim):
+    """
+    Rearrange tensor dimensions from cuda layout to reference layout, then directly call TriDao's ssd implementation
+    Arguments:
+        X/x: (D, L, C, H, B):(C*L, 1, L, D*C*L, H*D*C*L)
+        A/delta: (L, C, H, B):(1, L, C*L, H*C*L)
+        a: (H):(1)
+        B/C: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
+        D: (1, H):(0, 1) or (D, H):(1, D)
+        has_d: bool
+        d_has_hdim: bool
+    Return:
+        Y_out: (L, D, C, H, B):(1, C*L, L, D*C*L, H*D*C*L)
+        Fstate_out: (D, N, H, B):(N, 1, D*N, H*D*N)
+    """
+    assert x.dtype == a.dtype == delta.dtype == B.dtype == C.dtype
+
+    A = delta * a.view(1, 1, -1, 1)
+    X = x * delta.unsqueeze(0)
+
+    # Rearrange to match cutlass layout to tridao's layout
+    block_len = A.shape[0]
+    initial_states = None
+    # A: l c h b-> b c l h
+    A = A.permute(3, 1, 0, 2)
+    # X: p l c h b -> b c l h p
+    X = X.permute(4, 2, 1, 3, 0)
+    # B: l n c g b -> b c l g n
+    B = B.permute(4, 2, 0, 3, 1)
+    # C: l n c g b -> b c l g n
+    C = C.permute(4, 2, 0, 3, 1)
+    # X/A/B/C: b c l ... -> b (c l) ...
+    X, A, B, C = [x.reshape(x.shape[0], -1, *x.shape[3:]) for x in (X, A, B, C)]
+
+    # Ngroup (g to h) mapping
+    B_val, CL_val, G_val, N_val = B.shape
+    H_val = X.shape[2]
+    ngroup_ratio = H_val // G_val
+    # B/C: (B, CL, H, N)
+    h_to_g_mapping = torch.arange(H_val, device=B.device) // ngroup_ratio
+    B = B.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
+    C = C.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
+
+    ###################################################################
+    # Call reference implementation from Tri Dao ssd_minimal_discrete
+    Y, final_state = ssd_minimal_discrete_fp32_all(
+        X, A, B, C, block_len, initial_states
+    )
+    ###################################################################
+
+    if has_d:
+        D_val = Y.shape[3]
+        if not d_has_hdim:
+            D = D.expand(D_val, -1)
+        Y = Y + torch.einsum("bchp,ph->bchp", X, D)
+
+    # Rearrange to match tridao's layout to cutlass layout
+    # Y: b (c l) h p -> b c l h p
+    Y = Y.reshape(Y.shape[0], -1, block_len, Y.shape[2], Y.shape[3])
+    # Y: b c l h p -> l p c h b
+    Y = Y.permute(2, 4, 1, 3, 0)
+    # Fstate_out: b h p n -> p n h b
+    Fstate_out.copy_(final_state.permute(2, 3, 1, 0))
+    Y_out.copy_(Y)
+    return
+
+
+def ssd_reference_lowprecision_intermediates(
+    x, a, delta, B, C, Y_out, Fstate_out, intermediate_dtype, D, has_d, d_has_hdim
+):
+    """
+    Rearrange tensor dimensions from cuda layout to reference layout, then call a reduced intermediate dtype version of ssd implementation
+    Arguments:
+        X/x: (D, L, C, H, B):(C*L, 1, L, D*C*L, H*D*C*L)
+        A/delta: (L, C, H, B):(1, L, C*L, H*C*L)
+        a: (H):(1)
+        B/C: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
+        intermediate_dtype: input and intermediate data type
+        D: (1, H):(0, 1) or (D, H):(1, D)
+        has_d: bool
+        d_has_hdim: bool
+    Return:
+        Y_out: (L, D, C, H, B):(1, C*L, L, D*C*L, H*D*C*L)
+        Fstate_out: (D, N, H, B):(N, 1, D*N, H*D*N)
+    """
+    assert x.dtype == a.dtype == delta.dtype == B.dtype == C.dtype
+
+    A = delta * a.view(1, 1, -1, 1)
+
+    # Rearrange to match cutlass layout to tridao's layout
+    block_len = A.shape[0]
+    initial_states = None
+    # A: l c h b-> b c l h
+    A = A.permute(3, 1, 0, 2)
+    # delta: l c h b-> b c l h
+    delta = delta.permute(3, 1, 0, 2)
+    # x: p l c h b -> b c l h p
+    x = x.permute(4, 2, 1, 3, 0)
+    # B: l n c g b -> b c l g n
+    B = B.permute(4, 2, 0, 3, 1)
+    # C: l n c g b -> b c l g n
+    C = C.permute(4, 2, 0, 3, 1)
+    # x/A/delta/B/C: b c l ... -> b (c l) ...
+    x, A, delta, B, C = [
+        tensor.reshape(tensor.shape[0], -1, *tensor.shape[3:])
+        for tensor in (x, A, delta, B, C)
+    ]
+
+    # Ngroup (g to h) mapping
+    B_val, CL_val, G_val, N_val = B.shape
+    H_val = x.shape[2]
+    ngroup_ratio = H_val // G_val
+    # B/C: (B, CL, H, N)
+    h_to_g_mapping = torch.arange(H_val, device=B.device) // ngroup_ratio
+    B = B.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
+    C = C.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
+
+    # Type convert input tensors to input dtype (same as intermediate dtype)
+    x = x.to(intermediate_dtype).to(torch.float32)
+    A = A.to(intermediate_dtype).to(torch.float32)
+    delta = delta.to(intermediate_dtype).to(torch.float32)
+    B = B.to(intermediate_dtype).to(torch.float32)
+    C = C.to(intermediate_dtype).to(torch.float32)
+
+    #########################################################################
+    # Call reference implementation ssd_minimal_discrete_bf16_intermediates
+    Y, final_state = ssd_minimal_discrete_lowprecision_intermediates(
+        x, A, delta, B, C, block_len, intermediate_dtype, initial_states
+    )
+    #########################################################################
+
+    if has_d:
+        D = D.to(intermediate_dtype).to(torch.float32)
+        D_val = Y.shape[3]
+        if not d_has_hdim:
+            D = D.expand(D_val, -1)
+        Y = Y + torch.einsum("bchp,ph->bchp", x, D)
+
+    # Type convert output tensors to output dtype (same as intermediate dtype)
+    Y = Y.to(intermediate_dtype).to(torch.float32)
+    final_state = final_state.to(intermediate_dtype).to(torch.float32)
+
+    # Rearrange to match tridao's layout to cutlass layout
+    # Y: b (c l) h p -> b c l h p
+    Y = Y.reshape(Y.shape[0], -1, block_len, Y.shape[2], Y.shape[3])
+    # Y: b c l h p -> l p c h b
+    Y = Y.permute(2, 4, 1, 3, 0)
+    # Fstate_out: b h p n -> p n h b
+    Fstate_out.copy_(final_state.permute(2, 3, 1, 0))
+    Y_out.copy_(Y)
+    return
+
+
+def analyze_relative_diffs(actual, expected):
+    """
+    Print statistics of relative differences between actual and expected tensors
+    """
+    # Calculate relative differences
+    abs_diff = (actual - expected).abs()
+    rel_diff = abs_diff / (torch.maximum(expected.abs(), actual.abs()) + 0.00001)
+
+    total_elements = rel_diff.numel()
+
+    # Handle special cases first
+    nan_mask = torch.isnan(rel_diff)
+    inf_mask = torch.isinf(rel_diff)
+    nan_count = nan_mask.sum().item()
+    inf_count = inf_mask.sum().item()
+
+    # Find position and value of maximum relative difference
+    max_rel_diff = (
+        rel_diff[~nan_mask & ~inf_mask].max()
+        if (~nan_mask & ~inf_mask).any()
+        else float("nan")
+    )
+    max_rel_diff_pos = (
+        rel_diff[~nan_mask & ~inf_mask].argmax()
+        if (~nan_mask & ~inf_mask).any()
+        else -1
+    )
+
+    # Print max relative difference info
+    print("Maximum relative difference:")
+    print(f"Position: {max_rel_diff_pos}")
+    print(f"Value: {max_rel_diff:.6e}")
+    print(f"Actual value: {actual.flatten()[max_rel_diff_pos]}")
+    print(f"Expected value: {expected.flatten()[max_rel_diff_pos]}")
+    print(f"NaN values: {nan_count} ({100.0 * nan_count / total_elements:.2f}%)")
+    print(f"Inf values: {inf_count} ({100.0 * inf_count / total_elements:.2f}%)\n")
+
+    # Check different rtol thresholds
+    rtol_levels = [1e-5, 1e-4, 1e-3, 1e-2, 5e-02, 1e-01]
+
+    for i, rtol in enumerate(rtol_levels):
+        if i == 0:
+            mask = rel_diff <= rtol
+        else:
+            mask = (rel_diff <= rtol) & (rel_diff > rtol_levels[i - 1])
+
+        count = mask.sum().item()
+        percentage = (count / total_elements) * 100
+
+        if i == 0:
+            print(f"Elements with rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)")
+        else:
+            print(
+                f"Elements with {rtol_levels[i - 1]:.0e} < rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)"
+            )
+
+    # Print elements exceeding the largest rtol
+    mask = rel_diff > rtol_levels[-1]
+    count = mask.sum().item()
+    percentage = (count / total_elements) * 100
+    print(f"Elements with rtol > {rtol_levels[-1]:.0e}: {count} ({percentage:.2f}%)\n")
+
+
+def segsum(x):
+    """
+    More stable segment sum calculation.
+    x: b h c l
+    """
+    T = x.size(-1)
+    # x: b h c l -> b h c l l
+    x = x.unsqueeze(-1).expand(*x.shape, T)
+    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=-1)
+    x = x.masked_fill(~mask, 0)
+    x_segsum = torch.cumsum(x, dim=-2)
+    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
+    x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
+    return x_segsum
+
+
+def ssd_minimal_discrete_fp32_all(X, A, B, C, block_len, initial_states=None):
+    """
+    This is same with https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/ssd_minimal.py
+    (all accumulation and intermediate results in fp32)
+
+    Arguments:
+        X: (batch(B), length(C*L), n_heads(H), d_head(D))
+        A: (batch(B), length(C*L), n_heads(H))
+        B: (batch(B), length(C*L), n_heads(H), d_state(N))
+        C: (batch(B), length(C*L), n_heads(H), d_state(N))
+    Return:
+        Y: (batch(B), length(C*L), n_heads(H), d_head(D))
+        final_state: (B, H, D, N)
+    """
+    assert X.dtype == A.dtype == B.dtype == C.dtype
+    assert X.shape[1] % block_len == 0
+
+    # Rearrange into blocks/chunks
+    # X/A/B/C:b (c l) ... -> b c l ...
+    X, A, B, C = [
+        x.reshape(x.shape[0], -1, block_len, *x.shape[2:]) for x in (X, A, B, C)
+    ]
+
+    # A: b c l h -> b h c l
+    A = A.permute(0, 3, 1, 2)
+    # A_cumsum: (B, H, C, L)
+    A_cumsum = torch.cumsum(A, dim=-1)
+
+    # 1. Compute the output for each intra-chunk (diagonal blocks)
+    segsum_A = segsum(A)
+    L = torch.exp(segsum_A)
+    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)
+
+    # 2. Compute the state for each intra-chunk
+    # (right term of low-rank factorization of off-diagonal blocks; B terms)
+    decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
+    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
+
+    # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+    # (middle term of factorization of off-diag blocks; A terms)
+    if initial_states is None:
+        initial_states = torch.zeros_like(states[:, :1])
+    states = torch.cat([initial_states, states], dim=1)
+    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
+    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
+    states, final_state = new_states[:, :-1], new_states[:, -1]
+
+    # 4. Compute state -> output conversion per chunk
+    # (left term of low-rank factorization of off-diagonal blocks; C terms)
+    state_decay_out = torch.exp(A_cumsum)
+    Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, state_decay_out)
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    # Y: b c l h p -> b (c l) h p
+    Y = (Y_diag + Y_off).reshape(Y_diag.shape[0], -1, Y_diag.shape[3], Y_diag.shape[4])
+    return Y, final_state
+
+
+def ssd_minimal_discrete_lowprecision_intermediates(
+    X, A, delta, B, C, block_len, intermediate_dtype, initial_states=None
+):
+    """
+    This is adjusted from ssd_minimal_discrete_fp32_all, with exceptions:
+    1. accumulation in fp32 but intermediates Q/b_tmem/P are in intermediate_dtype
+    2. delta is not pre-multiplied with X, delta was applied to generate Q/b_tmem to match GPU implementation
+
+    Arguments:
+        X: (batch(B), length(C*L), n_heads(H), d_head(D))
+        A: (batch(B), length(C*L), n_heads(H))
+        delta: (batch(B), length(C*L), n_heads(H))
+        B: (batch(B), length(C*L), n_heads(H), d_state(N))
+        C: (batch(B), length(C*L), n_heads(H), d_state(N))
+    Return:
+        Y: (batch(B), length(C*L), n_heads(H), d_head(D))
+        final_state: (B, H, D, N)
+    """
+    assert X.dtype == A.dtype == B.dtype == C.dtype
+    assert X.shape[1] % block_len == 0
+
+    # Rearrange into blocks/chunks
+    # X/A/delta/B/C: b (c l) ... -> b c l ...
+    X, A, delta, B, C = [
+        x.reshape(x.shape[0], -1, block_len, *x.shape[2:]) for x in (X, A, delta, B, C)
+    ]
+
+    # A: b c l h -> b h c l
+    A = A.permute(0, 3, 1, 2)
+    # delta: b c l h -> b h c l
+    delta = delta.permute(0, 3, 1, 2)
+    # A_cumsum: (B, H, C, L)
+    A_cumsum = torch.cumsum(A, dim=-1)
+
+    # 1. Compute the output for each intra-chunk (diagonal blocks)
+    segsum_A = segsum(A)
+    L = torch.exp(segsum_A)
+    intra_acc_0 = torch.einsum("bclhn,bcshn->bclhs", C, B)
+    Q = torch.einsum("bclhs,bhcls,bhcs->bclhs", intra_acc_0, L, delta)
+    Y_diag = torch.einsum(
+        "bclhs,bcshp->bclhp", Q.to(intermediate_dtype).to(torch.float32), X
+    )
+
+    # 2. Compute the state for each intra-chunk
+    # (right term of low-rank factorization of off-diagonal blocks; B terms)
+    decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
+    b_tmem = torch.einsum("bclhn,bhcl,bhcl->bclhn", B, decay_states, delta)
+    states = torch.einsum(
+        "bclhn,bclhp->bchpn", b_tmem.to(intermediate_dtype).to(torch.float32), X
+    )
+
+    # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+    # (middle term of factorization of off-diag blocks; A terms)
+    if initial_states is None:
+        initial_states = torch.zeros_like(states[:, :1])
+    states = torch.cat([initial_states, states], dim=1)
+    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
+    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
+    states, final_state = new_states[:, :-1], new_states[:, -1]
+    final_state = final_state
+
+    # 4. Compute state -> output conversion per chunk
+    # (left term of low-rank factorization of off-diagonal blocks; C terms)
+    state_decay_out = torch.exp(A_cumsum)
+    Y_off_tmp = torch.einsum(
+        "bclhn,bchpn->bclhp", C, states.to(intermediate_dtype).to(torch.float32)
+    )
+    Y_off = torch.einsum("bclhp,bhcl->bclhp", Y_off_tmp, state_decay_out)
+
+    # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+    # Y: b c l h p -> b (c l) h p
+    Y = (Y_diag + Y_off).reshape(
+        Y_diag.shape[0], -1, Y_diag.shape[3], Y_diag.shape[4]
+    )  # b (c l) h p
+    return Y, final_state

examples/python/CuTeDSL/cute/blackwell/kernel/attention/mamba2_ssd/mamba2_ssd_tile_scheduler.py

@@ -0,0 +1,199 @@
+# 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.
+
+from typing import Tuple
+
+from cutlass.cutlass_dsl import (
+    Integer,
+    Int32,
+    min,
+    extract_mlir_values,
+    new_from_mlir_values,
+    dsl_user_op,
+)
+from cutlass._mlir import ir
+import cutlass.cute as cute
+from cutlass.utils import WorkTileInfo
+
+
+class Mamba2SSDTileSchedulerParams:
+    def __init__(
+        self,
+        problem_shape_ntiles: int,
+        eh: int,
+        ngroup_ratio: int,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        self.problem_shape_ntiles = problem_shape_ntiles
+        self.eh = eh
+        self.ngroup_ratio = ngroup_ratio
+        self._loc = loc
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in [self.problem_shape_ntiles, self.eh, self.ngroup_ratio]:
+            obj_values = extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        obj_list = []
+        for obj, n_items in zip(
+            [self.problem_shape_ntiles, self.eh, self.ngroup_ratio], self._values_pos
+        ):
+            obj_list.append(new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return Mamba2SSDTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
+
+    @dsl_user_op
+    def get_grid_shape(
+        self, max_active_clusters: Int32, *, loc=None, ip=None
+    ) -> Tuple[Integer, Integer, Integer]:
+        return (min(self.problem_shape_ntiles, max_active_clusters), 1, 1)
+
+
+class Mamba2SSDTileScheduler:
+    def __init__(
+        self,
+        params: Mamba2SSDTileSchedulerParams,
+        num_persistent_ctas: Int32,
+        current_work_linear_idx: Int32,
+        num_tiles_executed: Int32,
+    ):
+        self.params = params
+        self.num_persistent_ctas = num_persistent_ctas
+        self._current_work_linear_idx = current_work_linear_idx
+        self._num_tiles_executed = num_tiles_executed
+
+    def __extract_mlir_values__(self) -> list[ir.Value]:
+        values = extract_mlir_values(self.num_persistent_ctas)
+        values.extend(extract_mlir_values(self._current_work_linear_idx))
+        values.extend(extract_mlir_values(self._num_tiles_executed))
+        return values
+
+    def __new_from_mlir_values__(
+        self, values: list[ir.Value]
+    ) -> "Mamba2SSDTileScheduler":
+        assert len(values) == 3
+        new_num_persistent_ctas = new_from_mlir_values(
+            self.num_persistent_ctas, [values[0]]
+        )
+        new_current_work_linear_idx = new_from_mlir_values(
+            self._current_work_linear_idx, [values[1]]
+        )
+        new_num_tiles_executed = new_from_mlir_values(
+            self._num_tiles_executed, [values[2]]
+        )
+        return Mamba2SSDTileScheduler(
+            self.params,
+            new_num_persistent_ctas,
+            new_current_work_linear_idx,
+            new_num_tiles_executed,
+        )
+
+    # called by host
+    @staticmethod
+    @dsl_user_op
+    def create(
+        params: Mamba2SSDTileSchedulerParams,
+        block_idx: Tuple[Integer, Integer, Integer],
+        grid_dim: Tuple[Integer, Integer, Integer],
+        *,
+        loc=None,
+        ip=None,
+    ):
+        params = params
+
+        # Calculate the number of persistent clusters by dividing the total grid size
+        # by the number of CTAs per cluster
+        num_persistent_ctas = Int32(cute.size(grid_dim, loc=loc, ip=ip))
+
+        bidx, bidy, bidz = block_idx
+
+        # Initialize workload index equals to the cluster index in the grid
+        current_work_linear_idx = Int32(bidx)
+
+        # Initialize number of tiles executed to zero
+        num_tiles_executed = Int32(0)
+        return Mamba2SSDTileScheduler(
+            params,
+            num_persistent_ctas,
+            current_work_linear_idx,
+            num_tiles_executed,
+        )
+
+    # called by host
+    @staticmethod
+    def get_grid_shape(
+        params: Mamba2SSDTileSchedulerParams,
+        max_active_clusters: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Tuple[Integer, Integer, Integer]:
+        return params.get_grid_shape(max_active_clusters, loc=loc, ip=ip)
+
+    # private method
+    def _get_current_work_for_linear_idx(
+        self, current_work_linear_idx: Int32, *, loc=None, ip=None
+    ) -> WorkTileInfo:
+        is_valid = current_work_linear_idx < cute.size(
+            self.params.problem_shape_ntiles, loc=loc, ip=ip
+        )
+
+        eh_idx = current_work_linear_idx % self.params.eh
+        b_idx = current_work_linear_idx // self.params.eh
+        g_idx = eh_idx // self.params.ngroup_ratio
+        # cur_tile_coord is (b_idx, eh_idx, g_idx)
+        cur_tile_coord = tuple(Int32(x) for x in (b_idx, eh_idx, g_idx))
+
+        return WorkTileInfo(cur_tile_coord, is_valid)
+
+    @dsl_user_op
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        return self._get_current_work_for_linear_idx(
+            self._current_work_linear_idx, loc=loc, ip=ip
+        )
+
+    @dsl_user_op
+    def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo:
+        return self.get_current_work(loc=loc, ip=ip)
+
+    @dsl_user_op
+    def advance_to_next_work(self, *, advance_count: int = 1, loc=None, ip=None):
+        self._current_work_linear_idx += Int32(advance_count) * Int32(
+            self.num_persistent_ctas
+        )
+        self._num_tiles_executed += Int32(1)
+
+    @property
+    def num_tiles_executed(self) -> Int32:
+        return self._num_tiles_executed

examples/python/CuTeDSL/cute/blackwell/kernel/attention/mixed_input_fmha/prefill_helpers.py

@@ -0,0 +1,400 @@
+# 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.
+
+from typing import Tuple, Optional
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.nvgpu.tcgen05 as tcgen05
+import cutlass.pipeline as pipeline
+
+
+@cute.jit
+def load_qk(
+    iterations: int,
+    kv_step: cutlass.Int32,
+    k_args: Tuple,
+    scale_k_args: Optional[Tuple] = None,
+    q_args: Optional[Tuple] = None,
+) -> Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer]:
+    if cutlass.const_expr(q_args is not None):
+        tQgQ, tQsQ, tma_atom_q, load_q_producer = q_args
+    else:
+        tQgQ, tQsQ, tma_atom_q, load_q_producer = None, None, None, None
+    tKgK, tKsK, tma_atom_k, load_k_producer = k_args
+    tKgScaleK, tKsScaleK, tma_atom_scale_k, load_scale_k_producer = scale_k_args
+
+    scale_k_handle = load_scale_k_producer.acquire_and_advance()
+    cute.copy(
+        tma_atom_scale_k,
+        tKgScaleK[None, kv_step],
+        tKsScaleK[None, scale_k_handle.index],
+        tma_bar_ptr=scale_k_handle.barrier,
+    )
+    for iter in cutlass.range(iterations, unroll=1):
+        if cutlass.const_expr(q_args is not None):
+            q_handle = load_q_producer.acquire_and_advance()
+            cute.copy(
+                tma_atom_q,
+                tQgQ[None, iter],
+                tQsQ[None, q_handle.index],
+                tma_bar_ptr=q_handle.barrier,
+            )
+        k_handle = load_k_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_k,
+            tKgK[None, kv_step, iter],
+            tKsK[None, k_handle.index],
+            tma_bar_ptr=k_handle.barrier,
+        )
+    if cutlass.const_expr(q_args is not None):
+        return load_k_producer, load_scale_k_producer, load_q_producer
+    else:
+        return load_k_producer, load_scale_k_producer
+
+
+@cute.jit
+def load_v(
+    iterations: int,
+    kv_step: cutlass.Int32,
+    v_args: Tuple,
+    scale_v_args: Tuple,
+) -> pipeline.PipelineProducer:
+    tVgV, tVsV, tma_atom_v, load_v_producer = v_args
+    tScaleVgV, tScaleVsV, tma_atom_scale_v, load_scale_v_producer = scale_v_args
+    scale_v_handle = load_scale_v_producer.acquire_and_advance()
+    cute.copy(
+        tma_atom_scale_v,
+        tScaleVgV[None, kv_step],
+        tScaleVsV[None, scale_v_handle.index],
+        tma_bar_ptr=scale_v_handle.barrier,
+    )
+    for iter in cutlass.range(iterations, unroll=1):
+        v_handle = load_v_producer.acquire_and_advance()
+        cute.copy(
+            tma_atom_v,
+            tVgV[None, iter, kv_step],
+            tVsV[None, v_handle.index],
+            tma_bar_ptr=v_handle.barrier,
+        )
+    return load_v_producer, load_scale_v_producer
+
+
+@cute.jit
+def get_scale_smem_layout(
+    scale_granularity: int,
+    d_r: int,
+    mma_tiler: cute.Tile,
+    major_mode: tcgen05.OperandMajorMode,
+) -> Tuple[cute.Layout, cute.Tile]:
+    size_mn = mma_tiler[1] // 2  # 2cta by default
+    if cutlass.const_expr(major_mode == tcgen05.OperandMajorMode.MN):  # v
+        scale_tiler = (mma_tiler[2] * d_r,)
+        tma_view_layout = cute.make_layout(
+            (mma_tiler[2] * d_r),
+        )
+        assert scale_granularity % mma_tiler[1] == 0, (
+            "scale_granularity must be divisible by mma_tiler[1]"
+        )
+        rest_l = scale_granularity // mma_tiler[1]
+        s2r_view_layout = cute.make_layout(
+            (size_mn, mma_tiler[2], (rest_l, d_r)),
+            stride=(0, d_r, (0, 1)),
+        )
+    else:  # k
+        scale_tiler = (mma_tiler[1] * d_r,)
+        tma_view_layout = cute.make_layout((size_mn * d_r))
+        assert scale_granularity % mma_tiler[2] == 0, (
+            "scale_granularity must be divisible by mma_tiler[2]"
+        )
+        rest_l = scale_granularity // mma_tiler[2]
+        s2r_view_layout = cute.make_layout(
+            (size_mn, mma_tiler[2], (rest_l, d_r)),
+            stride=(d_r, 0, (0, 1)),
+        )
+    # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom
+    tma_view_smem_layout = cute.make_composed_layout(
+        cute.make_swizzle(0, 4, 3), 0, tma_view_layout
+    )
+    return tma_view_smem_layout, scale_tiler, s2r_view_layout
+
+
+@cute.jit
+def mma_qk(
+    iterations: int,
+    qk_tiled_mma: cute.TiledMma,
+    tensor_args: Tuple,
+    pipeline_args: Tuple,
+):
+    tStS, tSrQ, tSrK_trans = tensor_args
+    mma_s_producer, load_q_consumer, load_q_releaser, dequant_kv_consumer = (
+        pipeline_args
+    )
+    cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+    is_leader_cta = cta_rank_in_cluster % 2 == 0
+    if is_leader_cta:
+        s_handle = mma_s_producer.acquire_and_advance()
+        tStS_slice = tStS[None, None, None, s_handle.index]
+        qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+        for iter in cutlass.range(iterations, unroll=1):
+            if cutlass.const_expr(load_q_consumer is not None):
+                load_q_consumer.wait_and_advance()
+            tSrQ_slice = tSrQ[None, None, None, iter]
+            k_trans_handle = dequant_kv_consumer.wait_and_advance()
+            tSrK_trans_slice = tSrK_trans[None, None, None, k_trans_handle.index]
+            num_kphases = cute.size(tSrQ_slice, mode=[2])
+            for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
+                kphase_coord = (None, None, kphase_idx)
+                cute.gemm(
+                    qk_tiled_mma,
+                    tStS_slice,
+                    tSrQ_slice[kphase_coord],
+                    tSrK_trans_slice[kphase_coord],
+                    tStS_slice,
+                )
+                qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+            k_trans_handle.release()
+            if cutlass.const_expr(load_q_releaser is not None):
+                load_q_releaser.release()
+                load_q_releaser.advance()
+        s_handle.commit()
+    return mma_s_producer, load_q_consumer, dequant_kv_consumer
+
+
+@cute.jit
+def dequant_k(
+    iterations: int,
+    transform_warp_ids: Tuple,
+    dtype_args: Tuple,
+    tensor_args: Tuple,
+    pipeline_args: Tuple,
+):
+    (k_dtype, q_dtype) = dtype_args
+    (sOrig, sScale, sTrans) = tensor_args
+    (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args
+    tidx, _, _ = cute.arch.thread_idx()
+    THREADS_PER_WARP = 32
+    thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids))
+    r2s_copy_atom = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(), k_dtype, num_bits_per_copy=32
+    )
+    # Construct tiled_copy satisfying 16 contiguous elts per copy atom
+    r2s_tiled_copy = cute.make_cotiled_copy(
+        r2s_copy_atom,
+        cute.make_layout((256, 16), stride=(16, 1)),
+        sTrans[(None, None, None, 0)].layout,
+    )
+    thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx)
+    tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig)
+    tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans)
+    tOrOrig = cute.make_rmem_tensor_like(
+        cute.append(
+            tOsOrig[None, None, None, None, 0].layout,
+            cute.make_layout(
+                2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout)
+            ),
+        ),
+        k_dtype,
+    )
+    tTrTrans = cute.make_rmem_tensor_like(
+        cute.append(
+            tTsTrans[None, None, None, None, 0].layout,
+            cute.make_layout(
+                2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout)
+            ),
+        ),
+        q_dtype,
+    )
+    tSsScale = thr_r2s_tiled_copy.partition_S(sScale)
+    tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0])
+    scale_handle = load_scale_consumer.wait_and_advance()
+    cute.autovec_copy(
+        tSsScale[None, None, None, None, None, scale_handle.index], tSrScale
+    )
+    cute.arch.fence_view_async_shared()
+    scale_handle.release()
+    # prefetch iter = 0
+    kv_handle = load_kv_consumer.wait_and_advance()
+    cute.autovec_copy(
+        tOsOrig[None, None, None, None, kv_handle.index],
+        tOrOrig[None, None, None, None, 0],
+    )
+    transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype)
+    scale = cute.TensorSSA(
+        tSrScale[None, None, None, None, 0].load(),
+        transformed_tensor.shape,
+        q_dtype,
+    )
+    transformed_tensor = transformed_tensor * scale
+    tTrTrans[None, None, None, None, 0].store(transformed_tensor)
+    cute.arch.fence_view_async_shared()
+    kv_handle.release()
+    for iter in cutlass.range(1, iterations, unroll_full=True):
+        kv_trans_handle = dequant_kv_producer.acquire_and_advance()
+        cute.autovec_copy(
+            tTrTrans[None, None, None, None, (iter - 1) % 2],
+            tTsTrans[None, None, None, None, kv_trans_handle.index],
+        )
+        cute.arch.fence_view_async_shared()
+        kv_trans_handle.commit()
+        kv_handle = load_kv_consumer.wait_and_advance()
+        cute.autovec_copy(
+            tOsOrig[None, None, None, None, kv_handle.index],
+            tOrOrig[None, None, None, None, iter % 2],
+        )
+        transformed_tensor = (
+            tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype)
+        )
+        scale = cute.TensorSSA(
+            tSrScale[None, None, None, None, iter].load(),
+            transformed_tensor.shape,
+            q_dtype,
+        )
+        transformed_tensor = transformed_tensor * scale
+        tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor)
+        cute.arch.fence_view_async_shared()
+        kv_handle.release()
+    kv_trans_handle = dequant_kv_producer.acquire_and_advance()
+    cute.autovec_copy(
+        tTrTrans[None, None, None, None, (iterations - 1) % 2],
+        tTsTrans[None, None, None, None, kv_trans_handle.index],
+    )
+    cute.arch.fence_view_async_shared()
+    kv_trans_handle.commit()
+    return load_kv_consumer, load_scale_consumer, dequant_kv_producer
+
+
+@cute.jit
+def dequant_v(
+    iterations: int,
+    transform_warp_ids: Tuple,
+    dtype_args: Tuple,
+    tensor_args: Tuple,
+    pipeline_args: Tuple,
+):
+    (v_dtype, q_dtype) = dtype_args
+    (sOrig, sScale, sTrans) = tensor_args
+    (load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args
+    tidx, _, _ = cute.arch.thread_idx()
+    THREADS_PER_WARP = 32
+    thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids))
+    r2s_copy_atom = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(), v_dtype, num_bits_per_copy=32
+    )
+    # Construct tiled_copy satisfying 16 contiguous elts per copy atom
+    r2s_tiled_copy = cute.make_cotiled_copy(
+        r2s_copy_atom,
+        cute.make_layout((256, 16), stride=(16, 1)),
+        sTrans[(None, None, None, 0)].layout,
+    )
+    thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx)
+    tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig)
+    tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans)
+    # double buffer for better perf
+    tOrOrig = cute.make_rmem_tensor_like(
+        cute.append(
+            tOsOrig[None, None, None, None, 0].layout,
+            cute.make_layout(
+                2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout)
+            ),
+        ),
+        v_dtype,
+    )
+    tTrTrans = cute.make_rmem_tensor_like(
+        cute.append(
+            tTsTrans[None, None, None, None, 0].layout,
+            cute.make_layout(
+                2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout)
+            ),
+        ),
+        q_dtype,
+    )
+    tSsScale = thr_r2s_tiled_copy.partition_S(sScale)
+    tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0])
+    scale_v_handle = load_scale_consumer.wait_and_advance()
+    cute.autovec_copy(
+        tSsScale[None, None, None, None, None, scale_v_handle.index],
+        tSrScale,
+    )
+    cute.arch.fence_view_async_shared()
+    scale_v_handle.release()
+    # prefetch iter = 0
+    kv_handle = load_kv_consumer.wait_and_advance()
+    cute.autovec_copy(
+        tOsOrig[None, None, None, None, kv_handle.index],
+        tOrOrig[None, None, None, None, 0],
+    )
+    transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype)
+    scale = cute.TensorSSA(
+        tSrScale[None, None, None, None, 0].load(),
+        transformed_tensor.shape,
+        q_dtype,
+    )
+    transformed_tensor = transformed_tensor * scale
+    tTrTrans[None, None, None, None, 0].store(transformed_tensor)
+    cute.arch.fence_view_async_shared()
+    kv_handle.release()
+    for iter in cutlass.range(1, iterations, unroll_full=True):
+        kv_trans_handle = dequant_kv_producer.acquire_and_advance()
+        cute.autovec_copy(
+            tTrTrans[None, None, None, None, (iter - 1) % 2],
+            tTsTrans[None, None, None, None, kv_trans_handle.index],
+        )
+        cute.arch.fence_view_async_shared()
+        kv_trans_handle.commit()
+        kv_handle = load_kv_consumer.wait_and_advance()
+        cute.autovec_copy(
+            tOsOrig[None, None, None, None, kv_handle.index],
+            tOrOrig[None, None, None, None, iter % 2],
+        )
+        transformed_tensor = (
+            tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype)
+        )
+        scale = cute.TensorSSA(
+            tSrScale[
+                None,
+                None,
+                None,
+                None,
+                iter,
+            ].load(),
+            transformed_tensor.shape,
+            q_dtype,
+        )
+        transformed_tensor = transformed_tensor * scale
+        tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor)
+        cute.arch.fence_view_async_shared()
+        kv_handle.release()
+    kv_trans_handle = dequant_kv_producer.acquire_and_advance()
+    cute.autovec_copy(
+        tTrTrans[None, None, None, None, (iterations - 1) % 2],
+        tTsTrans[None, None, None, None, kv_trans_handle.index],
+    )
+    cute.arch.fence_view_async_shared()
+    kv_trans_handle.commit()
+    return load_kv_consumer, load_scale_consumer, dequant_kv_producer

examples/python/CuTeDSL/cute/blackwell/kernel/attention/mla/mla_helpers.py

@@ -0,0 +1,304 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+import cutlass
+import cutlass.cute as cute
+
+
+class MLAStaticTileSchedulerParams:
+    def __init__(
+        self,
+        is_persistent: bool,
+        problem_shape_b: cute.Int32,
+        problem_shape_s: cute.Int32,
+        cluster_shape_mnk: cute.Shape,
+        split_kv: cutlass.Int32,
+        *,
+        problem_shape_b_fdd: cute.FastDivmodDivisor = None,
+        problem_shape_s_fdd: cute.FastDivmodDivisor = None,
+        split_kv_fdd: cute.FastDivmodDivisor = None,
+        loc=None,
+        ip=None,
+    ):
+        """The static tile scheduler parameters prepared for MLA static tile scheduler.
+
+        :param is_persistent: Whether to use persistent kernel mode
+        :type is_persistent: bool
+        :param problem_shape_b: The shape of the problem
+        :type problem_shape_b: cute.Int32
+        :param problem_shape_s: The shape of the problem in sequence length Q dimension
+        :type problem_shape_s: cute.Int32
+        :param cluster_shape_mnk: The shape of the cluster
+        :type cluster_shape_mnk: cute.Shape
+        :param split_kv: The scalar factor for split KV
+        """
+        self.is_persistent = is_persistent
+        self.problem_shape_b = problem_shape_b
+        self.problem_shape_s = problem_shape_s
+        self.problem_shape_b_fdd = problem_shape_b_fdd
+        self.problem_shape_s_fdd = problem_shape_s_fdd
+        self.cluster_shape_mnk = cluster_shape_mnk
+        self.split_kv = split_kv
+        self.split_kv_fdd = split_kv_fdd
+        if cutlass.const_expr(problem_shape_b_fdd is None):
+            self.problem_shape_b_fdd = cute.fast_divmod_create_divisor(
+                problem_shape_b, loc=loc, ip=ip
+            )
+        if cutlass.const_expr(problem_shape_s_fdd is None):
+            self.problem_shape_s_fdd = cute.fast_divmod_create_divisor(
+                problem_shape_s, loc=loc, ip=ip
+            )
+        if cutlass.const_expr(split_kv_fdd is None):
+            self.split_kv_fdd = cute.fast_divmod_create_divisor(
+                split_kv, loc=loc, ip=ip
+            )
+        self.loc = loc
+        self.ip = ip
+
+    def __extract_mlir_values__(self):
+        values = cutlass.extract_mlir_values(self.problem_shape_b)
+        values += cutlass.extract_mlir_values(self.problem_shape_s)
+        values += cutlass.extract_mlir_values(self.split_kv)
+        values += cutlass.extract_mlir_values(self.problem_shape_b_fdd)
+        values += cutlass.extract_mlir_values(self.problem_shape_s_fdd)
+        values += cutlass.extract_mlir_values(self.split_kv_fdd)
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        problem_shape_b = cutlass.new_from_mlir_values(
+            self.problem_shape_b, (values[0],)
+        )
+        problem_shape_s = cutlass.new_from_mlir_values(
+            self.problem_shape_s, (values[1],)
+        )
+        split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2],))
+        problem_shape_b_fdd = cutlass.new_from_mlir_values(
+            self.problem_shape_b_fdd, (values[3],)
+        )
+        problem_shape_s_fdd = cutlass.new_from_mlir_values(
+            self.problem_shape_s_fdd, (values[4],)
+        )
+        split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, (values[5],))
+        return MLAStaticTileSchedulerParams(
+            self.is_persistent,
+            problem_shape_b,
+            problem_shape_s,
+            self.cluster_shape_mnk,
+            split_kv,
+            problem_shape_b_fdd=problem_shape_b_fdd,
+            problem_shape_s_fdd=problem_shape_s_fdd,
+            split_kv_fdd=split_kv_fdd,
+            loc=self.loc,
+        )
+
+
+def create_mla_static_tile_scheduler_params(
+    is_persistent: bool,
+    problem_shape_b: cute.Int32,
+    problem_shape_s: cute.Int32,
+    cluster_shape_mnk: cute.Shape,
+    split_kv: cutlass.Int32,
+) -> MLAStaticTileSchedulerParams:
+    return MLAStaticTileSchedulerParams(
+        is_persistent, problem_shape_b, problem_shape_s, cluster_shape_mnk, split_kv
+    )
+
+
+class WorkTileInfo:
+    def __init__(self, blk_coord: cute.Coord, is_valid: bool):
+        self.blk_coord = blk_coord
+        self.is_valid = cutlass.Boolean(is_valid)
+
+    def __extract_mlir_values__(self):
+        values = cutlass.extract_mlir_values(self.blk_coord)
+        values += cutlass.extract_mlir_values(self.is_valid)
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        new_tile_idx = cutlass.new_from_mlir_values(self.blk_coord, values[:-1])
+        new_is_valid_tile = cutlass.new_from_mlir_values(self.is_valid, [values[-1]])
+        return WorkTileInfo(new_tile_idx, new_is_valid_tile)
+
+    @property
+    def is_valid_tile(self) -> cutlass.Boolean:
+        return self.is_valid
+
+    @property
+    def tile_idx(self) -> cute.Coord:
+        return self.blk_coord
+
+
+class MLAStaticTileScheduler:
+    def __init__(
+        self,
+        params: MLAStaticTileSchedulerParams,
+        current_work_linear_idx: cutlass.Int32,
+        blk_coord: cute.Coord,
+        grid_shape: cute.Shape,
+        *,
+        is_valid: bool = True,
+        loc=None,
+        ip=None,
+    ):
+        """The static tile scheduler for MLA split kv kernel.
+        Based on `is_persistent`, it provides 2 modes for use:
+        - Persistent mode: Launch fixed blocks and reschedule the data blocks.
+        - Non-persistent mode: Launch dynamic blocks and exit when the current work is done.
+
+        :param params: The static tile scheduler parameters
+        :type params: MLAStaticTileSchedulerParams
+        :param current_work_linear_idx: The linear index of the current work
+        :type current_work_linear_idx: cutlass.Int32
+        :param blk_coord: The coordinate of the current work
+        :type blk_coord: cute.Coord
+        :param grid_shape: The shape of the grid
+        :type grid_shape: cute.Shape
+        :param is_valid: Whether the current work is valid
+        :type is_valid: bool
+        """
+        self.params = params
+        self.blk_coord = blk_coord
+        self.grid_shape = grid_shape
+        self.current_work_linear_idx = current_work_linear_idx
+        if params.is_persistent:
+            self.persistent_blk_layout = cute.make_layout(
+                (
+                    params.cluster_shape_mnk[0],
+                    params.problem_shape_s,
+                    params.problem_shape_b,
+                    params.split_kv,
+                ),
+                loc=loc,
+                ip=ip,
+            )
+            self.num_blocks = cute.size(self.persistent_blk_layout, loc=loc, ip=ip)
+            # Used for persistent scheduling
+            self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip)
+        else:
+            self.is_valid = is_valid
+        self.loc = loc
+        self.ip = ip
+
+    @staticmethod
+    def get_grid_shape(
+        params: MLAStaticTileSchedulerParams,
+        max_active_clusters: int,
+        *,
+        loc=None,
+        ip=None,
+    ) -> cute.Shape:
+        # called by host
+        grid_shape = (
+            params.cluster_shape_mnk[0],
+            params.problem_shape_b * params.problem_shape_s,
+            params.split_kv,
+        )
+        if params.is_persistent:
+            return (
+                cutlass.min(
+                    max_active_clusters * cute.size(params.cluster_shape_mnk),
+                    cute.size(grid_shape, loc=loc, ip=ip),
+                ),
+                1,
+                1,
+            )
+        else:
+            return grid_shape
+
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        is_valid = (
+            self.current_work_linear_idx < self.num_blocks
+            if self.params.is_persistent
+            else self.is_valid
+        )
+
+        if self.params.is_persistent:
+            current_work_cluster_batch, cluster_idx = (
+                self.current_work_linear_idx // self.params.cluster_shape_mnk[0],
+                self.current_work_linear_idx % self.params.cluster_shape_mnk[0],
+            )
+            current_work_s_batch, s_idx = divmod(
+                current_work_cluster_batch, self.params.problem_shape_s_fdd
+            )
+            current_work_b_batch, b_idx = divmod(
+                current_work_s_batch, self.params.problem_shape_b_fdd
+            )
+            _, split_kv_idx = divmod(current_work_b_batch, self.params.split_kv_fdd)
+
+            blk_coord = (cluster_idx, s_idx, b_idx, split_kv_idx)
+        else:
+            s_idx, b_idx = divmod(self.blk_coord[1], self.params.problem_shape_b_fdd)
+            blk_coord = (self.blk_coord[0], s_idx, b_idx, self.blk_coord[2])
+
+        return WorkTileInfo(blk_coord, is_valid)
+
+    def initial_work_tile_info(self, *, loc=None, ip=None):
+        return self.get_current_work(loc=loc, ip=ip)
+
+    def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None):
+        if self.params.is_persistent:
+            self.current_work_linear_idx += advance_count * self.num_persistent_sm
+        else:
+            self.is_valid = False
+
+    def __extract_mlir_values__(self):
+        values = cutlass.extract_mlir_values(self.params)
+        values.extend(cutlass.extract_mlir_values(self.current_work_linear_idx))
+        values.extend(cutlass.extract_mlir_values(self.blk_coord))
+        values.extend(cutlass.extract_mlir_values(self.grid_shape))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 13
+        new_params = cutlass.new_from_mlir_values(self.params, values[0:6])
+        new_current_work_linear_idx = cutlass.new_from_mlir_values(
+            self.current_work_linear_idx, [values[6]]
+        )
+        new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, values[7:10])
+        new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, values[10:])
+        return MLAStaticTileScheduler(
+            new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape
+        )
+
+
+def create_mla_static_tile_scheduler(
+    params: MLAStaticTileSchedulerParams,
+    blk_coord: cute.Coord,
+    grid_shape: cute.Shape,
+) -> MLAStaticTileScheduler:
+    return MLAStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape)
+
+
+LOG2_E = 1.4426950408889634074
+# avoid register indexing on array.
+MAX_SPLITS = 256
+
+
+def ceil_div(a: int, b: int) -> int:
+    return (a + b - 1) // b

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/README.md

@@ -0,0 +1,139 @@
+# CuTeDSL Distributed Examples
+
+This directory contains distributed examples using CuTeDSL with NVSHMEM for multi-GPU communication. Currently, we do not support to use NVSHMEM for any device side copy/put/get impl, only use the host side setup and allocations.
+
+## NVSHMEM Dependency
+
+These examples require two components:
+
+1. **NVSHMEM4Py** (`nvshmem4py-cu12` / `nvshmem4py-cu13`): A Python package that provides the official Python binding for NVIDIA's NVSHMEM. See the [NVSHMEM4Py Documentation](https://docs.nvidia.com/nvshmem/api/api/language_bindings/python/index.html).
+
+2. **NVSHMEM Library** (`nvidia-nvshmem-cu12` / `nvidia-nvshmem-cu13`): The underlying native library that contains the actual NVSHMEM implementation.
+
+### Overview
+
+**NVSHMEM4Py** (`nvshmem4py-cu12` / `nvshmem4py-cu13`) is a Python binding library that provides a Pythonic interface to NVSHMEM functionality. In these examples, we use it primarily for:
+
+- Allocating tensors that support peer-to-peer (P2P) communication across GPUs
+- Allocating multicast (MC) tensors that can leverage `multimem` instructions for efficient collective operations
+
+**nvidia-nvshmem** (`nvidia-nvshmem-cu12` / `nvidia-nvshmem-cu13`) is the underlying library that wraps NVSHMEM functions into dynamic libraries (`.so` files). NVSHMEM4Py dynamically loads and calls these libraries at runtime.
+
+### Installation
+
+For CUDA 12:
+```bash
+pip install nvshmem4py-cu12 nvidia-nvshmem-cu12
+```
+
+For CUDA 13:
+```bash
+pip install nvshmem4py-cu13 nvidia-nvshmem-cu13
+```
+
+> **Note:** `nvshmem4py` version >= 0.1.3 is recommended.
+
+### Key APIs Used
+
+We primarily use the following APIs from `nvshmem.core`:
+
+| API | Description |
+|-----|-------------|
+| `nvshmem.core.tensor(shape, dtype)` | Allocates a symmetric tensor that supports P2P communication |
+| `nvshmem.core.get_peer_tensor(tensor, pe)` | Returns a tensor handle for accessing the given tensor on a remote PE (processing element) |
+| `nvshmem.core.get_multicast_tensor(tensor)` | Returns a tensor that can be accessed using `multimem` instructions for efficient multicast operations |
+| `nvshmem.core.free_tensor(tensor)` | Explicitly frees the allocated symmetric memory |
+
+### Memory Management
+
+NVSHMEM requires **manual memory management**. Unlike PyTorch tensors that are garbage-collected automatically, NVSHMEM symmetric memory must be explicitly freed using `nvshmem.core.free_tensor()` to avoid memory leaks.
+
+Example:
+```python
+import nvshmem.core
+
+# init the environment
+# refer to the torchrun_uid_init_bcast() in example
+
+# Allocate symmetric tensor
+local_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+
+# Get peer tensors for P2P access
+tensor_list = [nvshmem.core.get_peer_tensor(local_tensor, rank) for rank in range(world_size)]
+
+# ... use tensors ...
+
+# Explicitly free memory when done
+for t in tensor_list:
+    nvshmem.core.free_tensor(t)
+
+# finalize the environment
+# refer to the torchrun_finalize() in example
+
+```
+
+## Multimem Instructions
+
+These examples demonstrate the use of NVIDIA's `multimem` PTX instructions for efficient multi-GPU collective operations. The `multimem` instructions operate on multicast (MC) addresses obtained via `nvshmem.core.get_multicast_tensor()`, enabling hardware-accelerated communication across multiple GPUs.
+
+### Why Multimem is Fast: NVLS (NVLink SHARP)
+
+The `multimem` instructions leverage **NVLS (NVLink SHARP)** technology to perform **in-network computation**. When multiple GPUs map the same symmetric memory region, `multimem` instructions can operate on a multicast address to perform hardware-accelerated reduction or broadcast operations directly in the NVLink/NVSwitch fabric, without requiring data to traverse to GPU memory first.
+
+**Key benefits:**
+- **In-network computation**: Reduction and broadcast operations happen in the NVSwitch hardware, not in GPU compute units
+- **Reduced memory traffic**: Data is processed in-flight within the interconnect, minimizing HBM bandwidth consumption
+- **Lower latency**: Single instruction replaces multiple loads/stores and arithmetic operations
+
+### Instruction Categories
+
+We use three types of `multimem` instructions in these examples:
+
+#### 1. `multimem.ld_reduce` - Reduction
+
+Reads data from a multicast address and returns the **reduced result** (e.g., sum) across all GPUs:
+
+```
+multimem.ld_reduce.sys.relaxed.global.add.v4.f32 {$0, $1, $2, $3}, [$4];
+```
+
+This instruction reads from a multicast address and performs a sum reduction (`.add`) across all GPUs that have mapped this address via NVLS.
+
+**Accumulator Precision**: For lower-precision data types, you can specify a higher accumulator precision to improve numerical accuracy:
+- **FP16 / BF16**: Can use FP32 accumulator (`.acc::f32`)
+- **FP8 (E4M3 / E5M2)**: Can use FP16 accumulator (`.acc::f16`)
+
+Example with FP16 using FP32 accumulator:
+```
+multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.f16x2 {$0, $1, $2, $3}, [$4];
+```
+
+#### 2. `multimem.st` - Broadcast via Store
+
+Stores data to a multicast address, which **broadcasts** the data to all participating GPUs:
+
+```
+multimem.st.sys.relaxed.global.v4.f32 [$1], {$2, $3, $4, $5};
+```
+
+This writes data to a multicast address, and the data becomes visible to all GPUs that have mapped this address via NVLS.
+
+#### 3. `multimem.red` - Broadcast via Atomic Reduction
+
+Performs an atomic reduction operation on a multicast address. This is commonly used for **signaling/synchronization** across GPUs:
+
+```
+multimem.red.release.sys.global.add.u32 [$0], 1;
+```
+
+This atomically adds a value to a multicast address. When used with synchronization patterns (e.g., spin locks), it enables efficient inter-GPU barriers where all GPUs can observe the updated value.
+
+## Future Work
+
+The `nvidia-nvshmem-cu12/cu13` packages include LLVM IR bitcode libraries that could potentially be integrated into CuTeDSL in the future. This would enable calling NVSHMEM functions directly from within CuTeDSL kernels, allowing for more fine-grained control over communication patterns at the kernel level.
+
+## References
+
+- [NVSHMEM4Py Documentation](https://docs.nvidia.com/nvshmem/api/api/language_bindings/python/index.html)
+- [NVSHMEM API Reference](https://docs.nvidia.com/nvshmem/api/api/language_bindings/python/index.html)
+- [multimem PTX instruction](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-multimem)

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_one_shot_lamport.py

@@ -0,0 +1,426 @@
+# 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 os
+import torch
+import argparse
+
+import numpy as np
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+import cuda.bindings.driver as cuda
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T
+from cutlass._mlir.dialects import vector
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+"""
+A Distributed One-Shot All-Reduce Example using CuTe DSL and fine-grained memory control. This is a mirrored version of the 
+existing tensorrt_llm kernel:
+https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu
+
+In Lamport terminology this is a classic flag-based busy-wait: every participant keeps polling the shared slot until the
+flag changes from the sentinel (negative zero) to real data, which indicates that the Lamport-style logical ordering has
+advanced and the payload is safe to consume.
+
+This example kernel demonstrates a one-shot all-reduce operation using the CuTe DSL with fine-grained memory control.
+It uses dedicated communication buffers for data exchange, and these buffers act as ping-pong buffers. During the 
+process, the kernel uses one buffer for communication and initializes the next buffer to all negative zeros.
+
+In this kernel, each thread is only responsible for 128bits of data. The kernel will write it's local data to every
+buffer at different ranks, then read the data from the local rank buffer. The buffer itself behaves as a barrier, 
+if kernel read negtive 0, then it means data are not ready or not visible yet so that the kernel will read the data again.
+
+If the input tensors from each device are not remotely accessible, this kernel can be used to perform the one-shot all-reduce 
+since it uses communication buffers for data exchange.
+
+The .SYS memory scope and .VOLATILE memory order are used to ensure that the data will be visible at the system scope.
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_one_shot_lamport.py --M 8192 --N 8192
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_one_shot_lamport.py \
+        --M 8192 --N 8192 --benchmark --warmup_iterations 2 --iterations 10
+"""
+
+
+PING_PONG_SIZE = 3
+
+
+class AllReduceOneShotLamportKernel:
+    @cute.jit
+    def __call__(
+        self,
+        rank: cutlass.Constexpr,
+        world_size: cutlass.Constexpr,
+        signal: cutlass.Int32,
+        local_input: cute.Tensor,
+        local_output: cute.Tensor,
+        buffers: list[cute.Tensor],
+        stream: cuda.CUstream,
+    ):
+        copy_bits = 128
+        dtype = local_input.element_type
+        vector_size = copy_bits // dtype.width
+
+        thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+        val_layout = cute.make_ordered_layout((1, vector_size), order=(1, 0))
+        tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+        grouped_buffers = [cute.group_modes(buffer, 0, 2) for buffer in buffers]
+        tiled_buffers = [
+            cute.zipped_divide(buffer, (tiler_mn, world_size, PING_PONG_SIZE))
+            for buffer in grouped_buffers
+        ]
+        tiled_input = cute.zipped_divide(local_input, tiler_mn)
+        tiled_output = cute.zipped_divide(local_output, tiler_mn)
+
+        self.kernel(
+            tiled_buffers,
+            tiled_input,
+            tiled_output,
+            thr_layout,
+            val_layout,
+            signal,
+            rank,
+        ).launch(
+            grid=[cute.size(tiled_input, mode=[1]), 1, 1],
+            block=[cute.size(tv_layout, mode=[0]), 1, 1],
+            stream=stream,
+        )
+
+    #  GPU device kernel
+    @cute.kernel
+    def kernel(
+        self,
+        buffers: list[cute.Tensor],
+        local_input: cute.Tensor,
+        local_output: cute.Tensor,
+        thr_layout: cute.Layout,
+        val_layout: cute.Layout,
+        signal: cutlass.Int32,
+        rank: cutlass.Constexpr,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+        ping = signal % 3
+        pong = (signal + 1) % 3
+
+        buffer_local = buffers[rank]
+        cta_coord = ((None, None), bidx)
+        local_tile_in = local_input[cta_coord]
+        local_tile_out = local_output[cta_coord]
+
+        ping_coord = (((None, None), None, ping), bidx)
+        pong_coord = (((None, None), None, pong), bidx)
+
+        read_buffer = buffer_local[ping_coord]
+        clear_buffer = buffer_local[pong_coord]
+
+        write_coord = (((None, None), rank, ping), bidx)
+        write_buffers = [buffer[write_coord] for buffer in buffers]
+
+        # assume all buffers have the same element type with input
+        copy_atom_load = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            buffers[0].element_type,
+            num_bits_per_copy=128,
+            memory_scope=cute.nvgpu.common.MemoryScope.SYS,
+            memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
+        )
+        copy_atom_store = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            buffers[0].element_type,
+            num_bits_per_copy=128,
+            memory_scope=cute.nvgpu.common.MemoryScope.SYS,
+            memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
+        )
+        tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+        thr_copy = tiled_copy.get_slice(tidx)
+
+        thr_write_buffer_list = [
+            thr_copy.partition_D(tensor) for tensor in write_buffers
+        ]
+        thr_read_buffer = thr_copy.partition_S(read_buffer)
+
+        thr_clear_buffer = thr_copy.partition_D(clear_buffer)
+
+        thr_in = thr_copy.partition_S(local_tile_in)
+        thr_out = thr_copy.partition_D(local_tile_out)
+
+        frg_in = cute.make_fragment_like(thr_in)
+        frg_clear = cute.make_fragment_like(thr_clear_buffer)
+        frg_acc = cute.make_fragment_like(thr_out)
+        frg_acc.fill(0.0)
+
+        # clear a next buffer to be all negtive 0
+        clear_tensor = frg_clear.load()
+        frg_size = cute.size(clear_tensor.shape)
+        neg0_i32_vec = cute.full_like(clear_tensor, 0x80000000, cutlass.Int32)
+        neg0_f32_vec = vector.bitcast(T.vector(frg_size, T.f32()), neg0_i32_vec)
+        neg0_f32_tensor = cute.TensorSSA(
+            neg0_f32_vec, clear_tensor.shape, cutlass.Float32
+        )
+        frg_clear.store(neg0_f32_tensor)
+        cute.copy(copy_atom_store, frg_clear, thr_clear_buffer)
+
+        # read local data to the register
+        cute.copy(copy_atom_load, thr_in, frg_in)
+
+        # write local data to every buffer at different ranks
+        for thr_write_buffer in thr_write_buffer_list:
+            cute.copy(copy_atom_store, frg_in, thr_write_buffer)
+
+        frg_in_vector_neg0_i32 = cute.full_like(
+            frg_in, cutlass.Int32(0x80000000), cutlass.Int32
+        )
+        frg_in_size = cute.size(frg_in.shape)
+
+        # loop over each buffer and accumulate the data
+        for i in cutlass.range_constexpr(len(buffers)):
+            read_coord = (None, 0, 0, i)
+            cute.copy(copy_atom_load, thr_read_buffer[read_coord], frg_in[None, 0, 0])
+            frg_vector = frg_in.load()
+            frg_vector_i32 = cute.TensorSSA(
+                vector.bitcast(T.vector(frg_in_size, T.i32()), frg_vector),
+                frg_in.shape,
+                cutlass.Int32,
+            )
+            isNotNeg0 = cute.all_(
+                cute.TensorSSA(
+                    frg_vector_i32 != frg_in_vector_neg0_i32,
+                    frg_in.shape,
+                    cutlass.Boolean,
+                )
+            )
+            # if the data is negtive 0, it means data are not ready or not visible yet, so we need to read the data again
+            while not isNotNeg0:
+                cute.copy(
+                    copy_atom_load, thr_read_buffer[read_coord], frg_in[None, 0, 0]
+                )
+                frg_vector = frg_in.load()
+                frg_vector_i32 = cute.TensorSSA(
+                    vector.bitcast(T.vector(frg_in_size, T.i32()), frg_vector),
+                    frg_in.shape,
+                    cutlass.Int32,
+                )
+                isNotNeg0 = cute.all_(
+                    cute.TensorSSA(
+                        frg_vector_i32 != frg_in_vector_neg0_i32,
+                        frg_in.shape,
+                        cutlass.Boolean,
+                    )
+                )
+            frg_acc.store(frg_in.load() + frg_acc.load())
+
+        cute.copy(copy_atom_store, frg_acc, thr_out)
+
+
+def run_all_reduce_one_shot(
+    M,
+    N,
+    warmup_iterations=2,
+    iterations=10,
+    skip_ref_check=False,
+    benchmark=True,
+):
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    if rank == 0:
+        print("\nRunning Elementwise Add test with:")
+        print(f"Tensor dimensions: [{M}, {N}]")
+        print(f"GPU count: {world_size}")
+
+    # init buffer tensors to be neg 0
+    local_buffer_tensor = nvshmem.core.tensor([PING_PONG_SIZE, world_size, M, N,], dtype=torch.float32).neg_()
+    buffer_tensor_list = [nvshmem.core.get_peer_tensor(local_buffer_tensor, rank).permute(2, 3, 1, 0) for rank in range(world_size)]
+    signal = cutlass.Int32(0)
+    input_tensor = torch.randn([M, N], device=f"cuda:{rank}")
+    output_tensor = torch.zeros([M, N], device=f"cuda:{rank}")
+    stream = cutlass.cuda.default_stream()
+    all_reduce_one_shot_lamport_kernel = AllReduceOneShotLamportKernel()
+
+    compiled_func = cute.compile(
+        all_reduce_one_shot_lamport_kernel,
+        rank,
+        world_size,
+        signal,
+        from_dlpack(input_tensor, assumed_align=32),
+        from_dlpack(output_tensor, assumed_align=32),
+        [from_dlpack(t, assumed_align=32) for t in buffer_tensor_list],
+        stream=stream,
+    )
+
+    if not skip_ref_check:
+        compiled_func(
+            signal,
+            from_dlpack(input_tensor, assumed_align=32),
+            from_dlpack(output_tensor, assumed_align=32),
+            [from_dlpack(t, assumed_align=32) for t in buffer_tensor_list],
+            stream,
+        )
+        if rank == 0:
+            print("Verifying results...")
+        dist.all_reduce(input_tensor, op=dist.ReduceOp.SUM)
+        dist.barrier(device_ids=[rank])
+        torch.testing.assert_close(input_tensor.cpu(), output_tensor.cpu())
+        if rank == 0:
+            print("Results verified successfully!")
+
+    for t in buffer_tensor_list:
+        nvshmem.core.free_tensor(t)
+
+    if not benchmark:
+        return
+
+    free_func_and_tensor_pairs = []
+    def add_free_func_and_tensor(free_func, tensor):
+        free_func_and_tensor_pairs.append((free_func, tensor))
+
+    def generate_tensors():
+        local_buffer = nvshmem.core.tensor([PING_PONG_SIZE, world_size, M, N,], dtype=torch.float32).neg_()
+        buffer_tensor_list = [nvshmem.core.get_peer_tensor(local_buffer, rank).permute(2, 3, 1, 0) for rank in range(world_size)]
+        input_tensor = torch.randn([M, N], device=f"cuda:{rank}")
+        output_tensor = torch.zeros([M, N], device=f"cuda:{rank}")
+
+        ja = testing.JitArguments(
+            cutlass.Int32(0),
+            from_dlpack(input_tensor, assumed_align=32),
+            from_dlpack(output_tensor, assumed_align=32),
+            [from_dlpack(t, assumed_align=32) for t in buffer_tensor_list],
+            stream=stream
+        )
+        for tensor in buffer_tensor_list:
+            add_free_func_and_tensor(nvshmem.core.free_tensor, tensor)
+        
+        return ja
+    avg_time_us = testing.benchmark(
+        compiled_func,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+
+    # Print execution results
+    if rank == 0:
+        print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
+        print(
+            f"Achieved memory throughput: {((world_size + 1) * output_tensor.numel() * 32 // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
+        )
+
+    for free_func, tensor in free_func_and_tensor_pairs:
+        free_func(tensor)
+
+def torchrun_uid_init_bcast():
+    """
+    Initialize NVSHMEM using UniqueID with `torchrun` as the launcher
+
+    It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting
+    """
+    # Set Torch device
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+
+    # nvshmem4py requires a cuda.core Device at init time
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    # Initialize torch.distributed process group
+    dist.init_process_group(
+        backend="cpu:gloo,cuda:nccl",
+    )
+
+    # Extract rank, nranks from process group
+    num_ranks = dist.get_world_size()
+
+    # Create an empty uniqueid for all ranks
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid")
+
+
+def torchrun_finalize():
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
+    )
+    parser.add_argument("--M", default=1024, type=int)
+    parser.add_argument("--N", default=1024, type=int)
+    parser.add_argument("--warmup_iterations", default=2, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+    args = parser.parse_args()
+
+    torchrun_uid_init_bcast()
+
+    run_all_reduce_one_shot(args.M, args.N, args.warmup_iterations, args.iterations, args.skip_ref_check, args.benchmark)
+
+    torchrun_finalize()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_simple.py

@@ -0,0 +1,312 @@
+# 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 os
+import time
+import importlib
+import argparse
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+from cutlass.cute.runtime import from_dlpack
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+"""
+A Distributed All-Reduce Addition Example using CuTe DSL and PyTorch Symmetric Memory.
+
+This example kernel demonstrates distributed all-reduce across multiple GPUs using the SIMT copy 
+of CuTe DSL and PyTorch's symmetric memory feature. Basic CuTe layout calculation is derived
+from the elementwise_add.py example.
+
+This kernel is a simple version of all-reduce. It will directly copy data from remote memory to 
+registers, then accumulate the data and finally store the accumulated data back to local global memory.
+If the input tensors from each device are remotely accessible, then this kernel can be used to perform the all-reduce.
+
+On the host side, we use `torch.distributed._symmetric_memory` to manage the symmetric memory. We use `symm_mem.empty`
+and `symm_mem.rendezvous` to create a symmetric tensor. Then we use `get_buffer` to get tensors that are accessible from all devices.
+In this way, we can hide the details of CUDA driver API calls to enable access to remote memory.
+
+.. code-block:: python
+
+    t = symm_mem.empty((M, N), device=torch.device(f"cuda:{rank}"))
+    hdl = symm_mem.rendezvous(t, dist.group.WORLD)
+    # get tensors from other devices from the symmetric memory
+    tensor_list = [hdl.get_buffer(rank, t.shape, t.dtype) for rank in range(world_size)]
+
+To run this example:
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_simple.py --M 1024 --N 512
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_simple.py \
+        --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 100
+"""
+
+
+@cute.kernel
+def all_reduce_simple_kernel(
+    inputs: list[cute.Tensor],
+    gOut: cute.Tensor,
+    thr_layout: cute.Layout,
+    val_layout: cute.Layout,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    # slice for CTAs
+    # logical id -> address
+    blk_coord = ((None, None), bidx)
+    local_tile_out = gOut[blk_coord]
+    local_tile_list = [t[blk_coord] for t in inputs]
+
+    assert all(t.element_type == inputs[0].element_type for t in inputs)
+
+    copy_atom_load = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        inputs[0].element_type,
+    )
+    copy_atom_store = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        inputs[0].element_type,
+    )
+    tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+    thr_copy = tiled_copy.get_slice(tidx)
+
+    thr_tensor_list = [thr_copy.partition_S(tensor) for tensor in local_tile_list]
+    thr_out = thr_copy.partition_D(local_tile_out)
+    frg_tensor_list = [cute.make_fragment_like(tensor) for tensor in thr_tensor_list]
+    frg_acc = cute.make_fragment_like(thr_out)
+    frg_acc.fill(0.0)
+
+    # load the frg at the same offset from all devices and accumulate the result in frg_acc
+    for thr, frg in zip(thr_tensor_list, frg_tensor_list):
+        cute.copy(copy_atom_load, thr, frg)
+        tmp = frg.load() + frg_acc.load()
+        frg_acc.store(tmp)
+
+    # copy from register memory to global memory
+    cute.copy(copy_atom_store, frg_acc, thr_out)
+
+
+@cute.jit
+def all_reduce_simple(
+    inputs: list[cute.Tensor], output: cute.Tensor, copy_bits: cutlass.Constexpr = 128
+):
+    dtype = inputs[0].element_type
+    vector_size = copy_bits // dtype.width
+
+    thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+    val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
+    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+    divided_inputs = [cute.zipped_divide(tensor, tiler_mn) for tensor in inputs]
+    gOut = cute.zipped_divide(output, tiler_mn)  # ((Tile),(Rest))
+    all_reduce_simple_kernel(
+        divided_inputs,
+        gOut,
+        thr_layout,
+        val_layout,
+    ).launch(
+        grid=[cute.size(gOut, mode=[1]), 1, 1],
+        block=[cute.size(tv_layout, mode=[0]), 1, 1],
+    )
+
+
+def run_all_reduce_simple(
+    M,
+    N,
+    warmup_iterations=2,
+    iterations=10,
+    skip_ref_check=False,
+    benchmark=True,
+):
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    if rank == 0:
+        print("\nRunning Elementwise Add test with:")
+        print(f"Tensor dimensions: [{M}, {N}]")
+        print(f"GPU count: {world_size}")
+
+    local_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+    local_tensor.random_(0, 100)
+    tensor_list = [nvshmem.core.get_peer_tensor(local_tensor, rank) for rank in range(world_size)]
+    output = torch.zeros((M, N), device=f"cuda:{rank}")
+    
+    if rank == 0:
+        print("Compiling kernel with cute.compile ...")
+    start_time = time.time()
+    compiled_func = cute.compile(all_reduce_simple, [from_dlpack(t) for t in tensor_list], from_dlpack(output))
+    compilation_time = time.time() - start_time
+    if rank == 0:
+        print(f"Compilation time: {compilation_time:.4f} seconds")
+        print("Executing vector add kernel...")
+
+    if not skip_ref_check:
+        dist.barrier(device_ids=[rank])
+        compiled_func([from_dlpack(t) for t in tensor_list], from_dlpack(output))
+        if rank == 0:
+            print("Verifying results...")
+        dist.barrier(device_ids=[rank])
+        torch.testing.assert_close(sum([t.cpu() for t in tensor_list]), output.cpu())
+        if rank == 0:
+            print("Results verified successfully!")
+    
+    for t in tensor_list:
+        nvshmem.core.free_tensor(t)
+
+    if not benchmark:
+        return
+
+    free_func_and_tensor_pairs = []
+    def add_free_func_and_tensor(free_func, tensor):
+        free_func_and_tensor_pairs.append((free_func, tensor))
+
+    def generate_tensors():
+        local_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+        local_tensor.random_(0, 100)
+        tensor_list = [nvshmem.core.get_peer_tensor(local_tensor, rank) for rank in range(world_size)]
+        output = torch.zeros((M, N), device=f"cuda:{rank}")
+
+        ja = testing.JitArguments(
+            [from_dlpack(t) for t in tensor_list],
+            from_dlpack(output),
+        )
+        for tensor in tensor_list:
+            add_free_func_and_tensor(nvshmem.core.free_tensor, tensor)
+        return ja
+   
+    avg_time_us = testing.benchmark(
+        compiled_func,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+
+    # Print execution results
+    if rank == 0:
+        print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
+        print(
+            f"Achieved memory throughput: {((world_size + 1) * output.numel() * 32 // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
+        )
+        print(f"First few elements of result: \n{output[:3, :3]}")
+
+    for free_func, tensor in free_func_and_tensor_pairs:
+        free_func(tensor)
+
+
+def torchrun_uid_init_bcast():
+    """
+    Initialize NVSHMEM using UniqueID with `torchrun` as the launcher
+
+    It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting
+    """
+    # Set Torch device
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+
+    # nvshmem4py requires a cuda.core Device at init time
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    # Initialize torch.distributed process group
+    dist.init_process_group(
+        backend="cpu:gloo,cuda:nccl",
+    )
+
+    # Extract rank, nranks from process group
+    num_ranks = dist.get_world_size()
+
+    # Create an empty uniqueid for all ranks
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid")
+
+
+def torchrun_finalize():
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
+    )
+    parser.add_argument("--M", default=1024, type=int)
+    parser.add_argument("--N", default=1024, type=int)
+    parser.add_argument("--warmup_iterations", default=2, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+
+    args = parser.parse_args()
+
+    torchrun_uid_init_bcast()
+
+    run_all_reduce_simple(args.M, args.N, args.warmup_iterations, args.iterations, args.skip_ref_check, args.benchmark)
+
+    torchrun_finalize()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_tma.py

@@ -0,0 +1,688 @@
+# 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.
+
+"""
+A Distributed All-Reduce Example using TMA (Tensor Memory Accelerator).
+
+This example demonstrates distributed all-reduce across multiple GPUs using TMA
+for data movement. It serves as a tutorial for TMA-based distributed operations,
+not as a performance-optimized implementation.
+
+Tensor Semantics:
+    - Input:  Logical shape (world_size, S), where S is the per-rank tensor size
+    - Output: Logical shape (world_size, S), each rank gets the sum of all inputs
+
+Kernel Parameters:
+    - input:  List of world_size tensors, each with shape S (accessible via NVSHMEM)
+    - output: Single tensor with shape S, using multicast address for broadcast
+
+Algorithm (Two-Shot):
+    1. Each CTA loads data from all ranks at its assigned tile position (TMA Load)
+    2. Accumulates the data locally in registers
+    3. Stores the result via TMA multicast (broadcasts to all ranks)
+    4. Cross-GPU barrier ensures completion before kernel exit
+
+Tile Assignment:
+    - Total tiles = ceil(S / elems_per_cta)
+    - Each rank processes ceil(total_tiles / world_size) CTAs
+    - CTA i on rank r processes global_tile_id = r * ctas_per_rank + i
+
+TMA Usage Notes (for tutorial purposes, not perf-optimal):
+    - Uses 1D TMA load to load from remote GPU memory via NVSHMEM addresses
+    - Uses 1D TMA load to store to multicast address for broadcasting to all ranks
+    - Supports any input shape by flattening to 1D and tiling linearly
+    - Pipeline with 2 stages overlaps TMA loads across ranks
+
+To run this example:
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8 examples/distributed/all_reduce_tma.py --shape 1024,1024
+    torchrun --nproc-per-node 8 examples/distributed/all_reduce_tma.py --shape 4,6,8,10,12
+"""
+
+import cutlass
+import cutlass.utils as utils
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass.cute.nvgpu import cpasync
+
+
+class AllReduceTmaKernel:
+    """
+    TMA-based distributed All-Reduce kernel.
+
+    This kernel performs an all-reduce operation across multiple GPUs using TMA
+    (Tensor Memory Accelerator) for efficient data movement.
+
+    Algorithm (Two-Shot):
+        1. Each CTA loads data from all ranks at its assigned tile position
+        2. Accumulates the data locally in registers
+        3. Stores the result via TMA multicast (broadcasts to all ranks)
+        4. Cross-GPU barrier ensures completion before kernel exit
+
+    The input/output tensors can be of any rank, as long as:
+        - All input tensors and output tensor share the same layout
+        - The layout is compact (no holes in memory)
+
+    We traverse the tensors linearly in codomain (physical offset) order,
+    which guarantees consistent logical coordinate access across all tensors.
+    """
+
+    _elems_per_cta: int = 128 * 128  # Elements processed per CTA
+    _tma_threads: int = 32
+    _consumer_threads: int = 128
+    _threads_per_cta: int = _tma_threads + _consumer_threads
+    _num_stages: int = 2
+
+    def __init__(self, dtype):
+        self.dtype = dtype
+
+        # SMEM layout shape (will be converted to Layout in JIT context)
+        self.smem_layout_shape = (self._elems_per_cta,)
+        self.tiler = (self._elems_per_cta,)
+
+        # TMA transaction bytes (computed from dtype size)
+        # dtype.width is in bits, divide by 8 to get bytes
+        self.tma_bytes = (dtype.width // 8) * self._elems_per_cta
+
+        # Dynamically create SharedStorage type based on dtype
+        elems = self._elems_per_cta
+        stages = self._num_stages
+
+        @cute.struct
+        class SharedStorage:
+            mbar_array: cute.struct.MemRange[cutlass.Int64, stages * 2]
+            smem_buffer: cute.struct.Align[
+                cute.struct.MemRange[dtype, elems * stages],  # stages 个 tile
+                128,
+            ]
+
+        self._SharedStorage = SharedStorage
+
+    @cute.jit
+    def __call__(
+        self,
+        input_tensors: list[cute.Tensor],
+        output_tensor_mc: cute.Tensor,
+        flag: cute.Tensor,
+        flag_mc: cute.Tensor,
+        local_rank: cutlass.Constexpr,
+        world_size: cutlass.Constexpr,
+    ):
+        """
+        Host-side JIT function: creates TMA descriptors and launches kernel.
+
+        Args:
+            input_tensors: List of input tensors from each rank (world_size tensors)
+            output_tensor_mc: Output tensor with multicast address
+            flag: Synchronization flag (local view)
+            flag_mc: Synchronization flag (multicast view)
+            local_rank: This rank's ID
+            world_size: Total number of ranks
+        """
+        # ======================================================================
+        # Layout validation
+        # ======================================================================
+        ref_layout = input_tensors[0].layout
+        ref_size = cute.size(ref_layout)
+        ref_cosize = cute.cosize(ref_layout)
+
+        # Check compact: size == cosize (no holes in memory)
+        assert ref_size == ref_cosize, (
+            f"Input tensor must be compact: size={ref_size}, cosize={ref_cosize}"
+        )
+        assert self.tma_bytes % 16 == 0, f"Not aligned to 16B, TMA should not be used."
+
+        # Check all input tensors have the same layout
+        for i in cutlass.range_constexpr(world_size):
+            assert input_tensors[i].layout == ref_layout, (
+                f"All input tensors must have the same layout. "
+                f"input_tensors[0].layout={ref_layout}, "
+                f"input_tensors[{i}].layout={input_tensors[i].layout}"
+            )
+
+        # Check output tensor has the same layout
+        assert output_tensor_mc.layout == ref_layout, (
+            f"Output tensor must have the same layout as input tensors. "
+            f"input layout={ref_layout}, output layout={output_tensor_mc.layout}"
+        )
+
+        # ======================================================================
+        # Extract tensor info
+        # ======================================================================
+        # Verify dtype matches
+        assert input_tensors[0].element_type == self.dtype, (
+            f"Input tensor dtype mismatch: expected {self.dtype}, "
+            f"got {input_tensors[0].element_type}"
+        )
+
+        total_elems = ref_size
+
+        # Flatten layout: treat tensor as 1D in codomain order
+        flat_layout = cute.make_layout((total_elems,))
+
+        # SMEM layout (created in JIT context)
+        smem_layout = cute.make_layout(self.smem_layout_shape)
+
+        # Create TMA load descriptors (one per rank)
+        tma_load_op = cpasync.CopyBulkTensorTileG2SOp()
+        tma_load_atoms = []
+        tma_load_tensors = []
+
+        for i in cutlass.range_constexpr(world_size):
+            flat_input = cute.make_tensor(input_tensors[i].iterator, flat_layout)
+            tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(
+                tma_load_op,
+                flat_input,
+                smem_layout,
+                self.tiler,
+            )
+            tma_load_atoms.append(tma_atom)
+            tma_load_tensors.append(tma_tensor)
+
+        # Create TMA store descriptor
+        tma_store_op = cpasync.CopyBulkTensorTileS2GOp()
+        flat_output = cute.make_tensor(output_tensor_mc.iterator, flat_layout)
+        tma_store_atom, tma_store_tensor = cpasync.make_tiled_tma_atom(
+            tma_store_op,
+            flat_output,
+            smem_layout,
+            self.tiler,
+        )
+
+        # Grid calculation
+        num_tiles_total = cute.ceil_div(total_elems, self._elems_per_cta)
+        ctas_per_rank = cute.ceil_div(num_tiles_total, world_size)
+
+        # SMEM size from SharedStorage
+        smem_bytes = self._SharedStorage.size_in_bytes()
+
+        # Launch kernel
+        self.kernel(
+            tma_load_atoms,
+            tma_load_tensors,
+            tma_store_atom,
+            tma_store_tensor,
+            flag,
+            flag_mc,
+            local_rank,
+            world_size,
+            num_tiles_total,
+            ctas_per_rank,
+        ).launch(
+            grid=[ctas_per_rank, 1, 1],
+            block=[self._threads_per_cta, 1, 1],
+            smem=smem_bytes,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        # TMA atoms and tensors for loading from each rank
+        tma_load_atoms: list[cute.CopyAtom],
+        tma_load_tensors: list[cute.Tensor],
+        # TMA atom and tensor for storing to multicast address
+        tma_store_atom: cute.CopyAtom,
+        tma_store_tensor: cute.Tensor,
+        # Synchronization flags
+        flag: cute.Tensor,
+        flag_mc: cute.Tensor,
+        # Rank info
+        local_rank: cutlass.Constexpr,
+        world_size: cutlass.Constexpr,
+        # Grid info for tile calculation
+        num_tiles_total: cutlass.Constexpr,
+        ctas_per_rank: cutlass.Constexpr,
+    ):
+        # ======================================================================
+        # Thread/Block indexing
+        # ======================================================================
+        tidx = cute.arch.thread_idx()[0]
+        bidx = cute.arch.block_idx()[0]
+        warp_idx = cute.arch.warp_idx()
+        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+
+        # ======================================================================
+        # SMEM allocation
+        # ======================================================================
+        staged_smem_layout = cute.make_layout((self._elems_per_cta, self._num_stages))
+
+        smem = utils.SmemAllocator()
+        storage = smem.allocate(self._SharedStorage)
+        mbar_ptr = storage.mbar_array.data_ptr()
+        staged_smem_tensor = storage.smem_buffer.get_tensor(staged_smem_layout)
+
+        # ======================================================================
+        # TMA Pipeline setup
+        # ======================================================================
+        producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+        consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, self._consumer_threads
+        )
+
+        tma_pipeline = pipeline.PipelineTmaAsync.create(
+            barrier_storage=mbar_ptr,
+            num_stages=self._num_stages,
+            producer_group=producer_group,
+            consumer_group=consumer_group,
+            tx_count=self.tma_bytes,
+            cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)),
+        )
+
+        global_tile_id = local_rank * ctas_per_rank + bidx
+
+        if global_tile_id < num_tiles_total:
+            # ======================================================================
+            # Warp 0: Producer - TMA Load from all ranks
+            # ======================================================================
+            if warp_idx == 0:
+                producer_state = pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Producer, self._num_stages
+                )
+
+                for rank_i in cutlass.range_constexpr(world_size):
+                    tma_pipeline.producer_acquire(producer_state)
+
+                    stage_idx = producer_state.index
+                    smem_tile = cute.slice_(staged_smem_tensor, (None, stage_idx))
+
+                    g_tensor_tiled = cute.zipped_divide(
+                        tma_load_tensors[rank_i], self.tiler
+                    )
+                    g_tile = g_tensor_tiled[(None,), global_tile_id]
+
+                    g_tile_flat = cute.group_modes(g_tile, 0, cute.rank(g_tile))
+                    s_tile_flat = cute.group_modes(smem_tile, 0, cute.rank(smem_tile))
+
+                    s_part, g_part = cute.nvgpu.cpasync.tma_partition(
+                        tma_load_atoms[rank_i],
+                        0,
+                        cute.make_layout(1),
+                        s_tile_flat,
+                        g_tile_flat,
+                    )
+
+                    cute.copy(
+                        tma_load_atoms[rank_i],
+                        g_part,
+                        s_part,
+                        tma_bar_ptr=tma_pipeline.producer_get_barrier(producer_state),
+                    )
+
+                    tma_pipeline.producer_commit(producer_state)
+                    producer_state.advance()
+
+            # ======================================================================
+            # Warp 1-4: Consumer - Load from smem, ADD, Store to smem
+            # ======================================================================
+            else:
+                consumer_tid = tidx - self._tma_threads
+
+                vec_size = 4
+                chunk_size = vec_size * self._consumer_threads
+
+                # ------------------------------------------------------------------
+                # Initialize accumulator using stage 0's layout
+                # ------------------------------------------------------------------
+                # (elems, stages) -> (elems,)
+                smem_tensor_wo_stage = cute.slice_(staged_smem_tensor, (None, 0))
+                # (elems,) -> ((thr_vec,), (num_chunks,))
+                smem_tensor_tiled_by_thr_vec = cute.zipped_divide(
+                    smem_tensor_wo_stage, (chunk_size,)
+                )
+                # ((thr_vec,), (num_chunks,)) -> (((vec, threads),), (num_chunks,))
+                smem_tensor_tiled_by_thr_vec_tiled_by_vec = cute.logical_divide(
+                    smem_tensor_tiled_by_thr_vec, (vec_size,)
+                )
+                # (((vec, threads),), (num_chunks,)) -> ((vec,), (num_chunks,))
+                per_thread_smem_tensor = cute.slice_(
+                    smem_tensor_tiled_by_thr_vec_tiled_by_vec,
+                    ((None, consumer_tid), None),
+                )
+
+                accum = cute.make_rmem_tensor(per_thread_smem_tensor.layout, self.dtype)
+                accum.fill(self.dtype(0.0))
+
+                # ------------------------------------------------------------------
+                # Main loop: load from SMEM and accumulate
+                # ------------------------------------------------------------------
+                consumer_state = pipeline.make_pipeline_state(
+                    pipeline.PipelineUserType.Consumer, self._num_stages
+                )
+
+                for rank_i in cutlass.range_constexpr(world_size):
+                    tma_pipeline.consumer_wait(consumer_state)
+
+                    stage_idx = consumer_state.index
+                    smem_tile = cute.slice_(staged_smem_tensor, (None, stage_idx))
+
+                    # (elems,) -> ((thr_vec,), (num_chunks,))
+                    smem_tiled_by_thr_vec = cute.zipped_divide(smem_tile, (chunk_size,))
+                    # ((thr_vec,), (num_chunks,)) -> (((vec, threads),), (num_chunks,))
+                    smem_tiled_by_thr_vec_tiled_by_vec = cute.logical_divide(
+                        smem_tiled_by_thr_vec, (vec_size,)
+                    )
+                    # (((vec, threads),), (num_chunks,)) -> ((vec,), (num_chunks,))
+                    per_thread_smem_view = cute.slice_(
+                        smem_tiled_by_thr_vec_tiled_by_vec,
+                        ((None, consumer_tid), None),
+                    )
+
+                    fragment = per_thread_smem_view.load()
+                    accum.store(accum.load() + fragment)
+
+                    tma_pipeline.sync_object_empty.arrive(
+                        consumer_state.index, tma_pipeline.consumer_mask
+                    )
+                    consumer_state.advance()
+
+                # Store accumulated result back to SMEM (stage 0)
+                per_thread_smem_tensor.store(accum.load())
+
+            # ======================================================================
+            # Sync point: all warps meet here
+            # ======================================================================
+            cute.arch.sync_threads()
+
+            # ======================================================================
+            # Warp 0: TMA Store to multicast output
+            # ======================================================================
+            if warp_idx == 0:
+                # Fence to ensure SMEM writes are visible
+                cute.arch.fence_proxy("async.shared", space="cta")
+
+                smem_tile_out = cute.slice_(staged_smem_tensor, (None, 0))
+
+                g_output_tiled = cute.zipped_divide(tma_store_tensor, self.tiler)
+                g_output_tile = g_output_tiled[(None,), global_tile_id]
+
+                g_out_flat = cute.group_modes(
+                    g_output_tile, 0, cute.rank(g_output_tile)
+                )
+                s_out_flat = cute.group_modes(
+                    smem_tile_out, 0, cute.rank(smem_tile_out)
+                )
+
+                s_part, g_part = cute.nvgpu.cpasync.tma_partition(
+                    tma_store_atom,
+                    0,
+                    cute.make_layout(1),
+                    s_out_flat,
+                    g_out_flat,
+                )
+
+                cute.copy(tma_store_atom, s_part, g_part)
+                cute.arch.cp_async_bulk_commit_group()
+                cute.arch.cp_async_bulk_wait_group(0)
+
+        # ==================================================================
+        # Cross-GPU barrier synchronization (thread 0 only)
+        # ==================================================================
+        if tidx == 0:
+            sm_id_linear = (
+                cute.arch.block_idx()[0]
+                + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0]
+                + cute.arch.block_idx()[2]
+                * cute.arch.grid_dim()[0]
+                * cute.arch.grid_dim()[1]
+            )
+
+            # Signal completion to all ranks
+            utils.distributed.multimem_red_add1(
+                flag_mc.iterator + sm_id_linear,
+                scope="sys",
+                order="release",
+            )
+
+            # The same idx ctas wait until all peer ranks' ctas complete
+            utils.distributed.spin_lock_atom_cas_relaxed_wait(
+                flag.iterator + sm_id_linear,
+                expected_val=world_size,
+                reset_val=0,
+                scope="sys",
+            )
+
+
+# =============================================================================
+# HOST-SIDE DRIVER CODE
+# =============================================================================
+
+import os
+import argparse
+import math
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+from cutlass.cute.runtime import from_dlpack
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+
+def torchrun_uid_init_bcast():
+    """Initialize NVSHMEM using UniqueID with torchrun as launcher."""
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    dist.init_process_group(backend="cpu:gloo,cuda:nccl")
+    num_ranks = dist.get_world_size()
+
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(
+        device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid"
+    )
+
+
+def torchrun_finalize():
+    """Finalize NVSHMEM and destroy process group."""
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def run_all_reduce_tma(
+    shape: tuple,
+    skip_ref_check: bool = False,
+):
+    """
+    Run the TMA-based All-Reduce kernel.
+
+    Args:
+        shape: Tensor shape tuple, e.g., (4, 6, 8, 10)
+        skip_ref_check: If True, skip reference result verification
+    """
+    local_rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    # Calculate total elements
+    total_elems = math.prod(shape)
+
+    if local_rank == 0:
+        print("\nRunning TMA All-Reduce test with:")
+        print(f"  Tensor shape: {shape}")
+        print(f"  Total elements: {total_elems}")
+        print(f"  GPU count: {world_size}")
+
+    # Allocate input tensor (symmetric memory, accessible from all ranks)
+    local_input_tensor = nvshmem.core.tensor(shape, dtype=torch.float32)
+    local_input_tensor.random_(0, 100)
+
+    # Get peer tensors (views into each rank's input)
+    peer_input_tensors = [
+        nvshmem.core.get_peer_tensor(local_input_tensor, r) for r in range(world_size)
+    ]
+
+    if local_rank == 0:
+        print(f"  Input tensor ptr: {local_input_tensor.data_ptr():#x}")
+
+    # Allocate output tensor with multicast address
+    local_output_tensor = nvshmem.core.tensor(shape, dtype=torch.float32)
+    local_output_tensor.fill_(0)
+    output_tensor_mc = nvshmem.core.get_multicast_tensor(
+        nvshmem.core.Teams.TEAM_NODE, local_output_tensor
+    )
+
+    # Allocate synchronization flags
+    # Flag size = ctas_per_rank (matches kernel's bidx indexing)
+    elems_per_cta = AllReduceTmaKernel._elems_per_cta
+    num_tiles = (total_elems + elems_per_cta - 1) // elems_per_cta
+    ctas_per_rank = (num_tiles + world_size - 1) // world_size
+    local_flag = nvshmem.core.tensor((ctas_per_rank,), dtype=torch.int32)
+    local_flag.fill_(0)
+    flag_mc = nvshmem.core.get_multicast_tensor(
+        nvshmem.core.Teams.TEAM_NODE, local_flag
+    )
+
+    if local_rank == 0:
+        print(f"  Number of tiles: {num_tiles}")
+        print(f"  CTAs per rank: {ctas_per_rank}")
+        print("Compiling kernel...")
+
+    # Create kernel instance and compile
+    kernel = AllReduceTmaKernel(cutlass.Float32)
+
+    compiled_func = cute.compile(
+        kernel,
+        [from_dlpack(t) for t in peer_input_tensors],
+        from_dlpack(output_tensor_mc),
+        from_dlpack(local_flag),
+        from_dlpack(flag_mc),
+        local_rank,
+        world_size,
+    )
+
+    if local_rank == 0:
+        print("Compilation successful!")
+
+    if not skip_ref_check:
+        if local_rank == 0:
+            print("Executing kernel...")
+
+        dist.barrier(device_ids=[local_rank])
+        compiled_func(
+            [from_dlpack(t) for t in peer_input_tensors],
+            from_dlpack(output_tensor_mc),
+            from_dlpack(local_flag),
+            from_dlpack(flag_mc),
+        )
+        dist.barrier(device_ids=[local_rank])
+
+        if local_rank == 0:
+            print("Verifying results...")
+
+        # Compute expected result: sum of all inputs
+        expected = sum([t.cpu() for t in peer_input_tensors])
+
+        # Compare with actual output
+        torch.testing.assert_close(expected, local_output_tensor.cpu())
+
+        if local_rank == 0:
+            print("Results verified successfully!")
+
+    # Cleanup
+    for i in range(world_size):
+        if i != local_rank:
+            nvshmem.core.free_tensor(peer_input_tensors[i])
+
+    nvshmem.core.free_tensor(output_tensor_mc)
+    nvshmem.core.free_tensor(flag_mc)
+    nvshmem.core.free_tensor(local_input_tensor)
+    nvshmem.core.free_tensor(local_output_tensor)
+    nvshmem.core.free_tensor(local_flag)
+
+
+def parse_shape(shape_str: str) -> tuple:
+    """
+    Parse shape string into tuple.
+    Examples:
+        "1024,1024" -> (1024, 1024)
+        "2,3,4,5,6,7,8" -> (2, 3, 4, 5, 6, 7, 8)
+    """
+    return tuple(int(x.strip()) for x in shape_str.split(","))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="TMA-based distributed all-reduce example"
+    )
+    parser.add_argument(
+        "--shape",
+        default="1024,1024",
+        type=str,
+        help="Tensor shape as comma-separated values, e.g., '1024,1024' or 4,6,8,10,12'",
+    )
+    parser.add_argument(
+        "--skip_ref_check",
+        action="store_true",
+        help="Skip reference result verification",
+    )
+
+    args = parser.parse_args()
+    shape = parse_shape(args.shape)
+
+    torchrun_uid_init_bcast()
+    run_all_reduce_tma(
+        shape=shape,
+        skip_ref_check=args.skip_ref_check,
+    )
+    torchrun_finalize()
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/distributed/all_reduce_two_shot_multimem.py

@@ -0,0 +1,393 @@
+# 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 os
+import time
+import argparse
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from cuda.core.experimental import Device
+from cuda.pathfinder import load_nvidia_dynamic_lib
+
+import cutlass
+import cutlass.utils as utils
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+import cutlass.torch as cutlass_torch
+from cutlass.cute.runtime import from_dlpack
+
+try:
+    import nvshmem.core
+except ImportError as exc:
+    raise ImportError(
+        "nvshmem4py is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvshmem4py-cu12\n"
+        "  For CUDA 13: pip install nvshmem4py-cu13\n"
+        "Note: nvshmem4py version >= 0.1.3 is recommended."
+    ) from None
+
+try:
+    load_nvidia_dynamic_lib("nvshmem_host")
+except RuntimeError as exc:
+    raise ImportError(
+        "nvshmem lib is required but not installed. Please install it using:\n"
+        "  For CUDA 12: pip install nvidia-nvshmem-cu12\n"
+        "  For CUDA 13: pip install nvidia-nvshmem-cu13\n"
+    ) from None
+
+
+"""
+A Distributed Two-Shot All-Reduce Example using CuTe DSL and PyTorch Symmetric Memory.
+
+This example kernel demonstrates how to leverage the multimem feature to do a two-shot all-reduce.
+The multimem instruction is operated on symmetric memory, it can offload the broadcast and reduce 
+to the Nvlink Switch so that the nvlink traffic will be reduced.
+
+When calling a 'multimem.ld_reduce addrA', the corresponding data from each remote device will be sent to the NVLS 
+and return the reduced data as result. And for 'multimem.st dataA addrA', the data will be sent to the NVLS once and 
+the data will be broadcast to each remote device. So the memory traffic and instruction count is reduced by 8 times 
+with multimem.
+
+In this example, we are using two-shot styled all-reduce which means each device computes a portion
+of data and stores them to each device. Compared to the one-shot styled all-reduce, the two-shot one can 
+maximize the performance of throughput. The input and output are symmetric memory so we don't need extra 
+communication buffers here. We use the `sm_wise_inter_gpu_multimem_barrier` to synchronize the data 
+between each device. It is to make sure that each device has done the data transfer.
+
+To run this example:
+
+.. code-block:: bash
+
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_two_shot_multimem.py --M 1024 --N 512
+    torchrun --nproc-per-node 8  examples/distributed/all_reduce_two_shot_multimem.py \
+        --M 1024 --N 1024 --benchmark --warmup_iterations 2 --iterations 100
+"""
+
+
+@cute.kernel
+def all_reduce_multimem_kernel(
+    gIn: cute.Tensor,
+    gOut: cute.Tensor,
+    flag: cute.Tensor,
+    flag_mc: cute.Tensor,
+    thr_layout: cute.Layout,
+    val_layout: cute.Layout,
+    local_rank: cutlass.Constexpr,
+    world_size: cutlass.Constexpr,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    bidx, _, _ = cute.arch.block_idx()
+
+    # slice for CTAs
+    # logical id -> address
+
+    num_ctas = cute.size(gIn, mode=[1])
+    chunk_size = num_ctas // world_size
+    blk_idx = local_rank * chunk_size + bidx
+
+    blk_coord = ((None, None), blk_idx)
+    local_tile_out = gOut[blk_coord]
+    local_tile_in = gIn[blk_coord]
+
+    assert gIn.element_type == gOut.element_type
+
+    copy_atom_load = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        gIn.element_type,
+        num_bits_per_copy=128,
+    )
+    tiled_copy = cute.make_tiled_copy_tv(copy_atom_load, thr_layout, val_layout)
+    thr_copy = tiled_copy.get_slice(tidx)
+
+    thr_in = thr_copy.partition_S(local_tile_in)
+    thr_out = thr_copy.partition_D(local_tile_out)
+
+    (_, rest_m), _, _ = thr_in.shape
+    (_, rest_m_stride), _, _ = thr_in.stride
+
+    for i in cutlass.range_constexpr(rest_m):
+        x, y, z, w = utils.distributed.multimem_ld_reduce_4xf32(
+            thr_in[(None, i), 0, 0].iterator
+        )
+        utils.distributed.multimem_st_4xb32(
+            thr_out[(None, i), 0, 0].iterator, x, y, z, w
+        )
+
+    # Ensure all threads in cta have finish issue multimem.ld_reduce and multimem.st instructions
+    cute.arch.sync_threads()
+
+    if tidx == 0:
+        # Linear id of current SM.
+        sm_id_linear = (
+            cute.arch.block_idx()[0]
+            + cute.arch.block_idx()[1] * cute.arch.grid_dim()[0]
+            + cute.arch.block_idx()[2]
+            * cute.arch.grid_dim()[0]
+            * cute.arch.grid_dim()[1]
+        )
+        # Release flag with sys scope
+        utils.distributed.multimem_red_add1(
+            flag_mc.iterator + sm_id_linear,
+            scope="sys",
+            order="release",
+        )
+        # Relaxed spin-lock wait flag with sys scope
+        utils.distributed.spin_lock_atom_cas_relaxed_wait(
+            flag.iterator + sm_id_linear,
+            expected_val=world_size,
+            reset_val=0,
+            scope="sys",
+        )
+
+@cute.jit
+def all_reduce_multimem(
+    mIn: cute.Tensor,
+    mOut: cute.Tensor,
+    flag: cute.Tensor,
+    flag_mc: cute.Tensor,
+    local_rank: cutlass.Constexpr,
+    world_size: cutlass.Constexpr,
+    copy_bits: cutlass.Constexpr = 128,
+):
+    dtype = mIn.element_type
+    vector_size = copy_bits // dtype.width
+
+    # we choose a 128x128 tile for a CTA
+    thr_layout = cute.make_ordered_layout((4, 32), order=(1, 0))
+    val_layout = cute.make_ordered_layout((32, vector_size), order=(1, 0))
+    tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
+
+    gIn = cute.zipped_divide(mIn, tiler_mn)
+    gOut = cute.zipped_divide(mOut, tiler_mn)
+
+    all_reduce_multimem_kernel(
+        gIn,
+        gOut,
+        flag,
+        flag_mc,
+        thr_layout,
+        val_layout,
+        local_rank,
+        world_size,
+    ).launch(
+        grid=[cute.size(gOut, mode=[1]) // world_size, 1, 1],
+        block=[cute.size(tv_layout, mode=[0]), 1, 1],
+    )
+
+
+def run_all_reduce_multimem(
+    M,
+    N,
+    warmup_iterations=2,
+    iterations=10,
+    skip_ref_check=False,
+    benchmark=True,
+):
+    local_rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    tile_m = 128
+    tile_n = 128
+
+    if local_rank == 0:
+        print("\nRunning Elementwise Add test with:")
+        print(f"Tensor dimensions: [{M}, {N}]")
+        print(f"GPU count: {world_size}")
+
+    local_input_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+    input_tensor = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_input_tensor)
+
+    local_output_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+    output_tensor = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_output_tensor)
+
+    local_flag = nvshmem.core.tensor((M*N//(tile_m*tile_n)), dtype=torch.int32)
+    flag_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_flag)
+    
+    if local_rank == 0:
+        print("Compiling kernel with cute.compile ...")
+    start_time = time.time()
+    compiled_func = cute.compile(
+        all_reduce_multimem,
+        from_dlpack(input_tensor),
+        from_dlpack(output_tensor),
+        from_dlpack(local_flag),
+        from_dlpack(flag_mc),
+        local_rank,
+        world_size,
+    )
+    compilation_time = time.time() - start_time
+    if local_rank == 0:
+        print(f"Compilation time: {compilation_time:.4f} seconds")
+        print("Executing all-reduce two shot multimem kernel...")
+
+    if not skip_ref_check:
+        dist.barrier(device_ids=[local_rank])
+        compiled_func(
+            from_dlpack(input_tensor),
+            from_dlpack(output_tensor),
+            from_dlpack(local_flag),
+            from_dlpack(flag_mc),
+        )
+        dist.barrier(device_ids=[local_rank])
+        if local_rank == 0:
+            print("Verifying results...")
+
+        local_buffers = [nvshmem.core.get_peer_tensor(local_input_tensor, local_rank) for local_rank in range(world_size)]
+        torch.testing.assert_close(sum([buffer.cpu() for buffer in local_buffers]), local_output_tensor.cpu())
+        if local_rank == 0:
+            print("Results verified successfully!")
+        for i in range(world_size):
+            if i != local_rank:
+                nvshmem.core.free_tensor(local_buffers[i])
+
+    # always free the multicast tensors first
+    nvshmem.core.free_tensor(input_tensor)
+    nvshmem.core.free_tensor(output_tensor)
+    nvshmem.core.free_tensor(flag_mc)
+    nvshmem.core.free_tensor(local_input_tensor)
+    nvshmem.core.free_tensor(local_output_tensor)
+    nvshmem.core.free_tensor(local_flag)
+
+    if not benchmark:
+        return
+    
+    free_func_and_tensor_pairs = []
+    def add_free_func_and_tensor(free_func, tensor):
+        free_func_and_tensor_pairs.append((free_func, tensor))
+
+    def generate_tensors():
+        local_input_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+        input_tensor_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_input_tensor)
+
+        local_output_tensor = nvshmem.core.tensor((M, N), dtype=torch.float32)
+        output_tensor_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_output_tensor)
+
+        local_flag = nvshmem.core.tensor((M*N//(tile_m*tile_n)), dtype=torch.int32)
+        flag_mc = nvshmem.core.get_multicast_tensor(nvshmem.core.Teams.TEAM_NODE, local_flag)
+
+        ja = testing.JitArguments(
+            from_dlpack(input_tensor_mc),
+            from_dlpack(output_tensor_mc),
+            from_dlpack(local_flag),
+            from_dlpack(flag_mc),
+        )
+        tensors_to_free = [input_tensor_mc, output_tensor_mc, flag_mc, local_input_tensor, local_output_tensor, local_flag]
+        for tensor in tensors_to_free:
+            add_free_func_and_tensor(nvshmem.core.free_tensor, tensor)
+        return ja
+
+    dist.barrier(device_ids=[local_rank])
+    avg_time_us = testing.benchmark(
+        compiled_func,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+    )
+    dist.barrier(device_ids=[local_rank])
+    torch.cuda.synchronize()
+
+    # Print execution results
+    if local_rank == 0:
+        print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
+        print(
+            f"Achieved memory throughput: {((world_size + 1) * output_tensor.numel() * 32 // 8) / (avg_time_us / 1e6) / 1e9:.2f} GB/s"
+        )
+    
+    for free_func, tensor in free_func_and_tensor_pairs:
+        free_func(tensor)
+    return
+
+
+def torchrun_uid_init_bcast():
+    """
+    Initialize NVSHMEM using UniqueID with `torchrun` as the launcher
+
+    It uses torch.distributed.broadcast on a NumPy array to handle the broadcasting
+    """
+    # Set Torch device
+    local_rank = int(os.environ['LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+
+    # nvshmem4py requires a cuda.core Device at init time
+    dev = Device(local_rank)
+    dev.set_current()
+    global stream
+    stream = dev.create_stream()
+
+    # Initialize torch.distributed process group
+    dist.init_process_group(
+        backend="cpu:gloo,cuda:nccl",
+    )
+
+    # Extract rank, nranks from process group
+    num_ranks = dist.get_world_size()
+
+    # Create an empty uniqueid for all ranks
+    uid = nvshmem.core.get_unique_id(empty=(local_rank != 0))
+    uid_bytes = uid._data.view(np.uint8).copy()
+    uid_tensor = torch.from_numpy(uid_bytes).cuda()
+    dist.broadcast(uid_tensor, src=0)
+    dist.barrier()
+    uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype)
+
+    nvshmem.core.init(device=dev, uid=uid, rank=local_rank, nranks=num_ranks, initializer_method="uid")
+
+
+def torchrun_finalize():
+    nvshmem.core.finalize()
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="example of elementwise add to demonstrate the numpy/pytorch as input for kernels"
+    )
+    parser.add_argument("--M", default=1024, type=int)
+    parser.add_argument("--N", default=1024, type=int)
+    parser.add_argument("--warmup_iterations", default=2, type=int)
+    parser.add_argument("--iterations", default=10, type=int)
+    parser.add_argument("--skip_ref_check", action="store_true")
+    parser.add_argument("--benchmark", action="store_true")
+
+    args = parser.parse_args()
+
+    torchrun_uid_init_bcast()
+
+    run_all_reduce_multimem(args.M, args.N, args.warmup_iterations, args.iterations, args.skip_ref_check, args.benchmark)
+
+    torchrun_finalize()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/python/CuTeDSL/cute/blackwell/kernel/mixed_input_gemm/mixed_input_host_utils.py

@@ -0,0 +1,506 @@
+# 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.
+
+from typing import Optional
+
+import torch
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.torch as cutlass_torch
+import cutlass.utils.mixed_input_helpers as mixed_input_utils
+from cutlass.cute.runtime import from_dlpack
+
+"""
+This file contains common host-side utilities for mixed-input GEMM.
+"""
+
+
+def create_cumsum_tensor(
+    num_groups: int,
+    fused_n: int,
+    alignment: int,
+    uniform_distribution: bool = False,
+) -> tuple[cute.Tensor, torch.Tensor]:
+    """
+    Create a tensor of shape (num_groups + 1) recording the cumulative sum of the elements in each group.
+    """
+    assert fused_n % alignment == 0, "fused_n must be divisible by alignment"
+    if uniform_distribution:
+        # keep a uniform distribution for debug and performance collection
+        group_counts = torch.tensor([fused_n // num_groups] * num_groups)
+    else:
+        # sample group sizes with equal probability for each group
+        probs = torch.ones(num_groups) / num_groups
+        group_sizes = torch.multinomial(probs, fused_n // alignment, replacement=True)
+        group_counts = torch.bincount(group_sizes, minlength=num_groups) * alignment
+    print(group_counts.tolist())
+
+    # Create cumulative sum
+    cumsum_torch = torch.cat([torch.tensor([0]), group_counts.cumsum(0)])
+    print(cumsum_torch.tolist())
+
+    cumsum_tensor, _ = cutlass_torch.cute_tensor_like(
+        cumsum_torch, cutlass.Int32, is_dynamic_layout=False
+    )
+
+    return cumsum_tensor, cumsum_torch.to("cpu")
+
+
+def create_i4_tensor_and_scale(
+    l: int,
+    m: int,
+    k: int,
+    is_m_major: bool,
+    dtype: type[cutlass.Numeric],
+    shuffle_a: bool,
+    scale_granularity_m: int,
+    scale_granularity_k: int,
+    is_dynamic_layout: bool = True,
+    init_config: tuple = (
+        cutlass_torch.TensorInitType.RANDOM,
+        cutlass_torch.RandomInitConfig(min_val=-7, max_val=6),
+    ),
+    divisibility: int = 16,
+    transformed_dtype: Optional[type[cutlass.Numeric]] = None,
+) -> tuple[
+    cute.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    cute.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+]:
+    """
+    Create quantized 4-bit tensor and corresponding scale tensor.
+    """
+    lb_4b = -8 if dtype == cutlass.Int4 else 0
+    up_4b = 7 if dtype == cutlass.Int4 else 15
+    if not (
+        init_config[0] == cutlass_torch.TensorInitType.RANDOM
+        or init_config[0] == cutlass_torch.TensorInitType.SCALAR
+    ):
+        raise ValueError(
+            "Only random and scalar initialization is supported for 4bit data type"
+        )
+
+    # Construct reference tensor in f32
+    ref_fp32 = cutlass_torch.matrix(l, m, k, is_m_major, cutlass.Float32, *init_config)
+    # Generate scale data and perform quantization
+    num_scales = k // scale_granularity_k
+    ref = ref_fp32.to(dtype=cutlass_torch.dtype(transformed_dtype)).reshape(
+        m, num_scales, scale_granularity_k, l
+    )
+    # Get elements with maximum absolute value to compute scaling factors
+    a_max = (
+        torch.maximum(ref / up_4b, ref / lb_4b)
+        if dtype == cutlass.Int4
+        else ref / up_4b
+    )
+    a_scales, _ = torch.max(a_max, dim=2, keepdim=True)
+    a_scale_inv = torch.where(a_scales == 0, 0, 1 / a_scales)
+    a_quant = ref * a_scale_inv
+    # Convert values to integer to avoid computation errors
+    a_quant = a_quant.to(dtype=torch.int32).reshape((m, k, l)).to(dtype=torch.float32)
+    # Construct cute scale tensor
+    a_scales = a_scales.random_(-3, 3).reshape((m, num_scales, l))
+    # Scale tensor is always m-major
+    a_scales = a_scales.permute(2, 1, 0).contiguous().permute(2, 1, 0).to(device="cuda")
+    if shuffle_a:
+        # shuffle within each group of 8 elements
+        perm = torch.tensor([0, 2, 1, 3, 4, 6, 5, 7], device=a_quant.device)
+        a_shuffled = (
+            a_quant.view(m, k // 8, 8, l)[:, :, perm, :]
+            .reshape(a_quant.shape)
+            .permute(2, 0, 1)
+            .contiguous()
+            .permute(1, 2, 0)
+        )
+        # Construct A quantized tensor
+        cute_a_quant_tensor, torch_a_quant_tensor = cutlass_torch.cute_tensor_like(
+            a_shuffled,
+            dtype,
+            is_dynamic_layout=is_dynamic_layout,
+            assumed_align=divisibility,
+        )
+    else:
+        # Construct A quantized tensor
+        cute_a_quant_tensor, torch_a_quant_tensor = cutlass_torch.cute_tensor_like(
+            a_quant,
+            dtype,
+            is_dynamic_layout=is_dynamic_layout,
+            assumed_align=divisibility,
+        )
+    cute_scale_tensor = from_dlpack(a_scales, assumed_align=divisibility)
+    for i, stride in enumerate(a_scales.stride()):
+        if stride == 1:
+            leading_dim = i
+            break
+    if is_dynamic_layout:
+        cute_scale_tensor = cute_scale_tensor.mark_layout_dynamic(
+            leading_dim=leading_dim
+        )
+
+    return (
+        cute_a_quant_tensor,
+        torch_a_quant_tensor,
+        a_quant.to("cpu"),
+        cute_scale_tensor,
+        a_scales,
+        a_scales.to("cpu"),
+    )
+
+
+def create_tensor_a(
+    l: int,
+    m: int,
+    k: int,
+    a_major: str,
+    a_dtype: type[cutlass.Numeric],
+    shuffle_a: bool,
+    scale_granularity_m: int = 0,
+    scale_granularity_k: int = 0,
+    transformed_dtype: Optional[type[cutlass.Numeric]] = None,
+) -> tuple[cute.Tensor, Optional[cute.Tensor], torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Create tensor A and scale tensor.
+    """
+    a_scale_tensor = None
+    a_scale_torch_cpu = None
+    if a_dtype in (cutlass.Int4,):
+        (
+            a_tensor,
+            a_torch_gpu,
+            a_torch_cpu,
+            a_scale_tensor,
+            a_scale_torch_gpu,
+            a_scale_torch_cpu,
+        ) = create_i4_tensor_and_scale(
+            l,
+            m,
+            k,
+            a_major == "m",
+            a_dtype,
+            shuffle_a,
+            scale_granularity_m,
+            scale_granularity_k,
+            divisibility=mixed_input_utils.get_divisibility(m if a_major == "m" else k),
+            transformed_dtype=transformed_dtype,
+        )
+    else:
+        a_torch_cpu = cutlass_torch.matrix(
+            l,
+            m,
+            k,
+            a_major == "m",
+            a_dtype,
+        )
+        a_tensor, _ = cutlass_torch.cute_tensor_like(
+            a_torch_cpu,
+            a_dtype,
+            is_dynamic_layout=True,
+            assumed_align=mixed_input_utils.get_divisibility(
+                m if a_major == "m" else k
+            ),
+        )
+    return a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu
+
+
+def create_tensors_for_contiguous_grouped_mixed_input_gemm(
+    l: int,
+    m: int,
+    n: int,
+    k: int,
+    a_major: str,
+    b_major: str,
+    c_major: str,
+    a_dtype: type[cutlass.Numeric],
+    b_dtype: type[cutlass.Numeric],
+    c_dtype: type[cutlass.Numeric],
+    shuffle_a: bool = False,
+    scale_granularity_m: int = 0,
+    scale_granularity_k: int = 0,
+    uniform_group_sizes: bool = False,
+) -> tuple:
+    """
+    Create all input and output tensors for the contiguous grouped mixed-input GEMM.
+    """
+    a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a(
+        l,
+        m,
+        k,
+        a_major,
+        a_dtype,
+        shuffle_a,
+        scale_granularity_m,
+        scale_granularity_k,
+        b_dtype,
+    )
+
+    # In GROUP mode, l specifies the number of groups. We'll fuse group into the n mode for tensor B and C.
+    # Batch mode will be set to 1.
+    num_groups = l
+    fused_n = n * num_groups
+    b_torch_cpu = cutlass_torch.matrix(
+        1,  # batch=1
+        fused_n,
+        k,
+        b_major == "n",
+        b_dtype,
+        cutlass_torch.TensorInitType.RANDOM,
+        cutlass_torch.RandomInitConfig(min_val=-10, max_val=10),
+    )
+    b_tensor, _ = cutlass_torch.cute_tensor_like(
+        b_torch_cpu,
+        b_dtype,
+        is_dynamic_layout=True,
+        assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k),
+    )
+
+    c_torch_cpu = cutlass_torch.matrix(
+        1,  # batch=1
+        m,
+        fused_n,
+        c_major == "m",
+        c_dtype,
+    )
+    c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
+        c_torch_cpu,
+        c_dtype,
+        is_dynamic_layout=True,
+        assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
+    )
+    c_tensor = c_tensor.mark_compact_shape_dynamic(
+        mode=(0 if c_major == "m" else 1),
+        stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1),
+        divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
+    )
+    # We need to ensure mode N satisfies 16B alignment for each group
+    alignment_n = 16 * 8 // b_dtype.width
+    cumsum_tensor, cumsum_torch = create_cumsum_tensor(
+        num_groups, fused_n, alignment_n, uniform_distribution=uniform_group_sizes
+    )
+
+    return (
+        a_tensor,
+        a_scale_tensor,
+        b_tensor,
+        cumsum_tensor,
+        c_tensor,
+        a_torch_cpu,
+        a_scale_torch_cpu,
+        b_torch_cpu,
+        cumsum_torch,
+        c_torch_gpu,
+    )
+
+
+def create_tensors_for_batched_mixed_input_gemm(
+    l: int,
+    m: int,
+    n: int,
+    k: int,
+    a_major: str,
+    b_major: str,
+    c_major: str,
+    a_dtype: type[cutlass.Numeric],
+    b_dtype: type[cutlass.Numeric],
+    c_dtype: type[cutlass.Numeric],
+    shuffle_a: bool = False,
+    scale_granularity_m: int = 0,
+    scale_granularity_k: int = 0,
+) -> tuple:
+    """
+    Create all input and output tensors for the batched mixed-input GEMM.
+    """
+    torch.manual_seed(2025)
+
+    a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a(
+        l,
+        m,
+        k,
+        a_major,
+        a_dtype,
+        shuffle_a,
+        scale_granularity_m,
+        scale_granularity_k,
+        b_dtype,
+    )
+
+    b_torch_cpu = cutlass_torch.matrix(
+        l,
+        n,
+        k,
+        b_major == "n",
+        b_dtype,
+        cutlass_torch.TensorInitType.RANDOM,
+        cutlass_torch.RandomInitConfig(min_val=-10, max_val=10),
+    )
+    c_torch_cpu = cutlass_torch.matrix(
+        l,
+        m,
+        n,
+        c_major == "m",
+        c_dtype,
+    )
+
+    b_tensor, _ = cutlass_torch.cute_tensor_like(
+        b_torch_cpu,
+        b_dtype,
+        is_dynamic_layout=True,
+        assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k),
+    )
+    c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
+        c_torch_cpu,
+        c_dtype,
+        is_dynamic_layout=True,
+        assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
+    )
+    c_tensor = c_tensor.mark_compact_shape_dynamic(
+        mode=(0 if c_major == "m" else 1),
+        stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1),
+        divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
+    )
+
+    return (
+        a_tensor,
+        a_scale_tensor,
+        b_tensor,
+        c_tensor,
+        a_torch_cpu,
+        a_scale_torch_cpu,
+        b_torch_cpu,
+        c_torch_gpu,
+    )
+
+
+def run_contiguous_grouped_ref_and_compare(
+    a_torch_cpu: torch.Tensor,
+    b_torch_cpu: torch.Tensor,
+    a_scale_torch_cpu: Optional[torch.Tensor],
+    cumsum_torch_cpu: torch.Tensor,
+    c_torch_gpu: torch.Tensor,
+    c_dtype: type[cutlass.Numeric],
+    tolerance: float,
+) -> None:
+    """
+    Compare kernel result with reference computation.
+    """
+    kernel_result = c_torch_gpu.cpu()
+    assert kernel_result.shape[2] == 1, "batch mode must be 1"
+    kernel_result = kernel_result.reshape(
+        kernel_result.shape[0], kernel_result.shape[1]
+    )
+    # Compute reference result
+    a_for_gemm = a_torch_cpu
+    if a_scale_torch_cpu is not None:
+        scale_shape = a_scale_torch_cpu.shape
+        a_shape = a_torch_cpu.shape
+        a_scale_torch_cpu = a_scale_torch_cpu.to(dtype=torch.float32).reshape(
+            scale_shape[0], scale_shape[1], 1, scale_shape[2]
+        )
+        a_torch_cpu = a_torch_cpu.to(dtype=torch.float32).reshape(
+            a_torch_cpu.shape[0], scale_shape[1], -1, a_torch_cpu.shape[2]
+        )
+        a_for_gemm = (a_torch_cpu * a_scale_torch_cpu).reshape(a_shape)
+    # A in (m, k, l), b in (n, k), c in (m, n)
+    assert cumsum_torch_cpu.shape[0] == a_for_gemm.shape[-1] + 1, (
+        "cumsum tensor must have one more element than a_for_gemm"
+    )
+    assert b_torch_cpu.shape[2] == 1, (
+        "b_torch_cpu must have a singleton dimension in the last position"
+    )
+    prev_idx = 0
+    ref = torch.zeros((a_for_gemm.shape[0], b_torch_cpu.shape[0]), dtype=torch.float32)
+    for group_idx in range(1, cumsum_torch_cpu.shape[0]):
+        # No computation for current group
+        if cumsum_torch_cpu[group_idx] == prev_idx:
+            continue
+        # Get A slice for current group
+        sliced_a = a_for_gemm[:, :, group_idx - 1]
+        # Get B slice for current group
+        sliced_b = b_torch_cpu[prev_idx : cumsum_torch_cpu[group_idx], :, 0]
+        sliced_ref = torch.einsum(
+            "mk,nk->mn",
+            sliced_a.to(dtype=torch.float32),
+            sliced_b.to(dtype=torch.float32),
+        )
+        ref[:, prev_idx : cumsum_torch_cpu[group_idx]] = sliced_ref
+        prev_idx = cumsum_torch_cpu[group_idx]
+    # Convert ref to c_dtype
+    _, ref_torch_gpu = cutlass_torch.cute_tensor_like(
+        ref, c_dtype, is_dynamic_layout=True, assumed_align=16
+    )
+    ref_result = ref_torch_gpu.cpu()
+
+    # Assert close results
+    torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05)
+
+
+def run_batched_mixed_input_ref_and_compare(
+    a_torch_cpu: torch.Tensor,
+    b_torch_cpu: torch.Tensor,
+    a_scale_torch_cpu: Optional[torch.Tensor],
+    c_torch_gpu: torch.Tensor,
+    c_dtype: type[cutlass.Numeric],
+    tolerance: float,
+) -> None:
+    """
+    Compare kernel result with reference computation.
+    """
+    kernel_result = c_torch_gpu.cpu()
+    # Compute reference result
+    if a_scale_torch_cpu is not None:
+        scale_shape = a_scale_torch_cpu.shape
+        a_shape = a_torch_cpu.shape
+        a_scale_torch_cpu = a_scale_torch_cpu.to(dtype=torch.float32).reshape(
+            scale_shape[0], scale_shape[1], 1, scale_shape[2]
+        )
+        a_torch_cpu = a_torch_cpu.to(dtype=torch.float32).reshape(
+            a_torch_cpu.shape[0], scale_shape[1], -1, a_torch_cpu.shape[2]
+        )
+        a_dequant = a_torch_cpu * a_scale_torch_cpu
+        ref = torch.einsum(
+            "mkl,nkl->mnl",
+            a_dequant.reshape(a_shape),
+            b_torch_cpu.to(dtype=torch.float32),
+        )
+    else:
+        ref = torch.einsum(
+            "mkl,nkl->mnl",
+            a_torch_cpu.to(dtype=torch.float32),
+            b_torch_cpu.to(dtype=torch.float32),
+        )
+    # Convert ref to c_dtype
+    _, ref_torch_gpu = cutlass_torch.cute_tensor_like(
+        ref, c_dtype, is_dynamic_layout=True, assumed_align=16
+    )
+    ref_result = ref_torch_gpu.cpu()
+
+    # Assert close results
+    torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05)

examples/python/CuTeDSL/cute/blackwell/kernel/moe/moe_persistent_scheduler.py

@@ -0,0 +1,695 @@
+# Copyright (c) 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.
+
+"""
+MoE Persistent Tile Scheduler
+
+A specialized tile scheduler for MoE (Mixture of Experts) grouped GEMM operations.
+This scheduler handles tile iteration across all experts, producing MoEWorkTileInfo
+(expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) for each tile.
+
+Scenarios:
+- 2Dx3D (Forward): A(tokens_sum, hidden) x B(experts, intermediate, hidden) -> C(tokens_sum, intermediate)
+- 2Dx2D (Backward): A(intermediate, tokens_sum) x B(hidden, tokens_sum) -> C(experts, intermediate, hidden)
+
+Key design principle:
+- Scheduler is ONLY responsible for tile iteration (tensor-agnostic, TMA-agnostic)
+- Domain conversion (fake tensor -> real expert tensor) is handled by MoESchedExtension
+- TMA descriptor management is handled by OnlineTensormapDescCreator
+- The kernel orchestrates all three components
+"""
+
+from typing import List, Tuple, Literal
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cutlass_dsl import (
+    Boolean,
+    Int32,
+    Integer,
+    extract_mlir_values,
+    new_from_mlir_values,
+    const_expr,
+    dsl_user_op,
+)
+from cutlass._mlir import ir
+
+# =============================================================================
+# Work Tile Info
+# =============================================================================
+
+class MoEWorkTileInfo:
+    """
+    Work tile information for MoE scheduler.
+    
+    Contains CTA-level tile information for executor warps:
+    - expert_idx: Which expert (-1 means invalid/done)
+    - tile_m_idx: CTA tile index along GEMM M dimension
+    - tile_n_idx: CTA tile index along GEMM N dimension  
+    - k_tile_cnt: Number of CTA tiles along K dimension
+    
+    Note: These are CTA-level indices, not cluster-level.
+    tile_l_idx is always 0 for MoE, executor can hardcode it.
+    
+    For 2Dx3D (Forward):
+        M = tokens_i (dynamic), N = intermediate (fixed), K = hidden (fixed)
+    
+    For 2Dx2D (Backward):
+        M = intermediate (fixed), N = hidden (fixed), K = tokens_i (dynamic)
+    """
+    
+    def __init__(
+        self,
+        expert_idx: Int32,      # -1 means invalid tile
+        tile_m_idx: Int32,
+        tile_n_idx: Int32,
+        k_tile_cnt: Int32,
+    ):
+        self.expert_idx = expert_idx
+        self.tile_m_idx = tile_m_idx
+        self.tile_n_idx = tile_n_idx
+        self.k_tile_cnt = k_tile_cnt
+    
+    @property
+    def is_valid_tile(self) -> Boolean:
+        """Check if this is a valid work tile (expert_idx >= 0)."""
+        return self.expert_idx >= Int32(0)
+    
+    def __extract_mlir_values__(self) -> List[ir.Value]:
+        values = extract_mlir_values(self.expert_idx)
+        values.extend(extract_mlir_values(self.tile_m_idx))
+        values.extend(extract_mlir_values(self.tile_n_idx))
+        values.extend(extract_mlir_values(self.k_tile_cnt))
+        return values
+    
+    def __new_from_mlir_values__(self, values: List[ir.Value]) -> "MoEWorkTileInfo":
+        assert len(values) == 4
+        return MoEWorkTileInfo(
+            expert_idx=new_from_mlir_values(self.expert_idx, [values[0]]),
+            tile_m_idx=new_from_mlir_values(self.tile_m_idx, [values[1]]),
+            tile_n_idx=new_from_mlir_values(self.tile_n_idx, [values[2]]),
+            k_tile_cnt=new_from_mlir_values(self.k_tile_cnt, [values[3]]),
+        )
+    
+    def to_rmem_tensor(self):
+        """Pack work tile info fields into an rmem tensor of shape (4,) for vectorized smem copy."""
+        rmem = cute.make_rmem_tensor((4,), Int32)
+        rmem[0] = self.expert_idx
+        rmem[1] = self.tile_m_idx
+        rmem[2] = self.tile_n_idx
+        rmem[3] = self.k_tile_cnt
+        return rmem
+    
+    @staticmethod
+    def from_rmem_tensor(rmem) -> "MoEWorkTileInfo":
+        """Unpack work tile info from an rmem tensor of shape (4,)."""
+        return MoEWorkTileInfo(
+            expert_idx=rmem[0],  # type: ignore[arg-type]
+            tile_m_idx=rmem[1],  # type: ignore[arg-type]
+            tile_n_idx=rmem[2],  # type: ignore[arg-type]
+            k_tile_cnt=rmem[3],  # type: ignore[arg-type]
+        )
+
+
+# =============================================================================
+# Scheduler Parameters
+# =============================================================================
+
+class MoEStaticSchedulerParams:
+    """
+    Parameters for MoE tile scheduler.
+    
+    Uses unified semantics for both scenarios:
+    - expert_shape: (expert_cnt, intermediate, hidden)
+    
+    For 2Dx3D: GEMM is (M=tokens_i, N=intermediate, K=hidden) per expert
+    For 2Dx2D: GEMM is (M=hidden, N=intermediate, K=tokens_i) per expert
+    
+    Tile hierarchy:
+    - cta_tile_shape_mnk: Single CTA tile shape (tile_m, tile_n, tile_k)
+    - cluster_shape_mn: CTAs per cluster (cluster_m, cluster_n)
+    - cluster_tile_shape_mn: Cluster tile shape = cta_tile_shape * cluster_shape
+    
+    This class is used both on host (for grid shape calculation) and on device
+    (stored in scheduler). Codegen-time constants (scenario, cta_tile_shape_mnk,
+    cluster_shape_mn) are NOT serialized to MLIR values.
+    """
+    
+    def __init__(
+        self,
+        scenario: Literal["2Dx3D", "2Dx2D"],
+        expert_shape: Tuple[int | Int32, int | Int32, int | Int32],  # (expert_cnt, intermediate, hidden)
+        cta_tile_shape_mnk: Tuple[int, int, int],  # (tile_m, tile_n, tile_k)
+        cluster_shape_mn: Tuple[int, int],         # (cluster_m, cluster_n)
+    ):
+        self.scenario = scenario
+        e, i, h = expert_shape
+        self.expert_cnt = e if isinstance(e, Int32) else Int32(e)
+        self.intermediate = i if isinstance(i, Int32) else Int32(i)
+        self.hidden = h if isinstance(h, Int32) else Int32(h)
+        self.cta_tile_shape_mnk = cta_tile_shape_mnk
+        self.cluster_shape_mn = cluster_shape_mn
+    
+    @property
+    def cluster_tile_m(self) -> int:
+        """Cluster tile size along M = cta_tile_m * cluster_m."""
+        return self.cta_tile_shape_mnk[0] * self.cluster_shape_mn[0]
+    
+    @property
+    def cluster_tile_n(self) -> int:
+        """Cluster tile size along N = cta_tile_n * cluster_n."""
+        return self.cta_tile_shape_mnk[1] * self.cluster_shape_mn[1]
+    
+    @property
+    def cta_tile_k(self) -> int:
+        """CTA tile size along K (same as cluster since cluster_k = 1)."""
+        return self.cta_tile_shape_mnk[2]
+    
+    def __extract_mlir_values__(self) -> List[ir.Value]:
+        """Only serialize runtime values, not codegen-time constants."""
+        values = []
+        values.extend(extract_mlir_values(self.expert_cnt))
+        values.extend(extract_mlir_values(self.intermediate))
+        values.extend(extract_mlir_values(self.hidden))
+        return values
+    
+    def __new_from_mlir_values__(self, values: List[ir.Value]) -> "MoEStaticSchedulerParams":
+        assert len(values) == 3
+        return MoEStaticSchedulerParams(
+            scenario=self.scenario,
+            expert_shape=(
+                new_from_mlir_values(self.expert_cnt, [values[0]]),
+                new_from_mlir_values(self.intermediate, [values[1]]),
+                new_from_mlir_values(self.hidden, [values[2]]),
+            ),
+            cta_tile_shape_mnk=self.cta_tile_shape_mnk,
+            cluster_shape_mn=self.cluster_shape_mn,
+        )
+    
+    @staticmethod
+    def get_grid_shape(
+        params: "MoEStaticSchedulerParams",
+        max_active_clusters: int,
+    ) -> Tuple[int, int, int]:
+        """
+        Compute grid shape for kernel launch.
+        
+        Since host doesn't know token distribution across experts,
+        we launch max_active_clusters and let device-side scheduler
+        determine which tiles are valid.
+        """
+        return (
+            params.cluster_shape_mn[0],
+            params.cluster_shape_mn[1],
+            max_active_clusters,
+        )
+
+
+# =============================================================================
+# Scheduler (Device-side)
+# =============================================================================
+
+class MoEStaticPersistentTileScheduler:
+    """
+    Persistent tile scheduler specialized for MoE grouped GEMM.
+
+    This scheduler is ONLY responsible for tile iteration. It does NOT know
+    about tensor types, TMA descriptors, or domain conversion. Those concerns
+    are handled by MoESchedExtension and OnlineTensormapDescCreator respectively.
+
+    Architecture:
+    - Scheduler warp: Holds scheduler instance, iterates tiles, broadcasts work_tile_info
+    - Executor warps: Read work_tile_info from smem, use MoESchedExtension for
+      domain conversion and TMA desc selection
+
+    The scheduler handles:
+    - 2Dx3D: Dynamic M per expert (from offs), fixed N (intermediate) and K (hidden)
+    - 2Dx2D: Fixed M (intermediate) and N (hidden), dynamic K per expert (reduction axis)
+
+    Usage (Scheduler warp):
+        scheduler = MoEStaticPersistentTileScheduler.create(params, offs, block_idx, grid_dim)
+        work_tile_info = scheduler.initial_work_tile_info()
+        # Broadcast work_tile_info to smem...
+
+        while work_tile_info.is_valid_tile:
+            # ... do work ...
+            work_tile_info = scheduler.advance_to_next_work()
+            # Broadcast work_tile_info to smem...
+
+    Usage (Executor warps - via MoESchedExtension):
+        # Read work_tile_info from smem...
+        real_a, desc_a = ext.get_gmem_tensor("a", tma_tensor_a, offs, work_tile_info)
+        real_b, desc_b = ext.get_gmem_tensor("b", tma_tensor_b, offs, work_tile_info)
+        real_c, desc_c = ext.get_gmem_tensor("c", tma_tensor_c, offs, work_tile_info)
+    """
+    
+    def __init__(
+        self,
+        # Params (contains scenario, expert_cnt, intermediate, hidden, tile/cluster shapes)
+        params: MoEStaticSchedulerParams,
+        # Runtime tensor for scheduling
+        offs: cute.Tensor,             # (experts,) cumsum of token counts
+        # Scheduling state
+        num_persistent_clusters: Int32,
+        current_work_linear_idx: Int32,
+        cta_id_in_cluster: cute.Coord,
+        # Expert tracking state (for O(1) advance within same expert)
+        current_expert_idx: Int32,
+        expert_tile_start: Int32,      # cumsum of tiles before current expert
+        expert_tile_end: Int32,        # cumsum of tiles including current expert
+    ):
+        self.params = params
+        self.offs = offs
+        self.num_persistent_clusters = num_persistent_clusters
+        self._current_work_linear_idx = current_work_linear_idx
+        self.cta_id_in_cluster = cta_id_in_cluster
+        # Expert tracking
+        self.current_expert_idx = current_expert_idx
+        self.expert_tile_start = expert_tile_start
+        self.expert_tile_end = expert_tile_end
+    
+    # =========================================================================
+    # Convenience accessors for params
+    # =========================================================================
+    
+    @property
+    def scenario(self) -> Literal["2Dx3D", "2Dx2D"]:
+        return self.params.scenario
+    
+    @property
+    def expert_cnt(self) -> Int32:
+        return self.params.expert_cnt
+    
+    @property
+    def intermediate(self) -> Int32:
+        return self.params.intermediate
+    
+    @property
+    def hidden(self) -> Int32:
+        return self.params.hidden
+    
+    @property
+    def cta_tile_shape_mnk(self) -> Tuple[int, int, int]:
+        return self.params.cta_tile_shape_mnk
+    
+    @property
+    def cluster_shape_mn(self) -> Tuple[int, int]:
+        return self.params.cluster_shape_mn
+    
+    @property
+    def cluster_tile_m(self) -> int:
+        return self.params.cluster_tile_m
+    
+    @property
+    def cluster_tile_n(self) -> int:
+        return self.params.cluster_tile_n
+    
+    @property
+    def cta_tile_k(self) -> int:
+        return self.params.cta_tile_k
+    
+    # =========================================================================
+    # MLIR value serialization (for SSA value passing in device code)
+    # =========================================================================
+    
+    def __extract_mlir_values__(self) -> List[ir.Value]:
+        values = []
+        # Params (only runtime values are extracted)
+        values.extend(extract_mlir_values(self.params))
+        # Runtime tensor for scheduling
+        values.extend(extract_mlir_values(self.offs))
+        # Scheduling state
+        values.extend(extract_mlir_values(self.num_persistent_clusters))
+        values.extend(extract_mlir_values(self._current_work_linear_idx))
+        values.extend(extract_mlir_values(self.cta_id_in_cluster))
+        # Expert tracking state
+        values.extend(extract_mlir_values(self.current_expert_idx))
+        values.extend(extract_mlir_values(self.expert_tile_start))
+        values.extend(extract_mlir_values(self.expert_tile_end))
+        return values
+    
+    def __new_from_mlir_values__(
+        self, values: List[ir.Value]
+    ) -> "MoEStaticPersistentTileScheduler":
+        idx = 0
+        
+        # Params (3 values: expert_cnt, intermediate, hidden)
+        new_params = new_from_mlir_values(self.params, values[idx:idx + 3])
+        idx += 3
+        
+        # Runtime tensor for scheduling (variable size)
+        offs_len = len(extract_mlir_values(self.offs))
+        new_offs = new_from_mlir_values(self.offs, values[idx:idx + offs_len])
+        idx += offs_len
+        
+        # Scheduling state
+        new_num_persistent_clusters = new_from_mlir_values(
+            self.num_persistent_clusters, [values[idx]]
+        )
+        idx += 1
+        new_current_work_linear_idx = new_from_mlir_values(
+            self._current_work_linear_idx, [values[idx]]
+        )
+        idx += 1
+        
+        # cta_id_in_cluster (3 values for Coord)
+        new_cta_id_in_cluster = new_from_mlir_values(
+            self.cta_id_in_cluster, values[idx:idx + 3]
+        )
+        idx += 3
+        
+        # Expert tracking state
+        new_current_expert_idx = new_from_mlir_values(
+            self.current_expert_idx, [values[idx]]
+        )
+        idx += 1
+        new_expert_tile_start = new_from_mlir_values(
+            self.expert_tile_start, [values[idx]]
+        )
+        idx += 1
+        new_expert_tile_end = new_from_mlir_values(
+            self.expert_tile_end, [values[idx]]
+        )
+        idx += 1
+        
+        return MoEStaticPersistentTileScheduler(
+            params=new_params,
+            offs=new_offs,
+            num_persistent_clusters=new_num_persistent_clusters,
+            current_work_linear_idx=new_current_work_linear_idx,
+            cta_id_in_cluster=new_cta_id_in_cluster,
+            current_expert_idx=new_current_expert_idx,
+            expert_tile_start=new_expert_tile_start,
+            expert_tile_end=new_expert_tile_end,
+        )
+    
+    # =========================================================================
+    # Factory method
+    # =========================================================================
+    
+    @staticmethod
+    @dsl_user_op
+    def create(
+        params: MoEStaticSchedulerParams,
+        offs: cute.Tensor,
+        block_idx: Tuple[Integer, Integer, Integer],
+        grid_dim: Tuple[Integer, Integer, Integer],
+        *,
+        loc=None,
+        ip=None,
+    ) -> "MoEStaticPersistentTileScheduler":
+        """
+        Create a MoE persistent tile scheduler.
+        
+        :param params: Scheduler parameters (from host)
+        :param offs: Cumsum tensor of token counts per expert, shape (experts,)
+        :param block_idx: CUDA block index
+        :param grid_dim: CUDA grid dimensions
+        """
+        num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size(
+            params.cluster_shape_mn, loc=loc, ip=ip
+        )
+        
+        bidx, bidy, bidz = block_idx
+        current_work_linear_idx = Int32(bidz)
+        
+        cta_id_in_cluster = (
+            Int32(bidx % params.cluster_shape_mn[0]),
+            Int32(bidy % params.cluster_shape_mn[1]),
+            Int32(0),
+        )
+        
+        # Initialize expert tracking to "before expert 0"
+        # The first call to _get_work_tile_for_linear_idx will advance to the correct expert
+        current_expert_idx = Int32(0)
+        expert_tile_start = Int32(0)
+        expert_tile_end = Int32(0)  # Will be computed on first access
+        
+        return MoEStaticPersistentTileScheduler(
+            params=params,
+            offs=offs,
+            num_persistent_clusters=num_persistent_clusters,
+            current_work_linear_idx=current_work_linear_idx,
+            cta_id_in_cluster=cta_id_in_cluster,
+            current_expert_idx=current_expert_idx,
+            expert_tile_start=expert_tile_start,
+            expert_tile_end=expert_tile_end,
+        )
+    
+    # =========================================================================
+    # Tile iteration methods
+    # =========================================================================
+    
+    @dsl_user_op
+    @cute.jit
+    def initial_work_tile_info(self, *, loc=None, ip=None) -> MoEWorkTileInfo:
+        """Get the initial work tile info."""
+        return self._get_work_tile_for_linear_idx(
+            self._current_work_linear_idx, loc=loc, ip=ip
+        )
+    
+    @dsl_user_op
+    @cute.jit
+    def advance_to_next_work(self, *, loc=None, ip=None) -> MoEWorkTileInfo:
+        """Advance to the next work tile and return its info."""
+        self._current_work_linear_idx += self.num_persistent_clusters
+        return self._get_work_tile_for_linear_idx(
+            self._current_work_linear_idx, loc=loc, ip=ip
+        )
+    
+    @dsl_user_op
+    @cute.jit
+    def _get_work_tile_for_linear_idx(
+        self, 
+        cluster_linear_idx: Int32, 
+        *, 
+        loc=None, 
+        ip=None
+    ) -> MoEWorkTileInfo:
+        """
+        Convert a linear cluster index to MoEWorkTileInfo.
+        
+        Uses cached expert tracking state for O(1) fast path when staying
+        within the same expert. Advances expert state when needed.
+        
+        Returns an invalid tile (expert_idx = -1) if cluster_linear_idx is out of range.
+        """
+        # Ensure expert tracking is initialized and up-to-date
+        self._advance_expert_to_contain(cluster_linear_idx, loc=loc, ip=ip)
+        
+        # Check if valid (still within expert range after advancing)
+        is_valid = self.current_expert_idx < self.expert_cnt
+
+        work_tile_info = MoEWorkTileInfo(
+            expert_idx=Int32(-1),
+            tile_m_idx=Int32(0),
+            tile_n_idx=Int32(0),
+            k_tile_cnt=Int32(0),
+        )
+        
+        if is_valid:
+            # Compute local cluster tile indices within current expert
+            local_idx = cluster_linear_idx - self.expert_tile_start
+            cluster_tile_m_idx, cluster_tile_n_idx = self._decompose_local_idx(
+                local_idx, self.current_expert_idx, loc=loc, ip=ip
+            )
+            
+            # Convert cluster tile indices to CTA tile indices
+            # cta_tile_idx = cluster_tile_idx * cluster_shape + cta_id_in_cluster
+            cta_tile_m_idx = (
+                cluster_tile_m_idx * self.cluster_shape_mn[0]
+                + self.cta_id_in_cluster[0]  # type: ignore[index]
+            )
+            cta_tile_n_idx = (
+                cluster_tile_n_idx * self.cluster_shape_mn[1]
+                + self.cta_id_in_cluster[1]  # type: ignore[index]
+            )
+            # Compute k_tile_cnt
+            k_tile_cnt = self._compute_k_tile_cnt(self.current_expert_idx, loc=loc, ip=ip)
+        
+            work_tile_info = MoEWorkTileInfo(
+                expert_idx=self.current_expert_idx,
+                tile_m_idx=cta_tile_m_idx,
+                tile_n_idx=cta_tile_n_idx,
+                k_tile_cnt=k_tile_cnt,
+            )
+        return work_tile_info
+    
+    @dsl_user_op
+    @cute.jit
+    def _advance_expert_to_contain(
+        self,
+        cluster_linear_idx: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> None:
+        """
+        Advance expert tracking state until current expert contains cluster_linear_idx,
+        or we run out of experts.
+        
+        Fast path: If already in correct expert, no work needed.
+        """
+        # Initialize expert_tile_end if this is the first call (expert_tile_end == 0)
+        if self.expert_tile_end == Int32(0):
+            tiles_for_expert_0 = self._compute_tiles_for_expert(Int32(0), loc=loc, ip=ip)
+            self.expert_tile_end = tiles_for_expert_0
+        
+        # Advance until cluster_linear_idx < expert_tile_end or no more experts
+        while cluster_linear_idx >= self.expert_tile_end and self.current_expert_idx < self.expert_cnt:
+            self.current_expert_idx = self.current_expert_idx + 1
+            self.expert_tile_start = self.expert_tile_end
+            
+            if self.current_expert_idx < self.expert_cnt:
+                tiles_for_expert = self._compute_tiles_for_expert(
+                    self.current_expert_idx, loc=loc, ip=ip
+                )
+                self.expert_tile_end = self.expert_tile_end + tiles_for_expert
+    
+    @dsl_user_op
+    @cute.jit
+    def _compute_tiles_for_expert(
+        self,
+        expert_idx: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Int32:
+        """Compute total cluster tiles for a given expert."""
+        if const_expr(self.scenario == "2Dx2D"):
+            # Fixed M=hidden, N=intermediate
+            cluster_tile_m_cnt = (self.hidden + self.cluster_tile_m - 1) // self.cluster_tile_m
+            cluster_tile_n_cnt = (self.intermediate + self.cluster_tile_n - 1) // self.cluster_tile_n
+            return cluster_tile_m_cnt * cluster_tile_n_cnt
+        else:  # 2Dx3D
+            # Variable M (tokens), fixed N
+            tokens_i = self.offs[expert_idx]
+            if expert_idx > 0:
+                tokens_i = tokens_i - self.offs[expert_idx - 1]  # type: ignore[operator]
+            cluster_tile_m_cnt = (
+                tokens_i + self.cluster_tile_m - 1  # type: ignore[operator]
+            ) // self.cluster_tile_m
+            cluster_tile_n_cnt = (
+                self.intermediate + self.cluster_tile_n - 1
+            ) // self.cluster_tile_n
+            return cluster_tile_m_cnt * cluster_tile_n_cnt
+    
+    @dsl_user_op
+    @cute.jit
+    def _decompose_local_idx(
+        self,
+        local_idx: Int32,
+        expert_idx: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Tuple[Int32, Int32]:
+        """
+        Decompose local cluster tile index within expert to (cluster_tile_m_idx, cluster_tile_n_idx).
+        
+        Uses "short side first" strategy: the shorter dimension changes faster.
+        This maximizes overlap between adjacent clusters for better L2 cache utilization.
+        
+        For example, if m_cnt=2, n_cnt=8:
+        - N is longer, so M changes faster: local_idx = n_idx * m_cnt + m_idx
+        - Linearization order: (0,0), (1,0), (0,1), (1,1), (0,2), (1,2), ...
+        """
+        # Get tile counts for M and N
+        cluster_tile_m_cnt, cluster_tile_n_cnt = self._get_cluster_tile_counts(
+            expert_idx, loc=loc, ip=ip
+        )
+        cluster_tile_m_idx = -1
+        cluster_tile_n_idx = -1
+        
+        # Short side first: shorter dimension changes faster
+        # If m_cnt <= n_cnt: m is shorter, m changes faster
+        #   local_idx = n_idx * m_cnt + m_idx
+        # If n_cnt < m_cnt: n is shorter, n changes faster
+        #   local_idx = m_idx * n_cnt + n_idx
+        if cluster_tile_m_cnt <= cluster_tile_n_cnt:
+            # M is shorter or equal, M changes faster
+            cluster_tile_m_idx = local_idx % cluster_tile_m_cnt
+            cluster_tile_n_idx = local_idx // cluster_tile_m_cnt
+        else:
+            # N is shorter, N changes faster
+            cluster_tile_n_idx = local_idx % cluster_tile_n_cnt
+            cluster_tile_m_idx = local_idx // cluster_tile_n_cnt
+        
+        return (cluster_tile_m_idx, cluster_tile_n_idx)
+    
+    @dsl_user_op
+    @cute.jit
+    def _get_cluster_tile_counts(
+        self,
+        expert_idx: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Tuple[Int32, Int32]:
+        """Get (cluster_tile_m_cnt, cluster_tile_n_cnt) for a given expert."""
+        if const_expr(self.scenario == "2Dx2D"):
+            # Fixed M=hidden, N=intermediate
+            cluster_tile_m_cnt = (self.hidden + self.cluster_tile_m - 1) // self.cluster_tile_m
+            cluster_tile_n_cnt = (self.intermediate + self.cluster_tile_n - 1) // self.cluster_tile_n
+        else:  # 2Dx3D
+            # Variable M (tokens), fixed N
+            tokens_i = self.offs[expert_idx]
+            if expert_idx > 0:
+                tokens_i = tokens_i - self.offs[expert_idx - 1]  # type: ignore[operator]
+            cluster_tile_m_cnt = (
+                tokens_i + self.cluster_tile_m - 1  # type: ignore[operator]
+            ) // self.cluster_tile_m
+            cluster_tile_n_cnt = (
+                self.intermediate + self.cluster_tile_n - 1
+            ) // self.cluster_tile_n
+        return (cluster_tile_m_cnt, cluster_tile_n_cnt)
+    
+    @dsl_user_op
+    @cute.jit
+    def _compute_k_tile_cnt(
+        self,
+        expert_idx: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ) -> Int32:
+        """
+        Compute the number of K tiles for this expert.
+        
+        2Dx3D: K = hidden (fixed) -> k_tile_cnt = ceil(hidden / cta_tile_k)
+        2Dx2D: K = tokens_i (variable) -> k_tile_cnt = ceil(tokens_i / cta_tile_k)
+        """
+        if const_expr(self.scenario == "2Dx3D"):
+            # K is hidden (fixed)
+            return (self.hidden + self.cta_tile_k - 1) // self.cta_tile_k
+        else:  # 2Dx2D
+            # K is tokens_i (variable per expert)
+            tokens_i = self.offs[expert_idx]
+            if expert_idx > cutlass.Int32(0):
+                tokens_i = tokens_i - self.offs[expert_idx - 1]  # type: ignore[operator]
+            return (tokens_i + self.cta_tile_k - 1) // self.cta_tile_k  # type: ignore[return-value, operator]

examples/python/CuTeDSL/cute/blackwell/kernel/moe/moe_sched_extension.py

@@ -0,0 +1,443 @@
+# Copyright (c) 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.
+
+"""
+MoE Scheduler Extension.
+
+Bridges the MoE tile scheduler (MoEStaticPersistentTileScheduler) with tensor-level
+domain conversion and TMA descriptor selection. This is the "glue" layer between:
+
+- Scheduler: produces MoEWorkTileInfo (expert_idx, tile_m, tile_n, k_tile_cnt)
+- OnlineTensormapDescCreator: builds/retrieves TMA descriptors from workspace
+- Kernel: orchestrates everything
+
+Different kernel types (grouped_mm, scaled_grouped_mm, etc.) provide their own
+MoESchedExtension subclass with kernel-specific domain conversion logic.
+
+Key design principles:
+- Unified interface: get_gmem_tensor() for all tensor types
+- Free implementation: no role-based templates, each subclass writes its own logic
+- Composable utilities: compute_expert_token_range, rewrite_tensor_shape, etc.
+  are available as tools but not mandatory
+
+Architecture:
+
+    Scheduler ──(produces)──> MoEWorkTileInfo
+                                    │
+                           expert_idx, tile_m, tile_n, k_cnt
+                                    │
+                                    v
+    Extension ──(uses)──> OnlineTensormapDescCreator
+        │                         │
+        │  get_gmem_tensor()      │  get_desc_ptr()
+        │  prefetch_for_expert()  │  construct_and_write()
+        │                         │
+        └── internal calls ───────┘
+
+    Kernel (caller): the only place that knows all three exist
+"""
+
+from abc import ABC, abstractmethod
+from typing import Literal, Tuple, Union
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.typing import Pointer
+from cutlass.cutlass_dsl import Int32
+
+from dataclasses import dataclass
+
+from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF
+from blackwell.kernel.moe.moe_utils import (
+    OnlineTensormapDescCreator,
+    tensormap_ptr_for_copy,
+    compute_expert_token_range,
+    rewrite_tensor_shape,
+    prefetch_tma_descriptor,
+)
+from blackwell.kernel.moe.moe_persistent_scheduler import MoEWorkTileInfo
+
+
+@dataclass(frozen=True)
+class MoESchedExtension(ABC):
+    """
+    Abstract base class for MoE scheduler extensions.
+
+    Bridges MoEWorkTileInfo with tensor-level domain conversion and TMA
+    descriptor selection. Each kernel type (grouped_mm, scaled_grouped_mm, etc.)
+    provides its own subclass with kernel-specific logic.
+
+    The extension:
+    - Holds a reference to an OnlineTensormapDescCreator for expert-wise desc retrieval
+    - Implements get_gmem_tensor() to convert MoE-view tensors to per-expert tensors
+    - Implements prefetch_for_expert() to prefetch expert-wise TMA descriptors
+
+    Subclasses are free to add any additional attributes in __init__ (scenario,
+    codegen configs, etc.) and implement get_gmem_tensor with arbitrary logic
+    per tensor_name. No role-based templates or rigid patterns are imposed.
+
+    Usage in kernel (caller):
+        ext = ConcreteSchedExtension(tensormap_ctor, scenario=...)
+
+        while work_tile_info.is_valid_tile:
+            real_a, desc_a = ext.get_gmem_tensor("a", tma_tensor_a, offs, work_tile_info)
+            real_b, desc_b = ext.get_gmem_tensor("b", tma_tensor_b, offs, work_tile_info)
+            # Use real_a, desc_a in cute.copy ...
+    """
+
+    def __init__(self, tensormap_ctor: OnlineTensormapDescCreator):
+        super().__init__()
+        self.tensormap_ctor = tensormap_ctor
+
+    @abstractmethod
+    def get_gmem_tensor(
+        self,
+        tensor_name: str,
+        gmem_tensor_in_moe_view: cute.Tensor,
+        offs: Union[cute.Tensor, Tuple[cute.Tensor, cute.Tensor]],
+        work_tile_info: MoEWorkTileInfo,
+    ) -> Tuple[cute.Tensor, "Pointer | None"]:
+        """
+        Convert an MoE-view tensor to the real per-expert tensor for the
+        current work tile, and return the appropriate TMA descriptor pointer.
+
+        The MoE-view tensor uses "fake" GEMM domain dimensions that span all
+        experts (e.g., fake_m = tokens_sum). This method slices/offsets it
+        to the current expert's actual region.
+
+        :param tensor_name: Identifies which tensor (e.g., "a", "b", "c", "sfa")
+        :param gmem_tensor_in_moe_view: Tensor in fake GEMM MNKL domain
+        :param offs: Either a single cumsum tensor (experts,), or a tuple of
+                     (offs_token, offs_padded) where offs_padded provides
+                     padded offsets for scale-factor domain conversion.
+        :param work_tile_info: Current work tile from the scheduler
+        :return: (real_tensor, tma_desc_ptr_or_none)
+                 - real_tensor: domain-offset and shape-rewritten tensor for this expert
+                 - tma_desc_ptr: expert-wise desc ptr (already converted for cute.copy),
+                   or None if the caller should use the global TMA descriptor
+        """
+        ...
+
+    @abstractmethod
+    def prefetch_for_expert(self, expert_idx: Int32) -> None:
+        """
+        Prefetch expert-wise TMA descriptors for the given expert.
+
+        Called when the scheduler advances to a new expert, allowing the TMA
+        descriptor cache to be warmed up before the descriptors are needed.
+
+        :param expert_idx: Index of the expert whose descriptors to prefetch
+        """
+        ...
+
+
+# =============================================================================
+# Grouped MM Extension
+# =============================================================================
+
+
+class GroupedMmSchedExtension(MoESchedExtension):
+    """
+    MoE scheduler extension for grouped_mm: handles tensors a, b, c.
+
+    Domain conversion logic per scenario:
+
+    2Dx3D:
+        A: (fake_m, k, 1)  -> offset fake_m by token_offset, global desc
+        B: (n, k, fake_l)  -> offset fake_l by expert_idx, global desc
+        C: (fake_m, n, 1)  -> rewrite shape only, expert-wise desc
+
+    2Dx2D:
+        A: (m, fake_k, 1)  -> rewrite shape only, expert-wise desc
+        B: (n, fake_k, 1)  -> rewrite shape only, expert-wise desc
+        C: (m, n, fake_l)  -> offset fake_l by expert_idx, global desc
+    """
+
+    def __init__(
+        self,
+        scenario: Literal["2Dx3D", "2Dx2D"],
+        tensormap_ctor: OnlineTensormapDescCreator,
+    ):
+        super().__init__(tensormap_ctor)
+        self.scenario = scenario
+
+    @cute.jit
+    def get_gmem_tensor(
+        self,
+        tensor_name: str,
+        gmem_tensor_in_moe_view: cute.Tensor,
+        offs: cute.Tensor,
+        work_tile_info: MoEWorkTileInfo,
+    ):
+        expert_idx = work_tile_info.expert_idx
+        token_offset, tokens_i = compute_expert_token_range(offs, expert_idx)
+
+        shape = gmem_tensor_in_moe_view.shape
+        c1 = cutlass.Int32(1)
+
+        if cutlass.const_expr(self.scenario == "2Dx3D"):
+            if cutlass.const_expr(tensor_name == "a"):
+                # A: (fake_m, k, 1) -> offset fake_m, global desc
+                real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view)
+                real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1))  # type: ignore[index]
+                return (real, None)
+            elif cutlass.const_expr(tensor_name == "b"):
+                # B: (n, k, fake_l) -> offset fake_l, global desc
+                real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view)
+                real = rewrite_tensor_shape(real, (shape[0], shape[1], c1))  # type: ignore[index]
+                return (real, None)
+            elif cutlass.const_expr(tensor_name == "c"):
+                # C: (fake_m, n, 1) -> expert-wise desc, no offset
+                real = rewrite_tensor_shape(
+                    gmem_tensor_in_moe_view,
+                    (tokens_i, shape[1], c1),  # type: ignore[index]
+                )
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("c", expert_idx)
+                )
+                return (real, desc)
+
+        elif cutlass.const_expr(self.scenario == "2Dx2D"):
+            if cutlass.const_expr(tensor_name == "a"):
+                # A: (m, fake_k, 1) -> expert-wise desc, no offset
+                real = rewrite_tensor_shape(
+                    gmem_tensor_in_moe_view,
+                    (shape[0], tokens_i, c1),  # type: ignore[index]
+                )
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("a", expert_idx)
+                )
+                return (real, desc)
+            elif cutlass.const_expr(tensor_name == "b"):
+                # B: (n, fake_k, 1) -> expert-wise desc, no offset
+                real = rewrite_tensor_shape(
+                    gmem_tensor_in_moe_view,
+                    (shape[0], tokens_i, c1),  # type: ignore[index]
+                )
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("b", expert_idx)
+                )
+                return (real, desc)
+            elif cutlass.const_expr(tensor_name == "c"):
+                # C: (m, n, fake_l) -> offset fake_l, global desc
+                real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view)
+                real = rewrite_tensor_shape(real, (shape[0], shape[1], c1))  # type: ignore[index]
+                return (real, None)
+        
+        raise ValueError("Invalid scenario or GEMM tensor name.")
+
+    @cute.jit
+    def prefetch_for_expert(self, expert_idx: Int32) -> None:
+        if cutlass.const_expr(self.scenario == "2Dx3D"):
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("c", expert_idx))
+        elif cutlass.const_expr(self.scenario == "2Dx2D"):
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("a", expert_idx))
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("b", expert_idx))
+        else:
+            raise ValueError("Invalid scenario.")
+        
+
+
+# =============================================================================
+# Scaled Grouped MM Extension (block-scaled MoE)
+# =============================================================================
+
+
+class ScaledGroupedMmSchedExtension(MoESchedExtension):
+    """
+    MoE scheduler extension for scaled_grouped_mm: handles a, b, c, sfa, sfb.
+
+    Extends GroupedMmSchedExtension with scale-factor tensor support.
+    SFA/SFB are passed as flat GEMM-domain tensors and atom-tiled per expert
+    via tile_atom_to_shape_SF.
+
+    The offs parameter is always a tuple (offs_token, offs_padded):
+    - offs_token: cumsum offsets in data (activation) domain
+    - offs_padded: cumsum offsets in scale-factor domain (padded to atom granularity)
+
+    sf_vec_size is obtained from self.tensormap_ctor.sf_vec_size.
+
+    Domain conversion logic per scenario:
+
+    2Dx3D:
+        A:   (fake_m, k, 1)               -> offset fake_m by token_offset, global desc
+        B:   (n, k, fake_l)               -> offset fake_l by expert_idx, global desc
+        C:   (fake_m, n, 1)               -> rewrite shape, expert-wise desc
+        SFA: (fake_m_pad, k_pad, 1)       -> offset fake_m_pad by padded_offset,
+                                              atom-tile, global desc
+        SFB: (n_pad, k_pad, fake_l)       -> offset fake_l by expert_idx,
+                                              atom-tile, global desc
+
+    2Dx2D:
+        A:   (m, fake_k, 1)               -> rewrite shape, expert-wise desc
+        B:   (n, fake_k, 1)               -> rewrite shape, expert-wise desc
+        C:   (m, n, fake_l)               -> offset fake_l by expert_idx, global desc
+        SFA: (m_pad, fake_k_pad, 1)       -> offset fake_k_pad by padded_offset,
+                                              atom-tile, expert-wise desc
+        SFB: (n_pad, fake_k_pad, 1)       -> offset fake_k_pad by padded_offset,
+                                              atom-tile, expert-wise desc
+    """
+
+    def __init__(
+        self,
+        scenario: Literal["2Dx3D", "2Dx2D"],
+        tensormap_ctor: OnlineTensormapDescCreator,
+    ):
+        super().__init__(tensormap_ctor)
+        self.scenario = scenario
+
+    @cute.jit
+    def get_gmem_tensor(
+        self,
+        tensor_name: str,
+        gmem_tensor_in_moe_view: cute.Tensor,
+        offs: Tuple[cute.Tensor, cute.Tensor],
+        work_tile_info: MoEWorkTileInfo,
+    ):
+        # Unpack the offs tuple
+        offs_token, offs_padded = offs
+
+        expert_idx = work_tile_info.expert_idx
+        token_offset, tokens_i = compute_expert_token_range(offs_token, expert_idx)
+        padded_offset, padded_size_i = compute_expert_token_range(
+            offs_padded, expert_idx
+        )
+
+        shape = gmem_tensor_in_moe_view.shape
+        stride = gmem_tensor_in_moe_view.stride
+        c1 = cutlass.Int32(1)
+        sf_vec_size = self.tensormap_ctor.sf_vec_size
+
+        if cutlass.const_expr(self.scenario == "2Dx3D"):
+            if cutlass.const_expr(tensor_name == "a"):
+                # A: (fake_m, k, 1) -> offset fake_m, global desc
+                real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view)
+                real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1))  # type: ignore[index]
+                return (real, None)
+
+            elif cutlass.const_expr(tensor_name == "b"):
+                # B: (n, k, fake_l) -> offset fake_l, global desc
+                real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view)
+                real = rewrite_tensor_shape(real, (shape[0], shape[1], c1))  # type: ignore[index]
+                return (real, None)
+
+            elif cutlass.const_expr(tensor_name == "c"):
+                # C: (fake_m, n, 1) -> expert-wise desc
+                real = rewrite_tensor_shape(
+                    gmem_tensor_in_moe_view,
+                    (tokens_i, shape[1], c1),  # type: ignore[index]
+                )
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("c", expert_idx)
+                )
+                return (real, desc)
+
+            elif cutlass.const_expr(tensor_name == "sfa"):
+                # SFA: (fake_m_pad, k_pad, 1) -> offset fake_m_pad, atom-tile, global desc
+                real = cute.domain_offset(
+                    (padded_offset, 0, 0), gmem_tensor_in_moe_view
+                )
+                per_expert_shape = (padded_size_i, shape[1], c1)  # type: ignore[index]
+                sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size)
+                real = cute.make_tensor(
+                    real.iterator, cute.make_layout(sf_layout.shape, stride=stride)
+                )
+                return (real, None)
+
+            elif cutlass.const_expr(tensor_name == "sfb"):
+                # SFB: (n_pad, k_pad, fake_l) -> offset fake_l, atom-tile, global desc
+                real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view)
+                per_expert_shape = (shape[0], shape[1], c1)  # type: ignore[index]
+                sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size)
+                real = cute.make_tensor(
+                    real.iterator, cute.make_layout(sf_layout.shape, stride=stride)
+                )
+                return (real, None)
+
+        elif cutlass.const_expr(self.scenario == "2Dx2D"):
+            if cutlass.const_expr(tensor_name == "a"):
+                # A: (m, fake_k, 1) -> expert-wise desc
+                real = rewrite_tensor_shape(
+                    gmem_tensor_in_moe_view,
+                    (shape[0], tokens_i, c1),  # type: ignore[index]
+                )
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("a", expert_idx)
+                )
+                return (real, desc)
+
+            elif cutlass.const_expr(tensor_name == "b"):
+                # B: (n, fake_k, 1) -> expert-wise desc
+                real = rewrite_tensor_shape(
+                    gmem_tensor_in_moe_view,
+                    (shape[0], tokens_i, c1),  # type: ignore[index]
+                )
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("b", expert_idx)
+                )
+                return (real, desc)
+
+            elif cutlass.const_expr(tensor_name == "c"):
+                # C: (m, n, fake_l) -> offset fake_l, global desc
+                real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view)
+                real = rewrite_tensor_shape(real, (shape[0], shape[1], c1))  # type: ignore[index]
+                return (real, None)
+
+            elif cutlass.const_expr(tensor_name == "sfa"):
+                # SFA: (m_pad, fake_k_pad, 1) -> offset fake_k_pad, atom-tile, expert-wise desc
+                per_expert_shape = (shape[0], padded_size_i, c1)  # type: ignore[index]
+                sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size)
+                real = rewrite_tensor_shape(gmem_tensor_in_moe_view, sf_layout.shape)
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("sfa", expert_idx)
+                )
+                return (real, desc)
+
+            elif cutlass.const_expr(tensor_name == "sfb"):
+                # SFB: (n_pad, fake_k_pad, 1) -> offset fake_k_pad, atom-tile, expert-wise desc
+                per_expert_shape = (shape[0], padded_size_i, c1)  # type: ignore[index]
+                sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size)
+                real = rewrite_tensor_shape(gmem_tensor_in_moe_view, sf_layout.shape)
+                desc = tensormap_ptr_for_copy(
+                    self.tensormap_ctor.get_desc_ptr("sfb", expert_idx)
+                )
+                return (real, desc)
+
+        raise ValueError("Invalid scenario or tensor name.")
+
+    @cute.jit
+    def prefetch_for_expert(self, expert_idx: Int32) -> None:
+        if cutlass.const_expr(self.scenario == "2Dx3D"):
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("c", expert_idx))
+        elif cutlass.const_expr(self.scenario == "2Dx2D"):
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("a", expert_idx))
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("b", expert_idx))
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("sfa", expert_idx))
+            prefetch_tma_descriptor(self.tensormap_ctor.get_desc_ptr("sfb", expert_idx))
+        else:
+            raise ValueError("Invalid scenario.")

examples/python/CuTeDSL/cute/blackwell/kernel/moe/moe_utils.py

@@ -0,0 +1,910 @@
+# Copyright (c) 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.
+
+"""
+Online TMA Descriptor Construction Utilities.
+
+Provides utilities for dynamically creating TMA descriptors at kernel runtime
+based on runtime-provided information (problem sizes, pointers, etc.).
+
+Key components:
+- OnlineTensormapDescCreator: Simplified ABC for TMA descriptor builders (2 abstract methods)
+- TensormapWorkspace: Helper for linear workspace layout of TMA descriptors
+- MoEGroupedGemmTensormapConstructor: TMA descriptor constructor for MoE Grouped GEMM
+- GeneralGroupedGemmTensormapConstructor: TMA descriptor constructor for general Grouped GEMM
+- Pointer utility functions (ptr_offset_bytes, gmem_ptr_to_generic, etc.)
+- tensormap_ptr_for_copy: Convert raw desc ptr to cute.copy-compatible type
+- compute_expert_token_range: Compute per-expert token offset and count from offs
+- rewrite_tensor_shape: Debug-friendly tensor shape rewrite utility
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional, Literal, Tuple, Union
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.typing import AddressSpace, Pointer
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cutlass_dsl import dsl_user_op, Int32
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm
+from cutlass._mlir.dialects import cute as _cute_ir
+from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
+from dataclasses import dataclass
+
+from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF
+
+TensormapDescBytes = 128
+
+# =============================================================================
+# Pointer Utilities
+# =============================================================================
+
+
+@dsl_user_op
+@cute.jit
+def spin_wait(
+    ptr: Pointer, condition, fail_sleep_cycles: int = 100, *, loc=None, ip=None
+) -> None:
+    """
+    Generic spin-wait.
+    Example usage:
+        # Wait until counter >= total_blocks
+        spin_wait(counter_ptr, lambda x: x >= total_blocks, fail_sleep_cycles=100)
+
+        # Wait until flag == 1
+        spin_wait(flag_ptr, lambda x: x == 1)
+    """
+    current = cute.arch.load(ptr, ptr.dtype, cop="cg", loc=loc, ip=ip)
+    while not condition(current):
+        # Load with L1 cache bypass (ld.global.cg)
+        if cutlass.const_expr(fail_sleep_cycles > 0):
+            cute.arch.nanosleep(sleep_time=fail_sleep_cycles, loc=loc, ip=ip)
+        current = cute.arch.load(ptr, ptr.dtype, cop="cg", loc=loc, ip=ip)
+
+
+@dsl_user_op
+def gmem_ptr_to_generic(
+    gmem_ptr: Pointer,
+    *,
+    loc: Optional[ir.Location] = None,
+    ip: Optional[ir.InsertionPoint] = None,
+) -> Pointer:
+    if gmem_ptr.memspace != AddressSpace.gmem:
+        raise ValueError(
+            f"gmem_ptr_to_generic requires pointer in gmem address space, "
+            f"got {gmem_ptr.memspace}"
+        )
+    # Get LLVM pointer and cast to generic address space
+    llvm_ptr = gmem_ptr.to_llvm_ptr(loc=loc, ip=ip)
+    generic_llvm_ptr = llvm.addrspacecast(
+        llvm.PointerType.get(AddressSpace.generic), llvm_ptr, loc=loc, ip=ip
+    )
+    # Create a new cute.Pointer with generic address space, preserving alignment
+    return cute.make_ptr(
+        gmem_ptr.dtype,
+        generic_llvm_ptr,
+        AddressSpace.generic,
+        assumed_align=gmem_ptr.alignment,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def generic_ptr_to_gmem(
+    generic_ptr: Pointer,
+    *,
+    loc: Optional[ir.Location] = None,
+    ip: Optional[ir.InsertionPoint] = None,
+) -> Pointer:
+    if generic_ptr.memspace != AddressSpace.generic:
+        raise ValueError(
+            f"generic_ptr_to_gmem requires pointer in generic address space, "
+            f"got {generic_ptr.memspace}"
+        )
+    # Get LLVM pointer and cast to gmem address space
+    llvm_ptr = generic_ptr.to_llvm_ptr(loc=loc, ip=ip)
+    gmem_llvm_ptr = llvm.addrspacecast(
+        llvm.PointerType.get(AddressSpace.gmem), llvm_ptr, loc=loc, ip=ip
+    )
+    # Create a new cute.Pointer with gmem address space, preserving alignment
+    return cute.make_ptr(
+        generic_ptr.dtype,
+        gmem_llvm_ptr,
+        AddressSpace.gmem,
+        assumed_align=generic_ptr.alignment,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def prefetch_tma_descriptor(tma_desc_ptr: Pointer, *, loc=None, ip=None) -> None:
+    """
+    Prefetch a TMA descriptor from global memory.
+
+    This function prefetches the TMA descriptor pointed to by tma_desc_ptr
+    into the TMA descriptor cache. The pointer must be in generic or global
+    address space. If a gmem pointer is passed, it will be automatically
+    converted to generic address space.
+
+    :param tma_desc_ptr: Pointer to the TMA descriptor in global or generic memory
+    :type tma_desc_ptr: Pointer
+    :raises ValueError: If pointer is not in generic or global address space
+    """
+    if tma_desc_ptr.memspace not in (AddressSpace.gmem, AddressSpace.generic):
+        raise ValueError(
+            f"prefetch_tma_descriptor requires pointer in gmem or generic address space, "
+            f"got {tma_desc_ptr.memspace}"
+        )
+    # Convert gmem pointer to generic if needed
+    if tma_desc_ptr.memspace == AddressSpace.gmem:
+        tma_desc_ptr = gmem_ptr_to_generic(tma_desc_ptr, loc=loc, ip=ip)
+    # Convert cute.Pointer to LLVM pointer for prefetch
+    llvm_ptr = tma_desc_ptr.to_llvm_ptr(loc=loc, ip=ip)
+    from cutlass.cute.arch.nvvm_wrappers import prefetch as nvvm_prefetch
+
+    nvvm_prefetch(llvm_ptr, tensormap=True, loc=loc, ip=ip)
+
+
+def ptr_offset_bytes(ptr: Pointer, byte_offset: int) -> Pointer:
+    """Offset a pointer by a given number of bytes."""
+    element_offset = byte_offset * 8 // ptr.dtype.width
+    return ptr + element_offset
+
+
+@dsl_user_op
+def tensormap_ptr_for_copy(raw_ptr: Pointer, *, loc=None, ip=None) -> Pointer:
+    """
+    Convert a raw TMA descriptor gmem pointer to the type expected by cute.copy.
+
+    cute.copy requires the tma_desc_ptr to be in generic address space and
+    recast to TmaDescriptorTiledType. This utility performs both conversions.
+
+    :param raw_ptr: Raw pointer to TMA descriptor in gmem
+    :type raw_ptr: Pointer
+    :return: Pointer compatible with cute.copy's tma_desc_ptr parameter
+    :rtype: Pointer
+    """
+    generic_ptr = gmem_ptr_to_generic(raw_ptr, loc=loc, ip=ip)
+    tma_desc_ptr_ty = _cute_ir.PtrType.get(
+        _cute_nvgpu_ir.TmaDescriptorTiledType.get(),
+        generic_ptr.memspace,
+        generic_ptr.alignment,
+    )
+    return _cute_ir.recast_iter(tma_desc_ptr_ty, generic_ptr.value)
+
+
+# =============================================================================
+# MoE Utilities
+# =============================================================================
+
+
+@dsl_user_op
+@cute.jit
+def compute_expert_token_range(
+    offs: cute.Tensor,
+    expert_idx: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> Tuple[Int32, Int32]:
+    """
+    Compute token offset and count for a given expert from the cumsum offs tensor.
+
+    :param offs: Cumulative sum tensor of token counts per expert, shape (experts,)
+    :param expert_idx: Index of the expert
+    :return: (token_offset, tokens_i) where token_offset is the start position
+             and tokens_i is the number of tokens for this expert
+    """
+    token_offset = Int32(0)
+    if expert_idx > Int32(0):
+        token_offset = offs[expert_idx - 1]  # type: ignore[assignment]
+    tokens_i = offs[expert_idx] - token_offset
+    return token_offset, tokens_i
+
+
+@dsl_user_op
+def rewrite_tensor_shape(
+    tensor: cute.Tensor,
+    new_shape: Tuple,
+    *,
+    loc=None,
+    ip=None,
+) -> cute.Tensor:
+    """
+    Rewrite tensor shape while keeping the same stride and iterator.
+
+    This is primarily for debug friendliness - shows the actual expert's shape
+    instead of the fake global shape. No runtime overhead as it becomes
+    dead code in non-debug builds.
+
+    :param tensor: Source tensor whose stride and iterator to preserve
+    :param new_shape: New shape to apply
+    :return: New tensor with the given shape but original stride and iterator
+    """
+    new_layout = cute.make_layout(new_shape, stride=tensor.stride, loc=loc, ip=ip)
+    return cute.make_tensor(tensor.iterator, new_layout, loc=loc, ip=ip)
+
+
+# =============================================================================
+# TMA Descriptor Workspace Helper
+# =============================================================================
+
+
+class TensormapWorkspace:
+    """
+    Helper for linear workspace layout of TMA descriptors.
+
+    Manages address calculation for a workspace buffer containing TMA descriptors
+    organized as: for each executor (e.g., expert or group), a fixed set of
+    named descriptor slots.
+
+    Layout: [slot_0_exec_0, slot_1_exec_0, ..., slot_0_exec_1, slot_1_exec_1, ...]
+
+    Example:
+        # 2Dx3D MoE: only C is expert-wise
+        workspace = TensormapWorkspace(workspace_ptr, ["c"])
+
+        # 2Dx2D MoE: A and B are expert-wise
+        workspace = TensormapWorkspace(workspace_ptr, ["a", "b"])
+
+        # General grouped GEMM: all three tensors
+        workspace = TensormapWorkspace(workspace_ptr, ["a", "b", "c"])
+    """
+
+    def __init__(self, workspace_ptr: Pointer, slot_names: list):
+        """
+        :param workspace_ptr: Pointer to the beginning of the workspace buffer
+        :param slot_names: Ordered list of tensor names, defining the slot layout
+                           per executor. e.g., ["a", "b", "c"]
+        """
+        self.workspace_ptr = workspace_ptr
+        self._name_to_slot = {name: i for i, name in enumerate(slot_names)}
+        self._slots_per_executor = len(slot_names)
+
+    @cute.jit
+    def get_ptr(self, tensor_name: str, executor_idx: Int32) -> Pointer:
+        """
+        Get the workspace pointer for a specific TMA descriptor.
+
+        :param tensor_name: Name of the tensor (must be one of the slot_names)
+        :param executor_idx: Index of the executor (e.g., group_idx or expert_idx)
+        :return: Aligned pointer to the TMA descriptor in workspace
+        """
+        if cutlass.const_expr(tensor_name not in self._name_to_slot):
+            raise ValueError(
+                f"Invalid tensor_name '{tensor_name}', "
+                f"expected one of {list(self._name_to_slot.keys())}"
+            )
+        slot = self._name_to_slot[tensor_name]
+        byte_offset = (
+            executor_idx * self._slots_per_executor + slot
+        ) * TensormapDescBytes
+        return ptr_offset_bytes(self.workspace_ptr, byte_offset).align(
+            TensormapDescBytes
+        )
+
+    @staticmethod
+    def size_bytes(num_slots: int, num_executors: int) -> int:
+        """
+        Calculate workspace size in bytes.
+
+        :param num_slots: Number of descriptor slots per executor
+        :param num_executors: Total number of executors (e.g., expert_cnt or group_cnt)
+        :return: Total workspace size in bytes
+        """
+        return num_slots * num_executors * TensormapDescBytes
+
+
+# =============================================================================
+# Online TMA Descriptor Creator (Abstract Base Class)
+# =============================================================================
+
+
+@dataclass(frozen=True)
+class OnlineTensormapDescCreator(ABC):
+    """
+    Abstract base class for building TMA descriptors online (at kernel runtime).
+
+    Subclasses store all needed parameters (both codegen-time configs and runtime
+    values) as explicit instance attributes in __init__. No dict-based APIs.
+
+    Subclasses must implement exactly 2 abstract methods:
+    - construct_and_write: Build TMA descriptor(s) for one executor and write to workspace
+    - get_desc_ptr: Return raw gmem pointer to a specific descriptor in workspace
+
+    To convert the raw pointer for use with cute.copy, callers should use the
+    standalone tensormap_ptr_for_copy() utility.
+    """
+
+    @abstractmethod
+    def construct_and_write(self, executor_idx: Int32, dependency=None) -> None:
+        """
+        Build TMA descriptor(s) for one executor and write to workspace.
+
+        :param executor_idx: Index of the executor (e.g., group_idx or expert_idx).
+            Semantics may vary by subclass when ``dependency`` is provided.
+        :param dependency: Optional pipeline consumer for inter-warp-group
+            synchronization. When provided, the subclass decides when to wait
+            (via ``dependency.wait_and_advance()``) and release. The subclass
+            also decides how to interpret ``executor_idx`` in this mode.
+        """
+        ...
+
+    @abstractmethod
+    def get_desc_ptr(self, tensor_name: str, executor_idx: Int32) -> Pointer:
+        """
+        Get the raw gmem pointer to a specific TMA descriptor in workspace.
+
+        :param tensor_name: Name identifying which tensor's descriptor
+        :param executor_idx: Index of the executor (e.g., group_idx or expert_idx)
+        :return: Raw pointer (gmem) to the TMA descriptor
+        """
+        ...
+
+
+# =============================================================================
+# MoE Grouped GEMM Tensormap Constructor
+# =============================================================================
+
+
+class MoEGroupedGemmTensormapConstructor(OnlineTensormapDescCreator):
+    """
+    Tensormap descriptor constructor for MoE Grouped GEMM (expert-wise descriptors only).
+
+    Non-expert-wise descriptors are passed directly at kernel launch.
+    This class only handles:
+    - 2Dx3D: C descriptors (expert-wise, to avoid write conflicts)
+    - 2Dx2D: A and B descriptors (expert-wise, tokens is reduction axis)
+
+    All parameters are stored as explicit instance attributes (no dicts).
+
+    Workspace layout:
+    - 2Dx3D: [C_0, C_1, ..., C_{n-1}]
+    - 2Dx2D: [A_0, A_1, ..., A_{n-1}, B_0, B_1, ..., B_{n-1}]
+    """
+
+    def __init__(
+        self,
+        scenario: Literal["2Dx3D", "2Dx2D"],
+        # Codegen-time configs
+        a_dtype,
+        b_dtype,
+        c_dtype,
+        a_smem_layout,
+        b_smem_layout,
+        epi_smem_layout,
+        a_tma_op,
+        b_tma_op,
+        c_tma_op,
+        tiled_mma,
+        mma_tiler,
+        cluster_layout_vmnk_shape,
+        epi_tile,
+        # Runtime params
+        a_tensor: cute.Tensor,  # fake GEMM domain A
+        b_tensor: cute.Tensor,  # fake GEMM domain B
+        c_tensor: cute.Tensor,  # fake GEMM domain C
+        offs: cute.Tensor,  # (experts,) cumsum
+        workspace_ptr: Pointer,
+    ) -> None:
+        super().__init__()
+        self.scenario = scenario
+        # Codegen-time configs
+        self.a_dtype = a_dtype
+        self.b_dtype = b_dtype
+        self.c_dtype = c_dtype
+        self.a_smem_layout = a_smem_layout
+        self.b_smem_layout = b_smem_layout
+        self.epi_smem_layout = epi_smem_layout
+        self.a_tma_op = a_tma_op
+        self.b_tma_op = b_tma_op
+        self.c_tma_op = c_tma_op
+        self.tiled_mma = tiled_mma
+        self.mma_tiler = mma_tiler
+        self.cluster_layout_vmnk_shape = cluster_layout_vmnk_shape
+        self.epi_tile = epi_tile
+        # Runtime params
+        self.a_tensor = a_tensor
+        self.b_tensor = b_tensor
+        self.c_tensor = c_tensor
+        self.offs = offs
+        # Workspace with scenario-specific slot layout
+        if scenario == "2Dx3D":
+            self.workspace = TensormapWorkspace(workspace_ptr, ["c"])
+        else:
+            self.workspace = TensormapWorkspace(workspace_ptr, ["a", "b"])
+
+    @staticmethod
+    def get_workspace_size(scenario: Literal["2Dx3D", "2Dx2D"], expert_cnt: int) -> int:
+        """Calculate workspace size in bytes for tensormap descriptors."""
+        if scenario == "2Dx3D":
+            return TensormapWorkspace.size_bytes(1, expert_cnt)  # only C
+        else:
+            return TensormapWorkspace.size_bytes(2, expert_cnt)  # A and B
+
+    @cute.jit
+    def get_desc_ptr(self, tensor_name: str, executor_idx: Int32) -> Pointer:
+        return self.workspace.get_ptr(tensor_name, executor_idx)
+
+    @cute.jit
+    def construct_and_write(self, executor_idx: Int32, dependency=None) -> None:
+        """
+        Create expert-wise tensormap descriptors for the given expert.
+
+        - 2Dx3D: Creates C descriptor for this expert
+        - 2Dx2D: Creates A and B descriptors for this expert
+        """
+        if cutlass.const_expr(self.scenario == "2Dx3D"):
+            self._construct_c_desc_2dx3d(executor_idx)
+        else:  # 2Dx2D
+            self._construct_ab_descs_2dx2d(executor_idx)
+
+    @cute.jit
+    def _construct_c_desc_2dx3d(self, expert_idx: Int32) -> None:
+        """
+        2Dx3D: Create expert-wise C descriptor.
+        C tensor: (fake_m, n, 1) = (tokens_sum, intermediate, 1)
+        Slice fake_m -> (tokens_i, intermediate, 1) per expert.
+        """
+        token_offset, tokens_i = compute_expert_token_range(self.offs, expert_idx)
+
+        c_ptr = self.c_tensor.iterator
+        c_stride = self.c_tensor.stride
+        intermediate = self.c_tensor.shape[1]  # type: ignore[index]
+
+        c1 = cutlass.Int32(1)
+        c0 = cutlass.Int32(0)
+
+        c_ptr_i = c_ptr + token_offset * c_stride[0]  # type: ignore[index]
+        c_layout_i = cute.make_layout(
+            (tokens_i, intermediate, c1),
+            stride=(c_stride[0], c_stride[1], c0),  # type: ignore[index]
+        )
+        c_tensor_i = cute.make_tensor(c_ptr_i, c_layout_i)
+
+        tma_atom_c, _ = cpasync.make_tiled_tma_atom(
+            self.c_tma_op,
+            c_tensor_i,
+            self.epi_smem_layout,
+            self.epi_tile,
+        )
+        cpasync.copy_tensormap(tma_atom_c, self.get_desc_ptr("c", expert_idx))
+
+    @cute.jit
+    def _construct_ab_descs_2dx2d(self, expert_idx: Int32) -> None:
+        """
+        2Dx2D: Create expert-wise A and B descriptors.
+        A: (m, fake_k, 1) -> slice to (m, tokens_i, 1)
+        B: (n, fake_k, 1) -> slice to (n, tokens_i, 1)
+        """
+        token_offset, tokens_i = compute_expert_token_range(self.offs, expert_idx)
+
+        c1 = cutlass.Int32(1)
+        c0 = cutlass.Int32(0)
+
+        # A tensor: (m, fake_k, 1) -> (m, tokens_i, 1)
+        a_ptr = self.a_tensor.iterator
+        a_stride = self.a_tensor.stride
+        a_m = self.a_tensor.shape[0]  # type: ignore[index]
+
+        a_ptr_i = a_ptr + token_offset * a_stride[1]  # type: ignore[index]
+        a_layout_i = cute.make_layout(
+            (a_m, tokens_i, c1),
+            stride=(a_stride[0], a_stride[1], c0),  # type: ignore[index]
+        )
+        a_tensor_i = cute.make_tensor(a_ptr_i, a_layout_i)
+
+        tma_atom_a, _ = cute.nvgpu.make_tiled_tma_atom_A(
+            self.a_tma_op,
+            a_tensor_i,
+            self.a_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk_shape,
+        )
+        cpasync.copy_tensormap(tma_atom_a, self.get_desc_ptr("a", expert_idx))
+
+        # B tensor: (n, fake_k, 1) -> (n, tokens_i, 1)
+        b_ptr = self.b_tensor.iterator
+        b_stride = self.b_tensor.stride
+        b_n = self.b_tensor.shape[0]  # type: ignore[index]
+
+        b_ptr_i = b_ptr + token_offset * b_stride[1]  # type: ignore[index]
+        b_layout_i = cute.make_layout(
+            (b_n, tokens_i, c1),
+            stride=(b_stride[0], b_stride[1], c0),  # type: ignore[index]
+        )
+        b_tensor_i = cute.make_tensor(b_ptr_i, b_layout_i)
+
+        tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B(
+            self.b_tma_op,
+            b_tensor_i,
+            self.b_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk_shape,
+        )
+        cpasync.copy_tensormap(tma_atom_b, self.get_desc_ptr("b", expert_idx))
+
+
+# =============================================================================
+# MoE Scaled Grouped GEMM Tensormap Constructor
+# =============================================================================
+
+
+class MoEScaledGroupedGemmTensormapConstructor(OnlineTensormapDescCreator):
+    """
+    Tensormap descriptor constructor for MoE Scaled Grouped GEMM (block-scaled).
+
+    .. py:attribute:: ChunkSize
+        :value: 128
+
+        Number of experts processed per chunk in the desc_init_kernel.
+        Must match the warp-group width (4 warps × 32 threads).
+
+    Extends MoEGroupedGemmTensormapConstructor with SFA/SFB descriptor support.
+
+    Expert-wise descriptors only — non-expert-wise descriptors are passed
+    directly at kernel launch.
+
+    Workspace layout:
+    - 2Dx3D: [C_0, C_1, ..., C_{n-1}]  (1 slot per expert)
+    - 2Dx2D: [A_0, B_0, SFA_0, SFB_0, A_1, B_1, SFA_1, SFB_1, ...]  (4 slots per expert)
+
+    :param scenario: "2Dx3D" or "2Dx2D"
+    :param sf_vec_size: Scale factor vector size (32 for MXFP8/MXFP4, 16 for NVFP4)
+    :param a_dtype: Data type for tensor A
+    :param b_dtype: Data type for tensor B
+    :param c_dtype: Data type for tensor C
+    :param sf_dtype: Data type for scale factors (SFA/SFB)
+    :param a_smem_layout: SMEM layout for A TMA
+    :param b_smem_layout: SMEM layout for B TMA
+    :param epi_smem_layout: SMEM layout for epilogue (C) TMA
+    :param sfa_smem_layout: SMEM layout for SFA TMA
+    :param sfb_smem_layout: SMEM layout for SFB TMA
+    :param a_tma_op: TMA operation for A
+    :param b_tma_op: TMA operation for B
+    :param c_tma_op: TMA operation for C (S2G store or reduce)
+    :param sfa_tma_op: TMA operation for SFA
+    :param sfb_tma_op: TMA operation for SFB
+    :param tiled_mma: TiledMma for A/B/SFA/C TMA atom construction
+    :param tiled_mma_sfb: TiledMma for SFB (separate due to 2CTA replication)
+    :param mma_tiler: MMA tiler shape (M, N, K)
+    :param mma_tiler_sfb: MMA tiler shape for SFB
+    :param cluster_layout_vmnk_shape: Cluster layout shape for A/B/SFA multicast
+    :param cluster_layout_sfb_vmnk_shape: Cluster layout shape for SFB multicast
+    :param epi_tile: Epilogue tile shape
+    :param a_tensor: Fake GEMM domain A tensor
+    :param b_tensor: Fake GEMM domain B tensor
+    :param c_tensor: Fake GEMM domain C tensor
+    :param sfa_tensor: Fake GEMM domain SFA tensor (atom-tiled layout)
+    :param sfb_tensor: Fake GEMM domain SFB tensor (atom-tiled layout)
+    :param offs: (experts,) cumsum offsets in data domain
+    :param offs_padded: (experts,) cumsum offsets in padded scale domain
+    :param workspace_ptr: Pointer to workspace for TMA descriptors
+    :param expert_cnt: Total number of experts
+    """
+
+    ChunkSize = 128
+
+    def __init__(
+        self,
+        scenario: Literal["2Dx3D", "2Dx2D"],
+        sf_vec_size: int,
+        # Codegen-time configs: dtypes
+        a_dtype,
+        b_dtype,
+        c_dtype,
+        sf_dtype,
+        # Codegen-time configs: SMEM layouts
+        a_smem_layout,
+        b_smem_layout,
+        epi_smem_layout,
+        sfa_smem_layout,
+        sfb_smem_layout,
+        # Codegen-time configs: TMA ops
+        a_tma_op,
+        b_tma_op,
+        c_tma_op,
+        sfa_tma_op,
+        sfb_tma_op,
+        # Codegen-time configs: MMA / cluster / tile
+        tiled_mma,
+        tiled_mma_sfb,
+        mma_tiler,
+        mma_tiler_sfb,
+        cluster_layout_vmnk_shape,
+        cluster_layout_sfb_vmnk_shape,
+        epi_tile,
+        # Runtime params
+        a_tensor: cute.Tensor,
+        b_tensor: cute.Tensor,
+        c_tensor: cute.Tensor,
+        sfa_tensor: cute.Tensor,
+        sfb_tensor: cute.Tensor,
+        offs: cute.Tensor,
+        offs_padded: cute.Tensor,
+        workspace_ptr: Pointer,
+        expert_cnt: Optional[Union[Int32, int]] = None,
+    ) -> None:
+        super().__init__()
+        self.scenario = scenario
+        self.sf_vec_size = sf_vec_size
+        # Dtypes
+        self.a_dtype = a_dtype
+        self.b_dtype = b_dtype
+        self.c_dtype = c_dtype
+        self.sf_dtype = sf_dtype
+        # SMEM layouts
+        self.a_smem_layout = a_smem_layout
+        self.b_smem_layout = b_smem_layout
+        self.epi_smem_layout = epi_smem_layout
+        self.sfa_smem_layout = sfa_smem_layout
+        self.sfb_smem_layout = sfb_smem_layout
+        # TMA ops
+        self.a_tma_op = a_tma_op
+        self.b_tma_op = b_tma_op
+        self.c_tma_op = c_tma_op
+        self.sfa_tma_op = sfa_tma_op
+        self.sfb_tma_op = sfb_tma_op
+        # MMA / cluster / tile
+        self.tiled_mma = tiled_mma
+        self.tiled_mma_sfb = tiled_mma_sfb
+        self.mma_tiler = mma_tiler
+        self.mma_tiler_sfb = mma_tiler_sfb
+        self.cluster_layout_vmnk_shape = cluster_layout_vmnk_shape
+        self.cluster_layout_sfb_vmnk_shape = cluster_layout_sfb_vmnk_shape
+        self.epi_tile = epi_tile
+        # Runtime params
+        self.a_tensor = a_tensor
+        self.b_tensor = b_tensor
+        self.c_tensor = c_tensor
+        self.sfa_tensor = sfa_tensor
+        self.sfb_tensor = sfb_tensor
+        self.offs = offs
+        self.offs_padded = offs_padded
+        self.expert_cnt = expert_cnt
+        # Workspace with scenario-specific slot layout
+        if scenario == "2Dx3D":
+            self.workspace = TensormapWorkspace(workspace_ptr, ["c"])
+        else:
+            self.workspace = TensormapWorkspace(workspace_ptr, ["a", "b", "sfa", "sfb"])
+
+    @staticmethod
+    def get_workspace_size(scenario: Literal["2Dx3D", "2Dx2D"], expert_cnt: int) -> int:
+        """Calculate workspace size in bytes for tensormap descriptors."""
+        if scenario == "2Dx3D":
+            return TensormapWorkspace.size_bytes(1, expert_cnt)  # C only
+        else:
+            return TensormapWorkspace.size_bytes(4, expert_cnt)  # A, B, SFA, SFB
+
+    @cute.jit
+    def get_desc_ptr(self, tensor_name: str, executor_idx: Int32) -> Pointer:
+        return self.workspace.get_ptr(tensor_name, executor_idx)
+
+    @cute.jit
+    def construct_and_write(self, lane_in_group: Int32, dependency=None) -> None:
+        """
+        Create expert-wise tensormap descriptors for all experts.
+
+        ``lane_in_group`` is the thread's position within its warp group
+        (0..ChunkSize-1). The method loops internally over all experts in
+        chunks of ``ChunkSize``, with two-phase pipeline synchronization
+        per chunk.
+
+        Per-chunk execution:
+
+        1. Phase 1: Build descriptors that do NOT depend on ``offs_padded``
+           (A/B for 2Dx2D, C for 2Dx3D). Overlaps with Group A's prefix sum.
+        2. Barrier: ``consumer.wait_and_advance()`` — all threads participate.
+        3. Phase 2: Build descriptors that depend on ``offs_padded``
+           (SFA/SFB for 2Dx2D). Reads padded offsets from SMEM buffer.
+        4. Release: ``handle.release()`` — all threads participate.
+
+        :param lane_in_group: Thread's position within the warp group (0..127).
+        :param dependency: ``(PipelineConsumer, smem_offs_padded)`` — the
+            consumer for mbarrier sync, and the SMEM tensor of shape
+            ``(ChunkSize + 1,)`` with layout ``[carry, offs_padded[0..127]]``.
+        """
+        consumer, smem_offs_padded = dependency
+        assert self.expert_cnt is not None
+        num_chunks = (self.expert_cnt + self.ChunkSize - 1) // self.ChunkSize
+
+        chunk_idx = cutlass.Int32(0)
+        while chunk_idx < num_chunks:
+            expert_idx = chunk_idx * self.ChunkSize + lane_in_group
+            in_bounds = expert_idx < self.expert_cnt
+
+            # Phase 1: non-dependent descriptors
+            if in_bounds:
+                if cutlass.const_expr(self.scenario == "2Dx2D"):
+                    self._construct_ab_descs_2dx2d(expert_idx)
+                else:
+                    self._construct_c_desc_2dx3d(expert_idx)
+
+            # All threads participate in barrier (fixed arrive count)
+            handle = consumer.wait_and_advance()
+
+            # Phase 2: dependent descriptors (read padded offsets from SMEM)
+            if in_bounds:
+                if cutlass.const_expr(self.scenario == "2Dx2D"):
+                    # smem_offs_padded layout: [carry, chunk[0], ..., chunk[127]]
+                    # padded_offset = smem[lane]     (prev expert's cumulative)
+                    # padded_end    = smem[lane + 1] (this expert's cumulative)
+                    padded_offset = smem_offs_padded[lane_in_group]
+                    padded_size_i = smem_offs_padded[lane_in_group + 1] - padded_offset
+                    self._construct_sf_descs_2dx2d_direct(
+                        expert_idx, padded_offset, padded_size_i
+                    )
+
+            # All threads release (fixed arrive count)
+            handle.release()
+
+            chunk_idx += 1
+
+    # -----------------------------------------------------------------
+    # 2Dx3D: C descriptor (same as MoEGroupedGemmTensormapConstructor)
+    # -----------------------------------------------------------------
+
+    @cute.jit
+    def _construct_c_desc_2dx3d(self, expert_idx: Int32) -> None:
+        """
+        2Dx3D: Create expert-wise C descriptor.
+        C: (fake_m, n, 1) -> slice to (tokens_i, n, 1) per expert.
+        """
+        token_offset, tokens_i = compute_expert_token_range(self.offs, expert_idx)
+        c1 = cutlass.Int32(1)
+
+        c_i = cute.domain_offset((token_offset, 0, 0), self.c_tensor)
+        c_i = rewrite_tensor_shape(c_i, (tokens_i, self.c_tensor.shape[1], c1))  # type: ignore[index]
+
+        tma_atom_c, _ = cpasync.make_tiled_tma_atom(
+            self.c_tma_op,
+            c_i,
+            self.epi_smem_layout,
+            self.epi_tile,
+        )
+        cpasync.copy_tensormap(tma_atom_c, self.get_desc_ptr("c", expert_idx))
+
+    # -----------------------------------------------------------------
+    # 2Dx2D: A, B descriptors (same as MoEGroupedGemmTensormapConstructor)
+    # -----------------------------------------------------------------
+
+    @cute.jit
+    def _construct_ab_descs_2dx2d(self, expert_idx: Int32) -> None:
+        """
+        2Dx2D: Create expert-wise A and B descriptors.
+        A: (m, fake_k, 1) -> slice to (m, tokens_i, 1)
+        B: (n, fake_k, 1) -> slice to (n, tokens_i, 1)
+        """
+        token_offset, tokens_i = compute_expert_token_range(self.offs, expert_idx)
+        c1 = cutlass.Int32(1)
+
+        # A: (m, fake_k, 1) -> domain_offset + rewrite shape
+        a_i = cute.domain_offset((0, token_offset, 0), self.a_tensor)
+        a_i = rewrite_tensor_shape(a_i, (self.a_tensor.shape[0], tokens_i, c1))  # type: ignore[index]
+
+        tma_atom_a, _ = cute.nvgpu.make_tiled_tma_atom_A(
+            self.a_tma_op,
+            a_i,
+            self.a_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk_shape,
+        )
+        cpasync.copy_tensormap(tma_atom_a, self.get_desc_ptr("a", expert_idx))
+
+        # B: (n, fake_k, 1) -> domain_offset + rewrite shape
+        b_i = cute.domain_offset((0, token_offset, 0), self.b_tensor)
+        b_i = rewrite_tensor_shape(b_i, (self.b_tensor.shape[0], tokens_i, c1))  # type: ignore[index]
+
+        tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B(
+            self.b_tma_op,
+            b_i,
+            self.b_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk_shape,
+        )
+        cpasync.copy_tensormap(tma_atom_b, self.get_desc_ptr("b", expert_idx))
+
+    # -----------------------------------------------------------------
+    # 2Dx2D: SFA, SFB descriptors (new for block-scaled)
+    # -----------------------------------------------------------------
+
+    @cute.jit
+    def _construct_sf_descs_2dx2d_direct(
+        self,
+        expert_idx: Int32,
+        padded_offset: Int32,
+        padded_size_i: Int32,
+    ) -> None:
+        """
+        2Dx2D: Create expert-wise SFA and SFB descriptors with pre-computed
+        padded offset and size.
+
+        This variant allows the caller to supply padded offsets from SMEM
+        (in desc_init_kernel) instead of reading from ``self.offs_padded`` in GMEM.
+        """
+        c1 = cutlass.Int32(1)
+
+        a_chunks_to_move = (
+            padded_offset
+            // self.sf_vec_size
+            * cute.size(self.sfa_tensor, mode=[0])
+            // 128
+        )
+        a_elems_to_move = (
+            cute.size(self.sfa_tensor, mode=[0]) * padded_offset // self.sf_vec_size
+        )
+        b_chunks_to_move = (
+            padded_offset
+            // self.sf_vec_size
+            * cute.size(self.sfb_tensor, mode=[0])
+            // 128
+        )
+        b_elems_to_move = (
+            cute.size(self.sfb_tensor, mode=[0]) * padded_offset // self.sf_vec_size
+        )
+
+        per_expert_sfa_shape = (self.sfa_tensor.shape[0], padded_size_i, c1)  # type: ignore[index]
+        sfa_layout_i = tile_atom_to_shape_SF(per_expert_sfa_shape, self.sf_vec_size)
+        sfa_i = cute.make_tensor(
+            self.sfa_tensor.iterator + a_elems_to_move, sfa_layout_i
+        )
+
+        tma_atom_sfa, _ = cute.nvgpu.make_tiled_tma_atom_A(
+            self.sfa_tma_op,
+            sfa_i,
+            self.sfa_smem_layout,
+            self.mma_tiler,
+            self.tiled_mma,
+            self.cluster_layout_vmnk_shape,
+            internal_type=cutlass.Uint64,
+        )
+        cpasync.copy_tensormap(tma_atom_sfa, self.get_desc_ptr("sfa", expert_idx))
+
+        per_expert_sfb_shape = (self.sfb_tensor.shape[0], padded_size_i, c1)  # type: ignore[index]
+        sfb_layout_i = tile_atom_to_shape_SF(per_expert_sfb_shape, self.sf_vec_size)
+        sfb_i = cute.make_tensor(
+            self.sfb_tensor.iterator + b_elems_to_move, sfb_layout_i
+        )
+
+        tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B(
+            self.sfb_tma_op,
+            sfb_i,
+            self.sfb_smem_layout,
+            self.mma_tiler_sfb,
+            self.tiled_mma_sfb,
+            self.cluster_layout_sfb_vmnk_shape,
+            internal_type=cutlass.Uint64,
+        )
+        cpasync.copy_tensormap(tma_atom_sfb, self.get_desc_ptr("sfb", expert_idx))

examples/python/CuTeDSL/cute/blackwell/kernel/reduce/reduce.py

@@ -0,0 +1,533 @@
+# 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.
+
+"""
+Hierarchical Reduction Utilities for CuTe-DSL Kernels
+=====================================================
+
+This module provides reusable reduction primitives for GPU kernels that need to
+reduce values across warps, thread blocks, and clusters (SM90+).
+
+Overview
+--------
+GPU reductions typically follow a hierarchical pattern:
+
+1. **Warp Reduction**: Threads within a warp reduce using shuffle instructions.
+   Use `cute.arch.warp_reduction()` from the CuTe-DSL library.
+
+2. **Block Reduction**: Multiple warps within a block reduce using shared memory.
+   Use `block_reduce()` from this module.
+
+3. **Cluster Reduction** (SM90+): Multiple CTAs in a cluster reduce using
+   distributed shared memory and mbarrier synchronization.
+   Use `cluster_reduce()` from this module.
+
+4. **Row Reduction**: Orchestrates all levels based on problem configuration.
+   Use `row_reduce()` from this module.
+
+Shared Memory Buffer Layout Assumptions
+---------------------------------------
+
+For `block_reduce`:
+    - Buffer shape: (rows_per_block, warps_per_row)
+    - Each warp's lane 0 writes its reduced value to buffer[row_idx, col_idx]
+    - Thread mapping: row_idx = warp_idx // warps_per_row
+                      col_idx = warp_idx % warps_per_row
+
+    Example for 8 warps, 2 rows, 4 warps per row:
+        Warp 0 -> buffer[0, 0]    Warp 4 -> buffer[1, 0]
+        Warp 1 -> buffer[0, 1]    Warp 5 -> buffer[1, 1]
+        Warp 2 -> buffer[0, 2]    Warp 6 -> buffer[1, 2]
+        Warp 3 -> buffer[0, 3]    Warp 7 -> buffer[1, 3]
+
+For `cluster_reduce`:
+    - Buffer shape: (rows_per_block, (warps_per_row, cluster_n))
+    - The second dimension is hierarchical: (local_warp_slot, cta_rank)
+    - Each CTA contributes to its own slot in the cluster dimension
+
+    Example for cluster_n=4, 2 warps per row:
+        CTA 0, Warp 0 -> buffer[row, (0, 0)]
+        CTA 0, Warp 1 -> buffer[row, (1, 0)]
+        CTA 1, Warp 0 -> buffer[row, (0, 1)]
+        CTA 1, Warp 1 -> buffer[row, (1, 1)]
+        ... etc for CTAs 2, 3
+
+Mbarrier Requirements (Cluster Reduction)
+-----------------------------------------
+For cluster reduction, the caller must:
+1. Allocate an mbarrier in shared memory
+2. Initialize it with `cute.arch.mbarrier_init(mbar_ptr, thread_count)`
+3. Pass the mbarrier pointer to `cluster_reduce()`
+
+The cluster_reduce function handles:
+- Setting up the expected transaction count
+- Performing async cross-CTA stores
+- Waiting for all stores to complete
+
+Usage Example
+-------------
+
+.. code-block:: python
+
+    from reduce import row_reduce, block_reduce, cluster_reduce
+
+    @cute.jit
+    def my_kernel(...):
+        # Allocate shared memory for reduction
+        # Shape depends on warps_per_row and cluster_n
+        if cluster_n > 1:
+            reduction_buffer = cute.make_smem_tensor(
+                cute.make_layout((rows_per_block, (warps_per_row, cluster_n))),
+                Float32
+            )
+        else:
+            reduction_buffer = cute.make_smem_tensor(
+                cute.make_layout((rows_per_block, warps_per_row)),
+                Float32
+            )
+
+        # Perform row reduction
+        result = row_reduce(
+            tensor_ssa,
+            cute.ReductionOp.ADD,
+            threads_per_row,
+            reduction_buffer,
+            mbar_ptr,
+            cluster_n,
+            init_val=Float32(0.0)
+        )
+
+References
+----------
+The cluster synchronization primitives (set_block_rank, store_shared_remote)
+are inspired by Quack: https://github.com/Dao-AILab/quack
+"""
+
+import operator
+from collections.abc import Callable
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32
+from cutlass._mlir.dialects import llvm
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+
+# =============================================================================
+# Inline PTX Operations for Cluster Communication
+# =============================================================================
+#
+# These operations enable cross-CTA communication within a cluster (SM90+).
+# They use inline PTX assembly for functionality not yet exposed in MLIR.
+# =============================================================================
+
+
+@dsl_user_op
+def set_block_rank(
+    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
+) -> Int32:
+    """
+    Map a shared memory pointer to the equivalent address in another CTA's
+    shared memory within the same cluster.
+
+    This uses the PTX `mapa.shared::cluster` instruction to translate a local
+    shared memory address to the corresponding address in a peer CTA's shared
+    memory space.
+
+    Args:
+        smem_ptr: Pointer to local shared memory
+        peer_cta_rank_in_cluster: Target CTA's rank within the cluster (0 to cluster_size-1)
+
+    Returns:
+        Int32 representing the mapped address in the peer CTA's shared memory
+
+    Note:
+        This operation requires SM90+ with cluster support enabled.
+        The cluster must be launched with the appropriate cluster dimensions.
+    """
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+    return Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
+            "mapa.shared::cluster.u32 $0, $1, $2;",
+            "=r,r,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+
+
+@dsl_user_op
+def store_shared_remote(
+    val: Float32,
+    smem_ptr: cute.Pointer,
+    mbar_ptr: cute.Pointer,
+    peer_cta_rank_in_cluster: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    """
+    Asynchronously store a Float32 value to shared memory on a remote CTA
+    within the cluster, with mbarrier completion tracking.
+
+    This uses the PTX `st.async.shared::cluster` instruction which:
+    1. Translates the local smem address to the peer CTA's address space
+    2. Performs an asynchronous store to the remote shared memory
+    3. Signals the mbarrier when the store completes
+
+    Args:
+        val: The Float32 value to store
+        smem_ptr: Pointer to the destination in local shared memory coordinates
+        mbar_ptr: Pointer to the mbarrier that tracks completion
+        peer_cta_rank_in_cluster: Target CTA's rank within the cluster
+
+    Note:
+        - The mbarrier must be initialized with the expected transaction byte count
+        - Use `cute.arch.mbarrier_arrive_and_expect_tx()` to set up the transaction
+        - Use `cute.arch.mbarrier_wait()` to wait for all stores to complete
+        - This operation requires SM90+ with cluster support enabled
+    """
+    remote_smem_ptr_i32 = set_block_rank(
+        smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    remote_mbar_ptr_i32 = set_block_rank(
+        mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
+    ).ir_value()
+    llvm.inline_asm(
+        None,
+        [remote_smem_ptr_i32, val.ir_value(loc=loc, ip=ip), remote_mbar_ptr_i32],
+        "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$0], $1, [$2];",
+        "r,f,r",
+        has_side_effects=True,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+    )
+
+
+@dsl_user_op
+def elem_pointer(x: cute.Tensor, coord, *, loc=None, ip=None) -> cute.Pointer:
+    """
+    Get a pointer to an element at the specified coordinate in a tensor.
+
+    This is useful for getting the shared memory address of a specific element
+    when performing cross-CTA stores in cluster reduction.
+
+    Args:
+        x: The tensor (typically a shared memory tensor)
+        coord: The coordinate tuple, can be hierarchical like (row, (col, cluster_idx))
+
+    Returns:
+        Pointer to the element at the specified coordinate
+    """
+    return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
+
+
+# =============================================================================
+# Block-Level Reduction
+# =============================================================================
+
+
+@cute.jit
+def block_reduce(
+    val: Float32,
+    op: Callable,
+    reduction_buffer: cute.Tensor,
+    init_val: Float32,
+) -> Float32:
+    """
+    Reduce values across all warps within a thread block using shared memory.
+
+    This function assumes each warp has already performed a warp-level reduction
+    and is contributing a single value (from lane 0). The function then:
+    1. Writes each warp's value to shared memory
+    2. Synchronizes the block
+    3. Performs a final warp reduction across the collected values
+
+    Args:
+        val: The warp-reduced value (only lane 0's value is used)
+        op: Binary reduction operator, e.g., `operator.add` or `cute.arch.fmax`
+        reduction_buffer: Shared memory tensor with shape (rows_per_block, warps_per_row)
+        init_val: Identity element for the reduction (0 for sum, -inf for max)
+
+    Returns:
+        The block-reduced result (same value across all threads)
+
+    Buffer Layout:
+        - Shape: (rows_per_block, warps_per_row)
+        - warps_per_row is inferred from reduction_buffer.shape[1]
+        - Thread mapping:
+            row_idx = warp_idx // warps_per_row
+            col_idx = warp_idx % warps_per_row
+
+    Example:
+        For a block with 8 warps processing 2 rows (4 warps per row):
+
+        .. code-block:: python
+
+            reduction_buffer = cute.make_smem_tensor(
+                cute.make_layout((2, 4)),  # 2 rows, 4 warps per row
+                Float32
+            )
+            result = block_reduce(warp_val, operator.add, reduction_buffer, Float32(0.0))
+    """
+    lane_idx = cute.arch.lane_idx()
+    warp_idx = cute.arch.warp_idx()
+    warps_per_row = cute.size(reduction_buffer.shape[1])
+    row_idx = warp_idx // warps_per_row
+    col_idx = warp_idx % warps_per_row
+
+    # Lane 0 of each warp writes its value to shared memory
+    if lane_idx == 0:
+        reduction_buffer[row_idx, col_idx] = val
+    cute.arch.barrier()
+
+    # All lanes participate in reading and reducing
+    # Only lanes < warps_per_row have valid data
+    block_reduce_val = init_val
+    if lane_idx < warps_per_row:
+        block_reduce_val = reduction_buffer[row_idx, lane_idx]
+    return cute.arch.warp_reduction(block_reduce_val, op)
+
+
+# =============================================================================
+# Cluster-Level Reduction (SM90+)
+# =============================================================================
+
+
+@cute.jit
+def cluster_reduce(
+    val: Float32,
+    op: Callable,
+    reduction_buffer: cute.Tensor,
+    mbar_ptr: cute.Pointer,
+    cluster_n: cutlass.Constexpr[int],
+    init_val: Float32,
+) -> Float32:
+    """
+    Reduce values across all CTAs within a cluster using distributed shared memory.
+
+    This function extends block reduction to work across multiple CTAs in a cluster
+    using asynchronous cross-CTA stores and mbarrier synchronization. It:
+    1. Sets up the mbarrier with expected transaction count
+    2. Asynchronously stores each warp's value to all peer CTAs
+    3. Waits for all stores to complete
+    4. Reduces across all collected values
+
+    Args:
+        val: The warp-reduced value (only lane 0's value is used for stores)
+        op: Binary reduction operator, e.g., `operator.add` or `cute.arch.fmax`
+        reduction_buffer: Shared memory tensor with hierarchical shape
+                          (rows_per_block, (warps_per_row, cluster_n))
+        mbar_ptr: Pointer to an initialized mbarrier in shared memory
+        cluster_n: Number of CTAs in the cluster (compile-time constant)
+        init_val: Identity element for the reduction (0 for sum, -inf for max)
+
+    Returns:
+        The cluster-reduced result (same value across all threads in all CTAs)
+
+    Buffer Layout:
+        - Shape: (rows_per_block, (warps_per_row, cluster_n))
+        - The second dimension is hierarchical:
+          - First level: warps_per_row (local warp slots)
+          - Second level: cluster_n (one slot per CTA in cluster)
+        - Access pattern: buffer[row_idx, (col_idx, cta_rank)]
+
+    Requirements:
+        - SM90+ with cluster support
+        - Mbarrier must be initialized before calling
+        - Kernel must be launched with appropriate cluster dimensions
+
+    Example:
+        For a cluster of 4 CTAs, each with 2 warps per row:
+
+        .. code-block:: python
+
+            # Allocate buffer with cluster dimension
+            reduction_buffer = cute.make_smem_tensor(
+                cute.make_layout((rows_per_block, (2, 4))),  # 2 warps, 4 CTAs
+                Float32
+            )
+
+            # Initialize mbarrier (once per kernel)
+            mbar = cute.make_smem_tensor(cute.make_layout((1,)), cute.arch.Mbarrier)
+            cute.arch.mbarrier_init(mbar.iterator, thread_count)
+
+            # Perform cluster reduction
+            result = cluster_reduce(
+                warp_val, operator.add, reduction_buffer,
+                mbar.iterator, cluster_n=4, init_val=Float32(0.0)
+            )
+    """
+    cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+    lane_idx = cute.arch.lane_idx()
+    warp_idx = cute.arch.warp_idx()
+
+    rows_per_block = reduction_buffer.shape[0]
+    warps_per_row = reduction_buffer.shape[1][0]
+
+    row_idx = warp_idx // warps_per_row
+    col_idx = warp_idx % warps_per_row
+
+    # Warp 0, lane 0 sets up mbarrier with expected transaction count
+    # Each warp sends cluster_n stores (one to each CTA), each store is 4 bytes
+    if warp_idx == 0:
+        with cute.arch.elect_one():
+            num_warps = rows_per_block * warps_per_row
+            expected_bytes = num_warps * cluster_n * 4  # 4 bytes per Float32
+            cute.arch.mbarrier_arrive_and_expect_tx(mbar_ptr, expected_bytes)
+
+    # Each lane < cluster_n writes to a different CTA's shared memory
+    # This distributes the warp's value to all CTAs in the cluster
+    if lane_idx < cluster_n:
+        store_shared_remote(
+            val,
+            elem_pointer(reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster))),
+            mbar_ptr,
+            peer_cta_rank_in_cluster=lane_idx,
+        )
+
+    # Wait for all cross-CTA stores to complete
+    cute.arch.mbarrier_wait(mbar_ptr, phase=0)
+
+    # Now each CTA has all values from all CTAs in the cluster
+    # Reduce across all collected values
+    num_total = warps_per_row * cluster_n
+    num_iter = cute.ceil_div(num_total, 32)
+
+    block_reduce_val = init_val
+    for i in cutlass.range_constexpr(num_iter):
+        idx = lane_idx + i * 32
+        if idx < num_total:
+            block_reduce_val = op(block_reduce_val, reduction_buffer[row_idx, idx])
+
+    return cute.arch.warp_reduction(block_reduce_val, op)
+
+
+# =============================================================================
+# Row Reduction (Orchestration Function)
+# =============================================================================
+
+
+@cute.jit
+def row_reduce(
+    x: cute.TensorSSA,
+    op: cute.ReductionOp,
+    threads_per_row: cutlass.Constexpr[int],
+    reduction_buffer: cute.Tensor,
+    mbar_ptr,
+    cluster_n: cutlass.Constexpr[int],
+    init_val: Float32,
+):
+    """
+    Perform hierarchical row reduction with automatic selection of reduction strategy.
+
+    This function orchestrates the full reduction pipeline:
+    1. Local reduction: Each thread reduces its portion of the row
+    2. Warp reduction: Threads within a warp reduce using shuffles
+    3. Block reduction: If needed, warps reduce using shared memory
+    4. Cluster reduction: If needed, CTAs reduce using distributed shared memory
+
+    The function automatically selects the appropriate reduction level based on
+    `threads_per_row` and `cluster_n`.
+
+    Args:
+        x: TensorSSA containing the values to reduce (in registers)
+        op: Reduction operation (cute.ReductionOp.ADD or cute.ReductionOp.MAX)
+        threads_per_row: Number of threads cooperating on each row (compile-time)
+        reduction_buffer: Shared memory tensor for block/cluster reduction
+        mbar_ptr: Mbarrier pointer (only used if cluster_n > 1)
+        cluster_n: Number of CTAs in cluster (1 for single-CTA reduction)
+        init_val: Identity element for the reduction
+
+    Returns:
+        The fully reduced result for each row
+
+    Reduction Strategy Selection:
+        - threads_per_row <= 32, cluster_n == 1: Warp reduction only
+        - threads_per_row > 32, cluster_n == 1: Warp + block reduction
+        - cluster_n > 1: Warp + cluster reduction (handles all cases)
+
+    Example:
+        .. code-block:: python
+
+            # Sum reduction across 128 threads per row, single CTA
+            result = row_reduce(
+                tensor_ssa,
+                cute.ReductionOp.ADD,
+                threads_per_row=128,
+                reduction_buffer=smem_buffer,
+                mbar_ptr=None,
+                cluster_n=1,
+                init_val=Float32(0.0)
+            )
+
+            # Max reduction across 256 threads per row, 4 CTAs in cluster
+            result = row_reduce(
+                tensor_ssa,
+                cute.ReductionOp.MAX,
+                threads_per_row=256,
+                reduction_buffer=smem_buffer,
+                mbar_ptr=mbar.iterator,
+                cluster_n=4,
+                init_val=Float32.neg_inf
+            )
+    """
+    # Step 1: Local reduction - each thread reduces its register values
+    local_val = x.reduce(op, init_val=init_val, reduction_profile=0)
+
+    # Map ReductionOp enum to binary operator for warp/block reductions
+    warp_op = {
+        cute.ReductionOp.ADD: operator.add,
+        cute.ReductionOp.MAX: cute.arch.fmax,
+    }[op]
+
+    # Step 2: Warp reduction
+    # If threads_per_row < 32, only use that many threads in the reduction
+    warp_width = min(threads_per_row, 32)
+    warp_val = cute.arch.warp_reduction(local_val, warp_op, threads_in_group=warp_width)
+
+    # Determine if we need additional reduction levels
+    warps_per_row = max(threads_per_row // 32, 1)
+
+    # Step 3 & 4: Block or cluster reduction (if needed)
+    if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
+        if cutlass.const_expr(cluster_n == 1):
+            # Single CTA: use block reduction
+            return block_reduce(warp_val, warp_op, reduction_buffer, init_val)
+        else:
+            # Multiple CTAs: use cluster reduction
+            return cluster_reduce(
+                warp_val, warp_op, reduction_buffer, mbar_ptr, cluster_n, init_val
+            )
+    else:
+        # Single warp handles entire row: warp reduction is sufficient
+        return warp_val
+

examples/python/CuTeDSL/cute/blackwell/kernel/rmsnorm/rmsnorm.py

@@ -0,0 +1,840 @@
+# 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 ctypes
+import functools
+import math
+from typing import Optional, Tuple, Union
+
+import cuda.bindings.driver as cuda
+import torch
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.testing as testing
+import cutlass.torch as cutlass_torch
+import cutlass.utils as utils
+from cutlass import Boolean, Float32, Int32, Int64
+from cutlass.cute.runtime import make_ptr
+
+# Support both direct execution and module import
+try:
+    from .reduce import row_reduce
+except ImportError:
+    from reduce import row_reduce
+
+"""
+RMSNorm: Root Mean Square Layer Normalization for Hopper & Blackwell (SM90+)
+====================================================================
+
+A high-performance RMSNorm implementation using CuTe DSL with cluster-based
+reduction for large hidden dimensions.
+
+RMSNorm computes: y = x / sqrt(mean(x²) + eps) * weight
+
+Key Features:
+-------------
+1. CLUSTER SYNCHRONIZATION (SM90+)
+   - Multiple CTAs cooperate to process large N dimensions
+   - Each CTA handles N/cluster_n elements, then reduces across the cluster
+   - Uses mbarrier for efficient cross-CTA synchronization
+
+2. ARCHITECTURE-SPECIFIC TUNING
+   - SM80 (Ampere): Single-CTA execution (cluster_n=1)
+   - SM90 (Hopper): Cluster support enabled for large N
+   - SM100 (Blackwell): Same as SM90
+
+3. VECTORIZED MEMORY ACCESS
+   - 128-bit vectorized loads/stores for optimal memory throughput
+   - TiledCopy abstraction for organized gmem↔smem↔rmem transfers
+
+Cluster Size Selection (FP16):
+------------------------------
+- N <= 16K: cluster_n = 1 (single CTA)
+- N <= 32K: cluster_n = 2
+- N <= 64K: cluster_n = 4
+- N <= 128K: cluster_n = 8
+- Larger: cluster_n = 16
+
+To run this example:
+
+.. code-block:: bash
+
+    python examples/python/CuTeDSL/blackwell/rmsnorm.py --M 2048 --N 4096 --dtype BFloat16
+    python examples/python/CuTeDSL/blackwell/rmsnorm.py --M 2048 --N 4096 --dtype BFloat16 --benchmark
+    python examples/python/CuTeDSL/blackwell/rmsnorm.py --M 2048 --N 32768 --dtype BFloat16 --benchmark
+
+To collect performance with NCU profiler:
+
+.. code-block:: bash
+
+    ncu python examples/python/CuTeDSL/blackwell/rmsnorm.py --M 2048 --N 4096 --dtype BFloat16 --skip_ref_check
+"""
+
+# =============================================================================
+# Architecture Detection
+# =============================================================================
+
+
+@functools.lru_cache(maxsize=16)
+def get_sm_version(device: Optional[Union[int, torch.device, str]] = None) -> int:
+    """Get the SM (compute capability) version of a CUDA device."""
+    if not torch.cuda.is_available():
+        return 80  # Default fallback
+    props = torch.cuda.get_device_properties(device)
+    return props.major * 10 + props.minor
+
+
+def supports_cluster() -> bool:
+    """Check if the current device supports cluster operations (SM90+)."""
+    return get_sm_version() >= 90
+
+
+# =============================================================================
+# Predicate Utility
+# =============================================================================
+
+
+@cute.jit
+def predicate_k(tXcX: cute.Tensor, limit: int) -> cute.Tensor:
+    """Create predicate tensor for bounds checking."""
+    tXpX = cute.make_rmem_tensor(
+        cute.make_layout(
+            (cute.size(tXcX, mode=[0, 1]), cute.size(tXcX, mode=[1]), cute.size(tXcX, mode=[2])),
+            stride=(cute.size(tXcX, mode=[2]), 0, 1),
+        ),
+        Boolean,
+    )
+    for rest_v in cutlass.range_constexpr(tXpX.shape[0]):
+        for rest_k in cutlass.range_constexpr(tXpX.shape[2]):
+            tXpX[rest_v, 0, rest_k] = cute.elem_less(tXcX[(0, rest_v), 0, rest_k][1], limit)
+    return tXpX
+
+
+# =============================================================================
+# RMSNorm Configuration Class
+# =============================================================================
+
+
+class RMSNormConfig:
+    """
+    Configuration for the RMSNorm kernel.
+
+    This class encapsulates all kernel configuration computed once at initialization,
+    following CuTe-DSL conventions from official CUTLASS examples.
+    """
+
+    COPY_BITS = 128  # 128-bit vectorized loads
+
+    def __init__(
+        self,
+        dtype: type[cutlass.Numeric],
+        N: int,
+        has_weight: bool = True,
+        sm_version: int | None = None,
+    ):
+        self.dtype = dtype
+        self.N = N
+        self.has_weight = has_weight
+        self.sm_version = sm_version if sm_version is not None else get_sm_version()
+
+        # Vector size for 128-bit loads
+        self.vec_size = self.COPY_BITS // dtype.width
+
+        # Compute cluster size (SM90+ only)
+        self.cluster_n = self._compute_cluster_n(N, dtype, self.sm_version)
+
+        # N per CTA for cluster case
+        self.N_per_cta = N // self.cluster_n
+
+        # Thread configuration using static methods
+        self.threads_per_row = self._compute_threads_per_row(self.N_per_cta)
+        self.num_threads = self._compute_num_threads(self.N_per_cta)
+
+        # Derived values
+        self.num_vec_blocks = max(
+            1, (self.N_per_cta // self.vec_size + self.threads_per_row - 1) // self.threads_per_row
+        )
+        self.rows_per_block = self.num_threads // self.threads_per_row
+        self.cols_per_tile = self.vec_size * self.num_vec_blocks * self.threads_per_row
+        self.warps_per_row = max(self.threads_per_row // 32, 1)
+
+    @staticmethod
+    def _compute_cluster_n(N: int, dtype: type[cutlass.Numeric], sm_version: int) -> int:
+        """Compute optimal cluster size based on N and architecture."""
+        if sm_version < 90:
+            return 1
+
+        if dtype.width == 16:  # FP16/BF16
+            if N <= 16 * 1024:
+                return 1
+            elif N <= 32 * 1024:
+                return 2
+            elif N <= 64 * 1024:
+                return 4
+            elif N <= 128 * 1024:
+                return 8
+            else:
+                return 16
+        else:  # FP32
+            if N <= 32 * 1024:
+                return 1
+            elif N <= 64 * 1024:
+                return 2
+            elif N <= 128 * 1024:
+                return 4
+            elif N <= 256 * 1024:
+                return 8
+            else:
+                return 16
+
+    @staticmethod
+    def _compute_threads_per_row(N_per_cta: int) -> int:
+        """Compute optimal threads per row based on N per CTA."""
+        if N_per_cta <= 64:
+            return 8
+        elif N_per_cta <= 128:
+            return 16
+        elif N_per_cta <= 3072:
+            return 32
+        elif N_per_cta <= 6144:
+            return 64
+        elif N_per_cta <= 16384:
+            return 128
+        else:
+            return 256
+
+    @staticmethod
+    def _compute_num_threads(N_per_cta: int) -> int:
+        """Compute total threads per block."""
+        return 128 if N_per_cta <= 16384 else 256
+
+    @staticmethod
+    def _make_tv_layout(
+        threads_per_row: int,
+        rows_per_block: int,
+        vec_size: int,
+        num_vec_blocks: int,
+    ) -> tuple:
+        """Create Thread-Value layout for coalesced vectorized memory access."""
+        shape = (
+            (threads_per_row, rows_per_block),
+            (vec_size, num_vec_blocks),
+        )
+        stride = (
+            (vec_size * rows_per_block, 1),
+            (rows_per_block, rows_per_block * vec_size * threads_per_row),
+        )
+        return shape, stride
+
+    def smem_size_in_bytes(self) -> int:
+        """Calculate shared memory requirement in bytes."""
+        tile_bytes = self.rows_per_block * self.cols_per_tile * (self.dtype.width // 8)
+        reduction_bytes = self.rows_per_block * self.warps_per_row * self.cluster_n * 4
+        mbar_bytes = 8 if self.cluster_n > 1 else 0
+        return tile_bytes + reduction_bytes + mbar_bytes
+
+
+# =============================================================================
+# RMSNorm Kernel Class
+# =============================================================================
+
+
+class RMSNormKernel:
+    """
+    RMSNorm kernel with cluster synchronization for large N.
+
+    Features:
+    - Cluster-based reduction for large N (SM90+)
+    - Multiple CTAs cooperate via mbarrier
+    - Single reduction (sum of squares) with cluster-level aggregation
+
+    Example:
+        >>> kernel = RMSNormKernel(cutlass.Float16, N=4096)
+        >>> kernel(x_ptr, w_ptr, o_ptr, M, eps, stream)
+    """
+
+    def __init__(
+        self,
+        dtype: cutlass.Numeric,
+        N: int,
+        has_weight: bool = True,
+        config: RMSNormConfig | None = None,
+    ):
+        # Use provided config or create new one
+        if config is not None:
+            self.cfg = config
+        else:
+            self.cfg = RMSNormConfig(dtype, N, has_weight)
+
+        # Expose key attributes for convenience
+        self.dtype = self.cfg.dtype
+        self.N = self.cfg.N
+        self.has_weight = self.cfg.has_weight
+        self.cluster_n = self.cfg.cluster_n
+
+    @cute.jit
+    def __call__(
+        self,
+        x_ptr: cute.Pointer,
+        w_ptr: cute.Pointer | None,
+        o_ptr: cute.Pointer,
+        M: Int32,
+        eps: Float32,
+        stream: cuda.CUstream,
+    ):
+        """Host function to launch the RMSNorm kernel."""
+        cfg = self.cfg
+
+        # Create CuTe tensors from raw pointers
+        mX = cute.make_tensor(
+            x_ptr,
+            cute.make_layout((M, cfg.N), stride=(cfg.N, 1)),
+        )
+        mO = cute.make_tensor(
+            o_ptr,
+            cute.make_layout((M, cfg.N), stride=(cfg.N, 1)),
+        )
+
+        if cutlass.const_expr(cfg.has_weight and w_ptr is not None):
+            mW = cute.make_tensor(
+                w_ptr,
+                cute.make_layout((cfg.N,), stride=(1,)),
+            )
+        else:
+            mW = None
+
+        # Create TV layout using static helper
+        tv_shape, tv_stride = RMSNormConfig._make_tv_layout(
+            cfg.threads_per_row,
+            cfg.rows_per_block,
+            cfg.vec_size,
+            cfg.num_vec_blocks,
+        )
+        tv_layout = cute.make_layout(tv_shape, stride=tv_stride)
+        tiler_mn = (cfg.rows_per_block, cfg.cols_per_tile)
+
+        self.kernel(mX, mW, mO, eps, tv_layout, tiler_mn).launch(
+            grid=[cute.ceil_div(M, cfg.rows_per_block), cfg.cluster_n, 1],
+            block=[cfg.num_threads, 1, 1],
+            cluster=[1, cfg.cluster_n, 1] if cutlass.const_expr(cfg.cluster_n > 1) else None,
+            smem=cfg.smem_size_in_bytes(),
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mX: cute.Tensor,
+        mW: cute.Tensor | None,
+        mO: cute.Tensor,
+        eps: Float32,
+        tv_layout: cute.Layout,
+        tiler_mn: cute.Shape,
+    ):
+        """Device kernel implementing RMSNorm with cluster support."""
+        cfg = self.cfg
+        tidx, _, _ = cute.arch.thread_idx()
+        bidx, _, _ = cute.arch.block_idx()
+
+        if cutlass.const_expr(cfg.cluster_n > 1):
+            cluster_y = cute.arch.block_idx()[1]
+        else:
+            cluster_y = cutlass.const_expr(0)
+
+        M = mX.shape[0]
+        threads_per_row = tv_layout.shape[0][0]
+        warps_per_row = max(threads_per_row // 32, 1)
+        rows_per_block = tiler_mn[0]
+
+        # =====================================================================
+        # Allocate shared memory
+        # =====================================================================
+        smem = utils.SmemAllocator()
+
+        sX = smem.allocate_tensor(
+            mX.element_type,
+            cute.make_ordered_layout(tiler_mn, order=(1, 0)),
+            byte_alignment=16,
+        )
+
+        if cutlass.const_expr(cfg.cluster_n == 1):
+            reduction_buffer = smem.allocate_tensor(
+                Float32,
+                cute.make_layout((rows_per_block, warps_per_row)),
+                byte_alignment=4,
+            )
+            mbar_ptr = None
+        else:
+            reduction_buffer = smem.allocate_tensor(
+                Float32,
+                cute.make_layout((rows_per_block, (warps_per_row, cfg.cluster_n))),
+                byte_alignment=4,
+            )
+            mbar_ptr = smem.allocate_array(Int64, num_elems=1)
+
+        # =====================================================================
+        # Initialize cluster
+        # =====================================================================
+        if cutlass.const_expr(cfg.cluster_n > 1):
+            if tidx == 0:
+                cute.arch.mbarrier_init(mbar_ptr, 1)
+            cute.arch.mbarrier_init_fence()
+            cute.arch.cluster_arrive_relaxed()
+            cute.arch.cluster_wait()
+
+        # =====================================================================
+        # Create identity tensor and partition
+        # =====================================================================
+        idX = cute.make_identity_tensor(mX.shape)
+
+        gX = cute.local_tile(mX, tiler_mn, (bidx, cluster_y))
+        gO = cute.local_tile(mO, tiler_mn, (bidx, cluster_y))
+        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
+
+        if cutlass.const_expr(cfg.has_weight and mW is not None):
+            mW_expanded_layout = cute.prepend(
+                mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,))
+            )
+            mW_2d = cute.make_tensor(mW.iterator, mW_expanded_layout)
+            gW = cute.local_tile(mW_2d, tiler_mn, (0, cluster_y))
+
+        # =====================================================================
+        # Create TiledCopy operations
+        # =====================================================================
+        copy_atom_load_async = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            mX.element_type,
+            num_bits_per_copy=RMSNormConfig.COPY_BITS,
+        )
+
+        copy_atom_load_W = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mX.element_type,
+            num_bits_per_copy=RMSNormConfig.COPY_BITS,
+        )
+
+        copy_atom_store = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mO.element_type,
+            num_bits_per_copy=RMSNormConfig.COPY_BITS,
+        )
+
+        tiled_copy_load = cute.make_tiled_copy(copy_atom_load_async, tv_layout, tiler_mn)
+        tiled_copy_W = cute.make_tiled_copy(copy_atom_load_W, tv_layout, tiler_mn)
+        tiled_copy_store = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn)
+
+        thr_copy_X = tiled_copy_load.get_slice(tidx)
+        thr_copy_W = tiled_copy_W.get_slice(tidx)
+        thr_copy_O = tiled_copy_store.get_slice(tidx)
+
+        # Partition tensors
+        tXgX = thr_copy_X.partition_S(gX)
+        tXsX = thr_copy_X.partition_D(sX)
+        tXgO = thr_copy_O.partition_D(gO)
+        tXcX = thr_copy_X.partition_S(cX)
+
+        # Register fragments
+        tXrX = cute.make_fragment_like(tXgX)
+        tXrO = cute.make_fragment_like(tXgO)
+
+        if cutlass.const_expr(cfg.has_weight and mW is not None):
+            tWgW = thr_copy_W.partition_S(gW)
+            tWrW = cute.make_fragment_like(tWgW)
+            tXrW = thr_copy_X.retile(tWrW)
+
+        # =====================================================================
+        # Bounds checking
+        # =====================================================================
+        tXpX = predicate_k(tXcX, limit=cfg.N)
+
+        row_coord = tXcX[(0, 0), 0, 0]
+        row_in_bounds = row_coord[0] < M
+
+        # =====================================================================
+        # Async copy global → shared
+        # =====================================================================
+        if row_in_bounds:
+            cute.copy(copy_atom_load_async, tXgX, tXsX, pred=tXpX)
+
+        cute.arch.cp_async_commit_group()
+
+        # Load weight while waiting
+        if cutlass.const_expr(cfg.has_weight and mW is not None):
+            tWpW = predicate_k(thr_copy_W.partition_S(cX), limit=cfg.N)
+            cute.copy(copy_atom_load_W, tWgW, tWrW, pred=tWpW)
+
+        cute.arch.cp_async_wait_group(0)
+
+        # =====================================================================
+        # Pass 1: Compute sum of squares with cluster reduction
+        # =====================================================================
+        cute.autovec_copy(tXsX, tXrX)
+        x = tXrX.load().to(Float32)
+
+        x_sq = x * x
+        sum_sq = row_reduce(
+            x_sq,
+            cute.ReductionOp.ADD,
+            threads_per_row,
+            reduction_buffer,
+            mbar_ptr,
+            cfg.cluster_n,
+            Float32(0.0),
+        )
+
+        # rstd = 1 / sqrt(mean(x²) + eps)
+        mean_sq = sum_sq / cfg.N
+        rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True)
+
+        # Sync after reduction
+        if cutlass.const_expr(cfg.cluster_n > 1):
+            cute.arch.cluster_arrive_relaxed()
+            cute.arch.cluster_wait()
+        else:
+            cute.arch.barrier()
+
+        # =====================================================================
+        # Pass 2: Normalize and output
+        # =====================================================================
+        cute.autovec_copy(tXsX, tXrX)
+        x = tXrX.load().to(Float32)
+
+        y = x * rstd
+
+        # Apply weight if present
+        if cutlass.const_expr(cfg.has_weight and mW is not None):
+            w = tXrW.load().to(Float32)
+            y = y * w
+
+        # Store to global memory
+        tXrO.store(y.to(cfg.dtype))
+
+        if row_in_bounds:
+            cute.copy(copy_atom_store, tXrO, tXgO, pred=tXpX)
+
+
+# =============================================================================
+# Kernel Compilation and Caching
+# =============================================================================
+
+# Mapping from torch dtype to cutlass dtype
+_torch_to_cutlass_dtype = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+}
+
+# Cache for compiled kernels
+_compile_cache: dict = {}
+
+
+def get_compiled_kernel(
+    dtype: type[cutlass.Numeric],
+    N: int,
+    has_weight: bool,
+    stream: cuda.CUstream,
+):
+    """
+    Get or compile the RMSNorm kernel for the given configuration.
+
+    Uses compilation cache to avoid recompiling for same (dtype, N, has_weight) tuples.
+
+    :param dtype: Data type (Float16, BFloat16, Float32)
+    :type dtype: type[cutlass.Numeric]
+    :param N: Hidden dimension size
+    :type N: int
+    :param has_weight: Whether weight is applied
+    :type has_weight: bool
+    :param stream: CUDA stream
+    :type stream: cuda.CUstream
+    :return: Compiled kernel function
+    """
+    key = (dtype, N, has_weight)
+    if key not in _compile_cache:
+        kernel_obj = RMSNormKernel(dtype, N, has_weight)
+
+        # Compile with representative arguments
+        compiled_kernel = cute.compile(
+            kernel_obj,
+            make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=16),  # x_ptr
+            make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
+            if has_weight
+            else None,  # w_ptr
+            make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=16),  # o_ptr
+            Int32(1),  # M (dummy)
+            Float32(1e-6),  # eps (dummy)
+            stream,
+        )
+
+        _compile_cache[key] = compiled_kernel
+
+    return _compile_cache[key]
+
+
+# =============================================================================
+# Tensor Creation Utilities
+# =============================================================================
+
+
+def create_tensors(
+    M: int,
+    N: int,
+    dtype: type[cutlass.Numeric],
+    has_weight: bool,
+) -> Tuple:
+    """Create input, weight, and output tensors for RMSNorm."""
+    torch.manual_seed(42)
+    torch_dtype = cutlass_torch.dtype(dtype)
+
+    x = torch.randn(M, N, device="cuda", dtype=torch_dtype)
+    weight = torch.randn(N, device="cuda", dtype=torch_dtype) if has_weight else None
+    out = torch.empty_like(x)
+
+    return x, weight, out
+
+
+def rmsnorm_ref(
+    x: torch.Tensor,
+    weight: torch.Tensor | None = None,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Reference RMSNorm implementation in PyTorch."""
+    x_f32 = x.float()
+    rms = torch.sqrt(torch.mean(x_f32**2, dim=-1, keepdim=True) + eps)
+    x_norm = x_f32 / rms
+    if weight is not None:
+        x_norm = x_norm * weight.float()
+    return x_norm.to(x.dtype)
+
+
+# =============================================================================
+# Run Function
+# =============================================================================
+
+
+def run(
+    M: int,
+    N: int,
+    dtype: type[cutlass.Numeric],
+    has_weight: bool = True,
+    eps: float = 1e-6,
+    tolerance: float = 1e-2,
+    warmup_iterations: int = 2,
+    iterations: int = 100,
+    skip_ref_check: bool = False,
+    benchmark: bool = False,
+) -> float:
+    """
+    Execute RMSNorm and optionally benchmark performance.
+
+    :param M: Number of rows (batch size * sequence length)
+    :type M: int
+    :param N: Hidden dimension size
+    :type N: int
+    :param dtype: Data type (Float16, BFloat16, Float32)
+    :type dtype: type[cutlass.Numeric]
+    :param has_weight: Whether to apply learnable weight
+    :type has_weight: bool
+    :param eps: Epsilon for numerical stability
+    :type eps: float
+    :param tolerance: Tolerance for correctness check
+    :type tolerance: float
+    :param warmup_iterations: Warmup iterations for benchmarking
+    :type warmup_iterations: int
+    :param iterations: Number of benchmark iterations
+    :type iterations: int
+    :param skip_ref_check: Skip reference correctness check
+    :type skip_ref_check: bool
+    :param benchmark: Enable benchmarking
+    :type benchmark: bool
+    :return: Execution time in microseconds (if benchmark=True, else 0)
+    :rtype: float
+    """
+    print("Running RMSNorm test with:")
+    print(f"  M: {M}, N: {N}")
+    print(f"  dtype: {dtype}")
+    print(f"  has_weight: {has_weight}")
+    print(f"  eps: {eps}")
+    print(f"  SM version: {get_sm_version()}")
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA GPU is required to run this example!")
+
+    # Get CUDA stream
+    torch_stream = torch.cuda.current_stream()
+    stream = cuda.CUstream(torch_stream.cuda_stream)
+
+    # Create tensors
+    x, weight, out = create_tensors(M, N, dtype, has_weight)
+
+    # Get configuration info
+    config = RMSNormConfig(dtype, N, has_weight)
+    print(f"  cluster_n: {config.cluster_n}")
+    print(f"  threads_per_row: {config.threads_per_row}")
+    print(f"  rows_per_block: {config.rows_per_block}")
+
+    # Get compiled kernel
+    compiled_kernel = get_compiled_kernel(dtype, N, has_weight, stream)
+
+    # Create pointers for kernel call
+    x_ptr = make_ptr(dtype, x.data_ptr())
+    w_ptr = make_ptr(dtype, weight.data_ptr()) if weight is not None else None
+    out_ptr = make_ptr(dtype, out.data_ptr())
+
+    # Run kernel and verify
+    if not skip_ref_check:
+        compiled_kernel(x_ptr, w_ptr, out_ptr, Int32(M), Float32(eps), stream)
+        torch.cuda.synchronize()
+
+        ref = rmsnorm_ref(x, weight, eps)
+        torch.testing.assert_close(out, ref, atol=tolerance, rtol=tolerance)
+        print("Correctness check passed!")
+
+    if not benchmark:
+        return 0.0
+
+    # Benchmark
+    print(f"\nBenchmarking with {warmup_iterations} warmup, {iterations} iterations...")
+
+    def generate_tensors():
+        x, weight, out = create_tensors(M, N, dtype, has_weight)
+        x_ptr = make_ptr(dtype, x.data_ptr())
+        w_ptr = make_ptr(dtype, weight.data_ptr()) if weight is not None else None
+        out_ptr = make_ptr(dtype, out.data_ptr())
+        return testing.JitArguments(x_ptr, w_ptr, out_ptr, Int32(M), Float32(eps), stream)
+
+    exec_time_us = testing.benchmark(
+        compiled_kernel,
+        workspace_generator=generate_tensors,
+        workspace_count=10,
+        warmup_iterations=warmup_iterations,
+        iterations=iterations,
+        stream=stream,
+    )
+
+    # Calculate throughput
+    torch_dtype = cutlass_torch.dtype(dtype)
+    bytes_per_elem = torch.tensor([], dtype=torch_dtype).element_size()
+    total_bytes = M * N * bytes_per_elem * 2  # read x, write out
+    if has_weight:
+        total_bytes += N * bytes_per_elem  # read weight (amortized across M)
+
+    throughput_gbps = (total_bytes / (exec_time_us / 1e6)) / 1e9
+
+    print(f"Kernel execution time: {exec_time_us:.2f} us")
+    print(f"Memory throughput: {throughput_gbps:.2f} GB/s")
+
+    return exec_time_us
+
+
+# =============================================================================
+# Main Entry Point
+# =============================================================================
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="RMSNorm kernel example for Blackwell (SM100)"
+    )
+
+    parser.add_argument("--M", type=int, default=2048, help="Number of rows")
+    parser.add_argument("--N", type=int, default=4096, help="Hidden dimension size")
+    parser.add_argument(
+        "--dtype",
+        type=cutlass.dtype,
+        default=cutlass.BFloat16,
+        help="Data type (Float16, BFloat16, Float32)",
+    )
+    parser.add_argument(
+        "--has_weight",
+        action="store_true",
+        default=True,
+        help="Apply learnable weight",
+    )
+    parser.add_argument(
+        "--no_weight",
+        action="store_true",
+        help="Disable learnable weight",
+    )
+    parser.add_argument(
+        "--eps",
+        type=float,
+        default=1e-6,
+        help="Epsilon for numerical stability",
+    )
+    parser.add_argument(
+        "--tolerance",
+        type=float,
+        default=1e-2,
+        help="Tolerance for correctness check",
+    )
+    parser.add_argument(
+        "--warmup_iterations",
+        type=int,
+        default=2,
+        help="Warmup iterations for benchmarking",
+    )
+    parser.add_argument(
+        "--iterations",
+        type=int,
+        default=100,
+        help="Number of benchmark iterations",
+    )
+    parser.add_argument(
+        "--skip_ref_check",
+        action="store_true",
+        help="Skip reference correctness check",
+    )
+    parser.add_argument(
+        "--benchmark",
+        action="store_true",
+        help="Enable benchmarking",
+    )
+
+    args = parser.parse_args()
+
+    # Handle weight flag
+    has_weight = args.has_weight and not args.no_weight
+
+    run(
+        M=args.M,
+        N=args.N,
+        dtype=args.dtype,
+        has_weight=has_weight,
+        eps=args.eps,
+        tolerance=args.tolerance,
+        warmup_iterations=args.warmup_iterations,
+        iterations=args.iterations,
+        skip_ref_check=args.skip_ref_check,
+        benchmark=args.benchmark,
+    )
+
+    print("PASS")
+