Repository navigation
GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations) - #1411
GLM-5.3-Flash on CUDA keeps conversation states on disk, every rank's part (--learn-conversations)#1411wojo wants to merge 7 commits into
Conversation
…on disk The CUDA Qwen hybrids refused --learn-conversations (cuda_refusal.zig). Their kept states are already host copies (cuda_reuse.zig): the recurrences and convolution tails, then each attention layer's key and value rows by head, in bf16 or fp8 with its scales. cuda_disk.zig lays that copy out as the rows store's image: 16 bytes of words (position, rotary shift, cache format), the recurrent state as the state's own part, and every head's rows as row parts, so 256-token blocks are shared between states and conversations as on Nemotron. A copy read back is a new host copy, restored as any other; one whose words name another position or format is refused. native/cuda_disk.zig's identity is split from Nemotron's probe, and a family that declares `kept_disk` (image, probe, facts) attaches the same store: the checkpoint, kernels, device, driver, the cache format and drafter taps, and a probe state's hash at the first chunk cut. Nemotron's identity bytes are unchanged; identityOf is public, as ashhart#1411 (GLM-5.3-Flash) factors it, and a module that opens the store itself (several ranks' parts) declares attach and is called in place of the probe. The payoff rule, the caps and the write-behind are the core's. A state whose prompt holds images stays in RAM (its image identities are not on disk). The qwen35 module imports the engine_api module the server links (build/cuda.zig), as the disk image names the rows store's types. Co-authored-by: gilby <880013+gilby@users.noreply.github.com>
|
Thanks @wojo. This is a careful port:
On main One change before we run it. After that we'll run it on our GB10 pair with NVFP4:
Two notes:
|
…on disk The CUDA Qwen hybrids refused --learn-conversations (cuda_refusal.zig). Their kept states are already host copies (cuda_reuse.zig): the recurrences and convolution tails, then each attention layer's key and value rows by head, in bf16 or fp8 with its scales. cuda_disk.zig lays that copy out as the rows store's image: 16 bytes of words (position, rotary shift, cache format), the recurrent state as the state's own part, and every head's rows as row parts, so 256-token blocks are shared between states and conversations as on Nemotron. A copy read back is a new host copy, restored as any other; one whose words name another position or format is refused. native/cuda_disk.zig's identity is split from Nemotron's probe, and a family that declares `kept_disk` (image, probe, facts) attaches the same store: the checkpoint, kernels, device, driver, the cache format and drafter taps, and a probe state's hash at the first chunk cut. Nemotron's identity bytes are unchanged; identityOf is public, as ashhart#1411 (GLM-5.3-Flash) factors it, and a module that opens the store itself (several ranks' parts) declares attach and is called in place of the probe. The payoff rule, the caps and the write-behind are the core's. A state whose prompt holds images stays in RAM (its image identities are not on disk). The qwen35 module imports the engine_api module the server links (build/cuda.zig), as the disk image names the rows store's types. Co-authored-by: gilby <880013+gilby@users.noreply.github.com>
|
Thanks @ashhart. I pushed c2b3583:
The seal check you suggested is written but not pushed yet. Noted on #1358: if it lands first, the first commit drops out. |
|
@ashhart I pushed the seal check you suggested, as 0d0084b.
One thing that predates this branch: a refusal inside On main that path already carries Nemotron's If you'd like those refusals to forget the state the way a failed file check does, I can do that in a follow-up and |
…on disk The CUDA Qwen hybrids refused --learn-conversations (cuda_refusal.zig). Their kept states are already host copies (cuda_reuse.zig): the recurrences and convolution tails, then each attention layer's key and value rows by head, in bf16 or fp8 with its scales. cuda_disk.zig lays that copy out as the rows store's image: 16 bytes of words (position, rotary shift, cache format), the recurrent state as the state's own part, and every head's rows as row parts, so 256-token blocks are shared between states and conversations as on Nemotron. A copy read back is a new host copy, restored as any other; one whose words name another position or format is refused. native/cuda_disk.zig's identity is split from Nemotron's probe, and a family that declares `kept_disk` (image, probe, facts) attaches the same store: the checkpoint, kernels, device, driver, the cache format and drafter taps, and a probe state's hash at the first chunk cut. Nemotron's identity bytes are unchanged; identityOf is public, as ashhart#1411 (GLM-5.3-Flash) factors it, and a module that opens the store itself (several ranks' parts) declares attach and is called in place of the probe. The payoff rule, the caps and the write-behind are the core's. A state whose prompt holds images stays in RAM (its image identities are not on disk). The qwen35 module imports the engine_api module the server links (build/cuda.zig), as the disk image names the rows store's types. Co-authored-by: gilby <880013+gilby@users.noreply.github.com>
|
Thanks @wojo, both asks are in, and both go red without their fix. Without the Main has moved since, and the branch no longer merges:
Agreed on handling the bare- After the rebase we'll run it on our GB10 pair with NVFP4. We'll cover the planned cells: eviction and read-back, restart, damaged files and the fresh twin. We'll add a conversation that returns while another stream decodes (read ahead), decode speed while a state stages, and a 60k-token conversation's first save. |
|
@ashhart Rebased onto main (5897e0b) and force-pushed as 96a156d (was 0d0084b).
GB10 pair, EXL3 at TP=2 (a local merge with commit 6's tree; commit 7 is test-only). Every case passes, with
All five disk reads in the run passed the seal check (the three above and the fillers' two). The description's table Two saves in flight. Each prompt here keeps two states, 2,048 tokens apart (16,384 and 18,432 for conversation A; Blocks in flight count twice. Both fixes belong in main's write path, so I haven't made them here: count a block that is still being written once Host. On the rebased head, the rows, CUDA, learn-recovery, read-ahead, snapshot-index, GLM snapshot, GLM expiry and |
|
Thanks @wojo for the rebase, and for the careful write-up of what the staging permit does to GLM's states. All three asks are in:
On a merge with current main, these all pass: On your question: please send both write-path fixes as a follow-up PR on main, not here. You're right that they belong in main's write path. Every family whose disk states go through the I/O worker shares that path, so each fix needs its own test that fails before it:
The second also changes when the engine stages, so it needs a decode-speed cell of its own. This PR doesn't depend on either fix: GLM's disk states are opt-in, and a skipped save costs a later first token, never a different one. Next we'll run it on our GB10 pair with NVFP4:
We'll also check that a Nemotron disk state written by main still reads back after the |
…on disk The CUDA Qwen hybrids refused --learn-conversations (cuda_refusal.zig). Their kept states are already host copies (cuda_reuse.zig): the recurrences and convolution tails, then each attention layer's key and value rows by head, in bf16 or fp8 with its scales. cuda_disk.zig lays that copy out as the rows store's image: 16 bytes of words (position, rotary shift, cache format), the recurrent state as the state's own part, and every head's rows as row parts, so 256-token blocks are shared between states and conversations as on Nemotron. A copy read back is a new host copy, restored as any other; one whose words name another position or format is refused. native/cuda_disk.zig's identity is split from Nemotron's probe, and a family that declares `kept_disk` (image, probe, facts) attaches the same store: the checkpoint, kernels, device, driver, the cache format and drafter taps, and a probe state's hash at the first chunk cut. Nemotron's identity bytes are unchanged; identityOf is public, as #1411 (GLM-5.3-Flash) factors it, and a module that opens the store itself (several ranks' parts) declares attach and is called in place of the probe. The payoff rule, the caps and the write-behind are the core's. A state whose prompt holds images stays in RAM (its image identities are not on disk). The qwen35 module imports the engine_api module the server links (build/cuda.zig), as the disk image names the rows store's types. Co-authored-by: gilby <880013+gilby@users.noreply.github.com>
|
Thanks @wojo. #1358, #1359 and #1413 have now landed. #1358 brings your first commit, @jschmied's rows change, onto main as 9330cc8, as you expected. Between them they also change Could you rebase onto current main? The first commit should drop out. Nothing else is needed from you. Once it's rebased, we'll run the cells we listed on our GB10 pair. |
|
@ashhart Rebased onto main (1b9a987) and force-pushed as fbea020 (was 96a156d). The first commit dropped out:
The two write-path fixes will come as their own PR on main, as you asked. Moving GLM onto |
|
@ashhart #1404 and #1424 landed after my rebase, so #1411 conflicted again. It's rebased onto main ( The two write-path fixes you asked for are #1548 (with a TP=2 decode cell), and the scatter refusal is #1543. |
|
Hi @wojo, a status update on #1411. We read the full diff. With The run is still owed for two reasons. Every cell needs two GB10s at once (TP=2). And main refuses Nothing is needed from you right now. Thanks for splitting the rows-store fixes out into #1548. |
|
@wojo, main has moved and this now conflicts with it. Could you rebase onto current main (a force-push is fine)? Then we'll run the flag-on and flag-off cells on our two GB10s at TP=2. |
|
@ashhart Rebased onto main (
|
|
Thanks @wojo for the rebase, and for the One thing before it lands. The receipt table ran on an older base, and this head changes If #2207 lands first, |
…re each byte sits in its lease The rows store's view of a GLM kept state, without the store: the words (format, position, head rows, ring bases, KDA copy bits), the KDA prefix, each DSA layer's latent rows, scales and whole pools by position, its rings and pools in progress as the state's own, the head's caches at a whole verify window (rows past the head's own read as zeros), then the tail. Copies run a page of the lease at a time. Nothing is wired to --learn-conversations yet.
|
@ashhart The table is rerun at TP=2 on our two GB10s. Since #2207 landed after your note, I rebased onto main first ( Setup: Mia's EXL3 checkpoint
|
…state on disk Rank 0's engine thread leases, fills, reads and seals each rank's copy through the follow link (cuda_disk_ranks: blank, put, take, seal, prefixes). The store's image is rank 0's image followed by each other rank's prefix, the one part a rank holds alone; the rest are the same bytes on every rank, checked at start by the identity probe. A blank state is refused by restore until its seal.
…sk state into that storage A family whose kept states live in preallocated storage it reuses (GLM's snapshot pages on CUDA) and cannot trim never had a disk read chosen: the selection counted its spare storage as fixed and its kept states as not reclaimable, so no state ever fit. Such a family's room for a read is now its spare storage plus what evicting unborrowed states frees, and a read is ready once the family says its spare storage takes the state, as a keep's admission already asks. Families with trim, or without reused storage, are unchanged.
…t, as before The disk read's early admission for a family whose spare storage takes the state now applies only to a family without trim, as the selection's capacity already does. A family with trim, such as Flash Next's or Qwen3.8-27B's pools on Metal, makes room for a read by trimming its spare storage and evicting, exactly as it did before, and a new test pins that its read stays within the budget when its reuses answers yes.
…he state's rows An index ring based at `base` holds the rows from `base` up to the state's position, and the image keeps at most a ring's rows of each. The seal now checks the words' main base against the state's position and the head's base against its rows, so an image whose words name a ring that starts past them, or too far back to be held, is refused before the blank state is made whole.
…/O worker GLM's store uses The write-behind gave GLM's store its own I/O worker, while the grid test wrote on the test's thread (sync_writes). It now runs on each, drained by write.drain, so a worker path that loses the grid's chunk starts fails the test.
…rn-conversations alone keeps its states Since ashhart#1413 CUDA serves --spill-gib and --checkpoint-slots, and --learn-conversations admits GLM, so --learn-conversations --spill-gib (or --checkpoint-slots) opened GLM and ignored the flag: its disk store writes every rank's part between passes and stages nothing at eviction. GLM now refuses both by name. Its store's spill stays off, since an eviction would gather on rank 0 mid-pass while the other ranks wait in the gate's all-reduce. docs/api.md lists every CUDA family that keeps states on disk.
Fixes #1312.
GLM-5.3-Flash on CUDA now keeps conversation states on disk through the rows store, with
--learn-conversations.After an eviction or a restart, the next turn reads the state back instead of prefilling the whole conversation again.
At TP>1 every rank's part goes to rank 0's disk (option (a) on #1312), so a state is read back on every rank or on
none of them.
Commits
jschmied's grid commit, which this PR carried first, landed with #1358 as 9330cc8, so the branch no longer carries it.
families/glm/cuda_disk_image.zig). It has the image'sparts at a position, maps each image byte to its place in the snapshot lease, and copies between host and lease.
The image is the snapshot plan's bytes in their order, so a disk restore writes exactly the bytes a RAM restore
reads.
--learn-conversationskeeps every rank's part of a kept state on disk.cuda_disk_ranks.zigadds one new follow step kind,disk, with opsblank,put,take,sealandprefixes. Rank 0's engine thread runs every op, as it does every other step. The bytes travel as the step'sids, a link buffer at a time.
cuda_snapshots.zigaddsblankAsandseal. A blank state is refused byrestoreuntil it is sealed.native/cuda_glm_disk.zigis therows.Imageadapter. The store's image is rank 0's image followed by eachother rank's KDA prefix, the only part a rank holds alone. The latent rows, rings, pooled keys, head and tail are
the same bytes on every rank.
refuses to open the store (
RanksDiffer) if they differ, so the "same bytes" premise is checked at every start,not assumed.
drafter, every rank's prefix size and the probe's hash, through main's
identityOf(from cuda qwen3_5: --learn-conversations keeps the 27B's and 35B's states on disk (#840) #1358), with GLM'sprofile
.{ .family = "glm5_next", .backend = "cuda" }since persisted states bind to the resolved profile values #2207.reuses) but cannot trim.The restore selection counted those spare pages as fixed, and the kept states as not reclaimable, so no disk
read ever fit.
A read is ready once
reusessays the spare storage takes the state, which is the question a keep's admissionalready asks.
trim, or without reused storage, are unchanged. The existing pooled test passes as before.readyreturn early whenever
reusessaid yes, which also reached Flash Next and Qwen3.8-27B on Metal (they define bothreusesandtrim). The early return now needstrim == null, asrows_select.capacityalready did, so a trimfamily's
readyis main's again.capacitywas already main's for them, and the branch nolonger changes
rows_storeorrows_write.trimandreuses, whose disk read takes new storage (as Qwen3.8-27B's disk blankdoes) and whose
reusesanswers yes after an eviction. With the gate it trims and stays within budget (held 205plus spare 25 of 240). Without it the read skips the trim (205 plus 120 of 240) and the test fails.
check on this PR. An index ring based at
baseholds the rows frombaseup tolen: the state's position forthe main ring, the head's rows for the head ring. The image keeps at most
ring_rowsof each ring.sealnowrequires
base <= lenandlen - base <= ring_rowsfor both rings. An image whose words (the bases and lengths itstores) fail is refused with
DiskWordsbefore the blank state changes, on the same path as the existing head andcopy-bit checks.
base + ring_rows, sobetween passes, where states are kept, both bounds hold.
passes, exactly
ring_rowsback. Without drafts the head's base and rows are both 0.worker, while the grid test wrote on the test's thread (
sync_writes). It now runs on each, drained bywrite.drain. WithstartsBelowdropped on the worker's path only, the worker case fails (rows_test.zig:277,the state's starts) and the sync case passes.
--spill-giband--checkpoint-slotsare refused. Since cuda: serve --spill-gib and --checkpoint-slots for Nemotron, the Qwen hybrids and Flash Next (#840, split from #1359) #1413 CUDA serves both flags, and--learn-conversationsadmits GLM, so on this base--learn-conversations --spill-gib(or--checkpoint-slots)opened GLM and ignored the flag. GLM's store writes every rank's part between passes and stages nothing at eviction.
GLM now refuses both by name, with tests (with the refusal removed, the
--learn-conversations --spill-gibcasefails at
cuda_refusal.zig:217). Its store's spill stays off: an eviction would gather on rank 0 mid-pass while theother ranks wait in the gate's all-reduce. docs/api.md now lists every CUDA family that keeps states on disk.
Receipts: two GB10s at TP=2
Setup:
--context 131072 --parallel 8 --prompt-cache-gib 6 --learn-conversations --learn-gib 30, greedy with thinking off.this run on this PR.
covers 18,432. Every row is turn 2's reply, compared by
token_sha.EXL3 checkpoint.
read ahead"draft": falseTF_GLM_KV_HOST=1.stdamaged, restartRowsHash), prefilled.blkdamaged, restartRowsHash), prefilled--prompt-cache-gib 0, no disk)its own 352 MiB in 1,678 ms. The 18,432-token state took 8 new blocks and its own 354 MiB in 592 ms.
TF_GLM_KV_HOST=1: the profile line readsregistered overrides: TF_GLM_KV_HOST=1 (identical, env). The stateswritten without it were found under the same identity, read back through the host-page path, and the next turn
resumed from RAM.
--learn-conversationsunset, same flags otherwise, fixed prompts so the two trees get identicalrequests): every reply has the same
token_shaon branch and main.Earlier runs (a 5897e0b-based build with read-ahead, and two with main's synchronous writer) passed the cases they
had; their tables are in this PR's history.
Main's staging permit and GLM's state sizes
ashhart pointed this out on this PR. (In the run above, on a757420, A's first save wrote both states. The figures
below are from the earlier 5897e0b run.) On main a save is staged whole between rounds, rank 1's
takesteps included,under a 1 GiB permit (
staging_bytes). A save that doesn't fit is skipped (imprint: saving N tokens skipped (StagingFull)) and not retried.GLM's states on this pair (TP=2,
--context 131072, MTP drafts), from that run:rank's KDA prefix, the index rings (capped) and the MTP head's caches; the 1 KiB a token is the head's caches.
What follows:
MiB) was written three times here; at 55,296 tokens the logged sizes add up to about 1,023 MiB before metadata,
within a MiB or so of the permit. Past that the first save is skipped, and so is every later save of that
conversation, since its blocks never reach disk.
states 2,048 tokens apart. The second was staged while the first was still being written, and was skipped all 5
times both were new (A's 18,432 twice; the 55,296 of both fillers and of the 55,942-token request). Its charge also
counted the blocks the first was still writing, since the store learns of them only when the first is published:
A's 18,432-token save was charged 564 MiB, of which 377 MiB was new, so the pair needed 1,103 MiB where 916 MiB
would have fit. A's 18,432-token state reached disk on its next turn (8 new blocks, 377 MiB). Until then the read
after eviction resumed from 16,384 tokens: first token 2.6 s, against 0.9-1.0 s with both states on disk.
extrapolated from a 131,072 context, not run), so at a 1M window no save of a longer conversation fits.
--spill-giboff, as for GLM, the stop's flush also stages the next savewithout waiting for the worker, so saves still queued at a stop can be skipped the same way.
TENSORFOLD_ROWS_SYNC=1selects main's synchronous writer, which has no staging permit; the earlier runs used it.Host tests (Zig 0.17.0, an M4 Mac)
test-cuda-host, learn recovery, snapshot index,GLM snapshots, GLM expiry, lane API, CUDA disk open, GLM disk charge, golden, and the GLM startup/pages/contract/copies
suites pass.
zig fmt --checkandgit diff --checkare clean, and the Linux aarch64 build passes.exl3_load's "a scan that fails an allocation frees every group prefix it held". It fails thesame way on main a757420 when the checkout's path is as long as ours, and passes from a shorter path.
bit no layer has, or ring bases that do not hold the state's rows (commit 5), and seal setting the plan's length.
if either bound were strict or if the main ring were checked against the head's rows.
rows_test.zig:277withstartsBelowdropped on the worker's path only; the synccase passes.
cuda_refusal.zig:217with GLM's refusal removed.zig build testfailed only in the 55 Metal library builds, on base and branch alike (no Metaltoolchain on this Mac). Not rerun since.
python3 tools/zig/lean_check.py: 425 problems on main and 425 on the branch (adds 0).Not covered yet
--spill-giband--checkpoint-slotsfor GLM on CUDA. Refused for now (commit 7). Taking them needs aneviction's staging moved out of the pass at TP>1, and a GPU run.
context). That is enough for the 428-548 ms reads above; a larger transfer buffer is a follow-up if it
is ever the bottleneck.
cover NVFP4.
@jschmied: your grid commit landed with #1358, thank you. The native side still opens GLM's image and store explicitly,
beside
keptDisk. Moving GLM ontokept_diskneedsengine_apiin the glm_cuda module, the prefix gather moved intoprobe, and commit 7's spill refusal kept by name (or a spill opt-out on the declaration), so I'd do it as a follow-up.By Robert Wojciechowski (wojo on GitHub, @rwojo on X)