[jschmied]

Found while porting the rest of #1156. On main a prompt on a pool lane runs the scalar paged attention partials with all of its chunk's rows in one launch. With 4096-row prompt chunks that overruns pm/pl/po, which hold 128 rows (most_verify):

• a 257-row prompt on a pool lane faults with CUDA_ERROR_ILLEGAL_ADDRESS;
• 129–256 rows overwrite pl/po without an error.

Repro on main: TF_QWEN_KV_PAGES=256 TF_QWEN_KV_PAGES_ALL=1, one 257-token prompt.

The fix runs the paged partials and their merge in launches of at most most_verify rows. A decode round is still one launch.

Gates on GB10 (Qwen3.6-35B-A3B and Qwen3.8-27B MLX 4-bit, bf16 and fp8 KV): flags unset == main on 37/37 greedy outputs per configuration, and with TF_QWEN_KV_PAGES=256 paging == main paging on 22/22 solo and resumed outputs. The every-stream-on-pages runs with 9k-token prompts are in the stacked #1357. zig build test -Dnvcc=/usr/local/cuda/bin/nvcc exits 0, lean shows 0 new.

[jschmied]

Follow-up checks on this PR, one by one:

• Row counts. New kernel-contract test qwen-page-slices (tests/cuda/contract_qwen_page_slices.zig) runs the paged partial and merge the way cuda_forward.zig now does, 128 rows at a time on shifted pointers, and compares the result with one launch over all rows and with the contiguous kernels. It covers one paged stream at 127, 128, 129, 255, 256, 257 and 4096 rows, plus two-stream batches of 127 to 4096 rows where one stream is paged and the other contiguous (some slices cross from one to the other), in bf16 and fp8. Every row matches byte for byte, and the guards around the 128-row partials hold. A planted control, slices that leave the row tables at row 0, comes out red. In the engine itself a prompt round always has exactly one segment (cuda_engine.zig refuses more) and decode rounds stay at 128 rows or fewer, so the mixed case above 128 rows is only reachable at kernel level.
• compute-sanitizer: qwen-page-slices small runs the 129 to 257-row cases. memcheck with leak check reports 0 errors and 0 bytes leaked; racecheck reports 0 hazards (18 cases each).
• Indexing. The partial and merge kernels take the row index from blockIdx (relative to the slice) and use it for Q, rowseg, rowanc, the partials and OUT. They take the segment from segs[rowseg[m]], and since rowseg holds absolute segment indices, the segment table needs no shift. One more thing to note: past_rows(P, m) compares that relative index with P[2], the round's total rows. That is only correct because a sliced (prompt) shape is exact (P[2] == total, so a relative index is always below it), and a padded graph-captured decode shape is never more than one slice. The test passes the engine's params with P[2] = total.
• Launches per attention layer for a 4096-row paged chunk: 32 partials and 32 merges, so 64 instead of the 2 (out of bounds) before. The head norms, rope and gate stay at 4. The contiguous tensor-core prompt path is 1 launch.

Gates: zig build test passes on the branch; lean has 0 new problems. No engine code changed on this level, only the test.

[ashhart]

Review of #1355–#1359 against main e9b69ed, in stack order.

• cuda qwen3_5: a paged prompt chunk's attention goes through most_verify rows at a time #1355: the fix is right: the partial-attention scratch holds 128 rows, and prompt chunks now go in slices of at most that, with offsets for the query, row tables and output. Please add a regression test (a launch-recording host test where a 257-row chunk issues 3 slices). Device check: TF_QWEN_KV_PAGES=256 TF_QWEN_KV_PAGES_ALL=1 with a 257-token prompt, which faults on main.
• cuda qwen3_5: an MTP stream's ring and state are part of its pool lane's run (#1156) #1356: reads correctly (the MTP ring is borrowed from the lane's run and released with the stream), but it has no tests. Please add a host test of the area layout and the borrow/release.
• cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357: off by default and exact (byte-identical default path), but the benefit is memory and there's no memory receipt, while with it on warm time-to-first-token is about 3.4× slower and cold 2–3×, because every stream takes the scalar paged path. It also adds a host synchronize on the engine stream at every kept mark, and cuda_prefix.zig has no tests. Please hold it until paged prompt chunks have a tensor-core path, or show peak device and pinned memory off vs on at --parallel 8.
• cuda qwen3_5: --learn-conversations keeps the 27B's and 35B's states on disk (#840) #1358: sound, but it doesn't compile on main without two pieces of cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357 (cuda_kv_pool.ownCache and the family host build's engine_api import). Please decouple it so it can land independently of cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357.
• cuda flashnext: --learn-conversations keeps Flash Next's states on disk (#840) #1359: two jobs in one PR: please split the Flash Next CUDA disk image from turning on CUDA spill. The spill flag helpers copy metal_spill.zig's slots() and refusal wording; move that into a backend-neutral module instead (gilby's flashnext: shared sessions spill evicted states to disk (--spill-gib) #1349 edits the same file, so coordinate there). Have the family declare kept_disk as cuda qwen3_5: --learn-conversations keeps the 27B's and 35B's states on disk (#840) #1358 does rather than special-casing Flash Next in cuda_disk.keptDisk, and say how Nemotron spill is gated on CUDA.

[BobClawblaw]

Tested on a discrete card under Windows: one RTX 4090 (24 GB, sm_89, driver 616.56), CUDA 13.3, MSVC host, Zig 0.17.0. Qwen3.8-27B MLX 4-bit with DFlash2, --context 8192 --parallel 3 (the most whole-window streams that fit 24 GB beside the drafter), bf16 and fp8 KV. Base 21db2c4c9 against this stack's current tip f3d545559 (#1357's head, so #1355 at 9ec466da6 and #1356 at 03e776237 included). Each arm is a fresh server; 8 greedy chats (24 to 4,616 prompt tokens, thinking off, 160 tokens), sent alone, with "draft": false, and three at a time, plus a 6,073-token two-turn chat.

The overrun reproduces on the base, and on this card it faults below 257 rows too:

• TF_QWEN_KV_PAGES=256 TF_QWEN_KV_PAGES_ALL=1, one 181-token prompt alone: cuStreamSynchronize: CUDA_ERROR_ILLEGAL_INSTRUCTION (715) and the server exits. The 24- and 31-token prompts before it were served.
• TF_QWEN_KV_PAGES=256 alone (only the streams beside the first on pages), bf16 and fp8: the first three-at-a-time group (24, 31 and 181 tokens) ends the same way.

At the tip neither faults: every paged arm serves all of its requests and the server is alive at the end. With paging on, its solo and undrafted replies equal the base's paging replies (16/16 where the base answered, bf16 and fp8).

[jschmied]

Added both.

Host test (commit "cuda qwen3_5: a paged chunk's slices are cuda_attention.pagedSlices, with a host test of their launches"): the slice loop moved unchanged from cuda_forward.zig into cuda_attention.pagedSlices, so a recorder can stand in for the engine. A 257-row paged chunk issues three partial + merge pairs: grid z / y = 128, 128, 1. Each pair runs on the queries, row table, ancestors and output of its slice's first row (offsets 0, 128 and 256 rows). The segment table and params stay unshifted. bf16 and fp8 both pick the paged partial. 1, 127, 128, 129, 256, 257 and 4096 rows issue one pair per started slice and merge every row exactly once. Mutation check: with the ancestor table left unshifted, the test goes red (expected 130048, found 128000). The merge launch keeps its compile-time signature check, and the launch-argument scanner (contract_q35_launch_args_host.zig) still sees it.

Device check (GB10, Qwen3.6-35B-A3B MLX-4bit, --parallel 2 --context 8192, no drafter, greedy; a prompt of exactly 257 tokens by /tokenize, 24 new tokens):

• main e9b69ed, TF_QWEN_KV_PAGES=256 TF_QWEN_KV_PAGES_ALL=1: cuStreamSynchronize: CUDA_ERROR_ILLEGAL_ADDRESS (700) on the 257-token prompt, and the server exits.
• this PR, same flags: 257, 1000 and 4096-token prompts are all served, and the server is alive at the end.
• the replies are byte-identical to this PR's contiguous run and to main's contiguous run, for all three lengths.

The kernel-contract test from the earlier commit (qwen-page-slices) is unchanged. zig build test passes and lean shows 0 new problems.

[jschmied]

Closing: since 0843339 (#1373) a paged prompt chunk takes the tensor-core prompt tiles over its page table, the same path as a contiguous one, so it no longer reaches the 128-row partial scratch this PR sliced for. Thanks for the review.