[jason-fxz]

Part of #893 (triton multi-token verify).

The triton backend sent every batch with more than one query row per request to the extend kernel, which has no KV split: one program per (request, query head) streams the whole context. Speculative verify (MTP, DFlash) and the MTP draft extend (which runs on triton even when the target uses fi) were therefore 6.6-60x slower than single-row decode at 32K-128K.

Changes:

• decode_paged_attention takes up to 8 causal query rows per request (qo_indptr, max_q_len). Stage 1 puts query rows x GQA heads on the M axis and keeps the KV split; each row is limited to its own position and window. Stage 2 merges per row.
• The triton backend sends batches with at most 8 rows per request and no multimodal blocks to it. CUDA graphs with rows_per_req > 1 get per-row split scratch.
• A program's query heads now divide the GQA group. This fixes wrong decode output on main for groups above 16 that are not a multiple of 16 (e.g. 48 query / 2 kv heads). The new test case fails on main.
• Multi-row stage-1 launches pass maxnreg=255 on CUDA. Without it, ptxas holds these tiles at 80-128 registers for occupancy and spills 8-16 B on sm_80 to sm_120. block_n and num_warps come from a sweep among the variants that do not spill on any of these archs.
• Single-row decode keeps main's code: same launch parameters, and the same registers and stack on sm_80/86/89/90/100/120.

End-to-end A/B (main vs this branch)

Qwen/Qwen3.6-35B-A3B bf16, bs 1, greedy, 256 output tokens, --moe-cache-size 10240. Each cell has 2 runs per side (order main, branch, branch, main); the runs agree within 0.3%.

config	32K main	32K branch	128K main	128K branch
bf16 KV, --attention-backend triton, MTP	59.7	243.3 (4.07x)	20.0	238.4 (11.92x)
--kv-cache-dtype fp8_per_token_head (triton), MTP	66.0	253.5 (3.84x)	19.9	252.6 (12.66x)
bf16 KV, default backend (fi), MTP	198.7	246.3 (1.24x)	121.0	236.2 (1.95x)
bf16 KV, --attention-backend triton, no spec	147.4	147.4	124.2	124.2

Output tokens/s. In the fi row the target model's attention stays on fi. The MTP draft layer's attention runs on triton unless the target uses fa (spec/mtp/proposer.py), so its (k+1)-row extend after each verify no longer uses the extend kernel: at 128K the verify step goes from 28.4 to 14.1 ms. Accept length is within 0.12 between main and branch in every MTP cell (3.31-3.54). The greedy outputs are identical at 128K. At 32K, 2 cells take another wording at token ~67 of 256, a near-tie; both texts are coherent. AIME-25 problem 1 pass@3 with fp8 KV + MTP on this branch: 70 / 70 / 70.

Server, as launched by the client (ft = python -m freetoken.cli; 132352 instead of 34048 at 128K):

ft serve --model Qwen/Qwen3.6-35B-A3B --moe-backend offload --max-running-requests 1 --max-seq-len-override 34048 \
    --memory-ratio 0.9 --cuda-graph-max-bs 1 --moe-hybrid-max-fetch -1 --kv-reserve-tokens 34048 --moe-cache-size 10240 \
    --speculative-algorithm mtp [--attention-backend triton | --kv-cache-dtype fp8_per_token_head]

The client is benchmarks/bench_decode_moe.py --context 32768 --greedy --cache 10240, plus two flags that it passes through to ft serve (--speculative-algorithm, --speculative-draft-model-path). The no-spec row runs on the script as it is on main.

Kernel time

Stage 1 + stage 2 in a CUDA graph with L2 flushed before each rep, RTX PRO 6000. Shape 16 query / 2 kv heads, head_dim 256 (Qwen3.6-35B-A3B full attention), bf16 KV, times in us.

ctx	bs	R=1 decode	R=4 before	R=4 after	R=4 FlashInfer	R=8 before	R=8 after	R=8 FlashInfer
32K	1	78	3713	85	88	3712	100	90
32K	16	804	6902	845	781	6915	922	782
128K	1	249	14820	254	258	14816	308	258
128K	16	2989	27473	3081	2888	27536	3381	2890

Before this PR, R > 1 ran the extend kernel. FlashInfer fa2 BatchPrefillWithPagedKVCacheWrapper is shown for reference. The R=1 code is unchanged. "after" is the interquartile mean; the other columns are the medians of an earlier run on the same GPU.

Tested on

NVIDIA RTX PRO 6000 Blackwell Server Edition 96 GB (sm_120), driver 595.91.07; Intel Xeon Platinum 8559C; torch 2.14.1+cu130, triton 3.8.0, flashinfer 0.6.18.post1. Checkpoint Qwen/Qwen3.6-35B-A3B.

• pytest tests/kernels/test_triton_attention.py tests/attention tests/spec tests/scheduler/test_spec_drain.py: all pass, except the 2 fa cases in test_capture_rows_per_req.py. Those fail on main too, because fa does not support sm_120.
• pytest tests -m "not slow": 91 failures, all of which fail the same way on main (ggml JIT 81, local-checkpoint quant configs 8, fa on sm_120 2).
• Register spills: every stage-1 and stage-2 kernel for 9 head shapes x bf16/fp8/int8 x 1-8 rows x bs 1/4/16/64 x window none/64, compiled with the launch specialization for sm_80/86/89/90/100/120. Each was checked with Triton's ptxas: no multi-row kernel spills or exceeds the GPU's shared memory. Single-row stage 1 has the same registers, stack and shared memory as on main, so its existing spills (bf16 window one-tile launches) are unchanged. One stage-2 config differs by 1-2 registers.

[jason-fxz]

Gemma-4 MTP at long context

With #895 merged, Gemma-4's assistant-draft MTP verifies k+1 rows on triton too. Its full-attention layers have head_dim 512 and 1 kv head, so on main those rows took the generic paged_attention, which has no KV split either. Main (58b3ad7) vs this branch merged onto it:

google/gemma-4-12B-it bf16 + unsloth's MTP/mtp-gemma-4-12b-it-BF16.gguf, bs 1, greedy, 256 output tokens, --cache-type naive, backend auto (resolves to triton). Output tokens/s; MTP k=3 cells are the mean of 2 runs per side, which agree within 1%.

config	8K main	8K branch	32K main	32K branch	128K main	128K branch
MTP k=3	102.7	149.5 (1.46x)	68.3	148.2 (2.17x)	27.5	128.8 (4.68x)
MTP k=7 (8 rows)	-	-	108.1	202.7 (1.87x)	42.3	169.6 (4.01x)
no spec	54.0	53.7	53.3	53.3	50.5	50.5

At 128K the k=3 verify step goes from 128.6 to 27.5 ms (a no-spec decode step is 20.1 ms), so MTP on main was slower than no spec there. The greedy outputs are identical between main and branch at every context and config. Accept length matches too, except k=7 at 128K (5.43 -> 5.64, same text).

ft serve --model google/gemma-4-12B-it --text-model-only --cache-type naive --max-running-requests 1 \
    --max-seq-len-override 133120 --num-tokens 135168 --memory-ratio 0.9 --cuda-graph-max-bs 1 \
    [--spec-algo mtp --spec-draft-model mtp-gemma-4-12b-it-BF16.gguf --spec-num-steps 3|7]