Repository navigation
GLM CUDA: the latent split across ranks (#1313 step 2), with a working reference at TP=2 #1431
Description
Version: native main 8c36359d5. Model: GLM-5.3-Flash at TP > 1. Follow-up to #1313, whose step 2 (the latent split across ranks) was left as design text when the issue closed with step 1 (#1322).
What happens
TP does not split the latent today. Every rank keeps every row: 13,056 bytes a token on each rank, 12,288 of them latent (#1142). At TP=2, the pool holds the same tokens as one Spark's memory would, so two GPUs buy no extra context.
A working reference
The GLM Zig engine in Mia's recipe and the native DeepSeek port both run a split latent ("KV sharding", from my GLM Python patch 0118). It keeps a conversation's latent rows on one owner part. The pieces:
- owner-computed decode;
- a prompt chunk's history fetch for rows a part does not hold;
- prompt borrowing (twins) at the first decode round;
- carries under memory pressure;
- each rank's own view in kept states and the disk tier.
Measured on two GB10s at TP=2, GLM:
- the pool grows 1.51M → 3.13M tokens in the same memory (
CONTEXT=1048576 PARALLEL=8, with a prompt cache and a spill tier); - replies identical to the unsharded server (39 of 39 cases; at 1M, 5 of 5 against the Python engine);
- decode within a few percent.
Questions
#1313's design moves queries rather than rows, on kv_pool's pages with DeepSeek's page book and owner rule. Before writing it:
- Do you want it on
kv_poolpages at any W (GLM-5.3-Flash CUDA at TP=2:--parallel 8at a 1M window is refused (90.48 GiB of windows, 2.33 GiB left), though KV pages would serve it #1313's design), or on the DeepSeekkv/split.zigmodel? - Would you take it in pieces, each off by default until the last, as GLM-5.3-Flash CUDA at TP=2:
--parallel 8at a 1M window is refused (90.48 GiB of windows, 2.33 GiB left), though KV pages would serve it #1313 question 4 asks?- the per-rank latent book and owner;
- owner decode;
- the history fetch;
- borrowing and carries.
- How should it interact with the prompt cache's byte order and with GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations) #1411's disk images? A sharded state's image holds each part's own rows.
Activity
ashhart commentedon Oct 27, 2026
Thanks, @wojo, for picking up #1313's step 2 with a working reference behind it. (1) On kv_pool pages at any W, as #1313 designed: since #1322 they are GLM's default from two ranks, so take DeepSeek's page book and owner rule as the model rather than adding a second pool. (2) Yes, in those four pieces, each off by default and leaving today's bytes unchanged when off; once whole, the split gives replies equal to the unsplit server at TP=2 and TP=3, with the drafted, resumed and concurrent checks, and any new launch goes in the kernel contract. Turning it on by default is a later call, once it is measured faster. (3) Each rank keeps and writes only its own rows, with the world size and owner map recorded; a restore into a different layout is refused by name, and resumed equals fresh across a save and restore. Please stack it on #1411 so the image format is settled once. The per-rank latent book and owner is the first PR.
wojo commentedon Oct 28, 2026
/atc claim
wojo commentedon Oct 28, 2026
@ashhart Thanks for the answer above. Before I write the first PR, here is its shape and the calls that are yours.
The per-rank latent book is designed against main (
ea01ae189) plus #1411 (fbea020a2), so it stacks as you asked.PR 1's shape:
kv_poolgains owner classes on DeepSeek's residue rule: page idsc, c + W, c + 2W, ...are classc,each with its own free stack. A lease takes all its ids from one class. One class is today's pool, id for id.
--parallel 8at a 1M window is refused (90.48 GiB of windows, 2.33 GiB left), though KV pages would serve it #1313'srule).
begincarries it in its unuseddepthword, and every other rank checks each id's residue(
ResidueMismatch).TF_GLM_KV_SPLIT=1a page's pooled keys and latent rows lie in two regions, each with its ownstride (derived from the weights) and its own page table. No kernel or launch is new: the latent launches take the
latent table and stride, the pooled-key launches the others. In PR 1 the latent table holds the same ids and every
rank still keeps every row, so replies, kept states and images equal off's. PR 2 then changes only that table's
contents and the region's size, together with owner decode.
--context 131072 --parallel 4serves (514 pages a class, 257 a window);--context 1048576 --parallel 8isrefused until PR 2 (1,028 pages a class against 2,049, computed). A request longer than a class is dropped by name,
never requeued.
Questions (my recommendation first):
kv_pool(recommended): one arena and one id space, no secondpool. The cost is a change to the pool the other families share, held by a pinned one-class run (same ids, same
order). The alternative is W
kv_poolpools inside GLM: core untouched, but two id spaces in GLM.onmeans in PR 1. Servable, with every row still on every rank (recommended): the book, owner, wire andboth tables run on GPUs before any row moves. The cost: in PR 1
onsaves no memory and caps a conversation atone class, and its start-up line says so. The alternative refuses
onat start until owner decode lands, sonothing of PR 1 runs on a GPU before PR 2.
--parallel 8at a 1M window is refused (90.48 GiB of windows, 2.33 GiB left), though KV pages would serve it #1313's most available pages, lowest rank on ties (recommended), kept in one function. Myreference used fewest live streams first, which evens the owners' decode work. I'd revisit it with PR 2's
numbers.
(1.89x the rows) and 4,864 at TP=3. Replicated: 7,424 (1.76x) and about 5,547, but the head's pass then needs no
exchange. PR 1's layout follows the answer: the head's latent part goes in the latent region or the replicated one.
world,layoutandownerwords toWordsnow(recommended): 48 to 72 bytes, the seal step 12 to 18 words, unsplit images written as
layout 0, owner 0, and arestore whose words name another layout refused by name. PR 4 then only fills them, and the format is settled
once. The cost is one more small commit on GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations) #1411 (about 30 lines with tests) under the cells you have queued. The
alternative leaves GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations) #1411 as it is, and PR 4 gives split images a format 2, refused by name in both directions.
(recommended). Or I hold it until GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations) #1411 lands. It is earlier review against a noisier diff.
glm-kv-pages-8e705835run the TP=3cells (the paged check and served exactness with the switch on, about 45 minutes)? Until then W=3 is host tests.
Unless you'd rather otherwise:
TF_GLM_KV_SPLIT(0/off,1/on), afterTF_DSV41_KV_SPLIT;--learn-conversations) are refused by name with the switch on untilPR 4, which takes borrowing, carries, kept states and images.
What we can run at TP=2 on our two GB10s, drafts on (about two hours):
glm-index-checkat the pooledstride;
glm-paged-checkwith the switch on, on both owners, pages out of order and NaN before use;--context 131072 --parallel 4underTF_KV_POISON=1: concurrent = alone, drafted ="draft": false, resumed = fresh, and the replies against main's;TP=3 needs a third GPU we don't have.
By Robert Wojciechowski (wojo on GitHub, @rwojo on X)