From 63976b9909283802fd2ead23c49a532ae78766cc Mon Sep 17 00:00:00 2001 From: Zhaohui Wang <173976389+GeoffreyWang1117@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:32:46 -0700 Subject: [PATCH] indexer: support non-power-of-two group sizes in the Triton kernels mm_bpr_kernel, reduce_rr_kernel and softmax_inplace_r_kernel walk the group axis with tl.arange(0, G) where G is taken straight from the tensor shape. Triton requires the length of a tl.arange to be a power of two, so any model whose GQA group size (num_attention_heads / num_key_value_heads) is not a power of two fails at kernel compile time: triton.compiler.errors.CompilationError: arange's range must be a power of 2 That rules out, among others, Qwen2.5-1.5B (12/2 = 6), Qwen2.5-7B (28/4 = 7) and Llama-3.2-3B (24/8 = 3) -- the indexer never gets to run at all on those. Fix: round the affected extent up to the next power of two for the tl.arange only, and mask the surplus lanes on both load and store. Pointer arithmetic and the Mean divisor keep using the real extent, so nothing changes for power-of-two shapes -- the mask is all-true and the generated code is equivalent. Padded lanes are filled with the identity element of the reduction (-inf for Max, +inf for Min, 0 for Mean / L2Norm / Sum) so they cannot influence the result, and they are never written back. Also fixes the fallback branch of the DIM == 2 reduce, which allocated its zero tile with the dim-1 extent instead of the dim-0 extent. Verified on an RTX 3090 against a torch reference: group sizes 1..8 for mm_bpr, all five ReduceTypes over both dims for six (D0, D1) shape pairs, and five shapes for softmax_inplace_r -- all within bf16 tolerance, with power-of-two shapes bit-for-bit unchanged from before. Repro script in the PR description. Not covered here, to keep the diff reviewable: mm_rrr / mm_rpr in the same file and the kernels under vortex_torch/cache/triton_kernels/ use the same tl.arange-on-raw-shape pattern. Happy to extend this PR or send a follow-up, whichever you prefer. --- .../indexer/triton_kernels/matmul_impl.py | 31 ++++++++++------- .../indexer/triton_kernels/reduce_impl.py | 33 ++++++++++++++----- .../indexer/triton_kernels/softmax_impl.py | 17 ++++++---- .../indexer/triton_kernels/utils_impl.py | 10 ++++++ 4 files changed, 64 insertions(+), 27 deletions(-) diff --git a/vortex_torch/indexer/triton_kernels/matmul_impl.py b/vortex_torch/indexer/triton_kernels/matmul_impl.py index 53f1ed80..de0a5b92 100644 --- a/vortex_torch/indexer/triton_kernels/matmul_impl.py +++ b/vortex_torch/indexer/triton_kernels/matmul_impl.py @@ -2,6 +2,7 @@ import triton import triton.language as tl from ..context import Context +from .utils_impl import next_pow2 @triton.jit def mm_bpr_kernel( @@ -14,7 +15,8 @@ def mm_bpr_kernel( winfo_y_lens, # int32 winfo_num_workloads, # int32* max_chunk_size: tl.constexpr, -G: tl.constexpr, +G: tl.constexpr, # real group size; may be any positive integer +G_PAD: tl.constexpr, # next_pow2(G); tl.arange needs a power-of-two length C: tl.constexpr, D: tl.constexpr, ): @@ -32,15 +34,16 @@ def mm_bpr_kernel( # Index vectors d_ptr = tl.arange(0, D) c_ptr = tl.arange(0, C) - g_ptr = tl.arange(0, G) + g_ptr = tl.arange(0, G_PAD) + g_mask = g_ptr < G idx_ptr = tl.arange(0, max_chunk_size) # Stride across B for x x_stride = G * D - # Persistent cache: the current x[x_idx] as a whole [G, D] tile + # Persistent cache: the current x[x_idx] as a whole [G_PAD, D] tile current_x_idx = tl.full((), -1, dtype=tl.int32) - x_i = tl.zeros((G, D), dtype=tl.float32) + x_i = tl.zeros((G_PAD, D), dtype=tl.float32) for i in range(start, end): @@ -48,9 +51,11 @@ def mm_bpr_kernel( x_idx_i32 = tl.load(winfo_x_indices + i).to(tl.int32) if x_idx_i32 != current_x_idx: x_base = (x_idx_i32 * x_stride).to(tl.int32) - # Load x_i: [G, D] (f32) + # Load x_i: [G_PAD, D] (f32). Rows g >= G are outside the + # tensor, so they are masked out and zero-filled; zeros contribute + # nothing to the dot product below and are never stored. x_offs = x_base + (g_ptr[:, None] * D + d_ptr[None, :]).to(tl.int32) - x_i = tl.load(x + x_offs).to(tl.float32) # f32 + x_i = tl.load(x + x_offs, mask=g_mask[:, None], other=0.0).to(tl.float32) # f32 current_x_idx = x_idx_i32 # Range of y rows for this workload @@ -74,14 +79,14 @@ def mm_bpr_kernel( rows_total: tl.constexpr = max_chunk_size * C y_rc = tl.reshape(y_tile, (rows_total, D)) # [RC, D], f32 - # Use x without transpose: x_i is [G, D] (f32) + # Use x without transpose: x_i is [G_PAD, D] (f32) # Elementwise multiply in f32, then cast to fp32 and reduce over D: - # [RC, 1, D] * [1, G, D] -> [RC, G, D] (f32), then sum over D -> [RC, G] (fp32) + # [RC, 1, D] * [1, G_PAD, D] -> [RC, G_PAD, D] (f32), then sum over D -> [RC, G_PAD] (fp32) prod_ = y_rc[:, None, :] * x_i[None, :, :] # f32 mult acc = tl.sum(prod_, 2) # fp32 reduction on D # Reshape back to [rows, C, G] and store (fp32) - o_i = tl.reshape(acc, (max_chunk_size, C, G)) # [rows, C, G], fp32 + o_i = tl.reshape(acc, (max_chunk_size, C, G_PAD)) # [rows, C, G_PAD], fp32 o_i = o_i.to(tl.bfloat16) # Linear output offset: row*C*G + c*G + g, where row starts at y_off offs_o = ( @@ -90,7 +95,7 @@ def mm_bpr_kernel( g_ptr[None, None, :] ).to(tl.int32) - tl.store(o + offs_o, o_i, mask=valid[:, None, None]) + tl.store(o + offs_o, o_i, mask=valid[:, None, None] & g_mask[None, None, :]) @@ -109,7 +114,8 @@ def mm_bpr( ctx.winfo_kv_lens, ctx.winfo_num_workloads, ctx.max_chunk_size, - x.shape[-2], y.shape[-2], x.shape[-1], num_warps=32, num_stages=1 + x.shape[-2], next_pow2(x.shape[-2]), y.shape[-2], x.shape[-1], + num_warps=32, num_stages=1 ) @@ -134,7 +140,8 @@ def _mm_bpr( winfo_y_lens, winfo_num_workloads, max_chunk_size, - x.shape[-2], y.shape[-2], x.shape[-1], num_warps=32, num_stages=1 + x.shape[-2], next_pow2(x.shape[-2]), y.shape[-2], x.shape[-1], + num_warps=32, num_stages=1 ) diff --git a/vortex_torch/indexer/triton_kernels/reduce_impl.py b/vortex_torch/indexer/triton_kernels/reduce_impl.py index a0f2aa10..ecefbe86 100644 --- a/vortex_torch/indexer/triton_kernels/reduce_impl.py +++ b/vortex_torch/indexer/triton_kernels/reduce_impl.py @@ -4,6 +4,7 @@ from ..context import Context from typing import Literal from ...utils import ReduceType +from .utils_impl import next_pow2 @triton.jit def reduce_rr_kernel( @@ -13,8 +14,10 @@ def reduce_rr_kernel( winfo_lens, # int32 winfo_num_workloads, # int32* max_chunk_size: tl.constexpr, -x_D0: tl.constexpr, -x_D1: tl.constexpr, +x_D0: tl.constexpr, # real extent of dim 0; may be any positive integer +x_D1: tl.constexpr, # real extent of dim 1; may be any positive integer +x_D0_PAD: tl.constexpr, # next_pow2(x_D0); tl.arange needs a power-of-two length +x_D1_PAD: tl.constexpr, # next_pow2(x_D1) DIM: tl.constexpr, REDUCE_TYPE: tl.constexpr ): @@ -31,8 +34,10 @@ def reduce_rr_kernel( end = start + per + (pid < r) idx_ptr = tl.arange(0, max_chunk_size) - dim0 = tl.arange(0, x_D0) - dim1 = tl.arange(0, x_D1) + dim0 = tl.arange(0, x_D0_PAD) + dim1 = tl.arange(0, x_D1_PAD) + dim0_mask = dim0 < x_D0 + dim1_mask = dim1 < x_D1 for i in range(start, end): @@ -45,7 +50,13 @@ def reduce_rr_kernel( dim0[None,:,None] * x_D1 + \ dim1[None, None, :] - x_i = tl.load(x_i_ptr, mask=valid[:,None,None], other=0.0).to(tl.float32) + # Lanes past the real extent are outside the tensor: mask them and + # fill with the identity element of the reduction, so they cannot + # affect the result. Mean divides by the real extent, not the padded + # one, so 0.0 is the right filler there too. + pad_val = -1e30 if REDUCE_TYPE == 1 else (1e30 if REDUCE_TYPE == 2 else 0.0) + load_mask = valid[:, None, None] & dim0_mask[None, :, None] & dim1_mask[None, None, :] + x_i = tl.load(x_i_ptr, mask=load_mask, other=pad_val).to(tl.float32) if DIM == 1: if REDUCE_TYPE == 0: @@ -64,11 +75,11 @@ def reduce_rr_kernel( x_i_reduce = tl.sum(x_i, axis=1) else: - x_i_reduce = tl.zeros((max_chunk_size, x_D1), dtype=tl.bfloat16) + x_i_reduce = tl.zeros((max_chunk_size, x_D1_PAD), dtype=tl.bfloat16) x_i_reduce = x_i_reduce.to(tl.bfloat16) o_i_ptr = o + x_off * x_D1 + idx_ptr[:, None] * x_D1 + dim1[None,:] - tl.store(o_i_ptr, x_i_reduce, mask=valid[:, None]) + tl.store(o_i_ptr, x_i_reduce, mask=valid[:, None] & dim1_mask[None, :]) elif DIM == 2: @@ -88,11 +99,11 @@ def reduce_rr_kernel( x_i_reduce = tl.sum(x_i, axis=2) else: - x_i_reduce = tl.zeros((max_chunk_size, x_D1), dtype=tl.float32) + x_i_reduce = tl.zeros((max_chunk_size, x_D0_PAD), dtype=tl.float32) x_i_reduce = x_i_reduce.to(tl.bfloat16) o_i_ptr = o + x_off * x_D0 + idx_ptr[:, None] * x_D0 + dim0[None,:] - tl.store(o_i_ptr, x_i_reduce, mask=valid[:, None]) + tl.store(o_i_ptr, x_i_reduce, mask=valid[:, None] & dim0_mask[None, :]) @@ -113,6 +124,8 @@ def reduce_rr( ctx.max_chunk_size, x.shape[-2], x.shape[-1], + next_pow2(x.shape[-2]), + next_pow2(x.shape[-1]), dim, reduce_type.value, num_warps=4, @@ -140,6 +153,8 @@ def _reduce_rr( max_chunk_size, x.shape[-2], x.shape[-1], + next_pow2(x.shape[-2]), + next_pow2(x.shape[-1]), dim, reduce_type.value, num_warps=4, diff --git a/vortex_torch/indexer/triton_kernels/softmax_impl.py b/vortex_torch/indexer/triton_kernels/softmax_impl.py index 8868286a..2ae213af 100644 --- a/vortex_torch/indexer/triton_kernels/softmax_impl.py +++ b/vortex_torch/indexer/triton_kernels/softmax_impl.py @@ -2,6 +2,7 @@ import triton import triton.language as tl from ..context import Context +from .utils_impl import next_pow2 @triton.jit def softmax_inplace_r_kernel( @@ -12,7 +13,8 @@ def softmax_inplace_r_kernel( eos: tl.constexpr, topk_val: tl.constexpr, x_D0: tl.constexpr, -x_D1: tl.constexpr, +x_D1: tl.constexpr, # real group size; may be any positive integer +x_D1_PAD: tl.constexpr, # next_pow2(x_D1); tl.arange needs a power-of-two length BLOCK_P: tl.constexpr = 256, ): pid = tl.program_id(0) @@ -32,13 +34,14 @@ def softmax_inplace_r_kernel( base_ptr = x + (start + bos) * (x_D0 * x_D1) d0_idx = tl.arange(0, x_D0) - d1_idx = tl.arange(0, x_D1) + d1_idx = tl.arange(0, x_D1_PAD) + d1_mask = d1_idx < x_D1 p_idx = tl.arange(0, BLOCK_P) # --- One-pass accumulation of (m, s) --- neg_inf = -1e30 - m = tl.full((x_D0, x_D1), neg_inf, dtype=tl.float32) - s = tl.zeros((x_D0, x_D1), dtype=tl.float32) + m = tl.full((x_D0, x_D1_PAD), neg_inf, dtype=tl.float32) + s = tl.zeros((x_D0, x_D1_PAD), dtype=tl.float32) for p in range(0, num_pages_to_compute, BLOCK_P): kp = tl.minimum(BLOCK_P, num_pages_to_compute - p) @@ -50,7 +53,7 @@ def softmax_inplace_r_kernel( + d1_idx[None, None, :] ).to(tl.int32) - mask = p_mask[:, None, None] + mask = p_mask[:, None, None] & d1_mask[None, None, :] slab = tl.load(base_ptr + offs, mask=mask, other=neg_inf).to(tl.float32) slab = slab * scale mc = tl.max(slab, axis=0) @@ -71,7 +74,7 @@ def softmax_inplace_r_kernel( + d1_idx[None, None, :] ).to(tl.int32) - mask = p_mask[:, None, None] + mask = p_mask[:, None, None] & d1_mask[None, None, :] slab = tl.load(base_ptr + offs, mask=mask, other=neg_inf).to(tl.float32) slab = slab * scale slab = tl.exp(slab - m[None, :, :]) / s[None, :, :] @@ -100,6 +103,7 @@ def softmax_inplace_r( ctx.topk_val, x.shape[-2], x.shape[-1], + next_pow2(x.shape[-1]), num_warps=4, num_stages=1 ) @@ -126,6 +130,7 @@ def _softmax_inplace_r( topk_val, x.shape[-2], x.shape[-1], + next_pow2(x.shape[-1]), num_warps=4, num_stages=1 ) \ No newline at end of file diff --git a/vortex_torch/indexer/triton_kernels/utils_impl.py b/vortex_torch/indexer/triton_kernels/utils_impl.py index e69de29b..25b110b3 100644 --- a/vortex_torch/indexer/triton_kernels/utils_impl.py +++ b/vortex_torch/indexer/triton_kernels/utils_impl.py @@ -0,0 +1,10 @@ +def next_pow2(n: int) -> int: + """Smallest power of two >= ``n``. + + Triton requires the length of a ``tl.arange`` to be a power of two. Kernels + that walk a tensor dimension with ``tl.arange`` therefore have to round that + dimension up to the next power of two and mask off the surplus lanes, while + still using the real extent for pointer arithmetic. + """ + assert n >= 1, f"next_pow2 expects n >= 1, got {n}" + return 1 << (n - 1).bit_length()