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

[Graffioh]

Summary

• Enable copied in-memory prefix checkpoints for Qwen continuous batching with paged attention. Full-prefill and disk snapshots remain disabled in paged mode.
• Keep token matching, boundary selection, pinning, LRU, and hit accounting in PrefixCache while Qwen owns checkpoint payloads, page-table restore, and recurrent state.
• Introduce an opaque checkpoint/ticket protocol and a move-only capture transaction so the scheduler never manipulates backend snapshot storage.
• Gather and scatter logical K/V rows through fresh sequence-owned pages, copy the matching recurrent slab, coalesce consecutive physical blocks, and publish replacements atomically.

Correctness review

The pre-draft review fixed these lifecycle hazards:

• stale restore cleanup now happens before busy or failed admission returns;
• unsupported engines use cold admission without deleting a valid checkpoint;
• malformed or mismatched events can abort only their own capture ticket;
• fatal step failures discard the pending destination instead of preserving potentially mutated state;
• failed replacement leaves the incumbent checkpoint intact;
• cache hits and MRU promotion are recorded only after a validated restore;
• exact-prompt checkpoints cannot mask a shorter restorable prefix when next-token logits are unavailable;
• tensor type, contiguity, extents, vector topology, recurrent slab shape, and per-layer tensor pairs are validated before destination writes.

Verification

• HIP Release build of dflash_server
• test_server_unit: 418 passed
• test_seq_slot_manager: 186 checks
• test_parallel_prefix_txn: 21 checks
• test_seq_engine_contract: 19 checks
• test_recurrent_snapshot: 2 passed
• git diff --check

Draft gates

• Measure capture and restore pause time plus snapshot memory at realistic long-prefix sizes while decode traffic is live.
• Run real Qwen GPU parity for a partial final page and a forced mid-message boundary, restoring into non-identity fresh pages.
• Decide the copied-checkpoint byte budget before broad production enablement.

This PR targets main as the AR ownership/mechanism base. Open speculative concurrency PRs #642 and #654 overlap in SeqEngine, Qwen, scheduler, and target-graph files; they should consume this opaque prefix-store boundary during their rebase rather than duplicate snapshot ownership.

[Graffioh]

@cubic-dev-ai review this PR

[cubic-dev-ai]

@cubic-dev-ai review this PR

@Graffioh I have started the AI code review. It will take a few minutes to complete.

[cubic-dev-ai]

All reported issues were addressed across 25 files

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

5 issues found across 26 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="server/src/server/prefix_cache.cpp">

<violation number="1" location="server/src/server/prefix_cache.cpp:412">
P2: When the backend saves below `target_cut_`, `InlineReservation::commit()` publishes `target_cut_` instead of the actual saved position. The classic path permits `saved_position < cache.snap_cut` and calls this overload, causing later lookups to request a checkpoint position the payload does not contain; commit with `commit_at(..., saved_position, ...)` and store the matching token prefix.</violation>
</file>

<file name="server/test/test_server_unit.cpp">

<violation number="1" location="server/test/test_server_unit.cpp:3888">
P3: The 1-second polling window for `returned_busy_before_restore` (1000 x 1ms) doubles as the pass/fail budget: if the scheduler's admission step takes longer than 1s on a slow or loaded machine, `TEST_ASSERT(returned_busy)` fails spuriously even though the restore-accounting behavior is correct. This makes the test flaky rather than just slow. Extend the wait window and only assert `returned_busy` after a sufficiently generous timeout, or poll until the flag is set with a larger ceiling.</violation>
</file>

<file name="server/src/common/concurrency/prefix_store.h">

<violation number="1" location="server/src/common/concurrency/prefix_store.h:62">
P2: When an engine returns a malformed restore result such as `{id, 0}`, `restore_attempted()` returns false, so the scheduler skips invalidation and protocol checks and leaves the stale `PrefixCache` entry live. Track restore attempts independently from `valid()` and reject malformed `restored`/`invalidated` references.</violation>
</file>

<file name="server/src/server/server_main.cpp">

<violation number="1" location="server/src/server/server_main.cpp:834">
P3: When `--max-concurrency > 1` is combined with `--prefix-cache-slots 0`, this startup branch claims copied checkpoints are enabled even though the cache is disabled. Include the prefix-cache capacity in the message condition or report that the cache is disabled.</violation>
</file>

<file name="server/src/qwen35/qwen35_target_graph.cpp">

<violation number="1" location="server/src/qwen35/qwen35_target_graph.cpp:2676">
P2: On host-accessible unified-memory backends, the snapshot buffer is allocated on `cache.backend`, but these transfers still use the host-oriented get/set async API. That can apply the wrong copy direction or fail when the destination/source is device-backed; use a backend-aware device copy path for same-backend snapshots, or keep paged snapshots on a true host backend.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P2: When the backend saves below target_cut_, InlineReservation::commit() publishes target_cut_ instead of the actual saved position. The classic path permits saved_position < cache.snap_cut and calls this overload, causing later lookups to request a checkpoint position the payload does not contain; commit with commit_at(..., saved_position, ...) and store the matching token prefix.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/server/prefix_cache.cpp, line 412:

