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

[GeoffreyWang1117]

Problem

Three indexer kernels walk the group axis with tl.arange sized directly from the tensor shape:

file	line
matmul_impl.py	g_ptr = tl.arange(0, G)
reduce_impl.py	dim0 = tl.arange(0, x_D0) / dim1 = tl.arange(0, x_D1)
softmax_impl.py	d1_idx = tl.arange(0, x_D1)

Triton requires the length of a tl.arange to be a power of two. G here is the GQA group size (num_qo_heads // num_kv_heads — see flow.py::run_indexer_virtual, which builds q_dummy as [1, group_size, head_dim]). So for any model whose group size is not a power of two, the indexer fails at kernel compile time and never runs at all:

    d_ptr   = tl.arange(0, D)
    c_ptr   = tl.arange(0, C)
    g_ptr   = tl.arange(0, G)
arange's range must be a power of 2

Models this rules out (numbers read from the published HF configs):

model	q heads	kv heads	G
Qwen2.5-1.5B, Qwen2.5-Coder-1.5B	12	2	6
Qwen2.5-7B, Qwen2.5-Coder-7B	28	4	7
Qwen2.5-14B, Qwen3-14B	40	8	5

The repo's own example (examples/verify_algo.py, Qwen3-1.7B) has G=2, so this doesn't surface until you point it at a different model.

Change

Round the affected extent up to the next power of two for the tl.arange only, and mask the surplus lanes on load and store. Pointer arithmetic and the Mean divisor keep using the real extent.

For power-of-two shapes the mask is all-true and the generated code is equivalent — nothing changes on the currently-supported path.

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.

next_pow2 lives in the (previously empty) indexer/triton_kernels/utils_impl.py so the three kernels share one definition.

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.

Verification

RTX 3090, torch 2.9.1, triton 3.5.1. The reproducer below covers mm_bpr for group sizes 1–8, reduce_rr for all five ReduceTypes over both dims across six (D0, D1) pairs, and softmax_inplace_r across five shapes — each against a torch reference. 73 checks total.

On v1 today — dies at the first non-power-of-two case:

mm_bpr  (G = num_qo_heads // num_kv_heads)
  [ok ] G=1                                err=2.637e-03
  [ok ] G=2                                err=2.637e-03
...
arange's range must be a power of 2

With this PR — 73/73 pass, all within bf16 tolerance:

mm_bpr  (G = num_qo_heads // num_kv_heads)
  [ok ] G=1                                err=2.637e-03
  [ok ] G=2                                err=2.637e-03
  [ok ] G=3                                err=2.724e-03
  [ok ] G=4                                err=2.523e-03
  [ok ] G=5                                err=2.523e-03
  [ok ] G=6                                err=2.523e-03
  [ok ] G=7                                err=2.523e-03
  [ok ] G=8                                err=2.193e-03

  ...
  [ok ] D0=3 D1=5 dim=1 Mean               err=3.257e-03
  [ok ] D0=3 D1=5 dim=1 Max                err=0.000e+00
  [ok ] D0=3 D1=5 dim=1 Min                err=0.000e+00
  [ok ] D0=7 D1=7 dim=2 Mean               err=2.687e-03
  [ok ] D0=7 D1=7 dim=2 Min                err=0.000e+00
  ...
softmax_inplace_r
  [ok ] D0=4 D1=8                          err=8.782e-04
  [ok ] D0=4 D1=3                          err=6.160e-04
  [ok ] D0=2 D1=7                          err=4.679e-04
  [ok ] D0=8 D1=5                          err=9.330e-04
  [ok ] D0=4 D1=6                          err=9.349e-04

all checks passed

Power-of-two shapes are bit-for-bit unchanged from before the patch.

Reproducer script

"""Reproducer: indexer Triton kernels vs. non-power-of-two GQA group sizes.

Checks mm_bpr / reduce_rr / softmax_inplace_r against a torch reference for a
range of group sizes. On upstream v1 every non-power-of-two case raises
`CompilationError: arange's range must be a power of 2`.

    python repro_non_pow2.py
"""
import torch
from vortex_torch.indexer.triton_kernels.matmul_impl import _mm_bpr
from vortex_torch.indexer.triton_kernels.reduce_impl import _reduce_rr
from vortex_torch.indexer.triton_kernels.softmax_impl import _softmax_inplace_r
from vortex_torch.utils import ReduceType

DEV, CHUNK, NSMS = "cuda", 8, 8
fails = []


def winfo(n_rows, chunk=CHUNK):
    offs, lens, o = [], [], 0
    while o < n_rows:
        L = min(chunk, n_rows - o)
        offs.append(o); lens.append(L); o += L
    t = lambda v: torch.tensor(v, dtype=torch.int32, device=DEV)
    return t(offs), t(lens), t([len(offs)])


def check(name, err, tol=3e-2):
    ok = err == err and err < tol           # NaN-safe
    print(f"  [{'ok ' if ok else 'FAIL'}] {name:<34} err={err:.3e}")
    if not ok:
        fails.append(name)


def mm(G, D=64, C=8, S=20):
    torch.manual_seed(0)
    x = torch.randn(2, G, D, device=DEV, dtype=torch.bfloat16)
    y = torch.randn(S, C, D, device=DEV, dtype=torch.bfloat16)
    o = torch.zeros(S, C, G, device=DEV, dtype=torch.bfloat16)
    offs, lens, nw = winfo(S)
    _mm_bpr(x, y, o, torch.arange(S, dtype=torch.int32, device=DEV),
            torch.zeros(len(lens), dtype=torch.int32, device=DEV),
            offs, lens, nw, CHUNK, NSMS)
    torch.cuda.synchronize()
    ref = torch.einsum("scd,gd->scg", y.float(), x[0].float())
    return (o.float() - ref).abs().max().item() / max(ref.abs().max().item(), 1.0)


def red(D0, D1, dim, rt, N=17):
    torch.manual_seed(0)
    x = torch.randn(N, D0, D1, device=DEV, dtype=torch.bfloat16)
    o = torch.zeros(N, D1 if dim == 1 else D0, device=DEV, dtype=torch.bfloat16)
    offs, lens, nw = winfo(N)
    _reduce_rr(x, o, dim, rt, offs, lens, nw, CHUNK, NSMS)
    torch.cuda.synchronize()
    xf = x.float()
    ref = {ReduceType.Mean: xf.mean, ReduceType.Max: xf.amax, ReduceType.Min: xf.amin,
           ReduceType.Sum: xf.sum}.get(rt)
    ref = xf.pow(2).sum(dim=dim).sqrt() if rt is ReduceType.L2Norm else ref(dim=dim)
    return (o.float() - ref).abs().max().item() / max(ref.abs().max().item(), 1.0)


def sm(D0, D1, B=2, P=12, scale=0.5):
    torch.manual_seed(0)
    x = torch.randn(B * P, D0, D1, device=DEV, dtype=torch.bfloat16)
    ref = torch.softmax(x.float().reshape(B, P, D0, D1) * scale, 1).reshape(B * P, D0, D1)
    indptr = torch.arange(0, B * P + 1, P, dtype=torch.int32, device=DEV)
    _softmax_inplace_r(x, 1, indptr, scale, 0, 0, 1, B)
    torch.cuda.synchronize()
    return (x.float() - ref).abs().max().item()


print(f"device: {torch.cuda.get_device_name(0)}\n")
print("mm_bpr  (G = num_qo_heads // num_kv_heads)")
for G in (1, 2, 3, 4, 5, 6, 7, 8):
    check(f"G={G}", mm(G))

print("\nreduce_rr")
for D0, D1 in ((4, 8), (3, 5), (3, 8), (8, 5), (7, 7), (1, 3)):
    for dim in (1, 2):
        for rt in ReduceType:
            check(f"D0={D0} D1={D1} dim={dim} {rt.name}", red(D0, D1, dim, rt))

print("\nsoftmax_inplace_r")
for D0, D1 in ((4, 8), (4, 3), (2, 7), (8, 5), (4, 6)):
    check(f"D0={D0} D1={D1}", sm(D0, D1))

print(f"\n{'FAILED: ' + ', '.join(fails) if fails else 'all checks passed'}")
raise SystemExit(1 if fails else 0)

tests/ is in .gitignore, so I kept this out of the tree rather than change your ignore rules. Happy to check it in as a real test if you'd like — just say where.

Scope

Deliberately not covered here, to keep the diff reviewable:

• mm_rrr / mm_rpr in the same file (six shape constexprs each) use the same pattern.
• The kernels under vortex_torch/cache/triton_kernels/ do too, though their dims are page/head layout rather than the GQA group, so they're less likely to be hit in practice.

Happy to extend this PR to cover those, or send a follow-up — whichever you prefer.