From 6b2d7b6016efa2f6d08c95efe7599ddea54d4dc7 Mon Sep 17 00:00:00 2001
From: ph0375 <wwwph410@vip.qq.com>
Date: Wed, 1 Apr 2026 11:30:43 +0800
Subject: [PATCH 01/41] extract_tile_strides

---
 .../TritonToTritonGPUPass.cpp                 |   7 +-
 .../experimental/tle/language/gpu/core.py     | 422 ++++++++++++++++++
 .../experimental/tle/language/gpu/semantic.py |  19 +-
 third_party/tle/dialect/include/IR/TleOps.td  |   7 +-
 third_party/tle/dialect/lib/IR/Ops.cpp        |  28 +-
 .../lib/Transforms/ExtractTileToLLVM.cpp      |  33 +-
 third_party/tle/triton_tle.cc                 |   8 +-
 7 files changed, 490 insertions(+), 34 deletions(-)

diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
index bf509ffe52..03c584e6ef 100644
--- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
@@ -882,12 +882,13 @@ class TleExtractTileOpPattern : public OpConversionPattern<tle::ExtractTileOp> {
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
+    auto stridesAttr = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
+    op->getAttr("strides"));
     auto newOp = rewriter.replaceOpWithNewOp<tle::ExtractTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getIndex());
-
+    op, retType, adaptor.getSrc(), adaptor.getIndex(), stridesAttr);
+   
     if (auto tileShapeAttr = op->getAttr("tile_shape"))
       newOp->setAttr("tile_shape", tileShapeAttr);
-
     addNamedAttrs(newOp, adaptor.getAttributes());
 
     return success();
diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py
index d913c33755..0a3603eb07 100644
--- a/python/triton/experimental/tle/language/gpu/core.py
+++ b/python/triton/experimental/tle/language/gpu/core.py
@@ -504,6 +504,428 @@ def tmacopy(
         return tmacopy(src, dst, direction, shape, offsets, _semantic)
 
 
+# ============================================================================
+# extract_tile helper functions (module-level, do not depend on @tl.builtin context)
+# ============================================================================
+
+
+def _try_unwrap_int(val):
+    """
+    Try to unwrap val as a Python int.
+    Supports: int, tl.constexpr(int), objects with .value attribute.
+    For runtime tl.tensor, returns None.
+    """
+    if isinstance(val, int):
+        return val
+    try:
+        v = tl._unwrap_if_constexpr(val)
+        if isinstance(v, int):
+            return v
+    except Exception:
+        pass
+    try:
+        if hasattr(val, 'value') and isinstance(val.value, int):
+            return val.value
+    except Exception:
+        pass
+    return None
+
+
+def _unwrap_tile_shape(tile_shape):
+    """Unwrap tile_shape (any form) as List[int], all elements must be compile-time constants."""
+    if hasattr(tile_shape, '__iter__') and not isinstance(tile_shape, str):
+        result = []
+        for s in tile_shape:
+            val = s
+            while hasattr(val, 'value'):
+                val = val.value
+            try:
+                val = tl._unwrap_if_constexpr(val)
+            except Exception:
+                pass
+            if not isinstance(val, int):
+                raise ValueError(
+                    f"All tile_shape dims must be int or tl.constexpr, got {type(val)} (original: {type(s)})")
+            result.append(val)
+        return result
+    else:
+        val = tile_shape
+        while hasattr(val, 'value'):
+            val = val.value
+        try:
+            val = tl._unwrap_if_constexpr(val)
+        except Exception:
+            pass
+        if not isinstance(val, int):
+            raise ValueError(f"tile_shape must be int/constexpr, got {type(val)}")
+        return [val]
+
+
+def _linearize_static_multidim_index(index_list, src_shape, tile_shape_ints, strides_ints=None):
+    """
+    Linearize multi-dimensional static index (row-major order).
+    index_list:      List[int]  tile coordinate in each dimension
+    src_shape:       List[int]  source tensor shape in each dimension
+    tile_shape_ints: List[int]  tile size in each dimension
+    Returns a linearized scalar int
+    """
+    rank = len(src_shape)
+    if len(index_list) != rank:
+        raise ValueError(f"Index rank {len(index_list)} must match source rank {rank}")
+    strides_eff = strides_ints if strides_ints else tile_shape_ints
+    grid = []
+    for i in builtins.range(rank):
+        remainder = (src_shape[i] - tile_shape_ints[i])
+        if remainder < 0 or remainder % strides_eff[i] != 0:
+            raise ValueError(f"(src-tile) not divisible by stride at dim {i}")
+        grid.append(remainder // strides_eff[i] + 1)
+
+    for i, v in builtins.enumerate(index_list):
+        if v < 0 or v >= grid[i]:
+            raise ValueError(f"Index[{i}]={v} out of bounds for grid size {grid[i]}")
+
+    # row major linearization
+    linear = 0
+    stride = 1
+    for i in builtins.reversed(builtins.range(rank)):
+        linear += index_list[i] * stride
+        stride *= grid[i]
+    return linear
+
+
+# ============================================================
+# dynamic multidim index linearization
+# ============================================================
+
+
+def _linearize_dynamic_multidim_index(index_tuple, src_shape, tile_shape_ints, _semantic, strides_ints=None):
+    """
+    Convert dynamic multi-dimensional tile index to linear index IR.
+    Example:
+        src_shape = [16,16]
+        tile_shape = [4,4]
+        tile grid = [4,4]
+        index = [i,j]
+        linear = i*4 + j
+    """
+
+    if any(not isinstance(s, int) for s in src_shape):
+        raise ValueError("Source shape must be static for dynamic multi-dim index")
+    # compute tile grid
+    strides_eff = strides_ints if strides_ints else tile_shape_ints
+    grid = []
+    for i, (s, t) in builtins.enumerate(builtins.zip(src_shape, tile_shape_ints)):
+        grid.append((s - t) // strides_eff[i] + 1)    
+    # compute strides
+    strides = [1] * len(grid)
+    acc = 1
+    for i in builtins.reversed(builtins.range(len(grid))):
+        strides[i] = acc
+        acc *= grid[i]
+
+    # Validation: index_tuple rank must match grid rank
+    if len(index_tuple) != len(grid):
+        raise ValueError(f"Dynamic multi-dim index rank {len(index_tuple)} does not match grid rank {len(grid)}")
+
+    linear_ir = None
+    for i, v in builtins.enumerate(index_tuple):
+        stride = strides[i]
+        if isinstance(v, tl.tensor):
+            term = v.handle
+        else:
+            iv = _try_unwrap_int(v)
+            if iv is None:
+                raise ValueError("Dynamic multidim index must contain tl.tensor or int")
+            term = _semantic._convert_to_ir_values([iv], require_i64=False)[0]
+        if stride != 1:
+            stride_ir = _semantic._convert_to_ir_values([stride], require_i64=False)[0]
+            term = _semantic.builder.create_mul(term, stride_ir)
+        if linear_ir is None:
+            linear_ir = term
+        else:
+            linear_ir = _semantic.builder.create_add(linear_ir, term)
+    return linear_ir
+
+
+@tl.builtin
+def extract_tile(
+    x: tl.tensor,
+    index,
+    tile_shape: tuple,
+    strides: tuple = None,
+    _semantic=None,
+) -> tl.tensor:
+    """
+    Extract a tile from a tensor at a given tile index.
+
+    Supported index forms:
+        1. Multi-dimensional static: tuple/list of int/constexpr (e.g. [1, 2])
+           → Linearized at compile time; uses register shuffle or SMEM path depending on alignment.
+        2. Scalar static: int or tl.constexpr
+           → Treated as already-linearized tile index (compile time constant).
+        3. Scalar dynamic: tl.tensor (scalar, i32/i64)
+           → Treated as a runtime linear tile index; always uses SMEM relay path.
+        4. Multi-dimensional dynamic: tuple/list containing tl.tensor (e.g. [i, j], i/j are tl.tensor)
+           → Automatically linearized at runtime in the frontend; supports mixed int/tl.tensor per axis.
+
+    For multi-dimensional dynamic index, the function will automatically compute the row-major linearized tile index as a dynamic IR expression, so users can pass [i, j, ...] directly.
+
+    Args:
+        x:          Source tensor (tl.tensor)
+        index:      Tile index (see above)
+        tile_shape: Tile shape in each dimension (must be compile-time constants)
+        _semantic:  Internal semantic analyzer (for lowering)
+
+    Returns:
+        Extracted tile tensor with shape = tile_shape
+    Raises:
+        ValueError: If index or shape is invalid
+        RuntimeError: If IR generation fails
+    """
+    strides_ints = None
+    if strides is not None:
+        strides_ints = _unwrap_tile_shape(strides)
+    # --- Parameter check ---
+    if not isinstance(x, tl.tensor):
+        raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
+
+    # --- Unwrap tile_shape (must all be compile-time constants) ---
+    tile_shape_ints = _unwrap_tile_shape(tile_shape)
+
+    src_shape = [tl._unwrap_if_constexpr(dim) for dim in x.type.shape]
+
+    # --- Parse index, three cases ---
+    #
+    #   Case A: tl.tensor → dynamic index, directly pass IR Value handle
+    #   Case B: tuple/list of int/constexpr → multi-dim static, linearize then go to Case C
+    #   Case C: scalar int/constexpr → static scalar
+    #
+    is_dynamic = False
+    index_value = None  # For static path: Python int
+    index_ir_handle = None  # For dynamic path: MLIR Value handle
+
+    if isinstance(index, tl.tensor):
+        # Case A: dynamic index, value known only at runtime
+        is_dynamic = True
+        index_ir_handle = index.handle
+    else:
+        # Try to unwrap, determine if multi-dim or scalar
+        index_unwrapped = index
+        try:
+            index_unwrapped = tl._unwrap_if_constexpr(index)
+        except Exception:
+            pass
+        try:
+            if hasattr(index_unwrapped, 'value'):
+                index_unwrapped = index_unwrapped.value
+        except Exception:
+            pass
+        if isinstance(index_unwrapped, (tuple, list, tl.tuple)):
+            # Case B: multi-dim index → unwrap each element, then linearize to scalar
+            has_tensor = any(isinstance(v, tl.tensor) for v in index_unwrapped)
+            if has_tensor:
+                # ====================================
+                # dynamic multidim index
+                # ====================================
+                index_ir_handle = _linearize_dynamic_multidim_index(index_unwrapped, src_shape, tile_shape_ints,
+                                                                    _semantic, strides_ints)
+                is_dynamic = True
+            else:
+                # ====================================
+                # static multidim index
+                # ====================================
+                idx_ints = []
+                for v in index_unwrapped:
+                    iv = _try_unwrap_int(v)
+                    if iv is None:
+                        raise ValueError("Multi-dim index must contain int/constexpr values.")
+                    idx_ints.append(iv)
+
+                if any(not isinstance(s, int) for s in src_shape):
+                    raise ValueError("Source shape must be static when using a multi-dim index")
+                index_value = _linearize_static_multidim_index(idx_ints, src_shape, tile_shape_ints, strides_ints)
+        else:
+            # Case C: scalar static index
+            scalar_int = _try_unwrap_int(index_unwrapped)
+            if scalar_int is not None:
+                index_value = scalar_int
+            else:
+                raise ValueError(f"index must be int, constexpr, tuple/list of int/constexpr, "
+                                 f"or a scalar tl.tensor; got {type(index)}")
+    # --- Basic dimension check ---
+    if len(tile_shape_ints) != len(src_shape):
+        raise ValueError(f"tile_shape rank ({len(tile_shape_ints)}) must match "
+                         f"source rank ({len(src_shape)})")
+
+    # --- Compile-time check for static index ---
+    if not is_dynamic:
+        strides_eff = strides_ints if strides_ints else tile_shape_ints
+        if all(isinstance(s, int) for s in src_shape):
+            for i, (s, t, st) in builtins.enumerate(
+                    builtins.zip(src_shape, tile_shape_ints, strides_eff)):
+                if (s - t) < 0 or (s - t) % st != 0:
+                    raise ValueError(
+                        f"(src-tile) not divisible by stride at dim {i}: "
+                        f"src={s}, tile={t}, stride={st}")
+            total_tiles = 1
+            for s, t, st in builtins.zip(src_shape, tile_shape_ints, strides_eff):
+                total_tiles *= (s - t) // st + 1 
+            if index_value < 0 or index_value >= total_tiles:
+                raise ValueError(f"index {index_value} out of range [0, {total_tiles})")
+
+        # Semantic validation (static path)
+        try:
+            from .semantic import TLESemantic
+            if isinstance(_semantic, TLESemantic):
+                _semantic.analyze_extract_tile_operation(x, index_value, tile_shape_ints, strides_ints)
+        except ImportError:
+            pass
+
+    # --- Generate MLIR IR ---
+    try:
+        if is_dynamic:
+            # Dynamic index: directly use the IR handle from the input tl.tensor
+            index_ir = index_ir_handle
+        else:
+            # Static index: encode compile-time constant as IR constant
+            index_ir = _semantic._convert_to_ir_values([index_value], require_i64=False)[0]
+
+        output = _semantic.builder.create_extract_tile(x.handle, index_ir, tile_shape_ints, strides_ints or tile_shape_ints)
+        block_type = tl.block_type(x.type.element_ty, tile_shape_ints)
+        return tl.tensor(output, block_type)
+    except Exception as e:
+        raise RuntimeError(f"Failed to create extract_tile operation: {str(e)}") from e
+
+
+@tl.builtin
+def insert_tile(
+    x: tl.tensor,
+    tile: tl.tensor,
+    index,
+    _semantic=None,
+) -> tl.tensor:
+    """
+    Insert a tile into source tensor.
+
+    index supports:
+      1. Multi-dim static index: list/tuple of int/constexpr (e.g. [i, j])
+      2. Scalar static index: int / tl.constexpr
+      3. Scalar dynamic index: tl.tensor (runtime value)
+    """
+    # Basic type checks for source and tile tensors.
+    if not isinstance(x, tl.tensor):
+        raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
+    if not isinstance(tile, tl.tensor):
+        raise ValueError(f"Tile must be tl.tensor, but got {type(tile)}")
+
+    # Shapes must be compile-time integers so tile-grid math stays static.
+    src_shape = [tl._unwrap_if_constexpr(dim) for dim in x.type.shape]
+    tile_shape = [tl._unwrap_if_constexpr(dim) for dim in tile.type.shape]
+    if any(not isinstance(dim, int) for dim in src_shape):
+        raise ValueError("Source shape must be static for insert_tile")
+    if any(not isinstance(dim, int) for dim in tile_shape):
+        raise ValueError("Tile shape must be static for insert_tile")
+    if len(src_shape) != len(tile_shape):
+        raise ValueError(f"Source rank ({len(src_shape)}) must match tile rank ({len(tile_shape)})")
+    if x.type.element_ty != tile.type.element_ty:
+        raise ValueError(f"Element type mismatch: source={x.type.element_ty}, tile={tile.type.element_ty}")
+
+    # Build per-dimension tile grid: how many tiles exist in each axis.
+    grid = []
+    for i, (src_dim, tile_dim) in enumerate(zip(src_shape, tile_shape)):
+        if tile_dim <= 0:
+            raise ValueError(f"Tile dimension {i} must be positive, got {tile_dim}")
+        if src_dim % tile_dim != 0:
+            raise ValueError(f"Source dimension {i}: {src_dim} must be divisible by tile dimension {tile_dim}")
+        grid.append(src_dim // tile_dim)
+
+    # Parse index: dynamic scalar tensor or static scalar/multi-dim.
+    is_dynamic = False
+    index_value = None
+    index_ir_handle = None
+
+    if isinstance(index, tl.tensor):
+        is_dynamic = True
+        index_ir_handle = index.handle
+    else:
+        index_unwrapped = index
+        try:
+            index_unwrapped = tl._unwrap_if_constexpr(index_unwrapped)
+        except Exception:
+            pass
+        try:
+            if hasattr(index_unwrapped, "value"):
+                index_unwrapped = index_unwrapped.value
+        except Exception:
+            pass
+
+        index_list = None
+        if isinstance(index_unwrapped, (tuple, list, tl.tuple)):
+            index_list = list(index_unwrapped)
+            if len(index_list) != len(src_shape):
+                raise ValueError(f"Index rank {len(index_list)} must match source rank {len(src_shape)}")
+            has_tensor = any(isinstance(v, tl.tensor) for v in index_list)
+            # ------------------------------------------------
+            # dynamic multi-dim index
+            # ------------------------------------------------
+            if has_tensor:
+                index_ir_handle = _linearize_dynamic_multidim_index(index_list, src_shape, tile_shape, _semantic)
+                is_dynamic = True
+            # ------------------------------------------------
+            # static multi-dim index
+            # ------------------------------------------------
+            else:
+                idx = []
+                for i, v in enumerate(index_list):
+                    iv = _try_unwrap_int(v)
+                    if iv is None:
+                        raise ValueError("Multi-dim index must contain int/constexpr values")
+                    if iv < 0 or iv >= grid[i]:
+                        raise ValueError(f"Index[{i}]={iv} out of bounds for tile grid (0~{grid[i]-1})")
+                    idx.append(iv)
+                index_value = _linearize_static_multidim_index(idx, src_shape, tile_shape)
+        else:
+            # Path B: scalar static index -> treat as already-linearized tile id.
+            scalar_int = _try_unwrap_int(index_unwrapped)
+            if scalar_int is None:
+                raise ValueError(f"index must be int, constexpr, tuple/list of int/constexpr, "
+                                 f"or a scalar tl.tensor; got {type(index)}")
+            index_value = scalar_int
+
+    # Static index checks + optional semantic pass.
+    if not is_dynamic:
+        if index_value < 0:
+            raise ValueError("Scalar index must be non-negative")
+
+        total_tiles = 1
+        for g in grid:
+            total_tiles *= g
+        if index_value >= total_tiles:
+            raise ValueError(f"Scalar index {index_value} out of bounds for total tiles {total_tiles}")
+
+        try:
+            from .semantic import TLESemantic
+            if isinstance(_semantic, TLESemantic):
+                _semantic.analyze_insert_tile_operation(x, tile, index_value)
+        except ImportError:
+            pass
+
+    # Lower to IR and construct output tensor with the source tensor type.
+    try:
+        if is_dynamic:
+            index_ir = index_ir_handle
+        else:
+            index_ir = _semantic._convert_to_ir_values([index_value], require_i64=False)[0]
+        output = _semantic.builder.create_insert_tile(
+            x.handle,
+            tile.handle,
+            index_ir,
+        )
+        return tl.tensor(output, x.type)
+    except Exception as e:
+        raise RuntimeError(f"Failed to create insert_tile operation: {str(e)}") from e
+
+
 def _expand_index_to_shape(index: tl.tensor, shape: Sequence[int], axis: int, _semantic) -> tl.tensor:
     idx = index
     for _ in builtins.range(axis):
diff --git a/python/triton/experimental/tle/language/gpu/semantic.py b/python/triton/experimental/tle/language/gpu/semantic.py
index 06c531bd99..4c0dccf6ed 100644
--- a/python/triton/experimental/tle/language/gpu/semantic.py
+++ b/python/triton/experimental/tle/language/gpu/semantic.py
@@ -96,7 +96,7 @@ def validate_local_pointer_buffer(self, buffer: tle.buffered_tensor) -> None:
         if not isinstance(buffer, tle.buffered_tensor):
             raise TLESemanticError(f"Buffer must be tle.buffered_tensor, but got {type(buffer)}", "local_ptr")
 
-    def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]]) -> None:
+    def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]], strides=None) -> None:
         """
 
         """
@@ -109,10 +109,17 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
             raise TLESemanticError("Index cannot be None", "extract_tile")
         if not tile_shape:
             raise TLESemanticError("tile_shape cannot be empty", "extract_tile")
+        src_shape = list(src.type.shape)         
 
         # Check 3: unpack and validate type
         tile_shape_unwrapped = [s.value if hasattr(s, 'value') else s for s in tile_shape]
-
+        strides_eff = strides if strides else tile_shape_unwrapped
+        # strides 合法性
+        if any(s <= 0 for s in strides_eff):
+            raise TLESemanticError("All strides must be positive", "extract_tile")
+        if len(strides_eff) != len(src_shape):
+            raise TLESemanticError("strides rank must match source rank", "extract_tile")  
+      
         # Check if every dim in tile_shape is int or constexpr-like
         if any(not isinstance(s, int) for s in tile_shape_unwrapped):
             raise TLESemanticError("All tile_shape dims must be int or constexpr", "extract_tile")
@@ -122,7 +129,7 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
             raise TLESemanticError("All tile_shape dims must be positive", "extract_tile")
 
         # Check 5: dimension match
-        src_shape = list(src.type.shape)
+   #     src_shape = list(src.type.shape)
 
         if len(tile_shape_unwrapped) != len(src_shape):
             raise TLESemanticError(
@@ -141,7 +148,7 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
                 raise TLESemanticError("All index values must be int or constexpr", "extract_tile")
             # Check tile grid out-of-bounds
             if all(isinstance(dim, int) for dim in src_shape):
-                grid = [src_dim // tile_shape_dim for src_dim, tile_shape_dim in zip(src_shape, tile_shape_unwrapped)]
+                grid = [(s - t) // st + 1 for s, t, st in zip(src_shape, tile_shape_unwrapped, strides_eff)]
                 for i, v in enumerate(idx):
                     if v < 0 or v >= grid[i]:
                         raise TLESemanticError(f"Index[{i}]={v} out of bounds for tile grid (0~{grid[i]-1})",
@@ -224,9 +231,9 @@ def validate_insert_tile_params(self, src: tl.tensor, tile: tl.tensor, index) ->
             if val >= total_tiles:
                 raise TLESemanticError(f"Linear index {val} out of bounds for total tiles {total_tiles}", "insert_tile")
 
-    def analyze_extract_tile_operation(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]]) -> None:
+    def analyze_extract_tile_operation(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]], strides=None) -> None:
         """Analyze extract_tile operation semantics"""
-        self.validate_extract_tile_params(src, index, tile_shape)
+        self.validate_extract_tile_params(src, index, tile_shape, strides)
 
     def analyze_insert_tile_operation(
         self,
diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td
index 777e47b5da..3bfe29bdf8 100644
--- a/third_party/tle/dialect/include/IR/TleOps.td
+++ b/third_party/tle/dialect/include/IR/TleOps.td
@@ -31,11 +31,12 @@ def Tle_ExtractTileOp : Tle_Op<"extract_tile", [Pure]> {
       3. Source and destination layouts must match at CTA tile level
   }];
 
-  let arguments = (ins TT_Tensor:$src, TT_IntLike:$index);
-
+  let arguments = (ins TT_Tensor:$src, TT_IntLike:$index,OptionalAttr<DenseI64ArrayAttr>:$strides);
   let results = (outs TT_Tensor:$result);
   let builders = [OpBuilder<(ins "Value":$src, "Value":$index,
-      "ArrayRef<int64_t>":$tileShape)>];
+      "ArrayRef<int64_t>":$tileShape,
+      CArg<"ArrayRef<int64_t>", "{}">:$strides)>];
+
   let assemblyFormat = [{
     $src `[` $index `]` attr-dict `:` qualified(type($src)) `,` qualified(type($index)) `->` qualified(type($result))
   }];
diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp
index f48a81d35e..2776162f72 100644
--- a/third_party/tle/dialect/lib/IR/Ops.cpp
+++ b/third_party/tle/dialect/lib/IR/Ops.cpp
@@ -25,13 +25,15 @@ constexpr int kClusterSharedMemoryAddressSpace = 7;
 // ExtractTileOp Builder
 // ============================================================================
 void ExtractTileOp::build(OpBuilder &builder, OperationState &state, Value src,
-                          Value index, ArrayRef<int64_t> tileShape) {
+                          Value index, ArrayRef<int64_t> tileShape, ArrayRef<int64_t> strides) {
   auto srcType = cast<RankedTensorType>(src.getType());
   auto resultType = RankedTensorType::get(tileShape, srcType.getElementType(),
                                           srcType.getEncoding());
   state.addOperands(src);
   state.addOperands(index);
   state.addAttribute("tile_shape", builder.getDenseI64ArrayAttr(tileShape));
+  SmallVector<int64_t> effectiveStrides(strides.empty() ? tileShape : strides);
+  state.addAttribute("strides", builder.getDenseI64ArrayAttr(effectiveStrides));
   state.addTypes(resultType);
 }
 
@@ -59,6 +61,12 @@ LogicalResult ExtractTileOp::verify() {
     for (auto v : denseArray64.asArrayRef())
       tileShape.push_back(v);
   }
+  SmallVector<int64_t> strides;
+  if (auto a = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
+          getOperation()->getAttr("strides")))
+    for (auto v : a.asArrayRef()) strides.push_back(v);
+  if (strides.empty()) strides = tileShape;  
+
 
   // ---- Basic checks required for both static and dynamic index ----
 
@@ -79,11 +87,12 @@ LogicalResult ExtractTileOp::verify() {
   for (size_t i = 0; i < srcShape.size(); ++i) {
     if (tileShape[i] <= 0)
       return emitOpError("tile_shape must be positive at dimension ") << i;
-    if (srcShape[i] % tileShape[i] != 0)
-      return emitOpError(
-                 "source shape must be divisible by tile_shape at dimension ")
-             << i << " (source=" << srcShape[i] << ", tile=" << tileShape[i]
-             << ")";
+    if (strides[i] <= 0)
+      return emitOpError("strides must be positive at dimension ") << i;
+    if ((srcShape[i] - tileShape[i]) < 0 ||
+        (srcShape[i] - tileShape[i]) % strides[i] != 0)
+      return emitOpError("(srcShape - tileShape) must be divisible by strides "
+                         "at dimension ") << i;
     if (dstShape[i] != tileShape[i])
       return emitOpError("result shape must equal tile_shape at dimension ")
              << i;
@@ -108,9 +117,9 @@ LogicalResult ExtractTileOp::verify() {
   SmallVector<int64_t> logicalGridShape(srcShape.size(), 0);
   int64_t totalTiles = 1;
   for (size_t i = 0; i < srcShape.size(); ++i) {
-    logicalGridShape[i] = srcShape[i] / tileShape[i];
+    logicalGridShape[i] = (srcShape[i] - tileShape[i]) / strides[i] + 1;
     totalTiles *= logicalGridShape[i];
-  }
+ }
 
   // Out-of-bounds check
   if (index < 0 || index >= totalTiles)
@@ -128,8 +137,7 @@ LogicalResult ExtractTileOp::verify() {
   // tile indices -> coordinate-level offsets
   SmallVector<int64_t> offsets(srcShape.size(), 0);
   for (size_t i = 0; i < srcShape.size(); ++i)
-    offsets[i] = tileIndices[i] * tileShape[i];
-
+    offsets[i] = tileIndices[i] * strides[i];
   // Boundary check
   if (offsets.size() != static_cast<size_t>(srcTy.getRank()))
     return emitError("offsets size must match tensor rank");
diff --git a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
index 9454ea913b..fe6975b3a5 100644
--- a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
+++ b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
@@ -29,6 +29,17 @@ static SmallVector<int64_t> getTileShape(ExtractTileOp op) {
   return ts;
 }
 
+static SmallVector<int64_t> getStrides(ExtractTileOp op) {
+  if (auto a = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
+          op->getAttr("strides"))) {
+    SmallVector<int64_t> s;
+    for (auto v : a.asArrayRef()) s.push_back(v);
+    return s;
+  }
+  // 向后兼容：没有 strides 属性时退化为 tile_shape
+  return getTileShape(op);
+}
+
 static std::optional<int64_t> getStaticIndex(ExtractTileOp op) {
   if (auto c = op->getOperand(1).getDefiningOp<mlir::arith::ConstantOp>())
     return mlir::cast<mlir::IntegerAttr>(c.getValue()).getInt();
@@ -39,18 +50,20 @@ static bool isCTATileAligned(ExtractTileOp op, int64_t linearIndex) {
   auto srcTy = cast<RankedTensorType>(op.getSrc().getType());
   auto srcShape = srcTy.getShape();
   auto tileShape = getTileShape(op);
+  auto strides   = getStrides(op);
   auto ctaTile = getShapePerCTATile(srcTy);
   int rank = srcShape.size();
   SmallVector<int64_t> logicalGrid(rank), tileCoords(rank);
   for (int i = 0; i < rank; ++i)
-    logicalGrid[i] = srcShape[i] / tileShape[i];
+    logicalGrid[i] = (srcShape[i] - tileShape[i]) / strides[i] + 1;
   int64_t remain = linearIndex;
   for (int i = rank - 1; i >= 0; --i) {
     tileCoords[i] = remain % logicalGrid[i];
     remain /= logicalGrid[i];
   }
   for (int i = 0; i < rank; ++i) {
-    int64_t off = tileCoords[i] * tileShape[i];
+   
+    int64_t off = tileCoords[i] * strides[i];
     if (tileShape[i] % (int64_t)ctaTile[i] != 0)
       return false;
     if (off % (int64_t)ctaTile[i] != 0)
@@ -73,6 +86,7 @@ lowerExtractTileStatic(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
   auto srcShape = srcTy.getShape(), dstShape = dstTy.getShape();
   auto tileShape = getTileShape(op);
   int rank = srcShape.size();
+  auto strides   = getStrides(op);
   auto vals = unpackLLElements(loc, adaptor.getSrc(), rewriter);
   auto shapePerCTATile = getShapePerCTATile(srcTy);
   auto srcCTAShape = multiDimElementwise<int64_t, unsigned>(
@@ -82,18 +96,18 @@ lowerExtractTileStatic(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
   SmallVector<int64_t> logicalGrid(rank), logicalCoords(rank),
       elementCoords(rank);
   for (int i = 0; i < rank; ++i)
-    logicalGrid[i] = srcShape[i] / tileShape[i];
+    logicalGrid[i] = (srcShape[i] - tileShape[i]) / strides[i] + 1;
   int64_t remain = linearIndex;
   for (int i = rank - 1; i >= 0; --i) {
     logicalCoords[i] = remain % logicalGrid[i];
     remain /= logicalGrid[i];
   }
   for (int i = 0; i < rank; ++i)
-    elementCoords[i] = logicalCoords[i] * tileShape[i];
+    elementCoords[i] = logicalCoords[i] * strides[i]; 
   auto firstTileCoord = multiDimElementwise<int64_t, unsigned>(
       elementCoords, shapePerCTATile, std::divides<unsigned>());
   auto srcCTAOrder = getCTATileOrder(srcTy),
-       dstCTAOrder = getCTATileOrder(dstTy);
+      dstCTAOrder = getCTATileOrder(dstTy);
   unsigned totalSrcCTAs = std::accumulate(
       srcCTAShape.begin(), srcCTAShape.end(), 1, std::multiplies<>());
   unsigned elemsPerCTA = ttg::getTotalElemsPerThread(srcTy) / totalSrcCTAs;
@@ -132,7 +146,7 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
   auto srcShape = srcTy.getShape(), dstShape = dstTy.getShape();
   auto tileShape = getTileShape(op);
   int rank = srcShape.size();
-
+  auto strides   = getStrides(op);
   MLIRContext *ctx = rewriter.getContext();
   auto i1Ty = rewriter.getIntegerType(1);
   auto i8Ty = rewriter.getIntegerType(8);
@@ -187,7 +201,7 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
   // ------------------------------------------------------------------
   SmallVector<int64_t> logicalGrid(rank), suffix(rank, 1);
   for (int d = 0; d < rank; ++d)
-    logicalGrid[d] = srcShape[d] / tileShape[d];
+    logicalGrid[d] = (srcShape[d] - tileShape[d]) / strides[d] + 1;
   for (int d = rank - 2; d >= 0; --d)
     suffix[d] = suffix[d + 1] * logicalGrid[d + 1];
 
@@ -204,11 +218,14 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)suffix[d]));
     Value gv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)logicalGrid[d]));
+    Value sv_stride = rewriter.create<LLVM::ConstantOp>(
+        loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d]));
     Value tv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d]));
+    
     Value coord = rewriter.create<LLVM::UDivOp>(loc, i32Ty, dynIndex, sv);
     coord = rewriter.create<LLVM::URemOp>(loc, i32Ty, coord, gv);
-    tileStartVals[d] = rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, tv);
+    tileStartVals[d] = rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, sv_stride);
     tileEndVals[d] =
         rewriter.create<LLVM::AddOp>(loc, i32Ty, tileStartVals[d], tv);
   }
diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc
index c77b116ede..a34bcced1e 100644
--- a/third_party/tle/triton_tle.cc
+++ b/third_party/tle/triton_tle.cc
@@ -80,12 +80,12 @@ void init_triton_tle_ir(py::module &&m) {
       .def(
           "create_extract_tile",
           [](TritonOpBuilder &self, Value &input,
-             // std::vector<int64_t> &offsets,
-             Value &index, std::vector<int64_t> &tileShape) -> Value {
-            auto op = self.create<tle::ExtractTileOp>(input, index, tileShape);
+            Value &index, std::vector<int64_t> &tileShape,std::vector<int64_t> strides) -> Value {
+            auto op = self.create<tle::ExtractTileOp>(input, index, tileShape, strides);
             return op.getResult();
           },
           py::arg("input"), py::arg("index"), py::arg("tileShape"),
+          py::arg("strides") = std::vector<int64_t>{},
           "Create extract_tile operation")
       .def(
           "create_insert_tile",
@@ -476,7 +476,7 @@ void init_triton_tle_passes(py::module &&m) {
                      tle::createTritonTleLowerExtractTile);
 
   ADD_PASS_WRAPPER_0("add_lower_insert_tile",
-                     tle::createTritonTleLowerInsertTile);
+                    tle::createTritonTleLowerInsertTile);
 }
 
 void init_tle_raw_ir(py::module &&m) {

From 3fbcb0821e897a02241591980f1a7632eb5e2213 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Mon, 13 Apr 2026 11:57:28 +0800
Subject: [PATCH 02/41] [TLE] Extend the strides parameter of insert_tile

---
 .../TritonToTritonGPUPass.cpp                 |   5 +-
 .../unit/test_insert_tile_dynamic_index.py    | 149 +++++++++++++++---
 .../tle/unit/test_insert_tile_static_index.py |  55 +++++++
 .../experimental/tle/language/gpu/core.py     |  34 +++-
 .../experimental/tle/language/gpu/semantic.py |  21 ++-
 third_party/tle/dialect/include/IR/TleOps.td  |   6 +-
 third_party/tle/dialect/lib/IR/Ops.cpp        |  42 ++++-
 .../lib/Transforms/InsertTileToLLVM.cpp       |  74 ++++++---
 third_party/tle/triton_tle.cc                 |   5 +-
 9 files changed, 326 insertions(+), 65 deletions(-)

diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
index 03c584e6ef..f1da3cb66c 100644
--- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
@@ -926,8 +926,11 @@ class TleInsertTileOpPattern : public OpConversionPattern<tle::InsertTileOp> {
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
+    auto stridesAttr = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
+        op->getAttr("strides"));
+
     auto newOp = rewriter.replaceOpWithNewOp<tle::InsertTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex());
+      op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(), stridesAttr);
 
     addNamedAttrs(newOp, adaptor.getAttributes());
 
diff --git a/python/test/tle/unit/test_insert_tile_dynamic_index.py b/python/test/tle/unit/test_insert_tile_dynamic_index.py
index 78135035ac..c374f3fd70 100644
--- a/python/test/tle/unit/test_insert_tile_dynamic_index.py
+++ b/python/test/tle/unit/test_insert_tile_dynamic_index.py
@@ -6,8 +6,22 @@
 
 
 @triton.jit
-def simple_insert_kernel(x_ptr, y_ptr, stride_xb, stride_xm, stride_xn, stride_ym, stride_yn, BLOCK_M: tl.constexpr,
-                         BLOCK_N: tl.constexpr, TILE_M: tl.constexpr, TILE_N: tl.constexpr):
+def simple_insert_dynamic_index_kernel(
+    x_ptr,
+    y_ptr,
+    index_ptr,
+    stride_xb,
+    stride_xm,
+    stride_xn,
+    stride_ym,
+    stride_yn,
+    stride_ib,
+    stride_ic,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+):
     # 1. Get 3D coordinates: z (layer/batch), m (row block), n (col block)
     pid_z = tl.program_id(0)
     pid_m = tl.program_id(1)
@@ -25,13 +39,51 @@ def simple_insert_kernel(x_ptr, y_ptr, stride_xb, stride_xm, stride_xn, stride_y
     y_ptrs = y_ptr + offs_tm[:, None] * stride_ym + offs_tn[None, :] * stride_yn
     small_tile = tl.load(y_ptrs)
 
-    # 4. Determine insertion position:
-    # Layer 0 inserts at top-left [0, 0], Layer 1 inserts at bottom-right [1, 1]
-    # Note: tle.insert_tile 'index' usually must be a constant or determined by static logic
-    if pid_z % 2 == 0:
-        res_tile = tle.insert_tile(bg_tile, small_tile, index=[0, 0])
-    else:
-        res_tile = tle.insert_tile(bg_tile, small_tile, index=[1, 1])
+    # 4. Runtime index per layer: index[pid_z] = [idx_m, idx_n]
+    idx_m = tl.load(index_ptr + pid_z * stride_ib + 0 * stride_ic)
+    idx_n = tl.load(index_ptr + pid_z * stride_ib + 1 * stride_ic)
+    res_tile = tle.insert_tile(bg_tile, small_tile, index=[idx_m, idx_n])
+
+    # 5. Store the resulting tile back to memory
+    tl.store(x_ptrs, res_tile)
+
+
+@triton.jit
+def simple_insert_dynamic_index_stride_kernel(
+    x_ptr,
+    y_ptr,
+    index_ptr,
+    stride_xb,
+    stride_xm,
+    stride_xn,
+    stride_ym,
+    stride_yn,
+    stride_ib,
+    stride_ic,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    STRIDE_M: tl.constexpr,
+    STRIDE_N: tl.constexpr,
+):
+    pid_z = tl.program_id(0)
+    pid_m = tl.program_id(1)
+    pid_n = tl.program_id(2)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    x_ptrs = x_ptr + pid_z * stride_xb + offs_m[:, None] * stride_xm + offs_n[None, :] * stride_xn
+    bg_tile = tl.load(x_ptrs)
+
+    offs_tm = tl.arange(0, TILE_M)
+    offs_tn = tl.arange(0, TILE_N)
+    y_ptrs = y_ptr + offs_tm[:, None] * stride_ym + offs_tn[None, :] * stride_yn
+    small_tile = tl.load(y_ptrs)
+
+    idx_m = tl.load(index_ptr + pid_z * stride_ib + 0 * stride_ic)
+    idx_n = tl.load(index_ptr + pid_z * stride_ib + 1 * stride_ic)
+    res_tile = tle.insert_tile(bg_tile, small_tile, index=[idx_m, idx_n], strides=(STRIDE_M, STRIDE_N))
 
     # 5. Store the resulting tile back to memory
     tl.store(x_ptrs, res_tile)
@@ -43,7 +95,7 @@ def simple_insert_kernel(x_ptr, y_ptr, stride_xb, stride_xm, stride_xn, stride_y
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_simple_insert_kernel_inserts_tiles_correctly():
+def test_simple_insert_kernel_with_dynamic_index():
     B = 2  # 2 layers (Z dimension)
     M, N = 32, 32  # 32x32 size per layer
     TM, TN = 16, 16  # The inserted small tile is 16x16
@@ -53,17 +105,25 @@ def test_simple_insert_kernel_inserts_tiles_correctly():
     # y is an all-99.0 2D small tile
     y = torch.ones((TM, TN), device="cuda", dtype=torch.float32) * 99.0
 
+    # Dynamic insertion indices per layer (no stride argument):
+    # Layer 0 -> [0, 0] => start at [0, 0]
+    # Layer 1 -> [1, 1] => start at [16, 16] (tile size is 16x16)
+    index = torch.tensor([[0, 0], [1, 1]], device="cuda", dtype=torch.int32)
+
     # Launch Kernel: B layers, each needs exactly 1x1 block (since M=32 and BLOCK_M=32)
     grid = (B, 1, 1)
 
-    simple_insert_kernel[grid](
+    simple_insert_dynamic_index_kernel[grid](
         x,
         y,
+        index,
         x.stride(0),
         x.stride(1),
         x.stride(2),
         y.stride(0),
         y.stride(1),
+        index.stride(0),
+        index.stride(1),
         BLOCK_M=32,
         BLOCK_N=32,
         TILE_M=16,
@@ -71,15 +131,58 @@ def test_simple_insert_kernel_inserts_tiles_correctly():
     )
 
     # --- Verification ---
-    # Layer 0 (Z=0): tile inserted at top-left [0:16, 0:16]
-    layer0_tl_mean = x[0, 0:16, 0:16].mean().item()
-    layer0_br_mean = x[0, 16:32, 16:32].mean().item()
-
-    # Layer 1 (Z=1): tile inserted at bottom-right [16:32, 16:32]
-    layer1_tl_mean = x[1, 0:16, 0:16].mean().item()
-    layer1_br_mean = x[1, 16:32, 16:32].mean().item()
-
-    assert layer0_tl_mean == pytest.approx(99.0)
-    assert layer0_br_mean == pytest.approx(0.0)
-    assert layer1_tl_mean == pytest.approx(0.0)
-    assert layer1_br_mean == pytest.approx(99.0)
+    expected = torch.zeros_like(x)
+    expected[0, 0:16, 0:16] = 99.0  # [idx_m, idx_n] = [0, 0]
+    expected[1, 16:32, 16:32] = 99.0  # [idx_m, idx_n] = [1, 1]
+
+    assert torch.allclose(x, expected)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.parametrize(
+    "SM,SN,idx_m1,idx_n1",
+    [
+        (4, 4, 3, 2),
+        (8, 8, 2, 1),
+        (16, 8, 1, 1),
+    ],
+)
+def test_simple_insert_kernel_with_dynamic_index_and_stride(SM, SN, idx_m1, idx_n1):
+    B = 2  # 2 layers (Z dimension)
+    M, N = 32, 32  # 32x32 size per layer
+    TM, TN = 16, 16  # The inserted small tile is 16x16
+
+    x = torch.zeros((B, M, N), device="cuda", dtype=torch.float32)
+    y = torch.ones((TM, TN), device="cuda", dtype=torch.float32) * 99.0
+
+    # Layer 0 is fixed at top-left; Layer 1 uses parameterized dynamic index.
+    index = torch.tensor([[0, 0], [idx_m1, idx_n1]], device="cuda", dtype=torch.int32)
+
+    grid = (B, 1, 1)
+
+    simple_insert_dynamic_index_stride_kernel[grid](
+        x,
+        y,
+        index,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        y.stride(0),
+        y.stride(1),
+        index.stride(0),
+        index.stride(1),
+        BLOCK_M=32,
+        BLOCK_N=32,
+        TILE_M=16,
+        TILE_N=16,
+        STRIDE_M=SM,
+        STRIDE_N=SN,
+    )
+
+    expected = torch.zeros_like(x)
+    expected[0, 0:16, 0:16] = 99.0
+    start_m = idx_m1 * SM
+    start_n = idx_n1 * SN
+    expected[1, start_m:start_m + TM, start_n:start_n + TN] = 99.0
+
+    assert torch.allclose(x, expected)
diff --git a/python/test/tle/unit/test_insert_tile_static_index.py b/python/test/tle/unit/test_insert_tile_static_index.py
index 1b5c38da06..adfe725502 100644
--- a/python/test/tle/unit/test_insert_tile_static_index.py
+++ b/python/test/tle/unit/test_insert_tile_static_index.py
@@ -5,6 +5,34 @@
 import pytest
 
 
+@triton.jit
+def insert_tile_stride_kernel(
+    x_ptr,
+    y_ptr,
+    out_ptr,
+    M: tl.constexpr,
+    N: tl.constexpr,
+    TM: tl.constexpr,
+    TN: tl.constexpr,
+    SM: tl.constexpr,
+    SN: tl.constexpr,
+    idx_m: tl.constexpr,
+    idx_n: tl.constexpr,
+):
+    offs_m = tl.arange(0, M)
+    offs_n = tl.arange(0, N)
+    x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :])
+
+    tile_m = tl.arange(0, TM)
+    tile_n = tl.arange(0, TN)
+    y = tl.load(y_ptr + tile_m[:, None] * TN + tile_n[None, :])
+
+    z = tle.insert_tile(x, y, index=[idx_m, idx_n], strides=(SM, SN))
+
+    tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], z)
+
+
+
 @triton.jit
 def insert_tile_kernel(
     x_ptr,
@@ -61,3 +89,30 @@ def test_insert_tile_static_index():
     print(out[128:132, 128:132].cpu().int())
 
     assert torch.allclose(out, expected)
+
+
+# 新增：专门测试 stride ≠ tile_shape 的 insert_tile
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for this test")
+def test_insert_tile_with_stride():
+    M, N = 256, 256
+    TM, TN = 64, 64
+    SM, SN = 32, 32  # stride < tile_shape
+    idx_m, idx_n = 2, 3
+
+    x = torch.arange(M * N, device="cuda", dtype=torch.float32).reshape(M, N)
+    y = (10000 + torch.arange(TM * TN, device="cuda", dtype=torch.float32)).reshape(TM, TN)
+    out = torch.empty_like(x)
+
+    print(f"Running insert_tile kernel with stride: x={M}x{N}, tile={TM}x{TN}, stride={SM}x{SN}, index=[{idx_m},{idx_n}]...")
+    insert_tile_stride_kernel[(1, )](x, y, out, M, N, TM, TN, SM, SN, idx_m, idx_n)
+    print("Kernel executed.\n")
+
+    expected = x.clone()
+    start_m = idx_m * SM
+    start_n = idx_n * SN
+    expected[start_m:start_m+TM, start_n:start_n+TN] = y
+
+    max_abs_diff = (out - expected).abs().max().item()
+    print(f"max_abs_diff = {max_abs_diff}")
+    assert torch.allclose(out, expected)
+    print("Test passed: insert_tile with stride updated the correct region.")
diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py
index 0a3603eb07..ed8e03ba84 100644
--- a/python/triton/experimental/tle/language/gpu/core.py
+++ b/python/triton/experimental/tle/language/gpu/core.py
@@ -802,6 +802,7 @@ def insert_tile(
     x: tl.tensor,
     tile: tl.tensor,
     index,
+    strides: tuple = None,
     _semantic=None,
 ) -> tl.tensor:
     """
@@ -811,7 +812,14 @@ def insert_tile(
       1. Multi-dim static index: list/tuple of int/constexpr (e.g. [i, j])
       2. Scalar static index: int / tl.constexpr
       3. Scalar dynamic index: tl.tensor (runtime value)
+
+    strides controls tile start-step per dimension. If omitted, defaults to
+    tile shape (non-overlapping tile grid, backward compatible behavior).
     """
+    strides_ints = None
+    if strides is not None:
+        strides_ints = _unwrap_tile_shape(strides)
+
     # Basic type checks for source and tile tensors.
     if not isinstance(x, tl.tensor):
         raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
@@ -831,13 +839,23 @@ def insert_tile(
         raise ValueError(f"Element type mismatch: source={x.type.element_ty}, tile={tile.type.element_ty}")
 
     # Build per-dimension tile grid: how many tiles exist in each axis.
+    # With explicit strides this can represent overlapping tiles.
+    strides_eff = strides_ints if strides_ints else tile_shape
+    if len(strides_eff) != len(src_shape):
+        raise ValueError(f"strides rank ({len(strides_eff)}) must match source rank ({len(src_shape)})")
+
     grid = []
-    for i, (src_dim, tile_dim) in enumerate(zip(src_shape, tile_shape)):
+    for i, (src_dim, tile_dim, stride_dim) in enumerate(zip(src_shape, tile_shape, strides_eff)):
         if tile_dim <= 0:
             raise ValueError(f"Tile dimension {i} must be positive, got {tile_dim}")
-        if src_dim % tile_dim != 0:
-            raise ValueError(f"Source dimension {i}: {src_dim} must be divisible by tile dimension {tile_dim}")
-        grid.append(src_dim // tile_dim)
+        if stride_dim <= 0:
+            raise ValueError(f"Stride dimension {i} must be positive, got {stride_dim}")
+        remainder = src_dim - tile_dim
+        if remainder < 0 or remainder % stride_dim != 0:
+            raise ValueError(
+                f"(src-tile) not divisible by stride at dim {i}: "
+                f"src={src_dim}, tile={tile_dim}, stride={stride_dim}")
+        grid.append(remainder // stride_dim + 1)
 
     # Parse index: dynamic scalar tensor or static scalar/multi-dim.
     is_dynamic = False
@@ -869,7 +887,8 @@ def insert_tile(
             # dynamic multi-dim index
             # ------------------------------------------------
             if has_tensor:
-                index_ir_handle = _linearize_dynamic_multidim_index(index_list, src_shape, tile_shape, _semantic)
+                index_ir_handle = _linearize_dynamic_multidim_index(index_list, src_shape, tile_shape, _semantic,
+                                                                    strides_ints)
                 is_dynamic = True
             # ------------------------------------------------
             # static multi-dim index
@@ -883,7 +902,7 @@ def insert_tile(
                     if iv < 0 or iv >= grid[i]:
                         raise ValueError(f"Index[{i}]={iv} out of bounds for tile grid (0~{grid[i]-1})")
                     idx.append(iv)
-                index_value = _linearize_static_multidim_index(idx, src_shape, tile_shape)
+                index_value = _linearize_static_multidim_index(idx, src_shape, tile_shape, strides_ints)
         else:
             # Path B: scalar static index -> treat as already-linearized tile id.
             scalar_int = _try_unwrap_int(index_unwrapped)
@@ -906,7 +925,7 @@ def insert_tile(
         try:
             from .semantic import TLESemantic
             if isinstance(_semantic, TLESemantic):
-                _semantic.analyze_insert_tile_operation(x, tile, index_value)
+                _semantic.analyze_insert_tile_operation(x, tile, index_value, strides_ints)
         except ImportError:
             pass
 
@@ -920,6 +939,7 @@ def insert_tile(
             x.handle,
             tile.handle,
             index_ir,
+            strides_ints or tile_shape,
         )
         return tl.tensor(output, x.type)
     except Exception as e:
diff --git a/python/triton/experimental/tle/language/gpu/semantic.py b/python/triton/experimental/tle/language/gpu/semantic.py
index 4c0dccf6ed..81543db3a3 100644
--- a/python/triton/experimental/tle/language/gpu/semantic.py
+++ b/python/triton/experimental/tle/language/gpu/semantic.py
@@ -169,7 +169,7 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
                     raise TLESemanticError(f"Linear index {val} out of bounds for total tiles {total_tiles}",
                                            "extract_tile")
 
-    def validate_insert_tile_params(self, src: tl.tensor, tile: tl.tensor, index) -> None:
+    def validate_insert_tile_params(self, src: tl.tensor, tile: tl.tensor, index, strides=None) -> None:
         """
 
         """
@@ -202,12 +202,20 @@ def validate_insert_tile_params(self, src: tl.tensor, tile: tl.tensor, index) ->
         if any(not isinstance(d, int) for d in src_shape_unwrapped):
             raise TLESemanticError("Source shape must be static integers for insert_tile", "insert_tile")
 
+        strides_eff = strides if strides else tile_shape_unwrapped
+        if len(strides_eff) != len(src_shape_unwrapped):
+            raise TLESemanticError("strides rank must match source rank", "insert_tile")
+
         grid = []
-        for i, (src_dim, tile_dim) in enumerate(zip(src_shape_unwrapped, tile_shape_unwrapped)):
-            if src_dim % tile_dim != 0:
+        for i, (src_dim, tile_dim, stride_dim) in enumerate(zip(src_shape_unwrapped, tile_shape_unwrapped, strides_eff)):
+            if stride_dim <= 0:
+                raise TLESemanticError(f"Stride dimension {i} must be positive, got {stride_dim}", "insert_tile")
+            remainder = src_dim - tile_dim
+            if remainder < 0 or remainder % stride_dim != 0:
                 raise TLESemanticError(
-                    f"Source dimension {i}: {src_dim} must be divisible by tile dimension {tile_dim}", "insert_tile")
-            grid.append(src_dim // tile_dim)
+                    f"(src-tile) not divisible by stride at dim {i}: src={src_dim}, tile={tile_dim}, stride={stride_dim}",
+                    "insert_tile")
+            grid.append(remainder // stride_dim + 1)
 
         # index checks
         if isinstance(index, (tuple, list)):
@@ -240,9 +248,10 @@ def analyze_insert_tile_operation(
         src: tl.tensor,
         tile: tl.tensor,
         index,
+        strides=None,
     ) -> None:
         """Analyze insert_tile operation semantics """
-        self.validate_insert_tile_params(src, tile, index)
+        self.validate_insert_tile_params(src, tile, index, strides)
 
     def analyze_alloc_operation(self, shape: Sequence[Union[int, any]], dtype: tl.dtype,
                                 layout: Optional[tle.shared_layout], storage: tle.scope) -> Tuple[List[int], tl.dtype]:
diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td
index 3bfe29bdf8..cea6cc9422 100644
--- a/third_party/tle/dialect/include/IR/TleOps.td
+++ b/third_party/tle/dialect/include/IR/TleOps.td
@@ -48,8 +48,12 @@ def Tle_InsertTileOp
     : Tle_Op<"insert_tile", [Pure,
                              DeclareOpInterfaceMethods<InferTypeOpInterface>,
 ]> {
-  let arguments = (ins TT_Tensor:$src, TT_Tensor:$tile, TT_IntLike:$index);
+  let arguments = (ins TT_Tensor:$src, TT_Tensor:$tile, TT_IntLike:$index,
+      OptionalAttr<DenseI64ArrayAttr>:$strides);
   let results = (outs TT_Tensor:$result);
+  let builders = [OpBuilder<(ins "Value":$src, "Value":$tile,
+      "Value":$index,
+      CArg<"ArrayRef<int64_t>", "{}">:$strides)>];
   let assemblyFormat = [{
     $src `[` $index `]` `=` $tile attr-dict `:` qualified(type($src)) `,` qualified(type($index)) `,` qualified(type($tile)) `->` qualified(type($result))
   }];
diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp
index 2776162f72..fefbafe1e9 100644
--- a/third_party/tle/dialect/lib/IR/Ops.cpp
+++ b/third_party/tle/dialect/lib/IR/Ops.cpp
@@ -356,6 +356,22 @@ void PipeReaderReleaseOp::getEffects(
 // ============================================================================
 // InsertTileOp Type Inference + Verification
 // ============================================================================
+void InsertTileOp::build(OpBuilder &builder, OperationState &state, Value src,
+                         Value tile, Value index, ArrayRef<int64_t> strides) {
+  auto srcType = cast<RankedTensorType>(src.getType());
+  auto tileType = cast<RankedTensorType>(tile.getType());
+  auto tileShape = tileType.getShape();
+
+  state.addOperands(src);
+  state.addOperands(tile);
+  state.addOperands(index);
+  SmallVector<int64_t> effectiveStrides(strides.begin(), strides.end());
+  if (effectiveStrides.empty())
+    effectiveStrides.assign(tileShape.begin(), tileShape.end());
+  state.addAttribute("strides", builder.getDenseI64ArrayAttr(effectiveStrides));
+  state.addTypes(srcType);
+}
+
 LogicalResult InsertTileOp::inferReturnTypes(
     [[maybe_unused]] MLIRContext *context,
     [[maybe_unused]] std::optional<Location> location, ValueRange operands,
@@ -403,6 +419,14 @@ LogicalResult InsertTileOp::verify() {
   auto tileShape = tileTy.getShape();
   auto dstShape = dstTy.getShape();
 
+  SmallVector<int64_t> strides;
+  if (auto a = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
+          getOperation()->getAttr("strides")))
+    for (auto v : a.asArrayRef())
+      strides.push_back(v);
+  if (strides.empty())
+    strides.assign(tileShape.begin(), tileShape.end());
+
   // --- Basic checks required for both static and dynamic index ---
 
   // Check 1: element types must match
@@ -423,17 +447,23 @@ LogicalResult InsertTileOp::verify() {
 
   // Check 4: tile_shape must be positive in each dimension and divide source
   // shape
+  if (strides.size() != srcShape.size())
+    return emitOpError("strides rank must match source rank");
+
   SmallVector<int64_t> logicalGridShape(srcShape.size(), 0);
   int64_t totalTiles = 1;
   for (size_t i = 0; i < srcShape.size(); ++i) {
     if (tileShape[i] <= 0)
       return emitOpError("tile shape must be positive at dimension ") << i;
-    if (srcShape[i] % tileShape[i] != 0)
-      return emitOpError(
-                 "source shape must be divisible by tile shape at dimension ")
+    if (strides[i] <= 0)
+      return emitOpError("strides must be positive at dimension ") << i;
+    if ((srcShape[i] - tileShape[i]) < 0 ||
+        (srcShape[i] - tileShape[i]) % strides[i] != 0)
+      return emitOpError("(source - tile) must be divisible by stride at "
+                         "dimension ")
              << i << " (source=" << srcShape[i] << ", tile=" << tileShape[i]
-             << ")";
-    logicalGridShape[i] = srcShape[i] / tileShape[i];
+             << ", stride=" << strides[i] << ")";
+    logicalGridShape[i] = (srcShape[i] - tileShape[i]) / strides[i] + 1;
     totalTiles *= logicalGridShape[i];
   }
 
@@ -471,7 +501,7 @@ LogicalResult InsertTileOp::verify() {
   // tile indices -> coordinate-level offsets
   SmallVector<int64_t> offsets(srcShape.size(), 0);
   for (size_t i = 0; i < srcShape.size(); ++i)
-    offsets[i] = tileIndices[i] * tileShape[i];
+    offsets[i] = tileIndices[i] * strides[i];
 
   // Boundary check: the full insertion region must be within the source
   for (size_t i = 0; i < srcShape.size(); ++i) {
diff --git a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp
index e0b5f08649..dbcac1ee4a 100644
--- a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp
+++ b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp
@@ -26,18 +26,32 @@ static std::optional<int64_t> getStaticIndex(InsertTileOp op) {
     return mlir::cast<mlir::IntegerAttr>(c.getValue()).getInt();
   return std::nullopt;
 }
+
+static SmallVector<int64_t> getStrides(InsertTileOp op) {
+  if (auto a = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
+          op->getAttr("strides"))) {
+    SmallVector<int64_t> s;
+    for (auto v : a.asArrayRef())
+      s.push_back(v);
+    return s;
+  }
+  auto tileTy = cast<RankedTensorType>(op.getTile().getType());
+  return SmallVector<int64_t>(tileTy.getShape().begin(), tileTy.getShape().end());
+}
+
 // Check if the tile to be inserted is CTA-aligned (for register shuffle path).
 static bool isCTATileAligned(InsertTileOp op, int64_t linearIndex) {
   auto srcTy = cast<RankedTensorType>(op.getSrc().getType());
   auto tileTy = cast<RankedTensorType>(op.getTile().getType());
   auto srcShape = srcTy.getShape();
   auto tileShape = tileTy.getShape();
+  auto strides = getStrides(op);
   auto ctaTile = getShapePerCTATile(srcTy);
   int rank = srcShape.size();
 
   SmallVector<int64_t> logicalGrid(rank), tileCoords(rank);
   for (int i = 0; i < rank; ++i)
-    logicalGrid[i] = srcShape[i] / tileShape[i];
+    logicalGrid[i] = (srcShape[i] - tileShape[i]) / strides[i] + 1;
 
   int64_t remain = linearIndex;
   for (int i = rank - 1; i >= 0; --i) {
@@ -46,7 +60,7 @@ static bool isCTATileAligned(InsertTileOp op, int64_t linearIndex) {
   }
 
   for (int i = 0; i < rank; ++i) {
-    int64_t off = tileCoords[i] * tileShape[i];
+    int64_t off = tileCoords[i] * strides[i];
     if (tileShape[i] % static_cast<int64_t>(ctaTile[i]) != 0)
       return false;
     if (off % static_cast<int64_t>(ctaTile[i]) != 0)
@@ -71,6 +85,7 @@ lowerInsertTileStatic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
 
   auto srcShape = srcTy.getShape();
   auto tileShape = tileTy.getShape();
+  auto strides = getStrides(op);
 
   // Compute CTA tile shapes and grid info.
   auto shapePerCTATile = getShapePerCTATile(srcTy);
@@ -83,9 +98,12 @@ lowerInsertTileStatic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
   SmallVector<int64_t> logicalGridShape(srcShape.size(), 0);
   // Check divisibility and compute logical grid shape.
   for (size_t i = 0; i < srcShape.size(); ++i) {
-    if (logicalTileShape[i] == 0 || srcShape[i] % logicalTileShape[i] != 0)
-      return op.emitError("source shape must be divisible by tile shape");
-    logicalGridShape[i] = srcShape[i] / logicalTileShape[i];
+    if (logicalTileShape[i] == 0 || strides[i] == 0)
+      return op.emitError("tile shape and strides must be non-zero");
+    if ((srcShape[i] - logicalTileShape[i]) < 0 ||
+        (srcShape[i] - logicalTileShape[i]) % strides[i] != 0)
+      return op.emitError("(source - tile) must be divisible by stride");
+    logicalGridShape[i] = (srcShape[i] - logicalTileShape[i]) / strides[i] + 1;
   }
   // Compute logical coordinates from linear index.
   SmallVector<int64_t> logicalCoords(srcShape.size(), 0);
@@ -97,7 +115,7 @@ lowerInsertTileStatic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
 
   SmallVector<int64_t> elementCoords(srcShape.size(), 0);
   for (size_t i = 0; i < srcShape.size(); ++i)
-    elementCoords[i] = logicalCoords[i] * logicalTileShape[i];
+    elementCoords[i] = logicalCoords[i] * strides[i];
 
   // Compute first tile coordinate in CTA space.
   auto firstTileCoordinate = multiDimElementwise<int64_t, unsigned>(
@@ -176,6 +194,7 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
   auto dstTy = cast<RankedTensorType>(op.getType());
   auto srcShape = srcTy.getShape();
   auto tileShape = tileTy.getShape();
+  auto strides = getStrides(op);
   int rank = srcShape.size();
 
   MLIRContext *ctx = rewriter.getContext();
@@ -226,7 +245,7 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
 
   SmallVector<int64_t> logicalGrid(rank), suffix(rank, 1);
   for (int d = 0; d < rank; ++d)
-    logicalGrid[d] = srcShape[d] / tileShape[d];
+    logicalGrid[d] = (srcShape[d] - tileShape[d]) / strides[d] + 1;
   for (int d = rank - 2; d >= 0; --d)
     suffix[d] = suffix[d + 1] * logicalGrid[d + 1];
 
@@ -244,11 +263,13 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)suffix[d]));
     Value gv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)logicalGrid[d]));
+    Value svStride = rewriter.create<LLVM::ConstantOp>(
+      loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d]));
     Value tv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d]));
     Value coord = rewriter.create<LLVM::UDivOp>(loc, i32Ty, dynIndex, sv);
     coord = rewriter.create<LLVM::URemOp>(loc, i32Ty, coord, gv);
-    tileStartVals[d] = rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, tv);
+    tileStartVals[d] = rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, svStride);
     tileEndVals[d] =
         rewriter.create<LLVM::AddOp>(loc, i32Ty, tileStartVals[d], tv);
   }
@@ -310,28 +331,43 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
       inRange = rewriter.create<LLVM::AndOp>(
           loc, rewriter.create<LLVM::AndOp>(loc, ge, lt), inRange);
 
-      Value tileShapeV = rewriter.create<LLVM::ConstantOp>(
-          loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d]));
-      Value tileLocalSafeV =
-          rewriter.create<LLVM::URemOp>(loc, i32Ty, globalCoordV, tileShapeV);
+      // Convert source global coord to tile-local coord.
+      Value tileLocalCoordV = rewriter.create<LLVM::SubOp>(
+          loc, i32Ty, globalCoordV, tileStartVals[d]);
 
       Value sb = rewriter.create<LLVM::ConstantOp>(
           loc, i32Ty,
           rewriter.getI32IntegerAttr((int32_t)(smemStrides[d] * elemBytes)));
       smemByteOffsetV = rewriter.create<LLVM::AddOp>(
           loc, i32Ty, smemByteOffsetV,
-          rewriter.create<LLVM::MulOp>(loc, i32Ty, tileLocalSafeV, sb));
+          rewriter.create<LLVM::MulOp>(loc, i32Ty, tileLocalCoordV, sb));
     }
 
+    // Only load from SMEM when the current source element falls in tile range.
+    // This mirrors extract_tile's conditional memory access style and avoids
+    // any out-of-range SMEM address materialization.
+    Block *cur = rewriter.getInsertionBlock();
+    Block *thenBlock = rewriter.splitBlock(cur, rewriter.getInsertionPoint());
+    Block *elseBlock = rewriter.splitBlock(thenBlock, thenBlock->begin());
+    Block *mergeBlock = rewriter.splitBlock(elseBlock, elseBlock->begin());
+    mergeBlock->addArgument(llvmElemTy, loc);
+
+    rewriter.setInsertionPointToEnd(cur);
+    rewriter.create<LLVM::CondBrOp>(loc, inRange, thenBlock, ValueRange{},
+                                    elseBlock, ValueRange{});
+
+    rewriter.setInsertionPointToStart(thenBlock);
     Value lp = rewriter.create<LLVM::GEPOp>(loc, smemPtrTy, i8Ty, smemBase,
                                             ValueRange{smemByteOffsetV},
                                             LLVM::GEPNoWrapFlags::inbounds);
-    Value tileLoaded =
-        rewriter.create<LLVM::LoadOp>(loc, llvmElemTy, lp, elemBytes);
-    // Overwrite source value with tile value if in range.
-    Value merged =
-        rewriter.create<LLVM::SelectOp>(loc, inRange, tileLoaded, srcVals[i]);
-    resultVals.push_back(merged);
+    Value tileLoaded = rewriter.create<LLVM::LoadOp>(loc, llvmElemTy, lp, elemBytes);
+    rewriter.create<LLVM::BrOp>(loc, ValueRange{tileLoaded}, mergeBlock);
+
+    rewriter.setInsertionPointToStart(elseBlock);
+    rewriter.create<LLVM::BrOp>(loc, ValueRange{srcVals[i]}, mergeBlock);
+
+    rewriter.setInsertionPointToStart(mergeBlock);
+    resultVals.push_back(mergeBlock->getArgument(0));
   }
 
   rewriter.create<NVVM::Barrier0Op>(loc);
diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc
index a34bcced1e..343d7239de 100644
--- a/third_party/tle/triton_tle.cc
+++ b/third_party/tle/triton_tle.cc
@@ -90,11 +90,12 @@ void init_triton_tle_ir(py::module &&m) {
       .def(
           "create_insert_tile",
           [](TritonOpBuilder &self, Value &input, Value &tile,
-             Value &index) -> Value {
-            auto op = self.create<tle::InsertTileOp>(input, tile, index);
+             Value &index, std::vector<int64_t> strides) -> Value {
+            auto op = self.create<tle::InsertTileOp>(input, tile, index, strides);
             return op.getResult();
           },
           py::arg("input"), py::arg("tile"), py::arg("index"),
+          py::arg("strides") = std::vector<int64_t>{},
           "Create insert_tile operation")
       // TLE-Struct
       .def("make_swizzled_shared_encoding_attr",

From d882b7d1d4e8db463006f870992b490a459906e1 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Tue, 21 Apr 2026 09:32:10 +0800
Subject: [PATCH 03/41] [TLE]Fix the issue of extract_tile and insert_tile
 allocating separate smem

---
 third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
index fe6975b3a5..a1bcf90a61 100644
--- a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
+++ b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
@@ -9,6 +9,7 @@
 #include "tle/dialect/include/Transforms/PatternTleToLLVM.h"
 #include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
+#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
 #include "llvm/Support/raw_ostream.h"

From 7254883edf2ee03f348b5159286ffb961f1ac00b Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 10:11:40 +0800
Subject: [PATCH 04/41] [TLE]Update the strides parameters of extract_tile and
 insert_tile, and add test examples optimized using these two primitives

---
 .../triton/experimental/tle/language/core.py  |  111 +-
 .../experimental/tle/language/gpu/core.py     |  441 ----
 python/tutorials/tle/05-glu.py                |  348 ++++
 python/tutorials/tle/06-2D_Depthwise_Conv.py  |  314 +++
 python/tutorials/tle/07-causal-conv1d.py      | 1767 +++++++++++++++++
 python/tutorials/tle/08-rope.py               |  707 +++++++
 6 files changed, 3210 insertions(+), 478 deletions(-)
 create mode 100644 python/tutorials/tle/05-glu.py
 create mode 100644 python/tutorials/tle/06-2D_Depthwise_Conv.py
 create mode 100644 python/tutorials/tle/07-causal-conv1d.py
 create mode 100644 python/tutorials/tle/08-rope.py

diff --git a/python/triton/experimental/tle/language/core.py b/python/triton/experimental/tle/language/core.py
index 03302bc60f..4fcec3c576 100644
--- a/python/triton/experimental/tle/language/core.py
+++ b/python/triton/experimental/tle/language/core.py
@@ -112,6 +112,10 @@ def cumsum(input, axis=0, reverse=False, dtype: tl.constexpr = None, _semantic=N
     return exclusive_sum, total_sum
 
 
+# ============================================================================
+# extract_tile helper functions 
+# ============================================================================
+
 def _try_unwrap_int(val):
     """
     Try to unwrap val as a Python int.
@@ -164,7 +168,7 @@ def _unwrap_tile_shape(tile_shape):
         return [val]
 
 
-def _linearize_static_multidim_index(index_list, src_shape, tile_shape_ints):
+def _linearize_static_multidim_index(index_list, src_shape, tile_shape_ints, strides_ints=None):
     """
     Linearize multi-dimensional static index (row-major order).
     index_list:      List[int]  tile coordinate in each dimension
@@ -175,12 +179,13 @@ def _linearize_static_multidim_index(index_list, src_shape, tile_shape_ints):
     rank = len(src_shape)
     if len(index_list) != rank:
         raise ValueError(f"Index rank {len(index_list)} must match source rank {rank}")
-
+    strides_eff = strides_ints if strides_ints else tile_shape_ints
     grid = []
     for i in builtins.range(rank):
-        if src_shape[i] % tile_shape_ints[i] != 0:
-            raise ValueError(f"Source dim {i} ({src_shape[i]}) not divisible by tile dim ({tile_shape_ints[i]})")
-        grid.append(src_shape[i] // tile_shape_ints[i])
+        remainder = (src_shape[i] - tile_shape_ints[i])
+        if remainder < 0 or remainder % strides_eff[i] != 0:
+            raise ValueError(f"(src-tile) not divisible by stride at dim {i}")
+        grid.append(remainder // strides_eff[i] + 1)
 
     for i, v in builtins.enumerate(index_list):
         if v < 0 or v >= grid[i]:
@@ -195,7 +200,12 @@ def _linearize_static_multidim_index(index_list, src_shape, tile_shape_ints):
     return linear
 
 
-def _linearize_dynamic_multidim_index(index_tuple, src_shape, tile_shape_ints, _semantic):
+# ============================================================
+# dynamic multidim index linearization
+# ============================================================
+
+
+def _linearize_dynamic_multidim_index(index_tuple, src_shape, tile_shape_ints, _semantic, strides_ints=None):
     """
     Convert dynamic multi-dimensional tile index to linear index IR.
     Example:
@@ -209,11 +219,10 @@ def _linearize_dynamic_multidim_index(index_tuple, src_shape, tile_shape_ints, _
     if any(not isinstance(s, int) for s in src_shape):
         raise ValueError("Source shape must be static for dynamic multi-dim index")
     # compute tile grid
+    strides_eff = strides_ints if strides_ints else tile_shape_ints
     grid = []
-    for s, t in builtins.zip(src_shape, tile_shape_ints):
-        if s % t != 0:
-            raise ValueError(f"Source dim {s} not divisible by tile dim {t}")
-        grid.append(s // t)
+    for i, (s, t) in builtins.enumerate(builtins.zip(src_shape, tile_shape_ints)):
+        grid.append((s - t) // strides_eff[i] + 1)    
     # compute strides
     strides = [1] * len(grid)
     acc = 1
@@ -250,6 +259,7 @@ def extract_tile(
     x: tl.tensor,
     index,
     tile_shape: tuple,
+    strides: tuple = None,
     _semantic=None,
 ) -> tl.tensor:
     """
@@ -257,13 +267,13 @@ def extract_tile(
 
     Supported index forms:
         1. Multi-dimensional static: tuple/list of int/constexpr (e.g. [1, 2])
-           -> Linearized at compile time; uses register shuffle or SMEM path depending on alignment.
+           → Linearized at compile time; uses register shuffle or SMEM path depending on alignment.
         2. Scalar static: int or tl.constexpr
-           -> Treated as already-linearized tile index (compile time constant).
+           → Treated as already-linearized tile index (compile time constant).
         3. Scalar dynamic: tl.tensor (scalar, i32/i64)
-           -> Treated as a runtime linear tile index; always uses SMEM relay path.
+           → Treated as a runtime linear tile index; always uses SMEM relay path.
         4. Multi-dimensional dynamic: tuple/list containing tl.tensor (e.g. [i, j], i/j are tl.tensor)
-           -> Automatically linearized at runtime in the frontend; supports mixed int/tl.tensor per axis.
+           → Automatically linearized at runtime in the frontend; supports mixed int/tl.tensor per axis.
 
     For multi-dimensional dynamic index, the function will automatically compute the row-major linearized tile index as a dynamic IR expression, so users can pass [i, j, ...] directly.
 
@@ -279,6 +289,9 @@ def extract_tile(
         ValueError: If index or shape is invalid
         RuntimeError: If IR generation fails
     """
+    strides_ints = None
+    if strides is not None:
+        strides_ints = _unwrap_tile_shape(strides)
     # --- Parameter check ---
     if not isinstance(x, tl.tensor):
         raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
@@ -290,9 +303,9 @@ def extract_tile(
 
     # --- Parse index, three cases ---
     #
-    #   Case A: tl.tensor -> dynamic index, directly pass IR Value handle
-    #   Case B: tuple/list of int/constexpr -> multi-dim static, linearize then go to Case C
-    #   Case C: scalar int/constexpr -> static scalar
+    #   Case A: tl.tensor → dynamic index, directly pass IR Value handle
+    #   Case B: tuple/list of int/constexpr → multi-dim static, linearize then go to Case C
+    #   Case C: scalar int/constexpr → static scalar
     #
     is_dynamic = False
     index_value = None  # For static path: Python int
@@ -315,14 +328,14 @@ def extract_tile(
         except Exception:
             pass
         if isinstance(index_unwrapped, (tuple, list, tl.tuple)):
-            # Case B: multi-dim index -> unwrap each element, then linearize to scalar
+            # Case B: multi-dim index → unwrap each element, then linearize to scalar
             has_tensor = any(isinstance(v, tl.tensor) for v in index_unwrapped)
             if has_tensor:
                 # ====================================
                 # dynamic multidim index
                 # ====================================
                 index_ir_handle = _linearize_dynamic_multidim_index(index_unwrapped, src_shape, tile_shape_ints,
-                                                                    _semantic)
+                                                                    _semantic, strides_ints)
                 is_dynamic = True
             else:
                 # ====================================
@@ -337,7 +350,7 @@ def extract_tile(
 
                 if any(not isinstance(s, int) for s in src_shape):
                     raise ValueError("Source shape must be static when using a multi-dim index")
-                index_value = _linearize_static_multidim_index(idx_ints, src_shape, tile_shape_ints)
+                index_value = _linearize_static_multidim_index(idx_ints, src_shape, tile_shape_ints, strides_ints)
         else:
             # Case C: scalar static index
             scalar_int = _try_unwrap_int(index_unwrapped)
@@ -353,21 +366,25 @@ def extract_tile(
 
     # --- Compile-time check for static index ---
     if not is_dynamic:
-        for i, (s, t) in builtins.enumerate(builtins.zip(src_shape, tile_shape_ints)):
-            if isinstance(s, int) and s % t != 0:
-                raise ValueError(f"Source dim {i} ({s}) not divisible by tile dim ({t})")
+        strides_eff = strides_ints if strides_ints else tile_shape_ints
         if all(isinstance(s, int) for s in src_shape):
+            for i, (s, t, st) in builtins.enumerate(
+                    builtins.zip(src_shape, tile_shape_ints, strides_eff)):
+                if (s - t) < 0 or (s - t) % st != 0:
+                    raise ValueError(
+                        f"(src-tile) not divisible by stride at dim {i}: "
+                        f"src={s}, tile={t}, stride={st}")
             total_tiles = 1
-            for s, t in builtins.zip(src_shape, tile_shape_ints):
-                total_tiles *= s // t
+            for s, t, st in builtins.zip(src_shape, tile_shape_ints, strides_eff):
+                total_tiles *= (s - t) // st + 1 
             if index_value < 0 or index_value >= total_tiles:
                 raise ValueError(f"index {index_value} out of range [0, {total_tiles})")
 
         # Semantic validation (static path)
         try:
-            from .gpu.semantic import TLESemantic
+            from .semantic import TLESemantic
             if isinstance(_semantic, TLESemantic):
-                _semantic.analyze_extract_tile_operation(x, index_value, tile_shape_ints)
+                _semantic.analyze_extract_tile_operation(x, index_value, tile_shape_ints, strides_ints)
         except ImportError:
             pass
 
@@ -380,7 +397,7 @@ def extract_tile(
             # Static index: encode compile-time constant as IR constant
             index_ir = _semantic._convert_to_ir_values([index_value], require_i64=False)[0]
 
-        output = _semantic.builder.create_extract_tile(x.handle, index_ir, tile_shape_ints)
+        output = _semantic.builder.create_extract_tile(x.handle, index_ir, tile_shape_ints, strides_ints or tile_shape_ints)
         block_type = tl.block_type(x.type.element_ty, tile_shape_ints)
         return tl.tensor(output, block_type)
     except Exception as e:
@@ -392,6 +409,7 @@ def insert_tile(
     x: tl.tensor,
     tile: tl.tensor,
     index,
+    strides: tuple = None,
     _semantic=None,
 ) -> tl.tensor:
     """
@@ -401,7 +419,14 @@ def insert_tile(
       1. Multi-dim static index: list/tuple of int/constexpr (e.g. [i, j])
       2. Scalar static index: int / tl.constexpr
       3. Scalar dynamic index: tl.tensor (runtime value)
+
+    strides controls tile start-step per dimension. If omitted, defaults to
+    tile shape (non-overlapping tile grid, backward compatible behavior).
     """
+    strides_ints = None
+    if strides is not None:
+        strides_ints = _unwrap_tile_shape(strides)
+
     # Basic type checks for source and tile tensors.
     if not isinstance(x, tl.tensor):
         raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
@@ -421,13 +446,23 @@ def insert_tile(
         raise ValueError(f"Element type mismatch: source={x.type.element_ty}, tile={tile.type.element_ty}")
 
     # Build per-dimension tile grid: how many tiles exist in each axis.
+    # With explicit strides this can represent overlapping tiles.
+    strides_eff = strides_ints if strides_ints else tile_shape
+    if len(strides_eff) != len(src_shape):
+        raise ValueError(f"strides rank ({len(strides_eff)}) must match source rank ({len(src_shape)})")
+
     grid = []
-    for i, (src_dim, tile_dim) in builtins.enumerate(builtins.zip(src_shape, tile_shape)):
+    for i, (src_dim, tile_dim, stride_dim) in enumerate(zip(src_shape, tile_shape, strides_eff)):
         if tile_dim <= 0:
             raise ValueError(f"Tile dimension {i} must be positive, got {tile_dim}")
-        if src_dim % tile_dim != 0:
-            raise ValueError(f"Source dimension {i}: {src_dim} must be divisible by tile dimension {tile_dim}")
-        grid.append(src_dim // tile_dim)
+        if stride_dim <= 0:
+            raise ValueError(f"Stride dimension {i} must be positive, got {stride_dim}")
+        remainder = src_dim - tile_dim
+        if remainder < 0 or remainder % stride_dim != 0:
+            raise ValueError(
+                f"(src-tile) not divisible by stride at dim {i}: "
+                f"src={src_dim}, tile={tile_dim}, stride={stride_dim}")
+        grid.append(remainder // stride_dim + 1)
 
     # Parse index: dynamic scalar tensor or static scalar/multi-dim.
     is_dynamic = False
@@ -459,21 +494,22 @@ def insert_tile(
             # dynamic multi-dim index
             # ------------------------------------------------
             if has_tensor:
-                index_ir_handle = _linearize_dynamic_multidim_index(index_list, src_shape, tile_shape, _semantic)
+                index_ir_handle = _linearize_dynamic_multidim_index(index_list, src_shape, tile_shape, _semantic,
+                                                                    strides_ints)
                 is_dynamic = True
             # ------------------------------------------------
             # static multi-dim index
             # ------------------------------------------------
             else:
                 idx = []
-                for i, v in builtins.enumerate(index_list):
+                for i, v in enumerate(index_list):
                     iv = _try_unwrap_int(v)
                     if iv is None:
                         raise ValueError("Multi-dim index must contain int/constexpr values")
                     if iv < 0 or iv >= grid[i]:
                         raise ValueError(f"Index[{i}]={iv} out of bounds for tile grid (0~{grid[i]-1})")
                     idx.append(iv)
-                index_value = _linearize_static_multidim_index(idx, src_shape, tile_shape)
+                index_value = _linearize_static_multidim_index(idx, src_shape, tile_shape, strides_ints)
         else:
             # Path B: scalar static index -> treat as already-linearized tile id.
             scalar_int = _try_unwrap_int(index_unwrapped)
@@ -494,9 +530,9 @@ def insert_tile(
             raise ValueError(f"Scalar index {index_value} out of bounds for total tiles {total_tiles}")
 
         try:
-            from .gpu.semantic import TLESemantic
+            from .semantic import TLESemantic
             if isinstance(_semantic, TLESemantic):
-                _semantic.analyze_insert_tile_operation(x, tile, index_value)
+                _semantic.analyze_insert_tile_operation(x, tile, index_value, strides_ints)
         except ImportError:
             pass
 
@@ -510,6 +546,7 @@ def insert_tile(
             x.handle,
             tile.handle,
             index_ir,
+            strides_ints or tile_shape,
         )
         return tl.tensor(output, x.type)
     except Exception as e:
diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py
index ed8e03ba84..b769797c34 100644
--- a/python/triton/experimental/tle/language/gpu/core.py
+++ b/python/triton/experimental/tle/language/gpu/core.py
@@ -504,447 +504,6 @@ def tmacopy(
         return tmacopy(src, dst, direction, shape, offsets, _semantic)
 
 
-# ============================================================================
-# extract_tile helper functions (module-level, do not depend on @tl.builtin context)
-# ============================================================================
-
-
-def _try_unwrap_int(val):
-    """
-    Try to unwrap val as a Python int.
-    Supports: int, tl.constexpr(int), objects with .value attribute.
-    For runtime tl.tensor, returns None.
-    """
-    if isinstance(val, int):
-        return val
-    try:
-        v = tl._unwrap_if_constexpr(val)
-        if isinstance(v, int):
-            return v
-    except Exception:
-        pass
-    try:
-        if hasattr(val, 'value') and isinstance(val.value, int):
-            return val.value
-    except Exception:
-        pass
-    return None
-
-
-def _unwrap_tile_shape(tile_shape):
-    """Unwrap tile_shape (any form) as List[int], all elements must be compile-time constants."""
-    if hasattr(tile_shape, '__iter__') and not isinstance(tile_shape, str):
-        result = []
-        for s in tile_shape:
-            val = s
-            while hasattr(val, 'value'):
-                val = val.value
-            try:
-                val = tl._unwrap_if_constexpr(val)
-            except Exception:
-                pass
-            if not isinstance(val, int):
-                raise ValueError(
-                    f"All tile_shape dims must be int or tl.constexpr, got {type(val)} (original: {type(s)})")
-            result.append(val)
-        return result
-    else:
-        val = tile_shape
-        while hasattr(val, 'value'):
-            val = val.value
-        try:
-            val = tl._unwrap_if_constexpr(val)
-        except Exception:
-            pass
-        if not isinstance(val, int):
-            raise ValueError(f"tile_shape must be int/constexpr, got {type(val)}")
-        return [val]
-
-
-def _linearize_static_multidim_index(index_list, src_shape, tile_shape_ints, strides_ints=None):
-    """
-    Linearize multi-dimensional static index (row-major order).
-    index_list:      List[int]  tile coordinate in each dimension
-    src_shape:       List[int]  source tensor shape in each dimension
-    tile_shape_ints: List[int]  tile size in each dimension
-    Returns a linearized scalar int
-    """
-    rank = len(src_shape)
-    if len(index_list) != rank:
-        raise ValueError(f"Index rank {len(index_list)} must match source rank {rank}")
-    strides_eff = strides_ints if strides_ints else tile_shape_ints
-    grid = []
-    for i in builtins.range(rank):
-        remainder = (src_shape[i] - tile_shape_ints[i])
-        if remainder < 0 or remainder % strides_eff[i] != 0:
-            raise ValueError(f"(src-tile) not divisible by stride at dim {i}")
-        grid.append(remainder // strides_eff[i] + 1)
-
-    for i, v in builtins.enumerate(index_list):
-        if v < 0 or v >= grid[i]:
-            raise ValueError(f"Index[{i}]={v} out of bounds for grid size {grid[i]}")
-
-    # row major linearization
-    linear = 0
-    stride = 1
-    for i in builtins.reversed(builtins.range(rank)):
-        linear += index_list[i] * stride
-        stride *= grid[i]
-    return linear
-
-
-# ============================================================
-# dynamic multidim index linearization
-# ============================================================
-
-
-def _linearize_dynamic_multidim_index(index_tuple, src_shape, tile_shape_ints, _semantic, strides_ints=None):
-    """
-    Convert dynamic multi-dimensional tile index to linear index IR.
-    Example:
-        src_shape = [16,16]
-        tile_shape = [4,4]
-        tile grid = [4,4]
-        index = [i,j]
-        linear = i*4 + j
-    """
-
-    if any(not isinstance(s, int) for s in src_shape):
-        raise ValueError("Source shape must be static for dynamic multi-dim index")
-    # compute tile grid
-    strides_eff = strides_ints if strides_ints else tile_shape_ints
-    grid = []
-    for i, (s, t) in builtins.enumerate(builtins.zip(src_shape, tile_shape_ints)):
-        grid.append((s - t) // strides_eff[i] + 1)    
-    # compute strides
-    strides = [1] * len(grid)
-    acc = 1
-    for i in builtins.reversed(builtins.range(len(grid))):
-        strides[i] = acc
-        acc *= grid[i]
-
-    # Validation: index_tuple rank must match grid rank
-    if len(index_tuple) != len(grid):
-        raise ValueError(f"Dynamic multi-dim index rank {len(index_tuple)} does not match grid rank {len(grid)}")
-
-    linear_ir = None
-    for i, v in builtins.enumerate(index_tuple):
-        stride = strides[i]
-        if isinstance(v, tl.tensor):
-            term = v.handle
-        else:
-            iv = _try_unwrap_int(v)
-            if iv is None:
-                raise ValueError("Dynamic multidim index must contain tl.tensor or int")
-            term = _semantic._convert_to_ir_values([iv], require_i64=False)[0]
-        if stride != 1:
-            stride_ir = _semantic._convert_to_ir_values([stride], require_i64=False)[0]
-            term = _semantic.builder.create_mul(term, stride_ir)
-        if linear_ir is None:
-            linear_ir = term
-        else:
-            linear_ir = _semantic.builder.create_add(linear_ir, term)
-    return linear_ir
-
-
-@tl.builtin
-def extract_tile(
-    x: tl.tensor,
-    index,
-    tile_shape: tuple,
-    strides: tuple = None,
-    _semantic=None,
-) -> tl.tensor:
-    """
-    Extract a tile from a tensor at a given tile index.
-
-    Supported index forms:
-        1. Multi-dimensional static: tuple/list of int/constexpr (e.g. [1, 2])
-           → Linearized at compile time; uses register shuffle or SMEM path depending on alignment.
-        2. Scalar static: int or tl.constexpr
-           → Treated as already-linearized tile index (compile time constant).
-        3. Scalar dynamic: tl.tensor (scalar, i32/i64)
-           → Treated as a runtime linear tile index; always uses SMEM relay path.
-        4. Multi-dimensional dynamic: tuple/list containing tl.tensor (e.g. [i, j], i/j are tl.tensor)
-           → Automatically linearized at runtime in the frontend; supports mixed int/tl.tensor per axis.
-
-    For multi-dimensional dynamic index, the function will automatically compute the row-major linearized tile index as a dynamic IR expression, so users can pass [i, j, ...] directly.
-
-    Args:
-        x:          Source tensor (tl.tensor)
-        index:      Tile index (see above)
-        tile_shape: Tile shape in each dimension (must be compile-time constants)
-        _semantic:  Internal semantic analyzer (for lowering)
-
-    Returns:
-        Extracted tile tensor with shape = tile_shape
-    Raises:
-        ValueError: If index or shape is invalid
-        RuntimeError: If IR generation fails
-    """
-    strides_ints = None
-    if strides is not None:
-        strides_ints = _unwrap_tile_shape(strides)
-    # --- Parameter check ---
-    if not isinstance(x, tl.tensor):
-        raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
-
-    # --- Unwrap tile_shape (must all be compile-time constants) ---
-    tile_shape_ints = _unwrap_tile_shape(tile_shape)
-
-    src_shape = [tl._unwrap_if_constexpr(dim) for dim in x.type.shape]
-
-    # --- Parse index, three cases ---
-    #
-    #   Case A: tl.tensor → dynamic index, directly pass IR Value handle
-    #   Case B: tuple/list of int/constexpr → multi-dim static, linearize then go to Case C
-    #   Case C: scalar int/constexpr → static scalar
-    #
-    is_dynamic = False
-    index_value = None  # For static path: Python int
-    index_ir_handle = None  # For dynamic path: MLIR Value handle
-
-    if isinstance(index, tl.tensor):
-        # Case A: dynamic index, value known only at runtime
-        is_dynamic = True
-        index_ir_handle = index.handle
-    else:
-        # Try to unwrap, determine if multi-dim or scalar
-        index_unwrapped = index
-        try:
-            index_unwrapped = tl._unwrap_if_constexpr(index)
-        except Exception:
-            pass
-        try:
-            if hasattr(index_unwrapped, 'value'):
-                index_unwrapped = index_unwrapped.value
-        except Exception:
-            pass
-        if isinstance(index_unwrapped, (tuple, list, tl.tuple)):
-            # Case B: multi-dim index → unwrap each element, then linearize to scalar
-            has_tensor = any(isinstance(v, tl.tensor) for v in index_unwrapped)
-            if has_tensor:
-                # ====================================
-                # dynamic multidim index
-                # ====================================
-                index_ir_handle = _linearize_dynamic_multidim_index(index_unwrapped, src_shape, tile_shape_ints,
-                                                                    _semantic, strides_ints)
-                is_dynamic = True
-            else:
-                # ====================================
-                # static multidim index
-                # ====================================
-                idx_ints = []
-                for v in index_unwrapped:
-                    iv = _try_unwrap_int(v)
-                    if iv is None:
-                        raise ValueError("Multi-dim index must contain int/constexpr values.")
-                    idx_ints.append(iv)
-
-                if any(not isinstance(s, int) for s in src_shape):
-                    raise ValueError("Source shape must be static when using a multi-dim index")
-                index_value = _linearize_static_multidim_index(idx_ints, src_shape, tile_shape_ints, strides_ints)
-        else:
-            # Case C: scalar static index
-            scalar_int = _try_unwrap_int(index_unwrapped)
-            if scalar_int is not None:
-                index_value = scalar_int
-            else:
-                raise ValueError(f"index must be int, constexpr, tuple/list of int/constexpr, "
-                                 f"or a scalar tl.tensor; got {type(index)}")
-    # --- Basic dimension check ---
-    if len(tile_shape_ints) != len(src_shape):
-        raise ValueError(f"tile_shape rank ({len(tile_shape_ints)}) must match "
-                         f"source rank ({len(src_shape)})")
-
-    # --- Compile-time check for static index ---
-    if not is_dynamic:
-        strides_eff = strides_ints if strides_ints else tile_shape_ints
-        if all(isinstance(s, int) for s in src_shape):
-            for i, (s, t, st) in builtins.enumerate(
-                    builtins.zip(src_shape, tile_shape_ints, strides_eff)):
-                if (s - t) < 0 or (s - t) % st != 0:
-                    raise ValueError(
-                        f"(src-tile) not divisible by stride at dim {i}: "
-                        f"src={s}, tile={t}, stride={st}")
-            total_tiles = 1
-            for s, t, st in builtins.zip(src_shape, tile_shape_ints, strides_eff):
-                total_tiles *= (s - t) // st + 1 
-            if index_value < 0 or index_value >= total_tiles:
-                raise ValueError(f"index {index_value} out of range [0, {total_tiles})")
-
-        # Semantic validation (static path)
-        try:
-            from .semantic import TLESemantic
-            if isinstance(_semantic, TLESemantic):
-                _semantic.analyze_extract_tile_operation(x, index_value, tile_shape_ints, strides_ints)
-        except ImportError:
-            pass
-
-    # --- Generate MLIR IR ---
-    try:
-        if is_dynamic:
-            # Dynamic index: directly use the IR handle from the input tl.tensor
-            index_ir = index_ir_handle
-        else:
-            # Static index: encode compile-time constant as IR constant
-            index_ir = _semantic._convert_to_ir_values([index_value], require_i64=False)[0]
-
-        output = _semantic.builder.create_extract_tile(x.handle, index_ir, tile_shape_ints, strides_ints or tile_shape_ints)
-        block_type = tl.block_type(x.type.element_ty, tile_shape_ints)
-        return tl.tensor(output, block_type)
-    except Exception as e:
-        raise RuntimeError(f"Failed to create extract_tile operation: {str(e)}") from e
-
-
-@tl.builtin
-def insert_tile(
-    x: tl.tensor,
-    tile: tl.tensor,
-    index,
-    strides: tuple = None,
-    _semantic=None,
-) -> tl.tensor:
-    """
-    Insert a tile into source tensor.
-
-    index supports:
-      1. Multi-dim static index: list/tuple of int/constexpr (e.g. [i, j])
-      2. Scalar static index: int / tl.constexpr
-      3. Scalar dynamic index: tl.tensor (runtime value)
-
-    strides controls tile start-step per dimension. If omitted, defaults to
-    tile shape (non-overlapping tile grid, backward compatible behavior).
-    """
-    strides_ints = None
-    if strides is not None:
-        strides_ints = _unwrap_tile_shape(strides)
-
-    # Basic type checks for source and tile tensors.
-    if not isinstance(x, tl.tensor):
-        raise ValueError(f"Source must be tl.tensor, but got {type(x)}")
-    if not isinstance(tile, tl.tensor):
-        raise ValueError(f"Tile must be tl.tensor, but got {type(tile)}")
-
-    # Shapes must be compile-time integers so tile-grid math stays static.
-    src_shape = [tl._unwrap_if_constexpr(dim) for dim in x.type.shape]
-    tile_shape = [tl._unwrap_if_constexpr(dim) for dim in tile.type.shape]
-    if any(not isinstance(dim, int) for dim in src_shape):
-        raise ValueError("Source shape must be static for insert_tile")
-    if any(not isinstance(dim, int) for dim in tile_shape):
-        raise ValueError("Tile shape must be static for insert_tile")
-    if len(src_shape) != len(tile_shape):
-        raise ValueError(f"Source rank ({len(src_shape)}) must match tile rank ({len(tile_shape)})")
-    if x.type.element_ty != tile.type.element_ty:
-        raise ValueError(f"Element type mismatch: source={x.type.element_ty}, tile={tile.type.element_ty}")
-
-    # Build per-dimension tile grid: how many tiles exist in each axis.
-    # With explicit strides this can represent overlapping tiles.
-    strides_eff = strides_ints if strides_ints else tile_shape
-    if len(strides_eff) != len(src_shape):
-        raise ValueError(f"strides rank ({len(strides_eff)}) must match source rank ({len(src_shape)})")
-
-    grid = []
-    for i, (src_dim, tile_dim, stride_dim) in enumerate(zip(src_shape, tile_shape, strides_eff)):
-        if tile_dim <= 0:
-            raise ValueError(f"Tile dimension {i} must be positive, got {tile_dim}")
-        if stride_dim <= 0:
-            raise ValueError(f"Stride dimension {i} must be positive, got {stride_dim}")
-        remainder = src_dim - tile_dim
-        if remainder < 0 or remainder % stride_dim != 0:
-            raise ValueError(
-                f"(src-tile) not divisible by stride at dim {i}: "
-                f"src={src_dim}, tile={tile_dim}, stride={stride_dim}")
-        grid.append(remainder // stride_dim + 1)
-
-    # Parse index: dynamic scalar tensor or static scalar/multi-dim.
-    is_dynamic = False
-    index_value = None
-    index_ir_handle = None
-
-    if isinstance(index, tl.tensor):
-        is_dynamic = True
-        index_ir_handle = index.handle
-    else:
-        index_unwrapped = index
-        try:
-            index_unwrapped = tl._unwrap_if_constexpr(index_unwrapped)
-        except Exception:
-            pass
-        try:
-            if hasattr(index_unwrapped, "value"):
-                index_unwrapped = index_unwrapped.value
-        except Exception:
-            pass
-
-        index_list = None
-        if isinstance(index_unwrapped, (tuple, list, tl.tuple)):
-            index_list = list(index_unwrapped)
-            if len(index_list) != len(src_shape):
-                raise ValueError(f"Index rank {len(index_list)} must match source rank {len(src_shape)}")
-            has_tensor = any(isinstance(v, tl.tensor) for v in index_list)
-            # ------------------------------------------------
-            # dynamic multi-dim index
-            # ------------------------------------------------
-            if has_tensor:
-                index_ir_handle = _linearize_dynamic_multidim_index(index_list, src_shape, tile_shape, _semantic,
-                                                                    strides_ints)
-                is_dynamic = True
-            # ------------------------------------------------
-            # static multi-dim index
-            # ------------------------------------------------
-            else:
-                idx = []
-                for i, v in enumerate(index_list):
-                    iv = _try_unwrap_int(v)
-                    if iv is None:
-                        raise ValueError("Multi-dim index must contain int/constexpr values")
-                    if iv < 0 or iv >= grid[i]:
-                        raise ValueError(f"Index[{i}]={iv} out of bounds for tile grid (0~{grid[i]-1})")
-                    idx.append(iv)
-                index_value = _linearize_static_multidim_index(idx, src_shape, tile_shape, strides_ints)
-        else:
-            # Path B: scalar static index -> treat as already-linearized tile id.
-            scalar_int = _try_unwrap_int(index_unwrapped)
-            if scalar_int is None:
-                raise ValueError(f"index must be int, constexpr, tuple/list of int/constexpr, "
-                                 f"or a scalar tl.tensor; got {type(index)}")
-            index_value = scalar_int
-
-    # Static index checks + optional semantic pass.
-    if not is_dynamic:
-        if index_value < 0:
-            raise ValueError("Scalar index must be non-negative")
-
-        total_tiles = 1
-        for g in grid:
-            total_tiles *= g
-        if index_value >= total_tiles:
-            raise ValueError(f"Scalar index {index_value} out of bounds for total tiles {total_tiles}")
-
-        try:
-            from .semantic import TLESemantic
-            if isinstance(_semantic, TLESemantic):
-                _semantic.analyze_insert_tile_operation(x, tile, index_value, strides_ints)
-        except ImportError:
-            pass
-
-    # Lower to IR and construct output tensor with the source tensor type.
-    try:
-        if is_dynamic:
-            index_ir = index_ir_handle
-        else:
-            index_ir = _semantic._convert_to_ir_values([index_value], require_i64=False)[0]
-        output = _semantic.builder.create_insert_tile(
-            x.handle,
-            tile.handle,
-            index_ir,
-            strides_ints or tile_shape,
-        )
-        return tl.tensor(output, x.type)
-    except Exception as e:
-        raise RuntimeError(f"Failed to create insert_tile operation: {str(e)}") from e
-
 
 def _expand_index_to_shape(index: tl.tensor, shape: Sequence[int], axis: int, _semantic) -> tl.tensor:
     idx = index
diff --git a/python/tutorials/tle/05-glu.py b/python/tutorials/tle/05-glu.py
new file mode 100644
index 0000000000..7635177fe5
--- /dev/null
+++ b/python/tutorials/tle/05-glu.py
@@ -0,0 +1,348 @@
+"""
+GLU with Triton (TLE Tutorial style) — autotune fixed
+======================================================
+
+This is test_122_v1 with the baseline BLOCK_D properly autotuned per
+(N, D, dtype). Two key fixes vs the previous attempts:
+
+  1. PRE-COMPILE pass: every BLOCK_D candidate is warmed/compiled BEFORE
+     do_bench is called for timing. This avoids the trap where the first
+     compile latency of a large BLOCK_D (like 1024) is counted into the
+     do_bench result, making it look slower than BLOCK_D=256 and forcing
+     a fallback to the initial value.
+
+  2. NO swallowed exceptions. If a BLOCK_D truly cannot compile (e.g.
+     out-of-shared-memory), the error is printed so it's visible.
+
+GLU(x) = x[:, :D] * sigmoid(x[:, D:])
+where x has shape (N, 2D).
+"""
+
+# %%
+# Setup
+# -----
+
+import argparse
+import math
+import sys
+
+import torch
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
+
+
+def _next_pow2(x: int) -> int:
+    return 1 if x <= 1 else 2 ** math.ceil(math.log2(x))
+
+
+# %%
+# Kernels (Triton baseline — test_122_v1 style: chunk + a/b pointers)
+# -------------------------------------------------------------------
+
+
+@triton.jit
+def glu_baseline(
+    a_ptr, b_ptr, out_ptr, D,
+    stride_an, stride_bn, stride_outn,
+    BLOCK_D: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_d = tl.program_id(1)
+
+    offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
+    mask_d = offs_d < D
+
+    a = tl.load(a_ptr + pid_n * stride_an + offs_d,
+                mask=mask_d, other=0.0).to(tl.float32)
+    b = tl.load(b_ptr + pid_n * stride_bn + offs_d,
+                mask=mask_d, other=0.0).to(tl.float32)
+
+    result = a * (1.0 / (1.0 + tl.exp(-b)))
+
+    tl.store(out_ptr + pid_n * stride_outn + offs_d,
+             result.to(out_ptr.dtype.element_ty), mask=mask_d)
+
+
+# %%
+# Kernels (TLE static extract_tile)
+# ---------------------------------
+
+
+@triton.jit
+def glu_extract_static(
+    x_ptr, out_ptr, D,
+    stride_xn, stride_outn,
+    D_P2:  tl.constexpr,
+    D2_P2: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+
+    offs = tl.arange(0, D2_P2)
+    mask = offs < (D * 2)
+    halo = tl.load(x_ptr + pid_n * stride_xn + offs,
+                   mask=mask, other=0.0)
+
+    a_tile = tle.extract_tile(halo, index=[0],
+                              tile_shape=[D_P2], strides=[D_P2])
+    b_tile = tle.extract_tile(halo, index=[1],
+                              tile_shape=[D_P2], strides=[D_P2])
+
+    a_f32     = a_tile.to(tl.float32)
+    b_f32     = b_tile.to(tl.float32)
+    sigmoid_b = 1.0 / (1.0 + tl.exp(-b_f32))
+    result    = a_f32 * sigmoid_b
+
+    offs_d = tl.arange(0, D_P2)
+    mask_d = offs_d < D
+    tl.store(out_ptr + pid_n * stride_outn + offs_d,
+             result.to(out_ptr.dtype.element_ty), mask=mask_d)
+
+
+# %%
+# Python wrappers
+# ---------------
+
+
+def triton_glu(x, BLOCK_D):
+    N, D2 = x.shape
+    D     = D2 // 2
+    out   = torch.empty((N, D), device=x.device, dtype=x.dtype)
+    a, b  = torch.chunk(x, 2, dim=-1)
+    grid  = (N, triton.cdiv(D, BLOCK_D))
+    glu_baseline[grid](
+        a, b, out, D,
+        a.stride(0), b.stride(0), out.stride(0),
+        BLOCK_D=BLOCK_D,
+    )
+    return out
+
+
+def tle_glu(x):
+    N, D2 = x.shape
+    D     = D2 // 2
+    out   = torch.empty((N, D), device=x.device, dtype=x.dtype)
+    d_p2  = _next_pow2(D)
+    d2_p2 = _next_pow2(D2)
+    glu_extract_static[(N,)](
+        x, out, D,
+        x.stride(0), out.stride(0),
+        D_P2=d_p2,
+        D2_P2=d2_p2,
+    )
+    return out
+
+
+def torch_glu(x):
+    return torch.nn.functional.glu(x, dim=-1)
+
+
+# %%
+# Baseline BLOCK_D autotune  ★ FIXED ★
+# ------------------------------------
+# Two-phase: (1) warm-compile all candidates, (2) timing pass.
+# Without phase 1, the first-compile time of BLOCK_D=1024 gets folded into
+# do_bench and the autotuner mis-elects BLOCK_D=256 as "fastest".
+
+
+_BLOCK_D_CANDIDATES = [64, 128, 256, 512, 1024]
+_block_d_cache: dict = {}
+
+
+def best_block_d(N: int, D: int, dtype: torch.dtype,
+                 verbose: bool = False) -> int:
+    """Pick the fastest BLOCK_D for the baseline on (N, D, dtype). Cached."""
+    key = (N, D, dtype)
+    if key in _block_d_cache:
+        return _block_d_cache[key]
+
+    x = _make_input(N, D, dtype)
+
+    # ── Phase 1: warm-compile every candidate ────────────────────────────
+    # This is the critical fix — drag every BLOCK_D through the JIT once
+    # so the timing phase below measures pure execution, not compilation.
+    valid_bds = []
+    for bd in _BLOCK_D_CANDIDATES:
+        try:
+            triton_glu(x, bd)
+            torch.cuda.synchronize()
+            valid_bds.append(bd)
+        except Exception as e:
+            if verbose:
+                print(f"  [autotune] N={N} D={D} {dtype} BLOCK_D={bd} "
+                      f"COMPILE FAIL ({type(e).__name__}: {e})")
+
+    if not valid_bds:
+        raise RuntimeError(f"no BLOCK_D worked for N={N} D={D} {dtype}")
+
+    # ── Phase 2: time each candidate fairly ──────────────────────────────
+    # NOTE: do_bench's return type depends on `quantiles`:
+    #   * quantiles=None  -> float (median ms)
+    #   * quantiles=[..]  -> list of floats  (some triton builds return
+    #                        a float when len==1; do not index)
+    # We use the no-quantiles form and let it return a single float.
+    best_bd, best_ms = valid_bds[0], float("inf")
+    for bd in valid_bds:
+        ms = triton.testing.do_bench(
+            lambda b=bd: triton_glu(x, b),
+            warmup=25, rep=100)
+        if verbose:
+            print(f"  [autotune] N={N} D={D} {dtype} BLOCK_D={bd:>4}  "
+                  f"{ms*1000:.3f} us")
+        if ms < best_ms:
+            best_ms, best_bd = ms, bd
+
+    if verbose:
+        print(f"  [autotune] -> picked BLOCK_D={best_bd}  "
+              f"({best_ms*1000:.3f} us)")
+
+    _block_d_cache[key] = best_bd
+    return best_bd
+
+
+# %%
+# Correctness check
+# -----------------
+
+
+def _get_dtype(name: str) -> torch.dtype:
+    name = name.lower()
+    if name in ("float16", "fp16", "half"):
+        return torch.float16
+    if name in ("float32", "fp32"):
+        return torch.float32
+    raise ValueError(f"unsupported dtype: {name}")
+
+
+def _make_input(N, D, dtype):
+    torch.manual_seed(0)
+    return torch.randn(N, 2 * D, device=DEVICE, dtype=dtype)
+
+
+def _check_static_feasibility(D: int) -> None:
+    d_p2 = _next_pow2(D)
+    if d_p2 != D:
+        raise ValueError(
+            f"D={D} is not pow2 (next_pow2={d_p2}); static extract_tile "
+            f"requires D == next_pow2(D)")
+    if d_p2 < 256:
+        raise ValueError(
+            f"D_P2={d_p2} < 256; cannot guarantee divisibility by "
+            f"ctaTile(<=128)")
+
+
+def run_correctness(N, D, dtype, BLOCK_D=None):
+    _check_static_feasibility(D)
+    x = _make_input(N, D, dtype)
+    if BLOCK_D is None:
+        BLOCK_D = best_block_d(N, D, dtype)
+
+    ref      = torch_glu(x.float()).to(dtype)
+    y_triton = triton_glu(x, BLOCK_D)
+    y_tle    = tle_glu(x)
+
+    atol = 1e-2 if dtype == torch.float16 else 1e-4
+    rtol = 1e-3 if dtype == torch.float16 else 1e-4
+    torch.testing.assert_close(y_triton.float(), ref.float(),
+                               rtol=rtol, atol=atol)
+    torch.testing.assert_close(y_tle.float(), ref.float(),
+                               rtol=rtol, atol=atol)
+    dt = "fp16" if dtype == torch.float16 else "fp32"
+    print(f"  [OK]  N={N} D={D} {dt}  BLOCK_D={BLOCK_D}")
+
+
+if "--only_unit_test" in sys.argv:
+    for _dtype in (torch.float32, torch.float16):
+        for _N, _D in [(1024, 256), (4096, 1024), (8192, 4096)]:
+            run_correctness(_N, _D, _dtype)
+    sys.exit(0)
+
+
+# %%
+# Benchmark (test_122_v1 perf_report style, three providers)
+# ----------------------------------------------------------
+
+
+_BENCH_PROVIDERS = ["triton", "tle", "torch"]
+_BENCH_NAMES     = ["Triton (baseline, autotuned)",
+                    "TLE (extract_tile static)",
+                    "Torch (F.glu)"]
+_BENCH_STYLES    = [("blue", "-"), ("orange", "-"), ("green", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["D"],
+        x_vals=[256, 512, 1024, 2048, 4096],
+        x_log=False,
+        line_arg="provider",
+        line_vals=_BENCH_PROVIDERS,
+        line_names=_BENCH_NAMES,
+        styles=_BENCH_STYLES,
+        ylabel="ms",
+        plot_name="tle-glu-performance",
+        args={"N": 4096},
+    ))
+def benchmark(N, D, provider, dtype_str):
+    dtype = _get_dtype(dtype_str)
+    _check_static_feasibility(D)
+
+    x = _make_input(N, D, dtype)
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "torch":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: torch_glu(x), quantiles=quantiles)
+    elif provider == "tle":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: tle_glu(x), quantiles=quantiles)
+    else:  # triton baseline with autotuned BLOCK_D
+        BLOCK_D = best_block_d(N, D, dtype)
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: triton_glu(x, BLOCK_D), quantiles=quantiles)
+
+    return ms, max_ms, min_ms
+
+
+# %%
+# Main
+# ----
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--N", type=int, default=4096,
+                        help="rows for correctness")
+    parser.add_argument("--D", type=int, default=1024,
+                        help="last-dim half-size for correctness "
+                             "(pow2, >=256)")
+    parser.add_argument("--dtype", type=str, default="float32",
+                        choices=["float16", "float32", "fp16", "fp32"])
+    parser.add_argument("--show_plots", action="store_true",
+                        help="show plots in benchmark")
+    parser.add_argument("--verbose_autotune", action="store_true",
+                        help="print autotune timing for each BLOCK_D")
+    args = parser.parse_args(argv)
+
+    dtype = _get_dtype(args.dtype)
+
+    # Optionally pre-warm autotune for the benchmark sizes with verbose
+    # output so the user can see what BLOCK_D actually got picked.
+    if args.verbose_autotune:
+        print("--- autotune (verbose) ---")
+        for D in [256, 512, 1024, 2048, 4096]:
+            best_block_d(4096, D, dtype, verbose=True)
+        print()
+
+    print("--- correctness ---")
+    run_correctness(args.N, args.D, dtype)
+    print()
+
+    benchmark.run(print_data=True, show_plots=args.show_plots,
+                  dtype_str=args.dtype.replace("float", "fp"))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/tutorials/tle/06-2D_Depthwise_Conv.py b/python/tutorials/tle/06-2D_Depthwise_Conv.py
new file mode 100644
index 0000000000..b4de39525b
--- /dev/null
+++ b/python/tutorials/tle/06-2D_Depthwise_Conv.py
@@ -0,0 +1,314 @@
+"""
+2D Depthwise Conv with Triton (TLE Tutorial style)
+==================================================
+
+This tutorial implements a 2D depthwise convolution and compares a Triton
+baseline kernel against a TLE kernel that uses ``tle.extract_tile`` for
+static halo reuse, against PyTorch's ``F.conv2d`` reference.
+
+Notes
+-----
+- Layout is HWC (channels-last) for both input and weight.
+- Kernel size K is odd; padding = 0 (valid conv). Output (H-K+1, W-K+1, C).
+- TLE uses a register-only halo path: load a (HALO_H_P2, HALO_W_P2, TC)
+  halo block once into registers, then ``tle.extract_tile`` with constexpr
+  indices for K*K reuse without any shared memory or barriers.
+- The benchmark focuses on K=5 with a 4x4 spatial tile.
+"""
+
+# %%
+# Setup
+# -----
+
+import argparse
+import math
+import sys
+
+import torch
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
+
+
+def _next_pow2(x: int) -> int:
+    return 1 if x <= 1 else 2 ** math.ceil(math.log2(x))
+
+
+# %%
+# Kernels (Triton baseline)
+# -------------------------
+
+
+@triton.jit
+def conv2d_baseline(
+    inp_ptr, wgt_ptr, out_ptr,
+    H, W, C, OH, OW,
+    KH: tl.constexpr, KW: tl.constexpr,
+    TILE_OH: tl.constexpr, TILE_OW: tl.constexpr, TILE_C: tl.constexpr,
+):
+    pid_oh = tl.program_id(0)
+    pid_ow = tl.program_id(1)
+    pid_c  = tl.program_id(2)
+    oh0 = pid_oh * TILE_OH
+    ow0 = pid_ow * TILE_OW
+    c0  = pid_c  * TILE_C
+
+    offs_c  = c0  + tl.arange(0, TILE_C)
+    offs_oh = oh0 + tl.arange(0, TILE_OH)
+    offs_ow = ow0 + tl.arange(0, TILE_OW)
+    c_mask  = offs_c < C
+
+    acc = tl.zeros((TILE_OH, TILE_OW, TILE_C), dtype=tl.float32)
+
+    for kh in tl.static_range(KH):
+        for kw in tl.static_range(KW):
+            w_ptrs = wgt_ptr + (kh * KW + kw) * C + offs_c
+            w = tl.load(w_ptrs, mask=c_mask, other=0.0)
+            ih = offs_oh + kh
+            iw = offs_ow + kw
+            ih_ok = (ih >= 0) & (ih < H)
+            iw_ok = (iw >= 0) & (iw < W)
+            inp_ptrs = (inp_ptr
+                        + ih[:, None, None] * (W * C)
+                        + iw[None, :, None] * C
+                        + offs_c[None, None, :])
+            mask = ih_ok[:, None, None] & iw_ok[None, :, None] & c_mask[None, None, :]
+            x = tl.load(inp_ptrs, mask=mask, other=0.0)
+            acc += x * w[None, None, :]
+
+    out_ptrs = (out_ptr
+                + offs_oh[:, None, None] * (OW * C)
+                + offs_ow[None, :, None] * C
+                + offs_c[None, None, :])
+    out_mask = ((offs_oh[:, None, None] < OH)
+                & (offs_ow[None, :, None] < OW)
+                & c_mask[None, None, :])
+    tl.store(out_ptrs, acc, mask=out_mask)
+
+
+# %%
+# Kernels (TLE static extract_tile)
+# ---------------------------------
+
+
+@triton.jit
+def conv2d_extract_static(
+    inp_ptr, wgt_ptr, out_ptr,
+    H, W, C, OH, OW,
+    KH: tl.constexpr, KW: tl.constexpr,
+    TILE_OH: tl.constexpr, TILE_OW: tl.constexpr, TILE_C: tl.constexpr,
+    HALO_H_P2: tl.constexpr, HALO_W_P2: tl.constexpr,
+    HALO_H:    tl.constexpr, HALO_W:    tl.constexpr,
+):
+    pid_oh = tl.program_id(0)
+    pid_ow = tl.program_id(1)
+    pid_c  = tl.program_id(2)
+    oh0 = pid_oh * TILE_OH
+    ow0 = pid_ow * TILE_OW
+    c0  = pid_c  * TILE_C
+
+    offs_c     = c0  + tl.arange(0, TILE_C)
+    c_mask     = offs_c < C
+    offs_ih_p2 = oh0 + tl.arange(0, HALO_H_P2)
+    offs_iw_p2 = ow0 + tl.arange(0, HALO_W_P2)
+
+    ih_ok = (offs_ih_p2 < oh0 + HALO_H) & (offs_ih_p2 >= 0) & (offs_ih_p2 < H)
+    iw_ok = (offs_iw_p2 < ow0 + HALO_W) & (offs_iw_p2 >= 0) & (offs_iw_p2 < W)
+
+    halo_ptrs = (inp_ptr
+                 + offs_ih_p2[:, None, None] * (W * C)
+                 + offs_iw_p2[None, :, None] * C
+                 + offs_c[None, None, :])
+    halo_mask = ih_ok[:, None, None] & iw_ok[None, :, None] & c_mask[None, None, :]
+    halo = tl.load(halo_ptrs, mask=halo_mask, other=0.0)
+
+    acc     = tl.zeros((TILE_OH, TILE_OW, TILE_C), dtype=tl.float32)
+    offs_oh = oh0 + tl.arange(0, TILE_OH)
+    offs_ow = ow0 + tl.arange(0, TILE_OW)
+
+    for kh in tl.static_range(KH):
+        for kw in tl.static_range(KW):
+            w_ptrs = wgt_ptr + (kh * KW + kw) * C + offs_c
+            w = tl.load(w_ptrs, mask=c_mask, other=0.0)
+            patch = tle.extract_tile(
+                halo,
+                index=[kh, kw, 0],
+                tile_shape=[TILE_OH, TILE_OW, TILE_C],
+                strides=[1, 1, 1],
+            )
+            acc += patch * w[None, None, :]
+
+    out_ptrs = (out_ptr
+                + offs_oh[:, None, None] * (OW * C)
+                + offs_ow[None, :, None] * C
+                + offs_c[None, None, :])
+    out_mask = ((offs_oh[:, None, None] < OH)
+                & (offs_ow[None, :, None] < OW)
+                & c_mask[None, None, :])
+    tl.store(out_ptrs, acc, mask=out_mask)
+
+
+# %%
+# Python wrappers
+# ---------------
+
+
+def _halo_dims(KH, KW, TOH, TOW):
+    hh = TOH + KH - 1
+    hw = TOW + KW - 1
+    return hh, hw, _next_pow2(hh), _next_pow2(hw)
+
+
+def triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC):
+    H, W, C = inp.shape
+    OH, OW = H - KH + 1, W - KW + 1
+    out = torch.empty((OH, OW, C), device=inp.device, dtype=inp.dtype)
+    grid = (math.ceil(OH / TOH), math.ceil(OW / TOW), math.ceil(C / TC))
+    conv2d_baseline[grid](
+        inp, wgt, out, H, W, C, OH, OW,
+        KH=KH, KW=KW, TILE_OH=TOH, TILE_OW=TOW, TILE_C=TC,
+    )
+    return out
+
+
+def tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC):
+    H, W, C = inp.shape
+    OH, OW = H - KH + 1, W - KW + 1
+    out = torch.empty((OH, OW, C), device=inp.device, dtype=inp.dtype)
+    hh, hw, hhp2, hwp2 = _halo_dims(KH, KW, TOH, TOW)
+    grid = (math.ceil(OH / TOH), math.ceil(OW / TOW), math.ceil(C / TC))
+    conv2d_extract_static[grid](
+        inp, wgt, out, H, W, C, OH, OW,
+        KH=KH, KW=KW, TILE_OH=TOH, TILE_OW=TOW, TILE_C=TC,
+        HALO_H_P2=hhp2, HALO_W_P2=hwp2,
+        HALO_H=hh, HALO_W=hw,
+    )
+    return out
+
+
+def torch_dwconv(inp, wgt, KH, KW):
+    C = inp.shape[2]
+    x = inp.permute(2, 0, 1).unsqueeze(0).float()
+    w = wgt.reshape(KH * KW, C).T.reshape(C, 1, KH, KW).float()
+    y = torch.nn.functional.conv2d(x, w, groups=C, padding=0)
+    return y.squeeze(0).permute(1, 2, 0).to(inp.dtype)
+
+
+# %%
+# Correctness check
+# -----------------
+
+
+def _get_dtype(name: str) -> torch.dtype:
+    name = name.lower()
+    if name in ("float16", "fp16", "half"):
+        return torch.float16
+    if name in ("float32", "fp32"):
+        return torch.float32
+    raise ValueError(f"unsupported dtype: {name}")
+
+
+def _make_input(H, W, C, KH, KW, dtype):
+    torch.manual_seed(0)
+    inp = torch.randn(H, W, C, device=DEVICE, dtype=dtype)
+    wgt = torch.randn(KH, KW, C, device=DEVICE, dtype=dtype)
+    return inp, wgt
+
+
+def run_correctness(H, W, C, KH, KW, TOH, TOW, TC, dtype):
+    inp, wgt = _make_input(H, W, C, KH, KW, dtype)
+    ref = torch_dwconv(inp, wgt, KH, KW)
+    y_triton = triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC)
+    y_tle    = tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC)
+    atol = 1e-2 if dtype == torch.float16 else 1e-3
+    torch.testing.assert_close(y_triton.float(), ref.float(), rtol=atol, atol=atol)
+    torch.testing.assert_close(y_tle.float(),    ref.float(), rtol=atol, atol=atol)
+    print("Correctness check passed (triton/tle).")
+
+
+if "--only_unit_test" in sys.argv:
+    _args = argparse.Namespace(H=128, W=128, C=64, KH=5, KW=5,
+                               TOH=4, TOW=4, TC=64, dtype="float32")
+    _dtype = _get_dtype(_args.dtype)
+    run_correctness(_args.H, _args.W, _args.C, _args.KH, _args.KW,
+                    _args.TOH, _args.TOW, _args.TC, _dtype)
+    sys.exit(0)
+
+
+# %%
+# Benchmark
+# ---------
+
+
+# Mirrors the FFT tutorial's "vary one dim, sweep providers" shape:
+# x-axis is spatial size H=W, y-axis is ms for {triton, tle, torch}.
+_BENCH_PROVIDERS = ["triton", "tle", "torch"]
+_BENCH_NAMES = ["Triton", "TLE", "Torch"]
+_BENCH_STYLES = [("blue", "-"), ("orange", "-"), ("green", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["H"],
+        x_vals=[112, 128, 256, 512],
+        x_log=False,
+        line_arg="provider",
+        line_vals=_BENCH_PROVIDERS,
+        line_names=_BENCH_NAMES,
+        styles=_BENCH_STYLES,
+        ylabel="ms",
+        plot_name="tle-dwconv-performance",
+        args={"KH": 5, "KW": 5, "TOH": 4, "TOW": 4, "TC": 64, "C": 64},
+    ))
+def benchmark(H, C, KH, KW, TOH, TOW, TC, provider, dtype):
+    inp, wgt = _make_input(H, H, C, KH, KW, dtype)
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "torch":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: torch_dwconv(inp, wgt, KH, KW), quantiles=quantiles)
+    elif provider == "tle":
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC), quantiles=quantiles)
+    else:
+        ms, min_ms, max_ms = triton.testing.do_bench(
+            lambda: triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC), quantiles=quantiles)
+
+    return ms, max_ms, min_ms
+
+
+# %%
+# Main
+# ----
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--H", type=int, default=128, help="spatial H for correctness")
+    parser.add_argument("--W", type=int, default=128, help="spatial W for correctness")
+    parser.add_argument("--C", type=int, default=64,  help="channels")
+    parser.add_argument("--KH", type=int, default=5)
+    parser.add_argument("--KW", type=int, default=5)
+    parser.add_argument("--TOH", type=int, default=4)
+    parser.add_argument("--TOW", type=int, default=4)
+    parser.add_argument("--TC",  type=int, default=64)
+    parser.add_argument("--dtype", type=str, default="float32",
+                        choices=["float16", "float32", "fp16", "fp32"])
+    parser.add_argument("--show_plots", action="store_true",
+                        help="show plots in benchmark")
+    args = parser.parse_args(argv)
+
+    dtype = _get_dtype(args.dtype)
+
+    check_H = min(args.H, 256)
+    check_W = min(args.W, 256)
+    run_correctness(check_H, check_W, args.C, args.KH, args.KW,
+                    args.TOH, args.TOW, args.TC, dtype)
+
+    benchmark.run(print_data=True, show_plots=args.show_plots, dtype=dtype)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/tutorials/tle/07-causal-conv1d.py b/python/tutorials/tle/07-causal-conv1d.py
new file mode 100644
index 0000000000..0cc7633ae0
--- /dev/null
+++ b/python/tutorials/tle/07-causal-conv1d.py
@@ -0,0 +1,1767 @@
+"""
+Causal Conv1d with TLE extract_tile
+====================================
+
+This tutorial compares a Triton baseline causal_conv1d kernel against a TLE
+optimized version that uses ``tle.extract_tile`` for register-only state reuse.
+
+Usage
+-----
+::
+
+    # correctness check only (default)
+    python python/tutorials/tle/07-causal-conv1d.py
+
+    # correctness + benchmark tables with speedup
+    python python/tutorials/tle/07-causal-conv1d.py --benchmark
+    
+Both ``causal_conv1d_fn`` (varlen forward) and ``causal_conv1d_update``
+(step/update) are included.
+"""
+
+# %%
+# Setup
+# -----
+
+import argparse
+import math
+import sys
+
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
+PAD_SLOT_ID = -1
+
+
+# %%
+# Kernels (baseline v1)
+# ---------------------
+
+@triton.jit()
+def _causal_conv1d_fwd_kernel_v1(
+    x_ptr,
+    w_ptr,
+    bias_ptr,
+    initial_states_ptr,
+    cache_indices_ptr,
+    has_initial_states_ptr,
+    query_start_loc_ptr,
+    batch_ptr,
+    token_chunk_offset_ptr,
+    block_idx_first_scheduled_token,
+    block_idx_last_scheduled_token,
+    initial_state_idx,
+    num_computed_tokens,
+    o_ptr,
+    dim: tl.constexpr,
+    seqlen: tl.int32,
+    num_cache_lines: tl.constexpr,
+    stride_x_dim: tl.constexpr,
+    stride_x_token: tl.constexpr,
+    stride_w_dim: tl.constexpr,
+    stride_w_width: tl.constexpr,
+    stride_istate_seq: tl.constexpr,
+    stride_istate_dim: tl.constexpr,
+    stride_istate_token: tl.constexpr,
+    stride_cache_indices: tl.constexpr,
+    stride_o_dim: tl.constexpr,
+    stride_o_token: tl.constexpr,
+    stride_block_m: tl.constexpr,
+    pad_slot_id: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    KERNEL_WIDTH: tl.constexpr,
+    SILU_ACTIVATION: tl.constexpr,
+    IS_APC_ENABLED: tl.constexpr,
+    USE_PAD_SLOT: tl.constexpr,
+    NP2_STATELEN: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    conv_states_ptr = initial_states_ptr
+    conv_state_indices_ptr = cache_indices_ptr
+    stride_conv_state_seq = stride_istate_seq
+    stride_conv_state_dim = stride_istate_dim
+    stride_conv_state_tok = stride_istate_token
+    state_len = KERNEL_WIDTH - 1
+
+    idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64)
+    chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0))
+
+    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    if idx_seq == pad_slot_id:
+        return
+
+    sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
+    sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1)
+    seqlen = sequence_end_index - sequence_start_index
+
+    B_size: tl.constexpr = stride_block_m * BLOCK_M
+
+    if IS_APC_ENABLED:
+        current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq)
+        current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
+        sequence_completed_index = tl.load(num_computed_tokens + idx_seq)
+
+        sequence_completed_offset_token = sequence_completed_index % B_size
+        seq_completed_offset = B_size - sequence_completed_offset_token
+        seq_end_offset = (seqlen - seq_completed_offset) % B_size
+        last_full_block_token_index = sequence_end_index - seq_end_offset
+        if seq_end_offset == 0:
+            last_full_block_token_index = last_full_block_token_index - B_size
+
+        n_block_to_fill = current_last_index - current_first_index
+
+        conv_state_init_index = tl.load(initial_state_idx + idx_seq)
+    else:
+        n_block_to_fill = 0
+        current_last_index = 0
+        conv_state_init_index = 0
+        current_first_index = 0
+        last_full_block_token_index = 0
+
+    token_offset = BLOCK_M * chunk_offset
+    segment_len = min(BLOCK_M, seqlen - token_offset)
+
+    x_base = (
+        x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
+    )
+
+    conv_states_input_coord = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
+    ).to(tl.int64)
+
+    if USE_PAD_SLOT:
+        if conv_states_input_coord == pad_slot_id:
+            return
+
+    conv_states_base = (
+        conv_states_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )
+
+    w_base = w_ptr + (idx_feats * stride_w_dim)
+
+    if chunk_offset == 0:
+        load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1)
+        if load_init_state:
+            prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok
+            mask_w = idx_feats < dim
+            if KERNEL_WIDTH == 2:
+                col0 = tl.load(prior_tokens, mask_w, 0.0)
+            if KERNEL_WIDTH == 3:
+                col1 = tl.load(prior_tokens, mask_w, 0.0)
+                col0 = tl.load(prior_tokens - 1 * stride_conv_state_tok, mask_w, 0.0)
+            if KERNEL_WIDTH == 4:
+                col2 = tl.load(prior_tokens, mask_w, 0.0)
+                col1 = tl.load(prior_tokens - 1 * stride_conv_state_tok, mask_w, 0.0)
+                col0 = tl.load(prior_tokens - 2 * stride_conv_state_tok, mask_w, 0.0)
+            if KERNEL_WIDTH == 5:
+                col3 = tl.load(prior_tokens, mask_w, 0.0)
+                col2 = tl.load(prior_tokens - 1 * stride_conv_state_tok, mask_w, 0.0)
+                col1 = tl.load(prior_tokens - 2 * stride_conv_state_tok, mask_w, 0.0)
+                col0 = tl.load(prior_tokens - 3 * stride_conv_state_tok, mask_w, 0.0)
+        else:
+            if KERNEL_WIDTH >= 2:
+                col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 3:
+                col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 4:
+                col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 5:
+                col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+
+        if state_len <= seqlen:
+            idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN)
+            x_ptrs = (
+                x_ptr
+                + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
+                + (idx_feats * stride_x_dim)[None, :]
+            )
+            mask_x = (
+                (idx_tokens_last >= 0)[:, None]
+                & (idx_tokens_last < seqlen)[:, None]
+                & (idx_feats < dim)[None, :]
+            )
+            loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+            idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+
+            conv_states_output_coord = tl.load(
+                conv_state_indices_ptr
+                + idx_seq * stride_cache_indices
+                + current_last_index
+            ).to(tl.int64)
+
+            conv_states_ptrs_target = (
+                conv_states_ptr
+                + (conv_states_output_coord * stride_conv_state_seq)
+                + (idx_feats * stride_conv_state_dim)
+            )[None, :] + (
+                idx_tokens_conv * stride_conv_state_tok
+            )[:, None]
+
+            mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+            tl.debug_barrier()
+            tl.store(conv_states_ptrs_target, loaded_x, mask)
+
+        else:
+            if load_init_state:
+                idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+
+                conv_states_ptrs_source = (
+                    conv_states_ptr
+                    + (conv_states_input_coord * stride_conv_state_seq)
+                    + (idx_feats * stride_conv_state_dim)[None, :]
+                    + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None]
+                )
+                mask = (
+                    (conv_states_input_coord < num_cache_lines)
+                    & ((idx_tokens_conv + seqlen) < state_len)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
+
+                VAL = state_len - seqlen
+                x_ptrs = (
+                    x_base[None, :]
+                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
+                )
+                mask_x = (
+                    (idx_tokens_conv - VAL >= 0)[:, None]
+                    & (idx_tokens_conv - VAL < seqlen)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+
+                tl.debug_barrier()
+                new_conv_state = tl.where(mask, conv_state, loaded_x)
+
+                conv_states_ptrs_target = (
+                    conv_states_base
+                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                )
+                mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+                tl.store(conv_states_ptrs_target, new_conv_state, mask)
+
+            else:
+                idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+                VAL = state_len - seqlen
+                x_ptrs = (
+                    x_base[None, :]
+                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
+                )
+                mask_x = (
+                    (idx_tokens_conv - VAL >= 0)[:, None]
+                    & (idx_tokens_conv - VAL < seqlen)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
+
+                conv_states_ptrs_target = (
+                    conv_states_base
+                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                )
+                mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+                tl.store(conv_states_ptrs_target, new_conv_state, mask)
+
+    else:
+        load_init_state = True
+        prior_tokens = x_base + (token_offset - 1) * stride_x_token
+        mask_w = idx_feats < dim
+        if KERNEL_WIDTH == 2:
+            col0 = tl.load(prior_tokens, mask_w, 0.0, cache_modifier=".ca")
+        if KERNEL_WIDTH == 3:
+            col1 = tl.load(prior_tokens, mask_w, 0.0, cache_modifier=".ca")
+            col0 = tl.load(prior_tokens - 1 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
+        if KERNEL_WIDTH == 4:
+            col2 = tl.load(prior_tokens, mask_w, 0.0, cache_modifier=".ca")
+            col1 = tl.load(prior_tokens - 1 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
+            col0 = tl.load(prior_tokens - 2 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
+        if KERNEL_WIDTH == 5:
+            col3 = tl.load(prior_tokens, mask_w, 0.0, cache_modifier=".ca")
+            col2 = tl.load(prior_tokens - 1 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
+            col1 = tl.load(prior_tokens - 2 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
+            col0 = tl.load(prior_tokens - 3 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
+
+        if (chunk_offset - 1) < n_block_to_fill:
+            idx_tokens_last = (
+                last_full_block_token_index
+                - (n_block_to_fill - chunk_offset) * B_size
+                - state_len
+            ) + tl.arange(0, NP2_STATELEN)
+            x_ptrs = (
+                x_ptr
+                + (idx_tokens_last * stride_x_token)[:, None]
+                + (idx_feats * stride_x_dim)[None, :]
+            )
+            mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[None, :]
+            loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+            idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+
+            conv_states_output_coord = tl.load(
+                conv_state_indices_ptr
+                + idx_seq * stride_cache_indices
+                + current_first_index
+                + (chunk_offset - 1)
+            ).to(tl.int64)
+
+            conv_states_ptrs_target = (
+                conv_states_ptr
+                + (conv_states_output_coord * stride_conv_state_seq)
+                + (idx_feats * stride_conv_state_dim)
+            )[None, :] + (
+                idx_tokens_conv * stride_conv_state_tok
+            )[:, None]
+
+            mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+            tl.debug_barrier()
+            tl.store(conv_states_ptrs_target, loaded_x, mask)
+
+    if HAS_BIAS:
+        bias = bias_ptr + idx_feats
+        mask_bias = idx_feats < dim
+        acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
+    else:
+        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    x_base_1d = x_base + token_offset * stride_x_token
+
+    mask_w = idx_feats < dim
+    if KERNEL_WIDTH >= 2:
+        w_ptrs = w_base + (0 * stride_w_width)
+        w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
+        w_ptrs = w_base + (1 * stride_w_width)
+        w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 3:
+        w_ptrs = w_base + (2 * stride_w_width)
+        w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 4:
+        w_ptrs = w_base + (3 * stride_w_width)
+        w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
+    mask_x_1d = idx_feats < dim
+    for idx_token in range(segment_len):
+        acc = acc_preload
+
+        matrix_w = w_col0
+        matrix_x = col0
+        for j in tl.static_range(KERNEL_WIDTH):
+            if KERNEL_WIDTH == 2:
+                if j == 1:
+                    matrix_w = w_col1
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 3:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 4:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+
+            acc += matrix_x * matrix_w
+
+        if KERNEL_WIDTH == 2:
+            col0 = matrix_x
+        elif KERNEL_WIDTH == 3:
+            col0 = col1
+            col1 = matrix_x
+        elif KERNEL_WIDTH == 4:
+            col0 = col1
+            col1 = col2
+            col2 = matrix_x
+
+        if SILU_ACTIVATION:
+            acc = acc / (1 + tl.exp(-acc))
+        mask_1d = (idx_token < segment_len) & (idx_feats < dim)
+        o_ptrs = (
+            o_ptr
+            + (sequence_start_index + token_offset + idx_token) * stride_o_token
+            + (idx_feats * stride_o_dim)
+        )
+        tl.store(o_ptrs, acc, mask=mask_1d)
+
+
+@triton.jit()
+def _causal_conv1d_update_kernel_v1(
+    x_ptr,
+    w_ptr,
+    bias_ptr,
+    conv_state_ptr,
+    conv_state_indices_ptr,
+    num_accepted_tokens_ptr,
+    query_start_loc_ptr,
+    block_idx_last_scheduled_token,
+    initial_state_idx,
+    o_ptr,
+    batch: int,
+    dim: tl.constexpr,
+    seqlen: tl.constexpr,
+    state_len: tl.constexpr,
+    num_cache_lines: tl.constexpr,
+    stride_x_seq: tl.constexpr,
+    stride_x_dim: tl.constexpr,
+    stride_x_token: tl.constexpr,
+    stride_w_dim: tl.constexpr,
+    stride_w_width: tl.constexpr,
+    stride_conv_state_seq: tl.constexpr,
+    stride_conv_state_dim: tl.constexpr,
+    stride_conv_state_tok: tl.constexpr,
+    stride_state_indices: tl.constexpr,
+    stride_o_seq: tl.constexpr,
+    stride_o_dim: tl.constexpr,
+    stride_o_token: tl.constexpr,
+    pad_slot_id: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    KERNEL_WIDTH: tl.constexpr,
+    SILU_ACTIVATION: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_APC_ENABLED: tl.constexpr,
+    IS_SPEC_DECODING: tl.constexpr,
+    NP2_STATELEN: tl.constexpr,
+    USE_PAD_SLOT: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    idx_seq = tl.program_id(0)
+    if idx_seq >= batch:
+        return
+
+    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    if IS_APC_ENABLED:
+        conv_state_init = tl.load(initial_state_idx + idx_seq)
+        current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
+    else:
+        conv_state_init = 0
+        current_last_index = 0
+
+    conv_states_input_coord = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
+    ).to(tl.int64)
+
+    if USE_PAD_SLOT:
+        if conv_states_input_coord == pad_slot_id:
+            return
+
+    if IS_VARLEN:
+        query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64)
+        query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64)
+        state_len = state_len - (seqlen - (query_end_index - query_start_index))
+        seqlen = query_end_index - query_start_index
+        x_offset = query_start_index * stride_x_token
+        o_offset = query_start_index * stride_o_token
+    else:
+        query_start_index = idx_seq * seqlen
+        query_end_index = query_start_index + seqlen
+        x_offset = idx_seq * stride_x_seq
+        o_offset = idx_seq * stride_o_seq
+
+    if query_start_index == query_end_index:
+        return
+
+    if IS_SPEC_DECODING:
+        conv_state_token_offset = (
+            tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
+        )
+    else:
+        conv_state_token_offset = 0
+
+    conv_states_base = (
+        conv_state_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )
+    mask_w = idx_feats < dim
+
+    prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok
+    if KERNEL_WIDTH >= 2:
+        col0 = tl.load(prior_tokens, mask_w, 0.0)
+    if KERNEL_WIDTH >= 3:
+        col1 = tl.load(prior_tokens + 1 * stride_conv_state_tok, mask_w, 0.0)
+    if KERNEL_WIDTH >= 4:
+        col2 = tl.load(prior_tokens + 2 * stride_conv_state_tok, mask_w, 0.0)
+    if KERNEL_WIDTH >= 5:
+        col3 = tl.load(prior_tokens + 3 * stride_conv_state_tok, mask_w, 0.0)
+    if KERNEL_WIDTH >= 6:
+        col4 = tl.load(prior_tokens + 4 * stride_conv_state_tok, mask_w, 0.0)
+
+    idx_tokens = tl.arange(0, NP2_STATELEN)
+
+    conv_state_ptrs_source = (
+        conv_state_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + conv_state_token_offset * stride_conv_state_tok
+        + (idx_feats * stride_conv_state_dim)[None, :]
+        + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[
+            :, None
+        ]
+    )
+    mask = (
+        (conv_states_input_coord < num_cache_lines)
+        & ((idx_tokens + seqlen) < state_len)[:, None]
+        & (idx_feats < dim)[None, :]
+    )
+    conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
+
+    VAL = state_len - seqlen
+    x_base = x_ptr + x_offset + (idx_feats * stride_x_dim)
+
+    x_ptrs = (
+        x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None]
+    )
+    mask_x = (
+        (idx_tokens - VAL >= 0)[:, None]
+        & (idx_tokens - VAL < seqlen)[:, None]
+        & (idx_feats < dim)[None, :]
+    )
+    loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+    tl.debug_barrier()
+
+    new_conv_state = tl.where(mask, conv_state, loaded_x)
+
+    conv_states_offset = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index
+    ).to(tl.int64)
+    conv_state_ptrs_target = (
+        conv_state_ptr
+        + (conv_states_offset * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )[None, :] + (
+        idx_tokens * stride_conv_state_tok
+    )[:, None]
+    mask_out = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
+    tl.store(conv_state_ptrs_target, new_conv_state, mask_out)
+
+    if HAS_BIAS:
+        bias = bias_ptr + idx_feats
+        mask_bias = idx_feats < dim
+        acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
+    else:
+        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    w_base = w_ptr + (idx_feats * stride_w_dim)
+    mask_w = idx_feats < dim
+    if KERNEL_WIDTH >= 2:
+        w_col0 = tl.load(w_base + (0 * stride_w_width), mask_w, other=0.0)
+        w_col1 = tl.load(w_base + (1 * stride_w_width), mask_w, other=0.0)
+    if KERNEL_WIDTH >= 3:
+        w_col2 = tl.load(w_base + (2 * stride_w_width), mask_w, other=0.0)
+    if KERNEL_WIDTH >= 4:
+        w_col3 = tl.load(w_base + (3 * stride_w_width), mask_w, other=0.0)
+    if KERNEL_WIDTH >= 5:
+        w_col4 = tl.load(w_base + (4 * stride_w_width), mask_w, other=0.0)
+    if KERNEL_WIDTH >= 6:
+        w_col5 = tl.load(w_base + (5 * stride_w_width), mask_w, other=0.0)
+
+    x_base_1d = x_base
+    mask_x_1d = idx_feats < dim
+
+    for idx_token in tl.range(seqlen):
+        acc = acc_preload
+
+        matrix_w = w_col0
+        matrix_x = col0
+        for j in tl.static_range(KERNEL_WIDTH):
+            if KERNEL_WIDTH == 2:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = tl.load(x_base_1d + idx_token * stride_x_token, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 3:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = tl.load(x_base_1d + idx_token * stride_x_token, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 4:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = tl.load(x_base_1d + idx_token * stride_x_token, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 5:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = col3
+                elif j == 4:
+                    matrix_w = w_col4
+                    matrix_x = tl.load(x_base_1d + idx_token * stride_x_token, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 6:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = col3
+                elif j == 4:
+                    matrix_w = w_col4
+                    matrix_x = col4
+                elif j == 5:
+                    matrix_w = w_col5
+                    matrix_x = tl.load(x_base_1d + idx_token * stride_x_token, mask=mask_x_1d)
+
+            acc += matrix_x * matrix_w
+
+        if KERNEL_WIDTH == 2:
+            col0 = matrix_x
+        elif KERNEL_WIDTH == 3:
+            col0 = col1
+            col1 = matrix_x
+        elif KERNEL_WIDTH == 4:
+            col0 = col1
+            col1 = col2
+            col2 = matrix_x
+        elif KERNEL_WIDTH == 5:
+            col0 = col1
+            col1 = col2
+            col2 = col3
+            col3 = matrix_x
+        elif KERNEL_WIDTH == 6:
+            col0 = col1
+            col1 = col2
+            col2 = col3
+            col3 = col4
+            col4 = matrix_x
+
+        if SILU_ACTIVATION:
+            acc = acc / (1 + tl.exp(-acc))
+        mask_1d = (idx_token < seqlen) & (idx_feats < dim)
+        o_ptrs = (
+            o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim)
+        )
+        tl.store(o_ptrs, acc, mask=mask_1d)
+
+
+# %%
+# Kernels (TLE v2)
+# ----------------
+@triton.jit()
+def _causal_conv1d_fwd_kernel_v2(  # continuous batching
+    # Pointers to matrices
+    x_ptr,  # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
+    w_ptr,  # (dim, width)
+    bias_ptr,
+    initial_states_ptr,  # conv_states_ptr
+    cache_indices_ptr,  # (batch, n_blocks + padding) The second dimension contains
+    # the block indices relevant for each sequence
+    # plus potential 0-padding at the beginning and at the end
+    has_initial_states_ptr,
+    query_start_loc_ptr,
+    batch_ptr,
+    token_chunk_offset_ptr,
+    block_idx_first_scheduled_token,  # (batch,)
+    block_idx_last_scheduled_token,  # (batch,)
+    initial_state_idx,  # (batch,)
+    num_computed_tokens,  # (batch,)
+    o_ptr,  # (dim, seqlen) - actually pointing to x_ptr
+    # Matrix dimensions
+    dim: tl.constexpr,
+    seqlen: tl.int32,  # cu_seqlen
+    num_cache_lines: tl.constexpr,  # added to support vLLM larger cache lines
+    # Strides
+    stride_x_dim: tl.constexpr,  # stride to get to next feature-value,
+    stride_x_token: tl.constexpr,  # stride to get to next token (same feature-index, same sequence-index)
+    stride_w_dim: tl.constexpr,  # stride to get to next dim-axis value
+    stride_w_width: tl.constexpr,  # stride to get to next width-axis value
+    stride_istate_seq: tl.constexpr,
+    stride_istate_dim: tl.constexpr,
+    stride_istate_token: tl.constexpr,
+    stride_cache_indices: tl.constexpr,
+    stride_o_dim: tl.constexpr,
+    stride_o_token: tl.constexpr,
+    stride_block_m: tl.constexpr,  # Stride block to align divided by BLOCK_M
+    # others
+    pad_slot_id: tl.constexpr,
+    # Meta-parameters
+    HAS_BIAS: tl.constexpr,
+    KERNEL_WIDTH: tl.constexpr,
+    SILU_ACTIVATION: tl.constexpr,
+    IS_APC_ENABLED: tl.constexpr,
+    USE_PAD_SLOT: tl.constexpr,
+    NP2_STATELEN: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    conv_states_ptr = initial_states_ptr
+    conv_state_indices_ptr = cache_indices_ptr
+    stride_conv_state_seq = stride_istate_seq
+    stride_conv_state_dim = stride_istate_dim
+    stride_conv_state_tok = stride_istate_token
+    state_len = (
+        KERNEL_WIDTH - 1
+    )  # can be passed via argument if it's not the same as this value
+
+    idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64)
+    chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0))
+
+    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    if idx_seq == pad_slot_id:
+        return
+
+    sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
+    sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1)
+    seqlen = sequence_end_index - sequence_start_index
+
+    B_size: tl.constexpr = stride_block_m * BLOCK_M
+
+    if IS_APC_ENABLED:
+        current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq)
+        current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
+        sequence_completed_index = tl.load(num_computed_tokens + idx_seq)
+
+        sequence_completed_offset_token = sequence_completed_index % B_size
+        seq_completed_offset = B_size - sequence_completed_offset_token
+        seq_end_offset = (seqlen - seq_completed_offset) % B_size
+        last_full_block_token_index = sequence_end_index - seq_end_offset
+        if seq_end_offset == 0:
+            last_full_block_token_index = last_full_block_token_index - B_size
+
+        n_block_to_fill = current_last_index - current_first_index
+
+        conv_state_init_index = tl.load(initial_state_idx + idx_seq)
+    else:
+        n_block_to_fill = 0
+        current_last_index = 0
+        conv_state_init_index = 0
+        current_first_index = 0
+        last_full_block_token_index = 0
+
+    token_offset = BLOCK_M * chunk_offset
+    segment_len = min(BLOCK_M, seqlen - token_offset)
+
+    x_base = (
+        x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
+    )  # [BLOCK_N,]
+
+    conv_states_input_coord = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
+    ).to(tl.int64)
+
+    if USE_PAD_SLOT:  # noqa
+        if conv_states_input_coord == pad_slot_id:
+            return
+
+    conv_states_base = (
+        conv_states_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )  # [BLOCK_N,]
+
+    w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
+
+
+    
+    # x_tile_start = token_offset - (KERNEL_WIDTH - 1)
+    # # idx_rows = tl.arange(0, X_TILE_M_P2)
+    # idx_rows = tl.arange(0, triton.next_power_of_2(BLOCK_M + KERNEL_WIDTH - 1))
+    # offs_tok = x_tile_start + idx_rows
+    # tile_ptrs = (
+    #         x_ptr
+    #         + (sequence_start_index + offs_tok)[:, None] * stride_x_token
+    #         + (idx_feats * stride_x_dim)[None, :]
+    #     )
+    # tile_mask = (
+    #         (offs_tok >= 0)[:, None]
+    #         & (offs_tok < seqlen)[:, None]
+    #         & (idx_feats < dim)[None, :]
+    #         & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None]
+    #     )
+    # tile_data = tl.load(
+    #         tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca"
+    #     )   
+
+    if chunk_offset == 0:
+        chunk = 0
+        load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1)
+        mask_w = idx_feats < dim
+        # ----------------------------------------------------------------
+        # tile_data load at chunk_offset == 0:
+        # read from token_offset (0), without shifting left by KERNEL_WIDTH - 1
+        # ----------------------------------------------------------------
+        idx_rows = tl.arange(0, triton.next_power_of_2(BLOCK_M + KERNEL_WIDTH - 1))
+        offs_tok = token_offset + idx_rows
+        tile_ptrs = (
+            x_ptr
+            + (sequence_start_index + offs_tok)[:, None] * stride_x_token
+            + (idx_feats * stride_x_dim)[None, :]
+        )
+        tile_mask = (
+            (offs_tok >= 0)[:, None]
+            & (offs_tok < seqlen)[:, None]
+            & (idx_feats < dim)[None, :]
+            & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None]
+        )
+        tile_data = tl.load(
+            tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca"
+        )
+
+        if load_init_state:
+            # ----------------------------------------------------------------
+            # [extract_tile] Load entire conv_state window in one tile load:
+            # Before: KERNEL_WIDTH-1 separate scalar loads per column
+            # After: 1 tile load [NP2_STATELEN, BLOCK_N], then slice with extract_tile
+            # ----------------------------------------------------------------
+            # Broadcast: conv_states_base[None, :]   [1, BLOCK_N]
+            #          + idx_state[:, None] * stride  [NP2_STATELEN, 1]
+            #          => [NP2_STATELEN, BLOCK_N]
+            idx_state = tl.arange(0, NP2_STATELEN)
+            conv_state_ptrs_full = (
+                conv_states_base[None, :]
+                + idx_state[:, None] * stride_conv_state_tok
+            )  # [NP2_STATELEN, BLOCK_N]
+            mask_state_full = (
+                (idx_state[:, None] < state_len) & (idx_feats[None, :] < dim)
+            )  # bool mask [NP2_STATELEN, BLOCK_N]
+            state_tile = tl.load(
+                conv_state_ptrs_full, mask_state_full, other=0.0
+            )  # [NP2_STATELEN, BLOCK_N]
+
+            if KERNEL_WIDTH == 2:
+                col0 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+            if KERNEL_WIDTH == 3:
+                col1 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+            if KERNEL_WIDTH == 4:
+                col2 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col1 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+            if KERNEL_WIDTH == 5:
+                col3 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col2 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col1 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 4, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+        else:
+            if KERNEL_WIDTH >= 2:
+                col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 3:
+                col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 4:
+                col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+            if KERNEL_WIDTH >= 5:
+                col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+
+        # STEP 2: write back conv_state
+        if state_len <= seqlen:
+            idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN)
+            x_ptrs = (
+                x_ptr
+                + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
+                + (idx_feats * stride_x_dim)[None, :]
+            )  # [NP2_STATELEN, BLOCK_N]
+            mask_x = (
+                (idx_tokens_last >= 0)[:, None]
+                & (idx_tokens_last < seqlen)[:, None]
+                & (idx_feats < dim)[None, :]
+            )
+            loaded_x = tl.load(x_ptrs, mask_x, 0.0)  # [NP2_STATELEN, BLOCK_N]
+
+            idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+
+            conv_states_output_coord = tl.load(
+                conv_state_indices_ptr
+                + idx_seq * stride_cache_indices
+                + current_last_index
+            ).to(tl.int64)
+
+            conv_states_ptrs_target = (
+                conv_states_ptr
+                + (conv_states_output_coord * stride_conv_state_seq)
+                + (idx_feats * stride_conv_state_dim)
+            )[None, :] + (
+                idx_tokens_conv * stride_conv_state_tok
+            )[:, None]
+
+            mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+            tl.debug_barrier()
+            tl.store(conv_states_ptrs_target, loaded_x, mask)
+
+        else:
+            if load_init_state:
+                idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+                # Shift right by seqlen to read the remaining cached history
+                conv_states_ptrs_source = (
+                    conv_states_ptr
+                    + (conv_states_input_coord * stride_conv_state_seq)
+                    + (idx_feats * stride_conv_state_dim)[None, :]
+                    + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None]
+                )  # [NP2_STATELEN, BLOCK_N]
+                mask = (
+                    (conv_states_input_coord < num_cache_lines)
+                    & ((idx_tokens_conv + seqlen) < state_len)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
+
+                VAL = state_len - seqlen
+                x_ptrs = (
+                    x_base[None, :]
+                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
+                )
+                mask_x = (
+                    (idx_tokens_conv - VAL >= 0)[:, None]
+                    & (idx_tokens_conv - VAL < seqlen)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                loaded_x = tl.load(x_ptrs, mask_x, 0.0)  # [NP2_STATELEN, BLOCK_N]
+                # Merge cached history (where mask is True) with new token data
+                # (where mask is False) to form the updated conv_state.
+
+                # NOTE: insert_tile approach was explored as an alternative to
+                # tl.where + tl.debug_barrier. The loop below shows the intended
+                # pattern — build new_conv_state row by row using extract_tile /
+                # insert_tile, avoiding the compiler bug in tl.where entirely.
+                # Currently kept as a comment for future reference; the code
+                # still uses tl.where with the required debug_barrier workaround.
+
+                tl.debug_barrier()  # required: tl.where bug when result of another tl.load
+                # Where mask is True: take cached history (left-shifted).
+                # Where mask is False: take new tokens from x.
+                new_conv_state = tl.where(
+                    mask, conv_state, loaded_x
+                )  # BUG in 'tl.where' requires a barrier before this
+
+                conv_states_ptrs_target = (
+                    conv_states_base
+                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                )  # [NP2_STATELEN, BLOCK_N]
+                mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+                tl.store(conv_states_ptrs_target, new_conv_state, mask)
+
+            else:  # load_init_state == False
+                idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+                VAL = state_len - seqlen
+                x_ptrs = (
+                    x_base[None, :]
+                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
+                )
+                mask_x = (
+                    (idx_tokens_conv - VAL >= 0)[:, None]
+                    & (idx_tokens_conv - VAL < seqlen)[:, None]
+                    & (idx_feats < dim)[None, :]
+                )
+                new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
+
+                conv_states_ptrs_target = (
+                    conv_states_base
+                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                )
+                mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+                tl.store(conv_states_ptrs_target, new_conv_state, mask)
+
+    else:  # chunk_offset > 0
+        load_init_state = True
+        chunk = 1
+        x_tile_start = token_offset - (KERNEL_WIDTH - 1)
+        idx_rows = tl.arange(0, triton.next_power_of_2(BLOCK_M + KERNEL_WIDTH - 1))
+        # idx_rows = tl.arange(0, X_TILE_M_P2)
+        offs_tok = x_tile_start + idx_rows
+        tile_ptrs = (
+            x_ptr
+            + (sequence_start_index + offs_tok)[:, None] * stride_x_token
+            + (idx_feats * stride_x_dim)[None, :]
+        )
+        tile_mask = (
+            (offs_tok >= 0)[:, None]
+            & (offs_tok < seqlen)[:, None]
+            & (idx_feats < dim)[None, :]
+            & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None]
+        )
+        tile_data = tl.load(
+            tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca"
+        )
+
+        if KERNEL_WIDTH == 2:
+            col0 = tl.reshape(
+                tle.extract_tile(tile_data, index=[0, 0], tile_shape=[1, BLOCK_N]),
+                [BLOCK_N],
+            )
+        if KERNEL_WIDTH == 3:
+            col0 = tl.reshape(
+                tle.extract_tile(tile_data, index=[0, 0], tile_shape=[1, BLOCK_N]),
+                [BLOCK_N],
+            )
+            col1 = tl.reshape(
+                tle.extract_tile(tile_data, index=[1, 0], tile_shape=[1, BLOCK_N]),
+                [BLOCK_N],
+            )
+        if KERNEL_WIDTH == 4:
+            col0 = tl.reshape(
+                tle.extract_tile(tile_data, index=[0, 0], tile_shape=[1, BLOCK_N]),
+                [BLOCK_N],
+            )
+            col1 = tl.reshape(
+                tle.extract_tile(tile_data, index=[1, 0], tile_shape=[1, BLOCK_N]),
+                [BLOCK_N],
+            )
+            col2 = tl.reshape(
+                tle.extract_tile(tile_data, index=[2, 0], tile_shape=[1, BLOCK_N]),
+                [BLOCK_N],
+            )
+        if (chunk_offset - 1) < n_block_to_fill:
+            idx_tokens_last = (
+                last_full_block_token_index
+                - (n_block_to_fill - chunk_offset) * B_size
+                - state_len
+            ) + tl.arange(0, NP2_STATELEN)
+            x_ptrs = (
+                x_ptr
+                + (idx_tokens_last * stride_x_token)[:, None]
+                + (idx_feats * stride_x_dim)[None, :]
+            )
+            mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[None, :]
+            loaded_x = tl.load(x_ptrs, mask_x, 0.0)
+            idx_tokens_conv = tl.arange(0, NP2_STATELEN)
+
+            conv_states_output_coord = tl.load(
+                conv_state_indices_ptr
+                + idx_seq * stride_cache_indices
+                + current_first_index
+                + (chunk_offset - 1)
+            ).to(tl.int64)
+
+            conv_states_ptrs_target = (
+                conv_states_ptr
+                + (conv_states_output_coord * stride_conv_state_seq)
+                + (idx_feats * stride_conv_state_dim)
+            )[None, :] + (
+                idx_tokens_conv * stride_conv_state_tok
+            )[:, None]
+
+            mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
+            tl.debug_barrier()
+            tl.store(conv_states_ptrs_target, loaded_x, mask)
+
+    if HAS_BIAS:
+        bias = bias_ptr + idx_feats
+        mask_bias = idx_feats < dim
+        acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
+    else:
+        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    mask_w = idx_feats < dim
+    if KERNEL_WIDTH >= 2:
+        w_ptrs = w_base + (0 * stride_w_width)
+        w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
+        w_ptrs = w_base + (1 * stride_w_width)
+        w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 3:
+        w_ptrs = w_base + (2 * stride_w_width)
+        w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 4:
+        w_ptrs = w_base + (3 * stride_w_width)
+        w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
+    
+    mask_x_1d = idx_feats < dim
+    for idx_token in range(segment_len):
+        acc = acc_preload
+
+        matrix_w = w_col0
+        matrix_x = col0
+        for j in tl.static_range(KERNEL_WIDTH):
+            if KERNEL_WIDTH == 2:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = tl.reshape(
+                        tle.extract_tile(
+                            tile_data,
+                            index=[idx_token + chunk*(KERNEL_WIDTH - 1), 0],
+                            tile_shape=[1, BLOCK_N],
+                        ),
+                        [BLOCK_N],
+                    )
+            elif KERNEL_WIDTH == 3:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = tl.reshape(
+                        tle.extract_tile(
+                            tile_data,
+                            index=[idx_token + chunk*(KERNEL_WIDTH - 1), 0],
+                            tile_shape=[1, BLOCK_N],
+                        ),
+                        [BLOCK_N],
+                    )
+            elif KERNEL_WIDTH == 4:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = tl.reshape(
+                        tle.extract_tile(
+                            tile_data,
+                            index=[idx_token + chunk*(KERNEL_WIDTH - 1), 0],
+                            tile_shape=[1, BLOCK_N],
+                        ),
+                        [BLOCK_N],
+                    )
+            acc += matrix_x * matrix_w  # [BLOCK_N]
+
+        if KERNEL_WIDTH == 2:
+            col0 = matrix_x
+        elif KERNEL_WIDTH == 3:
+            col0 = col1
+            col1 = matrix_x
+        elif KERNEL_WIDTH == 4:
+            col0 = col1
+            col1 = col2
+            col2 = matrix_x
+
+        if SILU_ACTIVATION:
+            acc = acc / (1 + tl.exp(-acc))
+        mask_1d = (idx_token < segment_len) & (idx_feats < dim)
+        o_ptrs = (
+            o_ptr
+            + (sequence_start_index + token_offset + idx_token) * stride_o_token
+            + (idx_feats * stride_o_dim)
+        )
+        tl.store(o_ptrs, acc, mask=mask_1d)
+
+
+@triton.jit()
+def _causal_conv1d_update_kernel_v2(
+    # Pointers to matrices
+    x_ptr,  # (batch, dim, seqlen)
+    w_ptr,  # (dim, width)
+    bias_ptr,
+    conv_state_ptr,
+    conv_state_indices_ptr,
+    num_accepted_tokens_ptr,
+    query_start_loc_ptr,  # (batch + 1)
+    block_idx_last_scheduled_token,  # (batch,)
+    initial_state_idx,  # (batch,)
+    o_ptr,  # (batch, dim, seqlen)
+    # Matrix dimensions
+    batch: int,
+    dim: tl.constexpr,
+    seqlen: tl.constexpr,
+    state_len: tl.constexpr,
+    num_cache_lines: tl.constexpr,
+    # Strides
+    stride_x_seq: tl.constexpr,
+    stride_x_dim: tl.constexpr,
+    stride_x_token: tl.constexpr,
+    stride_w_dim: tl.constexpr,
+    stride_w_width: tl.constexpr,
+    stride_conv_state_seq: tl.constexpr,
+    stride_conv_state_dim: tl.constexpr,
+    stride_conv_state_tok: tl.constexpr,
+    stride_state_indices: tl.constexpr,
+    stride_o_seq: tl.constexpr,
+    stride_o_dim: tl.constexpr,
+    stride_o_token: tl.constexpr,
+    # others
+    pad_slot_id: tl.constexpr,
+    # Meta-parameters
+    HAS_BIAS: tl.constexpr,
+    KERNEL_WIDTH: tl.constexpr,
+    SILU_ACTIVATION: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_APC_ENABLED: tl.constexpr,
+    IS_SPEC_DECODING: tl.constexpr,
+    NP2_STATELEN: tl.constexpr,
+    USE_PAD_SLOT: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    idx_seq = tl.program_id(0)
+    if idx_seq >= batch:
+        return
+
+    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    if IS_APC_ENABLED:
+        conv_state_init = tl.load(initial_state_idx + idx_seq)
+        current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
+    else:
+        conv_state_init = 0
+        current_last_index = 0
+
+    conv_states_input_coord = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
+    ).to(tl.int64)
+
+    if USE_PAD_SLOT:  # noqa
+        if conv_states_input_coord == pad_slot_id:
+            return
+
+    if IS_VARLEN:
+        query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64)
+        query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64)
+        state_len = state_len - (seqlen - (query_end_index - query_start_index))
+        seqlen = query_end_index - query_start_index
+        x_offset = query_start_index * stride_x_token
+        o_offset = query_start_index * stride_o_token
+    else:
+        query_start_index = idx_seq * seqlen
+        query_end_index = query_start_index + seqlen
+        x_offset = idx_seq * stride_x_seq
+        o_offset = idx_seq * stride_o_seq
+
+    if query_start_index == query_end_index:
+        return
+
+    if IS_SPEC_DECODING:
+        conv_state_token_offset = (
+            tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
+        )
+    else:
+        conv_state_token_offset = 0
+
+    # STEP 1: load conv_state tile and extract columns via extract_tile
+    conv_states_base = (
+        conv_state_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )
+    mask_w = idx_feats < dim
+
+
+    idx_state = tl.arange(0, NP2_STATELEN)
+    conv_state_ptrs_full = (
+        conv_states_base[None, :]
+        + (conv_state_token_offset + idx_state)[:, None] * stride_conv_state_tok
+    )  # [NP2_STATELEN, BLOCK_N]
+    mask_state_full = (idx_state[:, None] < state_len) & (idx_feats[None, :] < dim)
+    state_tile = tl.load(
+        conv_state_ptrs_full, mask_state_full, other=0.0
+    )  # [NP2_STATELEN, BLOCK_N]
+
+    if KERNEL_WIDTH >= 2:
+        col0 = tle.extract_tile(state_tile, index=[0, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+    if KERNEL_WIDTH >= 3:
+        col1 = tle.extract_tile(state_tile, index=[1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+    if KERNEL_WIDTH >= 4:
+        col2 = tle.extract_tile(state_tile, index=[2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+    if KERNEL_WIDTH >= 5:
+        col3 = tle.extract_tile(state_tile, index=[3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+    if KERNEL_WIDTH >= 6:
+        col4 = tle.extract_tile(state_tile, index=[4, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+
+    # STEP 2: build updated conv_state and write back
+    idx_tokens = tl.arange(0, NP2_STATELEN)
+
+    conv_state_ptrs_source = (
+        conv_state_ptr
+        + (conv_states_input_coord * stride_conv_state_seq)
+        + conv_state_token_offset * stride_conv_state_tok
+        + (idx_feats * stride_conv_state_dim)[None, :]
+        + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[
+            :, None
+        ]
+    )  # [NP2_STATELEN, BLOCK_N]
+    mask = (
+        (conv_states_input_coord < num_cache_lines)
+        & ((idx_tokens + seqlen) < state_len)[:, None]
+        & (idx_feats < dim)[None, :]
+    )
+    conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
+
+    VAL = state_len - seqlen
+    x_base = x_ptr + x_offset + (idx_feats * stride_x_dim)  # [BLOCK_N]
+    x_ptrs = (
+        x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None]
+    )  # [NP2_STATELEN, BLOCK_N]
+    mask_x = (
+        (idx_tokens - VAL >= 0)[:, None]
+        & (idx_tokens - VAL < seqlen)[:, None]
+        & (idx_feats < dim)[None, :]
+    )
+    loaded_x = tl.load(x_ptrs, mask_x, 0.0)  # [NP2_STATELEN, BLOCK_N]
+
+
+
+    tl.debug_barrier()
+
+    new_conv_state = tl.where(mask, conv_state, loaded_x)
+
+    conv_states_offset = tl.load(
+        conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index
+    ).to(tl.int64)
+    conv_state_ptrs_target = (
+        conv_state_ptr
+        + (conv_states_offset * stride_conv_state_seq)
+        + (idx_feats * stride_conv_state_dim)
+    )[None, :] + (
+        idx_tokens * stride_conv_state_tok
+    )[:, None]
+    mask_out = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
+    tl.store(conv_state_ptrs_target, new_conv_state, mask_out)
+
+    # STEP 3: initialize accumulator
+    if HAS_BIAS:
+        bias = bias_ptr + idx_feats
+        mask_bias = idx_feats < dim
+        acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
+    else:
+        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    # STEP 4: preload weights
+    w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
+    mask_w = idx_feats < dim
+    if KERNEL_WIDTH >= 2:
+        w_ptrs = w_base + (0 * stride_w_width)
+        w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
+        w_ptrs = w_base + (1 * stride_w_width)
+        w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 3:
+        w_ptrs = w_base + (2 * stride_w_width)
+        w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 4:
+        w_ptrs = w_base + (3 * stride_w_width)
+        w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 5:
+        w_ptrs = w_base + (4 * stride_w_width)
+        w_col4 = tl.load(w_ptrs, mask_w, other=0.0)
+    if KERNEL_WIDTH >= 6:
+        w_ptrs = w_base + (5 * stride_w_width)
+        w_col5 = tl.load(w_ptrs, mask_w, other=0.0)
+
+    x_base_1d = x_base  # [BLOCK_N]
+    mask_x_1d = idx_feats < dim
+
+    # STEP 5: per-token convolution
+    for idx_token in tl.range(seqlen):
+        acc = acc_preload
+
+        matrix_w = w_col0
+        matrix_x = col0
+        for j in tl.static_range(KERNEL_WIDTH):
+            if KERNEL_WIDTH == 2:
+                if j == 1:
+                    matrix_w = w_col1
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 3:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 4:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 5:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = col3
+                elif j == 4:
+                    matrix_w = w_col4
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+            elif KERNEL_WIDTH == 6:
+                if j == 1:
+                    matrix_w = w_col1
+                    matrix_x = col1
+                elif j == 2:
+                    matrix_w = w_col2
+                    matrix_x = col2
+                elif j == 3:
+                    matrix_w = w_col3
+                    matrix_x = col3
+                elif j == 4:
+                    matrix_w = w_col4
+                    matrix_x = col4
+                elif j == 5:
+                    matrix_w = w_col5
+                    x_ptrs_1d = x_base_1d + idx_token * stride_x_token
+                    matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
+
+            acc += matrix_x * matrix_w  # [BLOCK_N]
+
+        if KERNEL_WIDTH == 2:
+            col0 = matrix_x
+        elif KERNEL_WIDTH == 3:
+            col0 = col1
+            col1 = matrix_x
+        elif KERNEL_WIDTH == 4:
+            col0 = col1
+            col1 = col2
+            col2 = matrix_x
+        elif KERNEL_WIDTH == 5:
+            col0 = col1
+            col1 = col2
+            col2 = col3
+            col3 = matrix_x
+        elif KERNEL_WIDTH == 6:
+            col0 = col1
+            col1 = col2
+            col2 = col3
+            col3 = col4
+            col4 = matrix_x
+
+        if SILU_ACTIVATION:
+            acc = acc / (1 + tl.exp(-acc))
+        mask_1d = (idx_token < seqlen) & (idx_feats < dim)
+        o_ptrs = (
+            o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim)
+        )
+        tl.store(o_ptrs, acc, mask=mask_1d)
+
+
+
+# %%
+# Python wrappers
+# ---------------
+
+
+def _launch_fwd_kernel(kernel_fn, x, weight, bias, conv_states, query_start_loc,
+                       cache_indices, has_initial_state, activation, pad_slot_id):
+    if isinstance(activation, bool) and activation:
+        activation = "silu"
+
+    original_x_dtype = x.dtype
+    x = x.to(conv_states.dtype)
+    out = torch.empty_like(x)
+
+    seqlens = query_start_loc.diff().to("cpu")
+    MAX_NUM_PROGRAMS = 1024
+    batch_ptr = torch.full(
+        (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
+    )
+    token_chunk_offset_ptr = torch.full(
+        (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
+    )
+
+    dim, cu_seqlen = x.shape
+    _, width = weight.shape
+    state_len = width - 1
+    np2_statelen = triton.next_power_of_2(state_len)
+    BLOCK_M = 8
+    num_cache_lines = conv_states.size(0)
+
+    stride_x_dim, stride_x_token = x.stride()
+    stride_w_dim, stride_w_width = weight.stride()
+    stride_istate_seq, stride_istate_dim, stride_istate_token = conv_states.stride()
+    stride_o_dim, stride_o_token = out.stride()
+    stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0
+
+    def num_program(META, seqlens):
+        nums = -(-seqlens // META["BLOCK_M"])
+        tot = nums.sum().item()
+        mlist = np.repeat(np.arange(len(nums)), nums)
+        offsetlist = []
+        for idx, num in enumerate(nums):
+            offsetlist.extend(range(num))
+
+        if META["batch_ptr"].nelement() < len(mlist):
+            newlen = len(mlist) + 1
+            META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
+            META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
+
+        if META["batch_ptr"].nelement() >= len(mlist):
+            META["batch_ptr"][0:len(mlist)].copy_(torch.from_numpy(np.array(mlist)))
+            META["token_chunk_offset_ptr"][0:len(mlist)].copy_(
+                torch.from_numpy(np.array(offsetlist))
+            )
+
+        META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device)
+        META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(
+            META["x_ptr"].device
+        )
+        return tot
+
+    def grid(META):
+        return (num_program(META, seqlens), triton.cdiv(dim, META["BLOCK_N"]))
+
+    kernel_fn[grid](
+        x, weight, bias, conv_states, cache_indices, has_initial_state,
+        query_start_loc, batch_ptr, token_chunk_offset_ptr,
+        None, None, None, None, out,
+        dim, cu_seqlen, num_cache_lines,
+        stride_x_dim, stride_x_token, stride_w_dim, stride_w_width,
+        stride_istate_seq, stride_istate_dim, stride_istate_token,
+        stride_cache_indices, stride_o_dim, stride_o_token,
+        1, pad_slot_id,
+        HAS_BIAS=bias is not None,
+        KERNEL_WIDTH=width,
+        SILU_ACTIVATION=activation in ["silu", "swish"],
+        IS_APC_ENABLED=False,
+        USE_PAD_SLOT=pad_slot_id is not None,
+        NP2_STATELEN=np2_statelen,
+        BLOCK_M=BLOCK_M, BLOCK_N=256, num_stages=2,
+    )
+    return out.to(original_x_dtype)
+
+
+def causal_conv1d_fn_v1(x, weight, bias, conv_states, query_start_loc,
+                        cache_indices=None, has_initial_state=None,
+                        activation="silu", pad_slot_id=PAD_SLOT_ID):
+    return _launch_fwd_kernel(
+        _causal_conv1d_fwd_kernel_v1, x, weight, bias, conv_states,
+        query_start_loc, cache_indices, has_initial_state, activation, pad_slot_id
+    )
+
+
+def causal_conv1d_fn_v2(x, weight, bias, conv_states, query_start_loc,
+                        cache_indices=None, has_initial_state=None,
+                        activation="silu", pad_slot_id=PAD_SLOT_ID):
+    return _launch_fwd_kernel(
+        _causal_conv1d_fwd_kernel_v2, x, weight, bias, conv_states,
+        query_start_loc, cache_indices, has_initial_state, activation, pad_slot_id
+    )
+
+
+def _launch_update_kernel(kernel_fn, x, conv_state, weight, bias,
+                          activation, conv_state_indices, pad_slot_id):
+    if isinstance(activation, bool):
+        activation = "silu" if activation else None
+    elif activation is not None:
+        assert activation in ["silu", "swish"]
+
+    original_x_dtype = x.dtype
+    x = x.to(conv_state.dtype)
+    unsqueeze = x.dim() == 2
+    if unsqueeze:
+        x = x.unsqueeze(-1)
+    batch, dim, seqlen = x.shape
+    _, width = weight.shape
+    num_cache_lines, _, state_len = conv_state.size()
+    out = x
+
+    stride_w_dim, stride_w_width = weight.stride()
+    stride_x_seq, stride_x_dim, stride_x_token = x.stride()
+    stride_o_seq, stride_o_dim, stride_o_token = out.stride()
+    stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride()
+    stride_state_indices = conv_state_indices.stride(0) if conv_state_indices is not None else 0
+    state_len = width - 1
+    np2_statelen = triton.next_power_of_2(state_len)
+
+    def grid(META):
+        return (batch, triton.cdiv(dim, META["BLOCK_N"]))
+
+    kernel_fn[grid](
+        x, weight, bias, conv_state, conv_state_indices, None, None,
+        None, None, out,
+        batch, dim, seqlen, state_len, num_cache_lines,
+        stride_x_seq, stride_x_dim, stride_x_token,
+        stride_w_dim, stride_w_width,
+        stride_istate_seq, stride_istate_dim, stride_istate_token,
+        stride_state_indices,
+        stride_o_seq, stride_o_dim, stride_o_token,
+        pad_slot_id,
+        HAS_BIAS=bias is not None,
+        KERNEL_WIDTH=width,
+        SILU_ACTIVATION=activation in ["silu", "swish"],
+        IS_VARLEN=False,
+        IS_APC_ENABLED=False,
+        IS_SPEC_DECODING=False,
+        NP2_STATELEN=np2_statelen,
+        USE_PAD_SLOT=pad_slot_id is not None,
+        BLOCK_N=256,
+    )
+    if unsqueeze:
+        out = out.squeeze(-1)
+    return out.to(original_x_dtype)
+
+
+def causal_conv1d_update_v1(x, conv_state, weight, bias=None, activation=None,
+                            conv_state_indices=None, pad_slot_id=PAD_SLOT_ID):
+    return _launch_update_kernel(
+        _causal_conv1d_update_kernel_v1, x, conv_state, weight, bias,
+        activation, conv_state_indices, pad_slot_id
+    )
+
+
+def causal_conv1d_update_v2(x, conv_state, weight, bias=None, activation=None,
+                            conv_state_indices=None, pad_slot_id=PAD_SLOT_ID):
+    return _launch_update_kernel(
+        _causal_conv1d_update_kernel_v2, x, conv_state, weight, bias,
+        activation, conv_state_indices, pad_slot_id
+    )
+
+
+# %%
+# Correctness check
+# -----------------
+
+
+def _make_varlen_data(dim, total_seqlen, batch, width, dtype):
+    torch.manual_seed(0)
+    eos_pos = torch.randperm(total_seqlen - 1)[:batch - 1].sort().values
+    seqlens = torch.diff(torch.cat([
+        torch.tensor([-1], dtype=torch.int32),
+        eos_pos.to(dtype=torch.int32),
+        torch.tensor([total_seqlen - 1], dtype=torch.int32),
+    ]))
+    query_start_loc = torch.cat([
+        torch.tensor([0], dtype=torch.int32),
+        torch.cumsum(seqlens, dim=0).to(torch.int32),
+    ]).to(DEVICE)
+
+    x = torch.randn(dim, total_seqlen, device=DEVICE, dtype=dtype)
+    weight = torch.randn(dim, width, device=DEVICE, dtype=dtype)
+    bias = torch.randn(dim, device=DEVICE, dtype=dtype)
+    conv_states = torch.randn(batch, dim, width - 1, device=DEVICE, dtype=dtype)
+    cache_indices = torch.arange(batch, dtype=torch.int32, device=DEVICE)
+    has_initial_state = torch.ones(batch, dtype=torch.bool, device=DEVICE)
+    return x, weight, bias, conv_states, query_start_loc, cache_indices, has_initial_state
+
+
+def _make_update_data(dim, batch, width, dtype):
+    torch.manual_seed(0)
+    x = torch.randn(batch, dim, 1, device=DEVICE, dtype=dtype)
+    weight = torch.randn(dim, width, device=DEVICE, dtype=dtype)
+    bias = torch.randn(dim, device=DEVICE, dtype=dtype)
+    conv_state = torch.randn(batch, dim, width - 1, device=DEVICE, dtype=dtype)
+    conv_state_indices = torch.arange(batch, dtype=torch.int32, device=DEVICE)
+    return x, weight, bias, conv_state, conv_state_indices
+
+
+def run_correctness(dim=1024, total_seqlen=2048, batch=32, width=4, dtype=torch.bfloat16):
+    # varlen forward
+    x, weight, bias, conv_states, query_start_loc, cache_indices, has_initial_state = \
+        _make_varlen_data(dim, total_seqlen, batch, width, dtype)
+
+    out_v1 = causal_conv1d_fn_v1(
+        x.clone(), weight, bias, conv_states.clone(), query_start_loc,
+        cache_indices, has_initial_state, activation="silu"
+    )
+    out_v2 = causal_conv1d_fn_v2(
+        x.clone(), weight, bias, conv_states.clone(), query_start_loc,
+        cache_indices, has_initial_state, activation="silu"
+    )
+    torch.testing.assert_close(out_v1.float(), out_v2.float(), rtol=5e-2, atol=2e-1)
+    print("[OK] varlen forward correctness passed")
+
+    # update
+    x_u, weight_u, bias_u, conv_state_u, conv_state_indices_u = \
+        _make_update_data(dim, batch, width, dtype)
+
+    out_u_v1 = causal_conv1d_update_v1(
+        x_u.clone(), conv_state_u.clone(), weight_u, bias_u,
+        activation="silu", conv_state_indices=conv_state_indices_u
+    )
+    out_u_v2 = causal_conv1d_update_v2(
+        x_u.clone(), conv_state_u.clone(), weight_u, bias_u,
+        activation="silu", conv_state_indices=conv_state_indices_u
+    )
+    torch.testing.assert_close(out_u_v1.float(), out_u_v2.float(), rtol=5e-2, atol=2e-1)
+    print("[OK] update correctness passed")
+
+
+# %%
+# Benchmark
+# ---------
+
+# %%
+# Main
+# ----
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dim", type=int, default=1024)
+    parser.add_argument("--total_seqlen", type=int, default=2048)
+    parser.add_argument("--batch", type=int, default=32)
+    parser.add_argument("--width", type=int, default=4)
+    parser.add_argument("--dtype", type=str, default="bfloat16",
+                        choices=["float16", "bfloat16", "fp16", "bf16"])
+    parser.add_argument("--benchmark", action="store_true")
+    args = parser.parse_args(argv)
+
+    dtype = torch.bfloat16 if "bf" in args.dtype else torch.float16
+
+    run_correctness(args.dim, args.total_seqlen, args.batch, args.width, dtype)
+
+    if args.benchmark:
+        # ------------------------------------------------------------------
+        # You can sweep more configs here by editing the lists below.
+        # Each (dim, batch) pair produces one benchmark table.
+        # ------------------------------------------------------------------
+        VARLEN_CONFIGS = [(4096, 32), (8192, 128)]
+        UPDATE_CONFIGS = [(256,), (1024,)]  # batch values
+        for dim, batch in VARLEN_CONFIGS:
+            _run_benchmark_varlen(dim, batch, dtype)
+        for batch in UPDATE_CONFIGS:
+            _run_benchmark_update(batch[0], dtype)
+
+
+def _run_benchmark_varlen(dim, batch, dtype):
+    width = 4
+    x_vals = [2048, 4096, 8192, 16384]
+    print(f"\n--- Varlen benchmark  (dim={dim}, batch={batch}, width={width}, dtype={dtype}) ---")
+    print(f"{'total_seqlen':>12} | {'Baseline(ms)':>12} {'TLE(ms)':>10} {'Speedup':>8}")
+    print("-" * 52)
+    for total_seqlen in x_vals:
+        x, w, b, cs, qsl, ci, his = _make_varlen_data(dim, total_seqlen, batch, width, dtype)
+        ms_b, _, _ = triton.testing.do_bench(
+            lambda: causal_conv1d_fn_v1(x, w, b, cs, qsl, ci, his, "silu"),
+            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
+        ms_t, _, _ = triton.testing.do_bench(
+            lambda: causal_conv1d_fn_v2(x, w, b, cs, qsl, ci, his, "silu"),
+            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
+        sp = ms_b / ms_t if ms_t > 0 else 1.0
+        print(f"{total_seqlen:>12} | {ms_b:>12.4f} {ms_t:>10.4f} {sp:>7.2f}x")
+
+
+def _run_benchmark_update(batch, dtype):
+    width = 4
+    x_vals = [1024, 2048, 4096, 8192]
+    print(f"\n--- Update benchmark  (batch={batch}, width={width}, dtype={dtype}) ---")
+    print(f"{'dim':>6} | {'Baseline(ms)':>12} {'TLE(ms)':>10} {'Speedup':>8}")
+    print("-" * 46)
+    for dim in x_vals:
+        x, w, b, cs, ci = _make_update_data(dim, batch, width, dtype)
+        ms_b, _, _ = triton.testing.do_bench(
+            lambda: causal_conv1d_update_v1(x, cs, w, b, "silu", ci),
+            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
+        ms_t, _, _ = triton.testing.do_bench(
+            lambda: causal_conv1d_update_v2(x, cs, w, b, "silu", ci),
+            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
+        sp = ms_b / ms_t if ms_t > 0 else 1.0
+        print(f"{dim:>6} | {ms_b:>12.4f} {ms_t:>10.4f} {sp:>7.2f}x")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/tutorials/tle/08-rope.py b/python/tutorials/tle/08-rope.py
new file mode 100644
index 0000000000..8dcb37f89e
--- /dev/null
+++ b/python/tutorials/tle/08-rope.py
@@ -0,0 +1,707 @@
+"""
+Rotary Position Embedding with TLE extract_tile
+================================================
+
+This tutorial compares a Triton baseline RoPE kernel against a TLE
+optimized version that uses ``tle.extract_tile`` for register-only head reuse.
+
+**Key optimization**: Instead of loading each attention head's row from global
+memory one-by-one (v1), the TLE version loads **all heads as a single 2D block**
+``[NUM_HEADS, HEAD_DIM]`` into registers, then uses ``extract_tile`` to slice
+out individual head rows — eliminating redundant memory round-trips.
+
+Both **inplace** (modifies input tensors) and **outplace** (writes to new output
+tensors) variants are benchmarked.
+
+Usage
+-----
+::
+
+    # correctness check only (default)
+    python python/tutorials/tle/08-rope.py
+
+    # correctness + benchmark tables with speedup
+    python python/tutorials/tle/08-rope.py --benchmark
+
+The speedup is largest when **head_dim is small** (e.g. 128) and **num_heads is
+large** — the per-head loop overhead saved by the block load dominates.
+When head_dim is large (256+) and total work is large, arithmetic dominates and
+the gain shrinks to near zero.
+
+Both non-interleaved (default LLaMA-style) and interleaved (GPT-NeoX-style)
+RoPE variants are included.
+"""
+
+# %%
+# Setup
+# -----
+
+import argparse
+import math
+import random
+import sys
+
+import torch
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
+
+
+# %%
+# Helper: build cos/sin cache
+# ---------------------------
+
+def build_rope_cache(max_seq_len, head_dim, device, dtype=torch.float32, base=10000.0):
+    """Precompute cos/sin tables for all positions up to ``max_seq_len``."""
+    half_dim = head_dim // 2
+    theta = 1.0 / (base ** (torch.arange(0, half_dim, device=device, dtype=dtype) / half_dim))
+    positions = torch.arange(max_seq_len, device=device, dtype=dtype)
+    angles = torch.outer(positions, theta)
+    return angles.cos().contiguous(), angles.sin().contiguous()
+
+
+# %%
+# Kernels (baseline v1) — per-head loads
+# --------------------------------------
+# Each head's row is loaded from global memory with a separate ``tl.load``
+# inside the ``for off_h`` loop.  For ``N`` heads this means ``N`` loads
+# (plus ``N`` loads for the rotated element), even though the data for all
+# heads is contiguous in memory.
+
+@triton.jit
+def _rope_outplace_kernel_v1(
+    oq_ptr,         # (n_tokens, q_heads, head_dim)
+    ok_ptr,         # (n_tokens, k_heads, head_dim)
+    q_ptr,          # (n_tokens, q_heads, head_dim)
+    k_ptr,          # (n_tokens, k_heads, head_dim)
+    cos_ptr,        # (max_seq_len, head_dim // 2)
+    sin_ptr,        # (max_seq_len, head_dim // 2)
+    q_stride_s, q_stride_h, q_stride_d,
+    k_stride_s, k_stride_h, k_stride_d,
+    oq_stride_s, oq_stride_h, oq_stride_d,
+    ok_stride_s, ok_stride_h, ok_stride_d,
+    cos_stride_s, sin_stride_s,
+    seq_len,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    PADDED_HEAD_DIM: tl.constexpr,
+    ROTARY_INTERLEAVED: tl.constexpr,
+    MAX_POSITION_EMBEDDINGS: tl.constexpr,
+):
+    """Baseline outplace: per-head tl.load, store to separate output."""
+    s_id = tl.program_id(0)
+
+    pos_id = s_id % seq_len
+    cos_ptr += pos_id * cos_stride_s
+    sin_ptr += pos_id * sin_stride_s
+
+    tl.device_assert(pos_id < MAX_POSITION_EMBEDDINGS, "position id out of bound")
+
+    ordered_block = tl.arange(0, PADDED_HEAD_DIM)
+    mask = ordered_block < HEAD_DIM
+
+    if ROTARY_INTERLEAVED:
+        odd_mask = ordered_block % 2 == 0
+        rotated_block = tl.where(odd_mask, ordered_block + 1, ordered_block - 1)
+        sin_cos_block = ordered_block // 2
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(odd_mask, -sin, sin)
+    else:
+        rotated_block = (ordered_block + HEAD_DIM // 2) % HEAD_DIM
+        sin_cos_block = ordered_block % (HEAD_DIM // 2)
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(rotated_block < HEAD_DIM // 2, sin, -sin)
+
+    # --- Q: per-head load from q_ptr, store to oq_ptr ---
+    q_ptr += s_id * q_stride_s
+    oq_ptr += s_id * oq_stride_s
+    for off_h in range(0, NUM_Q_HEADS):
+        ordered_cols = off_h * q_stride_h + (ordered_block * q_stride_d)
+        rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
+        output_offs = off_h * oq_stride_h + (ordered_block * oq_stride_d)
+
+        q = tl.load(q_ptr + ordered_cols, mask=mask, other=0.0)          # ← load from GMEM
+        rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)  # ← load from GMEM
+        y = q * cos + rotated_q * sin
+        tl.store(oq_ptr + output_offs, y, mask=mask)
+
+    # --- K: per-head load from k_ptr, store to ok_ptr ---
+    k_ptr += s_id * k_stride_s
+    ok_ptr += s_id * ok_stride_s
+    for off_h in range(0, NUM_K_HEADS):
+        ordered_cols = off_h * k_stride_h + (ordered_block * k_stride_d)
+        rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
+        output_offs = off_h * ok_stride_h + (ordered_block * ok_stride_d)
+
+        k = tl.load(k_ptr + ordered_cols, mask=mask, other=0.0)          # ← load from GMEM
+        rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)  # ← load from GMEM
+        y = k * cos + rotated_k * sin
+        tl.store(ok_ptr + output_offs, y, mask=mask)
+
+
+@triton.jit
+def _rope_inplace_kernel_v1(
+    q_ptr,          # (n_tokens, q_heads, head_dim)
+    k_ptr,          # (n_tokens, k_heads, head_dim)
+    cos_ptr,        # (max_seq_len, head_dim // 2)
+    sin_ptr,        # (max_seq_len, head_dim // 2)
+    q_stride_s, q_stride_h, q_stride_d,
+    k_stride_s, k_stride_h, k_stride_d,
+    cos_stride_s, sin_stride_s,
+    seq_len,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    PADDED_HEAD_DIM: tl.constexpr,
+    ROTARY_INTERLEAVED: tl.constexpr,
+    MAX_POSITION_EMBEDDINGS: tl.constexpr,
+):
+    """Baseline inplace: per-head tl.load, store back to input."""
+    s_id = tl.program_id(0)
+
+    pos_id = s_id % seq_len
+    cos_ptr += pos_id * cos_stride_s
+    sin_ptr += pos_id * sin_stride_s
+
+    tl.device_assert(pos_id < MAX_POSITION_EMBEDDINGS, "position id out of bound")
+
+    ordered_block = tl.arange(0, PADDED_HEAD_DIM)
+    mask = ordered_block < HEAD_DIM
+
+    if ROTARY_INTERLEAVED:
+        odd_mask = ordered_block % 2 == 0
+        rotated_block = tl.where(odd_mask, ordered_block + 1, ordered_block - 1)
+        sin_cos_block = ordered_block // 2
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(odd_mask, -sin, sin)
+    else:
+        rotated_block = (ordered_block + HEAD_DIM // 2) % HEAD_DIM
+        sin_cos_block = ordered_block % (HEAD_DIM // 2)
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(rotated_block < HEAD_DIM // 2, sin, -sin)
+
+    # --- Q: per-head load, store inplace ---
+    q_ptr += s_id * q_stride_s
+    for off_h in range(0, NUM_Q_HEADS):
+        ordered_cols = off_h * q_stride_h + (ordered_block * q_stride_d)
+        rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
+
+        q = tl.load(q_ptr + ordered_cols, mask=mask, other=0.0)
+        rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)
+        y = q * cos + rotated_q * sin
+        tl.store(q_ptr + ordered_cols, y, mask=mask)
+
+    # --- K: per-head load, store inplace ---
+    k_ptr += s_id * k_stride_s
+    for off_h in range(0, NUM_K_HEADS):
+        ordered_cols = off_h * k_stride_h + (ordered_block * k_stride_d)
+        rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
+
+        k = tl.load(k_ptr + ordered_cols, mask=mask, other=0.0)
+        rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)
+        y = k * cos + rotated_k * sin
+        tl.store(k_ptr + ordered_cols, y, mask=mask)
+
+
+# %%
+# Kernels (TLE v2) — block load + extract_tile
+# --------------------------------------------
+# Instead of ``N`` individual loads, we load **all heads at once** as a 2D tile
+# ``[NUM_HEADS, PADDED_HEAD_DIM]`` into registers, then use ``tle.extract_tile``
+# to pull out each head's row.  The rotated element still requires a per-head
+# ``tl.load`` (it lives at a different offset), but the dominant load path is
+# coalesced into one bulk access.
+
+@triton.jit
+def _rope_outplace_kernel_v2(
+    oq_ptr,         # (n_tokens, q_heads, head_dim)
+    ok_ptr,         # (n_tokens, k_heads, head_dim)
+    q_ptr,          # (n_tokens, q_heads, head_dim)
+    k_ptr,          # (n_tokens, k_heads, head_dim)
+    cos_ptr,        # (max_seq_len, head_dim // 2)
+    sin_ptr,        # (max_seq_len, head_dim // 2)
+    q_stride_s, q_stride_h, q_stride_d,
+    k_stride_s, k_stride_h, k_stride_d,
+    oq_stride_s, oq_stride_h, oq_stride_d,
+    ok_stride_s, ok_stride_h, ok_stride_d,
+    cos_stride_s, sin_stride_s,
+    seq_len,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    PADDED_HEAD_DIM: tl.constexpr,
+    ROTARY_INTERLEAVED: tl.constexpr,
+    MAX_POSITION_EMBEDDINGS: tl.constexpr,
+):
+    """TLE outplace: block-load all heads, extract_tile per head."""
+    s_id = tl.program_id(0)
+
+    pos_id = s_id % seq_len
+    cos_ptr += pos_id * cos_stride_s
+    sin_ptr += pos_id * sin_stride_s
+
+    tl.device_assert(pos_id < MAX_POSITION_EMBEDDINGS, "position id out of bound")
+
+    ordered_block = tl.arange(0, PADDED_HEAD_DIM)
+    mask = ordered_block < HEAD_DIM
+
+    if ROTARY_INTERLEAVED:
+        odd_mask = ordered_block % 2 == 0
+        rotated_block = tl.where(odd_mask, ordered_block + 1, ordered_block - 1)
+        sin_cos_block = ordered_block // 2
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(odd_mask, -sin, sin)
+    else:
+        rotated_block = (ordered_block + HEAD_DIM // 2) % HEAD_DIM
+        sin_cos_block = ordered_block % (HEAD_DIM // 2)
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(rotated_block < HEAD_DIM // 2, sin, -sin)
+
+    # --- Q: block load all heads from q_ptr, extract_tile, store to oq_ptr ---
+    q_ptr += s_id * q_stride_s
+    oq_ptr += s_id * oq_stride_s
+    q_head_ids = tl.arange(0, NUM_Q_HEADS)
+    # ── KEY CHANGE: one bulk load for all heads ──
+    q_block = tl.load(
+        q_ptr + q_head_ids[:, None] * q_stride_h + ordered_block[None, :] * q_stride_d,
+        mask=mask[None, :],
+        other=0.0,
+    )  # shape: [NUM_Q_HEADS, PADDED_HEAD_DIM]
+
+    for off_h in range(0, NUM_Q_HEADS):
+        # ── KEY CHANGE: slice from register tile instead of GMEM load ──
+        q = tle.extract_tile(
+            q_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
+        ).reshape((PADDED_HEAD_DIM,))
+
+        rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
+        rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)
+        y = q * cos + rotated_q * sin
+        tl.store(oq_ptr + off_h * oq_stride_h + ordered_block * oq_stride_d, y, mask=mask)
+
+    # --- K: block load all heads from k_ptr, extract_tile, store to ok_ptr ---
+    k_ptr += s_id * k_stride_s
+    ok_ptr += s_id * ok_stride_s
+    k_head_ids = tl.arange(0, NUM_K_HEADS)
+    k_block = tl.load(
+        k_ptr + k_head_ids[:, None] * k_stride_h + ordered_block[None, :] * k_stride_d,
+        mask=mask[None, :],
+        other=0.0,
+    )  # shape: [NUM_K_HEADS, PADDED_HEAD_DIM]
+
+    for off_h in range(0, NUM_K_HEADS):
+        k = tle.extract_tile(
+            k_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
+        ).reshape((PADDED_HEAD_DIM,))
+
+        rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
+        rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)
+        y = k * cos + rotated_k * sin
+        tl.store(ok_ptr + off_h * ok_stride_h + ordered_block * ok_stride_d, y, mask=mask)
+
+
+@triton.jit
+def _rope_inplace_kernel_v2(
+    q_ptr,          # (n_tokens, q_heads, head_dim)
+    k_ptr,          # (n_tokens, k_heads, head_dim)
+    cos_ptr,        # (max_seq_len, head_dim // 2)
+    sin_ptr,        # (max_seq_len, head_dim // 2)
+    q_stride_s, q_stride_h, q_stride_d,
+    k_stride_s, k_stride_h, k_stride_d,
+    cos_stride_s, sin_stride_s,
+    seq_len,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    PADDED_HEAD_DIM: tl.constexpr,
+    ROTARY_INTERLEAVED: tl.constexpr,
+    MAX_POSITION_EMBEDDINGS: tl.constexpr,
+):
+    """TLE inplace: block-load all heads, then extract_tile per head."""
+    s_id = tl.program_id(0)
+
+    pos_id = s_id % seq_len
+    cos_ptr += pos_id * cos_stride_s
+    sin_ptr += pos_id * sin_stride_s
+
+    tl.device_assert(pos_id < MAX_POSITION_EMBEDDINGS, "position id out of bound")
+
+    ordered_block = tl.arange(0, PADDED_HEAD_DIM)
+    mask = ordered_block < HEAD_DIM
+
+    if ROTARY_INTERLEAVED:
+        odd_mask = ordered_block % 2 == 0
+        rotated_block = tl.where(odd_mask, ordered_block + 1, ordered_block - 1)
+        sin_cos_block = ordered_block // 2
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(odd_mask, -sin, sin)
+    else:
+        rotated_block = (ordered_block + HEAD_DIM // 2) % HEAD_DIM
+        sin_cos_block = ordered_block % (HEAD_DIM // 2)
+        cos = tl.load(cos_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + sin_cos_block, mask=mask, other=0.0).to(tl.float32)
+        sin = tl.where(rotated_block < HEAD_DIM // 2, sin, -sin)
+
+    # --- Q: block load all heads, then extract_tile ---
+    q_ptr += s_id * q_stride_s
+    q_head_ids = tl.arange(0, NUM_Q_HEADS)
+    q_block = tl.load(
+        q_ptr + q_head_ids[:, None] * q_stride_h + ordered_block[None, :] * q_stride_d,
+        mask=mask[None, :],
+        other=0.0,
+    )
+
+    for off_h in range(0, NUM_Q_HEADS):
+        q = tle.extract_tile(
+            q_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
+        ).reshape((PADDED_HEAD_DIM,))
+
+        rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
+        rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)
+        y = q * cos + rotated_q * sin
+        tl.store(q_ptr + off_h * q_stride_h + ordered_block * q_stride_d, y, mask=mask)
+
+    # --- K: block load all heads, then extract_tile ---
+    k_ptr += s_id * k_stride_s
+    k_head_ids = tl.arange(0, NUM_K_HEADS)
+    k_block = tl.load(
+        k_ptr + k_head_ids[:, None] * k_stride_h + ordered_block[None, :] * k_stride_d,
+        mask=mask[None, :],
+        other=0.0,
+    )
+
+    for off_h in range(0, NUM_K_HEADS):
+        k = tle.extract_tile(
+            k_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
+        ).reshape((PADDED_HEAD_DIM,))
+
+        rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
+        rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)
+        y = k * cos + rotated_k * sin
+        tl.store(k_ptr + off_h * k_stride_h + ordered_block * k_stride_d, y, mask=mask)
+
+
+# %%
+# Python wrappers
+# ---------------
+
+
+def _launch_rope_outplace(kernel, q, k, cos, sin, rotary_interleaved):
+    """Launch outplace kernel: returns new (q_embed, k_embed)."""
+    q_shape, k_shape = q.shape, k.shape
+    assert q.dim() == 4, f"Expected 4D input (B, S, H, D), got {q.shape}"
+    seq_len = q.shape[-3]
+
+    q = q.view(-1, q.shape[-2], q.shape[-1])
+    k = k.view(-1, k.shape[-2], k.shape[-1])
+
+    n_tokens, q_heads, head_dim = q.shape
+    padded_head_dim = max(triton.next_power_of_2(head_dim), 16)
+
+    q_embed = torch.empty_like(q)
+    k_embed = torch.empty_like(k)
+
+    grid = (n_tokens,)
+    kernel[grid](
+        q_embed, k_embed,
+        q, k, cos, sin,
+        q.stride(0), q.stride(1), q.stride(2),
+        k.stride(0), k.stride(1), k.stride(2),
+        q_embed.stride(0), q_embed.stride(1), q_embed.stride(2),
+        k_embed.stride(0), k_embed.stride(1), k_embed.stride(2),
+        cos.stride(0), sin.stride(0),
+        seq_len,
+        q_heads, k.shape[-2], head_dim, padded_head_dim,
+        rotary_interleaved,
+        MAX_POSITION_EMBEDDINGS=cos.shape[0],
+    )
+    return q_embed.view(q_shape), k_embed.view(k_shape)
+
+
+def _launch_rope_inplace(kernel, q, k, cos, sin, rotary_interleaved):
+    """Launch inplace kernel: modifies q, k in place."""
+    q_shape, k_shape = q.shape, k.shape
+    assert q.dim() == 4, f"Expected 4D input (B, S, H, D), got {q.shape}"
+    seq_len = q.shape[-3]
+
+    q = q.view(-1, q.shape[-2], q.shape[-1])
+    k = k.view(-1, k.shape[-2], k.shape[-1])
+
+    n_tokens, q_heads, head_dim = q.shape
+    padded_head_dim = max(triton.next_power_of_2(head_dim), 16)
+
+    grid = (n_tokens,)
+    kernel[grid](
+        q, k, cos, sin,
+        q.stride(0), q.stride(1), q.stride(2),
+        k.stride(0), k.stride(1), k.stride(2),
+        cos.stride(0), sin.stride(0),
+        seq_len,
+        q_heads, k.shape[-2], head_dim, padded_head_dim,
+        rotary_interleaved,
+        MAX_POSITION_EMBEDDINGS=cos.shape[0],
+    )
+    return q.view(q_shape), k.view(k_shape)
+
+
+def rope_outplace_v1(q, k, cos, sin, rotary_interleaved=False):
+    """Baseline RoPE (outplace)."""
+    return _launch_rope_outplace(_rope_outplace_kernel_v1, q, k, cos, sin, rotary_interleaved)
+
+def rope_outplace_v2(q, k, cos, sin, rotary_interleaved=False):
+    """TLE RoPE (outplace)."""
+    return _launch_rope_outplace(_rope_outplace_kernel_v2, q, k, cos, sin, rotary_interleaved)
+
+def rope_inplace_v1(q, k, cos, sin, rotary_interleaved=False):
+    """Baseline RoPE (inplace)."""
+    return _launch_rope_inplace(_rope_inplace_kernel_v1, q, k, cos, sin, rotary_interleaved)
+
+def rope_inplace_v2(q, k, cos, sin, rotary_interleaved=False):
+    """TLE RoPE (inplace)."""
+    return _launch_rope_inplace(_rope_inplace_kernel_v2, q, k, cos, sin, rotary_interleaved)
+
+
+# %%
+# Correctness check
+# -----------------
+
+
+def _assert_close_robust(name, a, b, rtol=5e-2, atol=2e-1):
+    """Robust comparison: pass if allclose OR quantile-based relaxed check passes.
+
+    Bfloat16 arithmetic is non-deterministic across different memory access
+    patterns (individual loads vs block loads).  A handful of elements near
+    zero can have large relative errors without indicating a real bug.
+    """
+    if torch.allclose(a, b, rtol=rtol, atol=atol):
+        return
+
+    diff = (a.float() - b.float()).abs()
+    denom = torch.maximum(a.float().abs(), b.float().abs()) + 1e-6
+    rel = diff / denom
+
+    flat_diff = diff.flatten()
+    flat_rel = rel.flatten()
+    max_samples = 1_000_000
+    if flat_diff.numel() > max_samples:
+        step = (flat_diff.numel() + max_samples - 1) // max_samples
+        flat_diff = flat_diff[::step]
+        flat_rel = flat_rel[::step]
+
+    abs_p999 = torch.quantile(flat_diff, 0.999).item()
+    rel_p999 = torch.quantile(flat_rel, 0.999).item()
+
+    if abs_p999 < 0.25 and rel_p999 < 0.25:
+        return  # robust check passed
+
+    raise AssertionError(
+        f"{name} mismatch: allclose failed and robust check failed; "
+        f"abs_p999={abs_p999:.6f}, rel_p999={rel_p999:.6f}"
+    )
+
+
+def check_correctness(batch=4, seq_len=256, q_heads=32, k_heads=8, head_dim=128,
+                      dtype=torch.bfloat16, rotary_interleaved=False):
+    """Verify v1 and v2 produce identical results (both inplace and outplace)."""
+    torch.manual_seed(0)
+
+    max_seq = seq_len + 1024
+    q = torch.randn(batch, seq_len, q_heads, head_dim, device=DEVICE, dtype=dtype)
+    k = torch.randn(batch, seq_len, k_heads, head_dim, device=DEVICE, dtype=dtype)
+    cos, sin = build_rope_cache(max_seq, head_dim, DEVICE)
+
+    # --- outplace ---
+    oq1, ok1 = rope_outplace_v1(q.clone(), k.clone(), cos, sin, rotary_interleaved)
+    oq2, ok2 = rope_outplace_v2(q.clone(), k.clone(), cos, sin, rotary_interleaved)
+    _assert_close_robust("outplace_q", oq1, oq2)
+    _assert_close_robust("outplace_k", ok1, ok2)
+
+    # --- inplace ---
+    iq1, ik1 = rope_inplace_v1(q.clone(), k.clone(), cos, sin, rotary_interleaved)
+    iq2, ik2 = rope_inplace_v2(q.clone(), k.clone(), cos, sin, rotary_interleaved)
+    _assert_close_robust("inplace_q", iq1, iq2)
+    _assert_close_robust("inplace_k", ik1, ik2)
+
+    inter = "interleaved" if rotary_interleaved else "non-interleaved"
+    print(f"[OK] Correctness passed ({inter}, {dtype}, dim={head_dim}) [outplace + inplace]")
+
+
+# %%
+# Benchmark
+# ---------
+# Uses the same rigorous methodology as the original rotary embedding benchmark:
+#   - Pre-allocated buffers with copy_ reset outside the timing window
+#   - L2 cache cleared between iterations
+#   - Multiple rounds with outlier trimming
+#   - Randomized v1/v2 measurement order
+
+BENCH_WARMUP = 20
+BENCH_REP = 200
+BENCH_ROUNDS = 5
+BENCH_TRIM = 0.20
+
+
+def _do_bench_rigorous(fn, reset_fn):
+    """Time ``fn`` with cache clearing and pre-allocated buffers."""
+    di = triton.runtime.driver.active.get_device_interface()
+    cache = triton.runtime.driver.active.get_empty_cache_for_benchmark()
+
+    for _ in range(BENCH_WARMUP):
+        reset_fn()
+        triton.runtime.driver.active.clear_cache(cache)
+        fn()
+
+    di.synchronize()
+
+    start_events = [di.Event(enable_timing=True) for _ in range(BENCH_REP)]
+    end_events = [di.Event(enable_timing=True) for _ in range(BENCH_REP)]
+
+    for i in range(BENCH_REP):
+        reset_fn()
+        triton.runtime.driver.active.clear_cache(cache)
+        di.synchronize()
+        start_events[i].record()
+        fn()
+        end_events[i].record()
+
+    di.synchronize()
+    times = [s.elapsed_time(e) for s, e in zip(start_events, end_events)]
+    t = torch.tensor(times, dtype=torch.float64)
+    return float(torch.quantile(t, 0.5).item())
+
+
+def _robust_bench(fn, reset_fn):
+    """Multiple rounds with outlier trimming for stable measurement."""
+    medians = []
+    for _ in range(BENCH_ROUNDS):
+        medians.append(_do_bench_rigorous(fn, reset_fn))
+
+    t = torch.tensor(medians, dtype=torch.float64)
+    if BENCH_ROUNDS >= 5:
+        lo = torch.quantile(t, BENCH_TRIM)
+        hi = torch.quantile(t, 1.0 - BENCH_TRIM)
+        keep = (t >= lo) & (t <= hi)
+        if keep.sum().item() < max(3, BENCH_ROUNDS // 2):
+            keep = torch.ones_like(t, dtype=torch.bool)
+    else:
+        keep = torch.ones_like(t, dtype=torch.bool)
+    return float(torch.quantile(t[keep], 0.5).item())
+
+
+def _run_benchmark_table(title, configs, dtype, rotary_interleaved,
+                         rope_v1_fn, rope_v2_fn, inplace):
+    """Print a benchmark table for a set of (batch, seq_len, q_heads, k_heads, head_dim)."""
+    mode = "inplace" if inplace else "outplace"
+    print(f"\n--- {title} [{mode}] (dtype={dtype}) ---")
+    print(f"{'batch':>6} {'seq':>6} {'qh':>4} {'kh':>4} {'dim':>5} | "
+          f"{'Baseline(ms)':>13} {'TLE(ms)':>10} {'Speedup':>8}")
+    print("-" * 72)
+
+    for batch, seq_len, q_heads, k_heads, head_dim in configs:
+        torch.manual_seed(0)
+        max_seq = seq_len + 1024
+        q_src = torch.randn(batch, seq_len, q_heads, head_dim, device=DEVICE, dtype=dtype)
+        k_src = torch.randn(batch, seq_len, k_heads, head_dim, device=DEVICE, dtype=dtype)
+        cos, sin = build_rope_cache(max_seq, head_dim, DEVICE)
+
+        # Pre-allocated buffers
+        q1, k1 = q_src.clone(), k_src.clone()
+        q2, k2 = q_src.clone(), k_src.clone()
+
+        def run_v1():
+            rope_v1_fn(q1, k1, cos, sin, rotary_interleaved)
+
+        def run_v2():
+            rope_v2_fn(q2, k2, cos, sin, rotary_interleaved)
+
+        def reset_v1():
+            q1.copy_(q_src); k1.copy_(k_src)
+
+        def reset_v2():
+            q2.copy_(q_src); k2.copy_(k_src)
+
+        # Randomize measurement order
+        runners = [("v1", run_v1, reset_v1), ("v2", run_v2, reset_v2)]
+        random.shuffle(runners)
+
+        stats = {}
+        for name, fn, rst in runners:
+            stats[name] = _robust_bench(fn, rst)
+
+        ms_v1, ms_v2 = stats["v1"], stats["v2"]
+        speedup = ms_v1 / ms_v2 if ms_v2 > 0 else 0.0
+        print(f"{batch:>6} {seq_len:>6} {q_heads:>4} {k_heads:>4} {head_dim:>5} | "
+              f"{ms_v1:>13.5f} {ms_v2:>10.5f} {speedup:>7.2f}×")
+
+
+# %%
+# Main
+# ----
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(
+        description="RoPE with TLE extract_tile — baseline vs optimized benchmark"
+    )
+    parser.add_argument("--benchmark", action="store_true",
+                        help="Run benchmark tables in addition to correctness check")
+    parser.add_argument("--dtype", type=str, default="bfloat16",
+                        choices=["float16", "bfloat16", "fp16", "bf16"],
+                        help="Data type for benchmark")
+    args = parser.parse_args(argv)
+
+    dtype = torch.bfloat16 if "bf" in args.dtype else torch.float16
+
+    # --- Correctness ---------------------------------------------------------
+    print("=" * 60)
+    print("Correctness checks")
+    print("=" * 60)
+    for interleaved in [False, True]:
+        for head_dim in [128, 256]:
+            check_correctness(head_dim=head_dim, dtype=dtype, rotary_interleaved=interleaved)
+
+    if not args.benchmark:
+        print("\n[INFO] Run with --benchmark to see performance tables.")
+        return
+
+    # --- Benchmark -----------------------------------------------------------
+    CONFIGS = [
+        # (batch, seq_len, q_heads, k_heads, head_dim)
+        # head_dim=128 — TLE shines
+        (1,   128,  32, 8,  128),
+        (1,  1024,  32, 8,  128),
+        (8,   128,  32, 8,  128),
+        (8,  1024,  32, 8,  128),
+        (32,  128,  32, 8,  128),
+        (32, 1024,  32, 8,  128),
+        (8,  1024,  16, 16, 128),
+        (32, 1024,  16, 16, 128),
+        # head_dim=256 — smaller gain
+        (1,   128,  32, 8,  256),
+        (8,   128,  32, 8,  256),
+        (32,  128,  32, 8,  256),
+        (32, 1024,  16, 16, 256),
+    ]
+
+    for interleaved in [False, True]:
+        ilabel = "interleaved (GPT-NeoX style)" if interleaved else "non-interleaved (LLaMA style)"
+        # outplace
+        _run_benchmark_table(ilabel, CONFIGS, dtype, interleaved,
+                             rope_outplace_v1, rope_outplace_v2, inplace=False)
+        # inplace
+        _run_benchmark_table(ilabel, CONFIGS, dtype, interleaved,
+                             rope_inplace_v1, rope_inplace_v2, inplace=True)
+
+
+if __name__ == "__main__":
+    main()

From fd6e15eb30149072fa86d931eaf927fba25562c3 Mon Sep 17 00:00:00 2001
From: liu-jiale <49688036+git-flyer@users.noreply.github.com>
Date: Wed, 3 Jun 2026 18:10:14 +0800
Subject: [PATCH 05/41] [FlagCICD] Add a step to test the accuracy of the
 FlagGems operator (#642)

---------

Co-authored-by: flagtree-bot <flagtree_ai@163.com>
---
 .../flagcicd-hopper-build-and-test.yml        | 244 +++++++-----------
 1 file changed, 92 insertions(+), 152 deletions(-)

diff --git a/.github/workflows/flagcicd-hopper-build-and-test.yml b/.github/workflows/flagcicd-hopper-build-and-test.yml
index b428a465f0..df4d6567a4 100644
--- a/.github/workflows/flagcicd-hopper-build-and-test.yml
+++ b/.github/workflows/flagcicd-hopper-build-and-test.yml
@@ -3,10 +3,8 @@ name: flagcicd-Hopper-Build-And-Test
 on:
   schedule:
     - cron: '0 21 * * *'
-  push:
-    branches: [ "triton_v3.4.x", "triton_v3.5.x", "triton_v3.6.x" ]
   pull_request:
-    branches: [ "triton_v3.4.x", "triton_v3.5.x", "triton_v3.6.x" ]
+    branches: [ "triton_v3.4.x", "triton_v3.5.x", "triton_v3.6.x"]
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
@@ -16,10 +14,17 @@ jobs:
   flagcicd-hopper-build-and-test:
     runs-on: [flagcicd-flagtree-nvidia3.6]
     container:
-      image: harbor.baai.ac.cn/flagtree/flagtree-3.6.x-py312-torch2.8.0a0_5228986c39.nv25.05-ubuntu24.04:202603
-      options: --uts=host --ipc=host --privileged  --ulimit stack=67108864 --ulimit memlock=-1 --security-opt seccomp=unconfined --gpus=all
+      image: harbor.baai.ac.cn/flagtree/flagtree-py312-torch2.8.0a0_5228986c39.nv25.05-ubuntu24.04:202605-3.6-base
+      options:
+        --uts=host
+        --ipc=host
+        --privileged
+        --ulimit stack=67108864
+        --ulimit memlock=-1
+        --security-opt seccomp=unconfined
+        --gpus=all
       volumes:
-        - /mnt/airs-business/cicd:/mnt/airs-business/cicd
+        - /mnt/airs-business/cicd/flagtree/flagcicd-flagGems-test/:/mnt/airs-business/cicd/flagtree/flagcicd-flagGems-test/
 
     if: ${{ github.repository == 'FlagTree/flagtree' || github.repository == 'flagos-ai/flagtree' }}
     steps:
@@ -50,165 +55,100 @@ jobs:
           echo "TARGET_BRANCH=$TARGET_BRANCH" >> $GITHUB_ENV
           echo "TARGET_BRANCH=$TARGET_BRANCH"
 
-      - name: FlagTree Build on NVidia (triton_v3.4.x branch)
-        if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.4.x' }}
-        shell: bash
-        run: |
-          set -x
-          pip uninstall -y triton
-          source ~/env-3.4.sh
-          env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true
-          MAX_JOBS=32 python3 -m pip install . --no-build-isolation
-
-      - name: FlagTree Build on NVidia (triton_v3.5.x branch)
-        if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.5.x' }}
-        shell: bash
-        run: |
-          set -x
-          pip uninstall -y triton
-          source ~/env-3.5.sh
-          env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true
-          MAX_JOBS=32 python3 -m pip install . --no-build-isolation
 
       - name: FlagTree Build on NVidia (triton_v3.6.x branch)
         if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.6.x' }}
         shell: bash
         run: |
           set -x
+          export TRITON_HOME="/root/"
           pip uninstall -y triton
           env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true
           MAX_JOBS=32 python3 -m pip install . --no-build-isolation
+          unset TRITON_HOME
 
-      - name: FlagTree Test on NVidia (triton_v3.4.x branch)
-        if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.4.x' }}
+      - name: FlagGems Accuracy Test on NVidia
+        if: ${{ steps.check_backend.outputs.should_skip != 'true' }}
         shell: bash
         run: |
-          set -x
-          # python unit test
-          python3 -m pytest -s python/test/unit/cuda
-          python3 -m pytest -s python/test/unit/instrumentation
-          python3 -m pytest -s python/test/unit/language \
-              --ignore=python/test/unit/language/test_line_info.py
-          python3 -m pytest -s python/test/unit/runtime
-          if [ -d "python/test/operators" ]; then
-            python3 -m pytest -s python/test/operators
+          # set -x
+          export TRITON_CACHE_DIR="/mnt/airs-business/cicd/flagtree/flagcicd-flagGems-test/nvidia/.triton/cache"
+          FLAGGEMS_BASE="/mnt/airs-business/cicd/flagtree/flagcicd-flagGems-test/nvidia"
+          FLAGGEMS_DIR="${FLAGGEMS_BASE}/FlagGems"
+          CONFIG_DIR="${FLAGGEMS_BASE}/config"
+          SCRIPTS_DIR="${FLAGGEMS_BASE}/../scripts"
+          IGNORE_LIST="${CONFIG_DIR}/ignore_list.txt"
+          OPERATORS_YAML="${FLAGGEMS_DIR}/conf/operators.yaml"
+
+          # echo "${FLAGGEMS_BASE}"   "${CONFIG_DIR}"   "${SCRIPTS_DIR}"
+
+          if [ ! -d "$FLAGGEMS_DIR" ]; then
+            echo "Error: FlagGems directory not found at $FLAGGEMS_DIR"
+            echo "Please clone FlagGems repository to the shared storage first."
+            exit 1
           fi
 
-      - name: FlagTree Test on NVidia (triton_v3.5.x branch)
-        if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.5.x' }}
-        shell: bash
-        run: |
-          set -x
-          source ~/env.sh
-          # python tutorials
-          python3 python/tutorials/01-vector-add.py --only_unit_test
-          python3 python/tutorials/02-fused-softmax.py --only_unit_test
-          python3 python/tutorials/03-matrix-multiplication.py --only_unit_test
-          python3 python/tutorials/04-low-memory-dropout.py --only_unit_test
-          python3 python/tutorials/05-layer-norm.py --only_unit_test
-          python3 python/tutorials/06-fused-attention.py --only_unit_test
-          python3 python/tutorials/07-extern-functions.py --only_unit_test
-          python3 python/tutorials/08-grouped-gemm.py --only_unit_test
-          python3 python/tutorials/09-persistent-matmul.py --only_unit_test
-          python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test
-          # python unit test
-          python3 -m pytest -s python/test/unit/cuda
-          python3 -m pytest -s python/test/unit/instrumentation
-          python3 -m pytest -s python/test/unit/language \
-              --ignore=python/test/unit/language/test_line_info.py
-          python3 -m pytest -s python/test/unit/runtime
-          if [ -d "python/test/operators" ]; then
-            python3 -m pytest -s python/test/operators
+          # Check if the file "operators.yaml" exists
+          if [ ! -f "$OPERATORS_YAML" ]; then
+            echo "Error: operators.yaml not found at $OPERATORS_YAML"
+            exit 1
           fi
-          # flagtree tle
-          # python tutorials
-          python3 python/tutorials/tle/01-fft.py
-          python3 python/tutorials/tle/02-moe_align_block_size.py
-          python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
-          python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
-          python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py --mode test
-          python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
-          # python unit test
-          python3 -m pytest -s python/test/tle/integration
-          python3 -m pytest -s python/test/tle/unit
-          # flagtree hints
-          # python tutorials
-          python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
-          # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test
-          python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test
-          python3 python/tutorials/hints/04/04-low-memory-dropout.py --only_unit_test
-          python3 python/tutorials/hints/05/05-layer-norm.py --only_unit_test
-          python3 python/tutorials/hints/06/06-fused-attention.py --only_unit_test
-          python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test
-          python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test
-          python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test
-          # flagtree tle raw
-          python3 python/tutorials/tle/raw/mlir/01-vector-add.py
-          python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py
-          python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py
-          python3 python/tutorials/tle/raw/mlir/04-hello-world.py
-          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton
-          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle
-          python3 python/tutorials/tle/raw/mlir/06-test-vassert.py
-
-      - name: FlagTree Test on NVidia (triton_v3.6.x branch)
-        if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.6.x' }}
-        shell: bash
-        run: |
-          set -x
-          # python tutorials
-          python3 python/tutorials/01-vector-add.py --only_unit_test
-          python3 python/tutorials/02-fused-softmax.py --only_unit_test
-          python3 python/tutorials/03-matrix-multiplication.py --only_unit_test
-          python3 python/tutorials/04-low-memory-dropout.py --only_unit_test
-          python3 python/tutorials/05-layer-norm.py --only_unit_test
-          python3 python/tutorials/06-fused-attention.py --only_unit_test
-          python3 python/tutorials/07-extern-functions.py --only_unit_test
-          python3 python/tutorials/08-grouped-gemm.py --only_unit_test
-          python3 python/tutorials/09-persistent-matmul.py --only_unit_test
-          python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test
-          # python unit test
-          python3 -m pytest -s python/test/unit/cuda
-          python3 -m pytest -s python/test/unit/instrumentation
-          python3 -m pytest -s python/test/unit/language \
-              --ignore=python/test/unit/language/test_line_info.py
-          python3 -m pytest -s python/test/unit/runtime
-          if [ -d "python/test/operators" ]; then
-            python3 -m pytest -s python/test/operators
+
+          # Install FlagGems dependencies
+          cd "$FLAGGEMS_DIR"
+          # pip install -r requirements/requirements_nvidia.txt
+          # pip list
+          pip install -v -e .  --no-build-isolation --no-deps
+          if [ $? -ne 0 ]; then
+              echo "Error: Failed to install FlagGems"
+              exit 1
+          fi
+
+          # Generate pytest marks parameters
+          MARKS=$(python3  "${SCRIPTS_DIR}/parse_operators.py"  "$OPERATORS_YAML")
+          if [ -z "$MARKS" ]; then
+            echo "Error: Failed to generate pytest marks"
+            exit 1
+          fi
+
+          # Print the first 200 characters of the MARKS string
+          # echo "Generated pytest marks: ${MARKS:0:200}..."
+
+          cd tests/
+          # pass the file path of the ignore list to the Python plugin
+          export FLAGGEMS_IGNORE_LIST_PATH="$IGNORE_LIST"
+
+          # Add the "scripts" directory to the Python search path
+          export PYTHONPATH="${SCRIPTS_DIR}:$PYTHONPATH"
+
+          # Build the Pytest command and use the -p parameter to load the dynamic_ignore plugin
+          PYTEST_CMD="python3 -m pytest -m \"$MARKS\" --ref=cpu --quick -p dynamic_ignore"
+
+          echo "=========================================="
+          echo "Running FlagGems accuracy tests"
+          echo "Quick mode: enabled"
+          echo "Reference device: cpu"
+          echo "=========================================="
+
+
+          set +e
+          bash -c "$PYTEST_CMD"
+          EXIT_CODE=$?
+          set -e
+
+          if [ $EXIT_CODE -eq 0 ]; then
+              echo "All tests passed"
+              echo "FlagGems accuracy tests completed successfully"
+              exit 0
+          elif [ $EXIT_CODE -eq 5 ]; then
+              # It might be a configuration error. An error should have been reported.
+              echo "Error: No tests were collected. Configuration issue suspected."
+              exit 1
+          elif [ $EXIT_CODE -eq 1 ]; then
+              echo "Warning: Some tests failed, tests exited with code $EXIT_CODE"
+              echo "FlagGems accuracy tests completed successfully (failures allowed)"
+              exit 0  # Test failure does not block CI
+          else
+              echo "Error: Tests exited with code $EXIT_CODE"
+              exit 1
           fi
-          # flagtree tle
-          # python tutorials
-          python3 python/tutorials/tle/01-fft.py
-          python3 python/tutorials/tle/02-moe_align_block_size.py
-          python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
-          python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
-          python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
-          # python unit test
-          python3 -m pytest -s python/test/tle/integration
-          python3 -m pytest -s python/test/tle/unit
-          # flagtree hints
-          # python tutorials
-          python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
-          # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test
-          python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test
-          python3 python/tutorials/hints/04/04-low-memory-dropout.py --only_unit_test
-          python3 python/tutorials/hints/05/05-layer-norm.py --only_unit_test
-          # python3 python/tutorials/hints/06/06-fused-attention.py --only_unit_test  # Error on 3.6
-          python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test
-          python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test
-          python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test
-          # flagtree tle raw
-          python3 python/tutorials/tle/raw/mlir/01-vector-add.py
-          python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py
-          python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py
-          python3 python/tutorials/tle/raw/mlir/04-hello-world.py
-          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton
-          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle
-          python3 python/tutorials/tle/raw/mlir/06-test-vassert.py
-          # flagtree tle cuda
-          # TODO: These tests are currently skipped because the CLANG environment variable cannot be set.
-          # python3 python/tutorials/tle/raw/cuda/01-vector-add.py
-          # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py
-          # python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py

From 397cf64da2c31e77a77fff0bece76ad0ffbf3365 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 10:34:12 +0800
Subject: [PATCH 06/41] [TLE]Modify the test file of glu

---
 python/tutorials/tle/05-glu.py | 33 ++-------------------------------
 1 file changed, 2 insertions(+), 31 deletions(-)

diff --git a/python/tutorials/tle/05-glu.py b/python/tutorials/tle/05-glu.py
index 7635177fe5..e228721b60 100644
--- a/python/tutorials/tle/05-glu.py
+++ b/python/tutorials/tle/05-glu.py
@@ -1,23 +1,3 @@
-"""
-GLU with Triton (TLE Tutorial style) — autotune fixed
-======================================================
-
-This is test_122_v1 with the baseline BLOCK_D properly autotuned per
-(N, D, dtype). Two key fixes vs the previous attempts:
-
-  1. PRE-COMPILE pass: every BLOCK_D candidate is warmed/compiled BEFORE
-     do_bench is called for timing. This avoids the trap where the first
-     compile latency of a large BLOCK_D (like 1024) is counted into the
-     do_bench result, making it look slower than BLOCK_D=256 and forcing
-     a fallback to the initial value.
-
-  2. NO swallowed exceptions. If a BLOCK_D truly cannot compile (e.g.
-     out-of-shared-memory), the error is printed so it's visible.
-
-GLU(x) = x[:, :D] * sigmoid(x[:, D:])
-where x has shape (N, 2D).
-"""
-
 # %%
 # Setup
 # -----
@@ -38,10 +18,6 @@ def _next_pow2(x: int) -> int:
     return 1 if x <= 1 else 2 ** math.ceil(math.log2(x))
 
 
-# %%
-# Kernels (Triton baseline — test_122_v1 style: chunk + a/b pointers)
-# -------------------------------------------------------------------
-
 
 @triton.jit
 def glu_baseline(
@@ -86,9 +62,9 @@ def glu_extract_static(
                    mask=mask, other=0.0)
 
     a_tile = tle.extract_tile(halo, index=[0],
-                              tile_shape=[D_P2], strides=[D_P2])
+                              tile_shape=[D_P2])
     b_tile = tle.extract_tile(halo, index=[1],
-                              tile_shape=[D_P2], strides=[D_P2])
+                              tile_shape=[D_P2])
 
     a_f32     = a_tile.to(tl.float32)
     b_f32     = b_tile.to(tl.float32)
@@ -260,11 +236,6 @@ def run_correctness(N, D, dtype, BLOCK_D=None):
     sys.exit(0)
 
 
-# %%
-# Benchmark (test_122_v1 perf_report style, three providers)
-# ----------------------------------------------------------
-
-
 _BENCH_PROVIDERS = ["triton", "tle", "torch"]
 _BENCH_NAMES     = ["Triton (baseline, autotuned)",
                     "TLE (extract_tile static)",

From f524ebaf281b291f18fe4abf4075a49e85b41ad0 Mon Sep 17 00:00:00 2001
From: flagtree-bot <flagtree_ai@163.com>
Date: Thu, 4 Jun 2026 02:42:00 +0000
Subject: [PATCH 07/41] Apply code-format changes

---
 .../TritonToTritonGPUPass.cpp                 |  15 +-
 .../tle/unit/test_insert_tile_static_index.py |   7 +-
 .../triton/experimental/tle/language/core.py  |  23 +-
 .../experimental/tle/language/gpu/core.py     |   1 -
 .../experimental/tle/language/gpu/semantic.py |  18 +-
 python/tutorials/tle/05-glu.py                | 144 ++--
 python/tutorials/tle/06-2D_Depthwise_Conv.py  | 160 ++--
 python/tutorials/tle/07-causal-conv1d.py      | 775 +++++++-----------
 python/tutorials/tle/08-rope.py               | 244 +++---
 third_party/tle/dialect/lib/IR/Ops.cpp        |  15 +-
 .../lib/Transforms/ExtractTileToLLVM.cpp      |  21 +-
 .../lib/Transforms/InsertTileToLLVM.cpp       |  11 +-
 third_party/tle/triton_tle.cc                 |  17 +-
 13 files changed, 671 insertions(+), 780 deletions(-)

diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
index f1da3cb66c..f7306536eb 100644
--- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
@@ -882,11 +882,11 @@ class TleExtractTileOpPattern : public OpConversionPattern<tle::ExtractTileOp> {
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
-    auto stridesAttr = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
-    op->getAttr("strides"));
+    auto stridesAttr =
+        mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(op->getAttr("strides"));
     auto newOp = rewriter.replaceOpWithNewOp<tle::ExtractTileOp>(
-    op, retType, adaptor.getSrc(), adaptor.getIndex(), stridesAttr);
-   
+        op, retType, adaptor.getSrc(), adaptor.getIndex(), stridesAttr);
+
     if (auto tileShapeAttr = op->getAttr("tile_shape"))
       newOp->setAttr("tile_shape", tileShapeAttr);
     addNamedAttrs(newOp, adaptor.getAttributes());
@@ -926,11 +926,12 @@ class TleInsertTileOpPattern : public OpConversionPattern<tle::InsertTileOp> {
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
-    auto stridesAttr = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
-        op->getAttr("strides"));
+    auto stridesAttr =
+        mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(op->getAttr("strides"));
 
     auto newOp = rewriter.replaceOpWithNewOp<tle::InsertTileOp>(
-      op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(), stridesAttr);
+        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(),
+        stridesAttr);
 
     addNamedAttrs(newOp, adaptor.getAttributes());
 
diff --git a/python/test/tle/unit/test_insert_tile_static_index.py b/python/test/tle/unit/test_insert_tile_static_index.py
index adfe725502..2fc28a5c59 100644
--- a/python/test/tle/unit/test_insert_tile_static_index.py
+++ b/python/test/tle/unit/test_insert_tile_static_index.py
@@ -32,7 +32,6 @@ def insert_tile_stride_kernel(
     tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], z)
 
 
-
 @triton.jit
 def insert_tile_kernel(
     x_ptr,
@@ -103,14 +102,16 @@ def test_insert_tile_with_stride():
     y = (10000 + torch.arange(TM * TN, device="cuda", dtype=torch.float32)).reshape(TM, TN)
     out = torch.empty_like(x)
 
-    print(f"Running insert_tile kernel with stride: x={M}x{N}, tile={TM}x{TN}, stride={SM}x{SN}, index=[{idx_m},{idx_n}]...")
+    print(
+        f"Running insert_tile kernel with stride: x={M}x{N}, tile={TM}x{TN}, stride={SM}x{SN}, index=[{idx_m},{idx_n}]..."
+    )
     insert_tile_stride_kernel[(1, )](x, y, out, M, N, TM, TN, SM, SN, idx_m, idx_n)
     print("Kernel executed.\n")
 
     expected = x.clone()
     start_m = idx_m * SM
     start_n = idx_n * SN
-    expected[start_m:start_m+TM, start_n:start_n+TN] = y
+    expected[start_m:start_m + TM, start_n:start_n + TN] = y
 
     max_abs_diff = (out - expected).abs().max().item()
     print(f"max_abs_diff = {max_abs_diff}")
diff --git a/python/triton/experimental/tle/language/core.py b/python/triton/experimental/tle/language/core.py
index 4fcec3c576..33aaf5dee8 100644
--- a/python/triton/experimental/tle/language/core.py
+++ b/python/triton/experimental/tle/language/core.py
@@ -113,9 +113,10 @@ def cumsum(input, axis=0, reverse=False, dtype: tl.constexpr = None, _semantic=N
 
 
 # ============================================================================
-# extract_tile helper functions 
+# extract_tile helper functions
 # ============================================================================
 
+
 def _try_unwrap_int(val):
     """
     Try to unwrap val as a Python int.
@@ -222,7 +223,7 @@ def _linearize_dynamic_multidim_index(index_tuple, src_shape, tile_shape_ints, _
     strides_eff = strides_ints if strides_ints else tile_shape_ints
     grid = []
     for i, (s, t) in builtins.enumerate(builtins.zip(src_shape, tile_shape_ints)):
-        grid.append((s - t) // strides_eff[i] + 1)    
+        grid.append((s - t) // strides_eff[i] + 1)
     # compute strides
     strides = [1] * len(grid)
     acc = 1
@@ -368,15 +369,13 @@ def extract_tile(
     if not is_dynamic:
         strides_eff = strides_ints if strides_ints else tile_shape_ints
         if all(isinstance(s, int) for s in src_shape):
-            for i, (s, t, st) in builtins.enumerate(
-                    builtins.zip(src_shape, tile_shape_ints, strides_eff)):
+            for i, (s, t, st) in builtins.enumerate(builtins.zip(src_shape, tile_shape_ints, strides_eff)):
                 if (s - t) < 0 or (s - t) % st != 0:
-                    raise ValueError(
-                        f"(src-tile) not divisible by stride at dim {i}: "
-                        f"src={s}, tile={t}, stride={st}")
+                    raise ValueError(f"(src-tile) not divisible by stride at dim {i}: "
+                                     f"src={s}, tile={t}, stride={st}")
             total_tiles = 1
             for s, t, st in builtins.zip(src_shape, tile_shape_ints, strides_eff):
-                total_tiles *= (s - t) // st + 1 
+                total_tiles *= (s - t) // st + 1
             if index_value < 0 or index_value >= total_tiles:
                 raise ValueError(f"index {index_value} out of range [0, {total_tiles})")
 
@@ -397,7 +396,8 @@ def extract_tile(
             # Static index: encode compile-time constant as IR constant
             index_ir = _semantic._convert_to_ir_values([index_value], require_i64=False)[0]
 
-        output = _semantic.builder.create_extract_tile(x.handle, index_ir, tile_shape_ints, strides_ints or tile_shape_ints)
+        output = _semantic.builder.create_extract_tile(x.handle, index_ir, tile_shape_ints, strides_ints
+                                                       or tile_shape_ints)
         block_type = tl.block_type(x.type.element_ty, tile_shape_ints)
         return tl.tensor(output, block_type)
     except Exception as e:
@@ -459,9 +459,8 @@ def insert_tile(
             raise ValueError(f"Stride dimension {i} must be positive, got {stride_dim}")
         remainder = src_dim - tile_dim
         if remainder < 0 or remainder % stride_dim != 0:
-            raise ValueError(
-                f"(src-tile) not divisible by stride at dim {i}: "
-                f"src={src_dim}, tile={tile_dim}, stride={stride_dim}")
+            raise ValueError(f"(src-tile) not divisible by stride at dim {i}: "
+                             f"src={src_dim}, tile={tile_dim}, stride={stride_dim}")
         grid.append(remainder // stride_dim + 1)
 
     # Parse index: dynamic scalar tensor or static scalar/multi-dim.
diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py
index b769797c34..d913c33755 100644
--- a/python/triton/experimental/tle/language/gpu/core.py
+++ b/python/triton/experimental/tle/language/gpu/core.py
@@ -504,7 +504,6 @@ def tmacopy(
         return tmacopy(src, dst, direction, shape, offsets, _semantic)
 
 
-
 def _expand_index_to_shape(index: tl.tensor, shape: Sequence[int], axis: int, _semantic) -> tl.tensor:
     idx = index
     for _ in builtins.range(axis):
diff --git a/python/triton/experimental/tle/language/gpu/semantic.py b/python/triton/experimental/tle/language/gpu/semantic.py
index 81543db3a3..c196926b87 100644
--- a/python/triton/experimental/tle/language/gpu/semantic.py
+++ b/python/triton/experimental/tle/language/gpu/semantic.py
@@ -96,7 +96,8 @@ def validate_local_pointer_buffer(self, buffer: tle.buffered_tensor) -> None:
         if not isinstance(buffer, tle.buffered_tensor):
             raise TLESemanticError(f"Buffer must be tle.buffered_tensor, but got {type(buffer)}", "local_ptr")
 
-    def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]], strides=None) -> None:
+    def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]],
+                                     strides=None) -> None:
         """
 
         """
@@ -109,7 +110,7 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
             raise TLESemanticError("Index cannot be None", "extract_tile")
         if not tile_shape:
             raise TLESemanticError("tile_shape cannot be empty", "extract_tile")
-        src_shape = list(src.type.shape)         
+        src_shape = list(src.type.shape)
 
         # Check 3: unpack and validate type
         tile_shape_unwrapped = [s.value if hasattr(s, 'value') else s for s in tile_shape]
@@ -118,8 +119,8 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
         if any(s <= 0 for s in strides_eff):
             raise TLESemanticError("All strides must be positive", "extract_tile")
         if len(strides_eff) != len(src_shape):
-            raise TLESemanticError("strides rank must match source rank", "extract_tile")  
-      
+            raise TLESemanticError("strides rank must match source rank", "extract_tile")
+
         # Check if every dim in tile_shape is int or constexpr-like
         if any(not isinstance(s, int) for s in tile_shape_unwrapped):
             raise TLESemanticError("All tile_shape dims must be int or constexpr", "extract_tile")
@@ -129,7 +130,8 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
             raise TLESemanticError("All tile_shape dims must be positive", "extract_tile")
 
         # Check 5: dimension match
-   #     src_shape = list(src.type.shape)
+
+#     src_shape = list(src.type.shape)
 
         if len(tile_shape_unwrapped) != len(src_shape):
             raise TLESemanticError(
@@ -207,7 +209,8 @@ def validate_insert_tile_params(self, src: tl.tensor, tile: tl.tensor, index, st
             raise TLESemanticError("strides rank must match source rank", "insert_tile")
 
         grid = []
-        for i, (src_dim, tile_dim, stride_dim) in enumerate(zip(src_shape_unwrapped, tile_shape_unwrapped, strides_eff)):
+        for i, (src_dim, tile_dim, stride_dim) in enumerate(zip(src_shape_unwrapped, tile_shape_unwrapped,
+                                                                strides_eff)):
             if stride_dim <= 0:
                 raise TLESemanticError(f"Stride dimension {i} must be positive, got {stride_dim}", "insert_tile")
             remainder = src_dim - tile_dim
@@ -239,7 +242,8 @@ def validate_insert_tile_params(self, src: tl.tensor, tile: tl.tensor, index, st
             if val >= total_tiles:
                 raise TLESemanticError(f"Linear index {val} out of bounds for total tiles {total_tiles}", "insert_tile")
 
-    def analyze_extract_tile_operation(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]], strides=None) -> None:
+    def analyze_extract_tile_operation(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]],
+                                       strides=None) -> None:
         """Analyze extract_tile operation semantics"""
         self.validate_extract_tile_params(src, index, tile_shape, strides)
 
diff --git a/python/tutorials/tle/05-glu.py b/python/tutorials/tle/05-glu.py
index e228721b60..2c0f631321 100644
--- a/python/tutorials/tle/05-glu.py
+++ b/python/tutorials/tle/05-glu.py
@@ -15,14 +15,18 @@
 
 
 def _next_pow2(x: int) -> int:
-    return 1 if x <= 1 else 2 ** math.ceil(math.log2(x))
-
+    return 1 if x <= 1 else 2**math.ceil(math.log2(x))
 
 
 @triton.jit
 def glu_baseline(
-    a_ptr, b_ptr, out_ptr, D,
-    stride_an, stride_bn, stride_outn,
+    a_ptr,
+    b_ptr,
+    out_ptr,
+    D,
+    stride_an,
+    stride_bn,
+    stride_outn,
     BLOCK_D: tl.constexpr,
 ):
     pid_n = tl.program_id(0)
@@ -31,15 +35,12 @@ def glu_baseline(
     offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
     mask_d = offs_d < D
 
-    a = tl.load(a_ptr + pid_n * stride_an + offs_d,
-                mask=mask_d, other=0.0).to(tl.float32)
-    b = tl.load(b_ptr + pid_n * stride_bn + offs_d,
-                mask=mask_d, other=0.0).to(tl.float32)
+    a = tl.load(a_ptr + pid_n * stride_an + offs_d, mask=mask_d, other=0.0).to(tl.float32)
+    b = tl.load(b_ptr + pid_n * stride_bn + offs_d, mask=mask_d, other=0.0).to(tl.float32)
 
     result = a * (1.0 / (1.0 + tl.exp(-b)))
 
-    tl.store(out_ptr + pid_n * stride_outn + offs_d,
-             result.to(out_ptr.dtype.element_ty), mask=mask_d)
+    tl.store(out_ptr + pid_n * stride_outn + offs_d, result.to(out_ptr.dtype.element_ty), mask=mask_d)
 
 
 # %%
@@ -49,32 +50,31 @@ def glu_baseline(
 
 @triton.jit
 def glu_extract_static(
-    x_ptr, out_ptr, D,
-    stride_xn, stride_outn,
-    D_P2:  tl.constexpr,
+    x_ptr,
+    out_ptr,
+    D,
+    stride_xn,
+    stride_outn,
+    D_P2: tl.constexpr,
     D2_P2: tl.constexpr,
 ):
     pid_n = tl.program_id(0)
 
     offs = tl.arange(0, D2_P2)
     mask = offs < (D * 2)
-    halo = tl.load(x_ptr + pid_n * stride_xn + offs,
-                   mask=mask, other=0.0)
+    halo = tl.load(x_ptr + pid_n * stride_xn + offs, mask=mask, other=0.0)
 
-    a_tile = tle.extract_tile(halo, index=[0],
-                              tile_shape=[D_P2])
-    b_tile = tle.extract_tile(halo, index=[1],
-                              tile_shape=[D_P2])
+    a_tile = tle.extract_tile(halo, index=[0], tile_shape=[D_P2])
+    b_tile = tle.extract_tile(halo, index=[1], tile_shape=[D_P2])
 
-    a_f32     = a_tile.to(tl.float32)
-    b_f32     = b_tile.to(tl.float32)
+    a_f32 = a_tile.to(tl.float32)
+    b_f32 = b_tile.to(tl.float32)
     sigmoid_b = 1.0 / (1.0 + tl.exp(-b_f32))
-    result    = a_f32 * sigmoid_b
+    result = a_f32 * sigmoid_b
 
     offs_d = tl.arange(0, D_P2)
     mask_d = offs_d < D
-    tl.store(out_ptr + pid_n * stride_outn + offs_d,
-             result.to(out_ptr.dtype.element_ty), mask=mask_d)
+    tl.store(out_ptr + pid_n * stride_outn + offs_d, result.to(out_ptr.dtype.element_ty), mask=mask_d)
 
 
 # %%
@@ -84,13 +84,18 @@ def glu_extract_static(
 
 def triton_glu(x, BLOCK_D):
     N, D2 = x.shape
-    D     = D2 // 2
-    out   = torch.empty((N, D), device=x.device, dtype=x.dtype)
-    a, b  = torch.chunk(x, 2, dim=-1)
-    grid  = (N, triton.cdiv(D, BLOCK_D))
+    D = D2 // 2
+    out = torch.empty((N, D), device=x.device, dtype=x.dtype)
+    a, b = torch.chunk(x, 2, dim=-1)
+    grid = (N, triton.cdiv(D, BLOCK_D))
     glu_baseline[grid](
-        a, b, out, D,
-        a.stride(0), b.stride(0), out.stride(0),
+        a,
+        b,
+        out,
+        D,
+        a.stride(0),
+        b.stride(0),
+        out.stride(0),
         BLOCK_D=BLOCK_D,
     )
     return out
@@ -98,13 +103,16 @@ def triton_glu(x, BLOCK_D):
 
 def tle_glu(x):
     N, D2 = x.shape
-    D     = D2 // 2
-    out   = torch.empty((N, D), device=x.device, dtype=x.dtype)
-    d_p2  = _next_pow2(D)
+    D = D2 // 2
+    out = torch.empty((N, D), device=x.device, dtype=x.dtype)
+    d_p2 = _next_pow2(D)
     d2_p2 = _next_pow2(D2)
-    glu_extract_static[(N,)](
-        x, out, D,
-        x.stride(0), out.stride(0),
+    glu_extract_static[(N, )](
+        x,
+        out,
+        D,
+        x.stride(0),
+        out.stride(0),
         D_P2=d_p2,
         D2_P2=d2_p2,
     )
@@ -122,13 +130,11 @@ def torch_glu(x):
 # Without phase 1, the first-compile time of BLOCK_D=1024 gets folded into
 # do_bench and the autotuner mis-elects BLOCK_D=256 as "fastest".
 
-
 _BLOCK_D_CANDIDATES = [64, 128, 256, 512, 1024]
 _block_d_cache: dict = {}
 
 
-def best_block_d(N: int, D: int, dtype: torch.dtype,
-                 verbose: bool = False) -> int:
+def best_block_d(N: int, D: int, dtype: torch.dtype, verbose: bool = False) -> int:
     """Pick the fastest BLOCK_D for the baseline on (N, D, dtype). Cached."""
     key = (N, D, dtype)
     if key in _block_d_cache:
@@ -161,9 +167,7 @@ def best_block_d(N: int, D: int, dtype: torch.dtype,
     # We use the no-quantiles form and let it return a single float.
     best_bd, best_ms = valid_bds[0], float("inf")
     for bd in valid_bds:
-        ms = triton.testing.do_bench(
-            lambda b=bd: triton_glu(x, b),
-            warmup=25, rep=100)
+        ms = triton.testing.do_bench(lambda b=bd: triton_glu(x, b), warmup=25, rep=100)
         if verbose:
             print(f"  [autotune] N={N} D={D} {dtype} BLOCK_D={bd:>4}  "
                   f"{ms*1000:.3f} us")
@@ -200,13 +204,11 @@ def _make_input(N, D, dtype):
 def _check_static_feasibility(D: int) -> None:
     d_p2 = _next_pow2(D)
     if d_p2 != D:
-        raise ValueError(
-            f"D={D} is not pow2 (next_pow2={d_p2}); static extract_tile "
-            f"requires D == next_pow2(D)")
+        raise ValueError(f"D={D} is not pow2 (next_pow2={d_p2}); static extract_tile "
+                         f"requires D == next_pow2(D)")
     if d_p2 < 256:
-        raise ValueError(
-            f"D_P2={d_p2} < 256; cannot guarantee divisibility by "
-            f"ctaTile(<=128)")
+        raise ValueError(f"D_P2={d_p2} < 256; cannot guarantee divisibility by "
+                         f"ctaTile(<=128)")
 
 
 def run_correctness(N, D, dtype, BLOCK_D=None):
@@ -215,16 +217,14 @@ def run_correctness(N, D, dtype, BLOCK_D=None):
     if BLOCK_D is None:
         BLOCK_D = best_block_d(N, D, dtype)
 
-    ref      = torch_glu(x.float()).to(dtype)
+    ref = torch_glu(x.float()).to(dtype)
     y_triton = triton_glu(x, BLOCK_D)
-    y_tle    = tle_glu(x)
+    y_tle = tle_glu(x)
 
     atol = 1e-2 if dtype == torch.float16 else 1e-4
     rtol = 1e-3 if dtype == torch.float16 else 1e-4
-    torch.testing.assert_close(y_triton.float(), ref.float(),
-                               rtol=rtol, atol=atol)
-    torch.testing.assert_close(y_tle.float(), ref.float(),
-                               rtol=rtol, atol=atol)
+    torch.testing.assert_close(y_triton.float(), ref.float(), rtol=rtol, atol=atol)
+    torch.testing.assert_close(y_tle.float(), ref.float(), rtol=rtol, atol=atol)
     dt = "fp16" if dtype == torch.float16 else "fp32"
     print(f"  [OK]  N={N} D={D} {dt}  BLOCK_D={BLOCK_D}")
 
@@ -235,12 +235,9 @@ def run_correctness(N, D, dtype, BLOCK_D=None):
             run_correctness(_N, _D, _dtype)
     sys.exit(0)
 
-
 _BENCH_PROVIDERS = ["triton", "tle", "torch"]
-_BENCH_NAMES     = ["Triton (baseline, autotuned)",
-                    "TLE (extract_tile static)",
-                    "Torch (F.glu)"]
-_BENCH_STYLES    = [("blue", "-"), ("orange", "-"), ("green", "-")]
+_BENCH_NAMES = ["Triton (baseline, autotuned)", "TLE (extract_tile static)", "Torch (F.glu)"]
+_BENCH_STYLES = [("blue", "-"), ("orange", "-"), ("green", "-")]
 
 
 @triton.testing.perf_report(
@@ -264,15 +261,12 @@ def benchmark(N, D, provider, dtype_str):
     quantiles = [0.5, 0.2, 0.8]
 
     if provider == "torch":
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: torch_glu(x), quantiles=quantiles)
+        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch_glu(x), quantiles=quantiles)
     elif provider == "tle":
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: tle_glu(x), quantiles=quantiles)
+        ms, min_ms, max_ms = triton.testing.do_bench(lambda: tle_glu(x), quantiles=quantiles)
     else:  # triton baseline with autotuned BLOCK_D
         BLOCK_D = best_block_d(N, D, dtype)
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: triton_glu(x, BLOCK_D), quantiles=quantiles)
+        ms, min_ms, max_ms = triton.testing.do_bench(lambda: triton_glu(x, BLOCK_D), quantiles=quantiles)
 
     return ms, max_ms, min_ms
 
@@ -284,17 +278,12 @@ def benchmark(N, D, provider, dtype_str):
 
 def main(argv=None):
     parser = argparse.ArgumentParser()
-    parser.add_argument("--N", type=int, default=4096,
-                        help="rows for correctness")
-    parser.add_argument("--D", type=int, default=1024,
-                        help="last-dim half-size for correctness "
-                             "(pow2, >=256)")
-    parser.add_argument("--dtype", type=str, default="float32",
-                        choices=["float16", "float32", "fp16", "fp32"])
-    parser.add_argument("--show_plots", action="store_true",
-                        help="show plots in benchmark")
-    parser.add_argument("--verbose_autotune", action="store_true",
-                        help="print autotune timing for each BLOCK_D")
+    parser.add_argument("--N", type=int, default=4096, help="rows for correctness")
+    parser.add_argument("--D", type=int, default=1024, help="last-dim half-size for correctness "
+                        "(pow2, >=256)")
+    parser.add_argument("--dtype", type=str, default="float32", choices=["float16", "float32", "fp16", "fp32"])
+    parser.add_argument("--show_plots", action="store_true", help="show plots in benchmark")
+    parser.add_argument("--verbose_autotune", action="store_true", help="print autotune timing for each BLOCK_D")
     args = parser.parse_args(argv)
 
     dtype = _get_dtype(args.dtype)
@@ -311,8 +300,7 @@ def main(argv=None):
     run_correctness(args.N, args.D, dtype)
     print()
 
-    benchmark.run(print_data=True, show_plots=args.show_plots,
-                  dtype_str=args.dtype.replace("float", "fp"))
+    benchmark.run(print_data=True, show_plots=args.show_plots, dtype_str=args.dtype.replace("float", "fp"))
 
 
 if __name__ == "__main__":
diff --git a/python/tutorials/tle/06-2D_Depthwise_Conv.py b/python/tutorials/tle/06-2D_Depthwise_Conv.py
index b4de39525b..bf29061f9d 100644
--- a/python/tutorials/tle/06-2D_Depthwise_Conv.py
+++ b/python/tutorials/tle/06-2D_Depthwise_Conv.py
@@ -33,7 +33,7 @@
 
 
 def _next_pow2(x: int) -> int:
-    return 1 if x <= 1 else 2 ** math.ceil(math.log2(x))
+    return 1 if x <= 1 else 2**math.ceil(math.log2(x))
 
 
 # %%
@@ -43,22 +43,31 @@ def _next_pow2(x: int) -> int:
 
 @triton.jit
 def conv2d_baseline(
-    inp_ptr, wgt_ptr, out_ptr,
-    H, W, C, OH, OW,
-    KH: tl.constexpr, KW: tl.constexpr,
-    TILE_OH: tl.constexpr, TILE_OW: tl.constexpr, TILE_C: tl.constexpr,
+    inp_ptr,
+    wgt_ptr,
+    out_ptr,
+    H,
+    W,
+    C,
+    OH,
+    OW,
+    KH: tl.constexpr,
+    KW: tl.constexpr,
+    TILE_OH: tl.constexpr,
+    TILE_OW: tl.constexpr,
+    TILE_C: tl.constexpr,
 ):
     pid_oh = tl.program_id(0)
     pid_ow = tl.program_id(1)
-    pid_c  = tl.program_id(2)
+    pid_c = tl.program_id(2)
     oh0 = pid_oh * TILE_OH
     ow0 = pid_ow * TILE_OW
-    c0  = pid_c  * TILE_C
+    c0 = pid_c * TILE_C
 
-    offs_c  = c0  + tl.arange(0, TILE_C)
+    offs_c = c0 + tl.arange(0, TILE_C)
     offs_oh = oh0 + tl.arange(0, TILE_OH)
     offs_ow = ow0 + tl.arange(0, TILE_OW)
-    c_mask  = offs_c < C
+    c_mask = offs_c < C
 
     acc = tl.zeros((TILE_OH, TILE_OW, TILE_C), dtype=tl.float32)
 
@@ -70,21 +79,13 @@ def conv2d_baseline(
             iw = offs_ow + kw
             ih_ok = (ih >= 0) & (ih < H)
             iw_ok = (iw >= 0) & (iw < W)
-            inp_ptrs = (inp_ptr
-                        + ih[:, None, None] * (W * C)
-                        + iw[None, :, None] * C
-                        + offs_c[None, None, :])
+            inp_ptrs = (inp_ptr + ih[:, None, None] * (W * C) + iw[None, :, None] * C + offs_c[None, None, :])
             mask = ih_ok[:, None, None] & iw_ok[None, :, None] & c_mask[None, None, :]
             x = tl.load(inp_ptrs, mask=mask, other=0.0)
             acc += x * w[None, None, :]
 
-    out_ptrs = (out_ptr
-                + offs_oh[:, None, None] * (OW * C)
-                + offs_ow[None, :, None] * C
-                + offs_c[None, None, :])
-    out_mask = ((offs_oh[:, None, None] < OH)
-                & (offs_ow[None, :, None] < OW)
-                & c_mask[None, None, :])
+    out_ptrs = (out_ptr + offs_oh[:, None, None] * (OW * C) + offs_ow[None, :, None] * C + offs_c[None, None, :])
+    out_mask = ((offs_oh[:, None, None] < OH) & (offs_ow[None, :, None] < OW) & c_mask[None, None, :])
     tl.store(out_ptrs, acc, mask=out_mask)
 
 
@@ -95,36 +96,44 @@ def conv2d_baseline(
 
 @triton.jit
 def conv2d_extract_static(
-    inp_ptr, wgt_ptr, out_ptr,
-    H, W, C, OH, OW,
-    KH: tl.constexpr, KW: tl.constexpr,
-    TILE_OH: tl.constexpr, TILE_OW: tl.constexpr, TILE_C: tl.constexpr,
-    HALO_H_P2: tl.constexpr, HALO_W_P2: tl.constexpr,
-    HALO_H:    tl.constexpr, HALO_W:    tl.constexpr,
+    inp_ptr,
+    wgt_ptr,
+    out_ptr,
+    H,
+    W,
+    C,
+    OH,
+    OW,
+    KH: tl.constexpr,
+    KW: tl.constexpr,
+    TILE_OH: tl.constexpr,
+    TILE_OW: tl.constexpr,
+    TILE_C: tl.constexpr,
+    HALO_H_P2: tl.constexpr,
+    HALO_W_P2: tl.constexpr,
+    HALO_H: tl.constexpr,
+    HALO_W: tl.constexpr,
 ):
     pid_oh = tl.program_id(0)
     pid_ow = tl.program_id(1)
-    pid_c  = tl.program_id(2)
+    pid_c = tl.program_id(2)
     oh0 = pid_oh * TILE_OH
     ow0 = pid_ow * TILE_OW
-    c0  = pid_c  * TILE_C
+    c0 = pid_c * TILE_C
 
-    offs_c     = c0  + tl.arange(0, TILE_C)
-    c_mask     = offs_c < C
+    offs_c = c0 + tl.arange(0, TILE_C)
+    c_mask = offs_c < C
     offs_ih_p2 = oh0 + tl.arange(0, HALO_H_P2)
     offs_iw_p2 = ow0 + tl.arange(0, HALO_W_P2)
 
     ih_ok = (offs_ih_p2 < oh0 + HALO_H) & (offs_ih_p2 >= 0) & (offs_ih_p2 < H)
     iw_ok = (offs_iw_p2 < ow0 + HALO_W) & (offs_iw_p2 >= 0) & (offs_iw_p2 < W)
 
-    halo_ptrs = (inp_ptr
-                 + offs_ih_p2[:, None, None] * (W * C)
-                 + offs_iw_p2[None, :, None] * C
-                 + offs_c[None, None, :])
+    halo_ptrs = (inp_ptr + offs_ih_p2[:, None, None] * (W * C) + offs_iw_p2[None, :, None] * C + offs_c[None, None, :])
     halo_mask = ih_ok[:, None, None] & iw_ok[None, :, None] & c_mask[None, None, :]
     halo = tl.load(halo_ptrs, mask=halo_mask, other=0.0)
 
-    acc     = tl.zeros((TILE_OH, TILE_OW, TILE_C), dtype=tl.float32)
+    acc = tl.zeros((TILE_OH, TILE_OW, TILE_C), dtype=tl.float32)
     offs_oh = oh0 + tl.arange(0, TILE_OH)
     offs_ow = ow0 + tl.arange(0, TILE_OW)
 
@@ -140,13 +149,8 @@ def conv2d_extract_static(
             )
             acc += patch * w[None, None, :]
 
-    out_ptrs = (out_ptr
-                + offs_oh[:, None, None] * (OW * C)
-                + offs_ow[None, :, None] * C
-                + offs_c[None, None, :])
-    out_mask = ((offs_oh[:, None, None] < OH)
-                & (offs_ow[None, :, None] < OW)
-                & c_mask[None, None, :])
+    out_ptrs = (out_ptr + offs_oh[:, None, None] * (OW * C) + offs_ow[None, :, None] * C + offs_c[None, None, :])
+    out_mask = ((offs_oh[:, None, None] < OH) & (offs_ow[None, :, None] < OW) & c_mask[None, None, :])
     tl.store(out_ptrs, acc, mask=out_mask)
 
 
@@ -167,8 +171,19 @@ def triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC):
     out = torch.empty((OH, OW, C), device=inp.device, dtype=inp.dtype)
     grid = (math.ceil(OH / TOH), math.ceil(OW / TOW), math.ceil(C / TC))
     conv2d_baseline[grid](
-        inp, wgt, out, H, W, C, OH, OW,
-        KH=KH, KW=KW, TILE_OH=TOH, TILE_OW=TOW, TILE_C=TC,
+        inp,
+        wgt,
+        out,
+        H,
+        W,
+        C,
+        OH,
+        OW,
+        KH=KH,
+        KW=KW,
+        TILE_OH=TOH,
+        TILE_OW=TOW,
+        TILE_C=TC,
     )
     return out
 
@@ -180,10 +195,23 @@ def tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC):
     hh, hw, hhp2, hwp2 = _halo_dims(KH, KW, TOH, TOW)
     grid = (math.ceil(OH / TOH), math.ceil(OW / TOW), math.ceil(C / TC))
     conv2d_extract_static[grid](
-        inp, wgt, out, H, W, C, OH, OW,
-        KH=KH, KW=KW, TILE_OH=TOH, TILE_OW=TOW, TILE_C=TC,
-        HALO_H_P2=hhp2, HALO_W_P2=hwp2,
-        HALO_H=hh, HALO_W=hw,
+        inp,
+        wgt,
+        out,
+        H,
+        W,
+        C,
+        OH,
+        OW,
+        KH=KH,
+        KW=KW,
+        TILE_OH=TOH,
+        TILE_OW=TOW,
+        TILE_C=TC,
+        HALO_H_P2=hhp2,
+        HALO_W_P2=hwp2,
+        HALO_H=hh,
+        HALO_W=hw,
     )
     return out
 
@@ -221,27 +249,23 @@ def run_correctness(H, W, C, KH, KW, TOH, TOW, TC, dtype):
     inp, wgt = _make_input(H, W, C, KH, KW, dtype)
     ref = torch_dwconv(inp, wgt, KH, KW)
     y_triton = triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC)
-    y_tle    = tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC)
+    y_tle = tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC)
     atol = 1e-2 if dtype == torch.float16 else 1e-3
     torch.testing.assert_close(y_triton.float(), ref.float(), rtol=atol, atol=atol)
-    torch.testing.assert_close(y_tle.float(),    ref.float(), rtol=atol, atol=atol)
+    torch.testing.assert_close(y_tle.float(), ref.float(), rtol=atol, atol=atol)
     print("Correctness check passed (triton/tle).")
 
 
 if "--only_unit_test" in sys.argv:
-    _args = argparse.Namespace(H=128, W=128, C=64, KH=5, KW=5,
-                               TOH=4, TOW=4, TC=64, dtype="float32")
+    _args = argparse.Namespace(H=128, W=128, C=64, KH=5, KW=5, TOH=4, TOW=4, TC=64, dtype="float32")
     _dtype = _get_dtype(_args.dtype)
-    run_correctness(_args.H, _args.W, _args.C, _args.KH, _args.KW,
-                    _args.TOH, _args.TOW, _args.TC, _dtype)
+    run_correctness(_args.H, _args.W, _args.C, _args.KH, _args.KW, _args.TOH, _args.TOW, _args.TC, _dtype)
     sys.exit(0)
 
-
 # %%
 # Benchmark
 # ---------
 
-
 # Mirrors the FFT tutorial's "vary one dim, sweep providers" shape:
 # x-axis is spatial size H=W, y-axis is ms for {triton, tle, torch}.
 _BENCH_PROVIDERS = ["triton", "tle", "torch"]
@@ -267,14 +291,13 @@ def benchmark(H, C, KH, KW, TOH, TOW, TC, provider, dtype):
     quantiles = [0.5, 0.2, 0.8]
 
     if provider == "torch":
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: torch_dwconv(inp, wgt, KH, KW), quantiles=quantiles)
+        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch_dwconv(inp, wgt, KH, KW), quantiles=quantiles)
     elif provider == "tle":
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC), quantiles=quantiles)
+        ms, min_ms, max_ms = triton.testing.do_bench(lambda: tle_dwconv(inp, wgt, KH, KW, TOH, TOW, TC),
+                                                     quantiles=quantiles)
     else:
-        ms, min_ms, max_ms = triton.testing.do_bench(
-            lambda: triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC), quantiles=quantiles)
+        ms, min_ms, max_ms = triton.testing.do_bench(lambda: triton_dwconv(inp, wgt, KH, KW, TOH, TOW, TC),
+                                                     quantiles=quantiles)
 
     return ms, max_ms, min_ms
 
@@ -288,24 +311,21 @@ def main(argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument("--H", type=int, default=128, help="spatial H for correctness")
     parser.add_argument("--W", type=int, default=128, help="spatial W for correctness")
-    parser.add_argument("--C", type=int, default=64,  help="channels")
+    parser.add_argument("--C", type=int, default=64, help="channels")
     parser.add_argument("--KH", type=int, default=5)
     parser.add_argument("--KW", type=int, default=5)
     parser.add_argument("--TOH", type=int, default=4)
     parser.add_argument("--TOW", type=int, default=4)
-    parser.add_argument("--TC",  type=int, default=64)
-    parser.add_argument("--dtype", type=str, default="float32",
-                        choices=["float16", "float32", "fp16", "fp32"])
-    parser.add_argument("--show_plots", action="store_true",
-                        help="show plots in benchmark")
+    parser.add_argument("--TC", type=int, default=64)
+    parser.add_argument("--dtype", type=str, default="float32", choices=["float16", "float32", "fp16", "fp32"])
+    parser.add_argument("--show_plots", action="store_true", help="show plots in benchmark")
     args = parser.parse_args(argv)
 
     dtype = _get_dtype(args.dtype)
 
     check_H = min(args.H, 256)
     check_W = min(args.W, 256)
-    run_correctness(check_H, check_W, args.C, args.KH, args.KW,
-                    args.TOH, args.TOW, args.TC, dtype)
+    run_correctness(check_H, check_W, args.C, args.KH, args.KW, args.TOH, args.TOW, args.TC, dtype)
 
     benchmark.run(print_data=True, show_plots=args.show_plots, dtype=dtype)
 
diff --git a/python/tutorials/tle/07-causal-conv1d.py b/python/tutorials/tle/07-causal-conv1d.py
index 0cc7633ae0..5f2eea045d 100644
--- a/python/tutorials/tle/07-causal-conv1d.py
+++ b/python/tutorials/tle/07-causal-conv1d.py
@@ -14,7 +14,7 @@
 
     # correctness + benchmark tables with speedup
     python python/tutorials/tle/07-causal-conv1d.py --benchmark
-    
+
 Both ``causal_conv1d_fn`` (varlen forward) and ``causal_conv1d_update``
 (step/update) are included.
 """
@@ -24,8 +24,6 @@
 # -----
 
 import argparse
-import math
-import sys
 
 import numpy as np
 import torch
@@ -36,11 +34,11 @@
 DEVICE = triton.runtime.driver.active.get_active_torch_device()
 PAD_SLOT_ID = -1
 
-
 # %%
 # Kernels (baseline v1)
 # ---------------------
 
+
 @triton.jit()
 def _causal_conv1d_fwd_kernel_v1(
     x_ptr,
@@ -127,23 +125,17 @@ def _causal_conv1d_fwd_kernel_v1(
     token_offset = BLOCK_M * chunk_offset
     segment_len = min(BLOCK_M, seqlen - token_offset)
 
-    x_base = (
-        x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
-    )
+    x_base = (x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim)
 
-    conv_states_input_coord = tl.load(
-        conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
-    ).to(tl.int64)
+    conv_states_input_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_cache_indices +
+                                      conv_state_init_index).to(tl.int64)
 
     if USE_PAD_SLOT:
         if conv_states_input_coord == pad_slot_id:
             return
 
-    conv_states_base = (
-        conv_states_ptr
-        + (conv_states_input_coord * stride_conv_state_seq)
-        + (idx_feats * stride_conv_state_dim)
-    )
+    conv_states_base = (conv_states_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                        (idx_feats * stride_conv_state_dim))
 
     w_base = w_ptr + (idx_feats * stride_w_dim)
 
@@ -168,42 +160,30 @@ def _causal_conv1d_fwd_kernel_v1(
                 col0 = tl.load(prior_tokens - 3 * stride_conv_state_tok, mask_w, 0.0)
         else:
             if KERNEL_WIDTH >= 2:
-                col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col0 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 3:
-                col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col1 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 4:
-                col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col2 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 5:
-                col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col3 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
 
         if state_len <= seqlen:
             idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN)
-            x_ptrs = (
-                x_ptr
-                + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
-                + (idx_feats * stride_x_dim)[None, :]
-            )
-            mask_x = (
-                (idx_tokens_last >= 0)[:, None]
-                & (idx_tokens_last < seqlen)[:, None]
-                & (idx_feats < dim)[None, :]
-            )
+            x_ptrs = (x_ptr + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] +
+                      (idx_feats * stride_x_dim)[None, :])
+            mask_x = ((idx_tokens_last >= 0)[:, None]
+                      & (idx_tokens_last < seqlen)[:, None]
+                      & (idx_feats < dim)[None, :])
             loaded_x = tl.load(x_ptrs, mask_x, 0.0)
             idx_tokens_conv = tl.arange(0, NP2_STATELEN)
 
-            conv_states_output_coord = tl.load(
-                conv_state_indices_ptr
-                + idx_seq * stride_cache_indices
-                + current_last_index
-            ).to(tl.int64)
+            conv_states_output_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_cache_indices +
+                                               current_last_index).to(tl.int64)
 
-            conv_states_ptrs_target = (
-                conv_states_ptr
-                + (conv_states_output_coord * stride_conv_state_seq)
-                + (idx_feats * stride_conv_state_dim)
-            )[None, :] + (
-                idx_tokens_conv * stride_conv_state_tok
-            )[:, None]
+            conv_states_ptrs_target = (conv_states_ptr + (conv_states_output_coord * stride_conv_state_seq) +
+                                       (idx_feats * stride_conv_state_dim))[None, :] + (idx_tokens_conv *
+                                                                                        stride_conv_state_tok)[:, None]
 
             mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
             tl.debug_barrier()
@@ -213,59 +193,38 @@ def _causal_conv1d_fwd_kernel_v1(
             if load_init_state:
                 idx_tokens_conv = tl.arange(0, NP2_STATELEN)
 
-                conv_states_ptrs_source = (
-                    conv_states_ptr
-                    + (conv_states_input_coord * stride_conv_state_seq)
-                    + (idx_feats * stride_conv_state_dim)[None, :]
-                    + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None]
-                )
-                mask = (
-                    (conv_states_input_coord < num_cache_lines)
-                    & ((idx_tokens_conv + seqlen) < state_len)[:, None]
-                    & (idx_feats < dim)[None, :]
-                )
+                conv_states_ptrs_source = (conv_states_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                                           (idx_feats * stride_conv_state_dim)[None, :] +
+                                           ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None])
+                mask = ((conv_states_input_coord < num_cache_lines)
+                        & ((idx_tokens_conv + seqlen) < state_len)[:, None]
+                        & (idx_feats < dim)[None, :])
                 conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
 
                 VAL = state_len - seqlen
-                x_ptrs = (
-                    x_base[None, :]
-                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
-                )
-                mask_x = (
-                    (idx_tokens_conv - VAL >= 0)[:, None]
-                    & (idx_tokens_conv - VAL < seqlen)[:, None]
-                    & (idx_feats < dim)[None, :]
-                )
+                x_ptrs = (x_base[None, :] + ((idx_tokens_conv - VAL) * stride_x_token)[:, None])
+                mask_x = ((idx_tokens_conv - VAL >= 0)[:, None]
+                          & (idx_tokens_conv - VAL < seqlen)[:, None]
+                          & (idx_feats < dim)[None, :])
                 loaded_x = tl.load(x_ptrs, mask_x, 0.0)
 
                 tl.debug_barrier()
                 new_conv_state = tl.where(mask, conv_state, loaded_x)
 
-                conv_states_ptrs_target = (
-                    conv_states_base
-                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
-                )
+                conv_states_ptrs_target = (conv_states_base + (idx_tokens_conv * stride_conv_state_tok)[:, None])
                 mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
                 tl.store(conv_states_ptrs_target, new_conv_state, mask)
 
             else:
                 idx_tokens_conv = tl.arange(0, NP2_STATELEN)
                 VAL = state_len - seqlen
-                x_ptrs = (
-                    x_base[None, :]
-                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
-                )
-                mask_x = (
-                    (idx_tokens_conv - VAL >= 0)[:, None]
-                    & (idx_tokens_conv - VAL < seqlen)[:, None]
-                    & (idx_feats < dim)[None, :]
-                )
+                x_ptrs = (x_base[None, :] + ((idx_tokens_conv - VAL) * stride_x_token)[:, None])
+                mask_x = ((idx_tokens_conv - VAL >= 0)[:, None]
+                          & (idx_tokens_conv - VAL < seqlen)[:, None]
+                          & (idx_feats < dim)[None, :])
                 new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
 
-                conv_states_ptrs_target = (
-                    conv_states_base
-                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
-                )
+                conv_states_ptrs_target = (conv_states_base + (idx_tokens_conv * stride_conv_state_tok)[:, None])
                 mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
                 tl.store(conv_states_ptrs_target, new_conv_state, mask)
 
@@ -289,34 +248,19 @@ def _causal_conv1d_fwd_kernel_v1(
             col0 = tl.load(prior_tokens - 3 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
 
         if (chunk_offset - 1) < n_block_to_fill:
-            idx_tokens_last = (
-                last_full_block_token_index
-                - (n_block_to_fill - chunk_offset) * B_size
-                - state_len
-            ) + tl.arange(0, NP2_STATELEN)
-            x_ptrs = (
-                x_ptr
-                + (idx_tokens_last * stride_x_token)[:, None]
-                + (idx_feats * stride_x_dim)[None, :]
-            )
+            idx_tokens_last = (last_full_block_token_index -
+                               (n_block_to_fill - chunk_offset) * B_size - state_len) + tl.arange(0, NP2_STATELEN)
+            x_ptrs = (x_ptr + (idx_tokens_last * stride_x_token)[:, None] + (idx_feats * stride_x_dim)[None, :])
             mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[None, :]
             loaded_x = tl.load(x_ptrs, mask_x, 0.0)
             idx_tokens_conv = tl.arange(0, NP2_STATELEN)
 
-            conv_states_output_coord = tl.load(
-                conv_state_indices_ptr
-                + idx_seq * stride_cache_indices
-                + current_first_index
-                + (chunk_offset - 1)
-            ).to(tl.int64)
-
-            conv_states_ptrs_target = (
-                conv_states_ptr
-                + (conv_states_output_coord * stride_conv_state_seq)
-                + (idx_feats * stride_conv_state_dim)
-            )[None, :] + (
-                idx_tokens_conv * stride_conv_state_tok
-            )[:, None]
+            conv_states_output_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_cache_indices +
+                                               current_first_index + (chunk_offset - 1)).to(tl.int64)
+
+            conv_states_ptrs_target = (conv_states_ptr + (conv_states_output_coord * stride_conv_state_seq) +
+                                       (idx_feats * stride_conv_state_dim))[None, :] + (idx_tokens_conv *
+                                                                                        stride_conv_state_tok)[:, None]
 
             mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
             tl.debug_barrier()
@@ -327,7 +271,7 @@ def _causal_conv1d_fwd_kernel_v1(
         mask_bias = idx_feats < dim
         acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
     else:
-        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        acc_preload = tl.zeros((BLOCK_N, ), dtype=tl.float32)
 
     x_base_1d = x_base + token_offset * stride_x_token
 
@@ -390,11 +334,8 @@ def _causal_conv1d_fwd_kernel_v1(
         if SILU_ACTIVATION:
             acc = acc / (1 + tl.exp(-acc))
         mask_1d = (idx_token < segment_len) & (idx_feats < dim)
-        o_ptrs = (
-            o_ptr
-            + (sequence_start_index + token_offset + idx_token) * stride_o_token
-            + (idx_feats * stride_o_dim)
-        )
+        o_ptrs = (o_ptr + (sequence_start_index + token_offset + idx_token) * stride_o_token +
+                  (idx_feats * stride_o_dim))
         tl.store(o_ptrs, acc, mask=mask_1d)
 
 
@@ -451,9 +392,8 @@ def _causal_conv1d_update_kernel_v1(
         conv_state_init = 0
         current_last_index = 0
 
-    conv_states_input_coord = tl.load(
-        conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
-    ).to(tl.int64)
+    conv_states_input_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init).to(
+        tl.int64)
 
     if USE_PAD_SLOT:
         if conv_states_input_coord == pad_slot_id:
@@ -476,17 +416,12 @@ def _causal_conv1d_update_kernel_v1(
         return
 
     if IS_SPEC_DECODING:
-        conv_state_token_offset = (
-            tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
-        )
+        conv_state_token_offset = (tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1)
     else:
         conv_state_token_offset = 0
 
-    conv_states_base = (
-        conv_state_ptr
-        + (conv_states_input_coord * stride_conv_state_seq)
-        + (idx_feats * stride_conv_state_dim)
-    )
+    conv_states_base = (conv_state_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                        (idx_feats * stride_conv_state_dim))
     mask_w = idx_feats < dim
 
     prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok
@@ -503,48 +438,30 @@ def _causal_conv1d_update_kernel_v1(
 
     idx_tokens = tl.arange(0, NP2_STATELEN)
 
-    conv_state_ptrs_source = (
-        conv_state_ptr
-        + (conv_states_input_coord * stride_conv_state_seq)
-        + conv_state_token_offset * stride_conv_state_tok
-        + (idx_feats * stride_conv_state_dim)[None, :]
-        + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[
-            :, None
-        ]
-    )
-    mask = (
-        (conv_states_input_coord < num_cache_lines)
-        & ((idx_tokens + seqlen) < state_len)[:, None]
-        & (idx_feats < dim)[None, :]
-    )
+    conv_state_ptrs_source = (conv_state_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                              conv_state_token_offset * stride_conv_state_tok +
+                              (idx_feats * stride_conv_state_dim)[None, :] +
+                              ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[:, None])
+    mask = ((conv_states_input_coord < num_cache_lines)
+            & ((idx_tokens + seqlen) < state_len)[:, None]
+            & (idx_feats < dim)[None, :])
     conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
 
     VAL = state_len - seqlen
     x_base = x_ptr + x_offset + (idx_feats * stride_x_dim)
 
-    x_ptrs = (
-        x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None]
-    )
-    mask_x = (
-        (idx_tokens - VAL >= 0)[:, None]
-        & (idx_tokens - VAL < seqlen)[:, None]
-        & (idx_feats < dim)[None, :]
-    )
+    x_ptrs = (x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None])
+    mask_x = ((idx_tokens - VAL >= 0)[:, None] & (idx_tokens - VAL < seqlen)[:, None] & (idx_feats < dim)[None, :])
     loaded_x = tl.load(x_ptrs, mask_x, 0.0)
     tl.debug_barrier()
 
     new_conv_state = tl.where(mask, conv_state, loaded_x)
 
-    conv_states_offset = tl.load(
-        conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index
-    ).to(tl.int64)
-    conv_state_ptrs_target = (
-        conv_state_ptr
-        + (conv_states_offset * stride_conv_state_seq)
-        + (idx_feats * stride_conv_state_dim)
-    )[None, :] + (
-        idx_tokens * stride_conv_state_tok
-    )[:, None]
+    conv_states_offset = tl.load(conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index).to(
+        tl.int64)
+    conv_state_ptrs_target = (conv_state_ptr + (conv_states_offset * stride_conv_state_seq) +
+                              (idx_feats * stride_conv_state_dim))[None, :] + (idx_tokens * stride_conv_state_tok)[:,
+                                                                                                                   None]
     mask_out = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
     tl.store(conv_state_ptrs_target, new_conv_state, mask_out)
 
@@ -553,7 +470,7 @@ def _causal_conv1d_update_kernel_v1(
         mask_bias = idx_feats < dim
         acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
     else:
-        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        acc_preload = tl.zeros((BLOCK_N, ), dtype=tl.float32)
 
     w_base = w_ptr + (idx_feats * stride_w_dim)
     mask_w = idx_feats < dim
@@ -655,9 +572,7 @@ def _causal_conv1d_update_kernel_v1(
         if SILU_ACTIVATION:
             acc = acc / (1 + tl.exp(-acc))
         mask_1d = (idx_token < seqlen) & (idx_feats < dim)
-        o_ptrs = (
-            o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim)
-        )
+        o_ptrs = (o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim))
         tl.store(o_ptrs, acc, mask=mask_1d)
 
 
@@ -716,9 +631,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
     stride_conv_state_seq = stride_istate_seq
     stride_conv_state_dim = stride_istate_dim
     stride_conv_state_tok = stride_istate_token
-    state_len = (
-        KERNEL_WIDTH - 1
-    )  # can be passed via argument if it's not the same as this value
+    state_len = (KERNEL_WIDTH - 1)  # can be passed via argument if it's not the same as this value
 
     idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64)
     chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0))
@@ -759,28 +672,20 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
     token_offset = BLOCK_M * chunk_offset
     segment_len = min(BLOCK_M, seqlen - token_offset)
 
-    x_base = (
-        x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
-    )  # [BLOCK_N,]
+    x_base = (x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim)  # [BLOCK_N,]
 
-    conv_states_input_coord = tl.load(
-        conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
-    ).to(tl.int64)
+    conv_states_input_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_cache_indices +
+                                      conv_state_init_index).to(tl.int64)
 
     if USE_PAD_SLOT:  # noqa
         if conv_states_input_coord == pad_slot_id:
             return
 
-    conv_states_base = (
-        conv_states_ptr
-        + (conv_states_input_coord * stride_conv_state_seq)
-        + (idx_feats * stride_conv_state_dim)
-    )  # [BLOCK_N,]
+    conv_states_base = (conv_states_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                        (idx_feats * stride_conv_state_dim))  # [BLOCK_N,]
 
     w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
 
-
-    
     # x_tile_start = token_offset - (KERNEL_WIDTH - 1)
     # # idx_rows = tl.arange(0, X_TILE_M_P2)
     # idx_rows = tl.arange(0, triton.next_power_of_2(BLOCK_M + KERNEL_WIDTH - 1))
@@ -798,7 +703,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
     #     )
     # tile_data = tl.load(
     #         tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca"
-    #     )   
+    #     )
 
     if chunk_offset == 0:
         chunk = 0
@@ -810,20 +715,13 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
         # ----------------------------------------------------------------
         idx_rows = tl.arange(0, triton.next_power_of_2(BLOCK_M + KERNEL_WIDTH - 1))
         offs_tok = token_offset + idx_rows
-        tile_ptrs = (
-            x_ptr
-            + (sequence_start_index + offs_tok)[:, None] * stride_x_token
-            + (idx_feats * stride_x_dim)[None, :]
-        )
-        tile_mask = (
-            (offs_tok >= 0)[:, None]
-            & (offs_tok < seqlen)[:, None]
-            & (idx_feats < dim)[None, :]
-            & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None]
-        )
-        tile_data = tl.load(
-            tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca"
-        )
+        tile_ptrs = (x_ptr + (sequence_start_index + offs_tok)[:, None] * stride_x_token +
+                     (idx_feats * stride_x_dim)[None, :])
+        tile_mask = ((offs_tok >= 0)[:, None]
+                     & (offs_tok < seqlen)[:, None]
+                     & (idx_feats < dim)[None, :]
+                     & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None])
+        tile_data = tl.load(tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca")
 
         if load_init_state:
             # ----------------------------------------------------------------
@@ -835,71 +733,64 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
             #          + idx_state[:, None] * stride  [NP2_STATELEN, 1]
             #          => [NP2_STATELEN, BLOCK_N]
             idx_state = tl.arange(0, NP2_STATELEN)
-            conv_state_ptrs_full = (
-                conv_states_base[None, :]
-                + idx_state[:, None] * stride_conv_state_tok
-            )  # [NP2_STATELEN, BLOCK_N]
+            conv_state_ptrs_full = (conv_states_base[None, :] + idx_state[:, None] * stride_conv_state_tok
+                                    )  # [NP2_STATELEN, BLOCK_N]
             mask_state_full = (
-                (idx_state[:, None] < state_len) & (idx_feats[None, :] < dim)
-            )  # bool mask [NP2_STATELEN, BLOCK_N]
-            state_tile = tl.load(
-                conv_state_ptrs_full, mask_state_full, other=0.0
-            )  # [NP2_STATELEN, BLOCK_N]
+                (idx_state[:, None] < state_len) & (idx_feats[None, :] < dim))  # bool mask [NP2_STATELEN, BLOCK_N]
+            state_tile = tl.load(conv_state_ptrs_full, mask_state_full, other=0.0)  # [NP2_STATELEN, BLOCK_N]
 
             if KERNEL_WIDTH == 2:
-                col0 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
             if KERNEL_WIDTH == 3:
-                col1 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
-                col0 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col1 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
             if KERNEL_WIDTH == 4:
-                col2 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
-                col1 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
-                col0 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col2 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
+                col1 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
             if KERNEL_WIDTH == 5:
-                col3 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
-                col2 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
-                col1 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
-                col0 = tle.extract_tile(state_tile, index=[state_len - 4, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+                col3 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
+                col2 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
+                col1 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
+                col0 = tle.extract_tile(state_tile, index=[state_len - 4, 0], tile_shape=[1, BLOCK_N]).reshape(
+                    (BLOCK_N, ))
         else:
             if KERNEL_WIDTH >= 2:
-                col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col0 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 3:
-                col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col1 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 4:
-                col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col2 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 5:
-                col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
+                col3 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
 
         # STEP 2: write back conv_state
         if state_len <= seqlen:
             idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN)
-            x_ptrs = (
-                x_ptr
-                + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
-                + (idx_feats * stride_x_dim)[None, :]
-            )  # [NP2_STATELEN, BLOCK_N]
-            mask_x = (
-                (idx_tokens_last >= 0)[:, None]
-                & (idx_tokens_last < seqlen)[:, None]
-                & (idx_feats < dim)[None, :]
-            )
+            x_ptrs = (x_ptr + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] +
+                      (idx_feats * stride_x_dim)[None, :])  # [NP2_STATELEN, BLOCK_N]
+            mask_x = ((idx_tokens_last >= 0)[:, None]
+                      & (idx_tokens_last < seqlen)[:, None]
+                      & (idx_feats < dim)[None, :])
             loaded_x = tl.load(x_ptrs, mask_x, 0.0)  # [NP2_STATELEN, BLOCK_N]
 
             idx_tokens_conv = tl.arange(0, NP2_STATELEN)
 
-            conv_states_output_coord = tl.load(
-                conv_state_indices_ptr
-                + idx_seq * stride_cache_indices
-                + current_last_index
-            ).to(tl.int64)
+            conv_states_output_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_cache_indices +
+                                               current_last_index).to(tl.int64)
 
-            conv_states_ptrs_target = (
-                conv_states_ptr
-                + (conv_states_output_coord * stride_conv_state_seq)
-                + (idx_feats * stride_conv_state_dim)
-            )[None, :] + (
-                idx_tokens_conv * stride_conv_state_tok
-            )[:, None]
+            conv_states_ptrs_target = (conv_states_ptr + (conv_states_output_coord * stride_conv_state_seq) +
+                                       (idx_feats * stride_conv_state_dim))[None, :] + (idx_tokens_conv *
+                                                                                        stride_conv_state_tok)[:, None]
 
             mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
             tl.debug_barrier()
@@ -909,29 +800,20 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
             if load_init_state:
                 idx_tokens_conv = tl.arange(0, NP2_STATELEN)
                 # Shift right by seqlen to read the remaining cached history
-                conv_states_ptrs_source = (
-                    conv_states_ptr
-                    + (conv_states_input_coord * stride_conv_state_seq)
-                    + (idx_feats * stride_conv_state_dim)[None, :]
-                    + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None]
-                )  # [NP2_STATELEN, BLOCK_N]
-                mask = (
-                    (conv_states_input_coord < num_cache_lines)
-                    & ((idx_tokens_conv + seqlen) < state_len)[:, None]
-                    & (idx_feats < dim)[None, :]
-                )
+                conv_states_ptrs_source = (conv_states_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                                           (idx_feats * stride_conv_state_dim)[None, :] +
+                                           ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None]
+                                           )  # [NP2_STATELEN, BLOCK_N]
+                mask = ((conv_states_input_coord < num_cache_lines)
+                        & ((idx_tokens_conv + seqlen) < state_len)[:, None]
+                        & (idx_feats < dim)[None, :])
                 conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
 
                 VAL = state_len - seqlen
-                x_ptrs = (
-                    x_base[None, :]
-                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
-                )
-                mask_x = (
-                    (idx_tokens_conv - VAL >= 0)[:, None]
-                    & (idx_tokens_conv - VAL < seqlen)[:, None]
-                    & (idx_feats < dim)[None, :]
-                )
+                x_ptrs = (x_base[None, :] + ((idx_tokens_conv - VAL) * stride_x_token)[:, None])
+                mask_x = ((idx_tokens_conv - VAL >= 0)[:, None]
+                          & (idx_tokens_conv - VAL < seqlen)[:, None]
+                          & (idx_feats < dim)[None, :])
                 loaded_x = tl.load(x_ptrs, mask_x, 0.0)  # [NP2_STATELEN, BLOCK_N]
                 # Merge cached history (where mask is True) with new token data
                 # (where mask is False) to form the updated conv_state.
@@ -946,35 +828,24 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                 tl.debug_barrier()  # required: tl.where bug when result of another tl.load
                 # Where mask is True: take cached history (left-shifted).
                 # Where mask is False: take new tokens from x.
-                new_conv_state = tl.where(
-                    mask, conv_state, loaded_x
-                )  # BUG in 'tl.where' requires a barrier before this
-
-                conv_states_ptrs_target = (
-                    conv_states_base
-                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
-                )  # [NP2_STATELEN, BLOCK_N]
+                new_conv_state = tl.where(mask, conv_state,
+                                          loaded_x)  # BUG in 'tl.where' requires a barrier before this
+
+                conv_states_ptrs_target = (conv_states_base + (idx_tokens_conv * stride_conv_state_tok)[:, None]
+                                           )  # [NP2_STATELEN, BLOCK_N]
                 mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
                 tl.store(conv_states_ptrs_target, new_conv_state, mask)
 
             else:  # load_init_state == False
                 idx_tokens_conv = tl.arange(0, NP2_STATELEN)
                 VAL = state_len - seqlen
-                x_ptrs = (
-                    x_base[None, :]
-                    + ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
-                )
-                mask_x = (
-                    (idx_tokens_conv - VAL >= 0)[:, None]
-                    & (idx_tokens_conv - VAL < seqlen)[:, None]
-                    & (idx_feats < dim)[None, :]
-                )
+                x_ptrs = (x_base[None, :] + ((idx_tokens_conv - VAL) * stride_x_token)[:, None])
+                mask_x = ((idx_tokens_conv - VAL >= 0)[:, None]
+                          & (idx_tokens_conv - VAL < seqlen)[:, None]
+                          & (idx_feats < dim)[None, :])
                 new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
 
-                conv_states_ptrs_target = (
-                    conv_states_base
-                    + (idx_tokens_conv * stride_conv_state_tok)[:, None]
-                )
+                conv_states_ptrs_target = (conv_states_base + (idx_tokens_conv * stride_conv_state_tok)[:, None])
                 mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
                 tl.store(conv_states_ptrs_target, new_conv_state, mask)
 
@@ -985,20 +856,13 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
         idx_rows = tl.arange(0, triton.next_power_of_2(BLOCK_M + KERNEL_WIDTH - 1))
         # idx_rows = tl.arange(0, X_TILE_M_P2)
         offs_tok = x_tile_start + idx_rows
-        tile_ptrs = (
-            x_ptr
-            + (sequence_start_index + offs_tok)[:, None] * stride_x_token
-            + (idx_feats * stride_x_dim)[None, :]
-        )
-        tile_mask = (
-            (offs_tok >= 0)[:, None]
-            & (offs_tok < seqlen)[:, None]
-            & (idx_feats < dim)[None, :]
-            & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None]
-        )
-        tile_data = tl.load(
-            tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca"
-        )
+        tile_ptrs = (x_ptr + (sequence_start_index + offs_tok)[:, None] * stride_x_token +
+                     (idx_feats * stride_x_dim)[None, :])
+        tile_mask = ((offs_tok >= 0)[:, None]
+                     & (offs_tok < seqlen)[:, None]
+                     & (idx_feats < dim)[None, :]
+                     & (idx_rows < (BLOCK_M + KERNEL_WIDTH - 1))[:, None])
+        tile_data = tl.load(tile_ptrs, mask=tile_mask, other=0.0, cache_modifier=".ca")
 
         if KERNEL_WIDTH == 2:
             col0 = tl.reshape(
@@ -1028,34 +892,19 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                 [BLOCK_N],
             )
         if (chunk_offset - 1) < n_block_to_fill:
-            idx_tokens_last = (
-                last_full_block_token_index
-                - (n_block_to_fill - chunk_offset) * B_size
-                - state_len
-            ) + tl.arange(0, NP2_STATELEN)
-            x_ptrs = (
-                x_ptr
-                + (idx_tokens_last * stride_x_token)[:, None]
-                + (idx_feats * stride_x_dim)[None, :]
-            )
+            idx_tokens_last = (last_full_block_token_index -
+                               (n_block_to_fill - chunk_offset) * B_size - state_len) + tl.arange(0, NP2_STATELEN)
+            x_ptrs = (x_ptr + (idx_tokens_last * stride_x_token)[:, None] + (idx_feats * stride_x_dim)[None, :])
             mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[None, :]
             loaded_x = tl.load(x_ptrs, mask_x, 0.0)
             idx_tokens_conv = tl.arange(0, NP2_STATELEN)
 
-            conv_states_output_coord = tl.load(
-                conv_state_indices_ptr
-                + idx_seq * stride_cache_indices
-                + current_first_index
-                + (chunk_offset - 1)
-            ).to(tl.int64)
-
-            conv_states_ptrs_target = (
-                conv_states_ptr
-                + (conv_states_output_coord * stride_conv_state_seq)
-                + (idx_feats * stride_conv_state_dim)
-            )[None, :] + (
-                idx_tokens_conv * stride_conv_state_tok
-            )[:, None]
+            conv_states_output_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_cache_indices +
+                                               current_first_index + (chunk_offset - 1)).to(tl.int64)
+
+            conv_states_ptrs_target = (conv_states_ptr + (conv_states_output_coord * stride_conv_state_seq) +
+                                       (idx_feats * stride_conv_state_dim))[None, :] + (idx_tokens_conv *
+                                                                                        stride_conv_state_tok)[:, None]
 
             mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
             tl.debug_barrier()
@@ -1066,7 +915,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
         mask_bias = idx_feats < dim
         acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
     else:
-        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        acc_preload = tl.zeros((BLOCK_N, ), dtype=tl.float32)
 
     mask_w = idx_feats < dim
     if KERNEL_WIDTH >= 2:
@@ -1080,7 +929,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
     if KERNEL_WIDTH >= 4:
         w_ptrs = w_base + (3 * stride_w_width)
         w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
-    
+
     mask_x_1d = idx_feats < dim
     for idx_token in range(segment_len):
         acc = acc_preload
@@ -1094,7 +943,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                     matrix_x = tl.reshape(
                         tle.extract_tile(
                             tile_data,
-                            index=[idx_token + chunk*(KERNEL_WIDTH - 1), 0],
+                            index=[idx_token + chunk * (KERNEL_WIDTH - 1), 0],
                             tile_shape=[1, BLOCK_N],
                         ),
                         [BLOCK_N],
@@ -1108,7 +957,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                     matrix_x = tl.reshape(
                         tle.extract_tile(
                             tile_data,
-                            index=[idx_token + chunk*(KERNEL_WIDTH - 1), 0],
+                            index=[idx_token + chunk * (KERNEL_WIDTH - 1), 0],
                             tile_shape=[1, BLOCK_N],
                         ),
                         [BLOCK_N],
@@ -1125,7 +974,7 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                     matrix_x = tl.reshape(
                         tle.extract_tile(
                             tile_data,
-                            index=[idx_token + chunk*(KERNEL_WIDTH - 1), 0],
+                            index=[idx_token + chunk * (KERNEL_WIDTH - 1), 0],
                             tile_shape=[1, BLOCK_N],
                         ),
                         [BLOCK_N],
@@ -1145,11 +994,8 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
         if SILU_ACTIVATION:
             acc = acc / (1 + tl.exp(-acc))
         mask_1d = (idx_token < segment_len) & (idx_feats < dim)
-        o_ptrs = (
-            o_ptr
-            + (sequence_start_index + token_offset + idx_token) * stride_o_token
-            + (idx_feats * stride_o_dim)
-        )
+        o_ptrs = (o_ptr + (sequence_start_index + token_offset + idx_token) * stride_o_token +
+                  (idx_feats * stride_o_dim))
         tl.store(o_ptrs, acc, mask=mask_1d)
 
 
@@ -1211,9 +1057,8 @@ def _causal_conv1d_update_kernel_v2(
         conv_state_init = 0
         current_last_index = 0
 
-    conv_states_input_coord = tl.load(
-        conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
-    ).to(tl.int64)
+    conv_states_input_coord = tl.load(conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init).to(
+        tl.int64)
 
     if USE_PAD_SLOT:  # noqa
         if conv_states_input_coord == pad_slot_id:
@@ -1236,89 +1081,61 @@ def _causal_conv1d_update_kernel_v2(
         return
 
     if IS_SPEC_DECODING:
-        conv_state_token_offset = (
-            tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
-        )
+        conv_state_token_offset = (tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1)
     else:
         conv_state_token_offset = 0
 
     # STEP 1: load conv_state tile and extract columns via extract_tile
-    conv_states_base = (
-        conv_state_ptr
-        + (conv_states_input_coord * stride_conv_state_seq)
-        + (idx_feats * stride_conv_state_dim)
-    )
+    conv_states_base = (conv_state_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                        (idx_feats * stride_conv_state_dim))
     mask_w = idx_feats < dim
 
-
     idx_state = tl.arange(0, NP2_STATELEN)
-    conv_state_ptrs_full = (
-        conv_states_base[None, :]
-        + (conv_state_token_offset + idx_state)[:, None] * stride_conv_state_tok
-    )  # [NP2_STATELEN, BLOCK_N]
+    conv_state_ptrs_full = (conv_states_base[None, :] +
+                            (conv_state_token_offset + idx_state)[:, None] * stride_conv_state_tok
+                            )  # [NP2_STATELEN, BLOCK_N]
     mask_state_full = (idx_state[:, None] < state_len) & (idx_feats[None, :] < dim)
-    state_tile = tl.load(
-        conv_state_ptrs_full, mask_state_full, other=0.0
-    )  # [NP2_STATELEN, BLOCK_N]
+    state_tile = tl.load(conv_state_ptrs_full, mask_state_full, other=0.0)  # [NP2_STATELEN, BLOCK_N]
 
     if KERNEL_WIDTH >= 2:
-        col0 = tle.extract_tile(state_tile, index=[0, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+        col0 = tle.extract_tile(state_tile, index=[0, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N, ))
     if KERNEL_WIDTH >= 3:
-        col1 = tle.extract_tile(state_tile, index=[1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+        col1 = tle.extract_tile(state_tile, index=[1, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N, ))
     if KERNEL_WIDTH >= 4:
-        col2 = tle.extract_tile(state_tile, index=[2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+        col2 = tle.extract_tile(state_tile, index=[2, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N, ))
     if KERNEL_WIDTH >= 5:
-        col3 = tle.extract_tile(state_tile, index=[3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+        col3 = tle.extract_tile(state_tile, index=[3, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N, ))
     if KERNEL_WIDTH >= 6:
-        col4 = tle.extract_tile(state_tile, index=[4, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N,))
+        col4 = tle.extract_tile(state_tile, index=[4, 0], tile_shape=[1, BLOCK_N]).reshape((BLOCK_N, ))
 
     # STEP 2: build updated conv_state and write back
     idx_tokens = tl.arange(0, NP2_STATELEN)
 
-    conv_state_ptrs_source = (
-        conv_state_ptr
-        + (conv_states_input_coord * stride_conv_state_seq)
-        + conv_state_token_offset * stride_conv_state_tok
-        + (idx_feats * stride_conv_state_dim)[None, :]
-        + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[
-            :, None
-        ]
-    )  # [NP2_STATELEN, BLOCK_N]
-    mask = (
-        (conv_states_input_coord < num_cache_lines)
-        & ((idx_tokens + seqlen) < state_len)[:, None]
-        & (idx_feats < dim)[None, :]
-    )
+    conv_state_ptrs_source = (conv_state_ptr + (conv_states_input_coord * stride_conv_state_seq) +
+                              conv_state_token_offset * stride_conv_state_tok +
+                              (idx_feats * stride_conv_state_dim)[None, :] +
+                              ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[:, None]
+                              )  # [NP2_STATELEN, BLOCK_N]
+    mask = ((conv_states_input_coord < num_cache_lines)
+            & ((idx_tokens + seqlen) < state_len)[:, None]
+            & (idx_feats < dim)[None, :])
     conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
 
     VAL = state_len - seqlen
     x_base = x_ptr + x_offset + (idx_feats * stride_x_dim)  # [BLOCK_N]
-    x_ptrs = (
-        x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None]
-    )  # [NP2_STATELEN, BLOCK_N]
-    mask_x = (
-        (idx_tokens - VAL >= 0)[:, None]
-        & (idx_tokens - VAL < seqlen)[:, None]
-        & (idx_feats < dim)[None, :]
-    )
+    x_ptrs = (x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None])  # [NP2_STATELEN, BLOCK_N]
+    mask_x = ((idx_tokens - VAL >= 0)[:, None] & (idx_tokens - VAL < seqlen)[:, None] & (idx_feats < dim)[None, :])
     loaded_x = tl.load(x_ptrs, mask_x, 0.0)  # [NP2_STATELEN, BLOCK_N]
 
-
-
     tl.debug_barrier()
 
     new_conv_state = tl.where(mask, conv_state, loaded_x)
 
-    conv_states_offset = tl.load(
-        conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index
-    ).to(tl.int64)
-    conv_state_ptrs_target = (
-        conv_state_ptr
-        + (conv_states_offset * stride_conv_state_seq)
-        + (idx_feats * stride_conv_state_dim)
-    )[None, :] + (
-        idx_tokens * stride_conv_state_tok
-    )[:, None]
+    conv_states_offset = tl.load(conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index).to(
+        tl.int64)
+    conv_state_ptrs_target = (conv_state_ptr + (conv_states_offset * stride_conv_state_seq) +
+                              (idx_feats * stride_conv_state_dim))[None, :] + (idx_tokens * stride_conv_state_tok)[:,
+                                                                                                                   None]
     mask_out = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
     tl.store(conv_state_ptrs_target, new_conv_state, mask_out)
 
@@ -1328,7 +1145,7 @@ def _causal_conv1d_update_kernel_v2(
         mask_bias = idx_feats < dim
         acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(tl.float32)
     else:
-        acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        acc_preload = tl.zeros((BLOCK_N, ), dtype=tl.float32)
 
     # STEP 4: preload weights
     w_base = w_ptr + (idx_feats * stride_w_dim)  # [BLOCK_N,]
@@ -1443,20 +1260,17 @@ def _causal_conv1d_update_kernel_v2(
         if SILU_ACTIVATION:
             acc = acc / (1 + tl.exp(-acc))
         mask_1d = (idx_token < seqlen) & (idx_feats < dim)
-        o_ptrs = (
-            o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim)
-        )
+        o_ptrs = (o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim))
         tl.store(o_ptrs, acc, mask=mask_1d)
 
 
-
 # %%
 # Python wrappers
 # ---------------
 
 
-def _launch_fwd_kernel(kernel_fn, x, weight, bias, conv_states, query_start_loc,
-                       cache_indices, has_initial_state, activation, pad_slot_id):
+def _launch_fwd_kernel(kernel_fn, x, weight, bias, conv_states, query_start_loc, cache_indices, has_initial_state,
+                       activation, pad_slot_id):
     if isinstance(activation, bool) and activation:
         activation = "silu"
 
@@ -1466,12 +1280,8 @@ def _launch_fwd_kernel(kernel_fn, x, weight, bias, conv_states, query_start_loc,
 
     seqlens = query_start_loc.diff().to("cpu")
     MAX_NUM_PROGRAMS = 1024
-    batch_ptr = torch.full(
-        (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
-    )
-    token_chunk_offset_ptr = torch.full(
-        (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
-    )
+    batch_ptr = torch.full((MAX_NUM_PROGRAMS, ), PAD_SLOT_ID, dtype=torch.int32, device=x.device)
+    token_chunk_offset_ptr = torch.full((MAX_NUM_PROGRAMS, ), PAD_SLOT_ID, dtype=torch.int32, device=x.device)
 
     dim, cu_seqlen = x.shape
     _, width = weight.shape
@@ -1501,59 +1311,71 @@ def num_program(META, seqlens):
 
         if META["batch_ptr"].nelement() >= len(mlist):
             META["batch_ptr"][0:len(mlist)].copy_(torch.from_numpy(np.array(mlist)))
-            META["token_chunk_offset_ptr"][0:len(mlist)].copy_(
-                torch.from_numpy(np.array(offsetlist))
-            )
+            META["token_chunk_offset_ptr"][0:len(mlist)].copy_(torch.from_numpy(np.array(offsetlist)))
 
         META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device)
-        META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(
-            META["x_ptr"].device
-        )
+        META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(META["x_ptr"].device)
         return tot
 
     def grid(META):
         return (num_program(META, seqlens), triton.cdiv(dim, META["BLOCK_N"]))
 
     kernel_fn[grid](
-        x, weight, bias, conv_states, cache_indices, has_initial_state,
-        query_start_loc, batch_ptr, token_chunk_offset_ptr,
-        None, None, None, None, out,
-        dim, cu_seqlen, num_cache_lines,
-        stride_x_dim, stride_x_token, stride_w_dim, stride_w_width,
-        stride_istate_seq, stride_istate_dim, stride_istate_token,
-        stride_cache_indices, stride_o_dim, stride_o_token,
-        1, pad_slot_id,
+        x,
+        weight,
+        bias,
+        conv_states,
+        cache_indices,
+        has_initial_state,
+        query_start_loc,
+        batch_ptr,
+        token_chunk_offset_ptr,
+        None,
+        None,
+        None,
+        None,
+        out,
+        dim,
+        cu_seqlen,
+        num_cache_lines,
+        stride_x_dim,
+        stride_x_token,
+        stride_w_dim,
+        stride_w_width,
+        stride_istate_seq,
+        stride_istate_dim,
+        stride_istate_token,
+        stride_cache_indices,
+        stride_o_dim,
+        stride_o_token,
+        1,
+        pad_slot_id,
         HAS_BIAS=bias is not None,
         KERNEL_WIDTH=width,
         SILU_ACTIVATION=activation in ["silu", "swish"],
         IS_APC_ENABLED=False,
         USE_PAD_SLOT=pad_slot_id is not None,
         NP2_STATELEN=np2_statelen,
-        BLOCK_M=BLOCK_M, BLOCK_N=256, num_stages=2,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=256,
+        num_stages=2,
     )
     return out.to(original_x_dtype)
 
 
-def causal_conv1d_fn_v1(x, weight, bias, conv_states, query_start_loc,
-                        cache_indices=None, has_initial_state=None,
+def causal_conv1d_fn_v1(x, weight, bias, conv_states, query_start_loc, cache_indices=None, has_initial_state=None,
                         activation="silu", pad_slot_id=PAD_SLOT_ID):
-    return _launch_fwd_kernel(
-        _causal_conv1d_fwd_kernel_v1, x, weight, bias, conv_states,
-        query_start_loc, cache_indices, has_initial_state, activation, pad_slot_id
-    )
+    return _launch_fwd_kernel(_causal_conv1d_fwd_kernel_v1, x, weight, bias, conv_states, query_start_loc,
+                              cache_indices, has_initial_state, activation, pad_slot_id)
 
 
-def causal_conv1d_fn_v2(x, weight, bias, conv_states, query_start_loc,
-                        cache_indices=None, has_initial_state=None,
+def causal_conv1d_fn_v2(x, weight, bias, conv_states, query_start_loc, cache_indices=None, has_initial_state=None,
                         activation="silu", pad_slot_id=PAD_SLOT_ID):
-    return _launch_fwd_kernel(
-        _causal_conv1d_fwd_kernel_v2, x, weight, bias, conv_states,
-        query_start_loc, cache_indices, has_initial_state, activation, pad_slot_id
-    )
+    return _launch_fwd_kernel(_causal_conv1d_fwd_kernel_v2, x, weight, bias, conv_states, query_start_loc,
+                              cache_indices, has_initial_state, activation, pad_slot_id)
 
 
-def _launch_update_kernel(kernel_fn, x, conv_state, weight, bias,
-                          activation, conv_state_indices, pad_slot_id):
+def _launch_update_kernel(kernel_fn, x, conv_state, weight, bias, activation, conv_state_indices, pad_slot_id):
     if isinstance(activation, bool):
         activation = "silu" if activation else None
     elif activation is not None:
@@ -1581,14 +1403,33 @@ def grid(META):
         return (batch, triton.cdiv(dim, META["BLOCK_N"]))
 
     kernel_fn[grid](
-        x, weight, bias, conv_state, conv_state_indices, None, None,
-        None, None, out,
-        batch, dim, seqlen, state_len, num_cache_lines,
-        stride_x_seq, stride_x_dim, stride_x_token,
-        stride_w_dim, stride_w_width,
-        stride_istate_seq, stride_istate_dim, stride_istate_token,
+        x,
+        weight,
+        bias,
+        conv_state,
+        conv_state_indices,
+        None,
+        None,
+        None,
+        None,
+        out,
+        batch,
+        dim,
+        seqlen,
+        state_len,
+        num_cache_lines,
+        stride_x_seq,
+        stride_x_dim,
+        stride_x_token,
+        stride_w_dim,
+        stride_w_width,
+        stride_istate_seq,
+        stride_istate_dim,
+        stride_istate_token,
         stride_state_indices,
-        stride_o_seq, stride_o_dim, stride_o_token,
+        stride_o_seq,
+        stride_o_dim,
+        stride_o_token,
         pad_slot_id,
         HAS_BIAS=bias is not None,
         KERNEL_WIDTH=width,
@@ -1605,20 +1446,16 @@ def grid(META):
     return out.to(original_x_dtype)
 
 
-def causal_conv1d_update_v1(x, conv_state, weight, bias=None, activation=None,
-                            conv_state_indices=None, pad_slot_id=PAD_SLOT_ID):
-    return _launch_update_kernel(
-        _causal_conv1d_update_kernel_v1, x, conv_state, weight, bias,
-        activation, conv_state_indices, pad_slot_id
-    )
+def causal_conv1d_update_v1(x, conv_state, weight, bias=None, activation=None, conv_state_indices=None,
+                            pad_slot_id=PAD_SLOT_ID):
+    return _launch_update_kernel(_causal_conv1d_update_kernel_v1, x, conv_state, weight, bias, activation,
+                                 conv_state_indices, pad_slot_id)
 
 
-def causal_conv1d_update_v2(x, conv_state, weight, bias=None, activation=None,
-                            conv_state_indices=None, pad_slot_id=PAD_SLOT_ID):
-    return _launch_update_kernel(
-        _causal_conv1d_update_kernel_v2, x, conv_state, weight, bias,
-        activation, conv_state_indices, pad_slot_id
-    )
+def causal_conv1d_update_v2(x, conv_state, weight, bias=None, activation=None, conv_state_indices=None,
+                            pad_slot_id=PAD_SLOT_ID):
+    return _launch_update_kernel(_causal_conv1d_update_kernel_v2, x, conv_state, weight, bias, activation,
+                                 conv_state_indices, pad_slot_id)
 
 
 # %%
@@ -1629,11 +1466,12 @@ def causal_conv1d_update_v2(x, conv_state, weight, bias=None, activation=None,
 def _make_varlen_data(dim, total_seqlen, batch, width, dtype):
     torch.manual_seed(0)
     eos_pos = torch.randperm(total_seqlen - 1)[:batch - 1].sort().values
-    seqlens = torch.diff(torch.cat([
-        torch.tensor([-1], dtype=torch.int32),
-        eos_pos.to(dtype=torch.int32),
-        torch.tensor([total_seqlen - 1], dtype=torch.int32),
-    ]))
+    seqlens = torch.diff(
+        torch.cat([
+            torch.tensor([-1], dtype=torch.int32),
+            eos_pos.to(dtype=torch.int32),
+            torch.tensor([total_seqlen - 1], dtype=torch.int32),
+        ]))
     query_start_loc = torch.cat([
         torch.tensor([0], dtype=torch.int32),
         torch.cumsum(seqlens, dim=0).to(torch.int32),
@@ -1663,14 +1501,10 @@ def run_correctness(dim=1024, total_seqlen=2048, batch=32, width=4, dtype=torch.
     x, weight, bias, conv_states, query_start_loc, cache_indices, has_initial_state = \
         _make_varlen_data(dim, total_seqlen, batch, width, dtype)
 
-    out_v1 = causal_conv1d_fn_v1(
-        x.clone(), weight, bias, conv_states.clone(), query_start_loc,
-        cache_indices, has_initial_state, activation="silu"
-    )
-    out_v2 = causal_conv1d_fn_v2(
-        x.clone(), weight, bias, conv_states.clone(), query_start_loc,
-        cache_indices, has_initial_state, activation="silu"
-    )
+    out_v1 = causal_conv1d_fn_v1(x.clone(), weight, bias, conv_states.clone(), query_start_loc, cache_indices,
+                                 has_initial_state, activation="silu")
+    out_v2 = causal_conv1d_fn_v2(x.clone(), weight, bias, conv_states.clone(), query_start_loc, cache_indices,
+                                 has_initial_state, activation="silu")
     torch.testing.assert_close(out_v1.float(), out_v2.float(), rtol=5e-2, atol=2e-1)
     print("[OK] varlen forward correctness passed")
 
@@ -1678,14 +1512,10 @@ def run_correctness(dim=1024, total_seqlen=2048, batch=32, width=4, dtype=torch.
     x_u, weight_u, bias_u, conv_state_u, conv_state_indices_u = \
         _make_update_data(dim, batch, width, dtype)
 
-    out_u_v1 = causal_conv1d_update_v1(
-        x_u.clone(), conv_state_u.clone(), weight_u, bias_u,
-        activation="silu", conv_state_indices=conv_state_indices_u
-    )
-    out_u_v2 = causal_conv1d_update_v2(
-        x_u.clone(), conv_state_u.clone(), weight_u, bias_u,
-        activation="silu", conv_state_indices=conv_state_indices_u
-    )
+    out_u_v1 = causal_conv1d_update_v1(x_u.clone(), conv_state_u.clone(), weight_u, bias_u, activation="silu",
+                                       conv_state_indices=conv_state_indices_u)
+    out_u_v2 = causal_conv1d_update_v2(x_u.clone(), conv_state_u.clone(), weight_u, bias_u, activation="silu",
+                                       conv_state_indices=conv_state_indices_u)
     torch.testing.assert_close(out_u_v1.float(), out_u_v2.float(), rtol=5e-2, atol=2e-1)
     print("[OK] update correctness passed")
 
@@ -1705,8 +1535,7 @@ def main(argv=None):
     parser.add_argument("--total_seqlen", type=int, default=2048)
     parser.add_argument("--batch", type=int, default=32)
     parser.add_argument("--width", type=int, default=4)
-    parser.add_argument("--dtype", type=str, default="bfloat16",
-                        choices=["float16", "bfloat16", "fp16", "bf16"])
+    parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float16", "bfloat16", "fp16", "bf16"])
     parser.add_argument("--benchmark", action="store_true")
     args = parser.parse_args(argv)
 
@@ -1720,7 +1549,7 @@ def main(argv=None):
         # Each (dim, batch) pair produces one benchmark table.
         # ------------------------------------------------------------------
         VARLEN_CONFIGS = [(4096, 32), (8192, 128)]
-        UPDATE_CONFIGS = [(256,), (1024,)]  # batch values
+        UPDATE_CONFIGS = [(256, ), (1024, )]  # batch values
         for dim, batch in VARLEN_CONFIGS:
             _run_benchmark_varlen(dim, batch, dtype)
         for batch in UPDATE_CONFIGS:
@@ -1735,12 +1564,10 @@ def _run_benchmark_varlen(dim, batch, dtype):
     print("-" * 52)
     for total_seqlen in x_vals:
         x, w, b, cs, qsl, ci, his = _make_varlen_data(dim, total_seqlen, batch, width, dtype)
-        ms_b, _, _ = triton.testing.do_bench(
-            lambda: causal_conv1d_fn_v1(x, w, b, cs, qsl, ci, his, "silu"),
-            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
-        ms_t, _, _ = triton.testing.do_bench(
-            lambda: causal_conv1d_fn_v2(x, w, b, cs, qsl, ci, his, "silu"),
-            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
+        ms_b, _, _ = triton.testing.do_bench(lambda: causal_conv1d_fn_v1(x, w, b, cs, qsl, ci, his, "silu"), warmup=10,
+                                             rep=100, quantiles=[0.5, 0.2, 0.8])
+        ms_t, _, _ = triton.testing.do_bench(lambda: causal_conv1d_fn_v2(x, w, b, cs, qsl, ci, his, "silu"), warmup=10,
+                                             rep=100, quantiles=[0.5, 0.2, 0.8])
         sp = ms_b / ms_t if ms_t > 0 else 1.0
         print(f"{total_seqlen:>12} | {ms_b:>12.4f} {ms_t:>10.4f} {sp:>7.2f}x")
 
@@ -1753,12 +1580,10 @@ def _run_benchmark_update(batch, dtype):
     print("-" * 46)
     for dim in x_vals:
         x, w, b, cs, ci = _make_update_data(dim, batch, width, dtype)
-        ms_b, _, _ = triton.testing.do_bench(
-            lambda: causal_conv1d_update_v1(x, cs, w, b, "silu", ci),
-            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
-        ms_t, _, _ = triton.testing.do_bench(
-            lambda: causal_conv1d_update_v2(x, cs, w, b, "silu", ci),
-            warmup=10, rep=100, quantiles=[0.5, 0.2, 0.8])
+        ms_b, _, _ = triton.testing.do_bench(lambda: causal_conv1d_update_v1(x, cs, w, b, "silu", ci), warmup=10,
+                                             rep=100, quantiles=[0.5, 0.2, 0.8])
+        ms_t, _, _ = triton.testing.do_bench(lambda: causal_conv1d_update_v2(x, cs, w, b, "silu", ci), warmup=10,
+                                             rep=100, quantiles=[0.5, 0.2, 0.8])
         sp = ms_b / ms_t if ms_t > 0 else 1.0
         print(f"{dim:>6} | {ms_b:>12.4f} {ms_t:>10.4f} {sp:>7.2f}x")
 
diff --git a/python/tutorials/tle/08-rope.py b/python/tutorials/tle/08-rope.py
index 8dcb37f89e..9fd6228ec1 100644
--- a/python/tutorials/tle/08-rope.py
+++ b/python/tutorials/tle/08-rope.py
@@ -37,9 +37,7 @@
 # -----
 
 import argparse
-import math
 import random
-import sys
 
 import torch
 import triton
@@ -48,15 +46,15 @@
 
 DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
-
 # %%
 # Helper: build cos/sin cache
 # ---------------------------
 
+
 def build_rope_cache(max_seq_len, head_dim, device, dtype=torch.float32, base=10000.0):
     """Precompute cos/sin tables for all positions up to ``max_seq_len``."""
     half_dim = head_dim // 2
-    theta = 1.0 / (base ** (torch.arange(0, half_dim, device=device, dtype=dtype) / half_dim))
+    theta = 1.0 / (base**(torch.arange(0, half_dim, device=device, dtype=dtype) / half_dim))
     positions = torch.arange(max_seq_len, device=device, dtype=dtype)
     angles = torch.outer(positions, theta)
     return angles.cos().contiguous(), angles.sin().contiguous()
@@ -70,19 +68,29 @@ def build_rope_cache(max_seq_len, head_dim, device, dtype=torch.float32, base=10
 # (plus ``N`` loads for the rotated element), even though the data for all
 # heads is contiguous in memory.
 
+
 @triton.jit
 def _rope_outplace_kernel_v1(
-    oq_ptr,         # (n_tokens, q_heads, head_dim)
-    ok_ptr,         # (n_tokens, k_heads, head_dim)
-    q_ptr,          # (n_tokens, q_heads, head_dim)
-    k_ptr,          # (n_tokens, k_heads, head_dim)
-    cos_ptr,        # (max_seq_len, head_dim // 2)
-    sin_ptr,        # (max_seq_len, head_dim // 2)
-    q_stride_s, q_stride_h, q_stride_d,
-    k_stride_s, k_stride_h, k_stride_d,
-    oq_stride_s, oq_stride_h, oq_stride_d,
-    ok_stride_s, ok_stride_h, ok_stride_d,
-    cos_stride_s, sin_stride_s,
+    oq_ptr,  # (n_tokens, q_heads, head_dim)
+    ok_ptr,  # (n_tokens, k_heads, head_dim)
+    q_ptr,  # (n_tokens, q_heads, head_dim)
+    k_ptr,  # (n_tokens, k_heads, head_dim)
+    cos_ptr,  # (max_seq_len, head_dim // 2)
+    sin_ptr,  # (max_seq_len, head_dim // 2)
+    q_stride_s,
+    q_stride_h,
+    q_stride_d,
+    k_stride_s,
+    k_stride_h,
+    k_stride_d,
+    oq_stride_s,
+    oq_stride_h,
+    oq_stride_d,
+    ok_stride_s,
+    ok_stride_h,
+    ok_stride_d,
+    cos_stride_s,
+    sin_stride_s,
     seq_len,
     NUM_Q_HEADS: tl.constexpr,
     NUM_K_HEADS: tl.constexpr,
@@ -125,7 +133,7 @@ def _rope_outplace_kernel_v1(
         rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
         output_offs = off_h * oq_stride_h + (ordered_block * oq_stride_d)
 
-        q = tl.load(q_ptr + ordered_cols, mask=mask, other=0.0)          # ← load from GMEM
+        q = tl.load(q_ptr + ordered_cols, mask=mask, other=0.0)  # ← load from GMEM
         rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)  # ← load from GMEM
         y = q * cos + rotated_q * sin
         tl.store(oq_ptr + output_offs, y, mask=mask)
@@ -138,7 +146,7 @@ def _rope_outplace_kernel_v1(
         rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
         output_offs = off_h * ok_stride_h + (ordered_block * ok_stride_d)
 
-        k = tl.load(k_ptr + ordered_cols, mask=mask, other=0.0)          # ← load from GMEM
+        k = tl.load(k_ptr + ordered_cols, mask=mask, other=0.0)  # ← load from GMEM
         rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)  # ← load from GMEM
         y = k * cos + rotated_k * sin
         tl.store(ok_ptr + output_offs, y, mask=mask)
@@ -146,13 +154,18 @@ def _rope_outplace_kernel_v1(
 
 @triton.jit
 def _rope_inplace_kernel_v1(
-    q_ptr,          # (n_tokens, q_heads, head_dim)
-    k_ptr,          # (n_tokens, k_heads, head_dim)
-    cos_ptr,        # (max_seq_len, head_dim // 2)
-    sin_ptr,        # (max_seq_len, head_dim // 2)
-    q_stride_s, q_stride_h, q_stride_d,
-    k_stride_s, k_stride_h, k_stride_d,
-    cos_stride_s, sin_stride_s,
+    q_ptr,  # (n_tokens, q_heads, head_dim)
+    k_ptr,  # (n_tokens, k_heads, head_dim)
+    cos_ptr,  # (max_seq_len, head_dim // 2)
+    sin_ptr,  # (max_seq_len, head_dim // 2)
+    q_stride_s,
+    q_stride_h,
+    q_stride_d,
+    k_stride_s,
+    k_stride_h,
+    k_stride_d,
+    cos_stride_s,
+    sin_stride_s,
     seq_len,
     NUM_Q_HEADS: tl.constexpr,
     NUM_K_HEADS: tl.constexpr,
@@ -219,19 +232,29 @@ def _rope_inplace_kernel_v1(
 # ``tl.load`` (it lives at a different offset), but the dominant load path is
 # coalesced into one bulk access.
 
+
 @triton.jit
 def _rope_outplace_kernel_v2(
-    oq_ptr,         # (n_tokens, q_heads, head_dim)
-    ok_ptr,         # (n_tokens, k_heads, head_dim)
-    q_ptr,          # (n_tokens, q_heads, head_dim)
-    k_ptr,          # (n_tokens, k_heads, head_dim)
-    cos_ptr,        # (max_seq_len, head_dim // 2)
-    sin_ptr,        # (max_seq_len, head_dim // 2)
-    q_stride_s, q_stride_h, q_stride_d,
-    k_stride_s, k_stride_h, k_stride_d,
-    oq_stride_s, oq_stride_h, oq_stride_d,
-    ok_stride_s, ok_stride_h, ok_stride_d,
-    cos_stride_s, sin_stride_s,
+    oq_ptr,  # (n_tokens, q_heads, head_dim)
+    ok_ptr,  # (n_tokens, k_heads, head_dim)
+    q_ptr,  # (n_tokens, q_heads, head_dim)
+    k_ptr,  # (n_tokens, k_heads, head_dim)
+    cos_ptr,  # (max_seq_len, head_dim // 2)
+    sin_ptr,  # (max_seq_len, head_dim // 2)
+    q_stride_s,
+    q_stride_h,
+    q_stride_d,
+    k_stride_s,
+    k_stride_h,
+    k_stride_d,
+    oq_stride_s,
+    oq_stride_h,
+    oq_stride_d,
+    ok_stride_s,
+    ok_stride_h,
+    ok_stride_d,
+    cos_stride_s,
+    sin_stride_s,
     seq_len,
     NUM_Q_HEADS: tl.constexpr,
     NUM_K_HEADS: tl.constexpr,
@@ -279,9 +302,7 @@ def _rope_outplace_kernel_v2(
 
     for off_h in range(0, NUM_Q_HEADS):
         # ── KEY CHANGE: slice from register tile instead of GMEM load ──
-        q = tle.extract_tile(
-            q_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
-        ).reshape((PADDED_HEAD_DIM,))
+        q = tle.extract_tile(q_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]).reshape((PADDED_HEAD_DIM, ))
 
         rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
         rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)
@@ -299,9 +320,7 @@ def _rope_outplace_kernel_v2(
     )  # shape: [NUM_K_HEADS, PADDED_HEAD_DIM]
 
     for off_h in range(0, NUM_K_HEADS):
-        k = tle.extract_tile(
-            k_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
-        ).reshape((PADDED_HEAD_DIM,))
+        k = tle.extract_tile(k_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]).reshape((PADDED_HEAD_DIM, ))
 
         rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
         rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)
@@ -311,13 +330,18 @@ def _rope_outplace_kernel_v2(
 
 @triton.jit
 def _rope_inplace_kernel_v2(
-    q_ptr,          # (n_tokens, q_heads, head_dim)
-    k_ptr,          # (n_tokens, k_heads, head_dim)
-    cos_ptr,        # (max_seq_len, head_dim // 2)
-    sin_ptr,        # (max_seq_len, head_dim // 2)
-    q_stride_s, q_stride_h, q_stride_d,
-    k_stride_s, k_stride_h, k_stride_d,
-    cos_stride_s, sin_stride_s,
+    q_ptr,  # (n_tokens, q_heads, head_dim)
+    k_ptr,  # (n_tokens, k_heads, head_dim)
+    cos_ptr,  # (max_seq_len, head_dim // 2)
+    sin_ptr,  # (max_seq_len, head_dim // 2)
+    q_stride_s,
+    q_stride_h,
+    q_stride_d,
+    k_stride_s,
+    k_stride_h,
+    k_stride_d,
+    cos_stride_s,
+    sin_stride_s,
     seq_len,
     NUM_Q_HEADS: tl.constexpr,
     NUM_K_HEADS: tl.constexpr,
@@ -362,9 +386,7 @@ def _rope_inplace_kernel_v2(
     )
 
     for off_h in range(0, NUM_Q_HEADS):
-        q = tle.extract_tile(
-            q_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
-        ).reshape((PADDED_HEAD_DIM,))
+        q = tle.extract_tile(q_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]).reshape((PADDED_HEAD_DIM, ))
 
         rotated_cols = off_h * q_stride_h + (rotated_block * q_stride_d)
         rotated_q = tl.load(q_ptr + rotated_cols, mask=mask, other=0.0)
@@ -381,9 +403,7 @@ def _rope_inplace_kernel_v2(
     )
 
     for off_h in range(0, NUM_K_HEADS):
-        k = tle.extract_tile(
-            k_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]
-        ).reshape((PADDED_HEAD_DIM,))
+        k = tle.extract_tile(k_block, index=[off_h, 0], tile_shape=[1, PADDED_HEAD_DIM]).reshape((PADDED_HEAD_DIM, ))
 
         rotated_cols = off_h * k_stride_h + (rotated_block * k_stride_d)
         rotated_k = tl.load(k_ptr + rotated_cols, mask=mask, other=0.0)
@@ -411,17 +431,33 @@ def _launch_rope_outplace(kernel, q, k, cos, sin, rotary_interleaved):
     q_embed = torch.empty_like(q)
     k_embed = torch.empty_like(k)
 
-    grid = (n_tokens,)
+    grid = (n_tokens, )
     kernel[grid](
-        q_embed, k_embed,
-        q, k, cos, sin,
-        q.stride(0), q.stride(1), q.stride(2),
-        k.stride(0), k.stride(1), k.stride(2),
-        q_embed.stride(0), q_embed.stride(1), q_embed.stride(2),
-        k_embed.stride(0), k_embed.stride(1), k_embed.stride(2),
-        cos.stride(0), sin.stride(0),
+        q_embed,
+        k_embed,
+        q,
+        k,
+        cos,
+        sin,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        q_embed.stride(0),
+        q_embed.stride(1),
+        q_embed.stride(2),
+        k_embed.stride(0),
+        k_embed.stride(1),
+        k_embed.stride(2),
+        cos.stride(0),
+        sin.stride(0),
         seq_len,
-        q_heads, k.shape[-2], head_dim, padded_head_dim,
+        q_heads,
+        k.shape[-2],
+        head_dim,
+        padded_head_dim,
         rotary_interleaved,
         MAX_POSITION_EMBEDDINGS=cos.shape[0],
     )
@@ -440,14 +476,25 @@ def _launch_rope_inplace(kernel, q, k, cos, sin, rotary_interleaved):
     n_tokens, q_heads, head_dim = q.shape
     padded_head_dim = max(triton.next_power_of_2(head_dim), 16)
 
-    grid = (n_tokens,)
+    grid = (n_tokens, )
     kernel[grid](
-        q, k, cos, sin,
-        q.stride(0), q.stride(1), q.stride(2),
-        k.stride(0), k.stride(1), k.stride(2),
-        cos.stride(0), sin.stride(0),
+        q,
+        k,
+        cos,
+        sin,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        cos.stride(0),
+        sin.stride(0),
         seq_len,
-        q_heads, k.shape[-2], head_dim, padded_head_dim,
+        q_heads,
+        k.shape[-2],
+        head_dim,
+        padded_head_dim,
         rotary_interleaved,
         MAX_POSITION_EMBEDDINGS=cos.shape[0],
     )
@@ -458,14 +505,17 @@ def rope_outplace_v1(q, k, cos, sin, rotary_interleaved=False):
     """Baseline RoPE (outplace)."""
     return _launch_rope_outplace(_rope_outplace_kernel_v1, q, k, cos, sin, rotary_interleaved)
 
+
 def rope_outplace_v2(q, k, cos, sin, rotary_interleaved=False):
     """TLE RoPE (outplace)."""
     return _launch_rope_outplace(_rope_outplace_kernel_v2, q, k, cos, sin, rotary_interleaved)
 
+
 def rope_inplace_v1(q, k, cos, sin, rotary_interleaved=False):
     """Baseline RoPE (inplace)."""
     return _launch_rope_inplace(_rope_inplace_kernel_v1, q, k, cos, sin, rotary_interleaved)
 
+
 def rope_inplace_v2(q, k, cos, sin, rotary_interleaved=False):
     """TLE RoPE (inplace)."""
     return _launch_rope_inplace(_rope_inplace_kernel_v2, q, k, cos, sin, rotary_interleaved)
@@ -504,14 +554,12 @@ def _assert_close_robust(name, a, b, rtol=5e-2, atol=2e-1):
     if abs_p999 < 0.25 and rel_p999 < 0.25:
         return  # robust check passed
 
-    raise AssertionError(
-        f"{name} mismatch: allclose failed and robust check failed; "
-        f"abs_p999={abs_p999:.6f}, rel_p999={rel_p999:.6f}"
-    )
+    raise AssertionError(f"{name} mismatch: allclose failed and robust check failed; "
+                         f"abs_p999={abs_p999:.6f}, rel_p999={rel_p999:.6f}")
 
 
-def check_correctness(batch=4, seq_len=256, q_heads=32, k_heads=8, head_dim=128,
-                      dtype=torch.bfloat16, rotary_interleaved=False):
+def check_correctness(batch=4, seq_len=256, q_heads=32, k_heads=8, head_dim=128, dtype=torch.bfloat16,
+                      rotary_interleaved=False):
     """Verify v1 and v2 produce identical results (both inplace and outplace)."""
     torch.manual_seed(0)
 
@@ -598,8 +646,7 @@ def _robust_bench(fn, reset_fn):
     return float(torch.quantile(t[keep], 0.5).item())
 
 
-def _run_benchmark_table(title, configs, dtype, rotary_interleaved,
-                         rope_v1_fn, rope_v2_fn, inplace):
+def _run_benchmark_table(title, configs, dtype, rotary_interleaved, rope_v1_fn, rope_v2_fn, inplace):
     """Print a benchmark table for a set of (batch, seq_len, q_heads, k_heads, head_dim)."""
     mode = "inplace" if inplace else "outplace"
     print(f"\n--- {title} [{mode}] (dtype={dtype}) ---")
@@ -625,10 +672,12 @@ def run_v2():
             rope_v2_fn(q2, k2, cos, sin, rotary_interleaved)
 
         def reset_v1():
-            q1.copy_(q_src); k1.copy_(k_src)
+            q1.copy_(q_src)
+            k1.copy_(k_src)
 
         def reset_v2():
-            q2.copy_(q_src); k2.copy_(k_src)
+            q2.copy_(q_src)
+            k2.copy_(k_src)
 
         # Randomize measurement order
         runners = [("v1", run_v1, reset_v1), ("v2", run_v2, reset_v2)]
@@ -650,13 +699,10 @@ def reset_v2():
 
 
 def main(argv=None):
-    parser = argparse.ArgumentParser(
-        description="RoPE with TLE extract_tile — baseline vs optimized benchmark"
-    )
+    parser = argparse.ArgumentParser(description="RoPE with TLE extract_tile — baseline vs optimized benchmark")
     parser.add_argument("--benchmark", action="store_true",
                         help="Run benchmark tables in addition to correctness check")
-    parser.add_argument("--dtype", type=str, default="bfloat16",
-                        choices=["float16", "bfloat16", "fp16", "bf16"],
+    parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float16", "bfloat16", "fp16", "bf16"],
                         help="Data type for benchmark")
     args = parser.parse_args(argv)
 
@@ -678,29 +724,27 @@ def main(argv=None):
     CONFIGS = [
         # (batch, seq_len, q_heads, k_heads, head_dim)
         # head_dim=128 — TLE shines
-        (1,   128,  32, 8,  128),
-        (1,  1024,  32, 8,  128),
-        (8,   128,  32, 8,  128),
-        (8,  1024,  32, 8,  128),
-        (32,  128,  32, 8,  128),
-        (32, 1024,  32, 8,  128),
-        (8,  1024,  16, 16, 128),
-        (32, 1024,  16, 16, 128),
+        (1, 128, 32, 8, 128),
+        (1, 1024, 32, 8, 128),
+        (8, 128, 32, 8, 128),
+        (8, 1024, 32, 8, 128),
+        (32, 128, 32, 8, 128),
+        (32, 1024, 32, 8, 128),
+        (8, 1024, 16, 16, 128),
+        (32, 1024, 16, 16, 128),
         # head_dim=256 — smaller gain
-        (1,   128,  32, 8,  256),
-        (8,   128,  32, 8,  256),
-        (32,  128,  32, 8,  256),
-        (32, 1024,  16, 16, 256),
+        (1, 128, 32, 8, 256),
+        (8, 128, 32, 8, 256),
+        (32, 128, 32, 8, 256),
+        (32, 1024, 16, 16, 256),
     ]
 
     for interleaved in [False, True]:
         ilabel = "interleaved (GPT-NeoX style)" if interleaved else "non-interleaved (LLaMA style)"
         # outplace
-        _run_benchmark_table(ilabel, CONFIGS, dtype, interleaved,
-                             rope_outplace_v1, rope_outplace_v2, inplace=False)
+        _run_benchmark_table(ilabel, CONFIGS, dtype, interleaved, rope_outplace_v1, rope_outplace_v2, inplace=False)
         # inplace
-        _run_benchmark_table(ilabel, CONFIGS, dtype, interleaved,
-                             rope_inplace_v1, rope_inplace_v2, inplace=True)
+        _run_benchmark_table(ilabel, CONFIGS, dtype, interleaved, rope_inplace_v1, rope_inplace_v2, inplace=True)
 
 
 if __name__ == "__main__":
diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp
index fefbafe1e9..df12e2f3d8 100644
--- a/third_party/tle/dialect/lib/IR/Ops.cpp
+++ b/third_party/tle/dialect/lib/IR/Ops.cpp
@@ -25,7 +25,8 @@ constexpr int kClusterSharedMemoryAddressSpace = 7;
 // ExtractTileOp Builder
 // ============================================================================
 void ExtractTileOp::build(OpBuilder &builder, OperationState &state, Value src,
-                          Value index, ArrayRef<int64_t> tileShape, ArrayRef<int64_t> strides) {
+                          Value index, ArrayRef<int64_t> tileShape,
+                          ArrayRef<int64_t> strides) {
   auto srcType = cast<RankedTensorType>(src.getType());
   auto resultType = RankedTensorType::get(tileShape, srcType.getElementType(),
                                           srcType.getEncoding());
@@ -64,9 +65,10 @@ LogicalResult ExtractTileOp::verify() {
   SmallVector<int64_t> strides;
   if (auto a = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
           getOperation()->getAttr("strides")))
-    for (auto v : a.asArrayRef()) strides.push_back(v);
-  if (strides.empty()) strides = tileShape;  
-
+    for (auto v : a.asArrayRef())
+      strides.push_back(v);
+  if (strides.empty())
+    strides = tileShape;
 
   // ---- Basic checks required for both static and dynamic index ----
 
@@ -92,7 +94,8 @@ LogicalResult ExtractTileOp::verify() {
     if ((srcShape[i] - tileShape[i]) < 0 ||
         (srcShape[i] - tileShape[i]) % strides[i] != 0)
       return emitOpError("(srcShape - tileShape) must be divisible by strides "
-                         "at dimension ") << i;
+                         "at dimension ")
+             << i;
     if (dstShape[i] != tileShape[i])
       return emitOpError("result shape must equal tile_shape at dimension ")
              << i;
@@ -119,7 +122,7 @@ LogicalResult ExtractTileOp::verify() {
   for (size_t i = 0; i < srcShape.size(); ++i) {
     logicalGridShape[i] = (srcShape[i] - tileShape[i]) / strides[i] + 1;
     totalTiles *= logicalGridShape[i];
- }
+  }
 
   // Out-of-bounds check
   if (index < 0 || index >= totalTiles)
diff --git a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
index a1bcf90a61..6bb0b14452 100644
--- a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
+++ b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp
@@ -9,7 +9,6 @@
 #include "tle/dialect/include/Transforms/PatternTleToLLVM.h"
 #include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
-#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
 #include "llvm/Support/raw_ostream.h"
@@ -34,7 +33,8 @@ static SmallVector<int64_t> getStrides(ExtractTileOp op) {
   if (auto a = mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(
           op->getAttr("strides"))) {
     SmallVector<int64_t> s;
-    for (auto v : a.asArrayRef()) s.push_back(v);
+    for (auto v : a.asArrayRef())
+      s.push_back(v);
     return s;
   }
   // 向后兼容：没有 strides 属性时退化为 tile_shape
@@ -51,7 +51,7 @@ static bool isCTATileAligned(ExtractTileOp op, int64_t linearIndex) {
   auto srcTy = cast<RankedTensorType>(op.getSrc().getType());
   auto srcShape = srcTy.getShape();
   auto tileShape = getTileShape(op);
-  auto strides   = getStrides(op);
+  auto strides = getStrides(op);
   auto ctaTile = getShapePerCTATile(srcTy);
   int rank = srcShape.size();
   SmallVector<int64_t> logicalGrid(rank), tileCoords(rank);
@@ -63,7 +63,7 @@ static bool isCTATileAligned(ExtractTileOp op, int64_t linearIndex) {
     remain /= logicalGrid[i];
   }
   for (int i = 0; i < rank; ++i) {
-   
+
     int64_t off = tileCoords[i] * strides[i];
     if (tileShape[i] % (int64_t)ctaTile[i] != 0)
       return false;
@@ -87,7 +87,7 @@ lowerExtractTileStatic(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
   auto srcShape = srcTy.getShape(), dstShape = dstTy.getShape();
   auto tileShape = getTileShape(op);
   int rank = srcShape.size();
-  auto strides   = getStrides(op);
+  auto strides = getStrides(op);
   auto vals = unpackLLElements(loc, adaptor.getSrc(), rewriter);
   auto shapePerCTATile = getShapePerCTATile(srcTy);
   auto srcCTAShape = multiDimElementwise<int64_t, unsigned>(
@@ -104,11 +104,11 @@ lowerExtractTileStatic(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
     remain /= logicalGrid[i];
   }
   for (int i = 0; i < rank; ++i)
-    elementCoords[i] = logicalCoords[i] * strides[i]; 
+    elementCoords[i] = logicalCoords[i] * strides[i];
   auto firstTileCoord = multiDimElementwise<int64_t, unsigned>(
       elementCoords, shapePerCTATile, std::divides<unsigned>());
   auto srcCTAOrder = getCTATileOrder(srcTy),
-      dstCTAOrder = getCTATileOrder(dstTy);
+       dstCTAOrder = getCTATileOrder(dstTy);
   unsigned totalSrcCTAs = std::accumulate(
       srcCTAShape.begin(), srcCTAShape.end(), 1, std::multiplies<>());
   unsigned elemsPerCTA = ttg::getTotalElemsPerThread(srcTy) / totalSrcCTAs;
@@ -147,7 +147,7 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
   auto srcShape = srcTy.getShape(), dstShape = dstTy.getShape();
   auto tileShape = getTileShape(op);
   int rank = srcShape.size();
-  auto strides   = getStrides(op);
+  auto strides = getStrides(op);
   MLIRContext *ctx = rewriter.getContext();
   auto i1Ty = rewriter.getIntegerType(1);
   auto i8Ty = rewriter.getIntegerType(8);
@@ -223,10 +223,11 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor,
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d]));
     Value tv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d]));
-    
+
     Value coord = rewriter.create<LLVM::UDivOp>(loc, i32Ty, dynIndex, sv);
     coord = rewriter.create<LLVM::URemOp>(loc, i32Ty, coord, gv);
-    tileStartVals[d] = rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, sv_stride);
+    tileStartVals[d] =
+        rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, sv_stride);
     tileEndVals[d] =
         rewriter.create<LLVM::AddOp>(loc, i32Ty, tileStartVals[d], tv);
   }
diff --git a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp
index dbcac1ee4a..7ada8d945e 100644
--- a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp
+++ b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp
@@ -36,7 +36,8 @@ static SmallVector<int64_t> getStrides(InsertTileOp op) {
     return s;
   }
   auto tileTy = cast<RankedTensorType>(op.getTile().getType());
-  return SmallVector<int64_t>(tileTy.getShape().begin(), tileTy.getShape().end());
+  return SmallVector<int64_t>(tileTy.getShape().begin(),
+                              tileTy.getShape().end());
 }
 
 // Check if the tile to be inserted is CTA-aligned (for register shuffle path).
@@ -264,12 +265,13 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
     Value gv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)logicalGrid[d]));
     Value svStride = rewriter.create<LLVM::ConstantOp>(
-      loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d]));
+        loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d]));
     Value tv = rewriter.create<LLVM::ConstantOp>(
         loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d]));
     Value coord = rewriter.create<LLVM::UDivOp>(loc, i32Ty, dynIndex, sv);
     coord = rewriter.create<LLVM::URemOp>(loc, i32Ty, coord, gv);
-    tileStartVals[d] = rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, svStride);
+    tileStartVals[d] =
+        rewriter.create<LLVM::MulOp>(loc, i32Ty, coord, svStride);
     tileEndVals[d] =
         rewriter.create<LLVM::AddOp>(loc, i32Ty, tileStartVals[d], tv);
   }
@@ -360,7 +362,8 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor,
     Value lp = rewriter.create<LLVM::GEPOp>(loc, smemPtrTy, i8Ty, smemBase,
                                             ValueRange{smemByteOffsetV},
                                             LLVM::GEPNoWrapFlags::inbounds);
-    Value tileLoaded = rewriter.create<LLVM::LoadOp>(loc, llvmElemTy, lp, elemBytes);
+    Value tileLoaded =
+        rewriter.create<LLVM::LoadOp>(loc, llvmElemTy, lp, elemBytes);
     rewriter.create<LLVM::BrOp>(loc, ValueRange{tileLoaded}, mergeBlock);
 
     rewriter.setInsertionPointToStart(elseBlock);
diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc
index 343d7239de..504c641802 100644
--- a/third_party/tle/triton_tle.cc
+++ b/third_party/tle/triton_tle.cc
@@ -79,9 +79,11 @@ void init_triton_tle_ir(py::module &&m) {
       // TLE-Lite
       .def(
           "create_extract_tile",
-          [](TritonOpBuilder &self, Value &input,
-            Value &index, std::vector<int64_t> &tileShape,std::vector<int64_t> strides) -> Value {
-            auto op = self.create<tle::ExtractTileOp>(input, index, tileShape, strides);
+          [](TritonOpBuilder &self, Value &input, Value &index,
+             std::vector<int64_t> &tileShape,
+             std::vector<int64_t> strides) -> Value {
+            auto op = self.create<tle::ExtractTileOp>(input, index, tileShape,
+                                                      strides);
             return op.getResult();
           },
           py::arg("input"), py::arg("index"), py::arg("tileShape"),
@@ -89,9 +91,10 @@ void init_triton_tle_ir(py::module &&m) {
           "Create extract_tile operation")
       .def(
           "create_insert_tile",
-          [](TritonOpBuilder &self, Value &input, Value &tile,
-             Value &index, std::vector<int64_t> strides) -> Value {
-            auto op = self.create<tle::InsertTileOp>(input, tile, index, strides);
+          [](TritonOpBuilder &self, Value &input, Value &tile, Value &index,
+             std::vector<int64_t> strides) -> Value {
+            auto op =
+                self.create<tle::InsertTileOp>(input, tile, index, strides);
             return op.getResult();
           },
           py::arg("input"), py::arg("tile"), py::arg("index"),
@@ -477,7 +480,7 @@ void init_triton_tle_passes(py::module &&m) {
                      tle::createTritonTleLowerExtractTile);
 
   ADD_PASS_WRAPPER_0("add_lower_insert_tile",
-                    tle::createTritonTleLowerInsertTile);
+                     tle::createTritonTleLowerInsertTile);
 }
 
 void init_tle_raw_ir(py::module &&m) {

From 6334df4ea3fdb6ba09e34b4dbe9077ffe7db3695 Mon Sep 17 00:00:00 2001
From: flagtree-bot <flagtree_ai@163.com>
Date: Thu, 4 Jun 2026 02:43:34 +0000
Subject: [PATCH 08/41] Apply code-format changes

---
 python/triton/experimental/tle/language/gpu/semantic.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/triton/experimental/tle/language/gpu/semantic.py b/python/triton/experimental/tle/language/gpu/semantic.py
index c196926b87..d0550efbe1 100644
--- a/python/triton/experimental/tle/language/gpu/semantic.py
+++ b/python/triton/experimental/tle/language/gpu/semantic.py
@@ -131,6 +131,7 @@ def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequen
 
         # Check 5: dimension match
 
+
 #     src_shape = list(src.type.shape)
 
         if len(tile_shape_unwrapped) != len(src_shape):

From a5866c99026bf39f71c9b08ac386e407b3d2f3fd Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 16:32:15 +0800
Subject: [PATCH 09/41] [TLE]fix some GCU backend problems

---
 python/tutorials/tle/07-causal-conv1d.py      | 10 ---
 .../TritonToGCU/TTSmemLowerToGCU.cpp          | 68 +++++++++++++------
 .../TritonToGCU/TritonGCULocalMemOptimize.cpp |  8 +--
 .../GCUConvertTritonToTritonGPU.cpp           |  6 +-
 4 files changed, 56 insertions(+), 36 deletions(-)

diff --git a/python/tutorials/tle/07-causal-conv1d.py b/python/tutorials/tle/07-causal-conv1d.py
index 5f2eea045d..f5adfa7d7b 100644
--- a/python/tutorials/tle/07-causal-conv1d.py
+++ b/python/tutorials/tle/07-causal-conv1d.py
@@ -154,7 +154,6 @@ def _causal_conv1d_fwd_kernel_v1(
                 col1 = tl.load(prior_tokens - 1 * stride_conv_state_tok, mask_w, 0.0)
                 col0 = tl.load(prior_tokens - 2 * stride_conv_state_tok, mask_w, 0.0)
             if KERNEL_WIDTH == 5:
-                col3 = tl.load(prior_tokens, mask_w, 0.0)
                 col2 = tl.load(prior_tokens - 1 * stride_conv_state_tok, mask_w, 0.0)
                 col1 = tl.load(prior_tokens - 2 * stride_conv_state_tok, mask_w, 0.0)
                 col0 = tl.load(prior_tokens - 3 * stride_conv_state_tok, mask_w, 0.0)
@@ -165,8 +164,6 @@ def _causal_conv1d_fwd_kernel_v1(
                 col1 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 4:
                 col2 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
-            if KERNEL_WIDTH >= 5:
-                col3 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
 
         if state_len <= seqlen:
             idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN)
@@ -242,7 +239,6 @@ def _causal_conv1d_fwd_kernel_v1(
             col1 = tl.load(prior_tokens - 1 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
             col0 = tl.load(prior_tokens - 2 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
         if KERNEL_WIDTH == 5:
-            col3 = tl.load(prior_tokens, mask_w, 0.0, cache_modifier=".ca")
             col2 = tl.load(prior_tokens - 1 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
             col1 = tl.load(prior_tokens - 2 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
             col0 = tl.load(prior_tokens - 3 * stride_x_token, mask_w, 0.0, cache_modifier=".ca")
@@ -755,8 +751,6 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                 col0 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape(
                     (BLOCK_N, ))
             if KERNEL_WIDTH == 5:
-                col3 = tle.extract_tile(state_tile, index=[state_len - 1, 0], tile_shape=[1, BLOCK_N]).reshape(
-                    (BLOCK_N, ))
                 col2 = tle.extract_tile(state_tile, index=[state_len - 2, 0], tile_shape=[1, BLOCK_N]).reshape(
                     (BLOCK_N, ))
                 col1 = tle.extract_tile(state_tile, index=[state_len - 3, 0], tile_shape=[1, BLOCK_N]).reshape(
@@ -770,9 +764,6 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
                 col1 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
             if KERNEL_WIDTH >= 4:
                 col2 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
-            if KERNEL_WIDTH >= 5:
-                col3 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
-
         # STEP 2: write back conv_state
         if state_len <= seqlen:
             idx_tokens_last = (seqlen - state_len) + tl.arange(0, NP2_STATELEN)
@@ -930,7 +921,6 @@ def _causal_conv1d_fwd_kernel_v2(  # continuous batching
         w_ptrs = w_base + (3 * stride_w_width)
         w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
 
-    mask_x_1d = idx_feats < dim
     for idx_token in range(segment_len):
         acc = acc_preload
 
diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
index 66654c7e30..dc28519d95 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
@@ -1023,18 +1023,19 @@ struct MemDescToPtrOpLowering
 // Convert a linear tile index to multi-dim tile coordinates (row-major).
 //
 // The source tensor is viewed as a grid of tiles:
-//   grid[d] = srcShape[d] / tileShape[d]
+//   grid[d] = (srcShape[d] - tileShape[d]) / strides[d] + 1
 // Then linearIdx is delinearized into per-dim tile coordinates.
-//   e.g. srcShape=[64,64], tileShape=[32,32] -> grid=[2,2]
+//   e.g. srcShape=[64,64], tileShape=[32,32], strides=[32,32] -> grid=[2,2]
 //        linearIdx=3 -> coords=[1,1] (bottom-right tile)
 static SmallVector<Value> delinearizeTileIndex(OpBuilder &b, Location loc,
                                                Value linearIdx,
                                                ArrayRef<int64_t> srcShape,
-                                               ArrayRef<int64_t> tileShape) {
+                                               ArrayRef<int64_t> tileShape,
+                                               ArrayRef<int64_t> strides) {
   unsigned rank = srcShape.size();
   SmallVector<int64_t> grid(rank);
   for (unsigned d = 0; d < rank; ++d)
-    grid[d] = srcShape[d] / tileShape[d];
+    grid[d] = (srcShape[d] - tileShape[d]) / strides[d] + 1;
 
   auto i32Ty = b.getI32Type();
   Value idx = linearIdx;
@@ -1062,20 +1063,22 @@ static SmallVector<Value> delinearizeTileIndex(OpBuilder &b, Location loc,
 // Convert a linear tile index to element-level offsets in the source tensor.
 //
 // First delinearizes the index into tile coordinates, then scales each
-// coordinate by the tile size in that dimension to get the element offset.
-//   e.g. srcShape=[64,64], tileShape=[32,32], linearIdx=3
+// coordinate by the stride in that dimension to get the element offset.
+//   e.g. srcShape=[64,64], tileShape=[32,32], strides=[32,32], linearIdx=3
 //        -> coords=[1,1] -> offsets=[32,32]
 //        meaning the tile starts at row 32, col 32 of the source.
 static SmallVector<Value> computeElemOffsets(OpBuilder &b, Location loc,
                                              Value linearIdx,
                                              ArrayRef<int64_t> srcShape,
-                                             ArrayRef<int64_t> tileShape) {
-  auto coords = delinearizeTileIndex(b, loc, linearIdx, srcShape, tileShape);
+                                             ArrayRef<int64_t> tileShape,
+                                             ArrayRef<int64_t> strides) {
+  auto coords = delinearizeTileIndex(b, loc, linearIdx, srcShape, tileShape,
+                                     strides);
   unsigned rank = srcShape.size();
   SmallVector<Value> offsets(rank);
   for (unsigned d = 0; d < rank; ++d) {
     auto cst = b.create<arith::ConstantIntOp>(
-        loc, static_cast<int32_t>(tileShape[d]), 32);
+        loc, static_cast<int32_t>(strides[d]), 32);
     offsets[d] = b.create<arith::MulIOp>(loc, coords[d], cst);
   }
   return offsets;
@@ -1128,13 +1131,14 @@ static SmallVector<Value, 4> computeCombinedOffsets(OpBuilder &b, Location loc,
 static Value
 emitSliceFromSmem(OpBuilder &rewriter, Location loc, Value smemBuf,
                   Value convertedIndex, ArrayRef<int64_t> srcShape,
-                  ArrayRef<int64_t> tileShape, RankedTensorType resultTensorTy,
-                  MemRefType outputType,
+                  ArrayRef<int64_t> tileShape, ArrayRef<int64_t> strides,
+                  RankedTensorType resultTensorTy, MemRefType outputType,
                   triton::gcu::FirstLastUserAnalysis &userAnalysis,
                   std::map<Operation *, Operation *> &replaced2Origin,
                   triton::gcu::PrivateTagPool &pTagPool, Operation *op) {
   auto elemOffsets =
-      computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape);
+      computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape,
+                         strides);
   auto lastUser = userAnalysis.getLastUser(op->getResults()[0]);
   auto output = syncAllocOp(rewriter, loc, lastUser, userAnalysis,
                             replaced2Origin, outputType);
@@ -1160,13 +1164,14 @@ static Value
 emitDesliceToSmem(OpBuilder &rewriter, Location loc, Value smemBuf,
                   Value convertedTile, Value convertedIndex,
                   ArrayRef<int64_t> srcShape, ArrayRef<int64_t> tileShape,
-                  RankedTensorType tileTensorTy,
+                  ArrayRef<int64_t> strides, RankedTensorType tileTensorTy,
                   RankedTensorType resultTensorTy,
                   triton::gcu::FirstLastUserAnalysis &userAnalysis,
                   std::map<Operation *, Operation *> &replaced2Origin,
                   triton::gcu::PrivateTagPool &pTagPool, Operation *op) {
   auto elemOffsets =
-      computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape);
+      computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape,
+                         strides);
   auto firstUser = userAnalysis.getFirstUser(op->getResults()[0]);
   auto lastUser = userAnalysis.getLastUser(op->getResults()[0]);
 
@@ -1235,6 +1240,17 @@ struct TleExtractTileLowering : SharedGenericConversionPattern {
       return failure();
     auto tileShape = tileShapeAttr.asArrayRef();
 
+    // Extract strides (defaults to tileShape for backward compatibility).
+    auto stridesAttr = op->getAttrOfType<DenseI64ArrayAttr>("strides");
+    SmallVector<int64_t> stridesVec;
+    if (stridesAttr) {
+      stridesVec.assign(stridesAttr.asArrayRef().begin(),
+                         stridesAttr.asArrayRef().end());
+    } else {
+      stridesVec.assign(tileShape.begin(), tileShape.end());
+    }
+    ArrayRef<int64_t> strides = stridesVec;
+
     auto srcTensorTy = cast<RankedTensorType>(op->getOperand(0).getType());
     auto srcShape = srcTensorTy.getShape();
     auto outputType =
@@ -1245,7 +1261,7 @@ struct TleExtractTileLowering : SharedGenericConversionPattern {
       auto output = syncAllocOp(rewriter, loc, lastUser, userAnalysis,
                                 replaced2Origin, outputType);
       auto elemOffsets = computeElemOffsets(rewriter, loc, convertedIndex,
-                                            srcShape, tileShape);
+                                            srcShape, tileShape, strides);
       auto totalNumElems = triton::gcu::getTotalElemsPerThread(resultTensorTy);
       auto defaultValue = triton::gcu::createConstantZero(
           rewriter, loc, resultTensorTy.getElementType());
@@ -1282,8 +1298,8 @@ struct TleExtractTileLowering : SharedGenericConversionPattern {
 
     auto output =
         emitSliceFromSmem(rewriter, loc, smemBuf, convertedIndex, srcShape,
-                          tileShape, resultTensorTy, outputType, userAnalysis,
-                          replaced2Origin, pTagPool, op);
+                          tileShape, strides, resultTensorTy, outputType,
+                          userAnalysis, replaced2Origin, pTagPool, op);
 
     leaveTritionOp(rewriter, op);
     rewriter.replaceOp(op, output);
@@ -1338,6 +1354,17 @@ struct TleInsertTileLowering : SharedGenericConversionPattern {
     }
     ArrayRef<int64_t> tileShape = tileShapeVec;
 
+    // Extract strides (defaults to tileShape for backward compatibility).
+    auto stridesAttr = op->getAttrOfType<DenseI64ArrayAttr>("strides");
+    SmallVector<int64_t> stridesVec;
+    if (stridesAttr) {
+      stridesVec.assign(stridesAttr.asArrayRef().begin(),
+                         stridesAttr.asArrayRef().end());
+    } else {
+      stridesVec.assign(tileShapeVec.begin(), tileShapeVec.end());
+    }
+    ArrayRef<int64_t> strides = stridesVec;
+
     auto srcTensorTy = cast<RankedTensorType>(op->getOperand(0).getType());
     auto srcShape = srcTensorTy.getShape();
 
@@ -1348,7 +1375,7 @@ struct TleInsertTileLowering : SharedGenericConversionPattern {
       auto output = syncAllocOp(rewriter, loc, lastUser, userAnalysis,
                                 replaced2Origin, outputType);
       auto elemOffsets = computeElemOffsets(rewriter, loc, convertedIndex,
-                                            srcShape, tileShape);
+                                            srcShape, tileShape, strides);
 
       auto cpTag = pTagPool.getPrivateSyncTagInfo(op);
       SmallVector<Value> zeroOffsets(
@@ -1397,8 +1424,9 @@ struct TleInsertTileLowering : SharedGenericConversionPattern {
 
     auto output =
         emitDesliceToSmem(rewriter, loc, smemBuf, convertedTile, convertedIndex,
-                          srcShape, tileShape, tileTensorTy, resultTensorTy,
-                          userAnalysis, replaced2Origin, pTagPool, op);
+                          srcShape, tileShape, strides, tileTensorTy,
+                          resultTensorTy, userAnalysis, replaced2Origin,
+                          pTagPool, op);
 
     leaveTritionOp(rewriter, op);
     rewriter.replaceOp(op, output);
diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TritonGCULocalMemOptimize.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TritonGCULocalMemOptimize.cpp
index c358339664..fddde4f10b 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TritonGCULocalMemOptimize.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TritonGCULocalMemOptimize.cpp
@@ -507,10 +507,10 @@ class FuseExtractTileSmemRelay : public RewritePattern {
     if (src.getDefiningOp()->use_empty())
       rewriter.eraseOp(src.getDefiningOp());
 
-    auto tileStridesAttr = op->getAttrOfType<DenseI64ArrayAttr>("tile_strides");
+    auto stridesAttr = op->getAttrOfType<DenseI64ArrayAttr>("strides");
     auto resultTy = op->getResult(0).getType();
     auto newOp = rewriter.create<triton::gcu::SliceFromLocalOp>(
-        loc, resultTy, smemAlloc, index, tileShapeAttr, tileStridesAttr);
+        loc, resultTy, smemAlloc, index, tileShapeAttr, stridesAttr);
     rewriter.replaceOp(op, newOp.getResult());
     return success();
   }
@@ -585,10 +585,10 @@ class FuseInsertTileSmemRelay : public RewritePattern {
     if (src.getDefiningOp()->use_empty())
       rewriter.eraseOp(src.getDefiningOp());
 
-    auto tileStridesAttr = op->getAttrOfType<DenseI64ArrayAttr>("tile_strides");
+    auto stridesAttr = op->getAttrOfType<DenseI64ArrayAttr>("strides");
     auto resultTy = op->getResult(0).getType();
     auto newOp = rewriter.create<triton::gcu::DesliceToLocalOp>(
-        loc, resultTy, smemAlloc, tile, index, tileShapeAttr, tileStridesAttr);
+        loc, resultTy, smemAlloc, tile, index, tileShapeAttr, stridesAttr);
     rewriter.replaceOp(op, newOp.getResult());
     return success();
   }
diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToTritonGPU/GCUConvertTritonToTritonGPU.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToTritonGPU/GCUConvertTritonToTritonGPU.cpp
index 969728a6c7..8ec1119604 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToTritonGPU/GCUConvertTritonToTritonGPU.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToTritonGPU/GCUConvertTritonToTritonGPU.cpp
@@ -792,7 +792,8 @@ class TleExtractTileOpPattern
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
     auto newOp = rewriter.replaceOpWithNewOp<triton::tle::ExtractTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getIndex());
+        op, retType, adaptor.getSrc(), adaptor.getIndex(),
+        op->getAttrOfType<DenseI64ArrayAttr>("strides"));
     addNamedAttrs(newOp, adaptor.getAttributes());
     return success();
   }
@@ -816,7 +817,8 @@ class TleInsertTileOpPattern
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
     auto newOp = rewriter.replaceOpWithNewOp<triton::tle::InsertTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex());
+        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(),
+        op->getAttrOfType<DenseI64ArrayAttr>("strides"));
     addNamedAttrs(newOp, adaptor.getAttributes());
     return success();
   }

From 0c5a41254c357af404117e742ea74eea0b375c6a Mon Sep 17 00:00:00 2001
From: flagtree-bot <flagtree_ai@163.com>
Date: Thu, 4 Jun 2026 08:34:54 +0000
Subject: [PATCH 10/41] Apply code-format changes

---
 .../TritonToGCU/TTSmemLowerToGCU.cpp          | 27 +++++++++----------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
index dc28519d95..a27fa11372 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
@@ -1072,8 +1072,8 @@ static SmallVector<Value> computeElemOffsets(OpBuilder &b, Location loc,
                                              ArrayRef<int64_t> srcShape,
                                              ArrayRef<int64_t> tileShape,
                                              ArrayRef<int64_t> strides) {
-  auto coords = delinearizeTileIndex(b, loc, linearIdx, srcShape, tileShape,
-                                     strides);
+  auto coords =
+      delinearizeTileIndex(b, loc, linearIdx, srcShape, tileShape, strides);
   unsigned rank = srcShape.size();
   SmallVector<Value> offsets(rank);
   for (unsigned d = 0; d < rank; ++d) {
@@ -1136,9 +1136,8 @@ emitSliceFromSmem(OpBuilder &rewriter, Location loc, Value smemBuf,
                   triton::gcu::FirstLastUserAnalysis &userAnalysis,
                   std::map<Operation *, Operation *> &replaced2Origin,
                   triton::gcu::PrivateTagPool &pTagPool, Operation *op) {
-  auto elemOffsets =
-      computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape,
-                         strides);
+  auto elemOffsets = computeElemOffsets(rewriter, loc, convertedIndex, srcShape,
+                                        tileShape, strides);
   auto lastUser = userAnalysis.getLastUser(op->getResults()[0]);
   auto output = syncAllocOp(rewriter, loc, lastUser, userAnalysis,
                             replaced2Origin, outputType);
@@ -1169,9 +1168,8 @@ emitDesliceToSmem(OpBuilder &rewriter, Location loc, Value smemBuf,
                   triton::gcu::FirstLastUserAnalysis &userAnalysis,
                   std::map<Operation *, Operation *> &replaced2Origin,
                   triton::gcu::PrivateTagPool &pTagPool, Operation *op) {
-  auto elemOffsets =
-      computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape,
-                         strides);
+  auto elemOffsets = computeElemOffsets(rewriter, loc, convertedIndex, srcShape,
+                                        tileShape, strides);
   auto firstUser = userAnalysis.getFirstUser(op->getResults()[0]);
   auto lastUser = userAnalysis.getLastUser(op->getResults()[0]);
 
@@ -1245,7 +1243,7 @@ struct TleExtractTileLowering : SharedGenericConversionPattern {
     SmallVector<int64_t> stridesVec;
     if (stridesAttr) {
       stridesVec.assign(stridesAttr.asArrayRef().begin(),
-                         stridesAttr.asArrayRef().end());
+                        stridesAttr.asArrayRef().end());
     } else {
       stridesVec.assign(tileShape.begin(), tileShape.end());
     }
@@ -1359,7 +1357,7 @@ struct TleInsertTileLowering : SharedGenericConversionPattern {
     SmallVector<int64_t> stridesVec;
     if (stridesAttr) {
       stridesVec.assign(stridesAttr.asArrayRef().begin(),
-                         stridesAttr.asArrayRef().end());
+                        stridesAttr.asArrayRef().end());
     } else {
       stridesVec.assign(tileShapeVec.begin(), tileShapeVec.end());
     }
@@ -1422,11 +1420,10 @@ struct TleInsertTileLowering : SharedGenericConversionPattern {
                                  replaced2Origin);
     }
 
-    auto output =
-        emitDesliceToSmem(rewriter, loc, smemBuf, convertedTile, convertedIndex,
-                          srcShape, tileShape, strides, tileTensorTy,
-                          resultTensorTy, userAnalysis, replaced2Origin,
-                          pTagPool, op);
+    auto output = emitDesliceToSmem(
+        rewriter, loc, smemBuf, convertedTile, convertedIndex, srcShape,
+        tileShape, strides, tileTensorTy, resultTensorTy, userAnalysis,
+        replaced2Origin, pTagPool, op);
 
     leaveTritionOp(rewriter, op);
     rewriter.replaceOp(op, output);

From 5f82345001bbdb272d6870dcaf51769debe7d4f3 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 17:26:07 +0800
Subject: [PATCH 11/41] [TLE]fix some GCU backends problems 2.

---
 .../TritonToGCU/TTSmemLowerToGCU.cpp          | 27 +++++++++++++++----
 .../TritonToTritonGPUPass.cpp                 |  6 +++--
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
index a27fa11372..2140428a81 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
@@ -1452,10 +1452,19 @@ struct SliceFromLocalOpLowering
     auto outputType =
         dyn_cast<MemRefType>(getTypeConverter()->convertType(resultTensorTy));
 
+    auto tileShapeArr = op.getTileShape();
+    auto stridesAttr = op.getTileStrides();
+    SmallVector<int64_t> stridesVec(
+        stridesAttr ? stridesAttr.asArrayRef().begin()
+                    : tileShapeArr.begin(),
+        stridesAttr ? stridesAttr.asArrayRef().end()
+                    : tileShapeArr.end());
+
     auto output = emitSliceFromSmem(
         rewriter, op.getLoc(), adaptor.getSrc(), adaptor.getIndex(),
-        op.getSrc().getType().getShape(), op.getTileShape(), resultTensorTy,
-        outputType, userAnalysis, replaced2Origin, pTagPool, op);
+        op.getSrc().getType().getShape(), tileShapeArr, stridesVec,
+        resultTensorTy, outputType, userAnalysis, replaced2Origin, pTagPool,
+        op);
 
     leaveTritionOp(rewriter, op.getOperation());
     rewriter.replaceOp(op, output);
@@ -1483,11 +1492,19 @@ struct DesliceToLocalOpLowering
     auto tileTensorTy = cast<RankedTensorType>(op.getTile().getType());
     auto resultTensorTy = cast<RankedTensorType>(op.getResult().getType());
 
+    auto tileShapeArr = op.getTileShape();
+    auto stridesAttr = op.getTileStrides();
+    SmallVector<int64_t> stridesVec(
+        stridesAttr ? stridesAttr.asArrayRef().begin()
+                    : tileShapeArr.begin(),
+        stridesAttr ? stridesAttr.asArrayRef().end()
+                    : tileShapeArr.end());
+
     auto output = emitDesliceToSmem(
         rewriter, op.getLoc(), adaptor.getSrc(), adaptor.getTile(),
-        adaptor.getIndex(), op.getSrc().getType().getShape(), op.getTileShape(),
-        tileTensorTy, resultTensorTy, userAnalysis, replaced2Origin, pTagPool,
-        op);
+        adaptor.getIndex(), op.getSrc().getType().getShape(), tileShapeArr,
+        stridesVec, tileTensorTy, resultTensorTy, userAnalysis,
+        replaced2Origin, pTagPool, op);
 
     leaveTritionOp(rewriter, op.getOperation());
     rewriter.replaceOp(op, output);
diff --git a/third_party/hcu/triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/third_party/hcu/triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
index b49871486d..ecb65948a9 100644
--- a/third_party/hcu/triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
+++ b/third_party/hcu/triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp
@@ -847,7 +847,8 @@ class TleExtractTileOpPattern : public OpConversionPattern<tle::ExtractTileOp> {
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
     auto newOp = rewriter.replaceOpWithNewOp<tle::ExtractTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getIndex());
+        op, retType, adaptor.getSrc(), adaptor.getIndex(),
+        op->getAttrOfType<DenseI64ArrayAttr>("strides"));
 
     if (auto tileShapeAttr = op->getAttr("tile_shape"))
       newOp->setAttr("tile_shape", tileShapeAttr);
@@ -890,7 +891,8 @@ class TleInsertTileOpPattern : public OpConversionPattern<tle::InsertTileOp> {
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
     auto newOp = rewriter.replaceOpWithNewOp<tle::InsertTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex());
+        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(),
+        op->getAttrOfType<DenseI64ArrayAttr>("strides"));
 
     addNamedAttrs(newOp, adaptor.getAttributes());
 

From 64e082cd651c780774f9dce892d9d0442360234c Mon Sep 17 00:00:00 2001
From: flagtree-bot <flagtree_ai@163.com>
Date: Thu, 4 Jun 2026 09:31:28 +0000
Subject: [PATCH 12/41] Apply code-format changes

---
 .../TritonToGCU/TTSmemLowerToGCU.cpp          | 26 ++++++++-----------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
index 2140428a81..78b578d728 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
@@ -1455,16 +1455,14 @@ struct SliceFromLocalOpLowering
     auto tileShapeArr = op.getTileShape();
     auto stridesAttr = op.getTileStrides();
     SmallVector<int64_t> stridesVec(
-        stridesAttr ? stridesAttr.asArrayRef().begin()
-                    : tileShapeArr.begin(),
-        stridesAttr ? stridesAttr.asArrayRef().end()
-                    : tileShapeArr.end());
+        stridesAttr ? stridesAttr.asArrayRef().begin() : tileShapeArr.begin(),
+        stridesAttr ? stridesAttr.asArrayRef().end() : tileShapeArr.end());
 
-    auto output = emitSliceFromSmem(
-        rewriter, op.getLoc(), adaptor.getSrc(), adaptor.getIndex(),
-        op.getSrc().getType().getShape(), tileShapeArr, stridesVec,
-        resultTensorTy, outputType, userAnalysis, replaced2Origin, pTagPool,
-        op);
+    auto output =
+        emitSliceFromSmem(rewriter, op.getLoc(), adaptor.getSrc(),
+                          adaptor.getIndex(), op.getSrc().getType().getShape(),
+                          tileShapeArr, stridesVec, resultTensorTy, outputType,
+                          userAnalysis, replaced2Origin, pTagPool, op);
 
     leaveTritionOp(rewriter, op.getOperation());
     rewriter.replaceOp(op, output);
@@ -1495,16 +1493,14 @@ struct DesliceToLocalOpLowering
     auto tileShapeArr = op.getTileShape();
     auto stridesAttr = op.getTileStrides();
     SmallVector<int64_t> stridesVec(
-        stridesAttr ? stridesAttr.asArrayRef().begin()
-                    : tileShapeArr.begin(),
-        stridesAttr ? stridesAttr.asArrayRef().end()
-                    : tileShapeArr.end());
+        stridesAttr ? stridesAttr.asArrayRef().begin() : tileShapeArr.begin(),
+        stridesAttr ? stridesAttr.asArrayRef().end() : tileShapeArr.end());
 
     auto output = emitDesliceToSmem(
         rewriter, op.getLoc(), adaptor.getSrc(), adaptor.getTile(),
         adaptor.getIndex(), op.getSrc().getType().getShape(), tileShapeArr,
-        stridesVec, tileTensorTy, resultTensorTy, userAnalysis,
-        replaced2Origin, pTagPool, op);
+        stridesVec, tileTensorTy, resultTensorTy, userAnalysis, replaced2Origin,
+        pTagPool, op);
 
     leaveTritionOp(rewriter, op.getOperation());
     rewriter.replaceOp(op, output);

From 4f9dd87078e290fb7989cf1f336e06de5691ae7f Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 18:32:06 +0800
Subject: [PATCH 13/41] [TLE]fix some GCU backend problems 3.

---
 python/tutorials/tle/04-cluster-gemm.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index 2c5c3e9b58..c8c184241c 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -485,13 +485,22 @@ def _autotune_config(
 ) -> KernelConfig:
     best_cfg = candidates[0]
     best_ms = float("inf")
+    any_success = False
     for cfg in candidates:
-        run_fn(cfg)
-        torch.cuda.synchronize()
-        ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)
+        try:
+            run_fn(cfg)
+            torch.cuda.synchronize()
+            ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)
+            any_success = True
+        except RuntimeError as e:
+            print(f"[autotune] {name}: skip cfg={cfg} ({e})")
+            continue
         if ms < best_ms:
             best_ms = ms
             best_cfg = cfg
+    if not any_success:
+        print(f"[autotune] {name}: all configs failed, giving up")
+        return None
     print(f"[autotune] {name}: best cfg={best_cfg} ms={best_ms:.3f} tflops={_tflops(M, N, K, best_ms):.2f}")
     return best_cfg
 
@@ -584,6 +593,9 @@ def main(argv: list[str] | None = None) -> None:
             args.tune_warmup,
             args.tune_rep,
         )
+        if remote_cfg is None:
+            print("SKIP: all cluster remote configs failed with misaligned address.")
+            return
 
     print(f"selected triton cfg: {triton_cfg}")
     print(f"selected remote cfg: {remote_cfg}")

From da21a7d9d07c867659bba939d69a263cf7e03fc7 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 19:28:14 +0800
Subject: [PATCH 14/41] [TLE] fix some GCU backend problems 4.

---
 python/tutorials/tle/04-cluster-gemm.py              |  5 +++++
 .../lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp  | 12 ++++++------
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index c8c184241c..894ecd86f4 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -494,6 +494,11 @@ def _autotune_config(
             any_success = True
         except RuntimeError as e:
             print(f"[autotune] {name}: skip cfg={cfg} ({e})")
+            # Clear the sticky CUDA error so subsequent kernels are not affected.
+            # cudaGetLastError is the only way to clear a sticky async error.
+            import ctypes
+            cuda = ctypes.CDLL("libcudart.so")
+            cuda.cudaGetLastError()
             continue
         if ms < best_ms:
             best_ms = ms
diff --git a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
index 78b578d728..0f3204477f 100644
--- a/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
+++ b/third_party/enflame/triton_gcu/triton_gcu400/lib/Conversion/TritonToGCU/TTSmemLowerToGCU.cpp
@@ -1453,10 +1453,10 @@ struct SliceFromLocalOpLowering
         dyn_cast<MemRefType>(getTypeConverter()->convertType(resultTensorTy));
 
     auto tileShapeArr = op.getTileShape();
-    auto stridesAttr = op.getTileStrides();
+    auto stridesOpt = op.getTileStrides();
     SmallVector<int64_t> stridesVec(
-        stridesAttr ? stridesAttr.asArrayRef().begin() : tileShapeArr.begin(),
-        stridesAttr ? stridesAttr.asArrayRef().end() : tileShapeArr.end());
+        stridesOpt.has_value() ? stridesOpt->begin() : tileShapeArr.begin(),
+        stridesOpt.has_value() ? stridesOpt->end() : tileShapeArr.end());
 
     auto output =
         emitSliceFromSmem(rewriter, op.getLoc(), adaptor.getSrc(),
@@ -1491,10 +1491,10 @@ struct DesliceToLocalOpLowering
     auto resultTensorTy = cast<RankedTensorType>(op.getResult().getType());
 
     auto tileShapeArr = op.getTileShape();
-    auto stridesAttr = op.getTileStrides();
+    auto stridesOpt = op.getTileStrides();
     SmallVector<int64_t> stridesVec(
-        stridesAttr ? stridesAttr.asArrayRef().begin() : tileShapeArr.begin(),
-        stridesAttr ? stridesAttr.asArrayRef().end() : tileShapeArr.end());
+        stridesOpt.has_value() ? stridesOpt->begin() : tileShapeArr.begin(),
+        stridesOpt.has_value() ? stridesOpt->end() : tileShapeArr.end());
 
     auto output = emitDesliceToSmem(
         rewriter, op.getLoc(), adaptor.getSrc(), adaptor.getTile(),

From 0d74e82aafb391e7f053adfa3041fd4533a9bf79 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 19:52:06 +0800
Subject: [PATCH 15/41] [TLE] Handle misaligned address errors gracefully in
 cluster-gemm autotuning

---
 python/tutorials/tle/04-cluster-gemm.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index 894ecd86f4..d3d01b4960 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -579,6 +579,14 @@ def main(argv: list[str] | None = None) -> None:
             args.tune_warmup,
             args.tune_rep,
         )
+        # Run triton correctness check before remote autotuning, because failed
+        # DSMEM kernels can leave CUDA in a state where subsequent kernels fail.
+        if args.check:
+            _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
+                        triton_cfg.num_stages)
+            ref = torch.matmul(a, b)
+            torch.testing.assert_close(c_triton, ref, atol=1e-1, rtol=1e-1)
+            print("triton correctness check: PASS")
         remote_cfg = _autotune_config(
             "cluster_tle_remote_gemm",
             REMOTE_TUNE_CONFIGS,
@@ -672,8 +680,13 @@ def main(argv: list[str] | None = None) -> None:
         print("remote lowering check: PASS")
 
     if args.check:
-        _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
-                    triton_cfg.num_stages)
+        ref = torch.matmul(a, b)
+        if not args.autotune:
+            # With --no-autotune, triton check wasn't done during autotuning.
+            _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
+                        triton_cfg.num_stages)
+            torch.testing.assert_close(c_triton, ref, atol=1e-1, rtol=1e-1)
+            print("triton correctness check: PASS")
         _run_cluster_remote(
             a,
             b,
@@ -684,10 +697,8 @@ def main(argv: list[str] | None = None) -> None:
             remote_cfg.num_warps,
             remote_cfg.num_stages,
         )
-        ref = torch.matmul(a, b)
-        torch.testing.assert_close(c_triton, ref, atol=1e-1, rtol=1e-1)
         torch.testing.assert_close(c_remote, ref, atol=1e-1, rtol=1e-1)
-        print("correctness check: PASS")
+        print("remote correctness check: PASS")
 
     run_triton = lambda: _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
                                      triton_cfg.num_stages)

From b1f8dac90bf1488ce79bee039200f83aa4d79115 Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Thu, 4 Jun 2026 20:31:46 +0800
Subject: [PATCH 16/41] fix 04-cluster-gemm problems

---
 python/tutorials/tle/04-cluster-gemm.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index d3d01b4960..ba09836bf8 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -581,11 +581,11 @@ def main(argv: list[str] | None = None) -> None:
         )
         # Run triton correctness check before remote autotuning, because failed
         # DSMEM kernels can leave CUDA in a state where subsequent kernels fail.
+        _ref = torch.matmul(a, b)
         if args.check:
             _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
                         triton_cfg.num_stages)
-            ref = torch.matmul(a, b)
-            torch.testing.assert_close(c_triton, ref, atol=1e-1, rtol=1e-1)
+            torch.testing.assert_close(c_triton, _ref, atol=1e-1, rtol=1e-1)
             print("triton correctness check: PASS")
         remote_cfg = _autotune_config(
             "cluster_tle_remote_gemm",
@@ -680,12 +680,12 @@ def main(argv: list[str] | None = None) -> None:
         print("remote lowering check: PASS")
 
     if args.check:
-        ref = torch.matmul(a, b)
         if not args.autotune:
-            # With --no-autotune, triton check wasn't done during autotuning.
+            # With --no-autotune, compute ref and run triton check.
+            _ref = torch.matmul(a, b)
             _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
                         triton_cfg.num_stages)
-            torch.testing.assert_close(c_triton, ref, atol=1e-1, rtol=1e-1)
+            torch.testing.assert_close(c_triton, _ref, atol=1e-1, rtol=1e-1)
             print("triton correctness check: PASS")
         _run_cluster_remote(
             a,
@@ -697,7 +697,7 @@ def main(argv: list[str] | None = None) -> None:
             remote_cfg.num_warps,
             remote_cfg.num_stages,
         )
-        torch.testing.assert_close(c_remote, ref, atol=1e-1, rtol=1e-1)
+        torch.testing.assert_close(c_remote, _ref, atol=1e-1, rtol=1e-1)
         print("remote correctness check: PASS")
 
     run_triton = lambda: _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,

From 9f14df64179ecf24dfa5aa0eaca33393e189686b Mon Sep 17 00:00:00 2001
From: lzllx123 <1803100521@qq.com>
Date: Fri, 5 Jun 2026 10:05:43 +0800
Subject: [PATCH 17/41] Revert 04-cluster-gemm.py to c7e6953

---
 python/tutorials/tle/04-cluster-gemm.py | 46 +++++--------------------
 1 file changed, 9 insertions(+), 37 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index ba09836bf8..2c5c3e9b58 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -485,27 +485,13 @@ def _autotune_config(
 ) -> KernelConfig:
     best_cfg = candidates[0]
     best_ms = float("inf")
-    any_success = False
     for cfg in candidates:
-        try:
-            run_fn(cfg)
-            torch.cuda.synchronize()
-            ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)
-            any_success = True
-        except RuntimeError as e:
-            print(f"[autotune] {name}: skip cfg={cfg} ({e})")
-            # Clear the sticky CUDA error so subsequent kernels are not affected.
-            # cudaGetLastError is the only way to clear a sticky async error.
-            import ctypes
-            cuda = ctypes.CDLL("libcudart.so")
-            cuda.cudaGetLastError()
-            continue
+        run_fn(cfg)
+        torch.cuda.synchronize()
+        ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)
         if ms < best_ms:
             best_ms = ms
             best_cfg = cfg
-    if not any_success:
-        print(f"[autotune] {name}: all configs failed, giving up")
-        return None
     print(f"[autotune] {name}: best cfg={best_cfg} ms={best_ms:.3f} tflops={_tflops(M, N, K, best_ms):.2f}")
     return best_cfg
 
@@ -579,14 +565,6 @@ def main(argv: list[str] | None = None) -> None:
             args.tune_warmup,
             args.tune_rep,
         )
-        # Run triton correctness check before remote autotuning, because failed
-        # DSMEM kernels can leave CUDA in a state where subsequent kernels fail.
-        _ref = torch.matmul(a, b)
-        if args.check:
-            _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
-                        triton_cfg.num_stages)
-            torch.testing.assert_close(c_triton, _ref, atol=1e-1, rtol=1e-1)
-            print("triton correctness check: PASS")
         remote_cfg = _autotune_config(
             "cluster_tle_remote_gemm",
             REMOTE_TUNE_CONFIGS,
@@ -606,9 +584,6 @@ def main(argv: list[str] | None = None) -> None:
             args.tune_warmup,
             args.tune_rep,
         )
-        if remote_cfg is None:
-            print("SKIP: all cluster remote configs failed with misaligned address.")
-            return
 
     print(f"selected triton cfg: {triton_cfg}")
     print(f"selected remote cfg: {remote_cfg}")
@@ -680,13 +655,8 @@ def main(argv: list[str] | None = None) -> None:
         print("remote lowering check: PASS")
 
     if args.check:
-        if not args.autotune:
-            # With --no-autotune, compute ref and run triton check.
-            _ref = torch.matmul(a, b)
-            _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
-                        triton_cfg.num_stages)
-            torch.testing.assert_close(c_triton, _ref, atol=1e-1, rtol=1e-1)
-            print("triton correctness check: PASS")
+        _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
+                    triton_cfg.num_stages)
         _run_cluster_remote(
             a,
             b,
@@ -697,8 +667,10 @@ def main(argv: list[str] | None = None) -> None:
             remote_cfg.num_warps,
             remote_cfg.num_stages,
         )
-        torch.testing.assert_close(c_remote, _ref, atol=1e-1, rtol=1e-1)
-        print("remote correctness check: PASS")
+        ref = torch.matmul(a, b)
+        torch.testing.assert_close(c_triton, ref, atol=1e-1, rtol=1e-1)
+        torch.testing.assert_close(c_remote, ref, atol=1e-1, rtol=1e-1)
+        print("correctness check: PASS")
 
     run_triton = lambda: _run_triton(a, b, c_triton, triton_cfg.bm, triton_cfg.bn, triton_cfg.bk, triton_cfg.num_warps,
                                      triton_cfg.num_stages)

From aaba5bcd68dafde388493e71165f14c02a6bf51d Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 06:24:25 +0000
Subject: [PATCH 18/41] Debug CI cluster GEMM failure

---
 .github/workflows/hopper-build-and-test.yml | 11 ++++++-
 python/tutorials/tle/04-cluster-gemm.py     | 32 ++++++++++++++++++++-
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 4084843f55..2a2db4dbbb 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -180,7 +180,16 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
+          echo "::group::TLE 04 debug: suspected BK=256 remote config"
+          PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 CUDA_LAUNCH_BLOCKING=1 \
+            python3 python/tutorials/tle/04-cluster-gemm.py \
+              --no-autotune --bm 32 --bn 512 --bk 256 --num-warps 4 --num-stages 2 \
+              --warmup 2 --rep 5
+          echo "::endgroup::"
+          echo "::group::TLE 04 debug: full default run"
+          PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 CUDA_LAUNCH_BLOCKING=1 \
+            python3 python/tutorials/tle/04-cluster-gemm.py
+          echo "::endgroup::"
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
           # python unit test
diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index 2c5c3e9b58..f36cccd467 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -14,9 +14,15 @@
 from __future__ import annotations
 
 import argparse
+import os
+import sys
 from dataclasses import dataclass
 from typing import Callable
 
+CI_DEBUG = os.environ.get("FLAGTREE_CI_DEBUG_CLUSTER_GEMM") == "1"
+if CI_DEBUG:
+    os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
+
 import torch
 import triton
 import triton.language as tl
@@ -25,6 +31,24 @@
 BLOCK_CLUSTER_MESH = tle.device_mesh({"block_cluster": [("cluster_x", 2)]})
 
 
+def _debug(msg: str) -> None:
+    if CI_DEBUG:
+        print(f"[ci-debug] {msg}", flush=True)
+
+
+def _print_runtime_debug(args: argparse.Namespace) -> None:
+    if not CI_DEBUG:
+        return
+    _debug(f"python={sys.version.split()[0]} executable={sys.executable}")
+    _debug(f"torch={torch.__version__} triton={getattr(triton, '__version__', 'unknown')}")
+    _debug(f"triton_file={triton.__file__}")
+    _debug(f"triton_cache_dir={os.environ.get('TRITON_CACHE_DIR', '<default>')}")
+    _debug(f"cuda_launch_blocking={os.environ.get('CUDA_LAUNCH_BLOCKING', '<unset>')}")
+    if torch.cuda.is_available():
+        _debug(f"cuda_device={torch.cuda.get_device_name()} capability={torch.cuda.get_device_capability()}")
+    _debug(f"args={args}")
+
+
 def _select_dot_k(bk: int) -> int:
     if bk % 32 == 0:
         return 32
@@ -485,10 +509,14 @@ def _autotune_config(
 ) -> KernelConfig:
     best_cfg = candidates[0]
     best_ms = float("inf")
-    for cfg in candidates:
+    for idx, cfg in enumerate(candidates, start=1):
+        _debug(f"{name}: candidate {idx}/{len(candidates)} launch cfg={cfg}")
         run_fn(cfg)
+        _debug(f"{name}: candidate {idx}/{len(candidates)} launch returned; synchronizing")
         torch.cuda.synchronize()
+        _debug(f"{name}: candidate {idx}/{len(candidates)} synchronized; benchmarking warmup={warmup} rep={rep}")
         ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)
+        _debug(f"{name}: candidate {idx}/{len(candidates)} bench_ms={ms:.6f}")
         if ms < best_ms:
             best_ms = ms
             best_cfg = cfg
@@ -537,6 +565,8 @@ def main(argv: list[str] | None = None) -> None:
         print(f"SKIP: {skip_reason}")
         return
 
+    _print_runtime_debug(args)
+
     torch.manual_seed(args.seed)
     dtype = torch.float16
     a = torch.randn((args.m, args.k), device="cuda", dtype=dtype)

From 70fc3f445a7616925ae650ba08dbade5958048a9 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 06:28:27 +0000
Subject: [PATCH 19/41] Fix debug import lint

---
 python/tutorials/tle/04-cluster-gemm.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index f36cccd467..d0abfbc123 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -23,10 +23,10 @@
 if CI_DEBUG:
     os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
 
-import torch
-import triton
-import triton.language as tl
-import triton.experimental.tle.language as tle
+import torch  # noqa: E402
+import triton  # noqa: E402
+import triton.language as tl  # noqa: E402
+import triton.experimental.tle.language as tle  # noqa: E402
 
 BLOCK_CLUSTER_MESH = tle.device_mesh({"block_cluster": [("cluster_x", 2)]})
 

From 6e254a55963116572f7b5d49d419e47466185ab9 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 07:01:03 +0000
Subject: [PATCH 20/41] Debug cold full cluster GEMM CI run

---
 .github/workflows/hopper-build-and-test.yml | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 2a2db4dbbb..86fdff0ec3 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -180,12 +180,6 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          echo "::group::TLE 04 debug: suspected BK=256 remote config"
-          PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 CUDA_LAUNCH_BLOCKING=1 \
-            python3 python/tutorials/tle/04-cluster-gemm.py \
-              --no-autotune --bm 32 --bn 512 --bk 256 --num-warps 4 --num-stages 2 \
-              --warmup 2 --rep 5
-          echo "::endgroup::"
           echo "::group::TLE 04 debug: full default run"
           PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 CUDA_LAUNCH_BLOCKING=1 \
             python3 python/tutorials/tle/04-cluster-gemm.py

From b7c8f2e5d4d375e1c917e18315c34c674922ba5c Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 07:50:02 +0000
Subject: [PATCH 21/41] Debug nonblocking cluster GEMM CI run

---
 .github/workflows/hopper-build-and-test.yml |  4 ++--
 python/tutorials/tle/04-cluster-gemm.py     | 12 +++++-------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 86fdff0ec3..0d354fe2ff 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -180,8 +180,8 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          echo "::group::TLE 04 debug: full default run"
-          PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 CUDA_LAUNCH_BLOCKING=1 \
+          echo "::group::TLE 04 debug: full default run without CUDA_LAUNCH_BLOCKING"
+          PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 \
             python3 python/tutorials/tle/04-cluster-gemm.py
           echo "::endgroup::"
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index d0abfbc123..0a03198f40 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -19,14 +19,12 @@
 from dataclasses import dataclass
 from typing import Callable
 
-CI_DEBUG = os.environ.get("FLAGTREE_CI_DEBUG_CLUSTER_GEMM") == "1"
-if CI_DEBUG:
-    os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
+import torch
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
 
-import torch  # noqa: E402
-import triton  # noqa: E402
-import triton.language as tl  # noqa: E402
-import triton.experimental.tle.language as tle  # noqa: E402
+CI_DEBUG = os.environ.get("FLAGTREE_CI_DEBUG_CLUSTER_GEMM") == "1"
 
 BLOCK_CLUSTER_MESH = tle.device_mesh({"block_cluster": [("cluster_x", 2)]})
 

From ce073c4e44c663c24d07c7a7da053c8d15850832 Mon Sep 17 00:00:00 2001
From: Liu Zhengliang <1803100521@qq.com>
Date: Fri, 5 Jun 2026 16:34:36 +0800
Subject: [PATCH 22/41] Update hopper-build-and-test.yml

---
 .github/workflows/hopper-build-and-test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 4084843f55..f6cb51dafb 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -122,7 +122,7 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py --mode test
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
@@ -180,7 +180,7 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
           # python unit test

From 51d1b1bd6c140370e837c3532df4b5b130bc9394 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 09:09:47 +0000
Subject: [PATCH 23/41] Test cluster GEMM with blocking launches

---
 .github/workflows/hopper-build-and-test.yml |  5 +---
 python/tutorials/tle/04-cluster-gemm.py     | 30 +--------------------
 2 files changed, 2 insertions(+), 33 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 0d354fe2ff..f7c27c65d6 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -180,10 +180,7 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          echo "::group::TLE 04 debug: full default run without CUDA_LAUNCH_BLOCKING"
-          PYTHONUNBUFFERED=1 FLAGTREE_CI_DEBUG_CLUSTER_GEMM=1 \
-            python3 python/tutorials/tle/04-cluster-gemm.py
-          echo "::endgroup::"
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
           # python unit test
diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index 0a03198f40..2c5c3e9b58 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -14,8 +14,6 @@
 from __future__ import annotations
 
 import argparse
-import os
-import sys
 from dataclasses import dataclass
 from typing import Callable
 
@@ -24,29 +22,9 @@
 import triton.language as tl
 import triton.experimental.tle.language as tle
 
-CI_DEBUG = os.environ.get("FLAGTREE_CI_DEBUG_CLUSTER_GEMM") == "1"
-
 BLOCK_CLUSTER_MESH = tle.device_mesh({"block_cluster": [("cluster_x", 2)]})
 
 
-def _debug(msg: str) -> None:
-    if CI_DEBUG:
-        print(f"[ci-debug] {msg}", flush=True)
-
-
-def _print_runtime_debug(args: argparse.Namespace) -> None:
-    if not CI_DEBUG:
-        return
-    _debug(f"python={sys.version.split()[0]} executable={sys.executable}")
-    _debug(f"torch={torch.__version__} triton={getattr(triton, '__version__', 'unknown')}")
-    _debug(f"triton_file={triton.__file__}")
-    _debug(f"triton_cache_dir={os.environ.get('TRITON_CACHE_DIR', '<default>')}")
-    _debug(f"cuda_launch_blocking={os.environ.get('CUDA_LAUNCH_BLOCKING', '<unset>')}")
-    if torch.cuda.is_available():
-        _debug(f"cuda_device={torch.cuda.get_device_name()} capability={torch.cuda.get_device_capability()}")
-    _debug(f"args={args}")
-
-
 def _select_dot_k(bk: int) -> int:
     if bk % 32 == 0:
         return 32
@@ -507,14 +485,10 @@ def _autotune_config(
 ) -> KernelConfig:
     best_cfg = candidates[0]
     best_ms = float("inf")
-    for idx, cfg in enumerate(candidates, start=1):
-        _debug(f"{name}: candidate {idx}/{len(candidates)} launch cfg={cfg}")
+    for cfg in candidates:
         run_fn(cfg)
-        _debug(f"{name}: candidate {idx}/{len(candidates)} launch returned; synchronizing")
         torch.cuda.synchronize()
-        _debug(f"{name}: candidate {idx}/{len(candidates)} synchronized; benchmarking warmup={warmup} rep={rep}")
         ms = triton.testing.do_bench(lambda: run_fn(cfg), warmup=warmup, rep=rep)
-        _debug(f"{name}: candidate {idx}/{len(candidates)} bench_ms={ms:.6f}")
         if ms < best_ms:
             best_ms = ms
             best_cfg = cfg
@@ -563,8 +537,6 @@ def main(argv: list[str] | None = None) -> None:
         print(f"SKIP: {skip_reason}")
         return
 
-    _print_runtime_debug(args)
-
     torch.manual_seed(args.seed)
     dtype = torch.float16
     a = torch.randn((args.m, args.k), device="cuda", dtype=dtype)

From 15734b78b970f7bafbc7b62ca8b112a67627615f Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 09:38:42 +0000
Subject: [PATCH 24/41] Run only cluster GEMM in Hopper debug

---
 .github/workflows/hopper-build-and-test.yml | 54 ---------------------
 1 file changed, 54 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index f7c27c65d6..9a92213eb5 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,58 +155,4 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          # python tutorials
-          python3 python/tutorials/01-vector-add.py --only_unit_test
-          python3 python/tutorials/02-fused-softmax.py --only_unit_test
-          python3 python/tutorials/03-matrix-multiplication.py --only_unit_test
-          python3 python/tutorials/04-low-memory-dropout.py --only_unit_test
-          python3 python/tutorials/05-layer-norm.py --only_unit_test
-          python3 python/tutorials/06-fused-attention.py --only_unit_test
-          python3 python/tutorials/07-extern-functions.py --only_unit_test
-          python3 python/tutorials/08-grouped-gemm.py --only_unit_test
-          python3 python/tutorials/09-persistent-matmul.py --only_unit_test
-          python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test
-          # python unit test
-          python3 -m pytest -s python/test/unit/cuda
-          python3 -m pytest -s python/test/unit/instrumentation
-          python3 -m pytest -s python/test/unit/language \
-              --ignore=python/test/unit/language/test_line_info.py
-          python3 -m pytest -s python/test/unit/runtime
-          if [ -d "python/test/operators" ]; then
-            python3 -m pytest -s python/test/operators
-          fi
-          # flagtree tle
-          # python tutorials
-          python3 python/tutorials/tle/01-fft.py
-          python3 python/tutorials/tle/02-moe_align_block_size.py
-          python3 python/tutorials/tle/03-topk.py
           CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
-          python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
-          python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
-          # python unit test
-          python3 -m pytest -s python/test/tle/integration
-          python3 -m pytest -s python/test/tle/unit
-          # flagtree hints
-          # python tutorials
-          python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
-          # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test
-          python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test
-          python3 python/tutorials/hints/04/04-low-memory-dropout.py --only_unit_test
-          python3 python/tutorials/hints/05/05-layer-norm.py --only_unit_test
-          # python3 python/tutorials/hints/06/06-fused-attention.py --only_unit_test  # Error on 3.6
-          python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test
-          python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test
-          python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test
-          # flagtree tle raw
-          python3 python/tutorials/tle/raw/mlir/01-vector-add.py
-          python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py
-          python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py
-          python3 python/tutorials/tle/raw/mlir/04-hello-world.py
-          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton
-          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle
-          python3 python/tutorials/tle/raw/mlir/06-test-vassert.py
-          # flagtree tle cuda
-          # TODO: These tests are currently skipped because the CLANG environment variable cannot be set.
-          # python3 python/tutorials/tle/raw/cuda/01-vector-add.py
-          # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py
-          # python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py

From b840cadf246e97d7a19219f6abbe5cf2d530c48b Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 09:48:59 +0000
Subject: [PATCH 25/41] Debug cluster GEMM without autotune

---
 .github/workflows/hopper-build-and-test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 9a92213eb5..0f5994cf88 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,4 +155,4 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py --no-autotune --warmup 1 --rep 1

From 3614cab3d634c747cbb47b1bda4ca38267d0ad43 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 09:58:18 +0000
Subject: [PATCH 26/41] Debug first remote cluster GEMM config

---
 .github/workflows/hopper-build-and-test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 0f5994cf88..faec112e8f 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,4 +155,4 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py --no-autotune --warmup 1 --rep 1
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py --no-autotune --bm 32 --bn 512 --bk 256 --num-warps 4 --num-stages 2 --warmup 1 --rep 1

From ad7bd5bfb3aa905faba9aaf75487179d73461376 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 10:19:25 +0000
Subject: [PATCH 27/41] Debug remote autotune only

---
 .github/workflows/hopper-build-and-test.yml | 47 ++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index faec112e8f..3af2508adf 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,4 +155,49 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py --no-autotune --bm 32 --bn 512 --bk 256 --num-warps 4 --num-stages 2 --warmup 1 --rep 1
+          CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY'
+          import importlib.util
+          import pathlib
+          import sys
+
+          import torch
+
+          path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py")
+          spec = importlib.util.spec_from_file_location("cluster_gemm_debug", path)
+          mod = importlib.util.module_from_spec(spec)
+          sys.modules[spec.name] = mod
+          spec.loader.exec_module(mod)
+
+          skip_reason = mod._cluster_remote_support_skip_reason()
+          if skip_reason is not None:
+              print(f"SKIP: {skip_reason}")
+              raise SystemExit(0)
+
+          M = N = K = 4096
+          torch.manual_seed(0)
+          a = torch.randn((M, K), device="cuda", dtype=torch.float16)
+          b = torch.randn((K, N), device="cuda", dtype=torch.float16)
+          c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16)
+
+          print("[debug] remote autotune only")
+          remote_cfg = mod._autotune_config(
+              "cluster_tle_remote_gemm_only",
+              mod.REMOTE_TUNE_CONFIGS,
+              lambda cfg: mod._run_cluster_remote(
+                  a,
+                  b,
+                  c_remote,
+                  cfg.bm,
+                  cfg.bn,
+                  cfg.bk,
+                  cfg.num_warps,
+                  cfg.num_stages,
+              ),
+              M,
+              N,
+              K,
+              1,
+              1,
+          )
+          print(f"[debug] remote autotune only selected {remote_cfg}")
+          PY

From c1bea1511f41f461f4bdf55b0ddd65e1e63b3959 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 10:29:10 +0000
Subject: [PATCH 28/41] Debug fixed config without lowering warmup

---
 .github/workflows/hopper-build-and-test.yml | 47 +--------------------
 1 file changed, 1 insertion(+), 46 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 3af2508adf..a46b1387b0 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,49 +155,4 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY'
-          import importlib.util
-          import pathlib
-          import sys
-
-          import torch
-
-          path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py")
-          spec = importlib.util.spec_from_file_location("cluster_gemm_debug", path)
-          mod = importlib.util.module_from_spec(spec)
-          sys.modules[spec.name] = mod
-          spec.loader.exec_module(mod)
-
-          skip_reason = mod._cluster_remote_support_skip_reason()
-          if skip_reason is not None:
-              print(f"SKIP: {skip_reason}")
-              raise SystemExit(0)
-
-          M = N = K = 4096
-          torch.manual_seed(0)
-          a = torch.randn((M, K), device="cuda", dtype=torch.float16)
-          b = torch.randn((K, N), device="cuda", dtype=torch.float16)
-          c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16)
-
-          print("[debug] remote autotune only")
-          remote_cfg = mod._autotune_config(
-              "cluster_tle_remote_gemm_only",
-              mod.REMOTE_TUNE_CONFIGS,
-              lambda cfg: mod._run_cluster_remote(
-                  a,
-                  b,
-                  c_remote,
-                  cfg.bm,
-                  cfg.bn,
-                  cfg.bk,
-                  cfg.num_warps,
-                  cfg.num_stages,
-              ),
-              M,
-              N,
-              K,
-              1,
-              1,
-          )
-          print(f"[debug] remote autotune only selected {remote_cfg}")
-          PY
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py --no-autotune --no-check-lowering --bm 32 --bn 512 --bk 256 --num-warps 4 --num-stages 2 --warmup 1 --rep 1

From 72bc8a509de721cd36ec86ffc68205608d449776 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 10:39:43 +0000
Subject: [PATCH 29/41] Debug remote fixed launch only

---
 .github/workflows/hopper-build-and-test.yml | 32 ++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index a46b1387b0..89ffade073 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,4 +155,34 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py --no-autotune --no-check-lowering --bm 32 --bn 512 --bk 256 --num-warps 4 --num-stages 2 --warmup 1 --rep 1
+          CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY'
+          import importlib.util
+          import pathlib
+          import sys
+
+          import torch
+
+          path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py")
+          spec = importlib.util.spec_from_file_location("cluster_gemm_debug", path)
+          mod = importlib.util.module_from_spec(spec)
+          sys.modules[spec.name] = mod
+          spec.loader.exec_module(mod)
+
+          skip_reason = mod._cluster_remote_support_skip_reason()
+          if skip_reason is not None:
+              print(f"SKIP: {skip_reason}")
+              raise SystemExit(0)
+
+          M = N = K = 4096
+          torch.manual_seed(0)
+          a = torch.randn((M, K), device="cuda", dtype=torch.float16)
+          b = torch.randn((K, N), device="cuda", dtype=torch.float16)
+          c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16)
+          cfg = mod.KernelConfig(32, 512, 256, 4, 2)
+
+          print(f"[debug] ptr mod256 a={a.data_ptr() % 256} b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}")
+          print(f"[debug] remote fixed only launch {cfg}")
+          mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages)
+          torch.cuda.synchronize()
+          print("[debug] remote fixed only launch: PASS")
+          PY

From 8725fef9ea835a02aa49a912dae91c768325925b Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 10:59:11 +0000
Subject: [PATCH 30/41] Debug remote autotune candidate progress

---
 .github/workflows/hopper-build-and-test.yml | 28 ++++++++++++++++-----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 89ffade073..d8a79c878c 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,6 +155,7 @@ jobs:
         run: |
           set -x
           source ~/env.sh
+          trap 'echo "[debug] last remote candidate:"; cat /tmp/remote-candidate.txt || true' ERR
           CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY'
           import importlib.util
           import pathlib
@@ -178,11 +179,26 @@ jobs:
           a = torch.randn((M, K), device="cuda", dtype=torch.float16)
           b = torch.randn((K, N), device="cuda", dtype=torch.float16)
           c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16)
-          cfg = mod.KernelConfig(32, 512, 256, 4, 2)
-
+          progress = pathlib.Path("/tmp/remote-candidate.txt")
           print(f"[debug] ptr mod256 a={a.data_ptr() % 256} b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}")
-          print(f"[debug] remote fixed only launch {cfg}")
-          mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages)
-          torch.cuda.synchronize()
-          print("[debug] remote fixed only launch: PASS")
+          for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS):
+              progress.write_text(f"{idx}: {cfg}\n")
+              mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages)
+              torch.cuda.synchronize()
+              ms = mod.triton.testing.do_bench(
+                  lambda: mod._run_cluster_remote(
+                      a,
+                      b,
+                      c_remote,
+                      cfg.bm,
+                      cfg.bn,
+                      cfg.bk,
+                      cfg.num_warps,
+                      cfg.num_stages,
+                  ),
+                  warmup=1,
+                  rep=1,
+              )
+              progress.write_text(f"{idx}: {cfg} done ms={ms}\n")
+          print("[debug] manual remote autotune loop: PASS")
           PY

From 790f435773707bb6a54aec5c2f3ca4dd9dddd75c Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 11:08:46 +0000
Subject: [PATCH 31/41] Debug second remote candidate alone

---
 .github/workflows/hopper-build-and-test.yml | 27 ++++-----------------
 1 file changed, 5 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index d8a79c878c..8bb81ffe39 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,7 +155,6 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          trap 'echo "[debug] last remote candidate:"; cat /tmp/remote-candidate.txt || true' ERR
           CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY'
           import importlib.util
           import pathlib
@@ -179,26 +178,10 @@ jobs:
           a = torch.randn((M, K), device="cuda", dtype=torch.float16)
           b = torch.randn((K, N), device="cuda", dtype=torch.float16)
           c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16)
-          progress = pathlib.Path("/tmp/remote-candidate.txt")
+          cfg = mod.KernelConfig(32, 512, 128, 4, 3)
           print(f"[debug] ptr mod256 a={a.data_ptr() % 256} b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}")
-          for idx, cfg in enumerate(mod.REMOTE_TUNE_CONFIGS):
-              progress.write_text(f"{idx}: {cfg}\n")
-              mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages)
-              torch.cuda.synchronize()
-              ms = mod.triton.testing.do_bench(
-                  lambda: mod._run_cluster_remote(
-                      a,
-                      b,
-                      c_remote,
-                      cfg.bm,
-                      cfg.bn,
-                      cfg.bk,
-                      cfg.num_warps,
-                      cfg.num_stages,
-                  ),
-                  warmup=1,
-                  rep=1,
-              )
-              progress.write_text(f"{idx}: {cfg} done ms={ms}\n")
-          print("[debug] manual remote autotune loop: PASS")
+          print(f"[debug] remote candidate1 fixed launch {cfg}")
+          mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages)
+          torch.cuda.synchronize()
+          print("[debug] remote candidate1 fixed launch: PASS")
           PY

From 7943cb312152d892697e877ad6309fca1be55681 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 11:18:45 +0000
Subject: [PATCH 32/41] Debug cluster GEMM slots for pipelining

---
 python/tutorials/tle/04-cluster-gemm.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index 2c5c3e9b58..d5cdd21446 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -290,7 +290,7 @@ def _run_cluster_remote(
     N = b.shape[1]
     dot_k = _select_remote_dot_k(bk)
     use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
-    a_slots = 2
+    a_slots = max(2, num_stages)
     # BK=64 with 2-stage pipeline benefits from MMA-friendly shared layout in
     # the remote A path; BK=32 keeps the existing fast path.
     use_nv_mma_smem_layout = (bk == 32) or (bk == 64 and num_stages <= 2)
@@ -337,7 +337,7 @@ def _verify_remote_lowering(
     dot_k = _select_remote_dot_k(bk)
     use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
     # Keep verifier config aligned with runtime launch config.
-    a_slots = 2
+    a_slots = max(2, num_stages)
     use_nv_mma_smem_layout = (bk == 32) or (bk == 64 and num_stages <= 2)
     compiled = _cluster_remote_gemm_kernel.warmup(
         a,
@@ -631,7 +631,7 @@ def main(argv: list[str] | None = None) -> None:
             DOT_K=_select_remote_dot_k(remote_cfg.bk),
             CLUSTER_SIZE=2,
             USE_MASK=((args.m % remote_cfg.bm != 0) or (args.n % remote_cfg.bn != 0) or (args.k % remote_cfg.bk != 0)),
-            A_SLOTS=2,
+            A_SLOTS=max(2, remote_cfg.num_stages),
             USE_NV_MMA_SMEM_LAYOUT=((remote_cfg.bk == 32) or (remote_cfg.bk == 64 and remote_cfg.num_stages <= 2)),
             grid=_grid_cluster_remote(args.m, args.n, remote_cfg.bm, remote_cfg.bn),
             num_ctas=1,

From c962a16cd8b53d75e1a422be2476f0af646237eb Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 11:28:01 +0000
Subject: [PATCH 33/41] Debug power-of-two remote slots

---
 python/tutorials/tle/04-cluster-gemm.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index d5cdd21446..0388105f1d 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -290,7 +290,7 @@ def _run_cluster_remote(
     N = b.shape[1]
     dot_k = _select_remote_dot_k(bk)
     use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
-    a_slots = max(2, num_stages)
+    a_slots = 4 if num_stages > 2 else 2
     # BK=64 with 2-stage pipeline benefits from MMA-friendly shared layout in
     # the remote A path; BK=32 keeps the existing fast path.
     use_nv_mma_smem_layout = (bk == 32) or (bk == 64 and num_stages <= 2)
@@ -337,7 +337,7 @@ def _verify_remote_lowering(
     dot_k = _select_remote_dot_k(bk)
     use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
     # Keep verifier config aligned with runtime launch config.
-    a_slots = max(2, num_stages)
+    a_slots = 4 if num_stages > 2 else 2
     use_nv_mma_smem_layout = (bk == 32) or (bk == 64 and num_stages <= 2)
     compiled = _cluster_remote_gemm_kernel.warmup(
         a,
@@ -631,7 +631,7 @@ def main(argv: list[str] | None = None) -> None:
             DOT_K=_select_remote_dot_k(remote_cfg.bk),
             CLUSTER_SIZE=2,
             USE_MASK=((args.m % remote_cfg.bm != 0) or (args.n % remote_cfg.bn != 0) or (args.k % remote_cfg.bk != 0)),
-            A_SLOTS=max(2, remote_cfg.num_stages),
+            A_SLOTS=4 if remote_cfg.num_stages > 2 else 2,
             USE_NV_MMA_SMEM_LAYOUT=((remote_cfg.bk == 32) or (remote_cfg.bk == 64 and remote_cfg.num_stages <= 2)),
             grid=_grid_cluster_remote(args.m, args.n, remote_cfg.bm, remote_cfg.bn),
             num_ctas=1,

From baea8f2fdff7be28fac2bd9a83c5890e7808885b Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 11:37:54 +0000
Subject: [PATCH 34/41] Verify full cluster GEMM with slot fix

---
 .github/workflows/hopper-build-and-test.yml | 31 +--------------------
 1 file changed, 1 insertion(+), 30 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 8bb81ffe39..9a92213eb5 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,33 +155,4 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 - <<'PY'
-          import importlib.util
-          import pathlib
-          import sys
-
-          import torch
-
-          path = pathlib.Path("python/tutorials/tle/04-cluster-gemm.py")
-          spec = importlib.util.spec_from_file_location("cluster_gemm_debug", path)
-          mod = importlib.util.module_from_spec(spec)
-          sys.modules[spec.name] = mod
-          spec.loader.exec_module(mod)
-
-          skip_reason = mod._cluster_remote_support_skip_reason()
-          if skip_reason is not None:
-              print(f"SKIP: {skip_reason}")
-              raise SystemExit(0)
-
-          M = N = K = 4096
-          torch.manual_seed(0)
-          a = torch.randn((M, K), device="cuda", dtype=torch.float16)
-          b = torch.randn((K, N), device="cuda", dtype=torch.float16)
-          c_remote = torch.empty((M, N), device="cuda", dtype=torch.float16)
-          cfg = mod.KernelConfig(32, 512, 128, 4, 3)
-          print(f"[debug] ptr mod256 a={a.data_ptr() % 256} b={b.data_ptr() % 256} c={c_remote.data_ptr() % 256}")
-          print(f"[debug] remote candidate1 fixed launch {cfg}")
-          mod._run_cluster_remote(a, b, c_remote, cfg.bm, cfg.bn, cfg.bk, cfg.num_warps, cfg.num_stages)
-          torch.cuda.synchronize()
-          print("[debug] remote candidate1 fixed launch: PASS")
-          PY
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py

From dfbccf2b95192ef5bd64c66101ecaa8d1028a219 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Fri, 5 Jun 2026 11:47:56 +0000
Subject: [PATCH 35/41] Restore Hopper workflow after debug

---
 .github/workflows/hopper-build-and-test.yml | 56 ++++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 9a92213eb5..4084843f55 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,4 +155,58 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
+          # python tutorials
+          python3 python/tutorials/01-vector-add.py --only_unit_test
+          python3 python/tutorials/02-fused-softmax.py --only_unit_test
+          python3 python/tutorials/03-matrix-multiplication.py --only_unit_test
+          python3 python/tutorials/04-low-memory-dropout.py --only_unit_test
+          python3 python/tutorials/05-layer-norm.py --only_unit_test
+          python3 python/tutorials/06-fused-attention.py --only_unit_test
+          python3 python/tutorials/07-extern-functions.py --only_unit_test
+          python3 python/tutorials/08-grouped-gemm.py --only_unit_test
+          python3 python/tutorials/09-persistent-matmul.py --only_unit_test
+          python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test
+          # python unit test
+          python3 -m pytest -s python/test/unit/cuda
+          python3 -m pytest -s python/test/unit/instrumentation
+          python3 -m pytest -s python/test/unit/language \
+              --ignore=python/test/unit/language/test_line_info.py
+          python3 -m pytest -s python/test/unit/runtime
+          if [ -d "python/test/operators" ]; then
+            python3 -m pytest -s python/test/operators
+          fi
+          # flagtree tle
+          # python tutorials
+          python3 python/tutorials/tle/01-fft.py
+          python3 python/tutorials/tle/02-moe_align_block_size.py
+          python3 python/tutorials/tle/03-topk.py
+          python3 python/tutorials/tle/04-cluster-gemm.py
+          python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
+          python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
+          # python unit test
+          python3 -m pytest -s python/test/tle/integration
+          python3 -m pytest -s python/test/tle/unit
+          # flagtree hints
+          # python tutorials
+          python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
+          # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test
+          python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test
+          python3 python/tutorials/hints/04/04-low-memory-dropout.py --only_unit_test
+          python3 python/tutorials/hints/05/05-layer-norm.py --only_unit_test
+          # python3 python/tutorials/hints/06/06-fused-attention.py --only_unit_test  # Error on 3.6
+          python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test
+          python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test
+          python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test
+          # flagtree tle raw
+          python3 python/tutorials/tle/raw/mlir/01-vector-add.py
+          python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py
+          python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py
+          python3 python/tutorials/tle/raw/mlir/04-hello-world.py
+          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton
+          python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle
+          python3 python/tutorials/tle/raw/mlir/06-test-vassert.py
+          # flagtree tle cuda
+          # TODO: These tests are currently skipped because the CLANG environment variable cannot be set.
+          # python3 python/tutorials/tle/raw/cuda/01-vector-add.py
+          # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py
+          # python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py

From db71b554109dd3abf9b24807e35d98e74ed1256d Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Mon, 8 Jun 2026 02:35:10 +0000
Subject: [PATCH 36/41] Debug Hopper TTGIR for remote slots

---
 .github/workflows/hopper-build-and-test.yml | 71 +++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 4084843f55..c5871f4568 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,6 +155,77 @@ jobs:
         run: |
           set -x
           source ~/env.sh
+          export CUDA_LAUNCH_BLOCKING=1
+          python3 - <<'PY'
+          import importlib.util
+          import sys
+          from pathlib import Path
+
+          import torch
+
+          path = Path("python/tutorials/tle/04-cluster-gemm.py")
+          spec = importlib.util.spec_from_file_location("cluster_gemm04", path)
+          mod = importlib.util.module_from_spec(spec)
+          sys.modules[spec.name] = mod
+          spec.loader.exec_module(mod)
+
+          M = N = K = 4096
+          bm, bn, bk, num_warps, num_stages = 32, 512, 128, 4, 3
+          a = torch.empty((M, K), device="cuda", dtype=torch.float16)
+          b = torch.empty((K, N), device="cuda", dtype=torch.float16)
+          c = torch.empty((M, N), device="cuda", dtype=torch.float16)
+          dot_k = mod._select_remote_dot_k(bk)
+          grid = mod._grid_cluster_remote(M, N, bm, bn)
+
+          print("[debug] device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
+          print("[debug] candidate", (bm, bn, bk, num_warps, num_stages), "grid", grid)
+
+          def compile_slots(slots):
+              compiled = mod._cluster_remote_gemm_kernel.warmup(
+                  a, b, c, M, N, K,
+                  a.stride(0), a.stride(1),
+                  b.stride(0), b.stride(1),
+                  c.stride(0), c.stride(1),
+                  mesh=mod.BLOCK_CLUSTER_MESH,
+                  BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
+                  CLUSTER_SIZE=2,
+                  USE_MASK=False,
+                  A_SLOTS=slots,
+                  USE_NV_MMA_SMEM_LAYOUT=False,
+                  grid=grid,
+                  num_ctas=1,
+                  num_warps=num_warps,
+                  num_stages=num_stages,
+              )
+              print(f"[debug] slots={slots} shared={getattr(compiled.metadata, 'shared', None)}")
+              print(f"===== TTGIR slots={slots} BEGIN =====")
+              print(compiled.asm.get("ttgir", ""))
+              print(f"===== TTGIR slots={slots} END =====")
+              return compiled
+
+          compile_slots(2)
+          compile_slots(4)
+
+          print("[debug] launching slots=2 candidate")
+          mod._cluster_remote_gemm_kernel[grid](
+              a, b, c, M, N, K,
+              a.stride(0), a.stride(1),
+              b.stride(0), b.stride(1),
+              c.stride(0), c.stride(1),
+              mesh=mod.BLOCK_CLUSTER_MESH,
+              BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
+              CLUSTER_SIZE=2,
+              USE_MASK=False,
+              A_SLOTS=2,
+              USE_NV_MMA_SMEM_LAYOUT=False,
+              num_ctas=1,
+              num_warps=num_warps,
+              num_stages=num_stages,
+          )
+          torch.cuda.synchronize()
+          print("[debug] slots=2 launch passed")
+          PY
+          exit 0
           # python tutorials
           python3 python/tutorials/01-vector-add.py --only_unit_test
           python3 python/tutorials/02-fused-softmax.py --only_unit_test

From ed0acaf28c8605673dddaffae7b937a72ca8de47 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Mon, 8 Jun 2026 02:56:06 +0000
Subject: [PATCH 37/41] Restore Hopper workflow after TTGIR debug

---
 .github/workflows/hopper-build-and-test.yml | 71 ---------------------
 1 file changed, 71 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index c5871f4568..4084843f55 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,77 +155,6 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          export CUDA_LAUNCH_BLOCKING=1
-          python3 - <<'PY'
-          import importlib.util
-          import sys
-          from pathlib import Path
-
-          import torch
-
-          path = Path("python/tutorials/tle/04-cluster-gemm.py")
-          spec = importlib.util.spec_from_file_location("cluster_gemm04", path)
-          mod = importlib.util.module_from_spec(spec)
-          sys.modules[spec.name] = mod
-          spec.loader.exec_module(mod)
-
-          M = N = K = 4096
-          bm, bn, bk, num_warps, num_stages = 32, 512, 128, 4, 3
-          a = torch.empty((M, K), device="cuda", dtype=torch.float16)
-          b = torch.empty((K, N), device="cuda", dtype=torch.float16)
-          c = torch.empty((M, N), device="cuda", dtype=torch.float16)
-          dot_k = mod._select_remote_dot_k(bk)
-          grid = mod._grid_cluster_remote(M, N, bm, bn)
-
-          print("[debug] device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
-          print("[debug] candidate", (bm, bn, bk, num_warps, num_stages), "grid", grid)
-
-          def compile_slots(slots):
-              compiled = mod._cluster_remote_gemm_kernel.warmup(
-                  a, b, c, M, N, K,
-                  a.stride(0), a.stride(1),
-                  b.stride(0), b.stride(1),
-                  c.stride(0), c.stride(1),
-                  mesh=mod.BLOCK_CLUSTER_MESH,
-                  BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
-                  CLUSTER_SIZE=2,
-                  USE_MASK=False,
-                  A_SLOTS=slots,
-                  USE_NV_MMA_SMEM_LAYOUT=False,
-                  grid=grid,
-                  num_ctas=1,
-                  num_warps=num_warps,
-                  num_stages=num_stages,
-              )
-              print(f"[debug] slots={slots} shared={getattr(compiled.metadata, 'shared', None)}")
-              print(f"===== TTGIR slots={slots} BEGIN =====")
-              print(compiled.asm.get("ttgir", ""))
-              print(f"===== TTGIR slots={slots} END =====")
-              return compiled
-
-          compile_slots(2)
-          compile_slots(4)
-
-          print("[debug] launching slots=2 candidate")
-          mod._cluster_remote_gemm_kernel[grid](
-              a, b, c, M, N, K,
-              a.stride(0), a.stride(1),
-              b.stride(0), b.stride(1),
-              c.stride(0), c.stride(1),
-              mesh=mod.BLOCK_CLUSTER_MESH,
-              BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
-              CLUSTER_SIZE=2,
-              USE_MASK=False,
-              A_SLOTS=2,
-              USE_NV_MMA_SMEM_LAYOUT=False,
-              num_ctas=1,
-              num_warps=num_warps,
-              num_stages=num_stages,
-          )
-          torch.cuda.synchronize()
-          print("[debug] slots=2 launch passed")
-          PY
-          exit 0
           # python tutorials
           python3 python/tutorials/01-vector-add.py --only_unit_test
           python3 python/tutorials/02-fused-softmax.py --only_unit_test

From ba677addf414ceed0883114e5bc28375147fd792 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Mon, 8 Jun 2026 03:00:21 +0000
Subject: [PATCH 38/41] Revert cluster GEMM slot workaround

---
 python/tutorials/tle/04-cluster-gemm.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tutorials/tle/04-cluster-gemm.py b/python/tutorials/tle/04-cluster-gemm.py
index 0388105f1d..2c5c3e9b58 100644
--- a/python/tutorials/tle/04-cluster-gemm.py
+++ b/python/tutorials/tle/04-cluster-gemm.py
@@ -290,7 +290,7 @@ def _run_cluster_remote(
     N = b.shape[1]
     dot_k = _select_remote_dot_k(bk)
     use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
-    a_slots = 4 if num_stages > 2 else 2
+    a_slots = 2
     # BK=64 with 2-stage pipeline benefits from MMA-friendly shared layout in
     # the remote A path; BK=32 keeps the existing fast path.
     use_nv_mma_smem_layout = (bk == 32) or (bk == 64 and num_stages <= 2)
@@ -337,7 +337,7 @@ def _verify_remote_lowering(
     dot_k = _select_remote_dot_k(bk)
     use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
     # Keep verifier config aligned with runtime launch config.
-    a_slots = 4 if num_stages > 2 else 2
+    a_slots = 2
     use_nv_mma_smem_layout = (bk == 32) or (bk == 64 and num_stages <= 2)
     compiled = _cluster_remote_gemm_kernel.warmup(
         a,
@@ -631,7 +631,7 @@ def main(argv: list[str] | None = None) -> None:
             DOT_K=_select_remote_dot_k(remote_cfg.bk),
             CLUSTER_SIZE=2,
             USE_MASK=((args.m % remote_cfg.bm != 0) or (args.n % remote_cfg.bn != 0) or (args.k % remote_cfg.bk != 0)),
-            A_SLOTS=4 if remote_cfg.num_stages > 2 else 2,
+            A_SLOTS=2,
             USE_NV_MMA_SMEM_LAYOUT=((remote_cfg.bk == 32) or (remote_cfg.bk == 64 and remote_cfg.num_stages <= 2)),
             grid=_grid_cluster_remote(args.m, args.n, remote_cfg.bm, remote_cfg.bn),
             num_ctas=1,

From 096f19d618e8c86f8d6d09c2a20d150ad221b35b Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Mon, 8 Jun 2026 03:02:26 +0000
Subject: [PATCH 39/41] Debug Hopper LLVM PTX for remote slots

---
 .github/workflows/hopper-build-and-test.yml | 81 +++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 4084843f55..4fd09c9147 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -155,6 +155,87 @@ jobs:
         run: |
           set -x
           source ~/env.sh
+          export CUDA_LAUNCH_BLOCKING=1
+          python3 -u - <<'PY'
+          import importlib.util
+          import sys
+          from pathlib import Path
+
+          import torch
+
+          path = Path("python/tutorials/tle/04-cluster-gemm.py")
+          spec = importlib.util.spec_from_file_location("cluster_gemm04", path)
+          mod = importlib.util.module_from_spec(spec)
+          sys.modules[spec.name] = mod
+          spec.loader.exec_module(mod)
+
+          M = N = K = 4096
+          bm, bn, bk, num_warps, num_stages = 32, 512, 128, 4, 3
+          a = torch.empty((M, K), device="cuda", dtype=torch.float16)
+          b = torch.empty((K, N), device="cuda", dtype=torch.float16)
+          c = torch.empty((M, N), device="cuda", dtype=torch.float16)
+          dot_k = mod._select_remote_dot_k(bk)
+          grid = mod._grid_cluster_remote(M, N, bm, bn)
+
+          print("[debug] device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
+          print("[debug] candidate", (bm, bn, bk, num_warps, num_stages), "grid", grid)
+          print("[debug] ptr_mod256", int(a.data_ptr() % 256), int(b.data_ptr() % 256), int(c.data_ptr() % 256))
+
+          def compile_slots(slots):
+              compiled = mod._cluster_remote_gemm_kernel.warmup(
+                  a, b, c, M, N, K,
+                  a.stride(0), a.stride(1),
+                  b.stride(0), b.stride(1),
+                  c.stride(0), c.stride(1),
+                  mesh=mod.BLOCK_CLUSTER_MESH,
+                  BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
+                  CLUSTER_SIZE=2,
+                  USE_MASK=False,
+                  A_SLOTS=slots,
+                  USE_NV_MMA_SMEM_LAYOUT=False,
+                  grid=grid,
+                  num_ctas=1,
+                  num_warps=num_warps,
+                  num_stages=num_stages,
+              )
+              llir = compiled.asm.get("llir", "")
+              ptx = compiled.asm.get("ptx", "")
+              print(f"[debug] slots={slots} shared={getattr(compiled.metadata, 'shared', None)} llir_bytes={len(llir)} ptx_bytes={len(ptx)}")
+              print(f"===== LLIR slots={slots} BEGIN =====")
+              print(llir)
+              print(f"===== LLIR slots={slots} END =====")
+              print(f"===== PTX slots={slots} BEGIN =====")
+              print(ptx)
+              print(f"===== PTX slots={slots} END =====")
+              return compiled
+
+          compile_slots(2)
+          compile_slots(4)
+
+          def launch(slots):
+              print(f"[debug] launching slots={slots} candidate")
+              mod._cluster_remote_gemm_kernel[grid](
+                  a, b, c, M, N, K,
+                  a.stride(0), a.stride(1),
+                  b.stride(0), b.stride(1),
+                  c.stride(0), c.stride(1),
+                  mesh=mod.BLOCK_CLUSTER_MESH,
+                  BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
+                  CLUSTER_SIZE=2,
+                  USE_MASK=False,
+                  A_SLOTS=slots,
+                  USE_NV_MMA_SMEM_LAYOUT=False,
+                  num_ctas=1,
+                  num_warps=num_warps,
+                  num_stages=num_stages,
+              )
+              torch.cuda.synchronize()
+              print(f"[debug] slots={slots} launch passed")
+
+          launch(4)
+          launch(2)
+          PY
+          exit 0
           # python tutorials
           python3 python/tutorials/01-vector-add.py --only_unit_test
           python3 python/tutorials/02-fused-softmax.py --only_unit_test

From 69e74e24ee6dbd35e3d6166bfa83bc20f90ef1f0 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Mon, 8 Jun 2026 03:35:43 +0000
Subject: [PATCH 40/41] Guard TLE shared store vectorization by alignment

---
 .../LoadStoreOpToLLVM.cpp                     | 38 +++++++++++--------
 .../test_tle_local_ptr_misaligned_store.mlir  | 26 +++++++++++++
 2 files changed, 49 insertions(+), 15 deletions(-)
 create mode 100644 third_party/tle/test/GPU/test_tle_local_ptr_misaligned_store.mlir

diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
index 1efeffad7d..af76eda86b 100644
--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
@@ -204,6 +204,14 @@ struct LoadStoreConversionBase {
   }
 
 #ifdef __TLE__
+  bool isTleSharedTensorPtr(Value ptr) const {
+    auto ptrTensorTy = dyn_cast<RankedTensorType>(ptr.getType());
+    auto ptrElemTy = ptrTensorTy
+                         ? dyn_cast<PointerType>(ptrTensorTy.getElementType())
+                         : PointerType();
+    return ptrElemTy && isSharedFamilyAddressSpace(ptrElemTy.getAddressSpace());
+  }
+
   unsigned getMaxVectorSizeByAlignment(Value ptr) const {
     auto tensorTy = dyn_cast<RankedTensorType>(ptr.getType());
     if (!tensorTy)
@@ -226,6 +234,18 @@ struct LoadStoreConversionBase {
 
     return std::min<unsigned>(128 / pointeeBitWidth, maxMultiple);
   }
+
+  unsigned getTleSharedPointerVectorSize(Value ptr, unsigned vec) const {
+    if (!isTleSharedTensorPtr(ptr))
+      return vec;
+
+    // AxisInfo contiguity and TLE layout hints may both propose vectorized
+    // shared-memory accesses. They are legal only up to the width whose first
+    // element alignment is proven by AxisInfo divisibility.
+    unsigned hint = tte::inferTlePointerLayoutVectorHint(ptr);
+    unsigned alignmentBound = getMaxVectorSizeByAlignment(ptr);
+    return std::min(std::max(vec, hint), alignmentBound);
+  }
 #endif
 
   unsigned getMaskAlignment(Value mask) const {
@@ -274,21 +294,8 @@ struct LoadOpConversion : public ConvertOpToLLVMPattern<triton::LoadOp>,
         typeConverter->convertType(getElementTypeOrSelf(op.getType()));
     unsigned vec = getVectorSize(ptr);
 #ifdef __TLE__
-    auto ptrTensorTy = dyn_cast<RankedTensorType>(ptr.getType());
-    auto ptrElemTy = ptrTensorTy
-                         ? dyn_cast<PointerType>(ptrTensorTy.getElementType())
-                         : PointerType();
-    bool isSharedTensorPtr =
-        ptrElemTy && isSharedFamilyAddressSpace(ptrElemTy.getAddressSpace());
-    if (!llMask && isSharedTensorPtr) {
-      // For TLE local/shared pointer chains, AxisInfo contiguity can be
-      // conservative on packed contiguous lanes. The layout hint may recover
-      // the lane grouping, but it is only legal up to the vector width whose
-      // first element alignment is proven by AxisInfo divisibility.
-      unsigned hint = tte::inferTlePointerLayoutVectorHint(ptr);
-      unsigned alignmentBound = getMaxVectorSizeByAlignment(ptr);
-      vec = std::max(vec, std::min(hint, alignmentBound));
-    }
+    if (!llMask)
+      vec = getTleSharedPointerVectorSize(ptr, vec);
 #endif
     unsigned numElems = getTotalElemsPerThread(ptr.getType());
     unsigned vecOrig = vec;
@@ -579,6 +586,7 @@ struct StoreOpConversion : public ConvertOpToLLVMPattern<triton::StoreOp>,
     const bool isClusterSharedPtr =
         inferPtrAddrSpace(ptrElems).value_or(1) ==
         static_cast<unsigned>(NVVM::NVVMMemorySpace::SharedCluster);
+    vec = getTleSharedPointerVectorSize(ptr, vec);
 #else
     auto ptrElems = unpackLLElements(loc, llPtr, rewriter);
     auto valueElems = unpackLLElements(loc, llValue, rewriter);
diff --git a/third_party/tle/test/GPU/test_tle_local_ptr_misaligned_store.mlir b/third_party/tle/test/GPU/test_tle_local_ptr_misaligned_store.mlir
new file mode 100644
index 0000000000..ff5383f8f6
--- /dev/null
+++ b/third_party/tle/test/GPU/test_tle_local_ptr_misaligned_store.mlir
@@ -0,0 +1,26 @@
+// RUN: triton-opt %s -pass-pipeline='builtin.module(allocate-shared-memory-nv{compute-capability=120 ptx-version=88}, tritongpu-global-scratch-memory-allocation, convert-triton-gpu-to-llvm{compute-capability=120 ptx-version=88}, canonicalize, cse, convert-nv-gpu-to-llvm, convert-warp-specialize-to-llvm, canonicalize, cse, symbol-dce, convert-nvvm-to-llvm)' | FileCheck %s
+
+#blocked = #ttg.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [16], order = [0]}>
+#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
+#smem = #ttg.shared_memory
+
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 16 : i32, ttg.target = "cuda:120", "ttg.threads-per-warp" = 32 : i32} {
+  // Layout-derived vector hints are legal only when AxisInfo divisibility proves
+  // that the first element is aligned for the requested vector width. This
+  // tensor-indexed local_ptr starts at offset 1, so it is only 4-byte aligned
+  // and must not lower to st.shared.v4.b32.
+  tt.func public @misaligned_tensor_indexed_local_ptr_store() {
+    %one = arith.constant dense<1> : tensor<128xi32, #blocked>
+    %offs0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #blocked>
+    %offs = arith.addi %offs0, %one : tensor<128xi32, #blocked>
+    %vals = arith.constant dense<0.000000e+00> : tensor<128xf32, #blocked>
+    %smem = ttg.local_alloc {tle.barrier_group = 0 : i64} : () -> !ttg.memdesc<256xf32, #shared, #smem, mutable>
+    %ptrs = "tle.local_pointers"(%smem, %offs) {tle.barrier_group = 0 : i64} : (!ttg.memdesc<256xf32, #shared, #smem, mutable>, tensor<128xi32, #blocked>) -> tensor<128x!tt.ptr<f32, 3>, #blocked>
+    tt.store %ptrs, %vals : tensor<128x!tt.ptr<f32, 3>, #blocked>
+    tt.return
+  }
+}
+
+// CHECK-LABEL: @misaligned_tensor_indexed_local_ptr_store
+// CHECK-NOT: st.shared.v4.b32
+// CHECK: st.shared.b32

From c84241a53e74a895ebe61ce02dbadb6d0834e833 Mon Sep 17 00:00:00 2001
From: sunnycase <sunnycase@live.cn>
Date: Mon, 8 Jun 2026 03:46:21 +0000
Subject: [PATCH 41/41] Restore Hopper workflow after alignment debug

---
 .github/workflows/hopper-build-and-test.yml | 85 +--------------------
 1 file changed, 2 insertions(+), 83 deletions(-)

diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml
index 4fd09c9147..f6cb51dafb 100644
--- a/.github/workflows/hopper-build-and-test.yml
+++ b/.github/workflows/hopper-build-and-test.yml
@@ -122,7 +122,7 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py --mode test
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
@@ -155,87 +155,6 @@ jobs:
         run: |
           set -x
           source ~/env.sh
-          export CUDA_LAUNCH_BLOCKING=1
-          python3 -u - <<'PY'
-          import importlib.util
-          import sys
-          from pathlib import Path
-
-          import torch
-
-          path = Path("python/tutorials/tle/04-cluster-gemm.py")
-          spec = importlib.util.spec_from_file_location("cluster_gemm04", path)
-          mod = importlib.util.module_from_spec(spec)
-          sys.modules[spec.name] = mod
-          spec.loader.exec_module(mod)
-
-          M = N = K = 4096
-          bm, bn, bk, num_warps, num_stages = 32, 512, 128, 4, 3
-          a = torch.empty((M, K), device="cuda", dtype=torch.float16)
-          b = torch.empty((K, N), device="cuda", dtype=torch.float16)
-          c = torch.empty((M, N), device="cuda", dtype=torch.float16)
-          dot_k = mod._select_remote_dot_k(bk)
-          grid = mod._grid_cluster_remote(M, N, bm, bn)
-
-          print("[debug] device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
-          print("[debug] candidate", (bm, bn, bk, num_warps, num_stages), "grid", grid)
-          print("[debug] ptr_mod256", int(a.data_ptr() % 256), int(b.data_ptr() % 256), int(c.data_ptr() % 256))
-
-          def compile_slots(slots):
-              compiled = mod._cluster_remote_gemm_kernel.warmup(
-                  a, b, c, M, N, K,
-                  a.stride(0), a.stride(1),
-                  b.stride(0), b.stride(1),
-                  c.stride(0), c.stride(1),
-                  mesh=mod.BLOCK_CLUSTER_MESH,
-                  BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
-                  CLUSTER_SIZE=2,
-                  USE_MASK=False,
-                  A_SLOTS=slots,
-                  USE_NV_MMA_SMEM_LAYOUT=False,
-                  grid=grid,
-                  num_ctas=1,
-                  num_warps=num_warps,
-                  num_stages=num_stages,
-              )
-              llir = compiled.asm.get("llir", "")
-              ptx = compiled.asm.get("ptx", "")
-              print(f"[debug] slots={slots} shared={getattr(compiled.metadata, 'shared', None)} llir_bytes={len(llir)} ptx_bytes={len(ptx)}")
-              print(f"===== LLIR slots={slots} BEGIN =====")
-              print(llir)
-              print(f"===== LLIR slots={slots} END =====")
-              print(f"===== PTX slots={slots} BEGIN =====")
-              print(ptx)
-              print(f"===== PTX slots={slots} END =====")
-              return compiled
-
-          compile_slots(2)
-          compile_slots(4)
-
-          def launch(slots):
-              print(f"[debug] launching slots={slots} candidate")
-              mod._cluster_remote_gemm_kernel[grid](
-                  a, b, c, M, N, K,
-                  a.stride(0), a.stride(1),
-                  b.stride(0), b.stride(1),
-                  c.stride(0), c.stride(1),
-                  mesh=mod.BLOCK_CLUSTER_MESH,
-                  BM=bm, BN=bn, BK=bk, DOT_K=dot_k,
-                  CLUSTER_SIZE=2,
-                  USE_MASK=False,
-                  A_SLOTS=slots,
-                  USE_NV_MMA_SMEM_LAYOUT=False,
-                  num_ctas=1,
-                  num_warps=num_warps,
-                  num_stages=num_stages,
-              )
-              torch.cuda.synchronize()
-              print(f"[debug] slots={slots} launch passed")
-
-          launch(4)
-          launch(2)
-          PY
-          exit 0
           # python tutorials
           python3 python/tutorials/01-vector-add.py --only_unit_test
           python3 python/tutorials/02-fused-softmax.py --only_unit_test
@@ -261,7 +180,7 @@ jobs:
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
-          python3 python/tutorials/tle/04-cluster-gemm.py
+          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
           # python unit test