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 diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index bf509ffe52..f7306536eb 100644 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -882,12 +882,13 @@ class TleExtractTileOpPattern : public OpConversionPattern { Type retType = op.getType().cloneWithEncoding(srcEnc); + auto stridesAttr = + mlir::dyn_cast_or_null(op->getAttr("strides")); auto newOp = rewriter.replaceOpWithNewOp( - 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(); @@ -925,8 +926,12 @@ class TleInsertTileOpPattern : public OpConversionPattern { Type retType = op.getType().cloneWithEncoding(srcEnc); + auto stridesAttr = + mlir::dyn_cast_or_null(op->getAttr("strides")); + auto newOp = rewriter.replaceOpWithNewOp( - 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..2fc28a5c59 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,33 @@ 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 +88,32 @@ 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/core.py b/python/triton/experimental/tle/language/core.py index 03302bc60f..33aaf5dee8 100644 --- a/python/triton/experimental/tle/language/core.py +++ b/python/triton/experimental/tle/language/core.py @@ -112,6 +112,11 @@ 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 +169,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 +180,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 +201,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 +220,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 +260,7 @@ def extract_tile( x: tl.tensor, index, tile_shape: tuple, + strides: tuple = None, _semantic=None, ) -> tl.tensor: """ @@ -257,13 +268,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 +290,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 +304,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 +329,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 +351,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 +367,23 @@ 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 +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) + 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,22 @@ 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 +493,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 +529,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 +545,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 06c531bd99..d0550efbe1 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]]) -> None: + def validate_extract_tile_params(self, src: tl.tensor, index, tile_shape: Sequence[Union[int, any]], + strides=None) -> None: """ """ @@ -109,9 +110,16 @@ 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): @@ -122,7 +130,9 @@ 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 +151,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})", @@ -162,7 +172,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: """ """ @@ -195,12 +205,21 @@ 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)): @@ -224,18 +243,20 @@ 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, 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/python/tutorials/tle/05-glu.py b/python/tutorials/tle/05-glu.py new file mode 100644 index 0000000000..2c0f631321 --- /dev/null +++ b/python/tutorials/tle/05-glu.py @@ -0,0 +1,307 @@ +# %% +# 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)) + + +@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]) + 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) + 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) + +_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..bf29061f9d --- /dev/null +++ b/python/tutorials/tle/06-2D_Depthwise_Conv.py @@ -0,0 +1,334 @@ +""" +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..f5adfa7d7b --- /dev/null +++ b/python/tutorials/tle/07-causal-conv1d.py @@ -0,0 +1,1582 @@ +""" +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 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: + 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 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: + 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: + 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) + # 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) + + 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..9fd6228ec1 --- /dev/null +++ b/python/tutorials/tle/08-rope.py @@ -0,0 +1,751 @@ +""" +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 random + +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() 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..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 @@ -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 delinearizeTileIndex(OpBuilder &b, Location loc, Value linearIdx, ArrayRef srcShape, - ArrayRef tileShape) { + ArrayRef tileShape, + ArrayRef strides) { unsigned rank = srcShape.size(); SmallVector 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 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 computeElemOffsets(OpBuilder &b, Location loc, Value linearIdx, ArrayRef srcShape, - ArrayRef tileShape) { - auto coords = delinearizeTileIndex(b, loc, linearIdx, srcShape, tileShape); + ArrayRef tileShape, + ArrayRef strides) { + auto coords = + delinearizeTileIndex(b, loc, linearIdx, srcShape, tileShape, strides); unsigned rank = srcShape.size(); SmallVector offsets(rank); for (unsigned d = 0; d < rank; ++d) { auto cst = b.create( - loc, static_cast(tileShape[d]), 32); + loc, static_cast(strides[d]), 32); offsets[d] = b.create(loc, coords[d], cst); } return offsets; @@ -1128,13 +1131,13 @@ static SmallVector computeCombinedOffsets(OpBuilder &b, Location loc, static Value emitSliceFromSmem(OpBuilder &rewriter, Location loc, Value smemBuf, Value convertedIndex, ArrayRef srcShape, - ArrayRef tileShape, RankedTensorType resultTensorTy, - MemRefType outputType, + ArrayRef tileShape, ArrayRef strides, + RankedTensorType resultTensorTy, MemRefType outputType, triton::gcu::FirstLastUserAnalysis &userAnalysis, std::map &replaced2Origin, triton::gcu::PrivateTagPool &pTagPool, Operation *op) { - auto elemOffsets = - computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape); + 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); @@ -1160,13 +1163,13 @@ static Value emitDesliceToSmem(OpBuilder &rewriter, Location loc, Value smemBuf, Value convertedTile, Value convertedIndex, ArrayRef srcShape, ArrayRef tileShape, - RankedTensorType tileTensorTy, + ArrayRef strides, RankedTensorType tileTensorTy, RankedTensorType resultTensorTy, triton::gcu::FirstLastUserAnalysis &userAnalysis, std::map &replaced2Origin, triton::gcu::PrivateTagPool &pTagPool, Operation *op) { - auto elemOffsets = - computeElemOffsets(rewriter, loc, convertedIndex, srcShape, tileShape); + auto elemOffsets = computeElemOffsets(rewriter, loc, convertedIndex, srcShape, + tileShape, strides); auto firstUser = userAnalysis.getFirstUser(op->getResults()[0]); auto lastUser = userAnalysis.getLastUser(op->getResults()[0]); @@ -1235,6 +1238,17 @@ struct TleExtractTileLowering : SharedGenericConversionPattern { return failure(); auto tileShape = tileShapeAttr.asArrayRef(); + // Extract strides (defaults to tileShape for backward compatibility). + auto stridesAttr = op->getAttrOfType("strides"); + SmallVector stridesVec; + if (stridesAttr) { + stridesVec.assign(stridesAttr.asArrayRef().begin(), + stridesAttr.asArrayRef().end()); + } else { + stridesVec.assign(tileShape.begin(), tileShape.end()); + } + ArrayRef strides = stridesVec; + auto srcTensorTy = cast(op->getOperand(0).getType()); auto srcShape = srcTensorTy.getShape(); auto outputType = @@ -1245,7 +1259,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 +1296,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 +1352,17 @@ struct TleInsertTileLowering : SharedGenericConversionPattern { } ArrayRef tileShape = tileShapeVec; + // Extract strides (defaults to tileShape for backward compatibility). + auto stridesAttr = op->getAttrOfType("strides"); + SmallVector stridesVec; + if (stridesAttr) { + stridesVec.assign(stridesAttr.asArrayRef().begin(), + stridesAttr.asArrayRef().end()); + } else { + stridesVec.assign(tileShapeVec.begin(), tileShapeVec.end()); + } + ArrayRef strides = stridesVec; + auto srcTensorTy = cast(op->getOperand(0).getType()); auto srcShape = srcTensorTy.getShape(); @@ -1348,7 +1373,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 zeroOffsets( @@ -1395,10 +1420,10 @@ struct TleInsertTileLowering : SharedGenericConversionPattern { replaced2Origin); } - auto output = - emitDesliceToSmem(rewriter, loc, smemBuf, convertedTile, convertedIndex, - srcShape, tileShape, 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); @@ -1427,10 +1452,17 @@ struct SliceFromLocalOpLowering auto outputType = dyn_cast(getTypeConverter()->convertType(resultTensorTy)); - auto output = emitSliceFromSmem( - rewriter, op.getLoc(), adaptor.getSrc(), adaptor.getIndex(), - op.getSrc().getType().getShape(), op.getTileShape(), resultTensorTy, - outputType, userAnalysis, replaced2Origin, pTagPool, op); + auto tileShapeArr = op.getTileShape(); + auto stridesOpt = op.getTileStrides(); + SmallVector stridesVec( + stridesOpt.has_value() ? stridesOpt->begin() : tileShapeArr.begin(), + stridesOpt.has_value() ? stridesOpt->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); leaveTritionOp(rewriter, op.getOperation()); rewriter.replaceOp(op, output); @@ -1458,11 +1490,17 @@ struct DesliceToLocalOpLowering auto tileTensorTy = cast(op.getTile().getType()); auto resultTensorTy = cast(op.getResult().getType()); + auto tileShapeArr = op.getTileShape(); + auto stridesOpt = op.getTileStrides(); + SmallVector stridesVec( + stridesOpt.has_value() ? stridesOpt->begin() : tileShapeArr.begin(), + stridesOpt.has_value() ? stridesOpt->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/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("tile_strides"); + auto stridesAttr = op->getAttrOfType("strides"); auto resultTy = op->getResult(0).getType(); auto newOp = rewriter.create( - 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("tile_strides"); + auto stridesAttr = op->getAttrOfType("strides"); auto resultTy = op->getResult(0).getType(); auto newOp = rewriter.create( - 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( - op, retType, adaptor.getSrc(), adaptor.getIndex()); + op, retType, adaptor.getSrc(), adaptor.getIndex(), + op->getAttrOfType("strides")); addNamedAttrs(newOp, adaptor.getAttributes()); return success(); } @@ -816,7 +817,8 @@ class TleInsertTileOpPattern Type retType = op.getType().cloneWithEncoding(srcEnc); auto newOp = rewriter.replaceOpWithNewOp( - op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex()); + op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(), + op->getAttrOfType("strides")); addNamedAttrs(newOp, adaptor.getAttributes()); return success(); } 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 { Type retType = op.getType().cloneWithEncoding(srcEnc); auto newOp = rewriter.replaceOpWithNewOp( - op, retType, adaptor.getSrc(), adaptor.getIndex()); + op, retType, adaptor.getSrc(), adaptor.getIndex(), + op->getAttrOfType("strides")); if (auto tileShapeAttr = op->getAttr("tile_shape")) newOp->setAttr("tile_shape", tileShapeAttr); @@ -890,7 +891,8 @@ class TleInsertTileOpPattern : public OpConversionPattern { Type retType = op.getType().cloneWithEncoding(srcEnc); auto newOp = rewriter.replaceOpWithNewOp( - op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex()); + op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(), + op->getAttrOfType("strides")); addNamedAttrs(newOp, adaptor.getAttributes()); 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(ptr.getType()); + auto ptrElemTy = ptrTensorTy + ? dyn_cast(ptrTensorTy.getElementType()) + : PointerType(); + return ptrElemTy && isSharedFamilyAddressSpace(ptrElemTy.getAddressSpace()); + } + unsigned getMaxVectorSizeByAlignment(Value ptr) const { auto tensorTy = dyn_cast(ptr.getType()); if (!tensorTy) @@ -226,6 +234,18 @@ struct LoadStoreConversionBase { return std::min(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, typeConverter->convertType(getElementTypeOrSelf(op.getType())); unsigned vec = getVectorSize(ptr); #ifdef __TLE__ - auto ptrTensorTy = dyn_cast(ptr.getType()); - auto ptrElemTy = ptrTensorTy - ? dyn_cast(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, const bool isClusterSharedPtr = inferPtrAddrSpace(ptrElems).value_or(1) == static_cast(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/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index 777e47b5da..cea6cc9422 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:$strides); let results = (outs TT_Tensor:$result); let builders = [OpBuilder<(ins "Value":$src, "Value":$index, - "ArrayRef":$tileShape)>]; + "ArrayRef":$tileShape, + CArg<"ArrayRef", "{}">:$strides)>]; + let assemblyFormat = [{ $src `[` $index `]` attr-dict `:` qualified(type($src)) `,` qualified(type($index)) `->` qualified(type($result)) }]; @@ -47,8 +48,12 @@ def Tle_InsertTileOp : Tle_Op<"insert_tile", [Pure, DeclareOpInterfaceMethods, ]> { - 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:$strides); let results = (outs TT_Tensor:$result); + let builders = [OpBuilder<(ins "Value":$src, "Value":$tile, + "Value":$index, + CArg<"ArrayRef", "{}">:$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 f48a81d35e..df12e2f3d8 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -25,13 +25,16 @@ constexpr int kClusterSharedMemoryAddressSpace = 7; // ExtractTileOp Builder // ============================================================================ void ExtractTileOp::build(OpBuilder &builder, OperationState &state, Value src, - Value index, ArrayRef tileShape) { + Value index, ArrayRef tileShape, + ArrayRef strides) { auto srcType = cast(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 effectiveStrides(strides.empty() ? tileShape : strides); + state.addAttribute("strides", builder.getDenseI64ArrayAttr(effectiveStrides)); state.addTypes(resultType); } @@ -59,6 +62,13 @@ LogicalResult ExtractTileOp::verify() { for (auto v : denseArray64.asArrayRef()) tileShape.push_back(v); } + SmallVector strides; + if (auto a = mlir::dyn_cast_or_null( + 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 +89,13 @@ 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,7 +120,7 @@ LogicalResult ExtractTileOp::verify() { SmallVector 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]; } @@ -128,8 +140,7 @@ LogicalResult ExtractTileOp::verify() { // tile indices -> coordinate-level offsets SmallVector 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(srcTy.getRank())) return emitError("offsets size must match tensor rank"); @@ -348,6 +359,22 @@ void PipeReaderReleaseOp::getEffects( // ============================================================================ // InsertTileOp Type Inference + Verification // ============================================================================ +void InsertTileOp::build(OpBuilder &builder, OperationState &state, Value src, + Value tile, Value index, ArrayRef strides) { + auto srcType = cast(src.getType()); + auto tileType = cast(tile.getType()); + auto tileShape = tileType.getShape(); + + state.addOperands(src); + state.addOperands(tile); + state.addOperands(index); + SmallVector 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, ValueRange operands, @@ -395,6 +422,14 @@ LogicalResult InsertTileOp::verify() { auto tileShape = tileTy.getShape(); auto dstShape = dstTy.getShape(); + SmallVector strides; + if (auto a = mlir::dyn_cast_or_null( + 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 @@ -415,17 +450,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 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]; } @@ -463,7 +504,7 @@ LogicalResult InsertTileOp::verify() { // tile indices -> coordinate-level offsets SmallVector 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/ExtractTileToLLVM.cpp b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp index 9454ea913b..6bb0b14452 100644 --- a/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp +++ b/third_party/tle/dialect/lib/Transforms/ExtractTileToLLVM.cpp @@ -29,6 +29,18 @@ static SmallVector getTileShape(ExtractTileOp op) { return ts; } +static SmallVector getStrides(ExtractTileOp op) { + if (auto a = mlir::dyn_cast_or_null( + op->getAttr("strides"))) { + SmallVector s; + for (auto v : a.asArrayRef()) + s.push_back(v); + return s; + } + // 向后兼容：没有 strides 属性时退化为 tile_shape + return getTileShape(op); +} + static std::optional getStaticIndex(ExtractTileOp op) { if (auto c = op->getOperand(1).getDefiningOp()) return mlir::cast(c.getValue()).getInt(); @@ -39,18 +51,20 @@ static bool isCTATileAligned(ExtractTileOp op, int64_t linearIndex) { auto srcTy = cast(op.getSrc().getType()); auto srcShape = srcTy.getShape(); auto tileShape = getTileShape(op); + auto strides = getStrides(op); auto ctaTile = getShapePerCTATile(srcTy); int rank = srcShape.size(); SmallVector 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 +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 vals = unpackLLElements(loc, adaptor.getSrc(), rewriter); auto shapePerCTATile = getShapePerCTATile(srcTy); auto srcCTAShape = multiDimElementwise( @@ -82,14 +97,14 @@ lowerExtractTileStatic(ExtractTileOp op, ExtractTileOp::Adaptor adaptor, SmallVector 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( elementCoords, shapePerCTATile, std::divides()); auto srcCTAOrder = getCTATileOrder(srcTy), @@ -132,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); MLIRContext *ctx = rewriter.getContext(); auto i1Ty = rewriter.getIntegerType(1); auto i8Ty = rewriter.getIntegerType(8); @@ -187,7 +202,7 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor, // ------------------------------------------------------------------ SmallVector 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 +219,15 @@ lowerExtractTileViaSMEM(ExtractTileOp op, ExtractTileOp::Adaptor adaptor, loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)suffix[d])); Value gv = rewriter.create( loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)logicalGrid[d])); + Value sv_stride = rewriter.create( + loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d])); Value tv = rewriter.create( loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d])); + Value coord = rewriter.create(loc, i32Ty, dynIndex, sv); coord = rewriter.create(loc, i32Ty, coord, gv); - tileStartVals[d] = rewriter.create(loc, i32Ty, coord, tv); + tileStartVals[d] = + rewriter.create(loc, i32Ty, coord, sv_stride); tileEndVals[d] = rewriter.create(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 e0b5f08649..7ada8d945e 100644 --- a/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp +++ b/third_party/tle/dialect/lib/Transforms/InsertTileToLLVM.cpp @@ -26,18 +26,33 @@ static std::optional getStaticIndex(InsertTileOp op) { return mlir::cast(c.getValue()).getInt(); return std::nullopt; } + +static SmallVector getStrides(InsertTileOp op) { + if (auto a = mlir::dyn_cast_or_null( + op->getAttr("strides"))) { + SmallVector s; + for (auto v : a.asArrayRef()) + s.push_back(v); + return s; + } + auto tileTy = cast(op.getTile().getType()); + return SmallVector(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(op.getSrc().getType()); auto tileTy = cast(op.getTile().getType()); auto srcShape = srcTy.getShape(); auto tileShape = tileTy.getShape(); + auto strides = getStrides(op); auto ctaTile = getShapePerCTATile(srcTy); int rank = srcShape.size(); SmallVector 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 +61,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(ctaTile[i]) != 0) return false; if (off % static_cast(ctaTile[i]) != 0) @@ -71,6 +86,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 +99,12 @@ lowerInsertTileStatic(InsertTileOp op, InsertTileOp::Adaptor adaptor, SmallVector 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 logicalCoords(srcShape.size(), 0); @@ -97,7 +116,7 @@ lowerInsertTileStatic(InsertTileOp op, InsertTileOp::Adaptor adaptor, SmallVector 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( @@ -176,6 +195,7 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor, auto dstTy = cast(op.getType()); auto srcShape = srcTy.getShape(); auto tileShape = tileTy.getShape(); + auto strides = getStrides(op); int rank = srcShape.size(); MLIRContext *ctx = rewriter.getContext(); @@ -226,7 +246,7 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor, SmallVector 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 +264,14 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor, loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)suffix[d])); Value gv = rewriter.create( loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)logicalGrid[d])); + Value svStride = rewriter.create( + loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)strides[d])); Value tv = rewriter.create( loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d])); Value coord = rewriter.create(loc, i32Ty, dynIndex, sv); coord = rewriter.create(loc, i32Ty, coord, gv); - tileStartVals[d] = rewriter.create(loc, i32Ty, coord, tv); + tileStartVals[d] = + rewriter.create(loc, i32Ty, coord, svStride); tileEndVals[d] = rewriter.create(loc, i32Ty, tileStartVals[d], tv); } @@ -310,28 +333,44 @@ lowerInsertTileViaSMEMDynamic(InsertTileOp op, InsertTileOp::Adaptor adaptor, inRange = rewriter.create( loc, rewriter.create(loc, ge, lt), inRange); - Value tileShapeV = rewriter.create( - loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)tileShape[d])); - Value tileLocalSafeV = - rewriter.create(loc, i32Ty, globalCoordV, tileShapeV); + // Convert source global coord to tile-local coord. + Value tileLocalCoordV = rewriter.create( + loc, i32Ty, globalCoordV, tileStartVals[d]); Value sb = rewriter.create( loc, i32Ty, rewriter.getI32IntegerAttr((int32_t)(smemStrides[d] * elemBytes))); smemByteOffsetV = rewriter.create( loc, i32Ty, smemByteOffsetV, - rewriter.create(loc, i32Ty, tileLocalSafeV, sb)); + rewriter.create(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(loc, inRange, thenBlock, ValueRange{}, + elseBlock, ValueRange{}); + + rewriter.setInsertionPointToStart(thenBlock); Value lp = rewriter.create(loc, smemPtrTy, i8Ty, smemBase, ValueRange{smemByteOffsetV}, LLVM::GEPNoWrapFlags::inbounds); Value tileLoaded = rewriter.create(loc, llvmElemTy, lp, elemBytes); - // Overwrite source value with tile value if in range. - Value merged = - rewriter.create(loc, inRange, tileLoaded, srcVals[i]); - resultVals.push_back(merged); + rewriter.create(loc, ValueRange{tileLoaded}, mergeBlock); + + rewriter.setInsertionPointToStart(elseBlock); + rewriter.create(loc, ValueRange{srcVals[i]}, mergeBlock); + + rewriter.setInsertionPointToStart(mergeBlock); + resultVals.push_back(mergeBlock->getArgument(0)); } rewriter.create(loc); 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, #blocked> + tt.store %ptrs, %vals : tensor<128x!tt.ptr, #blocked> + tt.return + } +} + +// CHECK-LABEL: @misaligned_tensor_indexed_local_ptr_store +// CHECK-NOT: st.shared.v4.b32 +// CHECK: st.shared.b32 diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index c77b116ede..504c641802 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -79,22 +79,26 @@ void init_triton_tle_ir(py::module &&m) { // TLE-Lite .def( "create_extract_tile", - [](TritonOpBuilder &self, Value &input, - // std::vector &offsets, - Value &index, std::vector &tileShape) -> Value { - auto op = self.create(input, index, tileShape); + [](TritonOpBuilder &self, Value &input, Value &index, + std::vector &tileShape, + std::vector strides) -> Value { + auto op = self.create(input, index, tileShape, + strides); return op.getResult(); }, py::arg("input"), py::arg("index"), py::arg("tileShape"), + py::arg("strides") = std::vector{}, "Create extract_tile operation") .def( "create_insert_tile", - [](TritonOpBuilder &self, Value &input, Value &tile, - Value &index) -> Value { - auto op = self.create(input, tile, index); + [](TritonOpBuilder &self, Value &input, Value &tile, Value &index, + std::vector strides) -> Value { + auto op = + self.create(input, tile, index, strides); return op.getResult(); }, py::arg("input"), py::arg("tile"), py::arg("index"), + py::arg("strides") = std::vector{}, "Create insert_tile operation") // TLE-Struct .def("make_swizzled_shared_encoding_attr",