diff --git a/.github/workflows/army-test.yml b/.github/workflows/army-test.yml new file mode 100644 index 000000000000..d37d60300f2d --- /dev/null +++ b/.github/workflows/army-test.yml @@ -0,0 +1,85 @@ +name: CI (army — paged + tiered KV) + +# MAD-137 CI integration. Hosted-runner-only: +# - Compile for the CPU backend so the mt:: + paged sources catch any +# cross-arch breakage immediately. +# - Run the mt:: unit tests + the integration tests with no model +# (they self-skip when LLAMACPP_TEST_MODELFILE is unset, but the +# binaries still have to build cleanly). +# +# Real-hardware testing on the army GPUs (R9700, 6900XT, 1070, RX 480) +# is run manually via scripts/test/run-army-matrix.sh on the dev boxes. +# Self-hosted runners were intentionally not wired up — they expose the +# dev boxes to inbound CI traffic, which we don't want. + +on: + workflow_dispatch: # manual trigger + push: + branches: [master, 'feat/MAD-*'] + paths: + - '.github/workflows/army-test.yml' + - 'src/llama-kv-cache-paged.{h,cpp}' + - 'src/memory-tier/**' + - 'tests/test-mt-*' + - 'tests/test-paged-*' + - 'tests/stress/**' + - 'scripts/test/**' + pull_request: + types: [opened, synchronize, reopened] + paths: + - '.github/workflows/army-test.yml' + - 'src/llama-kv-cache-paged.{h,cpp}' + - 'src/memory-tier/**' + - 'tests/test-mt-*' + - 'tests/test-paged-*' + - 'tests/stress/**' + - 'scripts/test/**' + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_LOG_COLORS: 1 + LLAMA_LOG_PREFIX: 1 + LLAMA_LOG_TIMESTAMPS: 1 + +jobs: + hosted-cpu-tests: + name: "hosted (cpu) — unit + integration smoke" + runs-on: ubuntu-24.04 + steps: + - name: Clone + uses: actions/checkout@v6 + + - name: ccache + uses: ggml-org/ccache-action@v1.2.21 + with: + key: army-hosted-cpu + evict-old-files: 1d + + - name: Configure + run: | + cmake -B build \ + -DGGML_CUDA=OFF -DGGML_HIP=OFF -DGGML_VULKAN=OFF \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLAMA_BUILD_TESTS=ON + + - name: "Build (mt:: tests + paged-* tests)" + run: | + cmake --build build --target \ + test-mt-quant \ + test-mt-block-pool \ + test-mt-block-table \ + test-mt-block-semantic-index \ + test-mt-tiered-thin \ + test-paged-lifecycle \ + test-paged-semantic \ + -j $(nproc) + + - name: "ctest — mt:: + paged-* (skips integration without model)" + working-directory: build + run: | + ctest -R 'test-mt-|test-paged-' --output-on-failure diff --git a/.gitignore b/.gitignore index 30d1baa81328..5f8c8f0b1c9b 100644 --- a/.gitignore +++ b/.gitignore @@ -184,6 +184,7 @@ llama.pc ggml/ggml-config.cmake ggml/ggml-version.cmake tests/libgguf-model-data.a +tests/results/ *.a tests/cmake_install.cmake Makefile diff --git a/README.md b/README.md index be23abcea67f..26ba9f6d8abb 100644 --- a/README.md +++ b/README.md @@ -523,6 +523,7 @@ To learn more about model quantization, [read this documentation](tools/quantize - [completion](tools/completion/README.md) - [server](tools/server/README.md) - [GBNF grammars](grammars/README.md) +- [Tiered KV cache](docs/memory-tier/) — paged + tiered + semantic prefetch for long-context / multi-agent serving (see [user guide](docs/memory-tier/USER-GUIDE.md), [operator runbook](docs/memory-tier/OPERATOR-RUNBOOK.md), [architecture](docs/memory-tier/ARCHITECTURE.md)) #### Development documentation diff --git a/common/arg.cpp b/common/arg.cpp index 41382976c867..37181195c1eb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1477,9 +1477,10 @@ common_params_context common_params_parser_init(common_params & params, llama_ex add_opt(common_arg( {"--kv-tier-paged-blocks"}, {"--no-kv-tier-paged-blocks"}, - "EXPERIMENTAL: enable mt:: paged attention KV cache (vLLM-style block-indexed layout). Routes attention through mt_paged_attention_kernel + mt_reshape_and_cache scatter on HIP/CUDA. Currently validated only on standard non-hybrid transformer models with per-attention-call ctx ≤ ~16k tokens — the kernel's smem footprint scales with ctx and overflows AMD's 64 KiB LDS limit beyond that (clear error logged at dispatch). Hybrid (DeltaNet/Mamba+attention) models compile through the paged path but produce incorrect attention output — needs further debugging. SWA models fall back to the regular kv cache. For multi-agent serving on hybrid models prefer --kv-tiered without this flag (software-only tier eviction; works at --parallel > 1).", + "MAD-134: paged-attention KV cache (vLLM-style block-indexed). When --kv-tiered is also set, this is auto-enabled by default (use --no-kv-tier-paged-blocks to opt out). Validated end-to-end on hybrid models (Qwen3.x family); non-hybrid + ctx > 16k may hit kernel LDS limits.", [](common_params & params, bool value) { params.kv_tier_paged_blocks = value; + params.kv_tier_paged_blocks_explicit = true; // MAD-134: user said something } ).set_env("LLAMA_ARG_KV_TIER_PAGED_BLOCKS").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( @@ -1492,6 +1493,39 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.kv_tier_paged_block_size = value; } ).set_env("LLAMA_ARG_KV_TIER_PAGED_BLOCK_SIZE").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--kv-tier-cold-resume"}, + {"--no-kv-tier-cold-resume"}, + "MAD-130: when true, skip O_TRUNC on cold-tier files at startup and load the in-memory index from " + "${ssd_path}/paged/instance-${INSTANCE_ID}/index.bin (written by the prior server's clean shutdown or " + "/slots/save). Lets the server resume cold-tier contents across a clean restart. Sidecar absent or " + "invalid → starts fresh with a warning. Default false (legacy behavior: fresh truncation).", + [](common_params & params, bool value) { + params.kv_tier_cold_resume = value; + } + ).set_env("LLAMA_ARG_KV_TIER_COLD_RESUME").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--instance-id"}, "ID", + "MAD-131: per-instance ID used to scope the cold-tier subdir (${ssd_path}/paged/instance-${ID}/) " + "and the per-instance lockfile. Lets multiple llama-server processes share one --kv-tier-ssd-path " + "without colliding. Default = process pid as a string. Use a stable ID (e.g. 'main-r9700') for " + "deterministic restarts with --kv-tier-cold-resume.", + [](common_params & params, const std::string & value) { + params.kv_tier_instance_id = value; + } + ).set_env("LLAMA_ARG_INSTANCE_ID").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--kv-tier-cold-budget-mb"}, "N", + "MAD-131: cap the cold-tier pool to N MiB total (across K+V × all attn layers). 0 = no cap " + "(size from --kv-tiered cold percentage). Use this to bound SSD wear per instance — e.g. " + "10000 (10 GiB) is reasonable for a consumer NVMe with ~600 TBW lifetime serving an army.", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("--kv-tier-cold-budget-mb must be >= 0"); + } + params.kv_tier_cold_budget_mb = value; + } + ).set_env("LLAMA_ARG_KV_TIER_COLD_BUDGET_MB").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.cpp b/common/common.cpp index 29b52c84286d..478194eedea2 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1547,8 +1547,23 @@ struct llama_context_params common_context_params_to_llama(const common_params & : params.kv_semantic_index.c_str(); cparams.kv_tier_semantic_threshold = params.kv_semantic_threshold; cparams.kv_tier_semantic_topk = params.kv_semantic_top_k; - cparams.kv_tier_paged_blocks = params.kv_tier_paged_blocks; + // MAD-134: auto-enable paged-blocks when --kv-tiered is set AND + // the user didn't explicitly choose either way. The army-goal use + // case (hybrid models + multi-agent serving) wants paged-on; the + // explicit-pct path stays opt-in via --no-kv-tier-paged-blocks. + bool effective_paged = params.kv_tier_paged_blocks; + if (!params.kv_tier_paged_blocks_explicit && params.kv_tiered_enabled && !params.kv_tier_paged_blocks) { + effective_paged = true; + LOG_INF("%s: auto-enabled --kv-tier-paged-blocks (--kv-tiered set; pass " + "--no-kv-tier-paged-blocks to opt out)\n", __func__); + } + cparams.kv_tier_paged_blocks = effective_paged; cparams.kv_tier_paged_block_size = params.kv_tier_paged_block_size; + cparams.kv_tier_cold_resume = params.kv_tier_cold_resume; + cparams.kv_tier_instance_id = params.kv_tier_instance_id.empty() + ? nullptr + : params.kv_tier_instance_id.c_str(); + cparams.kv_tier_cold_budget_mb = params.kv_tier_cold_budget_mb; return cparams; } diff --git a/common/common.h b/common/common.h index 8e9fd8c06c55..80dc6beb5c48 100644 --- a/common/common.h +++ b/common/common.h @@ -613,7 +613,11 @@ struct common_params { float kv_semantic_threshold = 0.65f; // minimum cosine similarity threshold for prefetch hints int kv_semantic_top_k = 5; // number of prefetch hints to return bool kv_tier_paged_blocks = false; // enable mt:: paged-attention KV cache (vLLM-style block-indexed); standard + hybrid models supported + bool kv_tier_paged_blocks_explicit = false; // MAD-134: true when user typed --kv-tier-paged-blocks or --no-...; false → auto-default applies int kv_tier_paged_block_size = 16; // tokens per block when paged_blocks is enabled (must be a power of 2; 16 matches vLLM) + bool kv_tier_cold_resume = false; // MAD-130: skip O_TRUNC on cold-tier files; load index sidecar from prior run + std::string kv_tier_instance_id; // MAD-131: per-instance ID for cold-tier subdir + lockfile (default: pid) + int kv_tier_cold_budget_mb = 0; // MAD-131: cap cold-pool size to N MiB (0 = no limit beyond percent-derived) std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/docs/memory-tier/ARCHITECTURE.md b/docs/memory-tier/ARCHITECTURE.md new file mode 100644 index 000000000000..cc6e95902575 --- /dev/null +++ b/docs/memory-tier/ARCHITECTURE.md @@ -0,0 +1,420 @@ +# Tiered KV cache — architecture + +How the paged + tiered + semantic-prefetch stack hangs together. This +doc is for the engineer modifying or extending the tier code; it is +not a "how do I configure the cache" doc (that's +[USER-GUIDE.md](USER-GUIDE.md)). + +The full epic context is in [Jira MAD-126](#jira-references) but Jira +isn't where engineers read documentation while debugging. The +authoritative source for design decisions is this file plus the per- +decision ADRs under [`adr/`](adr/). + +--- + +## The three-tier model + +``` + ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ + │ Hot (VRAM) │ ───▶ │ Warm (host) │ ───▶ │ Cold (SSD) │ + │ │ ◀─── │ │ ◀─── │ │ + └─────────────┘ └─────────────┘ └─────────────┘ + GPU buffers host RAM per-layer files at + sized by staging ${ssd_path}/paged/ + --kv-tiered buffers instance-${ID}/L*.{k,v}.bin + HOT% +``` + +A KV block — `block_size` tokens × `n_kv_heads` × `head_dim` of K and +V data per attention layer — lives in exactly one tier at any moment. +The tier is identified by the block's physical id range: + +- IDs `[0, n_blocks_total)` → GPU pool (hot) +- IDs `[n_blocks_total, n_blocks_total + n_warm_blocks)` → CPU pool (warm) +- Cold tier is keyed by `(seq_id, logical_block_idx)` rather than by + physical id — `KvtcStore` in `src/memory-tier/mt-kvtc-store.{h,cpp}` + maps those into per-layer file offsets. + +Movement between tiers is one-directional per call but fully bidirectional +in aggregate: + +| From | To | Trigger | Mechanism | +|---|---|---|---| +| Hot | Warm | `evict_lru_to_warm()` from `ensure_blocks_for` | `ggml_backend_tensor_get` into `warm_k_/warm_v_` host buffers | +| Warm | Hot | `restore_block_from_warm()` from semantic restore or kernel demand-fault | `ggml_backend_tensor_set` from host buffers back into the GPU layer tensor | +| Warm | Cold | `spill_one_to_cold()` when warm pool fills | `KvtcStore::write` to SSD; warm slot freed | +| Cold | Hot | `restore_semantic_paged()` cold-fault path | `KvtcStore::read` → host buffer → `ggml_backend_tensor_set` | +| Cold | (drop) | Cold pool full + new spill needed | `drop_oldest_cold_block()`; **data is lost**, table entry becomes `kInvalidBlockId` | + +Block contents are F16 (or whatever `--cache-type-k/v` is set to) on +hot; F16 in `warm_k_/v_` host RAM; **int4-with-scale** on cold. The +int4 cold compression is per-block: each block stores a single +`scale: f32` followed by `ceil(n_elts/2)` packed int4 nibbles. Cosine +similarity ≈ 0.99 against the unquantized baseline; this is acceptable +because cold blocks were already going to be re-attended-to in +combination with their original-precision neighbors. See +[`src/memory-tier/mt-quant.cpp`](../../src/memory-tier/mt-quant.cpp). + +--- + +## The paged-attention block model + +Adapted from vLLM's BlockTable / BlockPool design (Apache 2.0). The +key idea: tokens are not stored contiguously per sequence; they're +stored in **fixed-size physical blocks**, and a per-sequence +**logical→physical table** records which block holds which range. + +``` +seq 0 logical view physical pool view +┌────────┬────────┬────────┐ ┌────────┐ block 7 +│ tok0 │ tok16 │ tok32 │ ├────────┤ +│ ... │ ... │ ... │ │ tok0… │ block 0 +│ tok15 │ tok31 │ tok47 │ ├────────┤ +└────────┴────────┴────────┘ │ tok16… │ block 4 + ↓ ↓ ↓ ├────────┤ +table_[0] = [0, 4, 7] │ ... │ + └────────┘ +``` + +Why this matters: + +1. **Eviction is block-granular**, not byte-granular. The evict path + moves a single block (~16 tokens × per-token K/V row size) at a + time. That's a clean unit for `ggml_backend_tensor_get/set`. +2. **CoW is cheap**. `seq_cp` increments per-block refcounts in + `BlockPool` rather than copying bytes. A branched agent workflow + creating 10 conversation forks pays one block-table copy per fork + and zero K/V copies until a fork actually writes new tokens. See + ADR-A10. +3. **Holes are first-class**. A partial `seq_rm` can wipe one logical + block in the middle of a sequence by setting that table entry to + `kInvalidBlockId`. The paged-attn kernel reads `kInvalidBlockId` → + returns `-INFINITY` for the wiped positions → softmax weights → + zero contribution. No cache-side gymnastics required. + +The block size is fixed at construction (default 16 tokens). The +paged-attn kernel +([`ggml/src/ggml-cuda/mt_pagedattn.cu`](../../ggml/src/ggml-cuda/mt_pagedattn.cu)) +supports `(head_size, block_size) ∈ {(128, 16), (64, 16), (256, 16), +(128, 32)}` and K cache types `{F16, Q8_0, TURBO4_0}`. + +--- + +## Why hybrid+paged is THE primary path + +(See [adr/A1-hybrid-paged-primary.md](adr/A1-hybrid-paged-primary.md).) + +User direction 2026-05-10: "hybrid+paged is THE primary path; pure- +attention is least concern." All new tier features land on +`llama_kv_cache_paged`. Pure-attention via `mt::llama_memory_tiered` +stays passthrough; the legacy `llama_kv_cache_tiered` / +`server_tiered_cache` paths got removed in MAD-127. + +The `mt::llama_memory_tiered` wrapper survives in a thin form for two +specific responsibilities: + +1. **bge-small embed model ownership.** The wrapper holds a single + `EmbeddingModel` instance shared across server lifetime, and exposes + `embed_text(string) → vector` to whoever needs to compute a + fingerprint or query. +2. **Recurrent-state backup on hybrid models.** The recurrent half of + a hybrid model (Gated Delta Net, Mamba, etc.) uses + `llama_memory_recurrent`, whose `clear()` loses everything. The + tiered wrapper backs up the per-seq recurrent state into host RAM + on `seq_rm` so it can be restored later. + +Everything else — block management, eviction, persistence, semantic +prefetch — is in `llama_kv_cache_paged`. + +--- + +## Class hierarchy + +``` +llama_memory_i (interface) +└── llama_memory_hybrid (composition for hybrid models) + ├── mem_attn → llama_kv_cache_paged ◀── the active tier layer + └── mem_recr → llama_memory_recurrent (recurrent-only state) + +mt::llama_memory_tiered (thin wrapper, optional) +└── inner_ → llama_memory_hybrid (when wrapping a hybrid) + or → llama_kv_cache_paged (when wrapping pure paged) +``` + +For hybrid models routed through the tier stack, the runtime stack is: + +``` +server-side dispatch + ↓ +llama_context::decode + ↓ +llama_memory_hybrid (split into attn vs recr ubatches) + ├─→ mem_attn = llama_kv_cache_paged ──── paged-attn kernel dispatch + └─→ mem_recr = llama_memory_recurrent ── recurrent kernel dispatch +``` + +The `mt::llama_memory_tiered` wrapper sits **outside** this stack when +present — it intercepts `seq_rm` etc. for the recurrent backup +behavior, then delegates to its inner cache. + +The server-context helper `mt_get_paged_cache(llama_memory_i*)` peels +through three nestings (raw paged / hybrid / tiered+hybrid) to return +the underlying `llama_kv_cache_paged*`. See +[`tools/server/server-context.cpp`](../../tools/server/server-context.cpp) +near the `mt_get_paged_cache` definition. + +--- + +## The single-threading contract + +(See [adr/A4-single-threading.md](adr/A4-single-threading.md).) + +`llama_kv_cache_paged` and its `BlockPool` / `BlockTable` are NOT +internally locked. The server's main loop is the single mutator; HTTP +worker threads communicate with it via `server_queue`'s task channel +and never touch the cache directly. + +Concretely: + +- Tier-counter `_total` accessors (e.g. `evict_h2w_total()`) return + `uint64_t` from the cache. They're `volatile`-style monotonic counters + that the `/metrics` HTTP handler reads from a different thread — + this is the **only** cross-thread read on the cache. It's safe-ish + on x86_64 / aarch64 for monotonic 64-bit counters (atomic loads are + single instructions). +- Any future async path (semantic prefetch on a worker thread, bge- + small embed batched off the critical path, etc.) must gate on a real + concurrency design — not "by accident." +- Debug builds include a thread-id assertion (`check_thread_id_()` in + `llama_kv_cache_paged`) that traps if anyone other than the registered + main thread mutates the cache. + +This contract is documented in +[`src/llama-kv-cache-paged.h`](../../src/llama-kv-cache-paged.h) +near the class declaration and in +[`src/memory-tier/mt-block-pool.h`](../../src/memory-tier/mt-block-pool.h). + +--- + +## The persistence model + +(See [adr/A5-persistence-explicit.md](adr/A5-persistence-explicit.md).) + +The cache supports explicit save/restore via the server's `/slots/save` +and `/slots/restore` endpoints. Crash recovery is **not** automatic; +the model is "clean shutdown saves, restart restores." + +State written by `state_write()`: +- Block table per seq (logical→physical mapping). +- Hot-tier K/V tensors (or skip + reprefill — caller-configurable). +- Warm-tier K/V buffers. +- Cold-tier index sidecar (CIDX v1 magic = `0x58444943`). +- BlockSemanticIndex fingerprints (PSFI v1 magic = `0x49465350`). + +Format magic numbers and version bytes let future format changes be +detected and rejected (rather than silently corrupting state). + +The cold-tier sidecar is the recovery hinge: per-layer K/V files +contain the actual block data, but without the sidecar's +`(seq, lblock) → file_offset` mapping the cache can't read them back. +Hard crashes that miss the sidecar write produce orphan files — +`scripts/army/cleanup-cold.sh` removes them. + +--- + +## The multi-instance model + +(See [adr/A6-multi-instance.md](adr/A6-multi-instance.md).) + +Multiple `llama-server` processes can share one `--kv-tier-ssd-path` +without colliding because each writes cold-tier files under a +per-instance subdirectory: + +``` +${ssd_path}/paged/ +├── instance-main-r9700/ +│ ├── L0.k.bin +│ ├── L0.v.bin +│ ├── ... +│ ├── instance.lock # flock'd while the process is alive +│ └── index.bin # CIDX v1 sidecar +├── instance-main-6900xt/ +│ └── ... +``` + +`--instance-id` defaults to the process PID but is typically set +explicitly by the boot script for stable cold-resume across restarts. + +The `flock`-based lockfile prevents accidental double-start: a second +process trying to open the same instance subdir fails fast with a clear +error rather than silently corrupting the on-disk state. + +--- + +## The semantic prefetch model + +(See [adr/A2-fingerprint-at-prefill.md](adr/A2-fingerprint-at-prefill.md) +and [adr/A3-semantic-prefetch-only.md](adr/A3-semantic-prefetch-only.md).) + +Two paths share the same `BlockSemanticIndex` storage in +`src/memory-tier/mt-semantic.{h,cpp}`: + +### Write path: server-side prefill trigger + +When the server processes a prompt, after the prefill batch lands +and the block table is fully populated, the server walks the new seq's +**complete** logical blocks (any block with `n` < `block_size` of +its tokens written is skipped) and: + +1. Decodes the original tokens for that block back to text via + `slot.prompt.tokens` plus `common_token_to_piece`. +2. Calls `mt::llama_memory_tiered::embed_text(text)` to get an + L2-normalized 384-dim BGE-small embedding. +3. Calls `llama_kv_cache_paged::record_paged_block_fingerprint(seq, lblock, embedding, tier=Hot)`. + +CPU cost is ~5ms × n_complete_blocks per prefill, off the GPU +critical path. Skipped when the block already has a fingerprint +(the `has_paged_fingerprint(seq, lblock)` short-circuit). + +### Read path: server-side prefill query + +After every prefill the server also computes an embedding of the +**most recent** complete block (the "query" for purposes of semantic +recall) and calls +`llama_kv_cache_paged::restore_semantic_paged(seq, query_embedding, +top_k, threshold)`. The cache scores its stored fingerprints against +the query, picks the top-K above threshold, and faults each matching +block back from warm/cold to hot before kernel dispatch. + +### Why prefill-time and not eviction-time + +Eviction in the paged cache is internal — the server doesn't see +eviction events. Fingerprints written at eviction time would be +strictly *behind* the data they describe (the data has already been +evicted) and the timing is hard to control. Writing at prefill time +makes the fingerprint write a clean, predictable, additive operation +attached to the server's ordinary task lifecycle. + +### Why semantic doesn't drive eviction + +(See [adr/A3-semantic-prefetch-only.md](adr/A3-semantic-prefetch-only.md).) + +bge-small drives prefetch only, not eviction. Eviction stays on the +hybrid attention/recency/frequency policy in +`src/memory-tier/mt-eviction.{h,cpp}`. The reasons: +- Training-task mismatch: bge-small is trained for retrieval, not for + inference cache predictiveness. +- Hot-path latency: every eviction decision would need an embed call. +- Doesn't fix the structural problem (hot-pool fragmentation under + multi-seq load — that's the MAD-120 admission control's job). + +--- + +## Kernel dispatch + +The paged-attn kernel +([`ggml/src/ggml-cuda/mt_pagedattn.cu`](../../ggml/src/ggml-cuda/mt_pagedattn.cu)) +takes the layer's K and V tensors plus the block_table tensor and +context-lens / q-lens (per-batch tensors maintained by +`llama_kv_cache_paged::prepare_batch_tensors`). + +Dispatch table: + +``` +type_k: F16 | Q8_0 | TURBO4_0 +(head, block): + (128, 16) + (64, 16) + (256, 16) + (128, 32) +``` + +Aborts on unsupported tuples with a clear error. Hybrid models with +attention layers that fall outside this dispatch table cannot use +paged-attn until the kernel is extended. + +`kInvalidBlockTableEntry` (matches `mt::kInvalidBlockId`) is handled +in the kernel: any (seq, position) whose physical block id is the +sentinel returns `-INFINITY` as its attention logit. After softmax +this contributes zero weight to the attention output — equivalent to +"that token doesn't exist." Used by partial seq_rm and by cold-drop. + +--- + +## Eviction state machine + +``` + ┌────────────────────────────┐ + │ ensure_blocks_for(seq, N) │ + └─────────┬──────────────────┘ + │ + ┌───────────────────────▼─────────────────────────┐ + │ pool_.alloc_gpu() │ + └─┬─────────────────────────────────────────────┬──┘ + │ ok │ kInvalidBlockId + │ │ + ▼ ▼ + ┌────────────┐ ┌──────────────────────────┐ + │ DONE │ │ evict_lru_to_warm() │ + └────────────┘ └──┬─────────────────────┬──┘ + │ ok │ false (warm full) + ▼ │ + ┌──────────────┐ │ + │ retry alloc │ ▼ + └──┬───────────┘ ┌────────────────────┐ + │ ok │ spill_one_to_cold() │ + ▼ └──┬─────────────────┬┘ + ┌────────────┐ │ ok │ false (cold full) + │ DONE │ ▼ ▼ + └────────────┘ ┌────────────────┐ ┌─────────────────────────┐ + │ retry alloc │ │ drop_oldest_cold_block() │ + └────────────────┘ └──┬───────────────────────┘ + │ + ▼ + ┌──────────────────────────┐ + │ retry; on failure: false │ + └──────────────────────────┘ +``` + +The bottom escalation (`drop_oldest_cold_block` → data loss with +sentinel) is the **last resort**. The drop counter +(`paged_evict_cold_to_drop_total`) is the operator's signal to +re-tune sizing. + +--- + +## Where things live + +| File | What's in it | +|---|---| +| `src/llama-kv-cache-paged.h/.cpp` | The hot/warm-tier cache. Block table, eviction, semantic restore, state save/load, multi-instance lockfile. | +| `src/memory-tier/mt-block-pool.h/.cpp` | Physical block allocator (GPU + CPU pools, refcounting, watermark). | +| `src/memory-tier/mt-block-table.h/.cpp` | Per-seq logical→physical mapping. | +| `src/memory-tier/mt-semantic.h/.cpp` | `SemanticIndex` (chunk-level) + `BlockSemanticIndex` (per-block); save/load PSFI v1. | +| `src/memory-tier/mt-tiered.h/.cpp` | Thin `mt::llama_memory_tiered` wrapper: bge-small ownership, recurrent backup. | +| `src/memory-tier/mt-quant.h/.cpp` | int4 / int8 quant helpers including the per-block scaled int4 used for cold compression. | +| `src/memory-tier/mt-kvtc-store.h/.cpp` | Cold-tier file I/O. | +| `src/memory-tier/mt-eviction.h/.cpp` | `TokenMetadataStore` and the hybrid eviction policy. | +| `src/memory-tier/mt-mover-attn.h/.cpp` | Attention K/V mover for `mt::llama_memory_tiered`. | +| `src/memory-tier/mt-mover-recurrent.h/.cpp` | Recurrent state mover. | +| `src/memory-tier/mt-embed.h/.cpp` | `EmbeddingModel` wrapper around the bge-small gguf. | +| `src/memory-tier/mt-config.h/.cpp` | `TieredConfig` parsing. | +| `src/memory-tier/mt-capacity.h/.cpp` | `TierCapacityManager` for non-paged tiered (legacy). | +| `tools/server/server-context.cpp` | Server-side dispatch: prefill-time fingerprint write trigger, semantic restore call, /metrics extension, /slots tier residency, MAD-141 admission deadlock guard. | +| `ggml/src/ggml-cuda/mt_pagedattn.cu` | The paged-attn kernel + dispatch. | +| `tests/test-mt-*.cpp` | Unit tests for tier primitives. | +| `tests/test-paged-*.cpp` | Integration tests against `llama_kv_cache_paged`. | +| `tests/stress/stress-paged-multi-seq.py` | HTTP-driven stress driver. | +| `scripts/test/run-army-matrix.sh` | Per-device matrix runner. | +| `scripts/army/*.sh` | Boot scripts. | + +--- + +## Jira references + +- **Epic**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) + — production-quality paged + tiered + multi-seq for the agent army. +- **Children**: MAD-127 through MAD-138. See the Epic description for + the story map. +- **ADRs**: extracted from MAD-126's "Architecture decisions" section + into [`adr/`](adr/) — ten files A1-A10 covering the resolved design. diff --git a/docs/memory-tier/OPERATOR-RUNBOOK.md b/docs/memory-tier/OPERATOR-RUNBOOK.md new file mode 100644 index 000000000000..7164ea14de5c --- /dev/null +++ b/docs/memory-tier/OPERATOR-RUNBOOK.md @@ -0,0 +1,358 @@ +# Tiered KV cache — operator runbook + +Day-2 operations for the army stack: four `llama-server` instances +across four GPUs, each running paged + tiered with semantic prefetch. +This document is the playbook for "what do I do when X happens at 2am." + +For end-user docs see [USER-GUIDE.md](USER-GUIDE.md). For internals +see [ARCHITECTURE.md](ARCHITECTURE.md). + +--- + +## Topology + +| Instance ID | Box | GPU | Build | Model | +|---|---|---|---|---| +| `main-r9700` | main dev box | R9700 (gfx1201, 32 GB) | `build-hip` | Qwen3.6-27B-Q6_K | +| `main-6900xt` | main dev box | 6900XT (gfx1030, 16 GB) | `build-hip` | Qwen3.5-9B-TQ3_1S | +| `mad-lab-1070` | mad-lab-2026 | GTX 1070 (sm_61, 8 GB) | `build-army` (CUDA) | omnicoder-9b-q5_k_m | +| `mad-lab-rx480` | mad-lab-2026 | RX 480 (gfx803, 8 GB) | `build-rocm-gfx803` (in `rx480-army` docker) | omnicoder-9b-q5_k_m | + +Per-machine boot scripts live under `scripts/army/`: + +- `scripts/army/main.sh` — main dev box (r9700 + 6900xt) +- `scripts/army/mad-lab.sh` — mad-lab box (1070 + rx480) +- `scripts/army/cleanup-cold.sh` — wipes the cold-tier directory +- `scripts/army/army.service.example` — systemd template + +--- + +## Boot + +### Cold start + +```bash +# Main box +bash scripts/army/main.sh + +# mad-lab +ssh mad-lab-2026 'bash ~/GitHub/llama.cpp/scripts/army/mad-lab.sh' +``` + +Each script launches its instances in the background, writes pids and +log paths to stdout, and returns. Boot is considered complete when all +instances log `main: server is listening on http://…`. + +### Resume after a clean shutdown + +If the prior shutdown went through `/slots/save` (or systemd's +`ExecStop`), add `--kv-tier-cold-resume` to each instance's invocation. +The cold-tier files are re-opened and the cold index sidecar is +replayed. Without it the cold pool starts empty. + +### Hard restart (after crash) + +The cold-tier sidecar (`paged/instance-${ID}/index.bin`) is only +written on a graceful shutdown. After a crash the cold pool starts +empty regardless of `--kv-tier-cold-resume`; the per-layer files are +still on disk but with no live index they're inaccessible. If you want +clean disk hygiene, run `bash scripts/army/cleanup-cold.sh` before +restart to delete the orphan files. + +--- + +## Health checks + +### Per-instance + +```bash +# Aliveness +curl -sf http://127.0.0.1:11435/health | jq + +# Tier counters (Prometheus format; needs --metrics) +curl -s http://127.0.0.1:11435/metrics | grep -E 'paged_(evict|restore|semantic|seq|blocks)' + +# Per-seq tier residency +curl -s http://127.0.0.1:11435/slots | jq '.[] | {id, state, paged_tier}' +``` + +### Fleet-wide + +```bash +for port in 11435 11436 11437 11438; do + echo "=== :$port ===" + curl -sf "http://127.0.0.1:$port/health" >/dev/null && echo OK || echo DOWN +done +``` + +--- + +## Common alarms + +### Cold pool occupancy > 80% + +**Symptom**: `paged_blocks_capacity_cold` is fixed; the count of +unmapped (free) cold blocks is dropping toward zero. Visible by +diffing two `/metrics` snapshots, or by watching +`paged_evict_cold_to_drop_total` start to increment. + +**Action**: +1. Confirm the workload — is something pinning a lot of context that + isn't being released? Check `/slots` for tasks that have been + `is_processing=true` for an unusually long time. +2. If load is legitimate, raise `--kv-tier-cold-budget-mb` and restart + the affected instance with `--kv-tier-cold-resume` to preserve what + you have. +3. If load is a stuck task, kill the slot via the server's task API or + restart the instance. + +### `paged_evict_cold_to_drop_total` is non-zero + +**Symptom**: cold pool is full enough that the cache started **dropping +data** (not just spilling). Future queries against those positions get +`-INFINITY` logits — the kernel will try and silently produce nonsense. + +**Action**: this is a sizing failure. Either widen `COLD%`, raise +`--kv-tier-cold-budget-mb`, or reduce `-c` so the workload fits the +configured tiers. **Do not ignore this counter.** + +### Semantic hit rate < 10% + +**Symptom**: `paged_semantic_hits_total / paged_semantic_attempts_total` +< 0.10 across many requests. + +**Possible causes**: +- The workload doesn't revisit content (every turn is fresh prompt). + Semantic prefetch isn't useful here; consider disabling + `--kv-tier-semantic-index` to save the per-block embed cost. +- `--kv-tier-semantic-threshold` is too strict. Try lowering from 0.65 + to 0.55. +- The bge-small model failed to load. Check the boot log for + `mt::EmbeddingModel: failed to load model`. Common fix: re-download + the gguf or use a different embedding model. + +### Eviction rate > 100/sec sustained + +**Symptom**: `paged_evict_hot_to_warm_total` rate is climbing fast. + +**Action**: hot pool is undersized for the workload, or the hot +percentage is too low. Either raise `HOT%` in `--kv-tiered` or reduce +`--parallel` so each agent's working set has more room. Per-batch +eviction events are normal; sustained high rate hurts decode latency. + +### MAD-141: "paged KV admission could not fit a N-token request" + +**Symptom**: client receives HTTP 500 with this body. Server log shows +`MAD-141: paged admission stuck for 4 iters with no eviction progress`. + +**Action**: the prompt is genuinely too large to fit alongside the +live workload. Options: +- Caller sends a smaller prompt (chunk + use `cache_prompt: true`). +- Wait for sibling slots to release their blocks (they finish decoding, + client closes the connection). +- Operator: raise `-c` so the server has more total blocks, or reduce + `--parallel` so each slot has a larger share. + +This is a **graceful** failure mode (post-MAD-141 / commit `0d66d8aa3`). +Pre-fix builds would crash the server entirely; if you see a process +abort with `n_empty_consecutive > 3`, you're on a stale build. + +### Lockfile error on startup + +**Symptom**: `instance lock /…/instance.lock held by another process`. + +**Action**: another `llama-server` is already running with the same +`--instance-id` against the same `--kv-tier-ssd-path`. Either: +- Kill the previous instance (`pgrep -af "instance-id "` → + `kill `). +- Pick a different `--instance-id` for this one. +- If the previous instance crashed and left a stale lockfile, + `rm /…/instance.lock` then restart. Verify with `lsof` first that + no process actually holds it. + +--- + +## Manual interventions + +### Clean restart of a single instance + +```bash +INSTANCE=main-r9700 +PORT=11435 + +# Graceful shutdown via SIGTERM. Server will flush cold sidecar. +pkill -TERM -f "instance-id $INSTANCE" + +# Wait for socket to release. +while ss -ltn | grep -q ":$PORT"; do sleep 1; done + +# Restart with resume. +bash scripts/army/main.sh # or just the relevant chunk of it +``` + +### Cold-tier wipe (drop all SSD state) + +```bash +INSTANCE=main-r9700 +SSD_PATH=/var/lib/army/ssd + +# Stop the instance first. +pkill -TERM -f "instance-id $INSTANCE" + +# Delete this instance's cold-tier dir. +rm -rf "$SSD_PATH/paged/instance-$INSTANCE" + +# Restart without --kv-tier-cold-resume. +bash scripts/army/main.sh +``` + +For a fleet-wide cold wipe use `scripts/army/cleanup-cold.sh`. Stop +all instances first. + +### Fingerprint reset (drop semantic prefetch state) + +Fingerprints are stored alongside the live cache state and are dropped +on whole-seq `seq_rm` automatically. To force-drop without losing the +KV state, the cleanest path is: + +1. Save state via `/slots/save`. +2. Restart the instance (without `--kv-tier-semantic-index` if you + want to permanently disable semantic prefetch, or with it for a + fresh fingerprint generation). + +Per-block fingerprint deletion isn't exposed via HTTP; it requires a +custom server build that calls `paged_semantic_.clear()`. + +--- + +## Disaster recovery + +### Lost SSD (cold-tier dir disappeared) + +The hot and warm tiers are fine — they're in VRAM and host RAM. Cold +data is gone forever. The server will keep serving live requests; only +queries against positions that had been spilled to cold will read +sentinel logits (manifesting as garbage tokens for those positions). + +Action: +1. Recreate the directory and ensure permissions are correct. +2. Restart the instance **without** `--kv-tier-cold-resume`. Cold pool + starts fresh. +3. If client sessions need their cold-evicted context back, they'll + have to re-prefill it (i.e. resend the original long prompt). + +### Disk full on `--kv-tier-ssd-path` + +Cold writes start failing silently after the first ENOSPC. The cache +can't tell the difference between "wrote successfully" and "filesystem +silently dropped this." Eventually the server hits the cold pool's +configured size, can't drop further, and starts logging +`paged_evict_cold_to_drop_total` ticks. + +Action: +1. Free space on the disk (rotate logs, drop other temp data). +2. If urgent, reduce `--kv-tier-cold-budget-mb` and restart so the + server tries to use less. +3. Long-term: move `--kv-tier-ssd-path` to a larger disk. + +### GPU OOM on startup + +Server fails with `cudaMalloc failed: out of memory` or +`alloc_tensor_range: failed to allocate ROCm0 buffer`. + +Diagnose: +```bash +# AMD +rocm-smi --showmeminfo vram --showpids + +# NVIDIA +nvidia-smi +``` + +Common causes: +1. **Stale process holding VRAM**. Kill it (`kill -9 `). Common + on AMD with ungraceful shutdowns where the kernel module doesn't + reclaim immediately. +2. **HOT% sized too high** for the model's weight footprint. Reduce + `HOT%` in `--kv-tiered`, raise the warm/cold percentages. +3. **Other instance on the same GPU**. Confirm via `--device` mapping + that two `llama-server` instances aren't pointing at the same card. + +### Whole army down after kernel/driver update + +Symptoms: every instance fails to start with HIP/CUDA initialization +errors. + +Action: +1. Check `dmesg` for amdgpu/nvidia driver load failures. +2. Verify `rocm-smi` (AMD) and `nvidia-smi` (NVIDIA) both work. +3. Verify the binary's compiled-for arches still match the GPUs: + ``` + strings llama-server | grep -E 'gfx[0-9]+|sm_[0-9]+' + ``` + If a driver update bumped the GPU's minimum-supported arch beyond + what the binary targets, rebuild with the new arch in + `-DAMDGPU_TARGETS` / `-DCMAKE_CUDA_ARCHITECTURES`. + +--- + +## Per-device caveats (from MAD-136 verification) + +### R9700 (gfx1201, RDNA4) + +- Build with `-DAMDGPU_TARGETS=gfx1201` (or include it in a multi-arch + build). +- Stable on ROCm 6.4+; earlier ROCm versions don't recognize gfx1201. +- Decode rate baseline: ~22 t/s on Qwen3.6-27B-Q6_K with + `--parallel 1 -c 8192` and turbo4 KV. + +### 6900XT (gfx1030, RDNA2) + +- Build must include `-DAMDGPU_TARGETS=gfx1030`. **Silent crash** if + missing — HIP runtime emits "No compatible code objects found for + gfx1030" before llama.cpp's logger runs. Diagnose with + `AMD_LOG_LEVEL=3`. +- Runs as eGPU over Thunderbolt 3 in this fleet; PCIe bandwidth limits + apply but don't significantly affect inference (model load is the + one-time cost). + +### GTX 1070 (sm_61, Pascal) + +- Build with `-DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=61`. +- Compute capability 6.1 — no FP16 tensor cores, no bf16. Stick to + `f16` or `turbo4` cache types; `bf16` not supported by Pascal. + +### RX 480 (gfx803, Polaris) + +- **Vulkan path is ~10× slower than ROCm** on Polaris (1.9 vs 19.6 + t/s on omnicoder-9b in MAD-136 verification). Stick with ROCm in a + docker container (`rocm/dev-ubuntu-24.04:6.4-complete` base, env + `HSA_OVERRIDE_GFX_VERSION=8.0.3 PYTORCH_ROCM_ARCH=gfx803 + ROC_ENABLE_PRE_VEGA=1`). The `rx480-army` container on + `mad-lab-2026` is the working setup. +- Cannot run modern AMD-supported kernels natively; stays on the + custom container indefinitely (or until the box migrates to CachyOS). + +--- + +## Routine maintenance + +### Weekly + +- `df -h "$SSD_PATH"` — confirm cold-tier disk usage isn't drifting up. +- `journalctl --user -u army-* --since '7d ago' | grep ERROR` — scan + for warnings the alarms didn't catch. + +### Monthly + +- Rotate `/var/log/army/*.log` — these grow with verbosity 1 logs. +- Re-run `bash scripts/test/run-army-matrix.sh` against the live + branch to confirm regressions haven't snuck in. + +### After any branch merge that touches `src/llama-kv-cache-paged.*`, +`src/memory-tier/**`, `tools/server/server-context.cpp` + +- Rebuild affected boxes. +- Run `ctest -R 'test-mt-|test-paged-'` against each build. +- Run a 60-second smoke per device with the new binary before + promoting to production. diff --git a/docs/memory-tier/USER-GUIDE.md b/docs/memory-tier/USER-GUIDE.md new file mode 100644 index 000000000000..e1b10ee2702e --- /dev/null +++ b/docs/memory-tier/USER-GUIDE.md @@ -0,0 +1,242 @@ +# Tiered KV cache — user guide + +The tiered KV cache extends `llama-server` with a hot/warm/cold storage +hierarchy for the attention K/V state. It lets one server hold contexts +much larger than VRAM by spilling cold blocks to host RAM and SSD, with +optional semantic prefetch to bring the right blocks back when a query +revisits old material. + +This guide is for someone running `llama-server`. For the architecture +behind it, see [ARCHITECTURE.md](ARCHITECTURE.md). For the army-style +multi-instance operations setup, see +[OPERATOR-RUNBOOK.md](OPERATOR-RUNBOOK.md). + +--- + +## When to use it + +Turn on the tiered cache when **any** of these is true: + +- You want a context size larger than what fits in VRAM (e.g. 512k ctx + on a 32 GB card). +- You're running multiple agents on one server (`--parallel 4`+) and + want their contexts to coexist without OOM. +- You're running multiple `llama-server` instances on the same box and + want each to fail fast rather than fight over VRAM. +- You expect long sessions where parts of the context get revisited + (RAG-style queries, multi-turn agent workflows). + +## When *not* to use it + +Skip it when: + +- Your context fits in VRAM with comfortable headroom and you only run + one user. +- You're benchmarking raw decode throughput on a short prompt — the + tier machinery adds bookkeeping that doesn't pay off until the cache + actually fills. +- The model is pure-recurrent (no attention layers). The tier code is + attention-block-keyed; recurrent state lives elsewhere and isn't + paged. + +--- + +## Quick start + +The army-goal config — paged + tiered + semantic prefetch on a hybrid +model — is one extra line: + +```bash +./llama-server \ + -m /path/to/model.gguf \ + --device CUDA0 -ngl 99 \ + --parallel 4 -c 524288 --no-mmap \ + --kv-tier-paged-blocks --kv-tiered 25,75,0 \ + --cache-type-k turbo4 --cache-type-v turbo4 \ + --kv-tier-semantic-index /path/to/bge-small-en-v1.5-q8_0.gguf \ + --instance-id agent-1 +``` + +Read it as: "give this server 524k context split 25% hot / 75% warm / +0% cold, with paged-attn block management and bge-small semantic +prefetch, identified as `agent-1` so its cold-tier files don't collide +with sibling instances." + +For the same shape with a cold-tier slice on SSD: + +```bash +... --kv-tiered 25,25,50 \ + --kv-tier-ssd-path /var/llama/cache \ + --kv-tier-cold-budget-mb 8192 +``` + +--- + +## Flag reference + +### Tier-shape flags + +| Flag | Description | +|---|---| +| `--kv-tiered HOT,WARM,COLD` | Enable the tiered cache. Three percentages summing to 100. `25,75,0` = 25% VRAM, 75% host RAM, no SSD. `25,25,50` = SSD spillover. | +| `--kv-tier-paged-blocks` | Use vLLM-style block-indexed paged attention. Default-on for hybrid models when `--kv-tiered` is set; pass `--no-kv-tier-paged-blocks` to opt out. | +| `--kv-tier-paged-block-size N` | Tokens per paged block. Default 16; rarely needs changing. Must match the paged-attn kernel's supported block sizes (16 or 32). | +| `--kv-tier-total-ctx N` | Override the total ctx the tier sizing math uses. Defaults to `-c`. | + +### Cold-tier (SSD) flags + +| Flag | Description | +|---|---| +| `--kv-tier-ssd-path PATH` | Directory under which cold-tier files are written. Required when `COLD%` > 0. The server creates `paged/instance-${INSTANCE_ID}/` underneath. | +| `--kv-tier-cold-budget-mb N` | Cap the cold pool size to N MiB total (K+V across all attn layers). 0 = no cap (size from cold percentage). Use this to bound SSD wear. | +| `--kv-tier-cold-resume` | On startup, re-open the cold-tier files from the prior run instead of truncating. Pairs with `--instance-id`. Default: off (clean truncation). | +| `--no-kv-tier-cold-resume` | Explicit opt-out (overrides resume default if set elsewhere). | + +### Multi-instance flags + +| Flag | Description | +|---|---| +| `--instance-id ID` | Name this server instance. Used as the cold-tier subdir name (`paged/instance-${ID}/`) and to identify the lockfile. Default = the process pid. Required when running multiple instances against the same `--kv-tier-ssd-path`. | + +### Semantic prefetch flags + +| Flag | Description | +|---|---| +| `--kv-tier-semantic-index PATH` | Path to the embedding model (bge-small-en-v1.5 quantized to q8_0 is the tested config). When set, every block's contents are fingerprinted at prefill time; restore queries find semantically related blocks before kernel dispatch. | +| `--kv-tier-semantic-threshold F` | Cosine similarity threshold for semantic restore. Default 0.65. Higher = stricter match; lower = more aggressive prefetch. | +| `--kv-tier-semantic-topk N` | Maximum blocks restored per query. Default 5. | + +### Compression flags + +| Flag | Description | +|---|---| +| `--cache-type-k TYPE` / `--cache-type-v TYPE` | KV element type. `f16` (default), `q8_0`, or `turbo4`. `turbo4` is the int4-with-scale path used by the army stack and is ~4x smaller than f16 with cosine similarity ≈ 0.99 against the unquantized baseline. | + +### Eviction policy flags + +| Flag | Description | +|---|---| +| `--kv-tier-eviction-policy N` | Eviction policy id. The default Hybrid policy weights recency + attention attention magnitude + frequency. Reserved for tuning experiments. | +| `--kv-tier-attention-threshold F` | Lower bound on attention weight before a token is considered evictable. Default 0.10. | +| `--kv-tier-warm-device N` | Override the device id used to host warm-tier staging buffers. Default -1 (auto). | + +### Deprecated spellings (still accepted, will warn) + +| Old | New | +|---|---| +| `--kv-semantic-index` | `--kv-tier-semantic-index` | +| `--kv-semantic-threshold` | `--kv-tier-semantic-threshold` | +| `--kv-semantic-topk` | `--kv-tier-semantic-topk` | + +--- + +## Sizing guide + +The tier percentages multiply against `total_ctx × per_token_K_plus_V` +to size each pool. Pick from these starting points and tune from +metrics: + +| GPU class | Suggested split | Rationale | +|---|---|---| +| 32 GB+ (R9700, 4090, A6000) | `25,75,0` | Hot pool generous; warm absorbs spillover; SSD only if you hit the warm ceiling. | +| 16 GB (6900XT, 4080) | `25,75,0` or `25,50,25` | Same shape; add cold if your context budget really exceeds VRAM+RAM. | +| 8 GB (1070, 3060, RX 480) | `25,25,50` | Hot is small; lean on warm + cold. SSD path required. | +| <8 GB | Don't run paged here for prod | Bookkeeping overhead starts to hurt. | + +Rules of thumb: +- **Hot %** is what fits in VRAM with model weights still loaded. Don't + exceed `(VRAM_total − model_weight_bytes − ~1 GB)`. +- **Warm %** wants host RAM headroom; on a 64 GB box budget no more + than 32 GB for warm to leave room for the model loader and OS. +- **Cold %** needs writable SSD space at `--kv-tier-ssd-path`. Cap with + `--kv-tier-cold-budget-mb` to bound wear. + +The boot log prints the resulting allocation so you can sanity-check: + +``` +llama_kv_cache_paged: allocated 16/64 attn layers × 8.5 MiB (K+V) = 136.0 MiB total + (n_blocks=512, block_size=16, n_kv_heads=4, head_dim=256, type_k=turbo4, type_v=turbo4) +llama_kv_cache_paged: warm tier enabled — 384 host blocks × 17.0 KiB/block per layer × 16 attn layers = 102.0 MiB host warm storage +llama_kv_cache_paged: cold tier enabled — 256 blocks × 17.0 KiB/block (K+V) × 16 attn layers = 68.0 MiB on /var/llama/cache/paged/instance-agent-1 +``` + +--- + +## Health metrics + +Run with `--metrics` to enable the Prometheus endpoint at +`http://HOST:PORT/metrics`. The paged tier adds these keys: + +| Counter | Meaning | +|---|---| +| `llamacpp:paged_evict_hot_to_warm_total` | Blocks moved hot→warm. Going up = warm tier is engaging (good when ctx exceeds hot, bad if you sized hot too small). | +| `llamacpp:paged_evict_warm_to_cold_total` | Blocks moved warm→cold. Going up = warm full, spilling to SSD. | +| `llamacpp:paged_evict_cold_to_drop_total` | Blocks dropped (no recovery). Going up = your cold pool is also full; data is being lost. **Tune sizing.** | +| `llamacpp:paged_restore_warm_to_hot_total` | Successful warm→hot fault-ins. Going up = workload is revisiting recently-evicted blocks. | +| `llamacpp:paged_restore_cold_to_hot_total` | Successful cold→hot fault-ins. SSD reads. | +| `llamacpp:paged_seq_preempt_total` / `paged_seq_restore_total` | Whole-sequence preemptions (MAD-120 admission control). | +| `llamacpp:paged_semantic_attempts_total` | Times the semantic restore path was invoked. | +| `llamacpp:paged_semantic_hits_total` | Attempts that restored ≥ 1 block. | +| `llamacpp:paged_semantic_blocks_restored_total` | Total blocks restored via semantic prefetch. Hit rate = `hits / attempts`. | +| Gauge | | +| `llamacpp:paged_blocks_capacity_gpu` / `_warm` / `_cold` | Pool sizes in blocks. | +| `llamacpp:paged_fingerprints` | Live BGE-small fingerprints currently held. | + +The `/slots` endpoint extension shows per-seq tier residency +(hot/warm/cold block counts and live fingerprint count). + +--- + +## Troubleshooting + +### Server returns "paged KV admission could not fit a N-token request" + +The submitted prompt is too large to coexist with the live workload on +the hot pool. Two ways out: +- Send a smaller prompt (chunk into multiple turns with `cache_prompt: true`). +- Wait for other slots to drain (decode-finish releases their blocks). +- Increase `-c` so more total blocks are available, **or** reduce + `--parallel` so each slot has a larger share. + +### `paged_evict_cold_to_drop_total` is climbing + +Your cold pool is undersized for the workload. Either widen the cold +slice (raise `COLD%` in `--kv-tiered`), raise +`--kv-tier-cold-budget-mb`, or chunk requests so less old context lives +in the tier. + +### `paged_semantic_attempts_total` is zero + +Semantic prefetch fingerprints are written **at prefill time**, one per +complete (block-aligned) chunk of the prompt. If your prompts are +shorter than `block_size` (16 tokens), no fingerprints get written. +Either send longer prompts or increase prompt batch sizes so blocks +fill. + +### `paged_semantic_hits_total / paged_semantic_attempts_total < 0.30` + +The default 30% hit-rate target assumes cross-context queries. +Conversational workloads where every turn references only the most +recent N tokens won't benefit much from semantic prefetch — the LRU +path already keeps those hot. Lower +`--kv-tier-semantic-threshold` (e.g. 0.55) for more aggressive matches, +or accept that this workload doesn't need semantic prefetch. + +### "instance lock /…/instance.lock held by another process" + +Two `llama-server` instances tried to use the same `--instance-id` + +`--kv-tier-ssd-path` combination. Pick a unique `--instance-id` per +instance. + +### Server aborts with `n_empty_consecutive > 3` + +Pre-MAD-141 server. Pull the latest `feat/MAD-126-army-goal` branch +(or a release built after `0d66d8aa3`). The fix replaces the upstream +safety abort with a graceful per-slot `send_error()`. + +### Cold-tier load reports "sidecar … missing or invalid" + +The prior shutdown didn't write the cold-tier index sidecar — typical +on a crash or SIGKILL. The cold pool starts empty; existing per-layer +files are ignored on this run. Use the explicit `/slots/save` endpoint +before shutdown if you want the cold tier to survive a restart. diff --git a/docs/memory-tier/adr/A1-hybrid-paged-primary.md b/docs/memory-tier/adr/A1-hybrid-paged-primary.md new file mode 100644 index 000000000000..cb5e807da6f2 --- /dev/null +++ b/docs/memory-tier/adr/A1-hybrid-paged-primary.md @@ -0,0 +1,68 @@ +# A1. Hybrid+paged is THE primary path + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-127 (legacy path removal) + +## Context + +Three KV-cache implementations existed at the start of the army goal: + +1. `llama_kv_cache_paged` — vLLM-style block-indexed paged attention + for hybrid models (attention + recurrent layers). Active. +2. `llama_kv_cache_tiered` — non-paged tiered cache for pure-attention + models. Largely a parallel codebase to (1). +3. `server_tiered_cache` — server-side wrapper that brokered tier + movement for (2). Yet another parallel codebase. + +Maintaining three implementations of "tiered KV cache" forked the +work: every new tier feature had to be threaded through three +implementations, every bug fix required deciding which paths it +applied to, and the test surface tripled. + +The model landscape is shifting decisively toward hybrid architectures +(Qwen3.5/3.6 family, Mamba-attention hybrids, etc.). User direction +2026-05-10 was unambiguous: "hybrid+paged is THE primary path; pure- +attention is least concern. Basically no one is building pure- +attention models anymore." + +## Decision + +All new tier features land on `llama_kv_cache_paged`. The legacy paths +are removed: + +- `llama_kv_cache_tiered` deleted. +- `server_tiered_cache` deleted. +- `mt::llama_memory_tiered` retained but reduced to a thin shim. Its + surviving responsibilities are documented in + [A8](A8-bge-small-on-wrapper.md): bge-small embed model ownership + and recurrent-state backup on hybrid models. + +Pure-attention models still work — they route through +`llama_kv_cache_paged` directly without the hybrid wrapper. They just +don't get any special-cased treatment. + +## Consequences + +**Positive**: +- One codebase for tier management. Bug fixes apply uniformly. +- Test matrix shrinks. New features need one implementation, not + three. +- The `BlockPool` / `BlockTable` / `BlockSemanticIndex` primitives + serve every code path that needs paged-style block management. +- Eliminates the "which cache am I actually getting?" confusion for + operators and developers. + +**Negative**: +- Pure-attention models pay the (small) per-call indirection cost of + the paged dispatch even when they wouldn't strictly need it. +- The thin `mt::llama_memory_tiered` wrapper is a slight wart — + carrying it as a separate type for two narrow responsibilities + (bge-small + recurrent backup). A cleaner long-term move is to fold + those responsibilities into the model layer directly, but that's a + bigger refactor than the army goal needs. + +**Neutral**: +- Anyone holding stale references to `llama_kv_cache_tiered` or + `server_tiered_cache` will hit a clean compile error rather than + silently picking the wrong path. diff --git a/docs/memory-tier/adr/A10-real-seq-cp-cow.md b/docs/memory-tier/adr/A10-real-seq-cp-cow.md new file mode 100644 index 000000000000..9bb9aadd50a8 --- /dev/null +++ b/docs/memory-tier/adr/A10-real-seq-cp-cow.md @@ -0,0 +1,70 @@ +# A10. seq_cp CoW required for branching agent workflows + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-128 (paged-compat: real seq_cp with refcount-based CoW) + +## Context + +`llama_memory_i::seq_cp(src, dst, p0, p1)` copies positions +`[p0, p1)` from `src` to `dst`. Branching agent workflows depend on it +heavily: a planner clones the current conversation N times to explore +N hypothetical continuations, then either commits the best branch or +discards them all. + +The original `llama_kv_cache_paged::seq_cp` was a no-op that logged a +warning ("seq_cp not yet supported on paged cache"). Branching +workflows just didn't work — `dst` ended up empty and the model +re-prefilled from scratch on every branch, losing all the +conversation context. + +A literal byte-copy implementation would have worked but been wildly +expensive: an N-way branch with K-token shared context costs `N × K +× per_token_K_plus_V` extra bytes copied per fork. For a 128k-token +context with 4 branches, that's hundreds of MB of host↔device +copies per fork. + +## Decision + +Implement copy-on-write via the existing `BlockPool` refcount +machinery: + +- `seq_cp(src, dst, p0, p1)` walks `src`'s logical blocks covering + `[p0, p1)`. For each, the cache: + 1. Calls `pool_.bump_ref(physical_block_id)` — refcount++. + 2. Calls `table_.append_block(dst, physical_block_id)` — `dst`'s + logical sequence now points at the same physical block. +- `pool_.free_block` decrements refcount; only returns the block to + the free stack when refcount drops to 0. So a shared block stays + alive as long as ANY seq references it. + +When `dst` later writes new tokens to a shared block (the cell +positions diverge from `src`'s view), the cache detects shared +status via `pool_.refcount(physical) > 1` and allocates a fresh +block, memcpys the existing data into it, swaps the table entry, and +decrements the old block's refcount. That's the copy-on-write step. + +## Consequences + +**Positive**: +- Branching is essentially free at fork time — one block-table + copy + N refcount bumps. The N×K bytes get copied only when (and + if) a branch actually mutates that block. +- Multi-path planning workflows become viable on paged caches. +- The refcount machinery was already in `BlockPool` for other + reasons (the eviction path needed to know which blocks were live + vs. truly free); adding `bump_ref` and refcount-aware + `free_block` was minimal incremental work. + +**Negative**: +- Eviction has to skip refcount > 1 blocks during victim selection, + otherwise it'd evict a block another seq depends on. Current + behavior: if the LRU candidate is shared, skip and try the next. + Could thrash if every hot block is shared — pathological but + recoverable (operator sees `paged_evict_hot_to_warm_total` rate + spike). + +**Neutral**: +- The CoW write trigger lives inside `llama_kv_cache_paged`'s prefill + path. Anyone changing how new blocks get written needs to preserve + the `if (refcount > 1) clone-then-write` check. diff --git a/docs/memory-tier/adr/A2-fingerprint-at-prefill.md b/docs/memory-tier/adr/A2-fingerprint-at-prefill.md new file mode 100644 index 000000000000..88bb3d23f6b1 --- /dev/null +++ b/docs/memory-tier/adr/A2-fingerprint-at-prefill.md @@ -0,0 +1,81 @@ +# A2. Semantic fingerprint write trigger = prefill time + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-129 (server-side prefill trigger) +**Supersedes**: MAD-122 (eviction-time trigger; wrong layer) + +## Context + +The semantic prefetch system needs a per-block BGE-small embedding +("fingerprint") so that future queries can score similarity against +old blocks and pull the right ones back from warm/cold to hot. + +The original MAD-122 design wrote fingerprints at **eviction time** — +when a block was about to be moved out of hot, the embedding was +computed from the block's text and stored. + +This worked for `mt::llama_memory_tiered` (the legacy non-paged +tiered cache) where eviction was visible at the wrapper level. It +**did not** work for hybrid+paged because: + +1. Eviction in `llama_kv_cache_paged` is internal — the server doesn't + see eviction events. They happen inside `ensure_blocks_for` / + `evict_lru_to_warm` / MAD-120 admission. +2. The server-side proactive backup gating doesn't fire for + hybrid+paged either — its cap arithmetic uses full ctx since + `physical_attn_cells()` returns 0. +3. Even if both above were fixed, eviction-time fingerprinting writes + *after* the data is already mid-evict — racy and order-dependent. + +A smoke test of MAD-122 against the hybrid+paged path showed **zero +semantic activity** under a workload that should have produced many +fingerprints. The trigger never fired. + +## Decision + +The server fingerprints every complete (block-aligned) chunk of a new +seq's prompt at **prefill submission**, before block-table assignment. + +Specifically: after the server's `update_slots` builds the prefill +batch and the block table is populated, the server walks the new +seq's complete logical blocks and: + +1. Decodes the original tokens for that block via + `slot.prompt.tokens[]` + `common_token_to_piece`. +2. Calls `mt::llama_memory_tiered::embed_text(text)` for an + L2-normalized embedding. +3. Calls + `llama_kv_cache_paged::record_paged_block_fingerprint(seq, lblock, embedding, tier=Hot)`. + +Skipped when the block already has a fingerprint (the +`has_paged_fingerprint(seq, lblock)` short-circuit handles re-prefill +of a prior turn's accumulated context). + +CPU cost is ~5ms × n_complete_blocks per prefill, off the GPU +critical path. Fingerprint lifecycle is bound to block lifecycle — +they get dropped on whole-seq wipe; per-block removal happens +automatically when a block's table entry is cleared. + +## Consequences + +**Positive**: +- Trigger fires deterministically on every prefill, observable in + `paged_semantic_attempts_total`. +- Lives at the server layer where the original token text is still + available for decoding. The cache doesn't need to persist text. +- Decoupled from the cache's internal eviction policy — the cache can + evict however it wants without losing fingerprint coverage. + +**Negative**: +- Adds ~5ms × n_blocks per prefill of CPU work. For a 2048-token + prompt at block_size=16, that's ~640ms total. Off the GPU critical + path but visible in time-to-first-token. +- Sub-block-aligned prompts (< block_size new tokens) don't get + fingerprinted at all. Only matters for very short prompts. + +**Neutral**: +- The cache exposes `record_paged_block_fingerprint` and + `has_paged_fingerprint` as part of its public API. The wrapper / + embed path stays in `mt::llama_memory_tiered` + (see [A8](A8-bge-small-on-wrapper.md)). diff --git a/docs/memory-tier/adr/A3-semantic-prefetch-only.md b/docs/memory-tier/adr/A3-semantic-prefetch-only.md new file mode 100644 index 000000000000..ccd427c2abc0 --- /dev/null +++ b/docs/memory-tier/adr/A3-semantic-prefetch-only.md @@ -0,0 +1,69 @@ +# A3. Semantic drives prefetch only, not eviction + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-129 (semantic prefetch on `llama_kv_cache_paged`) + +## Context + +A natural extension of having a per-block embedding (BGE-small +fingerprint, see [A2](A2-fingerprint-at-prefill.md)) is to use it for +**eviction** decisions, not just prefetch — pick the block to evict +that's least semantically related to the current query. + +This was considered and rejected. + +## Decision + +bge-small drives **prefetch only**. Eviction stays on the existing +hybrid policy in `src/memory-tier/mt-eviction.{h,cpp}`, which weights +recency + attention magnitude + frequency without any semantic input. + +## Reasoning + +1. **Training-task mismatch.** BGE-small is trained for retrieval — + "given a query, find documents that answer it." That's a cousin + problem to "predict which cached blocks an LLM is about to attend + to," but only a cousin. There's no published evidence that + retrieval embeddings correlate well with attention-cache + predictiveness, and there's no ground-truth dataset to tune on. + +2. **Hot-path latency.** Every eviction decision would need an embed + call, or a lookup against pre-computed embeddings plus a cosine + computation per candidate. Eviction needs to be fast (<1ms) + because it sits in the prefill / decode critical path. Adding + embed lookups inflates the decision cost meaningfully. + +3. **Doesn't fix the structural problem.** The thing that actually + threatens correctness under multi-seq load is hot-pool + fragmentation when N agents collectively want more hot blocks + than exist (MAD-120's admission control problem). Smarter + eviction picks doesn't change that — admission control does. + +4. **Measurable risk if wrong.** If semantic-driven eviction + picks badly, decode quality degrades silently (the model attends + to misleading old context with wrong weights). Whereas + recency/frequency-driven eviction has known failure modes that + produce predictable degradation. + +## Consequences + +**Positive**: +- Eviction decisions stay simple, fast, and analyzable. The hybrid + policy in `mt-eviction.cpp` is ~200 lines of boring math. +- The semantic path is purely **additive** — turning it off (no + `--kv-tier-semantic-index`) leaves the cache fully functional. + +**Negative**: +- Some workloads where semantic-driven eviction would have been a + win don't get that win. Specifically: long-running agent + workflows where the "right" thing to evict is defined by what the + CURRENT query semantically isn't asking about. Re-evaluate this + decision if such workloads become important. + +**Neutral**: +- Two separate concerns now own two separate code paths: eviction in + `mt-eviction`, prefetch in `BlockSemanticIndex` + + `llama_kv_cache_paged::restore_semantic_paged`. Easier to reason + about, but anyone wanting to introduce semantic-driven eviction + later has a clear refactor surface. diff --git a/docs/memory-tier/adr/A4-single-threading.md b/docs/memory-tier/adr/A4-single-threading.md new file mode 100644 index 000000000000..9ff5ab1c6c15 --- /dev/null +++ b/docs/memory-tier/adr/A4-single-threading.md @@ -0,0 +1,69 @@ +# A4. Single-threading contract is explicit + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-132 (concurrency hardening) + +## Context + +`llama_kv_cache_paged` and its `BlockPool` / `BlockTable` / +`BlockSemanticIndex` companions hold mutable state. None of them have +internal synchronization (mutexes, atomics on counters, etc.). The +question was whether to add internal locking, document a single- +threading contract, or design for explicit concurrency. + +vLLM, the architectural reference, runs its scheduler single- +threaded for exactly the same reasons we're considering single- +threading: cache state machines are easier to reason about +sequentially, locking adds latency to the hot path, and the workload +(one batch per scheduler tick) doesn't naturally benefit from +parallelism inside the cache itself. + +## Decision + +Single-threading is the **explicit contract**. Document it in the +class headers; add debug-build assertions to catch violations. + +Concretely: + +- The server's main `update_slots` loop is the single mutator. HTTP + worker threads communicate with it via `server_queue`'s task + channel and never touch the cache directly. +- Tier-counter `_total` accessors return `uint64_t` from the cache. + These are read by the `/metrics` HTTP handler from a different + thread — this is the **only** sanctioned cross-thread cache + access. It's safe-ish on x86_64 and aarch64 for monotonic 64-bit + counters; aligned-uint64 atomic loads are single instructions + on those ISAs, so torn reads aren't a concern. +- Debug builds include `check_thread_id_()` in + `llama_kv_cache_paged` that traps if anyone other than the + registered main thread mutates the cache. +- Any future async work (semantic prefetch on a worker, batched + bge-small embed off the critical path) is gated on a real + concurrency design — not "by accident." + +## Consequences + +**Positive**: +- Reasoning about cache state is straightforward — no interleavings, + no data races, no lost-update bugs. +- No locking overhead in the hot path. Eviction, allocation, and + fingerprint scoring are all single-instruction-stream operations. +- Test surface stays manageable; we don't need to write race tests + for code paths that can't race by contract. + +**Negative**: +- Constrains future work. Any feature that wants to run on a worker + thread must either move the work outside the cache (compute on the + worker, hand the result back to the main thread to apply), or + carry a real concurrency-design proposal to add minimal locking. +- The cross-thread metrics read is technically a contract violation + — it works in practice but it's load-bearing on aligned-64-bit + atomicity. If the codebase ever ports to a 32-bit platform, the + counters need to be split or made atomic. + +**Neutral**: +- The contract is unenforceable at runtime in release builds. Anyone + introducing a multi-threaded mutation path will have to cross-check + the assertion in debug builds — relying on developer discipline to + do so. diff --git a/docs/memory-tier/adr/A5-persistence-explicit.md b/docs/memory-tier/adr/A5-persistence-explicit.md new file mode 100644 index 000000000000..63fcb3b972ea --- /dev/null +++ b/docs/memory-tier/adr/A5-persistence-explicit.md @@ -0,0 +1,78 @@ +# A5. Persistence is explicit save/restore, no implicit crash recovery + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-130 (state_write/state_read on paged + cold resume + fingerprint save/load) + +## Context + +The cache holds substantial state — block table, hot/warm/cold tier +contents, semantic fingerprints. Two persistence models were +candidates: + +1. **Implicit crash recovery**: write enough state continuously to + disk that a sudden process death leaves the cache in a recoverable + state on restart. Requires WAL-style logging or transactional + block writes. + +2. **Explicit save/restore**: serialize state on demand via the + server's `/slots/save` endpoint; restore via `/slots/restore` on + the next boot. A clean shutdown saves; a crash leaves nothing + recoverable. + +The tradeoffs run opposite directions on every axis: + +| Axis | Implicit | Explicit | +|---|---|---| +| Operator complexity | Low (just works) | Medium (call /slots/save) | +| Implementation complexity | High (WAL semantics, partial-write recovery, format versioning, double-write throughput cost) | Low (one serializer, one deserializer) | +| Crash semantics | "Best effort recovery — please report bugs" | "Crash = blank slate; clean shutdown = restored" | +| Performance | Continuous disk pressure | Zero disk pressure between saves | + +## Decision + +Explicit save/restore via the server's existing `/slots/save` and +`/slots/restore` endpoints. Crash recovery is **not** automatic. + +State written by `state_write()`: +- Block table per seq (logical→physical mapping; PAGS v1 magic). +- Hot-tier K/V tensors (or skip + reprefill — caller-configurable + via the `flags` argument). +- Warm-tier K/V buffers. +- Cold-tier index sidecar (CIDX v1 magic = `0x58444943`). +- BlockSemanticIndex fingerprints (PSFI v1 magic = `0x49465350`). + +Format magic numbers and version bytes let future format changes be +detected and rejected (rather than silently corrupting state). + +The cold-tier sidecar is the recovery hinge: per-layer K/V files +contain the actual block data, but without the sidecar's `(seq, +lblock) → file_offset` mapping the cache can't read them back. Hard +crashes that miss the sidecar write produce orphan files which +`scripts/army/cleanup-cold.sh` removes. + +## Consequences + +**Positive**: +- Implementation is straightforward. One serializer per state + category; one set of unit + integration tests. +- Zero steady-state disk overhead between saves. Operators control + when (and whether) to save. +- Format versioning gives a clean upgrade story. Adding a new field + bumps the version byte; old loaders refuse cleanly. + +**Negative**: +- Crash recovery is "start over." For long-running servers with + large cold tiers, that's a meaningful cost — the cold pool gets + rebuilt only as the workload re-prefills the same content. +- Operators must remember to call `/slots/save` (or wire it into + their service's `ExecStop`). A box that gets rebooted without that + call loses everything cold. + +**Neutral**: +- The cold-tier per-layer files survive on disk after a crash, but + without the sidecar they're unreadable. They're not corrupting + anything; they're just dead weight until cleaned up. +- A future "crash-recoverable" mode is not precluded — it would be a + separate feature on top of this one (probably a write-ahead log of + block-table mutations). diff --git a/docs/memory-tier/adr/A6-multi-instance.md b/docs/memory-tier/adr/A6-multi-instance.md new file mode 100644 index 000000000000..2c672160b3c2 --- /dev/null +++ b/docs/memory-tier/adr/A6-multi-instance.md @@ -0,0 +1,92 @@ +# A6. Multi-instance isolation = per-instance cold subdir + lockfile + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-131 (per-instance cold subdir + lockfile + SSD quota + per-machine boot scripts) + +## Context + +The army goal is four `llama-server` instances per box (one per GPU) +in the steady state. They share a host filesystem and may share a +single SSD configured via `--kv-tier-ssd-path`. Without isolation, +they would: + +1. Write to the same per-layer K/V files (`L0.k.bin`, `L0.v.bin`, + etc.). Different cache geometries between instances would corrupt + each other. +2. Race on the cold-tier sidecar. +3. Re-open each other's files on cold-resume, treating another + instance's persisted state as their own. + +Three isolation models were possible: + +1. **Disjoint paths**: each instance gets its own + `--kv-tier-ssd-path`. Operationally tedious — operators have to + provision N directories and remember which is which. + +2. **Shared path + namespacing**: one root path, every instance + writes under a per-instance subdirectory. Operators provision one + path; the server handles namespacing. + +3. **Single shared cache across instances**: one cache backing all + instances. Architecturally cleanest in some sense but requires + cross-process coordination (shared memory, IPC) that the + single-threading contract ([A4](A4-single-threading.md)) rules + out. + +## Decision + +Option 2: shared `--kv-tier-ssd-path` with per-instance namespacing. + +Each instance writes cold-tier files under +`${ssd_path}/paged/instance-${INSTANCE_ID}/`: + +``` +${ssd_path}/paged/ +├── instance-main-r9700/ +│ ├── L0.k.bin +│ ├── L0.v.bin +│ ├── ... +│ ├── instance.lock # flock'd while the process is alive +│ └── index.bin # CIDX v1 sidecar +├── instance-main-6900xt/ +│ └── ... +``` + +`--instance-id` defaults to the process PID but is typically set +explicitly by boot scripts so it's stable across restarts (which +matters for cold-resume — see [A5](A5-persistence-explicit.md)). + +A `flock`-based lockfile on `instance.lock` prevents accidental +double-start: a second process trying to open the same instance subdir +fails fast with a clear error rather than silently corrupting on-disk +state. + +Optional `--kv-tier-cold-budget-mb` caps each instance's cold pool +size in MiB (K+V across all attn layers) so a runaway instance can't +fill the shared disk and starve siblings. + +## Consequences + +**Positive**: +- One operator-provisioned path serves the whole army on a box. +- Per-instance isolation prevents cross-instance corruption by + construction. +- The lockfile gives clean error messages on configuration mistakes + (e.g. operator forgot to change `--instance-id` after copy-pasting + a boot command). +- Per-instance cold-resume works without any cross-process + coordination — each instance reads its own subdir. + +**Negative**: +- An operator who really wants single-shared-cache semantics (some + exotic deployment) doesn't get it. Acceptable because nobody's + asking. +- A stale lockfile from a hard crash blocks restart until cleaned. + The runbook includes the recovery procedure (verify with `lsof` + that no process holds it; then `rm`). + +**Neutral**: +- Disk-space accounting per instance becomes the operator's + responsibility. The server reports its own usage in metrics; total + disk usage across all instances requires `du` or equivalent. diff --git a/docs/memory-tier/adr/A7-paged-default-on.md b/docs/memory-tier/adr/A7-paged-default-on.md new file mode 100644 index 000000000000..9bc2a5c04ebf --- /dev/null +++ b/docs/memory-tier/adr/A7-paged-default-on.md @@ -0,0 +1,62 @@ +# A7. Auto-default `--kv-tier-paged-blocks` for hybrid models + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-134 (config ergonomics) + +## Context + +A hybrid model running with `--kv-tiered` and without +`--kv-tier-paged-blocks` falls into the legacy non-paged tiered +codepath. After [A1](A1-hybrid-paged-primary.md), that path is +neither maintained nor recommended for hybrid models. Most operators +who set `--kv-tiered` on a hybrid model **mean** to get paged + tiered +— forgetting `--kv-tier-paged-blocks` produces a worse cache +silently. + +Three options for the default behavior: + +1. **Keep current default off**. Operators who want paged must + remember the flag. Bad ergonomics; quiet performance regressions. + +2. **Default on for everyone (hybrid AND pure-attention)**. Risks + regressing pure-attention configurations that worked fine on the + non-paged path. + +3. **Default on for hybrid, leave pure-attention as-is**. Targets + the actual user — hybrid is what people actually run with + tiering. + +## Decision + +Option 3. When the model is hybrid AND `--kv-tiered` is set AND +`--kv-tier-paged-blocks` was NOT explicitly disabled, default-on. +Operators can opt out via `--no-kv-tier-paged-blocks` for +backwards-compat. + +The implementation uses a tristate-via-explicit-bool pattern in +`common_params`: a `bool kv_tier_paged_blocks_explicit` records +whether the user actually mentioned the flag. The default-on logic +runs in `common_context_params_to_llama` and only fires when the user +hasn't expressed an opinion either way. + +## Consequences + +**Positive**: +- Operator running the army-goal config doesn't need to remember the + paged flag. `--kv-tiered 25,75,0` on a hybrid model just works. +- Pure-attention deployments are untouched — no regression risk for + existing configurations that don't expect paged. +- The opt-out path (`--no-kv-tier-paged-blocks`) is still there for + edge cases or A/B comparisons. + +**Negative**: +- Adds one more piece of "automatic" behavior that operators have to + read about to fully understand. Mitigated by the boot-time log + line that prints the resolved configuration (whether paged is on + and why). + +**Neutral**: +- Anyone who explicitly set `--no-kv-tier-paged-blocks` keeps that + intent across all hybrid configurations. The default-on only + applies when the user hasn't expressed an opinion. diff --git a/docs/memory-tier/adr/A8-bge-small-on-wrapper.md b/docs/memory-tier/adr/A8-bge-small-on-wrapper.md new file mode 100644 index 000000000000..caacce43341c --- /dev/null +++ b/docs/memory-tier/adr/A8-bge-small-on-wrapper.md @@ -0,0 +1,68 @@ +# A8. bge-small ownership stays on `mt::llama_memory_tiered` + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-127 (wrapper thinning) + MAD-129 (server-side bridge) + +## Context + +After [A1](A1-hybrid-paged-primary.md), `mt::llama_memory_tiered` is a +thin shim — most of its old responsibilities moved into +`llama_kv_cache_paged`. Two responsibilities remained candidates for +the wrapper vs. the cache: + +1. **bge-small embed model ownership.** Loading the gguf, holding + the `EmbeddingModel` instance, exposing `embed_text(string) → + vector` to whoever needs to compute a fingerprint. + +2. **Recurrent-state backup on hybrid models.** Backing up the + per-seq recurrent state into host RAM on `seq_rm` so it can be + restored later (the recurrent half's `clear()` is destructive). + +Putting these on `llama_kv_cache_paged` would have made the cache the +single owner of all tier-related state. But it would also have +spread embed-model concerns into the cache header and forced the +cache to know about recurrent state — neither is its job. + +## Decision + +Both responsibilities stay on `mt::llama_memory_tiered`: + +- `mt::llama_memory_tiered::embed_text(text)` is the embed entry point. + Lazy-loads the gguf on first call (or warmup at construction — + see MAD-134). Returns L2-normalized 384-dim vectors. +- `mt::llama_memory_tiered::backup_seq_rm_recurrent` / + `restore_recurrent_from_warm` are the recurrent backup pair. + +The paged cache holds the resulting fingerprints in its own +`BlockSemanticIndex`. The server bridges the two via the existing +`mt_get_paged_cache(llama_get_memory(ctx))` accessor: it asks the +wrapper to compute an embedding, then hands the embedding to the +cache's `record_paged_block_fingerprint`. + +## Consequences + +**Positive**: +- The cache stays focused on its core responsibility (block + management, eviction, K/V tier movement). It doesn't know what an + embedding model is. +- The bge-small instance is shared across configurations — pure- + paged, hybrid+paged, and the legacy thin-wrapper paths can all + use it without each loading their own copy. +- Recurrent-state backup is a hybrid-specific concern; keeping it on + the wrapper means a pure-attention config with paged routing + doesn't carry the recurrent-backup code paths at all. + +**Negative**: +- Adds an indirection: server → wrapper.embed_text → cache.record. + Extra function call per fingerprint, negligible cost. +- The wrapper survives in a "thin but not gone" state, which is a + slight wart if you wanted A1 to fully kill it. Acceptable + tradeoff given the cleanliness benefit of keeping these two + concerns separate from the cache. + +**Neutral**: +- Anyone wanting to swap bge-small for a different embedding model + changes one place (the wrapper's `embed_text`). The cache and the + fingerprint store don't care which model produced the embedding, + as long as it's L2-normalized 384-dim. diff --git a/docs/memory-tier/adr/A9-real-partial-seq-rm.md b/docs/memory-tier/adr/A9-real-partial-seq-rm.md new file mode 100644 index 000000000000..fb40bef300f2 --- /dev/null +++ b/docs/memory-tier/adr/A9-real-partial-seq-rm.md @@ -0,0 +1,76 @@ +# A9. Speculative decoding requires real partial seq_rm + +**Status**: Accepted (2026-05-10) +**Decided in**: [MAD-126](https://mad-lab-ai.atlassian.net/browse/MAD-126) +**Implemented in**: MAD-128 (paged-compat: real seq_rm middle-range) + +## Context + +`llama_memory_i::seq_rm(seq, p0, p1)` removes positions `[p0, p1)` +from the seq's KV cache. The interface allows arbitrary middle-range +wipes; speculative decoding relies on this to roll back rejected +draft tokens (the verifier rejects token N, so positions `[N, end)` +get wiped to allow re-decode from a different sample). + +The original `llama_kv_cache_paged::seq_rm` implementation treated +**all** non-tail wipes as if they were tail-truncates (wiping +`[p0, end)` regardless of `p1`). This was a silent degradation: the +cache reported success but actually removed more data than asked. +With speculative decoding, that meant losing positions that the +verifier had **accepted** alongside the rejected ones — quietly +producing wrong outputs. + +Three choices: + +1. **Reject middle-range seq_rm** at the cache boundary, return false. + Forces speculative decoding to disable itself when paged is active. +2. **Implement real partial seq_rm**. Block-aligned wipes free whole + physical blocks; sub-block partial wipes leave holes in the + logical sequence (the kernel handles `kInvalidBlockId` by emitting + `-INFINITY` logits → zero attention weight; equivalent to the + token not existing). +3. **Convert to whole-seq wipe**. Simpler than (2) but breaks + conversation continuity — every speculative rollback would + blow away the entire prefill. + +## Decision + +Option 2: implement real partial seq_rm. The cache: + +- Block-aligned middle wipes free the wholly-covered physical block(s) + and reset the logical→physical table entry to `kInvalidBlockId`. +- Sub-block partial wipes (the rare case) leave the block's physical + storage in place but log a warning. The kernel still emits + `-INFINITY` for cells outside the live range thanks to per-block + context-len tracking. +- `seq_pos_max` updates to reflect the new logical extent. + +The `kInvalidBlockId` sentinel was already in the design from +[A1's](A1-hybrid-paged-primary.md) BlockTable scheme; partial seq_rm +just leverages it. + +## Consequences + +**Positive**: +- Speculative decoding works against paged caches without + degradation. +- Same plumbing supports `mt_record_paged_block_fingerprint`'s + re-fingerprint path (overwriting an existing fingerprint when the + block's content gets edited — the partial seq_rm wipes the cells, + the block_table entry stays, the fingerprint gets overwritten on + next prefill). +- The kernel's `kInvalidBlockId` handling is now load-bearing in + multiple places, which is fine because it was always part of the + paged design. + +**Negative**: +- Sub-block partial wipes (`p0 % block_size != 0` and `p1 % + block_size != 0` for the first/last block) currently log a warning + and leave the partially-wiped block's data in place. The kernel + copes via context-len, but it's a code-path that was harder to + test than block-aligned wipes. Keep an eye on it. + +**Neutral**: +- The implementation lives entirely in + `llama_kv_cache_paged::seq_rm`. The wrapper + (`mt::llama_memory_tiered::seq_rm`) just delegates. diff --git a/docs/memory-tier/adr/README.md b/docs/memory-tier/adr/README.md new file mode 100644 index 000000000000..18b0d6e95321 --- /dev/null +++ b/docs/memory-tier/adr/README.md @@ -0,0 +1,19 @@ +# Architecture Decision Records — tiered KV cache + +Each file documents one of the resolved design decisions from +Epic MAD-126. Format follows the standard ADR template (Context, +Decision, Consequences). The decisions are referenced by their +A-number throughout [`../ARCHITECTURE.md`](../ARCHITECTURE.md). + +| # | Title | +|---|---| +| [A1](A1-hybrid-paged-primary.md) | Hybrid+paged is THE primary path | +| [A2](A2-fingerprint-at-prefill.md) | Semantic fingerprint write trigger = prefill time, not eviction time | +| [A3](A3-semantic-prefetch-only.md) | Semantic drives prefetch only, not eviction | +| [A4](A4-single-threading.md) | Single-threading contract is explicit | +| [A5](A5-persistence-explicit.md) | Persistence is explicit save/restore, no implicit crash recovery | +| [A6](A6-multi-instance.md) | Multi-instance isolation = per-instance cold subdir + lockfile | +| [A7](A7-paged-default-on.md) | Auto-default `--kv-tier-paged-blocks` for hybrid models | +| [A8](A8-bge-small-on-wrapper.md) | bge-small ownership stays on `mt::llama_memory_tiered` | +| [A9](A9-real-partial-seq-rm.md) | Speculative decoding requires real partial seq_rm | +| [A10](A10-real-seq-cp-cow.md) | seq_cp CoW required for branching agent workflows | diff --git a/include/llama.h b/include/llama.h index c77704f49f06..8b79efbe3ae1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -405,6 +405,9 @@ extern "C" { int32_t kv_tier_semantic_topk; bool kv_tier_paged_blocks; // Phase 2a opt-in (off => current path) int32_t kv_tier_paged_block_size; // tokens/block (0 => default 16) + bool kv_tier_cold_resume; // MAD-130: skip O_TRUNC on cold-tier files; load index sidecar + const char * kv_tier_instance_id; // MAD-131: per-instance ID for cold subdir + lockfile (nullptr => pid) + int32_t kv_tier_cold_budget_mb; // MAD-131: cap cold pool to N MiB (0 => no limit) }; struct llama_model_tensor_override { diff --git a/scripts/army/README.md b/scripts/army/README.md new file mode 100644 index 000000000000..cc6fa8bbb8ad --- /dev/null +++ b/scripts/army/README.md @@ -0,0 +1,39 @@ +# Army boot scripts (MAD-131) + +Templates for spinning up the agent army across the per-machine targets: + +- `main.sh` — main box (R9700 + 6900XT). Two instances. +- `mad-lab.sh` — mad-lab box (1070 + RX 480 in ROCm 6.3 docker). Two instances. +- `cleanup-cold.sh` — walks the cold-tier root and removes orphaned `instance-*` subdirs. +- `army.service.example` — systemd unit template. + +These are templates. **Customize before use** — paths, model files, ports, and +GPU device IDs all need to match your local setup. + +## What "instance" means here + +Each `llama-server` process gets a unique `--instance-id` so multiple +processes can share one `--kv-tier-ssd-path` without colliding on +cold-tier files. The cache writes to: + +``` +${ssd_path}/paged/instance-${INSTANCE_ID}/ + L0.k.bin + L0.v.bin + ... + .lock (flock for double-start refusal) + .pid (informational; current pid) + index.bin (cold-index sidecar, written on /slots/save) +``` + +If two processes try to start with the same `--instance-id`, the second +will refuse with a clear error message naming the holder pid. + +## Operator runbook (short version) + +- Stop everything cleanly: `systemctl stop army@`. The cache + destructor releases the lock, so the next boot can reuse the same ID. +- After an unclean shutdown / crash: the lock is auto-released when the + process exits (kernel-level flock). The `.pid` file may be stale; + it's overwritten on next start. +- To wipe cold tier and start fresh: `./cleanup-cold.sh /path/to/ssd`. diff --git a/scripts/army/army.service.example b/scripts/army/army.service.example new file mode 100644 index 000000000000..d1ae16f91f61 --- /dev/null +++ b/scripts/army/army.service.example @@ -0,0 +1,33 @@ +# Example systemd unit template for army instances. +# Copy to /etc/systemd/system/army.service (or per-instance variants), +# customize Environment + ExecStart, then: +# +# systemctl daemon-reload +# systemctl enable army.service +# systemctl start army.service +# +# Logs: journalctl -u army.service -f + +[Unit] +Description=llama.cpp army instances +After=network-online.target +Wants=network-online.target + +[Service] +Type=simple +User=kmbandy +Group=kmbandy +Environment="LLAMA_BIN=/home/kmbandy/GitHub/llama.cpp/build-hip/bin/llama-server" +Environment="MODELS_DIR=/home/kmbandy/models" +Environment="SSD_PATH=/var/lib/army/ssd" +Environment="LOG_DIR=/var/log/army" +ExecStart=/home/kmbandy/GitHub/llama.cpp/scripts/army/main.sh +# 30s for graceful shutdown — gives the cache time to release its +# lockfile + (when wired) flush state via /slots/save. +TimeoutStopSec=30 +KillMode=mixed +Restart=on-failure +RestartSec=10 + +[Install] +WantedBy=multi-user.target diff --git a/scripts/army/cleanup-cold.sh b/scripts/army/cleanup-cold.sh new file mode 100755 index 000000000000..5b8bb08ce3e3 --- /dev/null +++ b/scripts/army/cleanup-cold.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# MAD-131: walk ${ssd_path}/paged and remove any instance-* subdir whose +# lockfile is no longer held (i.e. no live owner). Safe to run while +# other instances are alive — flock prevents touching live ones. + +set -euo pipefail + +SSD_PATH="${1:-${SSD_PATH:-/var/lib/army/ssd}}" +PAGED_DIR="$SSD_PATH/paged" + +if [[ ! -d "$PAGED_DIR" ]]; then + echo "cleanup-cold: $PAGED_DIR doesn't exist; nothing to do." + exit 0 +fi + +removed=0 +kept=0 +for dir in "$PAGED_DIR"/instance-*; do + [[ -d "$dir" ]] || continue + lock="$dir/.lock" + if [[ -f "$lock" ]]; then + # Try a non-blocking exclusive lock. If someone holds it (live + # llama-server with this instance ID), skip. flock returns 0 + # on success; we then immediately release. + if flock -nx "$lock" -c true 2>/dev/null; then + # We got the lock → no live holder. Safe to remove. + echo "cleanup-cold: removing orphaned $dir" + rm -rf "$dir" + ((removed++)) + else + echo "cleanup-cold: keeping $dir (live holder)" + ((kept++)) + fi + else + # No lockfile → safe to remove + echo "cleanup-cold: removing $dir (no lockfile)" + rm -rf "$dir" + ((removed++)) + fi +done + +echo "cleanup-cold: removed=$removed kept=$kept" diff --git a/scripts/army/mad-lab.sh b/scripts/army/mad-lab.sh new file mode 100755 index 000000000000..3b15f8503bc4 --- /dev/null +++ b/scripts/army/mad-lab.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +# Army boot script: mad-lab box (1070 Pascal + RX 480 Polaris in ROCm 6.3 docker). +# Customize paths, ports, and docker image name. + +set -euo pipefail + +LLAMA_BIN="${LLAMA_BIN:-$HOME/GitHub/llama.cpp/build-cuda/bin/llama-server}" +LLAMA_BIN_HIP="${LLAMA_BIN_HIP:-$HOME/GitHub/llama.cpp/build-hip-gfx803/bin/llama-server}" +MODELS_DIR="${MODELS_DIR:-$HOME/models}" +SSD_PATH="${SSD_PATH:-/var/lib/army/ssd}" +LOG_DIR="${LOG_DIR:-/var/log/army}" +BGE_SMALL="${BGE_SMALL:-$MODELS_DIR/bge-small-en-v1.5-q8_0.gguf}" +DOCKER_IMAGE="${DOCKER_IMAGE:-rocm-6.3-gfx803:latest}" + +mkdir -p "$SSD_PATH" "$LOG_DIR" + +# ── Instance 1: 1070 (Pascal CUDA, 8 GB) — Qwen3.5-9B-Omnicoder ─────── +INSTANCE=mad-lab-1070 +PORT=11437 +exec_log="$LOG_DIR/${INSTANCE}.log" + +"$LLAMA_BIN" \ + -m "$MODELS_DIR/Qwen3.5-9B-Omnicoder-Q5_K_S.gguf" \ + --device CUDA0 -ngl 99 \ + --parallel 4 -c 524288 \ + --no-mmap \ + --kv-tier-paged-blocks \ + --kv-tiered 25,25,50 \ + --cache-type-k turbo4 --cache-type-v turbo4 \ + --kv-tier-ssd-path "$SSD_PATH" \ + --kv-tier-semantic-index "$BGE_SMALL" \ + --instance-id "$INSTANCE" \ + --kv-tier-cold-budget-mb 8000 \ + --port "$PORT" \ + >> "$exec_log" 2>&1 & +PID_1070=$! +echo "[army-mad-lab] 1070 instance pid=$PID_1070 → $exec_log" + +# ── Instance 2: RX 480 (Polaris gfx803 in docker) — Qwen3.5-9B ──────── +# Run llama-server inside the gfx803 docker since native ROCm 7.x +# dropped Polaris support. The image is built with ROCm 6.3/6.4. +INSTANCE=mad-lab-rx480 +PORT=11438 +exec_log="$LOG_DIR/${INSTANCE}.log" + +docker run --rm \ + --device=/dev/kfd --device=/dev/dri \ + --group-add video \ + -v "$MODELS_DIR:/models:ro" \ + -v "$SSD_PATH:/ssd" \ + -v "$LOG_DIR:/logs" \ + -p $PORT:$PORT \ + "$DOCKER_IMAGE" \ + /opt/llama-server \ + -m /models/Qwen3.5-9B-Q5_K_S.gguf \ + -ngl 99 \ + --parallel 4 -c 524288 \ + --no-mmap \ + --kv-tier-paged-blocks \ + --kv-tiered 25,25,50 \ + --cache-type-k turbo4 --cache-type-v turbo4 \ + --kv-tier-ssd-path /ssd \ + --kv-tier-semantic-index "/models/bge-small-en-v1.5-q8_0.gguf" \ + --instance-id "$INSTANCE" \ + --kv-tier-cold-budget-mb 8000 \ + --port $PORT \ + >> "$exec_log" 2>&1 & +PID_RX480=$! +echo "[army-mad-lab] RX 480 docker pid=$PID_RX480 → $exec_log" + +wait $PID_1070 $PID_RX480 diff --git a/scripts/army/main.sh b/scripts/army/main.sh new file mode 100755 index 000000000000..00bf816ed169 --- /dev/null +++ b/scripts/army/main.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Army boot script: main box (R9700 + 6900XT). +# Customize the model paths, ports, and any per-machine specifics. + +set -euo pipefail + +# ── Paths ──────────────────────────────────────────────────────────── +LLAMA_BIN="${LLAMA_BIN:-$HOME/GitHub/llama.cpp/build-hip/bin/llama-server}" +MODELS_DIR="${MODELS_DIR:-$HOME/models}" +SSD_PATH="${SSD_PATH:-/var/lib/army/ssd}" +LOG_DIR="${LOG_DIR:-/var/log/army}" +BGE_SMALL="${BGE_SMALL:-$MODELS_DIR/bge-small-en-v1.5-q8_0.gguf}" + +mkdir -p "$SSD_PATH" "$LOG_DIR" + +# ── Wait for GPU + network ──────────────────────────────────────────── +# rocm-smi is a quick health check; bail early if HIP isn't up yet. +until rocm-smi >/dev/null 2>&1; do sleep 2; done + +# ── Instance 1: R9700 (32 GB RDNA4) — Qwen3.6-27B ───────────────────── +INSTANCE=main-r9700 +PORT=11435 +exec_log="$LOG_DIR/${INSTANCE}.log" + +"$LLAMA_BIN" \ + -m "$MODELS_DIR/Qwen3.6-27B-Q6_K.gguf" \ + --device ROCm0 -ngl 99 \ + --parallel 4 -c 524288 \ + --no-mmap \ + --kv-tier-paged-blocks \ + --kv-tiered 25,75,0 \ + --cache-type-k turbo4 --cache-type-v turbo4 \ + --kv-tier-ssd-path "$SSD_PATH" \ + --kv-tier-semantic-index "$BGE_SMALL" \ + --instance-id "$INSTANCE" \ + --kv-tier-cold-budget-mb 10000 \ + --port "$PORT" \ + >> "$exec_log" 2>&1 & +PID_R9700=$! +echo "[army-main] R9700 instance pid=$PID_R9700 → $exec_log" + +# ── Instance 2: 6900XT (16 GB RDNA2 eGPU) — gpt-oss-20B ─────────────── +INSTANCE=main-6900xt +PORT=11436 +exec_log="$LOG_DIR/${INSTANCE}.log" + +"$LLAMA_BIN" \ + -m "$MODELS_DIR/gpt-oss-20B-MXFP4.gguf" \ + --device ROCm1 -ngl 99 \ + --parallel 4 -c 262144 \ + --no-mmap \ + --kv-tier-paged-blocks \ + --kv-tiered 30,70,0 \ + --kv-tier-ssd-path "$SSD_PATH" \ + --kv-tier-semantic-index "$BGE_SMALL" \ + --instance-id "$INSTANCE" \ + --kv-tier-cold-budget-mb 5000 \ + --port "$PORT" \ + >> "$exec_log" 2>&1 & +PID_6900=$! +echo "[army-main] 6900XT instance pid=$PID_6900 → $exec_log" + +# Wait for both. systemd will SIGTERM us on shutdown; both children +# get the signal and clean up their lockfiles via the dtor. +wait $PID_R9700 $PID_6900 diff --git a/scripts/test/run-army-matrix.sh b/scripts/test/run-army-matrix.sh new file mode 100755 index 000000000000..a89bb521a1e7 --- /dev/null +++ b/scripts/test/run-army-matrix.sh @@ -0,0 +1,297 @@ +#!/usr/bin/env bash +# MAD-137: per-device test matrix runner for the army stack. +# +# For every device in the army manifest, run: +# 1. The mt:: + paged unit/integration test suite via ctest +# 2. The MAD-137 stress driver (bounded duration, low decode floor) +# +# Aggregates pass/fail into tests/results/army-matrix-.json so CI +# (or a human) can see at a glance what's green and what isn't. Stdout +# is a human summary; the JSON is the machine view. +# +# Each device is described by these env-vary-able knobs: +# +# DEVICE_NAME label used in the report +# SSH_HOST empty for "run locally"; otherwise the ssh target +# DOCKER_CONTAINER empty for host execution; otherwise `docker exec` +# wraps the commands +# REPO_PATH absolute path to the llama.cpp checkout +# BUILD_DIR build directory under REPO_PATH +# DEVICE_FLAG value passed to llama-server's --device +# MODEL_PATH gguf model path on the target box +# BGE_PATH optional bge-small gguf path (semantic prefetch) +# STRESS_DECODE_FLOOR decode tok/s threshold for the stress test +# STRESS_DURATION stress test duration in seconds +# STRESS_PARALLEL --n-agents for the stress test +# +# The default manifest below covers the four-GPU army: +# - R9700 + 6900XT on the main box (HIP) +# - 1070 on mad-lab-2026 (CUDA) +# - RX 480 on mad-lab-2026 inside the rx480-army docker container +# (HIP gfx803, ROCm 6.4) +# +# Usage: +# bash scripts/test/run-army-matrix.sh # full matrix +# DEVICES="r9700 1070" bash scripts/test/run-army-matrix.sh +# SKIP_STRESS=1 bash scripts/test/run-army-matrix.sh # unit/integration only + +set -uo pipefail + +REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)" +RESULTS_DIR="${REPO_ROOT}/tests/results" +mkdir -p "${RESULTS_DIR}" +REPORT_DATE="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +REPORT_PATH="${RESULTS_DIR}/army-matrix-${REPORT_DATE}.json" + +# ─── Device manifest ──────────────────────────────────────────────────── +# Each device is a function that exports the required env vars. Add new +# devices as new functions; the order list below picks which ones run. + +dev_r9700() { + DEVICE_NAME="r9700" + SSH_HOST="" + DOCKER_CONTAINER="" + REPO_PATH="${REPO_ROOT}" + BUILD_DIR="build-hip" + DEVICE_FLAG="ROCm0" + MODEL_PATH="${HOME}/models/Qwen3.6-27B-Q6_K.gguf" + BGE_PATH="${HOME}/models/bge-small-en-v1.5-q8_0.gguf" + STRESS_DECODE_FLOOR="${STRESS_DECODE_FLOOR_R9700:-15}" + STRESS_DURATION="${STRESS_DURATION:-60}" + STRESS_PARALLEL="${STRESS_PARALLEL:-2}" + INSTANCE_ID="matrix-r9700" +} + +dev_6900xt() { + DEVICE_NAME="6900xt" + SSH_HOST="" + DOCKER_CONTAINER="" + REPO_PATH="${REPO_ROOT}" + BUILD_DIR="build-hip" + DEVICE_FLAG="ROCm1" + MODEL_PATH="${HOME}/models/Qwen3.5-9B-TQ3_1S.gguf" + BGE_PATH="${HOME}/models/bge-small-en-v1.5-q8_0.gguf" + STRESS_DECODE_FLOOR="${STRESS_DECODE_FLOOR_6900XT:-15}" + STRESS_DURATION="${STRESS_DURATION:-60}" + STRESS_PARALLEL="${STRESS_PARALLEL:-2}" + INSTANCE_ID="matrix-6900xt" +} + +dev_1070() { + DEVICE_NAME="1070" + SSH_HOST="mad-lab-2026" + DOCKER_CONTAINER="" + REPO_PATH="/home/kmbandy/GitHub/llama.cpp" + BUILD_DIR="build-army" + DEVICE_FLAG="CUDA0" + MODEL_PATH="/home/kmbandy/models/omnicoder-9b-q5_k_m.gguf" + BGE_PATH="/home/kmbandy/models/bge-small-en-v1.5-q8_0.gguf" + STRESS_DECODE_FLOOR="${STRESS_DECODE_FLOOR_1070:-15}" + STRESS_DURATION="${STRESS_DURATION:-60}" + STRESS_PARALLEL="${STRESS_PARALLEL:-2}" + INSTANCE_ID="matrix-1070" +} + +dev_rx480() { + DEVICE_NAME="rx480" + SSH_HOST="mad-lab-2026" + DOCKER_CONTAINER="rx480-army" + REPO_PATH="/workspace/llama.cpp" # path INSIDE the container + BUILD_DIR="build-rocm-gfx803" + DEVICE_FLAG="ROCm0" + MODEL_PATH="/models/omnicoder-9b-q5_k_m.gguf" + BGE_PATH="/models/bge-small-en-v1.5-q8_0.gguf" + STRESS_DECODE_FLOOR="${STRESS_DECODE_FLOOR_RX480:-12}" + STRESS_DURATION="${STRESS_DURATION:-60}" + STRESS_PARALLEL="${STRESS_PARALLEL:-2}" + INSTANCE_ID="matrix-rx480" +} + +DEVICES_DEFAULT="r9700 6900xt 1070 rx480" +DEVICES="${DEVICES:-${DEVICES_DEFAULT}}" + +# ─── Command-execution wrapper — ssh / docker / local ─────────────────── + +# Print the prefix that wraps a command with `ssh ... -- "docker exec ..."` +# or `ssh ...` or nothing, depending on env vars. +build_runner_prefix() { + local prefix="" + if [[ -n "${SSH_HOST}" ]]; then + prefix="ssh ${SSH_HOST}" + if [[ -n "${DOCKER_CONTAINER}" ]]; then + prefix="${prefix} docker exec ${DOCKER_CONTAINER}" + fi + prefix="${prefix} bash -lc" + else + if [[ -n "${DOCKER_CONTAINER}" ]]; then + prefix="docker exec ${DOCKER_CONTAINER} bash -lc" + else + prefix="bash -lc" + fi + fi + echo "${prefix}" +} + +# Run a shell command (string) on the configured target. Returns its +# exit code; output captured to the named file. +run_remote() { + local cmd="$1" + local outfile="$2" + local prefix + prefix="$(build_runner_prefix)" + # shellcheck disable=SC2086 + ${prefix} "${cmd}" > "${outfile}" 2>&1 +} + +# ─── Per-device test phases ──────────────────────────────────────────── + +# Phase 1: ctest unit + integration suite. Gracefully skips integration +# tests that require LLAMACPP_TEST_MODELFILE — those just print the +# yellow "no model" warning and exit 0, which ctest treats as success. +run_unit_integration() { + local outfile="$1" + local cmd + cmd="cd ${REPO_PATH}/${BUILD_DIR} && \ + LLAMACPP_TEST_MODELFILE=${MODEL_PATH} \ + ctest -R 'test-mt-|test-paged-' --output-on-failure" + run_remote "${cmd}" "${outfile}" +} + +# Phase 2: MAD-137 stress driver. Always invoked from the *host's* repo +# checkout (REPO_ROOT) — even for remote devices, because the stress +# script is short and self-contained, and we want one canonical +# implementation. For ssh/docker targets the script is copied over +# first. +run_stress() { + local outfile="$1" + + # Stage the stress script onto the target. + local target_script="${REPO_PATH}/tests/stress/stress-paged-multi-seq.py" + if [[ -n "${SSH_HOST}" || -n "${DOCKER_CONTAINER}" ]]; then + # Best-effort: assume the stress driver was synced with the rest + # of the repo. If it's missing, the run will say so. + : + fi + + local stress_cmd + stress_cmd="python3 ${target_script} \ + --bin ${REPO_PATH}/${BUILD_DIR}/bin/llama-server \ + --model ${MODEL_PATH} \ + --device ${DEVICE_FLAG} \ + --ctx 8192 --n-agents ${STRESS_PARALLEL} \ + --tier 25,75,0 --instance-id ${INSTANCE_ID} \ + --duration ${STRESS_DURATION} \ + --decode-floor ${STRESS_DECODE_FLOOR} \ + --port 18095" + if [[ -n "${BGE_PATH}" ]]; then + stress_cmd="${stress_cmd} --bge-small ${BGE_PATH}" + fi + run_remote "${stress_cmd}" "${outfile}" +} + +# ─── Main loop ────────────────────────────────────────────────────────── + +# JSON aggregation — bash + a tiny python finalizer. +TMPJSON="$(mktemp)" +echo "[" > "${TMPJSON}" +first_entry=true + +emit() { + local sep="," + [[ "${first_entry}" == "true" ]] && sep="" + first_entry=false + printf '%s\n%s\n' "${sep}" "$1" >> "${TMPJSON}" +} + +overall_pass=true + +echo "──── Army test matrix — ${REPORT_DATE} ────" +echo + +for dev in ${DEVICES}; do + # Reset env vars before each device. + DEVICE_NAME=""; SSH_HOST=""; DOCKER_CONTAINER="" + REPO_PATH=""; BUILD_DIR=""; DEVICE_FLAG="" + MODEL_PATH=""; BGE_PATH="" + STRESS_DECODE_FLOOR=""; STRESS_DURATION=""; STRESS_PARALLEL="" + INSTANCE_ID="" + + case "${dev}" in + r9700) dev_r9700 ;; + 6900xt) dev_6900xt ;; + 1070) dev_1070 ;; + rx480) dev_rx480 ;; + *) + echo "skip: unknown device '${dev}'" + continue + ;; + esac + + location_label="local" + if [[ -n "${SSH_HOST}" ]]; then + location_label="ssh:${SSH_HOST}" + if [[ -n "${DOCKER_CONTAINER}" ]]; then + location_label="${location_label}+docker:${DOCKER_CONTAINER}" + fi + elif [[ -n "${DOCKER_CONTAINER}" ]]; then + location_label="docker:${DOCKER_CONTAINER}" + fi + echo "▶ ${DEVICE_NAME} (${location_label}, build=${BUILD_DIR}, dev=${DEVICE_FLAG})" + + unit_log="${RESULTS_DIR}/${DEVICE_NAME}-unit-${REPORT_DATE}.log" + stress_log="${RESULTS_DIR}/${DEVICE_NAME}-stress-${REPORT_DATE}.log" + + unit_status="skipped" + stress_status="skipped" + + # Phase 1: ctest + if run_unit_integration "${unit_log}"; then + unit_status="pass" + else + unit_status="fail" + overall_pass=false + fi + echo " unit/integration: ${unit_status} (log: ${unit_log})" + + # Phase 2: stress (skip with SKIP_STRESS=1) + if [[ "${SKIP_STRESS:-0}" != "1" ]]; then + if run_stress "${stress_log}"; then + stress_status="pass" + else + stress_status="fail" + overall_pass=false + fi + echo " stress: ${stress_status} (log: ${stress_log})" + fi + + # Per-device JSON entry. + json_entry=$(printf '{"device":"%s","location":"%s","build_dir":"%s","device_flag":"%s","unit":"%s","stress":"%s","unit_log":"%s","stress_log":"%s"}' \ + "${DEVICE_NAME}" \ + "${location_label}" \ + "${BUILD_DIR}" \ + "${DEVICE_FLAG}" \ + "${unit_status}" \ + "${stress_status}" \ + "${unit_log}" \ + "${stress_log}") + emit "${json_entry}" +done + +echo "]" >> "${TMPJSON}" + +# Finalize: wrap the per-device array in a top-level object with status. +overall_status="pass" +[[ "${overall_pass}" == "true" ]] || overall_status="fail" +python3 - "${TMPJSON}" "${REPORT_PATH}" "${overall_status}" "${REPORT_DATE}" <<'PY' +import json, sys +src, dst, overall, date = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4] +with open(src) as f: + devices = json.load(f) +out = {"date": date, "overall": overall, "devices": devices} +with open(dst, "w") as f: + json.dump(out, f, indent=2) +print(f"\nReport: {dst} (overall: {overall})") +PY +rm -f "${TMPJSON}" + +[[ "${overall_pass}" == "true" ]] diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c11d9dfeabcd..7bee837333f6 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -27,7 +27,6 @@ add_library(llama llama-kv-cache.cpp llama-kv-cache-iswa.cpp llama-kv-cache-paged.cpp - llama-kv-cache-tiered.cpp llama-memory.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5de13bb30fe7..626841488c15 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -300,6 +300,9 @@ llama_context::llama_context( /*.kv_tier_semantic_topk =*/ params.kv_tier_semantic_topk, /*.kv_tier_paged_blocks =*/ params.kv_tier_paged_blocks, /*.kv_tier_paged_block_size =*/ params.kv_tier_paged_block_size, + /*.kv_tier_cold_resume =*/ params.kv_tier_cold_resume, + /*.kv_tier_instance_id =*/ params.kv_tier_instance_id, + /*.kv_tier_cold_budget_mb =*/ params.kv_tier_cold_budget_mb, }; memory.reset(model.create_memory(params_mem, cparams)); @@ -3215,6 +3218,9 @@ llama_context_params llama_context_default_params() { /*.kv_tier_semantic_topk =*/ 5, /*.kv_tier_paged_blocks =*/ false, /*.kv_tier_paged_block_size =*/ 0, + /*.kv_tier_cold_resume =*/ false, + /*.kv_tier_instance_id =*/ nullptr, + /*.kv_tier_cold_budget_mb =*/ 0, }; return result; diff --git a/src/llama-eviction-policy.h b/src/llama-eviction-policy.h deleted file mode 100644 index 169f01c51ddc..000000000000 --- a/src/llama-eviction-policy.h +++ /dev/null @@ -1,235 +0,0 @@ -#pragma once - -#include "llama.h" -#include -#include -#include -#include -#include - -// Eviction policy types -enum llama_eviction_policy { - LRU, // Least Recently Used - LFU, // Least Frequently Used - ATTENTION, // Attention-based eviction - HYBRID // Attention + recency + frequency (default) -}; - -// Eviction score components -struct llama_eviction_score { - float attention_weight; // 0.0-1.0, lower = less important - float recency_weight; // Based on time since last access - float frequency_weight; // Based on access frequency - - // Default weights for hybrid policy - static constexpr float DEFAULT_ATTENTION_WEIGHT = 0.5f; - static constexpr float DEFAULT_RECENCY_WEIGHT = 0.3f; - static constexpr float DEFAULT_FREQUENCY_WEIGHT = 0.2f; - - // Configurable weights for hybrid policy - float att_w = DEFAULT_ATTENTION_WEIGHT; - float rec_w = DEFAULT_RECENCY_WEIGHT; - float freq_w = DEFAULT_FREQUENCY_WEIGHT; - - // Calculate combined eviction score (higher = more likely to evict) - float calculate() const { - return att_w * attention_weight + - rec_w * recency_weight + - freq_w * frequency_weight; - } - - // Calculate with custom weights - float calculate(float att_w, float rec_w, float freq_w) const { - return att_w * attention_weight + - rec_w * recency_weight + - freq_w * frequency_weight; - } -}; - -// Token metadata for eviction decisions -struct llama_token_metadata { - llama_pos position; - float attention_score; - std::time_t last_access; - uint32_t access_count; - - llama_token_metadata() : position(0), attention_score(0.0f), - last_access(std::time(nullptr)), access_count(0) {} - - llama_token_metadata(llama_pos pos, float att_score = 0.0f) - : position(pos), attention_score(att_score), - last_access(std::time(nullptr)), access_count(0) {} - - void record_access() { - access_count++; - last_access = std::time(nullptr); - } - - void update_attention_score(float score) { - // Exponential moving average for attention score - attention_score = 0.7f * attention_score + 0.3f * score; - } -}; - -// Token metadata store for eviction tracking -class llama_token_metadata_store { -public: - using pos_to_metadata = std::unordered_map; - - void record_access(llama_pos pos) { - auto it = metadata.find(pos); - if (it != metadata.end()) { - it->second.record_access(); - } else { - metadata.emplace(pos, llama_token_metadata(pos)); - } - } - - void update_attention(llama_pos pos, float attention_score) { - auto it = metadata.find(pos); - if (it != metadata.end()) { - it->second.update_attention_score(attention_score); - } - } - - void add_token(llama_pos pos, float initial_attention = 0.0f) { - if (metadata.find(pos) == metadata.end()) { - metadata.emplace(pos, llama_token_metadata(pos, initial_attention)); - } - } - - void remove_token(llama_pos pos) { - metadata.erase(pos); - } - - const llama_token_metadata* get_metadata(llama_pos pos) const { - auto it = metadata.find(pos); - return it != metadata.end() ? &it->second : nullptr; - } - - // Batch attention update for efficiency - void batch_update_attention(const std::vector& positions, - const std::vector& attention_scores) { - if (positions.size() != attention_scores.size()) { - return; - } - for (size_t i = 0; i < positions.size(); i++) { - update_attention(positions[i], attention_scores[i]); - } - } - - // Get tokens below attention threshold (candidates for eviction) - std::vector get_low_attention_tokens(float threshold) const { - std::vector result; - for (const auto& kv : metadata) { - if (kv.second.attention_score < threshold) { - result.push_back(kv.first); - } - } - return result; - } - - // Normalize attention scores across all tokens - void normalize_attention_scores() { - if (metadata.empty()) return; - - // Find max attention score - float max_score = 0.0f; - for (const auto& kv : metadata) { - max_score = std::max(max_score, kv.second.attention_score); - } - - if (max_score > 0.0f) { - for (auto& kv : metadata) { - kv.second.attention_score /= max_score; - } - } - } - - // Clear all metadata - void clear_all() { - metadata.clear(); - } - - // Get number of tracked tokens - size_t get_token_count() const { - return metadata.size(); - } - - // Get tokens sorted by eviction score (highest first) - std::vector get_eviction_candidates( - llama_eviction_policy policy, - uint32_t count, - float attention_weight = llama_eviction_score::DEFAULT_ATTENTION_WEIGHT, - float recency_weight = llama_eviction_score::DEFAULT_RECENCY_WEIGHT, - float frequency_weight = llama_eviction_score::DEFAULT_FREQUENCY_WEIGHT) const { - - std::vector> scored_tokens; - scored_tokens.reserve(metadata.size()); - - const auto now = std::time(nullptr); - - for (const auto& kv : metadata) { - const auto& pos = kv.first; - const auto& meta = kv.second; - float score = 0.0f; - - switch (policy) { - case LRU: - // Score based on recency (older = higher score) - score = float(now - meta.last_access); - break; - - case LFU: - // Score based on frequency (less frequent = higher score) - score = meta.access_count > 0 ? 1.0f / float(meta.access_count) : 1.0f; - break; - - case ATTENTION: - // Score based on attention (lower attention = higher score) - score = 1.0f - meta.attention_score; - break; - - case HYBRID: - default: - // Combined score - float att_score = 1.0f - meta.attention_score; - float rec_score = float(now - meta.last_access) / 60.0f; // Normalize to minutes - float freq_score = meta.access_count > 0 ? 1.0f / float(meta.access_count) : 1.0f; - score = att_score * attention_weight + - rec_score * recency_weight + - freq_score * frequency_weight; - break; - } - - scored_tokens.emplace_back(pos, score); - } - - // Sort by score (highest first) - std::sort(scored_tokens.begin(), scored_tokens.end(), - [](const auto& a, const auto& b) { return a.second > b.second; }); - - // Return top 'count' positions - std::vector result; - size_t n = std::min(count, uint32_t(scored_tokens.size())); - result.reserve(n); - for (size_t i = 0; i < n; ++i) { - result.push_back(scored_tokens[i].first); - } - - return result; - } - - size_t size() const { return metadata.size(); } - void clear() { metadata.clear(); } - -private: - pos_to_metadata metadata; -}; - -// Layer importance for per-layer eviction decisions -struct llama_layer_importance { - float importance; // 0.0-1.0, higher = more important to keep - - static constexpr float DEFAULT_IMPORTANCE = 1.0f; -}; diff --git a/src/llama-kv-cache-paged.cpp b/src/llama-kv-cache-paged.cpp index 94880987d7db..3cb49e436c84 100644 --- a/src/llama-kv-cache-paged.cpp +++ b/src/llama-kv-cache-paged.cpp @@ -2,16 +2,22 @@ #include "llama-impl.h" #include "llama-batch.h" +#include "llama-io.h" #include "llama-model.h" #include "llama-hparams.h" +#include "ggml.h" #include "ggml-backend.h" +#include "memory-tier/mt-quant.h" + #include #include #include +#include #include #include #include +#include #include #include #include @@ -59,6 +65,9 @@ llama_kv_cache_paged::llama_kv_cache_paged( uint32_t n_warm_blocks, uint32_t n_cold_blocks, std::string ssd_path, + bool cold_resume, + std::string instance_id, + uint32_t cold_budget_mb, layer_filter_cb filter, ggml_type type_k, ggml_type type_v) @@ -208,9 +217,77 @@ llama_kv_cache_paged::llama_kv_cache_paged( LLAMA_LOG_WARN("llama_kv_cache_paged: cold tier requested with empty ssd_path; disabling cold\n"); n_cold_blocks_ = 0; } else { - const std::string subdir = cold_path_ + "/paged"; - (void) ::mkdir(cold_path_.c_str(), 0700); - (void) ::mkdir(subdir.c_str(), 0700); + // MAD-131: per-instance subdir + lockfile. Default instance + // ID is the process pid as a string; operator override via + // --instance-id makes restarts deterministic. + instance_id_ = instance_id.empty() ? std::to_string(::getpid()) : instance_id; + + // MAD-131: budget cap on cold pool size. Convert MiB → blocks + // (across K+V per attention layer). Cap n_cold_blocks_ + // BEFORE building the file paths so the truncation/sizing + // arithmetic uses the capped count. + if (cold_budget_mb > 0) { + const size_t bytes_budget = (size_t) cold_budget_mb * 1024ull * 1024ull; + const size_t bytes_per_block_per_layer = k_bytes_per_block_ + v_bytes_per_block_; + if (bytes_per_block_per_layer == 0 || n_attn_layers == 0) { + LLAMA_LOG_WARN("llama_kv_cache_paged: cold-budget-mb=%u set but layer sizing " + "is unknown; ignoring cap\n", cold_budget_mb); + } else { + const size_t bytes_per_block = bytes_per_block_per_layer * n_attn_layers; + const uint32_t max_blocks_for_budget = (uint32_t) std::min( + UINT32_MAX, bytes_budget / bytes_per_block); + if (max_blocks_for_budget < n_cold_blocks_) { + LLAMA_LOG_INFO("llama_kv_cache_paged: cold-budget-mb=%u caps cold pool " + "to %u blocks (was %u)\n", + cold_budget_mb, max_blocks_for_budget, n_cold_blocks_); + n_cold_blocks_ = max_blocks_for_budget; + } + } + } + if (n_cold_blocks_ == 0) { + // Budget reduced cold to zero; skip file setup. + LLAMA_LOG_WARN("llama_kv_cache_paged: cold-budget-mb capped pool to 0 — " + "cold tier effectively disabled\n"); + goto cold_setup_done; + } + + const std::string subdir = + cold_path_ + "/paged/instance-" + instance_id_; + (void) ::mkdir(cold_path_.c_str(), 0700); + (void) ::mkdir((cold_path_ + "/paged").c_str(), 0700); + (void) ::mkdir(subdir.c_str(), 0700); + + // MAD-131: lockfile + .pid for double-start refusal. + // flock(LOCK_EX | LOCK_NB) gives clean OS-level mutual + // exclusion. The .pid file is informational — read it on + // lock failure to give the operator a useful diagnostic. + const std::string lock_path = subdir + "/.lock"; + const std::string pid_path = subdir + "/.pid"; + cold_lock_path_ = lock_path; + cold_pid_path_ = pid_path; + cold_lock_fd_ = ::open(lock_path.c_str(), O_RDWR | O_CREAT, 0600); + if (cold_lock_fd_ < 0) { + throw std::runtime_error("llama_kv_cache_paged: cannot open lockfile " + lock_path); + } + if (::flock(cold_lock_fd_, LOCK_EX | LOCK_NB) != 0) { + std::string holder_pid = "(unknown)"; + std::ifstream pf(pid_path); + if (pf) std::getline(pf, holder_pid); + ::close(cold_lock_fd_); + cold_lock_fd_ = -1; + throw std::runtime_error( + "llama_kv_cache_paged: instance " + instance_id_ + + " is already in use by pid " + holder_pid + + " (lockfile " + lock_path + " held). Use a different " + "--instance-id, or stop the holder, or rm the lockfile " + "if it's stale."); + } + // Write our pid into .pid so future failed-acquires can + // diagnose. Best-effort; not fatal if write fails. + { + std::ofstream pf(pid_path, std::ios::trunc); + if (pf) pf << ::getpid() << "\n"; + } cold_fd_k_.assign(n_layer, -1); cold_fd_v_.assign(n_layer, -1); @@ -218,12 +295,18 @@ llama_kv_cache_paged::llama_kv_cache_paged( const off_t v_file_bytes = (off_t) n_cold_blocks_ * (off_t) v_bytes_per_block_; uint64_t total_bytes = 0; bool ok = true; + // MAD-130: cold_resume=true → open WITHOUT O_TRUNC so existing + // bytes survive. We'll load the in-memory index from the + // sidecar after the layer files are open. + const int open_flags = cold_resume + ? (O_RDWR | O_CREAT) // preserve contents + : (O_RDWR | O_CREAT | O_TRUNC); // legacy: fresh for (uint32_t il = 0; il < n_layer && ok; ++il) { if (!layers_[il].k) continue; const std::string kpath = subdir + "/L" + std::to_string(il) + ".k.bin"; const std::string vpath = subdir + "/L" + std::to_string(il) + ".v.bin"; - const int fk = ::open(kpath.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600); - const int fv = ::open(vpath.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600); + const int fk = ::open(kpath.c_str(), open_flags, 0600); + const int fv = ::open(vpath.c_str(), open_flags, 0600); if (fk < 0 || fv < 0 || ::ftruncate(fk, k_file_bytes) < 0 || ::ftruncate(fv, v_file_bytes) < 0) { @@ -246,22 +329,43 @@ llama_kv_cache_paged::llama_kv_cache_paged( cold_fd_v_.clear(); n_cold_blocks_ = 0; } else { - // Free-stack of cold slot ids. + // Free-stack of cold slot ids. Initialize as all-free; + // the cold-resume sidecar load below will mark in-use + // slots and pull them off the free stack. cold_pool_free_.reserve(n_cold_blocks_); for (uint32_t i = n_cold_blocks_; i > 0; --i) { cold_pool_free_.push_back(i - 1); } cold_slot_for_.assign(n_seq_max_, std::vector()); + + // MAD-130: cold-resume — load the in-memory index from + // the sidecar. Sidecar absent → log warning but proceed + // (the on-disk bytes exist but we don't know what's + // in them; first-write of any slot will overwrite). + if (cold_resume) { + const std::string sidecar_path = subdir + "/index.bin"; + if (!load_cold_index_sidecar_(sidecar_path)) { + LLAMA_LOG_WARN("llama_kv_cache_paged: --kv-tier-cold-resume requested but " + "sidecar %s missing or invalid — starting with empty cold " + "index (existing file bytes ignored)\n", + sidecar_path.c_str()); + } + } + LLAMA_LOG_INFO("llama_kv_cache_paged: cold tier enabled — %u blocks × " - "%.1f KiB/block (K+V) × %u attn layers = %.1f MiB on %s\n", + "%.1f KiB/block (K+V) × %u attn layers = %.1f MiB on %s%s " + "(instance=%s)\n", n_cold_blocks_, (double)(k_bytes_per_block_ + v_bytes_per_block_) / 1024.0, n_attn_layers, (double) total_bytes / (1024.0 * 1024.0), - subdir.c_str()); + subdir.c_str(), + cold_resume ? " (resume mode)" : "", + instance_id_.c_str()); } } } +cold_setup_done:; // ── Input tensor (block_table) ── // @@ -301,10 +405,50 @@ llama_kv_cache_paged::llama_kv_cache_paged( static_assert((int32_t) mt::kInvalidBlockId == kInvalidBlockTableEntry, "mt::kInvalidBlockId must reinterpret to -1 in i32 for kernel compat"); +// MAD-132: enforce the single-threading contract. First mutator call +// captures std::this_thread::get_id(); subsequent calls assert match. +// Release builds compile this to a no-op via assert(). +// +// Method is const to allow calling from const methods that mutate +// internal caches (none today, but const-correctness for future). +// The captured_thread_id_ field is mutable for the same reason. +void llama_kv_cache_paged::check_thread_id_() const { +#ifndef NDEBUG + const std::thread::id self = std::this_thread::get_id(); + if (captured_thread_id_ == std::thread::id()) { + captured_thread_id_ = self; + } else { + assert(captured_thread_id_ == self && + "llama_kv_cache_paged: cross-thread mutation detected. " + "This cache is single-threaded; see header doc block " + "(MAD-132 / Epic A4). Add explicit synchronization or " + "use a worker queue if async access is required."); + } +#endif +} + +// MAD-132: timestamp helper (microseconds since epoch). Used for +// preempt-fairness victim selection. +static uint64_t now_us_() { + return (uint64_t) std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()).count(); +} + llama_kv_cache_paged::~llama_kv_cache_paged() { // MAD-121: close cold-tier fds. for (int fd : cold_fd_k_) if (fd >= 0) ::close(fd); for (int fd : cold_fd_v_) if (fd >= 0) ::close(fd); + // MAD-131: release the per-instance lockfile + remove .pid. flock + // releases automatically when the fd closes; .pid is informational + // and best-effort cleaned (next start with same instance_id will + // overwrite anyway). + if (cold_lock_fd_ >= 0) { + ::close(cold_lock_fd_); + cold_lock_fd_ = -1; + } + if (!cold_pid_path_.empty()) { + (void) ::unlink(cold_pid_path_.c_str()); + } if (buf_storage_) ggml_backend_buffer_free(buf_storage_); if (buf_inputs_) ggml_backend_buffer_free(buf_inputs_); if (ctx_storage_) ggml_free(ctx_storage_); @@ -312,6 +456,7 @@ llama_kv_cache_paged::~llama_kv_cache_paged() { } bool llama_kv_cache_paged::ensure_blocks_for(llama_seq_id seq_id, uint32_t n_new_tokens) { + check_thread_id_(); if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max_) return false; if (n_new_tokens == 0) return true; @@ -408,6 +553,7 @@ bool llama_kv_cache_paged::evict_block_to_warm(llama_seq_id seq_id, uint32_t log // free the GPU physical back to the pool. table_.swap_block(seq_id, logical_block, cpu_physical); pool_.free_block(gpu_physical); + ++evict_h2w_total_; // MAD-133 return true; } @@ -446,6 +592,7 @@ bool llama_kv_cache_paged::restore_block_from_warm(llama_seq_id seq_id, uint32_t table_.swap_block(seq_id, logical_block, gpu_physical); pool_.free_block(cpu_physical); + ++restore_w2h_total_; // MAD-133 return true; } @@ -501,6 +648,7 @@ bool llama_kv_cache_paged::fault_in_warm_blocks_for_batch(const llama_ubatch & u const uint32_t physical = table_.get_physical(sid, lb); if (physical == mt::kInvalidBlockId) continue; if (!pool_.is_gpu(physical)) continue; + if (pool_.refcount(physical) > 1) continue; // MAD-128: skip shared ++gpu_count; if (lb < oldest_lb) oldest_lb = lb; } @@ -570,6 +718,101 @@ bool llama_kv_cache_paged::fault_in_warm_blocks_for_batch(const llama_ubatch & u return true; } +// MAD-128: CoW any blocks being written this ubatch that are shared +// (refcount > 1 from a prior seq_cp). For each unique (seq, lblock) +// touched by a write in `ub`: if that physical's refcount > 1, allocate +// a fresh GPU block, copy K/V from the shared block into it, swap the +// seq's table entry, decrement the old block's refcount. +// +// Why this is needed: seq_cp shares physical blocks via refcount. If +// either sequence then writes to a shared block (e.g. partial last +// block at the share boundary), the kernel's K/V scatter would corrupt +// the OTHER sequence's view of that block. CoW preserves the invariant +// that each (seq, lblock) writable cell is uniquely owned. +// +// Called from apply_ubatch_to_state AFTER fault-in (so the block is +// GPU-resident before we copy) and BEFORE prepare_batch_tensors (so +// the uploaded block_table reflects the new physicals). +bool llama_kv_cache_paged::cow_writes_for_ubatch(const llama_ubatch & ub) { + if (ub.n_tokens == 0) return true; + + // Collect unique (seq, lblock) pairs touched by writes this ubatch. + // Using a set: small N (<= n_tokens), de-dup'd lookups. + std::vector> touched; + touched.reserve(ub.n_tokens); + for (uint32_t i = 0; i < ub.n_tokens; ++i) { + const llama_pos pos = ub.pos[i]; + if (pos < 0) continue; + llama_seq_id sid = 0; + if (ub.seq_id && ub.seq_id[i] && ub.n_seq_id && ub.n_seq_id[i] > 0) { + sid = ub.seq_id[i][0]; + } + if (sid < 0 || (uint32_t) sid >= n_seq_max_) continue; + touched.emplace_back(sid, (uint32_t) pos / block_size_); + } + std::sort(touched.begin(), touched.end()); + touched.erase(std::unique(touched.begin(), touched.end()), touched.end()); + + uint32_t n_cowed = 0; + for (auto [sid, lblock] : touched) { + if (lblock >= table_.num_blocks(sid)) continue; // freshly allocated by ensure_blocks_for; refcount=1 + const uint32_t physical = table_.get_physical(sid, lblock); + if (physical == mt::kInvalidBlockId) continue; + if (pool_.refcount(physical) <= 1) continue; // not shared + if (!pool_.is_gpu(physical)) continue; // CoW only on GPU-resident; warm/cold writes go through fault-in path + + // Allocate a fresh GPU block. If the pool is full, try evicting + // an LRU GPU block to warm to make room (only if warm enabled). + uint32_t new_phys = pool_.alloc_gpu(); + while (new_phys == mt::kInvalidBlockId && warm_enabled()) { + if (!evict_lru_to_warm()) break; + new_phys = pool_.alloc_gpu(); + } + if (new_phys == mt::kInvalidBlockId) { + LLAMA_LOG_ERROR("llama_kv_cache_paged::cow_writes_for_ubatch: GPU pool " + "exhausted attempting CoW for seq=%d lblock=%u — " + "write would corrupt shared block. Refusing batch.\n", + sid, lblock); + return false; + } + + // GPU-to-GPU copy via host bounce buffer (no native ggml D2D + // primitive; the get/set pair routes through host memory but the + // block is small — k+v_bytes_per_block, ~17 KiB at turbo4). + std::vector kbuf(k_bytes_per_block_); + std::vector vbuf(v_bytes_per_block_); + const size_t k_off_old = (size_t) physical * k_bytes_per_block_; + const size_t v_off_old = (size_t) physical * v_bytes_per_block_; + const size_t k_off_new = (size_t) new_phys * k_bytes_per_block_; + const size_t v_off_new = (size_t) new_phys * v_bytes_per_block_; + for (uint32_t il = 0; il < layers_.size(); ++il) { + const auto & layer = layers_[il]; + if (!layer.k) continue; + ggml_backend_tensor_get(layer.k, kbuf.data(), k_off_old, k_bytes_per_block_); + ggml_backend_tensor_get(layer.v, vbuf.data(), v_off_old, v_bytes_per_block_); + ggml_backend_tensor_set(layer.k, kbuf.data(), k_off_new, k_bytes_per_block_); + ggml_backend_tensor_set(layer.v, vbuf.data(), v_off_new, v_bytes_per_block_); + } + + // Swap seq's table entry to the fresh block; decrement old. + // The old block's refcount drops by 1 (from N to N-1). If N was + // 2, the old block now has refcount 1 (the OTHER seq still owns + // it, no one freed). If N was higher, more shares remain. + // Fingerprint follows the seq's logical block — same content, + // new physical, no fingerprint update needed. + table_.swap_block(sid, lblock, new_phys); + pool_.free_block(physical); + ++n_cowed; + } + + if (n_cowed > 0) { + LLAMA_LOG_DEBUG("llama_kv_cache_paged::cow_writes_for_ubatch: CoW'd %u " + "shared block(s) for upcoming writes (pool free: gpu=%zu)\n", + n_cowed, pool_.n_free_gpu()); + } + return true; +} + bool llama_kv_cache_paged::evict_lru_to_warm() { if (!warm_enabled()) return false; @@ -595,6 +838,12 @@ bool llama_kv_cache_paged::evict_lru_to_warm() { const uint32_t physical = table_.get_physical(sid, lb); if (physical == mt::kInvalidBlockId) continue; if (!pool_.is_gpu(physical)) continue; + // MAD-128: skip shared blocks (refcount > 1) — evicting them + // doesn't free GPU space (other seqs still hold the physical), + // so the eviction-retry caller would loop. Shared blocks stay + // GPU-resident until either the other seq frees its reference + // or its own write triggers CoW. + if (pool_.refcount(physical) > 1) continue; ++gpu_count; if (lb < oldest_gpu_lblock) oldest_gpu_lblock = lb; } @@ -654,6 +903,7 @@ int llama_kv_cache_paged::evict_seq_to_warm(llama_seq_id seq_id) { } ++moved; } + if (moved > 0) ++seq_preempt_total_; // MAD-133 LLAMA_LOG_DEBUG("evict_seq_to_warm: seq=%d evicted %d block(s)\n", seq_id, moved); return moved; } @@ -675,6 +925,7 @@ int llama_kv_cache_paged::restore_seq_from_warm(llama_seq_id seq_id) { } ++moved; } + if (moved > 0) ++seq_restore_total_; // MAD-133 LLAMA_LOG_DEBUG("restore_seq_from_warm: seq=%d restored %d block(s)\n", seq_id, moved); return moved; } @@ -684,8 +935,22 @@ bool llama_kv_cache_paged::evict_block_to_cold(llama_seq_id seq_id, uint32_t lbl if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max_) return false; if (lblock >= table_.num_blocks(seq_id)) return false; if (cold_pool_free_.empty()) { - LLAMA_LOG_WARN("evict_block_to_cold: cold pool full (%u in use); refusing\n", cold_in_use_); - return false; + // MAD-132: cold full → escalate by dropping the oldest cold + // block to make room. The dropped seq's K/V is gone (its table + // entry was already kInvalidBlockId from prior cold spill); + // future kernel reads of that block return -INFINITY logit (= + // zero attention contribution), same as middle-wipe holes. + // This is the last-resort policy before refusing the eviction + // and surfacing FAILED_PREPARE up the call chain. + if (!drop_oldest_cold_block()) { + LLAMA_LOG_WARN("evict_block_to_cold: cold pool full (%u in use) AND drop_oldest " + "failed; refusing eviction. Caller should fall back to keeping " + "the block in warm or returning a 503 to the client.\n", + cold_in_use_); + return false; + } + LLAMA_LOG_INFO("evict_block_to_cold: cold pool was full; dropped oldest cold block " + "to make room for seq=%d lblock=%u\n", seq_id, lblock); } const uint32_t phys = table_.get_physical(seq_id, lblock); @@ -701,6 +966,27 @@ bool llama_kv_cache_paged::evict_block_to_cold(llama_seq_id seq_id, uint32_t lbl std::vector kbuf(k_bytes_per_block_); std::vector vbuf(v_bytes_per_block_); + // MAD-135: int4 cold compression for F16 caches. Compressed slot + // layout: [4 bytes scale_f32][ceil(n/2) bytes packed int4]. The + // cold file slot is sized for raw bytes, so the compressed payload + // fits comfortably (4× smaller than F16). For Q8_0 / turbo4 the + // bytes are already byte-quantized at the cache layer; raw cold + // copy is correct + smaller than int4-on-quant would be. + const bool compress_k = (type_k_ == GGML_TYPE_F16); + const bool compress_v = (type_v_ == GGML_TYPE_F16); + std::vector tmp_k_f32, tmp_v_f32; + std::vector tmp_k_q4, tmp_v_q4; + if (compress_k) { + const size_t n_k = k_bytes_per_block_ / sizeof(uint16_t); // F16 element count + tmp_k_f32.resize(n_k); + tmp_k_q4.resize(sizeof(float) + (n_k + 1) / 2); // [scale][packed] + } + if (compress_v) { + const size_t n_v = v_bytes_per_block_ / sizeof(uint16_t); + tmp_v_f32.resize(n_v); + tmp_v_q4.resize(sizeof(float) + (n_v + 1) / 2); + } + for (uint32_t il = 0; il < layers_.size(); ++il) { const auto & layer = layers_[il]; if (!layer.k) continue; @@ -717,9 +1003,41 @@ bool llama_kv_cache_paged::evict_block_to_cold(llama_seq_id seq_id, uint32_t lbl std::memcpy(vbuf.data(), warm_v_[il].data() + v_off, v_bytes_per_block_); } - const ssize_t wk = ::pwrite(cold_fd_k_[il], kbuf.data(), k_bytes_per_block_, k_off_cold); - const ssize_t wv = ::pwrite(cold_fd_v_[il], vbuf.data(), v_bytes_per_block_, v_off_cold); - if (wk != (ssize_t) k_bytes_per_block_ || wv != (ssize_t) v_bytes_per_block_) { + // ── K write (compressed or raw) ── + ssize_t wk; + if (compress_k) { + const size_t n_k = tmp_k_f32.size(); + const ggml_fp16_t * src_f16 = reinterpret_cast(kbuf.data()); + for (size_t i = 0; i < n_k; ++i) tmp_k_f32[i] = ggml_fp16_to_fp32(src_f16[i]); + float scale = 0.0f; + mt::quantize_block_int4_with_scale(tmp_k_f32.data(), n_k, &scale, + tmp_k_q4.data() + sizeof(float)); + std::memcpy(tmp_k_q4.data(), &scale, sizeof(float)); + wk = ::pwrite(cold_fd_k_[il], tmp_k_q4.data(), tmp_k_q4.size(), k_off_cold); + if (wk != (ssize_t) tmp_k_q4.size()) wk = -1; + } else { + wk = ::pwrite(cold_fd_k_[il], kbuf.data(), k_bytes_per_block_, k_off_cold); + if (wk != (ssize_t) k_bytes_per_block_) wk = -1; + } + + // ── V write (compressed or raw) ── + ssize_t wv; + if (compress_v) { + const size_t n_v = tmp_v_f32.size(); + const ggml_fp16_t * src_f16 = reinterpret_cast(vbuf.data()); + for (size_t i = 0; i < n_v; ++i) tmp_v_f32[i] = ggml_fp16_to_fp32(src_f16[i]); + float scale = 0.0f; + mt::quantize_block_int4_with_scale(tmp_v_f32.data(), n_v, &scale, + tmp_v_q4.data() + sizeof(float)); + std::memcpy(tmp_v_q4.data(), &scale, sizeof(float)); + wv = ::pwrite(cold_fd_v_[il], tmp_v_q4.data(), tmp_v_q4.size(), v_off_cold); + if (wv != (ssize_t) tmp_v_q4.size()) wv = -1; + } else { + wv = ::pwrite(cold_fd_v_[il], vbuf.data(), v_bytes_per_block_, v_off_cold); + if (wv != (ssize_t) v_bytes_per_block_) wv = -1; + } + + if (wk < 0 || wv < 0) { LLAMA_LOG_WARN("evict_block_to_cold: pwrite failed at layer=%u cold_idx=%u (wk=%zd wv=%zd)\n", il, cold_idx, wk, wv); return false; @@ -731,6 +1049,7 @@ bool llama_kv_cache_paged::evict_block_to_cold(llama_seq_id seq_id, uint32_t lbl pool_.free_block(phys); mark_cold(cold_slot_for_, seq_id, lblock, cold_idx); ++cold_in_use_; + ++evict_w2c_total_; // MAD-133 return true; } @@ -750,17 +1069,77 @@ bool llama_kv_cache_paged::restore_block_from_cold(llama_seq_id seq_id, uint32_t std::vector kbuf(k_bytes_per_block_); std::vector vbuf(v_bytes_per_block_); + // MAD-135: dispatch on cache type for cold compression. Symmetric + // with evict_block_to_cold's compression policy. + const bool decompress_k = (type_k_ == GGML_TYPE_F16); + const bool decompress_v = (type_v_ == GGML_TYPE_F16); + std::vector tmp_k_f32, tmp_v_f32; + std::vector tmp_k_q4, tmp_v_q4; + if (decompress_k) { + const size_t n_k = k_bytes_per_block_ / sizeof(uint16_t); + tmp_k_f32.resize(n_k); + tmp_k_q4.resize(sizeof(float) + (n_k + 1) / 2); + } + if (decompress_v) { + const size_t n_v = v_bytes_per_block_ / sizeof(uint16_t); + tmp_v_f32.resize(n_v); + tmp_v_q4.resize(sizeof(float) + (n_v + 1) / 2); + } + for (uint32_t il = 0; il < layers_.size(); ++il) { const auto & layer = layers_[il]; if (!layer.k) continue; - const ssize_t rk = ::pread(cold_fd_k_[il], kbuf.data(), k_bytes_per_block_, k_off_cold); - const ssize_t rv = ::pread(cold_fd_v_[il], vbuf.data(), v_bytes_per_block_, v_off_cold); - if (rk != (ssize_t) k_bytes_per_block_ || rv != (ssize_t) v_bytes_per_block_) { - LLAMA_LOG_WARN("restore_block_from_cold: pread short at layer=%u cold_idx=%u (rk=%zd rv=%zd)\n", - il, cold_idx, rk, rv); - pool_.free_block(gpu_phys); - return false; + // ── K read + decompress ── + if (decompress_k) { + const ssize_t rk = ::pread(cold_fd_k_[il], tmp_k_q4.data(), tmp_k_q4.size(), k_off_cold); + if (rk != (ssize_t) tmp_k_q4.size()) { + LLAMA_LOG_WARN("restore_block_from_cold: pread short K at layer=%u cold_idx=%u (rk=%zd)\n", + il, cold_idx, rk); + pool_.free_block(gpu_phys); + return false; + } + float scale = 0.0f; + std::memcpy(&scale, tmp_k_q4.data(), sizeof(float)); + const size_t n_k = tmp_k_f32.size(); + mt::dequantize_block_int4_with_scale(tmp_k_q4.data() + sizeof(float), scale, + tmp_k_f32.data(), n_k); + ggml_fp16_t * dst_f16 = reinterpret_cast(kbuf.data()); + for (size_t i = 0; i < n_k; ++i) dst_f16[i] = ggml_fp32_to_fp16(tmp_k_f32[i]); + } else { + const ssize_t rk = ::pread(cold_fd_k_[il], kbuf.data(), k_bytes_per_block_, k_off_cold); + if (rk != (ssize_t) k_bytes_per_block_) { + LLAMA_LOG_WARN("restore_block_from_cold: pread short K at layer=%u cold_idx=%u (rk=%zd)\n", + il, cold_idx, rk); + pool_.free_block(gpu_phys); + return false; + } + } + + // ── V read + decompress ── + if (decompress_v) { + const ssize_t rv = ::pread(cold_fd_v_[il], tmp_v_q4.data(), tmp_v_q4.size(), v_off_cold); + if (rv != (ssize_t) tmp_v_q4.size()) { + LLAMA_LOG_WARN("restore_block_from_cold: pread short V at layer=%u cold_idx=%u (rv=%zd)\n", + il, cold_idx, rv); + pool_.free_block(gpu_phys); + return false; + } + float scale = 0.0f; + std::memcpy(&scale, tmp_v_q4.data(), sizeof(float)); + const size_t n_v = tmp_v_f32.size(); + mt::dequantize_block_int4_with_scale(tmp_v_q4.data() + sizeof(float), scale, + tmp_v_f32.data(), n_v); + ggml_fp16_t * dst_f16 = reinterpret_cast(vbuf.data()); + for (size_t i = 0; i < n_v; ++i) dst_f16[i] = ggml_fp32_to_fp16(tmp_v_f32[i]); + } else { + const ssize_t rv = ::pread(cold_fd_v_[il], vbuf.data(), v_bytes_per_block_, v_off_cold); + if (rv != (ssize_t) v_bytes_per_block_) { + LLAMA_LOG_WARN("restore_block_from_cold: pread short V at layer=%u cold_idx=%u (rv=%zd)\n", + il, cold_idx, rv); + pool_.free_block(gpu_phys); + return false; + } } const size_t k_off = (size_t) gpu_phys * k_bytes_per_block_; @@ -773,9 +1152,148 @@ bool llama_kv_cache_paged::restore_block_from_cold(llama_seq_id seq_id, uint32_t mark_cold(cold_slot_for_, seq_id, lblock, kInvalidColdIdx); cold_pool_free_.push_back(cold_idx); if (cold_in_use_ > 0) --cold_in_use_; + ++restore_c2h_total_; // MAD-133 return true; } +// MAD-125: BGE-small semantic prefetch — record + restore. + +void llama_kv_cache_paged::record_paged_block_fingerprint( + llama_seq_id seq_id, + uint32_t lblock, + std::vector embedding, + mt::SemanticIndex::Tier tier) { + paged_semantic_.add_fingerprint(seq_id, lblock, std::move(embedding), tier); +} + +uint32_t llama_kv_cache_paged::restore_semantic_paged( + llama_seq_id seq_id, + const std::vector & query_embedding, + int top_k, + float threshold) { + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max_) return 0; + + ++semantic_attempts_total_; // MAD-133 + auto hints = paged_semantic_.score(seq_id, query_embedding, top_k, threshold); + if (hints.empty()) return 0; + + uint32_t restored = 0; + uint32_t requested = 0; + uint32_t already_hot = 0; + uint32_t unmapped = 0; + uint32_t restore_fail = 0; + + for (const auto & h : hints) { + ++requested; + if (h.lblock >= table_.num_blocks(seq_id)) { ++unmapped; continue; } + const uint32_t phys = table_.get_physical(seq_id, h.lblock); + if (phys == mt::kInvalidBlockId) { ++unmapped; continue; } + + if (pool_.is_gpu(phys)) { ++already_hot; continue; } + + // Try warm first; if warm restore fails (e.g. block is in cold) + // fall back to cold restore. Both are no-ops if the tier isn't + // configured. + bool ok = warm_enabled() && restore_block_from_warm(seq_id, h.lblock); + if (!ok && cold_enabled()) { + ok = restore_block_from_cold(seq_id, h.lblock); + } + if (ok) ++restored; + else ++restore_fail; + } + + // MAD-122 acceptance criterion #5: hit-rate logging. The + // restored/requested ratio is the prefetch-effectiveness signal. + LLAMA_LOG_INFO("llama_kv_cache_paged::restore_semantic_paged: seq %d — %zu hints " + "(top_k=%d, threshold=%.2f), restored %u/%u " + "(hit-rate %.0f%%, %u already hot, %u unmapped, %u failed)\n", + seq_id, hints.size(), top_k, threshold, + restored, requested, + requested == 0 ? 0.0f : 100.0f * (float) restored / (float) requested, + already_hot, unmapped, restore_fail); + + if (restored > 0) { + ++semantic_hits_total_; // MAD-133 + semantic_blocks_restored_total_ += restored; // MAD-133 + } + return restored; +} + +// MAD-132: drop the oldest cold block — last-resort escalation when +// cold pool is full but eviction must continue. Walk cold_slot_for_ +// in seq-then-lblock order and pick the first in-use entry. This is +// "any" not strictly "oldest" — the cold tier doesn't track per-slot +// age. Acceptable for v1: the cold spillover is itself age-ordered +// (LRU warm → cold), so the lowest-indexed cold slots are typically +// the oldest evictions. True LRU tracking is a follow-up. +// +// Effect: the dropped block's data is gone. The owning seq's table +// entry stays kInvalidBlockId (already marked when the block was +// spilled to cold); future kernel reads contribute -INFINITY logit +// → 0 attention weight, same as middle-wipe holes. The seq sees a +// hole at that block position; correctness is preserved (no garbage +// reads), recall is degraded (the dropped K/V can't be attended to). +bool llama_kv_cache_paged::drop_oldest_cold_block() { + if (!cold_enabled()) return false; + + // Scan for any in-use cold slot. + for (uint32_t s = 0; s < (uint32_t) cold_slot_for_.size(); ++s) { + auto & row = cold_slot_for_[s]; + for (uint32_t lb = 0; lb < (uint32_t) row.size(); ++lb) { + if (row[lb] == kInvalidColdIdx) continue; + const uint32_t cold_idx = row[lb]; + row[lb] = kInvalidColdIdx; + cold_pool_free_.push_back(cold_idx); + if (cold_in_use_ > 0) --cold_in_use_; + ++evict_c2drop_total_; // MAD-133 + LLAMA_LOG_INFO("llama_kv_cache_paged::drop_oldest_cold_block: dropped " + "(seq=%u, lblock=%u, cold_idx=%u) — owning seq sees a " + "hole at this block (kernel handles via -INFINITY logit)\n", + s, lb, cold_idx); + return true; + } + } + return false; +} + +// MAD-133: count cold-resident blocks for a seq. +uint32_t llama_kv_cache_paged::n_blocks_cold_for(llama_seq_id seq_id) const { + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max_) return 0; + if ((size_t) seq_id >= cold_slot_for_.size()) return 0; + uint32_t cnt = 0; + for (uint32_t v : cold_slot_for_[seq_id]) { + if (v != kInvalidColdIdx) ++cnt; + } + return cnt; +} + +// MAD-132: idle-priority victim selection for whole-slot preemption. +// Picks the seq with the smallest (oldest) last_active_us, excluding +// any seq in exclude_seqs and any seq that has zero blocks (nothing +// to preempt). Returns -1 if no eligible victim exists. +llama_seq_id llama_kv_cache_paged::pick_preempt_victim( + const std::vector & exclude_seqs) const { + auto excluded = [&](llama_seq_id sid) { + for (auto e : exclude_seqs) if (e == sid) return true; + return false; + }; + + llama_seq_id victim = -1; + uint64_t oldest = UINT64_MAX; + for (uint32_t s = 0; s < n_seq_max_; ++s) { + const llama_seq_id sid = (llama_seq_id) s; + if (excluded(sid)) continue; + if (table_.num_blocks(sid) == 0) continue; // no blocks + if (n_gpu_blocks_for(sid) == 0) continue; // already preempted + const uint64_t last = seq_states_[s].last_active_us; + if (last < oldest) { + oldest = last; + victim = sid; + } + } + return victim; +} + bool llama_kv_cache_paged::evict_lru_warm_to_cold() { if (!cold_enabled() || cold_pool_free_.empty()) return false; @@ -903,6 +1421,15 @@ void llama_kv_cache_paged::prepare_batch_tensors() { } bool llama_kv_cache_paged::apply_ubatch_to_state(const llama_ubatch & ub) { + // MAD-133: snapshot counters for the per-batch tier_event log emitted + // at the end of this method. If any tier movement happened during + // this batch we emit one structured INF line; quiet otherwise. + const uint64_t pre_evict_h2w = evict_h2w_total_; + const uint64_t pre_evict_w2c = evict_w2c_total_; + const uint64_t pre_evict_drop = evict_c2drop_total_; + const uint64_t pre_restore_w2h = restore_w2h_total_; + const uint64_t pre_restore_c2h = restore_c2h_total_; + const uint64_t pre_seq_preempt = seq_preempt_total_; // Per-seq: count tokens, find max position. Sized to n_seq_max_ for // direct indexing — n_seq_max_ is bounded (typical 4–32). std::vector tokens_per_seq(n_seq_max_, 0); @@ -939,6 +1466,7 @@ bool llama_kv_cache_paged::apply_ubatch_to_state(const llama_ubatch & ub) { } // Commit per-seq state. + const uint64_t now = now_us_(); for (uint32_t s = 0; s < n_seq_max_; ++s) { h_q_lens_[s] = (int32_t) tokens_per_seq[s]; if (tokens_per_seq[s] > 0) { @@ -946,6 +1474,9 @@ bool llama_kv_cache_paged::apply_ubatch_to_state(const llama_ubatch & ub) { seq_states_[s].pos_max = max_pos_per_seq[s]; } if (seq_states_[s].pos_min < 0) seq_states_[s].pos_min = 0; + // MAD-132: stamp last-active so preempt-fairness picks + // genuinely idle seqs over actively-batched ones. + seq_states_[s].last_active_us = now; } } @@ -957,6 +1488,40 @@ bool llama_kv_cache_paged::apply_ubatch_to_state(const llama_ubatch & ub) { if (warm_enabled() && !fault_in_warm_blocks_for_batch(ub)) { return false; } + + // MAD-128: CoW any blocks being written this ubatch that are shared + // (refcount > 1 from a prior seq_cp). Must run AFTER fault-in (so + // the block is GPU-resident before we attempt to copy its data) and + // BEFORE prepare_batch_tensors (so the table snapshot reflects the + // CoW'd physicals). No-op when no shared blocks exist. + if (!cow_writes_for_ubatch(ub)) { + return false; + } + + // MAD-133: emit the per-batch tier_event log if any tier movement + // happened during this batch. Quiet for normal flow (no eviction = + // no log line); informative when something moved. + const uint64_t d_h2w = evict_h2w_total_ - pre_evict_h2w; + const uint64_t d_w2c = evict_w2c_total_ - pre_evict_w2c; + const uint64_t d_drop = evict_c2drop_total_ - pre_evict_drop; + const uint64_t d_w2h = restore_w2h_total_ - pre_restore_w2h; + const uint64_t d_c2h = restore_c2h_total_ - pre_restore_c2h; + const uint64_t d_preempt = seq_preempt_total_ - pre_seq_preempt; + if (d_h2w || d_w2c || d_drop || d_w2h || d_c2h || d_preempt) { + LLAMA_LOG_INFO("[mt::tier_event] instance=%s " + "evict_h2w=%llu evict_w2c=%llu evict_drop=%llu " + "restore_w2h=%llu restore_c2h=%llu preempt=%llu " + "pool_free_gpu=%zu pool_free_cpu=%zu cold_in_use=%u\n", + instance_id_.c_str(), + (unsigned long long) d_h2w, + (unsigned long long) d_w2c, + (unsigned long long) d_drop, + (unsigned long long) d_w2h, + (unsigned long long) d_c2h, + (unsigned long long) d_preempt, + pool_.n_free_gpu(), pool_.n_free_cpu(), cold_in_use_); + } + return true; } @@ -1057,6 +1622,7 @@ llama_memory_context_ptr llama_kv_cache_paged::init_update(llama_context * /*lct } void llama_kv_cache_paged::clear(bool /*data*/) { + check_thread_id_(); // Free all blocks back to the pool, reset table + seq states. for (uint32_t s = 0; s < n_seq_max_; ++s) { std::vector freed = table_.clear_seq((llama_seq_id) s); @@ -1067,6 +1633,7 @@ void llama_kv_cache_paged::clear(bool /*data*/) { } pool_.reset(); table_.reset(); + paged_semantic_.clear(); // MAD-125: drop all per-seq fingerprints std::fill(h_block_table_.begin(), h_block_table_.end(), kInvalidBlockTableEntry); std::fill(h_context_lens_.begin(), h_context_lens_.end(), 0); std::fill(h_q_lens_.begin(), h_q_lens_.end(), 0); @@ -1074,6 +1641,7 @@ void llama_kv_cache_paged::clear(bool /*data*/) { } bool llama_kv_cache_paged::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + check_thread_id_(); if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max_) return false; // Match the regular kv_cache convention: p0 < 0 → from 0; p1 < 0 → @@ -1082,6 +1650,7 @@ bool llama_kv_cache_paged::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p // produced position drift between the cache and the slot manager. if (p0 < 0) p0 = 0; if (p1 < 0) p1 = std::numeric_limits::max(); + if (p1 <= p0) return true; // empty range const llama_pos cur_max = seq_states_[seq_id].pos_max; @@ -1092,48 +1661,173 @@ bool llama_kv_cache_paged::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p if (bid != mt::kInvalidBlockId) pool_.free_block(bid); } seq_states_[seq_id] = seq_state{}; - LLAMA_LOG_DEBUG("llama_kv_cache_paged::seq_rm: whole-seq wipe seq=%d, freed %zu blocks\n", - seq_id, freed.size()); + // MAD-125: fingerprints from the prior task can't match the new + // one's K/V — wipe them so semantic restore doesn't fault in stale + // blocks. Mirrors the wrapper's whole-seq-wipe behavior. + const size_t n_finger = paged_semantic_.size(seq_id); + paged_semantic_.remove_seq(seq_id); + LLAMA_LOG_DEBUG("llama_kv_cache_paged::seq_rm: whole-seq wipe seq=%d, freed %zu blocks, dropped %zu paged-block fingerprints\n", + seq_id, freed.size(), n_finger); return true; } - // Partial wipe: tail truncation only for v1 (carving out middle - // ranges requires more bookkeeping; rare in practice). - if (p0 > cur_max) return true; // nothing to do - // Middle wipes (p1 < cur_max+1 with p0 > 0) aren't supported in v1. - // The server uses these for partial speculative-decode rollback; - // until we add support, treat as tail truncation from p0 (drops - // some valid positions but doesn't crash). - GGML_UNUSED(p1); - - // Compute new pos_max after truncation. - const llama_pos new_max = p0 - 1; - - // Drop blocks whose entire range is past p0. - const uint32_t keep_blocks = (uint32_t)(new_max + 1 + (llama_pos) block_size_ - 1) / block_size_; - if (keep_blocks > table_.num_blocks(seq_id)) return true; // nothing to drop - - // Walk blocks past keep_blocks and free. - while (table_.num_blocks(seq_id) > keep_blocks) { - std::vector tmp = table_.clear_seq(seq_id); // CAUTION: clears all - // Re-append the kept ones. - for (uint32_t i = 0; i < keep_blocks && i < tmp.size(); ++i) { - table_.append_block(seq_id, tmp[i]); - } - for (uint32_t i = keep_blocks; i < tmp.size(); ++i) { - if (tmp[i] != mt::kInvalidBlockId) pool_.free_block(tmp[i]); + if (p0 > cur_max) return true; // nothing in range + + // MAD-128: block-aligned partial wipe. Handles both tail truncate + // (p1 > cur_max) and middle wipe (p1 <= cur_max) uniformly. + // + // For each block whose [b_start, b_end) intersects [p0, p1): + // - Wholly covered (p0 <= b_start AND p1 >= b_end): free the + // physical, mark table entry as kInvalidBlockId, drop fingerprint. + // The mt_paged_attention_kernel handles invalid entries by + // contributing -INFINITY to attention logits → zero weight after + // softmax → no contribution. So freed blocks read as "no + // attention," not as garbage. + // - Partially overlapped (p0 inside the block OR p1 inside the + // block): leave the block alone. The unwiped slots within keep + // stale K/V that the kernel WILL read. Caller should round their + // range to block boundaries for clean wipes; we log a clear + // warning so sub-block partial wipes don't silently corrupt. + // + // Sub-block per-slot zeroing requires either a per-block valid-bitmask + // consulted by the kernel mask, or a layout-aware per-slot zero + // primitive. Both are deferred (separate ticket if a real consumer + // needs sub-block precision). + const llama_pos p1_clamped = std::min(p1, (llama_pos)(cur_max + 1)); + const uint32_t bsize = block_size_; + const uint32_t n_blocks_seq = table_.num_blocks(seq_id); + + uint32_t blocks_freed = 0; + uint32_t blocks_partial_skipped = 0; + + for (uint32_t lblock = 0; lblock < n_blocks_seq; ++lblock) { + const llama_pos b_start = (llama_pos) lblock * (llama_pos) bsize; + const llama_pos b_end = b_start + (llama_pos) bsize; + + if (b_end <= p0 || b_start >= p1_clamped) continue; // no overlap + + const bool wholly_covered = (p0 <= b_start) && (p1_clamped >= b_end); + if (wholly_covered) { + const uint32_t physical = table_.get_physical(seq_id, lblock); + if (physical != mt::kInvalidBlockId) { + pool_.free_block(physical); + table_.swap_block(seq_id, lblock, mt::kInvalidBlockId); + paged_semantic_.remove_block(seq_id, lblock); + ++blocks_freed; + } + } else { + ++blocks_partial_skipped; } - break; } - seq_states_[seq_id].pos_max = new_max; + if (blocks_partial_skipped > 0) { + LLAMA_LOG_WARN("llama_kv_cache_paged::seq_rm: range [%d,%d) is not " + "block-aligned (block_size=%u). %u block(s) wholly " + "freed; %u block(s) partially overlapped — those keep " + "stale K/V in the unwiped slots and the kernel will " + "attend to them. Round caller's range to block " + "boundaries for clean wipes.\n", + p0, p1, bsize, blocks_freed, blocks_partial_skipped); + } + + // Update pos_max iff the wipe touches the tail (p1 covers past cur_max). + // For middle wipes (p1_clamped <= cur_max), pos_max stays put — holes + // in the block table represent the wiped middle. For tail truncate, + // shrink pos_max to the position just before the wipe started. + if (p1 > cur_max) { + seq_states_[seq_id].pos_max = (llama_pos)(p0 - 1); + if (seq_states_[seq_id].pos_max < 0) seq_states_[seq_id].pos_min = -1; + } + + LLAMA_LOG_DEBUG("llama_kv_cache_paged::seq_rm: seq=%d range=[%d,%d) " + "blocks_freed=%u blocks_partial=%u pos_max=%d\n", + seq_id, p0, p1, blocks_freed, blocks_partial_skipped, + seq_states_[seq_id].pos_max); return true; } -void llama_kv_cache_paged::seq_cp(llama_seq_id /*src*/, llama_seq_id /*dst*/, - llama_pos /*p0*/, llama_pos /*p1*/) { - // CoW between sequences — not supported in v1. - LLAMA_LOG_WARN("llama_kv_cache_paged::seq_cp: not implemented in v1 — no-op\n"); +void llama_kv_cache_paged::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, + llama_pos p0, llama_pos p1) { + check_thread_id_(); + // MAD-128: block-aligned CoW. Wholly-covered blocks of src in [p0, p1) + // are SHARED into dst's table (refcount bumped in BlockPool). Future + // writes that target a shared block trigger CoW in + // cow_writes_for_ubatch (called from apply_ubatch_to_state). + // + // Sub-block partials are not shared (would require partial-block CoW + // which doesn't exist) — they're skipped with a warning, mirroring + // seq_rm's sub-block behavior. + if (seq_id_src < 0 || (uint32_t) seq_id_src >= n_seq_max_) return; + if (seq_id_dst < 0 || (uint32_t) seq_id_dst >= n_seq_max_) return; + if (seq_id_src == seq_id_dst) return; + if (p0 < 0) p0 = 0; + if (p1 < 0) p1 = std::numeric_limits::max(); + if (p1 <= p0) return; + + const llama_pos src_max = seq_states_[seq_id_src].pos_max; + if (p0 > src_max) return; // src has nothing in range + const llama_pos p1_clamped = std::min(p1, (llama_pos)(src_max + 1)); + + // Wipe dst's existing range first so its old blocks are properly + // refcount-released (recursive call uses the new partial seq_rm). + seq_rm(seq_id_dst, p0, p1_clamped); + + const uint32_t bsize = block_size_; + const uint32_t n_blocks_src = table_.num_blocks(seq_id_src); + + uint32_t blocks_shared = 0; + uint32_t blocks_partial_skipped = 0; + + for (uint32_t lblock = 0; lblock < n_blocks_src; ++lblock) { + const llama_pos b_start = (llama_pos) lblock * (llama_pos) bsize; + const llama_pos b_end = b_start + (llama_pos) bsize; + + if (b_end <= p0 || b_start >= p1_clamped) continue; + const bool wholly_covered = (p0 <= b_start) && (p1_clamped >= b_end); + if (!wholly_covered) { + ++blocks_partial_skipped; + continue; + } + + const uint32_t physical = table_.get_physical(seq_id_src, lblock); + if (physical == mt::kInvalidBlockId) continue; // hole in src + + // Bump refcount on src's physical and install in dst's table. + pool_.bump_ref(physical); + + // Pad dst's table with kInvalidBlockId for any logical gaps + // before this block, then either swap into an existing slot or + // append. + while (table_.num_blocks(seq_id_dst) < lblock) { + table_.append_block(seq_id_dst, mt::kInvalidBlockId); + } + if (lblock < table_.num_blocks(seq_id_dst)) { + table_.swap_block(seq_id_dst, lblock, physical); + } else { + table_.append_block(seq_id_dst, physical); + } + + ++blocks_shared; + } + + if (blocks_partial_skipped > 0) { + LLAMA_LOG_WARN("llama_kv_cache_paged::seq_cp: range [%d,%d) is not " + "block-aligned (block_size=%u). %u block(s) shared via " + "CoW; %u block(s) partially overlapped — those are NOT " + "shared. Round caller's range to block boundaries.\n", + p0, p1, bsize, blocks_shared, blocks_partial_skipped); + } + + // Update dst's pos_max if the copy extended its tail. + if (p1_clamped - 1 > seq_states_[seq_id_dst].pos_max) { + seq_states_[seq_id_dst].pos_max = p1_clamped - 1; + if (seq_states_[seq_id_dst].pos_min < 0) seq_states_[seq_id_dst].pos_min = 0; + } + + LLAMA_LOG_DEBUG("llama_kv_cache_paged::seq_cp: src=%d dst=%d range=[%d,%d) " + "blocks_shared=%u blocks_partial=%u\n", + seq_id_src, seq_id_dst, p0, p1, + blocks_shared, blocks_partial_skipped); } void llama_kv_cache_paged::seq_keep(llama_seq_id /*seq_id*/) { @@ -1168,15 +1862,460 @@ std::map llama_kv_cache_paged::memory_breakd return out; } -void llama_kv_cache_paged::state_write(llama_io_write_i & /*io*/, llama_seq_id /*seq_id*/, +// MAD-130: cold-tier sidecar persistence. +// +// File format (CIDX v1) at `${cold_path_}/paged/index.bin`: +// uint32 magic = 0x58444943 ("CIDX") +// uint32 version = 1 +// uint32 n_layers (sanity check vs current cache) +// uint32 n_cold_blocks (sanity check vs current cache) +// uint32 n_in_use +// For each in-use entry: +// int32 seq_id +// uint32 lblock +// uint32 cold_idx +// +// Save: called explicitly via save_cold_index_sidecar() (server +// shutdown handler or /slots/save). Atomic via write-to-tmp + rename. +// +// Load: called from ctor when cold_resume=true. On any error (missing +// file, bad magic, mismatched config), returns false and the caller +// proceeds with an empty cold index. + +namespace { +constexpr uint32_t kColdIndexMagic = 0x58444943; // "CIDX" +constexpr uint32_t kColdIndexVersion = 1; +} + +bool llama_kv_cache_paged::save_cold_index_sidecar() const { + if (!cold_enabled() || cold_path_.empty()) return true; // no-op + + const std::string subdir = cold_path_ + "/paged/instance-" + instance_id_; + const std::string sidecar = subdir + "/index.bin"; + const std::string tmp_path = sidecar + ".tmp"; + + std::ofstream f(tmp_path, std::ios::binary | std::ios::trunc); + if (!f) { + LLAMA_LOG_WARN("llama_kv_cache_paged::save_cold_index_sidecar: open(%s) failed\n", + tmp_path.c_str()); + return false; + } + + auto write_u32 = [&](uint32_t v) { f.write((const char *) &v, sizeof(v)); }; + auto write_i32 = [&](int32_t v) { f.write((const char *) &v, sizeof(v)); }; + + write_u32(kColdIndexMagic); + write_u32(kColdIndexVersion); + const uint32_t n_layers = (uint32_t) layers_.size(); + write_u32(n_layers); + write_u32(n_cold_blocks_); + write_u32(cold_in_use_); + + // Walk cold_slot_for_ and write each in-use (seq, lblock, cold_idx). + uint32_t written = 0; + for (uint32_t s = 0; s < (uint32_t) cold_slot_for_.size(); ++s) { + const auto & row = cold_slot_for_[s]; + for (uint32_t lb = 0; lb < (uint32_t) row.size(); ++lb) { + if (row[lb] == kInvalidColdIdx) continue; + write_i32((int32_t) s); + write_u32(lb); + write_u32(row[lb]); + ++written; + } + } + f.close(); + if (written != cold_in_use_) { + LLAMA_LOG_WARN("llama_kv_cache_paged::save_cold_index_sidecar: counted %u in-use entries " + "but cold_in_use_ tracker says %u — sidecar may be inconsistent\n", + written, cold_in_use_); + } + if (::rename(tmp_path.c_str(), sidecar.c_str()) != 0) { + LLAMA_LOG_WARN("llama_kv_cache_paged::save_cold_index_sidecar: rename %s -> %s failed\n", + tmp_path.c_str(), sidecar.c_str()); + return false; + } + LLAMA_LOG_INFO("llama_kv_cache_paged::save_cold_index_sidecar: wrote %u entries to %s\n", + written, sidecar.c_str()); + return true; +} + +bool llama_kv_cache_paged::load_cold_index_sidecar_(const std::string & path) { + std::ifstream f(path, std::ios::binary); + if (!f) return false; + + auto read_u32 = [&](uint32_t & v) -> bool { + f.read((char *) &v, sizeof(v)); + return (bool) f; + }; + auto read_i32 = [&](int32_t & v) -> bool { + f.read((char *) &v, sizeof(v)); + return (bool) f; + }; + + uint32_t magic = 0, version = 0, saved_n_layers = 0, saved_n_cold = 0, saved_n_in_use = 0; + if (!read_u32(magic) || !read_u32(version)) return false; + if (magic != kColdIndexMagic || version != kColdIndexVersion) { + LLAMA_LOG_WARN("load_cold_index_sidecar: bad header (magic=%08x ver=%u)\n", magic, version); + return false; + } + if (!read_u32(saved_n_layers) || !read_u32(saved_n_cold) || !read_u32(saved_n_in_use)) { + return false; + } + if (saved_n_layers != (uint32_t) layers_.size() || saved_n_cold != n_cold_blocks_) { + LLAMA_LOG_WARN("load_cold_index_sidecar: config mismatch (saved_n_layers=%u current=%zu, " + "saved_n_cold=%u current=%u) — rejecting sidecar\n", + saved_n_layers, layers_.size(), saved_n_cold, n_cold_blocks_); + return false; + } + + // Walk in-use entries. Mark each cold slot in cold_slot_for_, and + // rebuild cold_pool_free_ minus the in-use slots. + std::vector in_use(n_cold_blocks_, false); + for (uint32_t i = 0; i < saved_n_in_use; ++i) { + int32_t sid = 0; + uint32_t lb = 0, cold_idx = 0; + if (!read_i32(sid) || !read_u32(lb) || !read_u32(cold_idx)) return false; + if (sid < 0 || (uint32_t) sid >= n_seq_max_) continue; + if (cold_idx >= n_cold_blocks_) continue; + mark_cold(cold_slot_for_, (llama_seq_id) sid, lb, cold_idx); + in_use[cold_idx] = true; + } + + // Rebuild free stack — only push slots NOT in use. + cold_pool_free_.clear(); + cold_pool_free_.reserve(n_cold_blocks_); + for (uint32_t i = n_cold_blocks_; i > 0; --i) { + if (!in_use[i - 1]) cold_pool_free_.push_back(i - 1); + } + cold_in_use_ = saved_n_in_use; + + LLAMA_LOG_INFO("llama_kv_cache_paged::load_cold_index_sidecar: loaded %u in-use cold " + "entries (free pool: %zu)\n", + saved_n_in_use, cold_pool_free_.size()); + return true; +} + +// MAD-130: state persistence for the paged cache. +// +// On-disk format (PAGS v1): +// +// HEADER (76 bytes): +// magic uint32 = 0x53474150 ("PAGS" little-endian) +// version uint32 = 1 +// n_seq_max uint32 (config: must match on read) +// block_size uint32 (config: must match on read) +// n_blocks_total uint32 (config: GPU pool size; must match on read) +// n_warm_blocks uint32 (config: must match on read) +// n_cold_blocks uint32 (config: must match on read) +// n_layers uint32 (config: total, including filtered-out) +// k_bytes_per_block uint64 (config sanity) +// v_bytes_per_block uint64 (config sanity) +// type_k uint32 (ggml_type) +// type_v uint32 (ggml_type) +// flags uint32 (reserved) +// +// PER-SEQ SECTION (one per saved seq): +// seq_present uint8 (0 = end of seq list; 1 = entry follows) +// seq_id int32 +// pos_min int32 +// pos_max int32 +// num_blocks uint32 +// For each lblock in [0, num_blocks): +// tier_origin uint8 (0=hole, 1=hot, 2=warm, 3=cold) +// if hot or warm: +// For each restorable layer (filter applied at construction): +// K bytes (k_bytes_per_block) +// V bytes (v_bytes_per_block) +// if cold: +// cold_idx uint32 (slot in the per-layer .bin files) +// +// FINGERPRINT SECTION: +// n_fingerprints uint32 +// For each: seq_id (int32), lblock (uint32), tier (uint8), +// embedding_dim (uint32), floats[embedding_dim] +// +// Validation: header config fields must EXACTLY match the runtime +// cache (n_blocks_total, block_size, type_k, type_v, n_layers, +// k/v_bytes_per_block). Mismatch → throw runtime_error with a clear +// message; caller's /slots/restore returns the error to the client. + +namespace { +constexpr uint32_t kPagedStateMagic = 0x53474150; // "PAGS" +constexpr uint32_t kPagedStateVersion = 1; + +enum class PagedTierTag : uint8_t { + Hole = 0, + Hot = 1, + Warm = 2, + Cold = 3, +}; +} // namespace + +void llama_kv_cache_paged::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags /*flags*/) const { - // State persistence — punted to a later phase. - LLAMA_LOG_WARN("llama_kv_cache_paged::state_write: not implemented — state will not persist\n"); + // ── Header ── + io.write(&kPagedStateMagic, sizeof(kPagedStateMagic)); + io.write(&kPagedStateVersion, sizeof(kPagedStateVersion)); + io.write(&n_seq_max_, sizeof(n_seq_max_)); + io.write(&block_size_, sizeof(block_size_)); + io.write(&n_blocks_total_, sizeof(n_blocks_total_)); + io.write(&n_warm_blocks_, sizeof(n_warm_blocks_)); + io.write(&n_cold_blocks_, sizeof(n_cold_blocks_)); + const uint32_t n_layers = (uint32_t) layers_.size(); + io.write(&n_layers, sizeof(n_layers)); + const uint64_t k_bpb = (uint64_t) k_bytes_per_block_; + const uint64_t v_bpb = (uint64_t) v_bytes_per_block_; + io.write(&k_bpb, sizeof(k_bpb)); + io.write(&v_bpb, sizeof(v_bpb)); + const uint32_t type_k_u = (uint32_t) type_k_; + const uint32_t type_v_u = (uint32_t) type_v_; + io.write(&type_k_u, sizeof(type_k_u)); + io.write(&type_v_u, sizeof(type_v_u)); + const uint32_t flags_reserved = 0; + io.write(&flags_reserved, sizeof(flags_reserved)); + + // ── Per-seq section ── + auto write_one_seq = [&](llama_seq_id sid) { + if (sid < 0 || (uint32_t) sid >= n_seq_max_) return; + const uint32_t n_blocks_seq = table_.num_blocks(sid); + // Skip empty seqs (no blocks). + if (n_blocks_seq == 0) return; + + const uint8_t present = 1; + io.write(&present, sizeof(present)); + const int32_t sid_i32 = (int32_t) sid; + io.write(&sid_i32, sizeof(sid_i32)); + const int32_t pos_min = (int32_t) seq_states_[sid].pos_min; + const int32_t pos_max = (int32_t) seq_states_[sid].pos_max; + io.write(&pos_min, sizeof(pos_min)); + io.write(&pos_max, sizeof(pos_max)); + io.write(&n_blocks_seq, sizeof(n_blocks_seq)); + + std::vector kbuf(k_bytes_per_block_); + std::vector vbuf(v_bytes_per_block_); + + for (uint32_t lb = 0; lb < n_blocks_seq; ++lb) { + const uint32_t physical = table_.get_physical(sid, lb); + PagedTierTag tag = PagedTierTag::Hole; + if (physical != mt::kInvalidBlockId) { + tag = pool_.is_gpu(physical) ? PagedTierTag::Hot : PagedTierTag::Warm; + } else if (cold_enabled() && cold_slot(cold_slot_for_, sid, lb) != kInvalidColdIdx) { + tag = PagedTierTag::Cold; + } + const uint8_t tag_u8 = (uint8_t) tag; + io.write(&tag_u8, sizeof(tag_u8)); + + if (tag == PagedTierTag::Hot) { + const size_t k_off_gpu = (size_t) physical * k_bytes_per_block_; + const size_t v_off_gpu = (size_t) physical * v_bytes_per_block_; + for (uint32_t il = 0; il < layers_.size(); ++il) { + const auto & layer = layers_[il]; + if (!layer.k) continue; + ggml_backend_tensor_get(layer.k, kbuf.data(), k_off_gpu, k_bytes_per_block_); + ggml_backend_tensor_get(layer.v, vbuf.data(), v_off_gpu, v_bytes_per_block_); + io.write(kbuf.data(), k_bytes_per_block_); + io.write(vbuf.data(), v_bytes_per_block_); + } + } else if (tag == PagedTierTag::Warm) { + const uint32_t cpu_idx = physical - n_blocks_total_; + const size_t k_off_cpu = (size_t) cpu_idx * k_bytes_per_block_; + const size_t v_off_cpu = (size_t) cpu_idx * v_bytes_per_block_; + for (uint32_t il = 0; il < layers_.size(); ++il) { + if (!layers_[il].k) continue; + io.write(warm_k_[il].data() + k_off_cpu, k_bytes_per_block_); + io.write(warm_v_[il].data() + v_off_cpu, v_bytes_per_block_); + } + } else if (tag == PagedTierTag::Cold) { + const uint32_t cold_idx = cold_slot(cold_slot_for_, sid, lb); + io.write(&cold_idx, sizeof(cold_idx)); + } + // Hole: no extra bytes. + } + }; + + if (seq_id < 0) { + for (uint32_t s = 0; s < n_seq_max_; ++s) write_one_seq((llama_seq_id) s); + } else { + write_one_seq(seq_id); + } + const uint8_t end_marker = 0; + io.write(&end_marker, sizeof(end_marker)); + + // ── Fingerprint section ── + // BlockSemanticIndex doesn't expose iteration; for now we only + // serialize fingerprint COUNT here, with the actual data left to + // BlockSemanticIndex::save_to_disk (separate sidecar file). The + // sidecar approach matches the legacy SemanticIndex pattern and + // keeps state_write self-contained. + const uint32_t n_fingerprints_total = (uint32_t) paged_semantic_.size(); + io.write(&n_fingerprints_total, sizeof(n_fingerprints_total)); + // Fingerprints themselves are written via paged_semantic.save_to_disk, + // called by the server's /slots/save handler alongside this state_write. + + LLAMA_LOG_INFO("llama_kv_cache_paged::state_write: wrote %zu bytes " + "(seq_id=%d, n_fingerprints=%u)\n", + io.n_bytes(), seq_id, n_fingerprints_total); } -void llama_kv_cache_paged::state_read(llama_io_read_i & /*io*/, llama_seq_id /*seq_id*/, +void llama_kv_cache_paged::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags /*flags*/) { - LLAMA_LOG_WARN("llama_kv_cache_paged::state_read: not implemented — state will not load\n"); + check_thread_id_(); + // ── Header validation ── + uint32_t magic, version; + io.read(&magic, sizeof(magic)); + io.read(&version, sizeof(version)); + if (magic != kPagedStateMagic) { + throw std::runtime_error("llama_kv_cache_paged::state_read: bad magic — not a PAGS state file"); + } + if (version != kPagedStateVersion) { + throw std::runtime_error("llama_kv_cache_paged::state_read: unsupported version " + + std::to_string(version)); + } + + uint32_t saved_n_seq_max, saved_block_size, saved_n_blocks_total; + uint32_t saved_n_warm, saved_n_cold, saved_n_layers; + uint64_t saved_k_bpb, saved_v_bpb; + uint32_t saved_type_k, saved_type_v, saved_flags; + io.read(&saved_n_seq_max, sizeof(saved_n_seq_max)); + io.read(&saved_block_size, sizeof(saved_block_size)); + io.read(&saved_n_blocks_total, sizeof(saved_n_blocks_total)); + io.read(&saved_n_warm, sizeof(saved_n_warm)); + io.read(&saved_n_cold, sizeof(saved_n_cold)); + io.read(&saved_n_layers, sizeof(saved_n_layers)); + io.read(&saved_k_bpb, sizeof(saved_k_bpb)); + io.read(&saved_v_bpb, sizeof(saved_v_bpb)); + io.read(&saved_type_k, sizeof(saved_type_k)); + io.read(&saved_type_v, sizeof(saved_type_v)); + io.read(&saved_flags, sizeof(saved_flags)); + + auto require = [](bool cond, const char * what, auto saved, auto current) { + if (!cond) { + std::string msg = std::string("llama_kv_cache_paged::state_read: ") + what + + " mismatch — saved=" + std::to_string(saved) + + " current=" + std::to_string(current) + + " (state file from a different cache config)"; + throw std::runtime_error(msg); + } + }; + require(saved_n_seq_max == n_seq_max_, "n_seq_max", saved_n_seq_max, n_seq_max_); + require(saved_block_size == block_size_, "block_size", saved_block_size, block_size_); + require(saved_n_blocks_total == n_blocks_total_, "n_blocks_total", saved_n_blocks_total, n_blocks_total_); + require(saved_n_warm == n_warm_blocks_, "n_warm_blocks", saved_n_warm, n_warm_blocks_); + require(saved_n_cold == n_cold_blocks_, "n_cold_blocks", saved_n_cold, n_cold_blocks_); + require(saved_n_layers == (uint32_t) layers_.size(), "n_layers", saved_n_layers, (uint32_t) layers_.size()); + require(saved_k_bpb == (uint64_t) k_bytes_per_block_, "k_bytes_per_block", saved_k_bpb, (uint64_t) k_bytes_per_block_); + require(saved_v_bpb == (uint64_t) v_bytes_per_block_, "v_bytes_per_block", saved_v_bpb, (uint64_t) v_bytes_per_block_); + require(saved_type_k == (uint32_t) type_k_, "type_k", saved_type_k, (uint32_t) type_k_); + require(saved_type_v == (uint32_t) type_v_, "type_v", saved_type_v, (uint32_t) type_v_); + + std::vector kbuf(k_bytes_per_block_); + std::vector vbuf(v_bytes_per_block_); + + // ── Per-seq sections ── + uint32_t n_seqs_loaded = 0; + uint32_t n_blocks_loaded = 0; + while (true) { + uint8_t present = 0; + io.read(&present, sizeof(present)); + if (present == 0) break; + + int32_t saved_sid = 0; + io.read(&saved_sid, sizeof(saved_sid)); + // If caller asked for a specific seq_id, remap saved data into + // that slot (allows /slots/restore into a different slot ID). + const llama_seq_id load_sid = (seq_id < 0) ? (llama_seq_id) saved_sid : seq_id; + if (load_sid < 0 || (uint32_t) load_sid >= n_seq_max_) { + throw std::runtime_error("llama_kv_cache_paged::state_read: seq_id out of range"); + } + + // Wipe the existing seq before reading. Honors refcount via + // free_block (shared blocks decrement, free at 0). + seq_rm(load_sid, 0, std::numeric_limits::max()); + + int32_t pos_min = 0, pos_max = 0; + uint32_t n_blocks_seq = 0; + io.read(&pos_min, sizeof(pos_min)); + io.read(&pos_max, sizeof(pos_max)); + io.read(&n_blocks_seq, sizeof(n_blocks_seq)); + + for (uint32_t lb = 0; lb < n_blocks_seq; ++lb) { + uint8_t tag_u8 = 0; + io.read(&tag_u8, sizeof(tag_u8)); + const PagedTierTag tag = (PagedTierTag) tag_u8; + + if (tag == PagedTierTag::Hole) { + table_.append_block(load_sid, mt::kInvalidBlockId); + continue; + } + + if (tag == PagedTierTag::Hot) { + const uint32_t phys = pool_.alloc_gpu(); + if (phys == mt::kInvalidBlockId) { + throw std::runtime_error("llama_kv_cache_paged::state_read: GPU pool exhausted " + "while restoring hot block"); + } + const size_t k_off_gpu = (size_t) phys * k_bytes_per_block_; + const size_t v_off_gpu = (size_t) phys * v_bytes_per_block_; + for (uint32_t il = 0; il < layers_.size(); ++il) { + if (!layers_[il].k) continue; + io.read(kbuf.data(), k_bytes_per_block_); + io.read(vbuf.data(), v_bytes_per_block_); + ggml_backend_tensor_set(layers_[il].k, kbuf.data(), k_off_gpu, k_bytes_per_block_); + ggml_backend_tensor_set(layers_[il].v, vbuf.data(), v_off_gpu, v_bytes_per_block_); + } + table_.append_block(load_sid, phys); + } else if (tag == PagedTierTag::Warm) { + if (!warm_enabled()) { + throw std::runtime_error("llama_kv_cache_paged::state_read: warm tier disabled but " + "saved state has warm blocks"); + } + const uint32_t phys = pool_.alloc_cpu(); + if (phys == mt::kInvalidBlockId) { + throw std::runtime_error("llama_kv_cache_paged::state_read: CPU pool exhausted " + "while restoring warm block"); + } + const uint32_t cpu_idx = phys - n_blocks_total_; + const size_t k_off_cpu = (size_t) cpu_idx * k_bytes_per_block_; + const size_t v_off_cpu = (size_t) cpu_idx * v_bytes_per_block_; + for (uint32_t il = 0; il < layers_.size(); ++il) { + if (!layers_[il].k) continue; + io.read(warm_k_[il].data() + k_off_cpu, k_bytes_per_block_); + io.read(warm_v_[il].data() + v_off_cpu, v_bytes_per_block_); + } + table_.append_block(load_sid, phys); + } else if (tag == PagedTierTag::Cold) { + if (!cold_enabled()) { + throw std::runtime_error("llama_kv_cache_paged::state_read: cold tier disabled but " + "saved state has cold blocks"); + } + uint32_t cold_idx = 0; + io.read(&cold_idx, sizeof(cold_idx)); + if (cold_idx >= n_cold_blocks_) { + throw std::runtime_error("llama_kv_cache_paged::state_read: cold_idx out of range"); + } + // Mark as cold-resident; the actual bytes live in the + // cold-tier files (which require --kv-tier-cold-resume + // to survive a restart — separate part of MAD-130). + mark_cold(cold_slot_for_, load_sid, lb, cold_idx); + table_.append_block(load_sid, mt::kInvalidBlockId); + ++cold_in_use_; + } + ++n_blocks_loaded; + } + + seq_states_[load_sid].pos_min = pos_min; + seq_states_[load_sid].pos_max = pos_max; + ++n_seqs_loaded; + } + + // ── Fingerprint count (data lives in sidecar) ── + uint32_t n_fingerprints_meta = 0; + io.read(&n_fingerprints_meta, sizeof(n_fingerprints_meta)); + + LLAMA_LOG_INFO("llama_kv_cache_paged::state_read: loaded %u seq(s), %u block(s); " + "fingerprint sidecar reports %u entries (load via " + "BlockSemanticIndex::load_from_disk).\n", + n_seqs_loaded, n_blocks_loaded, n_fingerprints_meta); } // ─── llama_kv_cache_paged_context ─── diff --git a/src/llama-kv-cache-paged.h b/src/llama-kv-cache-paged.h index 33a7556280d7..e1a70b49ff2f 100644 --- a/src/llama-kv-cache-paged.h +++ b/src/llama-kv-cache-paged.h @@ -34,6 +34,38 @@ // appropriate logical block list. The earlier defensive assert is // gone; the cache works correctly with --parallel N for any N up to // n_seq_max_. +// +// THREADING CONTRACT (MAD-132 / Epic A4): +// +// This cache is SINGLE-THREADED. All mutator methods MUST be called +// from one thread (the server's update_slots main loop). The internal +// BlockPool, BlockTable, BlockSemanticIndex, cold_pool_free_, and +// per-seq state vectors have NO internal locking. +// +// Why this is safe today: the server's update_slots runs sequentially. +// --parallel N (multi-seq) processes N sequences inside one ubatch +// inside one main-thread iteration; the GPU runs them in parallel via +// the kernel's grid-y dimension over sequences, but the cache mutations +// that schedule the kernel all happen on one thread. +// +// What violates this: any future change that adds a worker thread +// touching the cache. Examples: +// - bge-small embedding on a CPU worker parallel to GPU compute +// - Async cold-tier I/O thread +// - Multi-threaded HTTP handler that touches the cache directly +// +// Each of these would silently corrupt the cache (BlockPool::alloc_gpu +// would race; two threads could pop the same block ID; kernel reads +// garbage at attention time). DEBUG builds catch the violation via +// captured_thread_id_ + check_thread_id_() — first call captures the +// thread, subsequent calls assert match. Release builds skip the check. +// +// To add real async support: either (a) a single-purpose worker queue +// where the worker enqueues "fingerprint this text" / "embed this query" +// and the main thread drains at safe points, OR (b) a redesign with +// fine-grained locking (mutex per pool, per-seq state shards). DO NOT +// add casual mutation from a second thread — the existing single-thread +// invariant is load-bearing for correctness. #include "llama-batch.h" #include "llama-graph.h" @@ -41,9 +73,12 @@ #include "memory-tier/mt-block-pool.h" #include "memory-tier/mt-block-table.h" +#include "memory-tier/mt-semantic.h" +#include #include #include +#include #include struct llama_model; @@ -89,6 +124,22 @@ class llama_kv_cache_paged : public llama_memory_i { // created. Ignored when n_cold_blocks == 0. The "/paged" // subdir is created on demand. std::string ssd_path = std::string(), + // MAD-130: when true, skip the O_TRUNC on cold-tier files + // and load the in-memory cold index from the sidecar file + // at `${ssd_path}/paged/instance-${INSTANCE_ID}/index.bin`. + // Lets the server resume cold-tier contents across a clean + // restart. Default false (legacy behavior: fresh truncation). + bool cold_resume = false, + // MAD-131: per-instance ID. Cold-tier files live under + // `${ssd_path}/paged/instance-${INSTANCE_ID}/`. Empty → + // use the process pid as a string. Required to allow + // multiple llama-server processes to share one ssd_path + // without colliding on cold-tier files. + std::string instance_id = std::string(), + // MAD-131: cap cold pool to this many MiB total (K+V across + // all attn layers). 0 = no cap (size from cold percentage). + // Lets operators bound SSD wear per instance. + uint32_t cold_budget_mb = 0, // Optional per-layer filter. Returns true for layers that should // get K/V allocation in the paged pool. Recurrent layers in // hybrid models should be filtered OUT (their state lives in @@ -255,17 +306,134 @@ class llama_kv_cache_paged : public llama_memory_i { uint32_t n_cold_blocks() const { return n_cold_blocks_; } bool cold_enabled() const { return n_cold_blocks_ > 0; } + // ─── MAD-133: tier-movement counters (monotonic since startup) ─── + // + // Read by the /metrics endpoint (paged_* keys). All uint64. No locking needed + // — single-thread mutator contract (Epic A4 / MAD-132). Reset on + // clear() + state_read. + uint64_t evict_h2w_total() const { return evict_h2w_total_; } + uint64_t evict_w2c_total() const { return evict_w2c_total_; } + uint64_t evict_c2drop_total() const { return evict_c2drop_total_; } + uint64_t restore_w2h_total() const { return restore_w2h_total_; } + uint64_t restore_c2h_total() const { return restore_c2h_total_; } + uint64_t seq_preempt_total() const { return seq_preempt_total_; } + uint64_t seq_restore_total() const { return seq_restore_total_; } + uint64_t semantic_attempts_total() const { return semantic_attempts_total_; } + uint64_t semantic_hits_total() const { return semantic_hits_total_; } + uint64_t semantic_blocks_restored_total() const { return semantic_blocks_restored_total_; } + + // ─── MAD-133: per-seq tier-residency accessors ─── + // + // n_blocks_hot_for / n_blocks_warm_for already exist (lines 217-218 + // — used by can_admit). New: cold + fingerprint counts for /slots. + uint32_t n_blocks_cold_for(llama_seq_id seq_id) const; + size_t n_fingerprints_for_seq(llama_seq_id seq_id) const { + return paged_semantic_.size(seq_id); + } + + // Per-instance ID accessor (set in ctor; immutable after). Used by + // /metrics + /slots to label paged_* output. + const std::string & instance_id() const { return instance_id_; } + + // MAD-132: drop the OLDEST cold block from cold_pool_free_ by + // unmapping the (seq, lblock) it was assigned to and freeing the + // cold_idx. Used as the final escalation when all three tiers are + // full (hot full, warm full, cold full): rather than refuse the + // request, drop the most-stale cold block to make room for the + // hot→warm→cold spillage. The dropped block's data is gone — the + // owning seq's table entry stays kInvalidBlockId; future reads + // contribute -INFINITY logit (= zero attention contribution) per + // the same kernel mechanism that handles middle-wipe holes. + // Returns true if a block was dropped, false if cold is empty. + bool drop_oldest_cold_block(); + + // MAD-132: pick a victim seq for whole-slot preemption using + // idle-priority fairness — the seq whose last apply_ubatch_to_state + // was longest ago is preferred. exclude_seqs lists seqs that must + // NOT be picked (e.g. the candidate trying to be admitted right + // now, or seqs already in the active batch). Returns the chosen + // seq_id (>= 0) on success, -1 if no eligible victim exists. + // Use evict_seq_to_warm() to actually carry out the preemption + // once the policy chose a target. + llama_seq_id pick_preempt_victim(const std::vector & exclude_seqs) const; + + // ─── MAD-125: BGE-small semantic prefetch ─── + // + // The cache holds an optional per-(seq, lblock) fingerprint store. + // Server populates it at backup time (one BGE embedding per block); + // at prefill arrival the server queries restore_semantic_paged with + // a query embedding, and this method scores against the fingerprints, + // picks the top-K matches above a cosine threshold, and faults each + // matching block back from warm/cold to hot before kernel dispatch. + // + // Lifecycle: fingerprints follow blocks. Whole-seq wipe (clear, + // seq_rm with full range) drops them; per-seq removal happens + // automatically. No FIFO cap — memory grows with active context. + + // Record the fingerprint for one paged block. embedding should be + // L2-normalized; tier annotates current location (informational). + void record_paged_block_fingerprint(llama_seq_id seq_id, + uint32_t lblock, + std::vector embedding, + mt::SemanticIndex::Tier tier); + + // Score the seq's paged-block fingerprints against query_embedding, + // restore the top-K above threshold from warm (and cold as fallback) + // back to hot. Returns the count of blocks actually restored. Logs + // hit-rate (restored / requested) for MAD-122 acceptance criterion. + uint32_t restore_semantic_paged(llama_seq_id seq_id, + const std::vector & query_embedding, + int top_k = 5, + float threshold = 0.65f); + + // Diagnostic: how many fingerprints currently held. + size_t n_paged_fingerprints() const { return paged_semantic_.size(); } + + // MAD-129: O(1) check whether (seq_id, lblock) already has a + // fingerprint. Used by the server's prefill-time write trigger to + // skip blocks already fingerprinted on prior turns. + bool has_paged_fingerprint(llama_seq_id seq_id, uint32_t lblock) const { + return paged_semantic_.has_fingerprint(seq_id, lblock); + } + + // MAD-130: persist the cold-tier index to a sidecar file at + // `${cold_path_}/paged/index.bin`. Called by the server on graceful + // shutdown or as part of /slots/save. Returns false on I/O error. + // No-op when cold tier is disabled. + bool save_cold_index_sidecar() const; + + // MAD-130: persist the BlockSemanticIndex (paged-block fingerprints) + // to a sidecar file at the given path. Thin forwarders so the + // server can save fingerprints alongside state_write without poking + // at private members. + bool save_paged_fingerprints(const std::string & path) const { + return paged_semantic_.save_to_disk(path); + } + bool load_paged_fingerprints(const std::string & path) { + return paged_semantic_.load_from_disk(path); + } + private: friend class llama_kv_cache_paged_context; // Per-seq tracking. seq_pos_max is the highest position written for // the seq (inclusive). seq_pos_min stays 0 unless seq_rm carved off - // the head. + // the head. last_active_us is the wall-clock microseconds at which + // apply_ubatch_to_state last saw tokens for this seq (MAD-132 + // preempt fairness; older = more eligible for preemption). struct seq_state { - llama_pos pos_min = -1; - llama_pos pos_max = -1; + llama_pos pos_min = -1; + llama_pos pos_max = -1; + uint64_t last_active_us = 0; }; + // MAD-132: thread-affinity capture for the single-threading contract. + // First mutator call captures std::this_thread::get_id(); subsequent + // calls assert match in DEBUG builds. Release builds: zero overhead + // (the assert macro compiles away to nothing). + mutable std::thread::id captured_thread_id_; + void check_thread_id_() const; + // Allocate and upload the block_table / context_lens / q_lens // tensors to the GPU for the current batch. Called by init_batch // before the graph runs. @@ -286,6 +454,24 @@ class llama_kv_cache_paged : public llama_memory_i { // No-op if warm tier is disabled. bool fault_in_warm_blocks_for_batch(const llama_ubatch & ub); + // MAD-130: load the cold-tier sidecar from disk and rebuild + // cold_slot_for_ + cold_pool_free_ + cold_in_use_. Called from the + // ctor when cold_resume=true and the sidecar exists. Returns false + // if the sidecar is missing/corrupt/mismatched-config — caller + // should treat as "start fresh." + bool load_cold_index_sidecar_(const std::string & path); + + // MAD-128: CoW any blocks being written this ubatch that are shared + // (refcount > 1 from a prior seq_cp). For each (seq, lblock) pair + // touched by a write in `ub`: if that physical block is shared, + // allocate a fresh GPU block, copy the existing data into it, swap + // the seq's table entry to point at the new block, and decrement + // the old block's refcount. After this returns, every write target + // is uniquely-owned and the kernel can write without corrupting + // other sequences sharing the same prefix. Returns false if a CoW + // alloc failed and eviction couldn't free a block. + bool cow_writes_for_ubatch(const llama_ubatch & ub); + const llama_model & model_; ggml_backend_buffer_type_t buft_; uint32_t n_blocks_total_; @@ -333,6 +519,30 @@ class llama_kv_cache_paged : public llama_memory_i { uint32_t n_cold_blocks_ = 0; uint32_t cold_in_use_ = 0; std::string cold_path_; + + // MAD-133: tier-movement counters (monotonic since startup, single- + // threaded so no atomics). Bumped at the relevant tier-move sites. + uint64_t evict_h2w_total_ = 0; + uint64_t evict_w2c_total_ = 0; + uint64_t evict_c2drop_total_ = 0; + uint64_t restore_w2h_total_ = 0; + uint64_t restore_c2h_total_ = 0; + uint64_t seq_preempt_total_ = 0; + uint64_t seq_restore_total_ = 0; + uint64_t semantic_attempts_total_ = 0; + uint64_t semantic_hits_total_ = 0; + uint64_t semantic_blocks_restored_total_ = 0; + + // MAD-131: per-instance subdir + flock-based double-start protection. + // instance_id_ defaults to the process pid as a string; --instance-id + // overrides for deterministic restarts. cold_lock_fd_ holds the + // OS-level lock on `${cold_path_}/paged/instance-${id}/.lock` for + // the lifetime of the cache; the .pid file is informational. Both + // get cleaned up by the destructor. + std::string instance_id_; + std::string cold_lock_path_; + std::string cold_pid_path_; + int cold_lock_fd_ = -1; std::vector cold_fd_k_; // [n_layers] fd for K cold file (-1 = none) std::vector cold_fd_v_; // [n_layers] fd for V cold file std::vector cold_pool_free_; // free cold_idx stack @@ -344,8 +554,9 @@ class llama_kv_cache_paged : public llama_memory_i { // Block-table machinery (independent from mt::, since this owns the // cache outright). Single-pool: GPU only for v1. - mt::BlockPool pool_; - mt::BlockTable table_; + mt::BlockPool pool_; + mt::BlockTable table_; + mt::BlockSemanticIndex paged_semantic_; // MAD-125 // Per-seq position tracking. std::vector seq_states_; diff --git a/src/llama-kv-cache-tiered.cpp b/src/llama-kv-cache-tiered.cpp deleted file mode 100644 index 51ae5e719839..000000000000 --- a/src/llama-kv-cache-tiered.cpp +++ /dev/null @@ -1,1145 +0,0 @@ -#include "llama-kv-cache-tiered.h" - -#include "llama-impl.h" -#include "llama-io.h" - -#include -#include -#include -#include -#include -#include -#include - -#ifdef GGML_USE_HIP -#include -#endif - -// -// llama_ssd_storage_format -// - -// Quantize float32 to int4 (4-bit) -std::vector llama_ssd_storage_format::quantize_int4(const float* data, uint32_t n_elements) { - std::vector result((n_elements + 1) / 2); // 2 values per byte - - for (uint32_t i = 0; i < n_elements; i += 2) { - // Quantize to 4-bit range [-8, 7] - int32_t v0 = std::max(-8, std::min(7, int(std::round(data[i] * 15.0f)))); - int32_t v1 = std::max(-8, std::min(7, int(std::round(data[i + 1] * 15.0f)))); - - // Pack into single byte - result[i / 2] = uint8_t((v0 & 0x0F) | ((v1 & 0x0F) << 4)); - } - - return result; -} - -// Quantize float32 to int8 (8-bit) -std::vector llama_ssd_storage_format::quantize_int8(const float* data, uint32_t n_elements) { - std::vector result(n_elements); - - for (uint32_t i = 0; i < n_elements; i++) { - // Quantize to int8 range [-128, 127] - result[i] = uint8_t(std::max(-128, std::min(127, int(std::round(data[i] * 127.0f)))) + 128); - } - - return result; -} - -// Dequantize int4 to float32 -bool llama_ssd_storage_format::dequantize_int4(const uint8_t* data, float* out, uint32_t n_elements) { - for (uint32_t i = 0; i < (n_elements + 1) / 2; i++) { - uint8_t byte = data[i]; - int32_t v0 = (byte & 0x0F) - 8; // Lower 4 bits - int32_t v1 = ((byte >> 4) & 0x0F) - 8; // Upper 4 bits - - out[2 * i] = float(v0) / 15.0f; - if (2 * i + 1 < n_elements) { - out[2 * i + 1] = float(v1) / 15.0f; - } - } - - return true; -} - -// Dequantize int8 to float32 -bool llama_ssd_storage_format::dequantize_int8(const uint8_t* data, float* out, uint32_t n_elements) { - for (uint32_t i = 0; i < n_elements; i++) { - out[i] = float(data[i] - 128) / 127.0f; - } - - return true; -} - -bool llama_ssd_storage_format::write(const std::string& path, - const float* k_data, - const float* v_data, - uint32_t n_tokens, - uint32_t n_layers, - uint32_t n_embd_k, - uint32_t n_embd_v, - llama_cache_compression compression) { - std::ofstream file(path, std::ofstream::binary); - if (!file) { - return false; - } - - // Calculate data size based on compression - size_t element_size = sizeof(float); - if (compression == COMPRESSION_INT4) { - element_size = sizeof(uint8_t) / 2; // 4-bit = 0.5 bytes per element - } else if (compression == COMPRESSION_INT8) { - element_size = sizeof(uint8_t); // 8-bit = 1 byte per element - } - - size_t total_elements = n_tokens * n_layers * (n_embd_k + n_embd_v); - size_t compressed_size = total_elements * element_size; - - // Write header - file_header header; - header.magic = MAGIC; - header.version = VERSION; - header.n_layers = n_layers; - header.n_embd_k = n_embd_k; - header.n_embd_v = n_embd_v; - header.n_tokens = n_tokens; - header.compression_type = uint32_t(compression); - header.layer_offset = 0; - header.attention_threshold = 0.0f; - header.index_offset = 0; // Will be filled later - header.index_size = 0; // Will be filled later - - file.write(reinterpret_cast(&header), sizeof(file_header)); - - // Write data with compression - if (compression == COMPRESSION_INT4) { - // Quantize and write K data - auto k_quant = quantize_int4(k_data, n_tokens * n_layers * n_embd_k); - uint32_t k_size = uint32_t(k_quant.size()); - file.write(reinterpret_cast(&k_size), sizeof(uint32_t)); - file.write(reinterpret_cast(k_quant.data()), k_quant.size()); - - // Quantize and write V data - auto v_quant = quantize_int4(v_data, n_tokens * n_layers * n_embd_v); - uint32_t v_size = uint32_t(v_quant.size()); - file.write(reinterpret_cast(&v_size), sizeof(uint32_t)); - file.write(reinterpret_cast(v_quant.data()), v_quant.size()); - } else if (compression == COMPRESSION_INT8) { - // Quantize and write K data - auto k_quant = quantize_int8(k_data, n_tokens * n_layers * n_embd_k); - uint32_t k_size = uint32_t(k_quant.size()); - file.write(reinterpret_cast(&k_size), sizeof(uint32_t)); - file.write(reinterpret_cast(k_quant.data()), k_quant.size()); - - // Quantize and write V data - auto v_quant = quantize_int8(v_data, n_tokens * n_layers * n_embd_v); - uint32_t v_size = uint32_t(v_quant.size()); - file.write(reinterpret_cast(&v_size), sizeof(uint32_t)); - file.write(reinterpret_cast(v_quant.data()), v_quant.size()); - } else { - // No compression - write raw floats - size_t k_size = n_tokens * n_layers * n_embd_k * sizeof(float); - size_t v_size = n_tokens * n_layers * n_embd_v * sizeof(float); - file.write(const_cast(reinterpret_cast(k_data)), k_size); - file.write(const_cast(reinterpret_cast(v_data)), v_size); - } - - return file.good(); -} - -bool llama_ssd_storage_format::read(const std::string& path, - float* k_data, - float* v_data, - uint32_t n_tokens, - uint32_t n_layers, - uint32_t n_embd_k, - uint32_t n_embd_v) { - std::ifstream file(path, std::ifstream::binary); - if (!file) { - return false; - } - - // Read and validate header - file_header header; - file.read(reinterpret_cast(&header), sizeof(file_header)); - - if (header.magic != MAGIC || header.version != VERSION) { - return false; - } - - // Read data based on compression type - if (header.compression_type == COMPRESSION_INT4) { - // Read K data - uint32_t k_size; - file.read(reinterpret_cast(&k_size), sizeof(uint32_t)); - std::vector k_quant(k_size); - file.read(reinterpret_cast(k_quant.data()), k_size); - dequantize_int4(k_quant.data(), k_data, n_tokens * n_layers * n_embd_k); - - // Read V data - uint32_t v_size; - file.read(reinterpret_cast(&v_size), sizeof(uint32_t)); - std::vector v_quant(v_size); - file.read(reinterpret_cast(v_quant.data()), v_size); - dequantize_int4(v_quant.data(), v_data, n_tokens * n_layers * n_embd_v); - } else if (header.compression_type == COMPRESSION_INT8) { - // Read K data - uint32_t k_size; - file.read(reinterpret_cast(&k_size), sizeof(uint32_t)); - std::vector k_quant(k_size); - file.read(reinterpret_cast(k_quant.data()), k_size); - dequantize_int8(k_quant.data(), k_data, n_tokens * n_layers * n_embd_k); - - // Read V data - uint32_t v_size; - file.read(reinterpret_cast(&v_size), sizeof(uint32_t)); - std::vector v_quant(v_size); - file.read(reinterpret_cast(v_quant.data()), v_size); - dequantize_int8(v_quant.data(), v_data, n_tokens * n_layers * n_embd_v); - } else { - // No compression - read raw floats - size_t k_size = n_tokens * n_layers * n_embd_k * sizeof(float); - size_t v_size = n_tokens * n_layers * n_embd_v * sizeof(float); - file.read(reinterpret_cast(k_data), k_size); - file.read(reinterpret_cast(v_data), v_size); - } - - return file.good(); -} - -bool llama_ssd_storage_format::build_index(const std::string& index_path, - const std::string& data_path, - uint32_t n_layers) { - // Build layer index for fast lookup - std::ifstream data_file(data_path, std::ifstream::binary); - if (!data_file) { - return false; - } - - // Read header to get file structure - file_header header; - data_file.read(reinterpret_cast(&header), sizeof(file_header)); - - // Build index entries for each layer - layer_index.clear(); - for (uint32_t layer = 0; layer < n_layers; layer++) { - layer_index_entry entry; - entry.layer = layer; - entry.file_offset = sizeof(file_header) + layer * (header.n_embd_k + header.n_embd_v) * sizeof(float); - entry.n_tokens = header.n_tokens; - entry.data_size = header.n_tokens * (header.n_embd_k + header.n_embd_v) * sizeof(float); - layer_index.push_back(entry); - } - - // Save index to file - return save_index(index_path); -} - -bool llama_ssd_storage_format::load_index(const std::string& index_path) { - std::ifstream file(index_path, std::ifstream::binary); - if (!file) { - return false; - } - - // Read number of entries - uint32_t n_entries; - file.read(reinterpret_cast(&n_entries), sizeof(uint32_t)); - - // Read each entry - layer_index.clear(); - for (uint32_t i = 0; i < n_entries; i++) { - layer_index_entry entry; - file.read(reinterpret_cast(&entry), sizeof(layer_index_entry)); - layer_index.push_back(entry); - } - - return file.good(); -} - -bool llama_ssd_storage_format::save_index(const std::string& index_path) { - std::ofstream file(index_path, std::ofstream::binary); - if (!file) { - return false; - } - - // Write number of entries - uint32_t n_entries = uint32_t(layer_index.size()); - file.write(reinterpret_cast(&n_entries), sizeof(uint32_t)); - - // Write each entry - for (const auto& entry : layer_index) { - file.write(reinterpret_cast(&entry), sizeof(layer_index_entry)); - } - - return file.good(); -} - -uint64_t llama_ssd_storage_format::get_layer_offset(uint32_t layer) const { - for (const auto& entry : layer_index) { - if (entry.layer == layer) { - return entry.file_offset; - } - } - return 0; // Not found -} - -bool llama_ssd_storage_format::has_layer(uint32_t layer) const { - for (const auto& entry : layer_index) { - if (entry.layer == layer) { - return true; - } - } - return false; -} - -// -// llama_kv_cache_tiered -// - -llama_kv_cache_tiered::llama_kv_cache_tiered( - const llama_model& model, - const llama_tier_config& config, - const std::string& ssd_path, - llama_eviction_policy eviction_policy, - llama_cache_compression compression, - float attention_threshold) - : config(config), - eviction_policy(eviction_policy), - ssd_path(ssd_path), - compression(compression), - attention_threshold(attention_threshold) { - stats.reset(); -} - -llama_kv_cache_tiered::~llama_kv_cache_tiered() { - for (auto & tl : kv_layers) { - delete[] tl.warm_k; - delete[] tl.warm_v; -#ifdef GGML_USE_HIP - if (tl.warm_k_dev) { hipFree(tl.warm_k_dev); } - if (tl.warm_v_dev) { hipFree(tl.warm_v_dev); } -#endif - } -} - -bool llama_kv_cache_tiered::init() { - // Create SSD directory if it doesn't exist - std::filesystem::path ssd_dir(ssd_path); - if (!std::filesystem::exists(ssd_dir)) { - std::filesystem::create_directories(ssd_dir); - } - - // Warm slot tracking — per-layer buffers allocated later in set_kv_layers_from_cache() - if (config.warm_capacity() > 0) { - warm_slots.assign(config.warm_capacity(), WarmSlot{}); - LLAMA_LOG_INFO("%s: warm tier: %u slots reserved (buffers wired after KV layers available)\n", - __func__, config.warm_capacity()); - } - if (config.cold_capacity() > 0) { - LLAMA_LOG_INFO("%s: cold tier: %u slots reserved (SSD path: %s)\n", - __func__, config.cold_capacity(), ssd_path.c_str()); - } - - return true; -} - -bool llama_kv_cache_tiered::resize(uint32_t new_total_ctx) { - std::lock_guard lock(mutex); - config.total_ctx = new_total_ctx; - return true; -} - -bool llama_kv_cache_tiered::set_eviction_policy(llama_eviction_policy policy) { - std::lock_guard lock(mutex); - eviction_policy = policy; - return true; -} - -bool llama_kv_cache_tiered::set_attention_threshold(float threshold) { - std::lock_guard lock(mutex); - attention_threshold = threshold; - return true; -} - -void llama_kv_cache_tiered::track_hot_range(uint32_t n_hot) { - std::lock_guard lock(mutex); - for (uint32_t p = 0; p < n_hot; p++) { - token_metadata.record_access((llama_pos)p); - } - stats.hot_tokens = n_hot; -} - -bool llama_kv_cache_tiered::evict_tokens(uint32_t n_tokens_to_evict, llama_cache_tier from_tier) { - std::lock_guard lock(mutex); - - auto candidates = token_metadata.get_eviction_candidates(eviction_policy, n_tokens_to_evict); - if (candidates.empty()) { - return false; - } - - // Apply semantic eviction weighting if we have a current query embedding - if (!current_query_emb.empty() && !fingerprints.empty()) { - std::vector> candidates_with_similarity; - - for (auto pos : candidates) { - float similarity = 0.0f; // Default similarity if no fingerprint found - - // Find fingerprint for this position - for (const auto& fp : fingerprints) { - bool found = false; - for (auto fp_pos : fp.positions) { - if (fp_pos == pos) { - found = true; - break; - } - } - if (found) { - // Compute cosine similarity (dot product since both are L2-normalized) - similarity = 0.0f; - for (size_t i = 0; i < current_query_emb.size() && i < fp.embedding.size(); i++) { - similarity += current_query_emb[i] * fp.embedding[i]; - } - break; - } - } - - candidates_with_similarity.emplace_back(pos, similarity); - } - - // Reorder candidates: positions with similarity < 0.3 evicted first, - // positions with similarity > 0.65 moved to back of eviction queue - std::sort(candidates_with_similarity.begin(), candidates_with_similarity.end(), - [](const auto& a, const auto& b) { - if (a.second > 0.65f) return false; // Keep high similarity at back - if (b.second > 0.65f) return true; // Keep high similarity at back - if (a.second < 0.3f) return true; // Evict low similarity first - if (b.second < 0.3f) return false; // Evict low similarity first - return a.second < b.second; // Otherwise by similarity - }); - - // Update candidates with reordered positions - candidates.clear(); - for (const auto& candidate : candidates_with_similarity) { - candidates.push_back(candidate.first); - } - - // Count semantic eviction saves - uint32_t saves = 0; - for (const auto& candidate : candidates_with_similarity) { - if (candidate.second > 0.65f) { - saves++; - } - } - stats.semantic_eviction_saves += saves; - } - - std::vector to_warm, to_cold; - - // Route to warm if slots available (RAM or GPU), otherwise cold - int free_warm = 0; - for (const auto & s : warm_slots) { - if (!s.occupied) free_warm++; - } - for (auto pos : candidates) { - if (free_warm > 0) { - to_warm.push_back(pos); - free_warm--; - } else { - to_cold.push_back(pos); - } - } - - // migrate_tokens handles null k_tensor/v_tensor as metadata-only (no-copy) path - if (!to_warm.empty()) { - migrate_tokens(to_warm, TIER_HOT, TIER_WARM); - stats.warm_tokens += (uint32_t)to_warm.size(); - stats.hot_tokens = stats.hot_tokens >= (uint32_t)to_warm.size() - ? stats.hot_tokens - (uint32_t)to_warm.size() : 0; - LLAMA_LOG_INFO("%s: evicted %zu hot->warm\n", __func__, to_warm.size()); - } - if (!to_cold.empty()) { - migrate_tokens(to_cold, TIER_HOT, TIER_COLD); - stats.cold_tokens += (uint32_t)to_cold.size(); - stats.hot_tokens = stats.hot_tokens >= (uint32_t)to_cold.size() - ? stats.hot_tokens - (uint32_t)to_cold.size() : 0; - LLAMA_LOG_INFO("%s: evicted %zu hot->cold\n", __func__, to_cold.size()); - } - - stats.eviction_count += candidates.size(); - return true; -} - -// Warm slot management — unconditional (works for RAM and GPU warm tiers) -int llama_kv_cache_tiered::warm_alloc_slot() { - for (int i = 0; i < (int)warm_slots.size(); i++) { - if (!warm_slots[i].occupied) { - warm_slots[i].occupied = true; - return i; - } - } - return -1; -} - -void llama_kv_cache_tiered::warm_free_slot(llama_pos pos) { - auto it = warm_pos_to_slot.find(pos); - if (it != warm_pos_to_slot.end()) { - int slot = it->second; - warm_slots[slot].occupied = false; - warm_slots[slot].pos = -1; - warm_pos_to_slot.erase(it); - } -} - -// Copy one token's K data from VRAM (or host) into warm slot — K is always contiguous -bool llama_kv_cache_tiered::warm_copy_to_host(uint32_t il, int slot, llama_pos pos) { - if (il >= kv_layers.size() || slot < 0 || (size_t)slot >= warm_slots.size()) return false; - const auto & tl = kv_layers[il]; - if (!tl.k || !tl.warm_k || !tl.v || !tl.warm_v) return false; - - const size_t elem_k = ggml_element_size(tl.k); - const size_t n_embd_k = (size_t)tl.k->ne[0]; - uint8_t * dst_k = tl.warm_k + (size_t)slot * tl.k_bytes; - - // K: contiguous row at pos — single copy -#ifdef GGML_USE_HIP - hipMemcpy(dst_k, (uint8_t *)tl.k->data + (size_t)pos * tl.k_bytes, tl.k_bytes, hipMemcpyDeviceToHost); -#else - memcpy(dst_k, (uint8_t *)tl.k->data + (size_t)pos * tl.k_bytes, tl.k_bytes); -#endif - - const size_t elem_v = ggml_element_size(tl.v); - const size_t n_embd_v = (size_t)tl.v->ne[0]; - uint8_t * dst_v = tl.warm_v + (size_t)slot * tl.v_bytes; - - if (!tl.v_trans) { - // V: also contiguous when not transposed -#ifdef GGML_USE_HIP - hipMemcpy(dst_v, (uint8_t *)tl.v->data + (size_t)pos * tl.v_bytes, tl.v_bytes, hipMemcpyDeviceToHost); -#else - memcpy(dst_v, (uint8_t *)tl.v->data + (size_t)pos * tl.v_bytes, tl.v_bytes); -#endif - } else { - // V transposed: element j at (j*kv_size + pos)*elem_v — gather into contiguous warm buf - // Use hipMemcpy2D (src column → dst row): copies n_embd_v elements, each elem_v bytes apart -#ifdef GGML_USE_HIP - hipMemcpy2D(dst_v, // dst (contiguous) - elem_v, // dst pitch - (uint8_t *)tl.v->data + (size_t)pos * elem_v, // src (column pos) - (size_t)tl.kv_size * elem_v, // src pitch (stride between rows) - elem_v, // width per row - n_embd_v, // number of rows - hipMemcpyDeviceToHost); -#else - for (size_t j = 0; j < n_embd_v; j++) { - memcpy(dst_v + j * elem_v, - (uint8_t *)tl.v->data + ((size_t)j * (size_t)tl.kv_size + (size_t)pos) * elem_v, - elem_v); - } -#endif - } - return true; -} - -// Restore one token's K/V from warm slot back into the hot VRAM tensor -bool llama_kv_cache_tiered::warm_copy_from_host(uint32_t il, int slot, llama_pos pos) { - if (il >= kv_layers.size() || slot < 0 || (size_t)slot >= warm_slots.size()) return false; - const auto & tl = kv_layers[il]; - if (!tl.k || !tl.warm_k || !tl.v || !tl.warm_v) return false; - - const size_t elem_k = ggml_element_size(tl.k); - const size_t n_embd_k = (size_t)tl.k->ne[0]; - const uint8_t * src_k = tl.warm_k + (size_t)slot * tl.k_bytes; - -#ifdef GGML_USE_HIP - hipMemcpy((uint8_t *)tl.k->data + (size_t)pos * tl.k_bytes, src_k, tl.k_bytes, hipMemcpyHostToDevice); -#else - memcpy((uint8_t *)tl.k->data + (size_t)pos * tl.k_bytes, src_k, tl.k_bytes); -#endif - - const size_t elem_v = ggml_element_size(tl.v); - const size_t n_embd_v = (size_t)tl.v->ne[0]; - const uint8_t * src_v = tl.warm_v + (size_t)slot * tl.v_bytes; - - if (!tl.v_trans) { -#ifdef GGML_USE_HIP - hipMemcpy((uint8_t *)tl.v->data + (size_t)pos * tl.v_bytes, src_v, tl.v_bytes, hipMemcpyHostToDevice); -#else - memcpy((uint8_t *)tl.v->data + (size_t)pos * tl.v_bytes, src_v, tl.v_bytes); -#endif - } else { - // Scatter: warm contiguous → V transposed column at pos -#ifdef GGML_USE_HIP - hipMemcpy2D((uint8_t *)tl.v->data + (size_t)pos * elem_v, // dst column pos - (size_t)tl.kv_size * elem_v, // dst pitch - src_v, // src (contiguous) - elem_v, // src pitch - elem_v, // width - n_embd_v, // height - hipMemcpyHostToDevice); -#else - for (size_t j = 0; j < n_embd_v; j++) { - memcpy((uint8_t *)tl.v->data + ((size_t)j * (size_t)tl.kv_size + (size_t)pos) * elem_v, - src_v + j * elem_v, - elem_v); - } -#endif - } - return true; -} - -#ifdef GGML_USE_HIP -// VRAM (R9700) → device VRAM (6900XT warm tier): K contiguous, V gather/scatter -bool llama_kv_cache_tiered::warm_copy_to_dev(uint32_t il, int slot, llama_pos pos) { - if (il >= kv_layers.size() || slot < 0 || (size_t)slot >= warm_slots.size()) return false; - const auto & tl = kv_layers[il]; - if (!tl.warm_k_dev || !tl.warm_v_dev || !tl.k || !tl.v) return false; - int prev_dev = 0; hipGetDevice(&prev_dev); hipSetDevice(config.warm_device); - // K: contiguous D2D copy - hipMemcpy((uint8_t *)tl.warm_k_dev + (size_t)slot * tl.k_bytes, - (uint8_t *)tl.k->data + (size_t)pos * tl.k_bytes, - tl.k_bytes, hipMemcpyDeviceToDevice); - const size_t ev = ggml_element_size(tl.v); - const size_t nv = (size_t)tl.v->ne[0]; // n_embd_v_gqa - if (!tl.v_trans) { - hipMemcpy((uint8_t *)tl.warm_v_dev + (size_t)slot * tl.v_bytes, - (uint8_t *)tl.v->data + (size_t)pos * tl.v_bytes, - tl.v_bytes, hipMemcpyDeviceToDevice); - } else { - // Gather column pos into contiguous warm buf - hipMemcpy2D((uint8_t *)tl.warm_v_dev + (size_t)slot * tl.v_bytes, ev, - (uint8_t *)tl.v->data + (size_t)pos * ev, - (size_t)tl.kv_size * ev, ev, nv, hipMemcpyDeviceToDevice); - } - hipSetDevice(prev_dev); - return true; -} - -bool llama_kv_cache_tiered::warm_copy_from_dev(uint32_t il, int slot, llama_pos pos) { - if (il >= kv_layers.size() || slot < 0 || (size_t)slot >= warm_slots.size()) return false; - const auto & tl = kv_layers[il]; - if (!tl.warm_k_dev || !tl.warm_v_dev || !tl.k || !tl.v) return false; - int prev_dev = 0; hipGetDevice(&prev_dev); hipSetDevice(config.warm_device); - hipMemcpy((uint8_t *)tl.k->data + (size_t)pos * tl.k_bytes, - (uint8_t *)tl.warm_k_dev + (size_t)slot * tl.k_bytes, - tl.k_bytes, hipMemcpyDeviceToDevice); - const size_t ev = ggml_element_size(tl.v); - const size_t nv = (size_t)tl.v->ne[0]; - if (!tl.v_trans) { - hipMemcpy((uint8_t *)tl.v->data + (size_t)pos * tl.v_bytes, - (uint8_t *)tl.warm_v_dev + (size_t)slot * tl.v_bytes, - tl.v_bytes, hipMemcpyDeviceToDevice); - } else { - hipMemcpy2D((uint8_t *)tl.v->data + (size_t)pos * ev, - (size_t)tl.kv_size * ev, - (uint8_t *)tl.warm_v_dev + (size_t)slot * tl.v_bytes, ev, - ev, nv, hipMemcpyDeviceToDevice); - } - hipSetDevice(prev_dev); - return true; -} -#endif // GGML_USE_HIP - -void llama_kv_cache_tiered::set_kv_layers_from_cache(llama_kv_cache * cache) { - if (!cache) return; - std::lock_guard lock(mutex); - - uint32_t n = cache->get_num_kv_layers(); - bool vtrans = cache->is_v_transposed(); - int64_t kvsz = (int64_t)cache->get_size(); - - kv_layers.resize(n); - for (uint32_t il = 0; il < n; il++) { - auto & tl = kv_layers[il]; - tl.k = cache->get_layer_k_raw(il); - tl.v = cache->get_layer_v_raw(il); - tl.v_trans = vtrans; - tl.kv_size = kvsz; - // Use nb[1] (the actual row stride) for per-token byte size so that - // block-quantized types (turbo4, Q4_K, etc.) are sized correctly. - // ggml_element_size returns the block size, not bytes-per-element, so - // ne[0]*element_size is wildly wrong for quantized caches. - if (tl.k) tl.k_bytes = tl.k->nb[1]; - if (tl.v) tl.v_bytes = tl.v->nb[1]; - } - - // (Re)allocate per-layer warm buffers using already-reserved slot count - uint32_t n_warm = (uint32_t)warm_slots.size(); - if (n_warm == 0 && config.warm_capacity() > 0) { - n_warm = config.warm_capacity(); - warm_slots.assign(n_warm, WarmSlot{}); - } - - for (auto & tl : kv_layers) { - delete[] tl.warm_k; - delete[] tl.warm_v; - tl.warm_k = tl.warm_v = nullptr; - // Skip RAM staging buffers when a GPU warm device is configured — the - // GPU VRAM path (warm_k_dev/warm_v_dev) is used instead, so allocating - // n_warm * kv_bytes of RAM here would just trigger OOM for large contexts. -#ifndef GGML_USE_HIP - if (n_warm > 0 && tl.k_bytes > 0) tl.warm_k = new uint8_t[n_warm * tl.k_bytes](); - if (n_warm > 0 && tl.v_bytes > 0) tl.warm_v = new uint8_t[n_warm * tl.v_bytes](); -#else - if (config.warm_device < 0) { - if (n_warm > 0 && tl.k_bytes > 0) tl.warm_k = new uint8_t[n_warm * tl.k_bytes](); - if (n_warm > 0 && tl.v_bytes > 0) tl.warm_v = new uint8_t[n_warm * tl.v_bytes](); - } -#endif - -#ifdef GGML_USE_HIP - if (config.warm_device >= 0 && n_warm > 0) { - hipSetDevice(config.warm_device); - if (tl.warm_k_dev) hipFree(tl.warm_k_dev); - if (tl.warm_v_dev) hipFree(tl.warm_v_dev); - tl.warm_k_dev = tl.warm_v_dev = nullptr; - if (tl.k_bytes > 0) hipMalloc(&tl.warm_k_dev, n_warm * tl.k_bytes); - if (tl.v_bytes > 0) hipMalloc(&tl.warm_v_dev, n_warm * tl.v_bytes); - } -#endif - } - - size_t total_mb = 0; - for (const auto & tl : kv_layers) - total_mb += (tl.k_bytes + tl.v_bytes) * n_warm; - total_mb /= (1024*1024); - -#ifdef GGML_USE_HIP - if (config.warm_device >= 0) { - LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm GPU: ~%zu MiB (%u slots x %u layers)\n", - __func__, n, (int)vtrans, (long long)kvsz, total_mb, n_warm, n); - } else { - LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm RAM: ~%zu MiB (%u slots x %u layers)\n", - __func__, n, (int)vtrans, (long long)kvsz, total_mb, n_warm, n); - } -#else - LLAMA_LOG_INFO("%s: wired %u KV layers (v_trans=%d, kv_size=%lld), warm RAM: ~%zu MiB (%u slots x %u layers)\n", - __func__, n, (int)vtrans, (long long)kvsz, total_mb, n_warm, n); -#endif -#ifdef GGML_USE_HIP - if (config.warm_device >= 0 && n_warm > 0 && !kv_layers.empty() && kv_layers[0].warm_k_dev) { - LLAMA_LOG_INFO("llama_kv_cache_tiered: ROCm%d KV warm buffer size = %6.2f MiB (%u slots x %u layers)\n", - config.warm_device, (float)total_mb, n_warm, n); - } -#endif -} - -bool llama_kv_cache_tiered::migrate_tokens(const std::vector& positions, - llama_cache_tier from_tier, - llama_cache_tier to_tier, - const ggml_tensor* /*unused_k*/, - const ggml_tensor* /*unused_v*/) { - auto start = std::chrono::high_resolution_clock::now(); - bool success = true; - - if (kv_layers.empty() || warm_slots.empty()) { - // Metadata-only path: layer tensors not yet wired (set_kv_layers_from_cache not called) - for (auto pos : positions) { - token_metadata.record_access(pos); - } - } else { - bool is_hot_to_warm = (from_tier == TIER_HOT && to_tier == TIER_WARM); - bool is_warm_to_cold = (from_tier == TIER_WARM && to_tier == TIER_COLD); - bool is_cold_to_hot = (from_tier == TIER_COLD && to_tier == TIER_HOT); - bool is_cold_to_warm = (from_tier == TIER_COLD && to_tier == TIER_WARM); - bool is_warm_to_hot = (from_tier == TIER_WARM && to_tier == TIER_HOT); - - if (is_hot_to_warm) { - // VRAM → RAM (or eGPU): copy K/V for every layer and every evicted position - for (auto pos : positions) { - int slot = warm_alloc_slot(); - if (slot < 0) { - LLAMA_LOG_WARN("%s: warm full, dropping pos %d (cold SSD TODO)\n", __func__, pos); - token_metadata.record_access(pos); - continue; - } - warm_slots[slot].pos = pos; - warm_pos_to_slot[pos] = slot; - - for (uint32_t il = 0; il < (uint32_t)kv_layers.size(); il++) { -#ifdef GGML_USE_HIP - if (kv_layers[il].warm_k_dev) { - success &= warm_copy_to_dev(il, slot, pos); - } else { - success &= warm_copy_to_host(il, slot, pos); - } -#else - success &= warm_copy_to_host(il, slot, pos); -#endif - } - token_metadata.record_access(pos); - } - if (success) - LLAMA_LOG_INFO("%s: hot→warm: %zu positions x %zu layers\n", - __func__, positions.size(), kv_layers.size()); - } else if (is_warm_to_hot) { - // RAM (or eGPU) → VRAM: restore - for (auto pos : positions) { - auto it = warm_pos_to_slot.find(pos); - if (it == warm_pos_to_slot.end()) continue; - int slot = it->second; - for (uint32_t il = 0; il < (uint32_t)kv_layers.size(); il++) { -#ifdef GGML_USE_HIP - if (kv_layers[il].warm_k_dev) { - success &= warm_copy_from_dev(il, slot, pos); - } else { - success &= warm_copy_from_host(il, slot, pos); - } -#else - success &= warm_copy_from_host(il, slot, pos); -#endif - } - warm_free_slot(pos); - token_metadata.record_access(pos); - } - } else if (is_warm_to_cold || is_cold_to_hot || is_cold_to_warm) { - // SSD paths: TODO — per-layer serialization - for (auto pos : positions) { - token_metadata.record_access(pos); - } - } - } - - auto end = std::chrono::high_resolution_clock::now(); - stats.total_migration_latency_us += - (double)std::chrono::duration_cast(end - start).count(); - return success; -} - -bool llama_kv_cache_tiered::batch_migrate_tokens(const std::vector& positions, - llama_cache_tier from_tier, - llama_cache_tier to_tier, - const ggml_tensor* k_tensor, - const ggml_tensor* v_tensor) { - auto start = std::chrono::high_resolution_clock::now(); - - // Delegate to migrate_tokens which handles the real per-layer data movement - return migrate_tokens(positions, from_tier, to_tier); -} - -bool llama_kv_cache_tiered::prefetch_tokens(const std::vector& positions, - llama_cache_tier target_tier) { - // Prefetch optimization: - // - Pre-load likely-needed tokens before they're requested - // - Use attention patterns to predict which tokens will be needed - // - Batch prefetch requests for efficiency - - // TODO: Implement attention-based prediction for prefetch - // TODO: Implement batch prefetch with priority queue - - return true; -} - -bool llama_kv_cache_tiered::contains_token(llama_pos pos) const { - std::lock_guard lock(mutex); - return token_metadata.get_metadata(pos) != nullptr; -} - -bool llama_kv_cache_tiered::load_token(llama_pos pos, llama_cache_tier target_tier) { - // Check if token is in cold tier and needs to be loaded - std::string ssd_file = get_ssd_file_path(pos); - if (std::filesystem::exists(ssd_file)) { - // Load from SSD - stats.cache_hits++; - return true; - } - stats.cache_misses++; - return false; -} - -llama_tier_stats llama_kv_cache_tiered::get_stats() const { - std::lock_guard lock(mutex); - return stats; -} - -void llama_kv_cache_tiered::reset_stats() { - std::lock_guard lock(mutex); - stats.reset(); -} - -void llama_kv_cache_tiered::add_fingerprint(const std::vector& positions, - const std::vector& embedding, - llama_cache_tier tier) { - std::lock_guard lock(mutex); - - semantic_fingerprint fp; - fp.positions = positions; - fp.embedding = embedding; - fp.tier = tier; - fp.turn = fingerprint_turn++; - - // Cap at 1000 entries - evict oldest when over cap - if (fingerprints.size() >= 1000) { - fingerprints.erase(fingerprints.begin()); - } - - fingerprints.push_back(fp); -} - -std::vector llama_kv_cache_tiered::score_fingerprints( - const std::vector& query_emb, int top_k, float threshold) const { - std::lock_guard lock(mutex); - - std::vector results; - - if (fingerprints.empty() || query_emb.empty()) { - return results; - } - - // Compute cosine similarity between query_emb and each fingerprint - // Since both are L2-normalized, cosine similarity = dot product - for (const auto& fp : fingerprints) { - // Compute dot product - float dot_product = 0.0f; - size_t min_size = std::min(query_emb.size(), fp.embedding.size()); - for (size_t i = 0; i < min_size; ++i) { - dot_product += query_emb[i] * fp.embedding[i]; - } - - // Check if above threshold - if (dot_product >= threshold) { - PrefetchHint hint; - hint.positions = fp.positions; - hint.score = dot_product; - hint.current_tier = fp.tier; - results.push_back(hint); - } - } - - // Sort by score descending - std::sort(results.begin(), results.end(), - [](const PrefetchHint& a, const PrefetchHint& b) { - return a.score > b.score; - }); - - // Return top_k results - if (static_cast(results.size()) > top_k) { - results.resize(top_k); - } - - return results; -} - -void llama_kv_cache_tiered::set_current_query_embedding(const std::vector& emb) { - std::lock_guard lock(mutex); - current_query_emb = emb; -} - -bool llama_kv_cache_tiered::save_fingerprints_to_disk(const std::string& path) { - std::lock_guard lock(mutex); - - std::ofstream file(path, std::ofstream::binary); - if (!file) { - return false; - } - - // Write number of entries - uint32_t n_entries = uint32_t(fingerprints.size()); - file.write(reinterpret_cast(&n_entries), sizeof(uint32_t)); - - // Write each fingerprint - for (const auto& fp : fingerprints) { - // Write positions count and positions - uint32_t n_pos = uint32_t(fp.positions.size()); - file.write(reinterpret_cast(&n_pos), sizeof(uint32_t)); - file.write(reinterpret_cast(fp.positions.data()), n_pos * sizeof(llama_pos)); - - // Write embedding dimension and embedding - uint32_t n_embd = uint32_t(fp.embedding.size()); - file.write(reinterpret_cast(&n_embd), sizeof(uint32_t)); - file.write(reinterpret_cast(fp.embedding.data()), n_embd * sizeof(float)); - - // Write tier and turn - file.write(reinterpret_cast(&fp.tier), sizeof(llama_cache_tier)); - file.write(reinterpret_cast(&fp.turn), sizeof(uint64_t)); - } - - return file.good(); -} - -bool llama_kv_cache_tiered::load_fingerprints_from_disk(const std::string& path) { - std::ifstream file(path, std::ifstream::binary); - if (!file) { - return false; - } - - // Read number of entries - uint32_t n_entries; - file.read(reinterpret_cast(&n_entries), sizeof(uint32_t)); - - // Clear existing fingerprints - fingerprints.clear(); - - // Read each fingerprint - for (uint32_t i = 0; i < n_entries; i++) { - semantic_fingerprint fp; - - // Read positions - uint32_t n_pos; - file.read(reinterpret_cast(&n_pos), sizeof(uint32_t)); - fp.positions.resize(n_pos); - file.read(reinterpret_cast(fp.positions.data()), n_pos * sizeof(llama_pos)); - - // Read embedding - uint32_t n_embd; - file.read(reinterpret_cast(&n_embd), sizeof(uint32_t)); - fp.embedding.resize(n_embd); - file.read(reinterpret_cast(fp.embedding.data()), n_embd * sizeof(float)); - - // Read tier and turn - file.read(reinterpret_cast(&fp.tier), sizeof(llama_cache_tier)); - file.read(reinterpret_cast(&fp.turn), sizeof(uint64_t)); - - fingerprints.push_back(fp); - } - - return file.good(); -} - -std::string llama_kv_cache_tiered::get_ssd_file_path(llama_pos pos) const { - // Generate unique file path for token position - return ssd_path + "/token_" + std::to_string(pos) + ".bin"; -} - -bool llama_kv_cache_tiered::save_to_ssd(const std::vector& positions, - const ggml_tensor* k_tensor, - const ggml_tensor* v_tensor, - bool is_device_data) { - // Save KV cache data to SSD using the SSD storage format - if (positions.empty() || k_tensor == nullptr || v_tensor == nullptr) { - return false; - } - - // Calculate number of tokens and layer info - uint32_t n_tokens = uint32_t(positions.size()); - uint32_t n_layers = 32; // Default placeholder - should come from model config - uint32_t n_embd_k = k_tensor->ne[0]; - uint32_t n_embd_v = v_tensor->ne[0]; - - // Calculate buffer size for device-to-host copy if needed - // This is a simplified calculation - actual implementation should use model-specific sizes - size_t buffer_size = n_tokens * n_embd_k * sizeof(float); - size_t buffer_size_v = n_tokens * n_embd_v * sizeof(float); - - // Use the SSD storage format to write data - llama_ssd_storage_format ssd; - std::string base_path = ssd_path + "/cache"; - -#ifdef GGML_USE_HIP - // If data is on device, we need to copy to host first - if (is_device_data) { - // Allocate host buffers for staging device data - std::vector k_host_buf(n_tokens * n_embd_k); - std::vector v_host_buf(n_tokens * n_embd_v); - - // Copy device data to host buffer using hipMemcpy - hipError_t err_k = hipMemcpy(k_host_buf.data(), - (const void*)k_tensor->data, - buffer_size, - hipMemcpyDeviceToHost); - hipError_t err_v = hipMemcpy(v_host_buf.data(), - (const void*)v_tensor->data, - buffer_size_v, - hipMemcpyDeviceToHost); - - if (err_k != hipSuccess || err_v != hipSuccess) { - return false; - } - - // Write K and V data for each position using host-staged data - for (auto pos : positions) { - std::string file_path = base_path + "_pos_" + std::to_string(pos) + ".bin"; - bool success = ssd.write(file_path, - k_host_buf.data(), - v_host_buf.data(), - n_tokens, n_layers, n_embd_k, n_embd_v, compression); - if (!success) { - return false; - } - } - return true; - } -#endif - - // Non-device data path (RAM/SSD) - use tensor data directly - const float* k_data = (const float*)k_tensor->data; - const float* v_data = (const float*)v_tensor->data; - - // Write K and V data for each position - for (auto pos : positions) { - std::string file_path = base_path + "_pos_" + std::to_string(pos) + ".bin"; - bool success = ssd.write(file_path, k_data, v_data, n_tokens, n_layers, n_embd_k, n_embd_v, compression); - if (!success) { - return false; - } - } - - return true; -} - -bool llama_kv_cache_tiered::load_from_ssd(const std::vector& positions, - ggml_tensor* k_tensor, - ggml_tensor* v_tensor, - bool to_device) { - // Load KV cache data from SSD using the SSD storage format - if (positions.empty() || k_tensor == nullptr || v_tensor == nullptr) { - return false; - } - - // Calculate number of tokens and layer info - uint32_t n_tokens = uint32_t(positions.size()); - uint32_t n_layers = 32; // Default placeholder - should come from model config - uint32_t n_embd_k = k_tensor->ne[0]; - uint32_t n_embd_v = v_tensor->ne[0]; - - // Calculate buffer size for host-to-device copy if needed - size_t buffer_size = n_tokens * n_embd_k * sizeof(float); - size_t buffer_size_v = n_tokens * n_embd_v * sizeof(float); - - // Use the SSD storage format to read data - llama_ssd_storage_format ssd; - std::string base_path = ssd_path + "/cache"; - -#ifdef GGML_USE_HIP - // If destination is device, we need to copy from host to device - if (to_device) { - // Allocate host buffer for staging before device copy - std::vector k_host_buf(n_tokens * n_embd_k); - std::vector v_host_buf(n_tokens * n_embd_v); - - // Read K and V data for each position into host buffer - for (auto pos : positions) { - std::string file_path = base_path + "_pos_" + std::to_string(pos) + ".bin"; - if (!ssd.read(file_path, k_host_buf.data(), v_host_buf.data(), n_tokens, n_layers, n_embd_k, n_embd_v)) { - return false; - } - } - - // Copy host data to device using hipMemcpy - hipError_t err_k = hipMemcpy((void*)k_tensor->data, - k_host_buf.data(), - buffer_size, - hipMemcpyHostToDevice); - hipError_t err_v = hipMemcpy((void*)v_tensor->data, - v_host_buf.data(), - buffer_size_v, - hipMemcpyHostToDevice); - - return (err_k == hipSuccess && err_v == hipSuccess); - } -#endif - - // Non-device path - read directly to tensor data - float* k_data = (float*)k_tensor->data; - float* v_data = (float*)v_tensor->data; - - // Read K and V data for each position - for (auto pos : positions) { - std::string file_path = base_path + "_pos_" + std::to_string(pos) + ".bin"; - if (!ssd.read(file_path, k_data, v_data, n_tokens, n_layers, n_embd_k, n_embd_v)) { - return false; - } - } - - return true; -} diff --git a/src/llama-kv-cache-tiered.h b/src/llama-kv-cache-tiered.h deleted file mode 100644 index 5011bec19134..000000000000 --- a/src/llama-kv-cache-tiered.h +++ /dev/null @@ -1,328 +0,0 @@ -#pragma once - -#include "llama.h" -#include "llama-kv-cache.h" -#include "llama-eviction-policy.h" - -#include -#include -#include -#include -#include -#include -#include - -#ifdef GGML_USE_HIP -#include -#endif - -// Tier types for KV cache -enum llama_cache_tier { - TIER_HOT, // VRAM - currently used context - TIER_WARM, // RAM - recently used context - TIER_COLD // SSD - less frequently accessed context -}; - -// Semantic fingerprint for KV cache tokens -struct semantic_fingerprint { - std::vector positions; // token positions this covers - std::vector embedding; // normalized embedding vector - llama_cache_tier tier; // TIER_WARM or TIER_COLD - uint64_t turn; // conversation turn when evicted -}; - -// Compression types for cold tier storage -enum llama_cache_compression { - COMPRESSION_NONE, - COMPRESSION_INT4, // 4-bit quantized (4x space reduction) - COMPRESSION_INT8, // 8-bit quantized (2x space reduction) - COMPRESSION_LZ4, // LZ4 compression - COMPRESSION_QUANTIZED // Model-native quantization -}; - -// Tier configuration -struct llama_tier_config { - float hot_percent; // Percentage of total context for hot tier (VRAM) - float warm_percent; // Percentage of total context for warm tier (RAM) - float cold_percent; // Percentage of total context for cold tier (SSD) - - uint32_t total_ctx; // Total context size - - uint32_t hot_capacity() const { return uint32_t(total_ctx * hot_percent / 100.0f); } - uint32_t warm_capacity() const { return uint32_t(total_ctx * warm_percent / 100.0f); } - uint32_t cold_capacity() const { return uint32_t(total_ctx * cold_percent / 100.0f); } - - // HIP device index for warm tier (-1 = RAM/SSD fallback, 1 = 6900XT) - int warm_device = -1; - - // Default configuration: 25% hot, 25% warm, 50% cold - static llama_tier_config default_config(uint32_t total_ctx) { - return {25.0f, 25.0f, 50.0f, total_ctx}; - } -}; - -// SSD storage format for cold tier -struct llama_ssd_storage_format { - // Magic header for validation - static constexpr uint32_t MAGIC = 0x4B565443; // "KVTC" - static constexpr uint32_t VERSION = 1; - - // File header structure - struct file_header { - uint32_t magic; - uint32_t version; - uint32_t n_layers; - uint32_t n_embd_k; - uint32_t n_embd_v; - uint32_t n_tokens; - uint32_t compression_type; - uint32_t layer_offset; // Starting token position for this layer - float attention_threshold; - uint64_t index_offset; // Offset to index section - uint32_t index_size; // Size of index in bytes - }; - - // Layer index entry for quick lookup - struct layer_index_entry { - uint32_t layer; - uint32_t n_tokens; - uint64_t file_offset; // Byte offset in file - uint32_t data_size; // Size of layer data in bytes - }; - - // Token range for a layer - struct token_range { - llama_pos start; - llama_pos end; - uint32_t n_tokens; - }; - - // Write KV cache to SSD - bool write(const std::string& path, - const float* k_data, - const float* v_data, - uint32_t n_tokens, - uint32_t n_layers, - uint32_t n_embd_k, - uint32_t n_embd_v, - llama_cache_compression compression = COMPRESSION_INT4); - - // Read KV cache from SSD - bool read(const std::string& path, - float* k_data, - float* v_data, - uint32_t n_tokens, - uint32_t n_layers, - uint32_t n_embd_k, - uint32_t n_embd_v); - - // Build index file for fast lookup - bool build_index(const std::string& index_path, - const std::string& data_path, - uint32_t n_layers); - - // Load index from file - bool load_index(const std::string& index_path); - - // Save index to file - bool save_index(const std::string& index_path); - - // Get layer data offset from index - uint64_t get_layer_offset(uint32_t layer) const; - - // Check if layer is in index - bool has_layer(uint32_t layer) const; - -private: - // Layer index for fast lookup - std::vector layer_index; - - // Helper for quantization - std::vector quantize_int4(const float* data, uint32_t n_elements); - std::vector quantize_int8(const float* data, uint32_t n_elements); - bool dequantize_int4(const uint8_t* data, float* out, uint32_t n_elements); - bool dequantize_int8(const uint8_t* data, float* out, uint32_t n_elements); -}; - -// Tier statistics for monitoring -struct llama_tier_stats { - uint32_t hot_tokens; - uint32_t warm_tokens; - uint32_t cold_tokens; - uint64_t eviction_count; - uint64_t cache_hits; - uint64_t cache_misses; - double total_migration_latency_us; - uint32_t semantic_prefetch_hits = 0; - uint32_t semantic_eviction_saves = 0; - - void reset() { - hot_tokens = 0; - warm_tokens = 0; - cold_tokens = 0; - eviction_count = 0; - cache_hits = 0; - cache_misses = 0; - total_migration_latency_us = 0.0; - semantic_prefetch_hits = 0; - semantic_eviction_saves = 0; - } - - double get_hit_rate() const { - uint64_t total = cache_hits + cache_misses; - return total > 0 ? double(cache_hits) / total : 0.0; - } -}; - -// Main tiered cache class -class llama_kv_cache_tiered { -public: - // Constructor - llama_kv_cache_tiered( - const llama_model& model, - const llama_tier_config& config, - const std::string& ssd_path, - llama_eviction_policy eviction_policy = HYBRID, - llama_cache_compression compression = COMPRESSION_INT4, - float attention_threshold = 0.1f); - - // Destructor - ~llama_kv_cache_tiered(); - - // Tier management - bool init(); - bool resize(uint32_t new_total_ctx); - bool set_eviction_policy(llama_eviction_policy policy); - bool set_attention_threshold(float threshold); - // Wire real KV layer tensors for actual data movement - void set_kv_layers_from_cache(llama_kv_cache * cache); - - // Eviction interface - bool evict_tokens(uint32_t n_tokens_to_evict, llama_cache_tier from_tier); - // Pre-register hot positions 0..n_hot-1 so eviction candidates are available - void track_hot_range(uint32_t n_hot); - - // Migration with tensor handles - requires actual KV cache tensor pointers - bool migrate_tokens(const std::vector& positions, - llama_cache_tier from_tier, - llama_cache_tier to_tier, - const ggml_tensor* k_tensor = nullptr, - const ggml_tensor* v_tensor = nullptr); - - // Migration optimization - bool batch_migrate_tokens(const std::vector& positions, - llama_cache_tier from_tier, - llama_cache_tier to_tier, - const ggml_tensor* k_tensor = nullptr, - const ggml_tensor* v_tensor = nullptr); - - bool prefetch_tokens(const std::vector& positions, - llama_cache_tier target_tier); - - // Token access - bool contains_token(llama_pos pos) const; - bool load_token(llama_pos pos, llama_cache_tier target_tier); - - // Prefetch hint for semantic similarity - struct PrefetchHint { - std::vector positions; - float score; - llama_cache_tier current_tier; - }; - - // Fingerprint management - void add_fingerprint(const std::vector& positions, - const std::vector& embedding, - llama_cache_tier tier); - bool save_fingerprints_to_disk(const std::string& path); - bool load_fingerprints_from_disk(const std::string& path); - - // Semantic similarity scoring - std::vector score_fingerprints( - const std::vector& query_emb, int top_k, float threshold) const; - void set_current_query_embedding(const std::vector& emb); - - // Statistics - llama_tier_stats get_stats() const; - void reset_stats(); - - // Configuration getters - const llama_tier_config& get_config() const { return config; } - llama_eviction_policy get_eviction_policy() const { return eviction_policy; } - float get_attention_threshold() const { return attention_threshold; } - -private: - // Tier capacities - llama_tier_config config; - - // Eviction policy - llama_eviction_policy eviction_policy; - - // SSD storage path - std::string ssd_path; - - // Compression type for cold tier - llama_cache_compression compression; - - // Attention threshold for eviction - float attention_threshold; - - // Token metadata store - llama_token_metadata_store token_metadata; - - // Statistics - llama_tier_stats stats; - - // Semantic fingerprints for evicted tokens - std::vector fingerprints; - uint64_t fingerprint_turn = 0; - - // Current query embedding for semantic eviction weighting - std::vector current_query_emb; - - // Thread safety - mutable std::mutex mutex; - - // Per-layer KV tensor info — populated by set_kv_layers_from_cache() - // K layout: [n_embd_k_gqa, kv_size, n_stream] — token p is contiguous at row p - // V layout: [n_embd_v_gqa, kv_size, n_stream] — when v_trans=true, token p's - // element j is at flat index (j*kv_size + p), requiring scatter/gather - struct TieredKVLayer { - ggml_tensor * k = nullptr; // raw K tensor in VRAM (or host) - ggml_tensor * v = nullptr; // raw V tensor in VRAM (or host) - bool v_trans = false; // true = V stored transposed (default) - int64_t kv_size = 0; // n_ctx_max (= k->ne[1]) - size_t k_bytes = 0; // bytes per token for K (contiguous row) - size_t v_bytes = 0; // bytes per token for V (de-transposed) - uint8_t * warm_k = nullptr; // warm_slots * k_bytes, host RAM - uint8_t * warm_v = nullptr; // warm_slots * v_bytes, host RAM (de-transposed) -#ifdef GGML_USE_HIP - void * warm_k_dev = nullptr; // optional: device (eGPU) warm K buffer - void * warm_v_dev = nullptr; // optional: device (eGPU) warm V buffer -#endif - }; - std::vector kv_layers; - - // Warm slot occupancy (shared across all layers) - struct WarmSlot { bool occupied = false; llama_pos pos = -1; }; - std::vector warm_slots; - std::unordered_map warm_pos_to_slot; - int warm_alloc_slot(); - void warm_free_slot(llama_pos pos); - - // Per-layer warm copy helpers: VRAM ↔ RAM (or VRAM ↔ device) - bool warm_copy_to_host(uint32_t il, int slot, llama_pos pos); - bool warm_copy_from_host(uint32_t il, int slot, llama_pos pos); -#ifdef GGML_USE_HIP - bool warm_copy_to_dev(uint32_t il, int slot, llama_pos pos); - bool warm_copy_from_dev(uint32_t il, int slot, llama_pos pos); -#endif - - // Helper methods - std::string get_ssd_file_path(llama_pos pos) const; - bool save_to_ssd(const std::vector& positions, const ggml_tensor* k_tensor, const ggml_tensor* v_tensor, bool is_device_data = false); - bool load_from_ssd(const std::vector& positions, ggml_tensor* k_tensor, ggml_tensor* v_tensor, bool to_device = false); - - // Tier capacity getters - uint32_t get_tier_capacity(llama_cache_tier tier) const; -}; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 192567510e2e..d3db1f2d3a5c 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -44,7 +44,10 @@ llama_memory_hybrid::llama_memory_hybrid( uint32_t paged_max_blocks_per_seq, uint32_t paged_n_warm_blocks, uint32_t paged_n_cold_blocks, - std::string paged_ssd_path) : + std::string paged_ssd_path, + bool paged_cold_resume, + std::string paged_instance_id, + uint32_t paged_cold_budget_mb) : hparams(model.hparams), mem_attn(paged_n_blocks > 0 ? nullptr : new llama_kv_cache( model, @@ -73,6 +76,9 @@ llama_memory_hybrid::llama_memory_hybrid( paged_n_warm_blocks, paged_n_cold_blocks, paged_ssd_path, + paged_cold_resume, + paged_instance_id, + paged_cold_budget_mb, // Filter out recurrent layers — paged cache only carries attention // K/V; recurrent state lives in mem_recr. Without this the cache // pre-allocates K/V for ALL layers (including recurrent ones that diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 51964dfd5594..f14536b7ce3c 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -56,7 +56,15 @@ class llama_memory_hybrid : public llama_memory_i { uint32_t paged_n_warm_blocks = 0, // MAD-121: cold-tier (SSD) capacity + path. 0 = cold disabled. uint32_t paged_n_cold_blocks = 0, - std::string paged_ssd_path = std::string()); + std::string paged_ssd_path = std::string(), + // MAD-130: when true, skip O_TRUNC on cold-tier files + // and load the in-memory index from the sidecar at + // ${paged_ssd_path}/paged/instance-${id}/index.bin. + bool paged_cold_resume = false, + // MAD-131: per-instance ID + cold-pool budget. Empty + // ID → use process pid. cold_budget_mb=0 → no cap. + std::string paged_instance_id = std::string(), + uint32_t paged_cold_budget_mb = 0); ~llama_memory_hybrid() = default; diff --git a/src/llama-memory.h b/src/llama-memory.h index de9214e457c3..b645ceea227a 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -39,6 +39,9 @@ struct llama_memory_params { int32_t kv_tier_semantic_topk; bool kv_tier_paged_blocks; // Phase 2a opt-in int32_t kv_tier_paged_block_size; // tokens/block (0 => default 16) + bool kv_tier_cold_resume; // MAD-130: skip O_TRUNC on cold-tier files; load index sidecar + const char * kv_tier_instance_id; // MAD-131: per-instance ID for cold subdir + lockfile (nullptr => pid) + int32_t kv_tier_cold_budget_mb; // MAD-131: cap cold pool to N MiB (0 => no limit) }; enum llama_memory_status { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index ed89938f1a80..b261ed56b241 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2176,7 +2176,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* paged_max_blocks_per_seq */ paged_max_bps, /* paged_n_warm_blocks */ paged_n_warm_blocks, /* paged_n_cold_blocks */ paged_n_cold_blocks, - /* paged_ssd_path */ paged_ssd_path); + /* paged_ssd_path */ paged_ssd_path, + /* paged_cold_resume */ params.kv_tier_cold_resume, + /* paged_instance_id */ params.kv_tier_instance_id ? std::string(params.kv_tier_instance_id) : std::string(), + /* paged_cold_budget_mb */ (uint32_t) std::max(0, params.kv_tier_cold_budget_mb)); } } else { llama_memory_i::layer_reuse_cb reuse = nullptr; diff --git a/src/memory-tier/mt-block-pool.cpp b/src/memory-tier/mt-block-pool.cpp index 8fe98529d8c7..9784d41811de 100644 --- a/src/memory-tier/mt-block-pool.cpp +++ b/src/memory-tier/mt-block-pool.cpp @@ -32,6 +32,9 @@ void BlockPool::init(uint32_t n_gpu, uint32_t n_cpu, float watermark) { cpu_free_.push_back(n_gpu + i); } + // MAD-128: per-block refcount, all start at 0 (free). + refcount_.assign((size_t) n_gpu + n_cpu, 0); + LLAMA_LOG_INFO("mt::BlockPool: init n_gpu=%u n_cpu=%u watermark=%.2f " "(reserve gpu=%u cpu=%u)\n", n_gpu, n_cpu, (double) watermark, @@ -42,6 +45,8 @@ uint32_t BlockPool::alloc_gpu() { if (gpu_free_.empty()) return kInvalidBlockId; const uint32_t id = gpu_free_.back(); gpu_free_.pop_back(); + assert(id < refcount_.size() && refcount_[id] == 0 && "alloc'd block had nonzero refcount"); + refcount_[id] = 1; return id; } @@ -49,15 +54,44 @@ uint32_t BlockPool::alloc_cpu() { if (cpu_free_.empty()) return kInvalidBlockId; const uint32_t id = cpu_free_.back(); cpu_free_.pop_back(); + assert(id < refcount_.size() && refcount_[id] == 0 && "alloc'd block had nonzero refcount"); + refcount_[id] = 1; return id; } +void BlockPool::bump_ref(uint32_t block_id) { + if (block_id == kInvalidBlockId) { + LLAMA_LOG_WARN("mt::BlockPool::bump_ref: kInvalidBlockId — ignoring\n"); + return; + } + assert(block_id < refcount_.size() && "bump_ref: block_id out of range"); + assert(refcount_[block_id] > 0 && "bump_ref: block is free (refcount==0)"); + ++refcount_[block_id]; +} + +uint32_t BlockPool::refcount(uint32_t block_id) const { + if (block_id == kInvalidBlockId) return 0; + if (block_id >= refcount_.size()) return 0; + return refcount_[block_id]; +} + void BlockPool::free_block(uint32_t block_id) { if (block_id == kInvalidBlockId) { LLAMA_LOG_WARN("mt::BlockPool::free_block: kInvalidBlockId — ignoring\n"); return; } + assert(block_id < refcount_.size() && "free_block: block_id out of range"); + if (refcount_[block_id] == 0) { + LLAMA_LOG_WARN("mt::BlockPool::free_block: double-free of block %u\n", block_id); + return; + } + --refcount_[block_id]; + if (refcount_[block_id] > 0) { + // Still has other references — don't return to free stack yet. + return; + } + if (is_gpu(block_id)) { assert(block_id < total_gpu_blocks_); gpu_free_.push_back(block_id); @@ -102,6 +136,8 @@ void BlockPool::reset() { for (uint32_t i = total_cpu_blocks_; i-- > 0; ) { cpu_free_.push_back(total_gpu_blocks_ + i); } + // MAD-128: zero all refcounts on whole-pool reset. + std::fill(refcount_.begin(), refcount_.end(), 0u); } } // namespace mt diff --git a/src/memory-tier/mt-block-pool.h b/src/memory-tier/mt-block-pool.h index f538a0c3c117..8bf3a0659559 100644 --- a/src/memory-tier/mt-block-pool.h +++ b/src/memory-tier/mt-block-pool.h @@ -55,12 +55,30 @@ class BlockPool { // first). Watermark is NOT enforced here — callers that care about // admission control should consult has_free_gpu_blocks() up-front. // alloc_*() always allocates if anything is free. + // + // The returned block has refcount=1. Use bump_ref() to share the + // block between sequences (CoW); free_block() decrements and only + // returns the block to the free stack when refcount drops to 0. uint32_t alloc_gpu(); uint32_t alloc_cpu(); - // Return a block to its pool. Idempotent on already-free blocks - // (logs a warning); double-free does NOT corrupt the pool. Asserts - // on out-of-range IDs. + // MAD-128: increment refcount on `block_id`. Used by seq_cp to share + // a physical block across two sequences without copying its bytes. + // The next free_block() call on this id decrements; the block goes + // back to the free stack only when the last reference is freed. + // Asserts on out-of-range or already-free IDs (refcount==0). + void bump_ref(uint32_t block_id); + + // MAD-128: query refcount. 0 means the block is free; 1 means + // single-owner; >1 means shared. Used by the CoW write trigger in + // llama_kv_cache_paged to decide whether a write needs to allocate + // a fresh block first. + uint32_t refcount(uint32_t block_id) const; + + // Return a block to its pool. Decrements refcount; only actually + // returns the block to the free stack when refcount drops to 0. + // Idempotent on already-free blocks (logs a warning); double-free + // does NOT corrupt the pool. Asserts on out-of-range IDs. void free_block(uint32_t block_id); // True if `block_id` is a GPU block (vs CPU). Inferred from the ID @@ -92,6 +110,12 @@ class BlockPool { std::vector gpu_free_; std::vector cpu_free_; + // MAD-128: per-block refcount (parallel to free stacks). Indexed + // by block_id (covers both GPU and CPU id ranges). 0 = free; 1 = + // single-owner; >1 = shared via seq_cp. alloc_*() sets to 1; + // bump_ref() increments; free_block() decrements (actual free at 0). + std::vector refcount_; + uint32_t total_gpu_blocks_ = 0; uint32_t total_cpu_blocks_ = 0; uint32_t watermark_gpu_ = 0; // reserve count, NOT a fraction diff --git a/src/memory-tier/mt-eviction.h b/src/memory-tier/mt-eviction.h index 05269491cc7d..132df1d92c3a 100644 --- a/src/memory-tier/mt-eviction.h +++ b/src/memory-tier/mt-eviction.h @@ -12,10 +12,8 @@ // Attention — tokens with lowest attention scores evict first // Hybrid — weighted combination of the above (default) // -// Threadsafe via internal mutex. Replaces the legacy -// llama_token_metadata_store in src/llama-eviction-policy.h. The -// numeric eviction-policy mapping (0..3) matches the legacy CLI for -// backward compat. +// Threadsafe via internal mutex. The numeric eviction-policy mapping +// (0..3) matches the long-standing CLI flag values. #include "mt-config.h" #include "llama.h" // llama_pos diff --git a/src/memory-tier/mt-mover-attn.h b/src/memory-tier/mt-mover-attn.h index 3c7a89bf4ee3..0b0f6340afb8 100644 --- a/src/memory-tier/mt-mover-attn.h +++ b/src/memory-tier/mt-mover-attn.h @@ -9,10 +9,8 @@ // evict_v VRAM[layer.v @ pos] -> host buf (gather if v_trans) // restore_v host buf -> VRAM[layer.v @ pos] (scatter if v_trans) // -// Synchronous via hipMemcpy / hipMemcpy2D — same pattern as the legacy -// pager and the legacy llama_kv_cache_tiered. Phase 2 follow-up will -// add async overlap with proper stream-event ordering once correctness -// is locked. +// Synchronous via hipMemcpy / hipMemcpy2D. Async overlap with proper +// stream-event ordering is a follow-up once correctness is locked. // // Layouts (ggml semantics): // K is always contiguous: token p starts at byte offset p * k_row_bytes. diff --git a/src/memory-tier/mt-mover-recurrent.h b/src/memory-tier/mt-mover-recurrent.h index efb40c32dfb3..3b4e6c5025cd 100644 --- a/src/memory-tier/mt-mover-recurrent.h +++ b/src/memory-tier/mt-mover-recurrent.h @@ -11,8 +11,9 @@ // of token count). For Qwen3.5-REAP-97B with 36 recurrent layers // this is ~149 MiB per seq. // -// New work in Phase 2 — the legacy llama_kv_cache_tiered did not tier -// recurrent state at all (B-K5 in the bug catalog). +// Recurrent state tiering is mt::-only — the paged cache itself only +// covers attention layers; recurrent state still flows through the +// wrapper for hybrid models. // // HIP-only. Stub for non-HIP builds (return false everywhere). diff --git a/src/memory-tier/mt-quant.cpp b/src/memory-tier/mt-quant.cpp index 6e2aafacd56f..6f965d9ee326 100644 --- a/src/memory-tier/mt-quant.cpp +++ b/src/memory-tier/mt-quant.cpp @@ -94,4 +94,56 @@ bool dequantize_int8(const uint8_t * src, float * dst, size_t n) { return true; } +// MAD-135: per-block int4 with explicit scale. +// +// Algorithm: +// 1. Compute scale = max(|x[i]|) over the block. If 0, set scale=1 +// (all output nibbles will encode 0 anyway). +// 2. For each x[i]: normalized = x[i] / scale; encode to int4. +// 3. Output: scale (one float), then ceil(n/2) packed bytes. +// +// On dequant: multiply each decoded value by scale to recover the +// original range. Quantization error is bounded by 1/(7 * 2) = 7.1% +// of the per-block max-abs (int4 has 16 levels, so half-step is ~7%). +// For cold-tier KV restoration this is acceptable — the original K/V +// values were going to be re-attended-to with their original-precision +// neighbors anyway, so per-block error doesn't compound across the +// attention sum. +bool quantize_block_int4_with_scale(const float * src, size_t n, + float * scale_out, uint8_t * dst) { + if (n == 0 || src == nullptr || scale_out == nullptr || dst == nullptr) return false; + + float scale = 0.0f; + for (size_t i = 0; i < n; ++i) { + const float a = std::fabs(src[i]); + if (a > scale) scale = a; + } + if (scale == 0.0f) scale = 1.0f; // all zero block; encoded nibbles will be zero anyway + + *scale_out = scale; + const float inv = 1.0f / scale; + + for (size_t i = 0; i < n; i += 2) { + const uint8_t lo = encode_int4(src[i] * inv); + const uint8_t hi = (i + 1 < n) ? encode_int4(src[i + 1] * inv) : 0u; + dst[i / 2] = (uint8_t)(lo | (hi << 4)); + } + return true; +} + +bool dequantize_block_int4_with_scale(const uint8_t * src, float scale_in, + float * dst, size_t n) { + if (n == 0) return true; + if (src == nullptr || dst == nullptr) return false; + + for (size_t i = 0; i < n; i += 2) { + const uint8_t byte = src[i / 2]; + dst[i] = decode_int4((uint8_t)(byte & 0x0F)) * scale_in; + if (i + 1 < n) { + dst[i + 1] = decode_int4((uint8_t)((byte >> 4) & 0x0F)) * scale_in; + } + } + return true; +} + } // namespace mt diff --git a/src/memory-tier/mt-quant.h b/src/memory-tier/mt-quant.h index cbf3eb0a7954..749817a87f59 100644 --- a/src/memory-tier/mt-quant.h +++ b/src/memory-tier/mt-quant.h @@ -6,11 +6,10 @@ // for inputs already in roughly [-1, +1]. Quality is "fine for cold // tier" — you would not use these for active inference. // -// IMPORTANT: the legacy llama_ssd_storage_format helpers in -// src/llama-kv-cache-tiered.cpp had a sign-bias bug — storing -8 -// round-tripped to 0 because the encoder packed `(v & 0x0F)` instead -// of `(v + 8) & 0x0F`. The new helpers here use the correct +8 bias -// so the full [-8, +7] range round-trips faithfully. +// IMPORTANT (history): an earlier int4 implementation had a sign-bias +// bug — storing -8 round-tripped to 0 because the encoder packed +// `(v & 0x0F)` instead of `(v + 8) & 0x0F`. These helpers use the +// correct +8 bias so the full [-8, +7] range round-trips faithfully. #include #include @@ -35,4 +34,22 @@ std::vector quantize_int8(const float * src, size_t n); // Decode `n` 8-bit signed values from `src` back to floats. bool dequantize_int8(const uint8_t * src, float * dst, size_t n); +// MAD-135: per-block int4 with explicit scale. Reads `n` floats from +// `src`, computes scale = max(abs(x[i])), normalizes each into +// [-1, +1], int4-quantizes. Output: `*scale_out` is the per-block +// scale; `dst` (size at least ceil(n/2)) holds packed int4 nibbles. +// scale==0 means the entire block is zero — quant bytes are still +// written (all zero). +// +// Cold-tier use case: bound the dynamic range without an external +// scale assumption. F16 KV values in attention can have magnitudes +// well outside [-1, +1]; per-block scale recovers a usable range. +bool quantize_block_int4_with_scale(const float * src, size_t n, + float * scale_out, uint8_t * dst); + +// Inverse: reads `*scale_in` + ceil(n/2) packed int4 bytes, +// dequantizes into `dst` (size at least n). +bool dequantize_block_int4_with_scale(const uint8_t * src, float scale_in, + float * dst, size_t n); + } // namespace mt diff --git a/src/memory-tier/mt-semantic.cpp b/src/memory-tier/mt-semantic.cpp index 3dc0a5e3f23f..37bf93111753 100644 --- a/src/memory-tier/mt-semantic.cpp +++ b/src/memory-tier/mt-semantic.cpp @@ -215,4 +215,214 @@ bool SemanticIndex::load_from_disk(const std::string & path) { return true; } +// --------------------------------------------------------------------------- +// BlockSemanticIndex — paged-block-keyed fingerprint store for MAD-122. +// --------------------------------------------------------------------------- + +void BlockSemanticIndex::add_fingerprint(llama_seq_id seq_id, + uint32_t lblock, + std::vector embedding, + SemanticIndex::Tier tier) { + std::lock_guard lk(mu_); + auto & seq_map = fps_[seq_id]; + auto & e = seq_map[lblock]; + e.embedding = std::move(embedding); + e.tier = tier; +} + +void BlockSemanticIndex::update_tier(llama_seq_id seq_id, uint32_t lblock, + SemanticIndex::Tier tier) { + std::lock_guard lk(mu_); + auto sit = fps_.find(seq_id); + if (sit == fps_.end()) return; + auto bit = sit->second.find(lblock); + if (bit == sit->second.end()) return; + bit->second.tier = tier; +} + +void BlockSemanticIndex::remove_block(llama_seq_id seq_id, uint32_t lblock) { + std::lock_guard lk(mu_); + auto sit = fps_.find(seq_id); + if (sit == fps_.end()) return; + sit->second.erase(lblock); + if (sit->second.empty()) fps_.erase(sit); +} + +bool BlockSemanticIndex::has_fingerprint(llama_seq_id seq_id, uint32_t lblock) const { + std::lock_guard lk(mu_); + auto sit = fps_.find(seq_id); + if (sit == fps_.end()) return false; + return sit->second.find(lblock) != sit->second.end(); +} + +void BlockSemanticIndex::remove_seq(llama_seq_id seq_id) { + std::lock_guard lk(mu_); + fps_.erase(seq_id); +} + +void BlockSemanticIndex::clear() { + std::lock_guard lk(mu_); + fps_.clear(); +} + +std::vector +BlockSemanticIndex::score(llama_seq_id seq_id, + const std::vector & query, + int top_k, + float threshold) const { + std::vector out; + if (query.empty() || top_k <= 0) return out; + + std::lock_guard lk(mu_); + auto sit = fps_.find(seq_id); + if (sit == fps_.end() || sit->second.empty()) return out; + + std::vector> scored; + scored.reserve(sit->second.size()); + for (const auto & kv : sit->second) { + const float s = dot(query, kv.second.embedding); + scored.emplace_back(s, kv.first); + } + + std::sort(scored.begin(), scored.end(), + [](const auto & a, const auto & b) { return a.first > b.first; }); + + out.reserve((size_t) top_k); + for (const auto & [s, lblock] : scored) { + if (s < threshold) break; + BlockHint h; + h.seq_id = seq_id; + h.lblock = lblock; + h.score = s; + h.tier = sit->second.at(lblock).tier; + out.push_back(std::move(h)); + if ((int) out.size() >= top_k) break; + } + return out; +} + +size_t BlockSemanticIndex::size() const { + std::lock_guard lk(mu_); + size_t n = 0; + for (const auto & kv : fps_) n += kv.second.size(); + return n; +} + +size_t BlockSemanticIndex::size(llama_seq_id seq_id) const { + std::lock_guard lk(mu_); + auto it = fps_.find(seq_id); + return it == fps_.end() ? 0 : it->second.size(); +} + +// --------------------------------------------------------------------------- +// MAD-130: persistence (PSFI v1 — Paged Semantic Fingerprint Index). +// +// File format: +// uint32 magic = 0x49465350 ("PSFI" little-endian) +// uint32 version = 1 +// uint32 n_seqs +// For each seq: +// int32 seq_id +// uint32 n_entries +// For each entry: +// uint32 lblock +// uint8 tier (0=Hot, 1=Warm, 2=Cold) +// uint32 embedding_dim +// float[embedding_dim] +// --------------------------------------------------------------------------- + +namespace { +constexpr uint32_t kBlockFingerprintFileMagic = 0x49465350; // "PSFI" +constexpr uint32_t kBlockFingerprintFileVersion = 1; +} + +bool BlockSemanticIndex::save_to_disk(const std::string & path) const { + std::lock_guard lk(mu_); + + std::ofstream f(path, std::ios::binary | std::ios::trunc); + if (!f) { + LLAMA_LOG_WARN("mt::BlockSemanticIndex::save: open(%s) failed\n", path.c_str()); + return false; + } + + auto write_u32 = [&](uint32_t v) { f.write((const char *) &v, sizeof(v)); }; + auto write_i32 = [&](int32_t v) { f.write((const char *) &v, sizeof(v)); }; + auto write_u8 = [&](uint8_t v) { f.write((const char *) &v, sizeof(v)); }; + + write_u32(kBlockFingerprintFileMagic); + write_u32(kBlockFingerprintFileVersion); + write_u32((uint32_t) fps_.size()); + + for (const auto & [sid, entries] : fps_) { + write_i32((int32_t) sid); + write_u32((uint32_t) entries.size()); + for (const auto & [lblock, entry] : entries) { + write_u32(lblock); + write_u8((uint8_t) entry.tier); + write_u32((uint32_t) entry.embedding.size()); + f.write((const char *) entry.embedding.data(), + (std::streamsize)(entry.embedding.size() * sizeof(float))); + } + } + return f.good(); +} + +bool BlockSemanticIndex::load_from_disk(const std::string & path) { + std::ifstream f(path, std::ios::binary); + if (!f) { + // Not an error — file may not exist on first run. + return false; + } + + auto read_u32 = [&](uint32_t & v) -> bool { + f.read((char *) &v, sizeof(v)); + return (bool) f; + }; + auto read_i32 = [&](int32_t & v) -> bool { + f.read((char *) &v, sizeof(v)); + return (bool) f; + }; + auto read_u8 = [&](uint8_t & v) -> bool { + f.read((char *) &v, sizeof(v)); + return (bool) f; + }; + + uint32_t magic = 0, version = 0; + if (!read_u32(magic) || !read_u32(version)) return false; + if (magic != kBlockFingerprintFileMagic || version != kBlockFingerprintFileVersion) { + LLAMA_LOG_WARN("mt::BlockSemanticIndex::load: bad header (magic=%08x ver=%u) in %s\n", + magic, version, path.c_str()); + return false; + } + + uint32_t n_seqs = 0; + if (!read_u32(n_seqs)) return false; + + decltype(fps_) loaded; + for (uint32_t s = 0; s < n_seqs; ++s) { + int32_t sid = 0; + uint32_t n_entries = 0; + if (!read_i32(sid) || !read_u32(n_entries)) return false; + + auto & seq_map = loaded[(llama_seq_id) sid]; + for (uint32_t e = 0; e < n_entries; ++e) { + uint32_t lblock = 0, emb_dim = 0; + uint8_t tier = 0; + if (!read_u32(lblock) || !read_u8(tier) || !read_u32(emb_dim)) return false; + + Entry entry; + entry.tier = (SemanticIndex::Tier) tier; + entry.embedding.resize(emb_dim); + f.read((char *) entry.embedding.data(), + (std::streamsize)(emb_dim * sizeof(float))); + if (!f) return false; + seq_map[lblock] = std::move(entry); + } + } + + std::lock_guard lk(mu_); + fps_ = std::move(loaded); + return true; +} + } // namespace mt diff --git a/src/memory-tier/mt-semantic.h b/src/memory-tier/mt-semantic.h index 4ec510de01a3..92dff939b441 100644 --- a/src/memory-tier/mt-semantic.h +++ b/src/memory-tier/mt-semantic.h @@ -29,6 +29,7 @@ #include #include #include +#include #include namespace mt { @@ -87,4 +88,88 @@ class SemanticIndex { uint64_t next_turn_ = 0; }; +// BlockSemanticIndex — paged-block-keyed fingerprint store for MAD-122. +// +// Parallel to SemanticIndex but keys fingerprints by (seq_id, logical_block_idx) +// so a query can find which paged blocks of a given seq are most semantically +// similar. The paged path stores at most one fingerprint per block (the 16-tok +// block size is small enough that one embedding per block is reasonable), +// matched 1:1 with BlockTable's lifecycle: blocks come and go with the seq; +// fingerprints are dropped on whole-seq wipe (mt::seq_rm with sentinel range). +// +// No FIFO cap: lifecycle is tied to BlockTable, so memory grows with the +// active context size rather than indefinitely. For the army goal (4 seqs × +// 8k blocks/seq × 384-dim fp32) the worst-case footprint is ~48 MiB — +// acceptable for a CPU-side index and well below the warm-tier staging cost. +class BlockSemanticIndex { +public: + struct BlockHint { + llama_seq_id seq_id = -1; + uint32_t lblock = 0; + float score = 0.0f; + SemanticIndex::Tier tier = SemanticIndex::Tier::Warm; + }; + + BlockSemanticIndex() = default; + + // Store the fingerprint for (seq_id, lblock). Overwrites any prior + // fingerprint at the same key — useful if a block is re-fingerprinted + // after a partial-range edit. embedding SHOULD be L2-normalized + // (caller's responsibility); scoring degrades to dot-product if it + // isn't. + void add_fingerprint(llama_seq_id seq_id, + uint32_t lblock, + std::vector embedding, + SemanticIndex::Tier tier); + + // Update only the tier annotation (e.g. when a block migrates + // hot→warm→cold). No-op if the (seq, lblock) isn't tracked. + void update_tier(llama_seq_id seq_id, uint32_t lblock, SemanticIndex::Tier tier); + + // Drop a single (seq, lblock) entry. No-op if not tracked. + void remove_block(llama_seq_id seq_id, uint32_t lblock); + + // MAD-129: O(1) check whether a fingerprint already exists for + // (seq_id, lblock). Used by the server's prefill-time write trigger + // to skip blocks that have already been fingerprinted (typical when + // a turn's prompt re-processes a prior turn's accumulated context). + bool has_fingerprint(llama_seq_id seq_id, uint32_t lblock) const; + + // Drop every fingerprint for `seq_id`. Called on whole-seq wipe. + void remove_seq(llama_seq_id seq_id); + + // Drop everything. + void clear(); + + // Score `seq_id`'s blocks against `query_embedding`. Returns up to + // `top_k` blocks with cosine similarity >= `threshold`, sorted by + // descending score. Blocks from other seqs are not considered — + // semantic prefetch is per-seq because cross-seq attention isn't a + // thing in the paged-attention model. + std::vector score(llama_seq_id seq_id, + const std::vector & query_embedding, + int top_k, + float threshold) const; + + // Diagnostics. + size_t size() const; + size_t size(llama_seq_id seq_id) const; + + // MAD-130: persistence. Format magic = "PSFI" v1. Saves all + // (seq_id, lblock, tier, embedding) tuples to the path. Returns + // false on I/O error. load_from_disk replaces the in-memory + // state with the file's contents (any prior data is dropped). + bool save_to_disk(const std::string & path) const; + bool load_from_disk(const std::string & path); + +private: + struct Entry { + std::vector embedding; + SemanticIndex::Tier tier = SemanticIndex::Tier::Warm; + }; + + mutable std::mutex mu_; + std::unordered_map> fps_; +}; + } // namespace mt diff --git a/src/memory-tier/mt-tiered.cpp b/src/memory-tier/mt-tiered.cpp index 8e12f826bfb3..633e579da0dc 100644 --- a/src/memory-tier/mt-tiered.cpp +++ b/src/memory-tier/mt-tiered.cpp @@ -67,30 +67,15 @@ llama_memory_tiered::llama_memory_tiered(llama_memory_ptr inner, tier_view_.recur_seqs.size(), tier_view_.empty() ? " (not tierable; will run as passthrough)" : ""); - // Phase 2a: paged-blocks scaffolding. When the flag is set, allocate - // BlockPool + BlockTable sized to the warm/cold tier capacities so - // we can start exercising allocation patterns. NO read/write paths - // touch these structures yet — Phase 2b/2c wire them in. This lets - // us validate that the structures behave under realistic sizing - // (right number of blocks, init logs make sense, no mem blowup) - // before changing any behavior. - if (cfg_.paged_blocks) { - const uint32_t bsize = cfg_.paged_block_size > 0 ? cfg_.paged_block_size : 16u; - // Hot tier corresponds to "GPU blocks" in vLLM terminology. - // Warm tier (RAM) is "CPU blocks." Cold tier (NVMe) doesn't get - // a block ID — it's KvtcStore-keyed and managed separately. - const uint32_t n_gpu = capacity_.hot_capacity() / bsize; - const uint32_t n_cpu = capacity_.warm_capacity() / bsize; - // Watermark default of 0.05 mirrors vLLM. Scaffolding only — - // no admission control consults it yet. - paged_pool_.init(n_gpu, n_cpu, /*watermark=*/0.05f); - paged_table_.init(n_seq_max_, bsize); - LLAMA_LOG_INFO("mt::llama_memory_tiered: paged-blocks scaffolding " - "ON (block_size=%u, n_gpu_blocks=%u, n_cpu_blocks=%u, " - "n_seq_max=%u) — no live wiring yet, Phase 2b will " - "start using them\n", - bsize, n_gpu, n_cpu, n_seq_max_); - } + // MAD-127: paged-blocks scaffolding removed. For the army-goal config + // (hybrid + paged) the active tier layer is llama_kv_cache_paged + // itself; the wrapper stays a thin shim for bge-small embedding + // ownership and recurrent-state backup. + + // MAD-134: warm the bge-small embed model at construction so the + // first user prompt doesn't pay the lazy-load cost. No-op when + // semantic_index isn't configured. + warmup_embed_(); } llama_memory_tiered::~llama_memory_tiered() { @@ -360,65 +345,6 @@ bool llama_memory_tiered::ensure_warm_staging() { return true; } -bool llama_memory_tiered::ensure_paged_warm_staging() { - if (paged_warm_initialized_) return true; - if (!cfg_.paged_blocks) return false; - - const uint32_t bsize = paged_table_.block_size(); - if (bsize == 0) { - LLAMA_LOG_WARN("mt::ensure_paged_warm_staging: paged_table_.block_size()=0; " - "did paged_table_.init() run?\n"); - return false; - } - - // Restorable layers — same selection criteria as warm_buf_. - size_t restorable_layers = 0; - for (const auto & c : tier_view_.attn_caches) { - if (!c.is_swa) restorable_layers += c.layers.size(); - } - if (restorable_layers == 0) { - LLAMA_LOG_WARN("mt::ensure_paged_warm_staging: no restorable attention " - "layers; paged backup is a no-op for this model\n"); - return false; - } - - // Per-block layout: for each restorable layer, allocate a K slab - // (block_size * k_row_bytes) followed by a V slab (block_size * - // v_row_bytes). Sum across all restorable layers gives the total - // bytes per block. Layers in the same order as ensure_warm_staging - // so the per-layer mover code can be reused later. - paged_layer_off_.resize(restorable_layers); - paged_layer_v_off_.resize(restorable_layers); - size_t cursor = 0; - size_t flat_idx = 0; - for (const auto & c : tier_view_.attn_caches) { - if (c.is_swa) continue; - for (const auto & a : c.layers) { - paged_layer_off_[flat_idx] = cursor; - paged_layer_v_off_[flat_idx] = cursor + (size_t) bsize * a.k_row_bytes; - cursor += (size_t) bsize * (a.k_row_bytes + a.v_row_bytes); - ++flat_idx; - } - } - paged_block_bytes_ = cursor; - - // Allocate the host buffer for ALL CPU blocks. n_cpu_blocks is the - // CPU-side pool size set in the constructor (warm_capacity / bsize). - const uint32_t n_cpu_blocks = paged_pool_.total_cpu_blocks(); - paged_warm_buf_.assign((size_t) n_cpu_blocks * paged_block_bytes_, 0); - - LLAMA_LOG_INFO("mt::ensure_paged_warm_staging: bsize=%u n_cpu_blocks=%u " - "block_bytes=%.1f KiB total=%.1f MiB across %zu restorable " - "layers\n", - bsize, n_cpu_blocks, - (double) paged_block_bytes_ / 1024.0, - (double) paged_warm_buf_.size() / (1024.0 * 1024.0), - restorable_layers); - - paged_warm_initialized_ = true; - return true; -} - std::vector llama_memory_tiered::embed_text(const std::string & text) { if (cfg_.semantic_index.empty()) { return {}; @@ -429,6 +355,27 @@ std::vector llama_memory_tiered::embed_text(const std::string & text) { return embed_model_->embed(text); } +// MAD-134: warm the bge-small model at construction so the first user +// prompt doesn't pay the lazy-load cost (~200ms on cold start). Called +// from the ctor right after embed_model_ would lazily come up. Logs +// the latency so operators can see it happened. Failures are non-fatal +// (the lazy path keeps working). +void llama_memory_tiered::warmup_embed_() { + if (cfg_.semantic_index.empty()) return; + const auto t0 = std::chrono::steady_clock::now(); + auto v = embed_text("warmup"); + const auto t1 = std::chrono::steady_clock::now(); + const auto ms = std::chrono::duration_cast(t1 - t0).count(); + if (v.empty()) { + LLAMA_LOG_WARN("mt::llama_memory_tiered: bge-small warmup returned empty embedding " + "(model load failed or degenerate input); semantic prefetch will lazy-init " + "on first real call instead\n"); + } else { + LLAMA_LOG_INFO("mt::llama_memory_tiered: bge-small warmup complete in %lldms (n_embd=%d)\n", + (long long) ms, (int) v.size()); + } +} + bool llama_memory_tiered::has_warm_recurrent(llama_seq_id seq_id) const { return warm_recur_buf_.find(seq_id) != warm_recur_buf_.end(); } @@ -485,11 +432,6 @@ bool llama_memory_tiered::restore_recurrent_from_warm(llama_seq_id seq_id) { // ---- public tier-restore API ---- bool llama_memory_tiered::has_warm(llama_seq_id seq_id, llama_pos position) const { - // Phase 2b-C: hard-branch to paged variant when flag is set. - if (cfg_.paged_blocks) { - return paged_has_warm(seq_id, position); - } - // "warm" in the API name is historical — this query returns true for // anything in warm OR cold for the given seq, since restore_from_warm // transparently demand-loads from cold via load_one_from_cold. @@ -498,26 +440,8 @@ bool llama_memory_tiered::has_warm(llama_seq_id seq_id, llama_pos position) cons || cold_positions_[seq_id].find(position) != cold_positions_[seq_id].end(); } -bool llama_memory_tiered::paged_has_warm(llama_seq_id seq_id, llama_pos position) const { - if (!paged_warm_initialized_) return false; - const uint32_t bsize = paged_table_.block_size(); - if (bsize == 0) return false; - const uint32_t lblock = (uint32_t) position / bsize; - const uint32_t physical = paged_table_.get_physical(seq_id, lblock); - if (physical == kInvalidBlockId) return false; - // GPU blocks are still in hot — caller doesn't need to restore. - // CPU blocks are in warm — restore-able. - // (Cold tier not yet wired into the paged path; comes in a later phase.) - return !paged_pool_.is_gpu(physical); -} - uint32_t llama_memory_tiered::restore_from_warm(llama_seq_id seq_id, const std::vector & positions) { - // Phase 2b-C: hard-branch to paged variant when flag is set. - if (cfg_.paged_blocks) { - return paged_restore_from_warm(seq_id, positions); - } - if (!inner_ || !warm_initialized_ || positions.empty()) return 0; if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max_) return 0; auto & seq_warm = warm_pos_to_slot_[seq_id]; @@ -601,92 +525,6 @@ uint32_t llama_memory_tiered::restore_from_warm(llama_seq_id s return restored; } -// Phase 2b-C: paged-blocks restore. For each requested position, looks -// up the physical block via paged_table_, copies K/V from -// paged_warm_buf_ back into the inner cache via mover_attn_.restore_*. -// CPU block stays mapped after restore — same chunk may be requested -// again by a related query (semantic prefetch). Phase 2c will validate -// parity by toggling cfg_.paged_blocks and confirming identical -// observable behavior. -uint32_t llama_memory_tiered::paged_restore_from_warm( - llama_seq_id seq_id, - const std::vector & positions) { - if (!inner_ || !paged_warm_initialized_ || positions.empty()) return 0; - - const uint32_t bsize = paged_table_.block_size(); - if (bsize == 0) return 0; - - // Group requested positions by their containing logical block. We - // restore at block granularity — the inner cache slots get tagged - // per position, but the block lookup happens once per block. - uint32_t restored = 0; - - for (auto pos : positions) { - if (pos < 0) continue; - const uint32_t lblock = (uint32_t) pos / bsize; - const uint32_t physical = paged_table_.get_physical(seq_id, lblock); - if (physical == kInvalidBlockId) continue; // never backed up - if (paged_pool_.is_gpu(physical)) continue; // still in hot, no-op - - const uint32_t cpu_block_idx = physical - paged_pool_.total_gpu_blocks(); - const uint8_t * block_base = paged_warm_buf_.data() + - (size_t) cpu_block_idx * paged_block_bytes_; - const uint32_t row_in_block = (uint32_t) pos % bsize; - - // Ask inner cache for a free slot tagged with this position. - const int inner_slot = inner_->mt_restore_tag_slot(seq_id, pos); - if (inner_slot < 0) { - LLAMA_LOG_WARN("mt::paged_restore_from_warm: no free slot in inner " - "cache for pos %d (seq %d, lblock=%u)\n", - pos, seq_id, lblock); - break; - } - - // Walk restorable layers in the same flat order as - // ensure_paged_warm_staging used to lay out paged_layer_off_. - bool ok = true; - size_t flat_idx = 0; - for (const auto & c : tier_view_.attn_caches) { - if (c.is_swa) continue; - for (const auto & layer : c.layers) { - const uint8_t * src_k = block_base - + paged_layer_off_[flat_idx] - + (size_t) row_in_block * layer.k_row_bytes; - const uint8_t * src_v = block_base - + paged_layer_v_off_[flat_idx] - + (size_t) row_in_block * layer.v_row_bytes; - if (!mover_attn_.restore_k(layer, src_k, inner_slot) || - !mover_attn_.restore_v(layer, src_v, inner_slot)) { - ok = false; - break; - } - ++flat_idx; - } - if (!ok) break; - } - - if (ok) { - ++restored; - } else { - LLAMA_LOG_WARN("mt::paged_restore_from_warm: mover failed restoring " - "pos %d into slot %d (lblock=%u, cpu_block=%u)\n", - pos, inner_slot, lblock, physical); - } - } - - if (restored > 0) { - capacity_.on_migrate(restored, - TierCapacityManager::Tier::Warm, - TierCapacityManager::Tier::Hot); - update_tier_state(); - LLAMA_LOG_INFO("mt::paged_restore_from_warm: restored %u positions " - "for seq %d (paged_pool free: gpu=%zu cpu=%zu)\n", - restored, seq_id, - paged_pool_.n_free_gpu(), paged_pool_.n_free_cpu()); - } - return restored; -} - void llama_memory_tiered::record_chunk_fingerprint(std::vector positions, std::vector embedding, SemanticIndex::Tier tier) { @@ -872,174 +710,9 @@ bool llama_memory_tiered::backup_seq_rm_recurrent(llama_seq_id seq_id) { // SWA caches are skipped: their distance mask hides any restored // position older than n_swa back from the head, so backing them up // would be wasted work. -// Phase 2b-C: paged-blocks variant. Block-keyed allocation via -// paged_pool_, block-keyed mapping via paged_table_, K/V written into -// paged_warm_buf_. Completely parallel to the position-keyed path — -// touches NONE of warm_pos_to_slot_ / evicted_to_warm_ / -// warm_insertion_order_ / warm_buf_. Phase 2c will validate by -// toggling cfg_.paged_blocks and confirming the same observable -// behavior (needle test still passes). -uint32_t llama_memory_tiered::paged_backup_seq_rm_range(llama_seq_id seq_id, - llama_pos p0, - llama_pos p1) { - if (!inner_ || p0 < 0 || p1 <= p0) return 0; - if (!ensure_paged_warm_staging()) return 0; - - const uint32_t bsize = paged_table_.block_size(); - - // Round the request to block boundaries: a backup range that ends - // mid-block can't fully evict that final block (some positions in - // it are still live). For Phase 2b-C we only back up FULLY-CONTAINED - // blocks — partial blocks at the tail are deferred to a later - // refinement (vLLM has the same constraint at the scheduler layer). - const uint32_t first_block_idx = (uint32_t) p0 / bsize; - const uint32_t last_block_idx = (uint32_t) p1 / bsize; // exclusive - if (last_block_idx <= first_block_idx) { - // Range is smaller than one block — nothing to back up at this - // granularity. Real eviction triggers always pass at least - // n_evict = 0.20 * cap which is many blocks, so this is rare. - return 0; - } - - // Snapshot inner cells so we can find which positions live in - // which physical hot slot for the K/V copy. Same pattern as the - // position-keyed path. - const auto fresh = inner_->make_tier_view(); - if (fresh.attn_caches.size() != tier_view_.attn_caches.size()) { - LLAMA_LOG_WARN("mt::paged_backup_seq_rm_range: cache topology changed " - "(%zu->%zu); skipping backup\n", - tier_view_.attn_caches.size(), fresh.attn_caches.size()); - return 0; - } - - // For each block in [first_block_idx, last_block_idx): - // - Skip if seq already has this logical block in paged_table_ - // (already backed up earlier) - // - Allocate a CPU block from paged_pool_ - // - Walk inner cells; for each cell whose pos falls in this - // block's range AND seq matches, copy K/V across all restorable - // layers into the CPU block buffer - // - Append the CPU block ID to paged_table_ for this seq - uint32_t backed_up_blocks = 0; - uint32_t backed_up_positions = 0; - const uint32_t n_logical_existing = paged_table_.num_blocks(seq_id); - - for (uint32_t lblock = first_block_idx; lblock < last_block_idx; ++lblock) { - // Already backed up for this seq? - if (lblock < n_logical_existing) { - const uint32_t existing = paged_table_.get_physical(seq_id, lblock); - if (existing != kInvalidBlockId && !paged_pool_.is_gpu(existing)) { - // Already mapped to a CPU block — skip. - continue; - } - } - - // Allocate a CPU block. - const uint32_t cpu_block = paged_pool_.alloc_cpu(); - if (cpu_block == kInvalidBlockId) { - LLAMA_LOG_DEBUG("mt::paged_backup_seq_rm_range: CPU pool empty at " - "lblock=%u; remaining blocks not backed up " - "(cold-spill not yet implemented in 2b-C)\n", - lblock); - break; - } - - // Walk inner cells for this block's position range, copy K/V. - const llama_pos block_start = (llama_pos) lblock * bsize; - const llama_pos block_end = block_start + (llama_pos) bsize; - - const uint32_t cpu_block_idx = cpu_block - paged_pool_.total_gpu_blocks(); - uint8_t * block_base = paged_warm_buf_.data() + - (size_t) cpu_block_idx * paged_block_bytes_; - - bool ok = true; - size_t flat_layer_idx = 0; - - for (size_t ci = 0; ci < tier_view_.attn_caches.size() && ok; ++ci) { - const auto & cache_view = tier_view_.attn_caches[ci]; - const auto & cells = fresh.attn_caches[ci].cells; - - if (cache_view.is_swa) continue; // not advancing flat_layer_idx - - // Within this block, find each cell whose (pos, seq) is in - // range. Slot index in the inner cache != position; we walk - // and match. O(kv_size_per_cache) per block — same cost - // shape as the position-keyed path. - for (uint32_t slot = 0; slot < cells.size() && ok; ++slot) { - const auto & cs = cells[slot]; - if (cs.pos < block_start || cs.pos >= block_end) continue; - if (cs.seq_id != seq_id) continue; - - const uint32_t row_in_block = (uint32_t)(cs.pos - block_start); - - // For each layer in this cache, copy K and V. - size_t layer_idx_in_block = flat_layer_idx; - for (const auto & layer : cache_view.layers) { - uint8_t * k_dst = block_base - + paged_layer_off_[layer_idx_in_block] - + (size_t) row_in_block * layer.k_row_bytes; - uint8_t * v_dst = block_base - + paged_layer_v_off_[layer_idx_in_block] - + (size_t) row_in_block * layer.v_row_bytes; - if (!mover_attn_.evict_k(layer, slot, k_dst) || - !mover_attn_.evict_v(layer, slot, v_dst)) { - LLAMA_LOG_WARN("mt::paged_backup_seq_rm_range: mover " - "failed at slot %u (pos %d, layer %zu); " - "abandoning this block\n", - slot, cs.pos, layer_idx_in_block); - ok = false; - break; - } - ++layer_idx_in_block; - } - - if (ok) ++backed_up_positions; - } - - flat_layer_idx += cache_view.layers.size(); - } - - if (ok) { - // Pad table with kInvalidBlockId for any logical gaps before - // this block (rare — only if backup is called non-contiguously). - while (paged_table_.num_blocks(seq_id) < lblock) { - paged_table_.append_block(seq_id, kInvalidBlockId); - } - // Either swap into an existing slot or append. - if (lblock < paged_table_.num_blocks(seq_id)) { - paged_table_.swap_block(seq_id, lblock, cpu_block); - } else { - paged_table_.append_block(seq_id, cpu_block); - } - ++backed_up_blocks; - } else { - paged_pool_.free_block(cpu_block); - } - } - - if (backed_up_blocks > 0) { - capacity_.on_migrate(backed_up_positions, - TierCapacityManager::Tier::Hot, - TierCapacityManager::Tier::Warm); - LLAMA_LOG_INFO("mt::paged_backup_seq_rm_range: seq=%d range=[%d,%d) " - "backed up %u blocks (%u positions); paged_pool free: " - "gpu=%zu cpu=%zu\n", - seq_id, p0, p1, backed_up_blocks, backed_up_positions, - paged_pool_.n_free_gpu(), paged_pool_.n_free_cpu()); - } - - return backed_up_positions; -} - uint32_t llama_memory_tiered::backup_seq_rm_range(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { - // Phase 2b-C: when paged_blocks flag is set, take the parallel - // block-keyed path. Both paths return positions-backed-up so the - // caller (server proactive trigger) treats them identically. - if (cfg_.paged_blocks) { - return paged_backup_seq_rm_range(seq_id, p0, p1); - } if (!inner_ || p0 < 0 || p1 <= p0) return 0; if (!ensure_warm_staging()) return 0; @@ -1371,35 +1044,11 @@ bool llama_memory_tiered::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 // n_seq_max=1 today so "for this seq" === "all of it." When mt:: // grows multi-seq tracking, scope these wipes by seq_id. if (p0 < 0 || p1 < 0) { - // Phase 2b-C-3: hard-branch to paged variant when flag is set. - // semantic_ / pressure_announced_ wipes are shared (not - // paged-specific) but the tier bookkeeping diverges. - if (cfg_.paged_blocks) { - // Block-keyed wipe. Per-seq scoping is real here (paged_table_ - // is indexed by seq_id), unlike the position-keyed path which - // wipes globally because its maps don't carry seq scope. - const size_t n_finger = semantic_.size(); - - std::vector freed = paged_table_.clear_seq(seq_id); - const size_t n_blocks_freed = freed.size(); - for (uint32_t bid : freed) { - if (bid != kInvalidBlockId) { - paged_pool_.free_block(bid); - } - } - - semantic_.clear(); - pressure_announced_ = false; - - if (n_blocks_freed + n_finger > 0) { - LLAMA_LOG_INFO("mt::seq_rm: paged whole-seq wipe for seq %d " - "— freed %zu blocks (paged_pool free: gpu=%zu " - "cpu=%zu), cleared %zu semantic fingerprints\n", - seq_id, n_blocks_freed, - paged_pool_.n_free_gpu(), paged_pool_.n_free_cpu(), - n_finger); - } - } else if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max_) { + // MAD-127: paged whole-seq wipe is now handled inside + // llama_kv_cache_paged::seq_rm (which the wrapper delegates to + // via inner_->seq_rm above). The wrapper only handles the + // non-paged tiered path's per-seq metadata cleanup below. + if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max_) { // Per-seq whole-seq wipe: drop only this seq's tier // metadata, freeing its warm slots back to the pool. // Other seqs' state stays intact. diff --git a/src/memory-tier/mt-tiered.h b/src/memory-tier/mt-tiered.h index 77dac40c2793..17993f33dfed 100644 --- a/src/memory-tier/mt-tiered.h +++ b/src/memory-tier/mt-tiered.h @@ -31,8 +31,6 @@ #include "mt-mover-recurrent.h" #include "mt-kvtc-store.h" #include "mt-semantic.h" -#include "mt-block-pool.h" -#include "mt-block-table.h" #include "llama-memory.h" @@ -210,6 +208,12 @@ class llama_memory_tiered : public llama_memory_i { // Cached from tier_view_ at init — cheap, no allocation. uint32_t physical_attn_cells() const; + // MAD-134: warm the bge-small embedding model at construction time + // so the first user prompt doesn't pay the lazy-load cost + // (~200ms cold). Called from the ctor when semantic_index is set. + // Failures are non-fatal — lazy path still works. + void warmup_embed_(); + // Compute an L2-normalized embedding for `text` via the embedding // model loaded from cfg_.semantic_index. Lazy-initializes the model // on first call. Returns empty vector if the model failed to load @@ -237,35 +241,6 @@ class llama_memory_tiered : public llama_memory_i { // (recurrent-only models — handled in 2d-recur). bool ensure_warm_staging(); - // Phase 2b-C: paged-blocks variant of ensure_warm_staging. Allocates - // paged_warm_buf_ + computes paged_layer_off_ / paged_layer_v_off_ - // for the block-keyed path. Only fires when cfg_.paged_blocks=true. - // Independent of warm_buf_ (parallel implementation under flag). - bool ensure_paged_warm_staging(); - - // Phase 2b-C: paged-blocks variant of backup_seq_rm_range. Allocates - // CPU blocks from paged_pool_, copies K/V from inner cache GPU into - // paged_warm_buf_, appends physical block IDs to paged_table_ for - // the seq. Returns the count of positions successfully backed up - // (= n_blocks_evicted * block_size if all rows succeed). Only - // called by backup_seq_rm_range when cfg_.paged_blocks=true. - uint32_t paged_backup_seq_rm_range(llama_seq_id seq_id, llama_pos p0, llama_pos p1); - - // Phase 2b-C: paged-blocks variant of restore_from_warm. For each - // requested position, looks up the physical block via paged_table_, - // copies K/V from paged_warm_buf_ back into the inner cache via - // mover_attn_.restore_*. The CPU block stays mapped in paged_table_ - // after restore (the same chunk may be requested again, e.g. by - // semantic prefetch on a related query). Phase 2c will validate - // that toggling cfg_.paged_blocks gives identical observable - // behavior to the position-keyed path. - uint32_t paged_restore_from_warm(llama_seq_id seq_id, - const std::vector & positions); - - // Phase 2b-C: paged-blocks variant of has_warm. Checks paged_table_ - // for a CPU-mapped physical block covering the position. - bool paged_has_warm(llama_seq_id seq_id, llama_pos position) const; - // Open KvtcStore on the configured SSD path the first time cold // is needed. Returns false if cold_capacity is 0 or KvtcStore::init // fails. @@ -314,29 +289,6 @@ class llama_memory_tiered : public llama_memory_i { KvtcStore store_; SemanticIndex semantic_; - // Phase 2a paged-blocks scaffolding. Allocated only when - // cfg_.paged_blocks=true; otherwise these stay default-constructed - // and the existing position-keyed paths are unaffected. Phase 2b+ - // wires them into the live read/write paths behind the same flag. - BlockPool paged_pool_; - BlockTable paged_table_; - - // Phase 2b-C paged warm-tier staging. Parallel to warm_buf_ — when - // cfg_.paged_blocks=true, the gated paged_backup / paged_restore - // paths allocate physical CPU blocks from paged_pool_ and copy - // K/V into this buffer. Layout per CPU block: - // block N starts at paged_warm_buf_ + (cpu_block_idx) * paged_block_bytes_ - // within block, layer L's K starts at +paged_layer_off_[L] - // layer L's V starts at +paged_layer_v_off_[L] - // K stride within layer: row_bytes (block_size rows of K) - // V stride within layer: row_bytes (block_size rows of V) - // Lazily allocated by ensure_paged_warm_staging() on first use. - std::vector paged_warm_buf_; - std::vector paged_layer_off_; // byte offset of layer L's K within a block - std::vector paged_layer_v_off_; // byte offset of layer L's V within a block - size_t paged_block_bytes_ = 0; // total bytes per block (all restorable layers, K+V) - bool paged_warm_initialized_ = false; - // Cached tier view captured at construction. Pointers stay stable // for the lifetime of inner_ — see mt-inner-access.h. InnerView tier_view_; diff --git a/tests/.gitignore b/tests/.gitignore index faacc273bb79..f0b02f7f98d0 100644 --- a/tests/.gitignore +++ b/tests/.gitignore @@ -1,6 +1,7 @@ * !*.* !snapshots/ +!stress/ *.o ggml-common.h **/*.swp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 518e727bb492..e2f1e28f3a06 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -265,6 +265,41 @@ if (NOT GGML_BACKEND_DL) llama_build_and_test(test-rope.cpp) endif() +# MAD-135: mt:: int4 quant round-trip tests +llama_build_and_test(test-mt-quant.cpp) + +# MAD-137: mt:: tier primitive unit tests +# These compile a small slice of src/memory-tier/ directly into the test +# binary because the classes aren't exported through the public llama +# API. src/ is added to the include path so internal headers (llama-impl.h) +# resolve. +llama_build_and_test(test-mt-block-pool.cpp ../src/memory-tier/mt-block-pool.cpp) +target_include_directories(test-mt-block-pool PRIVATE ${CMAKE_SOURCE_DIR}/src) + +llama_build_and_test(test-mt-block-table.cpp + ../src/memory-tier/mt-block-table.cpp + ../src/memory-tier/mt-block-pool.cpp) +target_include_directories(test-mt-block-table PRIVATE ${CMAKE_SOURCE_DIR}/src) + +llama_build_and_test(test-mt-block-semantic-index.cpp + ../src/memory-tier/mt-semantic.cpp) +target_include_directories(test-mt-block-semantic-index PRIVATE ${CMAKE_SOURCE_DIR}/src) + +llama_build_and_test(test-mt-tiered-thin.cpp + ../src/memory-tier/mt-semantic.cpp) +target_include_directories(test-mt-tiered-thin PRIVATE ${CMAKE_SOURCE_DIR}/src) + +# MAD-137: paged-cache integration test. Needs a real model (skips +# gracefully when LLAMACPP_TEST_MODELFILE / argv[1] are unset). Links +# the full llama lib because llama_kv_cache_paged is part of it. +llama_build_and_test(test-paged-lifecycle.cpp) +target_link_libraries(test-paged-lifecycle PRIVATE llama) +target_include_directories(test-paged-lifecycle PRIVATE ${CMAKE_SOURCE_DIR}/src) + +llama_build_and_test(test-paged-semantic.cpp) +target_link_libraries(test-paged-semantic PRIVATE llama) +target_include_directories(test-paged-semantic PRIVATE ${CMAKE_SOURCE_DIR}/src) + # libmtmd set(LLAMA_TEST_NAME test-mtmd-c-api) llama_build_and_test(test-mtmd-c-api.c) diff --git a/tests/stress/stress-paged-multi-seq.py b/tests/stress/stress-paged-multi-seq.py new file mode 100755 index 000000000000..ba81671278b9 --- /dev/null +++ b/tests/stress/stress-paged-multi-seq.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +"""MAD-137: paged-cache multi-seq stress driver. + +NOTE: this test is currently blocked by MAD-141 — the server's prefill +admission loop deadlocks under sustained workloads exceeding the hot +pool budget. The driver below correctly detects the symptom (HTTP errors ++ unadvanced tier counters) and reports a hard failure. Once MAD-141 +lands, the same invocation should reach the green path. + +Boots a real llama-server with --parallel N --kv-tier-paged-blocks and +drives N concurrent simulated agents at it. Each agent sends a long +initial prompt (mixed-locality references are emulated by interleaving +queries that target recently produced tokens with queries that recall +content from much earlier in the conversation). After the configured +duration, scrapes /metrics + /slots and asserts on: + + - no HTTP 5xx responses (proxy for "no OOM, no crash") + - paged_evict_hot_to_warm_total > 0 (warm tier engaged) + - paged_evict_warm_to_cold_total > 0 (cold tier engaged) [if cold > 0] + - paged_semantic_attempts_total > 0 (semantic path exercised) + - decode rate per agent above --decode-floor tok/s + - optionally: semantic hit rate > --hit-rate-floor (defaults to 0 because + hit-rate is content-dependent and only meaningful with curated workloads) + +The script is intentionally stdlib-only (urllib + threading + json) so it +can run inside the army's ROCm docker container (Python 3, no pip). + +Typical CI invocation (matrix runner sets the device-specific knobs): + + python3 stress-paged-multi-seq.py \\ + --bin /workspace/llama.cpp/build-army/bin/llama-server \\ + --model /models/omnicoder-9b-q5_k_m.gguf \\ + --device CUDA0 \\ + --ctx 32768 --n-agents 4 \\ + --tier 25,75,0 --instance-id stress-1070 \\ + --duration 60 --decode-floor 5 +""" + +from __future__ import annotations + +import argparse +import json +import os +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.request +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +# --------------------------------------------------------------------------- +# HTTP helpers (stdlib only) +# --------------------------------------------------------------------------- + +def _http_post_json(url: str, body: dict, timeout: float = 60.0) -> tuple[int, dict | None]: + data = json.dumps(body).encode("utf-8") + req = urllib.request.Request(url, data=data, + headers={"Content-Type": "application/json"}, + method="POST") + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + return resp.status, json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as e: + return e.code, None + except (urllib.error.URLError, TimeoutError, ConnectionError): + return 0, None + + +def _http_get_text(url: str, timeout: float = 5.0) -> str | None: + try: + with urllib.request.urlopen(url, timeout=timeout) as resp: + return resp.read().decode("utf-8") + except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, ConnectionError): + return None + + +def _parse_prom_counter(metrics_text: str, name: str) -> float | None: + """Parse a Prometheus-format counter/gauge value out of /metrics text.""" + full = f"llamacpp:{name}" + for line in metrics_text.splitlines(): + if line.startswith(full + " "): + try: + return float(line.split()[-1]) + except ValueError: + return None + return None + + +# --------------------------------------------------------------------------- +# Agent driver — one per simulated user +# --------------------------------------------------------------------------- + +@dataclass +class AgentStats: + agent_id: int + requests: int = 0 + http_errors: int = 0 + decode_tps_sum: float = 0.0 + prefill_tps_sum: float = 0.0 + decoded_total: int = 0 + transcript_tail: list[str] = field(default_factory=list) + + +def _seed_paragraph(agent_id: int, turn: int, locality: str) -> str: + """Synthesize a deterministic-but-distinct paragraph. The only thing + that matters for the cache is that text content varies enough to + occupy distinct fingerprint directions. Locality tags are recorded + in the prompt so the eventual server log is debuggable.""" + seed = f"agent{agent_id:02d}-turn{turn:04d}-{locality}" + # ~256 tokens per paragraph is a reasonable proxy at q5_k_m densities. + body = " ".join(f"{seed}-w{i:04d}" for i in range(256)) + return body + + +def _agent_loop(agent_id: int, base_url: str, ctx_per_agent: int, + stop_event: threading.Event, stats: AgentStats): + seq_id = agent_id + + # Backoff state: when the server returns an error or the request + # fails, sleep before retrying. Exponential up to 5s. Without this + # an agent in error state hammers the server thousands of times per + # second, which is both useless and makes the failure mode unreadable. + backoff_s = 0.5 + + def _post(body: dict, timeout: float) -> bool: + nonlocal backoff_s + code, resp = _http_post_json(f"{base_url}/completion", body, timeout=timeout) + stats.requests += 1 + if code != 200 or resp is None: + stats.http_errors += 1 + time.sleep(backoff_s) + backoff_s = min(5.0, backoff_s * 2.0) + return False + backoff_s = 0.5 + timings = resp.get("timings", {}) or {} + stats.prefill_tps_sum += float(timings.get("prompt_per_second", 0.0)) + stats.decode_tps_sum += float(timings.get("predicted_per_second", 0.0)) + stats.decoded_total += int(timings.get("predicted_n", 0)) + if len(stats.transcript_tail) < 32: + stats.transcript_tail.append((resp.get("content") or "")[:80]) + return True + + # Initial long context — pre-load the agent's KV with distinct content. + initial_paragraphs = max(1, ctx_per_agent // 256) + bootstrap = "\n\n".join( + _seed_paragraph(agent_id, i, "init") for i in range(initial_paragraphs) + ) + bootstrap += "\n\nQ: Repeat the agent identifier you saw at the start.\nA:" + _post({ + "prompt": bootstrap, + "n_predict": 16, + "temperature": 0.0, + "cache_prompt": True, + "id_slot": seq_id, + }, timeout=600.0) + + # Mixed-locality follow-ups until stop_event fires. + turn = 0 + while not stop_event.is_set(): + turn += 1 + # Alternate recent vs far-back queries. + locality = "recent" if (turn % 3) != 0 else "farback" + followup = ( + _seed_paragraph(agent_id, turn, locality) + + "\n\nQ: Summarize the prior section in one sentence.\nA:" + ) + _post({ + "prompt": followup, + "n_predict": 24, + "temperature": 0.0, + "cache_prompt": True, + "id_slot": seq_id, + }, timeout=300.0) + + +# --------------------------------------------------------------------------- +# Server lifecycle +# --------------------------------------------------------------------------- + +def _wait_for_listening(log_path: Path, timeout: float = 600.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + data = log_path.read_text(errors="replace") + if "main: server is listening" in data: + return True + if "model loading error" in data or "Aborted" in data or "GGML_ABORT" in data: + return False + except FileNotFoundError: + pass + time.sleep(0.5) + return False + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--bin", required=True, help="path to llama-server") + p.add_argument("--model", required=True) + p.add_argument("--device", default="", help="--device flag value (e.g. CUDA0, ROCm0)") + p.add_argument("--ctx", type=int, default=8192, + help="total context per server instance") + p.add_argument("--n-agents", type=int, default=4, + help="--parallel N (each agent gets ctx/N tokens)") + p.add_argument("--tier", default="25,75,0", + help="--kv-tiered HOT,WARM,COLD percentages") + p.add_argument("--instance-id", default="stress-test") + p.add_argument("--ssd-path", default="", + help="cold-tier directory; auto-tmp if unset") + p.add_argument("--cache-type-k", default="turbo4") + p.add_argument("--cache-type-v", default="turbo4") + p.add_argument("--bge-small", default="", + help="optional bge-small gguf for semantic prefetch") + p.add_argument("--port", type=int, default=18080) + p.add_argument("--ngl", type=int, default=99) + p.add_argument("--duration", type=int, default=60, + help="total stress duration in seconds") + p.add_argument("--ctx-per-agent", type=int, default=2048, + help="initial context tokens to pre-load per agent") + p.add_argument("--decode-floor", type=float, default=2.0, + help="minimum mean decode tok/s per agent for pass") + p.add_argument("--report-json", default="", + help="write final results json to this path") + args = p.parse_args() + + if not Path(args.bin).is_file(): + print(f"FAIL: --bin {args.bin} not found", file=sys.stderr) + return 2 + if not Path(args.model).is_file(): + print(f"FAIL: --model {args.model} not found", file=sys.stderr) + return 2 + + tmpdir = tempfile.mkdtemp(prefix="stress-paged-") + log_path = Path(tmpdir) / "server.log" + ssd_path = args.ssd_path or os.path.join(tmpdir, "ssd") + os.makedirs(ssd_path, exist_ok=True) + + cmd = [ + args.bin, + "-m", args.model, + "-c", str(args.ctx), + "--parallel", str(args.n_agents), + "-ngl", str(args.ngl), + "--no-mmap", + "--metrics", + "--kv-tier-paged-blocks", + "--kv-tiered", args.tier, + "--cache-type-k", args.cache_type_k, + "--cache-type-v", args.cache_type_v, + "--instance-id", args.instance_id, + "--kv-tier-ssd-path", ssd_path, + "--port", str(args.port), + ] + if args.device: + cmd += ["--device", args.device] + if args.bge_small: + cmd += ["--kv-tier-semantic-index", args.bge_small] + + print("LAUNCH:", " ".join(shlex.quote(c) for c in cmd)) + log_fh = log_path.open("wb") + proc = subprocess.Popen(cmd, stdout=log_fh, stderr=subprocess.STDOUT) + try: + if not _wait_for_listening(log_path): + print("FAIL: server did not reach 'main: server is listening' before timeout", + file=sys.stderr) + return 1 + + base_url = f"http://127.0.0.1:{args.port}" + + # Snapshot tier counters BEFORE the workload so we can compute deltas. + metrics_before = _http_get_text(f"{base_url}/metrics") or "" + before = { + "evict_h2w": _parse_prom_counter(metrics_before, "paged_evict_hot_to_warm_total"), + "evict_w2c": _parse_prom_counter(metrics_before, "paged_evict_warm_to_cold_total"), + "sem_attempts": _parse_prom_counter(metrics_before, "paged_semantic_attempts_total"), + "sem_hits": _parse_prom_counter(metrics_before, "paged_semantic_hits_total"), + } + + stop_event = threading.Event() + stats_per_agent: list[AgentStats] = [AgentStats(agent_id=i) for i in range(args.n_agents)] + threads = [ + threading.Thread( + target=_agent_loop, + args=(i, base_url, args.ctx_per_agent, stop_event, stats_per_agent[i]), + daemon=True, + ) + for i in range(args.n_agents) + ] + t0 = time.monotonic() + for t in threads: + t.start() + time.sleep(args.duration) + stop_event.set() + for t in threads: + t.join(timeout=60.0) + elapsed = time.monotonic() - t0 + + # Final metrics + metrics_after = _http_get_text(f"{base_url}/metrics") or "" + after = { + "evict_h2w": _parse_prom_counter(metrics_after, "paged_evict_hot_to_warm_total"), + "evict_w2c": _parse_prom_counter(metrics_after, "paged_evict_warm_to_cold_total"), + "sem_attempts": _parse_prom_counter(metrics_after, "paged_semantic_attempts_total"), + "sem_hits": _parse_prom_counter(metrics_after, "paged_semantic_hits_total"), + } + deltas = {k: (after[k] or 0) - (before[k] or 0) for k in after} + + # Per-agent decode mean (mean of per-request rates the server reports). + agent_decodes: list[float] = [] + total_errors = 0 + for s in stats_per_agent: + mean_decode = (s.decode_tps_sum / s.requests) if s.requests else 0.0 + agent_decodes.append(mean_decode) + total_errors += s.http_errors + + # Hit rate (only well-defined when sem_attempts > 0) + hit_rate = (deltas["sem_hits"] / deltas["sem_attempts"] + if deltas["sem_attempts"] > 0 else 0.0) + + # ─── Pass / fail ────────────────────────────────────────────────── + failures: list[str] = [] + if total_errors > 0: + failures.append(f"{total_errors} HTTP errors across all agents (no OOM expected)") + if deltas["evict_h2w"] <= 0: + failures.append("paged_evict_hot_to_warm_total did not advance (warm tier never engaged)") + # Cold spill is only expected when COLD% > 0 in the tier config. + cold_pct = int(args.tier.split(",")[2]) if len(args.tier.split(",")) >= 3 else 0 + if cold_pct > 0 and deltas["evict_w2c"] <= 0: + failures.append("paged_evict_warm_to_cold_total did not advance (cold tier configured but unused)") + if any(d < args.decode_floor for d in agent_decodes): + failures.append( + "decode rate floor not met: " + + ", ".join(f"agent{i}={d:.2f}<{args.decode_floor}" + for i, d in enumerate(agent_decodes) if d < args.decode_floor) + ) + + report: dict[str, Any] = { + "ok": not failures, + "failures": failures, + "elapsed_s": elapsed, + "n_agents": args.n_agents, + "tier": args.tier, + "device": args.device, + "decode_tps_per_agent": agent_decodes, + "decode_tps_floor": args.decode_floor, + "metrics_delta": deltas, + "hit_rate": hit_rate, + "total_http_errors": total_errors, + "log_path": str(log_path), + } + print(json.dumps(report, indent=2)) + if args.report_json: + Path(args.report_json).write_text(json.dumps(report, indent=2)) + + return 0 if not failures else 1 + + finally: + # Tear down the server. + try: + proc.send_signal(signal.SIGTERM) + proc.wait(timeout=15) + except subprocess.TimeoutExpired: + proc.kill() + log_fh.close() + # Leave tmpdir on failure for inspection; clean on success. + # (Caller can override via --report-json or by just inspecting tmpdir.) + if "ok" in locals() and locals().get("ok"): + shutil.rmtree(tmpdir, ignore_errors=True) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test-mt-block-pool.cpp b/tests/test-mt-block-pool.cpp new file mode 100644 index 000000000000..42b3b26d881f --- /dev/null +++ b/tests/test-mt-block-pool.cpp @@ -0,0 +1,222 @@ +// MAD-137: unit tests for mt::BlockPool — the physical block allocator +// that backs the paged tier (GPU pool + CPU pool, refcounting, watermark +// admission control). Bare main()/assert style; no real GPU. + +#include "../src/memory-tier/mt-block-pool.h" + +#undef NDEBUG +#include +#include + +using mt::BlockPool; +using mt::kInvalidBlockId; + +int main() { + // ─── Basic alloc / free cycle ─── + { + BlockPool p; + p.init(4, 4, 0.0f); + assert(p.n_free_gpu() == 4); + assert(p.n_free_cpu() == 4); + + const uint32_t b0 = p.alloc_gpu(); + assert(b0 != kInvalidBlockId); + assert(p.refcount(b0) == 1); + assert(p.n_free_gpu() == 3); + + p.free_block(b0); + assert(p.refcount(b0) == 0); + assert(p.n_free_gpu() == 4); + printf("test-mt-block-pool: basic alloc/free ok\n"); + } + + // ─── Allocation order — init pushes descending so pop returns 0 first ─── + { + BlockPool p; + p.init(4, 0, 0.0f); + const uint32_t a = p.alloc_gpu(); + const uint32_t b = p.alloc_gpu(); + const uint32_t c = p.alloc_gpu(); + assert(a == 0 && b == 1 && c == 2); + printf("test-mt-block-pool: deterministic alloc order ok\n"); + } + + // ─── GPU vs CPU id ranges + is_gpu() ─── + { + BlockPool p; + p.init(4, 4, 0.0f); + const uint32_t g = p.alloc_gpu(); + const uint32_t c = p.alloc_cpu(); + assert(g < 4); // GPU range [0, 4) + assert(c >= 4 && c < 8); // CPU range [4, 8) + assert(p.is_gpu(g)); + assert(!p.is_gpu(c)); + printf("test-mt-block-pool: gpu/cpu id ranges ok (gpu=%u cpu=%u)\n", g, c); + } + + // ─── Exhaustion — drain GPU, next alloc returns kInvalidBlockId, CPU unaffected ─── + { + BlockPool p; + p.init(2, 2, 0.0f); + const uint32_t a = p.alloc_gpu(); + const uint32_t b = p.alloc_gpu(); + assert(a != kInvalidBlockId && b != kInvalidBlockId); + assert(p.n_free_gpu() == 0); + + const uint32_t exhausted = p.alloc_gpu(); + assert(exhausted == kInvalidBlockId); + + // CPU pool independent: still allocates fine. + const uint32_t c = p.alloc_cpu(); + assert(c != kInvalidBlockId); + assert(p.n_free_cpu() == 1); + printf("test-mt-block-pool: exhaustion + multi-pool independence ok\n"); + } + + // ─── Refcount bump (CoW path) ─── + { + BlockPool p; + p.init(4, 0, 0.0f); + const uint32_t b = p.alloc_gpu(); + assert(p.refcount(b) == 1); + assert(p.n_free_gpu() == 3); + + p.bump_ref(b); + assert(p.refcount(b) == 2); + assert(p.n_free_gpu() == 3); // no change — still allocated + + p.free_block(b); + assert(p.refcount(b) == 1); + assert(p.n_free_gpu() == 3); // still NOT in free stack — second owner holds it + + p.free_block(b); + assert(p.refcount(b) == 0); + assert(p.n_free_gpu() == 4); // last ref dropped → returned to free stack + printf("test-mt-block-pool: bump_ref + multi-owner free ok\n"); + } + + // ─── Watermark admission — has_free_*_blocks() reserves a slice ─── + { + BlockPool p; + // 10 GPU blocks, 0.5 watermark → reserve = ceil(10 * 0.5) = 5. + // (Using 0.5f because it's exactly representable in float; values + // like 0.2f promote to double as 0.2000…0004, which trips ceil() + // and silently bumps the reserve by 1 — conservative behavior, + // but not what callers writing exact fractions might expect.) + // n_free_gpu() reports 10 (raw), has_free_gpu_blocks(N) is true + // only when (free - reserve) >= N → max admissible N = 5. + p.init(10, 0, 0.5f); + assert(p.n_free_gpu() == 10); + assert(p.has_free_gpu_blocks(5)); + assert(!p.has_free_gpu_blocks(6)); // would dip into reserve + + // Drain to reserve floor — has_free should refuse any further admit. + for (int i = 0; i < 5; ++i) { + const uint32_t id = p.alloc_gpu(); + assert(id != kInvalidBlockId); + } + assert(p.n_free_gpu() == 5); + assert(!p.has_free_gpu_blocks(1)); // exactly at watermark → false + + // alloc_*() ignores watermark — still allocates the reserved blocks. + for (int i = 0; i < 5; ++i) { + const uint32_t id = p.alloc_gpu(); + assert(id != kInvalidBlockId); + } + assert(p.n_free_gpu() == 0); + printf("test-mt-block-pool: watermark gating + alloc bypass ok\n"); + } + + // ─── Watermark with 0% disables the reserve ─── + { + BlockPool p; + p.init(4, 0, 0.0f); + assert(p.has_free_gpu_blocks(4)); + assert(!p.has_free_gpu_blocks(5)); + printf("test-mt-block-pool: 0%% watermark = no reserve ok\n"); + } + + // ─── Watermark applies independently to CPU pool ─── + { + BlockPool p; + p.init(0, 10, 0.5f); // reserve = ceil(10 * 0.5) = 5 + assert(p.has_free_cpu_blocks(5)); + assert(!p.has_free_cpu_blocks(6)); + // GPU side has 0 blocks → has_free returns false for any N. + assert(!p.has_free_gpu_blocks(1)); + printf("test-mt-block-pool: cpu-side watermark + empty-gpu side ok\n"); + } + + // ─── Double-free is idempotent (warn, no corruption) ─── + { + BlockPool p; + p.init(2, 0, 0.0f); + const uint32_t b = p.alloc_gpu(); + p.free_block(b); + assert(p.refcount(b) == 0); + const size_t free_after_first = p.n_free_gpu(); + p.free_block(b); // second free — should be a no-op + assert(p.refcount(b) == 0); + assert(p.n_free_gpu() == free_after_first); // didn't push twice + printf("test-mt-block-pool: double-free is idempotent ok\n"); + } + + // ─── kInvalidBlockId on every entry-point is a no-op ─── + { + BlockPool p; + p.init(2, 2, 0.0f); + p.bump_ref(kInvalidBlockId); // no crash + p.free_block(kInvalidBlockId); // no crash + assert(p.refcount(kInvalidBlockId) == 0); + // Pool state untouched. + assert(p.n_free_gpu() == 2); + assert(p.n_free_cpu() == 2); + printf("test-mt-block-pool: kInvalidBlockId paths safe ok\n"); + } + + // ─── refcount() for out-of-range ids returns 0 (safe accessor) ─── + { + BlockPool p; + p.init(2, 2, 0.0f); + assert(p.refcount(9999) == 0); + printf("test-mt-block-pool: out-of-range refcount safe ok\n"); + } + + // ─── reset() — restores all free counts and zeros all refcounts ─── + { + BlockPool p; + p.init(4, 4, 0.0f); + const uint32_t g = p.alloc_gpu(); + const uint32_t c = p.alloc_cpu(); + p.bump_ref(g); + p.bump_ref(c); + assert(p.n_free_gpu() == 3 && p.n_free_cpu() == 3); + assert(p.refcount(g) == 2 && p.refcount(c) == 2); + + p.reset(); + assert(p.n_free_gpu() == 4 && p.n_free_cpu() == 4); + assert(p.refcount(g) == 0 && p.refcount(c) == 0); + + // Post-reset alloc returns the same low IDs as a fresh init. + const uint32_t g2 = p.alloc_gpu(); + const uint32_t c2 = p.alloc_cpu(); + assert(g2 == 0 && c2 == 4); + printf("test-mt-block-pool: reset() restores all counts + refcounts ok\n"); + } + + // ─── total_*_blocks() reports init values, immutable across alloc ─── + { + BlockPool p; + p.init(7, 3, 0.1f); + assert(p.total_gpu_blocks() == 7); + assert(p.total_cpu_blocks() == 3); + (void) p.alloc_gpu(); + (void) p.alloc_cpu(); + assert(p.total_gpu_blocks() == 7); // doesn't shrink with alloc + assert(p.total_cpu_blocks() == 3); + printf("test-mt-block-pool: total_*_blocks() immutable ok\n"); + } + + printf("test-mt-block-pool: ALL PASS\n"); + return 0; +} diff --git a/tests/test-mt-block-semantic-index.cpp b/tests/test-mt-block-semantic-index.cpp new file mode 100644 index 000000000000..0e901f8969fd --- /dev/null +++ b/tests/test-mt-block-semantic-index.cpp @@ -0,0 +1,322 @@ +// MAD-137: unit tests for mt::BlockSemanticIndex — the per-(seq,lblock) +// fingerprint store + cosine-similarity scoring used for paged-block +// prefetch hints. Bare main()/assert style; no real GPU. + +#include "../src/memory-tier/mt-semantic.h" + +#undef NDEBUG +#include +#include +#include +#include +#include +#include +#include + +using mt::BlockSemanticIndex; +using mt::SemanticIndex; +using Tier = SemanticIndex::Tier; + +namespace { + +// Build an L2-normalized embedding from a free-form direction vector. +// Cosine similarity of the result with itself is 1.0 exactly. +std::vector normed(std::vector v) { + double sq = 0.0; + for (float x : v) sq += (double) x * x; + const double n = std::sqrt(sq); + if (n > 0.0) { + for (float & x : v) x = (float) ((double) x / n); + } + return v; +} + +// Unique tmp path so parallel test runs don't collide. Honors TMPDIR +// so sandboxed environments (and BSDs) get a writable directory. +std::string tmp_path(const char * tag) { + const char * dir = std::getenv("TMPDIR"); + if (dir == nullptr || dir[0] == '\0') dir = "/tmp"; + char buf[256]; + std::snprintf(buf, sizeof(buf), "%s/test-mt-bsi-%s-%d.bin", dir, tag, (int) ::getpid()); + return std::string(buf); +} + +} // namespace + +int main() { + // ─── Empty state — every accessor returns the safe zero ─── + { + BlockSemanticIndex idx; + assert(idx.size() == 0); + assert(idx.size(0) == 0); + assert(!idx.has_fingerprint(0, 0)); + const auto hints = idx.score(0, normed({1, 0, 0}), /* top_k */ 5, /* threshold */ 0.0f); + assert(hints.empty()); + printf("test-mt-block-semantic-index: empty state ok\n"); + } + + // ─── Basic add + has_fingerprint + size ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(/* seq */ 0, /* lblock */ 5, + normed({1, 0, 0}), Tier::Hot); + assert(idx.has_fingerprint(0, 5)); + assert(!idx.has_fingerprint(0, 6)); // different lblock + assert(!idx.has_fingerprint(1, 5)); // different seq + assert(idx.size() == 1); + assert(idx.size(0) == 1); + assert(idx.size(1) == 0); + printf("test-mt-block-semantic-index: add + has_fingerprint + size ok\n"); + } + + // ─── Overwrite at the same (seq, lblock) ─── + // Used after a partial-range edit re-fingerprints a block. + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 1, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(0, 1, normed({0, 1, 0}), Tier::Cold); + assert(idx.size() == 1); // not 2 — same key + + // Score should reflect the second embedding (now the "Y" direction) + // and the second tier (Cold). + const auto hints = idx.score(0, normed({0, 1, 0}), 5, 0.0f); + assert(hints.size() == 1); + assert(hints[0].lblock == 1); + assert(hints[0].tier == Tier::Cold); + assert(hints[0].score > 0.99f); // ~1.0 against same direction + printf("test-mt-block-semantic-index: overwrite at same key ok\n"); + } + + // ─── Per-seq scoping — score(seq) only sees that seq's blocks ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(1, 0, normed({1, 0, 0}), Tier::Hot); // same direction, different seq + idx.add_fingerprint(2, 0, normed({1, 0, 0}), Tier::Hot); + + const auto h0 = idx.score(0, normed({1, 0, 0}), 10, 0.0f); + assert(h0.size() == 1); + assert(h0[0].seq_id == 0); + + const auto h1 = idx.score(1, normed({1, 0, 0}), 10, 0.0f); + assert(h1.size() == 1); + assert(h1[0].seq_id == 1); + printf("test-mt-block-semantic-index: per-seq scoping ok\n"); + } + + // ─── Score: deterministic descending order ─── + { + BlockSemanticIndex idx; + // Three blocks at known angles to the query (1,0,0): + // block 0: (1,0,0) cos = 1.0 + // block 1: (cos45, sin45) cos ≈ 0.707 + // block 2: (0,1,0) cos = 0.0 + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(0, 1, normed({1, 1, 0}), Tier::Warm); + idx.add_fingerprint(0, 2, normed({0, 1, 0}), Tier::Cold); + + const auto hints = idx.score(0, normed({1, 0, 0}), 10, /* threshold */ -1.0f); + assert(hints.size() == 3); + assert(hints[0].lblock == 0); + assert(hints[1].lblock == 1); + assert(hints[2].lblock == 2); + // Scores strictly descending. + assert(hints[0].score > hints[1].score); + assert(hints[1].score > hints[2].score); + printf("test-mt-block-semantic-index: descending-score ordering ok\n"); + } + + // ─── Threshold filter — drops scores strictly below ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(0, 1, normed({1, 1, 0}), Tier::Warm); // ~0.707 + idx.add_fingerprint(0, 2, normed({0, 1, 0}), Tier::Cold); // ~0.0 + + // Threshold 0.5 keeps the first two only. + const auto hits = idx.score(0, normed({1, 0, 0}), 10, /* threshold */ 0.5f); + assert(hits.size() == 2); + assert(hits[0].lblock == 0); + assert(hits[1].lblock == 1); + printf("test-mt-block-semantic-index: threshold filter ok\n"); + } + + // ─── top_k cap ─── + { + BlockSemanticIndex idx; + for (uint32_t b = 0; b < 5; ++b) { + idx.add_fingerprint(0, b, normed({1, 0, 0}), Tier::Hot); + } + const auto hits = idx.score(0, normed({1, 0, 0}), /* top_k */ 3, -1.0f); + assert(hits.size() == 3); + printf("test-mt-block-semantic-index: top_k cap ok\n"); + } + + // ─── Edge cases for score() — empty query, top_k=0, unknown seq ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + assert(idx.score(0, /* empty query */ {}, 5, 0.0f).empty()); + assert(idx.score(0, normed({1, 0, 0}), /* top_k */ 0, 0.0f).empty()); + assert(idx.score(0, normed({1, 0, 0}), -1, 0.0f).empty()); + assert(idx.score(/* unknown seq */ 99, normed({1, 0, 0}), 5, 0.0f).empty()); + printf("test-mt-block-semantic-index: score edge cases ok\n"); + } + + // ─── update_tier — changes tier annotation; no-op on unknown key ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.update_tier(0, 0, Tier::Cold); + const auto hits = idx.score(0, normed({1, 0, 0}), 1, -1.0f); + assert(hits.size() == 1); + assert(hits[0].tier == Tier::Cold); + + // Unknown (seq, lblock): no crash, no state change. + idx.update_tier(99, 99, Tier::Hot); + idx.update_tier(0, 99, Tier::Hot); + assert(idx.size() == 1); + printf("test-mt-block-semantic-index: update_tier ok\n"); + } + + // ─── remove_block — drops one entry; seq stays if other entries remain ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(0, 1, normed({0, 1, 0}), Tier::Warm); + assert(idx.size() == 2); + + idx.remove_block(0, 0); + assert(!idx.has_fingerprint(0, 0)); + assert(idx.has_fingerprint(0, 1)); + assert(idx.size(0) == 1); + assert(idx.size() == 1); + + // Drop the last entry — seq map should be removed (size(seq)==0). + idx.remove_block(0, 1); + assert(idx.size() == 0); + assert(idx.size(0) == 0); + + // No-op on unknown. + idx.remove_block(99, 0); + idx.remove_block(0, 99); + assert(idx.size() == 0); + printf("test-mt-block-semantic-index: remove_block ok\n"); + } + + // ─── remove_seq — drops every entry for that seq, leaves others ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(0, 1, normed({0, 1, 0}), Tier::Hot); + idx.add_fingerprint(1, 0, normed({1, 0, 0}), Tier::Hot); + assert(idx.size() == 3); + + idx.remove_seq(0); + assert(idx.size() == 1); + assert(idx.size(0) == 0); + assert(idx.size(1) == 1); + + // No-op on unknown. + idx.remove_seq(42); + assert(idx.size() == 1); + printf("test-mt-block-semantic-index: remove_seq ok\n"); + } + + // ─── clear — wipes everything ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint(1, 0, normed({0, 1, 0}), Tier::Warm); + assert(idx.size() == 2); + + idx.clear(); + assert(idx.size() == 0); + assert(idx.size(0) == 0); + assert(idx.size(1) == 0); + // Index is reusable after clear. + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + assert(idx.size() == 1); + printf("test-mt-block-semantic-index: clear ok\n"); + } + + // ─── save_to_disk + load_from_disk round-trip (PSFI v1) ─── + { + const std::string path = tmp_path("rt"); + + BlockSemanticIndex original; + // Three seqs × multiple blocks × different tiers. + original.add_fingerprint(0, 0, normed({1, 0, 0, 0}), Tier::Hot); + original.add_fingerprint(0, 1, normed({0, 1, 0, 0}), Tier::Warm); + original.add_fingerprint(0, 2, normed({0, 0, 1, 0}), Tier::Cold); + original.add_fingerprint(7, 0, normed({0.5f, 0.5f, 0.5f, 0.5f}), Tier::Warm); + original.add_fingerprint(7, 4, normed({1, 1, 0, 0}), Tier::Hot); + + const bool save_ok = original.save_to_disk(path); + assert(save_ok); + + BlockSemanticIndex restored; + // Pre-load, restored should be empty. + assert(restored.size() == 0); + // Add a stray entry first to verify load REPLACES (not merges). + restored.add_fingerprint(99, 99, normed({1, 0, 0, 0}), Tier::Hot); + assert(restored.size() == 1); + + const bool load_ok = restored.load_from_disk(path); + assert(load_ok); + // The stray entry should be gone — load replaces in-memory state. + assert(!restored.has_fingerprint(99, 99)); + assert(restored.size() == 5); + assert(restored.size(0) == 3); + assert(restored.size(7) == 2); + + // Verify a specific (seq, lblock) round-trips its tier and embedding. + const auto hits = restored.score(0, normed({0, 1, 0, 0}), /* top_k */ 1, -1.0f); + assert(hits.size() == 1); + assert(hits[0].lblock == 1); + assert(hits[0].tier == Tier::Warm); + assert(hits[0].score > 0.99f); // exact direction match + + ::unlink(path.c_str()); + printf("test-mt-block-semantic-index: save/load round-trip ok\n"); + } + + // ─── load_from_disk on missing file returns false, leaves index alone ─── + { + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + const bool load_ok = idx.load_from_disk("/tmp/this-file-definitely-does-not-exist-mad137"); + assert(!load_ok); + // In-memory state should be untouched. + assert(idx.size() == 1); + assert(idx.has_fingerprint(0, 0)); + printf("test-mt-block-semantic-index: load missing file safe ok\n"); + } + + // ─── load_from_disk on wrong magic returns false ─── + { + const std::string path = tmp_path("badmagic"); + // Write a file that has the wrong magic. + FILE * f = std::fopen(path.c_str(), "wb"); + assert(f); + const uint32_t wrong_magic = 0xDEADBEEF; + const uint32_t version = 1; + std::fwrite(&wrong_magic, sizeof(wrong_magic), 1, f); + std::fwrite(&version, sizeof(version), 1, f); + std::fclose(f); + + BlockSemanticIndex idx; + idx.add_fingerprint(0, 0, normed({1, 0, 0}), Tier::Hot); + const bool load_ok = idx.load_from_disk(path); + assert(!load_ok); + // In-memory state untouched. + assert(idx.size() == 1); + assert(idx.has_fingerprint(0, 0)); + + ::unlink(path.c_str()); + printf("test-mt-block-semantic-index: load bad magic safe ok\n"); + } + + printf("test-mt-block-semantic-index: ALL PASS\n"); + return 0; +} diff --git a/tests/test-mt-block-table.cpp b/tests/test-mt-block-table.cpp new file mode 100644 index 000000000000..1b4edecba268 --- /dev/null +++ b/tests/test-mt-block-table.cpp @@ -0,0 +1,223 @@ +// MAD-137: unit tests for mt::BlockTable — the per-sequence +// logical→physical block mapping. Bare main()/assert style; no real GPU. + +#include "../src/memory-tier/mt-block-pool.h" // kInvalidBlockId +#include "../src/memory-tier/mt-block-table.h" + +#undef NDEBUG +#include +#include + +using mt::BlockTable; +using mt::kInvalidBlockId; + +int main() { + // ─── init() — accessors mirror the constructor args ─── + { + BlockTable t; + t.init(/* max_seqs */ 4, /* block_size */ 16); + assert(t.max_seqs() == 4); + assert(t.block_size() == 16); + // No appends yet — every seq has zero blocks. + for (uint32_t s = 0; s < 4; ++s) { + assert(t.num_blocks((int) s) == 0); + assert(t.get_physical((int) s, 0) == kInvalidBlockId); + } + printf("test-mt-block-table: init + zero state ok\n"); + } + + // ─── append + get_physical ─── + { + BlockTable t; + t.init(2, 16); + t.append_block(0, 100); + t.append_block(0, 101); + t.append_block(0, 102); + assert(t.num_blocks(0) == 3); + assert(t.get_physical(0, 0) == 100); + assert(t.get_physical(0, 1) == 101); + assert(t.get_physical(0, 2) == 102); + // Out-of-range logical_idx is a sentinel, not UB. + assert(t.get_physical(0, 3) == kInvalidBlockId); + printf("test-mt-block-table: append + get_physical ok\n"); + } + + // ─── get_physical_for_pos — pos / block_size division ─── + { + BlockTable t; + t.init(1, 16); + t.append_block(0, 200); + t.append_block(0, 201); + t.append_block(0, 202); + + // Block 0 covers pos [0, 16); block 1 covers [16, 32); etc. + assert(t.get_physical_for_pos(0, 0) == 200); + assert(t.get_physical_for_pos(0, 15) == 200); + assert(t.get_physical_for_pos(0, 16) == 201); + assert(t.get_physical_for_pos(0, 31) == 201); + assert(t.get_physical_for_pos(0, 32) == 202); + assert(t.get_physical_for_pos(0, 47) == 202); + // Past the live range → sentinel. + assert(t.get_physical_for_pos(0, 48) == kInvalidBlockId); + // Negative pos guard. + assert(t.get_physical_for_pos(0, -1) == kInvalidBlockId); + printf("test-mt-block-table: get_physical_for_pos arithmetic ok\n"); + } + + // ─── swap_block — returns old id, get_physical reflects new ─── + { + BlockTable t; + t.init(1, 16); + t.append_block(0, 300); + t.append_block(0, 301); + t.append_block(0, 302); + + const uint32_t old = t.swap_block(0, /* logical_idx */ 1, /* new */ 999); + assert(old == 301); + assert(t.get_physical(0, 1) == 999); + // Neighbors untouched. + assert(t.get_physical(0, 0) == 300); + assert(t.get_physical(0, 2) == 302); + assert(t.num_blocks(0) == 3); // swap doesn't change count + printf("test-mt-block-table: swap_block ok\n"); + } + + // ─── Non-contiguous mapping — swap_block(..., kInvalidBlockId) creates a hole ─── + // MAD-128 partial seq_rm uses this pattern: a logical block is wiped + // (its physical id replaced with kInvalidBlockId) but the surrounding + // blocks remain. The paged-attn kernel handles kInvalidBlockId by + // returning -INFINITY logits for the missing positions. + { + BlockTable t; + t.init(1, 16); + t.append_block(0, 400); + t.append_block(0, 401); + t.append_block(0, 402); + t.append_block(0, 403); + + const uint32_t evicted = t.swap_block(0, 2, kInvalidBlockId); + assert(evicted == 402); + + assert(t.get_physical(0, 0) == 400); + assert(t.get_physical(0, 1) == 401); + assert(t.get_physical(0, 2) == kInvalidBlockId); // the hole + assert(t.get_physical(0, 3) == 403); + + // num_blocks still counts the hole — the logical sequence length + // is unchanged; only the underlying physical mapping was wiped. + assert(t.num_blocks(0) == 4); + printf("test-mt-block-table: non-contiguous (hole) mapping ok\n"); + } + + // ─── clear_seq — returns the freed ids in append order, then empties ─── + { + BlockTable t; + t.init(2, 16); + t.append_block(0, 500); + t.append_block(0, 501); + t.append_block(0, 502); + // Throw a hole in to make sure it propagates through clear_seq too. + t.swap_block(0, 1, kInvalidBlockId); + + const auto freed = t.clear_seq(0); + assert(freed.size() == 3); + assert(freed[0] == 500); + assert(freed[1] == kInvalidBlockId); // the hole survives clearing + assert(freed[2] == 502); + + assert(t.num_blocks(0) == 0); + assert(t.get_physical(0, 0) == kInvalidBlockId); + printf("test-mt-block-table: clear_seq returns freed list + empties ok\n"); + } + + // ─── clear_seq on an already-empty sequence is a no-op ─── + { + BlockTable t; + t.init(2, 16); + const auto freed = t.clear_seq(0); + assert(freed.empty()); + printf("test-mt-block-table: clear_seq on empty seq ok\n"); + } + + // ─── Per-seq isolation — appends on one seq don't leak to others ─── + { + BlockTable t; + t.init(3, 16); + t.append_block(0, 600); + t.append_block(2, 700); + + assert(t.num_blocks(0) == 1); + assert(t.num_blocks(1) == 0); + assert(t.num_blocks(2) == 1); + assert(t.get_physical(0, 0) == 600); + assert(t.get_physical(1, 0) == kInvalidBlockId); + assert(t.get_physical(2, 0) == 700); + + // Clearing seq 0 doesn't touch seq 2. + t.clear_seq(0); + assert(t.num_blocks(0) == 0); + assert(t.num_blocks(2) == 1); + assert(t.get_physical(2, 0) == 700); + printf("test-mt-block-table: per-seq isolation ok\n"); + } + + // ─── Out-of-range seq queries return safe sentinels ─── + { + BlockTable t; + t.init(2, 16); + // Negative seq — every accessor must return safe values. + assert(t.get_physical(-1, 0) == kInvalidBlockId); + assert(t.get_physical_for_pos(-1, 0) == kInvalidBlockId); + assert(t.num_blocks(-1) == 0); + assert(t.clear_seq(-1).empty()); + + // Seq >= max_seqs — same. + assert(t.get_physical(99, 0) == kInvalidBlockId); + assert(t.get_physical_for_pos(99, 0) == kInvalidBlockId); + assert(t.num_blocks(99) == 0); + assert(t.clear_seq(99).empty()); + printf("test-mt-block-table: out-of-range seq queries safe ok\n"); + } + + // ─── reset() — wipes every sequence; the table is reusable ─── + { + BlockTable t; + t.init(3, 16); + t.append_block(0, 800); + t.append_block(1, 801); + t.append_block(2, 802); + assert(t.num_blocks(0) == 1 && t.num_blocks(1) == 1 && t.num_blocks(2) == 1); + + t.reset(); + for (uint32_t s = 0; s < 3; ++s) { + assert(t.num_blocks((int) s) == 0); + assert(t.get_physical((int) s, 0) == kInvalidBlockId); + } + // Accessors still report init values. + assert(t.max_seqs() == 3); + assert(t.block_size() == 16); + + // Table is reusable post-reset. + t.append_block(0, 900); + assert(t.get_physical(0, 0) == 900); + printf("test-mt-block-table: reset() + reusable ok\n"); + } + + // ─── Different block sizes resolve positions correctly ─── + { + BlockTable t; + t.init(1, 32); // larger blocks — block 0 covers pos [0, 32) + t.append_block(0, 700); + t.append_block(0, 701); + + assert(t.get_physical_for_pos(0, 0) == 700); + assert(t.get_physical_for_pos(0, 31) == 700); + assert(t.get_physical_for_pos(0, 32) == 701); + assert(t.get_physical_for_pos(0, 63) == 701); + assert(t.get_physical_for_pos(0, 64) == kInvalidBlockId); + printf("test-mt-block-table: block_size=32 arithmetic ok\n"); + } + + printf("test-mt-block-table: ALL PASS\n"); + return 0; +} diff --git a/tests/test-mt-quant.cpp b/tests/test-mt-quant.cpp new file mode 100644 index 000000000000..3b2e5d84b4b3 --- /dev/null +++ b/tests/test-mt-quant.cpp @@ -0,0 +1,144 @@ +// MAD-135: round-trip tests for mt::quantize_int4 / int8 + +// the per-block scaled int4 helpers used by the cold-tier +// compression in llama_kv_cache_paged. + +#include "../src/memory-tier/mt-quant.h" + +#undef NDEBUG +#include +#include +#include +#include +#include + +namespace { + +double max_abs_err(const std::vector & a, const std::vector & b) { + double m = 0.0; + const size_t n = std::min(a.size(), b.size()); + for (size_t i = 0; i < n; ++i) { + const double e = std::fabs((double) a[i] - (double) b[i]); + if (e > m) m = e; + } + return m; +} + +double cosine_sim(const std::vector & a, const std::vector & b) { + double dot = 0.0, na = 0.0, nb = 0.0; + const size_t n = std::min(a.size(), b.size()); + for (size_t i = 0; i < n; ++i) { + dot += (double) a[i] * b[i]; + na += (double) a[i] * a[i]; + nb += (double) b[i] * b[i]; + } + if (na == 0.0 || nb == 0.0) return 1.0; + return dot / (std::sqrt(na) * std::sqrt(nb)); +} + +// Generate Gaussian-distributed test data with given std, simulating +// real K/V magnitudes (~|x| up to 5-10× the std). +std::vector gauss(size_t n, float stddev, uint32_t seed) { + std::mt19937 rng(seed); + std::normal_distribution dist(0.0f, stddev); + std::vector v(n); + for (size_t i = 0; i < n; ++i) v[i] = dist(rng); + return v; +} + +} // anon + +int main() { + // ─── int4 (no scale) — input must be in [-1, +1] ─── + { + const std::vector src = {-1.0f, -0.5f, -0.125f, 0.0f, 0.125f, 0.5f, 1.0f}; + auto packed = mt::quantize_int4(src.data(), src.size()); + std::vector back(src.size()); + bool ok = mt::dequantize_int4(packed.data(), back.data(), src.size()); + assert(ok); + const double err = max_abs_err(src, back); + // 4 bits, 7 levels per side → ~1/7 ≈ 14% step; max-abs err < + // half-step ≈ 0.072. We give margin: 0.10. + printf("test-mt-quant: int4 simple max_abs_err=%.4f\n", err); + assert(err < 0.10); + } + + // ─── int8 (no scale) — input must be in [-1, +1] ─── + { + auto src = gauss(1024, 0.3f, 42); + for (auto & v : src) v = std::max(-1.0f, std::min(1.0f, v)); + auto packed = mt::quantize_int8(src.data(), src.size()); + std::vector back(src.size()); + bool ok = mt::dequantize_int8(packed.data(), back.data(), src.size()); + assert(ok); + const double err = max_abs_err(src, back); + printf("test-mt-quant: int8 max_abs_err=%.5f\n", err); + // 8 bits → step ≈ 1/127 ≈ 0.008; half-step ≈ 0.004. Margin 0.01. + assert(err < 0.01); + } + + // ─── int4 with per-block scale — handles arbitrary range ─── + { + // Realistic K/V magnitude: stddev=2.0 → range ~[-8, +8] typical. + auto src = gauss(1024, 2.0f, 123); + std::vector packed((src.size() + 1) / 2); + float scale = 0.0f; + bool ok = mt::quantize_block_int4_with_scale( + src.data(), src.size(), &scale, packed.data()); + assert(ok); + assert(scale > 0.0f); + + std::vector back(src.size()); + ok = mt::dequantize_block_int4_with_scale( + packed.data(), scale, back.data(), src.size()); + assert(ok); + + const double err = max_abs_err(src, back); + const double cs = cosine_sim(src, back); + printf("test-mt-quant: int4-with-scale n=%zu scale=%.3f max_abs_err=%.4f cosine_sim=%.5f\n", + src.size(), (double) scale, err, cs); + fflush(stdout); + // Step = 2*scale/14; half-step = scale/14 ≈ 0.07 * scale. + // For scale~6 (max-abs of stddev=2 Gaussian over 1024 samples), + // expected max err ≈ 0.43. Margin: 1.0. + assert(err < (double) scale * 0.15); + // Cosine similarity for int4 with per-block scale on Gaussian + // data: SNR ≈ 29 → cs ≈ 0.98. Threshold 0.97 leaves margin. + assert(cs > 0.97); + } + + // ─── Edge cases ─── + { + // All zeros: scale=1 by convention; round-trip is exact. + std::vector src(64, 0.0f); + std::vector packed(32); + float scale = 0.0f; + mt::quantize_block_int4_with_scale(src.data(), src.size(), &scale, packed.data()); + std::vector back(src.size()); + mt::dequantize_block_int4_with_scale(packed.data(), scale, back.data(), src.size()); + for (auto v : back) assert(v == 0.0f); + printf("test-mt-quant: all-zero round-trip ok\n"); + } + { + // Single non-zero element: scale = |x|; that element decodes + // back to ±scale. + std::vector src(16, 0.0f); + src[7] = -3.7f; + std::vector packed(8); + float scale = 0.0f; + mt::quantize_block_int4_with_scale(src.data(), src.size(), &scale, packed.data()); + std::vector back(src.size()); + mt::dequantize_block_int4_with_scale(packed.data(), scale, back.data(), src.size()); + // src[7] = -3.7 → encoded as -1.0 (fully saturating) → decoded as -scale = -3.7. + // (encode_int4 maps -1.0 to -7, decode_int4 maps -7 to -1.0; * 3.7 = -3.7.) + printf("test-mt-quant: single-value round-trip src[7]=%.3f back[7]=%.3f scale=%.3f\n", + (double) src[7], (double) back[7], (double) scale); + assert(std::fabs(back[7] - src[7]) < 0.01); + for (size_t i = 0; i < src.size(); ++i) { + if (i == 7) continue; + assert(std::fabs(back[i]) < scale * 0.10); // others stayed near zero + } + } + + printf("test-mt-quant: ALL PASS\n"); + return 0; +} diff --git a/tests/test-mt-tiered-thin.cpp b/tests/test-mt-tiered-thin.cpp new file mode 100644 index 000000000000..82c601322bd8 --- /dev/null +++ b/tests/test-mt-tiered-thin.cpp @@ -0,0 +1,264 @@ +// MAD-137: unit tests for mt::SemanticIndex — the chunk-level +// fingerprint store owned by mt::llama_memory_tiered. +// +// Notes on scope: +// +// The MAD-137 ticket asks us to "verify the wrapper is genuinely thin +// (no paged scaffolding)" and to round-trip embed_text + +// record_chunk_fingerprint + recurrent backup/restore. Those last two +// require a real model + a real inner cache and so live at the +// integration tier (test-paged-lifecycle / stress-paged-multi-seq). +// +// What's testable in isolation here is the chunk-level SemanticIndex +// that the wrapper composes — it's the slice that survived the MAD-127 +// thinning, so exercising it end-to-end is the right unit-level proof +// that the wrapper retained its public semantic surface after the +// paged scaffolding was carved out. +// +// The absence of paged-block surface is enforced by code review (no +// BlockPool / BlockTable / llama_kv_cache_paged references in +// mt-tiered.{h,cpp}); a unit test asserting on the absence of named +// members would just be brittle without catching anything semantic. + +#include "../src/memory-tier/mt-semantic.h" + +#undef NDEBUG +#include +#include +#include +#include +#include +#include +#include + +using mt::SemanticIndex; +using Tier = SemanticIndex::Tier; + +namespace { + +std::vector normed(std::vector v) { + double sq = 0.0; + for (float x : v) sq += (double) x * x; + const double n = std::sqrt(sq); + if (n > 0.0) { + for (float & x : v) x = (float) ((double) x / n); + } + return v; +} + +std::string tmp_path(const char * tag) { + const char * dir = std::getenv("TMPDIR"); + if (dir == nullptr || dir[0] == '\0') dir = "/tmp"; + char buf[256]; + std::snprintf(buf, sizeof(buf), "%s/test-mt-tiered-thin-%s-%d.bin", dir, tag, (int) ::getpid()); + return std::string(buf); +} + +} // namespace + +int main() { + // ─── Empty state ─── + { + SemanticIndex idx; + assert(idx.size() == 0); + assert(idx.score(normed({1, 0, 0}), 5, 0.0f).empty()); + printf("test-mt-tiered-thin: empty state ok\n"); + } + + // ─── add + size — every add increments size until the FIFO cap ─── + { + SemanticIndex idx; + idx.add_fingerprint({0, 1, 2}, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint({3, 4, 5}, normed({0, 1, 0}), Tier::Warm); + assert(idx.size() == 2); + printf("test-mt-tiered-thin: add + size ok\n"); + } + + // ─── score — descending order, returns positions + tier ─── + { + SemanticIndex idx; + // Three chunks at known angles to query (1,0,0): + // chunk A: (1,0,0) cos = 1.0 + // chunk B: (1,1,0) / √2 cos ≈ 0.707 + // chunk C: (0,1,0) cos = 0.0 + idx.add_fingerprint({100}, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint({200}, normed({1, 1, 0}), Tier::Warm); + idx.add_fingerprint({300}, normed({0, 1, 0}), Tier::Cold); + + const auto hits = idx.score(normed({1, 0, 0}), /* top_k */ 10, /* threshold */ -1.0f); + assert(hits.size() == 3); + assert(hits[0].positions.size() == 1 && hits[0].positions[0] == 100); + assert(hits[0].tier == Tier::Hot); + assert(hits[1].positions.size() == 1 && hits[1].positions[0] == 200); + assert(hits[1].tier == Tier::Warm); + assert(hits[2].positions.size() == 1 && hits[2].positions[0] == 300); + assert(hits[2].tier == Tier::Cold); + // Strictly descending. + assert(hits[0].score > hits[1].score); + assert(hits[1].score > hits[2].score); + printf("test-mt-tiered-thin: score returns positions + tier in descending order ok\n"); + } + + // ─── threshold filter ─── + { + SemanticIndex idx; + idx.add_fingerprint({1}, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint({2}, normed({1, 1, 0}), Tier::Warm); // ~0.707 + idx.add_fingerprint({3}, normed({0, 1, 0}), Tier::Cold); // ~0.0 + + const auto hits = idx.score(normed({1, 0, 0}), 10, /* threshold */ 0.5f); + assert(hits.size() == 2); + assert(hits[0].positions[0] == 1); + assert(hits[1].positions[0] == 2); + printf("test-mt-tiered-thin: threshold filter ok\n"); + } + + // ─── top_k cap + edge cases ─── + { + SemanticIndex idx; + for (int i = 0; i < 5; ++i) { + idx.add_fingerprint({(llama_pos) i}, normed({1, 0, 0}), Tier::Hot); + } + assert(idx.score(normed({1, 0, 0}), /* top_k */ 3, -1.0f).size() == 3); + assert(idx.score(normed({1, 0, 0}), /* top_k */ 0, -1.0f).empty()); + assert(idx.score(normed({1, 0, 0}), -1, -1.0f).empty()); + assert(idx.score(/* empty query */ {}, 5, -1.0f).empty()); + printf("test-mt-tiered-thin: top_k cap + edge cases ok\n"); + } + + // ─── FIFO eviction at kMaxFingerprints ─── + // Fill past the cap; verify size plateaus and the first-inserted + // entry is the one that gets evicted (lowest turn). + { + SemanticIndex idx; + const size_t cap = SemanticIndex::kMaxFingerprints; + + // Insert exactly the cap. The first one carries position 0 and + // a unique direction we can probe for later. + idx.add_fingerprint({/* pos */ 0}, normed({1, 0, 0, 0}), Tier::Hot); + for (size_t i = 1; i < cap; ++i) { + idx.add_fingerprint({(llama_pos) i}, normed({0, 1, 0, 0}), Tier::Warm); + } + assert(idx.size() == cap); + + // The oldest entry is still findable by its unique direction. + { + const auto hits = idx.score(normed({1, 0, 0, 0}), 1, /* threshold */ 0.99f); + assert(hits.size() == 1); + assert(hits[0].positions[0] == 0); + } + + // Insert one more — total should stay at cap, oldest must be gone. + idx.add_fingerprint({(llama_pos) cap}, normed({0, 0, 1, 0}), Tier::Cold); + assert(idx.size() == cap); + { + const auto hits = idx.score(normed({1, 0, 0, 0}), 1, 0.99f); + assert(hits.empty()); // the (1,0,0,0) entry was evicted + } + // The newest is findable. + { + const auto hits = idx.score(normed({0, 0, 1, 0}), 1, 0.99f); + assert(hits.size() == 1); + assert(hits[0].positions[0] == (llama_pos) cap); + } + printf("test-mt-tiered-thin: FIFO eviction at kMaxFingerprints ok\n"); + } + + // ─── clear + reuse — clear() resets size and the turn counter ─── + { + SemanticIndex idx; + idx.add_fingerprint({0}, normed({1, 0, 0}), Tier::Hot); + idx.add_fingerprint({1}, normed({0, 1, 0}), Tier::Warm); + assert(idx.size() == 2); + + idx.clear(); + assert(idx.size() == 0); + assert(idx.score(normed({1, 0, 0}), 5, -1.0f).empty()); + + // Index is reusable. + idx.add_fingerprint({0}, normed({1, 0, 0}), Tier::Hot); + const auto hits = idx.score(normed({1, 0, 0}), 1, 0.99f); + assert(hits.size() == 1); + printf("test-mt-tiered-thin: clear + reuse ok\n"); + } + + // ─── save_to_disk + load_from_disk round-trip (MTFI v1) ─── + { + const std::string path = tmp_path("rt"); + + SemanticIndex original; + original.add_fingerprint({10, 11, 12}, normed({1, 0, 0, 0}), Tier::Hot); + original.add_fingerprint({20, 21}, normed({0, 1, 0, 0}), Tier::Warm); + original.add_fingerprint({30}, normed({0, 0, 1, 0}), Tier::Cold); + + const bool save_ok = original.save_to_disk(path); + assert(save_ok); + + SemanticIndex restored; + // Pre-load: stray entry to verify load REPLACES (not merges). + restored.add_fingerprint({999}, normed({1, 0, 0, 0}), Tier::Hot); + assert(restored.size() == 1); + + const bool load_ok = restored.load_from_disk(path); + assert(load_ok); + assert(restored.size() == 3); // not 4 — replaced + + // Probe each direction; verify positions + tier survived. + { + const auto hits = restored.score(normed({1, 0, 0, 0}), 1, 0.99f); + assert(hits.size() == 1); + assert(hits[0].positions == std::vector({10, 11, 12})); + assert(hits[0].tier == Tier::Hot); + } + { + const auto hits = restored.score(normed({0, 1, 0, 0}), 1, 0.99f); + assert(hits.size() == 1); + assert(hits[0].positions == std::vector({20, 21})); + assert(hits[0].tier == Tier::Warm); + } + { + const auto hits = restored.score(normed({0, 0, 1, 0}), 1, 0.99f); + assert(hits.size() == 1); + assert(hits[0].positions == std::vector({30})); + assert(hits[0].tier == Tier::Cold); + } + + ::unlink(path.c_str()); + printf("test-mt-tiered-thin: save/load round-trip ok\n"); + } + + // ─── load_from_disk on missing file returns false ─── + { + SemanticIndex idx; + idx.add_fingerprint({1}, normed({1, 0, 0}), Tier::Hot); + const bool load_ok = idx.load_from_disk("/tmp/this-file-does-not-exist-mad137-tiered"); + assert(!load_ok); + // In-memory state untouched. + assert(idx.size() == 1); + printf("test-mt-tiered-thin: load missing file safe ok\n"); + } + + // ─── load_from_disk on wrong magic returns false ─── + { + const std::string path = tmp_path("badmagic"); + FILE * f = std::fopen(path.c_str(), "wb"); + assert(f); + const uint32_t wrong_magic = 0xCAFEBABE; + const uint32_t version = 1; + std::fwrite(&wrong_magic, sizeof(wrong_magic), 1, f); + std::fwrite(&version, sizeof(version), 1, f); + std::fclose(f); + + SemanticIndex idx; + idx.add_fingerprint({1}, normed({1, 0, 0}), Tier::Hot); + const bool load_ok = idx.load_from_disk(path); + assert(!load_ok); + assert(idx.size() == 1); + + ::unlink(path.c_str()); + printf("test-mt-tiered-thin: load bad magic safe ok\n"); + } + + printf("test-mt-tiered-thin: ALL PASS\n"); + return 0; +} diff --git a/tests/test-paged-lifecycle.cpp b/tests/test-paged-lifecycle.cpp new file mode 100644 index 000000000000..8bbaa13dc843 --- /dev/null +++ b/tests/test-paged-lifecycle.cpp @@ -0,0 +1,325 @@ +// MAD-137: integration test for the paged-cache lifecycle. +// +// Drives a real `llama_kv_cache_paged` instance through: +// - block allocation +// - hot→warm eviction and warm→hot restore (MAD-120) +// - cold spill via ssd_path (MAD-121) +// - whole-seq seq_rm (drops all blocks) +// - partial seq_rm (block-aligned wipe + sub-block warning path) +// - CoW seq_cp (refcount bump) +// - state_write → state_read round-trip on a fresh instance (MAD-130) +// - cold-resume from sidecar (MAD-130) +// +// The cache is constructed on the CPU backend so the test runs without a +// GPU. A real model is required to source hparams (n_layer, head_dim, +// n_kv_heads); follow the existing convention — pass the model path as +// argv[1] or set LLAMACPP_TEST_MODELFILE. If neither is present the test +// prints a warning and exits 0 (so CI without a model checkpoint stays +// green). + +#include "../src/llama-kv-cache-paged.h" +#include "../src/llama-model.h" +#include "../src/llama-io.h" +#include "get-model.h" +#include "llama.h" +#include "ggml-backend.h" + +#undef NDEBUG +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// ---- in-memory io_write_i / io_read_i for state_write/state_read ---- + +class MemoryWriter : public llama_io_write_i { +public: + void write(const void * src, size_t size) override { + const uint8_t * p = static_cast(src); + buf_.insert(buf_.end(), p, p + size); + } + void write_tensor(ggml_tensor * /*tensor*/, size_t /*offset*/, size_t size) override { + // For our test scenario the cache uses write() directly for its + // structural state; tensor writes don't appear. Append zeros so + // n_bytes() stays consistent if the path ever changes. + buf_.insert(buf_.end(), size, 0); + } + size_t n_bytes() override { return buf_.size(); } + + const std::vector & data() const { return buf_; } + +private: + std::vector buf_; +}; + +class MemoryReader : public llama_io_read_i { +public: + explicit MemoryReader(const std::vector & buf) : buf_(buf) {} + void read(void * dst, size_t size) override { + assert(pos_ + size <= buf_.size() && "MemoryReader underflow"); + std::memcpy(dst, buf_.data() + pos_, size); + pos_ += size; + } + void read_tensor(ggml_tensor * /*tensor*/, size_t /*offset*/, size_t size) override { + assert(pos_ + size <= buf_.size() && "MemoryReader underflow"); + pos_ += size; + } + size_t n_bytes() override { return pos_; } + +private: + const std::vector & buf_; + size_t pos_ = 0; +}; + +// Per-test directory under TMPDIR so parallel runs and sandboxed +// environments don't collide. +std::string tmp_dir(const char * tag) { + const char * dir = std::getenv("TMPDIR"); + if (dir == nullptr || dir[0] == '\0') dir = "/tmp"; + char buf[256]; + std::snprintf(buf, sizeof(buf), "%s/test-paged-lifecycle-%s-%d", dir, tag, (int) ::getpid()); + return std::string(buf); +} + +} // namespace + +int main(int argc, char * argv[]) { + char * model_path = get_model_or_exit(argc, argv); + + // --- Load the model on CPU only. We never run a forward pass; we + // only need hparams (n_layer, head_dim, n_kv_heads) and the model + // object itself for the cache constructor. --- + llama_backend_init(); + + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = 0; + mparams.use_mmap = true; + mparams.use_mlock = false; + mparams.vocab_only = false; + + llama_model * model = llama_model_load_from_file(model_path, mparams); + if (model == nullptr) { + fprintf(stderr, "test-paged-lifecycle: failed to load model %s — skipping\n", model_path); + llama_backend_free(); + return 0; + } + + // CPU backend buffer type — keeps everything host-resident so no GPU. + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + assert(buft != nullptr); + + // Small cache config so we can exercise eviction quickly: + // 4 GPU blocks × block_size=4 = 16 hot tokens + // 4 warm blocks (host) + // 4 cold blocks (SSD) + // 2 sequences, up to 8 logical blocks per seq + constexpr uint32_t kNBlocks = 4; + constexpr uint32_t kBlockSize = 4; + constexpr uint32_t kNSeqMax = 2; + constexpr uint32_t kMaxBlks = 16; // generous so cold-spill scenario fits + constexpr uint32_t kWarmBlocks = 4; + constexpr uint32_t kColdBlocks = 4; + + const std::string ssd_dir = tmp_dir("rt"); + ::mkdir(ssd_dir.c_str(), 0700); + + // ─── Construct cache, exhaust GPU pool, evict to warm, restore ─── + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + kWarmBlocks, /*n_cold_blocks=*/0, + /*ssd_path=*/std::string(), + /*cold_resume=*/false, + /*instance_id=*/"test-rt"); + + // Allocate every GPU block to seq 0 — exhausts the pool. + const uint32_t need_tokens = kNBlocks * kBlockSize; // 16 tokens → 4 blocks + const bool grew = cache.ensure_blocks_for(/*seq*/ 0, need_tokens); + assert(grew); + + // Next allocation request without freeing first should not be + // satisfiable from the GPU pool — but ensure_blocks_for can + // internally trigger evict_lru_to_warm to make room. With warm + // available it must succeed. + const bool grew_more = cache.ensure_blocks_for(/*seq*/ 1, kBlockSize); // 1 block + assert(grew_more && "warm tier should absorb the spillover"); + + // The eviction counter must have ticked. + assert(cache.evict_h2w_total() >= 1); + + printf("test-paged-lifecycle: alloc + hot→warm spill ok\n"); + + // Wipe seq 0 — every logical block returns to the pool. After + // this, GPU pool plus warm pool is fully refilled. + const bool removed = cache.seq_rm(/*seq*/ 0, /*p0=*/-1, /*p1=*/-1); + assert(removed); + printf("test-paged-lifecycle: whole-seq seq_rm ok\n"); + } + + // ─── Cold spill: warm full → spill oldest warm to cold ─── + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + kWarmBlocks, kColdBlocks, + ssd_dir, + /*cold_resume=*/false, + /*instance_id=*/"test-cold"); + + // Drive seq 0 long enough to overflow GPU + warm and force at + // least one cold spill: 4 hot + 4 warm = 8 blocks before cold + // engagement; the next ensure_blocks_for must spill. + const uint32_t total_tokens = (kNBlocks + kWarmBlocks + 1) * kBlockSize; + const bool grew = cache.ensure_blocks_for(/*seq*/ 0, total_tokens); + assert(grew); + assert(cache.evict_w2c_total() >= 1 && "warm→cold spill should have fired"); + printf("test-paged-lifecycle: warm→cold spill ok (w2c=%llu)\n", + (unsigned long long) cache.evict_w2c_total()); + } + + // ─── Partial seq_rm — block-aligned middle wipe ─── + // + // Note: pos_max-related accessors (seq_pos_max, seq_cp range copies) + // require the per-batch tensor population path to have run, which we + // can't trigger without spinning up a real llama_context + a graph. + // What this test verifies is that the seq_rm call doesn't crash and + // returns true — the freed-block bookkeeping is logged at info level + // and visible in the run output. Stress tests with real batches + // exercise the pos_max-driven assertions. + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, + /*n_cold_blocks=*/0, + std::string(), + false, "test-partrm"); + + const bool grew = cache.ensure_blocks_for(/*seq*/ 0, 3 * kBlockSize); + assert(grew); + + // Block-aligned middle wipe — covers the whole second block. + const bool ok = cache.seq_rm(/*seq*/ 0, /*p0=*/4, /*p1=*/8); + assert(ok); + printf("test-paged-lifecycle: partial seq_rm (middle block, no crash) ok\n"); + } + + // ─── seq_cp — verify the call doesn't crash on simple inputs ─── + // + // Same caveat as partial seq_rm: the full CoW range-copy semantics + // depend on pos_max, which only gets set during batch processing. + // What this verifies is the call signature + dispatch path runs. + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, + /*n_cold_blocks=*/0, + std::string(), + false, "test-cow"); + + const bool grew = cache.ensure_blocks_for(/*seq*/ 0, 2 * kBlockSize); + assert(grew); + + // Calling seq_cp with no live tokens (pos_max unset) is a no-op + // path; at minimum it must not crash and dst seq must remain + // unaffected. + cache.seq_cp(/*src*/ 0, /*dst*/ 1, /*p0=*/0, /*p1=*/2 * kBlockSize); + printf("test-paged-lifecycle: seq_cp no-crash ok\n"); + } + + // ─── state_write → state_read round-trip on structural state ─── + // + // After ensure_blocks_for, the cache has a populated block table + // even though pos_max is still -1 (no batch ran). state_write should + // serialize that structural state; state_read on a fresh cache + // should reconstruct enough that subsequent reads match. + { + llama_kv_cache_paged cache_a( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, + /*n_cold_blocks=*/0, + std::string(), + false, "test-state-a"); + + const bool grew = cache_a.ensure_blocks_for(/*seq*/ 0, 2 * kBlockSize); + assert(grew); + + MemoryWriter w; + cache_a.state_write(w, /*seq_id=*/0, /*flags=*/0); + assert(w.n_bytes() > 0 && "state_write produced empty buffer"); + + // Round-trip into a fresh cache. state_read should consume the + // exact byte count the writer produced (no underflow / overflow). + llama_kv_cache_paged cache_b( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, + /*n_cold_blocks=*/0, + std::string(), + false, "test-state-b"); + + MemoryReader r(w.data()); + cache_b.state_read(r, /*seq_id=*/0, /*flags=*/0); + assert(r.n_bytes() == w.data().size() && + "state_read should consume exactly the bytes state_write produced"); + + printf("test-paged-lifecycle: state_write/state_read round-trip ok (%zu bytes)\n", + w.data().size()); + } + + // ─── Cold-resume: second instance with same ssd_path + cold_resume=true ─── + // The point is just that the constructor accepts the resume flag and + // re-opens the cold-tier files without truncating; we don't validate + // recovered KV bytes here (that's stress-test territory). + { + const std::string resume_dir = tmp_dir("resume"); + ::mkdir(resume_dir.c_str(), 0700); + + // First instance writes some cold-tier data. + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + kWarmBlocks, kColdBlocks, + resume_dir, + /*cold_resume=*/false, + /*instance_id=*/"test-resume"); + // Force at least one cold spill. + const uint32_t total = (kNBlocks + kWarmBlocks + 1) * kBlockSize; + cache.ensure_blocks_for(/*seq*/ 0, total); + } + + // Second instance with cold_resume=true must construct without + // wiping the cold sidecar. + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + kWarmBlocks, kColdBlocks, + resume_dir, + /*cold_resume=*/true, + /*instance_id=*/"test-resume"); + // Live cold pool was rebuilt from the sidecar — n_cold_blocks + // is the configured size regardless. The accessor exists and + // returns a sane value. + assert(cache.n_cold_blocks() == kColdBlocks); + } + printf("test-paged-lifecycle: cold-resume ctor ok\n"); + } + + // ─── Cleanup ─── + llama_model_free(model); + llama_backend_free(); + printf("test-paged-lifecycle: ALL PASS\n"); + return 0; +} diff --git a/tests/test-paged-semantic.cpp b/tests/test-paged-semantic.cpp new file mode 100644 index 000000000000..d0f854f1b5c2 --- /dev/null +++ b/tests/test-paged-semantic.cpp @@ -0,0 +1,236 @@ +// MAD-137: integration test for the paged-cache semantic-prefetch path +// (MAD-125 / MAD-129 surface). +// +// Drives `llama_kv_cache_paged`'s fingerprint store + restore path: +// - record_paged_block_fingerprint writes per-block fingerprints +// - has_paged_fingerprint / n_paged_fingerprints reflect the writes +// - restore_semantic_paged scores them against a query and ticks +// the semantic_attempts counter +// - whole-seq seq_rm drops the seq's fingerprints +// - save/load_paged_fingerprints round-trips through PSFI v1 +// +// Constructs the cache on the CPU backend so no GPU is required. A real +// model is needed for hparams (n_layer, head_dim, n_kv_heads); skips +// gracefully when neither argv[1] nor LLAMACPP_TEST_MODELFILE is set. + +#include "../src/llama-kv-cache-paged.h" +#include "../src/llama-model.h" +#include "../src/memory-tier/mt-semantic.h" +#include "get-model.h" +#include "llama.h" +#include "ggml-backend.h" + +#undef NDEBUG +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// L2-normalize so cosine similarity == dot product. +std::vector normed(std::vector v) { + double sq = 0.0; + for (float x : v) sq += (double) x * x; + const double n = std::sqrt(sq); + if (n > 0.0) { + for (float & x : v) x = (float) ((double) x / n); + } + return v; +} + +std::string tmp_path(const char * tag) { + const char * dir = std::getenv("TMPDIR"); + if (dir == nullptr || dir[0] == '\0') dir = "/tmp"; + char buf[256]; + std::snprintf(buf, sizeof(buf), "%s/test-paged-semantic-%s-%d.bin", dir, tag, (int) ::getpid()); + return std::string(buf); +} + +} // namespace + +int main(int argc, char * argv[]) { + char * model_path = get_model_or_exit(argc, argv); + + llama_backend_init(); + + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = 0; + mparams.use_mmap = true; + + llama_model * model = llama_model_load_from_file(model_path, mparams); + if (model == nullptr) { + fprintf(stderr, "test-paged-semantic: failed to load model %s — skipping\n", model_path); + llama_backend_free(); + return 0; + } + + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + assert(buft != nullptr); + + constexpr uint32_t kNBlocks = 8; + constexpr uint32_t kBlockSize = 4; + constexpr uint32_t kNSeqMax = 2; + constexpr uint32_t kMaxBlks = 16; + constexpr uint32_t kWarmBlocks = 4; + + // ─── record / has / size — fingerprint ingestion ─── + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + kWarmBlocks, /*n_cold_blocks=*/0, + std::string(), + /*cold_resume=*/false, + /*instance_id=*/"test-record"); + + const bool grew = cache.ensure_blocks_for(/*seq*/ 0, 4 * kBlockSize); + assert(grew); + assert(cache.n_paged_fingerprints() == 0); + assert(!cache.has_paged_fingerprint(0, 0)); + + // Record a fingerprint per logical block, each in a unique + // direction so we can tell them apart later. + cache.record_paged_block_fingerprint(0, 0, normed({1, 0, 0, 0}), mt::SemanticIndex::Tier::Hot); + cache.record_paged_block_fingerprint(0, 1, normed({0, 1, 0, 0}), mt::SemanticIndex::Tier::Hot); + cache.record_paged_block_fingerprint(0, 2, normed({0, 0, 1, 0}), mt::SemanticIndex::Tier::Hot); + cache.record_paged_block_fingerprint(0, 3, normed({0, 0, 0, 1}), mt::SemanticIndex::Tier::Hot); + + assert(cache.n_paged_fingerprints() == 4); + assert(cache.has_paged_fingerprint(0, 0)); + assert(cache.has_paged_fingerprint(0, 3)); + assert(!cache.has_paged_fingerprint(0, 99)); + assert(!cache.has_paged_fingerprint(1, 0)); + printf("test-paged-semantic: record + has + size ok\n"); + + // ─── restore_semantic_paged — scores fingerprints, ticks counter ─── + // The blocks are still in hot (no eviction happened), so an + // actual warm→hot fault doesn't have anything to do; what we're + // verifying is that the call (a) doesn't crash, (b) ticks the + // semantic_attempts counter, and (c) the requested query matches + // a stored fingerprint in BlockSemanticIndex (visible by querying + // it directly through the existing accessor). + const uint64_t attempts_before = cache.semantic_attempts_total(); + const uint32_t restored = cache.restore_semantic_paged( + /*seq*/ 0, normed({1, 0, 0, 0}), /*top_k=*/2, /*threshold=*/0.5f); + // restored may be 0 — the matching block (lblock 0) is already + // in hot, so there's nothing to fault in. The point is the + // attempt was counted. + (void) restored; + assert(cache.semantic_attempts_total() == attempts_before + 1); + printf("test-paged-semantic: restore_semantic_paged attempts counter ticks ok\n"); + } + + // ─── Whole-seq seq_rm drops the seq's fingerprints ─── + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, /*n_cold_blocks=*/0, + std::string(), + false, "test-rm-fps"); + + const bool grew = cache.ensure_blocks_for(/*seq*/ 0, 2 * kBlockSize); + assert(grew); + cache.record_paged_block_fingerprint(0, 0, normed({1, 0, 0, 0}), mt::SemanticIndex::Tier::Hot); + cache.record_paged_block_fingerprint(0, 1, normed({0, 1, 0, 0}), mt::SemanticIndex::Tier::Hot); + // Seed a different seq's fingerprint to verify it survives. + const bool grew_b = cache.ensure_blocks_for(/*seq*/ 1, kBlockSize); + assert(grew_b); + cache.record_paged_block_fingerprint(1, 0, normed({0, 0, 1, 0}), mt::SemanticIndex::Tier::Hot); + assert(cache.n_paged_fingerprints() == 3); + + // Whole-seq wipe of seq 0. + const bool removed = cache.seq_rm(/*seq*/ 0, /*p0=*/-1, /*p1=*/-1); + assert(removed); + + // seq 0's fingerprints are gone; seq 1's survives. + assert(!cache.has_paged_fingerprint(0, 0)); + assert(!cache.has_paged_fingerprint(0, 1)); + assert(cache.has_paged_fingerprint(1, 0)); + assert(cache.n_paged_fingerprints() == 1); + printf("test-paged-semantic: whole-seq seq_rm drops only that seq's fingerprints ok\n"); + } + + // ─── save_paged_fingerprints / load_paged_fingerprints round-trip ─── + { + const std::string path = tmp_path("rt"); + + llama_kv_cache_paged cache_a( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, /*n_cold_blocks=*/0, + std::string(), + false, "test-fp-a"); + + const bool grew = cache_a.ensure_blocks_for(/*seq*/ 0, 3 * kBlockSize); + assert(grew); + cache_a.record_paged_block_fingerprint(0, 0, normed({1, 0, 0, 0}), mt::SemanticIndex::Tier::Hot); + cache_a.record_paged_block_fingerprint(0, 1, normed({0, 1, 0, 0}), mt::SemanticIndex::Tier::Warm); + cache_a.record_paged_block_fingerprint(0, 2, normed({0, 0, 1, 0}), mt::SemanticIndex::Tier::Cold); + assert(cache_a.n_paged_fingerprints() == 3); + + const bool save_ok = cache_a.save_paged_fingerprints(path); + assert(save_ok); + + // Load into a fresh cache that has its own (different) fingerprint + // pre-loaded; load must REPLACE rather than merge. + llama_kv_cache_paged cache_b( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + /*n_warm_blocks=*/0, /*n_cold_blocks=*/0, + std::string(), + false, "test-fp-b"); + const bool grew_b = cache_b.ensure_blocks_for(/*seq*/ 0, kBlockSize); + assert(grew_b); + cache_b.record_paged_block_fingerprint(7, 99, normed({1, 0, 0, 0}), mt::SemanticIndex::Tier::Hot); + assert(cache_b.n_paged_fingerprints() == 1); + + const bool load_ok = cache_b.load_paged_fingerprints(path); + assert(load_ok); + assert(cache_b.n_paged_fingerprints() == 3); // not 4 — load replaces + assert(cache_b.has_paged_fingerprint(0, 0)); + assert(cache_b.has_paged_fingerprint(0, 1)); + assert(cache_b.has_paged_fingerprint(0, 2)); + assert(!cache_b.has_paged_fingerprint(7, 99)); // stray entry was wiped + + ::unlink(path.c_str()); + printf("test-paged-semantic: save/load fingerprints round-trip ok\n"); + } + + // ─── Edge cases for restore_semantic_paged ─── + { + llama_kv_cache_paged cache( + *model, buft, + kNBlocks, kBlockSize, kNSeqMax, kMaxBlks, + kWarmBlocks, /*n_cold_blocks=*/0, + std::string(), + false, "test-edge"); + + // No fingerprints at all → restore returns 0 cleanly. + const uint32_t r0 = cache.restore_semantic_paged(0, normed({1, 0, 0, 0}), 5, 0.5f); + assert(r0 == 0); + + // Empty query → must not crash; returns 0. + const uint32_t r1 = cache.restore_semantic_paged(0, /* empty */ {}, 5, 0.5f); + assert(r1 == 0); + + // top_k = 0 → no work; returns 0. + const uint32_t r2 = cache.restore_semantic_paged(0, normed({1, 0, 0, 0}), 0, 0.5f); + assert(r2 == 0); + + // Unknown seq → no fingerprints for it; returns 0. + const uint32_t r3 = cache.restore_semantic_paged(99, normed({1, 0, 0, 0}), 5, 0.5f); + assert(r3 == 0); + + printf("test-paged-semantic: restore edge cases ok\n"); + } + + llama_model_free(model); + llama_backend_free(); + printf("test-paged-semantic: ALL PASS\n"); + return 0; +} diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index ecbbcc64b653..ffebebb4a1f7 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -15,8 +15,6 @@ add_library(${TARGET} STATIC server-common.h server-context.cpp server-context.h - server-tiered-cache.cpp - server-tiered-cache.h server-tools.cpp server-tools.h ) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index beed655494aa..f79db66289b2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -5,8 +5,9 @@ #include "server-http.h" #include "server-task.h" #include "server-queue.h" -#include "server-tiered-cache.h" #include "../src/llama-model.h" +#include "../src/llama-kv-cache-paged.h" +#include "../src/llama-memory-hybrid.h" #include "../src/memory-tier/mt-tiered.h" #include "build-info.h" @@ -22,11 +23,17 @@ #include #include #include +#include #include #include #include #include +#include +#include +#include +#include + // fix problem with std::min and std::max #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN @@ -40,6 +47,85 @@ using json = nlohmann::ordered_json; constexpr int HTTP_POLLING_SECONDS = 1; +// MAD-125: walk the active memory pointer chain and return the +// llama_kv_cache_paged at the bottom (if any). Three nestings to cover: +// - paged is the raw active memory (rare standalone test case) +// - paged is the attention member of llama_memory_hybrid (hybrid models) +// - llama_memory_hybrid is wrapped by mt::llama_memory_tiered (the +// tiered + paged + hybrid stack — Qwen3.x family with --kv-tiered +// and --kv-tier-paged-blocks). +// Returns nullptr when the active memory isn't running paged-blocks. +static llama_kv_cache_paged * mt_get_paged_cache(llama_memory_i * mem) { + if (!mem) return nullptr; + if (auto * p = dynamic_cast(mem)) return p; + if (auto * h = dynamic_cast(mem)) return h->get_mem_attn_paged(); + if (auto * t = dynamic_cast(mem)) { + return mt_get_paged_cache(t->inner_for_test()); + } + return nullptr; +} + +// MAD-122/125: write fingerprints for the token range [p0, p1). When a +// paged_cache is supplied we emit one BGE-small embedding per logical +// block via llama_kv_cache_paged::record_paged_block_fingerprint so +// query-time semantic prefetch can score at block granularity. When +// paged_cache is null we fall back to the legacy chunk-level path +// (one embedding for the whole range, position-keyed) on the tier +// wrapper — the dispatch happens at the call site so this function +// stays a single helper. +// +// mt_tier is required either way: it owns the embed_text / bge-small +// model. paged_cache is the destination; null = legacy path. +// +// Returns the number of fingerprints actually recorded. +static int mt_record_fingerprints_for_range( + mt::llama_memory_tiered * mt_tier, + llama_kv_cache_paged * paged_cache, + llama_context * ctx, + llama_seq_id seq_id, + const llama_tokens & toks, + int p0, + int p1, + uint32_t block_size) { + if (!mt_tier || p1 <= p0) return 0; + const int hi = std::min(p1, (int) toks.size()); + if (hi <= p0) return 0; + + if (!paged_cache) { + llama_tokens chunk(toks.begin() + p0, toks.begin() + hi); + const std::string text = common_detokenize(ctx, chunk, /*special=*/ false); + const auto emb = mt_tier->embed_text(text); + if (emb.empty()) return 0; + std::vector positions; + positions.reserve(hi - p0); + for (int i = p0; i < hi; ++i) positions.push_back((llama_pos) i); + mt_tier->record_chunk_fingerprint( + std::move(positions), emb, mt::SemanticIndex::Tier::Warm); + return 1; + } + + // Paged path: walk in block_size strides starting from a block-aligned + // floor. paged_backup is itself block-aligned so p0 should already be + // on a boundary, but rounding down keeps the lblock arithmetic clean + // even if a future caller passes an unaligned range. + const uint32_t bsize = block_size > 0 ? block_size : 16u; + const int aligned_p0 = (int)((uint32_t) p0 / bsize) * (int) bsize; + int n_recorded = 0; + for (int b = aligned_p0; b < hi; b += (int) bsize) { + const int chunk_hi = std::min(b + (int) bsize, hi); + if (chunk_hi <= b) continue; + llama_tokens chunk(toks.begin() + b, toks.begin() + chunk_hi); + const std::string text = common_detokenize(ctx, chunk, /*special=*/ false); + const auto emb = mt_tier->embed_text(text); + if (emb.empty()) continue; + const uint32_t lblock = (uint32_t) b / bsize; + paged_cache->record_paged_block_fingerprint( + seq_id, lblock, emb, mt::SemanticIndex::Tier::Warm); + ++n_recorded; + } + return n_recorded; +} + static void server_prompt_checkpoint_update(server_prompt_checkpoint & ckpt, llama_context * ctx, int id, int64_t n_tokens, bool on_device, llama_pos pos_min = -1, llama_pos pos_max = -1) { if (pos_min == -1) { pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx), id); @@ -122,6 +208,17 @@ struct server_slot { int32_t n_prompt_tokens_cache = 0; int32_t n_prompt_tokens_processed = 0; + // MAD-141: counts how many consecutive update_slots() iterations have + // preempted this slot via the MAD-120 admission gate without making + // any progress (evict_seq returned 0). Once it crosses + // kPagedPreemptDeadlockThreshold, the slot's request is unservable + // (the prompt simply doesn't fit alongside the rest of the live + // workload) and we fail it with a 503-equivalent rather than spin + // until the upstream "n_empty_consecutive > 3" safety abort fires. + // Reset on slot.reset() and any iteration where the slot DOES make + // progress through the prefill path. + int32_t paged_preempt_no_progress_count = 0; + size_t last_nl_pos = 0; std::string generated_text; @@ -208,7 +305,8 @@ struct server_slot { void reset() { SLT_DBG(*this, "%s", "\n"); - n_prompt_tokens_cache = 0; + n_prompt_tokens_cache = 0; + paged_preempt_no_progress_count = 0; // MAD-141 last_nl_pos = 0; generated_text = ""; @@ -704,7 +802,6 @@ struct server_context_impl { int n_empty_consecutive = 0; std::unique_ptr prompt_cache; - std::unique_ptr tiered_cache; server_metrics metrics; @@ -769,6 +866,36 @@ struct server_context_impl { SRV_INF("loading model '%s'\n", params.model.path.c_str()); + // MAD-134: validate tier-related config BEFORE model load so + // operators get fast clear failure instead of late crashes. + if (!params.kv_semantic_index.empty()) { + struct stat st; + if (::stat(params.kv_semantic_index.c_str(), &st) != 0) { + SRV_ERR("--kv-tier-semantic-index '%s' does not exist or is not " + "readable. Either provide a valid bge-small / nomic-embed gguf " + "file, or omit the flag to disable semantic prefetch.\n", + params.kv_semantic_index.c_str()); + return false; + } + } + if (params.kv_tiered_enabled && params.kv_tier_cold_pct > 0.0f && + !params.kv_tier_ssd_path.empty()) { + // Try to mkdir + create a test file to confirm writability. + const std::string test_dir = params.kv_tier_ssd_path; + (void) ::mkdir(test_dir.c_str(), 0700); // ok if exists + const std::string test_path = test_dir + "/.write_test_" + + std::to_string(::getpid()); + int fd = ::open(test_path.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600); + if (fd < 0) { + SRV_ERR("--kv-tier-ssd-path '%s' is not writable (errno %d: %s). " + "Pick a different path or fix permissions.\n", + test_dir.c_str(), errno, strerror(errno)); + return false; + } + ::close(fd); + ::unlink(test_path.c_str()); + } + params_base = params; if (params_base.kv_tiered_enabled && params_base.kv_tier_hot_pct > 0.0f && params_base.kv_tier_hot_pct < 100.0f) { // Save the original n_ctx so the cache layer can size pools @@ -886,7 +1013,25 @@ struct server_context_impl { if (!llama_memory_can_shift(llama_get_memory(ctx))) { if (params_base.ctx_shift) { params_base.ctx_shift = false; - SRV_WRN("%s\n", "ctx_shift is not supported by this context, it will be disabled"); + if (params_base.kv_tier_paged_blocks) { + // MAD-128: paged-blocks deliberately doesn't support + // position-shift in-place (would require re-indexing + // every block_table entry + per-block reindex of the + // GPU K/V layout). Instead: when a slot hits its + // context limit, the server stops it with + // STOP_TYPE_LIMIT and the caller resubmits as a fresh + // request — paged's prompt-cache + semantic prefetch + // recover most of the prefix on the next prefill. + SRV_WRN("%s\n", "ctx_shift disabled: --kv-tier-paged-blocks " + "is set; paged attention manages context via " + "tier movement, not in-place shift. Slots will " + "stop at n_ctx; clients should re-submit fresh " + "requests (prompt cache + semantic prefetch " + "will recover the prefix). Plan capacity " + "accordingly."); + } else { + SRV_WRN("%s\n", "ctx_shift is not supported by this context, it will be disabled"); + } } if (params_base.n_cache_reuse) { @@ -902,13 +1047,6 @@ struct server_context_impl { } } - // Initialize tiered cache if enabled - if (params_base.kv_tiered_enabled) { - tiered_cache = std::make_unique(params_base); - SRV_INF("tiered cache initialized, hot=%f%%, warm=%f%%, cold=%f%%\n", - params_base.kv_tier_hot_pct, params_base.kv_tier_warm_pct, params_base.kv_tier_cold_pct); - } - n_swa = params_base.swa_full ? 0 : llama_model_n_swa(model); // Necessary similarity of prompt for slot selection @@ -964,14 +1102,6 @@ struct server_context_impl { } } - // initialize tier manager for slot if enabled - if (params_base.kv_tiered_enabled) { - if (!tiered_cache->init_slot(i, *model, ctx)) { - SRV_WRN("failed to initialize tier manager for slot %d\n", i); - } else { - SLT_INF(slot, "tier manager initialized for slot %s\n", ""); - } - } SLT_INF(slot, "new slot, n_ctx = %d\n", slot.n_ctx); @@ -1482,23 +1612,14 @@ struct server_context_impl { // is the ONLY path that populates mt::'s warm + cold tiers // (otherwise the seq_rm-time backup hook never fires). // - // Dispatch order (precedence: newer mt:: rewrite over the older - // per-slot tiered_cache implementation): - // 1. If the active memory is mt::llama_memory_tiered, call its - // public backup_proactive — uses the real KV mover + KvtcStore - // pipeline. - // 2. Else fall back to the older tiered_cache->evict_from_slot - // (transformer paths that haven't migrated to mt::). - // - // **Threshold capacity**: for the mt:: path use the physical - // attention cache cell count, NOT slot.n_ctx. For hybrid models - // (Qwen3.6, DeepSeek V4) the attention KV cache is sized for a - // sliding window — much smaller than the user-facing context - // (recurrent layers carry the long context). Comparing against - // slot.n_ctx makes the trigger fire too late and the inner cache - // 500s with "failed to find free space in the KV cache" before - // we ever get a chance to evict. The old tiered_cache path keeps - // slot.n_ctx since its semantics haven't changed. + // **Threshold capacity**: use the physical attention cache cell + // count, NOT slot.n_ctx. For hybrid models (Qwen3.6, DeepSeek V4) + // the attention KV cache is sized for a sliding window — much + // smaller than the user-facing context (recurrent layers carry + // the long context). Comparing against slot.n_ctx makes the + // trigger fire too late and the inner cache 500s with "failed + // to find free space in the KV cache" before we ever get a + // chance to evict. if (params_base.kv_tiered_enabled) { const int n_tokens = slot.prompt.n_tokens(); @@ -1530,7 +1651,6 @@ struct server_context_impl { if (n_live_hot >= evict_threshold) { const int n_evict = std::max(1, (int)(cap * 0.20f)); - bool routed_to_mt = false; if (mt_tier) { // Eviction window starts at the cursor — the next // not-yet-backed-up positions. Cursor advances by the @@ -1557,23 +1677,23 @@ struct server_context_impl { // attention runs. Mirrors the context-shift-time // fingerprinting; only fires when --kv-tier-semantic-index // is set and the chunk isn't multimodal. + // + // MAD-122: under --kv-tier-paged-blocks the helper + // emits one fingerprint per logical block instead of + // one for the whole chunk, so query-time prefetch can + // score at block granularity. if (!params_base.kv_semantic_index.empty() && !slot.prompt.tokens.has_mtmd) { const auto & toks = slot.prompt.tokens.get_text_tokens(); - const int hi = std::min(p1, (int) toks.size()); - if (hi > p0) { - llama_tokens chunk(toks.begin() + p0, toks.begin() + hi); - const std::string text = common_detokenize(ctx, chunk, /*special=*/ false); - const auto emb = mt_tier->embed_text(text); - if (!emb.empty()) { - std::vector positions; - positions.reserve(hi - p0); - for (int i = p0; i < hi; ++i) positions.push_back((llama_pos) i); - mt_tier->record_chunk_fingerprint( - std::move(positions), emb, - mt::SemanticIndex::Tier::Warm); - SLT_INF(slot, "tier semantic: fingerprinted %d tokens [%d,%d) (%zu-dim) for proactive backup\n", - hi - p0, p0, hi, emb.size()); - } + llama_kv_cache_paged * paged_cache = params_base.kv_tier_paged_blocks + ? mt_get_paged_cache(llama_get_memory(ctx)) : nullptr; + const int n_fp = mt_record_fingerprints_for_range( + mt_tier, paged_cache, ctx, slot.id, toks, p0, p1, + (uint32_t) params_base.kv_tier_paged_block_size); + if (n_fp > 0) { + SLT_INF(slot, "tier semantic: %d %s fingerprint(s) [%d,%d) for proactive backup\n", + n_fp, + paged_cache ? "paged-block" : "chunk", + p0, p1); } } // Advance to the requested range end, not the count @@ -1585,17 +1705,6 @@ struct server_context_impl { SLT_DBG(slot, "proactive mt:: backup: %u/%d positions [%d,%d) at %d live / %u hot capacity\\n", backed_up, n_evict, p0, p1, n_live_hot, cap); } - routed_to_mt = true; - } - - if (!routed_to_mt && tiered_cache) { - auto* slot_tier = tiered_cache->get_slot_manager(slot.id); - if (slot_tier && slot_tier->initialized) { - if (tiered_cache->evict_from_slot(slot.id, n_evict, (uint32_t)n_tokens)) { - SLT_DBG(slot, "proactive tiered eviction: %d tokens at %d/%d hot capacity\\n", - n_evict, n_tokens, slot.n_ctx); - } - } } } } @@ -1606,8 +1715,20 @@ struct server_context_impl { slot.stop = STOP_TYPE_LIMIT; slot.has_next_token = false; - SLT_DBG(slot, "stopped due to running out of context capacity, prompt.n_tokens() = %d, task.n_tokens = %d, n_decoded = %d, n_ctx = %d\n", - slot.prompt.n_tokens(), slot.task->n_tokens(), slot.n_decoded, slot.n_ctx); + // MAD-128: clear log when paged hits the limit — the operator + // needs to know this is a "resubmit and rely on prompt cache" + // situation, not a hard failure. SLT_INF (not DBG) so it lands + // in the default log level. + if (params_base.kv_tier_paged_blocks) { + SLT_INF(slot, "paged: hit n_ctx limit (n_tokens=%d, n_ctx=%d) — " + "stopping slot. Client should re-submit as a fresh " + "request; the prompt cache + semantic prefetch will " + "recover the prefix.\n", + slot.prompt.n_tokens(), slot.n_ctx); + } else { + SLT_DBG(slot, "stopped due to running out of context capacity, prompt.n_tokens() = %d, task.n_tokens = %d, n_decoded = %d, n_ctx = %d\n", + slot.prompt.n_tokens(), slot.task->n_tokens(), slot.n_decoded, slot.n_ctx); + } } // check the limits @@ -2393,20 +2514,6 @@ struct server_context_impl { const int n_left = slot.prompt.n_tokens() - n_keep; const int n_discard = slot.task->params.n_discard ? slot.task->params.n_discard : (n_left / 2); - // Try tiered cache eviction before context shift - if (params_base.kv_tiered_enabled) { - auto* slot_tier = tiered_cache->get_slot_manager(slot.id); - if (slot_tier && slot_tier->initialized) { - // Evict tokens to tiered cache before context shift - int n_evict = std::min(n_discard, int(slot_tier->tiered_cache->get_config().cold_capacity())); - if (n_evict > 0) { - if (tiered_cache->evict_from_slot(slot.id, n_evict, (uint32_t)slot.prompt.n_tokens())) { - SLT_INF(slot, "tiered cache eviction: %d tokens\n", n_evict); - } - } - } - } - SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); // mt:: tier semantic fingerprint: capture an embedding of the @@ -2419,20 +2526,17 @@ struct server_context_impl { if (auto * mt_tier = dynamic_cast(llama_get_memory(ctx))) { if (!params_base.kv_semantic_index.empty() && !slot.prompt.tokens.has_mtmd) { const auto & toks = slot.prompt.tokens.get_text_tokens(); - const int hi = std::min(n_keep + n_discard, (int) toks.size()); - if (n_keep >= 0 && hi > n_keep) { - llama_tokens chunk(toks.begin() + n_keep, toks.begin() + hi); - const std::string text = common_detokenize(ctx, chunk, /*special=*/ false); - const auto emb = mt_tier->embed_text(text); - if (!emb.empty()) { - std::vector positions; - positions.reserve(hi - n_keep); - for (int i = n_keep; i < hi; ++i) positions.push_back((llama_pos) i); - mt_tier->record_chunk_fingerprint( - std::move(positions), emb, - mt::SemanticIndex::Tier::Warm); - SLT_INF(slot, "tier semantic: fingerprinted %d tokens (%zu-dim) for context shift\n", - hi - n_keep, emb.size()); + if (n_keep >= 0) { + llama_kv_cache_paged * paged_cache = params_base.kv_tier_paged_blocks + ? mt_get_paged_cache(llama_get_memory(ctx)) : nullptr; + const int n_fp = mt_record_fingerprints_for_range( + mt_tier, paged_cache, ctx, slot.id, toks, n_keep, n_keep + n_discard, + (uint32_t) params_base.kv_tier_paged_block_size); + if (n_fp > 0) { + SLT_INF(slot, "tier semantic: %d %s fingerprint(s) [%d,%d) for context shift\n", + n_fp, + paged_cache ? "paged-block" : "chunk", + n_keep, n_keep + n_discard); } } } @@ -2575,15 +2679,53 @@ struct server_context_impl { for (auto sid : paged_evicted_this_iter) { if (sid == slot.id) { already = true; break; } } + int n_evicted = 0; if (!already) { - int n_evicted = llama_memory_paged_evict_seq(mem_for_admit, slot.id); + n_evicted = llama_memory_paged_evict_seq(mem_for_admit, slot.id); SLT_INF(slot, "MAD-120 preempt (prefill): hot pool full, " "evicted %d block(s) to warm; will retry next iter\n", n_evicted); paged_evicted_this_iter.push_back(slot.id); } + + // MAD-141: deadlock break. evict_seq returns 0 when + // the slot has no GPU-resident blocks to give back — + // typically because nothing has been prefilled yet + // and the prompt itself is too large for the hot + // budget. Without this guard the slot loops forever + // until the upstream n_empty_consecutive safety + // abort fires and crashes the server. Track + // consecutive no-progress preempt iterations and + // fail the request cleanly once it's clearly stuck. + constexpr int32_t kPagedPreemptDeadlockThreshold = 4; + if (n_evicted <= 0) { + ++slot.paged_preempt_no_progress_count; + if (slot.paged_preempt_no_progress_count >= kPagedPreemptDeadlockThreshold) { + SLT_ERR(slot, + "MAD-141: paged admission stuck for %d iters with no eviction " + "progress (n_new_est=%u). Prompt does not fit alongside the " + "live workload. Failing request.\n", + slot.paged_preempt_no_progress_count, n_new_est); + send_error(slot, + string_format( + "paged KV admission could not fit a %u-token request " + "alongside the active workload after %d retries. " + "Reduce the prompt or wait for slots to drain.", + n_new_est, slot.paged_preempt_no_progress_count), + ERROR_TYPE_SERVER); + slot.release(); + // Releasing a deadlocked slot IS progress — reset + // the empty-batch streak so the upstream safety + // abort doesn't fire on this same iteration just + // because we haven't built a token batch yet. + n_empty_consecutive = 0; + continue; + } + } continue; } + // Slot was admitted — any prior no-progress streak ends here. + slot.paged_preempt_no_progress_count = 0; paged_admitted.push_back(slot.id); } @@ -2692,13 +2834,25 @@ struct server_context_impl { const std::string qtext = common_detokenize(ctx, q, /*special=*/ false); const auto qemb = mt_tier->embed_text(qtext); if (!qemb.empty()) { - const uint32_t restored = mt_tier->restore_semantic( - slot.id, qemb, - params_base.kv_semantic_top_k, - params_base.kv_semantic_threshold); + // MAD-122/125: paged-blocks routes through + // llama_kv_cache_paged (the active tier + // layer for hybrid+paged). Non-paged falls + // back to the chunk-level wrapper path. + llama_kv_cache_paged * paged_cache = params_base.kv_tier_paged_blocks + ? mt_get_paged_cache(llama_get_memory(ctx)) : nullptr; + const uint32_t restored = paged_cache + ? paged_cache->restore_semantic_paged( + slot.id, qemb, + params_base.kv_semantic_top_k, + params_base.kv_semantic_threshold) + : mt_tier->restore_semantic( + slot.id, qemb, + params_base.kv_semantic_top_k, + params_base.kv_semantic_threshold); if (restored > 0) { - SLT_INF(slot, "tier semantic: restored %u positions from warm via cosine search\n", - restored); + SLT_INF(slot, "tier semantic: restored %u positions/blocks via cosine search (%s path)\n", + restored, + paged_cache ? "paged-block" : "chunk"); } } } @@ -3047,73 +3201,62 @@ struct server_context_impl { slot.init_sampler(); SLT_INF(slot, "prompt processing done, n_tokens = %d, batch.n_tokens = %d\n", slot.prompt.n_tokens(), batch.n_tokens); - - // Get and log prefetch hints if semantic index is enabled - if (params_base.kv_tiered_enabled && tiered_cache->sem_enabled()) { - auto* slot_tier = tiered_cache->get_slot_manager(slot.id); - if (slot_tier && slot_tier->initialized) { - // Get the input text from the prompt tokens - std::string input_text; - const auto & input_tokens = slot.task->tokens; - for (size_t i = 0; i < input_tokens.size(); ++i) { - auto piece = common_token_to_piece(ctx, input_tokens[i]); - input_text += piece; - } - - // Get prefetch hints - auto hints = tiered_cache->get_prefetch_hints(slot.id, input_text, tiered_cache->semantic_top_k); - - // Log the hints - for (const auto& hint : hints) { - std::string positions_str; - if (!hint.positions.empty()) { - std::ostringstream oss; - oss << hint.positions.front() << ".." << hint.positions.back(); - positions_str = oss.str(); - } else { - positions_str = "empty"; - } - const char* tier_name = hint.current_tier == TIER_WARM ? "warm" : - hint.current_tier == TIER_COLD ? "cold" : "hot"; - SLT_INF(slot, "semantic prefetch: score=%.2f positions=[%s] tier=%s\n", - hint.score, positions_str.c_str(), tier_name); - } - - // Apply prefetch migration for WARM tier hints - for (const auto& hint : hints) { - if (hint.current_tier == TIER_WARM) { - // Migrate WARM tokens to HOT tier - if (tiered_cache->migrate_in_slot(slot.id, hint.positions, TIER_WARM, TIER_HOT)) { - SLT_INF(slot, "semantic prefetch: migrated %zu warm tokens to hot\n", hint.positions.size()); - auto* slot_mgr = tiered_cache->get_slot_manager(slot.id); - if (slot_mgr) { - slot_mgr->stats.semantic_prefetch_hits += (uint32_t)hint.positions.size(); - } - } else { - SLT_WRN(slot, "semantic prefetch: failed to migrate %zu warm tokens\n", hint.positions.size()); - } - } else if (hint.current_tier == TIER_COLD) { - // Stage cold→warm→hot: first bring cold to warm, then warm to hot - if (tiered_cache->migrate_in_slot(slot.id, hint.positions, TIER_COLD, TIER_WARM)) { - if (tiered_cache->migrate_in_slot(slot.id, hint.positions, TIER_WARM, TIER_HOT)) { - SLT_INF(slot, "semantic prefetch: staged %zu cold tokens to hot\n", hint.positions.size()); - auto* slot_mgr = tiered_cache->get_slot_manager(slot.id); - if (slot_mgr) { - slot_mgr->stats.semantic_prefetch_hits += (uint32_t)hint.positions.size(); - } - } else { - SLT_WRN(slot, "semantic prefetch: cold->warm ok but warm->hot failed for %zu tokens\n", hint.positions.size()); - } - } else { - SLT_WRN(slot, "semantic prefetch: failed to stage %zu cold tokens\n", hint.positions.size()); - } + + // MAD-129: prefill-time semantic fingerprint trigger. + // Walk the seq's COMPLETE blocks (skip the partial last + // block, which fills on the next prefill) and embed any + // that don't already have a fingerprint. Skip-already- + // fingerprinted via has_paged_fingerprint keeps + // multi-turn cost bounded — only NEW blocks (the + // accumulated assistant response from the prior turn) + // get embedded on each turn's prefill. + // + // Why here vs proactive-backup: the existing chunk- + // level trigger at the proactive-backup site (line + // ~1631) doesn't fire for hybrid+paged because the + // server's eviction threshold uses full ctx (cap + // arithmetic returns 0 for hybrid+paged) so the + // trigger never crosses. And eviction in the paged + // cache is internal — the server doesn't see those + // events anyway. Per Epic A2: write at prefill, not + // at eviction. + if (!params_base.kv_semantic_index.empty() && !slot.prompt.tokens.has_mtmd) { + llama_kv_cache_paged * paged_cache = params_base.kv_tier_paged_blocks + ? mt_get_paged_cache(llama_get_memory(ctx)) : nullptr; + auto * mt_tier = dynamic_cast(llama_get_memory(ctx)); + + if (paged_cache && mt_tier) { + const uint32_t bsize = (uint32_t) std::max(1, params_base.kv_tier_paged_block_size); + const auto & toks = slot.prompt.tokens.get_text_tokens(); + const int n_toks = (int) toks.size(); + const int n_complete_blocks = n_toks / (int) bsize; + + int n_new_fp = 0; + int n_skipped_existing = 0; + for (int lb = 0; lb < n_complete_blocks; ++lb) { + if (paged_cache->has_paged_fingerprint(slot.id, (uint32_t) lb)) { + ++n_skipped_existing; + continue; } + const int p0 = lb * (int) bsize; + const int p1 = p0 + (int) bsize; + llama_tokens chunk(toks.begin() + p0, toks.begin() + p1); + const std::string text = common_detokenize(ctx, chunk, /*special=*/ false); + const auto emb = mt_tier->embed_text(text); + if (emb.empty()) continue; + paged_cache->record_paged_block_fingerprint( + slot.id, (uint32_t) lb, emb, + mt::SemanticIndex::Tier::Hot); + ++n_new_fp; } - - // Set the current query embedding for semantic eviction weighting - auto embedding = tiered_cache->embed(input_text); - if (!embedding.empty()) { - slot_tier->tiered_cache->set_current_query_embedding(embedding); + + if (n_new_fp > 0 || n_skipped_existing > 0) { + SLT_INF(slot, "tier semantic: prefill fingerprint sweep — " + "%d new, %d already-fingerprinted, %d total complete blocks " + "(of %d total tokens, partial tail block of %d slots not " + "yet embedded)\n", + n_new_fp, n_skipped_existing, n_complete_blocks, + n_toks, n_toks % (int) bsize); } } } @@ -3945,6 +4088,48 @@ void server_routes::init_routes() { }}} }; + // MAD-133: Add paged-tier metrics when --kv-tier-paged-blocks + // is on. Reads counters directly from the live cache (single- + // thread contract — main thread is the only mutator; metrics + // endpoint runs on the HTTP thread but only READS volatile + // uint64s, which is safe-ish on x86/ARM64 for monotonic counters). + if (ctx_server.ctx) { + llama_kv_cache_paged * paged_cache = mt_get_paged_cache(llama_get_memory(ctx_server.ctx)); + if (paged_cache) { + auto add_counter = [&](const char * name, const char * help, uint64_t value) { + json m; + m["name"] = name; + m["help"] = help; + m["value"] = value; + all_metrics_def["counter"].push_back(m); + }; + auto add_gauge = [&](const char * name, const char * help, uint64_t value) { + json m; + m["name"] = name; + m["help"] = help; + m["value"] = value; + all_metrics_def["gauge"].push_back(m); + }; + + add_counter("paged_evict_hot_to_warm_total", "Hot→warm evictions", paged_cache->evict_h2w_total()); + add_counter("paged_evict_warm_to_cold_total", "Warm→cold evictions", paged_cache->evict_w2c_total()); + add_counter("paged_evict_cold_to_drop_total", "Cold-block drops (no recovery)", paged_cache->evict_c2drop_total()); + add_counter("paged_restore_warm_to_hot_total", "Warm→hot restores", paged_cache->restore_w2h_total()); + add_counter("paged_restore_cold_to_hot_total", "Cold→hot restores", paged_cache->restore_c2h_total()); + add_counter("paged_seq_preempt_total", "MAD-120 whole-seq preemptions", paged_cache->seq_preempt_total()); + add_counter("paged_seq_restore_total", "MAD-120 whole-seq restores", paged_cache->seq_restore_total()); + + add_counter("paged_semantic_attempts_total", "MAD-129 semantic restore attempts", paged_cache->semantic_attempts_total()); + add_counter("paged_semantic_hits_total", "MAD-129 semantic restore attempts that restored ≥1 block", paged_cache->semantic_hits_total()); + add_counter("paged_semantic_blocks_restored_total", "MAD-129 total blocks restored via semantic", paged_cache->semantic_blocks_restored_total()); + + add_gauge("paged_blocks_capacity_gpu", "GPU pool size (blocks)", paged_cache->n_blocks_total()); + add_gauge("paged_blocks_capacity_warm", "Warm pool size (blocks)", paged_cache->n_warm_blocks()); + add_gauge("paged_blocks_capacity_cold", "Cold pool size (blocks)", paged_cache->n_cold_blocks()); + add_gauge("paged_fingerprints", "MAD-129 paged-block fingerprints currently held", paged_cache->n_paged_fingerprints()); + } + } + // Add weight pager metrics if enabled if (ctx_server.model && ctx_server.model->weight_pager) { auto * pager = ctx_server.model->weight_pager.get(); @@ -4022,7 +4207,28 @@ void server_routes::init_routes() { } } - res->ok(res_task->slots_data); + // MAD-133: enrich each slot's JSON with paged-tier residency + // (blocks_hot / blocks_warm / blocks_cold / fingerprints). + // Skipped when paged isn't on or the slot's id is missing. + json slots_out = res_task->slots_data; + if (slots_out.is_array() && ctx_server.ctx) { + llama_kv_cache_paged * paged_cache = mt_get_paged_cache(llama_get_memory(ctx_server.ctx)); + if (paged_cache) { + for (auto & slot : slots_out) { + if (!slot.contains("id")) continue; + const llama_seq_id sid = slot["id"].get(); + if (sid < 0) continue; + slot["tier"] = { + {"blocks_hot", paged_cache->n_gpu_blocks_for(sid)}, + {"blocks_warm", paged_cache->n_warm_blocks_for(sid)}, + {"blocks_cold", paged_cache->n_blocks_cold_for(sid)}, + {"fingerprints", paged_cache->n_fingerprints_for_seq(sid)}, + }; + } + } + } + + res->ok(slots_out); return res; }; diff --git a/tools/server/server-tiered-cache.cpp b/tools/server/server-tiered-cache.cpp deleted file mode 100644 index eaf5ccff0615..000000000000 --- a/tools/server/server-tiered-cache.cpp +++ /dev/null @@ -1,356 +0,0 @@ -#include "server-tiered-cache.h" - -#include "server-common.h" -#include "llama.h" -#include "../../src/llama-memory-hybrid-iswa.h" -#include "../../src/llama-memory-hybrid.h" -#include "../../src/llama-kv-cache-iswa.h" -#include "../../src/llama-kv-cache.h" - -#include -#include -#include - -server_tiered_cache::server_tiered_cache(const common_params& params) - : params(params) { - // Check if tiered cache is enabled - enabled = params.kv_tiered_enabled; - - if (enabled) { - // Extract SSD path from params - ssd_path = params.kv_tier_ssd_path; - if (ssd_path.empty()) { - // Default SSD path if not specified - ssd_path = "./tiered-cache"; - } - - // Extract eviction policy - eviction_policy = llama_eviction_policy(params.kv_tier_eviction_policy); - - // Extract compression type - compression = llama_cache_compression(params.kv_tier_compression); - - // Extract attention threshold - attention_threshold = params.kv_tier_attention_threshold; - - // Extract semantic threshold and top_k - semantic_threshold = params.kv_semantic_threshold; - semantic_top_k = params.kv_semantic_top_k; - - stats_.reset(); - } - - // Load semantic embedding model if specified - if (!params.kv_semantic_index.empty()) { - auto sem_params = llama_model_default_params(); - sem_params.n_gpu_layers = 0; // CPU only - sem_params.progress_callback = NULL; - - sem_model = llama_model_load_from_file(params.kv_semantic_index.c_str(), sem_params); - if (!sem_model) { - SRV_ERR("failed to load semantic index model: %s\n", params.kv_semantic_index.c_str()); - return; - } - - auto ctx_params = llama_context_default_params(); - ctx_params.n_ctx = 512; - ctx_params.embeddings = true; - ctx_params.n_batch = 512; - ctx_params.n_ubatch = 512; - - sem_ctx = llama_init_from_model(sem_model, ctx_params); - if (!sem_ctx) { - LOG_ERR("failed to create context for semantic index model\n"); - llama_model_free(sem_model); - sem_model = nullptr; - return; - } - - SRV_INF("semantic KV index loaded: %s (CPU)\n", params.kv_semantic_index.c_str()); - } -} - -server_tiered_cache::~server_tiered_cache() { - // Cleanup all slot managers - std::lock_guard lock(mutex); - slot_managers.clear(); - - // Cleanup semantic embedding model - if (sem_ctx) { - llama_free(sem_ctx); - sem_ctx = nullptr; - } - if (sem_model) { - llama_model_free(sem_model); - sem_model = nullptr; - } -} - -std::vector server_tiered_cache::embed(const std::string & text) { - if (!sem_ctx || !sem_model) { - return {}; - } - - // Tokenize the input text - auto * vocab = llama_model_get_vocab(sem_model); - std::vector tokens(256); - int32_t n_tokens = llama_tokenize(vocab, text.c_str(), (int32_t)text.size(), - tokens.data(), (int32_t)tokens.capacity(), - true, true); - if (n_tokens < 0) { - // Tokenization failed, try without add_special - n_tokens = llama_tokenize(vocab, text.c_str(), (int32_t)text.size(), - tokens.data(), (int32_t)tokens.capacity(), - false, false); - if (n_tokens < 0) { - return {}; - } - } - tokens.resize(n_tokens); - - if (tokens.empty()) { - return {}; - } - - // Create a batch and decode to get embeddings - // Note: llama_batch_get_one borrows the tokens pointer — do NOT call llama_batch_free on it. - auto batch = llama_batch_get_one(tokens.data(), n_tokens); - int32_t result = llama_decode(sem_ctx, batch); - if (result != 0) { - SRV_WRN("embed: llama_decode failed with code %d\n", result); - return {}; - } - - // Extract embeddings for sequence 0 - float * embd_ptr = llama_get_embeddings_seq(sem_ctx, 0); - if (!embd_ptr) { - return {}; - } - - // Get embedding dimension - int32_t n_embd = llama_model_n_embd(sem_model); - std::vector embedding(embd_ptr, embd_ptr + n_embd); - - // L2-normalize the embedding vector - float norm = 0.0f; - for (int32_t i = 0; i < n_embd; ++i) { - norm += embedding[i] * embedding[i]; - } - norm = std::sqrt(norm); - if (norm > 1e-9f) { - for (int32_t i = 0; i < n_embd; ++i) { - embedding[i] /= norm; - } - } - - return embedding; -} - -bool server_tiered_cache::init_slot(int slot_id, const llama_model& model, struct llama_context * lctx) { - if (!enabled) { - return false; - } - - std::lock_guard lock(mutex); - - // Check if slot already has a manager - auto it = slot_managers.find(slot_id); - if (it != slot_managers.end()) { - // Already initialized - return true; - } - - // Store model pointer for detokenization - detokenize_model = &model; - - // Create new slot manager - slot_tier_manager manager; - - // Create tiered cache configuration - llama_tier_config config; - config.hot_percent = params.kv_tier_hot_pct; - config.warm_percent = params.kv_tier_warm_pct; - config.cold_percent = params.kv_tier_cold_pct; - config.total_ctx = params.kv_tier_total_ctx > 0 ? params.kv_tier_total_ctx : params.n_ctx; - config.warm_device = params.kv_warm_device; - - // Create tiered cache instance - manager.tiered_cache = std::make_unique( - model, - config, - ssd_path, - eviction_policy, - compression, - attention_threshold - ); - - // Initialize the tiered cache (warm slots reserved, buffers come after layer wiring) - if (!manager.tiered_cache->init()) { - return false; - } - - // Wire per-layer K/V tensor pointers for actual data movement. - // Handle both pure-attention (llama_kv_cache) and hybrid (llama_memory_hybrid_iswa) - // models like Qwen3-27B which wrap an iswa KV cache inside a hybrid container. - if (lctx) { - auto * mem = llama_get_memory(lctx); - auto * kv = dynamic_cast(mem); - auto * hybrid = dynamic_cast(mem); - auto * hybrid_iswa = dynamic_cast(mem); - if (kv) { - manager.tiered_cache->set_kv_layers_from_cache(kv); - } else if (hybrid && hybrid->get_mem_attn()) { - manager.tiered_cache->set_kv_layers_from_cache(hybrid->get_mem_attn()); - } else if (hybrid_iswa && hybrid_iswa->get_mem_attn()) { - manager.tiered_cache->set_kv_layers_from_cache(hybrid_iswa->get_mem_attn()->get_base()); - } else { - SRV_WRN("tiered cache: could not cast memory to llama_kv_cache for slot %d — metadata-only mode\n", slot_id); - } - } - - manager.initialized = true; - manager.stats.reset(); - slot_managers.emplace(slot_id, std::move(manager)); - - return true; -} - -server_tiered_cache::slot_tier_manager* server_tiered_cache::get_slot_manager(int slot_id) { - std::lock_guard lock(mutex); - - auto it = slot_managers.find(slot_id); - if (it == slot_managers.end()) { - return nullptr; - } - - return &it->second; -} - -bool server_tiered_cache::evict_from_slot(int slot_id, uint32_t n_tokens, uint32_t n_hot_positions) { - if (!enabled) { - return false; - } - - auto* manager = get_slot_manager(slot_id); - if (!manager || !manager->initialized) { - return false; - } - - // Populate metadata so eviction scoring has candidates to work with - if (n_hot_positions > 0) { - manager->tiered_cache->track_hot_range(n_hot_positions); - } - - // Evict tokens using the tiered cache - bool result = manager->tiered_cache->evict_tokens(n_tokens, TIER_HOT); - - if (result) { - // Update global stats - std::lock_guard lock(mutex); - stats_.total_evictions += n_tokens; - - // If semantic index is enabled, compute fingerprint for evicted tokens - if (sem_enabled() && detokenize_model && n_tokens > 0) { - // Create positions vector for the evicted tokens - std::vector positions(n_tokens); - for (uint32_t i = 0; i < n_tokens; i++) { - positions[i] = i; // Simplified - should track actual positions from tiered cache - } - - // Detokenize to get text for embedding - std::string text; - auto * vocab = llama_model_get_vocab(detokenize_model); - for (auto pos : positions) { - std::vector buf(64); - int len = llama_token_to_piece(vocab, pos, buf.data(), buf.size(), 0, false); - if (len > 0) { - text.append(buf.data(), len); - } - } - - // Compute embedding - auto embedding = embed(text); - if (!embedding.empty()) { - // Add fingerprint to tiered cache - manager->tiered_cache->add_fingerprint(positions, embedding, TIER_WARM); - - // Save fingerprints to disk - if (!fingerprints_path.empty()) { - manager->tiered_cache->save_fingerprints_to_disk(fingerprints_path + "/fingerprints.bin"); - } - } - } - } - - return result; -} - -bool server_tiered_cache::migrate_in_slot(int slot_id, - const std::vector& positions, - llama_cache_tier from_tier, - llama_cache_tier to_tier) { - if (!enabled) { - return false; - } - - auto* manager = get_slot_manager(slot_id); - if (!manager || !manager->initialized) { - return false; - } - - // Migrate tokens between tiers - bool result = manager->tiered_cache->migrate_tokens(positions, from_tier, to_tier); - - if (result) { - // Update global stats - std::lock_guard lock(mutex); - stats_.total_migrations += positions.size(); - } - - return result; -} - -llama_tier_stats server_tiered_cache::get_slot_stats(int slot_id) { - if (!enabled) { - return llama_tier_stats{}; - } - - auto* manager = get_slot_manager(slot_id); - if (!manager || !manager->initialized) { - return llama_tier_stats{}; - } - - return manager->tiered_cache->get_stats(); -} - -server_tiered_cache::global_stats server_tiered_cache::get_global_stats() { - std::lock_guard lock(mutex); - return stats_; -} - -void server_tiered_cache::reset_stats() { - std::lock_guard lock(mutex); - stats_.reset(); -} - -std::vector -server_tiered_cache::get_prefetch_hints(int slot_id, const std::string& input_text, int top_k) { - if (!sem_enabled()) { - return {}; - } - - // Embed the input text - auto embedding = embed(input_text); - if (embedding.empty()) { - return {}; - } - - // Get the slot manager - auto* manager = get_slot_manager(slot_id); - if (!manager || !manager->initialized) { - return {}; - } - - // Score fingerprints and return hints - return manager->tiered_cache->score_fingerprints(embedding, top_k, semantic_threshold); -} diff --git a/tools/server/server-tiered-cache.h b/tools/server/server-tiered-cache.h deleted file mode 100644 index 448c95015e01..000000000000 --- a/tools/server/server-tiered-cache.h +++ /dev/null @@ -1,101 +0,0 @@ -#pragma once - -#include "llama-kv-cache-tiered.h" -#include "server-common.h" - -#include -#include -#include - -// Tiered cache manager for server slots -struct server_tiered_cache { - // Per-slot tier manager - struct slot_tier_manager { - std::unique_ptr tiered_cache; - llama_tier_stats stats; - bool initialized = false; - - void reset() { - tiered_cache.reset(); - stats.reset(); - initialized = false; - } - }; - - server_tiered_cache() : enabled(false) {} - server_tiered_cache(const common_params& params); - ~server_tiered_cache(); - - // Initialize tier manager for a slot - bool init_slot(int slot_id, const llama_model& model, struct llama_context * lctx = nullptr); - - // Get tier manager for a slot - slot_tier_manager* get_slot_manager(int slot_id); - - // Evict tokens from a slot (with fingerprint computation if semantic index enabled) - bool evict_from_slot(int slot_id, uint32_t n_tokens, uint32_t n_hot_positions = 0); - - // Migrate tokens between tiers for a slot - bool migrate_in_slot(int slot_id, const std::vector& positions, - llama_cache_tier from_tier, llama_cache_tier to_tier); - - // Get statistics for a slot - llama_tier_stats get_slot_stats(int slot_id); - - // Get global statistics - struct global_stats { - uint64_t total_evictions = 0; - uint64_t total_migrations = 0; - uint64_t total_cache_hits = 0; - uint64_t total_cache_misses = 0; - double total_migration_latency_us = 0.0; - - void reset() { - total_evictions = 0; - total_migrations = 0; - total_cache_hits = 0; - total_cache_misses = 0; - total_migration_latency_us = 0.0; - } - }; - - global_stats get_global_stats(); - - // Reset all statistics - void reset_stats(); - - // Check if tiered cache is enabled - bool is_enabled() const { return enabled; } - - // Semantic embedding - bool sem_enabled() const { return sem_model != nullptr; } - std::vector embed(const std::string & text); - -private: - bool enabled = false; - std::unordered_map slot_managers; - global_stats stats_; - mutable std::mutex mutex; - - common_params params; - std::string ssd_path; - std::string fingerprints_path; - llama_eviction_policy eviction_policy; - llama_cache_compression compression; - float attention_threshold; - - // Semantic embedding model for KV fingerprints - struct llama_model * sem_model = nullptr; - struct llama_context * sem_ctx = nullptr; - - // Model pointer for detokenization (set during init_slot) - const llama_model* detokenize_model = nullptr; - - // Semantic threshold (private) - float semantic_threshold = 0.65f; - -public: - int semantic_top_k = 5; - std::vector - get_prefetch_hints(int slot_id, const std::string& input_text, int top_k = 5); -};