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
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.