[LuckyKo]

Problem

When running Strata with --layer-split auto across multiple GPUs, the engine divides model layers into stages (one per GPU). The existing session save/restore (--slot-save-path) only handled a single-stage (single-GPU) layout: it saved/restored one set of running state and K/V. On a multi-GPU split, the save was either refused or produced an incomplete file that couldn't be restored.

This made --slot-save-path unusable for the most common multi-GPU configuration, which is exactly where long-context session persistence matters most (the RAM cost of re-reading a 30K+ token prompt).

Solution

Introduce format v3 for session files that carries per-stage data. Version 2 is reserved for PR #2247's BF16 state codec; this build rejects it explicitly at the version gate.

• Write: put_payload serializes stage 0 (as before), then — only when stage_images is non-empty — a count followed by one block per later stage. Each block holds that stage's running state, K/V parts, and checkpoint. The file header version is set to kVersionStages = 3.
• Read: get_payload reads stage 0 identically to v1, then checks the version. For v3 files it reads the stage count (bounded by limits.max_stages) and each subsequent stage block, validating geometry and layer range per stage. A v1 file has no stage section; a v2 file is refused at the version gate (not silently misparsed as v1).
• Save guard: if conversation_checkpoints_split routes any checkpoint to cs.rest (incomplete for the split), the save is refused rather than silently dropping it.
• Backward compatible: v1 files read unchanged. A v3 file opened by old code is refused at the version check.

Read-limit hardening

SessionReadLimits gains max_stages and a per-stage stage_limits vector. The serve loop populates these from the live engine's stage carves before calling session_file_read, so a file from a different split (wrong stage count, wrong layer range, oversized K/V) is refused before any allocation — not after a partial read.

Changes by file

File	What
include/strata/core/conversation_file.hpp	kVersionStages = 3, SessionReadLimits::max_stages, stage_limits, SessionStageLimits struct
include/strata/core/conversation_snapshot.hpp	conversation_stage_read_limits() declaration
src/core/conversation_file.cpp	v3 write (put_payload), v3 read (get_payload), per-stage validation, explicit v2 rejection
src/core/conversation_state.cpp	conversation_stage_read_limits() implementation, conversation_session_read_limits updated to carry stage bounds
src/program/generate.cpp	Serve loop: build stage_limits from live stages, set max_stages, pass through to read; save path collects per-stage K/V sources; cs.rest guard; draft_stage_kv helper
src/core/conversation_file_test.cpp	Round-trip tests for v1 (no stages) and v3 (multi-stage), v2 rejection, v4 rejection, version-mismatch refusal, oversized-stage refusal
src/core/conversation_snapshot_test.cpp	conversation_stage_read_limits unit tests

Environment

Strata built from source, CUDA 13.4, Windows 11 / Debian 13
2× NVIDIA RTX 5060 Ti (16 GB) or 2× NVIDIA A2 (16 GB), 64 GB system RAM
Qwen3.8-Flash-Next Coder IQ3_XXS / UD-Q4_K_XL, up to 128K context

Validation

Revision: 94f55cd6 (rebased onto main e88ee39d, 0.2.0-beta)
Environment: CUDA 13.4, Windows 11, 2× NVIDIA RTX 5060 Ti 16 GB (sm_120), 64 GB RAM

Commands:

build-win-cuda13.bat
build-win-cuda13\conversation_file_test.exe
build-win-cuda13\conversation_snapshot_test.exe
build-win-cuda13\conversation_cache_test.exe
build-win-cuda13\conversation_memory_test.exe
build-win-cuda13\conversation_split_failure_test.exe
build-win-cuda13\session_save_reclaim_test.exe

Passed:

• conversation_file_test: 208 checks
• conversation_snapshot_test: 4027 checks
• conversation_cache_test: 4219 checks
• conversation_memory_test: 25 checks
• conversation_split_failure_test: 36 allocation-failure cases (stages 1–3)
• session_save_reclaim_test: 119 checks
• Smoke test: engine starts on dual-GPU layer-split, chat completion responds, session save/restore round-trip works

Failed: none

Not run:

