Pull requests / #3564

#3564 [Engine] Metal: working session cache, and the 12 GiB its arena was reserving for nothing

open · @drumblund · 0 comentarios · En GitHub

needs-format

Descripción

> Branch `metal-cache-only`, 4 commits on `origin/main` @ `9de04cd4`. Engine-side only — no `serve/` change.
> Status: built and measured on v0.1.44.1 (the engine this machine serves from); this branch is the same four
> commits resolved against `9de04cd4` and built there as `strata 0.1.46-metal`.

## Problem and result

Three things were broken on macOS/Metal, and they compound: the session cache could not be used at all.

**1. The memory guard could not read the RAM.** `sample_host_memory()` had a Windows branch and a `/proc`
branch, nothing else. On a Mac it returned nothing, so `conversation_memory_admit()` said no and both
`--conversation-cache-mib` (RAM parking) and session save/restore were refused outright:

```
strata generate: memory guard: the available RAM cannot be read; ... session save/restore are refused
```

**2. The Metal backend never published its K/V.** Parking and session files read `QsaState` descriptors
carved from an arena the Metal decode window never writes; its K/V lives in `Window::kq_/vq_/ksc_/vsc_`.
The descriptor's running-state sizes did not even match, so a SAVE died before writing a byte:

```
session save: conversation snapshot: invalid checkpoint running-state payload
   (tails 24,576 vs 18,432 bytes; block_pos 0 vs 48)
```

Had the sizes matched, the failure would have been worse than loud: the arena memory exists and reads
fine, it is simply all zeros, so the file would have saved successfully and restored into nothing.

**3. 12 GiB of VRAM was reserved for a copy nothing used.** The arena carved a full fp16 K/V pool per QSA
layer (1 GiB × 12 at 512K) plus a 128 MiB rope table that is already a no-op on this backend.

After, on the same machine (M5 Max, 128 GB, 512K context, IQ3_S):

| | before | after |
|---|---|---|
| memory guard | `the available RAM cannot be read` | `limit 128.0 GiB, available 75.3 GiB` |
| session SAVE at 512K | HTTP 500, `invalid checkpoint running-state payload` | 651 MB in 223 ms |
| session RESTORE at 512K | not reached | 27,174 tokens in 146 ms, `cached=26,987` of 26,992 |
| `VRAM free with everything loaded` | 30,227 MiB | 42,643 MiB |

The answer after a restore matches the one before it character for character, and matches the value
recomputed independently from the generated corpus (so it is the K/V being read, not a re-prefill
landing on the same greedy output).

## Related discussion

none — no issue was opened for this. Every behaviour change is behind `#if defined(STRATA_GPU_METAL)` or
a field that defaults to the current value, so CUDA/SYCL/HIP are byte-for-byte unchanged.

## Approach

