Skip to content

feat(laguna): pager-blob warm-restore survives eviction #451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
dusterbloom wants to merge 1 commit into
base: main
Choose a base branch
from
+426 −93
Draft

feat(laguna): pager-blob warm-restore survives eviction #451

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
260 changes: 245 additions & 15 deletions server/src/common/kvflash_pager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,12 @@

namespace dflash::common {

// Largest multiple of chunk_tokens that is <= pos.
// Used by the pooled-prefill snapshot path to align the save boundary.
inline int floor_to_chunk(int pos, int chunk_tokens) {
return (pos / chunk_tokens) * chunk_tokens;
}

struct KvFlashConfig {
int chunk_tokens = 64; // logical tokens per page
int pool_tokens = 0; // resident pool capacity (multiple of chunk_tokens)
Expand Down Expand Up @@ -161,6 +167,7 @@ class KvFlashPager {

// Drop all mappings and host backing (new request / cache reset).
// Cumulative stats are kept; the epoch advances so cached masks refill.
// Pins are cleared: the caller (backend) re-applies them after rebuild.
void reset() {
#ifdef KVFLASH_HAS_ASYNC_DMA
if (page_stream_) {
Expand All @@ -183,6 +190,7 @@ class KvFlashPager {
stats_.host_bytes = 0;
cur_chunk_ = 0;
epoch_++;
std::fill(pinned_.begin(), pinned_.end(), 0);
has_pending_page_in_ = false;
}

Expand All @@ -207,6 +215,53 @@ class KvFlashPager {
// Optional external relevance score; higher = keep. Falls back to LRU.
std::function<float(int /*chunk*/)> score_hook;

// ── Critical-chunk pinning ──────────────────────────────────────────
// Pinned chunks are never chosen as eviction victims (OR-ed on top of the
// sink/tail protections). Empty by default → byte-identical non-pin path.
// Pins are cleared by reset() and re-applied by the backend after each
// prefill/restore rebuild via apply_kvflash_pins().

// Pin logical token range [tok_lo, tok_hi) (inclusive of chunk boundaries).
// Maps to chunk range [tok_lo/chunk_tokens, tok_hi/chunk_tokens] inclusive.
// Best-effort deadlock guard: if (sink+tail+n_pinned+2) > n_blocks_ the
// span is refused with a one-line warning and the function returns without
// setting any pin.
void pin_range(int64_t tok_lo, int64_t tok_hi) {
if (!attached() || tok_lo > tok_hi) return;
const int c_lo = (int)(tok_lo / cfg_.chunk_tokens);
const int c_hi = (int)(tok_hi / cfg_.chunk_tokens);
if (c_lo > c_hi || c_lo < 0) return;
// Count currently-pinned chunks + the new ones.
int currently_pinned = 0;
for (int c = 0; c < (int)pinned_.size(); c++) {
if (pinned_[c]) currently_pinned++;
}
int new_pins = 0;
for (int c = c_lo; c <= c_hi; c++) {
if (c >= (int)pinned_.size() || !pinned_[c]) new_pins++;
}
// Deadlock guard: fixed protections + pinned + 2 (one evictable victim +
// one append-head) must fit in the pool.
if (cfg_.sink_chunks + cfg_.tail_window_chunks + currently_pinned + new_pins + 2 > n_blocks_) {
std::fprintf(stderr,
"[kvflash] pin_range [%lld,%lld] refused: "
"sink=%d tail=%d pinned=%d new=%d pool_blocks=%d — would deadlock eviction\n",
(long long)tok_lo, (long long)tok_hi,
cfg_.sink_chunks, cfg_.tail_window_chunks,
currently_pinned, new_pins, n_blocks_);
return;
}
if (c_hi + 1 > (int)pinned_.size()) pinned_.resize((size_t)c_hi + 1, 0);
for (int c = c_lo; c <= c_hi; c++) pinned_[(size_t)c] = 1;
}

bool is_pinned(int c) const {
return c >= 0 && c < (int)pinned_.size() && pinned_[(size_t)c];
}

// Clear all pins (also called by reset()).
void unpin_all() { std::fill(pinned_.begin(), pinned_.end(), 0); }

// Allocate slots for [kv_start, kv_start + n_tok) ahead of a forward
// step (evicting LRU/low-score chunks as needed). False — with a
// diagnostic — if the pool has no evictable block left.
Expand Down Expand Up @@ -323,20 +378,6 @@ class KvFlashPager {
}
return true;
}

// True iff every chunk intersecting [0, n_tok) is resident in its identity
// block (block_of(c) == c). Expresses "the logical prefix [0, n_tok) is a
// contiguous, identity-mapped, materialized span" — the exact precondition
// the non-paged tree-verify graph relies on. Stronger than is_identity(),
// which is also true for an empty / not-yet-materialized pager.
bool identity_prefix_covers(int n_tok) const {
if (n_tok <= 0) return true;
const int nc = (n_tok + cfg_.chunk_tokens - 1) / cfg_.chunk_tokens;
if (nc > (int)chunks_.size()) return false;
for (int c = 0; c < nc; c++)
if (chunks_[c].block != c) return false;
return true;
}
int block_of(int c) const {
return c < (int)chunks_.size() ? chunks_[c].block : -1;
}
Expand Down Expand Up @@ -397,7 +438,8 @@ class KvFlashPager {
const ChunkState & st = chunks_[c];
if (st.block < 0 && !st.on_host) continue; // never materialized
const bool prot = c < cfg_.sink_chunks ||
c > cur_chunk_ - 1 - cfg_.tail_window_chunks;
c > cur_chunk_ - 1 - cfg_.tail_window_chunks ||
is_pinned(c);
cands.push_back({c, prot ? 3.4e38f : score_hook(c)});
}
std::sort(cands.begin(), cands.end(),
Expand All @@ -418,11 +460,196 @@ class KvFlashPager {
return events;
}

// Snapshot resident+paged-out chunks into a flat byte blob (v1 format).
// max_chunks: if >= 0, serialize only chunks [0, max_chunks); the blob's
// nc header field encodes max_chunks so deserialize restores exactly that
// many chunks with the correct cur_pos. Pass -1 (default) to serialize
// all chunks.
//
// Layout (v1 / magic "KVFLASH1"):
// 8-byte magic
// 8×uint32 header: nc, chunk_tokens, n_head_kv, k_seg_bytes, v_seg_bytes,
// chunk_bytes, kv_k_type_u32, has_ledger(=1)
// nc × chunk_bytes_ — raw KV bytes (unchanged from v0)
// nc × 8 bytes — ledger section (has_ledger==1):
// uint8 was_resident (1 if chunk was in pool at serialize time)
// uint8 _pad[3]
// float score (qk_score; -INF if not set)
std::vector<uint8_t> serialize(int max_chunks = -1) const {
static constexpr uint64_t kMagic = 0x4b56464c41534831ULL; // "KVFLASH1"
const int nc = (max_chunks >= 0 && max_chunks < (int)chunks_.size())
? max_chunks
: (int)chunks_.size();
const size_t hdr = sizeof(uint64_t) + 8 * sizeof(uint32_t);
const size_t ledger_entry = 8; // uint8 was_resident + 3 pad + float score
const size_t total = hdr + (size_t)nc * chunk_bytes_ + (size_t)nc * ledger_entry;
std::vector<uint8_t> out;
out.resize(total, 0);
uint8_t * p = out.data();
std::memcpy(p, &kMagic, 8); p += 8;
auto w32 = [&](uint32_t v) { std::memcpy(p, &v, 4); p += 4; };
w32((uint32_t)nc);
w32((uint32_t)cfg_.chunk_tokens);
w32((uint32_t)n_head_kv_);
w32((uint32_t)k_seg_bytes_);
w32((uint32_t)v_seg_bytes_);
w32((uint32_t)chunk_bytes_);
// dtype guard: record the ggml_type enum for the K cache tensor
const uint32_t kv_k_type = (!attn_k_.empty() && attn_k_[0])
? (uint32_t)attn_k_[0]->type
: (uint32_t)GGML_TYPE_F16;
w32(kv_k_type);
w32(1u); // has_ledger = 1 always in v1 format
// KV bytes section
for (int c = 0; c < nc; ++c) {
uint8_t * dst = out.data() + hdr + (size_t)c * chunk_bytes_;
const ChunkState & st = chunks_[c];
if (st.block >= 0) {
// Resident: gather from pool tensors in fixed segment order
// (layer-major, K then V, head-minor) — matching copy_chunk.
uint8_t * q = dst;
for (size_t l = 0; l < attn_k_.size(); ++l) {
for (int kv = 0; kv < 2; ++kv) {
ggml_tensor * t = kv == 0 ? attn_k_[l] : attn_v_[l];
const size_t seg = kv == 0 ? k_seg_bytes_ : v_seg_bytes_;
for (int h = 0; h < n_head_kv_; ++h) {
const size_t off = (size_t)st.block * cfg_.chunk_tokens * t->nb[1]
+ (size_t)h * t->nb[2];
ggml_backend_tensor_get(t, q, off, seg);
q += seg;
}
}
}
} else if (st.on_host) {
// Host-backed: copy verbatim.
#ifdef KVFLASH_HAS_ASYNC_DMA
std::memcpy(dst, st.host_data, chunk_bytes_);
#else
std::memcpy(dst, st.host_data.data(), chunk_bytes_);
#endif
}
// else: never written — stays zero-filled from the resize above.
}
// Ledger section: was_resident + score per chunk
uint8_t * lp = out.data() + hdr + (size_t)nc * chunk_bytes_;
for (int c = 0; c < nc; ++c) {
const ChunkState & st = chunks_[c];
const uint8_t was_res = (st.block >= 0) ? 1u : 0u;
std::memcpy(lp, &was_res, 1); lp += 4; // 1 byte + 3 pad (out already zero)
std::memcpy(lp, &st.score, 4); lp += 4;
}
return out;
}

// Restore state from a blob produced by serialize(). Returns false on
// header mismatch (layout drift guard). Callers must rebuild slot masks.
//
// v1 blobs (magic "KVFLASH1"): ledger section is read back — was_resident
// controls whether a chunk gets a pool slot (resident) or stays host-backed
// (on_host only). Scores are restored into ChunkState.score.
//
// Ordering: pre-size chunks_ so each entry exists before slot_for() runs.
// For resident chunks: set host_data + on_host=true, then call slot_for().
// slot_for's recall branch sees on_host==true and calls copy_chunk(to_host=false).
// For non-resident chunks: set host_data + on_host=true, skip slot_for() so
// the chunk stays host-backed only (mirrors the original page-out state).
bool deserialize(const uint8_t * data, size_t n) {
static constexpr uint64_t kMagic = 0x4b56464c41534831ULL; // "KVFLASH1"
const size_t hdr = sizeof(uint64_t) + 8 * sizeof(uint32_t);
if (n < hdr) return false;
const uint8_t * p = data;
uint64_t magic = 0; std::memcpy(&magic, p, 8); p += 8;
if (magic != kMagic) return false;
auto r32 = [&]() { uint32_t v = 0; std::memcpy(&v, p, 4); p += 4; return v; };
const int nc = (int)r32();
const int ct = (int)r32();
const int nhkv = (int)r32();
const size_t kseg = (size_t)r32();
const size_t vseg = (size_t)r32();
const size_t cb = (size_t)r32();
const uint32_t k_type = r32(); // kv_k_type_u32
const uint32_t has_led = r32(); // has_ledger
// Dimension guard (same as before) + dtype guard.
if (ct != cfg_.chunk_tokens || nhkv != n_head_kv_ ||
kseg != k_seg_bytes_ || vseg != v_seg_bytes_ || cb != chunk_bytes_) {
return false;
}
// Dtype guard: refuse blobs whose KV dtype enum differs from the live cache.
const uint32_t live_k_type = (!attn_k_.empty() && attn_k_[0])
? (uint32_t)attn_k_[0]->type
: (uint32_t)GGML_TYPE_F16;
if (k_type != live_k_type) return false;

const size_t ledger_entry = 8;
const size_t kv_section = (size_t)nc * chunk_bytes_;
const size_t expected = hdr + kv_section + (has_led ? (size_t)nc * ledger_entry : 0);
if (n < expected) return false;

// Read ledger into a temp buffer before reset() clears state.
std::vector<uint8_t> ledger_was_res(nc, 1u); // default: treat as resident
std::vector<float> ledger_scores(nc, -std::numeric_limits<float>::infinity());
if (has_led) {
const uint8_t * lp = data + hdr + kv_section;
for (int c = 0; c < nc; ++c) {
std::memcpy(&ledger_was_res[c], lp, 1); lp += 4; // 1 byte + 3 pad
std::memcpy(&ledger_scores[c], lp, 4); lp += 4;
}
}

reset();
last_serialized_kv_k_type_ = k_type;
chunks_.resize(nc); // pre-size so slot_for doesn't resize again
for (int c = 0; c < nc; ++c) {
const uint8_t * src = data + hdr + (size_t)c * chunk_bytes_;
ChunkState & st = chunks_[c];
// Park bytes in host_data; slot_for's recall branch copies to pool.
#ifdef KVFLASH_HAS_ASYNC_DMA
if (!st.host_data) {
if (cudaMallocHost(&st.host_data, chunk_bytes_) != cudaSuccess) return false;
stats_.host_bytes += (int64_t)chunk_bytes_;
}
std::memcpy(st.host_data, src, chunk_bytes_);
#else
st.host_data.assign(src, src + chunk_bytes_);
#endif
st.on_host = true;
st.score = ledger_scores[c];
if (ledger_was_res[c]) {
// slot_for assigns a block and auto-recalls via copy_chunk.
slot_for((int64_t)c * cfg_.chunk_tokens);
}
// else: stays host-backed (on_host=true, block=-1), matching page_out state.
}
#ifdef KVFLASH_HAS_ASYNC_DMA
// Recalls above are async on page_stream_; settle before the pool is read.
if (has_pending_page_in_) synchronize_paging();
#endif
return true;
}

// ── Per-chunk score accessors (QK ledger) ─────────────────────────────
// set_chunk_score / chunk_score allow callers (the QK scorer wiring in
// qwen35_backend.cpp) to push per-chunk relevance scores into the pager's
// ledger before snapshot_save so they round-trip through serialize().
void set_chunk_score(int c, float s) {
if (c >= 0 && c < (int)chunks_.size()) chunks_[(size_t)c].score = s;
}
float chunk_score(int c) const {
if (c >= 0 && c < (int)chunks_.size()) return chunks_[(size_t)c].score;
return -std::numeric_limits<float>::infinity();
}

// The kv_k dtype enum read from the most-recently-deserialized blob.
// Useful for callers that want to confirm the restored dtype after
// deserialize() returns (e.g., Phase 2 QK pool rebuild).
uint32_t serialized_kv_k_type() const { return last_serialized_kv_k_type_; }

private:
struct ChunkState {
int block = -1; // pool block index, -1 = not resident
bool on_host = false; // backing store holds valid bytes
uint64_t last_use = 0;
float score = -std::numeric_limits<float>::infinity(); // QK relevance ledger
#ifdef KVFLASH_HAS_ASYNC_DMA
void * host_data = nullptr; // cudaMallocHost-pinned; allocated on first page_out
#else
Expand All @@ -441,6 +668,7 @@ class KvFlashPager {
if (chunks_[c].block < 0) continue;
if (c < cfg_.sink_chunks) continue;
if (c > cur_chunk_ - 1 - cfg_.tail_window_chunks) continue;
if (is_pinned(c)) continue;
if (score_hook) {
const float s = score_hook(c);
if (victim < 0 || s < v_score) { victim = c; v_score = s; }
Expand Down Expand Up @@ -566,11 +794,13 @@ class KvFlashPager {
std::vector<ChunkState> chunks_;
std::vector<int> free_blocks_;
std::vector<uint8_t> zero_buf_; // used by zero_block() in non-CUDA builds
std::vector<uint8_t> pinned_; // per-chunk pin flag; empty = no pins
KvFlashStats stats_;
size_t k_seg_bytes_ = 0, v_seg_bytes_ = 0, chunk_bytes_ = 0;
int n_blocks_ = 0, n_head_kv_ = 0, cur_chunk_ = 0;
uint64_t clock_ = 0;
uint64_t epoch_ = 0;
uint32_t last_serialized_kv_k_type_ = (uint32_t)GGML_TYPE_F16; // from most-recent deserialize

#ifdef KVFLASH_HAS_ASYNC_DMA
cudaStream_t page_stream_ = nullptr;
Expand Down
Loading
Loading