chore: import upstream snapshot with attribution

python/tvm/backend/cuda/lang/__init__.py

@@ -0,0 +1,70 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA-specific TIRx language helpers."""
+
+from importlib import import_module
+
+__all__ = [
+    "BaseTileScheduler",
+    "ClusterPersistentScheduler2D",
+    "FlashAttentionLPTScheduler",
+    "FlashAttentionLinearScheduler",
+    "GroupMajor3D",
+    "IndexedTripleTileScheduler",
+    "MBarrier",
+    "Pipeline",
+    "PipelineState",
+    "RankAwareGroupMajorTileScheduler",
+    "SMEMPool",
+    "SmemDescriptor",
+    "TCGen05Bar",
+    "TMABar",
+    "TMEMPool",
+    "TMEMStages",
+    "WarpRole",
+    "WarpgroupRole",
+]
+
+_HELPER_MODULES = {
+    "MBarrier": ".pipeline",
+    "Pipeline": ".pipeline",
+    "PipelineState": ".pipeline",
+    "BaseTileScheduler": ".tile_scheduler",
+    "ClusterPersistentScheduler2D": ".tile_scheduler",
+    "FlashAttentionLPTScheduler": ".tile_scheduler",
+    "FlashAttentionLinearScheduler": ".tile_scheduler",
+    "GroupMajor3D": ".tile_scheduler",
+    "IndexedTripleTileScheduler": ".tile_scheduler",
+    "RankAwareGroupMajorTileScheduler": ".tile_scheduler",
+    "SMEMPool": ".alloc_pool",
+    "TCGen05Bar": ".pipeline",
+    "TMABar": ".pipeline",
+    "TMEMPool": ".alloc_pool",
+    "TMEMStages": ".alloc_pool",
+    "WarpRole": ".warp_role",
+    "WarpgroupRole": ".warp_role",
+    "SmemDescriptor": ".smem_desc",
+}
+
+
+def __getattr__(name: str):
+    module_name = _HELPER_MODULES.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(module_name, __name__), name)
+    globals()[name] = value
+    return value

python/tvm/backend/cuda/lang/alloc_pool.py

@@ -0,0 +1,529 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""SMEM and TMEM bump-allocator pools for TIRX kernels."""
+
+from __future__ import annotations
+
+import functools
+import operator
+
+from tvm import DataType
+from tvm.tirx.layout import S, TCol, TileLayout, TLane
+
+# ---------------------------------------------------------------------------
+# ir_builder helpers — imported lazily to avoid circular deps at module level
+# ---------------------------------------------------------------------------
+
+_ir = None
+
+
+def _get_ir():
+    global _ir
+    if _ir is None:
+        from tvm.tirx.script.builder import ir as _mod
+
+        _ir = _mod
+    return _ir
+
+
+def _get_frame():
+    from tvm.tirx.script.builder import frame
+
+    return frame
+
+
+# ---------------------------------------------------------------------------
+# Shared utilities
+# ---------------------------------------------------------------------------
+
+_POOL_UNSET = object()
+
+
+def _default_tmem_layout(rows, cols):
+    return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)])
+
+
+def _emit_stmt(expr):
+    ir = _get_ir()
+    ir.add_to_parent(ir.evaluate(expr))
+
+
+def _shape_product(shape):
+    return functools.reduce(operator.mul, shape, 1)
+
+
+def _auto_swizzle_mode(dtype):
+    """Select the default MMA swizzle mode for a shared-memory allocation."""
+    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+    del dtype
+    return SwizzleMode.SWIZZLE_128B_ATOM
+
+
+def _swizzle_atom_bytes(swizzle_mode):
+    """Return the row width (in bytes) of one swizzle atom for *swizzle_mode*."""
+    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+    return {
+        SwizzleMode.SWIZZLE_NONE: 0,
+        SwizzleMode.SWIZZLE_32B_ATOM: 32,
+        SwizzleMode.SWIZZLE_64B_ATOM: 64,
+        SwizzleMode.SWIZZLE_128B_ATOM: 128,
+    }[swizzle_mode]
+
+
+def _suggest_swizzle_for_row_bytes(row_bytes):
+    """Pick the largest valid swizzle mode whose atom row fits within *row_bytes*."""
+
+    for atom_bytes, mode in (
+        (128, "SWIZZLE_128B_ATOM"),
+        (64, "SWIZZLE_64B_ATOM"),
+        (32, "SWIZZLE_32B_ATOM"),
+    ):
+        if row_bytes >= atom_bytes and row_bytes % atom_bytes == 0:
+            return mode
+    return "SWIZZLE_NONE"
+
+
+def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
+    """Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible.
+
+    ``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]``
+    over the last two logical dimensions of *shape*. If the row width or row count of
+    the request is smaller than (or not a multiple of) the atom, the underlying
+    ``Layout.tile_to`` lowers to a ``floordiv``/``floormod`` by zero and raises an
+    opaque internal "Divide by zero" diagnostic from ``tile_tile_ops.cc``. Catch the
+    misconfiguration here so callers see *what* is wrong and *how* to fix it.
+
+    Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
+    """
+    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+    if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
+        return
+
+    if len(shape) < 2:
+        raise ValueError(
+            f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; "
+            f"swizzled MMA layouts tile over the last two dims (rows, cols). "
+            f"Use swizzle_mode='none' for 1-D allocations."
+        )
+
+    # Only validate concrete int dims; symbolic dims fall through (the analyzer
+    # in C++ will still ICHECK on them, but at least we don't false-positive).
+    rows = shape[-2]
+    cols = shape[-1]
+    if not (isinstance(rows, int) and isinstance(cols, int)):
+        return
+
+    dtype_bytes = DataType(dtype).bits // 8
+    if dtype_bytes == 0:
+        # Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so
+        # use a fractional check expressed via bits.
+        col_bits = cols * DataType(dtype).bits
+        atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8
+        if col_bits < atom_bits or col_bits % atom_bits != 0:
+            row_bytes = col_bits // 8 if col_bits % 8 == 0 else col_bits / 8
+            atom_bytes = _swizzle_atom_bytes(swizzle_mode)
+            suggestion = _suggest_swizzle_for_row_bytes(col_bits // 8 if col_bits >= 8 else 0)
+            raise ValueError(
+                f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
+                f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
+                f"swizzle atom selected by {swizzle_mode.name}. "
+                f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
+                f"to a multiple of "
+                f"{(atom_bits + DataType(dtype).bits - 1) // DataType(dtype).bits} elements."
+            )
+    else:
+        row_bytes = cols * dtype_bytes
+        atom_bytes = _swizzle_atom_bytes(swizzle_mode)
+        if row_bytes < atom_bytes or row_bytes % atom_bytes != 0:
+            suggestion = _suggest_swizzle_for_row_bytes(row_bytes)
+            min_cols = atom_bytes // dtype_bytes
+            raise ValueError(
+                f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
+                f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
+                f"swizzle atom selected by {swizzle_mode.name}. "
+                f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
+                f"to a multiple of {min_cols} elements (>= {atom_bytes}B at {dtype})."
+            )
+
+    # Atom rows is always 8 (see ``mma_atom_shape`` in tma_utils.py).
+    atom_rows = 8
+    if rows < atom_rows or rows % atom_rows != 0:
+        raise ValueError(
+            f"alloc_mma shape={tuple(shape)} has shape[-2]={rows}, but the "
+            f"{swizzle_mode.name} atom requires shape[-2] to be a positive "
+            f"multiple of {atom_rows}. Use swizzle_mode='none', or widen shape[-2] "
+            f"to a multiple of {atom_rows}."
+        )
+
+
+# ---------------------------------------------------------------------------
+# TMEMStages
+# ---------------------------------------------------------------------------
+
+
+def _meta_class(cls):
+    """Apply @meta_class decorator from ir_builder."""
+    return _get_ir().meta_class(cls)
+
+
+@_meta_class
+class TMEMStages:
+    """Parse-time staged view over a TMEM buffer.
+
+    Parameters
+    ----------
+    buf : Buffer
+        The underlying TMEM buffer (e.g. f32 or f16 view).
+    col_start : int
+        First column of stage 0 in *buf*'s column space.
+    width : int
+        Number of columns per stage.
+    stages : int
+        Number of pipeline stages (default 1).
+    stride : int or None
+        Column distance between consecutive stages.  When *None* (default),
+        equals *width* (stages are packed back-to-back).
+    """
+
+    def __init__(self, buf, col_start, width, stages=1, stride=None):
+        self.buf = buf
+        self.col_start = col_start
+        self.width = width
+        self.stages = stages
+        self.stride = width if stride is None else stride
+
+    def _stage_base(self, stage):
+        return self.col_start + stage * self.stride
+
+    def __getitem__(self, item):
+        if isinstance(item, tuple):
+            assert len(item) == 2, "TMEMStages expects region[stage] or region[stage, start:stop]"
+            stage, col_slice = item
+            assert isinstance(col_slice, slice), "TMEMStages tuple indexing requires a slice"
+            base = self._stage_base(stage)
+            start = 0 if col_slice.start is None else col_slice.start
+            stop = self.width if col_slice.stop is None else col_slice.stop
+            return self.buf[:, base + start : base + stop : col_slice.step]
+        base = self._stage_base(item)
+        return self.buf[:, base : base + self.width]
+
+
+# ---------------------------------------------------------------------------
+# TMEMPool
+# ---------------------------------------------------------------------------
+
+
+@_meta_class
+class TMEMPool:
+    """Bump allocator over TMEM columns."""
+
+    def __init__(
+        self,
+        pool,
+        total_cols=512,
+        *,
+        cta_group=1,
+        alloc_warp=0,
+        dealloc_warp=None,
+        tmem_addr=None,
+        sync_after_alloc=True,
+    ):
+        # tcgen05 alloc/dealloc are warp-uniform PTX instructions: every lane
+        # in the chosen warp must participate, and exactly one warp in the
+        # CTA must execute them. The pool emits its own
+        # ``if warp_id() == target_warp: tcgen05.alloc(...)``
+        # guard, using the cta->warp scope id ``T.warp_id()``.
+        # NOTE: synccheck currently false-deadlocks on kernels that declare a
+        # second warp-scope id (cpusim binds only one warp var); the generated
+        # CUDA is equivalent to ``thread_rank() // 32 == target_warp``.
+        self.pool = pool
+        self.total_cols = total_cols
+        self.cta_group = cta_group
+        self.alloc_warp = alloc_warp
+        self.dealloc_warp = alloc_warp if dealloc_warp is None else dealloc_warp
+        self.sync_after_alloc = sync_after_alloc
+        self.offset = 0
+        self.max_offset = 0
+        self._committed = False
+        self._deallocated = False
+        self._addr_buf = pool.alloc([1], "uint32", align=4) if tmem_addr is None else tmem_addr
+
+    def _addr_slot(self):
+        try:
+            return self._addr_buf[0]
+        except TypeError:
+            return self._addr_buf
+
+    @property
+    def addr(self):
+        return self._addr_slot()
+
+    def _emit_warp_guard(self, target_warp, emit):
+        from tvm.script import tirx as T
+
+        warp_id = T.warp_id()
+        with T.If(warp_id == target_warp):
+            with T.Then():
+                emit()
+
+    def _resolve_cols(self, shape, dtype, cols, layout=None):
+        if cols is not None:
+            return cols
+        bits = DataType(dtype).bits
+        if layout is not None:
+            # span("TCol") is in *element* (buffer dtype) units; one TMEM cell
+            # holds 32 bits regardless of the element type.
+            tcol_elems = int(layout.span("TCol"))
+            tcol_bits = tcol_elems * bits
+            assert tcol_bits % 32 == 0, (
+                f"layout TCol span={tcol_elems} elems x {bits}b is not 32-bit aligned"
+            )
+            return tcol_bits // 32
+        assert len(shape) == 2, "TMEMPool.alloc() requires cols= for non-2D TMEM buffers"
+        total_bits = _shape_product(shape) * bits
+        rows = shape[0]
+        assert total_bits % (32 * rows) == 0, (
+            f"Cannot infer TMEM columns from shape={shape}, dtype={dtype!r}; "
+            "please pass cols= explicitly"
+        )
+        return total_bits // (32 * rows)
+
+    def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None):
+        """Allocate a TMEM buffer.
+
+        Parameters
+        ----------
+        shape, dtype, cols
+            Standard buffer shape / dtype / column count.
+        layout
+            Explicit ``TileLayout``. Mutually exclusive with ``datapath``.
+        datapath : str | None
+            Optional tcgen05 datapath letter (``"D"`` for M=128 full datapath,
+            ``"F"`` for M=64 non-``.ws`` scattered). When provided, the buffer's
+            layout is derived from ``tmem_datapath_layout(datapath, *shape)``
+            so the row index reflects the *physical* TMEM lane occupation
+            (PTX ISA §9.7.16.10.5). The downstream ``.16x*b`` / ``.32x32b``
+            dispatches structurally check this layout to catch mismatched
+            atoms (e.g. a ``.16x*b`` M=128 read against a Layout F buffer).
+            Defaults to ``None``, which means Layout D's identity row→lane
+            mapping — keep this for shape ``(128, X)`` buffers that hold
+            an M=128 MMA accumulator.
+        """
+        from tvm.tirx.layout import tmem_datapath_layout
+
+        if layout is not None and datapath is not None:
+            raise ValueError("TMEMPool.alloc: pass at most one of layout= and datapath=")
+        if datapath is not None:
+            assert len(shape) == 2, "TMEMPool.alloc: datapath= requires a 2-D shape"
+            layout = tmem_datapath_layout(datapath, shape[0], shape[1])
+
+        ir = _get_ir()
+        cols = self._resolve_cols(shape, dtype, cols, layout)
+        col_start = self.offset
+        col_end = col_start + cols
+        assert col_end <= self.total_cols, f"TMEM overflow: {col_end} > {self.total_cols}"
+        if layout is None:
+            assert len(shape) == 2, "TMEMPool.alloc() requires layout= for non-2D TMEM buffers"
+            layout = _default_tmem_layout(shape[0], shape[1])
+        res = ir.decl_buffer(shape, dtype, scope="tmem", allocated_addr=col_start, layout=layout)
+        self.offset = col_end
+        self.max_offset = max(self.max_offset, self.offset)
+        return res
+
+    def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1):
+        """Allocate a tcgen05 block-scaled SF TMEM buffer with an inferred layout.
+
+        ``shape`` last two dims are ``(rows, SF_K * sf_reuse)`` (the last dim is
+        what gemm dispatch iterates over). When ``shape`` has 3 dims, the first
+        is treated as a pipe-depth outer.
+        """
+        from tvm.backend.cuda.operator.tile_primitive.gemm_async.tcgen05 import sf_tmem_layout
+
+        if len(shape) == 2:
+            pipe_depth, rows, last = None, shape[0], shape[1]
+        elif len(shape) == 3:
+            pipe_depth, rows, last = shape[0], shape[1], shape[2]
+        else:
+            raise ValueError(
+                f"alloc_sf expects 2D (rows, SF_K*sf_reuse) or 3D "
+                f"(pipe_depth, rows, SF_K*sf_reuse); got shape={shape}"
+            )
+        assert last % sf_reuse == 0, (
+            f"alloc_sf: shape last dim {last} must be divisible by sf_reuse={sf_reuse}"
+        )
+        SF_K = last // sf_reuse
+        layout = sf_tmem_layout(
+            rows=rows, SF_K=SF_K, sf_per_mma=sf_per_mma, sf_reuse=sf_reuse, pipe_depth=pipe_depth
+        )
+        return self.alloc(shape, dtype, layout=layout)
+
+    def move_base_to(self, col):
+        self.offset = col
+        self.max_offset = max(self.max_offset, self.offset)
+
+    def commit(self):
+        assert not self._committed, "TMEMPool.commit() can only be called once"
+        from tvm.script import tirx as T
+
+        def emit_alloc():
+            _emit_stmt(
+                T.ptx.tcgen05.alloc(
+                    T.address_of(self.addr), n_cols=self.total_cols, cta_group=self.cta_group
+                )
+            )
+            if self.sync_after_alloc:
+                _emit_stmt(T.cuda.warp_sync())
+
+        self._emit_warp_guard(self.alloc_warp, emit_alloc)
+        self._committed = True
+
+    def dealloc(self):
+        assert self._committed, "TMEMPool.dealloc() called before commit()"
+        assert not self._deallocated, "TMEMPool.dealloc() can only be called once"
+        self._deallocated = True
+        from tvm.script import tirx as T
+
+        def emit_dealloc():
+            _emit_stmt(T.ptx.tcgen05.relinquish_alloc_permit(cta_group=self.cta_group))
+            _emit_stmt(
+                T.ptx.tcgen05.dealloc(self.addr, n_cols=self.total_cols, cta_group=self.cta_group)
+            )
+
+        self._emit_warp_guard(self.dealloc_warp, emit_dealloc)
+
+
+# ---------------------------------------------------------------------------
+# SMEMPool
+# ---------------------------------------------------------------------------
+
+
+@_meta_class
+class SMEMPool:
+    """Bump allocator over a contiguous shared memory region.
+
+    Parameters
+    ----------
+    ptr : Var or None, optional
+        If omitted, an ``alloc_buffer([0], "uint8", scope="shared.dyn")`` is
+        created automatically and ``commit()`` must be called after all
+        allocations to emit the size annotation.
+        If a ``Var`` is provided, the caller manages the backing buffer and
+        ``commit()`` is a no-op.
+    """
+
+    def __init__(self, ptr=_POOL_UNSET):
+        ir = _get_ir()
+        if ptr is _POOL_UNSET:
+            self.buf = ir.alloc_buffer([0], "uint8", scope="shared.dyn")
+            self.ptr = self.buf.data
+            self._owns_buffer = True
+        else:
+            self.buf = None
+            self.ptr = ptr
+            self._owns_buffer = False
+        self.offset = 0
+        self.max_offset = 0
+
+    def alloc(
+        self,
+        shape,
+        dtype="float32",
+        strides=None,
+        scope="shared.dyn",
+        align=0,
+        buffer_type="",
+        axis_separators=None,
+        layout="default",
+    ):
+        ir = _get_ir()
+        if align > 0:
+            self.offset = (self.offset + align - 1) // align * align
+        res = ir.decl_buffer(
+            shape,
+            dtype,
+            data=self.ptr,
+            strides=strides,
+            byte_offset=self.offset,
+            scope=scope,
+            align=align,
+            buffer_type=buffer_type,
+            axis_separators=axis_separators,
+            layout=layout,
+        )
+        # Advance in bits then round up to bytes so sub-byte dtypes (e.g.
+        # float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it
+        # at 0 (bits // 8) and silently overlapping the next allocation.
+        self.offset += (_shape_product(shape) * DataType(dtype).bits + 7) // 8
+        if self._owns_buffer:
+            self.max_offset = max(self.max_offset, self.offset)
+        return res
+
+    def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024):
+        """Allocate MMA-compatible shared memory with an inferred swizzle layout."""
+        from tvm.backend.cuda.operator.tile_primitive.tma_utils import (
+            SwizzleMode,
+            mma_shared_layout,
+        )
+
+        if isinstance(swizzle_mode, str):
+            if swizzle_mode == "auto":
+                swizzle_mode = _auto_swizzle_mode(dtype)
+            elif swizzle_mode == "none":
+                swizzle_mode = SwizzleMode.SWIZZLE_NONE
+            else:
+                raise ValueError(
+                    f"Unsupported swizzle_mode={swizzle_mode!r}; expected 'auto', 'none', "
+                    "or SwizzleMode"
+                )
+        _validate_mma_alloc_shape(shape, dtype, swizzle_mode)
+        layout = mma_shared_layout(dtype, swizzle_mode, shape)
+        return self.alloc(shape, dtype, align=align, layout=layout)
+
+    def move_base_to(self, offset):
+        self.offset = offset
+        if self._owns_buffer:
+            self.max_offset = max(self.max_offset, self.offset)
+
+    def commit(self, size=None):
+        """Emit pool size annotation into the IR.
+
+        Must be called after all ``alloc()`` / ``move_base_to()`` calls.
+
+        Parameters
+        ----------
+        size : int, optional
+            Explicit shared memory size in bytes.  When *None* (the default),
+            the high-water mark ``max_offset`` tracked by the allocator is used.
+        """
+        if not self._owns_buffer:
+            return
+        ir = _get_ir()
+        frame_mod = _get_frame()
+        resolved = size if size is not None else self.max_offset
+        assert resolved >= self.max_offset, (
+            f"Specified smem size ({resolved}) is smaller than "
+            f"the pool high-water mark ({self.max_offset})"
+        )
+        attr_frame = ir.attr(self.ptr, "tirx.pool_max_bytes", resolved)
+        if isinstance(attr_frame, frame_mod.AttrFrame):
+            from functools import partial
+
+            attr_frame.add_callback(partial(attr_frame.__exit__, None, None, None))
+            attr_frame.__enter__()

