[jschmied]

What this changes

Today about 25 snapshot files each have their own packing loop. Each one splits its live pieces where the pages
end and copies them, on CUDA, on Metal and on HIP. This PR moves that loop into core and puts GLM on it, on both
backends. Nothing changes in kernels, flags or allocation, and the copies are main's.

Two commits on main 8094ec962:

• bc41d9861 core: one walk for kept prompt states, GLM CUDA's snapshot pages on it

• 980f9a9c5 glm metal: kept-state copies and spills through the core walk

• Issue: Refs Metal: KV pages from the core pool #1293 (step 4, the walk first).

• Stacked on: main 8094ec962, two commits.

• atc overlap --json: one conflict, glm: TF_GLM_KV_LIVE_HALF=1 keeps the live half of the KDA state alone in a kept prompt state #2181 (volfco, TF_GLM_KV_LIVE_HALF) on glm/cuda_snapshot_layout.zig. Whichever lands second rebases; the walk's Part list takes a live-half part the same way. Every other shared file merges cleanly.

What moves to core (core/state_walk.zig, 164 lines with tests)

• Walk(Live, at) works over any live address: CUDA device pointers (Device) or Metal Refs.
  • Part{live, len, span, page, from} is one live region. span and page describe a sequence on pages;
    from skips the pages a page-sharing store already holds.
  • each(parts, ids, kept, ctx, hook) copies every part in order. It splits each one where a live page or a
    kept page ends, then calls hook(live, kept, n).
• Kept is where the kept bytes go: one buffer, or pages leased from a core kv_pool.Pool
  (lease, rounded, page).
• What stays with the family: its parts, the host scalars kept beside the bytes, rank agreement, disk shapes.
• What stays with the backend: the hook, one copy call.

The GLM pilot

