diff --git a/CLAUDE.md b/CLAUDE.md index a8d7c44..2191a4e 100644 --- a/CLAUDE.md +++ b/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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 1821229..5a3e166 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/apps/mlxforge.cpp b/apps/mlxforge.cpp index 5d8378a..3026f2a 100644 --- a/apps/mlxforge.cpp +++ b/apps/mlxforge.cpp @@ -73,12 +73,18 @@ void print_help() { " --max-waiting max queued requests (default 256)\n" " --kv-budget KV cache budget in bytes, 0 = unbounded (default 0)\n" " --kv-bits 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 prefix-pool block size in tokens (default 256)\n" + " --kv-pool prefix-pool RAM budget in bytes, 0 = unbounded (default 1 GiB)\n" + " --kv-spill-dir SSD spill dir for evicted prefix blocks (default off)\n" + " --kv-spill-bytes 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(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); diff --git a/apps/mlxforge_cli.cpp b/apps/mlxforge_cli.cpp index 9db2a2c..2fbdfca 100644 --- a/apps/mlxforge_cli.cpp +++ b/apps/mlxforge_cli.cpp @@ -15,6 +15,11 @@ // mlxforge-cli bench [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 [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 [--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 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 prefix; + while (static_cast(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 ids) { + auto req = std::make_shared(); + 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(t_first - t0).count(); + } + } + const double decode_s = + std::chrono::duration(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 ids = prefix; + const std::vector 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 [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). diff --git a/bindings/node/index.d.ts b/bindings/node/index.d.ts index 7631eea..db83cae 100644 --- a/bindings/node/index.d.ts +++ b/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 { diff --git a/bindings/node/src/addon.cc b/bindings/node/src/addon.cc index af3a3a2..d7b8b79 100644 --- a/bindings/node/src/addon.cc +++ b/bindings/node/src/addon.cc @@ -253,6 +253,7 @@ class EngineWrap : public Napi::ObjectWrap { 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(); if (o.Has("maxWaiting") && o.Get("maxWaiting").IsNumber()) @@ -261,6 +262,18 @@ class EngineWrap : public Napi::ObjectWrap { opts.kv_bits = o.Get("kvBits").As().Int32Value(); if (o.Has("kvGroupSize") && o.Get("kvGroupSize").IsNumber()) opts.kv_group_size = o.Get("kvGroupSize").As().Int32Value(); + if (o.Has("prefixCache") && o.Get("prefixCache").IsBoolean()) + opts.prefix_cache = o.Get("prefixCache").As().Value() ? 1 : 0; + if (o.Has("kvBlockSize") && o.Get("kvBlockSize").IsNumber()) + opts.kv_block_size = o.Get("kvBlockSize").As().Int32Value(); + if (o.Has("kvPoolBytes") && o.Get("kvPoolBytes").IsNumber()) + opts.kv_pool_bytes = o.Get("kvPoolBytes").As().Int64Value(); + if (o.Has("kvSpillDir") && o.Get("kvSpillDir").IsString()) { + spill_dir = o.Get("kvSpillDir").As().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().Int64Value(); } char* err = nullptr; diff --git a/doc/applications.md b/doc/applications.md index 9c43c00..02c6085 100644 --- a/doc/applications.md +++ b/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 diff --git a/doc/architecture.md b/doc/architecture.md index bb0aa6e..1d31d83 100644 --- a/doc/architecture.md +++ b/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`. | diff --git a/doc/embedding.md b/doc/embedding.md index fca9f3d..2fb9bf4 100644 --- a/doc/embedding.md +++ b/doc/embedding.md @@ -92,8 +92,12 @@ typedef struct mlxforge_request mlxforge_request; mlxforge_engine* eng = mlxforge_engine_create("mlx-community/Llama-3.2-1B-Instruct-4bit", /*opts=*/NULL, &err); // Or with extended options (ABI v6+): a quantized KV cache cuts the dominant -// growing allocation ~1.9x (8-bit, near-lossless) or ~3.6x (4-bit). -// mlxforge_engine_opts2 opts = { .struct_size = sizeof(opts), .kv_bits = 8 }; +// growing allocation ~1.9x (8-bit, near-lossless) or ~3.6x (4-bit), and the +// prefix cache (v7+) skips re-prefilling shared prompt prefixes — ~20x lower +// warm TTFT on a 2k-token system prompt, same greedy tokens. kv_spill_dir adds +// an SSD tier that survives engine restarts. +// mlxforge_engine_opts2 opts = { .struct_size = sizeof(opts), .kv_bits = 8, +// .prefix_cache = 1 }; // eng = mlxforge_engine_create2(spec, &opts, &err); while (!mlxforge_engine_ready(eng)) { /* model still loading on the worker thread */ } diff --git a/src/cache/batch_kv_cache.cpp b/src/cache/batch_kv_cache.cpp index dd420de..5f5b514 100644 --- a/src/cache/batch_kv_cache.cpp +++ b/src/cache/batch_kv_cache.cpp @@ -5,6 +5,8 @@ #include #include +#include "cache/block_pool.h" + #include "mlx/ops.h" #include "mlx/transforms.h" @@ -45,6 +47,41 @@ BatchKVCache BatchKVCache::from_single_sequence( return c; } +BatchKVCache BatchKVCache::from_prefix(int n_layers, + const std::vector>& blocks, + int len, KVQuantConfig qcfg) { + BatchKVCache c(n_layers, std::vector{0}, qcfg); // batch 1, no left padding + if (len <= 0 || blocks.empty()) return c; + + // Per layer, stitch the blocks into one [0, len) span and write it through + // the standard writer (one chunk at position 0) so capacity rounding and + // buffer layout match a normal prefill. `len` may stop short of the last + // block's end (the prompt's final token is always recomputed). + for (int l = 0; l < n_layers; ++l) { + std::vector k_in, v_in; + const std::size_t n_comp = blocks[0]->k[l].size(); + for (std::size_t i = 0; i < n_comp; ++i) { + std::vector kp, vp; + kp.reserve(blocks.size()); + vp.reserve(blocks.size()); + for (const auto& b : blocks) { + kp.push_back(b->k[l][i]); + vp.push_back(b->v[l][i]); + } + mx::array kj = kp.size() == 1 ? kp[0] : mx::concatenate(kp, /*axis=*/2); + mx::array vj = vp.size() == 1 ? vp[0] : mx::concatenate(vp, /*axis=*/2); + k_in.push_back(slice_seq(kj, 0, len)); + v_in.push_back(slice_seq(vj, 0, len)); + } + update_kv_components(c.keys_[l], k_in, /*prev=*/0, kStep); + update_kv_components(c.values_[l], v_in, /*prev=*/0, kStep); + } + c.idx_ = len; + c.offset_ = mx::array(&len, {1}, mx::int32); + mx::eval(c.offset_); + return c; +} + int BatchKVCache::s_cap() const { return keys_[0].empty() ? 0 : keys_[0][0].shape()[2]; } @@ -141,6 +178,24 @@ int scalar_int(const mx::array& a) { } } // namespace +std::pair, std::vector> BatchKVCache::fetch_row_components( + int layer, int row, int left_pad, int len) const { + auto row_slice = [&](const mx::array& c) { + const auto& s = c.shape(); + return mx::slice(c, {row, 0, left_pad, 0}, {row + 1, s[1], left_pad + len, s[3]}); + }; + std::vector k, v; + for (const auto& c : keys_[layer]) k.push_back(row_slice(c)); + for (const auto& c : values_[layer]) v.push_back(row_slice(c)); + return {std::move(k), std::move(v)}; +} + +std::vector BatchKVCache::left_padding_host() const { + mx::array c = mx::contiguous(left_padding_); + mx::eval(c); + return std::vector(c.data(), c.data() + c.size()); +} + void BatchKVCache::filter(const std::vector& keep) { mx::array idxs(keep.data(), {static_cast(keep.size())}, mx::int32); for (int l = 0; l < static_cast(keys_.size()); ++l) { diff --git a/src/cache/batch_kv_cache.h b/src/cache/batch_kv_cache.h index 3a0195a..f0fbb5b 100644 --- a/src/cache/batch_kv_cache.h +++ b/src/cache/batch_kv_cache.h @@ -20,6 +20,7 @@ // zero-filled pad regions dequantize to exactly 0 and are masked anyway. #pragma once +#include #include #include #include @@ -32,6 +33,8 @@ namespace mlxforge { namespace mx = mlx::core; +struct KVBlock; // cache/block_pool.h + class BatchKVCache { public: static constexpr int kStep = 256; @@ -52,6 +55,16 @@ class BatchKVCache { static BatchKVCache from_single_sequence( std::vector> kv_per_layer, int seq, int decode_offset); + // Build a batch-1 cache whose first `len` positions are the given prefix-pool + // blocks' K/V (consecutive from position 0; `len` may stop short of the last + // block's end — the prompt's final token is always recomputed). The blocks + // are written through the standard block-grow writer (update_kv_components) + // so the buffer layout matches a normal prefill chunk's; RoPE offset = len, + // no left padding. The suffix prefill then appends at idx == len. + static BatchKVCache from_prefix(int n_layers, + const std::vector>& blocks, + int len, KVQuantConfig qcfg = {}); + int batch_size() const { return batch_; } int idx() const { return idx_; } // populated sequence length (_idx) // Allocated capacity along the sequence axis (0 before the first write). @@ -89,6 +102,16 @@ class BatchKVCache { // inspection/tests — the model attends over the triplets directly). std::pair fetch_dequantized(int layer) const; + // One row's populated K/V component vectors for a layer, as + // (1, n_kv_heads, len, comp_dim) views over physical slots + // [left_pad, left_pad + len). For harvesting a finished row into the prefix + // pool (which materializes its own copies — these are lazy views). + std::pair, std::vector> fetch_row_components( + int layer, int row, int left_pad, int len) const; + + // Host copy of the per-row left padding (small, already materialized). + std::vector left_padding_host() const; + // Eviction: keep only the given batch rows (take on axis 0) across every // layer's K/V plus offset/left_padding, then shift off any common left // padding. `keep` indexes the current batch rows. diff --git a/src/cache/block_pool.cpp b/src/cache/block_pool.cpp new file mode 100644 index 0000000..7081ce7 --- /dev/null +++ b/src/cache/block_pool.cpp @@ -0,0 +1,61 @@ +#include "cache/block_pool.h" + +namespace mlxforge { + +uint64_t fnv1a(const void* data, std::size_t n, uint64_t seed) { + const auto* p = static_cast(data); + uint64_t h = seed; + for (std::size_t i = 0; i < n; ++i) { + h ^= p[i]; + h *= 1099511628211ull; + } + return h; +} + +uint64_t chain_hash(uint64_t prev, const int* ids, int n) { + // Fold the previous block's key in first so the result identifies the whole + // prefix, then the block's own ids. + uint64_t h = fnv1a(&prev, sizeof(prev)); + return fnv1a(ids, static_cast(n) * sizeof(int), h); +} + +std::size_t block_bytes(const KVBlock& b) { + std::size_t n = 0; + for (const auto& layer : b.k) + for (const auto& c : layer) n += c.nbytes(); + for (const auto& layer : b.v) + for (const auto& c : layer) n += c.nbytes(); + return n; +} + +std::shared_ptr BlockPool::get(uint64_t hash) { + auto it = map_.find(hash); + if (it == map_.end()) return nullptr; + lru_.splice(lru_.begin(), lru_, it->second.lru_it); // bump to most recent + return it->second.block; +} + +void BlockPool::insert(uint64_t hash, std::shared_ptr block) { + if (map_.count(hash) != 0) return; // immutable content; first write wins + // A block that alone exceeds the budget would just evict the whole pool and + // then be evicted itself on the next insert — don't admit it. + if (budget_ != 0 && block->bytes > budget_) return; + lru_.push_front(hash); + bytes_ += block->bytes; + map_.emplace(hash, Entry{std::move(block), lru_.begin()}); + evict_to_budget(); +} + +void BlockPool::evict_to_budget() { + if (budget_ == 0) return; + while (bytes_ > budget_ && !lru_.empty()) { + const uint64_t victim = lru_.back(); + auto it = map_.find(victim); + if (on_evict_) on_evict_(victim, it->second.block); + bytes_ -= it->second.block->bytes; + map_.erase(it); + lru_.pop_back(); + } +} + +} // namespace mlxforge diff --git a/src/cache/block_pool.h b/src/cache/block_pool.h new file mode 100644 index 0000000..75b936f --- /dev/null +++ b/src/cache/block_pool.h @@ -0,0 +1,95 @@ +// Block-pool ("paged") KV storage for prefix caching. +// +// The decode batch itself stays contiguous (BatchKVCache) — MLX has no paged +// SDPA kernel and the golden reference (mlx-lm) has none to gate one against. +// Pages live here instead: immutable KVBlocks of `block_size` tokens (all +// layers, all storage components), keyed by a chained hash of the token-id +// prefix. On admission matched blocks are gathered (copied) into the row's +// cache; on eviction a finished row's K/V is harvested back into the pool. +// Unified memory makes those copies cheap next to the prefill they replace. +// +// A block's key hashes its OWN ids chained onto the previous block's key, so a +// key identifies the entire token prefix up to the block's end — two prompts +// share a block only if they share every token before it. Keys are salted +// (model fingerprint + storage config) so persisted blocks can never cross +// models or quantization settings. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "mlx/array.h" + +namespace mlxforge { + +namespace mx = mlx::core; + +// Chain `n` token ids onto `prev` (FNV-1a-64). Deterministic across runs so +// hashes can key on-disk blocks. +uint64_t chain_hash(uint64_t prev, const int* ids, int n); + +// FNV-1a-64 of a byte string; used to derive the pool's salt (the chain seed) +// from the model fingerprint + storage config. +uint64_t fnv1a(const void* data, std::size_t n, uint64_t seed = 14695981039346656037ull); + +// One block_size-token span of cached K/V: per layer the same component vector +// the caches store (1 array dense fp16, 3 quantized), each +// (1, n_kv_heads, block_size, comp_dim). Immutable once pooled. +struct KVBlock { + std::vector> k; // [layer][component] + std::vector> v; + std::size_t bytes = 0; // summed component bytes (LRU budget accounting) +}; + +// Sum of all component buffer sizes — the block's RAM cost. +std::size_t block_bytes(const KVBlock& b); + +// hash -> KVBlock map with LRU eviction under a byte budget. Worker-thread-only +// (it holds MLX arrays, which are thread-bound); cross-thread visibility goes +// through the Worker's atomics, never through this object. +class BlockPool { + public: + // budget_bytes == 0 means "unbounded" (matching the kv_budget convention). + explicit BlockPool(std::size_t budget_bytes) : budget_(budget_bytes) {} + + // Look up a block, bumping it to most-recently-used. nullptr on miss. + std::shared_ptr get(uint64_t hash); + + // Membership test without touching recency (harvest uses it to skip slicing + // blocks that are already pooled). + bool contains(uint64_t hash) const { return map_.count(hash) != 0; } + + // Insert (no-op if present), then evict LRU entries until under budget. A + // single block larger than the whole budget is not admitted at all. In-flight + // references (shared_ptr held by an admission) survive eviction. + void insert(uint64_t hash, std::shared_ptr block); + + // Phase-2 spill hook: called with each evicted (hash, block) before it is + // dropped from RAM, so an SSD tier can persist it. + using EvictFn = std::function&)>; + void set_evict_hook(EvictFn fn) { on_evict_ = std::move(fn); } + + std::size_t bytes() const { return bytes_; } + std::size_t size() const { return map_.size(); } + std::size_t budget_bytes() const { return budget_; } + + private: + void evict_to_budget(); + + struct Entry { + std::shared_ptr block; + std::list::iterator lru_it; + }; + std::list lru_; // front = most recent, back = eviction victim + std::unordered_map map_; + std::size_t budget_; + std::size_t bytes_ = 0; + EvictFn on_evict_; +}; + +} // namespace mlxforge diff --git a/src/cache/block_store.cpp b/src/cache/block_store.cpp new file mode 100644 index 0000000..09e7196 --- /dev/null +++ b/src/cache/block_store.cpp @@ -0,0 +1,297 @@ +#include "cache/block_store.h" + +#include + +#include +#include +#include +#include +#include + +#include "core/logging.h" + +#include "mlx/ops.h" +#include "mlx/transforms.h" + +namespace fs = std::filesystem; + +namespace mlxforge { + +namespace { + +constexpr uint64_t kMagic = 0x3142564b46584c4dull; // "MLXFKVB1" little-endian +constexpr uint32_t kVersion = 1; + +// dtype codes in the file format. Only the two storage dtypes exist: dense and +// quantized scales/biases are fp16, quantized packed words are uint32. +uint32_t dtype_code(const mx::array& a) { + if (a.dtype() == mx::float16) return 0; + if (a.dtype() == mx::uint32) return 1; + throw std::logic_error("serialize_block: unexpected KV storage dtype"); +} + +template +void append(std::vector& out, const T& v) { + const char* p = reinterpret_cast(&v); + out.insert(out.end(), p, p + sizeof(T)); +} + +// Bounds-checked sequential reader over the serialized buffer. +struct Reader { + const char* p; + std::size_t left; + template + bool read(T& v) { + if (left < sizeof(T)) return false; + std::memcpy(&v, p, sizeof(T)); + p += sizeof(T); + left -= sizeof(T); + return true; + } +}; + +} // namespace + +std::vector serialize_block(const KVBlock& block, uint64_t salt) { + // Defensive contiguous+eval before data(): pooled blocks are already + // materialized contiguous copies, but a lazy view here would silently + // serialize the wrong elements. Order matches deserialize_block: per layer, + // K components then V components. + std::vector comps; + for (std::size_t l = 0; l < block.k.size(); ++l) { + for (const auto& c : block.k[l]) comps.push_back(mx::contiguous(c)); + for (const auto& c : block.v[l]) comps.push_back(mx::contiguous(c)); + } + mx::eval(comps); + + std::vector out; + out.reserve(64 + block_bytes(block)); + append(out, kMagic); + append(out, kVersion); + append(out, salt); + append(out, static_cast(block.k.size())); + append(out, static_cast(block.k.empty() ? 0 : block.k[0].size())); + append(out, static_cast(block.v.empty() ? 0 : block.v[0].size())); + for (const auto& c : comps) { + append(out, dtype_code(c)); + for (int d = 0; d < 4; ++d) append(out, static_cast(c.shape()[d])); + append(out, static_cast(c.nbytes())); + const char* p = c.data(); + out.insert(out.end(), p, p + c.nbytes()); + } + return out; +} + +std::shared_ptr deserialize_block(const std::vector& bytes, uint64_t salt) { + Reader r{bytes.data(), bytes.size()}; + uint64_t magic = 0, file_salt = 0; + uint32_t version = 0, n_layers = 0, k_comps = 0, v_comps = 0; + if (!r.read(magic) || magic != kMagic) return nullptr; + if (!r.read(version) || version != kVersion) return nullptr; + if (!r.read(file_salt) || file_salt != salt) return nullptr; + if (!r.read(n_layers) || !r.read(k_comps) || !r.read(v_comps)) return nullptr; + if (n_layers == 0 || n_layers > 4096 || k_comps > 3 || v_comps > 3) return nullptr; + + auto read_comp = [&](mx::array& out) { + uint32_t code = 0; + int32_t shape[4]; + uint64_t nbytes = 0; + if (!r.read(code) || code > 1) return false; + for (int d = 0; d < 4; ++d) { + if (!r.read(shape[d]) || shape[d] <= 0) return false; + } + if (!r.read(nbytes) || r.left < nbytes) return false; + const mx::Shape s{shape[0], shape[1], shape[2], shape[3]}; + const std::size_t elems = + static_cast(shape[0]) * shape[1] * shape[2] * shape[3]; + if (code == 0) { + if (nbytes != elems * sizeof(uint16_t)) return false; + out = mx::array(reinterpret_cast(r.p), s); + } else { + if (nbytes != elems * sizeof(uint32_t)) return false; + out = mx::array(reinterpret_cast(r.p), s); + } + r.p += nbytes; + r.left -= nbytes; + return true; + }; + + auto block = std::make_shared(); + block->k.resize(n_layers); + block->v.resize(n_layers); + for (uint32_t l = 0; l < n_layers; ++l) { + for (uint32_t i = 0; i < k_comps; ++i) { + mx::array c = mx::zeros({0}); + if (!read_comp(c)) return nullptr; + block->k[l].push_back(std::move(c)); + } + for (uint32_t i = 0; i < v_comps; ++i) { + mx::array c = mx::zeros({0}); + if (!read_comp(c)) return nullptr; + block->v[l].push_back(std::move(c)); + } + } + block->bytes = block_bytes(*block); + return block; +} + +BlockStore::BlockStore(std::string dir, std::size_t budget_bytes, uint64_t salt) + : dir_(std::move(dir)), budget_(budget_bytes), salt_(salt) { + fs::create_directories(dir_); + // Rescan surviving blocks, LRU-seeded by modification time so the budget + // evicts the stalest first. Unparseable names are ignored (foreign files). + struct Found { + fs::file_time_type mtime; + uint64_t hash; + std::size_t bytes; + }; + std::vector found; + for (const auto& e : fs::directory_iterator(dir_)) { + if (!e.is_regular_file() || e.path().extension() != ".kvb") continue; + uint64_t hash = 0; + try { + hash = std::stoull(e.path().stem().string(), nullptr, 16); + } catch (...) { + continue; + } + found.push_back({e.last_write_time(), hash, static_cast(e.file_size())}); + } + std::sort(found.begin(), found.end(), + [](const Found& a, const Found& b) { return a.mtime < b.mtime; }); + for (const Found& f : found) { + index_[f.hash] = Entry{f.bytes, ++stamp_}; + bytes_ += f.bytes; + } + log::info("block store: {} ({} blocks, {} bytes on disk)", dir_, index_.size(), bytes_); + writer_ = std::thread([this] { writer_loop(); }); +} + +BlockStore::~BlockStore() { + { + std::lock_guard lk(m_); + stop_ = true; + } + cv_.notify_all(); + if (writer_.joinable()) writer_.join(); +} + +std::string BlockStore::path_of(uint64_t hash) const { + char name[32]; + std::snprintf(name, sizeof(name), "%016llx.kvb", static_cast(hash)); + return dir_ + "/" + name; +} + +void BlockStore::put(uint64_t hash, std::vector bytes) { + { + std::lock_guard lk(m_); + if (index_.count(hash) != 0) return; // already on disk + queue_.emplace_back(hash, std::move(bytes)); + } + cv_.notify_one(); +} + +std::optional> BlockStore::get(uint64_t hash) { + { + std::lock_guard lk(m_); + auto it = index_.find(hash); + if (it == index_.end()) { + // Not on disk yet — it may still be in flight in the write queue (an + // immediate re-request right after a spill). Serve it from there so the + // async writer can never lose a hit. + for (const auto& [h, bytes] : queue_) { + if (h == hash) return bytes; + } + return std::nullopt; + } + it->second.stamp = ++stamp_; + } + std::ifstream f(path_of(hash), std::ios::binary | std::ios::ate); + if (!f) return std::nullopt; + std::vector bytes(static_cast(f.tellg())); + f.seekg(0); + f.read(bytes.data(), static_cast(bytes.size())); + if (!f) return std::nullopt; + return bytes; +} + +bool BlockStore::contains(uint64_t hash) const { + std::lock_guard lk(m_); + if (index_.count(hash) != 0) return true; + for (const auto& [h, _] : queue_) { + if (h == hash) return true; + } + return false; +} + +std::size_t BlockStore::bytes() const { + std::lock_guard lk(m_); + return bytes_; +} + +std::size_t BlockStore::size() const { + std::lock_guard lk(m_); + return index_.size(); +} + +void BlockStore::writer_loop() { + for (;;) { + std::pair> job; + { + std::unique_lock lk(m_); + cv_.wait(lk, [this] { return stop_ || !queue_.empty(); }); + if (queue_.empty()) return; // stop_ and drained + // Copy (don't pop): the entry must stay visible to get()/contains() + // while the file is being written, or an immediate re-request would + // land in the gap and silently miss. + job = queue_.front(); + } + // Write to a temp name then rename so a crash never leaves a torn block, + // and 0600 the file: the cache is conversation content. No fsync — this is + // a cache, not a durability contract. + const std::string path = path_of(job.first); + const std::string tmp = path + ".tmp"; + bool ok = false; + { + std::ofstream f(tmp, std::ios::binary | std::ios::trunc); + if (f) { + f.write(job.second.data(), static_cast(job.second.size())); + ok = static_cast(f); + } + } + if (ok) { + ::chmod(tmp.c_str(), 0600); + std::error_code ec; + fs::rename(tmp, path, ec); + ok = !ec; + } + if (!ok) { + log::warn("block store: failed to persist {}", path); + std::error_code ec; + fs::remove(tmp, ec); + } + + std::lock_guard lk(m_); + if (ok) { + index_[job.first] = Entry{job.second.size(), ++stamp_}; + bytes_ += job.second.size(); + evict_to_budget_locked(); + } + queue_.pop_front(); // the index entry (or the drop) is now authoritative + } +} + +void BlockStore::evict_to_budget_locked() { + if (budget_ == 0) return; + while (bytes_ > budget_ && index_.size() > 1) { + auto victim = index_.begin(); + for (auto it = index_.begin(); it != index_.end(); ++it) { + if (it->second.stamp < victim->second.stamp) victim = it; + } + std::error_code ec; + fs::remove(path_of(victim->first), ec); + bytes_ -= victim->second.bytes; + index_.erase(victim); + } +} + +} // namespace mlxforge diff --git a/src/cache/block_store.h b/src/cache/block_store.h new file mode 100644 index 0000000..435c6a7 --- /dev/null +++ b/src/cache/block_store.h @@ -0,0 +1,97 @@ +// SSD spill tier for the prefix-cache block pool. +// +// Blocks evicted from the RAM pool (BlockPool's evict hook) are serialized and +// written to one file per block (.kvb) under a spill directory; a pool +// miss loads the file back (PrefixCache's miss hook). The directory is +// rescanned at construction, so the prefix cache survives engine restarts. +// +// Threading (the MLX thread-bound rule): BlockStore itself never touches MLX +// arrays — its writer thread and file index handle only raw byte buffers. The +// array <-> bytes conversions (serialize_block / deserialize_block) must run on +// the worker thread, which owns every pooled array. Writes are asynchronous +// (queued to the writer thread); reads are synchronous on the caller's thread — +// an SSD read of a few MB replaces a far more expensive prefill. +// +// On-disk format (version 1), little-endian, one block per file: +// magic u64 ("MLXFKVB1"), version u32, salt u64, n_layers u32, +// k_comps u32, v_comps u32, then per layer, K components then V components: +// dtype u32 (0 = float16, 1 = uint32), shape i32[4], nbytes u64, raw bytes. +// The salt (model fingerprint + storage config + block size) is verified on +// load, and block keys are salted, so a file can never be revived for a +// different model or quantization setting. Files are created 0600 — the cache +// holds conversation content. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cache/block_pool.h" + +namespace mlxforge { + +// KVBlock -> bytes. Worker thread only (reads array buffers; defensively +// contiguous+eval'd first, per the data()-on-views gotcha). +std::vector serialize_block(const KVBlock& block, uint64_t salt); + +// bytes -> KVBlock. Worker thread only (creates MLX arrays). Returns nullptr on +// any mismatch (magic, version, salt) or truncation — the caller treats it as +// a plain miss, never an error. +std::shared_ptr deserialize_block(const std::vector& bytes, uint64_t salt); + +class BlockStore { + public: + // Creates `dir` if needed and rescans existing *.kvb files into the index + // (LRU-seeded by file modification time). budget_bytes == 0 means unbounded. + BlockStore(std::string dir, std::size_t budget_bytes, uint64_t salt); + ~BlockStore(); // drains queued writes, then joins the writer thread + + BlockStore(const BlockStore&) = delete; + BlockStore& operator=(const BlockStore&) = delete; + + // Queue a serialized block for writing (any thread; the writer thread does + // the file IO, then enforces the disk budget by deleting LRU files). + void put(uint64_t hash, std::vector bytes); + + // Read a block's bytes back (any thread, synchronous). Bumps LRU recency. + // nullopt on miss or unreadable file. + std::optional> get(uint64_t hash); + + bool contains(uint64_t hash) const; + std::size_t bytes() const; + std::size_t size() const; + + private: + void writer_loop(); + std::string path_of(uint64_t hash) const; + // Index mutation under m_; eviction scans for the smallest recency stamp + // (the index stays small enough that a linear scan beats bookkeeping). + void evict_to_budget_locked(); + + struct Entry { + std::size_t bytes = 0; + uint64_t stamp = 0; // recency: higher = more recently used + }; + + const std::string dir_; + const std::size_t budget_; + const uint64_t salt_; + + mutable std::mutex m_; + std::condition_variable cv_; + std::unordered_map index_; + std::size_t bytes_ = 0; + uint64_t stamp_ = 0; + std::deque>> queue_; + bool stop_ = false; + std::thread writer_; +}; + +} // namespace mlxforge diff --git a/src/cache/prefix_cache.cpp b/src/cache/prefix_cache.cpp new file mode 100644 index 0000000..52a4f31 --- /dev/null +++ b/src/cache/prefix_cache.cpp @@ -0,0 +1,84 @@ +#include "cache/prefix_cache.h" + +#include + +#include "cache/kv_quant.h" + +#include "mlx/ops.h" +#include "mlx/transforms.h" + +namespace mlxforge { + +PrefixCache::Match PrefixCache::match(const std::vector& ids) { + Match m; + const int bs = cfg_.block_size; + const int n_full = static_cast(ids.size()) / bs; + uint64_t h = cfg_.salt; + for (int b = 0; b < n_full; ++b) { + h = chain_hash(h, ids.data() + static_cast(b) * bs, bs); + std::shared_ptr blk = pool_.get(h); + if (!blk && on_miss_) { + blk = on_miss_(h); // SSD tier + // Re-promote: the revived block is hot again. (A block bigger than the + // whole pool budget stays un-pooled; the shared_ptr still serves this + // match.) + if (blk) pool_.insert(h, blk); + } + if (!blk) break; // keys chain, so the first miss ends every longer match + m.blocks.push_back(std::move(blk)); + } + m.tokens = std::max(0, std::min(static_cast(m.blocks.size()) * bs, + static_cast(ids.size()) - 1)); + return m; +} + +void PrefixCache::harvest(const std::vector& ids, int len, int n_layers, + const LayerFetch& fetch) { + const int bs = cfg_.block_size; + const int n_full = std::min(len, static_cast(ids.size())) / bs; + if (n_full == 0) return; + + // Chain the keys, keeping only blocks the pool doesn't already hold. + std::vector hashes; + std::vector fresh; + uint64_t h = cfg_.salt; + for (int b = 0; b < n_full; ++b) { + h = chain_hash(h, ids.data() + static_cast(b) * bs, bs); + hashes.push_back(h); + if (!pool_.contains(h)) fresh.push_back(b); + } + if (fresh.empty()) return; + + // Slice each fresh block out of the row and materialize it (a lazy slice + // would pin the whole batch buffer past the row's eviction), batching one + // eval over every new component. + std::vector> blocks(fresh.size()); + for (auto& b : blocks) { + b = std::make_shared(); + b->k.resize(n_layers); + b->v.resize(n_layers); + } + std::vector to_eval; + for (int l = 0; l < n_layers; ++l) { + auto [kc, vc] = fetch(l); + for (std::size_t j = 0; j < fresh.size(); ++j) { + const int start = fresh[j] * bs; + for (const auto& c : kc) { + blocks[j]->k[l].push_back(mx::contiguous(slice_seq(c, start, start + bs))); + to_eval.push_back(blocks[j]->k[l].back()); + } + for (const auto& c : vc) { + blocks[j]->v[l].push_back(mx::contiguous(slice_seq(c, start, start + bs))); + to_eval.push_back(blocks[j]->v[l].back()); + } + } + } + mx::eval(to_eval); + + for (std::size_t j = 0; j < fresh.size(); ++j) { + blocks[j]->bytes = block_bytes(*blocks[j]); + pool_.insert(hashes[fresh[j]], std::move(blocks[j])); + } +} + +} // namespace mlxforge diff --git a/src/cache/prefix_cache.h b/src/cache/prefix_cache.h new file mode 100644 index 0000000..db786c7 --- /dev/null +++ b/src/cache/prefix_cache.h @@ -0,0 +1,87 @@ +// Prefix cache: longest-prefix matching + harvest policy over a BlockPool. +// +// match() turns a prompt's token ids into the longest chain of consecutive +// cached full blocks; the worker seeds the row's cache with them and prefills +// only the suffix. harvest() seals a finished row's K/V back into the pool so +// the next turn of the same conversation (or the next request sharing the +// system prompt) hits. Both run on the worker thread only (MLX arrays are +// thread-bound; see BlockPool). +// +// At least the prompt's last token is always recomputed (cached_len is clamped +// to prompt_len - 1) so the admission still produces next-token logits — the +// same rule vLLM/SGLang use. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "cache/block_pool.h" + +namespace mlxforge { + +// Engine-wide prefix-cache setting (the pool stores the engine's single storage +// layout, so this cannot be per-request — same reasoning as KVQuantConfig). +struct PrefixCacheConfig { + bool enabled = false; + // Hashing/pooling granularity in tokens. Smaller blocks match shorter shared + // prefixes but cost more hash/slice overhead per token. Power of two in + // [16, 4096] (validated at engine creation). + int block_size = 256; + // RAM budget for pooled blocks (LRU beyond it); 0 = unbounded. + std::size_t pool_bytes = 1ull << 30; + // Salt folded into every block key: model fingerprint + storage config, so + // persisted blocks can never cross models or quantization settings. + uint64_t salt = 0; + // SSD spill tier: RAM-evicted blocks persist under this directory and are + // reloaded on a pool miss (also across engine restarts). Empty = no spill. + std::string spill_dir; + // Disk budget for spilled blocks (LRU-deleted beyond it); 0 = unbounded. + std::size_t spill_bytes = 0; +}; + +class PrefixCache { + public: + explicit PrefixCache(PrefixCacheConfig cfg) : cfg_(cfg), pool_(cfg.pool_bytes) {} + + struct Match { + std::vector> blocks; // consecutive from position 0 + int tokens = 0; // cached token count to reuse (<= blocks * block_size, >= 0) + }; + + // Longest run of consecutive cached full blocks covering ids[0..], clamped to + // ids.size() - 1 tokens. Hits bump the blocks' LRU recency. + Match match(const std::vector& ids); + + // Per-layer accessor a harvest caller provides: the row's populated K/V + // component vectors, each component (1, n_kv_heads, len, comp_dim) covering + // ids[0..len). Views are fine — harvest materializes its own copies. + using LayerFetch = + std::function, std::vector>(int layer)>; + + // Seal every full block of ids[0..len) not already pooled and insert it. + // Slices are materialized (mx::contiguous + one eval over all new blocks) so + // the pool never pins the batch cache's buffers. + void harvest(const std::vector& ids, int len, int n_layers, const LayerFetch& fetch); + + // Second-level lookup consulted when the RAM pool misses (the SSD tier): + // returns the revived block or nullptr. The revived block is re-promoted + // into the pool. Runs inside match() on the worker thread. + using MissFn = std::function(uint64_t)>; + void set_miss_fn(MissFn fn) { on_miss_ = std::move(fn); } + + const PrefixCacheConfig& config() const { return cfg_; } + std::size_t pool_bytes() const { return pool_.bytes(); } + std::size_t pool_blocks() const { return pool_.size(); } + BlockPool& pool() { return pool_; } // spill (evict) hook installation + + private: + PrefixCacheConfig cfg_; + BlockPool pool_; + MissFn on_miss_; +}; + +} // namespace mlxforge diff --git a/src/capi/mlxforge.cpp b/src/capi/mlxforge.cpp index 9156457..f837174 100644 --- a/src/capi/mlxforge.cpp +++ b/src/capi/mlxforge.cpp @@ -160,6 +160,18 @@ mlxforge_engine* mlxforge_engine_create2(const char* model_spec, if (covered(&opts->kv_bits + 1)) cfg.kv_bits = opts->kv_bits; if (covered(&opts->kv_group_size + 1) && opts->kv_group_size > 0) cfg.kv_group_size = opts->kv_group_size; + /* v7: prefix cache. kv_pool_bytes is zero-init-friendly: 0 keeps the + * engine default (1 GiB), negative means unbounded (engine 0). */ + if (covered(&opts->prefix_cache + 1)) cfg.prefix_cache = opts->prefix_cache != 0; + if (covered(&opts->kv_block_size + 1) && opts->kv_block_size > 0) + cfg.kv_block_size = opts->kv_block_size; + if (covered(&opts->kv_pool_bytes + 1) && opts->kv_pool_bytes != 0) + cfg.kv_pool_bytes = + opts->kv_pool_bytes < 0 ? 0 : static_cast(opts->kv_pool_bytes); + if (covered(&opts->kv_spill_dir + 1) && opts->kv_spill_dir && *opts->kv_spill_dir) + cfg.kv_spill_dir = opts->kv_spill_dir; + if (covered(&opts->kv_spill_bytes + 1) && opts->kv_spill_bytes > 0) + cfg.kv_spill_bytes = static_cast(opts->kv_spill_bytes); auto handle = std::make_unique(); handle->model_name = model_spec; diff --git a/src/capi/mlxforge.h b/src/capi/mlxforge.h index f9edf6a..9a3bce6 100644 --- a/src/capi/mlxforge.h +++ b/src/capi/mlxforge.h @@ -43,8 +43,11 @@ extern "C" { * accumulated as the request is drained and returned as an OpenAI-shaped * JSON array). * v6: mlxforge_engine_create2 + mlxforge_engine_opts2 (KV-cache quantization; - * opts2 carries struct_size so future fields append without a create3). */ -#define MLXFORGE_ABI_VERSION 6 + * opts2 carries struct_size so future fields append without a create3). + * v7: mlxforge_engine_opts2 prefix-cache fields (prefix_cache, kv_block_size, + * kv_pool_bytes, kv_spill_dir, kv_spill_bytes) — appended, struct_size- + * gated; no new symbols. */ +#define MLXFORGE_ABI_VERSION 7 typedef struct mlxforge_engine mlxforge_engine; typedef struct mlxforge_request mlxforge_request; @@ -125,12 +128,25 @@ mlxforge_engine* mlxforge_engine_create(const char* model_spec, * mlx-lm's QuantizedKVCache (8 is near-lossless at ~1.9x less cache memory; * 4 is ~3.6x). Unsupported setups (vision-language or hybrid Qwen3.5 models, * invalid bits/group sizes) FAIL engine creation with a clear *err — there is - * never a silent fp16 fallback. */ + * never a silent fp16 fallback. + * + * prefix_cache (v7+) enables prompt-prefix reuse (also engine-wide): finished + * prompts' KV is pooled in immutable kv_block_size-token blocks, and a later + * prompt sharing a token prefix skips that part of prefill (same greedy + * tokens, much lower time-to-first-token). kv_spill_dir adds an SSD tier: + * RAM-evicted blocks persist there and survive engine restarts. Vision- + * language and hybrid (Qwen3.5) models reject the option at creation. */ typedef struct { size_t struct_size; /* caller sets sizeof(mlxforge_engine_opts2) */ int max_waiting; /* max queued requests; <= 0 => default (256) */ int kv_bits; /* 0 = fp16 KV cache (default); 8 or 4 = quantized */ int kv_group_size; /* quantization group size; <= 0 => default (64) */ + /* ---- v7 ---- */ + int prefix_cache; /* 0 = off (default); non-zero = on */ + int kv_block_size; /* pool block size in tokens; <= 0 => default (256) */ + long long kv_pool_bytes; /* pool RAM budget; 0 => default (1 GiB); < 0 => unbounded */ + const char* kv_spill_dir; /* SSD spill directory; NULL/empty => no spill */ + long long kv_spill_bytes; /* spill disk budget; <= 0 => unbounded */ } mlxforge_engine_opts2; /* Create an engine with extended options (v6+). Identical contract to diff --git a/src/runtime/batching.cpp b/src/runtime/batching.cpp index d264e80..4311e37 100644 --- a/src/runtime/batching.cpp +++ b/src/runtime/batching.cpp @@ -51,4 +51,31 @@ PrefillResult prefill(const DecoderModel& model, const std::vector& prompt, + const std::vector>& blocks, + int cached_len, int step_size, KVQuantConfig kv_quant) { + const int n = static_cast(prompt.size()); + BatchKVCache cache = + BatchKVCache::from_prefix(model.config().n_layers, blocks, cached_len, kv_quant); + cache.eval_state(); // materialize the seeded storage before the forward + + // Suffix prefill, chunked like the cold path. The cache's offset/idx already + // sit at cached_len, so RoPE positions and the mask line up unchanged. + const int suffix = n - cached_len; + mx::array logits = mx::zeros({1, 1, model.config().vocab}, mx::float16); + for (int c = 0; c < suffix; c += step_size) { + const int m = std::min(step_size, suffix - c); + mx::array chunk(prompt.data() + cached_len + c, {1, m}, mx::int32); + logits = model.forward(chunk, cache); + cache.eval_state(); + } + + const int n_last = logits.shape()[1]; + const int vocab = logits.shape()[2]; + mx::array last = + mx::reshape(mx::slice(logits, {0, n_last - 1, 0}, {1, n_last, vocab}), {1, vocab}); + mx::eval(last); + return {std::move(cache), last, std::vector{0}}; +} + } // namespace mlxforge diff --git a/src/runtime/batching.h b/src/runtime/batching.h index 5f6a351..ea7b371 100644 --- a/src/runtime/batching.h +++ b/src/runtime/batching.h @@ -9,6 +9,7 @@ #include #include "cache/batch_kv_cache.h" +#include "cache/block_pool.h" #include "model/decoder_model.h" #include "scheduler/request.h" @@ -38,4 +39,14 @@ PrefillResult prefill(const DecoderModel& model, const std::vector& prompt, + const std::vector>& blocks, + int cached_len, int step_size = kPrefillStepSize, + KVQuantConfig kv_quant = {}); + } // namespace mlxforge diff --git a/src/runtime/engine.cpp b/src/runtime/engine.cpp index 38b9210..8c4a6f0 100644 --- a/src/runtime/engine.cpp +++ b/src/runtime/engine.cpp @@ -61,6 +61,43 @@ KVQuantConfig validate_kv_quant(const EngineConfig& ec, const ModelConfig& mc) { return {ec.kv_bits, ec.kv_group_size}; } +// Validate the prefix-cache request against the loaded model and return the +// Worker's PrefixCacheConfig. Same philosophy as validate_kv_quant: unsupported +// setups are hard errors, never a silent off. +// `model_name` is passed separately: by the time the Worker member initializes, +// EngineConfig::model_spec has already been moved into Engine::model_name_. +PrefixCacheConfig validate_prefix_cache(const EngineConfig& ec, const ModelConfig& mc, + const std::string& model_name) { + if (!ec.prefix_cache) { + if (!ec.kv_spill_dir.empty()) + throw std::runtime_error("kv_spill_dir requires prefix_cache to be enabled"); + return {}; + } + const int bs = ec.kv_block_size; + if (bs < 16 || bs > 4096 || (bs & (bs - 1)) != 0) + throw std::runtime_error("kv_block_size must be a power of two in [16, 4096]; got " + + std::to_string(bs)); + // Token-id hashing cannot identify image content / 3D positions, and the + // hybrid linear-attention state cannot be reconstructed at block boundaries. + if (mc.has_vision_tower()) + throw std::runtime_error("the prefix cache is not supported for vision-language models"); + if (mc.full_attention_interval > 0) + throw std::runtime_error("the prefix cache is not supported for hybrid (Qwen3.5) models"); + PrefixCacheConfig pc; + pc.enabled = true; + pc.block_size = bs; + pc.pool_bytes = ec.kv_pool_bytes; + pc.spill_dir = ec.kv_spill_dir; + pc.spill_bytes = ec.kv_spill_bytes; + // Salt every block key with the model identity + storage config so pooled + // (and, later, persisted) blocks can never cross models or settings. + const std::string fp = model_name + "|" + mc.model_type + "|" + std::to_string(mc.n_layers) + + "|" + std::to_string(ec.kv_bits) + "|" + std::to_string(ec.kv_group_size) + + "|" + std::to_string(bs); + pc.salt = fnv1a(fp.data(), fp.size()); + return pc; +} + } // namespace // Loads the model directory, config, and tokenizer metadata, but not weights. @@ -139,7 +176,7 @@ Engine::Engine(EngineConfig cfg, Loaded loaded) // constrained decoding. tok_ is initialized above and outlives worker_. // cfg_ is initialized above, so the KV-quant validation sees the model. worker_(make_factory(std::move(loaded.dir), loaded.is_gguf), &scheduler_, &tok_, - validate_kv_quant(cfg, cfg_)) { + validate_kv_quant(cfg, cfg_), validate_prefix_cache(cfg, cfg_, model_name_)) { // Configure the max waiting requests for the batch scheduler. scheduler_.set_max_waiting(cfg.max_waiting); diff --git a/src/runtime/engine.h b/src/runtime/engine.h index 94565f3..5a82498 100644 --- a/src/runtime/engine.h +++ b/src/runtime/engine.h @@ -33,6 +33,18 @@ struct EngineConfig { // head_dim) throw rather than silently falling back. int kv_bits = 0; int kv_group_size = 64; + // Prefix cache (engine-wide, like kv_bits: the pool stores the engine's + // single storage layout). Off by default. When on, finished rows' KV is + // harvested into a block pool and later prompts sharing a token prefix skip + // that part of prefill. Validated at construction: hybrid (Qwen3.5) and + // vision-language models are rejected (no golden gate for those paths yet). + bool prefix_cache = false; + int kv_block_size = 256; // pool granularity, power of two in [16, 4096] + std::size_t kv_pool_bytes = 1ull << 30; // pooled-KV RAM budget; 0 = unbounded + // SSD spill tier (requires prefix_cache): RAM-evicted blocks persist under + // this directory and survive engine restarts. Empty = no spill. + std::string kv_spill_dir; + std::size_t kv_spill_bytes = 0; // disk budget; 0 = unbounded }; // Per-call embedding options. The two int fields are tri-state: -1 means "use diff --git a/src/runtime/metrics.h b/src/runtime/metrics.h index f2c57da..db133a6 100644 --- a/src/runtime/metrics.h +++ b/src/runtime/metrics.h @@ -16,6 +16,14 @@ struct WorkerMetrics { double avg_ttft_ms = 0.0; // enqueue -> first token double avg_request_ms = 0.0; // enqueue -> finished double avg_tokens_per_second = 0.0; // aggregate: completion_tokens / gen_seconds + + // Prefix cache (all zero when the feature is off). + long prefix_hits = 0; // admissions served partly from the pool + long long prefix_tokens_reused = 0; // prompt tokens whose prefill was skipped + long long prefix_pool_bytes = 0; // pooled KV bytes currently held + long prefix_pool_blocks = 0; // pooled block count + long spill_writes = 0; // blocks spilled to the SSD tier + long spill_reads = 0; // blocks revived from the SSD tier }; } // namespace mlxforge diff --git a/src/runtime/worker.cpp b/src/runtime/worker.cpp index db5e124..93c2d36 100644 --- a/src/runtime/worker.cpp +++ b/src/runtime/worker.cpp @@ -5,6 +5,7 @@ #include #include +#include "cache/block_store.h" #include "core/logging.h" #include "model/qwen3_vl.h" #include "model/vision/vit.h" @@ -58,8 +59,9 @@ bool consume(Request& req, int& produced, int id, const TokenLogprob* lp) { } // namespace Worker::Worker(ModelFactory factory, Scheduler* scheduler, const Tokenizer* tok, - KVQuantConfig kv_quant) - : factory_(std::move(factory)), sched_(scheduler), tok_(tok), kv_quant_(kv_quant) {} + KVQuantConfig kv_quant, PrefixCacheConfig prefix) + : factory_(std::move(factory)), sched_(scheduler), tok_(tok), kv_quant_(kv_quant), + prefix_cfg_(prefix) {} Worker::~Worker() { stop(); } @@ -193,6 +195,32 @@ void Worker::stop() { void Worker::run() { log::info("worker: loading model..."); model_ = factory_(); // load the model on this thread so its arrays live here + // The prefix pool holds MLX arrays, so it lives (and dies) with this thread. + if (prefix_cfg_.enabled) { + prefix_ = std::make_unique(prefix_cfg_); + log::info("worker: prefix cache on (block={} pool={} bytes)", prefix_cfg_.block_size, + prefix_cfg_.pool_bytes); + if (!prefix_cfg_.spill_dir.empty()) { + // SSD tier: RAM-evicted blocks are serialized HERE (this thread owns the + // arrays) and queued to the store's byte-only writer thread; a pool miss + // synchronously revives the bytes into fresh worker-thread arrays. + block_store_ = std::make_unique(prefix_cfg_.spill_dir, prefix_cfg_.spill_bytes, + prefix_cfg_.salt); + prefix_->pool().set_evict_hook( + [this](uint64_t h, const std::shared_ptr& b) { + if (block_store_->contains(h)) return; + block_store_->put(h, serialize_block(*b, prefix_cfg_.salt)); + ++spill_writes_; + }); + prefix_->set_miss_fn([this](uint64_t h) -> std::shared_ptr { + std::optional> bytes = block_store_->get(h); + if (!bytes) return nullptr; + std::shared_ptr blk = deserialize_block(*bytes, prefix_cfg_.salt); + if (blk) ++spill_reads_; + return blk; + }); + } + } ready_.store(true); log::info("worker: model loaded, ready"); @@ -242,20 +270,51 @@ void Worker::run() { } void Worker::admit(const std::vector>& incoming) { - std::vector> prompts; - prompts.reserve(incoming.size()); - for (const auto& r : incoming) prompts.push_back(r->prompt_ids); + // Split on the prefix cache: matched requests prefill only their suffix + // (one by one — their cached lengths are heterogeneous), the rest share the + // batched cold prefill. + std::vector> cold; + std::vector, PrefixCache::Match>> warm; + for (const auto& r : incoming) { + if (prefix_) { + PrefixCache::Match m = prefix_->match(r->prompt_ids); + if (m.tokens > 0) { + warm.emplace_back(r, std::move(m)); + continue; + } + } + cold.push_back(r); + } - log::debug("worker: admitting {} request(s) (batch {} -> {})", incoming.size(), reqs_.size(), - reqs_.size() + incoming.size()); - PrefillResult pr = prefill(*model_, prompts, kPrefillStepSize, /*pad_id=*/0, kv_quant_); + log::debug("worker: admitting {} request(s), {} prefix-warm (batch {} -> {})", incoming.size(), + warm.size(), reqs_.size(), reqs_.size() + incoming.size()); - if (!cache_) { - cache_ = std::make_unique(std::move(pr.cache)); - } else { - cache_->merge(pr.cache); + auto adopt = [&](BatchKVCache&& fresh) { + if (!cache_) { + cache_ = std::make_unique(std::move(fresh)); + } else { + cache_->merge(fresh); + } + }; + + if (!cold.empty()) { + std::vector> prompts; + prompts.reserve(cold.size()); + for (const auto& r : cold) prompts.push_back(r->prompt_ids); + PrefillResult pr = prefill(*model_, prompts, kPrefillStepSize, /*pad_id=*/0, kv_quant_); + adopt(std::move(pr.cache)); + register_rows(cold, pr.last_logits); + } + for (auto& [r, m] : warm) { + PrefillResult pr = prefill_with_prefix(*model_, r->prompt_ids, m.blocks, m.tokens, + kPrefillStepSize, kv_quant_); + adopt(std::move(pr.cache)); + register_rows({r}, pr.last_logits); + ++prefix_hits_; + prefix_tokens_reused_ += m.tokens; + log::debug("worker: prefix hit ({} of {} prompt tokens reused)", m.tokens, + r->prompt_ids.size()); } - register_rows(incoming, pr.last_logits); } void Worker::register_rows(const std::vector>& incoming, @@ -417,7 +476,34 @@ void Worker::decode_step() { } } +void Worker::harvest_finished() { + if (!prefix_ || !cache_) return; + if (std::none_of(finished_.begin(), finished_.end(), [](char f) { return f != 0; })) return; + + const std::vector lp = cache_->left_padding_host(); + const int n_layers = model_->config().n_layers; + for (int b = 0; b < static_cast(finished_.size()); ++b) { + if (!finished_[b] || reqs_[b]->is_multimodal()) continue; + // Pool only the PROMPT's span of the row — prefill-produced K/V. The + // decode-produced K/V of generated tokens differs from a recompute by fp16 + // accumulation order (the decode-vs-recompute gap) and demonstrably flips + // later greedy choices, breaking the warm==cold gate. Multi-turn reuse + // still converges: the next turn's prompt contains this turn's generated + // text, so its (seeded) prefill recomputes that span once and pools it. + const int len = + std::min(cache_->idx() - lp[b], static_cast(reqs_[b]->prompt_ids.size())); + if (len < prefix_cfg_.block_size) continue; // nothing pool-able + prefix_->harvest(history_[b], len, n_layers, [&](int layer) { + return cache_->fetch_row_components(layer, b, lp[b], len); + }); + } + prefix_pool_bytes_.store(static_cast(prefix_->pool_bytes())); + prefix_pool_blocks_.store(static_cast(prefix_->pool_blocks())); +} + void Worker::evict_finished() { + harvest_finished(); // seal finished rows' K/V before filter() drops them + using ms = std::chrono::duration; using sec = std::chrono::duration; using us = std::chrono::duration; @@ -494,6 +580,13 @@ WorkerMetrics Worker::metrics() const { } const long long gen_us = gen_us_sum_.load(); if (gen_us > 0) m.avg_tokens_per_second = m.completion_tokens_total * 1e6 / gen_us; + + m.prefix_hits = prefix_hits_.load(); + m.prefix_tokens_reused = prefix_tokens_reused_.load(); + m.prefix_pool_bytes = prefix_pool_bytes_.load(); + m.prefix_pool_blocks = prefix_pool_blocks_.load(); + m.spill_writes = spill_writes_.load(); + m.spill_reads = spill_reads_.load(); return m; } diff --git a/src/runtime/worker.h b/src/runtime/worker.h index 5e62d89..081be28 100644 --- a/src/runtime/worker.h +++ b/src/runtime/worker.h @@ -17,6 +17,7 @@ #include #include "cache/batch_kv_cache.h" +#include "cache/prefix_cache.h" #include "model/decoder_model.h" #include "runtime/metrics.h" #include "scheduler/request.h" // Request, TokenLogprob @@ -28,6 +29,7 @@ namespace mlxforge { class Tokenizer; // for per-token byte strings used by grammar masking class VitEncoder; // lazily built for multimodal requests (borrows model weights) +class BlockStore; // SSD spill tier for the prefix cache (cache/block_store.h) class Worker { public: @@ -35,12 +37,13 @@ class Worker { // `tok` (optional) supplies the per-token byte strings used for constrained // decoding; when null, grammar-constrained requests fall back to unconstrained. - // `kv_quant` selects the decode cache's storage (dense fp16 by default); the - // Engine validates it against the model before construction. Defined - // out-of-line (with the destructor) because the unique_ptr member - // needs the complete type for cleanup. + // `kv_quant` selects the decode cache's storage (dense fp16 by default) and + // `prefix` the prefix-cache setting; the Engine validates both against the + // model before construction. Defined out-of-line (with the destructor) + // because the unique_ptr member needs the complete type for + // cleanup. Worker(ModelFactory factory, Scheduler* scheduler, const Tokenizer* tok = nullptr, - KVQuantConfig kv_quant = {}); + KVQuantConfig kv_quant = {}, PrefixCacheConfig prefix = {}); ~Worker(); Worker(const Worker&) = delete; @@ -64,8 +67,16 @@ class Worker { void run(); // the loop; the sole caller of MLX eval/async_eval // Prefill `incoming` and merge it into the decode batch (emitting each row's - // first token). + // first token). With the prefix cache on, requests whose prompt matches + // pooled blocks are admitted one-by-one via prefill_with_prefix (seeded + // cache + suffix-only prefill); the rest take the batched cold path. void admit(const std::vector>& incoming); + // Seal finished rows' PROMPT K/V into the prefix pool (called before the + // cache rows are filtered away). Prompt-only: decode-produced K/V differs + // from a recompute by fp16 accumulation order, so pooling it would break the + // warm==cold token gate. Skips multimodal rows — their K/V embeds image + // content and 3D positions that a token-id hash cannot identify. + void harvest_finished(); // Register freshly-admitted rows into the decode-batch state (already merged // into the cache) and sample each row's first token from `last_logits` (rows // aligned to the new tail). Shared by the text and multimodal admit paths. @@ -124,6 +135,12 @@ class Worker { Scheduler* sched_; const Tokenizer* tok_; // for per-token bytes (grammar masking); may be null KVQuantConfig kv_quant_; // decode-cache storage (dense when bits == 0) + PrefixCacheConfig prefix_cfg_; + // SSD tier; declared before prefix_ so the pool's hooks (which reference it) + // are destroyed first. Byte-only across threads; null when spill is off. + std::unique_ptr block_store_; + // Worker-thread-only (holds MLX arrays); null when the feature is off. + std::unique_ptr prefix_; std::vector token_bytes_; // id -> output bytes ("" for specials) bool token_bytes_built_ = false; std::unique_ptr model_; // constructed and owned on the worker thread @@ -152,6 +169,13 @@ class Worker { std::atomic ttft_us_sum_{0}; std::atomic gen_us_sum_{0}; std::atomic request_us_sum_{0}; + // Prefix-cache counters (worker writes after each admit/harvest). + std::atomic prefix_hits_{0}; + std::atomic prefix_tokens_reused_{0}; + std::atomic prefix_pool_bytes_{0}; + std::atomic prefix_pool_blocks_{0}; + std::atomic spill_writes_{0}; // blocks queued to the SSD tier + std::atomic spill_reads_{0}; // blocks revived from the SSD tier std::thread thread_; }; diff --git a/src/server/config.cpp b/src/server/config.cpp index c76e6a9..c76c995 100644 --- a/src/server/config.cpp +++ b/src/server/config.cpp @@ -57,7 +57,8 @@ ServerConfig ServerConfig::from_file(const std::string& path) { // Reject unknown keys up front so typos (e.g. "prot") fail loudly. static const std::set kKnownKeys = { - "model", "host", "port", "max_ctx", "max_waiting", "kv_budget", "kv_bits"}; + "model", "host", "port", "max_ctx", "max_waiting", "kv_budget", + "kv_bits", "prefix_cache", "kv_block", "kv_pool", "kv_spill_dir", "kv_spill_bytes"}; for (const auto& [key, _] : j.items()) { if (kKnownKeys.find(key) == kKnownKeys.end()) { throw std::runtime_error("config file: unknown key '" + key + "' in '" + path + "'"); @@ -91,6 +92,22 @@ ServerConfig ServerConfig::from_file(const std::string& path) { if (c.kv_bits != 0 && c.kv_bits != 4 && c.kv_bits != 8) throw std::runtime_error("config file: 'kv_bits' must be 0, 4, or 8"); } + if (j.contains("prefix_cache")) c.prefix_cache = require_type(j, "prefix_cache"); + if (j.contains("kv_block")) { + c.kv_block = require_type(j, "kv_block"); + if (c.kv_block <= 0) throw std::runtime_error("config file: 'kv_block' must be > 0"); + } + if (j.contains("kv_pool")) { + long long pool = require_type(j, "kv_pool"); + if (pool < 0) throw std::runtime_error("config file: 'kv_pool' must be >= 0"); + c.kv_pool_bytes = static_cast(pool); + } + if (j.contains("kv_spill_dir")) c.kv_spill_dir = require_type(j, "kv_spill_dir"); + if (j.contains("kv_spill_bytes")) { + long long spill = require_type(j, "kv_spill_bytes"); + if (spill < 0) throw std::runtime_error("config file: 'kv_spill_bytes' must be >= 0"); + c.kv_spill_bytes = static_cast(spill); + } return c; } @@ -134,6 +151,13 @@ ServerConfig ServerConfig::parse(const std::vector& args) { c.kv_budget_bytes = static_cast(env_long("MLXFORGE_KV_BUDGET", static_cast(c.kv_budget_bytes))); c.kv_bits = static_cast(env_long("MLXFORGE_KV_BITS", c.kv_bits)); + c.prefix_cache = env_long("MLXFORGE_PREFIX_CACHE", c.prefix_cache ? 1 : 0) != 0; + c.kv_block = static_cast(env_long("MLXFORGE_KV_BLOCK", c.kv_block)); + c.kv_pool_bytes = + static_cast(env_long("MLXFORGE_KV_POOL", static_cast(c.kv_pool_bytes))); + c.kv_spill_dir = env_or("MLXFORGE_KV_SPILL_DIR", c.kv_spill_dir); + c.kv_spill_bytes = static_cast( + env_long("MLXFORGE_KV_SPILL_BYTES", static_cast(c.kv_spill_bytes))); // Helper: extract value for a flag (accepts "--flag value" or "--flag=value") auto value_of = [&](const std::string& a, size_t& i) -> std::string { @@ -166,6 +190,16 @@ ServerConfig ServerConfig::parse(const std::vector& args) { c.kv_budget_bytes = static_cast(std::stoll(value_of(a, i))); else if (flag == "--kv-bits") c.kv_bits = std::stoi(value_of(a, i)); + else if (flag == "--prefix-cache") + c.prefix_cache = std::stoi(value_of(a, i)) != 0; + else if (flag == "--kv-block") + c.kv_block = std::stoi(value_of(a, i)); + else if (flag == "--kv-pool") + c.kv_pool_bytes = static_cast(std::stoll(value_of(a, i))); + else if (flag == "--kv-spill-dir") + c.kv_spill_dir = value_of(a, i); + else if (flag == "--kv-spill-bytes") + c.kv_spill_bytes = static_cast(std::stoll(value_of(a, i))); else throw std::runtime_error("unknown flag: " + flag); } diff --git a/src/server/config.h b/src/server/config.h index d759710..669c9d6 100644 --- a/src/server/config.h +++ b/src/server/config.h @@ -31,19 +31,34 @@ struct ServerConfig { // cache quantized (engine-wide; group size fixed at 64 for the server). int kv_bits = 0; + // Prefix cache: reuse pooled KV across requests sharing a token prefix. + bool prefix_cache = false; + // Prefix-pool block granularity in tokens (power of two; validated by the engine). + int kv_block = 256; + // Prefix-pool RAM budget in bytes. 0 = unbounded. + std::size_t kv_pool_bytes = 1ull << 30; + + // SSD spill tier for the prefix pool (requires prefix_cache): evicted blocks + // persist under this directory and survive restarts. Empty = no spill. + std::string kv_spill_dir; + // Disk budget for spilled blocks in bytes. 0 = unbounded. + std::size_t kv_spill_bytes = 0; + // Parses command line arguments (the model via -m/--model, plus optional flags as --flag value or --flag=value), // layering configuration sources by precedence (lowest to highest): // struct defaults < config file (-c/--config) < environment variables < CLI flags. // The config file is a JSON object (see from_file); CLI flags always override it. // 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. // Throws std::runtime_error if an unknown or malformed flag is encountered. static ServerConfig parse(const std::vector& args); // Loads and validates a JSON config file into a fully-populated ServerConfig, // with struct defaults filling any keys the file omits. Recognized keys // (snake_case): "model", "host", "port", "max_ctx", "max_waiting", "kv_budget", - // "kv_bits". + // "kv_bits", "prefix_cache", "kv_block", "kv_pool", "kv_spill_dir", + // "kv_spill_bytes". // Validates before applying: rejects unknown keys, wrong types, and out-of-range // values. Throws std::runtime_error (with the file path / offending key) on any // failure to open, parse, or validate. diff --git a/src/server/http_server.cpp b/src/server/http_server.cpp index 9b06f4c..04bcb68 100644 --- a/src/server/http_server.cpp +++ b/src/server/http_server.cpp @@ -421,6 +421,13 @@ void HttpServer::setup_routes() { {"avg_ttft_ms", m.avg_ttft_ms}, {"avg_request_ms", m.avg_request_ms}, {"avg_tokens_per_second", m.avg_tokens_per_second}}}, + {"prefix_cache", + {{"hits", m.prefix_hits}, + {"tokens_reused", m.prefix_tokens_reused}, + {"pool_bytes", m.prefix_pool_bytes}, + {"pool_blocks", m.prefix_pool_blocks}, + {"spill_writes", m.spill_writes}, + {"spill_reads", m.spill_reads}}}, }; res.set_content(body.dump(), "application/json"); }); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c816559..6093201 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -20,6 +20,9 @@ add_executable(mlxforge_tests cache/kv_quant_cache_test.cpp cache/kv_cache_quantized_test.cpp cache/kv_budget_test.cpp + cache/block_pool_test.cpp + cache/prefix_cache_test.cpp + cache/block_store_test.cpp model/quantized_sdpa_test.cpp sample/sampler_test.cpp sample/json_grammar_test.cpp @@ -29,7 +32,10 @@ add_executable(mlxforge_tests scheduler/scheduler_test.cpp scheduler/worker_test.cpp runtime/prefill_test.cpp + runtime/prefill_prefix_test.cpp scheduler/continuous_batch_test.cpp + scheduler/prefix_reuse_test.cpp + scheduler/prefix_spill_test.cpp runtime/bucketing_test.cpp scheduler/validation_test.cpp model/qwen3_test.cpp diff --git a/tests/cache/block_pool_test.cpp b/tests/cache/block_pool_test.cpp new file mode 100644 index 0000000..04c7d14 --- /dev/null +++ b/tests/cache/block_pool_test.cpp @@ -0,0 +1,123 @@ +// BlockPool pure-logic tests: chained prefix hashing, LRU eviction under a +// byte budget, and the Phase-2 evict hook — tiny synthetic tensors, no model. +#include + +#include +#include + +#include "cache/block_pool.h" + +#include "mlx/ops.h" +#include "mlx/transforms.h" + +using namespace mlxforge; +namespace mx = mlx::core; + +namespace { + +// A minimal pooled block: one layer, dense K/V of `tokens` positions. +std::shared_ptr make_block(int tokens, float fill) { + auto b = std::make_shared(); + b->k = {{mx::full({1, 2, tokens, 8}, fill, mx::float16)}}; + b->v = {{mx::full({1, 2, tokens, 8}, -fill, mx::float16)}}; + mx::eval(b->k[0][0], b->v[0][0]); + b->bytes = block_bytes(*b); + return b; +} + +} // namespace + +TEST_CASE("chain_hash identifies the whole prefix, not just the block") { + const std::vector a = {1, 2, 3, 4}; + const std::vector b = {5, 6, 7, 8}; + + // Deterministic for the same chain. + const uint64_t h1 = chain_hash(chain_hash(0, a.data(), 4), b.data(), 4); + const uint64_t h2 = chain_hash(chain_hash(0, a.data(), 4), b.data(), 4); + CHECK(h1 == h2); + + // The same block ids behind a different first block hash differently. + CHECK(chain_hash(chain_hash(0, b.data(), 4), b.data(), 4) != h1); + + // A different salt (seed) changes every key in the chain. + CHECK(chain_hash(chain_hash(99, a.data(), 4), b.data(), 4) != h1); + + // Single-token difference inside a block changes its key. + std::vector a2 = a; + a2[2] = 30; + CHECK(chain_hash(0, a2.data(), 4) != chain_hash(0, a.data(), 4)); +} + +TEST_CASE("block_bytes sums every component buffer") { + auto b = make_block(4, 1.0f); + // 2 tensors (K+V) * 1 layer * 1 component * (1*2*4*8) fp16 elements. + CHECK(b->bytes == 2 * 2 * 4 * 8 * sizeof(uint16_t)); +} + +TEST_CASE("BlockPool insert/get with LRU eviction under the byte budget") { + auto b = make_block(4, 1.0f); // 256 bytes each + BlockPool pool(/*budget_bytes=*/b->bytes * 2); + + pool.insert(1, make_block(4, 1.0f)); + pool.insert(2, make_block(4, 2.0f)); + CHECK(pool.size() == 2); + CHECK(pool.bytes() == b->bytes * 2); + + // Touch 1 so 2 becomes the LRU victim of the next insert. + CHECK(pool.get(1) != nullptr); + pool.insert(3, make_block(4, 3.0f)); + CHECK(pool.size() == 2); + CHECK(pool.get(2) == nullptr); // evicted + CHECK(pool.get(1) != nullptr); + CHECK(pool.get(3) != nullptr); +} + +TEST_CASE("BlockPool re-insert of an existing key is a no-op") { + BlockPool pool(/*budget_bytes=*/0); // unbounded + pool.insert(7, make_block(4, 1.0f)); + const std::size_t bytes = pool.bytes(); + pool.insert(7, make_block(4, 2.0f)); + CHECK(pool.size() == 1); + CHECK(pool.bytes() == bytes); + // First write wins (blocks are immutable; same key == same content). + mx::array k = pool.get(7)->k[0][0]; + mx::eval(k); + CHECK(static_cast(k.data()[0]) == doctest::Approx(1.0f)); +} + +TEST_CASE("BlockPool rejects a block larger than the whole budget") { + auto small = make_block(4, 1.0f); + BlockPool pool(small->bytes); + pool.insert(1, std::move(small)); + CHECK(pool.size() == 1); + pool.insert(2, make_block(8, 1.0f)); // 2x the budget: not admitted + CHECK(pool.size() == 1); + CHECK(pool.get(1) != nullptr); // and the existing entry survived +} + +TEST_CASE("BlockPool evict hook sees each victim before it is dropped") { + auto b = make_block(4, 1.0f); + BlockPool pool(b->bytes * 2); + std::vector evicted; + pool.set_evict_hook([&](uint64_t h, const std::shared_ptr&) { + evicted.push_back(h); + }); + pool.insert(1, make_block(4, 1.0f)); + pool.insert(2, make_block(4, 2.0f)); + pool.insert(3, make_block(4, 3.0f)); + pool.insert(4, make_block(4, 4.0f)); + CHECK(evicted == std::vector{1, 2}); +} + +TEST_CASE("in-flight references survive pool eviction") { + auto b = make_block(4, 5.0f); + BlockPool pool(b->bytes); + pool.insert(1, std::move(b)); + std::shared_ptr ref = pool.get(1); // an admission holding the block + pool.insert(2, make_block(4, 6.0f)); // evicts key 1 + CHECK(pool.get(1) == nullptr); + REQUIRE(ref != nullptr); // but the gathered copy is still usable + mx::array k = ref->k[0][0]; + mx::eval(k); + CHECK(static_cast(k.data()[0]) == doctest::Approx(5.0f)); +} diff --git a/tests/cache/block_store_test.cpp b/tests/cache/block_store_test.cpp new file mode 100644 index 0000000..17873dc --- /dev/null +++ b/tests/cache/block_store_test.cpp @@ -0,0 +1,178 @@ +// BlockStore pure-logic tests: serialize/deserialize byte round-trips (dense +// and quantized layouts), salt/corruption rejection, async write + reload, +// restart rescan, and the disk budget — temp dirs, no model. +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "cache/block_store.h" + +#include "mlx/ops.h" +#include "mlx/transforms.h" + +using namespace mlxforge; +namespace mx = mlx::core; +namespace fs = std::filesystem; + +namespace { + +constexpr uint64_t kSalt = 0xabcdef; + +// A scoped temp dir so failed runs don't leak files. +struct TempDir { + std::string path; + TempDir() { + char tmpl[] = "/tmp/mlxforge_block_store_XXXXXX"; + path = mkdtemp(tmpl); + } + ~TempDir() { + std::error_code ec; + fs::remove_all(path, ec); + } +}; + +std::shared_ptr dense_block(int n_layers, int tokens, float phase) { + auto b = std::make_shared(); + for (int l = 0; l < n_layers; ++l) { + mx::array base = mx::arange(static_cast(2 * tokens * 8)); + base = mx::sin(mx::add(mx::multiply(base, mx::array(0.37f)), mx::array(phase + l))); + mx::array k = mx::astype(mx::reshape(base, {1, 2, tokens, 8}), mx::float16); + mx::array v = mx::astype(mx::negative(mx::reshape(base, {1, 2, tokens, 8})), mx::float16); + mx::eval(k, v); + b->k.push_back({k}); + b->v.push_back({v}); + } + b->bytes = block_bytes(*b); + return b; +} + +std::shared_ptr quantized_block(int tokens) { + mx::array base = mx::arange(static_cast(2 * tokens * 64)); + base = mx::astype(mx::reshape(mx::sin(base), {1, 2, tokens, 64}), mx::float16); + auto b = std::make_shared(); + b->k.push_back(mx::quantize(base, 64, 8)); + b->v.push_back(mx::quantize(mx::negative(base), 64, 8)); + b->bytes = block_bytes(*b); + return b; +} + +bool blocks_equal(const KVBlock& a, const KVBlock& b) { + if (a.k.size() != b.k.size()) return false; + auto comps_equal = [](const std::vector& x, const std::vector& y) { + if (x.size() != y.size()) return false; + for (size_t i = 0; i < x.size(); ++i) { + mx::array eq = mx::array_equal(x[i], y[i]); + mx::eval(eq); + if (!eq.item()) return false; + } + return true; + }; + for (size_t l = 0; l < a.k.size(); ++l) { + if (!comps_equal(a.k[l], b.k[l]) || !comps_equal(a.v[l], b.v[l])) return false; + } + return true; +} + +// Wait until the async writer has persisted `n` blocks (bounded poll). +void wait_for_size(const BlockStore& store, std::size_t n) { + for (int i = 0; i < 500 && store.size() < n; ++i) + std::this_thread::sleep_for(std::chrono::milliseconds(2)); +} + +} // namespace + +TEST_CASE("serialize/deserialize round-trips a dense block exactly") { + auto b = dense_block(/*n_layers=*/3, /*tokens=*/16, 0.5f); + std::vector bytes = serialize_block(*b, kSalt); + std::shared_ptr back = deserialize_block(bytes, kSalt); + REQUIRE(back != nullptr); + CHECK(back->bytes == b->bytes); + CHECK(blocks_equal(*b, *back)); +} + +TEST_CASE("serialize/deserialize round-trips a quantized (triplet) block exactly") { + auto b = quantized_block(/*tokens=*/16); + std::vector bytes = serialize_block(*b, kSalt); + std::shared_ptr back = deserialize_block(bytes, kSalt); + REQUIRE(back != nullptr); + CHECK(back->k[0].size() == 3); + CHECK(blocks_equal(*b, *back)); +} + +TEST_CASE("deserialize rejects the wrong salt, corruption, and truncation") { + auto b = dense_block(1, 8, 0.1f); + std::vector bytes = serialize_block(*b, kSalt); + + CHECK(deserialize_block(bytes, kSalt + 1) == nullptr); // another model/config + + std::vector corrupt = bytes; + corrupt[0] ^= 0xff; // magic + CHECK(deserialize_block(corrupt, kSalt) == nullptr); + + std::vector truncated(bytes.begin(), bytes.begin() + bytes.size() / 2); + CHECK(deserialize_block(truncated, kSalt) == nullptr); + + CHECK(deserialize_block(std::vector{}, kSalt) == nullptr); +} + +TEST_CASE("BlockStore writes asynchronously and reads back the same bytes") { + TempDir dir; + auto b = dense_block(2, 16, 0.2f); + std::vector bytes = serialize_block(*b, kSalt); + BlockStore store(dir.path, /*budget=*/0, kSalt); + CHECK(!store.contains(11)); + store.put(11, bytes); + wait_for_size(store, 1); + REQUIRE(store.contains(11)); + auto back = store.get(11); + REQUIRE(back.has_value()); + CHECK(*back == bytes); + CHECK(store.bytes() == bytes.size()); +} + +TEST_CASE("BlockStore rescan revives the index across restarts") { + TempDir dir; + std::vector bytes = serialize_block(*dense_block(1, 8, 0.3f), kSalt); + { + BlockStore store(dir.path, 0, kSalt); + store.put(21, bytes); + store.put(22, bytes); + } // dtor drains the write queue + + // Foreign files in the dir are ignored by the rescan. + std::ofstream(dir.path + "/notes.txt") << "not a block"; + std::ofstream(dir.path + "/zzzz.kvb") << "bad name"; // unparseable hex is fine to skip... + + BlockStore revived(dir.path, 0, kSalt); + CHECK(revived.size() >= 2); + CHECK(revived.contains(21)); + CHECK(revived.contains(22)); + auto back = revived.get(21); + REQUIRE(back.has_value()); + CHECK(deserialize_block(*back, kSalt) != nullptr); +} + +TEST_CASE("BlockStore disk budget deletes the least recently used file") { + TempDir dir; + std::vector bytes = serialize_block(*dense_block(1, 8, 0.4f), kSalt); + BlockStore store(dir.path, /*budget=*/bytes.size() * 2, kSalt); + store.put(1, bytes); + store.put(2, bytes); + wait_for_size(store, 2); + CHECK(store.get(1).has_value()); // bump 1; 2 becomes the victim + store.put(3, bytes); + for (int i = 0; i < 500 && store.contains(2); ++i) + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + CHECK(store.contains(1)); + CHECK(!store.contains(2)); + CHECK(store.contains(3)); + CHECK(store.bytes() == bytes.size() * 2); +} diff --git a/tests/cache/prefix_cache_test.cpp b/tests/cache/prefix_cache_test.cpp new file mode 100644 index 0000000..ad9d322 --- /dev/null +++ b/tests/cache/prefix_cache_test.cpp @@ -0,0 +1,190 @@ +// PrefixCache pure-logic tests: longest-prefix matching, the last-token +// recompute clamp, harvest sealing/dedup, and seeding a BatchKVCache from +// matched blocks (dense and quantized) — tiny synthetic tensors, no model. +#include + +#include + +#include "cache/batch_kv_cache.h" +#include "cache/kv_quant.h" +#include "cache/prefix_cache.h" + +#include "mlx/ops.h" +#include "mlx/transforms.h" + +using namespace mlxforge; +namespace mx = mlx::core; + +namespace { + +constexpr int kH = 2, kD = 64; // head_dim a multiple of the quant group size + +// Deterministic varied values (constants would quantize trivially). +mx::array varied(int len, float phase) { + mx::array a = mx::arange(static_cast(kH * len * kD)); + a = mx::sin(mx::add(mx::multiply(a, mx::array(0.37f)), mx::array(phase))); + return mx::astype(mx::reshape(a, {1, kH, len, kD}), mx::float16); +} + +bool same(const mx::array& a, const mx::array& b) { + mx::array eq = mx::allclose(a, b, /*rtol=*/0.0, /*atol=*/1e-6); + mx::eval(eq); + return eq.item(); +} + +std::vector iota_ids(int n, int start = 0) { + std::vector ids(n); + for (int i = 0; i < n; ++i) ids[i] = start + i; + return ids; +} + +// A row's dense K/V (one component) per layer, len positions. +struct FakeRow { + std::vector k, v; // [layer], (1, kH, len, kD) + FakeRow(int n_layers, int len) { + for (int l = 0; l < n_layers; ++l) { + k.push_back(varied(len, 0.1f + l)); + v.push_back(varied(len, 0.2f + l)); + } + } + PrefixCache::LayerFetch fetch() const { + return [this](int l) { + return std::make_pair(std::vector{k[l]}, std::vector{v[l]}); + }; + } +}; + +} // namespace + +TEST_CASE("match on an empty pool reuses nothing") { + PrefixCache pc({true, 16, 1ull << 20, 0}); + PrefixCache::Match m = pc.match(iota_ids(64)); + CHECK(m.tokens == 0); + CHECK(m.blocks.empty()); +} + +TEST_CASE("harvest seals only full blocks; match clamps to len-1") { + const int bs = 16; + PrefixCache pc({true, bs, 1ull << 20, 0}); + FakeRow row(/*n_layers=*/2, /*len=*/40); // 2 full blocks + 8-token tail + pc.harvest(iota_ids(40), 40, 2, row.fetch()); + CHECK(pc.pool_blocks() == 2); // the partial tail is never sealed + + // A longer prompt sharing both blocks reuses all 32 cached tokens. + PrefixCache::Match m = pc.match(iota_ids(64)); + CHECK(m.tokens == 32); + CHECK(m.blocks.size() == 2); + // The pooled content is exactly the harvested span. + CHECK(same(m.blocks[0]->k[0][0], mx::slice(row.k[0], {0, 0, 0, 0}, {1, kH, bs, kD}))); + CHECK(same(m.blocks[1]->v[1][0], mx::slice(row.v[1], {0, 0, bs, 0}, {1, kH, 2 * bs, kD}))); + + // A prompt that IS the cached span still recomputes its last token. + m = pc.match(iota_ids(32)); + CHECK(m.tokens == 31); + CHECK(m.blocks.size() == 2); // both needed to cover [0, 31) +} + +TEST_CASE("a diverging block ends the match (keys chain over the whole prefix)") { + const int bs = 16; + PrefixCache pc({true, bs, 1ull << 20, 0}); + FakeRow row(1, 32); + pc.harvest(iota_ids(32), 32, 1, row.fetch()); + + std::vector ids = iota_ids(64); + ids[20] = 9999; // mutate inside block 1 + PrefixCache::Match m = pc.match(ids); + CHECK(m.tokens == bs); // block 0 still matches, block 1 no longer can + ids[3] = 9999; // mutate inside block 0: nothing matches + CHECK(pc.match(ids).tokens == 0); +} + +TEST_CASE("re-harvesting the same row is a no-op (dedup against the pool)") { + PrefixCache pc({true, 16, 1ull << 20, 0}); + FakeRow row(1, 32); + pc.harvest(iota_ids(32), 32, 1, row.fetch()); + const std::size_t bytes = pc.pool_bytes(); + pc.harvest(iota_ids(32), 32, 1, row.fetch()); + CHECK(pc.pool_bytes() == bytes); + CHECK(pc.pool_blocks() == 2); + + // Extending the row later seals only the new block. + FakeRow longer(1, 48); + pc.harvest(iota_ids(48), 48, 1, longer.fetch()); + CHECK(pc.pool_blocks() == 3); +} + +TEST_CASE("pool eviction under a small budget loses the oldest prefix") { + const int bs = 16; + FakeRow a(1, bs), b(1, bs); + PrefixCache pc({true, bs, 0, 0}); // measure one block first + pc.harvest(iota_ids(bs), bs, 1, a.fetch()); + const std::size_t one_block = pc.pool_bytes(); + + PrefixCache small({true, bs, one_block, 0}); + small.harvest(iota_ids(bs, 0), bs, 1, a.fetch()); + small.harvest(iota_ids(bs, 1000), bs, 1, b.fetch()); // different prefix -> evicts the first + CHECK(small.pool_blocks() == 1); + CHECK(small.match(iota_ids(bs + 1, 0)).tokens == 0); + CHECK(small.match(iota_ids(bs + 1, 1000)).tokens == bs); +} + +TEST_CASE("salt separates pools (a persisted block can never cross models)") { + FakeRow row(1, 16); + PrefixCache pc1({true, 16, 1ull << 20, 1}); + pc1.harvest(iota_ids(16), 16, 1, row.fetch()); + PrefixCache pc2({true, 16, 1ull << 20, 2}); + // Same ids, different salt: pc2's keys don't collide with pc1's content. + CHECK(pc2.match(iota_ids(17)).tokens == 0); + CHECK(pc1.match(iota_ids(17)).tokens == 16); +} + +TEST_CASE("BatchKVCache::from_prefix seeds a batch-1 cache from matched blocks") { + const int bs = 16, n_layers = 2, len = 48; + PrefixCache pc({true, bs, 1ull << 24, 0}); + FakeRow row(n_layers, len); + pc.harvest(iota_ids(len), len, n_layers, row.fetch()); + + PrefixCache::Match m = pc.match(iota_ids(len)); // clamps to 47 + CHECK(m.tokens == len - 1); + BatchKVCache cache = BatchKVCache::from_prefix(n_layers, m.blocks, m.tokens); + CHECK(cache.batch_size() == 1); + CHECK(cache.idx() == m.tokens); + CHECK(cache.s_cap() % BatchKVCache::kStep == 0); // standard block-grow rounding + + // RoPE offset == cached length; no left padding. + mx::array off = cache.offset(); + mx::eval(off); + CHECK(off.item() == m.tokens); + CHECK(cache.left_padding_host() == std::vector{0}); + + for (int l = 0; l < n_layers; ++l) { + auto [k, v] = cache.fetch(l); + CHECK(same(k, mx::slice(row.k[l], {0, 0, 0, 0}, {1, kH, m.tokens, kD}))); + CHECK(same(v, mx::slice(row.v[l], {0, 0, 0, 0}, {1, kH, m.tokens, kD}))); + } + + // Appending after the seed continues at the right position. + mx::array k1 = varied(1, 7.0f), v1 = varied(1, 8.0f); + auto [ks, vs] = cache.update_and_fetch(0, k1, v1); + CHECK(ks.shape()[2] == m.tokens + 1); +} + +TEST_CASE("from_prefix round-trips quantized (triplet) blocks") { + const int bs = 16, len = 32; + const KVQuantConfig qc{8, 64}; + PrefixCache pc({true, bs, 1ull << 24, 0}); + + // The harvested row stores triplets, exactly like a quantized BatchKVCache. + mx::array kd = varied(len, 0.3f), vd = varied(len, 0.6f); + std::vector kt = mx::quantize(kd, qc.group_size, qc.bits); + std::vector vt = mx::quantize(vd, qc.group_size, qc.bits); + pc.harvest(iota_ids(len), len, 1, [&](int) { return std::make_pair(kt, vt); }); + CHECK(pc.pool_blocks() == 2); + + PrefixCache::Match m = pc.match(iota_ids(len + 8)); + CHECK(m.tokens == len); + BatchKVCache cache = BatchKVCache::from_prefix(1, m.blocks, m.tokens, qc); + auto [k, v] = cache.fetch_dequantized(0); + CHECK(same(k, mx::dequantize(kt[0], kt[1], kt[2], qc.group_size, qc.bits))); + CHECK(same(v, mx::dequantize(vt[0], vt[1], vt[2], qc.group_size, qc.bits))); +} diff --git a/tests/capi/capi_test.cpp b/tests/capi/capi_test.cpp index e5c9089..b057090 100644 --- a/tests/capi/capi_test.cpp +++ b/tests/capi/capi_test.cpp @@ -320,3 +320,37 @@ TEST_CASE("C ABI constrained decoding forces well-formed JSON") { mlxforge_engine_free(eng); } + +TEST_CASE("C ABI v7 prefix cache: warm reuse keeps greedy output identical") { + if (!model_available()) { + MESSAGE("MLXFORGE_MODEL_DIR not present; skipping"); + return; + } + char* err = nullptr; + mlxforge_engine_opts2 opts = {}; + opts.struct_size = sizeof(opts); + opts.prefix_cache = 1; + opts.kv_block_size = 16; // the test prompt is short; default 256 would never hit + mlxforge_engine* eng = mlxforge_engine_create2(model_dir().c_str(), &opts, &err); + REQUIRE_MESSAGE(eng != nullptr, (err ? err : "engine_create2 failed")); + + // A prompt long enough to span full blocks (greedy, deterministic). + const char* prompt = + "Once upon a time in a quiet village by the sea, a young engineer set out " + "to build a tiny inference engine that could remember every prompt prefix " + "it had ever seen. Describe the first thing it cached."; + mlxforge_sampling s = {}; + s.max_tokens = 12; + + std::string first, second; + for (std::string* out : {&first, &second}) { + mlxforge_request* r = mlxforge_submit_text(eng, prompt, &s, &err); + REQUIRE_MESSAGE(r != nullptr, (err ? err : "submit failed")); + *out = drain(r); + mlxforge_request_free(r); + } + CHECK(first.size() > 0); + CHECK(first == second); // the warm (prefix-reused) run must not change tokens + + mlxforge_engine_free(eng); +} diff --git a/tests/runtime/prefill_prefix_test.cpp b/tests/runtime/prefill_prefix_test.cpp new file mode 100644 index 0000000..ce0e810 --- /dev/null +++ b/tests/runtime/prefill_prefix_test.cpp @@ -0,0 +1,114 @@ +// Seeded (prefix-cache) prefill against the cold path on the real model: a +// cache built from harvested blocks + suffix-only prefill must reproduce the +// cold prefill's K/V and next-token choice. The harvested K/V is a bit-copy of +// the cold run's, so K/V compares close and the next token compares exact — +// the warm==cold equivalence is this feature's golden gate (the cold path is +// already gated against mlx-lm). +#include + +#include + +#include "cache/prefix_cache.h" +#include "runtime/batching.h" +#include "support/model_fixture.h" +#include "support/reference.h" + +using namespace mlxforge::test; +namespace mx = mlx::core; + +namespace { + +// Fixture prompts are tiny; concatenate them into a prefix-cache-sized prompt. +std::vector long_prompt() { + std::vector p; + for (const char* name : {"prompt_2_ids.npy", "prompt_0_ids.npy", "prompt_1_ids.npy", + "prompt_2_ids.npy", "prompt_0_ids.npy", "prompt_1_ids.npy"}) { + std::vector ids = load_token_ids(name); + p.insert(p.end(), ids.begin(), ids.end()); + } + return p; // 46 tokens +} + +int argmax_id(const mx::array& logits_row) { + mx::array am = mx::argmax(logits_row, /*axis=*/-1); + mx::eval(am); + return static_cast(am.item()); +} + +} // namespace + +TEST_CASE("prefill_with_prefix matches the cold prefill (dense fp16)") { + if (!model_available()) { + MESSAGE("MLXFORGE_MODEL_DIR not present; skipping"); + return; + } + mlxforge::LlamaModel& model = shared_model(); + const int n_layers = model.config().n_layers; + const std::vector prompt = long_prompt(); + const int len = static_cast(prompt.size()); + + // Cold prefill, then harvest its rows into a fresh pool (block 16 so the + // short fixture prompt spans multiple blocks). + mlxforge::PrefillResult cold = mlxforge::prefill(model, {prompt}); + mlxforge::PrefixCache pc({true, /*block_size=*/16, 1ull << 30, /*salt=*/42}); + pc.harvest(prompt, len, n_layers, + [&](int l) { return cold.cache.fetch_row_components(l, 0, 0, len); }); + CHECK(pc.pool_blocks() == len / 16); + + mlxforge::PrefixCache::Match m = pc.match(prompt); + REQUIRE(m.tokens == (len / 16) * 16); + + mlxforge::PrefillResult warm = mlxforge::prefill_with_prefix(model, prompt, m.blocks, m.tokens); + + // Same next token (exact), same cache content (fp16-close; the suffix is + // recomputed in a different graph context, so raw-logit equality is not the + // gate — see the decode-vs-recompute gotcha). + CHECK(argmax_id(warm.last_logits) == argmax_id(cold.last_logits)); + CHECK(warm.cache.idx() == cold.cache.idx()); + for (int l = 0; l < n_layers; ++l) { + auto [ck, cv] = cold.cache.fetch(l); + auto [wk, wv] = warm.cache.fetch(l); + mlxforge::test::assert_close(wk, ck); + mlxforge::test::assert_close(wv, cv); + } +} + +TEST_CASE("prefill_with_prefix matches the cold prefill (quantized KV)") { + if (!model_available()) { + MESSAGE("MLXFORGE_MODEL_DIR not present; skipping"); + return; + } + mlxforge::LlamaModel& model = shared_model(); + const mlxforge::KVQuantConfig qc{8, 64}; + const int n_layers = model.config().n_layers; + const std::vector prompt = long_prompt(); + const int len = static_cast(prompt.size()); + + mlxforge::PrefillResult cold = + mlxforge::prefill(model, {prompt}, mlxforge::kPrefillStepSize, 0, qc); + mlxforge::PrefixCache pc({true, 16, 1ull << 30, 42}); + pc.harvest(prompt, len, n_layers, + [&](int l) { return cold.cache.fetch_row_components(l, 0, 0, len); }); + + mlxforge::PrefixCache::Match m = pc.match(prompt); + REQUIRE(m.tokens > 0); + mlxforge::PrefillResult warm = + mlxforge::prefill_with_prefix(model, prompt, m.blocks, m.tokens, + mlxforge::kPrefillStepSize, qc); + + // Quantized matmuls are fusion-context-sensitive (see the kv-quant gates): + // the recomputed suffix legitimately shifts within quantization noise, so the + // gate is the choice (argmax) plus exact reuse of the cached region — the + // pooled triplets must dequantize identically to the cold run's. + CHECK(argmax_id(warm.last_logits) == argmax_id(cold.last_logits)); + auto prefix_of = [&](const mx::array& a) { + const auto& s = a.shape(); + return mx::slice(a, {0, 0, 0, 0}, {s[0], s[1], m.tokens, s[3]}); + }; + for (int l = 0; l < n_layers; ++l) { + auto [ck, cv] = cold.cache.fetch_dequantized(l); + auto [wk, wv] = warm.cache.fetch_dequantized(l); + mlxforge::test::assert_close(prefix_of(wk), prefix_of(ck)); + mlxforge::test::assert_close(prefix_of(wv), prefix_of(cv)); + } +} diff --git a/tests/scheduler/prefix_reuse_test.cpp b/tests/scheduler/prefix_reuse_test.cpp new file mode 100644 index 0000000..b963991 --- /dev/null +++ b/tests/scheduler/prefix_reuse_test.cpp @@ -0,0 +1,125 @@ +// End-to-end prefix-cache gate through the continuous-batching worker: with +// the prefix cache on, a warm request (same or extended prompt) must produce +// the exact greedy stream of a cold solo run — reuse may only change speed, +// never tokens — and the worker's metrics must show the reuse happened. +#include + +#include +#include + +#include "core/config.h" +#include "core/weights.h" +#include "runtime/single_stream.h" +#include "runtime/worker.h" +#include "scheduler/request.h" +#include "scheduler/scheduler.h" +#include "support/model_fixture.h" +#include "support/reference.h" + +using namespace mlxforge::test; + +namespace { + +std::vector long_prompt() { + std::vector p; + for (const char* name : {"prompt_2_ids.npy", "prompt_0_ids.npy", "prompt_1_ids.npy", + "prompt_2_ids.npy", "prompt_0_ids.npy", "prompt_1_ids.npy"}) { + std::vector ids = load_token_ids(name); + p.insert(p.end(), ids.begin(), ids.end()); + } + return p; // 46 tokens: several 16-token blocks +} + +// Submit one greedy request and drain its stream (the worker harvests the row +// into the prefix pool when it finishes). +std::vector run_one(mlxforge::Scheduler& sched, const std::vector& prompt, + const std::vector& eos_ids, int max_tokens) { + auto req = std::make_shared(); + req->prompt_ids = prompt; + req->params.temperature = 0.0f; + req->max_tokens = max_tokens; + req->eos_ids = eos_ids; + REQUIRE(sched.submit(req)); + std::vector got; + int tok = 0; + while (req->tokens.pop(tok)) got.push_back(tok); + return got; +} + +} // namespace + +TEST_CASE("prefix-cache reuse reproduces the cold greedy stream exactly") { + if (!model_available()) { + MESSAGE("MLXFORGE_MODEL_DIR not present; skipping"); + return; + } + const std::string dir = model_dir(); + mlxforge::ModelConfig cfg = mlxforge::ModelConfig::from_file(dir + "/config.json"); + const int kMax = 16; + + const std::vector turn1 = long_prompt(); + // Solo expectations from the validated single-stream loop. Each turn's + // prompt extends the full prior conversation (prompt + answer + new text), + // the multi-turn shape prefix caching exists for. + mlxforge::LlamaModel& solo = shared_model(); + const std::vector expect1 = + mlxforge::greedy_generate(solo, turn1, kMax, cfg.eos_token_ids).tokens; + auto extend = [](std::vector conv, const std::vector& answer, const char* next) { + conv.insert(conv.end(), answer.begin(), answer.end()); + const std::vector extra = load_token_ids(next); + conv.insert(conv.end(), extra.begin(), extra.end()); + return conv; + }; + const std::vector turn2 = extend(turn1, expect1, "prompt_0_ids.npy"); + const std::vector expect2 = + mlxforge::greedy_generate(solo, turn2, kMax, cfg.eos_token_ids).tokens; + const std::vector turn3 = extend(turn2, expect2, "prompt_1_ids.npy"); + const std::vector expect3 = + mlxforge::greedy_generate(solo, turn3, kMax, cfg.eos_token_ids).tokens; + + // Worker with the prefix cache on (block 16: the fixture prompts are short; + // the engine's >= 16 floor still holds). + mlxforge::PrefixCacheConfig pcfg; + pcfg.enabled = true; + pcfg.block_size = 16; + pcfg.pool_bytes = 1ull << 30; + pcfg.salt = 7; + mlxforge::Scheduler sched; + mlxforge::Worker worker( + [dir] { + mlxforge::ModelConfig c = mlxforge::ModelConfig::from_file(dir + "/config.json"); + auto w = mlxforge::load_weights(dir, c); + return std::make_unique(std::move(c), std::move(w)); + }, + &sched, /*tok=*/nullptr, /*kv_quant=*/{}, pcfg); + worker.start(); + + // Cold: no pool content yet. + CHECK(run_one(sched, turn1, cfg.eos_token_ids, kMax) == expect1); + CHECK(worker.metrics().prefix_hits == 0); + CHECK(worker.metrics().prefix_pool_blocks > 0); // the finished row was harvested + + // Warm 1: identical prompt — served from the pool, same tokens. + CHECK(run_one(sched, turn1, cfg.eos_token_ids, kMax) == expect1); + // Warm 2: the multi-turn continuation. Its prefix hit covers turn 1's + // prompt blocks; harvesting turn 2's own (seeded) prefill then pools the + // conversation through turn 1's answer. + CHECK(run_one(sched, turn2, cfg.eos_token_ids, kMax) == expect2); + // Warm 3 reuses blocks spanning turn 1's *generated* text — pooled by turn + // 2's prompt prefill (decode-produced K/V itself is never pooled; see + // Worker::harvest_finished). + CHECK(run_one(sched, turn3, cfg.eos_token_ids, kMax) == expect3); + + const mlxforge::WorkerMetrics m = worker.metrics(); + CHECK(m.prefix_hits == 3); + const long long bs = 16; + // turn1 warm reuses turn1's full blocks; turn2 reuses the same; turn3 + // reuses turn2's full blocks (its prompt was pooled when turn2 finished). + const long long expected_reuse = (static_cast(turn1.size()) / bs) * bs * 2 + + (static_cast(turn2.size()) / bs) * bs; + CHECK(m.prefix_tokens_reused == expected_reuse); + CHECK(m.prefix_pool_bytes > 0); + CHECK(m.prefix_pool_blocks >= static_cast(turn3.size() / bs)); + + worker.stop(); +} diff --git a/tests/scheduler/prefix_spill_test.cpp b/tests/scheduler/prefix_spill_test.cpp new file mode 100644 index 0000000..caf4a7a --- /dev/null +++ b/tests/scheduler/prefix_spill_test.cpp @@ -0,0 +1,118 @@ +// End-to-end SSD spill gate: with a RAM pool deliberately too small to hold a +// prompt's blocks, reuse must round-trip through the spill tier (evict -> +// serialize -> SSD -> revive) and still reproduce the cold greedy stream +// exactly — including from a fresh Worker on the same spill dir (restart +// persistence). +#include + +#include + +#include +#include +#include +#include + +#include "core/config.h" +#include "core/weights.h" +#include "runtime/single_stream.h" +#include "runtime/worker.h" +#include "scheduler/request.h" +#include "scheduler/scheduler.h" +#include "support/model_fixture.h" +#include "support/reference.h" + +using namespace mlxforge::test; +namespace fs = std::filesystem; + +namespace { + +std::vector long_prompt() { + std::vector p; + for (const char* name : {"prompt_2_ids.npy", "prompt_0_ids.npy", "prompt_1_ids.npy", + "prompt_2_ids.npy", "prompt_0_ids.npy", "prompt_1_ids.npy"}) { + std::vector ids = load_token_ids(name); + p.insert(p.end(), ids.begin(), ids.end()); + } + return p; // 46 tokens: two full 16-token blocks +} + +std::vector run_one(mlxforge::Scheduler& sched, const std::vector& prompt, + const std::vector& eos_ids, int max_tokens) { + auto req = std::make_shared(); + req->prompt_ids = prompt; + req->params.temperature = 0.0f; + req->max_tokens = max_tokens; + req->eos_ids = eos_ids; + REQUIRE(sched.submit(req)); + std::vector got; + int tok = 0; + while (req->tokens.pop(tok)) got.push_back(tok); + return got; +} + +} // namespace + +TEST_CASE("prefix blocks spill to SSD, revive exactly, and survive a restart") { + if (!model_available()) { + MESSAGE("MLXFORGE_MODEL_DIR not present; skipping"); + return; + } + const std::string dir = model_dir(); + mlxforge::ModelConfig cfg = mlxforge::ModelConfig::from_file(dir + "/config.json"); + const int kMax = 16; + + const std::vector prompt = long_prompt(); + const std::vector expected = + mlxforge::greedy_generate(shared_model(), prompt, kMax, cfg.eos_token_ids).tokens; + + char tmpl[] = "/tmp/mlxforge_prefix_spill_XXXXXX"; + const std::string spill_dir = mkdtemp(tmpl); + + // One 16-token block of this model is n_layers * (K+V) * 8 heads * 64 dims + // * fp16 = 512 KiB; a 600 KB pool holds exactly one, forcing the other + // through the spill tier. + mlxforge::PrefixCacheConfig pcfg; + pcfg.enabled = true; + pcfg.block_size = 16; + pcfg.pool_bytes = 600'000; + pcfg.salt = 7; + pcfg.spill_dir = spill_dir; + auto factory = [dir] { + mlxforge::ModelConfig c = mlxforge::ModelConfig::from_file(dir + "/config.json"); + auto w = mlxforge::load_weights(dir, c); + return std::make_unique(std::move(c), std::move(w)); + }; + + { + mlxforge::Scheduler sched; + mlxforge::Worker worker(factory, &sched, nullptr, {}, pcfg); + worker.start(); + + CHECK(run_one(sched, prompt, cfg.eos_token_ids, kMax) == expected); // cold + // Harvest inserted 2 blocks into a 1-block pool: one was spilled. + CHECK(worker.metrics().spill_writes >= 1); + + // Warm: the evicted block must come back from the SSD tier. + CHECK(run_one(sched, prompt, cfg.eos_token_ids, kMax) == expected); + const mlxforge::WorkerMetrics m = worker.metrics(); + CHECK(m.prefix_hits == 1); + CHECK(m.spill_reads >= 1); + worker.stop(); + } // ~Worker drains the spill writer: everything queued is now on disk + + // Restart: a fresh worker on the same spill dir starts warm. + { + mlxforge::Scheduler sched; + mlxforge::Worker worker(factory, &sched, nullptr, {}, pcfg); + worker.start(); + CHECK(run_one(sched, prompt, cfg.eos_token_ids, kMax) == expected); + const mlxforge::WorkerMetrics m = worker.metrics(); + CHECK(m.prefix_hits == 1); + CHECK(m.prefix_tokens_reused > 0); + CHECK(m.spill_reads >= 1); + worker.stop(); + } + + std::error_code ec; + fs::remove_all(spill_dir, ec); +}