Skip to content

cuda: Gemma 4, Dory and Spark2.5 keep prompt states on disk, on one shared kept image (#840) - #3135

Open
jschmied wants to merge 5 commits into
ashhart:mainfrom
jschmied:disk-cells-next
Open

jschmied wants to merge 5 commits into
ashhart:mainfrom
jschmied:disk-cells-next

Conversation

@jschmied

@jschmied jschmied commented Nov 13, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

--learn-conversations and --spill-gib now open on Gemma 4 (E4B, 12B, 31B), Dory and Spark2.5 on CUDA. After a restart
or an eviction, a conversation resumes from disk instead of prefilling again. Both flags are opt-in: with neither set,
nothing changes and the served defaults stay as they are.

  • Shared code: zig/src/native/cuda_kept.zig holds the image of states kept in device buffers: the words, the copies,
    the probe and the key facts. Each family declares a view: Gemma 4 69 lines, Dory 54, Spark2.5 58 (without tests).
    Flash Next's family-local Kept moves onto it. Its probe bytes, seed and policy (Persistence: the key gaps #2163 leaves, for every family (closes #2131) #2287's resolved-policy key) are
    unchanged, and keptKey reads the policy from the view.
  • Keys: the shared identityOf / keptKey, extended by Kept.prompted with the tokenizer's and template's bytes and
    a drafter's checkpoint, plus the family's KV form: loops, ring, row, paging and the drafter's context. A stale key opens
    another namespace and runs a cold pass. Flash Next's key moves once (for the prompt files), so its old states miss once.
  • Host permit, from core: a Gemma 4 31B state at 9.3k tokens is 1.5 GiB, which is past the 1 GiB default, so every save
    was skipped as StagingFull. Kept.longest passes the served window as rows.Options.longest (Qwen3.8-27B on a Mac: the disk staging permit takes one whole state at its context #2388), so core prices
    one whole state within the RAM store's budget. It goes through native/cuda_disk.zig's open, beside cuda disk: a family's host staging permit for kept states #2337's
    staging hook. This PR adds no permit code of its own. A family with no window (Flash Next) keeps the default.
  • Gemma: a state over an image stays off disk, because only RAM can tell which image it was. The DFlash context is a
    second buffer.
  • Tails: each state keeps its ring tails in its own file, because the image is ordered by position, not by
    p % window. That is 176 MiB a state on Dory, so --learn-daily-gib buys fewer states than the block rows alone would.
  • TF_KV_SHARE together with --learn-conversations is already refused by name in native/cuda.zig. The refusal is
    unchanged here, so a page-store state never reaches the disk image.

Core vs backend: the store, writer, permit pricing and keying are core and shared, and this PR leaves them unchanged.
The CUDA-only parts are the shared image (native/cuda_kept.zig) and the three family views. Metal is untouched.

Refs #840. Overlap with open PRs on the same files: #2903 and #2956 (kotevcode) touch native/cuda_disk.zig's identity,
which is untouched here apart from keptKey's view lookup. #2392 (ours) touches gemma4/cuda_snap.zig; the hunks are
disjoint. Nobody else has opened Gemma 4, Dory or Spark2.5 on disk (searched open and closed PRs and main).

Proposed CHANGELOG line (for ashhart to place): CUDA: Gemma 4, Dory and Spark2.5 keep conversation states on disk with --learn-conversations and --spill-gib; Flash Next's disk states move onto the same shared image (#3135).

Checks, head 321c3d6 on main 6c5cef8

The PR head e213886 is 321c3d6 with only the PR number filled into docs/optimization-ledger.md; no code differs.

  • x86 (tfci): the CI set (test test-golden test-kernel-contract test-server-security test-pull-security fuzz-cortheon check-release-paths), check-native-cuda-host compile-cuda-tests, tools/pr_check.py, zig fmt --check,
    git diff --check: all 0. Lean: 345 on the base, 345 on the head, 0 new.
  • GB10: test-cuda-host and -Dtarget=aarch64-linux-gnu compile-cuda-tests: 0.
  • Host tests: test-cuda-kept-host (17 tests) covers the round trip over host memory, and a stale key or a new template
    or drafter missing through the real store. It also covers a partial copy failure draining what it queued, and core's
    permit taking a whole state at the window within the RAM budget (and capped by it). test-cuda-disk-open-host,
    test-disk-identity-host (Flash Next's policy key through the view), test-dory-host, test-spark25-host and
    test-gemma4-cuda-host run each family's parts against its own copy order, and check that the key moves.
  • No new kernel or device check, so no planted control applies. No new refusal: two existing messages now name the
    three families, and the test CUDA kept: Gemma 4, Dory and Spark2.5 take the store checks that a family with no
    view still refuses.

Receipt (head 321c3d6, binary md5 48b6b133; main 6c5cef8, md5 76b7f3ed)

GB10, sm_121, clock capped at 2,300 MHz, Zig 0.17.0. The page cache was dropped before each server. Flags:
--parallel 1 --context 16384, a 9.3-9.9k token prompt, greedy (temperature 0, seed 1). Each arm is a fresh server.

Exactness: all three turns gave equal token ids across main, a cold run with the flag on, and two resumed restarts
(4 models x 4 runs x 3 turns). Turn 1 after each restart resumed from disk (cached 9306-9779 tokens).

Turn after a restart, A-B-B-A (cold, disk, disk, cold), the mean of each pair:

model cold from disk change written
Gemma 4 E4B 2.013 s 0.123 s -93.9 % 414 MB
Gemma 4 31B 11.747 s 0.732 s -93.8 % 5.84 GB
Dory 2.563 s 0.345 s -86.5 % 1.65 GB
Spark2.5 4B, 8-bit affine 2.192 s 0.211 s -90.4 % 715 MB

Both disk arms beat both cold arms on every model (4 of 4). The replies' sha matched across all four arms.

Flag off against main (A-B-B-A: main, head, head, main, 256 tokens), head vs main means:

model TTFT decode tok/s sha
E4B 1.934 vs 1.929 s (+0.3 %) 102.63 vs 102.70 (-0.1 %) equal
31B 11.128 vs 11.127 s (+0.0 %) 21.38 vs 21.40 (-0.1 %) equal
Dory 2.318 vs 2.315 s (+0.2 %) 105.15 vs 105.29 (-0.1 %) equal
Spark2.5 2.060 vs 2.062 s (-0.1 %) 84.08 vs 83.89 (+0.2 %) equal

With the flag off, the head is slower than main on 3 of 4 TTFT cells and 3 of 4 decode cells, all by at most 0.3 %,
which is inside the run-to-run spread. With the flag off, the binary differs only by the unused path.

Known cost with the flag on. The turn that writes a save pays for it. The first cold turn costs more than with the
flag off: +4.3 % on E4B, +5.5 % on 31B, +10.7 % on Dory and +6.4 % on Spark2.5 (the ABBA cold arms against the off
arms). A warm turn that saves new states costs more as well: 31B 0.482 vs 0.335 s, Dory 0.190 vs 0.154 s. The cause is
that lane_host_rows.answer steps the writer between rounds, so a save's gather (DtoH plus a synchronize) stages on
the engine thread during the pass. The counterfactual was measured on 11-06 on the pre-rebase head. Dory with
--learn-daily-gib 0, which stops a save before its gather and keeps the store and the probe, came out level with
flag off (2.279 vs 2.284 s, six arms). 31B turns whose states were already on disk wrote nothing and matched main
(0.337 vs 0.335 s here). That writer schedule is core code shared with the sibling families (Qwen hybrids, Flash Next,
Nemotron) and is not changed here. Moving the gather off the engine thread would be a separate core PR. With the flag
off, nothing is paid.

Not covered on a GPU: Gemma 4 12B, Flash Next (host tests only; its key moves once), DFlash states on disk, and
--spill-gib eviction (host only). Concurrency 1/4/8 and seeded sampling were not run: all arms used --parallel 1
with greedy decoding. The store's concurrency paths are core's and are unchanged; ask and I will run c=4/8 arms.

Receipts: runs/r-{e4b,31b,dory,spk}/results.jsonl plus each arm's server log (md5 of the binary at its top). Tool
dc.py (ident: main, cold, resumed x2; abba: cold vs disk, then flag off vs main).

Ledger: docs/optimization-ledger.md row 8, gem4, spk25 and dory set to open.

…shhart#840)

native/cuda_kept.zig holds what Flash Next's kept_disk held: the image's words,
the copies, the probe. A family declares a view of its layout. Flash Next moves
onto it with its probe bytes and seed unchanged; its key gains the prompt files.

Gemma 4 keeps its states on disk with --learn-conversations and --spill-gib.
Keys bind the tokenizer's and template's bytes, a drafter's checkpoint and the
KV form. A state over an image stays off disk: RAM alone tells which image.
A stale key opens another namespace, a cold pass.
Their views over native/cuda_kept.zig: full caches as block rows, ring tails in
the state's own file, the loop count, ring, row and paging in the key.
…ough core's bound (ashhart#2388)

A Gemma 4 31B state at 9.3k tokens is 1.5 GiB, past the 1 GiB default: every save
was skipped as StagingFull. Kept passes the served window as rows.Options.longest,
so core prices one whole state within the RAM store's budget. Families without a
window (Flash Next) keep the default permit.

This branch has not been deployed

No deployments
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.

1 participant