feat(server): Anchored Prefix Caching for multi-turn agent workloads#492
Open
davidmroth wants to merge 20 commits into
Open
feat(server): Anchored Prefix Caching for multi-turn agent workloads#492davidmroth wants to merge 20 commits into
davidmroth wants to merge 20 commits into
Conversation
…l-gpu-dflash-safe-p2p-tool-use
RESTORE_CHAIN now honors snap= inline snapshots, prefix-cache entries are only confirmed when the daemon acks, and populated-slot tracking prevents ghost thick restores. Expose prefill_ms/decode metrics in usage.timings for benchmarks and the ai-platform Lucebox adapter.
Clamp conv prefix cache to one in-place slot when tool pins are enabled, always allocate slot 0 first, and force chain-restore snap refresh on the thick slot to avoid PrefixSnapshot OOM (which silently disabled speedups).
…fresh When a thick snapshot slot is refreshed at a deeper boundary, shallow hash→slot entries could restore the wrong cur_pos across conversations. Track committed cut depth per slot, evict stale keys on confirm, and skip lookup hits where key cut != slot depth.
Tool-split Python server (server_tools.py), thin KV snapshots, RESTORE_CHAIN daemon support, prefix-cache slot-depth validation, and benchmark docs/tests. Relocated from dflash/ to server/ to match the repo layout on main.
Thin tool KV snapshots (SNAPSHOT_THIN) + RESTORE_CHAIN for pinned tool definitions, Python server_tools.py orchestration, prefix-cache slot-depth validation (Python + C++), and daemon inline-snap on RESTORE_CHAIN. Validated on dual RTX 3090: ~18x incremental prefill on turn 3+, agent-after-tool ~3.6s wall-clock vs ~5s cold.
Reuse evicted thin-slot IDs within [slot_base, slot_base+pinned_slots), release ToolSlotCache reservations when SNAPSHOT_THIN fails, and fix a malformed apostrophe in the prefix-cache docstring. Add unit coverage.
Resolve add/add conflicts by keeping the Copilot-review fixes (bounded thin-slot LRU, release reservation on SNAPSHOT_THIN failure, docstring).
Explain the innovation in plain language: split tools out of the conversation path so PFlash can speed up chat without tool JSON in the way.
Qwen3.6 chat templates raise "No user query found in messages" when tools= is set without a user turn. Always slice the full rendered prompt at the tool/system boundary so Hermes webchat can pin tool KV and restore conversation cache.
…ackend The interactive/daemon test_dflash path allocated PrefixSnapshot buffers on the CUDA target backend. At large max_ctx (128K) with the model + live KV already resident, every snapshot allocation OOMed (ggml_backend_alloc_ctx_tensors failed for PrefixSnapshot), silently disabling the prefix cache. Route snapshots through create_snapshot_backend() (CPU RAM on discrete GPUs), matching the qwen35_backend, shard-IPC and layer-split daemons.
…cion - server_tools.py speaks the legacy inline daemon protocol (SNAPSHOT_THIN with kv ranges, RESTORE_CHAIN, '[snap] inline slot=' acks). test_dflash in daemon mode was dispatching to run_qwen35_daemon, which implements a different SNAPSHOT_THIN (SSM-only) and no RESTORE_CHAIN, so tool pinning and conversation-cache confirms silently never fired. Add DFLASH_LEGACY_DAEMON=1 to keep the inline loop for daemon mode. - _messages_for_adapter passed tool_call arguments through as JSON strings; Qwen templates iterate arguments with '| items' and raise 'Can only get item pairs from a mapping'. Coerce to dict for both dict-shaped and pydantic-shaped messages (mirrors _tokenize_prompt).
Thick prefix snapshots moved to the CPU snapshot backend (system RAM), so the OOM rationale for clamping prefix_cache_slots to 1 alongside pinned tool slots is gone. Agent stacks interleave multiple prompt families per turn; a single thick slot thrashes on every request and defeats the conversation cache.
Parse the daemon's per-request acceptance summary and per-step timing block into DaemonStdoutBus timings, so /v1/chat/completions responses report draft_accept_pct, avg_commit_per_step, and step_ms_* phase breakdowns alongside prefill/decode timings. This made the 700ms/step cross-GPU draft feature copy visible and is how decode regressions get diagnosed without shell access to the daemon.
Closed
5 tasks
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds “Anchored Prefix Caching” support aimed at multi-turn agent workloads by splitting the prompt into a reusable tool/schema anchor and a variable conversation tail, enabling fast restores (RESTORE_CHAIN) plus thin tool KV snapshots (SNAPSHOT_THIN). It also hardens prefix-cache correctness via slot-depth validation to prevent restoring into mismatched positions, and extends server-side telemetry to better observe prefill/step timing behavior.
Changes:
- Add tool-split orchestration (adapters, registry, daemon bridge) to pin tool KV separately and optionally compress only the conversation suffix.
- Implement prefix-cache slot-depth validation (Python + C++) to prevent stale hash→slot restores after in-place slot refreshes.
- Update the legacy daemon/test harness to support the tool-split protocol (CPU snapshot backend,
RESTORE_CHAINinlinesnap=parsing).
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| server/test/test_dflash.cpp | Adds legacy-daemon toggle, CPU snapshot backend for prefix snapshots, and RESTORE_CHAIN inline snap support. |
| server/src/server/prefix_cache.cpp | Adds slot-depth consistency validation during prefix-cache lookup to evict stale entries. |
| server/scripts/tool_split/base.py | Introduces core dataclasses/interfaces for tool-split adapters and daemon command planning. |
| server/scripts/tool_split/config.py | Adds env/CLI configuration for enabling tool-split, profile selection, and pinned tool slots. |
| server/scripts/tool_split/registry.py | Implements adapter registry plus user plugin discovery/loading. |
| server/scripts/tool_split/orchestrator.py | Builds split plans (restore chain vs restore vs cold) and manages pinned tool-slot LRU. |
| server/scripts/tool_split/daemon_bridge.py | Adds async IPC helpers to commit thin tool KV snapshots to daemon slots. |
| server/scripts/tool_split/qwen3.py | Implements Qwen3.x prompt splitting at the tool/system boundary. |
| server/scripts/tool_split/laguna.py | Adds a Laguna/Poolside splitter stub with marker-boundary fallback. |
| server/scripts/tool_split/example_user_plugin.py | Provides a template for user-defined tool-split adapters via plugin directory. |
| server/scripts/tool_split/init.py | Exposes tool-split public API and auto-registers built-in adapters. |
| server/scripts/tool_split/README.md | Documents how to enable tool-split, profiles, plugins, and daemon flow. |
| server/scripts/test_tool_split_adapters.py | Adds unit tests for registry, adapters, ToolSlotCache, and command formatting. |
| server/scripts/test_prefix_cache_slot_depth.py | Adds unit tests covering stale shallow-hash rejection and key eviction on deep refresh. |
| server/scripts/server_tools.py | Integrates tool-split into the OpenAI-compatible tool-aware server and daemon command composition. |
| server/scripts/prefix_cache.py | Adds slot-depth validation in Python prefix-cache lookup and related correctness plumbing. |
| server/docs/tool-split-goal.md | Adds a project goal doc describing tool-split success criteria and operational guidance. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+770
to
+772
| if req.tool_choice is not None: | ||
| kwargs["tool_choice"] = req.tool_choice | ||
| prompt = tokenizer.apply_chat_template(msgs, **kwargs) |
Comment on lines
+2563
to
+2571
| // Optional inline-snap suffix (same as RESTORE / bare prompt): | ||
| // snap=<pos>:<slot_id> | ||
| if (const char * sp = std::strstr(line.c_str(), "snap=")) { | ||
| if (std::sscanf(sp, "snap=%d:%d", &snap_pos, &snap_slot) != 2 | ||
| || snap_slot < 0 || snap_slot >= PREFIX_CACHE_SLOTS) { | ||
| std::fprintf(stderr, "[snap] bad inline-snap arg\n"); | ||
| snap_pos = -1; snap_slot = -1; | ||
| } | ||
| } |
Comment on lines
+1198
to
+1205
| // Prefix snapshots live in system RAM on discrete GPUs (CPU backend) so | ||
| // they never compete with model weights / KV cache for VRAM. Matches the | ||
| // qwen35_backend / shard-IPC / layer-split daemons. | ||
| ggml_backend_t snap_backend = dflash::common::create_snapshot_backend(target_backend); | ||
| if (!snap_backend) { | ||
| std::fprintf(stderr, "snapshot backend init failed\n"); | ||
| return 1; | ||
| } |
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
15 issues found across 17 files
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
…d hardening Defer tool-slot eviction until confirm(), purge stale mappings on abort, wire abort_full_snap through server_tools, fix qwen3 boundary off-by-one, add Anthropic stop_sequences handling, and harden config/registry/plugin paths.
Contributor
There was a problem hiding this comment.
2 issues found across 12 files (changes from recent commits).
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/scripts/server_tools.py">
<violation number="1" location="server/scripts/server_tools.py:958">
P1: Streaming request cancellations skip the snapshot-abort path because the new failure-handling blocks only catch `Exception`. In Python 3.8+, `asyncio.CancelledError` is a `BaseException`, so a cancelled streaming request (e.g., client disconnect) bypasses the `except Exception:` cleanup, runs only the file-cleanup `finally`, and leaves `full_snap_prep`/`snap_prep` neither confirmed nor aborted. That can leak a reserved prefix-cache slot or leave in-progress snapshot state visible to later requests.
Consider catching `asyncio.CancelledError` alongside `Exception` in any block that performs snapshot cleanup, so cancellation also triggers `_finalize_request_snaps(..., success=False)`.</violation>
<violation number="2" location="server/scripts/server_tools.py:959">
P1: In the streaming generator path, `prompt_bin` is unlinked in the `finally` block before `_finalize_request_snaps` runs on success. Because `PrefixCache.confirm_full_snap` copies the file with `shutil.copy2`, it needs the path to still exist. When a full snap was prepared (`full_snap_prep is not None`) and generation succeeds, the file is already gone by the time the snap is confirmed, so the cache commit silently fails. The non-streaming path already gets this right by finalizing before cleanup. I suggest deferring the success-path unlink until after snap finalization — for example, by only unlinking in `finally` when `not snap_ok` and adding an explicit unlink after the success-path `_finalize_request_snaps`.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| await bus.drain_timings() | ||
| snap_ok = True | ||
| except Exception: | ||
| await _finalize_request_snaps( |
Contributor
There was a problem hiding this comment.
P1: In the streaming generator path, prompt_bin is unlinked in the finally block before _finalize_request_snaps runs on success. Because PrefixCache.confirm_full_snap copies the file with shutil.copy2, it needs the path to still exist. When a full snap was prepared (full_snap_prep is not None) and generation succeeds, the file is already gone by the time the snap is confirmed, so the cache commit silently fails. The non-streaming path already gets this right by finalizing before cleanup. I suggest deferring the success-path unlink until after snap finalization — for example, by only unlinking in finally when not snap_ok and adding an explicit unlink after the success-path _finalize_request_snaps.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At server/scripts/server_tools.py, line 959:
<comment>In the streaming generator path, `prompt_bin` is unlinked in the `finally` block before `_finalize_request_snaps` runs on success. Because `PrefixCache.confirm_full_snap` copies the file with `shutil.copy2`, it needs the path to still exist. When a full snap was prepared (`full_snap_prep is not None`) and generation succeeds, the file is already gone by the time the snap is confirmed, so the cache commit silently fails. The non-streaming path already gets this right by finalizing before cleanup. I suggest deferring the success-path unlink until after snap finalization — for example, by only unlinking in `finally` when `not snap_ok` and adding an explicit unlink after the success-path `_finalize_request_snaps`.</comment>
<file context>
@@ -911,13 +950,29 @@ async def chat_completions(req: ChatRequest):
+ await bus.drain_timings()
+ snap_ok = True
+ except Exception:
+ await _finalize_request_snaps(
+ full_snap_prep=full_snap_prep,
+ snap_prep=snap_prep,
</file context>
| tokens = await asyncio.to_thread(list, _token_stream(r_pipe, gen_len)) | ||
| await bus.drain_timings() | ||
| snap_ok = True | ||
| except Exception: |
Contributor
There was a problem hiding this comment.
P1: Streaming request cancellations skip the snapshot-abort path because the new failure-handling blocks only catch Exception. In Python 3.8+, asyncio.CancelledError is a BaseException, so a cancelled streaming request (e.g., client disconnect) bypasses the except Exception: cleanup, runs only the file-cleanup finally, and leaves full_snap_prep/snap_prep neither confirmed nor aborted. That can leak a reserved prefix-cache slot or leave in-progress snapshot state visible to later requests.
Consider catching asyncio.CancelledError alongside Exception in any block that performs snapshot cleanup, so cancellation also triggers _finalize_request_snaps(..., success=False).
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At server/scripts/server_tools.py, line 958:
<comment>Streaming request cancellations skip the snapshot-abort path because the new failure-handling blocks only catch `Exception`. In Python 3.8+, `asyncio.CancelledError` is a `BaseException`, so a cancelled streaming request (e.g., client disconnect) bypasses the `except Exception:` cleanup, runs only the file-cleanup `finally`, and leaves `full_snap_prep`/`snap_prep` neither confirmed nor aborted. That can leak a reserved prefix-cache slot or leave in-progress snapshot state visible to later requests.
Consider catching `asyncio.CancelledError` alongside `Exception` in any block that performs snapshot cleanup, so cancellation also triggers `_finalize_request_snaps(..., success=False)`.</comment>
<file context>
@@ -911,13 +950,29 @@ async def chat_completions(req: ChatRequest):
+ tokens = await asyncio.to_thread(list, _token_stream(r_pipe, gen_len))
+ await bus.drain_timings()
+ snap_ok = True
+ except Exception:
+ await _finalize_request_snaps(
+ full_snap_prep=full_snap_prep,
</file context>
Suggested change
| except Exception: | |
| except (Exception, asyncio.CancelledError): |
Contributor
|
@davidmroth thanks for the huge contribution, we will review it soon. Meanwhile can you fix cubic P1s? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why this PR
Lucebox already had prefix caching; it didn't have Anchored Prefix Caching — caching the prompt as separate anchors instead of one monolithic prefix.
The problem
A typical agent prompt is ~20K tokens: system message, tool schemas, conversation history, and the latest user query. With monolithic prefix caching, the server treats that entire sequence as a single unit. If anything early in the prompt changes — or a different call type evicts the cache slot — KV for all ~20K tokens is recomputed. VRAM stores every token the same way; the waste is compute: re-running the transformer on thousands of positions whose KV was already valid on the previous turn.
Agent workloads make this worse:
flowchart TB subgraph mono["Monolithic prompt"] direction TB M1["System + tools + history + query"] end mono --> ALL["Recompute KV for entire prompt"]Anchored Prefix Caching
Anchored Prefix Caching separates the prompt into a cached anchor (system, tools, and history — restored from snapshot) and a delta tail (the variable user query or tool result — prefilled fresh each turn).
For agent apps, the anchor is further tiered (implemented today as tool-split):
SNAPSHOT_THIN), reusable across sessions that share a toolset.RESTORE_CHAINreloads tool + conversation anchors in one command, then prefills the delta.flowchart TB subgraph anchor["Cached anchor"] direction TB A1["System + tools + history"] end subgraph delta["Delta prefill"] D1["User query / tool result"] end anchor --> REST["Restore snapshot"] delta --> PREF["Prefill new tokens only"] REST --> OUT["Combined KV → decode"] PREF --> OUTVRAM still sees one linear token stream. The win is bookkeeping: which positions to reuse vs recompute.
Supporting fixes in this branch make Anchored Prefix Caching work at 131K context on discrete GPUs (right-sized CPU-RAM snapshots, legacy daemon protocol for
SNAPSHOT_THIN/RESTORE_CHAIN, tool-argument coercion for Qwen templating, removal of the obsolete single-slot VRAM clamp).Measured impact (dual RTX 3090, Qwen3.6-27B, 131K ctx, Hermes-shaped traffic)
Decode TPS is largely orthogonal (spec-decode / draft placement). This PR targets anchor restore + delta prefill — the path that determines whether agent chat feels instant or like a cold start every message.
Summary
SNAPSHOT_THIN) +RESTORE_CHAINserver_tools.py+tool_split/orchestrator (Qwen3 adapter; plugin surface for other templates)prefix_cache.py+ C++prefix_cache.cpp)snap=onRESTORE_CHAIN; CPU snapshot backend for right-sized slots at 128K+draft_accept_pct,avg_commit_per_step, per-stepstep_ms_*inusage.timingsTest plan
python3 server/scripts/test_prefix_cache_slot_depth.py(Python 3.10+)python3 server/scripts/test_tool_split_adapters.py