python/tvm/backend/cuda/lang/pipeline.py

@@ -0,0 +1,251 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Reusable pipeline state and mbarrier helpers for SM100 kernels.
+
+These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
+instances are automatically treated as meta values inside @T.prim_func.
+"""
+
+from tvm.script import tirx as T
+
+
+@T.meta_class
+class PipelineState:
+    """Tracks stage and phase for a software-pipelined ring buffer.
+
+    This class does not know anything about full/empty barriers. Use it when
+    the kernel manually waits/signals barriers, or when the stage/phase drives
+    a ring not wrapped in a ``Pipeline``.
+
+    Parameters
+    ----------
+    depth : int
+        Number of stages in the ring.
+    phase : int, optional
+        Initial phase. Omit when initialization should happen later.
+    """
+
+    def __init__(self, depth: int, phase=None):
+        self.stage = T.local_scalar("int32")
+        self.phase = T.local_scalar("int32")
+        self.depth = depth
+        if phase is not None:
+            self.init(phase)
+
+    @T.inline
+    def init(self, phase):
+        self.stage = 0
+        self.phase = phase
+
+    @T.inline
+    def advance(self):
+        if self.depth > 1:
+            self.stage = self.stage + 1
+            if self.stage == self.depth:
+                self.stage = 0
+                self.phase = self.phase ^ 1
+        else:
+            self.phase = self.phase ^ 1
+
+
+@T.meta_class
+class MBarrier:
+    """Mbarrier wrapper with regular ``mbarrier.arrive``.
+
+    Parameters
+    ----------
+    pool : SMEMPool
+        Shared memory pool allocator.
+    depth : int
+        Number of barrier slots (one per pipeline stage).
+    phase_offset : int
+        XORed into the phase bit on every ``wait`` / ``arrive``.
+    leader : Expr, optional
+        Boolean predicate selecting the single thread that runs
+        ``mbarrier.init``. Defaults to ``T.cuda.thread_rank() == 0`` --
+        thread 0 of the enclosing CTA, which always picks exactly one
+        thread regardless of which scope_id vars the caller declared.
+        Override only when you want a different CTA-local thread to do
+        the init.
+
+        Note: the default deliberately avoids ``T.warp_id()`` /
+        ``T.lane_id()``. Those introduce deferred ``cta->warp`` /
+        ``warp->thread`` ScopeIdDefs that the verifier cannot pin down
+        unless the kernel header declares the full warp/lane chain (e.g. a
+        single-CTA DSMEM kernel that only declares ``thread_id``). It also
+        avoids the synccheck false-deadlock on kernels that declare a
+        second warp-scope id. The generated CUDA is equivalent.
+    """
+
+    def __init__(self, pool, depth, phase_offset=0, leader=None):
+        self.buf = pool.alloc((depth,), "uint64", align=8)
+        self.depth = depth
+        self.phase_offset = phase_offset
+        self.leader = leader if leader is not None else (T.cuda.thread_rank() == 0)
+
+    @T.inline
+    def init(self, count):
+        if self.leader:
+            for i in T.unroll(self.depth):
+                T.ptx.mbarrier.init(self.buf.ptr_to([i]), count)
+
+    @T.inline
+    def wait(self, stage, phase):
+        # Blocks: ``mbarrier.try_wait`` loops internally until the phase flips,
+        # so this returns only once the barrier has completed.
+        T.ptx.mbarrier.try_wait(self.buf.ptr_to([stage]), phase ^ self.phase_offset)
+
+    @T.inline
+    def arrive(self, stage, cta_id=None, pred=None, count=None):
+        # Default: local-CTA arrive — emits the simple
+        # ``mbarrier.arrive.shared.b64`` form. To arrive on a remote
+        # CTA's mbarrier in a cluster kernel, callers must pass
+        # ``cta_id=`` explicitly (e.g. ``bar.arrive(stage, cta_id=0)``)
+        # or use ``MBarrier.remote_view(rank).arrive(stage)``. Defaulting
+        # the cross-CTA path was both surprising (``bar.arrive(stage)``
+        # silently ``mapa`` ed across the cluster) and a per-call cost
+        # of ~3 PTX ops on every single-CTA kernel.
+        #
+        # ``count`` (cross-CTA path only) emits the explicit arrival-count
+        # operand, i.e. ``mbarrier.arrive.shared::cluster.b64 _, [addr], count``.
+        # When ``None`` the implicit count-of-1 form is emitted. Passing
+        # ``count=1`` is semantically identical but spells the count explicitly.
+        if cta_id is None:
+            T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
+        else:
+            actual_pred = True if pred is None else pred
+            T.ptx.mbarrier.arrive(
+                self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred, count=count
+            )
+
+    def ptr_to(self, idx):
+        return self.buf.ptr_to(idx)
+
+    def remote_view(self, rank):
+        """Create a view of this barrier mapped to another CTA's shared memory.
+
+        Arrive-only: the returned view is built with ``object.__new__`` and
+        never copies ``self.leader``, so calling ``.init()`` on it would fail.
+        Use it solely to ``arrive`` on a remote CTA's mbarrier.
+        """
+        from tvm.ir import PointerType, PrimType
+        from tvm.tirx import Var as TIRVar
+
+        expr = T.reinterpret("handle", T.ptx.map_shared_rank(self.buf.ptr_to([0]), rank))
+        ptr = TIRVar("remote_mbar_ptr", PointerType(PrimType("uint64")))
+        T.Bind(expr, var=ptr)
+        buf = T.decl_buffer([self.depth], "uint64", data=ptr, scope="shared")
+        remote = object.__new__(type(self))
+        remote.buf = buf
+        remote.depth = self.depth
+        remote.phase_offset = self.phase_offset
+        return remote
+
+
+class TMABar(MBarrier):
+    """Barrier signaled by TMA (mbarrier.arrive.expect_tx).
+
+    When ``tx_count`` is None, falls back to a remote mbarrier.arrive
+    (matching MBarrier.arrive defaults).
+    """
+
+    @T.inline
+    def arrive(self, stage, tx_count=None, cta_id=None, pred=None):
+        # NOTE: this arrive() kwarg set intentionally differs from
+        # MBarrier.arrive (hardware necessity, LSP-incompatible by design).
+        # ``tx_count``: TMA byte count for ``mbarrier.arrive.expect_tx``.
+        # ``cta_id`` / ``pred``: forwarded to the underlying
+        # ``mbarrier.arrive`` (cluster path) when set; otherwise the
+        # arrive is local-CTA only. See ``MBarrier.arrive`` for the
+        # full default-local rationale.
+        if tx_count is not None:
+            T.ptx.mbarrier.arrive.expect_tx(self.buf.ptr_to([stage]), tx_count)
+        elif cta_id is None:
+            T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
+        else:
+            actual_pred = True if pred is None else pred
+            T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred)
+
+
+class TCGen05Bar(MBarrier):
+    """Barrier signaled by ``tcgen05`` commit.
+
+    The caller is responsible for ensuring only one thread issues the
+    commit, e.g. by wrapping the call in ``if T.ptx.elect_sync():``.
+    """
+
+    @T.inline
+    def arrive(self, stage, cta_group=1, cta_mask=None):
+        # NOTE: this arrive() kwarg set intentionally differs from
+        # MBarrier.arrive (hardware necessity, LSP-incompatible by design).
+        if cta_mask is None and cta_group == 1:
+            T.ptx.tcgen05.commit(self.buf.ptr_to([stage]))
+        else:
+            T.ptx.tcgen05.commit(self.buf.ptr_to([stage]), cta_group=cta_group, cta_mask=cta_mask)
+
+
+# Barrier-type tags accepted by Pipeline's ``full=`` / ``empty=`` arguments.
+_BAR_KINDS = {"tma": TMABar, "tcgen05": TCGen05Bar, "mbar": MBarrier}
+
+
+@T.meta_class
+class Pipeline:
+    """A full/empty mbarrier pair for a software-pipelined data flow.
+
+    Pass barrier-type tags and ``Pipeline`` constructs and ``init``\\ s the
+    barriers itself. Tags: ``"tma"`` (TMABar), ``"tcgen05"`` (TCGen05Bar),
+    ``"mbar"`` (MBarrier). The barrier type and arrival count of each event
+    stay explicit at the call site -- e.g. ``Pipeline(pool, n, full="tma",
+    empty="tcgen05", init_empty=NUM_CONSUMER)``.
+
+    Both signals are required: a ``Pipeline`` is a *pair*. For a one-way event
+    (a pure "X happened" signal with no slot to recycle) use a bare barrier
+    (``TMABar``/``TCGen05Bar``/``MBarrier``) directly -- it has no empty side.
+
+    Parameters
+    ----------
+    pool : SMEMPool
+        Shared memory pool allocator.
+    stages : int
+        Number of pipeline stages (barrier slots).
+    full, empty : str
+        Barrier-type tag for the full / empty signal (see above).
+    init_full, init_empty : int
+        Expected arrival count for the full / empty barrier.
+    empty_phase_offset : int
+        XORed into the empty barrier's phase bit on every wait / arrive.
+    leader : Expr, optional
+        Propagated to both barriers; defaults to thread 0 of the CTA.
+    """
+
+    def __init__(
+        self,
+        pool,
+        stages,
+        *,
+        full,
+        empty,
+        init_full=1,
+        init_empty=1,
+        empty_phase_offset=0,
+        leader=None,
+    ):
+        self.stages = stages
+        self.full = _BAR_KINDS[full](pool, stages, leader=leader)
+        self.full.init(init_full)
+        self.empty = _BAR_KINDS[empty](pool, stages, phase_offset=empty_phase_offset, leader=leader)
+        self.empty.init(init_empty)

python/tvm/backend/cuda/lang/smem_desc.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""SMEM matrix descriptor helper for tcgen05 / wgmma."""
+
+from tvm.backend.cuda.operator.tile_primitive.common import smem_desc_add_16B_offset
+from tvm.script import tirx as T
+
+
+@T.meta_class
+class SmemDescriptor:
+    """Encoded once via :meth:`init`, reused via :meth:`add_16B_offset`."""
+
+    def __init__(self):
+        self._buf = T.alloc_local([1], "uint64")
+
+    @property
+    def desc(self):
+        return self._buf[0]
+
+    @T.inline
+    def init(self, smem_ptr, ldo, sdo, swizzle):
+        T.ptx.tcgen05.encode_matrix_descriptor(
+            T.address_of(self._buf[0]), smem_ptr, ldo, sdo, swizzle
+        )
+
+    def add_16B_offset(self, offset):
+        return smem_desc_add_16B_offset(self._buf[0], offset)
+
+    def make_lo_uniform(self):
+        """Broadcast the lower 32 bits to all warp lanes via ``__shfl_sync``."""
+        func_name = "smem_desc_make_lo_uniform"
+        source_code = f"""
+__forceinline__ __device__ void {func_name}(uint64_t* desc) {{
+    SmemDescriptor* d = reinterpret_cast<SmemDescriptor*>(desc);
+    d->lo = __shfl_sync(0xffffffff, d->lo, 0);
+}}
+"""
+        return T.cuda.func_call(
+            func_name, T.address_of(self._buf[0]), source_code=source_code, return_type="void"
+        )