• CUDA (cuda_snapshot_layout.zig, cuda_snapshot_pages.zig, kept.zig):
  • Plan's parts are state_walk.Device.Part.
  • Pool.take is lease, and Pool.move is each with a cuMemcpyDtoDAsync hook.
  • TF_KV_SHARE with TF_KV_PAGES=on (prompt rows held as the store's pages) is the descriptor starting at from.
• Metal (snapshot.zig):
  • copy and spill build their parts over Walk(Ref, Ref.at).
  • The kept bytes stay in the state's own buffer, and the hook is the copy_u32 dispatch these functions already
    used.

Checks on this head

• Head commit: 980f9a9c5.

• Lean checker on the base and on this head: 425 and 425, 0 new.

• Copy oracle (GB10): main's Pool.move loop (from cuda_snapshot_pages.zig before this PR) against the core
  walk on 20,000 random plans (paged and unpaged parts, shuffled lease ids, both directions). Every copy is the same
  (dst, src, n) in the same order: 1,601,972 copies, 1,211,205,736,238 bytes, hash 3311f711a06ed690.

• GB10, on 980f9a9c5:

  • test-glm-snapshots-host 20/20, including the byte-for-byte restores through the host-RAM pages and through
    shuffled sequence pages.
  • test-glm-expiry 27/27, test-learn-recovery 503/503 (550/550 in that build).
  • zig build test RC 0; test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon
    RC 0 (server-security 24/24); git diff --check clean; lean 0 new (425 on the base, 425 on the head).

• Mac (M3 Pro), on 980f9a9c5: zig build test 669/669 steps, 9,033 of 9,099 tests passed, 66 skipped, 0 failed.
  check-glm-restore-host passes (main's glm/vision.zig import error is gone since 16f200f7d), and so do
  test-glm-expiry and 36 more test-/check-/compile-glm* host steps, among them scheduling 48/48, disk-charge
  47/47, snapshots 20/20 and window-contract 19/19.

• CI: the no-GPU set above is green on the GB10.

• Not run, and why: neither box has a GLM checkpoint, so the checkpoint gates did not run:

  • the GLM Metal resume gates (resumed == cold, a resident state spilled then restored);
  • tokens with the snapshot pages on and off on CUDA.

  The equivalence rests on the copy oracle and the host restores above.

Receipt

• Environment: TensorFold main 8094ec962 + this PR; Zig 0.17.0; GB10 (sm_121, CUDA 13.0) and an M3 Pro.
• Paths reached: GLM's kept prompt states on CUDA (copy and TF_KV_SHARE pages) and on Metal (copy and spill). The host gates and the copy oracle ran; the checkpoint paths could not.
• Exactness: the copy oracle (same copies, same order); the host restores byte for byte. No served token_sha without a GLM checkpoint.
• Decode speed and prompt speed: n/a. The same copies run in the same order, on the same streams; no checkpoint to time on.
• Licence of new weights: n/a.

Next (order from the RFC)

1. dory, kolibri1, spark25 on CUDA, with join and kv_prompts.Blob as a kept layout.
2. Liquid, once lfm3: serve LFM3-1.7B on CUDA: NVFP4 and FP8 weight-only, MLX 4-bit (stacked on #1513, #1492 and #1503) #1508 and lfm3 CUDA: serve the block FP8 export weight-only and the MLX 4-bit export as stored (stacked on #1508) #1634 land.
3. The five Metal "pieces, blit, wait" families, with one blit hook.
4. Nemotron Metal after metal nemotron: TF_KV_SHARE keeps prompt states as core pool pages (stacked on #1837) #1850, then GLM Metal's kept states as pool pages behind TF_KV_SHARE.
5. Flash Next, then the Qwen families.

Decided on #1293:

• The CUDA hook stays with the backend, so the walk is free of any driver (as here).
• join will be opt-in per family, in step 1 of the order above. A family that turns it on is gated on "bytes equal
  and resumed equal to fresh" rather than "copies equal".

The GLM resume gates are still open: neither of my boxes has a GLM checkpoint. As you suggested, we'll work out
that check from here.

[ashhart]

Thanks @jschmied. The copy oracle, the same 1.6 million copies in the same order over 20,000 plans, pins the walk well. Here are the GLM resume gates, main's binary against this PR's, greedy, with return_token_ids.

At TP=2 on two GB10s:

1. A three-turn chat at the default cache size: each turn's tokens and cached_tokens equal to main's.
2. Resumed equal to fresh at three marks: inside the first KV page, just past a page boundary, and near 32K.
3. The same resumed turns with MTP drafts, equal to "draft": false.
4. Cell 2 with TF_GLM_HOST_CACHE=1.
5. Cell 2 with TF_GLM_KV_SHARE=16, or TF_KV_SHARE=16 if glm cuda: TF_KV_SHARE keeps GLM's prompts in the shared prompt store, and GLM's own store goes #2211 lands first.
6. 1, 2 and 4 streams resuming at once, each equal to its solo reply.
7. States evicted past the budget, a survivor resumed equal to fresh, and the snapshot pool's free pages the same at start and stop.
8. Resumed time to first token and decode against main, with an A/A beside them.

On a Mac with a GLM checkpoint: resumed equal to cold, a spilled state restored equal to fresh, and two streams resuming at once, each equal to its solo reply.

You don't have these checkpoints, so we'll run the GB10 cells; that job is being written. It stays a draft until then. It also conflicts with #2181 on cuda_snapshot_layout.zig, as you say, so whichever lands second rebases.

[ashhart]

Thanks @jschmied. #2211 has since moved GLM's page-sharing prompts to TF_KV_SHARE. This still applies cleanly to today's main with lean at 425, and #2211 didn't touch any line you change. The walk is still GLM's only copy path for kept states. Under TF_KV_SHARE it now copies everything past the shared store's pages, with from coming from heldPages.

Two small things while it waits on the GLM resume gates:

1. The body still describes the shared pages as TF_GLM_KV_SHARE. That variable now stops the server at start, so please say TF_KV_SHARE (with TF_KV_PAGES=on) there.
2. A CHANGELOG line citing "(core: one walk for kept prompt states, GLM on it (CUDA and Metal) #2215, part of Metal: KV pages from the core pool #1293)", even though nothing a user sees changes.

For the record, #2211 changes our 5 Nov gate list in three places:

• The resume marks sit on GLM's keep grid (2,048, 4,096 and 32,768), since the in-page splits can't be restored over HTTP. Your host test and copy oracle cover those.
• The share cell is TF_KV_SHARE=16, with the idle check counting the store's pages.
• The concurrent-resume cell gets a store arm.

#2181 hasn't landed yet, so that conflict is still ahead for whichever lands second.

[jschmied]

Both done: the body now says TF_KV_SHARE with TF_KV_PAGES=on, and the CHANGELOG line is in, citing "(#2215, part of #1293)". It still waits on the GLM resume gates.