Pull requests / #3536

#3536 [Engine] metal: the session no longer holds unused QSA K/V pools (6.4 GiB resident at 262K context)

open · draft · @oceanio · 0 commentaires · Sur GitHub

BenchmarksServer & APIAMD / HIPNVIDIA / CUDAModels & quantsDocumentation

Description

## Problem and result
On the Mac (Metal) engine the decode window (`metal/src/engine/window.cpp`) keeps its own K/V and pooled indexer keys for the 12 QSA layers, sized by `--max-context`. It never reads or writes the session's QSA K/V pools, but `kv_plan` still sizes them for the whole context. On Metal a session buffer is one shared `MTLBuffer` in the residency set, so those pools are resident in full whether anything touches them or not: about 24.6 KB per context token.

Measured on main 0.2.1.4 (phys_footprint of the engine 5 s after it is loaded, Qwen3.8-Flash-Next IQ3_S, M5 Max):

| `--max-context` | main | this PR | saved |
|---|---:|---:|---:|
| 262144 | 16.99 / 16.98 GiB | 10.62 / 10.61 GiB | 6.37 GiB |
| 65536 | 9.27 / 9.27 GiB | 7.68 / 7.68 GiB | 1.59 GiB |

The same 6.4 GiB is still saved after 9 requests (24.91 → 18.53 GiB at 262K). Output and speed do not change (Validation).

## Related discussion
None. This is a focused fix. Context: on Metal, parking and session save/restore are refused on main already (#2502). This PR does not change that.

## Approach
- `src/core/layer.cpp`, `kv_plan`: under `#if defined(STRATA_GPU_METAL)` a QSA state gets 0 K/V slots and 0 pooled indexer rows. `qsa_state_init` then leaves the pool pointers and `idx_pooled` null, so they do not alias the page table or `block_pos` that follow. The page table, the indexer tail and dead rows, `block_pos` and the RoPE tables are carved as before.
- With `n_slots == 0`, a conversation snapshot refuses the state (`conversation_snapshot.cpp` `valid`) instead of reading pools that are not there. On main, the memory guard already refuses parking and session save/restore on Metal (the log says so in both builds below), so nothing that works today goes through this path.
- `docs/MAC.md`: one line saying that the engine's own memory grows with `--max-context`, with the numbers before and after.
- Non-Metal builds: unchanged. `unifdef -USTRATA_GPU_METAL src/core/layer.cpp` gives the same file before and after.
- No flag or setting: the pools were never used on Metal, so there is nothing to opt into.

## Validation
Environment: Apple M5 Max (128 GB), macOS 26.7.1, Apple clang 21.0.0. Both builds from the same checkout procedure (`cmake -S metal -B build-metal -G Ninja -DCMAKE_BUILD_TYPE=Release`, target `strata`): main `9de04cd4` and this branch (`9de04cd4` + 2 commits). 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`, through each tree's own `serve/server.py`. Runs alternate main / PR / main / PR at each context, 30 s apart.

**Passed:**
- Metal build: both build with the same 11 warnings, none in `layer.cpp`.
- `unifdef -USTRATA_GPU_METAL src/core/layer.cpp`: identical to main's (`cmp`).
- Memory: the table above (2 runs per build per context).
- Greedy output (temperature 0, up to 128 tokens): 8 requests per run. They are 4 short prompts, the same ~1.2K-token system prompt with two questions, a live follow-up turn (1,577 tokens reused), and a second follow-up rewound to the same point (1,577 reused). All 8 are identical between main and this PR, in both rounds, at 262K and at 65K, and identical across rounds within each build. A 9th request (a ~15K-token prompt carrying a per-run tag) read and answered in every run, with prompt reuse and checkpoints logged the same in both builds.
- Decode speed of the 4 short prompts (tok/s, round 1 / round 2):
  - 262K: main 93.1 94.5 91.9 94.6 / 93.3 93.6 91.7 95.1; PR 93.2 94.7 92.0 95.1 / 93.6 94.5 91.5 94.4
  - 65K: main 93.0 94.4 91.8 95.2 / 93.3 94.2 91.7 95.0; PR 93.2 94.1 91.6 94.5 / 92.9 93.6 92.1 95.3

**Failed:**
None observed.

**Not run:**
- CUDA, HIP and SYCL builds: no such toolchains here. `layer.cpp` is compiled by every backend, so this stays a draft. Could someone build it on one of them? The change is entirely inside `#if defined(STRATA_GPU_METAL)`, and the unifdef check above is my only evidence for the other backends.
- Other Macs, the SSD-streamed experts (Macs under ~96 GB), vision, `--max-context` above 262K.
- Long-context quality checks (needles, KL). The window's K/V is untouched, so I expect no change, but I did not run them for this PR.

## Limits
- The saving applies to the Metal engine only. It is per session: a feature that gives each request its own session would get the same saving per session.
- If session save/restore or parking on Metal are made to work later, they have to carry the decode window's own K/V rather than the session pools, which were never filled on Metal.

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

Sur le site

Liens install, modèles, releases.