From 902b99de1dea1bb8c641975562db214f01c7ff15 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Tue, 30 Jun 2026 15:48:01 +0200 Subject: [PATCH 1/2] feat(kvflash): pager serialize/deserialize + QK residency library Model-agnostic KV-paging primitives in server/src/common/, consumed by qwen35/qwen35moe, gemma4, and laguna backends. kvflash_pager.h: - serialize(max_chunks)/deserialize v1 format with ledger section (was_resident + qk_score + dtype), enabling pooled-prefix snapshot/restore - critical-chunk pinning + deadlock guard - floor_to_chunk helper - identity_prefix_covers query (preserved for existing consumers) kvflash_qk.h: - cosine QK residency scoring with seeded-fallback + overflow-safe seeded_n (limit defaults to 0, not n_chunks, to prevent OOB reads) Tests: test_kvflash_pager (serde + pinning + partial-serialize + ledger round-trip) and test_kvflash_qk (pure-math scoring). Both CPU-only. --- server/CMakeLists.txt | 6 + server/src/common/kvflash_pager.h | 273 +++++++++++++++++++++++-- server/src/common/kvflash_qk.h | 54 ++++- server/test/test_kvflash_pager.cpp | 318 +++++++++++++++++++++++++++++ server/test/test_kvflash_qk.cpp | 80 ++++++++ 5 files changed, 713 insertions(+), 18 deletions(-) create mode 100644 server/test/test_kvflash_pager.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 805d91c09..b5cf0699a 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -894,6 +894,12 @@ if(DFLASH27B_TESTS) target_link_libraries(test_kvflash PRIVATE CUDA::cudart) endif() endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kvflash_pager.cpp") + add_executable(test_kvflash_pager test/test_kvflash_pager.cpp) + target_include_directories(test_kvflash_pager PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(test_kvflash_pager PRIVATE dflash_common) + add_test(NAME kvflash_pager COMMAND test_kvflash_pager) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kvflash_qk.cpp") # Pure-math unit test (KVFLASH_QK_PURE_ONLY): no ggml, no GPU. add_executable(test_kvflash_qk test/test_kvflash_qk.cpp) diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index d61707d96..4d30c75c5 100644 --- a/server/src/common/kvflash_pager.h +++ b/server/src/common/kvflash_pager.h @@ -69,6 +69,12 @@ namespace dflash::common { +// Largest multiple of chunk_tokens that is <= pos. +// Used by the pooled-prefill snapshot path to align the save boundary. +inline int floor_to_chunk(int pos, int chunk_tokens) { + return (pos / chunk_tokens) * chunk_tokens; +} + struct KvFlashConfig { int chunk_tokens = 64; // logical tokens per page int pool_tokens = 0; // resident pool capacity (multiple of chunk_tokens) @@ -161,6 +167,7 @@ class KvFlashPager { // Drop all mappings and host backing (new request / cache reset). // Cumulative stats are kept; the epoch advances so cached masks refill. + // Pins are cleared: the caller (backend) re-applies them after rebuild. void reset() { #ifdef KVFLASH_HAS_ASYNC_DMA if (page_stream_) { @@ -183,6 +190,7 @@ class KvFlashPager { stats_.host_bytes = 0; cur_chunk_ = 0; epoch_++; + std::fill(pinned_.begin(), pinned_.end(), 0); has_pending_page_in_ = false; } @@ -207,6 +215,53 @@ class KvFlashPager { // Optional external relevance score; higher = keep. Falls back to LRU. std::function score_hook; + // ── Critical-chunk pinning ────────────────────────────────────────── + // Pinned chunks are never chosen as eviction victims (OR-ed on top of the + // sink/tail protections). Empty by default → byte-identical non-pin path. + // Pins are cleared by reset() and re-applied by the backend after each + // prefill/restore rebuild via apply_kvflash_pins(). + + // Pin logical token range [tok_lo, tok_hi) (inclusive of chunk boundaries). + // Maps to chunk range [tok_lo/chunk_tokens, tok_hi/chunk_tokens] inclusive. + // Best-effort deadlock guard: if (sink+tail+n_pinned+2) > n_blocks_ the + // span is refused with a one-line warning and the function returns without + // setting any pin. + void pin_range(int64_t tok_lo, int64_t tok_hi) { + if (!attached() || tok_lo > tok_hi) return; + const int c_lo = (int)(tok_lo / cfg_.chunk_tokens); + const int c_hi = (int)(tok_hi / cfg_.chunk_tokens); + if (c_lo > c_hi || c_lo < 0) return; + // Count currently-pinned chunks + the new ones. + int currently_pinned = 0; + for (int c = 0; c < (int)pinned_.size(); c++) { + if (pinned_[c]) currently_pinned++; + } + int new_pins = 0; + for (int c = c_lo; c <= c_hi; c++) { + if (c >= (int)pinned_.size() || !pinned_[c]) new_pins++; + } + // Deadlock guard: fixed protections + pinned + 2 (one evictable victim + + // one append-head) must fit in the pool. + if (cfg_.sink_chunks + cfg_.tail_window_chunks + currently_pinned + new_pins + 2 > n_blocks_) { + std::fprintf(stderr, + "[kvflash] pin_range [%lld,%lld] refused: " + "sink=%d tail=%d pinned=%d new=%d pool_blocks=%d — would deadlock eviction\n", + (long long)tok_lo, (long long)tok_hi, + cfg_.sink_chunks, cfg_.tail_window_chunks, + currently_pinned, new_pins, n_blocks_); + return; + } + if (c_hi + 1 > (int)pinned_.size()) pinned_.resize((size_t)c_hi + 1, 0); + for (int c = c_lo; c <= c_hi; c++) pinned_[(size_t)c] = 1; + } + + bool is_pinned(int c) const { + return c >= 0 && c < (int)pinned_.size() && pinned_[(size_t)c]; + } + + // Clear all pins (also called by reset()). + void unpin_all() { std::fill(pinned_.begin(), pinned_.end(), 0); } + // Allocate slots for [kv_start, kv_start + n_tok) ahead of a forward // step (evicting LRU/low-score chunks as needed). False — with a // diagnostic — if the pool has no evictable block left. @@ -323,20 +378,6 @@ class KvFlashPager { } return true; } - - // True iff every chunk intersecting [0, n_tok) is resident in its identity - // block (block_of(c) == c). Expresses "the logical prefix [0, n_tok) is a - // contiguous, identity-mapped, materialized span" — the exact precondition - // the non-paged tree-verify graph relies on. Stronger than is_identity(), - // which is also true for an empty / not-yet-materialized pager. - bool identity_prefix_covers(int n_tok) const { - if (n_tok <= 0) return true; - const int nc = (n_tok + cfg_.chunk_tokens - 1) / cfg_.chunk_tokens; - if (nc > (int)chunks_.size()) return false; - for (int c = 0; c < nc; c++) - if (chunks_[c].block != c) return false; - return true; - } int block_of(int c) const { return c < (int)chunks_.size() ? chunks_[c].block : -1; } @@ -367,6 +408,19 @@ class KvFlashPager { int resident_blocks() const { return n_blocks_ - (int)free_blocks_.size(); } int n_chunks() const { return (int)chunks_.size(); } + // True iff chunks [0, ceil(n_tok/chunk_tokens)) are all resident in their + // natural positions — i.e. the first n_tok tokens of KV are fully covered by + // the pool with no paging. Used by backends to decide if a prefix is + // bit-identical to a flat cache. + bool identity_prefix_covers(int n_tok) const { + if (n_tok <= 0) return true; + const int nc = (n_tok + cfg_.chunk_tokens - 1) / cfg_.chunk_tokens; + if (nc > (int)chunks_.size()) return false; + for (int c = 0; c < nc; c++) + if (chunks_[c].block != c) return false; + return true; + } + // Bumped on every residency change (alloc / page_out / page_in). // Callers cache the slot mask and refill only when the epoch moves. uint64_t epoch() const { return epoch_; } @@ -397,7 +451,8 @@ class KvFlashPager { const ChunkState & st = chunks_[c]; if (st.block < 0 && !st.on_host) continue; // never materialized const bool prot = c < cfg_.sink_chunks || - c > cur_chunk_ - 1 - cfg_.tail_window_chunks; + c > cur_chunk_ - 1 - cfg_.tail_window_chunks || + is_pinned(c); cands.push_back({c, prot ? 3.4e38f : score_hook(c)}); } std::sort(cands.begin(), cands.end(), @@ -418,11 +473,196 @@ class KvFlashPager { return events; } + // Snapshot resident+paged-out chunks into a flat byte blob (v1 format). + // max_chunks: if >= 0, serialize only chunks [0, max_chunks); the blob's + // nc header field encodes max_chunks so deserialize restores exactly that + // many chunks with the correct cur_pos. Pass -1 (default) to serialize + // all chunks. + // + // Layout (v1 / magic "KVFLASH1"): + // 8-byte magic + // 8×uint32 header: nc, chunk_tokens, n_head_kv, k_seg_bytes, v_seg_bytes, + // chunk_bytes, kv_k_type_u32, has_ledger(=1) + // nc × chunk_bytes_ — raw KV bytes (unchanged from v0) + // nc × 8 bytes — ledger section (has_ledger==1): + // uint8 was_resident (1 if chunk was in pool at serialize time) + // uint8 _pad[3] + // float score (qk_score; -INF if not set) + std::vector serialize(int max_chunks = -1) const { + static constexpr uint64_t kMagic = 0x4b56464c41534831ULL; // "KVFLASH1" + const int nc = (max_chunks >= 0 && max_chunks < (int)chunks_.size()) + ? max_chunks + : (int)chunks_.size(); + const size_t hdr = sizeof(uint64_t) + 8 * sizeof(uint32_t); + const size_t ledger_entry = 8; // uint8 was_resident + 3 pad + float score + const size_t total = hdr + (size_t)nc * chunk_bytes_ + (size_t)nc * ledger_entry; + std::vector out; + out.resize(total, 0); + uint8_t * p = out.data(); + std::memcpy(p, &kMagic, 8); p += 8; + auto w32 = [&](uint32_t v) { std::memcpy(p, &v, 4); p += 4; }; + w32((uint32_t)nc); + w32((uint32_t)cfg_.chunk_tokens); + w32((uint32_t)n_head_kv_); + w32((uint32_t)k_seg_bytes_); + w32((uint32_t)v_seg_bytes_); + w32((uint32_t)chunk_bytes_); + // dtype guard: record the ggml_type enum for the K cache tensor + const uint32_t kv_k_type = (!attn_k_.empty() && attn_k_[0]) + ? (uint32_t)attn_k_[0]->type + : (uint32_t)GGML_TYPE_F16; + w32(kv_k_type); + w32(1u); // has_ledger = 1 always in v1 format + // KV bytes section + for (int c = 0; c < nc; ++c) { + uint8_t * dst = out.data() + hdr + (size_t)c * chunk_bytes_; + const ChunkState & st = chunks_[c]; + if (st.block >= 0) { + // Resident: gather from pool tensors in fixed segment order + // (layer-major, K then V, head-minor) — matching copy_chunk. + uint8_t * q = dst; + for (size_t l = 0; l < attn_k_.size(); ++l) { + for (int kv = 0; kv < 2; ++kv) { + ggml_tensor * t = kv == 0 ? attn_k_[l] : attn_v_[l]; + const size_t seg = kv == 0 ? k_seg_bytes_ : v_seg_bytes_; + for (int h = 0; h < n_head_kv_; ++h) { + const size_t off = (size_t)st.block * cfg_.chunk_tokens * t->nb[1] + + (size_t)h * t->nb[2]; + ggml_backend_tensor_get(t, q, off, seg); + q += seg; + } + } + } + } else if (st.on_host) { + // Host-backed: copy verbatim. +#ifdef KVFLASH_HAS_ASYNC_DMA + std::memcpy(dst, st.host_data, chunk_bytes_); +#else + std::memcpy(dst, st.host_data.data(), chunk_bytes_); +#endif + } + // else: never written — stays zero-filled from the resize above. + } + // Ledger section: was_resident + score per chunk + uint8_t * lp = out.data() + hdr + (size_t)nc * chunk_bytes_; + for (int c = 0; c < nc; ++c) { + const ChunkState & st = chunks_[c]; + const uint8_t was_res = (st.block >= 0) ? 1u : 0u; + std::memcpy(lp, &was_res, 1); lp += 4; // 1 byte + 3 pad (out already zero) + std::memcpy(lp, &st.score, 4); lp += 4; + } + return out; + } + + // Restore state from a blob produced by serialize(). Returns false on + // header mismatch (layout drift guard). Callers must rebuild slot masks. + // + // v1 blobs (magic "KVFLASH1"): ledger section is read back — was_resident + // controls whether a chunk gets a pool slot (resident) or stays host-backed + // (on_host only). Scores are restored into ChunkState.score. + // + // Ordering: pre-size chunks_ so each entry exists before slot_for() runs. + // For resident chunks: set host_data + on_host=true, then call slot_for(). + // slot_for's recall branch sees on_host==true and calls copy_chunk(to_host=false). + // For non-resident chunks: set host_data + on_host=true, skip slot_for() so + // the chunk stays host-backed only (mirrors the original page-out state). + bool deserialize(const uint8_t * data, size_t n) { + static constexpr uint64_t kMagic = 0x4b56464c41534831ULL; // "KVFLASH1" + const size_t hdr = sizeof(uint64_t) + 8 * sizeof(uint32_t); + if (n < hdr) return false; + const uint8_t * p = data; + uint64_t magic = 0; std::memcpy(&magic, p, 8); p += 8; + if (magic != kMagic) return false; + auto r32 = [&]() { uint32_t v = 0; std::memcpy(&v, p, 4); p += 4; return v; }; + const int nc = (int)r32(); + const int ct = (int)r32(); + const int nhkv = (int)r32(); + const size_t kseg = (size_t)r32(); + const size_t vseg = (size_t)r32(); + const size_t cb = (size_t)r32(); + const uint32_t k_type = r32(); // kv_k_type_u32 + const uint32_t has_led = r32(); // has_ledger + // Dimension guard (same as before) + dtype guard. + if (ct != cfg_.chunk_tokens || nhkv != n_head_kv_ || + kseg != k_seg_bytes_ || vseg != v_seg_bytes_ || cb != chunk_bytes_) { + return false; + } + // Dtype guard: refuse blobs whose KV dtype enum differs from the live cache. + const uint32_t live_k_type = (!attn_k_.empty() && attn_k_[0]) + ? (uint32_t)attn_k_[0]->type + : (uint32_t)GGML_TYPE_F16; + if (k_type != live_k_type) return false; + + const size_t ledger_entry = 8; + const size_t kv_section = (size_t)nc * chunk_bytes_; + const size_t expected = hdr + kv_section + (has_led ? (size_t)nc * ledger_entry : 0); + if (n < expected) return false; + + // Read ledger into a temp buffer before reset() clears state. + std::vector ledger_was_res(nc, 1u); // default: treat as resident + std::vector ledger_scores(nc, -std::numeric_limits::infinity()); + if (has_led) { + const uint8_t * lp = data + hdr + kv_section; + for (int c = 0; c < nc; ++c) { + std::memcpy(&ledger_was_res[c], lp, 1); lp += 4; // 1 byte + 3 pad + std::memcpy(&ledger_scores[c], lp, 4); lp += 4; + } + } + + reset(); + last_serialized_kv_k_type_ = k_type; + chunks_.resize(nc); // pre-size so slot_for doesn't resize again + for (int c = 0; c < nc; ++c) { + const uint8_t * src = data + hdr + (size_t)c * chunk_bytes_; + ChunkState & st = chunks_[c]; + // Park bytes in host_data; slot_for's recall branch copies to pool. +#ifdef KVFLASH_HAS_ASYNC_DMA + if (!st.host_data) { + if (cudaMallocHost(&st.host_data, chunk_bytes_) != cudaSuccess) return false; + stats_.host_bytes += (int64_t)chunk_bytes_; + } + std::memcpy(st.host_data, src, chunk_bytes_); +#else + st.host_data.assign(src, src + chunk_bytes_); +#endif + st.on_host = true; + st.score = ledger_scores[c]; + if (ledger_was_res[c]) { + // slot_for assigns a block and auto-recalls via copy_chunk. + slot_for((int64_t)c * cfg_.chunk_tokens); + } + // else: stays host-backed (on_host=true, block=-1), matching page_out state. + } +#ifdef KVFLASH_HAS_ASYNC_DMA + // Recalls above are async on page_stream_; settle before the pool is read. + if (has_pending_page_in_) synchronize_paging(); +#endif + return true; + } + + // ── Per-chunk score accessors (QK ledger) ───────────────────────────── + // set_chunk_score / chunk_score allow callers (the QK scorer wiring in + // qwen35_backend.cpp) to push per-chunk relevance scores into the pager's + // ledger before snapshot_save so they round-trip through serialize(). + void set_chunk_score(int c, float s) { + if (c >= 0 && c < (int)chunks_.size()) chunks_[(size_t)c].score = s; + } + float chunk_score(int c) const { + if (c >= 0 && c < (int)chunks_.size()) return chunks_[(size_t)c].score; + return -std::numeric_limits::infinity(); + } + + // The kv_k dtype enum read from the most-recently-deserialized blob. + // Useful for callers that want to confirm the restored dtype after + // deserialize() returns (e.g., Phase 2 QK pool rebuild). + uint32_t serialized_kv_k_type() const { return last_serialized_kv_k_type_; } + private: struct ChunkState { int block = -1; // pool block index, -1 = not resident bool on_host = false; // backing store holds valid bytes uint64_t last_use = 0; + float score = -std::numeric_limits::infinity(); // QK relevance ledger #ifdef KVFLASH_HAS_ASYNC_DMA void * host_data = nullptr; // cudaMallocHost-pinned; allocated on first page_out #else @@ -441,6 +681,7 @@ class KvFlashPager { if (chunks_[c].block < 0) continue; if (c < cfg_.sink_chunks) continue; if (c > cur_chunk_ - 1 - cfg_.tail_window_chunks) continue; + if (is_pinned(c)) continue; if (score_hook) { const float s = score_hook(c); if (victim < 0 || s < v_score) { victim = c; v_score = s; } @@ -566,11 +807,13 @@ class KvFlashPager { std::vector chunks_; std::vector free_blocks_; std::vector zero_buf_; // used by zero_block() in non-CUDA builds + std::vector pinned_; // per-chunk pin flag; empty = no pins KvFlashStats stats_; size_t k_seg_bytes_ = 0, v_seg_bytes_ = 0, chunk_bytes_ = 0; int n_blocks_ = 0, n_head_kv_ = 0, cur_chunk_ = 0; uint64_t clock_ = 0; uint64_t epoch_ = 0; + uint32_t last_serialized_kv_k_type_ = (uint32_t)GGML_TYPE_F16; // from most-recent deserialize #ifdef KVFLASH_HAS_ASYNC_DMA cudaStream_t page_stream_ = nullptr; diff --git a/server/src/common/kvflash_qk.h b/server/src/common/kvflash_qk.h index f84bd85b0..aa39225da 100644 --- a/server/src/common/kvflash_qk.h +++ b/server/src/common/kvflash_qk.h @@ -42,12 +42,20 @@ struct KvFlashQkDims { // unscorable chunk always ranks last in reselect — never above a real chunk // whose query correlation is negative. A neutral 0.0 would let a no-info // chunk evict a genuinely low-relevance one. +// +// seeded: optional per-chunk fallback scores (n_chunks floats) from a +// prior turn's ledger. When pooled_keys[c] == nullptr and seeded != nullptr +// and seeded[c] != seeded_sentinel, seeded[c] is used instead of +// missing_score. Chunks with pooled keys always use the cosine path. +// seeded_sentinel defaults to -INF (the pager's unset-score marker). inline void kvflash_qk_chunk_scores( const std::vector & pooled_keys, const float * query, const KvFlashQkDims & d, std::vector & out, - float missing_score = -2.0f) { + float missing_score = -2.0f, + const float * seeded = nullptr, + float seeded_sentinel = -std::numeric_limits::infinity()) { const int group = d.n_q_heads / d.n_kv_heads; const int n_chunks = (int)pooled_keys.size(); out.assign((size_t)n_chunks, missing_score); @@ -83,6 +91,15 @@ inline void kvflash_qk_chunk_scores( } out[(size_t)c] = acc * inv_layers; // layer-MEAN (Phase-0 config) } + // Seeded fallback: for chunks with no pooled key, use the ledger score from + // a prior turn if it is not the sentinel (i.e. it was actually scored). + if (seeded) { + for (int c = 0; c < n_chunks; c++) { + if (!pooled_keys[(size_t)c] && seeded[c] != seeded_sentinel) { + out[(size_t)c] = seeded[c]; + } + } + } } } // namespace dflash::common @@ -103,6 +120,7 @@ class KvFlashQkPool { void reset(const KvFlashQkDims & d) { dims_ = d; keys_.clear(); + seeded_scores_.clear(); } const KvFlashQkDims & dims() const { return dims_; } bool has(int c) const { @@ -161,9 +179,35 @@ class KvFlashQkPool { return true; } + // Seed per-chunk fallback scores from the pager's ledger (Phase 2 restore). + // Called after KvFlashPager::deserialize() so the scorer can report prior- + // turn scores for restored chunks before the first re-pool pass. + // scores[c] == -INF means "never scored on the previous turn" (pager sentinel). + void seed_scores(const std::vector & scores) { + seeded_scores_ = scores; + } + + // Expose the seeded array (nc floats) for the scorer; nullptr if not seeded. + const float * seeded_scores_ptr() const { + return seeded_scores_.empty() ? nullptr : seeded_scores_.data(); + } + + // Rebuild seeded scores from a KvFlashPager after deserialize(). + // KvFlashPager accessors used: n_chunks(), chunk_score(c). + // Templated on pager type to avoid a hard dependency on kvflash_pager.h + // in this header (and to keep the pure test linkable). + template + void rebuild_pool_from_ledger(const Pager & pager) { + const int nc = pager.n_chunks(); + std::vector scores((size_t)nc); + for (int c = 0; c < nc; c++) scores[(size_t)c] = pager.chunk_score(c); + seed_scores(scores); + } + private: KvFlashQkDims dims_; - std::vector> keys_; // [chunk][L*Hkv*D], empty = missing + std::vector> keys_; // [chunk][L*Hkv*D], empty = missing + std::vector seeded_scores_; // per-chunk ledger scores post-restore }; // KvFlashScorer adapter: scores from the QkPool + the latest captured query @@ -186,7 +230,11 @@ class KvFlashTargetQkScorer : public KvFlashScorer { const int n_chunks = ((int)ids.size() + chunk_tokens - 1) / chunk_tokens; std::vector pk((size_t)n_chunks, nullptr); for (int c = 0; c < n_chunks; c++) pk[(size_t)c] = pool_->data(c); - kvflash_qk_chunk_scores(pk, query_.data(), d, out); + // Pass seeded scores from the restore ledger as fallback for chunks + // whose pooled keys have not been rebuilt yet (Phase 2 restore path). + kvflash_qk_chunk_scores(pk, query_.data(), d, out, + /*missing_score=*/-2.0f, + pool_->seeded_scores_ptr()); return true; } diff --git a/server/test/test_kvflash_pager.cpp b/server/test/test_kvflash_pager.cpp new file mode 100644 index 000000000..26bec6c14 --- /dev/null +++ b/server/test/test_kvflash_pager.cpp @@ -0,0 +1,318 @@ +// Unit tests for KvFlashPager serialize/deserialize round-trip and critical- +// chunk pinning, on a CPU ggml backend (no CUDA). +#include "../src/common/kvflash_pager.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { std::fprintf(stderr, "FAIL: %s\n", msg); std::exit(1); } +} + +// Small synthetic cache: head_dim=4, pool=512 tok, n_head_kv=2, 1 layer, +// chunk_tokens=64 -> 8 blocks. sink=1 + tail=4 -> min_pool = (1+4+2)*64=448 <= 512. +struct Harness { + ggml_context * ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_backend_buffer_t buf = nullptr; + std::vector k, v; + + Harness(int head_dim, int pool, int n_head_kv, int n_layer) { + backend = ggml_backend_cpu_init(); + ggml_init_params ip{}; + ip.mem_size = (size_t)(n_layer * 2 + 8) * ggml_tensor_overhead(); + ip.no_alloc = true; + ctx = ggml_init(ip); + for (int l = 0; l < n_layer; l++) { + ggml_tensor * kt = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_dim, pool, n_head_kv); + ggml_tensor * vt = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_dim, pool, n_head_kv); + k.push_back(kt); v.push_back(vt); + } + buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + } + ~Harness() { + if (buf) ggml_backend_buffer_free(buf); + if (ctx) ggml_free(ctx); + if (backend) ggml_backend_free(backend); + } + KvFlashConfig cfg(int pool, int chunk) const { + KvFlashConfig c; c.pool_tokens = pool; c.chunk_tokens = chunk; + c.sink_chunks = 1; c.tail_window_chunks = 4; return c; + } +}; + +static void test_serialize_roundtrip() { + const int head_dim = 4, pool = 512, nkv = 2, nlayer = 1, chunk = 64; + Harness A(head_dim, pool, nkv, nlayer); + KvFlashPager pa; + expect(pa.attach(A.cfg(pool, chunk), A.k, A.v), "attach A"); + // Map the whole pool and stamp a recognizable ramp into the K/V buffers. + expect(pa.alloc_span(0, pool), "alloc_span fills pool A"); + const size_t kbytes = ggml_nbytes(A.k[0]); + std::vector ramp(kbytes); + for (size_t i = 0; i < kbytes; i++) ramp[i] = (uint8_t)(i * 31 + 7); + ggml_backend_tensor_set(A.k[0], ramp.data(), 0, kbytes); + ggml_backend_tensor_set(A.v[0], ramp.data(), 0, kbytes); + + std::vector blob = pa.serialize(); + expect(!blob.empty(), "serialize produces a blob"); + + // Deserialize into a fresh pager over fresh (zeroed) tensors. + Harness B(head_dim, pool, nkv, nlayer); + KvFlashPager pb; + expect(pb.attach(B.cfg(pool, chunk), B.k, B.v), "attach B"); + expect(pb.deserialize(blob.data(), blob.size()), "deserialize succeeds"); + + // The restored pool must contain the same KV bytes as the source. + std::vector kb(kbytes), vb(kbytes); + ggml_backend_tensor_get(B.k[0], kb.data(), 0, kbytes); + ggml_backend_tensor_get(B.v[0], vb.data(), 0, kbytes); + expect(kb == ramp, "K restored byte-identical after round-trip"); + expect(vb == ramp, "V restored byte-identical after round-trip"); + + // A blob from a mismatched layout must be rejected (header validation). + KvFlashConfig bad = B.cfg(pool, chunk); + blob[8] ^= 0xFF; // corrupt a header byte + KvFlashPager pc; + Harness C(head_dim, pool, nkv, nlayer); + expect(pc.attach(C.cfg(pool, chunk), C.k, C.v), "attach C"); + expect(!pc.deserialize(blob.data(), blob.size()), "corrupt-header blob rejected"); + (void)bad; + std::printf("ok: serialize/deserialize round-trip + header guard\n"); +} + +// ── New tests for the pooled-prefill snapshot fix ──────────────────────── +// +// RED: floor_to_chunk helper and serialize(max_chunks=k) partial-serialize +// do NOT exist yet -> test_floor_to_chunk and test_serialize_partial MUST fail +// to compile / link / assert before the fix lands. + +static void test_floor_to_chunk() { + // floor_to_chunk(pos, chunk) must return the largest multiple of chunk_tokens + // that is <= pos — the chunk-aligned boundary used by the pooled-prefill + // snapshot path. + expect(dflash::common::floor_to_chunk(0, 64) == 0, "floor(0,64)==0"); + expect(dflash::common::floor_to_chunk(63, 64) == 0, "floor(63,64)==0"); + expect(dflash::common::floor_to_chunk(64, 64) == 64, "floor(64,64)==64"); + expect(dflash::common::floor_to_chunk(65, 64) == 64, "floor(65,64)==64"); + expect(dflash::common::floor_to_chunk(127, 64) == 64, "floor(127,64)==64"); + expect(dflash::common::floor_to_chunk(128, 64) == 128,"floor(128,64)==128"); + // snap_pos=34077, chunk=64: floor = 34048 (531*64) + expect(dflash::common::floor_to_chunk(34077, 64) == 34048, "floor(34077,64)==34048"); + std::printf("ok: floor_to_chunk\n"); +} + +static void test_serialize_partial() { + // Pool 8 chunks (8*64=512 tokens). Fill all 8, then partial-serialize + // only the first k=5 chunks. Deserialize into a fresh pager; assert: + // - exactly 5 chunks are restored + // - cur_pos == 5*64 == 320 + // - KV bytes for chunks 0..4 are bit-identical to source + // - the pager reports 5 chunks (not 8) + const int head_dim = 4, pool = 512, nkv = 2, nlayer = 1, chunk = 64; + const int k = 5; // only serialize first k chunks + Harness A(head_dim, pool, nkv, nlayer); + KvFlashPager pa; + expect(pa.attach(A.cfg(pool, chunk), A.k, A.v), "partial: attach A"); + expect(pa.alloc_span(0, pool), "partial: alloc_span fills pool A"); + + // Stamp a distinct ramp into each chunk's region so we can tell them apart. + const size_t bytes_per_chunk = (size_t)chunk * (size_t)head_dim * 2 * nkv; // F16=2 bytes + const size_t total_bytes = ggml_nbytes(A.k[0]); + std::vector ramp(total_bytes); + for (size_t i = 0; i < total_bytes; i++) ramp[i] = (uint8_t)((i * 37 + 11) & 0xFF); + ggml_backend_tensor_set(A.k[0], ramp.data(), 0, total_bytes); + ggml_backend_tensor_set(A.v[0], ramp.data(), 0, total_bytes); + + // Partial serialize: max_chunks=k means only chunks [0, k) go into the blob. + std::vector blob = pa.serialize(k); + expect(!blob.empty(), "partial: serialize(k) produces a blob"); + + // Deserialize into a fresh pager. + Harness B(head_dim, pool, nkv, nlayer); + KvFlashPager pb; + expect(pb.attach(B.cfg(pool, chunk), B.k, B.v), "partial: attach B"); + expect(pb.deserialize(blob.data(), blob.size()), "partial: deserialize succeeds"); + + // The pager must report exactly k chunks (not the full 8). + expect(pb.n_chunks() == k, "partial: n_chunks == k after restore"); + + // Verify KV bytes for chunks 0..k-1 are restored identically. + // Physical layout: chunk c is at pool slot c (identity for a fresh pager + // after alloc_span from position 0). The K tensor bytes at + // offset [c*chunk*head_dim*2 .. (c+1)*chunk*head_dim*2) for each head. + // We just verify the first k chunks' byte regions match the source. + std::vector kb(total_bytes, 0), vb(total_bytes, 0); + ggml_backend_tensor_get(B.k[0], kb.data(), 0, total_bytes); + ggml_backend_tensor_get(B.v[0], vb.data(), 0, total_bytes); + + // With identity slot assignment (chunk c -> block c), the first k*chunk rows + // of each head should be identical. head stride = pool*head_dim*2. + const size_t row_bytes = (size_t)head_dim * 2; // F16 + const size_t head_stride = (size_t)pool * row_bytes; // nb[2] stride + bool k_ok = true, v_ok = true; + for (int h = 0; h < nkv; h++) { + for (int r = 0; r < k * chunk; r++) { + const size_t off = (size_t)h * head_stride + (size_t)r * row_bytes; + for (size_t b = 0; b < row_bytes; b++) { + if (kb[off + b] != ramp[off + b]) { k_ok = false; break; } + if (vb[off + b] != ramp[off + b]) { v_ok = false; break; } + } + } + } + expect(k_ok, "partial: K bytes [0..k*chunk) restored bit-identical"); + expect(v_ok, "partial: V bytes [0..k*chunk) restored bit-identical"); + + std::printf("ok: serialize(max_chunks=%d) partial round-trip\n", k); +} + +// ── Ledger round-trip test (Phase 1: was_resident + qk_score + dtype_enum) ── +// +// RED before implementation: serialize() drops the ledger → round-trip +// asserts on score/residency fields will fail. +static void test_ledger_roundtrip() { + const int head_dim = 4, pool = 512, nkv = 2, nlayer = 1, chunk = 64; + // 8 blocks total (512/64). We'll fill 6 chunks, page-out 2 of them, and + // set distinct per-chunk scores so we can verify the score field round-trips. + Harness A(head_dim, pool, nkv, nlayer); + KvFlashPager pa; + KvFlashConfig cfg = A.cfg(pool, chunk); + expect(pa.attach(cfg, A.k, A.v), "ledger: attach A"); + + // Allocate 6 chunks (0..5) — stays within the 8-block pool. + expect(pa.alloc_span(0, 6 * chunk), "ledger: alloc 6 chunks"); + + // Stamp recognizable bytes. + const size_t kbytes = ggml_nbytes(A.k[0]); + std::vector ramp(kbytes); + for (size_t i = 0; i < kbytes; i++) ramp[i] = (uint8_t)(i * 17 + 3); + ggml_backend_tensor_set(A.k[0], ramp.data(), 0, kbytes); + ggml_backend_tensor_set(A.v[0], ramp.data(), 0, kbytes); + + // Assign per-chunk scores via the public score field on ChunkState. + // The pager exposes set_chunk_score() / chunk_score() for ledger access. + const float scores[6] = {1.0f, 0.5f, -0.25f, 2.0f, 0.0f, -1.0f}; + for (int c = 0; c < 6; c++) pa.set_chunk_score(c, scores[c]); + + // Page out chunk 1 and chunk 4 so they are host-backed (was_resident==false). + pa.page_out(1); + pa.page_out(4); + // After page_out, chunks 0,2,3,5 resident; 1,4 on_host. + + // Serialize all 6 chunks. + std::vector blob = pa.serialize(6); + expect(!blob.empty(), "ledger: serialize produces blob"); + + // Deserialize into a fresh pager. + Harness B(head_dim, pool, nkv, nlayer); + KvFlashPager pb; + expect(pb.attach(B.cfg(pool, chunk), B.k, B.v), "ledger: attach B"); + expect(pb.deserialize(blob.data(), blob.size()), "ledger: deserialize succeeds"); + + // --- dtype_enum guard --- + // The blob must record the dtype and the deserialized pager must expose it. + // For the harness we use GGML_TYPE_F16; serialize must record it. + expect(pb.serialized_kv_k_type() == (uint32_t)GGML_TYPE_F16, + "ledger: dtype_enum round-trips (F16)"); + + // --- was_resident round-trip --- + // Chunks 0,2,3,5 were resident → must be resident after restore. + // Chunks 1,4 were host-backed → must NOT be resident (just on_host) after restore. + expect( pb.is_resident(0), "ledger: chunk 0 resident after restore"); + expect(!pb.is_resident(1), "ledger: chunk 1 NOT resident after restore (was paged out)"); + expect( pb.is_resident(2), "ledger: chunk 2 resident after restore"); + expect( pb.is_resident(3), "ledger: chunk 3 resident after restore"); + expect(!pb.is_resident(4), "ledger: chunk 4 NOT resident after restore (was paged out)"); + expect( pb.is_resident(5), "ledger: chunk 5 resident after restore"); + + // --- qk_score round-trip --- + for (int c = 0; c < 6; c++) { + float got = pb.chunk_score(c); + expect(got == scores[c], "ledger: qk_score round-trips for chunk"); + } + + // --- KV bytes still round-trip for resident chunks --- + // Chunks 0,2,3,5 were resident when serialized (gathered from pool tensors). + // After restore they are resident again. Check a sample region. + std::vector kb(kbytes, 0); + ggml_backend_tensor_get(B.k[0], kb.data(), 0, kbytes); + // Chunk 0 occupies block 0 in a fresh pager (identity slot assignment). + // Verify first chunk's row bytes match the ramp. + const size_t row_bytes = (size_t)head_dim * 2; // F16 + const size_t head_stride = (size_t)pool * row_bytes; + bool bytes_ok = true; + for (int h = 0; h < nkv && bytes_ok; h++) { + for (int r = 0; r < chunk && bytes_ok; r++) { + const size_t off = (size_t)h * head_stride + (size_t)r * row_bytes; + for (size_t b = 0; b < row_bytes && bytes_ok; b++) { + if (kb[off + b] != ramp[off + b]) bytes_ok = false; + } + } + } + expect(bytes_ok, "ledger: chunk 0 KV bytes round-trip with new format"); + + // --- old-magic blob must be rejected --- + // Corrupt the first byte of the blob to simulate the old magic. + std::vector old_blob = blob; + old_blob[7] ^= 0x01; // flip the last byte of the magic field + Harness C(head_dim, pool, nkv, nlayer); + KvFlashPager pc; + expect(pc.attach(C.cfg(pool, chunk), C.k, C.v), "ledger: attach C"); + expect(!pc.deserialize(old_blob.data(), old_blob.size()), + "ledger: old/corrupted magic rejected"); + + std::printf("ok: ledger round-trip (was_resident + qk_score + dtype_enum)\n"); +} + +static void test_pinning() { + const int head_dim = 4, pool = 512, nkv = 2, nlayer = 1, chunk = 64; + Harness H(head_dim, pool, nkv, nlayer); + KvFlashPager p; + expect(p.attach(H.cfg(pool, chunk), H.k, H.v), "attach pin"); + + // No pins -> nothing pinned. + expect(!p.is_pinned(2), "pre: chunk 2 unpinned"); + + // Pin a single mid chunk (tokens 128..191 -> chunk 2). 8 blocks, sink1+tail4 + // +pinned1+2 = 8 <= 8 -> allowed. + p.pin_range(128, 191); + expect(p.is_pinned(2), "chunk 2 pinned"); + expect(!p.is_pinned(1), "chunk 1 not pinned"); + expect(!p.is_pinned(3), "chunk 3 not pinned"); + + // unpin clears it. + p.unpin_all(); + expect(!p.is_pinned(2), "unpin_all clears chunk 2"); + + // Deadlock guard: pinning the whole pool (8 chunks) would leave no evictable + // block (1+4+8+2 = 15 > 8) -> refused, nothing pinned. + p.pin_range(0, pool - 1); + expect(!p.is_pinned(0) && !p.is_pinned(4) && !p.is_pinned(7), + "over-pin refused by deadlock guard"); + + // reset() also clears pins. + p.pin_range(128, 191); + expect(p.is_pinned(2), "re-pin chunk 2"); + p.reset(); + expect(!p.is_pinned(2), "reset clears pins"); + std::printf("ok: pinning + deadlock guard + reset\n"); +} + +int main() { + test_serialize_roundtrip(); + test_pinning(); + test_floor_to_chunk(); + test_serialize_partial(); + test_ledger_roundtrip(); + std::printf("PASS: kvflash pager (serde + pinning + partial-serialize + ledger)\n"); + return 0; +} diff --git a/server/test/test_kvflash_qk.cpp b/server/test/test_kvflash_qk.cpp index e91e63bdd..f60a4f8ef 100644 --- a/server/test/test_kvflash_qk.cpp +++ b/server/test/test_kvflash_qk.cpp @@ -6,6 +6,11 @@ // aligned in both // 4. cosine: query/key magnitudes do not change scores // 5. missing pooled keys keep the missing_score sentinel +// +// Phase 2 (snapshot×ledger): seeded fallback scores +// 6. after restore, chunks with no pooled keys but a ledger score use that +// score (not missing_score=-2.0); chunks WITH pooled keys still use the +// cosine path; sentinel-valued seeded entries keep missing_score. #define KVFLASH_QK_PURE_ONLY #include "kvflash_qk.h" @@ -121,6 +126,81 @@ int main() { "fractional cosine propagates exactly"); } + // ── Phase 2: seeded fallback scores (restore×ledger) ────────────────── + // After a snapshot restore the QK pool has no entries: pk[c] == nullptr + // for every chunk. The seeded[] fallback array carries the per-chunk scores + // that were live at serialize time. For chunks without pooled keys, the + // scorer must use seeded[c] (when it is not the sentinel) instead of + // missing_score=-2.0. Chunks WITH pooled keys still go through the cosine + // path and seeded[] is ignored for them. + { + // 4 chunks, 2 have pooled keys (c0, c2), 2 are restore-only (c1, c3). + // Seeded scores: c1=0.7, c3=-0.3, c0/c2 sentinel (cosine path wins). + const float kSentinel = -std::numeric_limits::infinity(); + std::vector seeded(4, kSentinel); + seeded[1] = 0.7f; + seeded[3] = -0.3f; + + // c0: fully aligned key (cosine=1.0 with the query below) + // c2: half-aligned (cosine ~0.5) + KvFlashQkDims d2; + d2.n_layers = 2; d2.n_q_heads = 2; d2.n_kv_heads = 2; d2.head_dim = 4; + const size_t kstride2 = (size_t)d2.n_layers * d2.n_kv_heads * d2.head_dim; + const size_t qsize2 = (size_t)d2.n_layers * d2.n_q_heads * d2.head_dim; + + // Query: all heads / layers point at basis e0 + std::vector q2(qsize2, 0.0f); + for (size_t i = 0; i < (size_t)d2.n_layers * d2.n_q_heads; i++) { + q2[i * d2.head_dim + 0] = 1.0f; + } + + // c0: L2-normalized key e0 in both layers/heads -> cosine=1.0 + std::vector k0(kstride2, 0.0f); + for (int l = 0; l < d2.n_layers; l++) + for (int h = 0; h < d2.n_kv_heads; h++) + k0[((size_t)l * d2.n_kv_heads + h) * d2.head_dim + 0] = 1.0f; + + // c2: L2-normalized key e0 layer0 only, e2 layer1 -> cosine=0.5 + std::vector k2(kstride2, 0.0f); + k2[((size_t)0 * d2.n_kv_heads + 0) * d2.head_dim + 0] = 1.0f; + k2[((size_t)0 * d2.n_kv_heads + 1) * d2.head_dim + 0] = 1.0f; + k2[((size_t)1 * d2.n_kv_heads + 0) * d2.head_dim + 2] = 1.0f; + k2[((size_t)1 * d2.n_kv_heads + 1) * d2.head_dim + 2] = 1.0f; + + std::vector pk2(4, nullptr); + pk2[0] = k0.data(); // c0: pooled key present + // c1: no pooled key → seeded fallback 0.7 + pk2[2] = k2.data(); // c2: pooled key present + // c3: no pooled key → seeded fallback -0.3 + + std::vector out2; + kvflash_qk_chunk_scores(pk2, q2.data(), d2, out2, + /*missing_score=*/-2.0f, + seeded.data(), kSentinel); + + CHECK(out2.size() == 4, "phase2: output size == 4"); + CHECK(near(out2[0], 1.0f), "phase2: c0 pooled key -> cosine 1.0 (seeded ignored)"); + CHECK(near(out2[1], 0.7f), "phase2: c1 no pooled key -> seeded 0.7 (not -2.0)"); + CHECK(near(out2[2], 0.5f), "phase2: c2 pooled key -> cosine 0.5 (seeded ignored)"); + CHECK(near(out2[3], -0.3f), "phase2: c3 no pooled key -> seeded -0.3 (not -2.0)"); + + // A seeded sentinel must fall through to missing_score. + std::vector all_sentinel(4, kSentinel); + std::vector out3; + kvflash_qk_chunk_scores(pk2, q2.data(), d2, out3, + /*missing_score=*/-2.0f, + all_sentinel.data(), kSentinel); + CHECK(near(out3[1], -2.0f), "phase2: seeded=sentinel -> missing_score"); + CHECK(near(out3[3], -2.0f), "phase2: seeded=sentinel -> missing_score (c3)"); + + // Null seeded ptr must behave like the original (missing_score). + std::vector out4; + kvflash_qk_chunk_scores(pk2, q2.data(), d2, out4, + /*missing_score=*/-2.0f, + /*seeded=*/nullptr, kSentinel); + CHECK(near(out4[1], -2.0f), "phase2: null seeded ptr -> missing_score"); + } + std::printf("%s (%d failures)\n", g_fail == 0 ? "ALL PASS" : "FAILED", g_fail); return g_fail == 0 ? 0 : 1; } From f4382300d2d0e7297bede719cb445915a3a7ad94 Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Tue, 30 Jun 2026 16:40:15 +0200 Subject: [PATCH 2/2] fix(kvflash_qk): bound seeded fallback by seeded_n (P1 OOB) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit kvflash_qk_chunk_scores read seeded[c] for all c < n_chunks, but the seeded buffer comes from a restored ledger whose chunk count can be smaller than the current n_chunks (pool grew after restore) → out-of-bounds read. Fix: add an explicit seeded_n length parameter; chunks at/above seeded_n fall back to missing_score. Caller (KvFlashTargetQkScorer) now passes seeded_scores_n(). Test test_kvflash_qk gains a phase2b case that constructs a 3-entry seeded buffer over 6 chunks and asserts no OOB. Addresses cubic P1 on kvflash_qk.h seeded fallback. --- server/src/common/kvflash_qk.h | 18 ++++++++++++++---- server/test/test_kvflash_qk.cpp | 33 ++++++++++++++++++++++++++++++--- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/server/src/common/kvflash_qk.h b/server/src/common/kvflash_qk.h index aa39225da..7478ca1c5 100644 --- a/server/src/common/kvflash_qk.h +++ b/server/src/common/kvflash_qk.h @@ -55,7 +55,8 @@ inline void kvflash_qk_chunk_scores( std::vector & out, float missing_score = -2.0f, const float * seeded = nullptr, - float seeded_sentinel = -std::numeric_limits::infinity()) { + float seeded_sentinel = -std::numeric_limits::infinity(), + int seeded_n = 0) { const int group = d.n_q_heads / d.n_kv_heads; const int n_chunks = (int)pooled_keys.size(); out.assign((size_t)n_chunks, missing_score); @@ -93,9 +94,13 @@ inline void kvflash_qk_chunk_scores( } // Seeded fallback: for chunks with no pooled key, use the ledger score from // a prior turn if it is not the sentinel (i.e. it was actually scored). - if (seeded) { + // seeded_n bounds the valid range of the seeded array; a restored ledger + // can be shorter than the current n_chunks (pool grew after restore), so + // chunks at/above seeded_n fall back to missing_score instead of reading + // past the buffer. + if (seeded && seeded_n > 0) { for (int c = 0; c < n_chunks; c++) { - if (!pooled_keys[(size_t)c] && seeded[c] != seeded_sentinel) { + if (c < seeded_n && !pooled_keys[(size_t)c] && seeded[c] != seeded_sentinel) { out[(size_t)c] = seeded[c]; } } @@ -191,6 +196,9 @@ class KvFlashQkPool { const float * seeded_scores_ptr() const { return seeded_scores_.empty() ? nullptr : seeded_scores_.data(); } + // Number of valid entries in the seeded array (ledger chunk count at restore + // time). May be less than the current pool n_chunks after restore+growth. + int seeded_scores_n() const { return (int)seeded_scores_.size(); } // Rebuild seeded scores from a KvFlashPager after deserialize(). // KvFlashPager accessors used: n_chunks(), chunk_score(c). @@ -234,7 +242,9 @@ class KvFlashTargetQkScorer : public KvFlashScorer { // whose pooled keys have not been rebuilt yet (Phase 2 restore path). kvflash_qk_chunk_scores(pk, query_.data(), d, out, /*missing_score=*/-2.0f, - pool_->seeded_scores_ptr()); + pool_->seeded_scores_ptr(), + /*seeded_sentinel=*/-std::numeric_limits::infinity(), + /*seeded_n=*/pool_->seeded_scores_n()); return true; } diff --git a/server/test/test_kvflash_qk.cpp b/server/test/test_kvflash_qk.cpp index f60a4f8ef..67bbcd798 100644 --- a/server/test/test_kvflash_qk.cpp +++ b/server/test/test_kvflash_qk.cpp @@ -176,7 +176,8 @@ int main() { std::vector out2; kvflash_qk_chunk_scores(pk2, q2.data(), d2, out2, /*missing_score=*/-2.0f, - seeded.data(), kSentinel); + seeded.data(), kSentinel, + /*seeded_n=*/(int)seeded.size()); CHECK(out2.size() == 4, "phase2: output size == 4"); CHECK(near(out2[0], 1.0f), "phase2: c0 pooled key -> cosine 1.0 (seeded ignored)"); @@ -189,7 +190,8 @@ int main() { std::vector out3; kvflash_qk_chunk_scores(pk2, q2.data(), d2, out3, /*missing_score=*/-2.0f, - all_sentinel.data(), kSentinel); + all_sentinel.data(), kSentinel, + /*seeded_n=*/(int)all_sentinel.size()); CHECK(near(out3[1], -2.0f), "phase2: seeded=sentinel -> missing_score"); CHECK(near(out3[3], -2.0f), "phase2: seeded=sentinel -> missing_score (c3)"); @@ -197,8 +199,33 @@ int main() { std::vector out4; kvflash_qk_chunk_scores(pk2, q2.data(), d2, out4, /*missing_score=*/-2.0f, - /*seeded=*/nullptr, kSentinel); + /*seeded=*/nullptr, kSentinel, + /*seeded_n=*/0); CHECK(near(out4[1], -2.0f), "phase2: null seeded ptr -> missing_score"); + + // ── Phase 2b (P1 fix): seeded_n bounds the seeded read ────────────── + // A restored ledger can be SHORTER than the current n_chunks (pool grew + // after restore). seeded_n must clamp the read so chunks at/above + // seeded_n fall back to missing_score instead of reading past the buffer. + // RED (pre-fix): no seeded_n → reads seeded[3..5] out of bounds → UB. + // GREEN (post-fix): out[3..5] == missing_score. + { + // 6 chunks, none pooled; ledger only covered the first 3. + std::vector seeded_short = { 0.5f, 0.6f, 0.7f }; // length 3 + std::vector pk6(6, nullptr); + std::vector out6; + kvflash_qk_chunk_scores(pk6, q2.data(), d2, out6, + /*missing_score=*/-2.0f, + seeded_short.data(), kSentinel, + /*seeded_n=*/(int)seeded_short.size()); + CHECK(out6.size() == 6, "phase2b: output size == 6"); + CHECK(near(out6[0], 0.5f), "phase2b: c0 seeded 0.5"); + CHECK(near(out6[1], 0.6f), "phase2b: c1 seeded 0.6"); + CHECK(near(out6[2], 0.7f), "phase2b: c2 seeded 0.7"); + CHECK(near(out6[3], -2.0f), "phase2b: c3 beyond seeded_n -> missing_score (no OOB)"); + CHECK(near(out6[4], -2.0f), "phase2b: c4 beyond seeded_n -> missing_score (no OOB)"); + CHECK(near(out6[5], -2.0f), "phase2b: c5 beyond seeded_n -> missing_score (no OOB)"); + } } std::printf("%s (%d failures)\n", g_fail == 0 ? "ALL PASS" : "FAILED", g_fail);