[faresobeid]

Rewrites the indexed_attention kernels (Qwen3.8-Flash-Next QSA) as windowed GQA Triton kernels with a deterministic backward. The public API is unchanged: indexed_attention(query, key, value, indices, scale). select_indexed_blocks stays in TileLang.

On main, one TileLang CTA runs per query and KV head: the 12 heads of a KV group are padded to 16 MMA rows, each CTA gathers its ~2,051 selected K/V rows by itself, and the backward scatters dK/dV with fp32 atomics (not deterministic).

Change

• Plan (prime_kernels::indexed_attention_plan, new): groups 5 consecutive queries into a window and stores, per window, the sorted distinct keys with a 5-bit mask of which queries select each key, plus the transpose (per 32-key tile, the selecting windows in order). Built once per call and saved for backward; smaller than the indices main saves (60 MiB vs 128 MiB at 16K).
• Forward: one program per window and KV head; 64 MMA rows hold 5 queries x 12 heads; masked online softmax over the window's key union. Neighbouring queries share most blocks, so the union is 1.6-2.5x one query's selection, not 5x.
• Backward: dQ uses the same window loop. dK/dV run key-major: each 32-key tile walks its windows in fixed order and accumulates in registers. No atomics, no fp32 dK/dV buffers, bitwise deterministic.
• requires gains triton.

Numbers (H200, one Qwen3.8 QSA layer: 24 q / 2 KV heads, head_dim 256, 2,048-token budget; median ms, fwd+bwd incl. plan)

queries x keys	main	this PR	speedup	peak main / new
16K x 16K	116	31	3.69x	0.63 / 0.60 GiB
64K x 64K	352	191	1.84x	2.53 / 1.98 GiB
131K x 131K	669	448	1.49x	5.07 / 4.04 GiB
CP4 rank: 32K x 131K	163	134	1.22x	2.02 / 1.22 GiB
CP8 rank: 16K x 131K	84	65	1.28x	1.51 / 0.74 GiB

Indices come from prime-rl's SparseAttentionIndexer with random weights on locally correlated hidden states. The gain depends on window overlap: with a less local selection (union 3.8x at 64K) it is 1.08x.

Validation

• Against main's kernel (all shapes): out rel 0.13-0.16%, dq 0.11%, dk 0.11-0.14%, dv 0.02-0.08%; lse max diff 1.9e-6. Against an fp32 dense reference (8K, 16K), new and main have the same error.
• dq/dk/dv are bitwise identical across repeats at every shape.
• torch.compile(fullgraph=True, dynamic=True) is bitwise equal to eager. batch=2 matches two single calls.

Notes

• In prime-rl with full AC, the plan is rebuilt in the recompute (18 ms at 64K) unless prime_kernels::indexed_attention_plan is added to the AC save list. Using this in prime-rl needs a wheel release and a pin bump.

In a prime-rl RL run

prime-rl main 1882f7b98 on 1 node (8×H200). Random-weight 4-layer Qwen3.8 (indexed_attention runs in layer 3; the n-gram table is not frozen). Reverse-text task, debug algorithm, NCCL broadcast, vLLM TP4. Batch 64 × ~3.1K tokens (above the 2,048-token budget). 24 steps, no AC. A = main's kernel (prime-kernels 3101960), B = this PR (20685c4). Each arm is the built wheel on PYTHONPATH; a per-rank call counter confirmed the trainer used the intended kernel.

arm, context	mismatch KL, mean (min-max) over steps	loss	grad norm	fwd+bwd / step (median)
A, no CP (seq_len 4096)	0.0984 (0.093-0.104)	-0.0157	3.289	19.6 s / 28.6 s
B, no CP	0.0990 (0.093-0.105)	-0.0157	3.290	14.1 s / 23.3 s
A, CP4 (seq_len 16384)	0.0982 (0.093-0.104)	-0.0769	3.282	19.8 s / 29.1 s
B, CP4	0.0982 (0.093-0.103)	-0.0769	3.285	13.9 s / 23.2 s

Mismatch KL rises 0.093 → 0.104 in all four runs as token staleness grows (mean 11, max 23; max_staleness = 32). There is no difference between the kernels.

Exact check: same weights and batch, one trainer forward + backward (4 GPUs, full AC, two packed documents). Diffs of per-token logprobs, and of the layer-3 attention parameter gradients (relative L2):

	no CP, 8K/rank	CP4, 16K
A vs B	logprob max 0.21, mean 3.8e-3; grad 1.5-1.8%	max 0.23, mean 3.7e-3; grad 1.5-1.8%
A vs A (repeat)	max 0.21, mean 3.5e-3; grad 1.5-1.7%	max 0.22, mean 3.4e-3; grad 1.4-1.6%
B vs B (repeat)	0 (bitwise)	0 (bitwise)

A vs B differs by no more than main's kernel differs from itself on a repeat. With this PR, the whole trainer forward and backward is bitwise reproducible, with and without CP.