diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index d61707d96..e5602b368 100644 --- a/server/src/common/kvflash_pager.h +++ b/server/src/common/kvflash_pager.h @@ -84,6 +84,12 @@ struct KvFlashStats { }; class KvFlashPager { + struct AppendProtectRange { + uint64_t id = 0; + int begin = -1; + int end = -1; + }; + public: // `attn_k` / `attn_v` are the per-full-attention-layer cache tensors, // each [head_dim, pool_tokens, n_head_kv]. All must share dims/types @@ -184,6 +190,7 @@ class KvFlashPager { cur_chunk_ = 0; epoch_++; has_pending_page_in_ = false; + append_protect_ranges_.clear(); } // Zero every currently-free block. reset() drops mappings but leaves the @@ -207,6 +214,94 @@ class KvFlashPager { // Optional external relevance score; higher = keep. Falls back to LRU. std::function score_hook; + class AppendReservation { + public: + AppendReservation() = default; + AppendReservation(const AppendReservation &) = delete; + AppendReservation & operator=(const AppendReservation &) = delete; + AppendReservation(AppendReservation && other) noexcept { move_from(other); } + AppendReservation & operator=(AppendReservation && other) noexcept { + if (this != &other) { + release(); + move_from(other); + } + return *this; + } + ~AppendReservation() { release(); } + + bool ok() const { return ok_; } + int kv_start() const { return kv_start_; } + int n_tokens() const { return n_tok_; } + int slot(int i) const { + return i >= 0 && i < (int)slots_.size() ? slots_[(size_t)i] : -1; + } + + private: + friend class KvFlashPager; + AppendReservation(KvFlashPager & pager, int kv_start, int n_tok) + : pager_(&pager), + kv_start_(kv_start), + n_tok_(n_tok) { + if (n_tok <= 0 || pager.cfg_.chunk_tokens <= 0) { + pager_ = nullptr; + return; + } + protect_id_ = ++pager.append_protect_next_id_; + pager.append_protect_ranges_.push_back({ + protect_id_, + kv_start / pager.cfg_.chunk_tokens, + (kv_start + n_tok - 1) / pager.cfg_.chunk_tokens, + }); + slots_.resize((size_t)n_tok, -1); + for (int i = 0; i < n_tok; ++i) { + const int slot = pager.slot_for(kv_start + i); + if (slot < 0) { + std::fprintf(stderr, "[kvflash] no pool slot at pos %d " + "(pool %d exhausted)\n", + kv_start + i, pager.cfg_.pool_tokens); + release(); + return; + } + slots_[(size_t)i] = slot; + } + if (pager.has_pending_page_in_) pager.synchronize_paging(); + ok_ = true; + } + + void release() { + if (!pager_) return; + if (protect_id_ != 0) { + auto & ranges = pager_->append_protect_ranges_; + ranges.erase(std::remove_if(ranges.begin(), ranges.end(), + [&](const AppendProtectRange & r) { + return r.id == protect_id_; + }), ranges.end()); + } + pager_ = nullptr; + protect_id_ = 0; + ok_ = false; + } + + void move_from(AppendReservation & other) { + pager_ = other.pager_; + kv_start_ = other.kv_start_; + n_tok_ = other.n_tok_; + protect_id_ = other.protect_id_; + slots_ = std::move(other.slots_); + ok_ = other.ok_; + other.pager_ = nullptr; + other.protect_id_ = 0; + other.ok_ = false; + } + + KvFlashPager * pager_ = nullptr; + int kv_start_ = 0; + int n_tok_ = 0; + uint64_t protect_id_ = 0; + std::vector slots_; + bool ok_ = false; + }; + // 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. @@ -216,16 +311,12 @@ class KvFlashPager { // issued, page_stream_ is synchronised here so that recalled KV data is // resident before the next attention forward. bool alloc_span(int kv_start, int n_tok) { - for (int i = 0; i < n_tok; ++i) { - if (slot_for(kv_start + i) < 0) { - std::fprintf(stderr, "[kvflash] no pool slot at pos %d " - "(pool %d exhausted)\n", - kv_start + i, cfg_.pool_tokens); - return false; - } - } - if (has_pending_page_in_) synchronize_paging(); - return true; + if (n_tok <= 0) return true; + return reserve_append_span(kv_start, n_tok).ok(); + } + + AppendReservation reserve_append_span(int kv_start, int n_tok) { + return AppendReservation(*this, kv_start, n_tok); } // Physical pool slot for logical position `pos`. Allocates (and, when @@ -366,6 +457,10 @@ class KvFlashPager { const KvFlashStats & stats() const { return stats_; } int resident_blocks() const { return n_blocks_ - (int)free_blocks_.size(); } int n_chunks() const { return (int)chunks_.size(); } + int max_append_tokens() const { + const int chunks = n_blocks_ - cfg_.sink_chunks - cfg_.tail_window_chunks; + return std::max(1, chunks) * cfg_.chunk_tokens; + } // Bumped on every residency change (alloc / page_out / page_in). // Callers cache the slot mask and refill only when the epoch moves. @@ -393,17 +488,26 @@ class KvFlashPager { if (!score_hook) return 0; struct Cand { int c; float s; }; std::vector cands; + std::vector want(chunks_.size(), 0); + int kept = 0; for (int c = 0; c < (int)chunks_.size(); c++) { 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; - cands.push_back({c, prot ? 3.4e38f : score_hook(c)}); + c > cur_chunk_ - 1 - cfg_.tail_window_chunks || + append_protected(c); + if (prot) { + want[(size_t)c] = 1; + kept++; + continue; + } + cands.push_back({c, score_hook(c)}); } std::sort(cands.begin(), cands.end(), [](const Cand & a, const Cand & b) { return a.s > b.s; }); - std::vector want(chunks_.size(), 0); - for (int i = 0; i < (int)cands.size() && i < n_blocks_; i++) want[cands[i].c] = 1; + for (int i = 0; i < (int)cands.size() && kept < n_blocks_; i++, kept++) { + want[(size_t)cands[(size_t)i].c] = 1; + } int events = 0; for (int c = 0; c < (int)chunks_.size(); c++) { // out first: frees blocks @@ -430,6 +534,13 @@ class KvFlashPager { #endif }; + bool append_protected(int c) const { + for (const AppendProtectRange & r : append_protect_ranges_) { + if (c >= r.begin && c <= r.end) return true; + } + return false; + } + bool ensure_free_block() { if (!free_blocks_.empty()) return true; // Victim: unprotected resident chunk with the lowest score @@ -441,6 +552,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 (append_protected(c)) continue; if (score_hook) { const float s = score_hook(c); if (victim < 0 || s < v_score) { victim = c; v_score = s; } @@ -571,6 +683,8 @@ class KvFlashPager { int n_blocks_ = 0, n_head_kv_ = 0, cur_chunk_ = 0; uint64_t clock_ = 0; uint64_t epoch_ = 0; + uint64_t append_protect_next_id_ = 0; + std::vector append_protect_ranges_; #ifdef KVFLASH_HAS_ASYNC_DMA cudaStream_t page_stream_ = nullptr; @@ -746,4 +860,61 @@ inline bool kvflash_fill_rows_and_masks( return true; } +struct KvFlashPooledPrefillPlan { + KvFlashPager::AppendReservation reservation; + std::vector rows; + std::vector mask; + bool ok() const { return reservation.ok(); } +}; + +inline KvFlashPooledPrefillPlan kvflash_prepare_pooled_suffix_prefill( + KvFlashPager & pager, + int kv_start, int n_tok, int n_head_kv, + int mask_kv_dim, int mask_q_dim) { + constexpr uint16_t F16_ZERO = 0x0000, F16_NEG_INF = 0xFC00; + KvFlashPooledPrefillPlan plan; + if (n_tok <= 0 || n_head_kv <= 0 || + mask_kv_dim < pager.pool_tokens() || mask_q_dim < n_tok) { + std::fprintf(stderr, + "[kvflash] invalid pooled prefill plan kv_start=%d n_tok=%d " + "heads=%d mask=%dx%d pool=%d\n", + kv_start, n_tok, n_head_kv, mask_kv_dim, mask_q_dim, + pager.pool_tokens()); + return plan; + } + plan.reservation = pager.reserve_append_span(kv_start, n_tok); + if (!plan.reservation.ok()) return plan; + + plan.rows.resize((size_t)n_tok * n_head_kv); + for (int i = 0; i < n_tok; ++i) { + const int slot = plan.reservation.slot(i); + if (slot < 0 || slot >= mask_kv_dim) { + std::fprintf(stderr, + "[kvflash] pooled prefill slot missing pos=%d slot=%d mask_kv=%d\n", + kv_start + i, slot, mask_kv_dim); + plan.reservation = KvFlashPager::AppendReservation{}; + plan.rows.clear(); + plan.mask.clear(); + return plan; + } + for (int h = 0; h < n_head_kv; ++h) { + plan.rows[(size_t)h * n_tok + i] = slot; + } + } + + std::vector slot_pos((size_t)pager.pool_tokens(), -1); + pager.fill_slot_pos(slot_pos.data()); + plan.mask.assign((size_t)mask_kv_dim * mask_q_dim, F16_NEG_INF); + const int slots = std::min(mask_kv_dim, pager.pool_tokens()); + for (int q = 0; q < n_tok; ++q) { + const int abs_q = kv_start + q; + uint16_t * row = plan.mask.data() + (size_t)q * mask_kv_dim; + for (int s = 0; s < slots; ++s) { + const int p = slot_pos[(size_t)s]; + if (p >= 0 && p <= abs_q) row[s] = F16_ZERO; + } + } + return plan; +} + } // namespace dflash::common diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index 3eefe62c5..3e6d3a7a9 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -43,8 +43,11 @@ struct StepGraph { // size (stable topology for CUDA-graph replay); noise keys start here. int ctx_alloc = 0; ggml_tensor * hidden_input = nullptr; // lm-head projection only - // [n_tokens,n_head_kv] i64; step-invariant KV write (carries kv_start). Null on non-graph paths. + // [n_tokens,n_head_kv] i64; step-invariant KV write. Null on non-graph paths. ggml_tensor * kv_write_rows = nullptr; + // Non-pooled decode writes into a bucket-local cache view; callers upload + // logical_row - kv_write_row_base. Pooled KVFlash slots leave this at 0. + int kv_write_row_base = 0; // Output ggml_tensor * logits = nullptr; @@ -74,6 +77,7 @@ inline void step_graph_free(StepGraph & sg) { sg.hidden_input = nullptr; sg.parent_ids = nullptr; sg.kv_write_rows = nullptr; + sg.kv_write_row_base = 0; sg.logits = nullptr; sg.hidden_states = nullptr; sg.argmax_tokens = nullptr; diff --git a/server/src/internal.h b/server/src/internal.h index 6bf052f4d..3ed29120a 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -597,8 +597,11 @@ struct QwenGraphInputs { int fa_window = 0; // sliding window for FA layers: 0 = full attention bool last_token_logits_only = false; // if true, only compute logits for last token (prefill optimization) ggml_tensor * parent_ids = nullptr; // [n_tokens] i32; tree mode when non-null - // [n_tokens,n_head_kv] i64; non-null = step-invariant KV write via ggml_set_rows (carries kv_start). + // [n_tokens,n_head_kv] i64; non-null = step-invariant KV write via ggml_set_rows. ggml_tensor * kv_write_rows = nullptr; + // Base row subtracted by the caller before filling kv_write_rows. Non-zero + // selects a bucket-local cache destination for non-pooled decode. + int kv_write_row_base = 0; // Capture the LAST token's post-RoPE/post-rotation Q per full-attention // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. @@ -646,10 +649,11 @@ ggml_tensor * build_qwen35_layer( int kv_start, int n_tokens, bool capture, - int fa_window = 0, - ggml_tensor * q_tail_capture = nullptr, - int q_tail_start = 0, - ggml_tensor * kv_write_rows = nullptr); + int fa_window, + ggml_tensor * q_tail_capture, + int q_tail_start, + ggml_tensor * kv_write_rows, + int kv_write_row_base); // Overload that also exposes the MoE router selection tensor (if MoE layer). ggml_tensor * build_qwen35_layer( @@ -668,7 +672,8 @@ ggml_tensor * build_qwen35_layer( ggml_tensor * q_tail_capture, int q_tail_start, ggml_tensor ** moe_selected_out, - ggml_tensor * kv_write_rows = nullptr); + ggml_tensor * kv_write_rows, + int kv_write_row_base); QwenLayerPrefnOutputs build_qwen35_layer_prefn( ggml_context * ctx, @@ -681,8 +686,9 @@ QwenLayerPrefnOutputs build_qwen35_layer_prefn( ggml_tensor * attn_mask, int kv_start, int n_tokens, - int fa_window = 0, - ggml_tensor * kv_write_rows = nullptr, + int fa_window, + ggml_tensor * kv_write_rows, + int kv_write_row_base, bool skip_gdn_intermediate = true); } // namespace dflash::common diff --git a/server/src/qwen3/qwen3_drafter.cpp b/server/src/qwen3/qwen3_drafter.cpp index 3acc1775a..ddca4d397 100644 --- a/server/src/qwen3/qwen3_drafter.cpp +++ b/server/src/qwen3/qwen3_drafter.cpp @@ -374,7 +374,11 @@ static std::vector qwen35_score_and_compress( mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, align_up_i(kv_len, 32), align_up_i(n, 32)); ggml_set_input(mask); } - ggml_tensor * out = build_qwen35_layer(ctx, gf, w, cache, il, inp, pos, mask, start, n, false, 0); + ggml_tensor * out = build_qwen35_layer( + ctx, gf, w, cache, il, inp, pos, mask, start, n, + /*capture=*/false, /*fa_window=*/0, + /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, + /*kv_write_rows=*/nullptr, /*kv_write_row_base=*/0); ggml_tensor * dst = ggml_view_2d(ctx, act_out, hidden, n, act_out->nb[1], (size_t)start * act_out->nb[1]); if (ggml_nelements(out) != ggml_nelements(dst)) { std::fprintf(stderr, diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 842314b0d..0cd1fef8f 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -3,10 +3,21 @@ #include "ggml-alloc.h" #include +#include #include +#include namespace dflash::common { +static int qwen35_non_pooled_kvpad_max_row() { + const char * env = std::getenv("DFLASH_QWEN35_KVPAD_MAX_ROW"); + if (env) return std::atoi(env); + // q4_0 full-cache SET_ROWS is not reliable in the full long-context Qwen + // AR graph near 89k rows on CUDA. Keep the fast path below the observed + // cliff and use the exact legacy copy graph above it. + return 88000; +} + // ── build_layer_step ──────────────────────────────────────────── bool build_layer_step( @@ -81,7 +92,8 @@ bool build_layer_step( sg.inp_embed, sg.positions, sg.attn_mask, kv_start, n_tokens, capture, fa_window, /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, - sg.kv_write_rows); + sg.kv_write_rows, + sg.kv_write_row_base); if (!layer_out) return false; ggml_tensor * out_view = ggml_view_2d(sg.ctx, act_out, @@ -156,6 +168,7 @@ bool build_layer_prefn_step( sg.inp_embed, sg.positions, sg.attn_mask, kv_start, n_tokens, fa_window, sg.kv_write_rows, + sg.kv_write_row_base, /*skip_gdn_intermediate=*/true); if (!go.residual || !go.post) return false; sg.ffn_residual = go.residual; @@ -234,7 +247,9 @@ bool build_hybrid_full_layer_step( sg.inp_embed, sg.positions, sg.attn_mask, kv_start, n_tokens, /*capture=*/false, fa_window, /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, - &moe_selected); + &moe_selected, + /*kv_write_rows=*/nullptr, + /*kv_write_row_base=*/0); if (!layer_out) return false; // Use hidden_input as the layer output tensor (repurpose field) @@ -327,16 +342,37 @@ bool build_target_step( // region is reused by graph execution) and set_rows carries per-token // physical slots, so the slot-mapped write stays active for masked, // multi-token, and feature-capturing forwards (decode AND spec verify). + const bool non_pooled_step_kv_write = + n_tokens == 1 && fa_window == 0 && !with_mask && !capture; + const bool q4_full_cache_kv = + !kvflash_mask && + (cache.kv_k_type == GGML_TYPE_Q4_0 || cache.kv_v_type == GGML_TYPE_Q4_0); + const int non_pooled_kvpad_max_row = qwen35_non_pooled_kvpad_max_row(); + const bool non_pooled_step_kv_write_safe = + !q4_full_cache_kv || non_pooled_kvpad_max_row <= 0 || + kv_start < non_pooled_kvpad_max_row; + static std::atomic s_logged_q4_guard{false}; + if (non_pooled_step_kv_write && q4_full_cache_kv && + !non_pooled_step_kv_write_safe && + !s_logged_q4_guard.exchange(true, std::memory_order_relaxed)) { + std::fprintf(stderr, + "[qwen35] q4_0 non-pooled KV write fallback at row=%d " + "(DFLASH_QWEN35_KVPAD_MAX_ROW=%d)\n", + kv_start, non_pooled_kvpad_max_row); + } const bool use_kv_write_rows = !g_no_kvpad && !capture_delta_intermediate && (kvflash_mask ? (fa_window == 0) - : (n_tokens == 1 && fa_window == 0 && !with_mask && !capture)); + : (non_pooled_step_kv_write && non_pooled_step_kv_write_safe)); if (use_kv_write_rows) { sg.kv_write_rows = ggml_new_tensor_2d(sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv); ggml_set_name(sg.kv_write_rows, "kv_write_rows"); ggml_set_input(sg.kv_write_rows); + sg.kv_write_row_base = (!kvflash_mask && non_pooled_step_kv_write) + ? (kv_start & ~255) + : 0; } QwenGraphInputs gi{}; @@ -351,6 +387,7 @@ bool build_target_step( gi.fa_window = fa_window; gi.last_token_logits_only = last_token_logits_only; gi.kv_write_rows = sg.kv_write_rows; + gi.kv_write_row_base = sg.kv_write_row_base; gi.q_capture = capture_qk; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4c9a727da..4b60c2e01 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1012,12 +1012,11 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, prefill_last_logits_valid_ = false; // kvflash: a prompt that fits the pool prefills contiguously (identity - // mapping, normal chunking). A LARGER prompt switches to POOLED CHUNKED - // PREFILL: pager-chunk-sized batches whose KV rows are slot-mapped via - // set_rows, with a slot-space mask per chunk and live eviction as the - // pool fills (constant VRAM, linear time). Restore offsets are not - // supported in the pooled path (a relocated prefix cannot be restored - // identity-style in the first place). + // mapping, normal chunking). A LARGER prompt switches to POOLED SUFFIX + // PREFILL: large batches whose KV rows are slot-mapped via set_rows, with + // a slot-space causal mask and live eviction as the pool fills (constant + // VRAM, linear time). Restore offsets are not supported in the pooled path + // until snapshots serialize pager state. const bool kvf_paged = kvflash_active() && kv_offset + prompt_len > kvflash_tokens_ - kvflash_pager_.chunk_tokens(); if (kvf_paged && kv_offset != 0) { @@ -1029,15 +1028,21 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, return -1; } if (kvf_paged) { - prefill_ubatch = kvflash_pager_.chunk_tokens(); + prefill_ubatch = std::min(prefill_ubatch, kvflash_pager_.max_append_tokens()); + const int kvf_chunk = kvflash_pager_.chunk_tokens(); + if (kvf_chunk > 0 && prefill_ubatch > kvf_chunk) { + prefill_ubatch = (prefill_ubatch / kvf_chunk) * kvf_chunk; + } + prefill_ubatch = std::max(prefill_ubatch, kvf_chunk); kvflash_pager_.reset(); if (kvflash_qk_policy_) { kvflash_qk_pool_.reset(kvflash_qk_pool_.dims()); kvflash_qk_pooled_upto_ = 0; } std::printf("[kvflash] pooled prefill: %d tokens through a %d-token pool " - "(%d-token chunks, evicting)\n", - prompt_len, kvflash_tokens_, prefill_ubatch); + "(ubatch=%d, chunk=%d, evicting)\n", + prompt_len, kvflash_tokens_, prefill_ubatch, + kvflash_pager_.chunk_tokens()); std::fflush(stdout); } @@ -1060,56 +1065,57 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, // Chunked prefill std::vector embed_buf((size_t)hidden * prefill_ubatch); int committed = kv_offset; + auto save_inline_snapshot = [&](int cur_pos) { + if (snap_slot < 0 || snap_pos < 0 || snap_pos != cur_pos) return; + if (kvf_paged) { + std::fprintf(stderr, "[kvflash] boundary snapshot skipped: pooled " + "prefill relocates chunks\n"); + } else if (cur_pos > kv_offset) { + cache_.cur_pos = cur_pos; + if (snapshot_save(snap_slot)) { + std::printf("[snap] boundary slot=%d cur_pos=%d\n", + snap_slot, cur_pos); + std::fflush(stdout); + } + } else { + snap_pos = -1; + snap_slot = -1; + return; + } + snap_pos = -1; + snap_slot = -1; + }; + save_inline_snapshot(committed); for (int start = 0; start < prompt_len;) { const int kv_pos = kv_offset + start; int n_tokens = std::min(prefill_ubatch, prompt_len - start); - // FIX(bug2): do NOT shrink the prefill chunk to snap_pos. Shrinking - // realigns every subsequent chunk, changing GPU batch sizes vs the - // no-cache path -> FP-nondeterministic state divergence -> different - // greedy output on cache hits. Keep uniform chunks. When snap_pos falls - // inside this chunk, snapshot at the chunk START boundary kv_pos: the - // largest chunk boundary <= snap_pos. That stays (a) chunk-aligned, so - // the prefill is bit-identical to the no-cache path, and (b) strictly - // within the requested prefix, so a later request that shares only the - // system-prompt prefix still restores a valid cross-request hit. - // (Rounding UP would push the snapshot to prompt end -> the full prompt - // incl. the user message -> a different user msg restores garbage.) - if (snap_slot >= 0 && snap_pos >= 0 && - kv_pos <= snap_pos && snap_pos < kv_pos + n_tokens) { - if (kv_pos > kv_offset && !kvf_paged) { // skip degenerate / relocated - cache_.cur_pos = kv_pos; - if (snapshot_save(snap_slot)) { - std::printf("[snap] boundary slot=%d cur_pos=%d (req snap_pos=%d)\n", - snap_slot, kv_pos, snap_pos); - std::fflush(stdout); - } - } else if (kvf_paged) { - std::fprintf(stderr, "[kvflash] boundary snapshot skipped: pooled " - "prefill relocates chunks\n"); + if (!kvf_paged && prefill_ubatch > 0) { + const int phase = kv_pos % prefill_ubatch; + if (phase > 0) { + n_tokens = std::min(n_tokens, prefill_ubatch - phase); } + } + // Inline prefix-cache entries are exact by construction: the snapshot + // must physically materialize the requested chat boundary. If the + // boundary lands inside this ubatch, split once, save at snap_pos, then + // resume on the same absolute ubatch phase. A restored delta starting + // at snap_pos will take the same short first chunk and rejoin the + // regular ubatch grid at the next boundary. + if (!kvf_paged && snap_slot >= 0 && snap_pos > kv_pos && + snap_pos < kv_pos + n_tokens) { + n_tokens = snap_pos - kv_pos; + } else if (kvf_paged && snap_slot >= 0 && snap_pos >= kv_pos && + snap_pos < kv_pos + n_tokens) { + std::fprintf(stderr, "[kvflash] boundary snapshot skipped: pooled " + "prefill relocates chunks\n"); snap_pos = -1; snap_slot = -1; } const bool with_mask = kvf_paged || (cfg_.kq_stride_pad > KQ_MASK_PAD) || (n_tokens > 1); - // kvflash pooled prefill: allocate this chunk's slots up front - // (evicting the lowest-priority resident chunk once the pool fills). - std::vector kvf_slots; - if (kvf_paged) { - kvf_slots.resize((size_t)n_tokens); - bool ok = true; - for (int i = 0; i < n_tokens; i++) { - kvf_slots[(size_t)i] = kvflash_pager_.slot_for(kv_pos + i); - if (kvf_slots[(size_t)i] < 0) { ok = false; break; } - } - if (!ok) { - std::fprintf(stderr, "[kvflash] pooled prefill: slot alloc failed @%d\n", kv_pos); - set_last_error("kvflash: no evictable pool block"); - return -1; - } - } + KvFlashPooledPrefillPlan kvf_plan; // Prefill always uses full attention (fa_window=0) so that all // positions encode the complete context — critical for tool @@ -1132,15 +1138,21 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, std::fprintf(stderr, "[kvflash] pooled prefill requires the set_rows path\n"); return -1; } - // [n_tokens, n_head_kv] ne0-major (see verify_batch). - std::vector rows((size_t)n_tokens * w_.n_head_kv); - for (int h = 0; h < w_.n_head_kv; h++) { - for (int i = 0; i < n_tokens; i++) { - rows[(size_t)h * n_tokens + i] = kvf_slots[(size_t)i]; - } + if (!sg_.attn_mask) { + std::fprintf(stderr, "[kvflash] pooled prefill requires a slot mask\n"); + set_last_error("kvflash: pooled prefill mask missing"); + return -1; } - ggml_backend_tensor_set(sg_.kv_write_rows, rows.data(), 0, - sizeof(int64_t) * rows.size()); + kvf_plan = kvflash_prepare_pooled_suffix_prefill( + kvflash_pager_, kv_pos, n_tokens, w_.n_head_kv, + (int)sg_.attn_mask->ne[0], (int)sg_.attn_mask->ne[1]); + if (!kvf_plan.ok()) { + std::fprintf(stderr, "[kvflash] pooled prefill: plan failed @%d\n", kv_pos); + set_last_error("kvflash: pooled prefill plan failed"); + return -1; + } + ggml_backend_tensor_set(sg_.kv_write_rows, kvf_plan.rows.data(), 0, + sizeof(int64_t) * kvf_plan.rows.size()); } // Embed @@ -1164,32 +1176,8 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, // Mask — full attention during prefill (no windowing) if (sg_.attn_mask && kvf_paged) { - // Slot-space mask (same recipe as verify_batch): row q attends - // (a) the slots of resident chunks holding positions < kv_pos - // and (b) this chunk's own slots, causally. - constexpr uint16_t F16_ZERO = 0x0000, F16_NEG_INF = 0xFC00; - const size_t kvd = (size_t)sg_.attn_mask->ne[0]; - const int q_pad = (int)sg_.attn_mask->ne[1]; - std::vector mask_buf((size_t)kvd * q_pad, F16_NEG_INF); - const int ct = kvflash_pager_.chunk_tokens(); - for (int c = 0; c < kvflash_pager_.n_chunks(); c++) { - const int blk = kvflash_pager_.block_of(c); - if (blk < 0) continue; - for (int i = 0; i < ct; i++) { - if ((int64_t)c * ct + i >= kv_pos) break; - mask_buf[(size_t)blk * ct + i] = F16_ZERO; - } - } - for (int q = 1; q < n_tokens; q++) { - std::memcpy(mask_buf.data() + (size_t)q * kvd, mask_buf.data(), kvd * 2); - } - for (int q = 0; q < n_tokens; q++) { - for (int i = 0; i <= q; i++) { - mask_buf[(size_t)q * kvd + kvf_slots[(size_t)i]] = F16_ZERO; - } - } - ggml_backend_tensor_set(sg_.attn_mask, mask_buf.data(), 0, - sizeof(uint16_t) * mask_buf.size()); + ggml_backend_tensor_set(sg_.attn_mask, kvf_plan.mask.data(), 0, + sizeof(uint16_t) * kvf_plan.mask.size()); } else if (sg_.attn_mask) { const int win_start = 0; const int kv_len = kv_pos + n_tokens - win_start; @@ -1221,6 +1209,7 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, committed = kv_pos + n_tokens; cache_.cur_pos = committed; + save_inline_snapshot(committed); // QK policy: pool the post-RoPE keys of chunks this batch sealed // (they are resident — sealed inside the protected tail window). @@ -1249,11 +1238,8 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, } } - // End-of-prefill snapshot: scoped disk-cache saves (auto/fixed policy) - // request snap_pos == prompt end, which never falls inside a chunk so the - // boundary branch above cannot fire. Taking the snapshot here changes - // nothing about the prefill computation; it only persists the final state - // (cache_.cur_pos == committed). + // Defensive end-of-prefill snapshot for callers that request exactly the + // final committed position and did not already consume the inline save. if (snap_slot >= 0 && snap_pos == committed) { if (snapshot_save(snap_slot)) { std::printf("[snap] end-of-prefill slot=%d cur_pos=%d\n", @@ -1557,16 +1543,25 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, } // Fill kv_write_rows with this step's cache slot for set_rows: - // the logical position directly, or its pool slot in kvflash mode. + // a pooled physical slot, or a bucket-local logical row in non-pool mode. if (sg_.kv_write_rows) { const int n_head_kv = w_.n_head_kv; - const int64_t slot = pool ? (int64_t)kvflash_pager_.slot_for(committed) - : (int64_t)committed; - if (pool && slot < 0) { - std::fprintf(stderr, "[kvflash] no pool slot at pos %d " - "(pool %d exhausted)\n", - committed, kvflash_tokens_); - set_last_error("kvflash: no evictable pool block"); + const int64_t slot = pool + ? (int64_t)kvflash_pager_.slot_for(committed) + : (int64_t)committed - (int64_t)sg_.kv_write_row_base; + if (slot < 0 || (!pool && slot >= 256)) { + if (pool) { + std::fprintf(stderr, "[kvflash] no pool slot at pos %d " + "(pool %d exhausted)\n", + committed, kvflash_tokens_); + } else { + std::fprintf(stderr, "[qwen35] invalid bucket-local KV row " + "(pos=%d base=%d row=%lld)\n", + committed, sg_.kv_write_row_base, + (long long)slot); + } + set_last_error(pool ? "kvflash: no evictable pool block" + : "qwen35: invalid bucket-local KV row"); return false; } std::vector row_vals(n_head_kv, slot); diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index ea4684e09..807efe283 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -36,6 +36,7 @@ #include "qwen35_ops.h" #include "qwen35moe_ffn.h" +#include #include #include #include @@ -527,6 +528,8 @@ static ggml_tensor * build_swiglu_ffn(ggml_context * ctx, ggml_tensor * cur, // over [0..kv_start + n_tokens). // // kv_write_rows: non-null selects the step-invariant ggml_set_rows KV write; null = legacy ggml_cpy. +// kv_write_row_base: for non-pooled decode, set_rows targets a 256-row cache +// view starting here, and callers upload bucket-local row ids. static ggml_tensor * build_full_attn_block( ggml_context * ctx, ggml_cgraph * gf, @@ -547,6 +550,7 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * q_tail_capture = nullptr, int q_tail_start = 0, ggml_tensor * kv_write_rows = nullptr, + int kv_write_row_base = 0, ggml_tensor ** q_fa_out = nullptr // post-RoPE/post-rotation Q [head_dim, n_tokens, n_head] ) { const int head_dim = w.n_embd_head_k; @@ -639,8 +643,25 @@ static ggml_tensor * build_full_attn_block( // Step-invariant: constant dst pointer, idx carries kv_start. set_rows needs contiguous src. ggml_tensor * Kcur_cont = ggml_is_contiguous(Kcur_T) ? Kcur_T : ggml_cont(ctx, Kcur_T); ggml_tensor * Vcur_cont = ggml_is_contiguous(Vcur_T) ? Vcur_T : ggml_cont(ctx, Vcur_T); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache_k, Kcur_cont, kv_write_rows)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache_v, Vcur_cont, kv_write_rows)); + ggml_tensor * cache_k_write = cache_k; + ggml_tensor * cache_v_write = cache_v; + int row_span = (int) cache_k->ne[1]; + if (kv_write_row_base > 0) { + const int cache_rows = (int) cache_k->ne[1]; + GGML_ASSERT(kv_write_row_base < cache_rows); + row_span = std::min(256, cache_rows - kv_write_row_base); + GGML_ASSERT(row_span > 0); + cache_k_write = ggml_view_3d(ctx, cache_k, + head_dim, row_span, n_head_kv, + cache_k->nb[1], cache_k->nb[2], + (size_t)cache_k->nb[1] * (size_t)kv_write_row_base); + cache_v_write = ggml_view_3d(ctx, cache_v, + head_dim, row_span, n_head_kv, + cache_v->nb[1], cache_v->nb[2], + (size_t)cache_v->nb[1] * (size_t)kv_write_row_base); + } + ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache_k_write, Kcur_cont, kv_write_rows)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, cache_v_write, Vcur_cont, kv_write_rows)); } else { // Legacy: kv_start as literal view offset (not step-invariant; prefill/verify/non-graph). ggml_tensor * k_slot = ggml_view_3d(ctx, cache_k, @@ -1044,7 +1065,8 @@ static ggml_tensor * build_single_layer( ggml_tensor * q_tail_capture = nullptr, int q_tail_start = 0, ggml_tensor ** moe_selected_out = nullptr, - ggml_tensor * kv_write_rows = nullptr) + ggml_tensor * kv_write_rows = nullptr, + int kv_write_row_base = 0) { const int hidden = w.n_embd; const float eps = w.rms_eps; @@ -1070,7 +1092,8 @@ static ggml_tensor * build_single_layer( cache.kv_k_rotated, fa_window, q_tail_capture, q_tail_start, - kv_write_rows); + kv_write_rows, + kv_write_row_base); } else { int dn_idx = 0; for (int il = 0; il < layer_idx; il++) { @@ -1191,6 +1214,7 @@ QwenGraphOutputs build_qwen35_graph( /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, in.kv_write_rows, + in.kv_write_row_base, want_q_cap ? &q_fa : nullptr); if (want_q_cap && q_fa) { // Last token's Q, all heads: src [head_dim, 1, n_head] view of @@ -1334,12 +1358,13 @@ ggml_tensor * build_qwen35_layer( int fa_window, ggml_tensor * q_tail_capture, int q_tail_start, - ggml_tensor * kv_write_rows) + ggml_tensor * kv_write_rows, + int kv_write_row_base) { return build_single_layer(ctx, gf, w, cache, layer_idx, inp, positions, attn_mask, kv_start, n_tokens, capture, fa_window, q_tail_capture, q_tail_start, nullptr, - kv_write_rows); + kv_write_rows, kv_write_row_base); } ggml_tensor * build_qwen35_layer( @@ -1358,12 +1383,13 @@ ggml_tensor * build_qwen35_layer( ggml_tensor * q_tail_capture, int q_tail_start, ggml_tensor ** moe_selected_out, - ggml_tensor * kv_write_rows) + ggml_tensor * kv_write_rows, + int kv_write_row_base) { return build_single_layer(ctx, gf, w, cache, layer_idx, inp, positions, attn_mask, kv_start, n_tokens, capture, fa_window, q_tail_capture, q_tail_start, moe_selected_out, - kv_write_rows); + kv_write_rows, kv_write_row_base); } QwenLayerPrefnOutputs build_qwen35_layer_prefn( @@ -1379,6 +1405,7 @@ QwenLayerPrefnOutputs build_qwen35_layer_prefn( int n_tokens, int fa_window, ggml_tensor * kv_write_rows, + int kv_write_row_base, bool skip_gdn_intermediate) { QwenLayerPrefnOutputs out{}; const float eps = w.rms_eps; @@ -1401,7 +1428,8 @@ QwenLayerPrefnOutputs build_qwen35_layer_prefn( cache.kv_k_rotated, fa_window, /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, - kv_write_rows); + kv_write_rows, + kv_write_row_base); } else { int dn_idx = 0; for (int il = 0; il < layer_idx; il++) { diff --git a/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp b/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp index b81e8b87d..58ac8607b 100644 --- a/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp +++ b/server/src/qwen35moe/qwen35moe_pipelined_decode.cpp @@ -38,6 +38,8 @@ void CachedPrefnGraph::free() { positions = nullptr; kv_write_rows = nullptr; kv_win = 0; + kv_write_row_base = 0; + kv_write_pooled = false; } @@ -74,6 +76,7 @@ static bool build_cached_deltanet_prefn( out.inp_embed, /*positions=*/nullptr, /*attn_mask=*/nullptr, kv_start, /*n_tokens=*/1, /*fa_window=*/0, /*kv_write_rows=*/nullptr, + /*kv_write_row_base=*/0, /*skip_gdn_intermediate=*/true); if (!go.residual || !go.post) { out.free(); return false; } @@ -116,6 +119,7 @@ static bool build_cached_attn_prefn( TargetCache & cache, int layer_idx, int kv_win, // 256-aligned FA span to bake in + bool pooled_kvflash, int kq_stride_pad) { ggml_gallocr_t keep_alloc = out.alloc; // reuse allocator across rebuilds @@ -142,6 +146,8 @@ static bool build_cached_attn_prefn( out.kv_write_rows = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I64, 1, w.n_head_kv); ggml_set_name(out.kv_write_rows, "kv_write_rows"); ggml_set_input(out.kv_write_rows); + out.kv_write_pooled = pooled_kvflash; + out.kv_write_row_base = pooled_kvflash ? 0 : ((kv_win - 1) & ~255); // Bake the FA span: build with kv_start = kv_win-1 so win_len_padded == kv_win. // The actual write/read positions come from kv_write_rows + positions DATA. @@ -151,6 +157,7 @@ static bool build_cached_attn_prefn( out.inp_embed, out.positions, /*attn_mask=*/nullptr, /*kv_start=*/kv_win - 1, /*n_tokens=*/1, /*fa_window=*/0, out.kv_write_rows, + out.kv_write_row_base, /*skip_gdn_intermediate=*/true); if (!go.residual || !go.post) { out.free(); return false; } @@ -623,20 +630,25 @@ bool pipelined_decode_one_token( bool attn_cached_ok = false; if (is_attn && !g_no_kvpad) { auto & cpg = state.cached_prefn[(size_t)il]; + const bool pooled_kvflash = kv_slot >= 0; // Clamp the baked FA span to the cache tensor's physical capacity: // with kvflash the tensors are pool-sized, so the window stops // growing at the pool (and the cached graph never rebuilds again). const int kv_phys = (int)cache.attn_k[0]->ne[1]; const int kv_win_needed = std::min(((kv_pos + 1) + 255) & ~255, kv_phys); - if (!cpg.valid() || cpg.kv_win < kv_win_needed) { + if (!cpg.valid() || cpg.kv_win < kv_win_needed || + cpg.kv_write_pooled != pooled_kvflash) { if (!build_cached_attn_prefn(cpg, backend, w, cache, il, - kv_win_needed, kq_stride_pad)) { + kv_win_needed, pooled_kvflash, + kq_stride_pad)) { std::fprintf(stderr, "[pipelined] cached attn prefn build failed (layer %d); dyn fallback\n", il); } } - attn_cached_ok = cpg.valid() && cpg.kv_win >= kv_win_needed; + attn_cached_ok = cpg.valid() && + cpg.kv_win >= kv_win_needed && + cpg.kv_write_pooled == pooled_kvflash; } if (attn_cached_ok) { @@ -644,7 +656,16 @@ bool pipelined_decode_one_token( ggml_backend_tensor_copy_async(backend, backend, state.gpu_state.act_cur, cpg.inp_embed); int32_t pos4[4] = {kv_pos, kv_pos, kv_pos, 0}; ggml_backend_tensor_set_async(backend, cpg.positions, pos4, 0, sizeof(pos4)); - std::vector row_vals((size_t)w.n_head_kv, (int64_t)kv_row); + const int64_t write_row = (kv_slot >= 0) + ? (int64_t)kv_row + : (int64_t)kv_row - (int64_t)cpg.kv_write_row_base; + if (write_row < 0 || (kv_slot < 0 && write_row >= 256)) { + std::fprintf(stderr, + "[pipelined] invalid bucket-local KV row (pos=%d base=%d row=%lld)\n", + kv_pos, cpg.kv_write_row_base, (long long)write_row); + return false; + } + std::vector row_vals((size_t)w.n_head_kv, write_row); ggml_backend_tensor_set_async(backend, cpg.kv_write_rows, row_vals.data(), 0, sizeof(int64_t) * row_vals.size()); diff --git a/server/src/qwen35moe/qwen35moe_pipelined_decode.h b/server/src/qwen35moe/qwen35moe_pipelined_decode.h index 4045f727d..9fa3c1c3e 100644 --- a/server/src/qwen35moe/qwen35moe_pipelined_decode.h +++ b/server/src/qwen35moe/qwen35moe_pipelined_decode.h @@ -38,6 +38,8 @@ struct CachedPrefnGraph { ggml_tensor * positions = nullptr; // [4] I32 input (attn layers only) ggml_tensor * kv_write_rows = nullptr; // [1, n_head_kv] I64 input (attn layers only) int kv_win = 0; // FA span baked into the graph (attn layers; 256-aligned) + int kv_write_row_base = 0; // non-pooled local-row base for attn set_rows + bool kv_write_pooled = false; // true when rows are KVFlash physical slots ggml_tensor * ffn_post = nullptr; // output: post-norm hidden state ggml_tensor * ffn_residual = nullptr; // output: pre-FFN residual ggml_tensor * moe_selected = nullptr; // output: selected expert IDs @@ -57,12 +59,16 @@ struct CachedPrefnGraph { moe_selected = o.moe_selected; moe_weights = o.moe_weights; positions = o.positions; kv_write_rows = o.kv_write_rows; kv_win = o.kv_win; + kv_write_row_base = o.kv_write_row_base; + kv_write_pooled = o.kv_write_pooled; o.ctx = nullptr; o.gf = nullptr; o.alloc = nullptr; o.inp_embed = nullptr; o.ffn_post = nullptr; o.ffn_residual = nullptr; o.moe_selected = nullptr; o.moe_weights = nullptr; o.positions = nullptr; o.kv_write_rows = nullptr; o.kv_win = 0; + o.kv_write_row_base = 0; + o.kv_write_pooled = false; } return *this; } diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 8d4bdab24..bf7ed0648 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -2461,6 +2461,7 @@ void HttpServer::worker_loop() { int cache_slot = full_cache_hit_slot; int prefix_len = full_cache_hit_len; bool using_restore = (cache_slot >= 0); + bool inline_restore_hit = false; if (!using_restore) { auto [full_slot, full_len] = prefix_cache_.lookup_full(req.prompt_tokens); if (full_slot >= 0) { @@ -2477,10 +2478,11 @@ void HttpServer::worker_loop() { // Inline prefix cache: check for cached turn-boundary KV state if no // exact full-prompt cache was available. if (!using_restore) { - auto [inline_slot, inline_len] = prefix_cache_.lookup(effective_prompt); - cache_slot = inline_slot; - prefix_len = inline_len; + auto inline_hit = prefix_cache_.lookup(effective_prompt); + cache_slot = inline_hit.slot; + prefix_len = inline_hit.snapshot_len; using_restore = (cache_slot >= 0); + inline_restore_hit = using_restore; } // Disk prefix cache: try disk if memory missed. @@ -2636,14 +2638,19 @@ void HttpServer::worker_loop() { // pointless snapshot copy-in. if (using_restore) { const int snap_len = backend_.snapshot_cur_pos(cache_slot); - if (snap_len > (int)effective_prompt.size()) { + if (snap_len <= 0 || snap_len > (int)effective_prompt.size()) { std::fprintf(stderr, - "[pc] slot=%d snapshot pos=%d > prompt=%zu — treating as miss\n", + "[pc] slot=%d invalid snapshot pos=%d prompt=%zu — treating as miss\n", cache_slot, snap_len, effective_prompt.size()); cache_slot = -1; prefix_len = 0; using_restore = false; disk_hit = false; + inline_restore_hit = false; + } else { + // Always account and restore from physical backend truth. Inline + // entries may have a newer logical key_len than snapshot_len. + prefix_len = snap_len; } } @@ -2711,7 +2718,11 @@ void HttpServer::worker_loop() { snap_cut = prepared.second; } bool snap_prepared = (snap_slot >= 0); + bool snap_slot_preexisting = false; + int snap_slot_old_pos = 0; if (snap_prepared) { + snap_slot_preexisting = backend_.snapshot_used(snap_slot); + snap_slot_old_pos = backend_.snapshot_cur_pos(snap_slot); gen_req.snap_slot = snap_slot; gen_req.snap_pos = snap_cut; } @@ -2914,21 +2925,57 @@ void HttpServer::worker_loop() { // Confirm or abort the inline snapshot. if (snap_prepared) { - if (completion_tokens > 0 && visible_output_seen && !client_disconnected && - backend_.snapshot_used(snap_slot)) { - prefix_cache_.confirm_inline_snap(snap_slot, snap_cut, effective_prompt); - // Track for shutdown save. - slot_tokens_[snap_slot] = std::vector( - effective_prompt.begin(), effective_prompt.begin() + snap_cut); - // Save to disk cache if threshold met. - if (!disk_cache_.disabled()) { - disk_cache_.learn_layout(snap_slot); - if (disk_policy.mode == DiskPrefixCacheMode::Full) { - disk_cache_.save(snap_slot, effective_prompt); + const bool can_commit_inline = + completion_tokens > 0 && visible_output_seen && !client_disconnected; + auto alias_to_inline_restore = [&]() { + if (!can_commit_inline || !inline_restore_hit || cache_slot < 0 || + !backend_.snapshot_used(cache_slot)) { + return; + } + const int alias_pos = backend_.snapshot_cur_pos(cache_slot); + if (alias_pos > 0 && alias_pos <= snap_cut) { + prefix_cache_.alias_inline_snap(cache_slot, snap_cut, alias_pos, + effective_prompt); + } + }; + + if (can_commit_inline && backend_.snapshot_used(snap_slot)) { + const int saved_pos = backend_.snapshot_cur_pos(snap_slot); + if (saved_pos <= 0 || saved_pos > snap_cut) { + prefix_cache_.abort_inline_snap(snap_slot); + alias_to_inline_restore(); + } else { + const bool refreshed_snapshot = + !snap_slot_preexisting || saved_pos != snap_slot_old_pos; + const bool reused_restore_slot = + inline_restore_hit && snap_slot == cache_slot; + if (!refreshed_snapshot && !reused_restore_slot) { + prefix_cache_.abort_inline_snap(snap_slot); + alias_to_inline_restore(); + } else if (reused_restore_slot && saved_pos < snap_cut) { + prefix_cache_.alias_inline_snap(snap_slot, snap_cut, saved_pos, + effective_prompt); + } else { + prefix_cache_.confirm_inline_snap(snap_slot, snap_cut, saved_pos, + effective_prompt); + // Track the exact physical snapshot prefix for shutdown + // saves. Shorter pooled snapshots are published only + // under their saved length, never as the later boundary. + slot_tokens_[snap_slot] = std::vector( + effective_prompt.begin(), effective_prompt.begin() + saved_pos); + // Do not persist alias-style entries as if the slot + // physically contained the longer prepared boundary. + if (!disk_cache_.disabled() && saved_pos == snap_cut) { + disk_cache_.learn_layout(snap_slot); + if (disk_policy.mode == DiskPrefixCacheMode::Full) { + disk_cache_.save(snap_slot, effective_prompt); + } + } } } } else { prefix_cache_.abort_inline_snap(snap_slot); + alias_to_inline_restore(); } } diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 420733672..022efd618 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -7,6 +7,7 @@ #include #include #include +#include namespace dflash::common { @@ -115,7 +116,10 @@ std::vector find_all_boundaries(const std::vector & ids, } } found: - if (next_match < 0) break; + if (next_match < 0) { + cursor = end_idx + 1; + continue; + } int boundary = next_match + next_len; out.push_back(boundary); cursor = boundary; @@ -192,6 +196,17 @@ PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) std::fprintf(stderr, "[pc] enabled: cap=%d family=%s\n", cap_, markers_.family.c_str()); } +PrefixCache::PrefixCache(int cap, ChatMarkers markers) + : cap_(std::min(cap, MAX_SLOTS)), markers_(std::move(markers)) +{ + if (cap_ <= 0) { + disabled_ = true; + cap_ = 0; + return; + } + disabled_ = false; +} + // ── LRU helpers ───────────────────────────────────────────────────────── int PrefixCache::find_entry(const PrefixHash & h) const { @@ -208,6 +223,110 @@ void PrefixCache::move_to_end(int idx) { entries_.push_back(std::move(e)); } +void PrefixCache::erase_inline_at(int idx) { + if (idx < 0 || idx >= (int)entries_.size()) return; + entries_.erase(entries_.begin() + idx); + publish_inline_counts(); +} + +void PrefixCache::erase_inline_slot(int slot) { + for (int i = (int)entries_.size() - 1; i >= 0; --i) { + if (entries_[(size_t)i].slot == slot) { + erase_inline_at(i); + } + } +} + +void PrefixCache::evict_pending_inline() { + if (!has_pending_evict_) return; + erase_inline_slot(pending_evict_slot_); + pending_evict_slot_ = -1; + has_pending_evict_ = false; +} + +bool PrefixCache::inline_slot_in_use(int slot) const { + for (const auto & e : entries_) { + if (e.slot == slot) return true; + } + return false; +} + +int PrefixCache::count_inline_slots() const { + bool seen[MAX_SLOTS] = {}; + int count = 0; + for (const auto & e : entries_) { + if (e.slot < 0 || e.slot >= MAX_SLOTS || seen[e.slot]) continue; + seen[e.slot] = true; + count++; + } + return count; +} + +int PrefixCache::select_inline_evict_slot() const { + if (entries_.empty()) return 0; + + // Eviction frees a physical slot, not one logical cache key. Aliases that + // share a slot must be considered as a group; otherwise an alias leaf can + // evict the shared slot and drop the ancestor entry the policy meant to keep. + for (const auto & candidate : entries_) { + const int slot = candidate.slot; + bool protects_other_slot = false; + for (const auto & own : entries_) { + if (own.slot != slot) continue; + for (const auto & other : entries_) { + if (other.slot == slot) continue; + if (is_strict_prefix(own.ids, other.ids)) { + protects_other_slot = true; + break; + } + } + if (protects_other_slot) break; + } + if (!protects_other_slot) return slot; + } + + return entries_.front().slot; +} + +void PrefixCache::publish_inline_counts() { + entries_size_count_.store((int64_t)entries_.size(), std::memory_order_relaxed); + inline_slot_count_.store((int64_t)count_inline_slots(), std::memory_order_relaxed); +} + +void PrefixCache::insert_inline_entry(int slot, int target_cut, int snapshot_len, + const std::vector & prompt_ids, + bool replace_slot_entries) { + if (slot < 0 || target_cut <= 0 || snapshot_len <= 0 || + target_cut > (int)prompt_ids.size() || snapshot_len > target_cut) { + std::fprintf(stderr, + "[pc] ignoring invalid inline entry slot=%d key_len=%d snapshot_len=%d prompt=%zu\n", + slot, target_cut, snapshot_len, prompt_ids.size()); + return; + } + + auto key = hash_prefix(prompt_ids.data(), target_cut); + int existing = find_entry(key); + if (existing >= 0) { + erase_inline_at(existing); + } + + if (replace_slot_entries) { + // A new physical snapshot changes this slot's KV contents. Drop any + // older logical aliases to the slot before publishing the replacement. + for (int i = (int)entries_.size() - 1; i >= 0; --i) { + if (entries_[(size_t)i].slot == slot) { + std::fprintf(stderr, + "[pc] dropping stale entry for reused slot=%d\n", slot); + erase_inline_at(i); + } + } + } + + std::vector ids(prompt_ids.begin(), prompt_ids.begin() + target_cut); + entries_.push_back({key, slot, snapshot_len, std::move(ids)}); + publish_inline_counts(); +} + int PrefixCache::find_full_entry(const PrefixHash & h) const { for (int i = 0; i < (int)full_entries_.size(); i++) { if (full_entries_[i].hash == h) return i; @@ -224,30 +343,36 @@ void PrefixCache::move_full_to_end(int idx) { // ── Inline prefix cache ───────────────────────────────────────────────── -std::pair PrefixCache::lookup(const std::vector & prompt_ids) { - if (disabled_) return {-1, 0}; - - auto boundaries = find_all_boundaries(prompt_ids, markers_); - int best_slot = -1, best_len = 0; - - for (int cut : boundaries) { - auto key = hash_prefix(prompt_ids.data(), cut); - int idx = find_entry(key); - if (idx >= 0) { - if (cut > best_len) { - best_slot = entries_[idx].slot; - best_len = cut; - } - move_to_end(idx); +PrefixCache::InlineLookup PrefixCache::lookup(const std::vector & prompt_ids) { + if (disabled_) return {}; + + InlineLookup best; + int best_idx = -1; + + // Entries are stored only for exact physical snapshot lengths. Most are + // chat-boundary prefixes, but KVFlash may save at the previous pool chunk + // boundary. Restoring that exact shorter prefix and prefilling the suffix is + // still correct, so lookup must consider stored prefixes directly. + for (int idx = 0; idx < (int)entries_.size(); ++idx) { + const auto & e = entries_[(size_t)idx]; + const int key_len = (int)e.ids.size(); + if (key_len <= best.key_len || key_len > (int)prompt_ids.size()) continue; + if (std::equal(e.ids.begin(), e.ids.end(), prompt_ids.begin())) { + best.slot = e.slot; + best.key_len = key_len; + best.snapshot_len = e.snapshot_len; + best_idx = idx; } } - if (best_slot >= 0) { + if (best.slot >= 0) { + move_to_end(best_idx); lifetime_hits_.fetch_add(1, std::memory_order_relaxed); - std::fprintf(stderr, "[pc] lookup hit slot=%d prefix_len=%d (of %zu total)\n", - best_slot, best_len, prompt_ids.size()); + std::fprintf(stderr, + "[pc] lookup hit slot=%d key_len=%d snapshot_len=%d (of %zu total)\n", + best.slot, best.key_len, best.snapshot_len, prompt_ids.size()); } - return {best_slot, best_len}; + return best; } std::pair PrefixCache::prepare_inline_snap( @@ -257,96 +382,116 @@ std::pair PrefixCache::prepare_inline_snap( auto candidates = find_all_boundaries(prompt_ids, markers_); if (candidates.empty()) return {-1, 0}; - // Best cache point: second-to-last boundary (last completed assistant turn). - int target_cut = candidates.size() >= 2 - ? candidates[candidates.size() - 2] - : candidates.back(); + // Snapshot the newest completed boundary so replay-style conversations keep + // growing the cached prefix instead of reusing a stale ancestor forever. + int target_cut = candidates.back(); auto key = hash_prefix(prompt_ids.data(), target_cut); if (find_entry(key) >= 0) return {-1, 0}; // already cached - int slot; - if ((int)entries_.size() >= cap_) { - // At capacity — reserve a slot without evicting yet. Prefix-aware: prefer - // the oldest leaf so shared ancestor prefixes (reused by later branches) - // stay resident. entries_ is already in LRU order (front = oldest). - std::vector *> ids_lru; - ids_lru.reserve(entries_.size()); - for (const auto & e : entries_) ids_lru.push_back(&e.ids); - int victim = select_inline_evict_victim(ids_lru); - pending_evict_key_ = entries_[victim].hash; + int slot = -1; + if (count_inline_slots() >= cap_) { + // At physical-slot capacity: reserve an entire slot group without + // evicting yet. The group selector protects ancestors that are still + // shared by entries on other physical slots. + slot = select_inline_evict_slot(); + pending_evict_slot_ = slot; has_pending_evict_ = true; - slot = entries_[victim].slot; - if (victim != 0) { - std::fprintf(stderr, - "[pc] prefix-aware evict: victim idx=%d (len=%zu) kept oldest " - "ancestor (len=%zu)\n", - victim, entries_[victim].ids.size(), entries_.front().ids.size()); - } + std::fprintf(stderr, "[pc] prefix-aware evict: reserved slot=%d\n", slot); } else { - slot = next_slot_; - next_slot_ = (next_slot_ + 1) % cap_; + for (int tries = 0; tries < cap_; ++tries) { + int candidate = next_slot_; + next_slot_ = (next_slot_ + 1) % cap_; + if (!inline_slot_in_use(candidate)) { + slot = candidate; + break; + } + } + if (slot < 0) { + // Defensive fallback for inconsistent accounting. + slot = select_inline_evict_slot(); + pending_evict_slot_ = slot; + has_pending_evict_ = true; + } else { + pending_evict_slot_ = -1; + has_pending_evict_ = false; + } + } + if (!has_pending_evict_) { + pending_evict_slot_ = -1; has_pending_evict_ = false; } return {slot, target_cut}; } -void PrefixCache::confirm_inline_snap(int slot, int target_cut, +void PrefixCache::confirm_inline_snap(int slot, int target_cut, int snapshot_len, const std::vector & prompt_ids) { if (disabled_) return; + if (slot < 0 || target_cut <= 0 || snapshot_len <= 0 || + target_cut > (int)prompt_ids.size() || + snapshot_len > target_cut || snapshot_len > (int)prompt_ids.size()) { + std::fprintf(stderr, + "[pc] refusing inline-snap slot=%d key_len=%d snapshot_len=%d prompt=%zu\n", + slot, target_cut, snapshot_len, prompt_ids.size()); + abort_inline_snap(slot); + return; + } // Evict the reserved entry (if any). - if (has_pending_evict_) { - int idx = find_entry(pending_evict_key_); - if (idx >= 0) { - entries_.erase(entries_.begin() + idx); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - } - has_pending_evict_ = false; + evict_pending_inline(); + + if (snapshot_len == target_cut) { + insert_inline_entry(slot, target_cut, snapshot_len, prompt_ids, + /*replace_slot_entries=*/true); + std::fprintf(stderr, + "[pc] inline-snap committed slot=%d key_len=%d snapshot_len=%d\n", + slot, target_cut, snapshot_len); + } else { + insert_inline_entry(slot, snapshot_len, snapshot_len, prompt_ids, + /*replace_slot_entries=*/true); + std::fprintf(stderr, + "[pc] inline-snap committed slot=%d key_len=%d snapshot_len=%d " + "(requested_key_len=%d)\n", + slot, snapshot_len, snapshot_len, target_cut); } +} - // The new snapshot replaces whatever this slot previously held. Drop any - // other entries still pointing at the slot: their hashes describe a - // different (or shorter) token stream than the new snapshot, and a later - // restore through them would attach mismatched KV. Stale entries arise - // when an aborted snap burns a round-robin next_slot_ step and a later - // confirm wraps onto a slot with a live entry (PR #370 repro). - for (int i = (int)entries_.size() - 1; i >= 0; --i) { - if (entries_[(size_t)i].slot == slot) { - std::fprintf(stderr, - "[pc] dropping stale entry for reused slot=%d\n", slot); - entries_.erase(entries_.begin() + i); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - } - } +void PrefixCache::alias_inline_snap(int slot, int target_cut, int snapshot_len, + const std::vector & prompt_ids) { + if (disabled_) return; - auto key = hash_prefix(prompt_ids.data(), target_cut); - std::vector ids(prompt_ids.begin(), prompt_ids.begin() + target_cut); - entries_.push_back({key, slot, std::move(ids)}); - entries_size_count_.fetch_add(1, std::memory_order_relaxed); - std::fprintf(stderr, "[pc] inline-snap committed slot=%d prefix_len=%d\n", - slot, target_cut); + // A failed prepared snap may have reserved an eviction victim. Release that + // reservation. Do not publish the longer logical key for a shorter physical + // snapshot: restore must materialize exactly the key length by construction. + evict_pending_inline(); + if (slot >= 0 && snapshot_len > 0 && snapshot_len <= target_cut && + snapshot_len <= (int)prompt_ids.size()) { + insert_inline_entry(slot, snapshot_len, snapshot_len, prompt_ids, + /*replace_slot_entries=*/false); + std::fprintf(stderr, + "[pc] inline-snap alias committed slot=%d key_len=%d snapshot_len=%d " + "(requested_key_len=%d)\n", + slot, snapshot_len, snapshot_len, target_cut); + } else { + std::fprintf(stderr, + "[pc] inline-snap alias skipped slot=%d key_len=%d snapshot_len=%d\n", + slot, target_cut, snapshot_len); + } } void PrefixCache::abort_inline_snap(int /*slot*/) { if (disabled_) return; - if (has_pending_evict_) { - int idx = find_entry(pending_evict_key_); - if (idx >= 0) { - entries_.erase(entries_.begin() + idx); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - } - has_pending_evict_ = false; - } + evict_pending_inline(); } void PrefixCache::mark_all_cleared() { if (disabled_) return; int n = (int)entries_.size(); entries_.clear(); - entries_size_count_.store(0, std::memory_order_relaxed); + publish_inline_counts(); next_slot_ = 0; + pending_evict_slot_ = -1; has_pending_evict_ = false; std::fprintf(stderr, "[pc] all-cleared — dropped %d LRU entries\n", n); } @@ -465,7 +610,7 @@ void PrefixCache::abort_full_snap(int /*slot*/) { PrefixCache::InlineStats PrefixCache::stats() const { if (disabled_) return {0, 0, 0}; return {cap_, - (int)entries_size_count_.load(std::memory_order_relaxed), + (int)inline_slot_count_.load(std::memory_order_relaxed), lifetime_hits_.load(std::memory_order_relaxed)}; } diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 675df911c..4e48b0318 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -76,6 +76,7 @@ class PrefixCache { // cap = number of prefix-cache slots (0 disables). PrefixCache(int cap, const Tokenizer & tokenizer); + PrefixCache(int cap, ChatMarkers markers); bool disabled() const { return disabled_; } @@ -84,16 +85,32 @@ class PrefixCache { // ── Inline prefix cache ───────────────────────────────────────── - // Look up the longest cached prefix. Returns (slot, prefix_len) or (-1, 0). - std::pair lookup(const std::vector & prompt_ids); + struct InlineLookup { + int slot = -1; + // Logical prompt boundary that matched the cache key. + int key_len = 0; + // Physical KV length saved in the backend snapshot. Inline entries are + // keyed by the exact saved prefix, so key_len == snapshot_len even when + // the originally requested chat boundary was slightly later. + int snapshot_len = 0; + }; + + // Look up the longest cached prefix. + InlineLookup lookup(const std::vector & prompt_ids); // Prepare an inline snapshot. Returns (slot, target_cut) or (-1, 0). std::pair prepare_inline_snap(const std::vector & prompt_ids); // Confirm after daemon successfully saved the snapshot. - void confirm_inline_snap(int slot, int target_cut, + void confirm_inline_snap(int slot, int target_cut, int snapshot_len, const std::vector & prompt_ids); + // Release a prepared inline snapshot without publishing the longer logical + // boundary. If the shorter physical snapshot is itself a valid prompt + // prefix, it may be published under that exact shorter key. + void alias_inline_snap(int slot, int target_cut, int snapshot_len, + const std::vector & prompt_ids); + // Abort if the snapshot failed. void abort_inline_snap(int slot); @@ -153,11 +170,12 @@ class PrefixCache { struct LruEntry { PrefixHash hash; int slot; + int snapshot_len; std::vector ids; // prefix tokens [0, target_cut) for prefix-aware eviction }; std::vector entries_; int next_slot_ = 0; - PrefixHash pending_evict_key_{}; + int pending_evict_slot_ = -1; bool has_pending_evict_ = false; // Full-cache state @@ -186,12 +204,23 @@ class PrefixCache { // data race per the C++ memory model. Bump these alongside every // push_back / erase / clear so the public introspection counters // stay well-defined. (Codex r1 P2 follow-up.) - std::atomic entries_size_count_{0}; // mirrors entries_.size() + std::atomic entries_size_count_{0}; // mirrors logical entries_.size() + std::atomic inline_slot_count_{0}; // mirrors distinct physical inline slots std::atomic full_entries_size_count_{0}; // mirrors full_entries_.size() // Helpers int find_entry(const PrefixHash & h) const; void move_to_end(int idx); + void erase_inline_at(int idx); + void erase_inline_slot(int slot); + void evict_pending_inline(); + bool inline_slot_in_use(int slot) const; + int count_inline_slots() const; + int select_inline_evict_slot() const; + void publish_inline_counts(); + void insert_inline_entry(int slot, int target_cut, int snapshot_len, + const std::vector & prompt_ids, + bool replace_slot_entries); int find_full_entry(const PrefixHash & h) const; void move_full_to_end(int idx); }; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 534d553d4..e00fd1e69 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -1353,6 +1353,190 @@ static void test_find_boundaries_empty() { TEST_ASSERT(bounds.empty()); } +static ChatMarkers synthetic_chat_markers() { + ChatMarkers markers; + markers.family = "synthetic"; + markers.sys_role_prefix = {1}; + markers.end_msg_seqs = {{2}}; + markers.next_role_starts = {{3}}; + return markers; +} + +static void test_find_boundaries_skips_unmatched_content_markers() { + auto markers = synthetic_chat_markers(); + std::vector ids = { + 1, 100, 2, 3, + 200, 2, 901, 902, 903, 904, 905, 201, 2, 3, 202 + }; + + auto bounds = find_all_boundaries(ids, markers); + TEST_ASSERT(bounds.size() == 2); + TEST_ASSERT(bounds[0] == 4); + TEST_ASSERT(bounds[1] == 14); +} + +static void test_prefix_cache_prepares_newest_boundary() { + auto markers = synthetic_chat_markers(); + PrefixCache cache(2, markers); + std::vector ids = {1, 100, 2, 3, 200, 2, 3, 300}; + + auto prep = cache.prepare_inline_snap(ids); + TEST_ASSERT(prep.first == 0); + TEST_ASSERT(prep.second == 7); +} + +static void test_prefix_cache_lookup_grows_chain_from_exact_prefix() { + auto markers = synthetic_chat_markers(); + PrefixCache cache(2, markers); + std::vector ids = {1, 100, 2, 3, 200, 2, 3, 300}; + + auto prep = cache.prepare_inline_snap(ids); + cache.confirm_inline_snap(prep.first, prep.second, prep.second, ids); + + std::vector appended = { + 1, 100, 2, 3, 200, 2, 3, 300, 2, 3, 400 + }; + auto hit = cache.lookup(appended); + TEST_ASSERT(hit.slot == prep.first); + TEST_ASSERT(hit.key_len == 7); + TEST_ASSERT(hit.snapshot_len == 7); + + std::vector sibling = { + 1, 100, 2, 3, 201, 2, 3, 300, 2, 3, 400 + }; + auto miss = cache.lookup(sibling); + TEST_ASSERT(miss.slot == -1); + TEST_ASSERT(miss.key_len == 0); + TEST_ASSERT(miss.snapshot_len == 0); +} + +static void test_prefix_cache_rejects_shorter_snapshot_alias() { + auto markers = synthetic_chat_markers(); + PrefixCache cache(4, markers); + std::vector base = {1, 100, 2, 3, 200}; + + auto base_prep = cache.prepare_inline_snap(base); + TEST_ASSERT(base_prep.first == 0); + TEST_ASSERT(base_prep.second == 4); + cache.confirm_inline_snap(base_prep.first, base_prep.second, + base_prep.second, base); + + std::vector longer = {1, 100, 2, 3, 200, 2, 3, 300}; + auto longer_prep = cache.prepare_inline_snap(longer); + TEST_ASSERT(longer_prep.first == 1); + TEST_ASSERT(longer_prep.second == 7); + + // Simulate a restored path where no new physical snapshot was materialized. + // The cache must not publish the longer logical boundary as a hit to the + // shorter physical snapshot: key_len must equal snapshot_len by construction. + cache.abort_inline_snap(longer_prep.first); + cache.alias_inline_snap(base_prep.first, longer_prep.second, + base_prep.second, longer); + + std::vector appended = { + 1, 100, 2, 3, 200, 2, 3, 300, 2, 3, 400 + }; + auto hit = cache.lookup(appended); + TEST_ASSERT(hit.slot == base_prep.first); + TEST_ASSERT(hit.key_len == base_prep.second); + TEST_ASSERT(hit.snapshot_len == 4); + TEST_ASSERT(hit.key_len == hit.snapshot_len); + + auto base_hit = cache.lookup(base); + TEST_ASSERT(base_hit.slot == base_prep.first); + TEST_ASSERT(base_hit.key_len == 4); + TEST_ASSERT(base_hit.snapshot_len == 4); +} + +static void test_prefix_cache_commits_shorter_snapshot_prefix() { + auto markers = synthetic_chat_markers(); + PrefixCache cache(2, markers); + std::vector ids = {1, 100, 2, 3, 200, 2, 3, 300}; + + auto prep = cache.prepare_inline_snap(ids); + TEST_ASSERT(prep.first == 0); + TEST_ASSERT(prep.second == 7); + // KVFlash snapshots can land at the previous pool chunk boundary rather + // than the requested chat boundary. Reuse is still correct if the cache key + // is the exact saved prefix, followed by suffix prefill. + cache.confirm_inline_snap(prep.first, prep.second, 6, ids); + + std::vector appended = { + 1, 100, 2, 3, 200, 2, 3, 300, 2, 3, 400 + }; + auto hit = cache.lookup(appended); + TEST_ASSERT(hit.slot == prep.first); + TEST_ASSERT(hit.key_len == 6); + TEST_ASSERT(hit.snapshot_len == 6); + TEST_ASSERT(hit.key_len == hit.snapshot_len); + + std::vector sibling = { + 1, 100, 2, 3, 201, 2, 3, 300, 2, 3, 400 + }; + auto miss = cache.lookup(sibling); + TEST_ASSERT(miss.slot == -1); + TEST_ASSERT(miss.key_len == 0); + TEST_ASSERT(miss.snapshot_len == 0); +} + +static void test_prefix_cache_rejects_oversized_confirm() { + auto markers = synthetic_chat_markers(); + PrefixCache cache(2, markers); + std::vector ids = {1, 100, 2, 3, 200, 2, 3, 300}; + + auto prep = cache.prepare_inline_snap(ids); + TEST_ASSERT(prep.first == 0); + TEST_ASSERT(prep.second == 7); + + cache.confirm_inline_snap(prep.first, prep.second, 8, ids); + + std::vector appended = { + 1, 100, 2, 3, 200, 2, 3, 300, 2, 3, 400 + }; + auto hit = cache.lookup(appended); + TEST_ASSERT(hit.slot == -1); + TEST_ASSERT(hit.key_len == 0); + TEST_ASSERT(hit.snapshot_len == 0); + + auto retry = cache.prepare_inline_snap(ids); + TEST_ASSERT(retry.first >= 0); + TEST_ASSERT(retry.second == prep.second); + cache.abort_inline_snap(retry.first); +} + +static void test_prefix_cache_alias_eviction_preserves_shared_ancestor_slot() { + auto markers = synthetic_chat_markers(); + PrefixCache cache(2, markers); + std::vector base = {1, 100, 2, 3, 200}; + + auto base_prep = cache.prepare_inline_snap(base); + TEST_ASSERT(base_prep.first == 0); + TEST_ASSERT(base_prep.second == 4); + cache.confirm_inline_snap(base_prep.first, base_prep.second, + base_prep.second, base); + + std::vector branch_a = {1, 100, 2, 3, 200, 2, 3, 300}; + cache.alias_inline_snap(base_prep.first, 7, 4, branch_a); + + std::vector branch_b = {1, 100, 2, 3, 201, 2, 3, 301}; + auto branch_b_prep = cache.prepare_inline_snap(branch_b); + TEST_ASSERT(branch_b_prep.first == 1); + TEST_ASSERT(branch_b_prep.second == 7); + cache.confirm_inline_snap(branch_b_prep.first, branch_b_prep.second, + branch_b_prep.second, branch_b); + + std::vector branch_c = {1, 100, 2, 3, 202, 2, 3, 302}; + auto branch_c_prep = cache.prepare_inline_snap(branch_c); + TEST_ASSERT(branch_c_prep.first == branch_b_prep.first); + TEST_ASSERT(branch_c_prep.first != base_prep.first); + cache.abort_inline_snap(branch_c_prep.first); + + auto base_hit = cache.lookup(base); + TEST_ASSERT(base_hit.slot == base_prep.first); + TEST_ASSERT(base_hit.key_len == 4); + TEST_ASSERT(base_hit.snapshot_len == 4); +} + // ── Prefix-aware eviction policy (model-free) ─────────────────────────── static void test_evict_empty_is_zero() { @@ -1974,6 +2158,98 @@ static void test_kvflash_pager_identity_sync_contract() { TEST_ASSERT(local.slot_of(127) == 127); } +static void test_kvflash_alloc_span_preserves_pending_suffix_chunks() { + KvFlashConfig cfg; + cfg.chunk_tokens = 4; + cfg.pool_tokens = 10 * cfg.chunk_tokens; + cfg.sink_chunks = 1; + cfg.tail_window_chunks = 4; + + KvFlashPager pager; + TEST_ASSERT(pager.attach(cfg, {}, {})); + TEST_ASSERT(pager.alloc_span(0, cfg.pool_tokens)); + + pager.score_hook = [](int c) { + return c < 10 ? 1000.0f - (float)c : (float)c; + }; + + const int suffix_start = cfg.pool_tokens; + const int suffix_tokens = + (10 - cfg.sink_chunks - cfg.tail_window_chunks) * cfg.chunk_tokens; + TEST_ASSERT(pager.alloc_span(suffix_start, suffix_tokens)); + for (int p = suffix_start; p < suffix_start + suffix_tokens; ++p) { + TEST_ASSERT_MSG(pager.slot_of(p) >= 0, "suffix position stayed resident"); + } +} + +static void test_kvflash_append_reservation_nested_release_keeps_live_span() { + KvFlashConfig cfg; + cfg.chunk_tokens = 4; + cfg.pool_tokens = 6 * cfg.chunk_tokens; + cfg.sink_chunks = 1; + cfg.tail_window_chunks = 0; + + KvFlashPager pager; + TEST_ASSERT(pager.attach(cfg, {}, {})); + TEST_ASSERT(pager.alloc_span(0, cfg.pool_tokens)); + + auto first = pager.reserve_append_span(cfg.pool_tokens, 2 * cfg.chunk_tokens); + TEST_ASSERT(first.ok()); + auto second = pager.reserve_append_span(cfg.pool_tokens + 2 * cfg.chunk_tokens, + 2 * cfg.chunk_tokens); + TEST_ASSERT(second.ok()); + + first = KvFlashPager::AppendReservation{}; + + pager.score_hook = [](int c) { + return (c == 8 || c == 9) ? -1000.0f : 1000.0f + (float)c; + }; + TEST_ASSERT(pager.reselect() >= 0); + for (int p = second.kv_start(); p < second.kv_start() + second.n_tokens(); ++p) { + TEST_ASSERT_MSG(pager.slot_of(p) >= 0, "live reservation stayed resident"); + } +} + +static void test_kvflash_pooled_suffix_prefill_plan_slot_mask() { + KvFlashConfig cfg; + cfg.chunk_tokens = 4; + cfg.pool_tokens = 10 * cfg.chunk_tokens; + cfg.sink_chunks = 1; + cfg.tail_window_chunks = 4; + + KvFlashPager pager; + TEST_ASSERT(pager.attach(cfg, {}, {})); + TEST_ASSERT(pager.alloc_span(0, cfg.pool_tokens)); + + const int kv_start = cfg.pool_tokens; + const int n_tok = 12; + const int n_head_kv = 3; + const int mask_kv_dim = cfg.pool_tokens + 8; + auto plan = kvflash_prepare_pooled_suffix_prefill( + pager, kv_start, n_tok, n_head_kv, + mask_kv_dim, 32); + TEST_ASSERT(plan.ok()); + + TEST_ASSERT(plan.rows.size() == (size_t)n_tok * n_head_kv); + for (int h = 0; h < n_head_kv; ++h) { + for (int i = 0; i < n_tok; ++i) { + TEST_ASSERT(plan.rows[(size_t)h * n_tok + i] == pager.slot_of(kv_start + i)); + } + } + + std::vector slot_pos((size_t)cfg.pool_tokens, -1); + pager.fill_slot_pos(slot_pos.data()); + constexpr uint16_t F16_ZERO = 0x0000, F16_NEG_INF = 0xFC00; + for (int q = 0; q < n_tok; ++q) { + const int abs_q = kv_start + q; + for (int s = 0; s < cfg.pool_tokens; ++s) { + const uint16_t got = plan.mask[(size_t)q * mask_kv_dim + s]; + const bool want_open = slot_pos[(size_t)s] >= 0 && slot_pos[(size_t)s] <= abs_q; + TEST_ASSERT(got == (want_open ? F16_ZERO : F16_NEG_INF)); + } + } +} + static void test_layer_split_kvflash_history_contract() { std::vector history; layer_split_kvflash_sync_history(history, {1, 2, 3}, 0); @@ -4454,6 +4730,13 @@ int main() { RUN_TEST(test_hash_prefix_different_lengths); RUN_TEST(test_hash_prefix_empty); RUN_TEST(test_find_boundaries_empty); + RUN_TEST(test_find_boundaries_skips_unmatched_content_markers); + RUN_TEST(test_prefix_cache_prepares_newest_boundary); + RUN_TEST(test_prefix_cache_lookup_grows_chain_from_exact_prefix); + RUN_TEST(test_prefix_cache_rejects_shorter_snapshot_alias); + RUN_TEST(test_prefix_cache_commits_shorter_snapshot_prefix); + RUN_TEST(test_prefix_cache_rejects_oversized_confirm); + RUN_TEST(test_prefix_cache_alias_eviction_preserves_shared_ancestor_slot); RUN_TEST(test_evict_empty_is_zero); RUN_TEST(test_evict_single_is_zero); RUN_TEST(test_evict_chain_keeps_ancestors); @@ -4512,6 +4795,9 @@ int main() { RUN_TEST(test_target_shard_plan_mixed_backend_split); RUN_TEST(test_target_shard_plan_rejects_bad_local_backend); RUN_TEST(test_kvflash_pager_identity_sync_contract); + RUN_TEST(test_kvflash_alloc_span_preserves_pending_suffix_chunks); + RUN_TEST(test_kvflash_append_reservation_nested_release_keeps_live_span); + RUN_TEST(test_kvflash_pooled_suffix_prefill_plan_slot_mask); RUN_TEST(test_layer_split_kvflash_history_contract); RUN_TEST(test_backend_precision_cuda_sm_policy); RUN_TEST(test_backend_precision_hip_arch_policy);