Skip to content

[Engine] restore a session file without holding its K/V in RAM - #3251

Open
pspranger-throw wants to merge 1 commit into
Niko1221:mainfrom
pspranger-throw:pr/streaming-restore
Open

pspranger-throw wants to merge 1 commit into
Niko1221:mainfrom
pspranger-throw:pr/streaming-restore

Conversation

@pspranger-throw

@pspranger-throw pspranger-throw commented Oct 19, 2026 •

Copy link
Copy Markdown

Supersedes #1269, which was closed in the 0.2.0-beta tracker reset. That PR cannot be reopened after its branch was rebased while closed (GitHub returns 422 on reopen); the diff and the validation below are unchanged.

Problem and result

The RESTORE path reads the whole snapshot into host memory before it writes anything, so its peak is about the size of the file. On a machine that is also serving the model, a deep session is larger than the RAM left free: a 198,830-token IQ3_S session is a 3.05 GiB file, with roughly 5 GiB available while the model is loaded, and the server answers 503 ("when the RAM to read the file is not there", DETAILS.md). Shallow sessions restore fine, which makes the limit easy to miss.

After this change the restore is two passes over the same file, format v1 unchanged: the read pass validates the whole file without materializing the K/V, and the apply pass re-reads it and copies the K/V into the authoritative pools in blocks of at most 16 MiB. Peak RAM is two 16 MiB buffers plus the running-state metadata, so a session of any depth restores under the same floor.

Related discussion

PR #1269. Niko, 2026-10-17: "It avoids the 503 on deep sessions by restoring in blocks, which is a memory fix and not a speed change, so it is outside this round. It needs a rebase and we would take it as a fix."

The branch is rebased onto main (e88ee39d, head 69fb7912) and retested on 0.2.0-beta. Rebase-only: the diff is unchanged from the reviewed patch (8 files, 561 insertions / 50 deletions).

Approach

  • Read pass: the existing parser and validation run unchanged (header, identity, size, geometry, every count against this session's limits, payload hash), but the K/V bytes are read through one 16 MiB block, hashed, and not retained. Each layer's five part sizes are recorded for the apply pass.
  • Apply pass: re-reads the file and copies each K/V part into its authoritative pool through ConversationKvTarget. It is bound to the read pass: a file whose K/V layer count or any part size differs is refused before that layer or part is applied. The payload hash is recomputed and compared to the trailer at the end.
  • The first applied block is the commit gate: a failure before it refuses cleanly (the live session is exactly as it was); after it the engine fail-stops rather than decode from bytes it cannot vouch for.
  • The residency refill after the pool writes is kept per layer (kv_stream_reset / kv_ring_restore); without it the modes that keep K/V in VRAM read stale slots.
  • session_file_read is untouched, so callers that hold the K/V in RAM behave exactly as before.

Validation

Passed:

  • Full engine build on 0.2.0-beta source rebased onto main (e88ee39d), CUDA 13.4, archs 75;86 (368 targets).
  • conversation_file_test: 223/223 checks. All 8 conversation suites pass under ctest: conversation_cache, conversation_split_failure, conversation_memory, conversation_file, session_save_reclaim, conversation_snapshot, conversation_validation, conversation_transfer.
  • Live drill on 0.2.0-beta (engine 0.1.45.3), IQ3_S, --peer-device 1, 153K session, direct server on loopback:
    • cold prefill: 153,277 tokens in 106.1 s (1,445 tok/s);
    • SAVE: 153,426 tokens / 2.58 GB in 1.57 s;
    • RESTORE (verb): 2.68 s; next turn 153,426 / 153,465 tokens reused (99.975 %);
    • second SAVE: 1.54 s; fresh engine process + RESTORE: 2.79 s, with MemAvailable 5,658 -> 5,308 MiB (the 2.58 GB payload was not materialized); next turn 153,538 / 153,574 tokens reused (99.977 %);
    • no SERR, no FATAL, no refusals. Peer tier active throughout: 3,339 experts / 6.34 GiB on the RTX 2060S.
  • Prior 0.1.45 production drill (same code, older base): 153K session save 1.80 s / restore 2.66 s / 99.977 % reuse; a 198,830-token (3.05 GiB) session restored in 3.4-3.5 s, next turn 99.99 % reused in 1.8 s against ~177 s cold.

Failed:

  • none observed.

Not run:

  • HIP, SYCL and Metal builds, and Windows/macOS; only the CUDA/Linux build was run here.

Limits

  • The file is now read twice (read pass + apply pass). The second read normally comes from the page cache; a cold file on slow storage pays for two reads. Measured restore time stays in the previous band (2.68-2.79 s for 2.58 GB here; 2.66 s on 0.1.45).
  • The commit gate is the first applied block rather than a fully validated in-RAM image; the fail-closed rule is unchanged, but the point of no return moves earlier in the restore.
  • Format v1 is unchanged; no migration is needed.

The RESTORE path materialized the whole snapshot in host memory, so a deep
session could not be restored on a machine that serves it: upstream's own
documented 503 'when the RAM to read the file is not there'.  The restore is
now two passes over the same file, format v1 unchanged:

- the read pass parses and validates the whole file (running state, K/V
  headers, every count against this session's limits, the payload hash)
  without materializing the K/V; it reports each layer's part sizes.
- the apply pass re-reads the file and copies the K/V into the authoritative
  pools a 16 MiB block at a time.  It is bound to the read pass: a file whose
  layer count or any part size differs is refused before that part is applied,
  and the payload hash is compared to the trailer at the end.  The first
  applied block is the gate - before it a failure refuses cleanly, after it
  the engine fail-stops rather than decode from bytes it cannot vouch for.

Peak RAM is two 16 MiB buffers and the running-state metadata, so a session
of any depth restores under the same floor.  Measured on a 198,830-token
IQ3_S session: 3.05 GiB restored in 3.4-3.5 s, next turn 99.99% reused.
The residency refill after the pool writes is mandatory for the modes that
keep K/V in VRAM and runs per layer (kv_stream_reset / kv_ring_restore).

(cherry picked from commit d438ca5)
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