• --batch / parallel slots
• Restore with a different GPU split than the save
• Full-disk / near-full-VRAM restore ([Engine] avoid CUDA OOM during conversation-cache restore #2409)
• Linux (previously validated on 2× A2 sm_86 by @Eranicus30666, 23K and 94K tokens)

[qiuzan2001]

One control-flow bug from reading the branch (not built here): in the restore block of src/program/generate.cpp (around line 10637), the continue after refuse(err, strata::core::SessionError::io) sits inside for (size_t k = 0; k < n_st; ++k), so it only continues that loop. After a conversation_stage_read_limits failure the engine has already sent its refusal, then still calls session_file_read with a short stage_limits and can answer a second time (or restore), which would desynchronize the server's reader. A flag plus break, as the validate loop beside it does, would fix it. Also, #2247 bumps the same session file to version 2 for its running-state codec (kVersionCoded = 2), while this PR's kVersionStages is 2 as well, so the two v2 layouts collide whichever lands second.

[LuckyKo]

One control-flow bug from reading the branch (not built here): in the restore block of src/program/generate.cpp (around line 10637), the continue after refuse(err, strata::core::SessionError::io) sits inside for (size_t k = 0; k < n_st; ++k), so it only continues that loop. After a conversation_stage_read_limits failure the engine has already sent its refusal, then still calls session_file_read with a short stage_limits and can answer a second time (or restore), which would desynchronize the server's reader. A flag plus break, as the validate loop beside it does, would fix it. Also, #2247 bumps the same session file to version 2 for its running-state codec (kVersionCoded = 2), while this PR's kVersionStages is 2 as well, so the two v2 layouts collide whichever lands second.

Thanks for the careful read — both points are valid. Fixed in 75e7df6:

1. Control-flow bug (the continue inside the stage-limits loop)

Replaced with the same flag+break pattern the validate loop beside it uses:

bool stage_limit_bad = false;
for (size_t k = 0; k < n_st; ++k) {
...
if (!conversation_stage_read_limits(...)) {
stage_limit_bad = true;
break;
}
...
}
if (stage_limit_bad) { refuse(err, SessionError::io); continue; }
The outer continue now correctly skips to the next request in the serve loop instead of falling through to session_file_read.

2. Version collision with --session-state-codec bf16: a saved/parked conversation's running state at half the bytes #2247

Bumped kVersionStages from 2 → 3. The version gate now explicitly rejects v2:

if (version != kVersion && version != kVersionStages) {
// v2 is reserved for #2247's state codec; this build does not implement it.
error = "session file: unsupported version " + std::to_string(version); return false;
}
This avoids the fall-through hazard where a v2 (coded-state) file would be parsed by the v1 grammar. Once #2247 merges and adds its own v2 read path, this check can be relaxed to accept all three.

Also addressed while in here:

Guard on save: if conversation_checkpoints_split routes any checkpoint to cs.rest (incomplete for the split), the save is refused rather than silently dropping it from the file.
Stale "v2" references in public header comments → "v3".
Test: explicit v2-rejection case added alongside the existing v4-rejection case.
Happy to coordinate on the version numbering if you'd prefer a different scheme (e.g., a flags field instead of sequential integers).

[Eranicus30666]

Tested on Linux with two NVIDIA A2 (Ampere, sm_86), a setup the PR description does not cover (it was checked on 2x RTX 5060 Ti, Windows, 6K tokens). This is a test report, not a code review.

Setup: Strata v0.1.45 (477ac13c) + this PR (75e7df6b), built from source for sm_86 with CUDA 13.4.92, driver 615.71.09, Debian 13. 2x NVIDIA A2 16 GB, --layer-split auto, Qwen3.8-Flash-Next UD-Q4_K_XL, --max-context 131072, --kv int8, MTP on, one slot (no --batch), --slot-save-path set.

Merge: the branch does not merge into current main. One conflict, in src/program/generate.cpp: both sides added code at the same spot (upstream added CheckpointPool and StageThreads, the PR adds draft_stage_kv). I kept both; it builds without other changes.

Without the PR (v0.1.45), POST /slots/0?action=save on the two cards answers 400: session files do not support --layer-split.

With the PR, one run each, engine restarted between save and restore:

	23,150 tokens	94,435 tokens
Reading the prompt from scratch	44.8 s	147 s
Save	590,627,512 bytes, 2.16 s	1,679,847,048 bytes, 5.46 s
Restore after an engine restart	1.43 s	3.71 s
Next turn after restore	33 tokens read (23,150 of 23,183 reused), 3.1 s	33 tokens read (94,435 of 94,468 reused), 3.9 s

• At 23K the reply after the restore is identical to the reply from the live session.
• At 94K it is the same sentence with one more clause at the end. Could be ordinary numeric noise (I also saw a reply differ between a warm session and a cold read at 23K), but I ran it once, so I cannot say.
• Invalid files: a truncated file, random bytes and an empty file are each refused with 400 (kind: invalid), a missing file with 404, and a normal request right after each one worked.
• VRAM after a restore was the same as after the save.
• Unit tests on the merged tree (CPU-only build): conversation_file_test (227 checks), conversation_cache_test, conversation_memory_test, conversation_split_failure_test and session_save_reclaim_test pass. I could not build conversation_snapshot_test in that configuration.

Not tested: --batch or parallel slots, restoring a file saved with a different split (only the unit tests cover it), a full disk, a restore while VRAM is nearly full (#2409), Windows, other models.

Test run and write-up with Claude.

[LuckyKo]

Tested on Linux with two NVIDIA A2 (Ampere, sm_86), a setup the PR description does not cover (it was checked on 2x RTX 5060 Ti, Windows, 6K tokens). This is a test report, not a code review.

Setup: Strata v0.1.45 (477ac13c) + this PR (75e7df6b), built from source for sm_86 with CUDA 13.4.92, driver 615.71.09, Debian 13. 2x NVIDIA A2 16 GB, --layer-split auto, Qwen3.8-Flash-Next UD-Q4_K_XL, --max-context 131072, --kv int8, MTP on, one slot (no --batch), --slot-save-path set.

Merge: the branch does not merge into current main. One conflict, in src/program/generate.cpp: both sides added code at the same spot (upstream added CheckpointPool and StageThreads, the PR adds draft_stage_kv). I kept both; it builds without other changes.

Without the PR (v0.1.45), POST /slots/0?action=save on the two cards answers 400: session files do not support --layer-split.

With the PR, one run each, engine restarted between save and restore:

23,150 tokens 94,435 tokens
Reading the prompt from scratch 44.8 s 147 s
Save 590,627,512 bytes, 2.16 s 1,679,847,048 bytes, 5.46 s
Restore after an engine restart 1.43 s 3.71 s
Next turn after restore 33 tokens read (23,150 of 23,183 reused), 3.1 s 33 tokens read (94,435 of 94,468 reused), 3.9 s

• At 23K the reply after the restore is identical to the reply from the live session.
• At 94K it is the same sentence with one more clause at the end. Could be ordinary numeric noise (I also saw a reply differ between a warm session and a cold read at 23K), but I ran it once, so I cannot say.
• Invalid files: a truncated file, random bytes and an empty file are each refused with 400 (kind: invalid), a missing file with 404, and a normal request right after each one worked.
• VRAM after a restore was the same as after the save.
• Unit tests on the merged tree (CPU-only build): conversation_file_test (227 checks), conversation_cache_test, conversation_memory_test, conversation_split_failure_test and session_save_reclaim_test pass. I could not build conversation_snapshot_test in that configuration.

Not tested: --batch or parallel slots, restoring a file saved with a different split (only the unit tests cover it), a full disk, a restore while VRAM is nearly full (#2409), Windows, other models.

Test run and write-up with Claude.

Thanks for the thorough cross-platform test! A few notes:

Merge conflict — resolved now; I just rebased onto current main (5e1ec82). The conflict was exactly what you hit (CheckpointPool/StageThreads vs draft_stage_kv), and the branch should merge cleanly going forward.

94K one-clause difference — sounds like decode non-determinism as you suspected. The restore is byte-identical (VRAM matches, same token reuse count: 94,435 of 94,468 reused), so the KV state is correct; the divergence is in attention numerics between a fresh 147s prompt pass and a restored state computed incrementally over many turns. Your observation that warm vs cold at 23K also differed once supports this — it's floating-point accumulation order, not a serialization bug. If you can reproduce it deterministically (same seed, --temp 0), that'd be worth a separate issue, but I wouldn't block this PR on it.

PR title/description — updating now to say v3 (the code and commit already did; the GitHub page was stale).

[Niko1221]

Thanks for the pull request.

We're resetting the tracker for 0.2.0-beta. There are more than 1,300 open issues and pull requests, too many for us to review properly, so we're closing everything that isn't being worked on right now. Closed here does not mean rejected.

To reopen it takes about two minutes:

1. Rebase on main if it no longer applies, and re-run your checks on 0.2.0-beta.
2. Start the title with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other].
3. Fill in the pull request template: the problem and result, and the validation (what you ran, on which revision, what you saw).
4. Comment /reopen. A check runs and either reopens the pull request or lists what is still missing.

The rules are in CONTRIBUTING.md.

[LuckyKo]

/reopen

[github-actions]

Not reopened yet. Still missing:

• Say in Validation which revision you tested: rebase on main (or 0.2.0-beta) and re-run your checks.

Edit the pull request and comment /reopen again. Details: CONTRIBUTING.md

[LuckyKo]

/reopen

[github-actions]

Not reopened yet. Still missing:

• Fill in "Validation": the exact commands, the revision and environment tested, what you saw. Keep Passed, Failed and Not run apart.

Edit the pull request and comment /reopen again. Details: CONTRIBUTING.md

[LuckyKo]

/reopen

[github-actions]

The check passed, but GitHub would not reopen it (Validation Failed: ). If the branch was deleted, restore it or open a new pull request and link this one.