<comment>When the backend saves below `target_cut_`, `InlineReservation::commit()` publishes `target_cut_` instead of the actual saved position. The classic path permits `saved_position < cache.snap_cut` and calls this overload, causing later lookups to request a checkpoint position the payload does not contain; commit with `commit_at(..., saved_position, ...)` and store the matching token prefix.</comment>

<file context>
@@ -314,23 +355,125 @@ std::pair<int, int> PrefixCache::lookup(const std::vector<int32_t> & prompt_ids)
+        return false;
+    }
+    return cache_->commit_inline_reservation(
+        *this, prompt_ids, target_cut_, resident_bytes, protect);
+}
+
</file context>

[cubic-dev-ai]

P2: When an engine returns a malformed restore result such as {id, 0}, restore_attempted() returns false, so the scheduler skips invalidation and protocol checks and leaves the stale PrefixCache entry live. Track restore attempts independently from valid() and reject malformed restored/invalidated references.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/common/concurrency/prefix_store.h, line 62:

<comment>When an engine returns a malformed restore result such as `{id, 0}`, `restore_attempted()` returns false, so the scheduler skips invalidation and protocol checks and leaves the stale `PrefixCache` entry live. Track restore attempts independently from `valid()` and reject malformed `restored`/`invalidated` references.</comment>

<file context>
@@ -0,0 +1,84 @@
+    uint64_t restore_elapsed_us = 0;
+
+    bool restore_attempted() const {
+        return restored.valid() || invalidated.valid();
+    }
+};
</file context>

[cubic-dev-ai]

P2: On host-accessible unified-memory backends, the snapshot buffer is allocated on cache.backend, but these transfers still use the host-oriented get/set async API. That can apply the wrong copy direction or fail when the destination/source is device-backed; use a backend-aware device copy path for same-backend snapshots, or keep paged snapshots on a true host backend.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/qwen35/qwen35_target_graph.cpp, line 2676:

<comment>On host-accessible unified-memory backends, the snapshot buffer is allocated on `cache.backend`, but these transfers still use the host-oriented get/set async API. That can apply the wrong copy direction or fail when the destination/source is device-backed; use a backend-aware device copy path for same-backend snapshots, or keep paged snapshots on a true host backend.</comment>

<file context>
@@ -2593,6 +2600,376 @@ bool restore_target_cache(const PrefixSnapshot & snap, TargetCache & cache) {
+                (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);
</file context>

[cubic-dev-ai]

P3: The 1-second polling window for returned_busy_before_restore (1000 x 1ms) doubles as the pass/fail budget: if the scheduler's admission step takes longer than 1s on a slow or loaded machine, TEST_ASSERT(returned_busy) fails spuriously even though the restore-accounting behavior is correct. This makes the test flaky rather than just slow. Extend the wait window and only assert returned_busy after a sufficiently generous timeout, or poll until the flag is set with a larger ceiling.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/test/test_server_unit.cpp, line 3888:

<comment>The 1-second polling window for `returned_busy_before_restore` (1000 x 1ms) doubles as the pass/fail budget: if the scheduler's admission step takes longer than 1s on a slow or loaded machine, `TEST_ASSERT(returned_busy)` fails spuriously even though the restore-accounting behavior is correct. This makes the test flaky rather than just slow. Extend the wait window and only assert `returned_busy` after a sufficiently generous timeout, or poll until the flag is set with a larger ceiling.</comment>

<file context>
@@ -3550,6 +3705,317 @@ struct MockBackend : ModelBackend {
+        SchedulerTestHarness::run(server, backend.engine);
+    });
+
+    for (int i = 0; i < 1000 &&
+         !backend.engine.returned_busy_before_restore.load(
+             std::memory_order_relaxed); ++i) {
</file context>

[cubic-dev-ai]

P3: When --max-concurrency > 1 is combined with --prefix-cache-slots 0, this startup branch claims copied checkpoints are enabled even though the cache is disabled. Include the prefix-cache capacity in the message condition or report that the cache is disabled.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/server/server_main.cpp, line 834:

<comment>When `--max-concurrency > 1` is combined with `--prefix-cache-slots 0`, this startup branch claims copied checkpoints are enabled even though the cache is disabled. Include the prefix-cache capacity in the message condition or report that the cache is disabled.</comment>

<file context>
@@ -798,19 +826,29 @@ int main(int argc, char ** argv) {
-            "[server] --paged-attention disables prefix/prefill snapshots "
-            "until their format stores page tables\n");
-        sconfig.prefix_cache_cap = 0;
+        if (bargs.max_concurrency > 1) {
+            std::fprintf(stderr,
+                "[server] concurrent paged serving enables copied in-memory "
</file context>

Suggested change

	if (bargs.max_concurrency > 1) {
	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] single-sequence --paged-attention still disables "
	"prefix snapshots\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 leaves copied in-memory "
	"prefix checkpoints disabled (--prefix-cache-slots 0)\n");
	}
	} else {