From cdb0cd7594f9037b7ecbb90176e2ee413a166bbf Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:26:26 +0200 Subject: [PATCH] feat(laguna): pager-blob snapshot/restore-consume for warm-restore after eviction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LagunaCacheSnapshot gains is_pooled + kvflash_blob fields mirroring PrefixSnapshot in qwen35. snapshot_save() serializes the pager blob when the pool has relocated chunks instead of refusing. restore_and_generate_impl() deserializes the blob and prefills only the suffix [snap_pos, prompt_len) — the restore-consume path that survives eviction. Non-pooled (identity) saves use the existing GPU-tensor copy unchanged. kvflash_pager.h: forward-port serialize/deserialize/ledger/pinning from the perf/qwen35-decode-fuse-elementwise branch so laguna can call them. The header is arch-agnostic; the impl is shared with qwen35 once merged to main. Validated: -fsyntax-only clean (laguna_backend.cpp, single-TU). Full CUDA build + NIAH recall gate deferred (resource contention). --- server/src/common/kvflash_pager.h | 260 +++++++++++++++++++++++++-- server/src/laguna/laguna_backend.cpp | 251 ++++++++++++++++++-------- server/src/laguna/laguna_internal.h | 8 + 3 files changed, 426 insertions(+), 93 deletions(-) diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index d61707d96..b298bc19a 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; } @@ -397,7 +438,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 +460,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 +668,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 +794,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/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index f4de8fd8e..73499cabf 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -318,16 +318,37 @@ bool LagunaBackend::ensure_slot(int slot) { } bool LagunaBackend::snapshot_save(int slot) { - // kvflash: snapshots copy rows assuming identity layout, which breaks - // after the first page-out relocates a chunk. - if (kvflash_active() && !kvflash_pager_.is_identity()) { - std::fprintf(stderr, "[kvflash] snapshot skipped: pool has relocated " - "chunks (page-table serialization not implemented)\n"); - return false; - } if (!ensure_slot(slot)) return false; + LagunaCacheSnapshot & snap = snapshots_[slot]; + // kvflash: when the pool has relocated chunks (is_identity() == false) the + // GPU-tensor copy path silently captures wrong rows. Serialize the pager + // blob instead — it carries every chunk's raw bytes regardless of pool slot. + const bool use_blob = kvflash_active() && + (cache_.cur_pos > kvflash_tokens_ || !kvflash_pager_.is_identity()); + if (use_blob) { + std::vector blob = kvflash_pager_.serialize(); + // Save non-KV state (recurrent state if any, cur_pos, last_tok) via the + // existing helper with skip_kv semantics: pass an empty dummy attn list + // by temporarily zeroing the cache attn vectors then restoring. + // Simpler: just save the scalar fields ourselves and call the helper with + // skip_kv so it doesn't touch attn tensors. + // laguna_snapshot_save does not have a skip_kv flag, so we fall through + // to save non-KV scalars (cur_pos, last_tok) directly on the snapshot. + snap.kvflash_blob = std::move(blob); + snap.is_pooled = true; + snap.cur_pos = cache_.cur_pos; + snap.last_tok = cache_.last_tok; + snap.used = true; + std::fprintf(stderr, "[kvflash] snapshot_save slot=%d pooled blob=%zu B " + "cur_pos=%d\n", slot, + snap.kvflash_blob.size(), snap.cur_pos); + return true; + } + // Identity layout: GPU-tensor copy is safe. + snap.is_pooled = false; + snap.kvflash_blob.clear(); if (!laguna_snapshot_save(cache_, snap_backend_, w_.n_layer, - w_.n_head_kv, w_.head_dim, snapshots_[slot])) { + w_.n_head_kv, w_.head_dim, snap)) { std::fprintf(stderr, "[snap] save slot=%d: %s\n", slot, dflash27b_last_error()); return false; @@ -945,23 +966,35 @@ GenerateResult LagunaBackend::generate_impl(const GenerateRequest & req, result.prefill_s = std::chrono::duration(t_pf1 - t_pf0).count(); // ── Inline snapshot (if requested) ── - // kvflash: snapshots copy rows [0, snap_pos) assuming identity layout, - // which holds until the first page-out relocates a chunk. - if (kvflash_active() && req.snap_slot >= 0 && - !kvflash_pager_.is_identity()) { - std::fprintf(stderr, "[kvflash] snapshot skipped: pool has relocated " - "chunks (page-table serialization not implemented)\n"); - } else if (req.snap_slot >= 0 && req.snap_pos > 0 && req.snap_pos <= N) { - if (ensure_slot(req.snap_slot) && - laguna_snapshot_save(cache_, snap_backend_, w_.n_layer, - w_.n_head_kv, w_.head_dim, snapshots_[req.snap_slot])) { - snapshots_[req.snap_slot].cur_pos = req.snap_pos; - // Record the last committed token so an exact-hit restore can - // re-embed it even when handed a zero-filled restore-only prompt. - snapshots_[req.snap_slot].last_tok = req.prompt[req.snap_pos - 1]; - std::printf("[snap] inline slot=%d cur_pos=%d\n", - req.snap_slot, req.snap_pos); + if (ensure_slot(req.snap_slot)) { + const int snap_tok = req.prompt[req.snap_pos - 1]; + LagunaCacheSnapshot & snap_dst = snapshots_[req.snap_slot]; + const bool pool_relocated = kvflash_active() && + (cache_.cur_pos > kvflash_tokens_ || !kvflash_pager_.is_identity()); + if (pool_relocated) { + // Serialize pager blob up to the snap boundary (chunk-aligned). + const int cfg_chunk = kvflash_pager_.chunk_tokens(); + const int max_chunks = req.snap_pos / cfg_chunk; + std::vector blob = kvflash_pager_.serialize(max_chunks > 0 ? max_chunks : -1); + snap_dst.kvflash_blob = std::move(blob); + snap_dst.is_pooled = true; + snap_dst.cur_pos = req.snap_pos; + snap_dst.last_tok = snap_tok; + snap_dst.used = true; + std::printf("[snap] inline pooled slot=%d cur_pos=%d blob=%zu B\n", + req.snap_slot, req.snap_pos, snap_dst.kvflash_blob.size()); + } else { + snap_dst.is_pooled = false; + snap_dst.kvflash_blob.clear(); + if (laguna_snapshot_save(cache_, snap_backend_, w_.n_layer, + w_.n_head_kv, w_.head_dim, snap_dst)) { + snap_dst.cur_pos = req.snap_pos; + snap_dst.last_tok = snap_tok; + std::printf("[snap] inline slot=%d cur_pos=%d\n", + req.snap_slot, req.snap_pos); + } + } std::fflush(stdout); } } @@ -1111,15 +1144,38 @@ GenerateResult LagunaBackend::restore_and_generate_impl(int slot, sampler_rng_.seed(sampler_.seed); } - if (!laguna_snapshot_restore(snapshots_[slot], cache_)) { - std::fprintf(stderr, "[snap] RESTORE slot=%d: %s\n", - slot, dflash27b_last_error()); - result.error = "restore"; - return result; + const LagunaCacheSnapshot & snap_ref = snapshots_[slot]; + const bool snap_pooled = snap_ref.is_pooled; + const int N = (int)req.prompt.size(); + + if (snap_pooled) { + // Pager-blob restore: the GPU-tensor snapshot was NOT taken (attn_k/v + // are stale/nullptr for pooled saves). Restore only the scalar fields + // (cur_pos, last_tok) from the snapshot and rebuild the pager from the + // serialized blob so KV rows are accessible without re-prefill. + cache_.cur_pos = snap_ref.cur_pos; + cache_.last_tok = snap_ref.last_tok; + if (!kvflash_active()) { + // Pooled save requires kvflash; if it disappeared between save and + // restore, fall through to an error rather than silently misbehave. + result.error = "kvflash: pooled snapshot requires kvflash active"; + return result; + } + if (!kvflash_pager_.deserialize(snap_ref.kvflash_blob.data(), + snap_ref.kvflash_blob.size())) { + result.error = "kvflash: pager deserialize failed"; + return result; + } + } else { + if (!laguna_snapshot_restore(snap_ref, cache_)) { + std::fprintf(stderr, "[snap] RESTORE slot=%d: %s\n", + slot, dflash27b_last_error()); + result.error = "restore"; + return result; + } } const int prefix_len = cache_.cur_pos; - const int N = (int)req.prompt.size(); if (N < prefix_len) { std::fprintf(stderr, "[snap] RESTORE prompt shorter than cached prefix (%d < %d)\n", N, prefix_len); @@ -1127,16 +1183,8 @@ GenerateResult LagunaBackend::restore_and_generate_impl(int slot, return result; } - // kvflash: restore lands rows identity-mapped; the full prompt (prefix + - // diff) must fit the pool. Rebuild the pager mapping over the prefix. - if (kvflash_active() && - N > kvflash_tokens_ - kvflash_pager_.chunk_tokens()) { - std::fprintf(stderr, "[kvflash] restore prompt (%d) exceeds pool %d; " - "raise --kvflash\n", N, kvflash_tokens_); - result.error = "kvflash: prompt exceeds resident pool"; - return result; - } - if (kvflash_active()) { + // Non-pooled: rebuild pager over the prefix (identity layout, rows intact). + if (!snap_pooled && kvflash_active()) { kvflash_pager_.reset(); if (!kvflash_alloc_span(0, prefix_len)) { result.error = "kvflash_slot"; @@ -1145,46 +1193,93 @@ GenerateResult LagunaBackend::restore_and_generate_impl(int slot, } const KvFlashPager * kvf = kvflash_active() ? &kvflash_pager_ : nullptr; - // Re-prefill diff tokens (or last cached token when diff is empty). - bool restore_only = false; - if (prefix_len == N) { - if (prefix_len <= 0) { result.error = "empty_diff"; return result; } - cache_.cur_pos = prefix_len - 1; - restore_only = true; - } - const int kv_start = cache_.cur_pos; - const int diff_n = N - kv_start; - - // On an exact full-prompt hit (restore_only, diff_n == 1) the caller may - // hand us a zero-filled restore-only prompt (pflash full-cache hit, where - // effective_prompt is a dummy buffer sized to the cached/compressed length). - // req.prompt[kv_start] would then be token 0, corrupting the final logits. - // Re-embed the real last committed token recorded in the snapshot instead. - std::vector diff_ids; - const int32_t * diff_src = req.prompt.data() + kv_start; - if (restore_only && snapshots_[slot].last_tok >= 0) { - diff_ids.assign(1, snapshots_[slot].last_tok); - diff_src = diff_ids.data(); - } - - std::vector embed_diff((size_t)diff_n * w_.n_embd); - if (!w_.embedder.embed(diff_src, diff_n, embed_diff.data())) { - result.error = "embed_prefill"; - return result; - } - std::vector last_logits; bool ok = true; - const int n_chunks = (diff_n + args_.chunk - 1) / args_.chunk; - for (int c = 0; c < n_chunks && ok; ++c) { - const int off = c * args_.chunk; - const int n_tok = std::min(args_.chunk, diff_n - off); - const int starts = kv_start + off; - ok = kvflash_alloc_span(starts, n_tok) && - laguna_step(backend_, w_, cache_, - embed_diff.data() + (size_t)off * w_.n_embd, - n_tok, starts, no_mask, last_logits, kvf, - /*capture=*/true); + + if (snap_pooled && kvflash_active()) { + // Restore-consume path: pager already holds the deserialized prefix KV + // for [0, prefix_len). Prefill only the suffix [prefix_len, N). + // This is the warm-restore win — no re-prefill of the cached prefix. + std::fprintf(stderr, + "[kvflash] laguna restore-consume: snap_pos=%d prompt=%d " + "(suffix prefill only)\n", prefix_len, N); + std::fflush(stderr); + + if (N > prefix_len) { + // Suffix prefill: embed tokens [prefix_len, N) and step from prefix_len. + const int suffix_n = N - prefix_len; + std::vector embed_sfx((size_t)suffix_n * w_.n_embd); + if (!w_.embedder.embed(req.prompt.data() + prefix_len, suffix_n, + embed_sfx.data())) { + result.error = "embed_prefill"; + return result; + } + cache_.cur_pos = prefix_len; // pager already seeded; do NOT reset + const int n_sfx_chunks = (suffix_n + args_.chunk - 1) / args_.chunk; + for (int c = 0; c < n_sfx_chunks && ok; ++c) { + const int off = c * args_.chunk; + const int n_tok = std::min(args_.chunk, suffix_n - off); + const int starts = prefix_len + off; + ok = kvflash_alloc_span(starts, n_tok) && + laguna_step(backend_, w_, cache_, + embed_sfx.data() + (size_t)off * w_.n_embd, + n_tok, starts, no_mask, last_logits, kvf, + /*capture=*/true); + } + } else { + // Exact cache hit (N == prefix_len): re-embed last token to get fresh + // logits. Use the recorded last_tok if the prompt is zero-filled. + int32_t last_id = (snap_ref.last_tok >= 0) ? snap_ref.last_tok + : req.prompt[prefix_len - 1]; + std::vector embed_last((size_t)w_.n_embd); + if (!w_.embedder.embed(&last_id, 1, embed_last.data())) { + result.error = "embed_prefill"; + return result; + } + cache_.cur_pos = prefix_len - 1; + ok = kvflash_alloc_span(cache_.cur_pos, 1) && + laguna_step(backend_, w_, cache_, + embed_last.data(), 1, cache_.cur_pos, + no_mask, last_logits, kvf, /*capture=*/true); + } + } else { + // Non-pooled (or kvflash disabled): re-prefill diff tokens as before. + bool restore_only = false; + if (prefix_len == N) { + if (prefix_len <= 0) { result.error = "empty_diff"; return result; } + cache_.cur_pos = prefix_len - 1; + restore_only = true; + } + const int kv_start = cache_.cur_pos; + const int diff_n = N - kv_start; + + // On an exact full-prompt hit (restore_only, diff_n == 1) the caller may + // hand us a zero-filled restore-only prompt (pflash full-cache hit, where + // effective_prompt is a dummy buffer sized to the cached/compressed length). + std::vector diff_ids; + const int32_t * diff_src = req.prompt.data() + kv_start; + if (restore_only && snap_ref.last_tok >= 0) { + diff_ids.assign(1, snap_ref.last_tok); + diff_src = diff_ids.data(); + } + + std::vector embed_diff((size_t)diff_n * w_.n_embd); + if (!w_.embedder.embed(diff_src, diff_n, embed_diff.data())) { + result.error = "embed_prefill"; + return result; + } + + const int n_chunks = (diff_n + args_.chunk - 1) / args_.chunk; + for (int c = 0; c < n_chunks && ok; ++c) { + const int off = c * args_.chunk; + const int n_tok = std::min(args_.chunk, diff_n - off); + const int starts = kv_start + off; + ok = kvflash_alloc_span(starts, n_tok) && + laguna_step(backend_, w_, cache_, + embed_diff.data() + (size_t)off * w_.n_embd, + n_tok, starts, no_mask, last_logits, kvf, + /*capture=*/true); + } } if (!ok) { result.error = "prefill"; return result; } diff --git a/server/src/laguna/laguna_internal.h b/server/src/laguna/laguna_internal.h index 2ad75e7a9..2c8c26b92 100644 --- a/server/src/laguna/laguna_internal.h +++ b/server/src/laguna/laguna_internal.h @@ -226,6 +226,14 @@ struct LagunaCacheSnapshot { // not recorded (older/inline snapshots) -> restore falls back to req.prompt. int last_tok = -1; bool used = false; + // Pager-blob snapshot: when is_pooled=true the KV cache rows are NOT copied + // into attn_k/attn_v (skip_kv path); instead the full pager state is + // serialized to kvflash_blob via KvFlashPager::serialize(). On restore, + // deserialize() rebuilds the resident pool and restore-consume skips + // re-prefilling the already-snapped prefix. This survives eviction-relocation + // because the blob carries every chunk's raw bytes regardless of pool slot. + bool is_pooled = false; + std::vector kvflash_blob; // pager-serialized KV; valid iff is_pooled }; bool laguna_snapshot_alloc(const LagunaTargetCache & cache,