Repository navigation
GLM CUDA: the KV page pool crowds out the prompt cache, so long chats prefill again every turn #1510
Description
GLM CUDA sizes its KV page pool first, for the window and the streams, and the prompt cache gets no more than what the pool leaves. When the pool takes most of the memory, the cache can no longer hold a long conversation's state (on this checkpoint about 131 MiB plus 12.45 MiB per 1k tokens), so each turn of a long chat prefills the whole conversation again, even while most of the pool's pages sit idle. Asking for more with --prompt-cache-gib doesn't help, because the pool is sized before the cache.
For example, at TP=3 with --parallel 4 and the default 1,048,576-token window, the pool takes 6,147 pages (38.27 GiB, "enough for one at a 1048576-token window") and leaves the cache 497 MiB, so nothing past about 29k tokens is kept. A 64k chat's second turn logs kept nothing at 63488 tokens: 922 MiB passes the 497 MiB budget and prefills again (36.0 s; at 128k, 75.3 s), and --prompt-cache-gib 8 still gives 0.5 GiB. With CONTEXT=262144 the pool is 12.77 GiB, the cache 16.0 GiB, and the second turn resumes (0.39 s and 0.73 s).
What I think is going on. The cache is one arena allocated at start (glm/cuda_snapshots.zig prepare), and each kept state copies into it the rows its stream's pages already hold: every DSA layer's latent and pooled keys and the MTP layer's (glm/cuda_snapshot_layout.zig Plan), besides the KDA state, the rings and the tail. So most of a long state's bytes duplicate rows the pool has just held, and the pool's idle pages can't serve as cache.
Where I'd take it. Keep a kept state's attention pages in the pool instead of copying them, and copy only the fixed part into a slot. The pool's store would hold the pages until the pool needs them, then reclaim them with the entry. A resumed stream would adopt them at the head of its table and copy back only the fixed part (about 131 MiB a rank at TP=3: the KDA state, the rings and the tail). core/kv_pool.zig has the pieces (retain, adopt, share, reclaim), and HIP's Nemotron and Qwen3.5 already use them through core/prompt_radix.zig, where conversations with the same prefix share its pages. CUDA doesn't yet: Nemotron CUDA copies its KV rows as GLM does, and prompt_radix assumes lookahead 0 and an unplanned grid, while GLM's cache is planned with lookahead 1. I see two ways there:
- A, GLM alone. Keep the flat store and have GLM's save, restore and drop retain, adopt and release the pages themselves, without sharing between streams. As far as I read, with marks on the 2,048-row grid a resumed stream never writes a page that holds rows below the mark, so it needs no copy-on-write.
- B, the radix store on CUDA. Teach the CUDA lane host
prompt_radix, with the grid and the lookahead it doesn't take yet, so GLM shares pages as the HIP families do.
Which would you take, or is B already coming? I'd start with A unless it is. Either way I'd add a GPU check: keep, recycle the free pages, resume to the cold run's token ids, and reclaim and miss cleanly.
Meanwhile. With --parallel above 1, a smaller --context leaves the cache room, as above. And for anyone measuring this: on this engine and checkpoint, tools/bench_cache.py's repeated phase, like RigMark's immediate replay, kept nothing. Both send a completion of token ids, and as far as I read, GLM's planned cache gets no marks for one (history_len, rewind_len and user_end stay 0). The same prompt sent twice as a chat message does resume.
Runs
Three GB10 at TP=3 over a switchless RoCE ring, XSFP4, main 1b4f2f979 built with nvcc 13.3 for sm_121, MTP drafts, one user message of filler text sent twice as a chat (16 tokens, temperature 0).
--parallel 4, window |
pool | prompt cache | 64k: 1st, 2nd ttft | 128k: 1st, 2nd ttft |
|---|---|---|---|---|
| 1,048,576 | 6,147 pages, 38.27 GiB | 497 MiB | 36.07 s, 36.04 s (refused) | 75.54 s, 75.33 s (refused) |
| 262,144 | 2,052 pages, 12.77 GiB | 16.0 GiB | 35.70 s, 0.386 s (resumed at 63,488) | 39.16 s (resumed at 63,488 from the 64k state), 0.725 s (resumed at 126,976) |
States as the cache lines count them, by where they were kept: 6,144 tokens 208 MiB, 30,720 tokens 514 MiB (both at --parallel 1, where the cache is 16.0 GiB), 63,488 tokens 922 MiB, 126,976 tokens 1,714 MiB.
At --parallel 1 a chat sent twice resumes: 8,025 tokens at 6,144 (ttft 4.59 → 1.03 s), 32,025 tokens at 30,720 (14.21 → 0.76 s). The same prompts sent through bench_cache.py keep nothing (resumed at 0, kept 0, no refusals) at 8k, 32k and 128k.
Written with Weave (My Symbiotic AI)
Activity
ashhart commentedon Oct 28, 2026
Thanks @Bizuayeu, a clear diagnosis, and the runs make the cost plain: 36 s against 0.39 s for a 64k chat's second turn.
We'd take B, the radix store on CUDA. Kept states belong in the core, and #1357 from @jschmied (open) is already bringing
core/prompt_radix.zigto the CUDA lane host for the Qwen hybrids. GLM should share that store rather than grow page retention of its own. That means teachingprompt_radixa planned grid and a lookahead of 1, with host tests showing that HIP's families and #1357's Qwen path keep their marks and bytes. Please coordinate the CUDA host side with @jschmied there.Your GPU check is the right one. Please add resumed == fresh with MTP drafts on, concurrent == solo, and decode speed against main at TP=2 and TP=3. We can run TP=2 on a GB10 pair.
On the token-id completions keeping nothing: a planned cache gets no marks without message starts, as #1334 found for the Qwen hybrids. It deserves its own issue if you'd like
bench_cache.py's repeated phase to count on GLM.jschmied commentedon Oct 28, 2026
Happy to coordinate, @Bizuayeu. #1357 brings
core/prompt_radix.zigto the CUDA lane host for the Qwen hybrids: kept states at chat marks, the share keyed by the prompt's tokens, and the pool lanes sized to--parallelwhen every stream runs on pages (as of7db39e45b). Its host tests intest-qwen35-share-faults-hostpin the Qwen path's marks and bytes, so a planned grid and a lookahead of 1 can be checked against them.To keep us from colliding: I won't touch
prompt_radixbeyond #1357's needs. If you take the planned grid and the lookahead inprompt_radix, I'll rebase #1357 on it, or the other way round, whichever lands first. I'm glad to review the CUDA host side when you have a branch.nullburn commentedon Oct 28, 2026
Same class of failure on a different CUDA family, in case it widens the report: DeepSeek-V4.1-Flash-1017 (dsv41) on 2× GB10, engine
tensorfold-native 1.0.2(recipe27741e1, imagetensorfold-dsv41:exl3-tp2-928adbd0e57d),TF_DS_POOL=2097152,CONTEXT=524288,--parallel 5,TF_DS_KV_SHARD=1,--spill-gib 64,--prompt-cache-gibunset.What we see (one process, 3.5 days uptime):
prompt cachepinned at 497 MiB of 512 MiB, 0 MiB spare;keptandevictedadvance 1:1 (1775 / 1714); 79 entries.spill_points558 /spill_blocks8992 a rank (~11 GB on/spill),spill_evicted_turns1193 — butprompt_cache_disk_lookups2963 vsprompt_cache_disk_hits33 (all early),spill_reads33,restores_begun33,restored_tokens4.66M.spill_compatible_misses0,restore_wait_timeouts0,restores_deferred_expired0,crc/size/read_failed0. So the lookup is consulted on every request and almost never issues a read.Two questions:
--prompt-cache-gib 2be ignored whileTF_DS_POOL=2097152stands? Our startup line ispool 2097152 positions, fetch 1048576 a part; compressed planes 1060 MiB a rank (1780 MiB replicated); arena 4285 MiB used— at a 524,288 window the pool looks ~2× the derived size, so we would rather shrink the pool first if that is what unlocks the cache.Deployment-side record with the full metric dump is on the recipe repo: MiaAI-Lab/DeepSeek-v4.1-Flash-1017-2x-DGX-Sparks#75.
Bizuayeu commentedon Oct 28, 2026
#1510 (comment)
B it is. I'll build on #1357: teach
prompt_radixa planned grid and a lookahead of 1 for GLM's CUDA path, with host tests that HIP's families and #1357's Qwen path keep their marks and bytes, then the GPU checks you list (resumed equal to fresh with MTP drafts on, concurrent equal to solo, and decode against main at TP=2 and TP=3). @jschmied, thank you; I'll keep to what GLM needs inprompt_radixand rebase on whichever lands first, and I'll post the branch here early so you can see where it touches #1357's host side.Bizuayeu commentedon Oct 28, 2026
#1510 (comment)
The
prompt_radixhalf is on a branch: two commits on main, no GLM code, and none of #1357's files touched.With lookahead 0 and no grid every path stays as it was: the golden run, and on main
test-kv-pool28/28 andzig build test430/430 steps; stacked on #1357's65ef4ffcc,test-qwen35-share-faults-hostpasses 14/14. GLM's CUDA side comes next, along the lines of #1357'scuda_prefix.zig.jschmied commentedon Oct 28, 2026
@Bizuayeu I read the
prompt_radixhalf (5d64bc05e,11fc92b49), and it looks right to me from #1357's side:child,common,split,have,extends);max(1, lookahead)tokens to prefill;deepesttakes only states on the step;Merged onto #1357's head
65ef4ffcc(no conflicts),test-kv-poolandtest-qwen35-share-faults-hostpass 42/42, against 28/28 for #1357 alone. The difference is your new tests.One small point, not blocking:
grid % page == 0is checked only bystd.debug.assertinStore.init, so a release build would accept a planned grid off the page. Today'slookahead == 0check worked the same way. If a family's rules can come from a config, an error there would be safer.Whichever lands first, I'll rebase #1357 on it.
Bizuayeu commentedon Oct 28, 2026
#1510 (comment)
Thanks for reading it and for the stacked run. Done in
59e416ddeon the same branch:Store.initnow returnserror.BadPromptRadixRulesfor a zero page or a planned grid that isn't a multiple of the page, with a test for each (the old assert fails it).ashhart commentedon Oct 28, 2026
Thanks @nullburn, that does look like the same class. The numbers fit: a 497 MiB cache holds DeepSeek V4.1's kept state up to about 26k tokens, and 26,624 is 13 of its 2,048-row prompt steps. Within a turn, the stream resumes from its own live state, so it looks fine. After an idle gap the next request falls back to the deepest state the cache could keep, and prefills the rest again.
What we don't know yet is whether DeepSeek V4.1 sizes its page pool before the cache the way GLM does, which would mean a larger
--prompt-cache-gibcan't help either. We'll check its budget path on main and report back here. The fix in this issue starts with GLM, through the coreprompt_radixstore, and DeepSeek V4.1 would follow on the same store. One note: 1.0.2 is well behind. If you can, please try 1.1.8 with the same settings and say whether the idle-gap resume point changes.ashhart commentedon Oct 28, 2026
@nullburn, we checked main. DeepSeek V4.1's cache works differently from GLM's, so this isn't the same clamp:
--prompt-cache-gibto its session store, so on main neither 2 nor 0 changes anything there. Its own settings areTF_DSV41_SESSION_RAM_MIB(host snapshots, default 256 MiB),TF_DSV41_SESSION_RESIDENTandTF_DSV41_SESSION_RESIDENT_GIB(idle conversations and their pages), andTF_DSV41_POOL_TOKENSfor the pool. Kept pages compete with other requests inside that pool.Your build's
TF_DS_*settings and its 512 MiB counter don't match main, so we can't tell from main why the newest full state left RAM in your runs. In that fork, the useful things to log are the value its prompt store receives, the bytes charged per state, and the reason the newest full state is dropped. We'll fix the--prompt-cache-gibgap on main.Bizuayeu commentedon Oct 28, 2026
#1510 (comment)
The GLM half is on a branch: on #1357's head
65ef4ffcc, theprompt_radixhalf's three commits, then GLM's two (63f4f8495on one rank,7739c20f8for TP). WithTF_GLM_KV_SHARE=nandTF_KV_PAGES=on, GLM keeps prompts as the pool's pages inprompt_radix's store, with at most n states past them; unset, every path stays as it was.Resumeevery rank takes.TF_KV_POISONon the other ranks leaves their rows alone, and the snapshot step says whether states leave the pages out, so every rank's snapshot has rank 0's shape. Rank 0 sends its drops before the begin step that maps a freed page, asbeginOnalready flushes first.own_prompt_cache), as for cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357's Qwen.Host tests (Linux x86_64, on
7739c20f8):test-kv-pool28/28,test-cuda-host260/262 steps (3,801 of 3,811 tests, 9 skipped),zig build test425/428 steps (1,656 of 1,674, 17 skipped),compile-cuda-tests10/10,compile-config-server7/7,zig fmt --check, andtools/zig/lean_check.pywith nothing new; in both suites the one failure is Qwen3.5's timing-dependent cost-identity test, which passes in other runs. The new GLM tests fail with their property taken out: a NaN fill over the adopted pages, aPlanthat ignores the shared pages, a keep that names one page more, and another rank's map that fills the cache's pages.GPU check at TP=3 on our ring of three GB10 (
0cc5cd348, the same tree as7739c20f8; PARALLEL=4,TF_KV_PAGES=on, the 1M window, MTP drafts on). One binary served withTF_GLM_KV_SHAREunset (main's flat store) and=4, launched unset, on, on, unset; in each, one chat message of filler text sent fresh, again, then twice at once, at 64k and 128k tokens:TF_GLM_KV_SHARE=4bench_openai)GPU check at TP=2 on our pair of GB10 (the same binary; PARALLEL=4,
TF_KV_PAGES=on, the 524k window, MTP drafts on), launched and sent the same way:TF_GLM_KV_SHARE=4bench_openai)ashhart commentedon Oct 30, 2026
Thanks @Bizuayeu. This is what we asked for: the GLM half on
prompt_radix, resumed equal to fresh and concurrentequal to solo in every launch, and the 64k repeat going from 37 s to 1.8 s at TP=3 with decode level.
Next on our side: review both branches and run TP=2 on our GB10 pair, as offered. Your TP=2 table shows main's flat
store has room at the 524k window, so the pair is a no-regression check, not the case the fix is for.
Next on yours: open the
prompt_radixPR, then GLM's two commits stacked on it. Whichever of it and #1357 landssecond rebases on the other. #1357 passed its GB10 cells and lands today.
@nullburn: the DeepSeek V4.1 side is now its own issue, #1625. We'll report there what we find in its budget path, and
whether
--prompt-cache-gibapplies to it.Bizuayeu commentedon Oct 30, 2026
#1510 (comment)
Thanks. The
prompt_radixPR is up as #1659, CI green. GLM's two commits come next, stacked on it, once #1357 lands.Bizuayeu commentedon Oct 31, 2026
#1510 (comment)
GLM's half is up as #1732, on main now that #1357 and #1659 are both in. It has a third commit: the start-up warning below 4 ×
--parallelstates, as #1357 added for Qwen.ashhart commentedon Oct 31, 2026
Status: #1659 is on main. It teaches
prompt_radixa lookahead and a planned grid. A store given rules it can't keepnow fails with
error.BadPromptRadixRulesin every build. No family turns these rules on yet, so every path on mainkeeps its marks, pages and bytes. #1357 is on main too.
Left: GLM's half, #1732 (opt-in,
TF_GLM_KV_SHARE), which brings the 64k repeat at TP=3 from 37 s to 1.8 s in@Bizuayeu's runs. We still have to review it and run the TP=2 no-regression check on our pair. Thanks @Bizuayeu and
@jschmied.
glm cuda: TF_GLM_HOST_CACHE keeps kept prompt states in pinned host m…
glm cuda: TF_GLM_HOST_CACHE keeps kept prompt states in pinned host m…
ashhart commentedon Nov 1, 2026
Thanks @Bizuayeu for the diagnosis and for building both halves, and @jschmied for coordinating with #1357. This is fixed on main, in 1.2.1. #1659 taught the shared prompt store a lookahead and a planned grid (in 1.2.0). #1732 adds
TF_GLM_KV_SHARE=nwithTF_KV_PAGES=on, so GLM keeps prompts as the KV pool's pages: at TP=3 with the 1M window, a 64k message sent again resumes in about 1.7 s instead of 36 s. #1736 addsTF_GLM_HOST_CACHE=1for discrete cards. Both are opt-in; unset, nothing changes. The DeepSeek V4.1 side, reported here by @nullburn, continues on #1625. Closing.