Skip to content

GLM CUDA: the KV page pool crowds out the prompt cache, so long chats prefill again every turn #1510

Closed

Description

@Bizuayeu

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

  1. ashhart commented on Oct 28, 2026

    @ashhart
    Owner

    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.zig to the CUDA lane host for the Qwen hybrids. GLM should share that store rather than grow page retention of its own. That means teaching prompt_radix a 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.

  2. jschmied commented on Oct 28, 2026

    @jschmied
    Contributor

    Happy to coordinate, @Bizuayeu. #1357 brings core/prompt_radix.zig to 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 --parallel when every stream runs on pages (as of 7db39e45b). Its host tests in test-qwen35-share-faults-host pin 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_radix beyond #1357's needs. If you take the planned grid and the lookahead in prompt_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.

  3. nullburn commented on Oct 28, 2026

    @nullburn
    Contributor

    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 (recipe 27741e1, image tensorfold-dsv41:exl3-tp2-928adbd0e57d), TF_DS_POOL=2097152, CONTEXT=524288, --parallel 5, TF_DS_KV_SHARD=1, --spill-gib 64, --prompt-cache-gib unset.

    What we see (one process, 3.5 days uptime):

    • prompt cache pinned at 497 MiB of 512 MiB, 0 MiB spare; kept and evicted advance 1:1 (1775 / 1714); 79 entries.
    • A long conversation resumes fine within a turn (resumed 260096 → 262144 → 266240 …), but the first request after an idle gap (5–11 min, no other request reached the engine — checked in the container log) resumes at 26,624 tokens (13 × 2048) and re-prefills ~240K (TTFT 137–156 s). The next request in that turn resumes correctly again.
    • Spill is on and writing — spill_points 558 / spill_blocks 8992 a rank (~11 GB on /spill), spill_evicted_turns 1193 — but prompt_cache_disk_lookups 2963 vs prompt_cache_disk_hits 33 (all early), spill_reads 33, restores_begun 33, restored_tokens 4.66M. spill_compatible_misses 0, restore_wait_timeouts 0, restores_deferred_expired 0, crc/size/read_failed 0. So the lookup is consulted on every request and almost never issues a read.

    Two questions:

    1. Is the dsv41 prompt cache also sized from what the pool leaves (as you describe for GLM), i.e. would --prompt-cache-gib 2 be ignored while TF_DS_POOL=2097152 stands? Our startup line is pool 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.
    2. Are 1.1.8's spill changes — write-on-idle-unload (Prompt cache spill: conversations held in RAM are lost at idle unload #1471), I/O-worker writes (Prompt cache spill: evictions write to disk on the engine thread and pause decoding streams #1453), queued read-ahead (Prompt cache spill: restoring a conversation from disk reads on the engine thread and pauses decoding streams #1464) — ported to dsv41, or is dsv41 still on the 1.0.x spill path? We are on 1.0.2, six releases behind 1.1.8.

    Deployment-side record with the full metric dump is on the recipe repo: MiaAI-Lab/DeepSeek-v4.1-Flash-1017-2x-DGX-Sparks#75.

  4. Bizuayeu commented on Oct 28, 2026

    @Bizuayeu
    ContributorAuthor

    We'd take B, the radix store on CUDA. [...] That means teaching prompt_radix a 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.

    #1510 (comment)

    B it is. I'll build on #1357: teach prompt_radix a 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 in prompt_radix and rebase on whichever lands first, and I'll post the branch here early so you can see where it touches #1357's host side.

  5. Bizuayeu commented on Oct 28, 2026

    @Bizuayeu
    ContributorAuthor

    If you take the planned grid and the lookahead in prompt_radix, I'll rebase #1357 on it, or the other way round, whichever lands first.

    #1510 (comment)

    The prompt_radix half is on a branch: two commits on main, no GLM code, and none of #1357's files touched.

    • The first pins today's behaviour: a golden run of seven passes under the rules HIP's families and cuda qwen3_5: the prompt cache keeps KV pages on the device (TF_QWEN_KV_SHARE, #1156) #1357's Qwen path use (lookahead 0, no grid, the tail kept), checking each pass's resume point and marks and the tree's states, pages, bytes and counts.
    • The second teaches the store a lookahead and a planned grid. The lookahead goes into the tree, not only the snapshot: GLM keeps its MTP rows on pages, and row at−1 reads token at, so a page followed by a different token holds different rows. A page is matched and shared only while the prompt repeats the lookahead's tokens past it, so each node keeps those tokens after its last page. A planned store cuts and resumes only on its grid, which must be a multiple of the page.

    With lookahead 0 and no grid every path stays as it was: the golden run, and on main test-kv-pool 28/28 and zig build test 430/430 steps; stacked on #1357's 65ef4ffcc, test-qwen35-share-faults-host passes 14/14. GLM's CUDA side comes next, along the lines of #1357's cuda_prefix.zig.

  6. jschmied commented on Oct 28, 2026

    @jschmied
    Contributor

    @Bizuayeu I read the prompt_radix half (5d64bc05e, 11fc92b49), and it looks right to me from #1357's side:

    • the lookahead is part of every page match (child, common, split, have, extends);
    • a resume leaves max(1, lookahead) tokens to prefill;
    • deepest takes only states on the step;
    • with lookahead 0 and no grid, every path stays as it was.

    Merged onto #1357's head 65ef4ffcc (no conflicts), test-kv-pool and test-qwen35-share-faults-host pass 42/42, against 28/28 for #1357 alone. The difference is your new tests.

    One small point, not blocking: grid % page == 0 is checked only by std.debug.assert in Store.init, so a release build would accept a planned grid off the page. Today's lookahead == 0 check 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.

  7. Bizuayeu commented on Oct 28, 2026

    @Bizuayeu
    ContributorAuthor

    One small point, not blocking: grid % page == 0 is checked only by std.debug.assert in Store.init, so a release build would accept a planned grid off the page. [...] If a family's rules can come from a config, an error there would be safer.

    #1510 (comment)

    Thanks for reading it and for the stacked run. Done in 59e416dde on the same branch: Store.init now returns error.BadPromptRadixRules for 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).

  8. ashhart commented on Oct 28, 2026

    @ashhart
    Owner

    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-gib can'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 core prompt_radix store, 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.

  9. ashhart commented on Oct 28, 2026

    @ashhart
    Owner

    @nullburn, we checked main. DeepSeek V4.1's cache works differently from GLM's, so this isn't the same clamp:

    • It creates the KV pool first, but keeps complete KV pages by reference inside that pool rather than copying the conversation into a second prompt-cache arena. A long kept state doesn't have to fit in a leftover copied-cache budget.
    • There is an option gap, though: the DeepSeek server doesn't pass --prompt-cache-gib to its session store, so on main neither 2 nor 0 changes anything there. Its own settings are TF_DSV41_SESSION_RAM_MIB (host snapshots, default 256 MiB), TF_DSV41_SESSION_RESIDENT and TF_DSV41_SESSION_RESIDENT_GIB (idle conversations and their pages), and TF_DSV41_POOL_TOKENS for 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-gib gap on main.

  10. Bizuayeu commented on Oct 28, 2026

    @Bizuayeu
    ContributorAuthor

    GLM should share that store rather than grow page retention of its own. [...] Please add resumed == fresh with MTP drafts on, concurrent == solo, and decode speed against main at TP=2 and TP=3.

    #1510 (comment)

    The GLM half is on a branch: on #1357's head 65ef4ffcc, the prompt_radix half's three commits, then GLM's two (63f4f8495 on one rank, 7739c20f8 for TP). With TF_GLM_KV_SHARE=n and TF_KV_PAGES=on, GLM keeps prompts as the pool's pages in prompt_radix's store, with at most n states past them; unset, every path stays as it was.

    • Marks sit on the planned grid (2,048 rows, the prompt step) over 512-row pages, with a lookahead of 1: the MTP head's row at a mark reads the token after it.
    • A kept state holds its prefix's pages by reference in the store. Its snapshot keeps only what lies past them: the arena's head, the rings, the two pooled-key rows each paged cache has past the mark, and the MTP head's rows past it. A resumed pass writes nothing below its mark (the pooled keys are written in whole pools from the step's head), so nothing is copied on write.
    • A resume adopts the deepest kept state's pages at the head of the new stream's lease and restores the rest at the prompt pass's start, through the same Resume every rank takes.
    • At TP the tree and the pages stay rank 0's. The begin step says how many of the table's first pages are the cache's, so TF_KV_POISON on 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, as beginOn already flushes first.
    • With it on, the server makes no flat store for GLM (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-pool 28/28, test-cuda-host 260/262 steps (3,801 of 3,811 tests, 9 skipped), zig build test 425/428 steps (1,656 of 1,674, 17 skipped), compile-cuda-tests 10/10, compile-config-server 7/7, zig fmt --check, and tools/zig/lean_check.py with 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, a Plan that 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 as 7739c20f8; PARALLEL=4, TF_KV_PAGES=on, the 1M window, MTP drafts on). One binary served with TF_GLM_KV_SHARE unset (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:

    unset (as main) TF_GLM_KV_SHARE=4
    64k message again (s) 37.07, 37.06 1.77, 1.77
    128k message again (s) 75.63, 75.55 2.14, 2.14
    64k twice at once (s) 71.3 to 74.4 2.31 to 2.32
    128k twice at once (s) 115.7 to 151.4 3.39 to 3.44
    decode (tok/s, bench_openai) 51.17, 50.13 51.07, 51.06
    • Resumed == fresh and concurrent == solo: every reply's tokens are the same in all four launches (rank 0's done-line hash: 16 replies at 64k, 6 of them resumed, and 16 at 128k, 8 resumed).
    • Unset, every repeat prefills again (cached 0), as this issue reports. On, a repeat resumes at the last mark below the prompt (63,488 of 64,030 tokens, 126,976 of 128,030), and the 128k message's first send resumes from the 64k one's state (40.4 s against 75.7 s).

    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:

    unset (as main) TF_GLM_KV_SHARE=4
    64k message again (s) 2.13, 2.13 2.12, 2.12
    128k message again (s) 2.47, 2.48 2.45, 2.46
    64k twice at once (s) 3.25 to 3.31 3.24 to 3.30
    128k twice at once (s) 3.65 to 3.66 3.58 to 3.59
    decode (tok/s, bench_openai) 40.18, 40.06 40.05, 40.15
    • Resumed == fresh and concurrent == solo: every reply's tokens are the same in all four launches (16 replies at 64k, 12 of them resumed, and 16 at 128k, all resumed).
    • At this window main's flat store has room (4.8 GiB, the two states in 2.7 GiB), so unset resumes too, at the same marks, and the two match. The clamp this issue reports shows at TP=3's 1M window.
  11. ashhart commented on Oct 30, 2026

    @ashhart
    Owner

    Thanks @Bizuayeu. This is what we asked for: the GLM half on prompt_radix, resumed equal to fresh and concurrent
    equal 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_radix PR, then GLM's two commits stacked on it. Whichever of it and #1357 lands
    second 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-gib applies to it.

  12. Bizuayeu commented on Oct 30, 2026

    @Bizuayeu
    ContributorAuthor

    Next on yours: open the prompt_radix PR, then GLM's two commits stacked on it. Whichever of it and #1357 lands second rebases on the other.

    #1510 (comment)

    Thanks. The prompt_radix PR is up as #1659, CI green. GLM's two commits come next, stacked on it, once #1357 lands.

  13. Bizuayeu commented on Oct 31, 2026

    @Bizuayeu
    ContributorAuthor

    Next on yours: open the prompt_radix PR, then GLM's two commits stacked on it. Whichever of it and #1357 lands second rebases on the other.

    #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 × --parallel states, as #1357 added for Qwen.

  14. ashhart commented on Oct 31, 2026

    @ashhart
    Owner

    Status: #1659 is on main. It teaches prompt_radix a lookahead and a planned grid. A store given rules it can't keep
    now fails with error.BadPromptRadixRules in every build. No family turns these rules on yet, so every path on main
    keeps 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.

  15. ashhart commented on Nov 1, 2026

    @ashhart
    Owner

    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=n with TF_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 adds TF_GLM_HOST_CACHE=1 for discrete cards. Both are opt-in; unset, nothing changes. The DeepSeek V4.1 side, reported here by @nullburn, continues on #1625. Closing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions