diff --git a/README.md b/README.md index 005318eb8..f455c273a 100644 --- a/README.md +++ b/README.md @@ -363,6 +363,7 @@ When compression is on, multi-turn continuations automatically use **FlowKV**: a | `DFLASH27B_KV_TQ3=1` | (default) | Preset TQ3_0 K+V (3.5 bpv, fits 256K @ 24 GB) | | `DFLASH27B_KV_Q4=1` | off | Q4_0 K+V (4.5 bpv, legacy, ~128K ceiling) | | `--prefix-cache-slots N` | — | Live prefix-cache slot count | +| `--concurrent-prefix-cache-max-mib N` | `4096` | Resident RAM limit for copied concurrent paged checkpoints; `0` is unlimited. | | `DFLASH_PREFIX_CACHE_SLOTS=N` | `32` | Container-entrypoint equivalent of `--prefix-cache-slots`; the native binary itself uses the CLI flag. | | `DFLASH_PREFILL_CACHE_SLOTS=N` | `0` | Container-entrypoint equivalent of `--prefill-cache-slots`; the native binary itself uses the CLI flag. | | `--kv-cache-dir ` | — | Persist prefix cache to disk | diff --git a/docs/specs/props-endpoint.md b/docs/specs/props-endpoint.md index e4238df3d..60385fc8c 100644 --- a/docs/specs/props-endpoint.md +++ b/docs/specs/props-endpoint.md @@ -363,15 +363,36 @@ enabled, fields carry the runtime configuration: ```json "prefix_cache": { - "capacity": 0, - "in_use": 0, - "lifetime_hits": 0 + "capacity": 0, + "in_use": 0, + "lifetime_hits": 0, + "agent_turn_enabled": false, + "max_resident_bytes": 4294967296, + "resident_bytes": 0, + "budget_skips": 0, + "capture_attempts": 0, + "capture_failures": 0, + "capture_stall_ms_total": 0.0, + "capture_stall_ms_max": 0.0, + "restore_attempts": 0, + "restore_invalidations": 0, + "restore_stall_ms_total": 0.0, + "restore_stall_ms_max": 0.0 } ``` The inline prefix cache (system-prompt KV reuse). Same atomic / non-strictly-consistent semantics as `full_cache` (§4.7). -`capacity = 0` means the cache is disabled. +`capacity = 0` means the cache is disabled. `max_resident_bytes` and +`resident_bytes` cover committed copied checkpoints used by concurrent paged +serving; a maximum of `0` means unlimited. `budget_skips` counts captures +declined because no single eligible LRU entry could make enough room. + +The capture and restore counters expose synchronous time spent copying +checkpoint state on the scheduler thread. `*_attempts` include successful and +unsuccessful operations, totals are cumulative milliseconds, and maxima are +the largest single measured operation. A failed capture increments +`capture_failures`; an unusable restore increments `restore_invalidations`. ### 4.13 `reasoning` @@ -609,9 +630,21 @@ version increments. "threshold": null }, "prefix_cache": { - "capacity": 0, - "in_use": 0, - "lifetime_hits": 0 + "capacity": 0, + "in_use": 0, + "lifetime_hits": 0, + "agent_turn_enabled": false, + "max_resident_bytes": 4294967296, + "resident_bytes": 0, + "budget_skips": 0, + "capture_attempts": 0, + "capture_failures": 0, + "capture_stall_ms_total": 0.0, + "capture_stall_ms_max": 0.0, + "restore_attempts": 0, + "restore_invalidations": 0, + "restore_stall_ms_total": 0.0, + "restore_stall_ms_max": 0.0 }, "reasoning": { "default": null, diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5f0cd92bb..d5c8d6980 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1434,6 +1434,15 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_seq_engine_contract) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_parallel_prefix_txn.cpp") + # Host-only ownership tests for the scheduler/cache capture ticket. + add_executable(test_parallel_prefix_txn + test/test_parallel_prefix_txn.cpp) + target_include_directories(test_parallel_prefix_txn PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_parallel_prefix_txn) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_batch_plan.cpp") # Pure-host tests for model-neutral token-budget/FIFO planning. add_executable(test_seq_batch_plan test/test_seq_batch_plan.cpp) diff --git a/server/docs/PREFIX_CACHE.md b/server/docs/PREFIX_CACHE.md index 2d095050d..20bc662a0 100644 --- a/server/docs/PREFIX_CACHE.md +++ b/server/docs/PREFIX_CACHE.md @@ -195,6 +195,7 @@ free_snapshot_backend(snap_backend_, compute_backend_); // then backend | Server flag | Default | Description | |-------------|---------|-------------| | `--prefix-cache-slots N` | 32 | Max turn-boundary prefix cache slots | +| `--concurrent-prefix-cache-max-mib N` | 4096 | Resident RAM limit for copied concurrent paged checkpoints; `0` is unlimited | | `--prefill-cache-slots N` | 0 | Max exact full-prompt prefill cache slots | | `--skip-park` | false | Skip parking draft model during compress | @@ -202,8 +203,16 @@ free_snapshot_backend(snap_backend_, compute_backend_); // then backend With right-sized, CPU-resident snapshots the limiting resource is **system RAM**, not VRAM. Each slot costs approximately `cur_pos × 5 KB` (for Qwen3.5-27B Q8_0 KV), -so 32 slots with an average prefix of 2000 tokens ≈ 320 MB of system RAM — negligible -on most workstations. +so 32 slots with an average prefix of 2000 tokens use about 320 MB of system RAM. + +Concurrent paged serving measures the exact backend allocation required for +each checkpoint before copying it. The cache keeps committed checkpoints under +`--concurrent-prefix-cache-max-mib`: when necessary it replaces one eligible least-recently-used +entry, and if no single eligible entry can make enough room it skips the new +checkpoint without disturbing the committed cache. The configured limit covers +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. | Scenario | Typical prefix length | Recommended cap | |----------|----------------------|-----------------| diff --git a/server/src/common/concurrency/prefix_store.h b/server/src/common/concurrency/prefix_store.h new file mode 100644 index 000000000..17f43778d --- /dev/null +++ b/server/src/common/concurrency/prefix_store.h @@ -0,0 +1,92 @@ +// Model-neutral checkpoint protocol for continuous-batching prefix reuse. +// +// The scheduler owns token-prefix lookup and LRU policy. A SeqEngine receives +// only opaque checkpoint identities and logical token positions. The concrete +// engine owns checkpoint payloads and cache-layout-specific copies. Copied +// pages today and a shared-page/radix engine later use the same protocol. + +#pragma once + +#include +#include +#include + +namespace dflash::common { + +struct PrefixStoreRef { + uint64_t id = 0; + int tokens = 0; + + bool empty() const { return id == 0 && tokens == 0; } + bool valid() const { return id != 0 && tokens > 0; } +}; + +inline bool operator==(PrefixStoreRef a, PrefixStoreRef b) { + return a.id == b.id && a.tokens == b.tokens; +} + +inline bool operator!=(PrefixStoreRef a, PrefixStoreRef b) { + return !(a == b); +} + +struct PrefixCaptureTicket { + uint64_t id = 0; + PrefixStoreRef checkpoint; + + bool valid() const { return id != 0 && checkpoint.valid(); } +}; + +inline bool operator==(const PrefixCaptureTicket & a, + const PrefixCaptureTicket & b) { + return a.id == b.id && a.checkpoint == b.checkpoint; +} + +inline bool operator!=(const PrefixCaptureTicket & a, + const PrefixCaptureTicket & b) { + return !(a == b); +} + +struct PrefixStorePlan { + PrefixStoreRef restore; + PrefixCaptureTicket capture; +}; + +struct PrefixStoreAdmission { + PrefixStoreRef restored; + PrefixStoreRef invalidated; + PrefixCaptureTicket capture; + // True only after the engine begins validating/copying a requested + // restore. Keep this independent from the result references: malformed + // references must still count as attempts and be rejected by the caller. + bool restore_attempted = false; + // Wall time spent validating/copying a requested restore. Non-zero for + // both successful restores and invalidations so operators can see stalls. + uint64_t restore_elapsed_us = 0; + + bool malformed_restore_state() const { + const bool has_result = !restored.empty() || !invalidated.empty(); + return restore_attempted != has_result || + (!restored.empty() && !restored.valid()) || + (!invalidated.empty() && !invalidated.valid()); + } +}; + +struct PrefixStoreEvent { + enum class Status { + none, + saved, + failed, + }; + + Status status = Status::none; + PrefixCaptureTicket ticket; + std::string error; + // Actual committed payload size. Present only for `saved`. + size_t bytes = 0; + // Wall time spent in the capture attempt, including a failed copy. + uint64_t elapsed_us = 0; + + bool attempted() const { return status != Status::none; } +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa4dbba35..4e036d581 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -55,6 +55,7 @@ #include #include "common/sampler.h" +#include "prefix_store.h" namespace dflash::common { @@ -161,6 +162,7 @@ class SeqEngine { Status status = Status::failed; int slot = -1; std::string error; + PrefixStoreAdmission prefix_store; }; // Admit one request into a free slot and queue its prompt for chunked @@ -178,6 +180,23 @@ class SeqEngine { const std::vector & prompt, const SamplerCfg & sampler) = 0; + // Optional prefix-checkpoint admission. Unsupported engines remain on + // cold admission and never receive a plan from the scheduler. + virtual bool supports_prefix_store() const { return false; } + // Conservative resident-byte estimate for one checkpoint. Returning zero + // means the engine cannot safely participate in a configured byte budget. + virtual size_t estimate_prefix_store_bytes(int) const { return 0; } + virtual AdmitResult admit_with_prefix( + uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler, + const PrefixStorePlan &) { + return admit(request_id, prompt, sampler); + } + + // Release engine-owned payload without touching server policy metadata. + virtual void discard_prefix_store(PrefixStoreRef) {} + struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position @@ -205,6 +224,10 @@ class SeqEngine { int32_t token = -1; // Present only for failed. std::string error; + // A capture ending on this successfully-computed prefill boundary. + // Capture failure does not fail generation: the scheduler invalidates + // the reserved cache entry and continues the request cold. + PrefixStoreEvent prefix_store; }; // One scheduler iteration owns both kinds of logical work. `decode` must @@ -318,6 +341,35 @@ inline std::string validate_step_result( if (output.status == PrefillStatus::failed && (output.token >= 0 || output.error.empty())) return "failed prefill has invalid payload"; + const PrefixStoreEvent & store = output.prefix_store; + if (store.status == PrefixStoreEvent::Status::none) { + if (store.ticket.id != 0 || + store.ticket.checkpoint.id != 0 || + store.ticket.checkpoint.tokens != 0 || + !store.error.empty() || + store.bytes != 0 || store.elapsed_us != 0) + return "inactive prefix capture carries payload"; + } else { + if (store.status != PrefixStoreEvent::Status::saved && + store.status != PrefixStoreEvent::Status::failed) + return "prefix capture has an unknown status"; + if (!store.ticket.valid()) + return "prefix capture has an invalid ticket"; + if (output.status == PrefillStatus::failed) + return "failed prefill carries a prefix capture"; + if (store.status == PrefixStoreEvent::Status::saved && + !store.error.empty()) + return "saved prefix capture carries an error"; + if (store.status == PrefixStoreEvent::Status::saved && + store.bytes == 0) + return "saved prefix capture omits its byte size"; + if (store.status == PrefixStoreEvent::Status::failed && + store.error.empty()) + return "failed prefix capture omits its error"; + if (store.status == PrefixStoreEvent::Status::failed && + store.bytes != 0) + return "failed prefix capture carries committed bytes"; + } prefill_seen[(size_t)output.slot] = 1; } diff --git a/server/src/internal.h b/server/src/internal.h index e7eaaa898..01e8cd062 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -524,11 +524,12 @@ struct PrefixSnapshot { ggml_context * ctx = nullptr; ggml_backend_buffer_t buf = nullptr; - // Phase B: thin-mode snapshots cover only a KV-position range. - bool is_thin = false; - int kv_start = 0; // inclusive (only meaningful when is_thin) - int kv_end = 0; // exclusive (only meaningful when is_thin) - // When is_thin == true: + // Snapshot payload shape; one value avoids impossible flag combinations. + enum class Layout { empty, dense, thin, paged }; + Layout layout = Layout::empty; + int kv_start = 0; // inclusive (only meaningful for Layout::thin) + int kv_end = 0; // exclusive (only meaningful for Layout::thin) + // For Layout::thin: // - attn_k_snap[i] / attn_v_snap[i] are sized // [HEAD_DIM, kv_end-kv_start, N_HEAD_KV] (smaller than cache). // - ssm_state_snap, conv_state_snap, target_feat_snap are NOT @@ -554,6 +555,45 @@ bool restore_target_cache(const PrefixSnapshot & snap, TargetCache & cache); // Free the snapshot's GPU buffers. 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. +size_t estimate_paged_target_cache_snapshot_bytes( + const TargetCache & cache, int token_count); + +// 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. +bool snapshot_paged_target_cache( + const TargetCache & cache, + int seq_slot, + const std::vector & block_table, + int block_size, + int token_count, + PrefixSnapshot & snap); + +// Atomically replace a paged snapshot. The incumbent remains valid when +// allocation, layout validation, or any staged copy fails. +bool replace_paged_target_cache( + const TargetCache & cache, + int seq_slot, + const std::vector & block_table, + int block_size, + int token_count, + PrefixSnapshot & destination); + +// 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. +bool restore_paged_target_cache( + const PrefixSnapshot & snap, + TargetCache & cache, + int seq_slot, + const std::vector & block_table, + int block_size); + // Thin snapshot: capture only KV slice [kv_start, kv_end). // SSM/conv/target_feat are not preserved (caller chains thin entries // onto a thick base via restore_target_cache_chain). diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 0a54761f9..66c706ad2 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -15,6 +15,7 @@ #include "internal.h" #include +#include #include #include #include @@ -50,6 +51,158 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( return result; } +size_t Qwen35SeqEngine::estimate_prefix_store_bytes(int tokens) const { + return estimate_paged_target_cache_snapshot_bytes(b_.cache_, tokens); +} + +int Qwen35SeqEngine::checkpoint_index(PrefixStoreRef checkpoint) const { + if (!checkpoint.valid() || + checkpoint.id > (uint64_t)Qwen35Backend::PREFIX_SLOTS) return -1; + return (int)checkpoint.id - 1; +} + +void Qwen35SeqEngine::discard_prefix_store(PrefixStoreRef checkpoint) { + const int index = checkpoint_index(checkpoint); + if (index >= 0 && + b_.prefix_snapshots_[index].cur_pos == checkpoint.tokens) { + free_prefix_snapshot(b_.prefix_snapshots_[index]); + } +} + +bool Qwen35SeqEngine::arm_capture( + int slot, PrefixCaptureTicket ticket, int restored_tokens) { + if (slot < 0 || slot >= slots_.slot_count()) return false; + const Qwen35Slot & sequence = slots_.slot(slot); + if (!ticket.valid() || checkpoint_index(ticket.checkpoint) < 0 || + ticket.checkpoint.tokens <= restored_tokens || + ticket.checkpoint.tokens > sequence.prompt_len) return false; + slots_.slot(slot).pending_capture = ticket; + return true; +} + +SeqEngine::AdmitResult Qwen35SeqEngine::admit_with_prefix( + uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler, + const PrefixStorePlan & plan) { + AdmitResult result = slots_.admit(request_id, prompt, sampler); + if (result.status != AdmitResult::Status::admitted) return result; + + const int slot = result.slot; + slots_.slot(slot).pending_capture = {}; + bool restored = false; + if (plan.restore.valid()) { + const int restore_index = checkpoint_index(plan.restore); + const bool metadata_valid = + restore_index >= 0 && + plan.restore.tokens < (int)prompt.size(); + PrefixSnapshot * snap = metadata_valid + ? &b_.prefix_snapshots_[restore_index] : nullptr; + const auto restore_started = std::chrono::steady_clock::now(); + if (snap && snap->ctx && + snap->layout == PrefixSnapshot::Layout::paged && + snap->cur_pos == plan.restore.tokens) { + Qwen35SlotManager::PrefillChunk seeded = + slots_.seed_restored_prefix(slot, plan.restore.tokens); + PagedKvSequenceSnapshot sequence; + const bool pool_ok = seeded.ok && + pool_.sequence(slots_.slot(slot).handle, sequence) == + PagedKvStatus::Ok; + const bool table_ok = pool_ok && upload_block_table_delta( + slot, seeded.first_new_block, seeded.new_blocks.data(), + seeded.new_blocks.size()); + restored = table_ok && restore_paged_target_cache( + *snap, b_.cache_, slot, sequence.block_table, + (int)pool_.block_size()); + } + const uint64_t restore_elapsed_us = + (uint64_t)std::chrono::duration_cast( + std::chrono::steady_clock::now() - restore_started).count(); + + if (!restored) { + discard_prefix_store(plan.restore); + slots_.retire(slot); + result = slots_.admit(request_id, prompt, sampler); + result.prefix_store.invalidated = plan.restore; + if (result.status == AdmitResult::Status::admitted) { + reset_recurrent_slot(b_.cache_, result.slot); + } else { + result.error = + "cold admission failed after stale prefix restore"; + } + } else { + result.prefix_store.restored = plan.restore; + std::fprintf(stderr, + "[parallel-pc] restored checkpoint=%llu seq_slot=%d " + "tokens=%d time_ms=%.1f\n", + (unsigned long long)plan.restore.id, slot, + plan.restore.tokens, (double)restore_elapsed_us / 1000.0); + } + result.prefix_store.restore_attempted = true; + result.prefix_store.restore_elapsed_us = restore_elapsed_us; + if (result.status != AdmitResult::Status::admitted) return result; + } else { + reset_recurrent_slot(b_.cache_, slot); + } + + const int admitted_slot = result.slot; + if (!result.prefix_store.invalidated.valid() && + arm_capture( + admitted_slot, plan.capture, + result.prefix_store.restored.tokens)) { + result.prefix_store.capture = plan.capture; + } + return result; +} + +PrefixStoreEvent Qwen35SeqEngine::capture_prefix( + int slot, PrefixCaptureTicket ticket) { + PrefixStoreEvent event; + event.ticket = ticket; + event.status = PrefixStoreEvent::Status::failed; + const int checkpoint = checkpoint_index(ticket.checkpoint); + if (!ticket.valid() || checkpoint < 0 || + !slots_.is_prefilling(slot) || + slots_.slot(slot).cur_pos != ticket.checkpoint.tokens) { + event.error = "invalid prefix capture boundary"; + return event; + } + PagedKvSequenceSnapshot sequence; + if (pool_.sequence(slots_.slot(slot).handle, sequence) != + PagedKvStatus::Ok || + sequence.kv_seq_len != (uint32_t)ticket.checkpoint.tokens) { + event.error = "prefix capture page table is incomplete"; + return event; + } + const auto capture_started = std::chrono::steady_clock::now(); + PrefixSnapshot & snapshot = b_.prefix_snapshots_[checkpoint]; + if (!replace_paged_target_cache( + b_.cache_, slot, sequence.block_table, + (int)pool_.block_size(), ticket.checkpoint.tokens, snapshot)) { + event.elapsed_us = + (uint64_t)std::chrono::duration_cast( + std::chrono::steady_clock::now() - capture_started).count(); + event.error = dflash27b_last_error(); + if (event.error.empty()) { + event.error = "paged prefix capture failed"; + } + return event; + } + event.status = PrefixStoreEvent::Status::saved; + event.elapsed_us = + (uint64_t)std::chrono::duration_cast( + std::chrono::steady_clock::now() - capture_started).count(); + event.bytes = snapshot.buf + ? ggml_backend_buffer_get_size(snapshot.buf) : 0; + std::fprintf(stderr, + "[parallel-pc] saved checkpoint=%llu seq_slot=%d tokens=%d " + "bytes=%zu time_ms=%.1f\n", + (unsigned long long)ticket.checkpoint.id, slot, + ticket.checkpoint.tokens, event.bytes, + (double)event.elapsed_us / 1000.0); + return event; +} + int32_t Qwen35SeqEngine::sample_graph_row( int slot, int logits_row, const int32_t * cached_argmax, std::vector * logits_scratch) { @@ -125,6 +278,14 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( stage.kv_pos = seq.cur_pos; stage.chunk = std::min( max_tokens, seq.prompt_len - stage.kv_pos); + const PrefixCaptureTicket capture = slots_.slot(slot).pending_capture; + if (capture.valid() && + stage.kv_pos < capture.checkpoint.tokens) { + // Recurrent state is only checkpoint-consistent after compute, so a + // selected capture boundary must be the exact end of this graph slice. + stage.chunk = std::min( + stage.chunk, capture.checkpoint.tokens - stage.kv_pos); + } if (stage.chunk <= 0) return PrefillStage{}; stage.commit = stage.kv_pos + stage.chunk >= seq.prompt_len; @@ -522,6 +683,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const int slot = plan.prefills[i].slot; PrefillOutput out; out.slot = slot; + const PrefixCaptureTicket capture = slots_.slot(slot).pending_capture; + if (capture.valid() && + slots_.slot(slot).cur_pos == capture.checkpoint.tokens) { + out.prefix_store = capture_prefix(slot, capture); + slots_.slot(slot).pending_capture = {}; + } if (prefills[i].commit) { out.status = PrefillOutput::Status::completed; out.token = sample_graph_row( diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index d9391784c..e897d148f 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -51,16 +51,26 @@ class Qwen35SeqEngine final : public SeqEngine { long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), long_prefill_threshold_(std::max(1, long_prefill_threshold)), idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), - prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), - slots_(pool, max_ctx), scratch_row_(scratch_row) {} + prefill_quantum_(std::max(1, prefill_quantum)), pool_(pool), + b_(backend), slots_(pool, max_ctx), scratch_row_(scratch_row) {} int slot_count() const override { return slots_.slot_count(); } int max_context() const override { return slots_.max_context(); } + bool supports_prefix_store() const override { return true; } + size_t estimate_prefix_store_bytes(int tokens) const override; + + void discard_prefix_store(PrefixStoreRef checkpoint) override; AdmitResult admit(uint64_t request_id, const std::vector & prompt, const SamplerCfg & sampler) override; + AdmitResult admit_with_prefix( + uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler, + const PrefixStorePlan & plan) override; + StepResult step(const StepPlan & plan) override; StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; @@ -112,7 +122,13 @@ class Qwen35SeqEngine final : public SeqEngine { int32_t sample_graph_row(int slot, int logits_row, const int32_t * cached_argmax = nullptr, std::vector * logits_scratch = nullptr); + PrefixStoreEvent capture_prefix( + int slot, PrefixCaptureTicket ticket); + bool arm_capture( + int slot, PrefixCaptureTicket ticket, int restored_tokens); + int checkpoint_index(PrefixStoreRef checkpoint) const; + PagedKvPool & pool_; Qwen35Backend & b_; Qwen35SlotManager slots_; int64_t scratch_row_ = 0; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index b8dcc36ec..dde1e6e22 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -223,6 +223,16 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( return out; } +Qwen35SlotManager::PrefillChunk Qwen35SlotManager::seed_restored_prefix( + int slot, int restored_tokens) { + if (!is_prefilling(slot) || slots_[(size_t)slot].cur_pos != 0 || + restored_tokens <= 0 || + restored_tokens >= slots_[(size_t)slot].prompt_len) { + return {}; + } + return append_prefill(slot, restored_tokens); +} + void Qwen35SlotManager::commit_prefill(int slot) { if (!is_prefilling(slot)) return; Qwen35Slot & s = slots_[(size_t)slot]; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 1da009f69..664ae050f 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -42,6 +42,7 @@ struct Qwen35Slot { int cur_pos = 0; SamplerCfg sampler; std::mt19937_64 rng{0x9E3779B97F4A7C15ull}; + PrefixCaptureTicket pending_capture; // Penalty history is recorded as fed rather than sampled: the scheduler // may override a sample before the model consumes it. std::vector sample_history; @@ -87,6 +88,13 @@ class Qwen35SlotManager { // admitted prompt is guaranteed not to wait on another sequence. PrefillChunk append_prefill(int slot, int n_tokens); + // Materialize a copied checkpoint into this request's freshly-reserved + // physical pages. Valid only before ordinary prefill has advanced. The + // returned rows and block-table delta are the destinations into which the + // engine scatters the checkpoint's logical K/V rows. The slot remains in + // prefill so the uncached suffix can continue normally. + PrefillChunk seed_restored_prefix(int slot, int restored_tokens); + // Record a finished prefill and expose the slot to decode. void commit_prefill(int slot); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 963ded27a..f274c7ae1 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -845,8 +845,9 @@ bool Qwen35Backend::snapshot_save(int slot) { static bool warned = false; if (!warned) { std::fprintf(stderr, - "[paged-attention] prefix snapshots are disabled until the " - "snapshot format stores block tables\n"); + "[paged-attention] the classic single-sequence snapshot API " + "is unavailable; continuous batching uses copied paged " + "checkpoints\n"); warned = true; } return false; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index e4615cbc2..6b29a642e 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -2389,7 +2389,9 @@ bool snapshot_target_cache(const TargetWeights & w, // Reuse existing buffer if shapes match (same cur_pos); otherwise reallocate. // Right-sized KV tensors use [head_dim, cur_pos, n_head_kv] — orders of // magnitude smaller than [head_dim, max_ctx, n_head_kv] for short prefixes. - const bool needs_alloc = (snap.ctx == nullptr) || (snap.cur_pos != snap_pos); + const bool needs_alloc = snap.ctx == nullptr || + snap.layout != PrefixSnapshot::Layout::dense || + snap.cur_pos != snap_pos; if (needs_alloc) { free_prefix_snapshot(snap); @@ -2504,11 +2506,16 @@ bool snapshot_target_cache(const TargetWeights & w, snap.kv_k_type = cache.kv_k_type; snap.max_ctx = cache.max_ctx; snap.target_feat_cap = cache.target_feat_cap; + snap.layout = PrefixSnapshot::Layout::dense; return true; } bool restore_target_cache(const PrefixSnapshot & snap, TargetCache & cache) { + if (snap.layout != PrefixSnapshot::Layout::dense) { + set_last_error("restore_target_cache: snapshot is not dense"); + return false; + } if (cache.n_seq_slots > 1) { set_last_error("restore_target_cache: multi-slot caches are unsupported"); return false; @@ -2593,6 +2600,379 @@ bool restore_target_cache(const PrefixSnapshot & snap, TargetCache & cache) { return true; } +namespace { + +size_t blocks_for_prefix(int tokens, int block_size) { + return tokens <= 0 ? 0 : + ((size_t)tokens + (size_t)block_size - 1) / (size_t)block_size; +} + +size_t contiguous_block_run( + const std::vector & blocks, + size_t first, + size_t block_count) { + size_t run = 1; + while (first + run < block_count && + (uint64_t)blocks[first + run] == + (uint64_t)blocks[first + run - 1] + 1) { + ++run; + } + return run; +} + +bool paged_tensor_layout_matches( + const ggml_tensor * dense, const ggml_tensor * paged, + int tokens) { + return dense && paged && dense->type == paged->type && + dense->ne[0] == paged->ne[0] && dense->ne[1] == tokens && + dense->ne[2] == paged->ne[2] && dense->ne[3] == paged->ne[3] && + ggml_is_contiguous(dense) && ggml_is_contiguous(paged); +} + +bool paged_rows_fit( + const ggml_tensor * tensor, + const std::vector & blocks, + int block_size, + int tokens) { + if (!tensor || block_size <= 0 || tokens <= 0 || + blocks.size() < blocks_for_prefix(tokens, block_size)) { + return false; + } + for (size_t logical = 0; + logical < blocks_for_prefix(tokens, block_size); ++logical) { + const uint64_t first = + (uint64_t)blocks[logical] * (uint64_t)block_size; + const int remaining = tokens - (int)logical * block_size; + const uint64_t count = (uint64_t)std::min(block_size, remaining); + if (first + count > (uint64_t)tensor->ne[1]) return false; + } + return true; +} + +ggml_backend_buffer_type_t paged_snapshot_buffer_type() { + // Paged gather/scatter passes snapshot tensor data to get/set as host + // staging, so even unified-memory compute backends need true CPU storage. + return ggml_backend_cpu_buffer_type(); +} + +enum class PagedCopyDirection { gather, scatter }; + +void copy_paged_tensor( + ggml_backend_t backend, + ggml_tensor * dense, ggml_tensor * paged, + const std::vector & blocks, + int block_size, int tokens, PagedCopyDirection direction) { + const size_t block_count = blocks_for_prefix(tokens, block_size); + for (int head = 0; head < (int)paged->ne[2]; ++head) { + for (size_t logical = 0; logical < block_count;) { + const size_t run = + contiguous_block_run(blocks, logical, block_count); + const int logical_row = (int)logical * block_size; + const int count = std::min( + (int)run * block_size, tokens - logical_row); + const size_t bytes = (size_t)count * paged->nb[1]; + const size_t paged_row = + (size_t)blocks[logical] * (size_t)block_size; + const size_t paged_offset = + (size_t)head * paged->nb[2] + paged_row * paged->nb[1]; + const size_t dense_offset = + (size_t)head * dense->nb[2] + + (size_t)logical_row * dense->nb[1]; + if (direction == PagedCopyDirection::gather) { + ggml_backend_tensor_get_async( + backend, paged, (char *)dense->data + dense_offset, + paged_offset, bytes); + } else { + ggml_backend_tensor_set_async( + backend, paged, (const char *)dense->data + dense_offset, + paged_offset, bytes); + } + logical += run; + } + } +} + +bool recurrent_slab_matches( + const ggml_tensor * dense, const ggml_tensor * slotted, + int n_slots, int slot_axis) { + if (!dense || !slotted || n_slots < 1 || slot_axis < 0 || + slot_axis >= GGML_MAX_DIMS || dense->type != slotted->type || + !ggml_is_contiguous(dense) || !ggml_is_contiguous(slotted) || + slotted->ne[slot_axis] != n_slots || + dense->ne[slot_axis] != 1) return false; + for (int axis = 0; axis < GGML_MAX_DIMS; ++axis) { + if (axis != slot_axis && dense->ne[axis] != slotted->ne[axis]) + return false; + } + return ggml_nbytes(dense) == + ggml_nbytes(slotted) / (size_t)n_slots; +} + +bool paged_cache_pairs_complete(const TargetCache & cache) { + if (cache.attn_k.size() != cache.attn_v.size() || + cache.ssm_state.size() != cache.conv_state.size()) return false; + for (size_t i = 0; i < cache.attn_k.size(); ++i) { + if ((cache.attn_k[i] == nullptr) != + (cache.attn_v[i] == nullptr)) return false; + } + for (size_t i = 0; i < cache.ssm_state.size(); ++i) { + if ((cache.ssm_state[i] == nullptr) != + (cache.conv_state[i] == nullptr)) return false; + } + return true; +} + +bool create_paged_snapshot_layout( + const TargetCache & cache, int token_count, PrefixSnapshot & snap) { + const int total_tensors = 2 * (int)cache.attn_k.size() + + 2 * (int)cache.ssm_state.size(); + ggml_init_params params{}; + params.mem_size = + (size_t)(total_tensors + 16) * ggml_tensor_overhead(); + params.no_alloc = true; + snap.ctx = ggml_init(params); + if (!snap.ctx) return false; + + 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); + snap.conv_state_snap.assign(cache.conv_state.size(), nullptr); + for (size_t i = 0; i < cache.attn_k.size(); ++i) { + if (!cache.attn_k[i]) continue; + snap.attn_k_snap[i] = ggml_new_tensor_3d( + snap.ctx, cache.attn_k[i]->type, cache.attn_k[i]->ne[0], + token_count, cache.attn_k[i]->ne[2]); + snap.attn_v_snap[i] = ggml_new_tensor_3d( + snap.ctx, cache.attn_v[i]->type, cache.attn_v[i]->ne[0], + token_count, cache.attn_v[i]->ne[2]); + char name[64]; + std::snprintf(name, sizeof(name), "snap_cache_k_%zu", i); + ggml_set_name(snap.attn_k_snap[i], name); + std::snprintf(name, sizeof(name), "snap_cache_v_%zu", i); + ggml_set_name(snap.attn_v_snap[i], name); + } + for (size_t i = 0; i < cache.ssm_state.size(); ++i) { + if (!cache.ssm_state[i]) continue; + const ggml_tensor * ssm = cache.ssm_state[i]; + const ggml_tensor * conv = cache.conv_state[i]; + snap.ssm_state_snap[i] = ggml_new_tensor_3d( + snap.ctx, ssm->type, ssm->ne[0], ssm->ne[1], ssm->ne[2]); + snap.conv_state_snap[i] = ggml_new_tensor_2d( + snap.ctx, conv->type, conv->ne[0], conv->ne[1]); + char name[64]; + std::snprintf(name, sizeof(name), "snap_ssm_state_%zu", i); + ggml_set_name(snap.ssm_state_snap[i], name); + std::snprintf(name, sizeof(name), "snap_conv_state_%zu", i); + ggml_set_name(snap.conv_state_snap[i], name); + } + return true; +} + +bool paged_snapshot_matches( + const TargetCache & cache, const PrefixSnapshot & snap, + const std::vector & blocks, int block_size, int tokens) { + if (!paged_cache_pairs_complete(cache) || + 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()) return false; + + for (size_t i = 0; i < cache.attn_k.size(); ++i) { + const bool present = cache.attn_k[i] != nullptr; + if ((snap.attn_k_snap[i] != nullptr) != present || + (snap.attn_v_snap[i] != nullptr) != present) return false; + if (present && + (!paged_tensor_layout_matches( + snap.attn_k_snap[i], cache.attn_k[i], tokens) || + !paged_tensor_layout_matches( + snap.attn_v_snap[i], cache.attn_v[i], tokens) || + !paged_rows_fit(cache.attn_k[i], blocks, block_size, tokens) || + !paged_rows_fit(cache.attn_v[i], blocks, block_size, tokens))) { + return false; + } + } + for (size_t i = 0; i < cache.ssm_state.size(); ++i) { + const bool present = cache.ssm_state[i] != nullptr; + if ((snap.ssm_state_snap[i] != nullptr) != present || + (snap.conv_state_snap[i] != nullptr) != present) return false; + if (present && + (!recurrent_slab_matches( + snap.ssm_state_snap[i], cache.ssm_state[i], + cache.n_seq_slots, /*slot_axis=*/3) || + !recurrent_slab_matches( + snap.conv_state_snap[i], cache.conv_state[i], + cache.n_seq_slots, /*slot_axis=*/2))) return false; + } + return true; +} + +} // namespace + +size_t estimate_paged_target_cache_snapshot_bytes( + const TargetCache & cache, + int token_count) { + 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; + const size_t bytes = ggml_backend_alloc_ctx_tensors_from_buft_size( + layout.ctx, + paged_snapshot_buffer_type()); + free_prefix_snapshot(layout); + return bytes; +} + +bool snapshot_paged_target_cache( + const TargetCache & cache, + int seq_slot, + const std::vector & block_table, + int block_size, + int token_count, + PrefixSnapshot & snap) { + 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 || + block_table.size() < blocks_for_prefix(token_count, block_size) || + !paged_cache_pairs_complete(cache)) { + set_last_error("snapshot_paged_target_cache: invalid arguments"); + return false; + } + 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(); + if (needs_alloc) { + free_prefix_snapshot(snap); + if (!create_paged_snapshot_layout(cache, token_count, snap)) { + set_last_error("paged PrefixSnapshot ggml_init failed"); + return false; + } + snap.buf = ggml_backend_alloc_ctx_tensors_from_buft( + snap.ctx, paged_snapshot_buffer_type()); + if (!snap.buf) { + set_last_error("paged PrefixSnapshot buffer allocation failed"); + free_prefix_snapshot(snap); + return false; + } + } + + if (!paged_snapshot_matches( + cache, snap, block_table, block_size, token_count)) { + set_last_error("paged snapshot layout mismatch"); + free_prefix_snapshot(snap); + return false; + } + + // Validate the entire topology before submitting any transfers. Then queue + // every logical K/V run and recurrent slab and synchronize once, avoiding + // one device-wide scheduler stall per tensor/run. + for (size_t i = 0; i < cache.attn_k.size(); ++i) { + if (!cache.attn_k[i]) continue; + copy_paged_tensor( + cache.backend, snap.attn_k_snap[i], cache.attn_k[i], + block_table, block_size, token_count, + PagedCopyDirection::gather); + copy_paged_tensor( + cache.backend, snap.attn_v_snap[i], cache.attn_v[i], + block_table, block_size, token_count, + PagedCopyDirection::gather); + } + for (size_t i = 0; i < cache.ssm_state.size(); ++i) { + if (!cache.ssm_state[i]) continue; + const size_t ssm_bytes = ggml_nbytes(snap.ssm_state_snap[i]); + const size_t conv_bytes = ggml_nbytes(snap.conv_state_snap[i]); + ggml_backend_tensor_get_async( + cache.backend, + cache.ssm_state[i], snap.ssm_state_snap[i]->data, + (size_t)seq_slot * ssm_bytes, ssm_bytes); + ggml_backend_tensor_get_async( + cache.backend, + cache.conv_state[i], snap.conv_state_snap[i]->data, + (size_t)seq_slot * conv_bytes, conv_bytes); + } + 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; +} + +bool replace_paged_target_cache( + const TargetCache & cache, + int seq_slot, + const std::vector & block_table, + int block_size, + int token_count, + PrefixSnapshot & destination) { + PrefixSnapshot candidate; + if (!snapshot_paged_target_cache( + cache, seq_slot, block_table, block_size, token_count, + candidate)) { + free_prefix_snapshot(candidate); + return false; + } + using std::swap; + swap(candidate, destination); + free_prefix_snapshot(candidate); + return true; +} + +bool restore_paged_target_cache( + const PrefixSnapshot & snap, + TargetCache & cache, + int seq_slot, + const std::vector & block_table, + int block_size) { + if (!snap.ctx || snap.layout != PrefixSnapshot::Layout::paged || + !cache.backend || + cache.n_seq_slots < 1 || + seq_slot < 0 || seq_slot >= cache.n_seq_slots || block_size <= 0 || + snap.cur_pos <= 0 || snap.cur_pos > cache.max_ctx || + snap.max_ctx != cache.max_ctx || snap.kv_k_type != cache.kv_k_type || + block_table.size() < blocks_for_prefix(snap.cur_pos, block_size) || + !paged_snapshot_matches( + cache, snap, block_table, block_size, snap.cur_pos)) { + set_last_error("restore_paged_target_cache: incompatible checkpoint"); + return false; + } + for (size_t i = 0; i < cache.attn_k.size(); ++i) { + if (!cache.attn_k[i]) continue; + copy_paged_tensor( + cache.backend, snap.attn_k_snap[i], cache.attn_k[i], + block_table, block_size, snap.cur_pos, + PagedCopyDirection::scatter); + copy_paged_tensor( + cache.backend, snap.attn_v_snap[i], cache.attn_v[i], + block_table, block_size, snap.cur_pos, + PagedCopyDirection::scatter); + } + for (size_t i = 0; i < cache.ssm_state.size(); ++i) { + if (!cache.ssm_state[i]) continue; + const size_t ssm_bytes = ggml_nbytes(snap.ssm_state_snap[i]); + const size_t conv_bytes = ggml_nbytes(snap.conv_state_snap[i]); + ggml_backend_tensor_set_async( + cache.backend, + cache.ssm_state[i], snap.ssm_state_snap[i]->data, + (size_t)seq_slot * ssm_bytes, ssm_bytes); + ggml_backend_tensor_set_async( + cache.backend, + cache.conv_state[i], snap.conv_state_snap[i]->data, + (size_t)seq_slot * conv_bytes, conv_bytes); + } + ggml_backend_synchronize(cache.backend); + return true; +} + void free_prefix_snapshot(PrefixSnapshot & snap) { if (snap.buf) { ggml_backend_buffer_free(snap.buf); snap.buf = nullptr; } if (snap.ctx) { ggml_free(snap.ctx); snap.ctx = nullptr; } @@ -2605,7 +2985,7 @@ void free_prefix_snapshot(PrefixSnapshot & snap) { snap.kv_k_type = GGML_TYPE_COUNT; snap.max_ctx = 0; snap.target_feat_cap = 0; - snap.is_thin = false; + snap.layout = PrefixSnapshot::Layout::empty; snap.kv_start = 0; snap.kv_end = 0; } @@ -2631,7 +3011,7 @@ bool snapshot_target_cache_thin(const TargetWeights & w, // Lazy alloc; if snap was already a THIN with same range, reuse. bool needs_alloc = (snap.ctx == nullptr) || - !snap.is_thin || + snap.layout != PrefixSnapshot::Layout::thin || snap.kv_start != kv_start || snap.kv_end != kv_end; if (needs_alloc) { @@ -2691,7 +3071,7 @@ bool snapshot_target_cache_thin(const TargetWeights & w, ggml_backend_tensor_set(dv, bufv.data(), v_dst, v_strip); } } - snap.is_thin = true; + snap.layout = PrefixSnapshot::Layout::thin; snap.kv_start = kv_start; snap.kv_end = kv_end; snap.cur_pos = kv_end; @@ -2706,7 +3086,7 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, TargetCache & cache) { // Step 1: restore thick base if provided. if (thick) { - if (thick->is_thin) { + if (thick->layout == PrefixSnapshot::Layout::thin) { set_last_error("restore_chain: 'thick' arg is actually a thin snapshot"); return false; } @@ -2716,8 +3096,8 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, int max_kv_end = cache.cur_pos; for (int t = 0; t < n_thins; t++) { const PrefixSnapshot * thin = thins[t]; - if (!thin->is_thin) { - set_last_error("restore_chain: 'thin' arg has is_thin=false"); + if (thin->layout != PrefixSnapshot::Layout::thin) { + set_last_error("restore_chain: 'thin' arg has the wrong layout"); return false; } if (thin->kv_k_type != cache.kv_k_type || diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 20a758880..2e12d094e 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -876,10 +876,25 @@ json build_props_body(const ServerConfig & config, }}, {"pflash", pflash}, {"prefix_cache", { - {"capacity", pcs.capacity}, - {"in_use", pcs.in_use}, - {"lifetime_hits", pcs.lifetime_hits}, + {"capacity", pcs.capacity}, + {"in_use", pcs.in_use}, + {"lifetime_hits", pcs.lifetime_hits}, {"agent_turn_enabled", config.agent_turn_cache}, + {"max_resident_bytes", pcs.max_resident_bytes}, + {"resident_bytes", pcs.resident_bytes}, + {"budget_skips", pcs.budget_skips}, + {"capture_attempts", pcs.capture_attempts}, + {"capture_failures", pcs.capture_failures}, + {"capture_stall_ms_total", + (double)pcs.capture_stall_us_total / 1000.0}, + {"capture_stall_ms_max", + (double)pcs.capture_stall_us_max / 1000.0}, + {"restore_attempts", pcs.restore_attempts}, + {"restore_invalidations", pcs.restore_invalidations}, + {"restore_stall_ms_total", + (double)pcs.restore_stall_us_total / 1000.0}, + {"restore_stall_ms_max", + (double)pcs.restore_stall_us_max / 1000.0}, }}, {"full_cache", { {"enabled", pcfs.enabled}, @@ -1111,7 +1126,9 @@ HttpServer::HttpServer(ModelBackend & backend, , tokenizer_(tokenizer) , config_(config) , chat_format_(ChatFormat::QWEN3) // default, overridden by arch - , prefix_cache_(config.prefix_cache_cap, tokenizer) + , prefix_cache_(config.prefix_cache_cap, tokenizer, + config.concurrent_paged_prefix_cache + ? config.concurrent_prefix_cache_max_bytes : 0) , disk_cache_({config.disk_cache_dir, config.disk_cache_budget_mb * (size_t)(1024 * 1024), config.disk_cache_min_tokens, @@ -3275,7 +3292,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( // invalidating both ownership tables is unambiguous. forget_inline_slot_metadata(cache.cache_slot); backend_.snapshot_free(cache.cache_slot); - prefix_cache_.abort_inline_snap(cache.cache_slot); + prefix_cache_.invalidate_inline_snap(cache.cache_slot); prefix_cache_.abort_full_snap(cache.cache_slot); } cache.cache_slot = -1; @@ -3335,13 +3352,13 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( // requests prefer the reusable system/tool boundary; otherwise an // enabled exact full-prompt cache retains its existing priority. auto prepare_inline = [&]() { - const auto prepared_snapshot = prefix_cache_.prepare_inline_snap( + cache.snap_reservation = prefix_cache_.reserve_inline_snap( effective_prompt, cache.using_restore ? cache.prefix_len : 0, prefer_tools_boundary, forced_cut); - cache.snap_slot = prepared_snapshot.first; - cache.snap_cut = prepared_snapshot.second; + cache.snap_slot = cache.snap_reservation.slot(); + cache.snap_cut = cache.snap_reservation.target_cut(); }; auto prepare_full = [&]() { const auto & full_key = cache.full_snap_key_effective @@ -3373,7 +3390,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( // checkpoint; preserving the current hit is better than invalidating // it before restore starts. if (cache.using_restore && cache.snap_slot == cache.cache_slot) { - prefix_cache_.cancel_inline_snap(cache.snap_slot); + cache.snap_reservation.cancel(); cache.snap_slot = -1; cache.snap_cut = 0; } @@ -3410,7 +3427,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( void HttpServer::finalize_generation_cache( const ParsedRequest & req, const PreparedPrompt & prepared, - const GenerationCacheState & cache, const GenerateResult & result, + GenerationCacheState & cache, const GenerateResult & result, int completion_tokens, bool visible_output_seen, bool client_disconnected) { const auto & effective_prompt = prepared.tokens; @@ -3447,27 +3464,28 @@ void HttpServer::finalize_generation_cache( std::fprintf(stderr, "[pc] inline snapshot requested=%d saved=%d slot=%d\n", cache.snap_cut, saved_position, cache.snap_slot); - prefix_cache_.confirm_inline_snap( - cache.snap_slot, cache.snap_cut, effective_prompt); - // Track for shutdown save. The key may be stricter than a - // Qwen chunk-aligned snapshot, which is safe: matching the - // longer token prefix necessarily matches saved KV rows. + cache.snap_reservation.commit_at( + effective_prompt, saved_position); + // Track the same prefix published by the in-memory cache. + // Some backends may save short of the requested cut, so the + // shutdown key must not claim rows the snapshot lacks. slot_tokens_[cache.snap_slot] = std::vector( effective_prompt.begin(), - effective_prompt.begin() + cache.snap_cut); + effective_prompt.begin() + saved_position); if (!disk_cache_.disabled()) { disk_cache_.learn_layout(cache.snap_slot); if (cache.disk_policy.mode == DiskPrefixCacheMode::Full) { - disk_cache_.save(cache.snap_slot, effective_prompt); + disk_cache_.save( + cache.snap_slot, slot_tokens_[cache.snap_slot]); } } } else { backend_.snapshot_free(cache.snap_slot); - prefix_cache_.abort_inline_snap(cache.snap_slot); + cache.snap_reservation.abort(); } } else { backend_.snapshot_free(cache.snap_slot); - prefix_cache_.abort_inline_snap(cache.snap_slot); + cache.snap_reservation.abort(); } } @@ -3602,13 +3620,14 @@ void HttpServer::remember_agent_turn( } const int canonical_end = (int) canonical_tokens.size(); - const auto pending = prefix_cache_.prepare_inline_snap( + auto reservation = prefix_cache_.reserve_inline_snap( canonical_tokens, source_pos, false, canonical_end); - if (pending.first < 0 || pending.second != canonical_end) return; + if (!reservation.active() || + reservation.target_cut() != canonical_end) return; - const int slot = pending.first; + const int slot = reservation.slot(); if (slot == source_slot) { - prefix_cache_.cancel_inline_snap(slot); + reservation.cancel(); return; } forget_inline_slot_metadata(slot); @@ -3627,7 +3646,7 @@ void HttpServer::remember_agent_turn( const int saved_pos = replay_result.ok() && backend_.snapshot_used(slot) ? backend_.snapshot_cur_pos(slot) : 0; if (saved_pos > source_pos && saved_pos <= canonical_end) { - prefix_cache_.confirm_inline_snap(slot, saved_pos, canonical_tokens); + reservation.commit_at(canonical_tokens, saved_pos); canonical_tokens.resize((size_t) saved_pos); slot_tokens_[slot] = std::move(canonical_tokens); agent_turn_cache_slots_.insert(slot); @@ -3638,7 +3657,7 @@ void HttpServer::remember_agent_turn( canonical_end - source_pos); } else { backend_.snapshot_free(slot); - prefix_cache_.abort_inline_snap(slot); + reservation.abort(); } } diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index e5ca7a28b..18d9dd0e1 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -92,6 +92,11 @@ struct ServerConfig { bool enable_cors = true; std::string model_name = "dflash"; int prefix_cache_cap = 32; // prefix cache slots (0 disables) + // Resident system-memory budget for copied paged checkpoints. The + // scheduler enforces it only when concurrent paged prefix storage is + // active. Zero means unlimited. + size_t concurrent_prefix_cache_max_bytes = (size_t)4 * 1024 * 1024 * 1024; + bool concurrent_paged_prefix_cache = false; int prefill_cache_cap = 0; // full-prompt/prefill cache slots (0 disables) // Extend the existing prefix cache through generated tool-call turns. bool agent_turn_cache = false; @@ -365,6 +370,8 @@ class HttpServer { } private: + friend struct SchedulerTestHarness; + // Client thread: read HTTP request, parse, enqueue job, wait. void handle_client(SocketHandle fd); @@ -406,6 +413,7 @@ class HttpServer { // When DiffPin rewrote tokens, full-cache keys must use // prepared.tokens (effective), not req.prompt_tokens. bool full_snap_key_effective = false; + PrefixCache::InlineReservation snap_reservation; int snap_slot = -1; int snap_cut = 0; bool snap_prepared = false; @@ -416,7 +424,7 @@ class HttpServer { GenerateRequest & generate_request); void finalize_generation_cache( const ParsedRequest & req, const PreparedPrompt & prepared, - const GenerationCacheState & cache, const GenerateResult & result, + GenerationCacheState & cache, const GenerateResult & result, int completion_tokens, bool visible_output_seen, bool client_disconnected); void remember_agent_turn( diff --git a/server/src/server/parallel_prefix_txn.h b/server/src/server/parallel_prefix_txn.h new file mode 100644 index 000000000..d1afad99b --- /dev/null +++ b/server/src/server/parallel_prefix_txn.h @@ -0,0 +1,99 @@ +// Move-only ownership for one continuous-batching prefix capture. +// +// Policy metadata and checkpoint payload have different owners. This object +// keeps their resolution ordered and makes every early exit cancel exactly +// one untouched reservation. A malformed saved outcome discards only this +// transaction's payload and metadata; it never trusts an event-supplied id. + +#pragma once + +#include "common/concurrency/prefix_store.h" + +#include +#include +#include + +namespace dflash::common { + +template +class BasicPrefixCaptureTxn { +public: + enum class Resolution { + inactive, + saved, + failed, + mismatched, + }; + + BasicPrefixCaptureTxn() = default; + + BasicPrefixCaptureTxn( + Reservation reservation, Engine & engine, + PrefixCaptureTicket ticket) + : reservation_(std::move(reservation)), engine_(&engine), + ticket_(ticket) {} + + ~BasicPrefixCaptureTxn() { cancel(); } + + BasicPrefixCaptureTxn(const BasicPrefixCaptureTxn &) = delete; + BasicPrefixCaptureTxn & operator=( + const BasicPrefixCaptureTxn &) = delete; + BasicPrefixCaptureTxn(BasicPrefixCaptureTxn &&) noexcept = default; + BasicPrefixCaptureTxn & operator=( + BasicPrefixCaptureTxn &&) noexcept = default; + + bool active() const { + return engine_ && ticket_.valid() && reservation_.active() && + ticket_.checkpoint == PrefixStoreRef{ + (uint64_t)reservation_.slot() + 1, + reservation_.target_cut()}; + } + + Resolution resolve( + const PrefixStoreEvent & event, + const std::vector & prompt) { + if (!active()) return Resolution::inactive; + if (!event.attempted() || event.ticket != ticket_) { + if (event.status == PrefixStoreEvent::Status::saved) { + discard_saved(); + } else { + cancel(); + } + return Resolution::mismatched; + } + if (event.status == PrefixStoreEvent::Status::saved) { + const bool committed = reservation_.commit(prompt, event.bytes); + clear(); + return committed ? Resolution::saved : Resolution::failed; + } + if (event.status == PrefixStoreEvent::Status::failed) { + cancel(); + return Resolution::failed; + } + cancel(); + return Resolution::mismatched; + } + + void cancel() { + reservation_.cancel(); + clear(); + } + +private: + void discard_saved() { + engine_->discard_prefix_store(ticket_.checkpoint); + reservation_.abort(); + clear(); + } + + void clear() { + engine_ = nullptr; + ticket_ = {}; + } + + Reservation reservation_; + Engine * engine_ = nullptr; + PrefixCaptureTicket ticket_; +}; + +} // namespace dflash::common diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 5ad6bcb5d..218eb6f94 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -220,8 +220,10 @@ int select_inline_snapshot_boundary(const std::vector & boundaries, // ─── PrefixCache ──────────────────────────────────────────────────────── -PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) +PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer, + size_t max_resident_bytes) : cap_(std::min(cap, MAX_CACHE_SLOTS)) + , max_resident_bytes_(max_resident_bytes) { if (cap_ <= 0) { disabled_ = true; @@ -235,7 +237,13 @@ PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) return; } disabled_ = false; - std::fprintf(stderr, "[pc] enabled: cap=%d family=%s\n", cap_, markers_.family.c_str()); + if (max_resident_bytes_ > 0) { + std::fprintf(stderr, "[pc] enabled: cap=%d family=%s resident_budget=%zu MiB\n", + cap_, markers_.family.c_str(), max_resident_bytes_ / (1024 * 1024)); + } else { + std::fprintf(stderr, "[pc] enabled: cap=%d family=%s\n", + cap_, markers_.family.c_str()); + } } // ── LRU helpers ───────────────────────────────────────────────────────── @@ -247,6 +255,23 @@ int PrefixCache::find_entry(const PrefixHash & h) const { return -1; } +int PrefixCache::find_slot_entry(int slot) const { + for (int i = 0; i < (int)entries_.size(); ++i) { + if (entries_[(size_t)i].slot == slot) return i; + } + return -1; +} + +void PrefixCache::erase_inline_entry(int idx) { + if (idx < 0 || idx >= (int)entries_.size()) return; + const size_t bytes = entries_[(size_t)idx].resident_bytes; + resident_bytes_ = bytes <= resident_bytes_ ? resident_bytes_ - bytes : 0; + resident_bytes_count_.store( + (uint64_t)resident_bytes_, std::memory_order_relaxed); + entries_.erase(entries_.begin() + idx); + entries_size_count_.fetch_sub(1, std::memory_order_relaxed); +} + void PrefixCache::move_to_end(int idx) { if (idx < 0 || idx >= (int)entries_.size()) return; auto e = std::move(entries_[idx]); @@ -270,14 +295,30 @@ 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}; +std::pair PrefixCache::lookup( + const std::vector & prompt_ids) { + return lookup_impl( + prompt_ids, (int)prompt_ids.size(), /*record_hit=*/true); +} + +std::pair PrefixCache::lookup_candidate( + const std::vector & prompt_ids, + int max_prefix_tokens) { + return lookup_impl(prompt_ids, max_prefix_tokens, /*record_hit=*/false); +} + +std::pair PrefixCache::lookup_impl( + const std::vector & prompt_ids, + int max_prefix_tokens, + bool record_hit) { + if (disabled_ || max_prefix_tokens <= 0) return {-1, 0}; auto boundaries = find_all_boundaries(prompt_ids, markers_); int best_slot = -1, best_len = 0; int best_idx = -1; for (int cut : boundaries) { + if (cut > max_prefix_tokens) continue; auto key = hash_prefix(prompt_ids.data(), cut); int idx = find_entry(key); if (idx >= 0) { @@ -288,8 +329,7 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) std::fprintf(stderr, "[pc] lookup stale slot=%d key_cut=%d committed=%d — evicting\n", entries_[idx].slot, cut, committed); - entries_.erase(entries_.begin() + idx); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); + erase_inline_entry(idx); continue; } if (cut > best_len) { @@ -305,7 +345,8 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) for (int i = 0; i < (int)entries_.size(); ++i) { const auto & e = entries_[(size_t)i]; const int len = (int)e.ids.size(); - if (len <= best_len || len > (int)prompt_ids.size()) continue; + if (len <= best_len || len > (int)prompt_ids.size() || + len > max_prefix_tokens) continue; if (!std::equal(e.ids.begin(), e.ids.end(), prompt_ids.begin())) { continue; } @@ -314,23 +355,125 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) best_idx = i; } - if (best_idx >= 0) { - move_to_end(best_idx); + if (best_idx >= 0 && record_hit) + record_inline_hit(best_slot, best_len, prompt_ids.size()); + return {best_slot, best_len}; +} + +void PrefixCache::record_inline_hit( + int slot, int prefix_len, size_t prompt_len) { + for (int i = 0; i < (int)entries_.size(); ++i) { + if (entries_[(size_t)i].slot != slot || + (int)entries_[(size_t)i].ids.size() != prefix_len) continue; + move_to_end(i); 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 prefix_len=%d (of %zu total)\n", + slot, prefix_len, prompt_len); + return; } - return {best_slot, best_len}; } -std::pair PrefixCache::prepare_inline_snap( +PrefixCache::InlineReservation::InlineReservation( + PrefixCache * cache, uint64_t id, int slot, int target_cut, + PrefixHash victim, bool has_victim, bool protect) + : cache_(cache), id_(id), slot_(slot), target_cut_(target_cut), + victim_(victim), has_victim_(has_victim), protect_(protect) {} + +PrefixCache::InlineReservation::~InlineReservation() { cancel(); } + +PrefixCache::InlineReservation::InlineReservation( + InlineReservation && other) noexcept { + take(std::move(other)); +} + +PrefixCache::InlineReservation & +PrefixCache::InlineReservation::operator=( + InlineReservation && other) noexcept { + if (this != &other) { + cancel(); + take(std::move(other)); + } + return *this; +} + +bool PrefixCache::InlineReservation::active() const { + return cache_ && cache_->inline_reservation_active(id_); +} + +bool PrefixCache::InlineReservation::commit( + const std::vector & prompt_ids, + size_t resident_bytes, bool protect) { + if (!active()) { + clear(); + return false; + } + return cache_->commit_inline_reservation( + *this, prompt_ids, target_cut_, resident_bytes, protect); +} + +bool PrefixCache::InlineReservation::commit_at( + const std::vector & prompt_ids, int committed_cut, + size_t resident_bytes, bool protect) { + if (!active()) { + clear(); + return false; + } + return cache_->commit_inline_reservation( + *this, prompt_ids, committed_cut, resident_bytes, protect); +} + +void PrefixCache::InlineReservation::cancel() { + if (cache_) cache_->release_inline_reservation(id_); + clear(); +} + +void PrefixCache::InlineReservation::abort() { + if (active()) { + cache_->abort_inline_reservation(*this); + } else { + clear(); + } +} + +void PrefixCache::InlineReservation::clear() { + cache_ = nullptr; + id_ = 0; + slot_ = -1; + target_cut_ = 0; + victim_ = {}; + has_victim_ = false; + protect_ = false; +} + +void PrefixCache::InlineReservation::take(InlineReservation && other) { + cache_ = other.cache_; + id_ = other.id_; + slot_ = other.slot_; + target_cut_ = other.target_cut_; + victim_ = other.victim_; + has_victim_ = other.has_victim_; + protect_ = other.protect_; + other.clear(); +} + +bool PrefixCache::inline_reservation_active(uint64_t id) const { + return id != 0 && active_inline_reservation_ == id; +} + +void PrefixCache::release_inline_reservation(uint64_t id) { + if (inline_reservation_active(id)) active_inline_reservation_ = 0; +} + +PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( const std::vector & prompt_ids, int restored_prefix_len, bool prefer_tools_boundary, - int forced_cut) { - if (disabled_) return {-1, 0}; + int forced_cut, + InlineSnapshotSize estimate_bytes) { + if (disabled_ || active_inline_reservation_ != 0) return {}; - auto candidates = find_all_boundaries(prompt_ids, markers_); + const auto candidates = find_all_boundaries(prompt_ids, markers_); int target_cut = 0; bool forced = false; if (forced_cut > restored_prefix_len && @@ -341,117 +484,226 @@ std::pair PrefixCache::prepare_inline_snap( target_cut = select_inline_snapshot_boundary( candidates, restored_prefix_len, prefer_tools_boundary); } - if (target_cut <= 0) return {-1, 0}; - - auto key = hash_prefix(prompt_ids.data(), target_cut); - if (find_entry(key) >= 0) return {-1, 0}; // already cached + if (target_cut <= 0) return {}; - // Protect the tools head pin for tool-heavy requests so multi-chat deepen - // snaps cannot thrash the ~18k system+tools KV away. PPP forced cuts are - // the stable tools/identity span and stay protected as well. - pending_protect_ = prefer_tools_boundary && - (forced || - (!candidates.empty() && target_cut == candidates.front())); + const auto key = hash_prefix(prompt_ids.data(), target_cut); + if (find_entry(key) >= 0) return {}; - int slot; + const bool protect = prefer_tools_boundary && + (forced || (!candidates.empty() && target_cut == candidates.front())); + PrefixHash victim_key{}; + bool has_victim = false; + int slot = -1; 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. Skip protected tools pins when an unprotected leaf exists. std::vector *> ids_lru; std::vector protected_lru; ids_lru.reserve(entries_.size()); protected_lru.reserve(entries_.size()); - for (const auto & e : entries_) { - ids_lru.push_back(&e.ids); - protected_lru.push_back(e.protect); + for (const auto & entry : entries_) { + ids_lru.push_back(&entry.ids); + protected_lru.push_back(entry.protect); } - int victim = select_inline_evict_victim(ids_lru, &protected_lru); - pending_evict_key_ = entries_[victim].hash; - has_pending_evict_ = true; - slot = entries_[victim].slot; - if (victim != 0 || entries_[victim].protect) { + const int victim = select_inline_evict_victim( + ids_lru, &protected_lru); + victim_key = entries_[(size_t)victim].hash; + has_victim = true; + slot = entries_[(size_t)victim].slot; + if (victim != 0 || entries_[(size_t)victim].protect) { std::fprintf(stderr, "[pc] prefix-aware evict: victim idx=%d protect=%d (len=%zu) " "kept oldest ancestor (len=%zu)\n", - victim, (int)entries_[victim].protect, - entries_[victim].ids.size(), entries_.front().ids.size()); + victim, (int)entries_[(size_t)victim].protect, + entries_[(size_t)victim].ids.size(), + entries_.front().ids.size()); } } else { slot = next_slot_; next_slot_ = (next_slot_ + 1) % cap_; - has_pending_evict_ = false; } - return {slot, target_cut}; -} - -void PrefixCache::confirm_inline_snap(int slot, int target_cut, - const std::vector & prompt_ids, - bool protect) { - if (disabled_) 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); + if (max_resident_bytes_ > 0) { + const size_t estimated_bytes = estimate_bytes + ? estimate_bytes(target_cut) : 0; + const auto fits = [&](int victim_idx) { + if (estimated_bytes == 0 || + estimated_bytes > max_resident_bytes_) return false; + const size_t freed = victim_idx >= 0 + ? entries_[(size_t)victim_idx].resident_bytes : 0; + const size_t after_free = freed <= resident_bytes_ + ? resident_bytes_ - freed : 0; + return after_free <= max_resident_bytes_ && + estimated_bytes <= max_resident_bytes_ - after_free; + }; + + const int planned_victim = has_victim + ? find_entry(victim_key) : find_slot_entry(slot); + if (!fits(planned_victim)) { + const auto is_leaf = [&](int candidate) { + for (int i = 0; i < (int)entries_.size(); ++i) { + if (i != candidate && + is_strict_prefix(entries_[(size_t)candidate].ids, + entries_[(size_t)i].ids)) { + return false; + } + } + return true; + }; + int victim = -1; + for (int pass = 0; pass < 4 && victim < 0; ++pass) { + const bool require_leaf = pass < 2; + const bool allow_protected = pass == 1 || pass == 3; + for (int i = 0; i < (int)entries_.size(); ++i) { + if (!fits(i)) continue; + if (require_leaf && !is_leaf(i)) continue; + if (!allow_protected && entries_[(size_t)i].protect) + continue; + victim = i; + break; + } + } + if (victim >= 0) { + victim_key = entries_[(size_t)victim].hash; + has_victim = true; + const int redirected_slot = entries_[(size_t)victim].slot; + std::fprintf(stderr, + "[pc] resident budget redirects capture slot=%d -> %d " + "estimate=%zu resident=%zu budget=%zu\n", + slot, redirected_slot, estimated_bytes, resident_bytes_, + max_resident_bytes_); + slot = redirected_slot; + } else { + const uint64_t skipped = + budget_skips_.fetch_add(1, std::memory_order_relaxed) + 1; + if (skipped == 1 || skipped % 64 == 0) { + std::fprintf(stderr, + "[pc] resident budget skips capture estimate=%zu " + "resident=%zu budget=%zu (skips=%llu)\n", + estimated_bytes, resident_bytes_, max_resident_bytes_, + (unsigned long long)skipped); + } + return {}; + } } - has_pending_evict_ = false; } - // 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). + uint64_t id = next_inline_reservation_++; + if (id == 0) id = next_inline_reservation_++; + active_inline_reservation_ = id; + return InlineReservation( + this, id, slot, target_cut, victim_key, has_victim, protect); +} + +void PrefixCache::replace_inline_entry( + int slot, int target_cut, + const std::vector & prompt_ids, + bool protect, size_t resident_bytes) { 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); + erase_inline_entry(i); } } - const bool protect_entry = protect || pending_protect_; - pending_protect_ = false; - - 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), protect_entry}); + const 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), protect, resident_bytes}); entries_size_count_.fetch_add(1, std::memory_order_relaxed); + resident_bytes_ += resident_bytes; + resident_bytes_count_.store( + (uint64_t)resident_bytes_, std::memory_order_relaxed); std::fprintf(stderr, - "[pc] inline-snap committed slot=%d prefix_len=%d protect=%d\n", - slot, target_cut, (int)protect_entry); + "[pc] inline-snap committed slot=%d prefix_len=%d protect=%d " + "bytes=%zu resident=%zu\n", + slot, target_cut, (int)protect, resident_bytes, resident_bytes_); } -void PrefixCache::abort_inline_snap(int slot) { - if (disabled_) return; - // The HTTP layer clears the reserved backend slot before generation. Any - // metadata still pointing at it is therefore invalid, whether the slot was - // selected through the explicit eviction path or through a round-robin - // hole left by an earlier aborted reservation. +bool PrefixCache::commit_inline_reservation( + InlineReservation & reservation, + const std::vector & prompt_ids, + int committed_cut, size_t resident_bytes, bool protect) { + if (!reservation.active() || committed_cut <= 0 || + committed_cut > reservation.target_cut_ || + committed_cut > (int)prompt_ids.size()) { + reservation.abort(); + return false; + } + if (reservation.has_victim_) { + const int victim = find_entry(reservation.victim_); + if (victim >= 0) erase_inline_entry(victim); + } + replace_inline_entry( + reservation.slot_, committed_cut, prompt_ids, + protect || reservation.protect_, resident_bytes); + release_inline_reservation(reservation.id_); + reservation.clear(); + return true; +} + +void PrefixCache::abort_inline_reservation( + InlineReservation & reservation) { + if (!reservation.active()) { + reservation.clear(); + return; + } for (int i = (int)entries_.size() - 1; i >= 0; --i) { - if (entries_[(size_t)i].slot == slot) { - entries_.erase(entries_.begin() + i); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); + if (entries_[(size_t)i].slot == reservation.slot_) { + erase_inline_entry(i); } } - has_pending_evict_ = false; - pending_protect_ = false; + release_inline_reservation(reservation.id_); + reservation.clear(); +} + +void PrefixCache::confirm_inline_snap( + int slot, int target_cut, + const std::vector & prompt_ids, + bool protect, size_t resident_bytes) { + if (disabled_ || slot < 0 || target_cut <= 0 || + target_cut > (int)prompt_ids.size()) return; + if (active_inline_reservation_ != 0) { + std::fprintf(stderr, + "[pc] direct commit refused while a reservation is active\n"); + return; + } + replace_inline_entry( + slot, target_cut, prompt_ids, protect, resident_bytes); } -void PrefixCache::cancel_inline_snap(int slot) { +void PrefixCache::invalidate_inline_snap(int slot) { if (disabled_) return; - if (has_pending_evict_) { - const int idx = find_entry(pending_evict_key_); - if (idx >= 0 && entries_[idx].slot != slot) return; + for (int i = (int)entries_.size() - 1; i >= 0; --i) { + if (entries_[(size_t)i].slot == slot) erase_inline_entry(i); + } +} + +static void update_atomic_max(std::atomic & value, + uint64_t candidate) { + uint64_t current = value.load(std::memory_order_relaxed); + while (current < candidate && + !value.compare_exchange_weak( + current, candidate, std::memory_order_relaxed)) { } - has_pending_evict_ = false; - pending_protect_ = false; +} + +void PrefixCache::record_capture_attempt(uint64_t elapsed_us, bool success) { + capture_attempts_.fetch_add(1, std::memory_order_relaxed); + if (!success) { + capture_failures_.fetch_add(1, std::memory_order_relaxed); + } + capture_stall_us_total_.fetch_add(elapsed_us, std::memory_order_relaxed); + update_atomic_max(capture_stall_us_max_, elapsed_us); +} + +void PrefixCache::record_restore_attempt(uint64_t elapsed_us, bool restored) { + restore_attempts_.fetch_add(1, std::memory_order_relaxed); + if (!restored) { + restore_invalidations_.fetch_add(1, std::memory_order_relaxed); + } + restore_stall_us_total_.fetch_add(elapsed_us, std::memory_order_relaxed); + update_atomic_max(restore_stall_us_max_, elapsed_us); } void PrefixCache::mark_all_cleared() { @@ -459,9 +711,10 @@ void PrefixCache::mark_all_cleared() { int n = (int)entries_.size(); entries_.clear(); entries_size_count_.store(0, std::memory_order_relaxed); + resident_bytes_ = 0; + resident_bytes_count_.store(0, std::memory_order_relaxed); next_slot_ = 0; - has_pending_evict_ = false; - pending_protect_ = false; + active_inline_reservation_ = 0; std::fprintf(stderr, "[pc] all-cleared — dropped %d LRU entries\n", n); } @@ -586,10 +839,30 @@ 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), - lifetime_hits_.load(std::memory_order_relaxed)}; + InlineStats out{}; + if (disabled_) return out; + out.capacity = cap_; + out.in_use = + (int)entries_size_count_.load(std::memory_order_relaxed); + out.lifetime_hits = lifetime_hits_.load(std::memory_order_relaxed); + out.max_resident_bytes = (uint64_t)max_resident_bytes_; + out.resident_bytes = + resident_bytes_count_.load(std::memory_order_relaxed); + out.budget_skips = budget_skips_.load(std::memory_order_relaxed); + out.capture_attempts = capture_attempts_.load(std::memory_order_relaxed); + out.capture_failures = capture_failures_.load(std::memory_order_relaxed); + out.capture_stall_us_total = + capture_stall_us_total_.load(std::memory_order_relaxed); + out.capture_stall_us_max = + capture_stall_us_max_.load(std::memory_order_relaxed); + out.restore_attempts = restore_attempts_.load(std::memory_order_relaxed); + out.restore_invalidations = + restore_invalidations_.load(std::memory_order_relaxed); + out.restore_stall_us_total = + restore_stall_us_total_.load(std::memory_order_relaxed); + out.restore_stall_us_max = + restore_stall_us_max_.load(std::memory_order_relaxed); + return out; } PrefixCache::FullStats PrefixCache::full_stats() const { diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index e92bde87c..f1d64828d 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -94,7 +94,8 @@ class PrefixCache { static constexpr int MAX_CACHE_SLOTS = MAX_SLOTS - 1; // cap = number of prefix-cache slots (0 disables). - PrefixCache(int cap, const Tokenizer & tokenizer); + PrefixCache(int cap, const Tokenizer & tokenizer, + size_t max_resident_bytes = 0); bool disabled() const { return disabled_; } @@ -105,32 +106,82 @@ class PrefixCache { // Look up the longest cached prefix. Returns (slot, prefix_len) or (-1, 0). std::pair lookup(const std::vector & prompt_ids); + // Side-effect-free candidate for engines that must validate payloads. + std::pair lookup_candidate( + const std::vector & prompt_ids, + int max_prefix_tokens); + + // Promote and count only after an engine restored this checkpoint. + void record_inline_hit( + int slot, int prefix_len, size_t prompt_len); + + + class InlineReservation { + public: + InlineReservation() = default; + ~InlineReservation(); + + InlineReservation(const InlineReservation &) = delete; + InlineReservation & operator=(const InlineReservation &) = delete; + InlineReservation(InlineReservation && other) noexcept; + InlineReservation & operator=(InlineReservation && other) noexcept; + + bool active() const; + int slot() const { return slot_; } + int target_cut() const { return target_cut_; } + + // Commit after the engine saved the payload, cancel before the target + // slot was touched, or abort after a failed write invalidated it. + bool commit(const std::vector & prompt_ids, + size_t resident_bytes = 0, bool protect = false); + bool commit_at(const std::vector & prompt_ids, + int committed_cut, size_t resident_bytes = 0, + bool protect = false); + void cancel(); + void abort(); + + private: + friend class PrefixCache; + InlineReservation(PrefixCache * cache, uint64_t id, int slot, + int target_cut, PrefixHash victim, bool has_victim, + bool protect); + void clear(); + void take(InlineReservation && other); + + PrefixCache * cache_ = nullptr; + uint64_t id_ = 0; + int slot_ = -1; + int target_cut_ = 0; + PrefixHash victim_{}; + bool has_victim_ = false; + bool protect_ = false; + }; + + using InlineSnapshotSize = std::function; - // Prepare an inline snapshot. `restored_prefix_len` prevents reserving a - // slot for a boundary already covered by the restored snapshot. - // `prefer_tools_boundary` selects the system/tools head first (see - // select_inline_snapshot_boundary). When `forced_cut` > restored, that - // cut is used instead (PPP pin_end, including mid-message LCP cuts). - // Returns (slot, target_cut) or (-1, 0). - std::pair prepare_inline_snap( + // Select a boundary, destination, and optional budget victim as one owned + // operation. At most one reservation can be live; destroying it cancels + // without changing committed metadata. + InlineReservation reserve_inline_snap( const std::vector & prompt_ids, int restored_prefix_len = 0, bool prefer_tools_boundary = false, - int forced_cut = 0); + int forced_cut = 0, + InlineSnapshotSize estimate_bytes = {}); - // Confirm after daemon successfully saved the snapshot. - // `protect` marks the entry non-evictable by unprotected traffic (tool pin). + // Commit an already-materialized snapshot without a reservation. Used by + // cache import/bootstrap paths and tests. void confirm_inline_snap(int slot, int target_cut, const std::vector & prompt_ids, - bool protect = false); + bool protect = false, + size_t resident_bytes = 0); - // Abort if the snapshot failed. - void abort_inline_snap(int slot); + // Remove committed metadata for an engine-invalidated checkpoint. + void invalidate_inline_snap(int slot); - // Cancel before the backend slot is touched (for example when the selected - // destination is also the snapshot being restored). Unlike abort, this - // preserves the existing entry and only drops the pending reservation. - void cancel_inline_snap(int slot); + // Record synchronous scheduler stalls caused by copied checkpoints. + void record_capture_attempt(uint64_t elapsed_us, bool success); + void record_restore_attempt(uint64_t elapsed_us, bool restored); // Drop all entries (e.g., after OOM recovery). void mark_all_cleared(); @@ -159,6 +210,17 @@ class PrefixCache { int capacity; int in_use; int64_t lifetime_hits; + uint64_t max_resident_bytes; + uint64_t resident_bytes; + uint64_t budget_skips; + uint64_t capture_attempts; + uint64_t capture_failures; + uint64_t capture_stall_us_total; + uint64_t capture_stall_us_max; + uint64_t restore_attempts; + uint64_t restore_invalidations; + uint64_t restore_stall_us_total; + uint64_t restore_stall_us_max; }; struct FullStats { bool enabled; @@ -190,13 +252,14 @@ class PrefixCache { int slot; std::vector ids; // prefix tokens [0, target_cut) for prefix-aware eviction bool protect = false; // sticky tools-boundary pin + size_t resident_bytes = 0; }; - // Pending protect flag for the in-flight reservation (applied on confirm). - bool pending_protect_ = false; std::vector entries_; int next_slot_ = 0; - PrefixHash pending_evict_key_{}; - bool has_pending_evict_ = false; + uint64_t active_inline_reservation_ = 0; + uint64_t next_inline_reservation_ = 1; + size_t max_resident_bytes_ = 0; + size_t resident_bytes_ = 0; // Full-cache state bool full_disabled_ = true; @@ -215,6 +278,16 @@ class PrefixCache { // tearing across the daemon thread's increments. Relaxed ordering // is sufficient — no synchronization with other state required. std::atomic lifetime_hits_{0}; // inline cache hits + std::atomic resident_bytes_count_{0}; + std::atomic budget_skips_{0}; + std::atomic capture_attempts_{0}; + std::atomic capture_failures_{0}; + std::atomic capture_stall_us_total_{0}; + std::atomic capture_stall_us_max_{0}; + std::atomic restore_attempts_{0}; + std::atomic restore_invalidations_{0}; + std::atomic restore_stall_us_total_{0}; + std::atomic restore_stall_us_max_{0}; std::atomic full_lifetime_hits_{0}; // full-compress cache hits std::atomic full_disk_bytes_{0}; // best-effort snapshot of disk usage // Atomic mirrors of `entries_.size()` and `full_entries_.size()`. @@ -229,7 +302,24 @@ class PrefixCache { // Helpers int find_entry(const PrefixHash & h) const; + int find_slot_entry(int slot) const; + void erase_inline_entry(int idx); void move_to_end(int idx); + std::pair lookup_impl( + const std::vector & prompt_ids, + int max_prefix_tokens, + bool record_hit); + bool inline_reservation_active(uint64_t id) const; + void release_inline_reservation(uint64_t id); + bool commit_inline_reservation(InlineReservation & reservation, + const std::vector & prompt_ids, + int committed_cut, size_t resident_bytes, + bool protect); + void abort_inline_reservation(InlineReservation & reservation); + void replace_inline_entry(int slot, int target_cut, + const std::vector & prompt_ids, + bool protect, size_t resident_bytes); + int find_full_entry(const PrefixHash & h) const; void move_full_to_end(int idx); }; diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index e22a4cf3f..5d82080ee 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -10,6 +10,7 @@ #include "http_server.h" #include "common/concurrency/seq_engine.h" +#include "parallel_prefix_txn.h" #include #include @@ -25,6 +26,9 @@ namespace { // engine slot id returned from admit(), so scheduler and engine agree on // which engine-owned state record a request owns. This remains the one // external phase: sockets stay here, prompt/KV/sampler/progress stay in Qwen. +using PrefixCaptureTxn = BasicPrefixCaptureTxn< + PrefixCache::InlineReservation, SeqEngine>; + struct SchedSlot { ServerJob * job = nullptr; SocketHandle fd = kInvalidSocket; @@ -34,6 +38,8 @@ struct SchedSlot { std::chrono::steady_clock::time_point started_at{}; std::chrono::steady_clock::time_point decode_started_at{}; double prefill_s = 0.0; + int cached_prefix_tokens = 0; + PrefixCaptureTxn cache_capture; int n_gen_cap = 0; int completion_tokens = 0; bool client_disconnected = false; @@ -92,6 +98,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // Replaces the O(n_slots) scan that was called 2-3× per iteration. int live_slots = 0; + // Capture tickets are never reused during this scheduler run. + uint64_t next_prefix_capture_id = 1; + int published_live_count = -1; int published_prefill_count = -1; auto publish_live_count = [&]() { @@ -240,6 +249,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { SchedSlot & s = slots[(size_t)idx]; if (!s.job) return; const ParsedRequest & req = s.job->req; + s.cache_capture.cancel(); // Stop monitor-thread heartbeats before queuing terminal frames. stop_job_stream(s.job, &s.send_buffer); const double decode_s = std::chrono::duration( @@ -248,9 +258,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { GenTimings gen_timings{ s.prefill_s, decode_s, - /*cache_hit=*/false, - /*cached_prefix_tokens=*/0, - /*prefilled_tokens=*/prompt_tokens, + /*cache_hit=*/s.cached_prefix_tokens > 0, + /*cached_prefix_tokens=*/s.cached_prefix_tokens, + /*prefilled_tokens=*/prompt_tokens - s.cached_prefix_tokens, /*effective_prompt_tokens=*/prompt_tokens, }; @@ -258,8 +268,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { PerfRecord perf; perf.prompt_tokens = (int)req.prompt_tokens.size(); perf.completion_tokens = s.completion_tokens; + const int computed_prompt_tokens = + prompt_tokens - s.cached_prefix_tokens; perf.prefill_tok_s = s.prefill_s > 0.0 - ? (double)req.prompt_tokens.size() / s.prefill_s : 0.0; + ? (double)computed_prompt_tokens / s.prefill_s : 0.0; perf.decode_tok_s = decode_s > 0.0 ? (double)s.completion_tokens / decode_s : 0.0; status_.record_perf(perf); @@ -446,12 +458,133 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { start_job_stream(job); } - // Admission only claims the slot and queues the prompt. Prefill - // advances one chunk per engine step alongside live decode. - auto ar = engine.admit(next_request_id, req.prompt_tokens, - req.sampler); - if (ar.status == SeqEngine::AdmitResult::Status::busy) + // PrefixCache owns token policy; SeqEngine owns checkpoint payloads. + // Unsupported engines never receive a plan, so cold fallback cannot + // invalidate a valid entry. + PrefixStorePlan prefix_plan; + PrefixCaptureTxn prepared_capture; + PrefixCache::InlineReservation capture_reservation; + int restore_policy_slot = -1; + const bool prefix_supported = + engine.supports_prefix_store() && !prefix_cache_.disabled(); + if (prefix_supported) { + const auto hit = prefix_cache_.lookup_candidate( + req.prompt_tokens, + (int)req.prompt_tokens.size() - 1); + if (hit.first >= 0 && hit.second > 0 && + hit.second < (int)req.prompt_tokens.size()) { + restore_policy_slot = hit.first; + prefix_plan.restore = { + (uint64_t)hit.first + 1, hit.second}; + } + + capture_reservation = prefix_cache_.reserve_inline_snap( + req.prompt_tokens, + prefix_plan.restore.valid() + ? prefix_plan.restore.tokens : 0, + /*prefer_tools_boundary=*/!req.tools.empty(), + req.pin_end_token, + [&engine](int target_cut) { + return engine.estimate_prefix_store_bytes(target_cut); + }); + if (capture_reservation.active()) { + const uint64_t capture_id = next_prefix_capture_id++; + if (next_prefix_capture_id == 0) + next_prefix_capture_id = 1; + prefix_plan.capture.id = capture_id; + prefix_plan.capture.checkpoint = { + (uint64_t)capture_reservation.slot() + 1, + capture_reservation.target_cut()}; + if (prefix_plan.restore.valid() && + prefix_plan.capture.checkpoint == prefix_plan.restore) { + capture_reservation.cancel(); + prefix_plan.capture = {}; + } + } + } + if (prefix_plan.capture.valid()) { + prepared_capture = PrefixCaptureTxn( + std::move(capture_reservation), engine, + prefix_plan.capture); + } + + const PrefixStorePlan requested_plan = prefix_plan; + auto ar = prefix_supported + ? engine.admit_with_prefix( + next_request_id, req.prompt_tokens, req.sampler, + requested_plan) + : engine.admit( + next_request_id, req.prompt_tokens, req.sampler); + + std::string prefix_protocol_error; + const PrefixStoreAdmission & prefix = ar.prefix_store; + if (prefix.restore_attempted) { + prefix_cache_.record_restore_attempt( + prefix.restore_elapsed_us, prefix.restored.valid()); + } + if (prefix.malformed_restore_state()) { + prefix_protocol_error = + "engine returned malformed prefix restore state"; + // A malformed outcome does not tell us whether the engine-owned + // payload is still usable. Drop both sides of the requested + // checkpoint so a later lookup cannot retry stale metadata. + prepared_capture.cancel(); + if (requested_plan.restore.valid() && restore_policy_slot >= 0) { + engine.discard_prefix_store(requested_plan.restore); + prefix_cache_.invalidate_inline_snap(restore_policy_slot); + } + } + if (prefix.invalidated.valid()) { + if (prefix.invalidated != requested_plan.restore || + restore_policy_slot < 0) { + prefix_protocol_error = + "engine invalidated an unrequested prefix checkpoint"; + } else { + // Cancel this admission's reservation before removing the + // stale restore metadata. The invalidation itself must not + // clear a capture reservation owned by another live request. + // Qwen already discarded the engine-owned payload. + prepared_capture.cancel(); + prefix_cache_.invalidate_inline_snap(restore_policy_slot); + } + } + if (prefix.restored.valid() && + prefix.restored != requested_plan.restore) { + engine.discard_prefix_store(prefix.restored); + prefix_protocol_error = + "engine restored an unrequested prefix checkpoint"; + } + if (prefix.restored.valid() && prefix.invalidated.valid()) { + prefix_protocol_error = + "engine both restored and invalidated one checkpoint"; + } + if (prefix.capture.valid() && + prefix.capture != requested_plan.capture) { + prefix_protocol_error = + "engine accepted an unrequested prefix capture"; + } + if (ar.status != SeqEngine::AdmitResult::Status::admitted && + (prefix.restored.valid() || prefix.capture.valid())) { + prefix_protocol_error = + "failed admission returned accepted prefix state"; + } + if (!prefix_protocol_error.empty()) { + prepared_capture.cancel(); + if (ar.status == SeqEngine::AdmitResult::Status::admitted) + engine.retire(ar.slot); + ar.status = SeqEngine::AdmitResult::Status::failed; + ar.error = "prefix admission protocol violation: " + + prefix_protocol_error; + } else if (prefix.capture != requested_plan.capture) { + // Rejected capture: no payload was touched, so preserve any + // incumbent selected by the cache's transactional eviction. + prepared_capture.cancel(); + } + + // Invalid restore cleanup deliberately precedes both branches below. + if (ar.status == SeqEngine::AdmitResult::Status::busy) { return AdmissionDisposition::Deferred; + } if (ar.status != SeqEngine::AdmitResult::Status::admitted) { std::fprintf(stderr, "[server] admit failed: %s\n", ar.error.c_str()); @@ -470,6 +603,11 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { return AdmissionDisposition::Retired; } next_request_id++; + if (prefix.restored.valid() && restore_policy_slot >= 0) { + prefix_cache_.record_inline_hit( + restore_policy_slot, prefix.restored.tokens, + req.prompt_tokens.size()); + } SchedSlot & s = slots[(size_t)ar.slot]; s = SchedSlot{}; @@ -479,6 +617,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.admission_order = next_admission_order++; s.started_at = started_at; s.decode_started_at = started_at; // sane on prefill failure + s.cached_prefix_tokens = prefix.restored.tokens; + s.cache_capture = std::move(prepared_capture); s.n_gen_cap = std::min( n_gen_cap, engine.max_context() - (int)req.prompt_tokens.size() + 1); @@ -677,6 +817,32 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { if (out.slot < 0 || out.slot >= n_slots) continue; SchedSlot & s = slots[(size_t)out.slot]; if (!s.job) continue; + if (out.prefix_store.attempted()) { + prefix_cache_.record_capture_attempt( + out.prefix_store.elapsed_us, + out.prefix_store.status == + PrefixStoreEvent::Status::saved); + using Resolution = PrefixCaptureTxn::Resolution; + const Resolution resolution = s.cache_capture.resolve( + out.prefix_store, s.job->req.prompt_tokens); + if (resolution == Resolution::failed) { + std::fprintf(stderr, + "[parallel-pc] capture failed checkpoint=%llu: %s\n", + (unsigned long long) + out.prefix_store.ticket.checkpoint.id, + out.prefix_store.error.c_str()); + } else if (resolution == Resolution::mismatched || + resolution == Resolution::inactive) { + // The transaction has already aborted only its own + // destination. Never act on an event-supplied checkpoint. + std::fprintf(stderr, + "[parallel-pc] capture ticket mismatch id=%llu " + "checkpoint=%llu\n", + (unsigned long long)out.prefix_store.ticket.id, + (unsigned long long) + out.prefix_store.ticket.checkpoint.id); + } + } if (out.status == PrefillStatus::failed) { s.failed = true; s.error = out.error; diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index b0d3fe9da..b13f00db1 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -122,6 +123,9 @@ static void print_usage(const char * prog) { " (default: sized from available device memory)\n" " --model-name Model name for /v1/models (default: dflash)\n" " --prefix-cache-slots Prefix cache slots (default: 32, 0 disables)\n" + " --concurrent-prefix-cache-max-mib \n" + " Resident RAM limit for copied concurrent paged\n" + " checkpoints (default: 4096; 0 unlimited)\n" " --agent-turn-cache Extend prefix caching through generated tool calls\n" " --prefill-cache-slots Full prompt/prefill cache slots (default: 0)\n" " --fast-rollback Enable speculative fast rollback (default: on)\n" @@ -419,6 +423,30 @@ int main(int argc, char ** argv) { sconfig.model_name = argv[++i]; } else if (std::strcmp(argv[i], "--prefix-cache-slots") == 0 && i + 1 < argc) { sconfig.prefix_cache_cap = std::atoi(argv[++i]); + } else if (std::strcmp( + argv[i], "--concurrent-prefix-cache-max-mib") == 0) { + if (i + 1 >= argc) { + std::fprintf(stderr, + "[server] --concurrent-prefix-cache-max-mib requires " + "a value\n"); + return 2; + } + const char * value = argv[++i]; + const char * end = value + std::strlen(value); + uint64_t mib = 0; + const auto parsed = std::from_chars(value, end, mib); + constexpr uint64_t bytes_per_mib = 1024ull * 1024ull; + if (parsed.ec != std::errc{} || parsed.ptr != end || + mib > (uint64_t)std::numeric_limits::max() / + bytes_per_mib) { + std::fprintf(stderr, + "[server] --concurrent-prefix-cache-max-mib must be a " + "non-negative " + "integer that fits in addressable memory\n"); + return 2; + } + sconfig.concurrent_prefix_cache_max_bytes = + (size_t)(mib * bytes_per_mib); } else if (std::strcmp(argv[i], "--agent-turn-cache") == 0) { sconfig.agent_turn_cache = true; } else if (std::strcmp(argv[i], "--prefill-cache-slots") == 0 && i + 1 < argc) { @@ -798,19 +826,35 @@ int main(int argc, char ** argv) { } } - // Paged decode owns its K/V through a block table that the snapshot format - // cannot describe yet, so the caches it would restore into are turned off. - // This rewrites ServerConfig rather than rejecting the launch, which is why - // it lives here and not in the gate. + // Continuous Qwen serving supports copied in-memory prefix checkpoints: + // restore allocates fresh pages and scatters logical K/V through the new + // sequence's block table. Exact-prefill and disk snapshots still use the + // classic single-sequence format and remain disabled in paged mode. if (bargs.paged_attention) { - std::fprintf(stderr, - "[server] --paged-attention disables prefix/prefill snapshots " - "until their format stores page tables\n"); - sconfig.prefix_cache_cap = 0; + if (bargs.max_concurrency > 1) { + if (sconfig.prefix_cache_cap > 0) { + std::fprintf(stderr, + "[server] concurrent paged serving enables copied in-memory " + "prefix checkpoints; full-prefill and disk caches remain disabled\n"); + } else { + std::fprintf(stderr, + "[server] concurrent paged serving: prefix checkpoints are " + "disabled; full-prefill and disk caches remain disabled\n"); + } + } else { + std::fprintf(stderr, + "[server] single-sequence --paged-attention still disables " + "prefix snapshots\n"); + sconfig.prefix_cache_cap = 0; + } sconfig.prefill_cache_cap = 0; sconfig.disk_cache_dir.clear(); sconfig.disk_cache_policy.mode = DiskPrefixCacheMode::Off; } + sconfig.concurrent_paged_prefix_cache = + bargs.paged_attention && bargs.max_concurrency > 1 && + sconfig.prefix_cache_cap > 0; + if (sconfig.agent_turn_cache && bargs.paged_attention) { std::fprintf(stderr, "[server] --agent-turn-cache is not yet supported with " @@ -1236,6 +1280,17 @@ int main(int argc, char ** argv) { } std::fprintf(stderr, "[server] │ ddtree_budget = %d\n", bargs.ddtree_budget); std::fprintf(stderr, "[server] │ prefix_cache = %d slots\n", sconfig.prefix_cache_cap); + if (sconfig.concurrent_paged_prefix_cache) { + if (sconfig.concurrent_prefix_cache_max_bytes == 0) { + std::fprintf(stderr, + "[server] │ prefix_cache_ram= unlimited\n"); + } else { + std::fprintf(stderr, + "[server] │ prefix_cache_ram= %zu MiB resident limit\n", + sconfig.concurrent_prefix_cache_max_bytes / + (1024 * 1024)); + } + } std::fprintf(stderr, "[server] │ agent_turn_cache= %s\n", sconfig.agent_turn_cache ? "ON" : "off"); std::fprintf(stderr, "[server] │ prefill_cache = %d slots\n", sconfig.prefill_cache_cap); diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 91e96604e..40bd7a7b8 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -2574,7 +2574,8 @@ int main(int argc, char ** argv) { if (thick_slot_local != -1 && (thick_slot_local < 0 || thick_slot_local >= PREFIX_CACHE_SLOTS || prefix_snapshots[thick_slot_local].ctx == nullptr - || prefix_snapshots[thick_slot_local].is_thin)) { + || prefix_snapshots[thick_slot_local].layout == + PrefixSnapshot::Layout::thin)) { std::fprintf(stderr, "[snap] RESTORE_CHAIN bad thick slot=%d\n", thick_slot_local); stream_emit(-1); continue; @@ -2597,7 +2598,8 @@ int main(int argc, char ** argv) { int id = (int)id_l; if (id < 0 || id >= PREFIX_CACHE_SLOTS || prefix_snapshots[id].ctx == nullptr - || !prefix_snapshots[id].is_thin) { + || prefix_snapshots[id].layout != + PrefixSnapshot::Layout::thin) { std::fprintf(stderr, "[snap] RESTORE_CHAIN bad thin slot=%d\n", id); thin_parse_ok = false; break; } diff --git a/server/test/test_parallel_prefix_txn.cpp b/server/test/test_parallel_prefix_txn.cpp new file mode 100644 index 000000000..1b8aec56c --- /dev/null +++ b/server/test/test_parallel_prefix_txn.cpp @@ -0,0 +1,212 @@ +#include "server/parallel_prefix_txn.h" +#include "host_check.h" + +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +namespace { + +struct FakeReservationState { + std::vector cancelled; + std::vector aborted; + std::vector committed; + size_t committed_bytes = 0; +}; + +class FakeReservation { +public: + FakeReservation() = default; + explicit FakeReservation( + FakeReservationState & state, int slot = 6, int cut = 16) + : state_(&state), slot_(slot), cut_(cut) {} + ~FakeReservation() { cancel(); } + + FakeReservation(const FakeReservation &) = delete; + FakeReservation & operator=(const FakeReservation &) = delete; + FakeReservation(FakeReservation && other) noexcept { + take(std::move(other)); + } + FakeReservation & operator=(FakeReservation && other) noexcept { + if (this != &other) { + cancel(); + take(std::move(other)); + } + return *this; + } + + bool active() const { return state_ != nullptr; } + int slot() const { return slot_; } + int target_cut() const { return cut_; } + + bool commit(const std::vector &, size_t bytes, bool = false) { + if (!active()) return false; + state_->committed.push_back(slot_); + state_->committed_bytes = bytes; + clear(); + return true; + } + + void cancel() { + if (!active()) return; + state_->cancelled.push_back(slot_); + clear(); + } + + void abort() { + if (!active()) return; + state_->aborted.push_back(slot_); + clear(); + } + +private: + void clear() { + state_ = nullptr; + slot_ = -1; + cut_ = 0; + } + + void take(FakeReservation && other) { + state_ = other.state_; + slot_ = other.slot_; + cut_ = other.cut_; + other.clear(); + } + + FakeReservationState * state_ = nullptr; + int slot_ = -1; + int cut_ = 0; +}; + +struct FakeEngine { + std::vector discarded; + + void discard_prefix_store(PrefixStoreRef checkpoint) { + discarded.push_back(checkpoint); + } +}; + +using Txn = BasicPrefixCaptureTxn; + +PrefixCaptureTicket ticket(uint64_t id = 11) { + PrefixCaptureTicket value; + value.id = id; + value.checkpoint = {7, 16}; + return value; +} + +PrefixStoreEvent event( + PrefixStoreEvent::Status status, PrefixCaptureTicket value) { + PrefixStoreEvent out; + out.status = status; + out.ticket = value; + if (status == PrefixStoreEvent::Status::saved) { + out.bytes = 4096; + } + if (status == PrefixStoreEvent::Status::failed) { + out.error = "copy failed"; + } + return out; +} + +} // namespace + +int main() { + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(std::is_move_constructible_v); + static_assert(std::is_move_assignable_v); + + { + FakeReservationState state; + FakeEngine engine; + { + Txn first(FakeReservation(state), engine, ticket()); + Txn owner(std::move(first)); + CHECK(!first.active()); + CHECK(owner.active()); + } + CHECK(state.cancelled == std::vector({6})); + CHECK(state.aborted.empty()); + CHECK(engine.discarded.empty()); + } + + { + FakeReservationState state; + FakeEngine engine; + { + Txn destination( + FakeReservation(state), engine, ticket(/*id=*/11)); + Txn source( + FakeReservation(state), engine, ticket(/*id=*/12)); + destination = std::move(source); + CHECK(!source.active()); + CHECK(destination.active()); + CHECK(state.cancelled == std::vector({6})); + CHECK(state.aborted.empty()); + CHECK(engine.discarded.empty()); + } + CHECK(state.cancelled == std::vector({6, 6})); + } + + { + FakeReservationState state; + FakeEngine engine; + Txn txn(FakeReservation(state), engine, ticket()); + CHECK(txn.resolve( + event(PrefixStoreEvent::Status::saved, ticket()), + std::vector(16, 1)) == + Txn::Resolution::saved); + CHECK(!txn.active()); + CHECK(state.committed_bytes == 4096); + CHECK(state.committed == std::vector({6})); + CHECK(state.cancelled.empty()); + CHECK(state.aborted.empty()); + CHECK(engine.discarded.empty()); + } + + { + FakeReservationState state; + FakeEngine engine; + Txn txn(FakeReservation(state), engine, ticket()); + CHECK(txn.resolve( + event(PrefixStoreEvent::Status::failed, ticket()), + std::vector(16, 1)) == + Txn::Resolution::failed); + CHECK(state.cancelled == std::vector({6})); + CHECK(state.aborted.empty()); + CHECK(engine.discarded.empty()); + } + + { + FakeReservationState state; + FakeEngine engine; + Txn txn(FakeReservation(state), engine, ticket()); + PrefixCaptureTicket unrelated = ticket(/*id=*/99); + unrelated.checkpoint = {63, 16}; + CHECK(txn.resolve( + event(PrefixStoreEvent::Status::saved, unrelated), + std::vector(16, 1)) == + Txn::Resolution::mismatched); + CHECK(state.cancelled.empty()); + CHECK(state.aborted == std::vector({6})); + CHECK(engine.discarded == std::vector({{7, 16}})); + } + + { + FakeReservationState state; + FakeEngine engine; + Txn txn(FakeReservation(state), engine, ticket()); + txn.cancel(); + CHECK(state.cancelled == std::vector({6})); + CHECK(state.aborted.empty()); + CHECK(engine.discarded.empty()); + } + + std::printf("OK test_parallel_prefix_txn (%d checks)\n", g_checks); + return 0; +} diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 6fb5de4d6..257b0d443 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -10,8 +10,14 @@ using namespace CppUnitTestFramework; +using dflash::common::PrefixSnapshot; using dflash::common::TargetCache; +using dflash::common::estimate_paged_target_cache_snapshot_bytes; +using dflash::common::free_prefix_snapshot; +using dflash::common::replace_paged_target_cache; +using dflash::common::restore_paged_target_cache; using dflash::common::restore_ssm_state; +using dflash::common::snapshot_paged_target_cache; using dflash::common::snapshot_ssm_state; namespace { @@ -101,3 +107,137 @@ TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { ggml_free(ctx); ggml_backend_free(backend); } + +TEST_CASE(RecurrentSnapshotFixture, copied_paged_prefix_uses_fresh_pages) { + 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 = 12 * 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"); + } + + 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); + 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"); + } + + std::vector keys((size_t)ggml_nelements(key)); + std::vector values((size_t)ggml_nelements(value)); + std::vector states((size_t)ggml_nelements(ssm)); + std::vector convs((size_t)ggml_nelements(conv)); + for (size_t i = 0; i < keys.size(); ++i) { + keys[i] = (float)i + 1.0f; + values[i] = (float)i + 1001.0f; + } + for (size_t i = 0; i < states.size(); ++i) states[i] = (float)i + 2001.0f; + for (size_t i = 0; i < convs.size(); ++i) convs[i] = (float)i + 3001.0f; + set_tensor(key, keys); + set_tensor(value, values); + set_tensor(ssm, states); + set_tensor(conv, convs); + + 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}; + + const size_t estimated_bytes = + estimate_paged_target_cache_snapshot_bytes( + cache, /*token_count=*/20); + CHECK(estimated_bytes > 0); + + PrefixSnapshot snap; + const std::vector source_blocks = {2, 0}; + CHECK(snapshot_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/20, snap)); + CHECK(snap.layout == PrefixSnapshot::Layout::paged && + snap.cur_pos == 20); + CHECK(ggml_backend_buffer_get_size(snap.buf) == estimated_bytes); + // Paged copies use get/set with snapshot tensor data as host staging. + // Keep that storage on a true CPU buffer, including on unified-memory + // compute backends. + CHECK(ggml_backend_buffer_get_type(snap.buf) == + ggml_backend_cpu_buffer_type()); + + // An incomplete per-layer pair is invalid topology and must not replace + // the committed payload. + cache.attn_v[0] = nullptr; + CHECK(!replace_paged_target_cache( + cache, /*seq_slot=*/1, source_blocks, + /*block_size=*/16, /*token_count=*/20, snap)); + cache.attn_v[0] = value; + CHECK(snap.layout == PrefixSnapshot::Layout::paged && + snap.cur_pos == 20); + + // A failed shape-changing replacement must leave the incumbent payload + // intact. Block 4 begins beyond this cache's 64 physical rows, so the + // staged copy fails after allocating a differently-sized candidate. + const std::vector invalid_source_blocks = {4, 0}; + CHECK(!replace_paged_target_cache( + cache, /*seq_slot=*/1, invalid_source_blocks, + /*block_size=*/16, /*token_count=*/17, snap)); + CHECK(snap.layout == PrefixSnapshot::Layout::paged && + snap.cur_pos == 20); + + set_tensor(key, std::vector(keys.size(), 0.0f)); + set_tensor(value, std::vector(values.size(), 0.0f)); + set_tensor(ssm, std::vector(states.size(), 0.0f)); + set_tensor(conv, std::vector(convs.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_keys = get_tensor(key); + const auto restored_values = get_tensor(value); + for (int head = 0; head < 2; ++head) { + for (int logical = 0; logical < 20; ++logical) { + const int source_row = logical < 16 ? 32 + logical : logical - 16; + const int destination_row = logical < 16 ? 16 + logical : 48 + logical - 16; + for (int element = 0; element < 2; ++element) { + const size_t src = (size_t)head * 128 + + (size_t)source_row * 2 + element; + const size_t dst = (size_t)head * 128 + + (size_t)destination_row * 2 + element; + CHECK(restored_keys[dst] == keys[src]); + CHECK(restored_values[dst] == values[src]); + } + } + } + const auto restored_states = get_tensor(ssm); + const auto restored_convs = get_tensor(conv); + CHECK(std::equal(states.begin() + states.size() / 2, states.end(), + restored_states.begin())); + CHECK(std::equal(convs.begin() + convs.size() / 2, convs.end(), + restored_convs.begin())); + + free_prefix_snapshot(snap); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index be3f031b4..a416fb7f3 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -250,6 +250,102 @@ static bool mentions(const std::vector & violations, } int main() { + { + FakeSeqEngine engine(1); + PrefixStorePlan plan; + plan.restore = {3, 8}; + plan.capture.id = 1; + plan.capture.checkpoint = {4, 16}; + const auto result = engine.admit_with_prefix( + 1, std::vector{1, 2, 3}, SamplerCfg{}, plan); + CHECK(result.status == SeqEngine::AdmitResult::Status::admitted); + CHECK(!engine.supports_prefix_store()); + engine.retire(result.slot); + } + + { + PrefixStoreAdmission admission; + CHECK(!admission.malformed_restore_state()); + admission.restored = {3, 0}; + CHECK(admission.malformed_restore_state()); + admission.restored = {3, 8}; + CHECK(admission.malformed_restore_state()); + admission.restore_attempted = true; + CHECK(!admission.malformed_restore_state()); + admission.restored = {}; + CHECK(admission.malformed_restore_state()); + admission.invalidated = {3, -1}; + CHECK(admission.malformed_restore_state()); + } + + { + SeqEngine::StepPlan plan; + plan.prefills.push_back({0, 8}); + SeqEngine::StepResult result; + SeqEngine::PrefillOutput output; + output.slot = 0; + output.prefix_store.status = PrefixStoreEvent::Status::saved; + output.prefix_store.ticket.id = 7; + output.prefix_store.ticket.checkpoint = {3, 8}; + output.prefix_store.bytes = 4096; + result.prefills.push_back(output); + CHECK(validate_step_result(plan, result, 1).empty()); + + result.prefills[0].prefix_store.error = "saved with an error"; + CHECK(validate_step_result(plan, result, 1).find( + "saved prefix capture") != std::string::npos); + + result.prefills[0].prefix_store.error.clear(); + result.prefills[0].prefix_store.bytes = 0; + CHECK(validate_step_result(plan, result, 1).find( + "omits its byte size") != std::string::npos); + + result.prefills[0].prefix_store.status = + PrefixStoreEvent::Status::failed; + result.prefills[0].prefix_store.error = "copy failed"; + result.prefills[0].prefix_store.bytes = 4096; + CHECK(validate_step_result(plan, result, 1).find( + "committed bytes") != std::string::npos); + + result.prefills[0].prefix_store = {}; + result.prefills[0].prefix_store.elapsed_us = 1; + CHECK(validate_step_result(plan, result, 1).find( + "inactive prefix capture") != std::string::npos); + + result.prefills[0].prefix_store = {}; + result.prefills[0].prefix_store.ticket.id = 7; + CHECK(validate_step_result(plan, result, 1).find( + "inactive prefix capture") != std::string::npos); + + result.prefills[0].prefix_store = {}; + result.prefills[0].prefix_store.ticket.checkpoint.id = 3; + CHECK(validate_step_result(plan, result, 1).find( + "inactive prefix capture") != std::string::npos); + + result.prefills[0].prefix_store = {}; + result.prefills[0].prefix_store.ticket.checkpoint.tokens = 8; + CHECK(validate_step_result(plan, result, 1).find( + "inactive prefix capture") != std::string::npos); + + result.prefills[0].prefix_store.status = + static_cast(99); + result.prefills[0].prefix_store.ticket = output.prefix_store.ticket; + CHECK(validate_step_result(plan, result, 1).find( + "unknown status") != std::string::npos); + + result.prefills[0].prefix_store.status = + PrefixStoreEvent::Status::saved; + result.prefills[0].prefix_store.ticket = output.prefix_store.ticket; + result.prefills[0].prefix_store.bytes = output.prefix_store.bytes; + result.prefills[0].prefix_store.error.clear(); + result.prefills[0].status = + SeqEngine::PrefillOutput::Status::failed; + result.prefills[0].error = "prefill failed"; + CHECK(validate_step_result(plan, result, 1).find( + "failed prefill carries a prefix capture") != + std::string::npos); + } + for (const int slots : {2, 4}) { FakeSeqEngine engine(slots); const auto violations = check_seq_engine_contract(engine); diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index a64329e5f..18ade92a9 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -8,6 +8,7 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" +#include #include #include @@ -39,6 +40,43 @@ static bool is_busy(const SeqEngine::AdmitResult & result) { } int main() { + // A copied prefix checkpoint resumes into freshly-owned pages. The + // checkpoint never lends its original physical blocks to the new request: + // the slot consumes its own admission reservation, exposes the resulting + // block-table delta, and starts the remaining prefill at the restored + // logical position. + { + PagedKvPool pool(/*physical_block_count=*/8, /*max_sequences=*/2, + /*block_size=*/16); + Qwen35SlotManager mgr(pool, /*max_ctx=*/64); + const std::vector prompt = prompt_tokens(40); + + auto first = admit(mgr, 1, prompt, greedy_sampler()); + CHECK(is_admitted(first)); + auto original = mgr.append_prefill(first.slot, 24); + CHECK(original.ok && original.rows.front() == 0 && + original.rows.back() == 23); + + auto restored = admit(mgr, 2, prompt, greedy_sampler()); + CHECK(is_admitted(restored)); + auto seeded = mgr.seed_restored_prefix(restored.slot, 24); + CHECK(seeded.ok && seeded.rows.size() == 24); + for (int64_t row : seeded.rows) { + CHECK(std::find(original.rows.begin(), original.rows.end(), row) == + original.rows.end()); + } + CHECK(mgr.slot(restored.slot).cur_pos == 24); + CHECK(mgr.slot(restored.slot).prefilling()); + CHECK(seeded.first_new_block == 0 && seeded.new_blocks.size() == 2); + mgr.retire(first.slot); + + auto suffix = mgr.append_prefill(restored.slot, 16); + CHECK(suffix.ok && suffix.rows.size() == 16); + CHECK(mgr.slot(restored.slot).cur_pos == 40); + mgr.commit_prefill(restored.slot); + CHECK(mgr.slot(restored.slot).decoding()); + } + // 8 blocks x 16 tokens = 128 pool tokens, 2 slots, per-seq max_ctx 64. { PagedKvPool pool(/*physical_block_count=*/8, /*max_sequences=*/2, diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 68aaf9648..b1949f3db 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -20,6 +20,7 @@ #include "server/api_types.h" #include "server/http_server.h" #include "server/chat_template.h" +#include "common/concurrency/seq_engine.h" #include "common/sampler.h" #include "common/backend_precision.h" #include "common/backend_ipc.h" @@ -52,6 +53,7 @@ #include #include #include +#include #include #include #include @@ -78,6 +80,52 @@ std::vector normalize_chat_messages( const json & messages, ApiFormat format, ToolMemory & tool_memory); + +struct SchedulerTestHarness { + static PrefixCache & prefix_cache(HttpServer & server) { + return server.prefix_cache_; + } + + static void enqueue(HttpServer & server, ServerJob * job) { + server.enqueue(job); + } + + static void run(HttpServer & server, SeqEngine & engine) { + server.scheduler_loop(engine); + } + + static void stop(HttpServer & server) { + server.stopping_.store(true, std::memory_order_relaxed); + server.queue_cv_.notify_all(); + } + + static void finalize_inline_snapshot( + HttpServer & server, const std::vector & prompt, + PrefixCache::InlineReservation reservation, + int slot, int requested_cut) { + ParsedRequest req; + req.prompt_tokens = prompt; + HttpServer::PreparedPrompt prepared; + prepared.tokens = prompt; + HttpServer::GenerationCacheState cache; + cache.snap_reservation = std::move(reservation); + cache.snap_slot = slot; + cache.snap_cut = requested_cut; + cache.snap_prepared = true; + GenerateResult result; + result.error.reset(); + server.finalize_generation_cache( + req, prepared, cache, result, + /*completion_tokens=*/1, + /*visible_output_seen=*/true, + /*client_disconnected=*/false); + } + + static const std::vector & slot_tokens( + const HttpServer & server, int slot) { + return server.slot_tokens_.at(slot); + } +}; } namespace { @@ -2089,6 +2137,140 @@ TEST_CASE(ServerUnitFixture, test_prefix_cache_reserves_disk_staging_slot) { unlink(path.c_str()); } +TEST_CASE(ServerUnitFixture, test_prefix_cache_records_only_validated_restore) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + PrefixCache cache(2, tokenizer); + const std::vector prompt = { + 1, 100, 3, 101, 4, 102, 2, 3, 103, 4, + }; + cache.confirm_inline_snap(0, 8, prompt); + cache.confirm_inline_snap(1, 10, prompt); + + const auto candidate = cache.lookup_candidate( + prompt, (int)prompt.size() - 1); + TEST_ASSERT(candidate.first == 0); + TEST_ASSERT(candidate.second == 8); + TEST_ASSERT(cache.stats().lifetime_hits == 0); + + cache.record_inline_hit( + candidate.first, candidate.second, prompt.size()); + TEST_ASSERT(cache.stats().lifetime_hits == 1); + + // The classic path still accepts/counts an exact snapshot. Concurrent + // admission requests a strict prefix because snapshots do not store the + // next-token logits required for an empty suffix. + const auto exact = cache.lookup(prompt); + TEST_ASSERT(exact.first == 1); + TEST_ASSERT(exact.second == 10); + TEST_ASSERT(cache.stats().lifetime_hits == 2); + + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, test_restore_invalidation_preserves_pending_pin) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + PrefixCache cache(2, tokenizer); + const std::vector first = {1, 100}; + const std::vector stale = {1, 200}; + const std::vector pinned = {1, 300}; + const std::vector replacement = {1, 400}; + const std::vector next = {1, 500}; + cache.confirm_inline_snap(0, 2, first); + cache.confirm_inline_snap(1, 2, stale); + + // Reserve the oldest slot for a protected tool-prefix capture, then + // invalidate an unrelated restore while that reservation is in flight. + auto prepared = cache.reserve_inline_snap( + pinned, /*restored_prefix_len=*/0, + /*prefer_tools_boundary=*/true, /*forced_cut=*/2); + TEST_ASSERT(prepared.slot() == 0); + TEST_ASSERT(prepared.target_cut() == 2); + auto blocked = cache.reserve_inline_snap( + next, /*restored_prefix_len=*/0, + /*prefer_tools_boundary=*/false, /*forced_cut=*/2); + TEST_ASSERT(!blocked.active()); + cache.invalidate_inline_snap(/*slot=*/1); + TEST_ASSERT(prepared.commit(pinned)); + + // Refill the unrelated slot. The next eviction must choose this + // unprotected entry, proving invalidation did not clear the pending pin. + cache.confirm_inline_snap(1, 2, replacement); + auto victim = cache.reserve_inline_snap( + next, /*restored_prefix_len=*/0, + /*prefer_tools_boundary=*/false, /*forced_cut=*/2); + TEST_ASSERT(victim.slot() == 1); + TEST_ASSERT(victim.target_cut() == 2); + victim.cancel(); + + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, test_prefix_cache_resident_budget_and_stall_stats) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + PrefixCache cache(3, tokenizer, /*max_resident_bytes=*/300); + const std::vector pinned = {1, 100}; + const std::vector ordinary = {1, 200}; + const std::vector replacement = {1, 300}; + const std::vector oversized = {1, 400}; + + auto prepared = cache.reserve_inline_snap( + pinned, 0, /*prefer_tools_boundary=*/true, /*forced_cut=*/2, + [](int) { return 100; }); + TEST_ASSERT(prepared.slot() == 0); + TEST_ASSERT(prepared.commit(pinned, /*resident_bytes=*/100)); + + prepared = cache.reserve_inline_snap( + ordinary, 0, /*prefer_tools_boundary=*/false, /*forced_cut=*/2, + [](int) { return 100; }); + TEST_ASSERT(prepared.slot() == 1); + TEST_ASSERT(prepared.commit(ordinary, /*resident_bytes=*/100)); + + // A third slot exists, but its 150-byte checkpoint would exceed the + // resident ceiling. Replace the oldest unprotected leaf (slot 1) while + // preserving the protected tools pin in slot 0. + prepared = cache.reserve_inline_snap( + replacement, 0, /*prefer_tools_boundary=*/false, /*forced_cut=*/2, + [](int) { return 150; }); + TEST_ASSERT(prepared.slot() == 1); + TEST_ASSERT(prepared.commit(replacement, /*resident_bytes=*/150)); + + auto stats = cache.stats(); + TEST_ASSERT(stats.in_use == 2); + TEST_ASSERT(stats.max_resident_bytes == 300); + TEST_ASSERT(stats.resident_bytes == 250); + TEST_ASSERT(cache.lookup(pinned).first == 0); + + prepared = cache.reserve_inline_snap( + oversized, 0, /*prefer_tools_boundary=*/false, /*forced_cut=*/2, + [](int) { return 301; }); + TEST_ASSERT(!prepared.active()); + cache.record_capture_attempt(/*elapsed_us=*/1500, /*success=*/false); + cache.record_restore_attempt(/*elapsed_us=*/2500, /*restored=*/false); + stats = cache.stats(); + TEST_ASSERT(stats.in_use == 2); + TEST_ASSERT(stats.resident_bytes == 250); + TEST_ASSERT(stats.budget_skips == 1); + TEST_ASSERT(stats.capture_attempts == 1); + TEST_ASSERT(stats.capture_failures == 1); + TEST_ASSERT(stats.capture_stall_us_total == 1500); + TEST_ASSERT(stats.capture_stall_us_max == 1500); + TEST_ASSERT(stats.restore_attempts == 1); + TEST_ASSERT(stats.restore_invalidations == 1); + TEST_ASSERT(stats.restore_stall_us_total == 2500); + TEST_ASSERT(stats.restore_stall_us_max == 2500); + + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_canonical_turn_matches_replay_checkpoint) { TEST_ASSERT(http_detail::canonical_turn_matches_checkpoint( {1, 2, 3}, {1, 2, 9, 4}, 2)); @@ -3550,6 +3732,485 @@ struct MockBackend : ModelBackend { void shutdown() override {} }; +struct ShortInlineSnapshotBackend : MockBackend { + int saved_slot = -1; + int saved_position = 0; + + bool snapshot_used(int slot) const override { + return slot == saved_slot && saved_position > 0; + } + int snapshot_cur_pos(int slot) const override { + return snapshot_used(slot) ? saved_position : 0; + } +}; + +TEST_CASE(ServerUnitFixture, + test_inline_snapshot_finalization_uses_actual_saved_position) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + ShortInlineSnapshotBackend backend; + ServerConfig config; + config.prefix_cache_cap = 2; + HttpServer server(backend, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + const std::vector prompt = {1, 100, 3, 101}; + auto reservation = cache.reserve_inline_snap( + prompt, /*restored_prefix_len=*/0, + /*prefer_tools_boundary=*/false, /*forced_cut=*/4); + TEST_ASSERT(reservation.active()); + TEST_ASSERT(reservation.target_cut() == 4); + + backend.saved_slot = reservation.slot(); + backend.saved_position = 3; + SchedulerTestHarness::finalize_inline_snapshot( + server, prompt, std::move(reservation), backend.saved_slot, + /*requested_cut=*/4); + + const auto hit = cache.lookup(prompt); + TEST_ASSERT(hit.first == backend.saved_slot); + TEST_ASSERT(hit.second == backend.saved_position); + TEST_ASSERT(SchedulerTestHarness::slot_tokens( + server, backend.saved_slot) == + std::vector(prompt.begin(), prompt.begin() + 3)); + + unlink(path.c_str()); +} + +#if !defined(_WIN32) +class SchedulerPrefixEngine final : public SeqEngine { +public: + SchedulerPrefixEngine() : slots_(2) {} + + int slot_count() const override { return (int)slots_.size(); } + int max_context() const override { return 64; } + bool supports_prefix_store() const override { return true; } + size_t estimate_prefix_store_bytes(int) const override { return 256; } + bool token_is_eos(int32_t token) const override { return token == 2; } + StepPlanLimits step_plan_limits(int) const override { + return {/*max_prefill_sequences=*/2, + /*max_prefill_tokens_per_sequence=*/64, + /*max_prefill_tokens_total=*/128, + /*prefill_allocation_quantum=*/64}; + } + + AdmitResult admit( + uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler) override { + return admit_with_prefix( + request_id, prompt, sampler, PrefixStorePlan{}); + } + + AdmitResult admit_with_prefix( + uint64_t, + const std::vector & prompt, + const SamplerCfg &, + const PrefixStorePlan & plan) override { + AdmitResult result; + if (prompt.empty()) { + result.error = "empty prompt"; + return result; + } + if (defer_restore.load(std::memory_order_relaxed) && + plan.restore.valid()) { + returned_busy_before_restore.store(true, std::memory_order_relaxed); + result.status = AdmitResult::Status::busy; + result.error = "restore admission deferred"; + return result; + } + int chosen = -1; + for (int i = 0; i < (int)slots_.size(); ++i) { + if (!slots_[(size_t)i].active) { + chosen = i; + break; + } + } + if (chosen < 0) { + result.status = AdmitResult::Status::busy; + result.error = "all slots live"; + return result; + } + + Slot & slot = slots_[(size_t)chosen]; + slot.active = true; + slot.prefilling = true; + slot.capture = {}; + result.status = AdmitResult::Status::admitted; + result.slot = chosen; + + if (plan.restore.valid()) { + result.prefix_store.restore_attempted = true; + result.prefix_store.restore_elapsed_us = 2500; + if (unrequested_restore.load(std::memory_order_relaxed)) { + result.prefix_store.restored = { + plan.restore.id + 1, plan.restore.tokens}; + } else if (malformed_restore.load(std::memory_order_relaxed)) { + result.prefix_store.restored = {plan.restore.id, 0}; + } else if (plan.restore == PrefixStoreRef{2, 2}) { + saw_stale_restore = true; + discarded.push_back(plan.restore); + result.prefix_store.invalidated = plan.restore; + } else { + result.prefix_store.restored = plan.restore; + } + } + if (!result.prefix_store.invalidated.valid() && + plan.capture.valid()) { + saw_capture = true; + slot.capture = plan.capture; + result.prefix_store.capture = plan.capture; + } + return result; + } + + StepResult step(const StepPlan & plan) override { + StepResult result; + for (const StepInput & input : plan.decode) { + result.decode.push_back({input.slot, 2, false, {}}); + } + for (const PrefillSlice & slice : plan.prefills) { + if (slice.slot < 0 || slice.slot >= (int)slots_.size() || + !slots_[(size_t)slice.slot].active || + !slots_[(size_t)slice.slot].prefilling) { + result.error = "invalid prefill"; + result.prefills.clear(); + result.decode.clear(); + return result; + } + Slot & slot = slots_[(size_t)slice.slot]; + slot.prefilling = false; + PrefillOutput output; + output.slot = slice.slot; + output.status = PrefillOutput::Status::completed; + output.token = 2; + if (slot.capture.valid()) { + output.prefix_store.status = PrefixStoreEvent::Status::saved; + output.prefix_store.ticket = slot.capture; + output.prefix_store.bytes = 256; + output.prefix_store.elapsed_us = 1500; + slot.capture = {}; + } + result.prefills.push_back(std::move(output)); + } + return result; + } + + void retire(int slot) override { + if (slot >= 0 && slot < (int)slots_.size()) { + slots_[(size_t)slot] = Slot{}; + } + } + + void discard_prefix_store(PrefixStoreRef checkpoint) override { + discarded.push_back(checkpoint); + } + + bool saw_capture = false; + bool saw_stale_restore = false; + std::atomic defer_restore{false}; + std::atomic returned_busy_before_restore{false}; + std::atomic unrequested_restore{false}; + std::atomic malformed_restore{false}; + std::vector discarded; + +private: + struct Slot { + bool active = false; + bool prefilling = false; + PrefixCaptureTicket capture; + }; + std::vector slots_; +}; + +struct SchedulerPrefixBackend : MockBackend { + SchedulerPrefixEngine engine; + SeqEngine * seq_engine() override { return &engine; } +}; + +TEST_CASE(ServerUnitFixture, + test_scheduler_counts_restore_only_after_engine_attempts_it) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + SchedulerPrefixBackend backend; + backend.engine.defer_restore.store(true, std::memory_order_relaxed); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; + config.concurrent_prefix_cache_max_bytes = 1024; + config.concurrent_paged_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(backend, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + cache.confirm_inline_snap( + /*slot=*/0, /*target_cut=*/2, {1, 100}, false, 128); + + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + + ServerJob job; + job.fd = sockets[0]; + job.req.format = ApiFormat::OPENAI_CHAT; + job.req.prompt_tokens = {1, 100, 999}; + job.req.max_output = 1; + job.req.stream = false; + job.req.model = "scheduler-test"; + job.req.response_id = "restore-after-busy"; + + SchedulerTestHarness::enqueue(server, &job); + std::thread scheduler([&] { + SchedulerTestHarness::run(server, backend.engine); + }); + + const auto busy_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!backend.engine.returned_busy_before_restore.load( + std::memory_order_relaxed) && + std::chrono::steady_clock::now() < busy_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + const bool returned_busy = + backend.engine.returned_busy_before_restore.load( + std::memory_order_relaxed); + const auto pre_attempt_stats = cache.stats(); + backend.engine.defer_restore.store(false, std::memory_order_relaxed); + + std::unique_lock lock(job.mu); + const bool done = job.cv.wait_for( + lock, std::chrono::seconds(5), [&] { return job.done; }); + lock.unlock(); + SchedulerTestHarness::stop(server); + scheduler.join(); + + close(sockets[0]); + close(sockets[1]); + unlink(path.c_str()); + + TEST_ASSERT(done); + TEST_ASSERT(returned_busy); + TEST_ASSERT(pre_attempt_stats.restore_attempts == 0); + TEST_ASSERT(pre_attempt_stats.restore_invalidations == 0); + TEST_ASSERT(pre_attempt_stats.restore_stall_us_total == 0); + TEST_ASSERT(pre_attempt_stats.restore_stall_us_max == 0); + + const auto stats = cache.stats(); + TEST_ASSERT(stats.restore_attempts == 1); + TEST_ASSERT(stats.restore_invalidations == 0); + TEST_ASSERT(stats.restore_stall_us_total == 2500); +} + +TEST_CASE(ServerUnitFixture, + test_scheduler_rejects_malformed_restore_and_drops_stale_entry) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + SchedulerPrefixBackend backend; + backend.engine.malformed_restore.store(true, std::memory_order_relaxed); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; + config.concurrent_prefix_cache_max_bytes = 1024; + config.concurrent_paged_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(backend, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + cache.confirm_inline_snap( + /*slot=*/0, /*target_cut=*/2, {1, 100}, false, 128); + + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + + ServerJob job; + job.fd = sockets[0]; + job.req.format = ApiFormat::OPENAI_CHAT; + job.req.prompt_tokens = {1, 100, 999}; + job.req.max_output = 1; + job.req.stream = false; + job.req.model = "scheduler-test"; + job.req.response_id = "malformed-restore"; + + SchedulerTestHarness::enqueue(server, &job); + std::thread scheduler([&] { + SchedulerTestHarness::run(server, backend.engine); + }); + + std::unique_lock lock(job.mu); + const bool done = job.cv.wait_for( + lock, std::chrono::seconds(5), [&] { return job.done; }); + lock.unlock(); + SchedulerTestHarness::stop(server); + scheduler.join(); + + close(sockets[0]); + close(sockets[1]); + unlink(path.c_str()); + + TEST_ASSERT(done); + TEST_ASSERT(cache.lookup_candidate({1, 100, 999}, 2).first == -1); + TEST_ASSERT(backend.engine.discarded == + std::vector({{1, 2}})); + const auto stats = cache.stats(); + TEST_ASSERT(stats.restore_attempts == 1); + TEST_ASSERT(stats.restore_invalidations == 1); + TEST_ASSERT(stats.restore_stall_us_total == 2500); +} + +TEST_CASE(ServerUnitFixture, + test_scheduler_discards_unrequested_restored_checkpoint) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + SchedulerPrefixBackend backend; + backend.engine.unrequested_restore.store(true, std::memory_order_relaxed); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; + config.concurrent_prefix_cache_max_bytes = 1024; + config.concurrent_paged_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(backend, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + cache.confirm_inline_snap( + /*slot=*/0, /*target_cut=*/2, {1, 100}, false, 128); + + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + + ServerJob job; + job.fd = sockets[0]; + job.req.format = ApiFormat::OPENAI_CHAT; + job.req.prompt_tokens = {1, 100, 999}; + job.req.max_output = 1; + job.req.stream = false; + job.req.model = "scheduler-test"; + job.req.response_id = "unrequested-restore"; + + SchedulerTestHarness::enqueue(server, &job); + std::thread scheduler([&] { + SchedulerTestHarness::run(server, backend.engine); + }); + + std::unique_lock lock(job.mu); + const bool done = job.cv.wait_for( + lock, std::chrono::seconds(5), [&] { return job.done; }); + lock.unlock(); + SchedulerTestHarness::stop(server); + scheduler.join(); + + close(sockets[0]); + close(sockets[1]); + unlink(path.c_str()); + + TEST_ASSERT(done); + TEST_ASSERT(backend.engine.discarded == + std::vector({{2, 2}})); + TEST_ASSERT(cache.lookup_candidate({1, 100, 999}, 2).first == 0); +} + +TEST_CASE(ServerUnitFixture, + test_scheduler_stale_restore_preserves_other_protected_capture) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + SchedulerPrefixBackend backend; + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; + config.concurrent_prefix_cache_max_bytes = 1024; + config.concurrent_paged_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(backend, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + cache.confirm_inline_snap( + /*slot=*/0, /*target_cut=*/2, {1, 100}, false, 128); + cache.confirm_inline_snap( + /*slot=*/1, /*target_cut=*/2, {1, 200}, false, 128); + + int capture_sockets[2] = {-1, -1}; + int restore_sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, capture_sockets) == 0); + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, restore_sockets) == 0); + + ServerJob capture_job; + capture_job.fd = capture_sockets[0]; + capture_job.req.format = ApiFormat::OPENAI_CHAT; + capture_job.req.prompt_tokens = {1, 300, 999}; + capture_job.req.max_output = 1; + capture_job.req.stream = false; + capture_job.req.model = "scheduler-test"; + capture_job.req.response_id = "capture"; + capture_job.req.pin_end_token = 2; + capture_job.req.tools = json::array( + {{{"type", "function"}, + {"function", {{"name", "probe"}, {"parameters", json::object()}}}}}); + + ServerJob restore_job; + restore_job.fd = restore_sockets[0]; + restore_job.req.format = ApiFormat::OPENAI_CHAT; + restore_job.req.prompt_tokens = {1, 200, 999}; + restore_job.req.max_output = 1; + restore_job.req.stream = false; + restore_job.req.model = "scheduler-test"; + restore_job.req.response_id = "restore"; + + SchedulerTestHarness::enqueue(server, &capture_job); + SchedulerTestHarness::enqueue(server, &restore_job); + std::thread scheduler([&] { + SchedulerTestHarness::run(server, backend.engine); + }); + + const auto wait_done = [](ServerJob & job) { + std::unique_lock lock(job.mu); + return job.cv.wait_for(lock, std::chrono::seconds(5), + [&] { return job.done; }); + }; + const bool capture_done = wait_done(capture_job); + const bool restore_done = wait_done(restore_job); + SchedulerTestHarness::stop(server); + scheduler.join(); + + close(capture_sockets[0]); + close(capture_sockets[1]); + close(restore_sockets[0]); + close(restore_sockets[1]); + unlink(path.c_str()); + + TEST_ASSERT(capture_done); + TEST_ASSERT(restore_done); + TEST_ASSERT(backend.engine.saw_capture); + TEST_ASSERT(backend.engine.saw_stale_restore); + auto stats = cache.stats(); + TEST_ASSERT(stats.in_use == 1); + TEST_ASSERT(stats.resident_bytes == 256); + TEST_ASSERT(stats.capture_attempts == 1); + TEST_ASSERT(stats.capture_failures == 0); + TEST_ASSERT(stats.capture_stall_us_total == 1500); + TEST_ASSERT(stats.restore_attempts == 1); + TEST_ASSERT(stats.restore_invalidations == 1); + TEST_ASSERT(stats.restore_stall_us_total == 2500); + + // Refill the stale slot. The next capture must evict this unprotected + // entry, proving the real scheduler preserved the other request's pin. + cache.confirm_inline_snap( + /*slot=*/1, /*target_cut=*/2, {1, 400}, false, 128); + auto victim = cache.reserve_inline_snap( + {1, 500}, 0, /*prefer_tools_boundary=*/false, /*forced_cut=*/2, + [](int) { return 128; }); + TEST_ASSERT(victim.slot() == 1); + victim.cancel(); +} +#endif + struct MockBatchCompressBackend : MockBackend { int compress_calls = 0; @@ -4786,6 +5447,8 @@ TEST_CASE(ServerUnitFixture, test_sampler_needs_logit_processing) { TEST_CASE(ServerUnitFixture, test_server_config_cache_defaults) { ServerConfig cfg; TEST_ASSERT(cfg.prefix_cache_cap == 32); + TEST_ASSERT(cfg.concurrent_prefix_cache_max_bytes == (size_t)4 * 1024 * 1024 * 1024); + TEST_ASSERT(!cfg.concurrent_paged_prefix_cache); TEST_ASSERT(cfg.prefill_cache_cap == 0); } @@ -5039,6 +5702,18 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { TEST_ASSERT(rt["chunk"].get() == 512); TEST_ASSERT(rt["target_device"].get() == "auto:0"); TEST_ASSERT(rt["draft_device"].get() == "auto:0"); + const json & pc_props = body["prefix_cache"]; + TEST_ASSERT(pc_props.contains("max_resident_bytes")); + TEST_ASSERT(pc_props.contains("resident_bytes")); + TEST_ASSERT(pc_props.contains("budget_skips")); + TEST_ASSERT(pc_props.contains("capture_attempts")); + TEST_ASSERT(pc_props.contains("capture_failures")); + TEST_ASSERT(pc_props.contains("capture_stall_ms_total")); + TEST_ASSERT(pc_props.contains("capture_stall_ms_max")); + TEST_ASSERT(pc_props.contains("restore_attempts")); + TEST_ASSERT(pc_props.contains("restore_invalidations")); + TEST_ASSERT(pc_props.contains("restore_stall_ms_total")); + TEST_ASSERT(pc_props.contains("restore_stall_ms_max")); TEST_ASSERT(rt["continuous_batching"]["admission_coalesce_ms"] .get() == 20); TEST_ASSERT(body["pflash"]["draft_residency"].get() == "persistent");