Prefix cache + paged (block-pool) KV storage with SSD spill tier

CLAUDE.md

@@ -164,6 +164,25 @@ reference/.venv/bin/python reference/dump_ref.py
   affect kernel accumulation order. Engine-wide setting, default off;
   vision/hybrid models are rejected at engine creation (no silent fp16
   fallback).
+- **The prefix cache harvests PROMPT K/V only — never decode-produced K/V.**
+  Decode-with-cache K/V differs from a recompute by fp16 accumulation order
+  (the decode-vs-recompute gap below) and demonstrably flips later greedy
+  choices; prefill-produced K/V is the proven exact-stable class, so pooling
+  only it keeps the feature's gate (warm == cold, token-exact) sound.
+  Multi-turn reuse still converges — the next turn's prompt contains the prior
+  answer as text and pools after its own (seeded) prefill. The pool
+  (`cache/block_pool`) stores immutable blocks keyed by a salted chain hash;
+  matched blocks seed a batch-1 cache via `BatchKVCache::from_prefix`, written
+  through the standard `update_kv_components` writer so buffer layout matches a
+  cold prefill (strides are load-bearing, see kv-quant above). Harvest
+  materializes copies (`mx::contiguous` + eval) — a lazy slice would pin the
+  whole batch buffer. The SSD tier (`cache/block_store`) is byte-only across
+  threads (worker does all array<->bytes conversion); its writer keeps the
+  queue front visible to `get()`/`contains()` until the file lands, and its
+  serialize order (per layer, K then V components) is gated by the exact-token
+  spill test — an order mismatch produced silent garbage. Engine-wide opt-in;
+  vision/hybrid models and spill-without-prefix-cache are rejected at engine
+  creation. Multimodal rows are never harvested or matched.
 - **Qwen3-VL interleaved M-RoPE can't use `fast::rope`** (it takes a 1D offset,
   not 3D `(t,h,w)` positions). `Qwen3VLModel` hand-rolls a half-split rotation
   with a per-frequency t/h/w selector; text tokens have `t==h==w` so it reduces to

CMakeLists.txt

@@ -69,6 +69,9 @@ add_library(mlxforge_core STATIC
   src/cache/batch_kv_cache.cpp
   src/cache/kv_quant.cpp
   src/cache/kv_budget.cpp
+  src/cache/block_pool.cpp
+  src/cache/prefix_cache.cpp
+  src/cache/block_store.cpp
   src/sample/sampler.cpp
   src/sample/json_grammar.cpp
   src/runtime/single_stream.cpp

apps/mlxforge.cpp

@@ -73,12 +73,18 @@ void print_help() {
       "      --max-waiting <N>  max queued requests (default 256)\n"
       "      --kv-budget <B>    KV cache budget in bytes, 0 = unbounded (default 0)\n"
       "      --kv-bits <N>      KV cache quantization: 0 = fp16, 8 or 4 (default 0)\n"
+      "      --prefix-cache <0|1>  reuse pooled KV across shared prompt prefixes (default 0)\n"
+      "      --kv-block <N>     prefix-pool block size in tokens (default 256)\n"
+      "      --kv-pool <B>      prefix-pool RAM budget in bytes, 0 = unbounded (default 1 GiB)\n"
+      "      --kv-spill-dir <D> SSD spill dir for evicted prefix blocks (default off)\n"
+      "      --kv-spill-bytes <B>  spill-dir disk budget in bytes, 0 = unbounded (default 0)\n"
       "  -h, --help             show this help and exit\n"
       "\n"
       "The model may be given via -m or the config file's \"model\" key.\n"
       "Config precedence (low to high): defaults < config file < env vars < CLI flags.\n"
       "Env vars: MLXFORGE_HOST, MLXFORGE_PORT, MLXFORGE_MAX_CTX, MLXFORGE_MAX_WAITING, "
-      "MLXFORGE_KV_BUDGET, MLXFORGE_KV_BITS.");
+      "MLXFORGE_KV_BUDGET, MLXFORGE_KV_BITS, MLXFORGE_PREFIX_CACHE, MLXFORGE_KV_BLOCK, "
+      "MLXFORGE_KV_POOL, MLXFORGE_KV_SPILL_DIR, MLXFORGE_KV_SPILL_BYTES.");
   std::fflush(stdout);
 }
 
@@ -131,6 +137,11 @@ int main(int argc, char** argv) {
     ec.model_spec = sc.model_dir;
     ec.max_waiting = sc.max_waiting;
     ec.kv_bits = sc.kv_bits;
+    ec.prefix_cache = sc.prefix_cache;
+    ec.kv_block_size = sc.kv_block;
+    ec.kv_pool_bytes = sc.kv_pool_bytes;
+    ec.kv_spill_dir = sc.kv_spill_dir;
+    ec.kv_spill_bytes = sc.kv_spill_bytes;
     engine = std::make_unique<mlxforge::Engine>(std::move(ec));
   } catch (const std::exception& e) {
     mlxforge::log::error("model error: {}", e.what());
@@ -167,8 +178,9 @@ int main(int argc, char** argv) {
   std::signal(SIGTERM, on_signal);
 
   // Info log: server has started, print bind details and config bounds.
-  mlxforge::log::info("mlxforge serving on http://{}:{} (max_ctx={} max_waiting={} kv_bits={})",
-                      sc.host, sc.port, sc.max_ctx, sc.max_waiting, sc.kv_bits);
+  mlxforge::log::info(
+      "mlxforge serving on http://{}:{} (max_ctx={} max_waiting={} kv_bits={} prefix_cache={})",
+      sc.host, sc.port, sc.max_ctx, sc.max_waiting, sc.kv_bits, sc.prefix_cache);
 
   // Run the server's request loop (blocks until stop() is called).
   server.listen(sc.host, sc.port);

apps/mlxforge_cli.cpp

@@ -15,6 +15,11 @@
 //   mlxforge-cli bench <model> [max_tokens] [runs]
 //     - Repeatable throughput benchmark over a fixed prompt: one discarded warmup run, then `runs`
 //       timed runs (defaults: max_tokens=128, runs=3) reporting time-to-first-token and decode tok/s.
+//   mlxforge-cli bench-prefix <model> [prefix_tokens] [runs]
+//     - Prefix-cache benchmark (defaults: prefix_tokens=2048, runs=3): builds an Engine with the
+//       prefix cache on, then measures TTFT for one COLD request and `runs` WARM requests that share
+//       a prefix_tokens-long prompt prefix (distinct tails). Reports the cold/warm speedup and the
+//       engine's reuse metrics; warm decode tok/s shows reuse leaves throughput unchanged.
 //   mlxforge-cli embed <model> <text> [--last|--mean] [--eos] [--instruct "..."] [--no-normalize]
 //     - Embeds text and prints the (by default unit-normalized) vector. With no flags the model
 //       self-selects its convention (a Qwen3-Embedding checkpoint uses last-token pooling + a
@@ -44,6 +49,7 @@
 #include "runtime/engine.h"
 #include "runtime/multimodal_stream.h"
 #include "runtime/single_stream.h"
+#include "scheduler/request.h"
 #include "tokenizer/tokenizer.h"
 #include "vision/image_decode.h"
 
@@ -322,6 +328,95 @@ int run_bench(const std::string& spec, int max_tokens, int runs) {
   return 0;
 }
 
+// Prefix-cache benchmark: build a real Engine (the exact library path, prefix
+// cache on) and measure time-to-first-token for prompts sharing a long prefix —
+// the shared-system-prompt scenario the cache exists for. One discarded warmup
+// absorbs Metal kernel compilation; the COLD run then prefills the shared
+// prefix from scratch, and each WARM run reuses its pooled blocks with a unique
+// tail (a distinct "user question"). Greedy, EOS disabled, fixed max_tokens, so
+// runs are comparable; decode tok/s is reported to show reuse does not change
+// steady-state throughput.
+int run_bench_prefix(const std::string& spec, int prefix_tokens, int runs) {
+  mlxforge::EngineConfig cfg;
+  cfg.model_spec = spec;
+  cfg.prefix_cache = true;
+  mlxforge::Engine engine(cfg);
+  while (!engine.ready()) std::this_thread::sleep_for(std::chrono::milliseconds(20));
+  const mlxforge::Tokenizer& tok = engine.tokenizer();
+
+  // Build a shared prefix of ~prefix_tokens ids by repeating a paragraph.
+  const std::vector<int> para =
+      tok.encode("The ocean covers most of the planet, and its slow currents move heat between "
+                 "the equator and the poles, shaping weather on every continent. ");
+  std::vector<int> prefix;
+  while (static_cast<int>(prefix.size()) < prefix_tokens)
+    prefix.insert(prefix.end(), para.begin(), para.end());
+  prefix.resize(prefix_tokens);
+
+  const int kMaxTokens = 32;
+  // Submit a prompt through the scheduler (the continuous-batching path the
+  // server and bindings use) and return {ttft_ms, decode tok/s}.
+  auto timed_run = [&](std::vector<int> ids) {
+    auto req = std::make_shared<mlxforge::Request>();
+    req->prompt_ids = std::move(ids);
+    req->params.temperature = 0.0f;
+    req->max_tokens = kMaxTokens;  // eos_ids stays empty: fixed-length runs
+    const auto t0 = std::chrono::steady_clock::now();
+    engine.scheduler().submit(req);
+    double ttft_ms = 0.0;
+    auto t_first = t0;
+    int produced = 0, tk = 0;
+    while (req->tokens.pop(tk)) {
+      if (produced++ == 0) {
+        t_first = std::chrono::steady_clock::now();
+        ttft_ms = std::chrono::duration<double, std::milli>(t_first - t0).count();
+      }
+    }
+    const double decode_s =
+        std::chrono::duration<double>(std::chrono::steady_clock::now() - t_first).count();
+    const double tps = (produced > 1 && decode_s > 0) ? (produced - 1) / decode_s : 0.0;
+    return std::make_pair(ttft_ms, tps);
+  };
+  auto with_tail = [&](int i) {
+    std::vector<int> ids = prefix;
+    const std::vector<int> tail =
+        tok.encode("Question " + std::to_string(i) + ": summarize the key point briefly.");
+    ids.insert(ids.end(), tail.begin(), tail.end());
+    return ids;
+  };
+
+  mlxforge::log::info("bench-prefix: prefix={} tokens, max_tokens={}, warmup=1, runs={}",
+                      prefix.size(), kMaxTokens, runs);
+
+  // Warmup (discarded): an UNRELATED prompt — it absorbs Metal kernel
+  // compilation but shares no prefix, so the cold run below stays cold.
+  timed_run(tok.encode("A completely unrelated warmup prompt about gardening tools."));
+  const auto [cold_ttft, cold_tps] = timed_run(with_tail(0));
+  std::printf("  cold:      ttft %8.1f ms   decode %.1f tok/s\n", cold_ttft, cold_tps);
+  std::fflush(stdout);
+
+  double warm_sum = 0.0, warm_min = 1e300, warm_max = 0.0, tps_sum = 0.0;
+  for (int i = 1; i <= runs; ++i) {
+    const auto [ttft, tps] = timed_run(with_tail(i));
+    warm_sum += ttft;
+    tps_sum += tps;
+    warm_min = std::min(warm_min, ttft);
+    warm_max = std::max(warm_max, ttft);
+    std::printf("  warm %d/%d:  ttft %8.1f ms   decode %.1f tok/s\n", i, runs, ttft, tps);
+    std::fflush(stdout);
+  }
+
+  const mlxforge::WorkerMetrics m = engine.metrics();
+  const double warm_mean = warm_sum / runs;
+  std::printf("\nttft       cold %.1f ms   warm mean %.1f ms (min %.1f, max %.1f)   speedup %.1fx\n",
+              cold_ttft, warm_mean, warm_min, warm_max,
+              warm_mean > 0 ? cold_ttft / warm_mean : 0.0);
+  std::printf("decode     cold %.1f tok/s   warm mean %.1f tok/s\n", cold_tps, tps_sum / runs);
+  std::printf("reuse      hits %ld   tokens reused %lld   pool %ld blocks / %lld bytes\n",
+              m.prefix_hits, m.prefix_tokens_reused, m.prefix_pool_blocks, m.prefix_pool_bytes);
+  return 0;
+}
+
 // Embedding smoke harness: build a real Engine (so this exercises the exact
 // library path — detection of embedding defaults, instruction wrap, EOS append,
 // pooling, normalize) and print the resulting vector to stdout. This is the
@@ -423,6 +518,16 @@ int main(int argc, char** argv) {
     const int runs = argc >= 5 ? std::stoi(argv[4]) : 3;
     return run_bench(argv[2], max_tokens, runs);
   }
+  if (cmd == "bench-prefix") {
+    // Prefix-cache benchmark: cold vs warm TTFT over a shared prompt prefix.
+    if (argc < 3) {
+      std::fprintf(stderr, "usage: mlxforge-cli bench-prefix <model_dir> [prefix_tokens] [runs]\n");
+      return 2;
+    }
+    const int prefix_tokens = argc >= 4 ? std::stoi(argv[3]) : 2048;
+    const int runs = argc >= 5 ? std::stoi(argv[4]) : 3;
+    return run_bench_prefix(argv[2], prefix_tokens, runs);
+  }
   if (cmd == "embed") {
     // Embed text and print the vector. Flags override the model's detected
     // defaults (a Qwen3-Embedding checkpoint self-selects last-token + EOS).

bindings/node/index.d.ts

@@ -12,6 +12,25 @@ export interface EngineOptions {
   kvBits?: 0 | 4 | 8;
   /** Quantization group size (default 64; must divide the model's head_dim). */
   kvGroupSize?: number;
+  /**
+   * Prefix cache (engine-wide, default off): finished prompts' KV is pooled
+   * in fixed-size token blocks, and a later prompt sharing a token prefix
+   * (system prompt, multi-turn history) skips that part of prefill — same
+   * greedy tokens, much lower time-to-first-token. Vision-language and hybrid
+   * (Qwen3.5) models fail engine creation rather than silently ignoring it.
+   */
+  prefixCache?: boolean;
+  /** Prefix-pool block size in tokens (default 256; power of two, 16..4096). */
+  kvBlockSize?: number;
+  /** Prefix-pool RAM budget in bytes (default 1 GiB; negative = unbounded). */
+  kvPoolBytes?: number;
+  /**
+   * SSD spill directory: RAM-evicted prefix blocks persist here and survive
+   * engine restarts. Unset = no spill. Requires prefixCache.
+   */
+  kvSpillDir?: string;
+  /** Spill-directory disk budget in bytes (0/unset = unbounded). */
+  kvSpillBytes?: number;
 }
 
 export interface SamplingOptions {

bindings/node/src/addon.cc

@@ -253,6 +253,7 @@ class EngineWrap : public Napi::ObjectWrap<EngineWrap> {
 
     mlxforge_engine_opts2 opts = {};
     opts.struct_size = sizeof(opts);
+    std::string spill_dir;  // must outlive the create call (opts borrows it)
     if (info.Length() >= 2 && info[1].IsObject()) {
       Napi::Object o = info[1].As<Napi::Object>();
       if (o.Has("maxWaiting") && o.Get("maxWaiting").IsNumber())
@@ -261,6 +262,18 @@ class EngineWrap : public Napi::ObjectWrap<EngineWrap> {
         opts.kv_bits = o.Get("kvBits").As<Napi::Number>().Int32Value();
       if (o.Has("kvGroupSize") && o.Get("kvGroupSize").IsNumber())
         opts.kv_group_size = o.Get("kvGroupSize").As<Napi::Number>().Int32Value();
+      if (o.Has("prefixCache") && o.Get("prefixCache").IsBoolean())
+        opts.prefix_cache = o.Get("prefixCache").As<Napi::Boolean>().Value() ? 1 : 0;
+      if (o.Has("kvBlockSize") && o.Get("kvBlockSize").IsNumber())
+        opts.kv_block_size = o.Get("kvBlockSize").As<Napi::Number>().Int32Value();
+      if (o.Has("kvPoolBytes") && o.Get("kvPoolBytes").IsNumber())
+        opts.kv_pool_bytes = o.Get("kvPoolBytes").As<Napi::Number>().Int64Value();
+      if (o.Has("kvSpillDir") && o.Get("kvSpillDir").IsString()) {
+        spill_dir = o.Get("kvSpillDir").As<Napi::String>().Utf8Value();
+        opts.kv_spill_dir = spill_dir.c_str();
+      }
+      if (o.Has("kvSpillBytes") && o.Get("kvSpillBytes").IsNumber())
+        opts.kv_spill_bytes = o.Get("kvSpillBytes").As<Napi::Number>().Int64Value();
     }
 
     char* err = nullptr;

doc/applications.md

@@ -58,6 +58,11 @@ with environment-variable fallbacks (`server/config`):
 | `--max-waiting` | `MLXFORGE_MAX_WAITING` | `256` | Bounded waiting queue → `429` on overflow. |
 | `--kv-budget` | `MLXFORGE_KV_BUDGET` | `0` (unbounded) | KV-memory admission budget in bytes. |
 | `--kv-bits` | `MLXFORGE_KV_BITS` | `0` (fp16) | KV-cache quantization: `8` (~1.9× less cache memory, near-lossless) or `4` (~3.6×). Unsupported models (vision, hybrid Qwen3.5) fail startup rather than silently falling back. |
+| `--prefix-cache` | `MLXFORGE_PREFIX_CACHE` | `0` (off) | Reuse pooled KV across requests sharing a token prefix (system prompts, multi-turn). Same greedy tokens, ~20× lower warm TTFT on a 2k-token prefix. Unsupported models (vision, hybrid Qwen3.5) fail startup. |
+| `--kv-block` | `MLXFORGE_KV_BLOCK` | `256` | Prefix-pool block size in tokens (power of two, 16–4096). |
+| `--kv-pool` | `MLXFORGE_KV_POOL` | `1 GiB` | Prefix-pool RAM budget in bytes (LRU beyond it); `0` = unbounded. |
+| `--kv-spill-dir` | `MLXFORGE_KV_SPILL_DIR` | off | SSD spill directory: RAM-evicted prefix blocks persist here and survive restarts. |
+| `--kv-spill-bytes` | `MLXFORGE_KV_SPILL_BYTES` | `0` (unbounded) | Disk budget for the spill directory. |
 
 ### Logging
 
@@ -226,7 +231,8 @@ print(final.choices[0].message.content)
 
 ## The CLI: `mlxforge-cli`
 
-Subcommands (the main ones below; `embed` and `bench` also exist).
+Subcommands (the main ones below; `embed`, `bench`, and `bench-prefix` — the
+prefix-cache cold-vs-warm TTFT benchmark — also exist).
 
 ### `generate` — single-stream greedy generation
 

doc/architecture.md

@@ -204,6 +204,52 @@ The per-token budget figure adjusts accordingly: a K-or-V head row is
 `head_dim × bits/8` packed bytes plus a fp16 scale and bias per group (D=64/g=64:
 68 B at 8-bit, 36 B at 4-bit, vs 128 B fp16).
 
+## Prefix cache (block-pool KV storage, optional SSD tier)
+
+`--prefix-cache 1` (engine option `prefix_cache`; default off) reuses KV across
+requests that share a token prefix — the shared-system-prompt and multi-turn
+shapes. On a 2048-token shared prefix the warm TTFT drops ~20× (see
+`mlxforge-cli bench-prefix`); decode throughput is unchanged.
+
+The design is **gather-on-admit, not paged attention**: MLX has no paged SDPA
+kernel and mlx-lm has no paged reference to gate one against, so the decode
+batch stays the contiguous left-padded `BatchKVCache`. The *pages* live in a
+pool instead:
+
+- **BlockPool** (`cache/block_pool`): immutable `kv_block_size`-token blocks
+  (default 256; all layers, dense or quantized component vectors), keyed by a
+  salted **chain hash** of the token-id prefix — a key identifies the entire
+  prefix up to the block's end. LRU-evicted under `--kv-pool` bytes.
+- **PrefixCache** (`cache/prefix_cache`): longest-chain match (clamped to
+  `prompt_len - 1`: the last prompt token is always recomputed so admission
+  still yields next-token logits) + harvest policy. On admission, matched
+  blocks seed a batch-1 cache (`BatchKVCache::from_prefix`, written through
+  the standard block-grow writer so buffer layout matches a cold prefill) and
+  only the suffix is prefilled (`prefill_with_prefix`); the row then merges
+  into the decode batch like any single-row admission.
+- **Harvest is prompt-only.** When a row finishes, only its *prompt* span is
+  sealed into the pool. Decode-produced K/V differs from a recompute by fp16
+  accumulation order (the decode-vs-recompute gap) and demonstrably flips
+  later greedy choices; prefill-produced K/V is the proven exact-stable class,
+  so warm == cold stays token-exact. Multi-turn reuse still converges: the
+  next turn's prompt contains the prior answer as text, so its (seeded)
+  prefill recomputes that span once and pools it. Multimodal rows are never
+  harvested (a token-id hash cannot identify image content / 3D positions).
+- **SSD tier** (`cache/block_store`, `--kv-spill-dir`): RAM-evicted blocks are
+  serialized on the worker thread and written by a byte-only IO thread (one
+  0600 file per block, salted-hash name, tmp+rename); a pool miss revives the
+  file synchronously. The directory is rescanned at startup, so the cache
+  survives restarts; the salt (model fingerprint + storage config + block
+  size) is verified on load, so a block can never be revived for a different
+  model or quantization setting.
+
+Engine-wide, like `kv_bits` (the pool stores one storage layout); hybrid
+(Qwen3.5) and vision-language models reject the option at engine creation. The
+gate is **warm == cold**: reuse may change speed, never tokens — no new
+mlx-lm fixtures are needed because the cold path is already golden-gated and
+prefix reuse is an engine-internal equivalence property
+(`tests/scheduler/prefix_reuse_test.cpp`, `tests/scheduler/prefix_spill_test.cpp`).
+
 ## Module map
 
 Source lives under `src/`, grouped by responsibility. Tests mirror the module
@@ -221,6 +267,9 @@ path under `tests/`.
 | `cache/batch_kv_cache` | Batched, left-padded KV cache: `update_and_fetch`, `filter` (evict), `merge` (admit), `pad_dummies` (bucketing). |
 | `cache/kv_quant` | Quantized-KV shared types (`KVQuantConfig`, triplets) + the block-grow component writer both caches use. |
 | `cache/kv_budget` | KV memory projection / admission gate (fp16 and quantized accounting). |
+| `cache/block_pool` | Prefix-cache pages: immutable KV blocks keyed by a salted chain hash, LRU under a byte budget. |
+| `cache/prefix_cache` | Longest-prefix block matching + prompt-only harvest over the pool. |
+| `cache/block_store` | SSD spill tier: byte-only writer thread, salted/versioned block files, restart rescan. |
 | `model/sdpa` | Cache-aware SDPA dispatch: dense fast kernel vs the hand-rolled quantized path (mlx-lm port). |
 | `sample/sampler` | greedy / temperature / top-k / top-p, all as MLX graph ops. |
 | `scheduler/request` | The `Request` struct and the bounded, blocking `TokenQueue`. |