Repository navigation
qwen CUDA: streams on pages by default (256 positions, the engine's own too), with the paged prompt and decode at the window path's speed - #2011
Conversation
|
Thanks @jschmied. An independent review on current main (d8b80bd):
Applying it:
Nothing was built or run here. Is paged output the same as windowed output?The attention reads are exact by construction. Only the address each key and value row is read from changes. The
The rest of the paged path isn't new. The page pool, the per-stream page runs, the engine stream on pages The author's receipts (one GB10; the 27B with DFlash2 and the 35B-A3B with MTP; their own harness and a local
Not shown: a resumed turn against a fresh one (the author says so), and any chip other than sm_121. Memory admission
Defaults and the window path
Speed
Contract
ScopeOnly Qwen 3.5/3.6 CUDA files, the shared CUDA prefill-attention kernel (its other instantiations get only Small points (not blocking)
Device cells
|
… 16-byte chunk hpattn_kernel and hpattn8_kernel staged each chunk through PT[row / PTOK], a runtime divide and a dependent table load per chunk. An item's 32 keys sit in one page whenever PTOK is a multiple of 32, so its first page is read once and a row's offset is row - lo; a row past that page (PTOK 48) still looks up its own. Addresses only: the same bytes land in the same slots.
…, with the page size a constant TF_QWEN_KV_PAGES's default page holds 256 positions. hpattn_p256_kernel is hpattn_kernel with PTOK a compile-time 256: an item's keys sit in one page (static_assert), so it reads the table once and finds its base once, with no runtime divide or modulo. cuda_prompt.zig takes it for bf16 pages of 256; every other page size keeps hpattn_kernel. Addresses only: the same bytes land in the same slots. The kernel contract checks it against contiguous rows twice, with a swapped page table as its control, and registers its usage.
…e, not once per key Bf16Tiles::fetch and Fp8Tiles::fetch found each key's page with a runtime pos / ptok and pos % ptok. A thread's keys run in order, so it keeps the page it is in (its first position and base) and looks a new one up only when a key leaves it. Addresses only: the same bytes are copied. On a GB10 this brings paged decode on Qwen3.8-27B to the window path's speed (attn_partial was 0.327 against 0.266 ms a call before).
TF_QWEN_KV_PAGES unset now pages at 256 positions (a divisor of the 4096-row prompt chunk, as TF_QWEN_KV_SHARE needs) and TF_QWEN_KV_PAGES_ALL unset puts the engine's stream on pages too. 0 keeps a window each (PAGES) or the engine's window (ALL). Beside --slide, whose learner reads whole windows, an unset TF_QWEN_KV_PAGES keeps windows; set explicitly it is still refused. Both are still read through the profile registry, which supplies no value. Its reporting follows the new default: a count entry names what its reader takes when unset (256 for TF_QWEN_KV_PAGES), and TF_QWEN_KV_PAGES_ALL takes the `off` syntax, on unless exactly 0, so the startup line reports 0 and not 256 or 1.
…les, so its first row is looked up once a page, not once a tile
|
Thanks @ashhart. Before closing the gap, I ran an A/A on main 68906a4 to separate the paging cost from run-to-run noise. Main's binary ran twice (mainA, mainB). This PR, rebased on 68906a4 plus three fixes (below), also ran twice: with pages (head), and with 35B (MTP). Mean of 4 paired rounds, [min, max]; + means slower:
What the A/A shows:
Where the 35B's warm-turn cost goes. I compared CPU-sampled nsys traces of one warm turn, paged against windowed: 701.6 against 699.0 ms, so +2.6 ms.
Choices for the zeroing. I'd like your pick before building one. My proposal is 1:
Expected with 1: the 35B warm turn goes from about +0.65 % to about +0.3 %. Even with 1, the copies and paged attention leave the 35B warm turn about +0.3 % slower, and decode at 16k needs more rounds to settle. So the cell-by-cell "no default cell slower" bar is likely met only with 4, or with 1 plus the copy work. #2043 tracks that copy work, where a kernel writing to mapped host memory measured 6.4 against 11.3 ms at 32k. The three fixes the head runs with are local for now, rebased on 68906a4:
I'll push them with whichever choice you pick. |
|
A correction to my last comment. I measured the zeroing before building anything: a head binary with Mean of 4 paired rounds, [min, max]; + means slower:
What this changes:
Next: before claiming any cell is a paging cost, I'm finding out why two runs of the same binary differ by up to 2 %. I'll post the cell-by-cell table again once the noise floor is explained and tight enough to read a 0.3 % gap. Separately, I had a static audit run on the three buffers: log, taps and drafter ring. Every read of them is bounded by a counter that the lane resets at open, and covers only bytes the stream has written. The window path already relies on this, since it zeroes them only once at init. I'll keep it for later and won't change it in this PR. |
|
Thanks @jschmied. Running an A/A was the right step, and so was dropping the zeroing once it measured nothing. Our call is option 4. The 35B's warm turn was slower in all four rounds of your first run, and a slower default is a no. So please:
Your noise work and #2043 can then make the 35B default its own follow-up. |
…ids in shared memory (up to 768 pages), hpattn_p256_kernel past that
…page, not one run a page and piece
…; a dense one pages at 256 positions TF_QWEN_KV_PAGES unset now pages a dense Qwen3.5 (the 27B) and keeps a MoE (the 35B-A3B) on windows: the 35B's warm turn measured slower on pages, the 27B's no different. The reader takes the MoE bit from the loaded geometry, so a malformed value is still refused before the load. The profile registry reports a count against what this family takes unset (unset_in), and TF_QWEN_KV_PAGES_ALL takes 1 or 0, refusing anything else by name.
c4eb274 to
4e43c68
Compare
|
Thanks @ashhart. Items 1 to 3 are in the pushed head
Here is item 4, the cell-by-cell A/B on that head with an A/A beside it. Setup:
Mean of 6 paired rounds, [min, max], ±2 SE; + means slower: Qwen3.6-35B-A3B, MTP (windows by default)
Qwen3.8-27B, DFlash2 (pages by default)
What it shows:
Proposal: pages opt-in for both models for now. This PR would then:
Making pages the default would move to #2043, once its copy and attention work closes the 0.3 % (its integrated-GPU save and restore measured 6.4 against 11.3 ms at 32k). Another option is to keep the 27B default and accept 0.3 % for the admission gain, but that's against the rule as you stated it. I can push the opt-in change today, with the CHANGELOG line reworded. Which do you prefer? One caveat: some of my host builds overlapped the 35B run's middle rounds. Its numbers are level all the same. |
|
Thanks @jschmied. Locking the clock and dropping the page cache turned the A/A from ±2% into ±0.1% on prompts, and that's what shows the 27B's cost: about 0.3% on 8k and 32k prompts and the warm turn, slower in all six rounds. So yes, pages opt-in for both models for now, as you propose: keep the faster paged path and the MoE fallback here, make |
…set off again, both as on main TF_QWEN_KV_PAGES unset keeps a window each on the 27B and the 35B-A3B alike: the 27B measured about 0.3 % slower on 8k and 32k prompts and the warm turn on pages. The reader still takes the MoE bit from the loaded geometry, so a per-model default is one constant (unset_pages) once ashhart#2043 closes the gap. With pages unset there is no pool, so --parallel reserves a window a stream as on main; set, it is a cap and a request can wait for pages. TF_QWEN_KV_PAGES_ALL is on only at exactly 1 again, so setting only TF_QWEN_KV_PAGES gives what it gave on main. The profile registry, its startup line and docs/profiles.md are main's again.
…oE), _ALL unset off, --parallel a cap only with pages test-qwen35-pages-default-host runs them, part of test-cuda-host. A planted wrong admit count fails the cap test.
|
Thanks @ashhart. Pushed as
Host tests ( Startup lines on the GB10 ( 35B, default: 35B, 27B, default: 27B, The defaults match main. Both default lines are the same as main
|
|
Thanks @jschmied. Pages opt-in for both models, |
|
Thanks @jschmied. The first GB10 runs are in, both against main with Qwen3.8-27B MLX 4-bit, balanced order:
0 of 20 prompt rows differ in either arm, and both read GREEN. The 16-stream overcommit run and the pages contract are next, and it lands if they hold. It'll need a rebase by then; we'll do that at prep. |
|
The 16-stream overcommit run on a GB10, main against this PR with the same model, balanced order: GREEN. Aggregate at 16 streams:
The pages contract is the last run, and it's on the box now. It lands if that holds. |
|
The pages contract on a GB10, this PR against its base: 553 of 553 kernels pass gates 1-4, 359 of 359 planted controls red, and the same 190 limit problems as main, with no limit line new. That was the last run. It goes in the queue, and we'll do the rebase at prep. Thanks @jschmied. |
What this changes
Qwen 3.5/3.6 on CUDA now pages by default. With
TF_QWEN_KV_PAGESunset, every stream takes 256-position pages from the core pool, including the engine's own (TF_QWEN_KV_PAGES_ALLunset). Setting either to0keeps today's windows.So
--parallelbecomes a cap rather than a reservation. A request holds its prompt and reply in pages, and waits when the pool is full. Two kernel changes keep the paged path at the window path's speed:the prompt attention gets a kernel with the page size fixed at 256;
the decode tiles and the prompt rows look a page up once, rather than once per key or 16-byte chunk.
Issue: Refs Proposal: one KV admission ledger for unpaged streams: a request reserves its prompt and reply, not a whole window per --parallel stream (#781 follow-up) #1943 (the "Qwen 3.5/3.6 paged by default" step).
Stacked on: main. Four commits on 071cefe; the branch merges cleanly with ef9fa70.
cuda prefill_attention: a paged item's page looked up once, not every 16-byte chunkcuda prefill_attention: hpattn_p256_kernel for pages of 256 positions, with the page size a constantqwen35 attention tiles: a decode fetch looks its page up once per page, not once per keyqwen CUDA: streams on pages by default, the engine's own tooatc overlap --json:families/qwen3_5/cuda_native.zigwith fix(cuda): allocate Qwen3.5 scratch lazily #1879 (lazy scratch) and qwen3_5 cuda: a chat's next turn resumes at the previous turn's history (message starts for the planned cache) #1335 (resume at history), and git reports conflicts with both. This PR's lines there are thepaging(o.slide)call, two comments and two error texts, so whichever lands second has a small rebase.core/contract/qwen_whole_records.zigoverlaps cuda: use MT2 FP8 attention on RTX PRO 4500 #1731 cleanly.Behaviour:
Unset: 256-position pages for every stream. 256 divides the 4096-row prompt chunk, as
TF_QWEN_KV_SHAREneeds.Beside
--slide: an unsetTF_QWEN_KV_PAGESkeeps windows, since the learner reads whole windows. Set explicitly beside--slide, it's still refused.The profile registry (feat: add stable profile registry and scoped startup reporting #1944) keeps supplying no value, and its startup line follows the new default:
unset = 256forTF_QWEN_KV_PAGES);TF_QWEN_KV_PAGES_ALLtakes a newoffsyntax (on unless exactly0).So
TF_QWEN_KV_PAGES=0andTF_QWEN_KV_PAGES_ALL=0are reported as overrides, while256and1are not.docs/profiles.mdsays so.Checks on this head
zig build testpasses.test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon.git diff --checkis clean.--slideand at0, the registry'sunsetandoffreporting, and the startup line.Receipt
--parallel 4 --context 40960unless stated.qwen3_5andqwen3_5_moe. bf16 and fp8 KV cache, drafted and--no-drafts. HIP and Metal are untouched.TF_GATE_LOGITS, a local hook that isn't in this diff.--parallel 8 --context 32768.--no-drafts. Drafted replies equal--no-draftsreplies.--parallel 16 --context 131072with the prompt cache off:max_tokens16 is the start of its text at 48;--parallel 16needs 143.38 GiB for its streams at a 131072-token window, and the CUDA memory budget leaves 87.99 GiB: serve--parallel 9".Main's first 27B round was a cold outlier; its second matches this PR within 0.1 % on prompts.
The remaining cost, from nsys on one 35B 32k prompt (window path against paged, twice):
So the 35B pays about 0.5 % on prompts and on 4 streams, and the 27B is level. In exchange, a request reserves its prompt and reply instead of a whole window, as the 16-stream run shows.
Kernel contract
hpattn_p256_kernelreplaceshpattn_kernelfor bf16 pages of 256 positions. It's checked against contiguous rows twice. Its control, a swapped page table, fails as planted.hpattn_kernelandhpattn8_kernel(one page lookup per item), and the decode tilesBf16Tiles::fetchandFp8Tiles::fetch(one lookup per page). The addresses change, not the bytes.hpattn_kernel(251 and 233).tf-cuda-test qwen-pageson the GB10 passes 26 of 26, with every control RED as planted for bf16 and fp8.Checklist
"draft": false, and concurrent equals solo. Resumed turns (the 2-turn prompts and the warm turn) equal main's resumed turns; I didn't compare a resumed turn with a fresh one separately.