python/tvm/backend/cuda/lang/tile_scheduler.py

@@ -0,0 +1,945 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Reusable tile scheduler helpers for TIR tests/kernels.
+
+These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
+instances are automatically treated as meta values inside @T.prim_func.
+"""
+
+from tvm.backend.cuda.lang.pipeline import Pipeline, PipelineState
+from tvm.script import tirx as T
+
+
+@T.meta_class
+class BaseTileScheduler:
+    """Base class for tile schedulers with common state and macros."""
+
+    def __init__(self, prefix: str):
+        self.m_idx = T.local_scalar("int32")
+        self.n_idx = T.local_scalar("int32")
+        self.linear_idx = T.local_scalar("int32")
+
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        # To be implemented by subclasses
+        pass
+
+    @T.inline
+    def init(self, linear_init):
+        self.linear_idx = linear_init
+        self.update_current_m_n_idx(linear_init)
+
+    @T.inline
+    def next_tile(self, step):
+        self.linear_idx = self.linear_idx + step
+        self.update_current_m_n_idx(self.linear_idx)
+
+    def valid(self, total_tiles):
+        return self.linear_idx < total_tiles
+
+
+class ClusterPersistentScheduler2D(BaseTileScheduler):
+    """
+    Tile scheduler for cluster-based persistent kernels.
+
+    Distributes a 2D tile grid across persistent clusters using group-major ordering
+    for L2 cache locality. Each cluster starts at its cluster_id and strides by
+    num_clusters to process tiles.
+
+    Tile Ordering (group-major for L2 locality):
+    - Tiles are grouped into "L2 groups" of `l2_group_size` rows
+    - Within a group, tiles are visited in column-major order within the group
+    - Groups are processed in row-major order
+
+    Example with 4x4 tiles, l2_group_size=2:
+        Group 0 (rows 0-1):  0  2  4  6
+                             1  3  5  7
+        Group 1 (rows 2-3):  8 10 12 14
+                             9 11 13 15
+
+    Serpentine Mode (serpentine=True):
+    - Uses CUTLASS-style 2D block swizzle with serpentine traversal
+    - Grid is divided into swizzle_size x swizzle_size blocks
+    - Within each block, tiles are visited in row-major order
+    - Blocks are traversed in serpentine order (even block-rows forward, odd backward)
+    - This provides better L2 locality by reusing both A and B tiles
+
+    Example with 4x4 tiles, swizzle_size=2, serpentine=True:
+        Block layout:
+          Block(0,0)  Block(0,1)
+          Block(1,0)  Block(1,1)
+
+        Tile numbering with serpentine:
+               n=0  n=1  n=2  n=3
+          m=0   0    1   14   15
+          m=1   2    3   12   13
+          m=2   4    5   10   11
+          m=3   6    7    8    9
+
+        Traversal: Block(0,0) -> Block(1,0) -> Block(1,1) -> Block(0,1)
+                   (serpentine: down in col 0, then up in col 1)
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for TIR variable names
+    num_m_tiles : int | T.ExprLike
+        Total number of tiles in M dimension (can be runtime expression)
+    num_n_tiles : int
+        Total number of tiles in N dimension
+    num_clusters : int
+        Number of persistent clusters (determines stride)
+    l2_group_size : int
+        Number of M-tile rows per L2 locality group (default: 8)
+        When serpentine=True, this is used as swizzle_size for 2D blocks
+    cluster_m : int
+        Cluster dimension in M for hierarchical scheduling (default: 1)
+    cluster_n : int
+        Cluster dimension in N for hierarchical scheduling (default: 1)
+    serpentine : bool
+        If True, use CUTLASS-style 2D block swizzle with serpentine traversal (default: False)
+
+    Attributes
+    ----------
+    m_idx : T.local_scalar
+        Current M tile index (output)
+    n_idx : T.local_scalar
+        Current N tile index (output)
+    work_idx : T.local_scalar
+        Global work item index for this cluster
+    tile_count : T.local_scalar
+        Number of tiles processed by this cluster so far
+
+    Usage
+    -----
+    ```python
+    scheduler = ClusterPersistentScheduler2D(
+        "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES,
+        num_clusters=NUM_CLUSTERS, l2_group_size=8
+    )
+    scheduler.init(cluster_id)  # cluster_id = cta_idx // CLUSTER_SIZE
+
+    while scheduler.valid():
+        m = T.meta_var(scheduler.m_idx)  # current M tile
+        n = T.meta_var(scheduler.n_idx)  # current N tile
+        # ... process tile (m, n) ...
+        scheduler.next_tile()
+    ```
+
+    Examples
+    --------
+    Example 1: Basic persistent kernel
+    ```
+    num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2
+    cluster_m=1, cluster_n=1 (default, no tile subdivision)
+
+    Group-major tile numbering (l2_group_size=2):
+           n=0  n=1  n=2  n=3
+      m=0   0    2    4    6   ┐ L2 group 0
+      m=1   1    3    5    7   ┘
+      m=2   8   10   12   14   ┐ L2 group 1
+      m=3   9   11   13   15   ┘
+
+    Work distribution (cluster starts at cluster_id, strides by num_clusters=3):
+      cluster 0: work_idx 0,3,6,9,12,15  -> tiles 0,3,6,9,12,15
+      cluster 1: work_idx 1,4,7,10,13    -> tiles 1,4,7,10,13
+      cluster 2: work_idx 2,5,8,11,14    -> tiles 2,5,8,11,14
+
+    Tile grid (which cluster handles each tile):
+           n=0  n=1  n=2  n=3
+      m=0   C0   C2   C1   C0   ┐ L2 group 0
+      m=1   C1   C0   C2   C1   ┘
+      m=2   C2   C1   C0   C2   ┐ L2 group 1
+      m=3   C0   C2   C1   C0   ┘
+
+    Tile sequence per cluster (in execution order):
+      cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3)
+      cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2)
+      cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3)
+    ```
+
+    Example 2: 2SM GEMM (typical B200 config)
+    ```
+    M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1
+    => M_TILES=8, N_TILES=4
+    => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity)
+
+    Scheduler params:
+      num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8
+      cluster_m=1, cluster_n=1
+
+    Key: Scheduler outputs CLUSTER-level tiles.
+         All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler.
+         CTAs differentiate via cluster_rank (computed OUTSIDE scheduler):
+           cluster_rank = cta_idx % CLUSTER_SIZE
+           cb_m = cluster_rank % CLUSTER_M   # 0 or 1 for 2SM
+           cb_n = cluster_rank // CLUSTER_M  # 0 for 2SM
+
+    Final CTA tile:
+      cta_m = m_idx * CLUSTER_M + cb_m
+      cta_n = n_idx * CLUSTER_N + cb_n
+
+    Example: cluster 5 gets scheduler tile (1,2)
+      CTA rank=0 (cb_m=0): actual tile (2,2)
+      CTA rank=1 (cb_m=1): actual tile (3,2)
+    ```
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        num_m_tiles,
+        num_n_tiles: int,
+        num_clusters: int,
+        l2_group_size: int = 8,
+        cluster_m: int = 1,
+        cluster_n: int = 1,
+        serpentine: bool = False,
+    ):
+        super().__init__(prefix)
+        self._num_m_tiles = num_m_tiles
+        self._num_n_tiles = num_n_tiles
+        self._num_clusters = num_clusters
+        self._l2_group_size = l2_group_size
+        self._cluster_m = cluster_m
+        self._cluster_n = cluster_n
+        self._serpentine = serpentine
+
+        # Rename internal state for clarity
+        self.work_idx = self.linear_idx  # alias: global work item index
+        self.tile_count = T.local_scalar("int32")
+        self.tile_idx = self.tile_count  # alias for backward compatibility
+
+        is_static_m = isinstance(num_m_tiles, int)
+
+        # Number of tile columns after accounting for cluster_n
+        n_tile_cols = (num_n_tiles + cluster_n - 1) // cluster_n
+        self._N_TILE_COLS = n_tile_cols
+
+        if is_static_m:
+            self._M_TILE_ROWS = (num_m_tiles + cluster_m - 1) // cluster_m
+            self._FULL_GROUPS = self._M_TILE_ROWS // l2_group_size
+        else:
+            # Dynamic expressions for runtime M
+            self._M_TILE_ROWS = T.truncdiv(self._num_m_tiles + self._cluster_m - 1, self._cluster_m)
+            self._FULL_GROUPS = T.truncdiv(self._M_TILE_ROWS, self._l2_group_size)
+
+        self._TAIL_ROWS = self._M_TILE_ROWS - self._FULL_GROUPS * l2_group_size
+        self._TOTAL_TILES = self._M_TILE_ROWS * n_tile_cols * cluster_m * cluster_n
+
+        # For serpentine mode: precompute block counts
+        if serpentine:
+            self._N_BLOCKS = n_tile_cols // l2_group_size  # full blocks in N
+            self._M_BLOCKS = (
+                self._M_TILE_ROWS // l2_group_size
+                if is_static_m
+                else T.truncdiv(self._M_TILE_ROWS, l2_group_size)
+            )
+            self._BLOCK_SIZE = l2_group_size * l2_group_size  # tiles per block
+            self._FULL_BLOCK_TILES = self._M_BLOCKS * self._N_BLOCKS * self._BLOCK_SIZE
+            # Residual tiles (not covered by full blocks)
+            self._RESIDUAL_N = n_tile_cols - self._N_BLOCKS * l2_group_size
+            self._RESIDUAL_M = self._M_TILE_ROWS - self._M_BLOCKS * l2_group_size
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, work_idx):
+        """Convert global work index to (m_idx, n_idx) tile coordinates."""
+        CLUSTER_M = T.meta_var(self._cluster_m)
+        CLUSTER_N = T.meta_var(self._cluster_n)
+
+        # Extract hierarchical cluster-local offsets
+        cluster_m_offset = T.meta_var(work_idx % CLUSTER_M)
+        t = T.meta_var(work_idx // CLUSTER_M)
+        cluster_n_offset = T.meta_var(t % CLUSTER_N)
+        tile_linear = T.meta_var(t // CLUSTER_N)
+
+        @T.inline
+        def set_tile_coords(tile_row, tile_col):
+            self.m_idx = tile_row * CLUSTER_M + cluster_m_offset
+            self.n_idx = tile_col * CLUSTER_N + cluster_n_offset
+
+        if self._serpentine:
+            self._update_serpentine(tile_linear, set_tile_coords)
+        else:
+            self._update_group_major(tile_linear, set_tile_coords)
+
+    def _update_group_major(self, tile_linear, set_tile_coords):
+        """Group-major ordering with parse-time pruning of statically-dead branches.
+
+        The TIR script parser does not constant-fold ``if False: ...``, so a
+        Python-literal ``FULL_GROUPS == 0`` would otherwise produce
+        ``T.bitwise_and(T.bool(False), tile_linear < 0)`` IR plus the dead
+        then-leg.  Branch in plain Python here and only invoke the inline
+        emitter that can actually fire.
+        """
+        full_zero = isinstance(self._FULL_GROUPS, int) and self._FULL_GROUPS == 0
+        tail_zero = isinstance(self._TAIL_ROWS, int) and self._TAIL_ROWS == 0
+        if full_zero and tail_zero:
+            self._gm_emit_zero(set_tile_coords)
+        elif full_zero:
+            self._gm_emit_tail_only(tile_linear, set_tile_coords)
+        elif tail_zero:
+            self._gm_emit_full_only(tile_linear, set_tile_coords)
+        else:
+            self._gm_emit_full_and_tail(tile_linear, set_tile_coords)
+
+    @T.inline
+    def _gm_emit_zero(self, set_tile_coords):
+        set_tile_coords(0, 0)
+
+    @T.inline
+    def _gm_emit_full_only(self, tile_linear, set_tile_coords):
+        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
+        GROUP_SIZE = T.meta_var(self._l2_group_size)
+        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
+        if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
+            group_id: T.let = tile_linear // GROUP_SPAN
+            within_group: T.let = tile_linear % GROUP_SPAN
+            tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
+            tile_col: T.let = within_group // GROUP_SIZE
+            set_tile_coords(tile_row, tile_col)
+        else:
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def _gm_emit_tail_only(self, tile_linear, set_tile_coords):
+        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
+        TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
+        GROUP_SIZE = T.meta_var(self._l2_group_size)
+        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
+        if TAIL_ROWS > 0:
+            rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
+            tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
+            tile_col: T.let = rem // TAIL_ROWS
+            set_tile_coords(tile_row, tile_col)
+        else:
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def _gm_emit_full_and_tail(self, tile_linear, set_tile_coords):
+        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
+        TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
+        GROUP_SIZE = T.meta_var(self._l2_group_size)
+        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
+        if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
+            group_id: T.let = tile_linear // GROUP_SPAN
+            within_group: T.let = tile_linear % GROUP_SPAN
+            tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
+            tile_col: T.let = within_group // GROUP_SIZE
+            set_tile_coords(tile_row, tile_col)
+        elif TAIL_ROWS > 0:
+            rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
+            tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
+            tile_col: T.let = rem // TAIL_ROWS
+            set_tile_coords(tile_row, tile_col)
+        else:
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def _update_serpentine(self, tile_linear, set_tile_coords):
+        """CUTLASS-style 2D block swizzle with serpentine traversal.
+
+        Algorithm:
+        1. Divide grid into swizzle_size x swizzle_size blocks
+        2. Within each block, visit tiles in row-major order
+        3. Blocks are traversed column by column (along N)
+        4. Within each column of blocks, use serpentine:
+           - Even columns: top to bottom
+           - Odd columns: bottom to top
+
+        This maximizes L2 reuse for both A and B matrices.
+        """
+        S = T.meta_var(self._l2_group_size)  # swizzle_size
+        M_BLOCKS = T.meta_var(self._M_BLOCKS)
+        N_BLOCKS = T.meta_var(self._N_BLOCKS)
+        BLOCK_SIZE = T.meta_var(self._BLOCK_SIZE)  # S * S
+        FULL_BLOCK_TILES = T.meta_var(self._FULL_BLOCK_TILES)
+        M_TILE_ROWS = T.meta_var(self._M_TILE_ROWS)
+        T.meta_var(self._N_TILE_COLS)
+        RESIDUAL_N = T.meta_var(self._RESIDUAL_N)
+        RESIDUAL_M = T.meta_var(self._RESIDUAL_M)
+
+        # Check if we're in the full block region
+        if (M_BLOCKS > 0) & (N_BLOCKS > 0) & (tile_linear < FULL_BLOCK_TILES):
+            # Which block (in linear order along columns of blocks)
+            block_linear: T.let = tile_linear // BLOCK_SIZE
+            within_block: T.let = tile_linear % BLOCK_SIZE
+
+            # Block column and row
+            block_col: T.let = block_linear // M_BLOCKS
+            block_row_raw: T.let = block_linear % M_BLOCKS
+
+            # Serpentine: odd columns go bottom-to-top
+            block_row: T.let = T.Select(
+                block_col % 2 == 0,
+                block_row_raw,
+                M_BLOCKS - 1 - block_row_raw
+            )
+
+            # Position within block (row-major within block)
+            local_row: T.let = within_block // S
+            local_col: T.let = within_block % S
+
+            tile_row: T.let = block_row * S + local_row
+            tile_col: T.let = block_col * S + local_col
+            set_tile_coords(tile_row, tile_col)
+
+        elif RESIDUAL_N > 0:
+            # Residual tiles in the rightmost partial column of blocks
+            # These are tiles where n >= N_BLOCKS * S
+            rem: T.let = tile_linear - FULL_BLOCK_TILES
+
+            # First handle the right residual strip (full M height, partial N width)
+            right_strip_tiles: T.let = M_TILE_ROWS * RESIDUAL_N
+            if rem < right_strip_tiles:
+                # Row-major within the right strip
+                tile_row: T.let = rem // RESIDUAL_N
+                tile_col: T.let = N_BLOCKS * S + (rem % RESIDUAL_N)
+                set_tile_coords(tile_row, tile_col)
+            elif RESIDUAL_M > 0:
+                # Bottom residual strip (already covered in right strip overlap)
+                # This handles corner case - shouldn't normally reach here
+                # as right strip already covers full M height
+                set_tile_coords(0, 0)
+            else:
+                set_tile_coords(0, 0)
+
+        elif RESIDUAL_M > 0:
+            # Bottom residual strip only (no right residual)
+            rem: T.let = tile_linear - FULL_BLOCK_TILES
+            bottom_strip_tiles: T.let = RESIDUAL_M * (N_BLOCKS * S)
+            if rem < bottom_strip_tiles:
+                tile_row: T.let = M_BLOCKS * S + (rem % RESIDUAL_M)
+                tile_col: T.let = rem // RESIDUAL_M
+                set_tile_coords(tile_row, tile_col)
+            else:
+                set_tile_coords(0, 0)
+        else:
+            # Fallback
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def init(self, cluster_id):
+        """Initialize scheduler for a given cluster.
+
+        Parameters
+        ----------
+        cluster_id : int
+            The cluster's index (typically cta_idx // CLUSTER_SIZE)
+        """
+        self.linear_idx = cluster_id
+        self.tile_count = 0
+        self.update_current_m_n_idx(cluster_id)
+
+    @T.inline
+    def next_tile(self):
+        """Advance to the next tile for this cluster."""
+        self.linear_idx = self.linear_idx + self._num_clusters
+        self.tile_count = self.tile_count + 1
+        self.update_current_m_n_idx(self.linear_idx)
+
+    @T.inline
+    def next_tile_stride(self, stride: int):
+        """Advance by a custom stride (for non-standard scheduling)."""
+        self.linear_idx = self.linear_idx + stride
+        self.tile_count = self.tile_count + 1
+        self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        """Check if this cluster has more tiles to process."""
+        return self.linear_idx < self._TOTAL_TILES
+
+
+class GroupMajor3D(BaseTileScheduler):
+    """
+    3D grouped-row scheduler (M,N,K) with tail handling on M.
+
+    Args
+    ----
+    prefix: str
+    m_tiles: int | T Expr   # tiles along M (static or runtime)
+    n_tiles: int                # tiles along N (static)
+    k_tiles: int                # tiles along K (static)
+    group_rows: int             # rows per group along M
+    step: int = 1               # default stride for next_tile()
+    """
+
+    def __init__(
+        self, prefix: str, m_tiles, n_tiles: int, k_tiles: int, group_rows: int, step: int = 1
+    ):
+        super().__init__(prefix)
+        self._step = step
+        self.tile_idx = T.local_scalar("int32")
+        self.k_idx = T.local_scalar("int32")
+
+        # ---- constants / primexprs baked once ----
+        self._G = group_rows
+        self._N = n_tiles
+        self._K = k_tiles
+
+        if isinstance(m_tiles, int):
+            self._GROUPS = m_tiles // group_rows
+            self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
+            self._SAFE_FINAL_ROWS = max(self._FINAL_ROWS, 1)
+            self._GROUP_SIZE = group_rows * n_tiles * k_tiles
+            self._TOTAL = m_tiles * n_tiles * k_tiles
+        else:
+            self._GROUPS = T.truncdiv(m_tiles, group_rows)
+            self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
+            self._SAFE_FINAL_ROWS = T.max(self._FINAL_ROWS, 1)
+            self._GROUP_SIZE = self._G * self._N * self._K
+            self._TOTAL = m_tiles * n_tiles * k_tiles
+
+        # handy composites used in macro
+        self._FULL_BOUND = self._GROUPS * self._GROUP_SIZE
+        self._HAS_FULL = self._GROUPS > 0
+        self._HAS_TAIL = self._FINAL_ROWS > 0
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        # full-group formulas
+        full_m: T.let = T.floordiv(linear_idx, self._GROUP_SIZE) * self._G + T.floormod(
+            linear_idx, self._G
+        )
+        full_n: T.let = T.floormod(T.floordiv(linear_idx, self._G), self._N)
+        full_k: T.let = T.floordiv(T.floormod(linear_idx, self._GROUP_SIZE), self._G * self._N)
+
+        # tail formulas (relative to FULL_BOUND)
+        # Use _SAFE_FINAL_ROWS (max(FINAL_ROWS, 1)) to avoid divide-by-zero when there is no tail
+        rem: T.let = linear_idx - self._FULL_BOUND
+        tail_m: T.let = self._GROUPS * self._G + T.floormod(rem, self._SAFE_FINAL_ROWS)
+        tail_n: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS) % self._N
+        tail_k: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS * self._N)
+
+        # choose phase
+        if self._HAS_FULL & (linear_idx < self._FULL_BOUND):
+            self.m_idx = full_m
+            self.n_idx = full_n
+            self.k_idx = full_k
+        elif self._HAS_TAIL:
+            self.m_idx = tail_m
+            self.n_idx = tail_n
+            self.k_idx = tail_k
+        else:
+            self.m_idx = 0
+            self.n_idx = 0
+            self.k_idx = 0
+
+    @T.inline
+    def init(self, linear_init):
+        self.linear_idx = linear_init
+        self.tile_idx = 0
+        self.update_current_m_n_idx(linear_init)
+
+    @T.inline
+    def next_tile(self):
+        self.linear_idx = self.linear_idx + self._step
+        self.tile_idx = self.tile_idx + 1
+        self.update_current_m_n_idx(self.linear_idx)
+
+    @T.inline
+    def next_tile_stride(self, stride: int):
+        self.linear_idx = self.linear_idx + stride
+        self.tile_idx = self.tile_idx + 1
+        self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        return self.linear_idx < self._TOTAL
+
+
+class RankAwareGroupMajorTileScheduler(BaseTileScheduler):
+    """
+    Group-major scheduler that applies a rank-aware remapping (remote rows first).
+    Kept as a thin adapter because it depends on NVSHMEM rank at device-side.
+    """
+
+    def __init__(
+        self, prefix: str, m_clusters: int, n_clusters: int, group_size: int, world_size: int
+    ):
+        super().__init__(prefix)
+        self._m_clusters = m_clusters
+        self._n_clusters = n_clusters
+        self._group_size = group_size
+        self._world_size = world_size
+
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        my_rank: T.let = T.nvshmem.my_pe()
+        remote_m_clusters: T.let = self._m_clusters - self._m_clusters // self._world_size
+        group_rows: T.let = (remote_m_clusters // self._group_size) * self._group_size
+        final_rows: T.let = remote_m_clusters - group_rows
+        group_repeat: T.let = self._group_size * self._n_clusters
+        if linear_idx < group_rows * self._n_clusters and group_rows > 0:
+            self.m_idx = (
+                (linear_idx // group_repeat) * self._group_size
+                + (linear_idx % self._group_size)
+                + (my_rank + 1) * self._m_clusters // self._world_size
+            ) % self._m_clusters
+            self.n_idx = (linear_idx % group_repeat) // self._group_size
+        elif linear_idx < remote_m_clusters * self._n_clusters:
+            remainder_idx: T.let = linear_idx - group_rows * self._n_clusters
+            self.m_idx = (
+                group_rows
+                + remainder_idx % final_rows
+                + (my_rank + 1) * self._m_clusters // self._world_size
+            ) % self._m_clusters
+            self.n_idx = remainder_idx // final_rows
+        else:
+            remainder_idx: T.let = linear_idx - remote_m_clusters * self._n_clusters
+            self.m_idx = (
+                remote_m_clusters
+                + remainder_idx % (self._m_clusters // self._world_size)
+                + (my_rank + 1) * self._m_clusters // self._world_size
+            ) % self._m_clusters
+            self.n_idx = remainder_idx // (self._m_clusters // self._world_size)
+
+    @T.inline
+    def next_tile(self, stride: int):
+        self.linear_idx = self.linear_idx + stride
+        self.update_current_m_n_idx(self.linear_idx)
+
+    def valid(self):
+        return self.linear_idx < self._m_clusters * self._n_clusters
+
+
+class IndexedTripleTileScheduler(BaseTileScheduler):
+    """Scheduler that maps linear_idx to (b_idx, h_idx, q_idx) via index lists."""
+
+    def __init__(self, prefix: str, b_indices, h_indices, q_indices, tiles_indptr):
+        super().__init__(prefix)
+        self.b_indices = b_indices
+        self.h_indices = h_indices
+        self.q_indices = q_indices
+        self.tiles_indptr = tiles_indptr
+        self.q_idx = T.local_scalar("int32")
+        self.h_idx = T.local_scalar("int32")
+        self.b_idx = T.local_scalar("int32")
+        self.linear_lim = T.local_scalar("int32")
+
+    @T.inline
+    def _load(self):
+        self.q_idx = self.q_indices[self.linear_idx]
+        self.h_idx = self.h_indices[self.linear_idx]
+        self.b_idx = self.b_indices[self.linear_idx]
+
+    @T.inline
+    def init(self, sm):
+        self.linear_idx = self.tiles_indptr[sm]
+        self.linear_lim = self.tiles_indptr[sm + 1]
+        self._load()
+
+    @T.inline
+    def next_tile(self):
+        self.linear_idx = self.linear_idx + 1
+        self._load()
+
+    def valid(self):
+        return self.linear_idx < self.linear_lim
+
+
+class FlashAttentionLinearScheduler(BaseTileScheduler):
+    """Linear 3D scheduler for flash attention (batch, head, m_block).
+
+    Used for non-causal attention with simple linear decomposition.
+    Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:
+        batch = linear_idx // (num_heads * num_m_blocks)
+        head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks
+        m_block = linear_idx % num_m_blocks
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for TIR variable names
+    num_batches : int
+        Number of batches
+    num_heads : int
+        Number of KV heads
+    num_m_blocks : int
+        Number of Q blocks (M dimension tiles)
+    num_ctas : int
+        Number of CTAs for persistent kernel stride
+    """
+
+    def __init__(
+        self, prefix: str, num_batches: int, num_heads: int, num_m_blocks: int, num_ctas: int
+    ):
+        super().__init__(prefix)
+        self._num_batches = num_batches
+        self._num_heads = num_heads
+        self._num_m_blocks = num_m_blocks
+        self._num_ctas = num_ctas
+        self._total_tasks = num_batches * num_heads * num_m_blocks
+
+        # Output indices
+        self.batch_idx = T.local_scalar("int32")
+        self.head_idx = T.local_scalar("int32")
+        self.m_block_idx = T.local_scalar("int32")
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        """Convert linear index to (batch, head, m_block) coordinates."""
+        NUM_HEADS = T.meta_var(self._num_heads)
+        NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
+        HEAD_M_PRODUCT = T.meta_var(NUM_HEADS * NUM_M_BLOCKS)
+
+        self.batch_idx = linear_idx // HEAD_M_PRODUCT
+        self.head_idx = (linear_idx % HEAD_M_PRODUCT) // NUM_M_BLOCKS
+        self.m_block_idx = linear_idx % NUM_M_BLOCKS
+
+    @T.inline
+    def init(self, cta_id):
+        """Initialize scheduler with CTA ID."""
+        self.linear_idx = cta_id
+        self.update_current_m_n_idx(cta_id)
+
+    @T.inline
+    def next_tile(self):
+        """Advance to next tile by striding by num_ctas."""
+        self.linear_idx = self.linear_idx + self._num_ctas
+        self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        """Check if there are more tiles to process."""
+        return self.linear_idx < self._total_tasks
+
+
+class FlashAttentionLPTScheduler(BaseTileScheduler):
+    """LPT scheduler with L2 swizzle for causal flash attention.
+
+    Processes high-work Q blocks (with more KV blocks to attend to) first using
+    Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle
+    for better cache locality across batch*head dimensions.
+
+    The LPT aspect comes from reversing m_block order: lower Q blocks have more
+    KV blocks to process due to causal masking, so processing them first balances load.
+
+    The scheduler is only applied to non-persistent kernels.
+
+    L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for TIR variable names
+    num_batches : int
+        Number of batches
+    num_heads : int
+        Number of KV heads
+    num_m_blocks : int
+        Number of Q blocks (M dimension tiles)
+    num_ctas : int
+        Number of CTAs (should equal total_tasks for causal)
+    l2_swizzle : int
+        L2 swizzle factor for cache locality
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        num_batches: int,
+        num_heads: int,
+        num_m_blocks: int,
+        l2_swizzle: int,
+        num_ctas: int | None = None,
+    ):
+        super().__init__(prefix)
+        self._num_batches = num_batches
+        self._num_heads = num_heads
+        self._num_m_blocks = num_m_blocks
+        self._l2_swizzle = l2_swizzle
+        self._num_ctas = num_ctas
+        self._total_tasks = num_batches * num_heads * num_m_blocks
+
+        # Derived constants for L2 swizzle
+        self._num_hb = num_batches * num_heads
+        self._l2_major = l2_swizzle * num_m_blocks
+        self._num_hb_quotient = self._num_hb // l2_swizzle
+
+        # Output indices
+        self.batch_idx = T.local_scalar("int32")
+        self.head_idx = T.local_scalar("int32")
+        self.m_block_idx = T.local_scalar("int32")
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        """Convert linear index to (batch, head, m_block) with LPT + L2 swizzle."""
+        L2_SWIZZLE = T.meta_var(self._l2_swizzle)
+        L2_MAJOR = T.meta_var(self._l2_major)
+        NUM_HB_QUOTIENT = T.meta_var(self._num_hb_quotient)
+        NUM_HB = T.meta_var(self._num_hb)
+        NUM_HEADS = T.meta_var(self._num_heads)
+        NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
+
+        # L2 swizzle decomposition
+        bidhb: T.let = linear_idx // L2_MAJOR
+        l2_mod: T.let = linear_idx % L2_MAJOR
+
+        # Handle residual section (last partial swizzle group)
+        num_hb_remainder: T.let = T.max(NUM_HB % L2_SWIZZLE, 1)
+        m_block_raw: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod // L2_SWIZZLE, l2_mod // num_hb_remainder)  # noqa: E501
+        bidhb_residual: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod % L2_SWIZZLE, l2_mod % num_hb_remainder)  # noqa: E501
+        bidhb_actual: T.let = bidhb * L2_SWIZZLE + bidhb_residual
+
+        self.batch_idx = bidhb_actual // NUM_HEADS
+        self.head_idx = bidhb_actual % NUM_HEADS
+
+        # LPT: Reverse block order so high-work blocks are processed first
+        self.m_block_idx = (NUM_M_BLOCKS - 1) - m_block_raw
+
+    @T.inline
+    def init(self, cta_id):
+        """Initialize scheduler with CTA ID."""
+        self.linear_idx = cta_id
+        self.update_current_m_n_idx(cta_id)
+
+    @T.inline
+    def next_tile(self):
+        """Advance to the next tile.
+
+        Single-tile mode (``num_ctas=None``, the default): each CTA owns one
+        task; terminate. Persistent mode (``num_ctas=N``): stride by N, like
+        :class:`FlashAttentionLinearScheduler`, while keeping the LPT + L2
+        swizzle index mapping.
+        """
+        if self._num_ctas is None:
+            self.linear_idx = self._total_tasks
+        else:
+            self.linear_idx = self.linear_idx + self._num_ctas
+            self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        """Check if there are more tiles to process."""
+        return self.linear_idx < self._total_tasks
+
+
+class _CLCWorker(ClusterPersistentScheduler2D):
+    """Per-role CLC handle: IS-A ClusterPersistentScheduler2D (so m_idx / n_idx work as
+    usual) plus the role-local barrier phase and handshake. A coord-free role (e.g. an
+    MMA warp consuming whatever a loader staged) arms the loop with reset() not init().
+    """
+
+    def __init__(self, clc, prefix):
+        super().__init__(
+            prefix,
+            num_m_tiles=clc._num_m_tiles,
+            num_n_tiles=clc._num_n_tiles,
+            num_clusters=clc._num_m_tiles * clc._num_n_tiles,
+            l2_group_size=clc._l2_group_size,
+        )
+        self._clc = clc
+        self._sa = PipelineState(1, 0)
+        self._done = T.local_scalar("int32")
+        self._nxt = T.local_scalar("uint32")
+
+    @T.inline
+    def reset(self):
+        self._done = 0
+
+    @T.inline
+    def init(self, cluster_id):
+        # Explicit base call: TVMScript's parser has no zero-arg super().
+        ClusterPersistentScheduler2D.init(self, cluster_id)
+        self._done = 0
+
+    def valid(self):
+        return self._done == 0
+
+    @T.inline
+    def consume(self):
+        # Single-elected-thread scope: wait for the handle, decode, release the slot.
+        self._clc.sched_arr.full.wait(0, self._sa.phase)
+        self._sa.advance()
+        self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
+        self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)
+
+    @T.inline
+    def consume_wg(self, wg_id, warp_id, lane_id):
+        # Warpgroup scope: all threads decode; one elected lane releases the slot.
+        self._clc.sched_arr.full.wait(0, self._sa.phase)
+        self._sa.advance()
+        self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
+        T.cuda.warpgroup_sync(wg_id + 1)
+        if (warp_id == 0) & (lane_id == 0):
+            self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)
+
+    @T.inline
+    def advance_coords(self):
+        if self._nxt != 0xFFFFFFFF:
+            self.update_current_m_n_idx(self._nxt // self._clc._cta_group)
+
+    @T.inline
+    def mark_done_if_drained(self):
+        if self._nxt == 0xFFFFFFFF:
+            self._done = 1
+
+
+@T.meta_class
+class ClusterLaunchControlScheduler:
+    """Blackwell Cluster Launch Control (CLC) tile scheduler.
+
+    A scheduler warp runs ``run_scheduler`` (issues ``try_cancel`` to steal the next
+    cluster); worker roles each take a ``worker()`` handle and pull the stolen tile
+    through the shared smem handshake. Owns the CLC smem: the 16B response handle, the
+    arrival barrier (handle ready), and the finished barrier (slot consumed;
+    ``finish_arrivals`` arrivals per round). Tile-coord mapping is delegated to
+    ``ClusterPersistentScheduler2D`` (group-major L2 ordering).
+    """
+
+    def __init__(self, pool, num_m_tiles, num_n_tiles, l2_group_size, cta_group, finish_arrivals):
+        self._num_m_tiles = num_m_tiles
+        self._num_n_tiles = num_n_tiles
+        self._l2_group_size = l2_group_size
+        self._cta_group = cta_group
+        self.sched_arr = Pipeline(pool, 1, full="tma", empty="mbar", init_empty=1)
+        self.sched_fin = Pipeline(pool, 1, full="mbar", empty="mbar", init_empty=finish_arrivals)
+        self.clc_handle = pool.alloc((4,), "uint32", align=16)
+        self._s_done = T.local_scalar("int32")
+        self._s_nxt = T.local_scalar("uint32")
+
+    def worker(self, prefix):
+        return _CLCWorker(self, prefix)
+
+    @T.inline
+    def run_scheduler(self, cbx):
+        # cta0 drives try_cancel; both CTAs expect_bytes + consume the handle so the
+        # finished-barrier count is met and the slot can be reissued.
+        if T.ptx.elect_sync():
+            sa = PipelineState(1, 0)
+            sf = PipelineState(1, 1)
+            self._s_done = 0
+            while self._s_done == 0:
+                if cbx == 0:
+                    self.sched_fin.empty.wait(0, sf.phase)
+                    sf.advance()
+                    T.ptx.clc_try_cancel(
+                        T.address_of(self.clc_handle[0]), T.address_of(self.sched_arr.full.buf[0])
+                    )
+                self.sched_arr.full.arrive(0, 16)  # expect_bytes for the 16B handle
+                self.sched_arr.full.wait(0, sa.phase)
+                sa.advance()
+                self._s_nxt = T.ptx.clc_query_cancel(T.address_of(self.clc_handle[0]))
+                self.sched_fin.empty.arrive(0, cta_id=0, pred=True)
+                if self._s_nxt == 0xFFFFFFFF:
+                    self._s_done = 1

python/tvm/backend/cuda/lang/warp_role.py

@@ -0,0 +1,144 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Warp role helpers for SM100 kernels.
+
+Simplifies the common pattern of dispatching warps to named roles
+with register budgets.
+
+Example::
+
+    # Declare roles
+    tma_warp = WarpRole(warp_id, 1, regs=48)
+    store_warp = WarpRole(warp_id, 2, regs=48)
+    mma_warp = WarpRole(warp_id, 0, regs=232, increase=True)
+
+    # Use with context manager
+    with tma_warp:
+        # TMA load code
+    with store_warp:
+        # TMA store code
+    with mma_warp:
+        # MMA compute code
+"""
+
+from tvm.script import tirx as T
+
+
+class WarpRole:
+    """A warp-level role that guards a block of code by warp_id comparison
+    with optional register budget.
+
+    Generates::
+
+        if <warp_id_var> == <warp_id_val>:
+            T.ptx.setmaxnreg(<increase>, <regs>)  # if regs specified
+            <user code>
+
+    The ``if`` guard narrows the active set to the single warp; individual
+    tile-primitive calls inside ``<user code>`` carry their own exec scope via
+    a scope-namespace prefix (e.g. ``Tx.warp.copy(...)``).
+
+    Parameters
+    ----------
+    warp_id_var : Var
+        The warp_id variable (from ``T.warp_id(...)``).
+    warp_id_val : int
+        Which warp index this role corresponds to.
+    regs : int, optional
+        Register budget (passed to ``T.ptx.setmaxnreg``).
+        If None, no setmaxnreg is emitted.
+    increase : bool
+        Direction for ``setmaxnreg`` (default False = decrease).
+    """
+
+    def __init__(self, warp_id_var, warp_id_val, regs=None, increase=False):
+        self.warp_id_var = warp_id_var
+        self.warp_id_val = warp_id_val
+        self.regs = regs
+        self.increase = increase
+
+    def __enter__(self):
+        self._if_frame = T.If(self.warp_id_var == self.warp_id_val)
+        self._if_frame.__enter__()
+        self._then_frame = T.Then()
+        self._then_frame.__enter__()
+        if self.regs is not None:
+            T.evaluate(T.ptx.setmaxnreg(self.increase, self.regs))
+        return self
+
+    def __exit__(self, *exc):
+        self._then_frame.__exit__(*exc)
+        self._if_frame.__exit__(*exc)
+        return False
+
+
+class WarpgroupRole:
+    """A warpgroup-level role that guards by wg_id comparison,
+    with optional register budget.
+
+    Generates (single wg_id)::
+
+        if <wg_id_var> == <wg_id_val>:
+            T.ptx.setmaxnreg(<increase>, <regs>)  # if regs specified
+            <user code>
+
+    Generates (range of wg_ids, e.g. ``wg_id_val=(0, 2)``)::
+
+        if 0 <= <wg_id_var> and <wg_id_var> < 2:
+            T.ptx.setmaxnreg(<increase>, <regs>)
+            <user code>
+
+    The ``if`` guard narrows the active set to the target warpgroup(s);
+    individual tile-primitive calls inside ``<user code>`` carry their own exec
+    scope via a scope-namespace prefix (e.g. ``Tx.wg.copy(...)``).
+
+    Parameters
+    ----------
+    wg_id_var : Var
+        The warpgroup_id variable (from ``T.warpgroup_id(...)``).
+    wg_id_val : int or tuple[int, int]
+        Which warpgroup index (int) or range ``(start, stop)`` this role
+        corresponds to.
+    regs : int, optional
+        Register budget.
+    increase : bool
+        Direction for ``setmaxnreg`` (default False = decrease).
+    """
+
+    def __init__(self, wg_id_var, wg_id_val, regs=None, increase=False):
+        self.wg_id_var = wg_id_var
+        self.wg_id_val = wg_id_val
+        self.regs = regs
+        self.increase = increase
+
+    def __enter__(self):
+        if isinstance(self.wg_id_val, tuple):
+            start, stop = self.wg_id_val
+            self._if_frame = T.If(start <= self.wg_id_var and self.wg_id_var < stop)
+        else:
+            self._if_frame = T.If(self.wg_id_var == self.wg_id_val)
+        self._if_frame.__enter__()
+        self._then_frame = T.Then()
+        self._then_frame.__enter__()
+        if self.regs is not None:
+            T.evaluate(T.ptx.setmaxnreg(self.increase, self.regs))
+        return self
+
+    def __exit__(self, *exc):
+        self._then_frame.__exit__(*exc)
+        self._if_frame.__exit__(*exc)
+        return False