Pull requests / #3541

#3541 [Engine] metal: the conversation cache and session files keep the decode window's own K/V (#2502)

open · draft · @oceanio · 0 Kommentare · Auf GitHub

BenchmarksServer & APIAMD / HIPNVIDIA / CUDAModels & quantsSecurityDocumentationWindowsLinux

Beschreibung

## Problem and result
On the Mac (Metal) engine the conversation cache never keeps a conversation, and session files cannot be saved (#2502). A short side request, such as a client's title call, pushes the long conversation out, and its next turn reads the whole prompt again. There are two causes:

1. The memory guard has free-memory readings for Windows and Linux only. On a Mac it logs `memory guard: the available RAM cannot be read; parking and session save/restore are refused` at every start, and refuses every park and session save/restore.
2. A conversation snapshot copies the session's QSA K/V pools, but the Metal decode window keeps its own K/V and running state and never uses those pools. The checkpoint check also expects CUDA's running-state layout, so parking fails with `invalid checkpoint running-state payload`. After #3536 the session no longer has the pools at all.

With this PR, parking and session files work on the Mac and carry the decode window's own state. Measured on an M5 Max with Qwen3.8-Flash-Next IQ3_S: a 28,785-token conversation is parked as 798 MB in 49 ms and restored in 41 ms, and its follow-up decodes the same tokens as the conversation that was never parked. As a session file it is 679 MB, saved in 0.16 s and restored in a new engine process in 0.10 s, with the same follow-up.

## Related discussion
- #2502 (bug report, both causes; closed in the 0.2.0 reset).
- Earlier fixes, also closed in the reset: #2503 by @joesinvestments (cause 1) and #2631 by @PiMatyja (cause 2, parking only). This PR covers both causes, and differs in that:
  - the macOS reading is the Metal engine's own (free and purgeable pages and the clean file cache, as its streaming arena uses). Other programs' inactive anonymous memory is not counted, because macOS frees it only by compressing or swapping it. #2503 counts inactive pages;
  - session files work too. #2631 refuses them on Metal. Their fingerprint carries the window's layout, so a file from a build with another layout is refused instead of misread;
  - it builds on #3536, so a snapshot takes the window's K/V instead of pools the session no longer has.
  If either author would rather continue their own PR, I'm happy to rebase mine on theirs or close it.

## Approach
Under `STRATA_GPU_METAL` only:
- `src/core/conversation_memory.cpp`: `sample_host_memory()` takes `host_memory_now()` (the Metal engine's free-memory reading) on macOS.
- `metal/src/engine/window.cpp`:
  - `SlotState` gathers the window's per-conversation state (a refactor, the same buffers);
  - `SlotState::parts / layer_cells / cells_of` walk that state as byte ranges in a fixed order;
  - `Window::kv_layer` and `checkpoint_sizes` hand the snapshot code the K/V cells `[0, upto)` of each QSA layer and of the MTP draft layer, plus the running state's sizes.
  - The draft layer's cache uses the same per-cell layout.
- `src/core/conversation_snapshot.cpp`: a snapshot's K/V layer is the window's own (format key `0x4D570001`, page size 1), copied with `memcpy` (unified memory, the window idle). A K/V layer in any other format is refused.
- `src/core/conversation_state.cpp`: the checkpoint and session-read sizes come from the window's running state.
- `src/program/generate.cpp`: the session fingerprint names backend `metal-window`, the window's layout, and the MTP GGUF the window actually loads.
- Tests: `tools/metal_snapshot_test.py` and `tools/session_restart_test.py` are new. `tools/parking_test.py` reads the engine's own log path.
- Docs: MAC.md (Limits: the conversation cache and session files, measured) and DETAILS.md (the macOS memory reading, Metal in the fingerprint).

Cost: about 15.4 KB per token plus 118 MB per running state (the live one and each checkpoint kept with it). Nothing changes unless `--conversation-cache-mib` is set or a session file is saved or restored. With `STRATA_GPU_METAL` undefined, `unifdef` gives the same four shared C++ files as main's, apart from blank lines.

## Validation
Environment: Apple M5 Max (128 GB), macOS 26.7.1, Apple clang 21.0.0. Model: Qwen3.8-Flash-Next GSQ-RCO IQ3_S (SC117's abliterated build, two-file layout), resident, `--expert-cache 64 --prefill auto --spec 2 --mtp <rt> --mmap-experts`. Revision: this branch's tip, on #3536 (main `9de04cd4` + 2 commits).

**Passed:**
- Metal build (`cmake --build build-metal`) at every one of the 7 commits; no new warnings.
- `unifdef -USTRATA_GPU_METAL` on the four shared C++ files (conversation_memory, conversation_snapshot, conversation_state, generate): identical to main's, apart from blank lines.
- `tools/metal_snapshot_test.py --max-context 8192`: all checks pass.
  - The cache stays on (`conversation_cache_mib=1024`), and a second conversation parks the first.
  - SAVE answers `SAVED`. RESTORE of a non-session file answers `SERR invalid`, and RESTORE of the saved file answers `RESTORED 35 ...`.
  - The engine keeps serving.
- `tools/parking_test.py --extra "--conversation-cache-mib 8192 --conversation-cache-slots 4"`:
  - 3,682-token conversation: follow-up identical (99 of 99 tokens);
  - 28,692-token conversation (`--repeat 470`): follow-up identical (95 of 95). Engine log: `parked 28785 tokens in 48.8 ms ... bytes=797571752`, `restored 28785 tokens (live) in 41.1 ms`. The second park of the conversation reused 439 MB of K/V it still held.
- `tools/session_restart_test.py --extra "--max-context 65536"`:
  - 3,785 tokens: `SAVED ... 294433688 73.5`, `RESTORED ... 39.4` in a new engine;
  - 28,785 tokens: `SAVED ... 679433728 160.5`, `RESTORED ... 103.2`.
  - At both sizes the follow-up is identical (99/99 and 95/95), and the restored engine reused the whole conversation.
- `tools/session_restart_test.py --restore-exe <main's engine> --expect-refused`: main refuses this build's file cleanly (`SERR io 0 session fingerprint: ...`), and the engine keeps serving.
- Default path against main at `--max-context 262144`, through each tree's own `serve/server.py`, without `--conversation-cache-mib`:
  - greedy output identical on 8 requests (4 short prompts, a shared system prompt with two questions, a follow-up turn and a rewind, each reusing 1,577 tokens);
  - footprint 10.61 GiB (main 16.98: #3536's saving);
  - decode 92.0-95.3 tok/s (main 92.0-95.4).

**Failed:**
None observed.

**Not run:**
- CUDA, HIP and SYCL builds: no such toolchains here, and the shared files are compiled by every backend, so this is a draft. All changes are inside `#if defined(STRATA_GPU_METAL)`. Could someone build it on one of them?
- The refusal of a file from a Mac build with another window layout was not re-run on this branch. Main's refusal above comes from its own fingerprint (it looks for the CUDA layout's MTP files), not from the layout key. It was checked before the split: a file from an earlier build of this change was refused, and the engine kept serving.
- Other Macs, the SSD-streamed experts, vision with parking, contexts above 262K, a file from the PC engine restored on the Mac.

## Limits
- Builds on #3536: until that merges, this PR also shows its 2 commits.
- Snapshots copy the window's state with `memcpy` while the window is idle (unified memory). They do not overlap with decoding.
- Only the solo conversation is snapshotted. Batch slots on Metal (#2715) are a separate change.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Mehr auf der Site

Links zu Install, Modellen, Releases.