diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 9f62ddb08..31f8a1da1 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -77,6 +77,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MMID_TELEMETRY` | unset | DEBUG: report MUL_MAT_ID dispatch, MMVQ variant, and per-node graph compatibility. | | `DFLASH_KVFLASH` | unset | Prefer the CLI: `--kvflash` (token count or `auto`). | | `DFLASH_PREFIX_CACHE_SLOTS` | 32 | Container-entrypoint equivalent of `--prefix-cache-slots`; not read directly by the native binary. | +| `DFLASH_PREFIX_NO_FEAT` | unset | DEBUG/A-B: =1 drops the drafter feature-ring slab from concurrent prefix checkpoints (speculative `--max-concurrency` only) to measure the payload's cost and exercise the cold-ring restore path. | | `DFLASH_PREFILL_CACHE_SLOTS` | 0 | Container-entrypoint equivalent of `--prefill-cache-slots`; not read directly by the native binary. | | `DFLASH_PREFILL_POOL_TRIM_TOKENS` | unset | OPT-IN: trim cached allocations from legacy CUDA/HIP device pools at completed Qwen3.5 prefill chunk boundaries after each configured token interval. Intended for long, shape-changing prefills on non-VMM devices; each trim synchronizes the target backend and retires captured graphs. | | `DFLASH_SPLIT_FAST_ROLLBACK` | unset | OPT-IN: exact F32 checkpoints and replay-free rollback for local qwen35 target layer splits. Prefer `--target-split-fast-rollback`; adds checkpoint VRAM (~1.65 GiB for the measured Qwen3.6-27B q=16 split). | @@ -292,6 +293,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_PREFILL_POOL_TRIM_TOKENS` - qwen35_backend.cpp (OPT-IN: trim legacy device pools during long prefills) - `DFLASH_PREFILL_TIMING` - qwen35_backend.cpp (DEBUG: per-ubatch prefill build/alloc/compute timing) - `DFLASH_PREFIX_CACHE_SLOTS` - scripts/entrypoint.sh (maps to `--prefix-cache-slots`) +- `DFLASH_PREFIX_NO_FEAT` - qwen35_seq_engine.cpp (DEBUG/A-B: =1 drops the drafter feature-ring slab from concurrent prefix checkpoints) - `DFLASH_QWEN35MOE_CACHE_SLOTS` - qwen35moe_backend.cpp - `DFLASH_QWEN35MOE_HOTNESS` - qwen35moe_backend.cpp - `DFLASH_QWEN35MOE_NEXT_PLACEMENT_OUT` - qwen35moe_backend.cpp diff --git a/server/docs/PREFIX_CACHE.md b/server/docs/PREFIX_CACHE.md index 53120f7ff..a4bb4bc22 100644 --- a/server/docs/PREFIX_CACHE.md +++ b/server/docs/PREFIX_CACHE.md @@ -221,6 +221,25 @@ resident committed checkpoint buffers. During an atomic replacement, the new buffer and the selected victim can coexist briefly, so transient process memory can exceed the limit by up to one checkpoint. +When the engine runs speculative continuous batching (a local same-device +DFlash drafter under `--max-concurrency`), eligible checkpoints additionally +carry the slot's drafter feature-ring slab so a restored prefix keeps the +draft window warm — the resident estimate charges for it. AR-only engines +keep the leaner KV + recurrent payload. A capture is ineligible while its +slot is still floored by a featureless restore (below), and the bench/debug +switch `DFLASH_PREFIX_NO_FEAT=1` drops the slab entirely (see +ENVIRONMENT.md). + +A checkpoint without a feature payload still restores correctly under +speculation. The engine then treats ring rows below the restore cut as +untrusted: the drafter rebuilds its K/V window only from rows the +occupying sequence wrote itself (the restore cut plus everything +generated afterward), so proposals are never conditioned on a previous +slot occupant's features. Such a slot also refrains from blessing those +rows into later feature-bearing checkpoints — until it has written a full +ring worth of positions past the restore cut, after which every row is +self-written and the payload resumes. + | Scenario | Typical prefix length | Recommended cap | |----------|----------------------|-----------------| | Single-user chat | 200–2000 tokens | 16–32 | diff --git a/server/src/internal.h b/server/src/internal.h index 810ee7f75..e4c6471ac 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -623,6 +623,11 @@ struct PrefixSnapshot { // [HEAD_DIM, kv_end-kv_start, N_HEAD_KV] (smaller than cache). // - ssm_state_snap, conv_state_snap, target_feat_snap are NOT // allocated (THIN snapshots are KV-only). + // For Layout::paged: + // - target_feat_snap is optional. Speculating engines capture the + // sequence slot's live drafter feature-ring slab so a restored prefix + // keeps the draft window warm; engines without speculation leave it + // null and pay nothing. }; // Snapshot the slim state of `cache` into `snap`. KV tensors are RIGHT-SIZED @@ -647,21 +652,28 @@ void free_prefix_snapshot(PrefixSnapshot & snap); // Exact CPU-buffer allocation size for the dense checkpoint layout used by // snapshot_paged_target_cache(). Returns zero when the cache topology or token // count is invalid. This lets the scheduler enforce a resident-memory budget -// before allocating or copying a checkpoint. +// before allocating or copying a checkpoint. `with_target_feat` must match the +// capture call so the estimate covers the optional drafter feature payload. size_t estimate_paged_target_cache_snapshot_bytes( - const TargetCache & cache, int token_count); + const TargetCache & cache, int token_count, bool with_target_feat = false); // Capture one live sequence from a multi-slot paged cache. Attention rows are // gathered through `block_table` into dense logical order in the copied // snapshot; recurrent state is copied only from `seq_slot`'s slab. The page // table itself is intentionally not retained: every restore owns fresh pages. +// When `with_target_feat` is set and the cache owns a drafter feature ring, +// the slot's live ring slab is captured verbatim (ring slots are absolute +// positions mod cap, so the first min(token_count, cap) slab rows are +// self-describing on restore). Engines without speculation pass false and +// keep the checkpoint at KV + recurrent state only. bool snapshot_paged_target_cache( const TargetCache & cache, int seq_slot, const std::vector & block_table, int block_size, int token_count, - PrefixSnapshot & snap); + PrefixSnapshot & snap, + bool with_target_feat = false); // Atomically replace a paged snapshot. The incumbent remains valid when // allocation, layout validation, or any staged copy fails. @@ -671,11 +683,14 @@ bool replace_paged_target_cache( const std::vector & block_table, int block_size, int token_count, - PrefixSnapshot & destination); + PrefixSnapshot & destination, + bool with_target_feat = false); // Restore a copied paged snapshot into fresh destination pages and one // recurrent-state slab. `block_table` describes the destination sequence and -// must cover snap.cur_pos logical tokens. +// must cover snap.cur_pos logical tokens. A checkpoint that carries a +// drafter feature payload also restores the slot's feature-ring slab; a +// checkpoint without one restores KV + recurrent state only. bool restore_paged_target_cache( const PrefixSnapshot & snap, TargetCache & cache, diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 6becfdaf8..2c68605de 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -11,6 +11,7 @@ #include "graph_builders.h" #include "attn_masks.h" #include "prefill_helpers.h" +#include "common/chain_rollback_policy.h" #include "common/concurrency/chain_spec_shapes.h" #include "common/dflash2_head.h" #include "common/sampler.h" @@ -64,6 +65,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( slot_draft_kv_.resize(static_cast(n_slots)); seq_lens_.assign(static_cast(n_slots), 0); reserve_growth_.assign(static_cast(n_slots), 0); + slot_ring_valid_from_.assign(static_cast(n_slots), 0); fixed_chain_ready_ = fixed_chain_.enabled && fixed_chain_.width > 1 && fixed_chain_.width <= 16 && @@ -264,6 +266,13 @@ Qwen35SeqEngine::prepare_chain_drafts( return std::nullopt; } lanes.push_back({i, input.slot, input.token, state, mirror}); + // Ring floor from a featureless restore: rows below it belong to a + // previous occupant, so seed the append cursor past them — they stay + // unappended (slot_pos -1) and masked out of the draft context. + const int32_t ring_floor = + input.slot < static_cast(slot_ring_valid_from_.size()) + ? slot_ring_valid_from_[static_cast(input.slot)] : 0; + if (state->next_pos < ring_floor) state->next_pos = ring_floor; if (!draft_kv_begin_step( *state, b_.dw_, b_.draft_backend_, *mirror, slots_.slot(input.slot).cur_pos)) { @@ -380,17 +389,37 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { reset_recurrent_slot(b_.cache_, result.slot); - if (result.slot >= 0 && - result.slot < static_cast(slot_draft_kv_.size()) && - slot_draft_kv_[static_cast(result.slot)]) { - draft_kv_reset(*slot_draft_kv_[static_cast(result.slot)]); - } + reset_slot_draft_state(result.slot); } return result; } +void Qwen35SeqEngine::reset_slot_draft_state(int slot) { + if (slot >= 0 && slot < static_cast(slot_draft_kv_.size()) && + slot_draft_kv_[static_cast(slot)]) { + draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); + } + if (slot >= 0 && slot < static_cast(slot_ring_valid_from_.size())) { + slot_ring_valid_from_[static_cast(slot)] = 0; + } +} + +// DFLASH_PREFIX_NO_FEAT=1 drops the drafter feature slab from concurrent +// prefix checkpoints — a bench/debug knob that measures the payload's +// contribution and exercises the cold-ring restore path. +static bool prefix_feat_payload_disabled() { + static const bool off = env_flag_enabled("DFLASH_PREFIX_NO_FEAT"); + return off; +} + size_t Qwen35SeqEngine::estimate_prefix_store_bytes(int tokens) const { - return estimate_paged_target_cache_snapshot_bytes(b_.cache_, tokens); + // Speculating engines also carry the slot's drafter feature-ring slab in + // each checkpoint; the estimate must charge for it so the resident-byte + // budget stays honest. + const bool with_target_feat = + fixed_chain_ready_ && !prefix_feat_payload_disabled(); + return estimate_paged_target_cache_snapshot_bytes( + b_.cache_, tokens, with_target_feat); } int Qwen35SeqEngine::checkpoint_index(PrefixStoreRef checkpoint) const { @@ -428,6 +457,12 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit_with_prefix( const int slot = result.slot; slots_.slot(slot).pending_capture = {}; + // A restored prefix repopulates the slot's target_feat slab below; a + // cold admission leaves whatever the prefill graph will write. Either + // way the previous occupant's drafter K/V window must not survive — + // its next_pos/slot_pos bookkeeping would otherwise feed stale rows to + // the new sequence's drafts. + reset_slot_draft_state(slot); bool restored = false; if (plan.restore.valid()) { const int restore_index = checkpoint_index(plan.restore); @@ -464,17 +499,30 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit_with_prefix( result.prefix_store.invalidated = plan.restore; if (result.status == AdmitResult::Status::admitted) { reset_recurrent_slot(b_.cache_, result.slot); + reset_slot_draft_state(result.slot); } else { result.error = "cold admission failed after stale prefix restore"; } } else { result.prefix_store.restored = plan.restore; + // A checkpoint without the feature slab leaves ring rows below + // the restore cut holding a previous occupant's features. Floor + // them so the drafter never bulk-appends rows this sequence did + // not write (unpopulated slots stay masked out of the draft + // context) and a later capture does not bless them. + if (snap && !snap->target_feat_snap && + slot < static_cast(slot_ring_valid_from_.size())) { + slot_ring_valid_from_[static_cast(slot)] = + plan.restore.tokens; + } std::fprintf(stderr, "[parallel-pc] restored checkpoint=%llu seq_slot=%d " - "tokens=%d time_ms=%.1f\n", + "tokens=%d feat=%d time_ms=%.1f\n", (unsigned long long)plan.restore.id, slot, - plan.restore.tokens, (double)restore_elapsed_us / 1000.0); + plan.restore.tokens, + snap && snap->target_feat_snap ? 1 : 0, + (double)restore_elapsed_us / 1000.0); } result.prefix_store.restore_attempted = true; result.prefix_store.restore_elapsed_us = restore_elapsed_us; @@ -514,9 +562,22 @@ PrefixStoreEvent Qwen35SeqEngine::capture_prefix( } const auto capture_started = std::chrono::steady_clock::now(); PrefixSnapshot & snapshot = b_.prefix_snapshots_[checkpoint]; + // A slot whose ring still holds foreign rows below a featureless restore + // must not bless them into a feature payload. Ring rows are keyed by + // position % cap, so the foreign region is fully overwritten once the + // sequence has written `cap` positions past the cut and the payload + // resumes. + const int32_t ring_floor = + slot < static_cast(slot_ring_valid_from_.size()) + ? slot_ring_valid_from_[static_cast(slot)] : 0; + const bool with_target_feat = + fixed_chain_ready_ && !prefix_feat_payload_disabled() && + (ring_floor == 0 || + ticket.checkpoint.tokens >= ring_floor + b_.cache_.target_feat_cap); if (!replace_paged_target_cache( b_.cache_, slot, sequence.block_table, - (int)pool_.block_size(), ticket.checkpoint.tokens, snapshot)) { + (int)pool_.block_size(), ticket.checkpoint.tokens, snapshot, + with_target_feat)) { event.elapsed_us = (uint64_t)std::chrono::duration_cast( std::chrono::steady_clock::now() - capture_started).count(); @@ -1622,10 +1683,7 @@ bool Qwen35SeqEngine::restore_kv(int slot, std::string & error) { // state together — so it must be reset exactly as at admission. if (!slots_.resume_recompute(slot)) return false; reset_recurrent_slot(b_.cache_, slot); - if (slot < static_cast(slot_draft_kv_.size()) && - slot_draft_kv_[static_cast(slot)]) { - draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); - } + reset_slot_draft_state(slot); return true; } std::vector blocks; @@ -1651,10 +1709,7 @@ bool Qwen35SeqEngine::evict_kv(int slot, int32_t pending_token, void Qwen35SeqEngine::retire(int slot) { offload_.discard(slot); if (!slots_.is_active(slot)) return; - if (slot >= 0 && slot < static_cast(slot_draft_kv_.size()) && - slot_draft_kv_[static_cast(slot)]) { - draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); - } + reset_slot_draft_state(slot); slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 0558e2281..20fb55da6 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -168,6 +168,19 @@ class Qwen35SeqEngine final : public SeqEngine { bool arm_capture( int slot, PrefixCaptureTicket ticket, int restored_tokens); int checkpoint_index(PrefixStoreRef checkpoint) const; + // Invalidate the slot's drafter K/V window so the next speculative step + // bulk-appends from the feature ring. No-op when speculation is off. + void reset_slot_draft_state(int slot); + + // Per-slot lower bound on feature-ring rows the occupying sequence wrote + // itself. A featureless checkpoint restore leaves ring rows below the + // restored cut holding a previous occupant's features; the floor keeps + // the drafter from bulk-appending them (they stay masked out of the + // draft context) and keeps captures from claiming them as a feature + // payload until the sequence has written cap positions past the cut and + // every ring row is self-written again. Zero means the whole ring slab + // is self-written. + std::vector slot_ring_valid_from_; PagedKvPool & pool_; Qwen35Backend & b_; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index f0cd8ef87..891f76261 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -3162,9 +3162,10 @@ bool paged_cache_pairs_complete(const TargetCache & cache) { } bool create_paged_snapshot_layout( - const TargetCache & cache, int token_count, PrefixSnapshot & snap) { + const TargetCache & cache, int token_count, int feat_rows, + PrefixSnapshot & snap) { const int total_tensors = 2 * (int)cache.attn_k.size() + - 2 * (int)cache.ssm_state.size(); + 2 * (int)cache.ssm_state.size() + (feat_rows > 0 ? 1 : 0); ggml_init_params params{}; params.mem_size = (size_t)(total_tensors + 16) * ggml_tensor_overhead(); @@ -3172,6 +3173,21 @@ bool create_paged_snapshot_layout( snap.ctx = ggml_init(params); if (!snap.ctx) return false; + // Live drafter feature-ring slab: verbatim ring slots [0, feat_rows). + // Ring position→slot mapping is absolute, so the same byte range + // round-trips through capture and restore for this token count. + snap.target_feat_snap = nullptr; + snap.target_feat_cap = 0; + if (feat_rows > 0 && cache.target_feat) { + snap.target_feat_snap = ggml_new_tensor_2d( + snap.ctx, cache.target_feat->type, cache.target_feat->ne[0], + feat_rows); + if (snap.target_feat_snap) { + ggml_set_name(snap.target_feat_snap, "snap_target_feat"); + snap.target_feat_cap = cache.target_feat_cap; + } + } + snap.attn_k_snap.assign(cache.attn_k.size(), nullptr); snap.attn_v_snap.assign(cache.attn_v.size(), nullptr); snap.ssm_state_snap.assign(cache.ssm_state.size(), nullptr); @@ -3242,21 +3258,60 @@ bool paged_snapshot_matches( snap.conv_state_snap[i], cache.conv_state[i], cache.n_seq_slots, /*slot_axis=*/2))) return false; } + // Optional drafter feature payload: when a checkpoint carries the ring + // slab, it must describe this cache's per-slot ring exactly — the live + // window for `tokens` committed tokens is the first min(tokens, cap) + // slab columns of a [fc_in, cap * n_seq_slots (+ dead row)] ring. + if (snap.target_feat_snap) { + const ggml_tensor * feat = cache.target_feat; + if (!feat || cache.target_feat_cap <= 0 || + !ggml_is_matrix(feat) || + !ggml_is_matrix(snap.target_feat_snap) || + snap.target_feat_cap != cache.target_feat_cap || + feat->ne[1] < (int64_t)cache.target_feat_cap * + cache.n_seq_slots || + snap.target_feat_snap->type != feat->type || + snap.target_feat_snap->ne[0] != feat->ne[0] || + snap.target_feat_snap->ne[1] != + std::min(tokens, cache.target_feat_cap) || + !ggml_is_contiguous(snap.target_feat_snap) || + !ggml_is_contiguous(feat)) { + return false; + } + } return true; } +// Slab rows a speculating engine captures for `token_count` committed tokens: +// the live ring window, bounded by the per-slot ring capacity. Zero when the +// caller keeps the checkpoint AR-only or the cache has no feature ring. +int paged_feat_rows( + const TargetCache & cache, int token_count, bool with_target_feat) { + if (!with_target_feat || !cache.target_feat || + cache.target_feat_cap <= 0) { + return 0; + } + return std::min(token_count, cache.target_feat_cap); +} + } // namespace size_t estimate_paged_target_cache_snapshot_bytes( const TargetCache & cache, - int token_count) { + int token_count, + bool with_target_feat) { if (cache.n_seq_slots < 1 || token_count <= 0 || token_count > cache.max_ctx || !paged_cache_pairs_complete(cache)) { return 0; } PrefixSnapshot layout; - if (!create_paged_snapshot_layout(cache, token_count, layout)) return 0; + if (!create_paged_snapshot_layout( + cache, token_count, paged_feat_rows(cache, token_count, + with_target_feat), + layout)) { + return 0; + } const size_t bytes = ggml_backend_alloc_ctx_tensors_from_buft_size( layout.ctx, paged_snapshot_buffer_type()); @@ -3270,7 +3325,8 @@ bool snapshot_paged_target_cache( const std::vector & block_table, int block_size, int token_count, - PrefixSnapshot & snap) { + PrefixSnapshot & snap, + bool with_target_feat) { if (!cache.backend || cache.n_seq_slots < 1 || seq_slot < 0 || seq_slot >= cache.n_seq_slots || token_count <= 0 || token_count > cache.max_ctx || block_size <= 0 || @@ -3279,16 +3335,22 @@ bool snapshot_paged_target_cache( set_last_error("snapshot_paged_target_cache: invalid arguments"); return false; } + const int feat_rows = + paged_feat_rows(cache, token_count, with_target_feat); const bool needs_alloc = !snap.ctx || snap.layout != PrefixSnapshot::Layout::paged || snap.cur_pos != token_count || snap.attn_k_snap.size() != cache.attn_k.size() || snap.attn_v_snap.size() != cache.attn_v.size() || snap.ssm_state_snap.size() != cache.ssm_state.size() || - snap.conv_state_snap.size() != cache.conv_state.size(); + snap.conv_state_snap.size() != cache.conv_state.size() || + (snap.target_feat_snap != nullptr) != (feat_rows > 0) || + (snap.target_feat_snap && + snap.target_feat_snap->ne[1] != feat_rows); if (needs_alloc) { free_prefix_snapshot(snap); - if (!create_paged_snapshot_layout(cache, token_count, snap)) { + if (!create_paged_snapshot_layout( + cache, token_count, feat_rows, snap)) { set_last_error("paged PrefixSnapshot ggml_init failed"); return false; } @@ -3335,13 +3397,20 @@ bool snapshot_paged_target_cache( cache.conv_state[i], snap.conv_state_snap[i]->data, (size_t)seq_slot * conv_bytes, conv_bytes); } + // The slot's live drafter feature ring is one contiguous slab of + // target_feat columns — a single range copy. + if (snap.target_feat_snap) { + ggml_backend_tensor_get_async( + cache.backend, cache.target_feat, snap.target_feat_snap->data, + (size_t)seq_slot * (size_t)cache.target_feat_cap * + cache.target_feat->nb[1], + ggml_nbytes(snap.target_feat_snap)); + } ggml_backend_synchronize(cache.backend); snap.cur_pos = token_count; snap.last_tok = -1; snap.kv_k_type = cache.kv_k_type; snap.max_ctx = cache.max_ctx; - snap.target_feat_cap = 0; - snap.target_feat_snap = nullptr; snap.layout = PrefixSnapshot::Layout::paged; return true; } @@ -3352,11 +3421,12 @@ bool replace_paged_target_cache( const std::vector & block_table, int block_size, int token_count, - PrefixSnapshot & destination) { + PrefixSnapshot & destination, + bool with_target_feat) { PrefixSnapshot candidate; if (!snapshot_paged_target_cache( cache, seq_slot, block_table, block_size, token_count, - candidate)) { + candidate, with_target_feat)) { free_prefix_snapshot(candidate); return false; } @@ -3408,6 +3478,15 @@ bool restore_paged_target_cache( cache.conv_state[i], snap.conv_state_snap[i]->data, (size_t)seq_slot * conv_bytes, conv_bytes); } + // Restore the slot's drafter feature-ring slab when the checkpoint + // carries it; feat-less checkpoints keep the AR-only payload. + if (snap.target_feat_snap) { + ggml_backend_tensor_set_async( + cache.backend, cache.target_feat, snap.target_feat_snap->data, + (size_t)seq_slot * (size_t)cache.target_feat_cap * + cache.target_feat->nb[1], + ggml_nbytes(snap.target_feat_snap)); + } ggml_backend_synchronize(cache.backend); return true; } diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index e83a8536b..11e401042 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -397,3 +397,207 @@ TEST_CASE(RecurrentSnapshotFixture, copied_paged_prefix_uses_fresh_pages) { ggml_free(ctx); ggml_backend_free(backend); } + +TEST_CASE(RecurrentSnapshotFixture, copied_paged_prefix_carries_draft_features) { + ggml_backend_t backend = ggml_backend_cpu_init(); + CHECK(backend != nullptr); + if (!backend) SKIP("CPU backend is unavailable"); + + ggml_init_params params{}; + params.mem_size = 16 * ggml_tensor_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + CHECK(ctx != nullptr); + if (!ctx) { + ggml_backend_free(backend); + SKIP("could not initialize ggml context"); + } + + const int feat_cap = 8; // per-slot ring width + ggml_tensor * key = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2, 64, 2); + ggml_tensor * value = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2, 64, 2); + ggml_tensor * ssm = + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, 2, 2); + ggml_tensor * conv = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3, 2, 2); + // Concurrent ring: [fc_in, cap * n_seq_slots + 1] — one slab per slot + // plus the dead padding row. + ggml_tensor * target_feat = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 6, feat_cap * 2 + 1); + ggml_backend_buffer_t buffer = + ggml_backend_alloc_ctx_tensors(ctx, backend); + CHECK(buffer != nullptr); + if (!buffer) { + ggml_free(ctx); + ggml_backend_free(backend); + SKIP("could not allocate CPU backend tensors"); + } + + TargetCache cache; + cache.backend = backend; + cache.max_ctx = 64; + cache.n_seq_slots = 2; + cache.kv_k_type = GGML_TYPE_F32; + cache.attn_k = {key}; + cache.attn_v = {value}; + cache.ssm_state = {ssm}; + cache.conv_state = {conv}; + cache.target_feat = target_feat; + cache.target_feat_cap = feat_cap; + + // Slot 1's ring slab: each live column tagged by its ring slot. + const int64_t fc_in = target_feat->ne[0]; + std::vector feat((size_t)ggml_nelements(target_feat), -7.0f); + for (int64_t col = feat_cap; col < 2 * feat_cap; ++col) { + for (int64_t e = 0; e < fc_in; ++e) { + feat[(size_t)col * fc_in + e] = 4000.0f + (float)(col - feat_cap); + } + } + set_tensor(target_feat, feat); + + const std::vector source_blocks = {2, 0}; + + // The resident-byte estimate must charge for the feature payload so the + // scheduler budget covers what a speculating engine actually copies. + const size_t ar_bytes = estimate_paged_target_cache_snapshot_bytes( + cache, /*token_count=*/20, /*with_target_feat=*/false); + const size_t spec_bytes = estimate_paged_target_cache_snapshot_bytes( + cache, /*token_count=*/20, /*with_target_feat=*/true); + CHECK(ar_bytes > 0); + CHECK(spec_bytes > ar_bytes); + + PrefixSnapshot snap; + CHECK(snapshot_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/20, snap, + /*with_target_feat=*/true)); + CHECK(snap.target_feat_snap != nullptr); + CHECK(snap.target_feat_cap == feat_cap); + CHECK(snap.target_feat_snap->ne[1] == feat_cap); + CHECK(ggml_backend_buffer_get_size(snap.buf) == spec_bytes); + + const auto snap_feat = get_tensor(snap.target_feat_snap); + for (int64_t col = 0; col < feat_cap; ++col) { + for (int64_t e = 0; e < fc_in; ++e) { + CHECK(snap_feat[(size_t)col * fc_in + e] == + 4000.0f + (float)col); + } + } + + // Restore into slot 0 of a cleared cache: the slot's live ring slab + // returns verbatim, while untouched ring slots and the dead row stay + // zeroed. + set_tensor(target_feat, std::vector(feat.size(), 0.0f)); + const std::vector destination_blocks = {1, 3}; + CHECK(restore_paged_target_cache( + snap, cache, /*seq_slot=*/0, destination_blocks, + /*block_size=*/16)); + const auto restored_feat = get_tensor(target_feat); + for (int64_t col = 0; col < feat_cap; ++col) { + for (int64_t e = 0; e < fc_in; ++e) { + CHECK(restored_feat[(size_t)col * fc_in + e] == + 4000.0f + (float)col); + } + } + for (int64_t col = feat_cap; col < target_feat->ne[1]; ++col) { + for (int64_t e = 0; e < fc_in; ++e) { + CHECK(restored_feat[(size_t)col * fc_in + e] == 0.0f); + } + } + free_prefix_snapshot(snap); + + // Short prefixes copy only live rows, including when restoring into a + // nonzero slot. Neither the unused tail nor the dead row is overwritten. + set_tensor(target_feat, feat); + CHECK(replace_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/3, snap, + /*with_target_feat=*/true)); + CHECK(snap.target_feat_snap->ne[1] == 3); + CHECK(ggml_backend_buffer_get_size(snap.buf) == + estimate_paged_target_cache_snapshot_bytes(cache, 3, true)); + set_tensor(target_feat, std::vector(feat.size(), -9.0f)); + CHECK(restore_paged_target_cache( + snap, cache, /*seq_slot=*/1, destination_blocks, + /*block_size=*/16)); + const auto short_feat = get_tensor(target_feat); + for (int64_t col = 0; col < target_feat->ne[1]; ++col) { + const float expected = col >= feat_cap && col < feat_cap + 3 + ? 4000.0f + (float)(col - feat_cap) : -9.0f; + for (int64_t e = 0; e < fc_in; ++e) { + CHECK(short_feat[(size_t)col * fc_in + e] == expected); + } + } + + // Reject an extra payload plane before the slab copy: checking only + // width/rows/contiguity would allow ggml_nbytes() to cross slot boundaries. + ggml_tensor * saved_feat = snap.target_feat_snap; + ggml_tensor * extra_plane = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, fc_in, 3, 2); + snap.target_feat_snap = extra_plane; + CHECK(!restore_paged_target_cache( + snap, cache, /*seq_slot=*/0, destination_blocks, + /*block_size=*/16)); + CHECK(get_tensor(target_feat) == short_feat); + snap.target_feat_snap = saved_feat; + + // An incompatible ring cannot replace the incumbent checkpoint. + cache.target_feat_cap = feat_cap + 1; + CHECK(!replace_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/3, snap, + /*with_target_feat=*/true)); + CHECK(snap.target_feat_snap == saved_feat); + CHECK(snap.cur_pos == 3); + cache.target_feat_cap = feat_cap; + free_prefix_snapshot(snap); + + // AR-only checkpoints keep the lean payload: no feature tensor, and a + // restore leaves the cache's ring untouched. + PrefixSnapshot ar_snap; + CHECK(snapshot_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/20, ar_snap)); + CHECK(ar_snap.target_feat_snap == nullptr); + CHECK(ar_snap.target_feat_cap == 0); + CHECK(ggml_backend_buffer_get_size(ar_snap.buf) == ar_bytes); + set_tensor(target_feat, std::vector(feat.size(), -3.0f)); + CHECK(restore_paged_target_cache( + ar_snap, cache, /*seq_slot=*/0, destination_blocks, + /*block_size=*/16)); + const auto untouched_feat = get_tensor(target_feat); + for (float v : untouched_feat) { + CHECK(v == -3.0f); + } + free_prefix_snapshot(ar_snap); + + // A checkpoint carrying features cannot restore into a cache that has no + // ring — the payload would be silently dropped otherwise. + CHECK(snapshot_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/20, snap, + /*with_target_feat=*/true)); + CHECK(snap.target_feat_snap != nullptr); + cache.target_feat = nullptr; + CHECK(!restore_paged_target_cache( + snap, cache, /*seq_slot=*/0, destination_blocks, + /*block_size=*/16)); + free_prefix_snapshot(snap); + cache.target_feat = target_feat; + + // Requesting the payload on a ring-less cache degrades to the lean + // layout rather than failing the capture. + cache.target_feat = nullptr; + CHECK(snapshot_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/20, snap, + /*with_target_feat=*/true)); + CHECK(snap.target_feat_snap == nullptr); + free_prefix_snapshot(snap); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +}