Skip to content

glm cuda: TF_GLM_KV_SHARE keeps prompts as the pool's pages in prompt_radix's store - #1732

Merged
ashhart merged 4 commits into
ashhart:mainfrom
Bizuayeu:pr/glm-prompt-pages
Oct 31, 2026
Merged

ashhart merged 4 commits into
ashhart:mainfrom
Bizuayeu:pr/glm-prompt-pages

Conversation

@Bizuayeu

@Bizuayeu Bizuayeu commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

TF_GLM_KV_SHARE=n with TF_KV_PAGES=on keeps GLM's prompts as the pool's pages in prompt_radix's store, the store TF_QWEN_KV_SHARE uses, with at most n states past them. Unset, every path stays as it was. This is the GLM half of #1510's option B, on the rules #1659 landed. On GB10 at TP=3 with the 1M window, the pool leaves main's flat store about 450 MiB at the default prompt cache, less than a 64k state's 922 MiB, so a message sent again prefills from the start:

unset (as main) TF_GLM_KV_SHARE=16
64k message sent again (s) 36.39, 36.46 1.73, 1.75
128k message sent again (s) 72.14, 72.27 2.05, 2.06
64k sent twice at once (s) 70.1 to 73.1 2.27 to 2.28
128k sent twice at once (s) 111.5 to 144.7 3.35 to 3.45
136k of token ids sent again, ttft (s) 83.3 to 83.5 0.69 to 0.71
decode (tok/s, bench_openai) 51.13, 51.06 51.09, 51.10

At TP=2 with the 524k window main's flat store has room, so unset resumes the chats too and the two match. Only token-id prompts gain, as the flat store keeps no marks for them:

unset (as main) TF_GLM_KV_SHARE=16
64k message sent again (s) 2.10, 2.11 2.08, 2.09
128k message sent again (s) 2.40, 2.39 2.36, 2.37
64k sent twice at once (s) 3.21 to 3.28 3.17 to 3.23
128k sent twice at once (s) 3.50 to 3.51 3.43
136k of token ids sent again, ttft (s) 100.8 to 101.0 0.85, 0.86
decode (tok/s, bench_openai) 40.21, 40.13 40.14, 40.20

What this changes

Three commits in families/glm, with their build and host-test glue, on the rules #1659 gave prompt_radix.

  • 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, as 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 Qwen. Set without TF_KV_PAGES=on, the engine refuses before the weights load.
  • Below 4 × --parallel states the server warns at start and names the value to set, as TF_QWEN_KV_SHARE does.

Choosing n. n bounds the states the tree keeps, whatever pages are free, and start-up takes n times the state at the window's last mark for their snapshots. More conversations going back and forth than n miss: on an earlier head of this branch, six 8k chats sent twice at TP=2 resumed none at n=4 (5.5 s each), and all six at n=16 and on the flat store (1.2 s each), and at TP=3 n=16 resumed the six in 1.04 s against 4.65 s fresh. Over the launches above, the lowest MemAvailable by rank was 9, 10 and 15 GiB at n=16 against 11, 11 to 12 and 17 GiB unset at TP=3, and 13 and 14 GiB against 11 to 12 and 12 GiB at TP=2, where the flat store takes 4.7 GiB. Kept pages do not hold streams back: on that earlier head at n=4 and TP=2, with four 128k chats filling the tree to 992 of 1,028 pages, a fifth started at once and matched unset's ttft (86.43 and 86.48 s against 86.67 s), a chat the reclaim had evicted missed cleanly when sent again (cached 0), and decode with the tree that full stayed within the runs' spread (four launches against four, −0.15%, the ranges overlap).

#1411. Its disk images hold a snapshot's bytes, and a state here leaves its pages out, so the two do not combine until an image carries the pages. They also conflict in cuda.zig and cuda_follow.zig: whichever lands second rebases.

