Repository navigation
CUDA: Qwen3.5/3.6 attention caches on the core KV pool (#781 follow-up) #1156
Description
Summary
Each Qwen3.5/3.6 stream on CUDA reserves a full-context window of attention keys and values when it is admitted. On the 27B at --context 32768 that is 2 GiB a stream (16 attention layers × 4 KV heads × 256 × bf16), whatever the request's length. Nemotron moved to the core pool in 1.1.0 (#781/#935). This issue proposes the same for the qwen3_5 / qwen3_5_moe CUDA engine:
- the first stream keeps its reserved window;
- every other stream takes pages from one pool, sized to its prompt + max_tokens + one window;
- admission goes through
fits/LaneMemory, as Nemotron's does.
Overlap check
- Proposal: GLM-5.3-Flash on CUDA takes its streams' KV pages from the core pool, and kept prompt states copy only their rows #1142 (GLM-5.3-Flash pages): GLM only.
- Proposal: several models on one Mac under one memory budget #1027: Mac multi-model budget.
- sessions: session memory for paged-KV families (page pool, prefix trie, RAM residency budget, streamed NVMe tier) (DeepSeek-V4.1 stack 4/14) #637: closed.
- qwen27 CUDA: an opt-in 4-bit attention cache (262K on a 24 GB card) — agreeing on its gate #799 (
--kv-cache q4): touches the qwen KV formats, not addressing. The two would only meet in the fp8/q4 kernels, which this proposal leaves contiguous.
No open claim on qwen3_5 paging.
Design
Kernels
qwen3_5 uses its own attention kernels (qwen35.cu: rope's store, the split-K tf_q35_attn_partial, and keep's tf_q35k_compact), not the shared tree/window attention. Seg gains three fields: pt (page table), ps (page stride) and ptok (positions a page).
One helper, kv_row(), gives every row address. With pt == 0 it is the old expression, so contiguous mode is unchanged.
Only the address changes; the per-key summation order does not. Paged sums are therefore byte-equal to contiguous sums for any page size, not just multiples of the chunk width.
Host
cuda_pages.zig:layout(): the fixed part is state, log, tap ring and table; a token costskvBytes(1).Pages: one arena, akv.Pool, and lanes ofCache.initPaged.
- Prompt-cache saves and restores copy page runs into the same host layout, so kept states move freely between paged and contiguous caches.
src/native/cuda_kv_pool.served():- a wrapping drafter's per-stream bytes (DFlash2, MTP) join the fixed part;
- a family with no free stream pays its first window out of the room before pages are planned.
- It is a no-op for Nemotron.
- The pool is off by default, behind
TF_QWEN_KV_PAGES=<positions a page>.TF_QWEN_KV_PAGES_ALL=1puts every stream on pages (checks). --kv-cache fp8pages too:qwen35_kv8.cugoes through the same kind of helper (row8: codes and scale), byte-equal to contiguous fp8.- Warm-up: reservation warms the 2..N-stream shapes on lanes opened for one row, so the first concurrent round runs nothing cold.
- Recurrent state from the pool: a lane's GDN recurrences and tails, rollback copy, verify log and tap ring are one run of consecutive pages, taken when its stream opens and returned when it closes. Core
kv_poolgainsLayout.run_bytes(default 0, so Nemotron is unchanged) andPool.claimRun/releaseRun/findRun. Runs are taken from the top of the arena and KV pages from the bottom; fragmentation can only delay admission. An idle lane holds only its page table. --slideis still refused with the pool on (its learner reads whole windows).
Branch: jschmied:qwen-paged-kv, six commits on d1018b31: 3f5b87d kernels + contract, d5189b3 pool, 6037621 warm-up, 5c09685 fp8, 9a99d1f comments, 1ac2075 pooled recurrent state.
Verification (GB10, sm_121, Qwen3.8-27B MLX-4bit)
Contract: tf-cuda-test qwen-pages, also run in kernel-contract
Setup (bf16 and fp8):
- 6 shapes, with cached lengths up to 9,000 and pages of 16, 48, 128, 256 and 512 positions.
- Two streams' pages are shuffled through one arena.
Gates passed:
- Paged == contiguous for:
- the rotated queries,
- the attention rows,
- rope's stores,
- keep's moves.
- Three repeats are identical.
- The float64 bound holds.
- Guards and inputs are intact.
Both controls are red:
- a swapped table over correctly placed pages;
- a planted row.
Engine
| Arm | Chat / 8k prefill bytes | Decode tok/s | Gates (drafted==plain / concurrent==solo / resumed==fresh) |
|---|---|---|---|
| contiguous (this branch) | == main 3/3, 3/3 | 13.75 / 13.83 | — |
| page test, 16-position reversed table | == contiguous 3/3, 3/3 | 13.83 / 14.00 | 8/8, 8/8, True (DFlash2, --parallel 8) |
| page test, 256 | == contiguous 3/3, 3/3 | — | — |
pool, 256, --parallel 4 |
== contiguous 3/3, 3/3 | 13.88 / 14.05; c=4 48.6 code, 48.2 chat (contiguous 48.4, 46.9) | — |
pool, 256, --parallel 8 + DFlash2 |
— | — | 8/8, 8/8, True; outputs == page test 8/8 + 8/8 |
pool, pages only (TF_QWEN_KV_PAGES_ALL=1) |
== contiguous 3/3, 3/3 | — | 8/8, 8/8, True (state restored into a lane; 8 requests on 7 lanes, one waited) |
| fp8 contiguous (new kernels) | == previous fp8 kernels 3/3, 3/3 | 14.0 / 13.7 (was 14.05 / 13.6); c=4 48.9 / 47.5 (was 49.2 / 47.7) | 8/8, 8/8, True |
| fp8 page test 16, fp8 pool 256 | == fp8 contiguous 3/3, 3/3 | pool: 14.0 / 13.7; c=4 49.2 / 47.7 | pages only: 8/8, 8/8, True, == fp8 contiguous |
| pool with pooled state, 256 | == contiguous 3/3, 3/3 (also pages only) | 13.9 / 14.1; c=4 48.8 / 48.3 | pages only, bf16 and fp8: 8/8, 8/8, True, == earlier runs |
First concurrent round on a fresh server (8 requests, --parallel 8): TTFT max 0.97 s warmed, 0.98 s unwarmed, 0.99 s contiguous. Qwen captures no graphs, so cold shapes cost little; the warm-up is there for consistency.
Memory at --parallel 4, 32k context, 27B:
| Other streams' memory | |
|---|---|
| main | three 2.5 GiB windows + their state, held from start |
| pool (d5189b3) | 387 pages (6.05 GiB) taken as needed + 558 MiB fixed a lane |
| pool + pooled state (1ac2075) | 492 pages (7.69 GiB, KV and state) taken as needed, 0 MiB fixed a lane (the table only; a drafter's own buffers, 169 MiB with DFlash2, stay fixed) |
Unit tests and style
zig build testpasses.lean_checkadds no findings.
Open points
- A drafter's per-stream buffers (DFlash2 ring, MTP cache) are still fixed: the same run mechanism fits them.
- The first stream keeps a reserved window (the learner and the drafter calibration use it).
- The default: on for Qwen once soaked, as Nemotron's
TF_KV_PAGES.
Gotchas
- Off by default (
TF_QWEN_KV_PAGES). The first stream still reserves a full window. --slideis refused with the pool on.- Measured on the MLX 27B only. The MoE engine (35B) takes the same pool, but it wasn't run here.
- The branch is off
d1018b31and doesn't include qwen3_5 cuda: Qwen3.8-27B ModelOpt and compressed-tensors NVFP4 checkpoints, weight-only and in their own math #1145 or qwen CUDA: TF_DRAFT_NODES above 15 serves; a stop or cancel lands between prompt chunks #1078-qwen CUDA: long-prompt prefill 5.4x faster (4096-row chunks, tensor-core attention, the delta-rule chain) #1080. I expect textual overlap inqwen3_5/cuda_*.zigwith both, and with qwen27 CUDA: an opt-in 4-bit attention cache (262K on a 24 GB card) — agreeing on its gate #799's--kv-cache q4.
Question for @ashhart: one PR with the six commits, or kernels + contract first (3f5b87dd), then the pool?
Activity
jschmied commentedon Oct 21, 2026
Fix pushed to the branch:
51f83540on top of1ac20753. At1ac20753, with a drafter andTF_QWEN_KV_PAGESon, the pool was planned before the drafter attached. Each lane's DFlash2 tap ring (~100 MB) was then mapped past its run, into the next one, or past the arena for the top lane. The gates still passed by luck.51f83540plans the pool after the drafter attaches and refuses a run too short for its parts. The DFlash2 ring is now part of the lane's run: 23 MiB fixed per lane instead of 169 MiB with DFlash2. Same gates, all passing: bytes equal contiguous; pages-only bf16 and fp8 at--parallel 8give drafted==plain 8/8, concurrent==solo 8/8, resumed==fresh. Speed is unchanged: DFlash2 decode 60.5 / 44.6 tok/s contiguous, 59.9 / 44.4 pooled.jschmied commentedon Oct 21, 2026
Two updates on
qwen-paged-kv.ce5b78d7, a speed fix. Running the MoE (Qwen3.6-35B-A3BMLX 4-bit with its MTP layer) showed the branch 2.7-3.6 % slower than main at one stream, even with the pool off. The cause: the attention kernels tested for a page table once per key inside the loop. They now test once per block, and the contiguous path computes row addresses exactly as main does. Alternating main/branch runs afterwards: 82.7/74.5 and 84.0/75.2 tok/s (code/chat) against main's 84.2/75.0 and 83.3/75.1, inside run-to-run spread. Bytes are unchanged: the 35B (bf16 and fp8) and the 27B still match main 3/3 chat + 3/3 8k prefill, andqwen-pagespasses with its controls red.The MoE on the pool (closes the "not run" gotcha). At
--parallel 4, 32k context:--parallel 8, all pages: drafted==plain 8/8, concurrent==solo 8/8, resumed==fresh, bf16 and fp8.Still open: MTP's per-stream buffers stay fixed per lane. The target's tap ring is now counted in both the run and the fixed part, which over-reserves a little but is never wrong.
jschmied commentedon Oct 21, 2026
Flash Next on the same core pool: #1169 (its attention caches through a page table, on main) and #1170 (its streams beside the first on pool pages, stacked on this issue's branch at
ce5b78d7and #1169).jschmied commentedon Oct 21, 2026
Two more commits on
qwen-paged-kv, now at0566be8a.ac90d6b8: MTP's per-stream buffers in the lane's run. On the 35B with its MTP layer, a pool lane's MTP ring and newest-row state are part of its run. At--parallel 4that leaves 0 MiB fixed per lane ("3 more taking 582 pages (2.84 GiB)"). Bytes equal main and contiguous, and the gates pass with every stream on pages, in bf16 and fp8.0566be8a: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE=n, off by default).--parallel 8. Every turn's reply is equal across today, share and no prompt cache (80/80 on the 27B + DFlash2 and on the 35B + MTP). Drafted==plain, concurrent==solo and resumed==fresh pass with share on.Today's host cache hit its 8 GiB budget in both runs, so with a bigger budget the gap is larger. First-token time is within noise (median per-turn change −0.14 s and −0.03 s).
build/cuda.ziggives the qwen35 module the server'sengine_api, so both share one instance of the radix store.src/native/cuda.zigskips the host prompt cache when the family keeps prompts itself (cuda_kv_pool.ownCache).Speed A/Bs on this branch need a fresh window-cost table per start; see #1173.
ashhart commentedon Oct 22, 2026
Thanks @jschmied for the Qwen3.5/3.6 pages and for keeping this issue so clearly up to date. Your eight commits on
qwen-paged-kvthroughce5b78d7are carried, kept as yours, onto the KV pool line we are building for the next release, together with fixes of ours for what our review found there: pool pages and runs go back only once the engine's stream has drained, the arena and lane tables are zeroed on the engine's stream, owners are allocated before their fills, and the drafter's ring is counted in a stream's bytes. We'll pick upac90d6b8(MTP's per-stream buffers in the lane's run) and0566be8a(TF_QWEN_KV_SHARE, off by default) next, with the same gates, and we'll take #1169 and #1170 as the Flash Next part. No need to rebase anything; if you push more, just say here which commits are new.- Synthema, Ash's AI Assistant.
jschmied commentedon Oct 26, 2026
The rest of this issue is now three stacked PRs on main
21db2c4c9:ac90d6b8). The 35B goes from 64 to 0 MiB fixed a lane.TF_QWEN_KV_SHARE(was0566be8a), off by default. It saves memory but is slower today: every stream goes on pages, and paged prompts take the scalar attention instead of the tensor-core prompt path (27B warm turn 2: 0.75 → 2.2 s). A paged tensor-core prompt kernel would remove that; I'm working on one.One more thing seen on main with
TF_QWEN_KV_PAGESalone, not changed by these: a concurrent stream's bytes depend on whether its batch holds a paged stream. On the 27B, 12–14 of 15 concurrent outputs equal a solo; with every stream on pages, 15/15.jschmied commentedon Oct 26, 2026
/atc presence 22a2bc6f6c4d4ebf40aa36fbe95cd8a5a7dad45252827688d6bf9d3dfc7fe35f
ashhart commentedon Oct 28, 2026
Status: #1356 shipped in 1.1.8. #1355 was closed once #1373 (in 1.1.7) moved paged prompt chunks onto the tensor-core tiles. #1373 should also cover the paged-batch bytes you saw on 26 Oct, since a paged stream's attention now runs the same tiles as a whole window. Could you confirm that on your next run? #1357 is the remaining piece.
jschmied commentedon Oct 28, 2026
Confirmed on a GB10: #1373 fixed the paged-batch bytes. On main a concurrent stream's reply no longer depends on whether its batch holds a paged stream.
Setup. Main
d0c9dc9fb(1.1.8 + 1),zig build native, CUDA 13.0. Qwen3.8-27B MLX-4bit + DFlash2, bf16 KV,--context 16384, a fresh server per arm. 12 prompts: 8 short questions (11–15 tokens) and 4 filler logs (1,201 / 2,804 / 4,413 / 6,003 tokens). 96 new tokens each, greedy. Each prompt runs alone first, then all 12 are sent at once in two rounds, in forward and in reverse order. With--parallel 8or4the extra requests queue, so the batches mix short and long prompts and window and page lanes.--parallel 8--parallel 4TF_QWEN_KV_PAGES=256: concurrent == soloTF_QWEN_KV_PAGES=256 TF_QWEN_KV_PAGES_ALL=1: concurrent == soloPAGES, andPAGES_ALL)The solo replies are also the same at
--parallel 8and4(12/12).Control. I ran the same suite on the 26 Oct build, from before #1373 (the one behind my earlier counts), at
--parallel 8withTF_QWEN_KV_PAGES=256. 14 of 24 concurrent replies matched their solo run, the same failure as on 26 Oct. Its solo replies equal main's (12/12). So this suite does catch the old difference, and main no longer has it.