Proposal: GLM-5.3-Flash on CUDA takes its streams' KV pages from the core pool, and kept prompt states copy only their rows #1142

#781 gave stream memory one owner in core, and #935 piloted it on Nemotron CUDA. GLM-5.3-Flash on CUDA still keeps a full window for every stream, and each kept prompt state copies that whole window. This proposes the GLM slice of #781's follow-ups. We are starting on it in our fork and will open a draft PR that links here, so your answers to the questions below can still change its shape. Line numbers are on main at d1018b31 (1.1.1).

What happens now

• A stream's sequence is one device arena sized for the window (zig/src/families/glm/cuda_state.zig:269). Each of the 11 MLA layers and the MTP layer keeps capacity rows of 512 bf16 latent values (:286), plus one pooled index key for every 4 tokens. Index keys and gates sit in 4,096-row rings since GLM-5.3-Flash on two CUDA ranks: 16 streams at a 40960-token window #984.
• A slot makes its arena at e.capacity on its first bind (cuda_engine.zig:242). Admission prices every stream at that arena (stream_bytes = seqBytes(), cuda_native.zig:221; native/cuda_memory.zig:149).
• TP does not split the latent: every rank computes kv_a whole (cuda_state.zig:61) and keeps the same rows. A token costs 13,056 bytes on each rank: 12 caches × (1,024 + 64).
• At a 1,048,576-token window a stream takes 12.9 GiB a rank, so TP=3 holds 38 GiB for one conversation. These numbers are computed from the config, not measured.
• A kept prompt state copies the whole arena, whatever its position: snap_bytes does not use at (cuda_native.zig:234), and keepAs copies arena.buf.len (cuda_engine.zig:385). At a 1M window, one kept state takes 12.9 GiB a rank. The cache gets what the streams leave, up to 16 GiB by default (native/cuda_memory.zig:163), so it holds one state at most.
• On one Mac, GLM already shares one MLA cache across its streams (GLM-5.3-Flash on one Mac: the streams share one MLA cache pool, so --parallel 3 serves the full 1M window #639): a request takes a region for its prompt and max_tokens, and kept prefixes hold their blocks. With GLM_POOL=0, kept states copy only their prefixes (GLM: a kept state copies its prefixes out of its slot before the slot writes below it #644), and the two-Mac pair does the same. CUDA has neither.

Proposal

Both pieces below go in one PR. Piece 1 also covers the engine's own full-window sequence, whose kept states stay copies.

1. Kept states copy only their rows (no kernel change)

• snap_bytes(at) returns the fixed part plus at rows. keepAs and restoreState copy each cache's first at rows, the pooled keys up to at, the ring rows a resume reads, and the KDA state. Metal's GLM_POOL=0 path copies its prefixes the same way (GLM: a kept state copies its prefixes out of its slot before the slot writes below it #644, about 18 KB a token there, since Metal keeps index keys and gates for every token). On CUDA the rings keep the copy to 13,056 bytes a token.
• Rows past at are never read (cuda_state.zig:311), so a resumed reply keeps its bits.
• At TP=3, a 100k-token conversation's kept state takes 1.38 GB instead of 13.8 GB (computed).
• Checks, at TP=2 and TP=3:
  • resumed == fresh: tools/zig/reuse_served.py record, replayed with --prompt-cache-gib 0
  • drafted == plain
  • the kept bytes follow at

2. GLM declares a kv_pool.Layout

This is #781's contract, with #935's Nemotron pilot as the model:

• The kernels that read or write the caches take a page table, the way Nemotron's take pt and the elements between pages:
  • writes: tf_glm_lat_write and tf_glm_pool_keys
  • pooled keys: tf_glm_scores
  • latent, for dense chunks and the sparse gather: tf_glm_chunks_16 and tf_glm_chunks_32
• The MTP layer's caches go through the same kernels (cuda_forward.zig:547). A null table keeps contiguous rows.
• The copy store stays, with piece 1's copies, as KV pages: one core owner with reservation leases, Nemotron CUDA pilot (#781) #935 kept it.
• The family reads the same switches as Nemotron: TF_KV_PAGES=0 keeps a window for every stream, and TF_KV_POISON=1 fills taken pages with NaN.

What is new beside #935: ranks

• GLM on CUDA serves at TP=2 and TP=3. Rank 0 admits by the smallest rank's room (peer_room), and the other ranks replay rank 0's steps (families/glm/cuda_follow.zig:13). No step carries page ids today.
• Proposed: every rank builds the same pool from rank 0's plan, and rank 0 sends a stream's page ids with the step that takes them. A follower writes them to its table and never allocates pages on its own.

What stays as it is

• A stream's bits do not depend on its pages. The dense pass reads the same 512-key chunks and merges them in the same order, and the sparse pass gathers the same rows. drafted == plain, resumed == fresh and concurrent == solo at TP=2 and TP=3 check this, with shuffled and NaN-filled pages as in KV pages: one core owner with reservation leases, Nemotron CUDA pilot (#781) #935.
• KV stays bf16 (GLM-5.3-Flash and full GLM-5.3 on the native CUDA engine (tracker) #593's gate).

Out of scope

• Growing a stream past its reservation, and preemption (as Proposal: streams take memory from one core pool as they grow, instead of a full window each at load #781 left them).
• FP8 latent rows, splitting the latent across ranks, and a disk tier.

Questions

1. We plan to build on the core pool. If you would rather have a CUDA port of Metal's shared MLA pool (GLM-5.3-Flash on one Mac: the streams share one MLA cache pool, so --parallel 3 serves the full 1M window #639), please say so early. The port needs no kernel change. It brings contiguous regions with their fragmentation and relocation copies, and a second pool owner beside core.
2. Page size: 512 rows, as for Nemotron?
3. Page ids across ranks: should rank 0 send them with its steps, or should every rank replay the same pool from the same step order?
4. The engine's own sequence: a full window charged once, as Nemotron keeps it, or pages too? GLM on CUDA captures no graphs.

@ashhart, @jschmied: we are taking the GLM CUDA side, and the draft PR will link here.

Written with Weave (My Symbiotic AI)