Pull requests / #2590
#2590 session: support multi-GPU layer-split save/restore (format v2)
open · @LuckyKo · 0 Kommentare · Auf GitHub
Server & APIMulti-GPUNVIDIA / CUDAModels & quantsWindows
Beschreibung
# session: support multi-GPU layer-split save/restore (format v2) ## 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 v2** for session files that carries per-stage data: - **Write:** `put_payload` now 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 bumped to `kVersionStages = 2`. - **Read:** `get_payload` reads stage 0 identically to v1, then checks the version. For v2 files it reads the stage count (bounded by `limits.max_stages`) and each subsequent stage block, validating geometry and layer range per stage. - **Backward compatible:** a v1 file has no stage section; `stage_images` stays empty and the engine behaves exactly as before. A v2 file opened by old code is refused at the version check (no silent corruption). ### 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`, `SessionReadLimits::max_stages`, `stage_limits`, `SessionStageLimits` struct | | `include/strata/core/conversation_snapshot.hpp` | `conversation_stage_read_limits()` declaration | | `src/core/conversation_file.cpp` | v2 write (`put_payload`), v2 read (`get_payload`), per-stage validation | | `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 | | `src/core/conversation_file_test.cpp` | Round-trip tests for v1 (no stages) and v2 (multi-stage), version-mismatch refusal, oversized-stage refusal | | `src/core/conversation_snapshot_test.cpp` | `conversation_stage_read_limits` unit tests | ## Testing - All existing session save/restore tests pass unchanged (v1 path untouched). - New tests cover: v2 round-trip with 2 and 3 stages, v1→v2 version gate, stage-count mismatch refusal, per-stage layer-range validation, K/V size bounds per stage. - Manually verified on a 2×RTX 5060 Ti (sm_120) with `--layer-split auto`: save a 6K-token multi-turn session, kill the engine, restart, restore — conversation resumes with full context and correct expert routing. ## Environment - Strata built from source, CUDA 13.4, Windows 11 - 2× NVIDIA RTX 5060 Ti (16 GB each), 64 GB system RAM - Qwen3.8-Flash-Next Coder IQ3_XXS, 128K context
Mehr auf der Site
Links zu Install, Modellen, Releases.