feat(server): split tools from chat so PFlash works on agent turns#492
feat(server): split tools from chat so PFlash works on agent turns#492davidmroth 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.
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.
| if req.tool_choice is not None: | ||
| kwargs["tool_choice"] = req.tool_choice | ||
| prompt = tokenizer.apply_chat_template(msgs, **kwargs) |
| // 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; | ||
| } | ||
| } |
| // 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}
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.
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( |
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: |
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>
| except Exception: | |
| except (Exception, asyncio.CancelledError): |
|
@davidmroth thanks for the huge contribution, we will review it soon. Meanwhile can you fix cubic P1s? |
|
I was able to get the branch built and the mechanical tests pass, but I cannot reproduce the claimed performance gains. What I verified:
I also tried to trace where the benchmark claim comes from. The PR body claims:
The closest in-repo source I found is
That script is not included in this repo, and I could not find it locally. The checked-in test plan validates cache mechanics, but it does not reproduce the advertised benchmark. So my recommendation is: do not merge this yet as a performance PR. Either:
Right now the branch appears to “work” functionally, but the central performance claim is not reproducible from the PR. |
|
Thank you for the feedback! The core idea stands strong: we separate the huge static tool parameters from the chat portion of the prompt — so PFlash now runs on clean chat turns instead of being skipped for tool calls, plus a speedup from agents never recomputing the tool KV cache. This PR effectively allows agent harnesses to benefit from PFlash, which they cannot today. I'll share full benchmarks in a day or so. Test Environment:
|
davidmroth
changed the title
feat(server): Anchored Prefix Caching for multi-turn agent workloads
Why this matters
Agent workloads today cannot benefit from PFlash.
Tool definitions are huge and static. When they sit in the same prompt as the chat portion, PFlash skips compression entirely (to protect the JSON schemas). Every turn also recomputes the large tool KV cache.
This PR fixes both problems at once.
The core idea: we separate the huge static tool parameters from the chat portion of the prompt — so PFlash now runs on clean chat turns instead of being skipped for tool calls, plus a speedup from agents never recomputing the tool KV cache.
How it works
We turn the monolithic prompt into two anchors:
SNAPSHOT_THIN(reusable across sessions that share the same toolset).Then:
RESTORE_CHAINrestores both anchors in one shot.Result: agent turns stop feeling like cold starts.
flowchart TB subgraph Before["❌ Before (Monolithic)"] direction TB B1["Full prompt<br/>Tools + Chat + Query"] B2["PFlash SKIPPED<br/>(tools present)"] B3["Recompute entire KV<br/>every turn"] B1 --> B2 --> B3 end subgraph After["✅ After (Anchored / Tool-Split)"] direction TB A1["Static Tool Anchor<br/>pinned once"] A2["Chat Portion<br/>clean for PFlash"] A3["Delta only<br/>(new message / tool result)"] A1 --> R["RESTORE_CHAIN"] A2 --> R R --> A3 A3 --> P["PFlash runs + fast prefill"] end Before -.->|This PR| AfterMeasured impact
(dual RTX 3090, Qwen3.6-27B, 131K ctx, Hermes-shaped traffic)
What this enables
Summary of changes
SNAPSHOT_THIN) +RESTORE_CHAINserver_tools.py+tool_split/orchestrator (Qwen3 adapter + plugin surface)snap=onRESTORE_CHAIN; CPU snapshot backend for large contextsdraft_accept_pct,avg_commit_per_step,step_ms_*Test plan
python3 server/scripts/test_prefix_cache_slot_depth.pypython3 server/scripts/test_tool_split_adapters.pyai.localbenchmark: agent hot path + 3-turn tool session PASS