**A backend says what it owns, instead of the snapshot code assuming.** `QsaState::page_size` (0 = the
artifact's granule, so every existing backend is byte-for-byte unchanged) lets a backend say its pools are
cell-major. The Metal window sets it to 1, and at a granule of one cell `[page][kv_head][page_size][head_dim]`
and `[cell][kv_head][head_dim]` are the same bytes in the same order — so parking, session files and the
state hash read the real K/V through the existing code path, with no second copy of any of it. That is the
whole trick; `Window::bind_session` publishes the pointers once, at `window_init`, before the first request.

**Two running-state fields were not in the shared format.** The indexer tail was saved as four rows per
layer where the arena keeps three (`idx_block - 1` — the cell that completes a block is consumed by the
pool instead of being kept, exactly as `src/kernels/cuda/qsa.cu` notes), and `block_pos` was cleared
because this backend derives it from the cell count and keeps no copy; it is reconstructed from the same
count. The draft layer's descriptor was never initialised at all on this backend, so `bind_session`'s
sibling binds the MTP window's own buffers, and `MtpDrafter::load` brings the window up (idempotent) so a
RESTORE that arrives before the first prefill has a descriptor to size the file's draft part against.

**The pools are skipped at compile time, not at run time.** With the descriptors published, the arena's
copy is a second set that nothing writes and nothing reads; `qsa_pools_in_arena` is a `constexpr` false on
Metal only, applied to `qsa_state_bytes()` and `qsa_state_init()` in pairs so offsets stay aligned.

Compatibility: CUDA/SYCL/HIP builds keep `page_size = 0`, so their layout, sizes and file format are
unchanged. Session files written on Metal declare `page_size = 1` and a CUDA engine refuses them with
`incompatible K/V geometry` rather than misreading them.

## Validation

**Passed:**

*Environment for all of the below:* macOS 26.5.2, M5 Max 128 GB, Xcode 26.6, engine `0.1.44.1-metal`,
IQ3_S 79 GB, `cmake --build build-metal` (incremental, ~5 s per rebuild).

- Memory guard, 512K config: before `the available RAM cannot be read; ... session save/restore are refused`,
  after `limit 128.0 GiB (source: meminfo), available 60.4 → 75.3 GiB`.
- 64K config, SAVE → kill -9 → restart → RESTORE, 4,250-token conversation:
  `saved 291,371,100 B in 525 ms`, `restored 3,617 tokens in 77 ms`, `cached=3,548`, output identical.
- 512K config, 26,958-token conversation: `parked 27,023 tokens in 120.7 ms`,
  `saved 651,131,596 B in 223.3 ms`, `restored 27,174 tokens in 146.0 ms (read+check 115.1)`,
  `cached=26,987`, output identical and equal to the value recomputed from the corpus.
- RAM parking round trip at 512K: switching back to a parked conversation went from 2,306 ms to 54 ms.
- `VRAM free with everything loaded`: 30,227 → 42,643 MiB (+12,416), matching the computed
  `131,072 × 2 × 4 × 256 × 2 × 2 B × 12 layers + 128 MiB` to the MiB.
- The branch builds on current `main` (`9de04cd4`): `cmake -S metal -B build-check-metal -DMETAL_CPP_DIR=…`
  then `ninja strata` → `strata 0.1.46-metal`, 4,542,424 B, with `bind_session` / `mtp_kv_state` present
  (`nm`). This is the Metal project only; the engine sources that this PR touches are the ones compiled.

**Failed:** none observed.

**Not run:**

- CUDA, HIP and SYCL builds and their test suites — no such hardware here. The guard is
  `#if defined(STRATA_GPU_METAL)` and `page_size` defaults to 0, but I could not compile those targets.
- Windows.
- The root project's `strata_kernels_cpu` target on arm64 (a clean `cmake -S .` stops at `-mavx512f` …
  on `arm64-apple-darwin`). Unrelated to this PR - `metal/` is its own CMake project and is what was built.
- A run of the 512K configuration against the `0.1.46-metal` binary: the measurements below are from the
  same four commits built into `0.1.44.1-metal`, the engine this machine serves from. Upstream's only
  Metal change between the two is `qsa_prompt.metal` (an M1 Max thread-group fix).

## Limits

- **Metal session files do not travel to CUDA.** They are written with `page_size = 1`; a CUDA engine
  refuses them (`incompatible K/V geometry`). That is the intended behaviour, not a bug to fix here.
- **`idx_pooled`'s moving spare row is not written by this backend.** Metal scores from `idx_dead`; the
  row may therefore hold a stale value in a Metal file. Neither backend's scoring reads it, but a future
  caller that does would need it.
- **The first request after a restore still pays ~8 s at 512K** — Metal's one-time pipeline capture in that
  process, unrelated to the restore (a second identical request is 0.5 s). Worth knowing before reading a
  wall-clock number as the restore cost.
- **The parked pool is RAM, not free.** `--conversation-cache-mib` is a ceiling; a 280K-token conversation
  is ~4.5 GB (`236 MB + tokens × 15.3 KB`, measured at two sizes). On a memory-tight machine, lower the
  ceiling before anything else.
**The title must start with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other], for example "[Engine] fix decode stop on long prompts" or "[GUI] keep the model list after a tab switch".**

<!--
Read CONTRIBUTING.md first. Lead with the problem and the result, not a file-by-file list.
Placeholders below are not test results: replace them with what you observed, or write "none".
-->

## Problem and result
<!-- What triggered this change. How it behaves before and after. -->

## Related discussion
<!-- The issue URL and the maintainer's agreement on scope, when the change needs one (features, refactors, defaults). Write "none" if not applicable. -->

## Approach
<!-- Why this fix. Compatibility effects. Keep unrelated cleanup out of the PR. -->

## Validation
<!-- Exact commands, the revision and the environment tested, what you saw. Keep the three lists apart. -->

**Passed:**
<!-- command, revision/environment, observed result, log -->

**Failed:**
<!-- command and failure, or "none observed" -->

**Not run:**
<!-- check, why, and what stays unverified (for example CUDA, HIP, SYCL or Metal builds you could not do) -->

## Limits
<!-- Known limitations or remaining work. -->

En el sitio

Enlaces a install, modelos, releases.