indexer: support non-power-of-two group sizes in the Triton kernels

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,25 +34,28 @@ 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):
         # Select x for this workload
         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
     )
 
 

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, 

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
     )

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