feat(pflash): cross-request prefix cache for agent loops composing with pFlash

Merged
DU
6mo ago
+2896-107

TL;DR

40K NIAH @ RTX 3090, pFlash + this PR's full-result cache: cold ~9s (beats published 13.5s @ 64K bench by ~1s when linear-scaled to 40K), warm turns 0.4–0.7s (~50× speedup). Phase A/B prefix cache alone (no pFlash) on a 5-turn ~2K agent loop: cold 5.09s → warm 0.7–1.5s; deeper-hit verified each turn (1206 → 1294 prefix_len).

Summary

Cross-request prefix cache for the dflash daemon, composed cleanly with pFlash speculative-prefill. Today every HTTP request to dflash/scripts/server.py (or server_tools.py) destroys and recreates the target KV cache, paying full prefill cost per turn. For agent loops where each turn appends ~400 tokens to a growing conversation, that's O(total_context) prefill on every tool round-trip. This PR makes it O(new_tokens) for prefixes the daemon has already seen — and skips both compress + prefill entirely when an identical long prompt is repeated and pFlash is on.

Six commits, landed together (current branch tip = e216877):

Phase A — single-snapshot prefix cache (c32d9d7)

  • PrefixSnapshot struct + snapshot_target_cache / restore_target_cache / free_prefix_snapshot primitives in dflash/src/.
  • Daemon protocol: SNAPSHOT <slot> / RESTORE <slot> <prompt> <n_gen> / FREE_SNAPSHOT <slot> / LIST_SLOTS.
  • Python PrefixCache (LRU dict) with Qwen chat-template boundary auto-detect, hash key includes kv_k_type + fa_window, integrated at the 4 _tokenize_prompt callsites.

Phase B — multi-slot LRU + inline-snap (2ebaf57)

  • snapshot_target_cache_thin / restore_target_cache_chain (per-strip H2D+D2H staging via ggml_backend_tensor_get/set because ggml_backend_tensor_copy refuses tensor views with mismatched layouts — verified with spike_thin_copy.cpp).
  • Daemon: SNAPSHOT_THIN <slot> <kv_start> <kv_end> / RESTORE_CHAIN <thick> <thins,csv> <prompt> <n_gen>, plus the inline-snap snap=<pos>:<slot> suffix that fires SNAPSHOT during the actual response prefill (the daemon splits the ubatch at <pos>, snapshots, then continues prefill — no separate n_gen=0 prefill pass).
  • find_all_boundaries returns every Qwen chat-template turn break; lookup returns longest match; prepare_inline_snap/confirm_inline_snap always pick the second-to-last boundary so we cache one turn behind the user's current prompt.
  • cache.last_tok field on TargetCache + bridge in the daemon prefill loop so a RESTORE that covers the entire prompt seeds decode from the snapshot's last token rather than -1.

Phase B fix — defer LRU eviction (7e143f9)

  • LRU strand on aborted requests: prepare_inline_snap reserves the slot id without removing the old key. confirm_inline_snap evicts + inserts atomically. abort_inline_snap for explicit cancellation. Old entry survives if the request fails mid-stream.
  • bench_agent_loop.py — real Claude Code session JSONL replay.

Full-result compression cache + auto-tune (7fd3939)

The Phase A/B prefix cache hashes prompt-id PREFIXES at Qwen chat-template boundaries. With pFlash compression on, the compressed token stream loses chat-template structure, so find_all_boundaries() returns empty, no snapshot is taken, and lookups always miss. Verified failure on a 40K NIAH prompt: 3 turns each pay full ~10s drafter compress + ~10s target prefill, t1=18.9s t2=26.1s t3=29.5s with replies identical (silent degradation from cache disabled to cache no-op).

Adds an alternate path keyed on the raw prompt token ids (markers intact), storing the snapshot slot AND the compressed cur_bin file on disk:

  • New methods on PrefixCache: init_full_cache(full_cap, cache_dir), lookup_full(prompt_ids), prepare_full_snap, confirm_full_snap, abort_full_snap. Separate slot pool [prefix_cap, prefix_cap+full_cap) so the two caches don't fight for the daemon's 8-slot hard cap.
  • On hit: skip compress entirely, RESTORE <slot> <cached_cur_bin> <n_gen>; the daemon's Phase B cache.last_tok bridge seeds decode when the snapshot covers the full cur_bin (no daemon changes needed).
  • New --prefill-cache-slots flag (default 4).

Also auto-sets DFLASH_FP_USE_BSA=1 and DFLASH_FP_ALPHA=0.85 in server.py / server_tools.py when --prefill-compression != off (matches the bench-tuned defaults from pflash/README.md). Without these the drafter falls through to the WMMA fallback at default ALPHA=0.12, roughly tripling cold-start TTFT. setdefault so explicit user overrides still win.

Review fixes (b002090)

