feat(server): add prefix cache to paged concurrency (qwen AR)

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 <path>` | — | Persist prefix cache to disk |

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,

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)

server/docs/PREFIX_CACHE.md

@@ -195,15 +195,24 @@ 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 |
 
 ### Choosing `--prefix-cache-slots`
 
 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 |
 |----------|----------------------|-----------------|

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 <cstddef>
+#include <cstdint>
+#include <string>
+
+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

server/src/common/concurrency/seq_engine.h

@@ -55,6 +55,7 @@
 #include <vector>
 
 #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<int32_t> & 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<int32_t> & 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;
     }
 

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<uint32_t> & 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<uint32_t> & 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<uint32_t> & 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).