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 |
…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.
…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.
|
@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. |
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).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.case reads ahead. The build that ran (221d3722c, a local merge of main, not pushed) had the source of jschmied's
commit and commits 1-5; commit 6 changes only a test. The branch has since been rebased onto dd0a9cc (after cuda qwen3_5: --learn-conversations keeps the 27B's and 35B's states on disk (#840) #1358,
cuda flashnext: --learn-conversations keeps Flash Next's states on disk (#840) #1359, cuda: serve --spill-gib and --checkpoint-slots for Nemotron, the Qwen hybrids and Flash Next (#840, split from #1359) #1413, glm cuda: a prompt chunk's indexer rows split across the ranks, their picks all-gathered #1404 and glm cuda: a decode step's select gives each row a 1,024-thread block #1424) and gained commit 7 (the two refusals); that head was built and host-tested, not run
on GPUs.
commit and commits 1-3 on the earlier base, and 2454d86f0, jschmied's commit and commits 1-5 on 3f408b6; neither
pushed), passed the cases they had (no read-ahead case): disk reads in 393-452 ms, first token 0.89-0.96 s. There the
18,432-token state was on disk; here it was skipped at first. See Main's staging permit below.
EXL3 checkpoint.
token_shabd5a8cc0e246.StagingFull) while that one was in flight; see Main's staging permit below.token_sha. Turn 2 also kept the 18,432-token state, which was written this time (8 new blocks, 377 MiB).read ahead; the scatter and rank 1's part follow on the engine thread, untimed). Turn 2's own prefill 0.59 s; its first token 3.8 s, 2.5 s of it queued beside the other request. Sametoken_sha, and the other request finished.token_sha..stpayload damaged, then restarta disk state failed its check (RowsHash). The turn prefills (15.3 s), sametoken_sha, server keeps serving..blkdamaged, then restarta disk block failed its check (RowsHash). The turn prefills (15.7 s), sametoken_sha.--prompt-cache-gib 0, no disktoken_sha: a disk read gives the reply a fresh pass gives.Main's staging permit and GLM's state sizes
ashhart pointed this out on this PR. 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 the run above: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-rows-host,test-cuda-host,test-learn-recovery,test-rows-prefetch-host,test-snapshot-index-host,test-glm-snapshots-host,test-glm-expiry,test-lane-api-host,test-cuda-disk-open-hostandtest-goldenpass: 295/295 steps, 4,132/4,146 tests (14skipped).
zig fmt --checkandgit diff --checkare clean, and the Linux aarch64 build passes.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 446 and 499 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)