Cubic flagged 10 issues on 7fd3939; 8 fixed, 2 justified as non-bugs:

  • P1: prompt_bin tempfile leaked on full-cache hits (9 callsites in server.py + server_tools.py).
  • P1: restore_target_cache now validates snap layer counts + cur_pos range against the cache before iterating (prevents OOB on a stale snapshot from a different daemon run).
  • P2: await_reply removed canceled futures from _waiters in try/finally (was leaking on every timeout).
  • P2: init_full_cache idempotency guard.
  • P2: maybe_snapshot failure path calls abort_inline_snap (was leaving _pending_evict_key set, corrupting later eviction decisions).
  • P2: reset_target_cache now clears last_tok = -1 (was leaking previous request's argmax across resets).
  • P2: snapshot_target_cache_thin rejects kv_end > cache.cur_pos (was capturing uninitialized KV).
  • P2: test_server_prefix_cache.py includes DRAFT in skip preconditions.
  • Justified non-bugs: spike_thin_copy.cpp filename is misleading but the file's purpose IS to verify the host-staging fallback (because ggml_backend_tensor_copy refuses mismatched layouts). test_full_compress_cache.py ungated ok_needle_t1 is intentional — needle retrieval is a function of pFlash compression quality (--prefill-keep-ratio + scoring), not cache correctness; gating on it would couple cache tests to compression quality.

Review round 2 (e216877)

Cubic flagged 1 P1: silent KV-corruption window. prepare_inline_snap's at-cap branch peeked the LRU and stashed old_key in _pending_evict_key without removing it — by design, so a failed request leaves the old entry intact. But every failure path that reaches abort_inline_snap happens AFTER the SNAPSHOT command was already issued, so the daemon may well have committed the snapshot to that slot before the ack timed out. In that case the slot now holds NEW-prompt KV but old_key → slot still lives in entries → next lookup(old_key) returns the slot, daemon RESTOREs it, decode resumes from the wrong KV. Two-part fix: abort_inline_snap now drops _pending_evict_key from entries (conservative — lose one valid cache entry on false abort, never return wrong-KV restore); maybe_snapshot failure path explicitly issues FREE_SNAPSHOT slot (best-effort, 2s ack) before aborting so the daemon slot doesn't dangle.

Numbers

Phase A/B (no pFlash)

test_multi_turn_prefix_cache.py (5 turns, ~2K-token system prompt, deeper-hit assertion):

  • Turn 1 (cold): 5.09s
  • Turns 2-5 warm: 0.70 / 1.13 / 1.32 / 1.54s — 14-30% of cold.
  • Cache walks deeper across turns: 1206 → 1234 → 1262 → 1294 prefix_len.

test_server_prefix_cache.py (3 turns, 2K system prompt):

  • Turn 1 cold: 8.03s. Turns 2-3 warm: 0.39 / 0.37s — 5% of cold.

Full-result cache + pFlash composition

test_full_compress_cache.py (40K NIAH prompt, RTX 3090, --prefill-compression auto --prefill-keep-ratio 0.05 --prefix-cache-slots 4 --prefill-cache-slots 4):

  • Turn 1 cold: 8.9-12.4s (variance band, ~9s typical with auto-tune)
  • Turns 2-3 warm: 0.4-0.7s — ~50× speedup
  • Replies identical across all 3 turns
  • [pc] full-cache hit slot=4 for turns 2 and 3 in server log
  • [snap] restored slot=4 cur_pos=2104 on daemon side (no prefill executed)
  • Beats the published pflash/README.md 64K bench (13.5s @ 64K → ~10s @ 40K linear) by ~1s when scaled.

Drafter stage breakdown with auto-tune on (BSA + α=0.85): forward+score 3.0s (vs ~10s on WMMA fallback), [qwen3-0.6b-fp] FP phase 0.7s (vs 7.5s without BSA).

bench_agent_loop.py (real Claude Code session JSONL)

  • Real-session prompts are too short for the prefix cache alone to dominate; cold and warm turns 2-5 are within noise (~0.8-1s either way), and the only 6.28× signal is turn 1 (OS page-cache + CUDA warmup). Cache wins where prefill is genuinely expensive — long system prompts, long conversation history. Full-result cache (with pFlash on) is where the biggest agent-loop wins live.

Constraints

  • Greedy-only: SSM/conv state captured in snapshots is deterministic on greedy + locked-clock setups. If a sampler is added in the future, the cache must be disabled. Hash key already incorporates kv_k_type/fa_window so a flag change is automatically a cache miss.
  • Daemon slot cap = 8 (hard array limit). Python LRU defaults to --prefix-cache-slots=4 and --prefill-cache-slots=4; combined cap is clamped to 8 (full-cache yields if remaining is too tight).
  • Per-slot memory: prefix cache ~233 MB at --max-ctx=4096 Q8_0 (KV ~36 MB + SSM ~150 MB + conv ~6 MB + target_feat ring ~41 MB). Full cache stores the same per-slot state plus a small cur_bin file on disk (/tmp/dflash-pflash-cache/{hex_digest}.bin, ~8 KB per slot for 2K compressed tokens). Default cap = 4+4 → ~1.9 GB on top of the model.
  • Daemon is single-threaded; concurrent requests are serialised under daemon_lock in the Python server, so snap-bus + LRU don't need additional locking.

Test plan

  • dflash/test/test_prefix_cache.cpp (Phase A correctness — first 32 generated tokens match between cache-hit and cache-miss paths)
  • dflash/scripts/test_server_prefix_cache.py (Phase A python integration: 3 chat completions, turns 2-3 < 10% of turn 1)
  • dflash/scripts/test_multi_turn_prefix_cache.py (Phase B: 5-turn agent loop, deeper-hit assertion + no-regression vs cold gate)
  • dflash/scripts/test_full_compress_cache.py (Full-result cache + pFlash compose: 40K NIAH, turns 2-3 < 30% of turn 1, log-confirmed full-cache hits)
  • dflash/scripts/bench_agent_loop.py (real Claude Code session JSONL replay, slots=0 vs slots=4)
  • cubic review round 1: 10 findings, 8 fixed + 2 justified (b002090)
  • cubic review round 2: 1 P1 finding fixed (e216877)
  • Existing bench_he.py etc. unchanged with --prefix-cache-slots 0 --prefill-compression off (no behavior change vs current main)
  • Build verified clean: cmake -DCMAKE_CUDA_ARCHITECTURES=86 -DDFLASH27B_ENABLE_BSA=ON --build build --target test_dflash (sm_86, BSA enabled)

Phase C (sliding-window KV growth) and Phase E (tool-loop incremental tokenization) remain as follow-ups.