Receipt

  • Environment: TensorFold fb84de61c, these three commits on main 19372fb36; de41646f9 is the same on d80e5f2ef but for the two lines that met glm cuda: at three ranks or more, TF_GLM_ROCE_MULTI=1 sends the small all-gathers over fabric/tp's multirank RoCE #1524's. GLM-5.3-Flash XSFP4 at 88736f916f89, PARALLEL=4, KV pages on, the drafts on; three GB10 in a ring on NCCL (TP=3, 1M window) and two (TP=2, 524k window). Launched unset, on, on, unset at each TP.
  • Exactness: every reply's tokens are the same in all launches at each TP (rank 0's done-line hash): 140 replies at TP=3, 18 of them resumed, and 140 at TP=2, 32 resumed. tools/bench_concurrent.py --alone --serial on two prompts at 1 and 4 streams, four launches at each TP: 40 of 40 equal alone and 16 of 16 equal "draft": false at both. One more launch at each TP, n=16 under TF_KV_POISON=1 on every rank, gave the same hashes on its 35 replies, 9 of them resumed, and its concurrent cells were equal too, so no rank read a row before its write.
  • Decode speed before and after, from tools/bench_openai.py: the tables above. Eight more launches at TP=2 on an earlier head of this branch (n=4 against unset, four of each, 40.424 against 40.349 tok/s with nothing kept, 40.195 against 40.140 with a 64k state held): −0.19% and −0.14%, with standard errors of 0.12% and 0.13%. The ranges overlap, and holding a state does not change the gap.
  • Prompt speed before and after, the cold first request's ttft_s from tools/bench_cache.py --require-cold on 136k tokens: 83.40 and 83.46 s unset, 83.39 and 84.13 s on at TP=3; 101.01 and 100.82 s unset, 100.94 and 100.91 s on at TP=2.
  • Tests run (Linux x86_64, at de41646f9): test-kv-pool 28/28, test-qwen35-share-faults-host 14/14, compile-cuda-tests 10/10, compile-config-server 7/7, zig fmt --check, and tools/zig/lean_check.py with nothing new (426 against main's 426). test-cuda-host (3,987 tests passed, 9 skipped) and zig build test (1,808 passed, 18 skipped) pass but for cuda_cost_identity's stable-reads drafter test, outside these files: run 40 times, it fails 11 times here and 11 times on main, as two same-sized writes on this Linux often share a 4 ms clock tick and leave one stamp. Under CPU load zig build test passed whole, 445/445 steps.
  • Not run, and why: Metal and HIP devices, since the change is GLM's CUDA path.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (A resume adopts pages and restores a state before the pass. The tokens after it go through the rounds as before, with the drafts on.)
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (bench_concurrent.py --alone --serial at 1 and 4 streams and the done-line hashes, at TP=3 and TP=2, on and unset, as above.)
  • No precision traded for speed. If the bits change, the description says which sums change and why. (No bits change: a resume restores the rows a fresh pass wrote, and the hashes match.)
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (No GPU kernel is added or changed.)
  • Prompt processing is no slower than the last release. (Unset, the paths are main's. With it on, the cold ttft is 83.39 and 84.13 s against 83.40 and 83.46 s at TP=3, and 100.94 and 100.91 s against 101.01 and 100.82 s at TP=2, as above.)
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. (CUDA on GB10 at TP=3 and TP=2, host tests on Linux. Metal and HIP devices not run, as above.)
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (On main they do not build, as openWith, Plan's shared and Prefix are not there. With their property taken out here, one at a time, the test named for it fails: the poison reaching the adopted pages on rank 0 or on a follower, a follower taking more kept pages than its table, a refused resume keeping the cache's pages, a snapshot ignoring shared, the store told one page too many, a lookahead of 0, and the warning's factor at 3. The last run with all of them back passes. None needs a device.)
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

@russellmm

Copy link
Copy Markdown
Contributor

As I am building Mia GLM-5.3-Flash model to work with my 2 RTX6000s I pulled this PR in during my testing. Working fine for me.
Screenshot From 2026-10-30 20-48-56

Bizuayeu and others added 4 commits October 31, 2026 18:13
…_radix's store

With TF_GLM_KV_SHARE=n and TF_KV_PAGES=on, GLM keeps a pass's states at its marks 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), in the
core's radix store, at most n states past the pages. A state holds its prefix's pages by reference; its snapshot keeps
only what lies past them (Plan.shared): 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, since the pooled keys are
written in whole pools from the step's head, so nothing is copied on write. A resume adopts the deepest state's pages at
the head of the stream's lease (Pages.openWith, where TF_KV_POISON's NaN spares them) and restores the rest through the
prompt pass's Resume. With it on the server makes no flat store for GLM (own_prompt_cache), as for the Qwen hybrids.
One rank here; unset, every path is as before.
…s pages, its snapshots step their absence

The tree and the pages stay rank 0's. The begin step carries how many of the table's first pages are the cache's, so
another rank's map leaves their rows out of TF_KV_POISON's fill and refuses a count past the table. The snapshots step
carries whether states leave the pool's 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. The one-rank refusal goes.
Six 8k conversations sent twice at --parallel 4 resumed none at n=4 and all at n=16, with pages to spare: the store
limits states as well as pages, and each new conversation evicted the state the next turn wanted. The server now logs a
warning below 4 x --parallel and names the value to set, as ashhart#1357 does for TF_QWEN_KV_SHARE.
@ashhart
ashhart force-pushed the pr/glm-prompt-pages branch from 3495cda to 682cbd2 Compare October 31, 2026 18:13
@ashhart
ashhart merged commit ab4b5d8 into ashhart:main Oct 31, 2026
9 checks passed
@ashhart

ashhart commented Oct 31, 2026

Copy link
Copy Markdown
Owner

Thank you @Bizuayeu. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants