diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 6b35df216..b9fba8b99 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -619,6 +619,12 @@ if(DFLASH27B_TESTS) target_include_directories(test_kvflash_placement PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) add_test(NAME kvflash_placement COMMAND test_kvflash_placement) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kvflash_pager.cpp") + add_executable(test_kvflash_pager test/test_kvflash_pager.cpp) + target_include_directories(test_kvflash_pager PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(test_kvflash_pager PRIVATE dflash_common) + add_test(NAME kvflash_pager COMMAND test_kvflash_pager) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_bandit_integration.cpp") add_executable(test_bandit_integration test/test_bandit_integration.cpp) target_include_directories(test_bandit_integration PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) @@ -796,18 +802,30 @@ if(DFLASH27B_TESTS) target_include_directories(test_kvflash_qk PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kvflash_pool_sizing.cpp") - # Pure unit test for kvflash_pool_from_env: ggml headers only, no ggml link, no GPU. + # Pure unit test for kvflash_pool_from_env: no ggml link, no GPU. + # kvflash_pager.h includes ggml.h for tensor types used by the pager + # (not by pool_from_env itself), so the ggml include path is required + # even though no ggml symbols are linked. add_executable(test_kvflash_pool_sizing test/test_kvflash_pool_sizing.cpp) target_include_directories(test_kvflash_pool_sizing PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS} - ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include - ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) add_test(NAME kvflash_pool_sizing COMMAND test_kvflash_pool_sizing) endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_restore_delta.cpp") add_executable(test_restore_delta test/test_restore_delta.cpp) target_include_directories(test_restore_delta PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_warm_prefill.cpp") + add_executable(test_warm_prefill test/test_warm_prefill.cpp) + target_include_directories(test_warm_prefill PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + add_test(NAME warm_prefill COMMAND test_warm_prefill) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_spec_gate.cpp") + add_executable(test_spec_gate test/test_spec_gate.cpp) + target_include_directories(test_spec_gate PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + add_test(NAME spec_gate COMMAND test_spec_gate) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash.cpp") add_executable(test_dflash test/test_dflash.cpp) target_include_directories(test_dflash PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index d61707d96..8320fb897 100644 --- a/server/src/common/kvflash_pager.h +++ b/server/src/common/kvflash_pager.h @@ -161,6 +161,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 +184,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 +209,55 @@ 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) half-open. + // Maps to chunk range [c_lo, c_hi] inclusive where c_hi covers the last + // token in the range (tok_hi - 1), not tok_hi itself, so an exact boundary + // at tok_hi does not pin the next chunk beyond the range. + // 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 - 1) / 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. @@ -397,7 +448,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,6 +470,101 @@ class KvFlashPager { return events; } + // Snapshot all resident+paged-out chunks into a flat byte blob. + // Layout: 8-byte magic, header fields (6×uint32), then for each logical + // chunk c in [0, n_chunks): chunk_bytes_ bytes in for_each_segment order. + std::vector serialize() const { + static constexpr uint64_t kMagic = 0x4b56464c41534800ULL; // "KVFLASH\0" + const int nc = (int)chunks_.size(); + const size_t hdr = sizeof(uint64_t) + 6 * sizeof(uint32_t); + std::vector out; + out.resize(hdr + (size_t)nc * chunk_bytes_, 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_); + 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. + uint8_t * q = dst; + for_each_segment(st.block, [&](ggml_tensor * t, size_t off, size_t seg) { + ggml_backend_tensor_get(t, q, off, seg); + q += seg; + }); + } else if (st.on_host) { + // Host-backed: copy verbatim. host_data is a raw pinned pointer + // under async DMA, a std::vector otherwise. +#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. + } + return out; + } + + // Restore state from a blob produced by serialize(). Returns false on + // header mismatch (layout drift guard). Callers must rebuild slot masks. + // + // Ordering: pre-size chunks_ so each entry exists before slot_for() runs. + // We 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) itself — + // no extra copy_chunk needed (that would double-write). + bool deserialize(const uint8_t * data, size_t n) { + static constexpr uint64_t kMagic = 0x4b56464c41534800ULL; + const size_t hdr = sizeof(uint64_t) + 6 * 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(); + if (ct != cfg_.chunk_tokens || nhkv != n_head_kv_ || + kseg != k_seg_bytes_ || vseg != v_seg_bytes_ || cb != chunk_bytes_) { + return false; + } + if (n < hdr + (size_t)nc * chunk_bytes_) return false; + reset(); + 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. + // host_data is a raw pinned pointer under async DMA, a vector otherwise. +#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; + // slot_for assigns a block and auto-recalls via copy_chunk(to_host=false). + slot_for((int64_t)c * cfg_.chunk_tokens); + } +#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; + } + private: struct ChunkState { int block = -1; // pool block index, -1 = not resident @@ -441,6 +588,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; } @@ -451,6 +599,24 @@ class KvFlashPager { return victim >= 0 && page_out(victim); } + // Walk every (tensor, head) segment for `block` in the fixed layout order + // (layer-major, K then V, head-minor). fn(tensor, byte_offset, seg_bytes). + // Used by serialize/deserialize (synchronous); copy_chunk has its own + // async-DMA path below. + template + void for_each_segment(int block, Fn&& fn) const { + 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)block * cfg_.chunk_tokens * t->nb[1] + (size_t)h * t->nb[2]; + fn(t, off, seg); + } + } + } + } + // Move one chunk between pool slots and host backing. // Segment order is fixed (layer-major, K then V, head-minor). // When KVFLASH_HAS_ASYNC_DMA: transfers are issued on page_stream_ @@ -566,6 +732,7 @@ 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; diff --git a/server/src/common/kvflash_qk.h b/server/src/common/kvflash_qk.h index f84bd85b0..a7e3ef445 100644 --- a/server/src/common/kvflash_qk.h +++ b/server/src/common/kvflash_qk.h @@ -47,7 +47,10 @@ inline void kvflash_qk_chunk_scores( const float * query, const KvFlashQkDims & d, std::vector & out, - float missing_score = -2.0f) { + float missing_score = -2.0f, + const float * seeded = nullptr, + float seeded_sentinel = -std::numeric_limits::infinity(), + int seeded_n = -1) { const int group = d.n_q_heads / d.n_kv_heads; const int n_chunks = (int)pooled_keys.size(); out.assign((size_t)n_chunks, missing_score); @@ -83,6 +86,19 @@ inline void kvflash_qk_chunk_scores( } out[(size_t)c] = acc * inv_layers; // layer-MEAN (Phase-0 config) } + // Seeded fallback: for chunks with no pooled key, use the ledger score from + // a prior turn if it is not the sentinel (i.e. it was actually scored). + // seeded_n bounds the valid range of the seeded array; chunks beyond it + // (n_chunks > seeded array length) fall back to missing_score safely. + if (seeded) { + const int seeded_limit = (seeded_n >= 0) ? seeded_n : n_chunks; + for (int c = 0; c < n_chunks; c++) { + if (!pooled_keys[(size_t)c] && c < seeded_limit && + seeded[c] != seeded_sentinel) { + out[(size_t)c] = seeded[c]; + } + } + } } } // namespace dflash::common diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index e8bddebd2..16a2fc547 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -1066,6 +1066,27 @@ static bool eval_moe_hybrid_ffn_batched_core( if (cl >= 0) { cold_sel[i] = cl; cold_wts[i] = selected_weights[i]; fp_has_cold = true; } } } + // Dummy slots (wts==0) may alias a real hot expert's local ID per token → + // ids_to_sorted_host drops entries → ASSERT in slow ggml_mul_mat_id path. + for (int t = 0; t < n_tokens; ++t) { + const int base = t * n_used; + int32_t next = 0; + for (int s = 0; s < n_used; ++s) { + if (hot_wts[base + s] > 0.0f) continue; + // Bounded search: at most n_hot_init probes. If every ID in + // [0, n_hot_init) is already taken by another slot we break and + // keep `next` as-is (duplicate), which is safe — the zero-weight + // slot is ignored by ids_to_sorted_host anyway. + int tries = 0; + while (tries < n_hot_init && + [&]{ for (int k=0; k= n_hot_init) next = 0; + ++tries; + } + hot_sel[base + s] = next++; + if (next >= n_hot_init) next = 0; + } + } CachedHotBatchedGraph & hg = storage.hot_batched_mixed[n_tokens]; const bool hg_ok = (hg.valid() && hg.n_tokens == n_tokens) @@ -1145,6 +1166,23 @@ static bool eval_moe_hybrid_ffn_batched_core( } } } + // Dummy slots (wts==0) may alias a real hot expert's local ID per token → + // ids_to_sorted_host drops entries → ASSERT in slow ggml_mul_mat_id path. + for (int t = 0; t < n_tokens; ++t) { + const int base = t * n_used; + int32_t next = 0; + for (int s = 0; s < n_used; ++s) { + if (hot_wts[base + s] > 0.0f) continue; + int tries = 0; + while (tries < n_hot_init && + [&]{ for (int k=0; k= n_hot_init) next = 0; + ++tries; + } + hot_sel[base + s] = next++; + if (next >= n_hot_init) next = 0; + } + } // ── Step 2: Build and run hot GPU graph (includes shared expert always) ── std::vector hot_partial((size_t)n_embd * (size_t)n_tokens, 0.0f); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4e376950e..6d0ac62a9 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1020,15 +1020,38 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, // identity-style in the first place). const bool kvf_paged = kvflash_active() && kv_offset + prompt_len > kvflash_tokens_ - kvflash_pager_.chunk_tokens(); - if (kvf_paged && kv_offset != 0) { - std::fprintf(stderr, - "[kvflash] restored prefix (%d) + prompt (%d) exceeds pool %d; " - "pooled prefill requires a fresh request\n", - kv_offset, prompt_len, kvflash_tokens_); - set_last_error("kvflash: restore + pooled prefill unsupported"); - return -1; - } - if (kvf_paged) { + // restore-consume: kv_offset != 0 means the caller passed a suffix after + // deserializing the prefix KV; allow it in the pooled path without reset. + const bool kvf_restore_suffix = kvf_paged && kv_offset != 0; + if (kvf_restore_suffix) { + // Suffix pooled prefill: pager was seeded by deserialize(); just set + // ubatch size and skip reset. Log to distinguish from the cold path. + prefill_ubatch = kvflash_pager_.chunk_tokens(); + // Verify kv_offset is chunk-aligned (snapshot boundary contract). + // On misalignment, fall back to a full cold re-prefill rather than + // aborting the request — the KV restored from the snapshot is stale. + if (kv_offset % prefill_ubatch != 0) { + std::fprintf(stderr, + "[kvflash] restore-consume: kv_offset=%d not chunk-aligned " + "(chunk_tokens=%d) — falling back to full re-prefill\n", + kv_offset, prefill_ubatch); + kv_offset = 0; + 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 (fallback): %d tokens through a %d-token pool " + "(%d-token chunks, evicting)\n", + prompt_len, kvflash_tokens_, prefill_ubatch); + std::fflush(stdout); + } else { + std::printf("[kvflash] restore-consume suffix: offset=%d suffix=%d " + "pool=%d chunk=%d\n", + kv_offset, prompt_len, kvflash_tokens_, prefill_ubatch); + std::fflush(stdout); + } + } else if (kvf_paged) { prefill_ubatch = kvflash_pager_.chunk_tokens(); kvflash_pager_.reset(); if (kvflash_qk_policy_) { diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 742aaa329..5450c59e5 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -190,8 +190,24 @@ bool create_target_cache_partial(const TargetWeights & w, } } - constexpr int TARGET_FEAT_CAP_DEFAULT = 4096; - out.target_feat_cap = std::min(max_ctx, TARGET_FEAT_CAP_DEFAULT); + // Feature ring cap. The drafter needs the last `cap` target hidden states + // as context. A cap smaller than the context makes the ring WRAP and the + // drafter reads stale features → acceptance collapses (0.1% at 27K with the + // old 4096 default). Default to the full reserved context so the ring never + // wraps; lower via DFLASH_FEAT_RING_CAP to save VRAM (ring = cap*fc_in*2B). + int target_feat_cap_default = max_ctx; + if (const char * e = std::getenv("DFLASH_FEAT_RING_CAP")) { + char * endp = nullptr; + long v = std::strtol(e, &endp, 10); + if (endp == e || *endp != '\0' || v <= 0) { + std::fprintf(stderr, + "[dflash] DFLASH_FEAT_RING_CAP='%s' is invalid or non-positive; " + "using default max_ctx=%d\n", e, max_ctx); + } else { + target_feat_cap_default = (int)std::min(v, (long)max_ctx); + } + } + out.target_feat_cap = std::min(max_ctx, target_feat_cap_default); if (allocate_target_feat) { const int fc_in = w.n_capture_layers * w.n_embd; out.target_feat = ggml_new_tensor_2d(out.base_ctx, GGML_TYPE_BF16, fc_in, out.target_feat_cap); diff --git a/server/test/test_kvflash_pager.cpp b/server/test/test_kvflash_pager.cpp new file mode 100644 index 000000000..ec8b71c90 --- /dev/null +++ b/server/test/test_kvflash_pager.cpp @@ -0,0 +1,132 @@ +// Unit tests for KvFlashPager serialize/deserialize round-trip and critical- +// chunk pinning, on a CPU ggml backend (no CUDA). +#include "../src/common/kvflash_pager.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static void expect(bool cond, const char * msg) { + if (!cond) { std::fprintf(stderr, "FAIL: %s\n", msg); std::exit(1); } +} + +// Small synthetic cache: head_dim=4, pool=512 tok, n_head_kv=2, 1 layer, +// chunk_tokens=64 -> 8 blocks. sink=1 + tail=4 -> min_pool = (1+4+2)*64=448 <= 512. +struct Harness { + ggml_context * ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_backend_buffer_t buf = nullptr; + std::vector k, v; + + Harness(int head_dim, int pool, int n_head_kv, int n_layer) { + backend = ggml_backend_cpu_init(); + ggml_init_params ip{}; + ip.mem_size = (size_t)(n_layer * 2 + 8) * ggml_tensor_overhead(); + ip.no_alloc = true; + ctx = ggml_init(ip); + for (int l = 0; l < n_layer; l++) { + ggml_tensor * kt = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_dim, pool, n_head_kv); + ggml_tensor * vt = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_dim, pool, n_head_kv); + k.push_back(kt); v.push_back(vt); + } + buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + } + ~Harness() { + if (buf) ggml_backend_buffer_free(buf); + if (ctx) ggml_free(ctx); + if (backend) ggml_backend_free(backend); + } + KvFlashConfig cfg(int pool, int chunk) const { + KvFlashConfig c; c.pool_tokens = pool; c.chunk_tokens = chunk; + c.sink_chunks = 1; c.tail_window_chunks = 4; return c; + } +}; + +static void test_serialize_roundtrip() { + const int head_dim = 4, pool = 512, nkv = 2, nlayer = 1, chunk = 64; + Harness A(head_dim, pool, nkv, nlayer); + KvFlashPager pa; + expect(pa.attach(A.cfg(pool, chunk), A.k, A.v), "attach A"); + // Map the whole pool and stamp a recognizable ramp into the K/V buffers. + expect(pa.alloc_span(0, pool), "alloc_span fills pool A"); + const size_t kbytes = ggml_nbytes(A.k[0]); + std::vector ramp(kbytes); + for (size_t i = 0; i < kbytes; i++) ramp[i] = (uint8_t)(i * 31 + 7); + ggml_backend_tensor_set(A.k[0], ramp.data(), 0, kbytes); + ggml_backend_tensor_set(A.v[0], ramp.data(), 0, kbytes); + + std::vector blob = pa.serialize(); + expect(!blob.empty(), "serialize produces a blob"); + + // Deserialize into a fresh pager over fresh (zeroed) tensors. + Harness B(head_dim, pool, nkv, nlayer); + KvFlashPager pb; + expect(pb.attach(B.cfg(pool, chunk), B.k, B.v), "attach B"); + expect(pb.deserialize(blob.data(), blob.size()), "deserialize succeeds"); + + // The restored pool must contain the same KV bytes as the source. + std::vector kb(kbytes), vb(kbytes); + ggml_backend_tensor_get(B.k[0], kb.data(), 0, kbytes); + ggml_backend_tensor_get(B.v[0], vb.data(), 0, kbytes); + expect(kb == ramp, "K restored byte-identical after round-trip"); + expect(vb == ramp, "V restored byte-identical after round-trip"); + + // A blob from a mismatched layout must be rejected (header validation). + KvFlashConfig bad = B.cfg(pool, chunk); + blob[8] ^= 0xFF; // corrupt a header byte + KvFlashPager pc; + Harness C(head_dim, pool, nkv, nlayer); + expect(pc.attach(C.cfg(pool, chunk), C.k, C.v), "attach C"); + expect(!pc.deserialize(blob.data(), blob.size()), "corrupt-header blob rejected"); + (void)bad; + std::printf("ok: serialize/deserialize round-trip + header guard\n"); +} + +static void test_pinning() { + const int head_dim = 4, pool = 512, nkv = 2, nlayer = 1, chunk = 64; + Harness H(head_dim, pool, nkv, nlayer); + KvFlashPager p; + expect(p.attach(H.cfg(pool, chunk), H.k, H.v), "attach pin"); + + // No pins -> nothing pinned. + expect(!p.is_pinned(2), "pre: chunk 2 unpinned"); + + // Pin a single mid chunk (tokens 128..191 -> chunk 2). 8 blocks, sink1+tail4 + // +pinned1+2 = 8 <= 8 -> allowed. + p.pin_range(128, 191); + expect(p.is_pinned(2), "chunk 2 pinned"); + expect(!p.is_pinned(1), "chunk 1 not pinned"); + expect(!p.is_pinned(3), "chunk 3 not pinned"); + + // unpin clears it. + p.unpin_all(); + expect(!p.is_pinned(2), "unpin_all clears chunk 2"); + + // Deadlock guard: pinning the whole pool (8 chunks) would leave no evictable + // block (1+4+8+2 = 15 > 8) -> refused, nothing pinned. + p.pin_range(0, pool - 1); + expect(!p.is_pinned(0) && !p.is_pinned(4) && !p.is_pinned(7), + "over-pin refused by deadlock guard"); + + // reset() also clears pins. + p.pin_range(128, 191); + expect(p.is_pinned(2), "re-pin chunk 2"); + p.reset(); + expect(!p.is_pinned(2), "reset clears pins"); + std::printf("ok: pinning + deadlock guard + reset\n"); +} + +int main() { + test_serialize_roundtrip(); + test_pinning(); + std::printf("PASS: kvflash pager (serde + pinning)\n"); + return 0; +}