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()