diff --git a/server/docs/tool-split-goal.md b/server/docs/tool-split-goal.md new file mode 100644 index 000000000..bec3f09f3 --- /dev/null +++ b/server/docs/tool-split-goal.md @@ -0,0 +1,96 @@ +# Tool-split — project goal + +## The idea + +Agents that use tools (read a file, run a command, search code) should feel **fast after the first turn** — not like every message is starting from scratch. + +Today, tool definitions get mixed into the same prompt as the conversation. That hurts **PFlash**, the system that compresses and speeds up long chat history. When tools and chat share one blob of text, PFlash has to fight through tool JSON it was never meant to optimize. You pay for the tools again and again, and the conversation speedups never fully kick in. + +**Tool-split** is the fix: **pull tools out of the conversation path.** + +- Tool schemas live in their own pinned memory (thin KV slots). +- The chat history stays clean for PFlash and prefix cache. +- Tools no longer drag down the algorithm that makes multi-turn chat fast. + +Pay the full cost once. After that — especially after a tool result comes back — the agent should feel snappy. + +## What users should feel + +- **Turn 1 (cold)** — full cost once (tools + first message). Expected. +- **Later turns with a little new text** — a few seconds, not another cold start. +- **After a tool result** — the common “continue” path must be fast and reliable. + +If the cache is “working” but people still wait 15 seconds per message, we have not succeeded. + +## How it works (short version) + +1. **Pin tools separately** — tool schemas sit in thin snapshot slots (`SNAPSHOT_THIN`). No re-prefill of hundreds of tool tokens every turn. +2. **Cache conversation alone** — chat history uses `RESTORE_CHAIN` + prefix cache so PFlash can focus on what it does best. +3. **Respect VRAM** — on 2×24GB, `DFLASH_PREFIX_CACHE_SLOTS=1` is a conservative default (one conversation slot, updated in place). Thick prefix snapshots live in CPU RAM, so additional slots are viable when `prefix + tool_pins ≤ DAEMON_MAX_SLOTS`. Extra thick snapshots can still OOM and **silently kill every speedup**. + +Stack: `model-runner-v4` → lucebox (:8080) → ai-platform proxy (:8000) + +## Practical success criteria + +| What users care about | Target | How we measure | +|----------------------|--------|----------------| +| Incremental turn latency | **Wall-clock ≪ turn 1** when only a few new tokens are added | `elapsed_s` on benchmark turn 3 / agent-after-tool | +| Time to start generating | **Prefill ≪ cold** on cached turns | `usage.timings.prefill_ms` | +| Reliability | Speedup works every session, not 1-in-3 after OOM | `inline-snap committed` ≥ 1, no `inline snap failed` in logs | +| Agent hot path | **After tool result**, response in **< 4s** typical | `agent_after_tool` benchmark phase | +| Correctness | Multi-turn tools complete, no `bad thick slot` | 3+ turn session completes | + +### Validated on ai.local (when cache is active) + +| Turn | elapsed | prefill_ms | Notes | +|------|---------|------------|-------| +| 1 cold | ~8s | ~2570 | full prompt + tool pin | +| 2 (bigger delta) | ~6s | ~2240 | still prefills new messages — **not the main win** | +| 3 (tiny delta) | **~3.7s** | **~120** | **21× prefill speedup** — this is the usable win | + +**Key insight:** Speedup tracks **how many new prompt tokens** you add. Small follow-ups feel fast; huge new user messages still cost prefill. That is expected — agents usually add short tool results or short replies between turns. + +## What “done” is not + +- Faster decode tok/s on tool turns (decode is short; **prefill** is what we cache). +- Benchmark-only wins that do not show up as lower `elapsed_s` for incremental turns. +- Production-ready without soak tests, baked image, and CI gate. + +## One-line summary + +**Split tools out so PFlash can speed up the conversation — pay full prefill once; every small follow-up (especially after tool results) should feel snappy.** + +## Infrastructure reference + +| Item | Path / detail | +|------|----------------| +| Server | `user@host`, `/path/to/projects/` | +| Patch scripts | `model-runner-v4/lucebox-patch/dflash/scripts/` | +| Daemon binary | `lucebox-hub/build/test_dflash` | +| Benchmark | `model-runner-v4/scripts/benchmark-tool-split.py` | +| Goal doc | `server/docs/tool-split-goal.md` | + +## Key env + +```bash +DFLASH_TOOL_SPLIT_ENABLED=1 +DFLASH_PREFIX_CACHE_SLOTS=1 # conservative default for 2×24GB; not a hard cap +DFLASH_TOOL_SPLIT_PINNED_SLOTS=2 +DFLASH_LAYER_SPLIT=0 +``` + +## How to verify practical speed + +1. Use **`tools`** in the request (tool-split is off without them). +2. **Restart lucebox** after deploy (clears stale GPU snapshot slots). +3. Run `benchmark-tool-split.py` — check **elapsed_s** and **prefill_ms**, not decode tok/s. +4. Logs must show `thick=0`, `inline-snap committed`, and **no** `inline snap failed`. +5. Compare **turn 3** or **agent_after_tool** to turn 1 — that is the user-visible win. + +## Remaining work for production-grade practical speed + +- [x] Agent-realistic benchmark phase (user → tool_call → tool result → continue) +- [ ] Bake patch + binary into image (no host-mount drift) +- [ ] CI gate: incremental turn `elapsed_s` < 4s, `prefill_ms` < 500ms +- [ ] Soak test: 50 sessions without daemon death or OOM +- [ ] Merge `lucebox-hub` PR and sync `model-runner-v4` defaults diff --git a/server/scripts/prefix_cache.py b/server/scripts/prefix_cache.py new file mode 100644 index 000000000..200258345 --- /dev/null +++ b/server/scripts/prefix_cache.py @@ -0,0 +1,1029 @@ +"""Phase A: single-point prefix cache. + +Auto-detects the system-prompt boundary in token id streams via Qwen chat +template markers, hashes prefixes, and maintains an LRU map of hash → daemon +slot id. Daemon owns slot buffers; Python is the index. + +Usage: + bus = DaemonStdoutBus(daemon_proc.stdout) + bus.start(loop) + + pc = PrefixCache( + daemon_stdin=daemon_proc.stdin, + await_reply=bus.await_reply, + daemon_lock=lock, + tokenizer=tokenizer, + cap=4, + ) + await pc.startup_sync() # free orphaned slots from a previous daemon run + + # Per request (caller holds daemon_lock): + hit = pc.lookup(prompt_ids, kv_k_type, fa_window) # (slot_id, prefix_len) or None + if hit: + slot, prefix_len = hit + # send "RESTORE " instead of bare line + ... + else: + # send bare " " + ... + # after daemon finishes, snapshot for future cache hits: + await pc.maybe_snapshot(prompt_ids, kv_k_type, fa_window) + +Option 3 — full-compress-result cache: + When pFlash compression is enabled, the prefix-cache path above silently + no-ops because compressed tokens lack Qwen chat-template markers. The + full-cache path caches the compressed cur_bin keyed on the ORIGINAL raw + prompt token IDs, so that an identical long prompt sent a second time skips + BOTH the drafter compression dance AND the target prefill. + + full_hit = pc.lookup_full(prompt_ids) + if full_hit: + slot, cached_cur_bin, cur_ids_len = full_hit + cmd_line = f"RESTORE {slot} {cached_cur_bin} {gen_len}\\n" + else: + cur_bin, cur_ids = _maybe_compress(...) + if cur_bin != prompt_bin: # compression actually fired + prep = pc.prepare_full_snap(prompt_ids) + if prep: + slot, _ = prep + cmd_line = f"{cur_bin} {gen_len} snap={len(cur_ids)}:{slot}\\n" + # ...after response completes: + pc.confirm_full_snap(slot, prompt_ids, cur_bin, len(cur_ids)) + # on exception: + pc.abort_full_snap(slot) +""" +import asyncio +import hashlib +import os +import re +import shutil +import struct +from collections import OrderedDict +from pathlib import Path + + +# --------------------------------------------------------------------------- +# DaemonStdoutBus +# --------------------------------------------------------------------------- + +class DaemonStdoutBus: + """Owns the read loop on daemon stdout. + + Lines that start with a registered prefix are routed to the waiting + coroutine; everything else is printed as a log (with noise filtering). + """ + + # Prefixes that are too spammy to print in normal operation. + _SUPPRESS_PREFIXES = ( + "[step ", "[timing]", "[dflash]", "[prompt]", + "[prefill]", "[migrate]", "[dbg ", " ", + ) + + _PREFILL_TIMING_RE = re.compile( + r"\[prefill\](?: layer-seg)? (\d+) tokens in ([\d.]+) s" + ) + _DFLASH_TIMING_RE = re.compile( + r"\[dflash\] generated (\d+) tokens in ([\d.]+) s\s+->\s+([\d.]+) tok/s" + ) + _DFLASH_ACCEPT_RE = re.compile( + r"\[dflash\] (\d+) draft steps, accepted=(\d+)/(\d+) \(([\d.]+)% per step\), " + r"avg commit/step=([\d.]+)" + ) + + _STEP_TIMING_LINE_RE = re.compile(r"^ ([a-z_]+|----- sum)\s+([\d.]+)$") + + def __init__(self, stdout): + self.stdout = stdout + self._waiters: list[tuple[str, asyncio.Future]] = [] + self._task: asyncio.Task | None = None + self._request_inline_slot: int | None = None + self._timings: dict[str, float | int] = {} + self._in_step_timing_block = False + + def begin_request(self) -> None: + """Reset per-request daemon telemetry (inline snap ack, timings, etc.).""" + self._request_inline_slot = None + self._timings = {} + + def inline_snap_slot(self) -> int | None: + """Slot id if daemon emitted ``[snap] inline slot=N`` this request.""" + return self._request_inline_slot + + def request_timings(self) -> dict[str, float | int]: + """Daemon-reported prefill/decode timings for the active request.""" + return dict(self._timings) + + def _parse_timing_line(self, decoded: str) -> None: + # Per-step timing block: "[timing] per-step averages ..." followed by + # indented " name ms" lines. Captured into step_ms_* keys. + if decoded.startswith("[timing] per-step averages"): + self._in_step_timing_block = True + return + if self._in_step_timing_block: + m = self._STEP_TIMING_LINE_RE.match(decoded) + if m: + key = m.group(1).replace("----- ", "") + self._timings[f"step_ms_{key}"] = float(m.group(2)) + return + self._in_step_timing_block = False + m = self._PREFILL_TIMING_RE.search(decoded) + if m: + self._timings["prefill_tokens"] = int(m.group(1)) + self._timings["prefill_ms"] = round(float(m.group(2)) * 1000.0, 2) + return + m = self._DFLASH_TIMING_RE.search(decoded) + if m: + self._timings["completion_tokens"] = int(m.group(1)) + gen_s = float(m.group(2)) + self._timings["decode_ms"] = round(gen_s * 1000.0, 2) + self._timings["decode_tokens_per_sec"] = round(float(m.group(3)), 2) + return + m = self._DFLASH_ACCEPT_RE.search(decoded) + if m: + self._timings["draft_steps"] = int(m.group(1)) + self._timings["draft_accept_pct"] = float(m.group(4)) + self._timings["avg_commit_per_step"] = float(m.group(5)) + + def start(self, loop: asyncio.AbstractEventLoop) -> None: + self._task = loop.create_task(self._run()) + + async def _run(self) -> None: + loop = asyncio.get_running_loop() + while True: + line = await loop.run_in_executor(None, self.stdout.readline) + if not line: + # Daemon exited — wake all waiters with an error. + for _, fut in self._waiters: + if not fut.done(): + fut.set_exception(EOFError("daemon stdout closed")) + self._waiters.clear() + return + decoded = line.decode("utf-8", errors="replace").rstrip() + + self._parse_timing_line(decoded) + + if decoded.startswith("[snap] inline slot="): + try: + slot_s = decoded.split("slot=", 1)[1].split()[0] + self._request_inline_slot = int(slot_s) + except (IndexError, ValueError): + pass + + # Try to satisfy a waiter first. + matched = False + for i, (prefix, fut) in enumerate(self._waiters): + if decoded.startswith(prefix) and not fut.done(): + fut.set_result(decoded) + self._waiters.pop(i) + matched = True + break + + if not matched: + # Log line — suppress very noisy prefixes. + if decoded and not any(decoded.startswith(p) for p in self._SUPPRESS_PREFIXES): + print(f" [daemon] {decoded}", flush=True) + + async def drain_timings(self, timeout: float = 3.0) -> None: + """Wait until daemon emits prefill + decode timing lines (or timeout).""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + has_prefill = self._timings.get("prefill_ms") is not None + has_decode = ( + self._timings.get("decode_ms") is not None + or self._timings.get("decode_tokens_per_sec") is not None + ) + if has_prefill and has_decode: + return + await asyncio.sleep(0.005) + + async def drain_inline_snap(self, timeout: float = 10.0) -> int | None: + """Wait until the read loop consumes ``[snap] inline slot=N`` (or timeout).""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + if self._request_inline_slot is not None: + return self._request_inline_slot + await asyncio.sleep(0.005) + return self._request_inline_slot + + async def await_reply(self, prefix: str, timeout: float = 10.0) -> str: + """Block until daemon emits a line starting with *prefix*.""" + loop = asyncio.get_running_loop() + fut: asyncio.Future[str] = loop.create_future() + entry = (prefix, fut) + self._waiters.append(entry) + try: + return await asyncio.wait_for(fut, timeout=timeout) + finally: + # On timeout / cancellation the matcher loop never popped us; + # remove ourselves so _waiters doesn't grow without bound. + try: self._waiters.remove(entry) + except ValueError: pass + + +# --------------------------------------------------------------------------- +# Qwen chat template helpers +# --------------------------------------------------------------------------- + +def _qwen_marker_ids(tokenizer): + """Resolve <|im_end|>, <|im_start|>, and 'system' token ids.""" + im_end = tokenizer.encode("<|im_end|>", add_special_tokens=False) + im_start = tokenizer.encode("<|im_start|>", add_special_tokens=False) + system_t = tokenizer.encode("system", add_special_tokens=False) + if len(im_end) != 1 or len(im_start) != 1: + raise ValueError( + f"Expected single-token chat markers; got " + f"im_end={im_end} im_start={im_start}" + ) + return im_end[0], im_start[0], system_t[0] if len(system_t) == 1 else None + + +def _resolve_chat_markers(tokenizer): + """Return a marker spec for the prefix-cache boundary detector. + + Supports two chat-template families used by the dflash daemon: + + - Qwen3.x: single-token ``<|im_end|>`` / ``<|im_start|>`` markers, + ``system`` role keyword. Used by Qwen3.5/3.6-27B target. + - Laguna-XS.2 (Poolside): XML-style ```` / ```` / + ```` / ```` / ```` / ```` + markers. Each tokenizes to a 4-6 token sequence under byte-level + BPE. + + Returns a dict with token-sequence patterns: + family : str + sys_role_prefix : tuple[int] pattern that opens the system role + end_msg_seqs : list[tuple] any of these closes a message + next_role_starts : list[tuple] any of these opens the next role + """ + qe = tokenizer.encode("<|im_end|>", add_special_tokens=False) + qs = tokenizer.encode("<|im_start|>", add_special_tokens=False) + if len(qe) == 1 and len(qs) == 1: + sys_t = tokenizer.encode("system", add_special_tokens=False) + sys_seq = tuple(qs) + (tuple(sys_t) if len(sys_t) == 1 else ()) + return { + "family": "qwen", + "sys_role_prefix": sys_seq, + "end_msg_seqs": [tuple(qe)], + "next_role_starts": [tuple(qs)], + } + + lstart_sys = tokenizer.encode("", add_special_tokens=False) + lend_sys = tokenizer.encode("", add_special_tokens=False) + lstart_usr = tokenizer.encode("", add_special_tokens=False) + lend_usr = tokenizer.encode("", add_special_tokens=False) + lstart_ast = tokenizer.encode("", add_special_tokens=False) + lend_ast = tokenizer.encode("", add_special_tokens=False) + if all(x for x in (lstart_sys, lend_sys, lstart_usr, lend_usr, + lstart_ast, lend_ast)): + return { + "family": "laguna", + "sys_role_prefix": tuple(lstart_sys), + "end_msg_seqs": [tuple(lend_sys), tuple(lend_usr), tuple(lend_ast)], + "next_role_starts": [tuple(lstart_usr), tuple(lstart_ast), + tuple(lstart_sys)], + } + + raise ValueError( + f"Could not resolve chat markers for this tokenizer: " + f"qwen im_end={qe} im_start={qs}; laguna seqs missing" + ) + + +def _seq_at(ids, idx, seq): + """Return True iff ids[idx:idx+len(seq)] == seq (and bounds OK).""" + if idx < 0 or idx + len(seq) > len(ids): + return False + for k, t in enumerate(seq): + if ids[idx + k] != t: + return False + return True + + +def _find_first_seq(ids, seq, start=0): + """Index of first occurrence of `seq` in ids[start:], or -1.""" + if not seq: + return -1 + head = seq[0] + n = len(ids); m = len(seq) + i = start + while i + m <= n: + if ids[i] == head and _seq_at(ids, i, seq): + return i + i += 1 + return -1 + + +def _find_first_seq_any(ids, seqs, start=0): + """Position of the earliest match among `seqs` in ids[start:], or (-1, None).""" + best_idx = -1 + best_seq = None + for s in seqs: + idx = _find_first_seq(ids, s, start) + if idx >= 0 and (best_idx < 0 or idx < best_idx): + best_idx = idx + best_seq = s + return best_idx, best_seq + + +def find_prefix_boundary_markers(ids, markers): + """Multi-token-sequence variant of find_prefix_boundary(). + + `markers` is the dict returned by _resolve_chat_markers. The boundary + is the index right after the FIRST next-role start sequence that + follows the system message: i.e., ids[:boundary] = system header. + + Returns -1 if the system role isn't found. + """ + sys_seq = markers["sys_role_prefix"] + end_seqs = markers["end_msg_seqs"] + next_seqs = markers["next_role_starts"] + + sys_idx = _find_first_seq(ids, sys_seq) + if sys_idx < 0: + return -1 + after_sys = sys_idx + len(sys_seq) + + end_idx, end_seq = _find_first_seq_any(ids, end_seqs, after_sys) + if end_idx < 0: + return -1 + after_end = end_idx + len(end_seq) + + # Allow up to 4 separator tokens (whitespace) between message-end and + # the next role-start sequence. + for skip in range(0, 5): + probe = after_end + skip + for s in next_seqs: + if _seq_at(ids, probe, s): + return probe + len(s) + return -1 + + +def find_all_boundaries_markers(ids, markers): + """Multi-token-sequence variant of find_all_boundaries().""" + sys_seq = markers["sys_role_prefix"] + end_seqs = markers["end_msg_seqs"] + next_seqs = markers["next_role_starts"] + + out = [] + sys_idx = _find_first_seq(ids, sys_seq) + if sys_idx < 0: + return out + + cursor = sys_idx + len(sys_seq) + while True: + end_idx, end_seq = _find_first_seq_any(ids, end_seqs, cursor) + if end_idx < 0: + break + after_end = end_idx + len(end_seq) + next_match = -1 + next_len = 0 + for skip in range(0, 5): + probe = after_end + skip + for s in next_seqs: + if _seq_at(ids, probe, s): + next_match = probe + next_len = len(s) + break + if next_match >= 0: + break + if next_match < 0: + break + boundary = next_match + next_len + out.append(boundary) + cursor = boundary + return out + + +def find_prefix_boundary(ids, im_end_id, im_start_id, system_token_id): + """Return the index AFTER the FIRST end-of-system-message marker, or -1. + + Qwen's chat template renders to: + + <|im_start|>system\\nCONTENT<|im_end|>\\n<|im_start|>user\\n... + + so a `\\n` token sits BETWEEN ``<|im_end|>`` and the next ``<|im_start|>``. + We allow up to 2 intervening tokens (covers `\\n` and similar separators). + + The cacheable prefix is the SYSTEM message: from index 0 through and + including the ``<|im_start|>`` that begins the next role. Subsequent turns + sharing this system message hash to the same key. + + Returns the index right after that ``<|im_start|>``, so ``ids[:boundary]`` + is the cached state and ``ids[boundary:]`` is the per-request suffix. + Returns -1 if there is no recognizable system message. + """ + # Find the first <|im_start|>system sequence. + sys_idx = -1 + for i in range(len(ids) - 1): + if ids[i] == im_start_id: + if system_token_id is None or ids[i + 1] == system_token_id: + sys_idx = i + break + if sys_idx < 0: + return -1 + + # Find the FIRST <|im_end|> after sys_idx, then locate the next <|im_start|> + # within a small lookahead (handles a single-token newline separator). + for i in range(sys_idx + 1, len(ids)): + if ids[i] == im_end_id: + for j in range(i + 1, min(i + 3, len(ids))): + if ids[j] == im_start_id: + return j + 1 # boundary is one past <|im_start|> + return -1 # malformed — im_end without subsequent im_start + return -1 + + +def find_all_boundaries(ids, im_end_id, im_start_id, system_token_id): + """Return ascending list of candidate cut points for multi-slot caching. + + Each cut point is the index AFTER an ``<|im_start|>`` that begins a new + role's content. The first cut is the system-prompt boundary (same as + ``find_prefix_boundary``); subsequent cuts are at every following + ``<|im_end|>`` + ``<|im_start|>`` pair. + + Returns an empty list if no recognizable system message is found. + """ + boundaries = [] + + # Locate the opening <|im_start|>system token. + sys_idx = -1 + for i in range(len(ids) - 1): + if ids[i] == im_start_id: + if system_token_id is None or ids[i + 1] == system_token_id: + sys_idx = i + break + if sys_idx < 0: + return boundaries + + # Walk forward from sys_idx: every time we see <|im_end|> followed + # (within 2 tokens) by <|im_start|>, record the position just after + # that <|im_start|> as a cache cut-point. + i = sys_idx + 1 + while i < len(ids): + if ids[i] == im_end_id: + found_start = False + for j in range(i + 1, min(i + 3, len(ids))): + if ids[j] == im_start_id: + boundaries.append(j + 1) + i = j + 1 + found_start = True + break + if not found_start: + break + else: + i += 1 + return boundaries + + +def hash_prefix(prefix_ids, kv_k_type, fa_window): + """Stable SHA-1 (truncated 16 B) of (token ids, kv type, fa window).""" + h = hashlib.sha1() + h.update(struct.pack(" str`` — provided by + ``DaemonStdoutBus.await_reply``. + daemon_lock: + ``asyncio.Lock`` that serialises all stdin writes + stdout reads. + Callers must acquire it before calling ``lookup`` and hold it through + any subsequent ``RESTORE`` / ``SNAPSHOT`` IPC. + tokenizer: + HuggingFace tokenizer (used only to resolve Qwen chat marker ids). + cap: + Maximum number of snapshot slots. 0 disables the cache entirely. + log_prefix: + String prepended to cache-hit/miss log lines. + """ + + # Daemon-side hard cap (PREFIX_CACHE_SLOTS in test_dflash.cpp). Any + # configured cap > this is silently clamped down — exceeding it would + # cause silent SNAPSHOT failures on slots ≥ 8. + DAEMON_MAX_SLOTS = 8 + + def __init__(self, *, daemon_stdin, await_reply, daemon_lock, + tokenizer, kv_k_type: str, fa_window: int, + cap: int = 4, log_prefix: str = "[pc]"): + self.stdin = daemon_stdin + self._await_reply = await_reply + self.lock = daemon_lock + self.log_prefix = log_prefix + # Cache key fields — fixed at daemon spawn (env vars passed through). + # Mismatched values across turns are not possible within one server + # process, but they're still part of the hash so a daemon restart + # with different flags doesn't return stale state. + self.kv_k_type = kv_k_type + self.fa_window = fa_window + + if cap > self.DAEMON_MAX_SLOTS: + print(f"{log_prefix} cap={cap} exceeds daemon limit " + f"({self.DAEMON_MAX_SLOTS}); clamping", flush=True) + cap = self.DAEMON_MAX_SLOTS + self.cap = cap + + if cap <= 0: + self.disabled = True + return + self.disabled = False + + self.entries: OrderedDict[bytes, int] = OrderedDict() # hash → slot_id + self._slot_prefix_len: dict[int, int] = {} # slot → committed cut depth + self.next_slot = 0 + try: + self.markers = _resolve_chat_markers(tokenizer) + except ValueError as e: + print(f"{log_prefix} disabling: {e}", flush=True) + self.disabled = True + self.cap = 0 + return + print(f"{log_prefix} chat markers: family={self.markers['family']} " + f"sys_seq={list(self.markers['sys_role_prefix'])[:6]}… " + f"end_seqs={[list(s)[:4] for s in self.markers['end_msg_seqs']]} " + f"next_seqs={[list(s)[:4] for s in self.markers['next_role_starts']]}", + flush=True) + # Pending eviction: set by prepare_inline_snap when at cap; the old + # entry is NOT removed until confirm_inline_snap succeeds. This ensures + # that if the request aborts before confirm runs, the old entry survives + # and the daemon slot count stays consistent. + self._pending_evict_key: bytes | None = None + # Slots known to hold KV in the daemon (inline snap ack or full snap). + self._populated_slots: set[int] = set() + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def boundary(self, ids: list[int]) -> int: + """Return first boundary (system-prompt end), or -1. Legacy helper.""" + if self.disabled: + return -1 + return find_prefix_boundary_markers(ids, self.markers) + + def _all_boundaries(self, ids: list[int]) -> list[int]: + """Return all candidate cache cut-points in ascending order.""" + return find_all_boundaries_markers(ids, self.markers) + + def lookup(self, prompt_ids: list[int]) -> tuple[int, int] | None: + """Return ``(slot_id, prefix_len)`` for the LONGEST cached prefix, or ``None``. + + Iterates all block-aligned turn boundaries in ``prompt_ids``, checks + each against the LRU index, and returns the deepest (longest) match. + + The caller must already hold ``daemon_lock`` before inspecting the + returned slot, since the slot id may be evicted by a concurrent + request otherwise. + """ + if self.disabled: + return None + candidates = self._all_boundaries(prompt_ids) + best: tuple[int, int] | None = None # (slot_id, prefix_len) + for cut in candidates: + key = hash_prefix(prompt_ids[:cut], self.kv_k_type, self.fa_window) + if key not in self.entries: + continue + slot = self.entries[key] + committed = self._slot_prefix_len.get(slot) + if committed is not None and committed != cut: + # Slot was refreshed in-place at a deeper boundary; an older + # shallow hash→slot mapping would restore the wrong cur_pos. + print(f"{self.log_prefix} lookup stale slot={slot} " + f"key_cut={cut} committed={committed} — evicting", + flush=True) + del self.entries[key] + continue + if best is None or cut > best[1]: + best = (slot, cut) + self.entries.move_to_end(key) # mark fresh + if best is not None and best[0] not in self._populated_slots: + print(f"{self.log_prefix} lookup skip slot={best[0]} " + f"(not populated in daemon)", flush=True) + best = None + if best is not None: + print(f"{self.log_prefix} lookup hit slot={best[0]} prefix_len={best[1]} " + f"(of {len(prompt_ids)} total)", flush=True) + return best + + def slot_populated(self, slot: int) -> bool: + return slot in self._populated_slots + + def mark_slot_populated(self, slot: int) -> None: + self._populated_slots.add(slot) + + def finish_inline_snap( + self, + snap_prep: tuple[int, int] | None, + prompt_ids: list[int], + *, + inline_slot: int | None, + ) -> None: + """Confirm or abort an inline snap reservation based on daemon ack.""" + if not snap_prep: + return + slot, target_cut = snap_prep + if inline_slot == slot: + self.confirm_inline_snap(slot, target_cut, prompt_ids) + else: + self.abort_inline_snap(slot) + + def prepare_inline_snap( + self, + prompt_ids: list[int], + *, + reuse_slot: int | None = None, + ) -> tuple[int, int] | None: + """Pick a target boundary + slot for inline snapshot during the next + request. Returns ``(slot_id, target_cut)`` or ``None`` if no + snapshot is needed (e.g. boundary already cached). + + Caller must: + 1. Append ``snap=:`` to the daemon command + that runs the actual response (bare prompt OR ``RESTORE``). + 2. After the daemon emits ``[snap] inline slot=N cur_pos=M`` + during prefill, call ``confirm_inline_snap(slot_id, target_cut, + prompt_ids)`` to register the entry in the LRU. + + For an agent loop that monotonically grows conversation history, the + most valuable cache point is "end of the most recent completed + assistant message" — i.e., the second-to-last `<|im_start|>` + boundary. The LAST boundary is the current turn's opening, whose + content hasn't been generated yet. + """ + if self.disabled: + return None + candidates = self._all_boundaries(prompt_ids) + if not candidates: + return None + target_cut = candidates[-2] if len(candidates) >= 2 else candidates[-1] + + target_key = hash_prefix(prompt_ids[:target_cut], + self.kv_k_type, self.fa_window) + if target_key in self.entries: + self.entries.move_to_end(target_key) + return None # already cached + + # Pick slot: when at cap, reserve the LRU slot WITHOUT evicting yet. + # The actual eviction is deferred to confirm_inline_snap so that if the + # request aborts before confirm runs, the old entry survives and the + # daemon slot count stays consistent. + if reuse_slot is not None and reuse_slot in self._populated_slots: + slot = reuse_slot + self._pending_evict_key = None + elif not self._populated_slots: + # First successful snap always lands in slot 0 so later turns + # refresh in-place (no second full PrefixSnapshot alloc). + slot = 0 + self._pending_evict_key = None + elif len(self.entries) >= self.cap: + # Peek at LRU without removing. + old_key = next(iter(self.entries)) + slot = self.entries[old_key] + self._pending_evict_key = old_key + else: + slot = self.next_slot + self.next_slot = (self.next_slot + 1) % self.cap + self._pending_evict_key = None + + return (slot, target_cut) + + def confirm_inline_snap(self, slot: int, target_cut: int, + prompt_ids: list[int]) -> None: + """Register an inline snapshot in the LRU after the daemon has + successfully fired ``[snap] inline``. Called from the caller after + the actual response stream completes. + + If prepare_inline_snap reserved a slot by displacing an LRU entry, + the eviction happens HERE (atomically with the insert), so an aborted + request that never reaches confirm leaves the old entry intact. + """ + if self.disabled: + return + # Atomically evict the reserved old entry (if any) and insert the new one. + if self._pending_evict_key is not None: + self.entries.pop(self._pending_evict_key, None) + self._pending_evict_key = None + key = hash_prefix(prompt_ids[:target_cut], + self.kv_k_type, self.fa_window) + # In-place slot refresh leaves stale shallow hash→slot entries; drop them. + stale_keys = [k for k, s in self.entries.items() + if s == slot and k != key] + for k in stale_keys: + del self.entries[k] + self.entries[key] = slot + self._slot_prefix_len[slot] = target_cut + self._populated_slots.add(slot) + print(f"{self.log_prefix} inline-snap committed slot={slot} " + f"prefix_len={target_cut}", flush=True) + + def abort_inline_snap(self, slot: int) -> None: + """Release the reservation made by prepare_inline_snap. + + At-cap case: prepare_inline_snap peeked at the LRU (old_key -> slot) + and stashed old_key in _pending_evict_key WITHOUT removing it. We + cannot tell from here whether the daemon already committed the + snapshot to ``slot`` before the failure was observed: + - If it didn't: old_key -> slot is still semantically valid and + we should keep it. + - If it did: slot now holds the NEW prompt's KV, so old_key + -> slot is stale and a future lookup would return data that + doesn't match the key. + Without daemon-side query we conservatively assume the worst and + drop old_key from the LRU. We accept losing one valid cache entry + in exchange for never returning a wrong-KV restore. Callers that + know the daemon did NOT process the snap (e.g. early validation + failure before any send) should evict only the pending key — but + in practice, every failure path that calls this happens AFTER the + daemon command was issued, so the conservative drop is correct. + """ + if self.disabled: + return + if self._pending_evict_key is not None: + self.entries.pop(self._pending_evict_key, None) + self._pending_evict_key = None + # reuse_slot path: no pending evict, but daemon may have committed new KV. + stale_keys = [k for k, s in self.entries.items() if s == slot] + for k in stale_keys: + del self.entries[k] + self._slot_prefix_len.pop(slot, None) + + # ------------------------------------------------------------------ + # Option 3: full-compress-result cache + # ------------------------------------------------------------------ + # When pFlash compression is enabled the existing prefix-cache path above + # silently no-ops (compressed tokens lack Qwen chat-template markers so + # find_all_boundaries returns []). The full-cache path solves this by + # caching the compressed cur_bin keyed on the ORIGINAL raw prompt_ids so + # that an identical long prompt sent a second time skips BOTH the drafter + # dance AND the target prefill. + # + # Slot allocation: prefix-cache uses slots [0, cap); full-cache uses slots + # [cap, cap + full_cap). Both are initialised at PrefixCache construction + # time; the daemon cap (8) is shared, so prefix_cap + full_cap <= 8. + # ------------------------------------------------------------------ + + def init_full_cache(self, full_cap: int, + cache_dir: str | None = None) -> None: + """Initialise the full-cache pool. Must be called once after __init__ + if you want Option 3 to be active. Idempotent if called again with + the same parameters. + + Parameters + ---------- + full_cap: + Number of daemon slots reserved for full-cache entries. + prefix_cap (self.cap) + full_cap must not exceed DAEMON_MAX_SLOTS. + cache_dir: + Directory to persist cur_bin files across requests. + Defaults to /tmp/dflash-pflash-cache/. + """ + if self.disabled or full_cap <= 0: + self._full_cap = 0 + self._full_disabled = True + return + + # Idempotency guard: a second call would otherwise reset full_entries + + # slot allocator and orphan any cur_bin files already on disk. + if not getattr(self, "_full_disabled", True): + return + + remaining = self.DAEMON_MAX_SLOTS - self.cap + if full_cap > remaining: + print(f"{self.log_prefix} full-cache cap={full_cap} would exceed " + f"daemon limit (prefix uses {self.cap}); clamping to {remaining}", + flush=True) + full_cap = remaining + if full_cap <= 0: + self._full_cap = 0 + self._full_disabled = True + return + + self._full_cap = full_cap + self._full_disabled = False + # Slots used by the full-cache start AFTER the prefix-cache slots. + self._full_slot_base = self.cap + self._full_next_slot = 0 # relative; absolute = _full_slot_base + _full_next_slot + # LRU map: prompt_ids_hash -> (absolute_slot, cached_cur_bin_path, cur_ids_len) + self.full_entries: OrderedDict[bytes, tuple[int, str, int]] = OrderedDict() + # Pending eviction: the LRU entry reserved for the next confirm. + self._full_pending_evict_key: bytes | None = None + self._full_pending_evict_path: str | None = None + + cache_dir_path = Path(cache_dir) if cache_dir else Path("/tmp/dflash-pflash-cache") + cache_dir_path.mkdir(parents=True, exist_ok=True) + self._full_cache_dir = cache_dir_path + print(f"{self.log_prefix} full-cache enabled: cap={full_cap} " + f"slots=[{self._full_slot_base},{self._full_slot_base + full_cap}) " + f"dir={cache_dir_path}", flush=True) + + def lookup_full(self, prompt_ids: list[int]) -> tuple[int, str, int] | None: + """Exact-match on full prompt_ids hash (keyed on raw, pre-compression ids). + + Returns ``(slot, cached_cur_bin_path, cur_ids_len)`` on hit, else None. + The cur_bin_path points to a file in the persistent cache dir that the + caller passes directly to the daemon as a RESTORE command's second arg. + + Caller must hold daemon_lock before inspecting the returned slot. + """ + if getattr(self, "_full_disabled", True): + return None + key = hash_prefix(prompt_ids, self.kv_k_type, self.fa_window) + entry = self.full_entries.get(key) + if entry is None: + return None + slot, cur_bin_path, cur_ids_len = entry + # Verify the cached file still exists (could have been deleted externally). + if not Path(cur_bin_path).exists(): + self.full_entries.pop(key, None) + return None + self.full_entries.move_to_end(key) # mark fresh in LRU + print(f"{self.log_prefix} full-cache hit slot={slot} " + f"cur_ids_len={cur_ids_len} key={key.hex()[:8]}", flush=True) + return slot, cur_bin_path, cur_ids_len + + def prepare_full_snap(self, prompt_ids: list[int]) -> tuple[int, int] | None: + """Reserve a daemon slot for the full-prefill snapshot. + + Returns ``(absolute_slot, 0)`` — the second element is a placeholder; + the real target_pos (== len(cur_ids)) is supplied by the caller to + ``confirm_full_snap``. Returns None if full-cache is disabled or the + prompt is already cached. + """ + if getattr(self, "_full_disabled", True): + return None + key = hash_prefix(prompt_ids, self.kv_k_type, self.fa_window) + if key in self.full_entries: + self.full_entries.move_to_end(key) + return None # already cached + + # Pick a slot, deferring eviction until confirm succeeds. + if len(self.full_entries) >= self._full_cap: + old_key = next(iter(self.full_entries)) + old_slot, old_path, _ = self.full_entries[old_key] + self._full_pending_evict_key = old_key + self._full_pending_evict_path = old_path + abs_slot = old_slot + else: + abs_slot = self._full_slot_base + self._full_next_slot + self._full_next_slot = (self._full_next_slot + 1) % self._full_cap + self._full_pending_evict_key = None + self._full_pending_evict_path = None + + return abs_slot, 0 # 0 is a placeholder; real pos passed to confirm + + def confirm_full_snap(self, slot: int, prompt_ids: list[int], + cur_bin_src: str | Path, cur_ids_len: int) -> None: + """Persist cur_bin_src into the cache dir and register the entry. + + ``cur_bin_src`` is the path to the tempfile written by _maybe_compress; + its content is copied (not moved, to keep the original available for the + daemon) into the persistent cache dir before registering. + + Atomically evicts the LRU entry (and its on-disk file) if one was + reserved by prepare_full_snap. + """ + if getattr(self, "_full_disabled", True): + return + + key = hash_prefix(prompt_ids, self.kv_k_type, self.fa_window) + dest = self._full_cache_dir / (key.hex() + ".bin") + + try: + shutil.copy2(str(cur_bin_src), str(dest)) + except OSError as exc: + print(f"{self.log_prefix} full-cache: failed to copy cur_bin " + f"({cur_bin_src} -> {dest}): {exc}", flush=True) + # Don't evict the old entry — leave cache consistent. + self._full_pending_evict_key = None + self._full_pending_evict_path = None + return + + # Atomically evict the reserved entry (if any) and insert new one. + if self._full_pending_evict_key is not None: + evicted_path = self._full_pending_evict_path + self.full_entries.pop(self._full_pending_evict_key, None) + if evicted_path: + Path(evicted_path).unlink(missing_ok=True) + self._full_pending_evict_key = None + self._full_pending_evict_path = None + + self.full_entries[key] = (slot, str(dest), cur_ids_len) + print(f"{self.log_prefix} full-cache committed slot={slot} " + f"cur_ids_len={cur_ids_len} key={key.hex()[:8]}", flush=True) + + def abort_full_snap(self, slot: int) -> None: + """Cancel a prepare_full_snap reservation without registering anything. + + Clears the pending eviction so the old LRU entry is not evicted. If the + daemon may have committed KV to *slot* before the failure, also drop + stale hash→slot mappings (conservative, same as abort_inline_snap). + """ + if getattr(self, "_full_disabled", True): + return + if self._full_pending_evict_key is not None: + self.full_entries.pop(self._full_pending_evict_key, None) + self._full_pending_evict_key = None + self._full_pending_evict_path = None + stale_keys = [k for k, (s, _, _) in self.full_entries.items() if s == slot] + for k in stale_keys: + self.full_entries.pop(k, None) + + # Legacy out-of-band snapshot (kept for backward-compatibility tests + # that call it directly; new code uses prepare_inline_snap + + # confirm_inline_snap so the snapshot rides on the actual response). + async def maybe_snapshot(self, prompt_ids: list[int], + token_stream_consumer=None) -> None: + if self.disabled: + return + prep = self.prepare_inline_snap(prompt_ids) + if prep is None: + return + slot, cut = prep + + import os, struct, tempfile + fd, tmp_path = tempfile.mkstemp(suffix="_prefix.bin") + with os.fdopen(fd, "wb") as f: + for t in prompt_ids[:cut]: + f.write(struct.pack(" None: + self.stdin.write(line.encode("utf-8")) + self.stdin.flush() diff --git a/server/scripts/server_tools.py b/server/scripts/server_tools.py new file mode 100644 index 000000000..a666b5107 --- /dev/null +++ b/server/scripts/server_tools.py @@ -0,0 +1,1872 @@ +""" +OpenAI-compatible HTTP server on top of test_dflash, **with tool-calling support**. + +Patched fork of scripts/server.py that: + 1. Accepts the OpenAI `tools` array in ChatRequest. + 2. Renders tools into the prompt via Qwen's chat template (`tools=...`). + 3. Parses `` blocks out + of the model output and returns them as proper OpenAI `tool_calls`. + 4. Supports `role: "tool"` and assistant `tool_calls` in input messages so + multi-turn agent loops round-trip correctly. + +Streaming behavior: + - Content tokens are streamed as `delta.content` until a `` opener + is detected; the rest of the response is then buffered, parsed at the end + of generation, and emitted as a single final `delta.tool_calls` chunk with + `finish_reason: "tool_calls"`. + - If no tool call appears in the output, behavior is identical to the + upstream server. + +Greedy decoding still applies (verify path is greedy-only). `temperature` and +`top_p` are accepted but ignored, matching upstream. + +When ``tools`` is non-empty, ``enable_thinking`` is forced off in the chat template +(Qwen3 thinking prefill otherwise dominates and the model rarely emits ````). + +Run: + pip install fastapi uvicorn transformers + python3 scripts/server_tools.py --port 8000 +""" +import argparse +import asyncio +import json +import os +import re +import struct +import subprocess +import sys +import tempfile +import time +import uuid +from pathlib import Path +from typing import Any, AsyncIterator + +from fastapi import FastAPI +from fastapi.responses import JSONResponse, StreamingResponse + +from _prefill_hook import ( + PrefillConfig, add_cli_flags, config_from_args, + compress_text_via_daemon, +) +from pydantic import BaseModel, ConfigDict +from starlette.concurrency import iterate_in_threadpool +from transformers import AutoTokenizer + +from prefix_cache import DaemonStdoutBus, PrefixCache +from tool_split.config import ToolSplitConfig, add_cli_flags as add_tool_split_flags +from tool_split.config import config_from_env_and_args as tool_split_config_from_args +from tool_split.base import ToolRequestContext +from tool_split.daemon_bridge import commit_pending_tool_snap +from tool_split.orchestrator import ToolSplitOrchestrator +from tool_split.registry import resolve_adapter as resolve_tool_split_adapter + + +ROOT = Path(__file__).resolve().parent.parent +DEFAULT_TARGET = Path(os.environ.get( + "DFLASH_TARGET", + str(ROOT / "models" / "Qwen3.6-27B-Q4_K_M.gguf"), +)) +DEFAULT_DRAFT_ROOT = ROOT / "models" / "draft" +DEFAULT_BIN = ROOT / "build" / ("test_dflash" + (".exe" if sys.platform == "win32" else "")) +DEFAULT_BUDGET = 22 +MODEL_NAME = "luce-dflash" + +# Passed through to apply_chat_template only (see server.py — avoid arbitrary kwargs). +_ALLOWED_TEMPLATE_KWARGS = frozenset({"enable_thinking", "add_generation_prompt", "tools"}) + + +def _extra_daemon_has_target_sharding(extra: list[str] | None) -> bool: + """True if we spawn test_dflash with multi-GPU target layer split.""" + if not extra: + return False + return any(tok.startswith("--target-gpus") for tok in extra) + + +# Architecture strings in `general.architecture` of the GGUF (see server.py). +_QWEN35_ARCHES = {"qwen35", "qwen36"} +_LAGUNA_ARCHES = {"laguna"} + +_QWEN35_FAMILY_TOKENIZERS = { + "Qwen3.5-27B": "Qwen/Qwen3.5-27B", + "Qwen3.6-27B": "Qwen/Qwen3.6-27B", +} +_LAGUNA_FAMILY_TOKENIZERS = { + "Laguna-XS.2": "poolside/Laguna-XS.2", + "Laguna-XS": "poolside/Laguna-XS.2", + "laguna-xs2": "poolside/Laguna-XS.2", +} + + +def _read_gguf_str(reader, key: str) -> str | None: + f = reader.fields.get(key) + if f is None or not f.data: + return None + import numpy as np + p = f.parts[f.data[0]] + if not isinstance(p, np.ndarray): + return None + try: + return bytes(p).decode("utf-8", errors="replace") + except Exception: + return None + + +def _arch_from_gguf(gguf_path: Path) -> str: + try: + from gguf import GGUFReader # type: ignore + r = GGUFReader(str(gguf_path)) + v = _read_gguf_str(r, "general.architecture") + return v.lower() if v else "unknown" + except Exception: + return "unknown" + + +def _tokenizer_id_from_gguf(gguf_path: Path) -> str: + default = "Qwen/Qwen3.5-27B" + try: + from gguf import GGUFReader # type: ignore + r = GGUFReader(str(gguf_path)) + arch = (_read_gguf_str(r, "general.architecture") or "").lower() + family = _LAGUNA_FAMILY_TOKENIZERS if arch in _LAGUNA_ARCHES else _QWEN35_FAMILY_TOKENIZERS + if arch in _LAGUNA_ARCHES: + default = next(iter(_LAGUNA_FAMILY_TOKENIZERS.values())) + for key in ("general.basename", "general.name"): + val = _read_gguf_str(r, key) + if val is None: + continue + for known, repo in family.items(): + if known.lower() in val.lower(): + return repo + except Exception: + pass + return default + + +def resolve_draft(root: Path) -> Path: + for st in root.rglob("model.safetensors"): + return st + raise FileNotFoundError(f"no model.safetensors under {root}") + + +# ─── pydantic schemas ────────────────────────────────────────────── + +class ToolCallFunction(BaseModel): + name: str + arguments: str # JSON string per OpenAI spec + + +class ToolCall(BaseModel): + id: str | None = None + type: str = "function" + function: ToolCallFunction + + +class ChatMessage(BaseModel): + role: str + content: Any | None = None # str, list, or null when tool_calls present + name: str | None = None + tool_call_id: str | None = None + tool_calls: list[ToolCall] | None = None + + +class ToolDef(BaseModel): + type: str = "function" + function: dict # {name, description, parameters: {...JSON schema...}} + + +class ChatRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + model: str = MODEL_NAME + messages: list[ChatMessage] + stream: bool = False + max_tokens: int = 512 + # OpenAI-compatible alias (newer clients send this instead of max_tokens). + max_completion_tokens: int | None = None + temperature: float | None = None + top_p: float | None = None + tools: list[ToolDef] | None = None + tool_choice: Any | None = None # "auto" | "none" | {"type":"function",...} | {"type":"required"} + chat_template_kwargs: dict | None = None # e.g. {"enable_thinking": false} + stop: Any | None = None # str or list[str] + stream_options: dict | None = None # e.g. {"include_usage": true} + + +class AnthropicMessage(BaseModel): + role: str + # Anthropic allows either a plain string or a list of content blocks. + content: str | list[dict] + + +class AnthropicMessagesRequest(BaseModel): + model: str = MODEL_NAME + max_tokens: int + messages: list[AnthropicMessage] + system: str | list[dict] | None = None + stream: bool = False + temperature: float | None = None + top_p: float | None = None + stop_sequences: list[str] | None = None + + +# ─── tool-call parser ────────────────────────────────────────────── + +# Qwen3.6 chat template emits: +# +# +# +# VALUE +# +# ... +# +# +# Parsers ported from vLLM (Apache-2.0) for behavioral parity with +# `--reasoning-parser qwen3` and `--tool-call-parser qwen3_coder`: +# vllm/reasoning/qwen3_reasoning_parser.py +# vllm/tool_parsers/qwen3coder_tool_parser.py +# Core algorithms reproduced without vLLM runtime dependencies. + +TOOL_CALL_COMPLETE_RE = re.compile(r"(.*?)", re.DOTALL) +TOOL_CALL_FUNCTION_RE = re.compile( + r"| by using +# next or end-of-string as a terminator. +TOOL_CALL_PARAMETER_RE = re.compile( + r"|(?=)|$)", + re.DOTALL, +) +TOOL_OPEN_TAG = "" + +# Qwen3.6 chat template wraps the model's CoT inside .... +# The template typically prefills `\n` into the prompt (headless mode) +# so only `` appears in generated output; older templates emit both. +THINK_OPEN_TAG = "" +THINK_CLOSE_TAG = "" + + +def normalize_stop(stop) -> list[str]: + """Coerce OpenAI's stop field (str | list[str] | None) to list[str].""" + if not stop: + return [] + if isinstance(stop, str): + return [stop] + return [s for s in stop if isinstance(s, str) and s] + + +def first_stop_match(text: str, stops: list[str]) -> int: + """Return the earliest index where any stop sequence appears, or -1.""" + best = -1 + for s in stops: + i = text.find(s) + if i != -1 and (best == -1 or i < best): + best = i + return best + + +def trim_at_stop(text: str, stops: list[str]) -> tuple[str, str | None]: + """Trim *text* at the first stop sequence. Returns (trimmed, matched_stop).""" + if not stops: + return text, None + best_i = -1 + matched: str | None = None + for s in stops: + i = text.find(s) + if i != -1 and (best_i == -1 or i < best_i): + best_i = i + matched = s + if best_i == -1: + return text, None + return text[:best_i], matched + + +def parse_reasoning(text: str, thinking_enabled: bool = True) -> tuple[str, str | None]: + """Port of vLLM's Qwen3ReasoningParser.extract_reasoning. + + Handles the three Qwen3.x thinking flavors: + 1. Paired: `...` both in generated output. + 2. Headless: template prefilled `\\n` into the prompt, model + only emits `......`. + 3. Disabled: user passed `chat_template_kwargs: {enable_thinking: false}`. + Template still emits `\\n\\n\\n\\n` but into the prompt; + the model output is pure content and contains no tags. + + If the output was truncated mid-thinking (no `` seen and + `thinking_enabled=True`), returns `("", full_output_as_reasoning)` — + matching vLLM's convention. + + Returns (cleaned_content, reasoning_content). + """ + # Strip if the model emitted it itself (older templates). + parts = text.partition(THINK_OPEN_TAG) + rest = parts[2] if parts[1] else parts[0] + if THINK_CLOSE_TAG not in rest: + if thinking_enabled: + # No close tag — assume truncated; everything is reasoning. + return "", (rest.strip() or None) + else: + # Thinking disabled — output is pure content. + return rest.strip(), None + reasoning, _, content = rest.partition(THINK_CLOSE_TAG) + return content.strip(), (reasoning.strip() or None) + + +def _find_tool_properties(tools, function_name): + """Helper matching vLLM's `find_tool_properties`: returns the parameters + dict for a given function name, or {} if not found. + Accepts pydantic ToolDef instances or plain dicts. + """ + for t in tools or []: + fn = t.function if hasattr(t, "function") else t.get("function", {}) + if hasattr(fn, "model_dump"): + fn = fn.model_dump() + if fn.get("name") == function_name: + params = fn.get("parameters", {}) + if isinstance(params, dict): + return params.get("properties", {}) + return {} + + +def _convert_param_value(param_value: str, param_name: str, param_config: dict, + func_name: str): + """Port of vLLM's _convert_param_value. Coerces stringified XML values + to their JSON-schema type (int/float/bool/object/array/string).""" + import ast + if param_value.lower() == "null": + return None + if param_name not in param_config: + return param_value + cfg = param_config[param_name] + if isinstance(cfg, dict) and "type" in cfg: + ptype = str(cfg["type"]).strip().lower() + elif isinstance(cfg, dict) and "anyOf" in cfg: + ptype = "object" + else: + ptype = "string" + if ptype in ("string", "str", "text", "varchar", "char", "enum"): + return param_value + if any(ptype.startswith(p) for p in ("int", "uint", "long", "short", "unsigned")): + try: return int(param_value) + except (ValueError, TypeError): return param_value + if ptype.startswith("num") or ptype.startswith("float"): + try: + f = float(param_value) + return f if f - int(f) != 0 else int(f) + except (ValueError, TypeError): + return param_value + if ptype in ("boolean", "bool", "binary"): + return param_value.lower() == "true" + # object / array / dict / list + if (ptype in ("object", "array", "arr") + or ptype.startswith("dict") or ptype.startswith("list")): + try: return json.loads(param_value) + except (json.JSONDecodeError, TypeError, ValueError): pass + try: return ast.literal_eval(param_value) + except (ValueError, SyntaxError, TypeError): return param_value + + +def parse_tool_calls(text: str, tools=None) -> tuple[str, list[dict]]: + """Port of Qwen3CoderToolParser._parse_xml_function_call (non-streaming). + + Handles Qwen3.x's `...VAL + ...` XML. Uses vLLM's improved + parameter regex that tolerates unclosed tags. When `tools` + is provided, each parameter value is coerced to its JSON-schema type. + + Returns (cleaned_content, tool_calls_list). + """ + tool_calls: list[dict] = [] + cleaned_parts: list[str] = [] + cursor = 0 + for m in TOOL_CALL_COMPLETE_RE.finditer(text): + cleaned_parts.append(text[cursor:m.start()]) + cursor = m.end() + body = m.group(1) + fn_match = TOOL_CALL_FUNCTION_RE.search(body) + if not fn_match: + continue + fn_text = fn_match.group(1) or fn_match.group(2) or "" + end_idx = fn_text.find(">") + if end_idx == -1: + continue + function_name = fn_text[:end_idx].strip() + params_region = fn_text[end_idx + 1:] + param_config = _find_tool_properties(tools, function_name) + args: dict = {} + for match_text in TOOL_CALL_PARAMETER_RE.findall(params_region): + eq_idx = match_text.find(">") + if eq_idx == -1: + continue + k = match_text[:eq_idx].strip() + v = match_text[eq_idx + 1:] + if v.startswith("\n"): v = v[1:] + if v.endswith("\n"): v = v[:-1] + args[k] = _convert_param_value(v, k, param_config, function_name) + tool_calls.append({ + "id": "call_" + uuid.uuid4().hex[:24], + "type": "function", + "function": { + "name": function_name, + "arguments": json.dumps(args, ensure_ascii=False), + }, + }) + cleaned_parts.append(text[cursor:]) + return "".join(cleaned_parts).strip(), tool_calls + + +# ─── app ─────────────────────────────────────────────────────────── + +def build_app(target: Path, draft: Path | None, bin_path: Path, budget: int, + max_ctx: int, tokenizer: AutoTokenizer, stop_ids: set[int], + prefill_cfg: PrefillConfig | None = None, + drafter_tokenizer: AutoTokenizer | None = None, + prefix_cache_slots: int = 4, + prefill_cache_slots: int = 4, + arch: str = "qwen35", + extra_daemon_args: list[str] | None = None, + tool_split_cfg: ToolSplitConfig | None = None, + tool_split: ToolSplitOrchestrator | None = None) -> FastAPI: + import asyncio + if _extra_daemon_has_target_sharding(extra_daemon_args): + if prefix_cache_slots > 0 or prefill_cache_slots > 0: + print( + " [cfg] target-gpus sharding: disabling prefix/full cache " + "(daemon SNAPSHOT/RESTORE not implemented for this mode)", + flush=True, + ) + prefix_cache_slots = 0 + prefill_cache_slots = 0 + if tool_split is not None: + print( + " [cfg] target-gpus sharding: disabling tool-split " + "(SNAPSHOT_THIN / RESTORE_CHAIN unsupported)", + flush=True, + ) + tool_split = None + app = FastAPI(title="Luce DFlash OpenAI server (tool-aware)") + daemon_lock = asyncio.Lock() + + r_pipe, w_pipe = os.pipe() + if sys.platform == "win32": + import msvcrt + os.set_inheritable(w_pipe, True) + stream_fd_val = int(msvcrt.get_osfhandle(w_pipe)) + else: + stream_fd_val = w_pipe + + bin_abs = str(Path(bin_path).resolve()) + dll_dir = str(Path(bin_abs).parent / "bin") + env = {**os.environ} + if sys.platform == "win32": + env["PATH"] = dll_dir + os.pathsep + str(Path(bin_abs).parent) + os.pathsep + env.get("PATH", "") + + if arch in _LAGUNA_ARCHES: + cmd = [bin_abs, str(target), "--daemon", + f"--max-ctx={max_ctx}", + f"--stream-fd={stream_fd_val}"] + else: + if draft is None: + raise SystemExit("qwen35 arch requires --draft model.safetensors") + cmd = [bin_abs, str(target), str(draft), "--daemon", + "--fast-rollback", "--ddtree", f"--ddtree-budget={budget}", + f"--max-ctx={max_ctx}", + f"--stream-fd={stream_fd_val}"] + if extra_daemon_args: + cmd.extend(extra_daemon_args) + if sys.platform == "win32": + daemon_proc = subprocess.Popen(cmd, close_fds=False, env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, bufsize=0) + else: + daemon_proc = subprocess.Popen(cmd, pass_fds=(w_pipe,), env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, bufsize=0) + os.close(w_pipe) + + bus = DaemonStdoutBus(daemon_proc.stdout) + # Mirror server.py: resolve effective KV-K type + FA window from env so + # they participate in the prefix-cache hash key. + def _resolve_kv_k_type(): + kv = "q8_0" + if os.environ.get("DFLASH27B_KV_F16", "0") != "0": + kv = "f16" + if os.environ.get("DFLASH27B_KV_Q4", "0") != "0": + kv = "q4_0" + if os.environ.get("DFLASH27B_KV_TQ3", "0") != "0": + kv = "tq3_0" + if os.environ.get("DFLASH27B_KV_K"): + kv = os.environ["DFLASH27B_KV_K"].lower() + return kv + _fa_window = int(os.environ.get("DFLASH27B_FA_WINDOW", 2048)) + prefix_cache = PrefixCache( + daemon_stdin=daemon_proc.stdin, + await_reply=bus.await_reply, + daemon_lock=daemon_lock, + tokenizer=tokenizer, + kv_k_type=_resolve_kv_k_type(), + fa_window=_fa_window, + cap=prefix_cache_slots, + ) + # Option 3: full-compress-result cache. Only meaningful when pFlash + # compression is enabled. Uses a separate slot range [prefix_cap, ...). + if prefill_cfg is not None and prefill_cache_slots > 0: + prefix_cache.init_full_cache(prefill_cache_slots) + + async def _finish_inline_snap( + snap_prep: tuple[int, int] | None, + prompt_ids: list[int], + ) -> None: + if snap_prep: + await bus.drain_inline_snap() + prefix_cache.finish_inline_snap( + snap_prep, prompt_ids, inline_slot=bus.inline_snap_slot()) + + def _abort_full_snap_if_needed(full_snap_prep) -> None: + if full_snap_prep is not None: + fslot, _ = full_snap_prep + prefix_cache.abort_full_snap(fslot) + + async def _finalize_request_snaps( + *, + full_snap_prep, + snap_prep: tuple[int, int] | None, + prompt_ids: list[int], + cur_bin: Path, + cur_ids: list[int] | None, + success: bool, + ) -> None: + if full_snap_prep is not None: + if success and cur_ids is not None: + fslot, _ = full_snap_prep + prefix_cache.confirm_full_snap( + fslot, prompt_ids, cur_bin, len(cur_ids)) + else: + _abort_full_snap_if_needed(full_snap_prep) + elif success: + await _finish_inline_snap(snap_prep, prompt_ids) + elif snap_prep: + prefix_cache.abort_inline_snap(snap_prep[0]) + + def _build_usage( + prompt_len: int, + completion_tokens: int, + *, + conv_prefix_len: int | None = None, + ) -> dict: + usage = { + "prompt_tokens": prompt_len, + "completion_tokens": completion_tokens, + "total_tokens": prompt_len + completion_tokens, + } + timings: dict[str, float | int] = {} + raw = bus.request_timings() + for key, val in raw.items(): + if isinstance(val, (int, float)) and val > 0: + timings[key] = val + if conv_prefix_len is not None and conv_prefix_len > 0: + timings["prefix_len"] = conv_prefix_len + if timings: + usage["timings"] = timings + return usage + + def _compose_daemon_cmd( + cur_bin: Path, + gen_len: int, + prompt_ids: list[int], + *, + full_hit, + compression_fired: bool, + full_snap_prep, + cur_ids: list[int] | None, + tool_ctx: ToolRequestContext | None, + ) -> tuple[str, object | None, int | None]: + """Build daemon stdin line, optional inline snap_prep, conv prefix len.""" + if full_hit is not None: + slot, cached_cur_bin, cached_len = full_hit + return f"RESTORE {slot} {cached_cur_bin} {gen_len}\n", None, cached_len + if compression_fired: + if full_snap_prep is not None: + fslot, _ = full_snap_prep + return f"{cur_bin} {gen_len} snap={len(cur_ids)}:{fslot}\n", None, None + return f"{cur_bin} {gen_len}\n", None, None + + hit = prefix_cache.lookup(prompt_ids) + conv_prefix_len = hit[1] if hit else None + reuse = hit[0] if hit else None + snap_prep = prefix_cache.prepare_inline_snap(prompt_ids, reuse_slot=reuse) + + if tool_split and tool_ctx and tool_ctx.tool_slot_hit is not None: + plan = tool_split.build_plan( + split=tool_ctx.split, + tools_fingerprint=tool_ctx.fingerprint, + prompt_bin=cur_bin, + prompt_len=len(prompt_ids), + tool_slot_hit=tool_ctx.tool_slot_hit, + conv_hit=hit, + snap_prep=snap_prep, + pending_tool_snap=tool_ctx.pending_tool_snap, + ) + # On chain restore with a thick conv slot, always refresh that + # slot in-place — never allocate a new prefix snapshot while thin + # tool KV is also resident (OOM on 24GB cards). + if ( + snap_prep + and plan.conv_restore_slot is not None + and plan.conv_restore_slot >= 0 + ): + slot, cut = snap_prep + if slot != plan.conv_restore_slot: + snap_prep = (plan.conv_restore_slot, cut) + cmd = tool_split.format_daemon_command(plan, gen_len) + print( + f" [tool-split] RESTORE_CHAIN thick={plan.conv_restore_slot} " + f"thin={plan.thin_slot_ids} prompt_tokens={len(prompt_ids)}", + flush=True, + ) + return cmd, snap_prep, conv_prefix_len + + if hit: + slot, _prefix_len = hit + cmd = f"RESTORE {slot} {cur_bin} {gen_len}" + else: + cmd = f"{cur_bin} {gen_len}" + if snap_prep: + cmd += f" snap={snap_prep[1]}:{snap_prep[0]}" + return cmd + "\n", snap_prep, conv_prefix_len + + async def _commit_tool_snap_if_needed(tool_ctx: ToolRequestContext | None) -> None: + if not tool_split or not tool_ctx or not tool_ctx.pending_tool_snap: + return + if not tool_ctx.fingerprint: + return + slot, kv_end = tool_ctx.pending_tool_snap + await commit_pending_tool_snap( + orchestrator=tool_split, + daemon_stdin=daemon_proc.stdin, + await_reply=bus.await_reply, + fingerprint=tool_ctx.fingerprint, + slot=slot, + kv_end=kv_end, + ) + + @app.on_event("startup") + async def _startup(): + import asyncio + bus.start(asyncio.get_running_loop()) + await prefix_cache.startup_sync() + + @app.get("/health") + def health(): + alive = daemon_proc.poll() is None + if not alive: + return JSONResponse( + {"status": "error", "detail": "daemon exited"}, + status_code=503, + ) + return {"status": "ok"} + + @app.get("/v1/models") + def list_models(): + return {"object": "list", + "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "luce"}]} + + def _maybe_compress_tool_chat(req: "ChatRequest", prompt_bin: Path, + prompt_len: int, started_in_thinking: bool + ) -> tuple[Path, int, bool]: + """If prefill is on and the request has no tools and the last user + message is long, run the daemon compress + re-tokenise. Returns + (bin, prompt_len, started_in_thinking) — the last is recomputed when + compression fires, otherwise passed through. + + When ``tool_split`` is active and ``req.tools`` is non-empty, tool + definitions stay in the template while only conversation text is + eligible for PFlash (same last-user-message compress path). + """ + if not prefill_cfg or not prefill_cfg.enabled: + return prompt_bin, prompt_len, started_in_thinking + tools_present = bool(req.tools) + if tools_present and not tool_split: + # Legacy: monolithic prompt — compressing mangles tool JSON in-stream. + return prompt_bin, prompt_len, started_in_thinking + + compress_len = prompt_len + if tools_present and tool_split: + try: + split = tool_split.split_request( + tokenizer, req.messages, req.tools, + chat_template_kwargs=req.chat_template_kwargs, + ) + if tool_split.conversation_compressible(split): + compress_len = split.conversation_len + except Exception as exc: + print(f" [tool-split] split failed, skipping conv compress: {exc}", + flush=True) + + if not prefill_cfg.should_compress(compress_len) or drafter_tokenizer is None: + return prompt_bin, prompt_len, started_in_thinking + + last_user = next((m for m in reversed(req.messages) if m.role == "user"), None) + if last_user is None or not isinstance(last_user.content, str): + return prompt_bin, prompt_len, started_in_thinking + + compressed_text = compress_text_via_daemon( + daemon_stdin=daemon_proc.stdin, + r_pipe=r_pipe, + drafter_tokenizer=drafter_tokenizer, + cfg=prefill_cfg, + prompt_text=last_user.content, + skip_park=prefill_cfg.skip_park, + ) + + new_msgs = [] + compressed_emitted = False + for m in req.messages: + if m is last_user and not compressed_emitted: + new_msgs.append({"role": "user", "content": compressed_text}) + compressed_emitted = True + else: + d = {"role": m.role} + if m.content is not None: + d["content"] = m.content + new_msgs.append(d) + + kwargs: dict = { + "tokenize": False, + "add_generation_prompt": True, + "enable_thinking": False, + } + kwargs.update( + {k: v for k, v in (req.chat_template_kwargs or {}).items() + if k in _ALLOWED_TEMPLATE_KWARGS}, + ) + if req.tools: + kwargs["tools"] = [t.model_dump() for t in req.tools] + kwargs["enable_thinking"] = False + prompt = tokenizer.apply_chat_template(new_msgs, **kwargs) + new_started_in_thinking = bool(re.search(r"\s*$", prompt)) + ids = tokenizer.encode(prompt, add_special_tokens=False) + fd, path = tempfile.mkstemp(suffix=".bin") + with os.fdopen(fd, "wb") as f: + for t in ids: + f.write(struct.pack(" tuple[Path, bool]: + """Returns (prompt_bin_path, started_in_thinking). started_in_thinking + is True when the chat template prefilled \\n at the end of the + prompt — the model's first emitted tokens are reasoning content.""" + # Convert pydantic messages to dicts the chat template expects. + msgs: list[dict] = [] + for m in req.messages: + d: dict = {"role": m.role} + if m.content is not None: + d["content"] = m.content + if m.name is not None: + d["name"] = m.name + if m.tool_call_id is not None: + d["tool_call_id"] = m.tool_call_id + if m.tool_calls is not None: + # The Qwen template walks tool_calls[i].function.{name, arguments} + d["tool_calls"] = [] + for tc in m.tool_calls: + args = tc.function.arguments + # Template expects arguments as a dict, not a JSON string. + if isinstance(args, str): + try: + args_obj = json.loads(args) + except (json.JSONDecodeError, ValueError): + args_obj = {"_raw": args} + else: + args_obj = args + d["tool_calls"].append({ + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": args_obj}, + }) + msgs.append(d) + + tools_arg = None + if req.tools: + # OpenAI-shaped tool defs (type + function{name,description,parameters}). + tools_arg = [t.model_dump() for t in req.tools] + + # Mirror server.py: default enable_thinking=False. Qwen3's template default is + # often True, which pre-fills and the model rambles in + # "reasoning" instead of emitting XML — looks like "ignores tools". + kwargs: dict = { + "tokenize": False, + "add_generation_prompt": True, + "enable_thinking": False, + } + kwargs.update( + {k: v for k, v in (req.chat_template_kwargs or {}).items() + if k in _ALLOWED_TEMPLATE_KWARGS}, + ) + if tools_arg: + kwargs["tools"] = tools_arg + # Thinking + tool XML is a bad combo for Qwen3.x; never let merged kwargs + # re-enable thinking while tools are mounted. + kwargs["enable_thinking"] = False + if req.tool_choice is not None: + kwargs["tool_choice"] = req.tool_choice + prompt = tokenizer.apply_chat_template(msgs, **kwargs) + # Did the template prefill `\n` at the end? Then streaming should + # start in reasoning mode. + started_in_thinking = bool(re.search(r"\s*$", prompt)) + ids = tokenizer.encode(prompt, add_special_tokens=False) + fd, path = tempfile.mkstemp(suffix=".bin") + tmp = Path(path) + with os.fdopen(fd, "wb") as f: + for t in ids: + f.write(struct.pack("= n_gen: + hit_stop = True + + def _max_gen_tokens(req: ChatRequest) -> int: + if req.max_completion_tokens is not None: + return req.max_completion_tokens + return req.max_tokens + + @app.post("/v1/chat/completions") + async def chat_completions(req: ChatRequest): + prompt_bin, started_in_thinking = _tokenize_prompt(req) + prompt_len = prompt_bin.stat().st_size // 4 + + # Read back token ids for cache key (cheap — file is small). + raw = prompt_bin.read_bytes() + prompt_ids = [struct.unpack_from("` must not turn + # the whole completion into reasoning-only. + cleaned, reasoning = parse_reasoning( + cleaned, thinking_enabled=started_in_thinking) + + msg: dict = {"role": "assistant"} + finish_reason = "stop" + if reasoning: + msg["reasoning_content"] = reasoning + if tool_calls: + msg["content"] = cleaned if cleaned else None + msg["tool_calls"] = tool_calls + finish_reason = "tool_calls" + else: + msg["content"] = cleaned + + return JSONResponse({ + "id": completion_id, + "object": "chat.completion", + "created": created, + "model": MODEL_NAME, + "choices": [{ + "index": 0, + "message": msg, + "finish_reason": finish_reason, + }], + "usage": _build_usage( + prompt_len, len(tokens), conv_prefix_len=conv_prefix_len), + }) + + async def _stream_response(req, prompt_bin, prompt_ids, gen_len, completion_id, + created, started_in_thinking, lock, + full_hit=None, full_snap_prep=None, + compression_fired=False, cur_ids=None, + tool_ctx: ToolRequestContext | None = None): + # prompt_bin may be cur_bin (the compressed bin) when coming from the + # compression or full-cache-hit path; prompt_len is derived from it. + prompt_len = prompt_bin.stat().st_size // 4 if full_hit is None else ( + full_hit[2] # cached_cur_ids_len + ) + include_usage = bool(req.stream_options and req.stream_options.get("include_usage")) + def chunk(delta_obj, finish=None): + return {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": MODEL_NAME, + "choices": [{"index": 0, "delta": delta_obj, + "finish_reason": finish}]} + + async def sse() -> AsyncIterator[str]: + async with lock: + cmd_line, snap_prep, conv_prefix_len = _compose_daemon_cmd( + prompt_bin, gen_len, prompt_ids, + full_hit=full_hit, + compression_fired=compression_fired, + full_snap_prep=full_snap_prep, + cur_ids=cur_ids, + tool_ctx=tool_ctx, + ) + bus.begin_request() + daemon_proc.stdin.write(cmd_line.encode("utf-8")) + daemon_proc.stdin.flush() + + yield f"data: {json.dumps(chunk({'role': 'assistant'}))}\n\n" + + # State machine: mode ∈ {'reasoning', 'content', 'tool_buffer'} + mode = "reasoning" if started_in_thinking else "content" + window = "" # holdback buffer for tag detection + tool_buffer = "" + stops = normalize_stop(req.stop) + # Holdback must cover longest tag AND longest stop sequence. + tag_holdback = max(len(THINK_OPEN_TAG), len(THINK_CLOSE_TAG), len(TOOL_OPEN_TAG)) + stop_holdback = max((len(s) for s in stops), default=0) + HOLDBACK = max(tag_holdback, stop_holdback) + completion_tokens = 0 + stop_hit = False + + def emit_delta(text, kind): + """kind: 'content' or 'reasoning_content'""" + if not text: + return None + return f"data: {json.dumps(chunk({kind: text}))}\n\n" + + snap_ok = False + try: + async for tok_id in iterate_in_threadpool(_token_stream(r_pipe, gen_len)): + completion_tokens += 1 + piece = tokenizer.decode([tok_id]) + window += piece + + # Stop-sequence check on the visible (content/reasoning) stream. + if stops and mode != "tool_buffer": + si = first_stop_match(window, stops) + if si != -1: + window = window[:si] + stop_hit = True + # Flush truncated remainder per current mode. + kind = "reasoning_content" if mode == "reasoning" else "content" + out = emit_delta(window, kind) + if out: yield out + window = "" + break + + # Process state transitions until no more tags found in window. + while True: + if mode == "tool_buffer": + tool_buffer += window + window = "" + break + + # Look for the next tag of interest based on mode. + if mode == "reasoning": + idx = window.find(THINK_CLOSE_TAG) + if idx != -1: + pre = window[:idx] + out = emit_delta(pre, "reasoning_content") + if out: yield out + window = window[idx + len(THINK_CLOSE_TAG):] + mode = "content" + continue + # No close tag yet. Stream all but holdback. + if len(window) > HOLDBACK: + safe = window[:-HOLDBACK] + out = emit_delta(safe, "reasoning_content") + if out: yield out + window = window[-HOLDBACK:] + break # need more tokens + + else: # mode == "content" + think_idx = window.find(THINK_OPEN_TAG) + tool_idx = window.find(TOOL_OPEN_TAG) + # Pick the earliest tag that actually appears. + hits = [(i, t) for i, t in + ((think_idx, "think"), (tool_idx, "tool")) if i != -1] + if hits: + hits.sort() + idx, which = hits[0] + pre = window[:idx] + out = emit_delta(pre, "content") + if out: yield out + if which == "think": + window = window[idx + len(THINK_OPEN_TAG):] + mode = "reasoning" + else: # tool + tool_buffer = window[idx:] + window = "" + mode = "tool_buffer" + continue + if len(window) > HOLDBACK: + safe = window[:-HOLDBACK] + out = emit_delta(safe, "content") + if out: yield out + window = window[-HOLDBACK:] + break # need more tokens + + if stop_hit: + finish_reason = "stop" + yield f"data: {json.dumps(chunk({}, finish=finish_reason))}\n\n" + if include_usage: + usage_chunk = {"id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": MODEL_NAME, "choices": [], + "usage": _build_usage( + prompt_len, completion_tokens, + conv_prefix_len=conv_prefix_len)} + yield f"data: {json.dumps(usage_chunk)}\n\n" + yield "data: [DONE]\n\n" + await bus.drain_timings() + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=prompt_bin, + cur_ids=cur_ids, + success=True, + ) + await _commit_tool_snap_if_needed(tool_ctx) + if full_hit is None: + try: prompt_bin.unlink() + except Exception: pass + return + + # Generation done. Flush remaining window per current mode. + if mode == "reasoning" and window: + out = emit_delta(window, "reasoning_content") + if out: yield out + elif mode == "content" and window: + out = emit_delta(window, "content") + if out: yield out + elif mode == "tool_buffer": + tool_buffer += window + window = "" + + finish_reason = "stop" + if mode == "tool_buffer": + cleaned_after, tool_calls = parse_tool_calls(tool_buffer, tools=req.tools) + if tool_calls: + if cleaned_after: + out = emit_delta(cleaned_after, "content") + if out: yield out + tc_delta_list = [{ + "index": i, "id": tc["id"], "type": "function", + "function": {"name": tc["function"]["name"], + "arguments": tc["function"]["arguments"]}, + } for i, tc in enumerate(tool_calls)] + yield f"data: {json.dumps(chunk({'tool_calls': tc_delta_list}))}\n\n" + finish_reason = "tool_calls" + else: + # Unclosed — emit raw as content fallback. + out = emit_delta(tool_buffer, "content") + if out: yield out + snap_ok = True + except Exception: + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=prompt_bin, + cur_ids=cur_ids, + success=False, + ) + raise + finally: + if full_hit is None: + try: prompt_bin.unlink() + except Exception: pass + + await bus.drain_timings() + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=prompt_bin, + cur_ids=cur_ids, + success=snap_ok, + ) + await _commit_tool_snap_if_needed(tool_ctx) + + yield f"data: {json.dumps(chunk({}, finish=finish_reason))}\n\n" + if include_usage: + usage_chunk = { + "id": completion_id, "object": "chat.completion.chunk", + "created": created, "model": MODEL_NAME, + "choices": [], + "usage": _build_usage( + prompt_len, completion_tokens, + conv_prefix_len=conv_prefix_len), + } + yield f"data: {json.dumps(usage_chunk)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") + + # ── Anthropic Messages API ────────────────────────────────────── + # Mirrors the OpenAI endpoint but formatted for the Anthropic SDK + # (Claude Code, Anthropic clients). Tool calling NOT forwarded here + # yet — agent CLIs that want tools should use /v1/chat/completions. + + def _anthropic_text_from_content(content) -> str: + if isinstance(content, str): + return content + parts = [] + for b in content: + if isinstance(b, dict) and b.get("type") == "text": + parts.append(b.get("text", "")) + return "".join(parts) + + def _tokenize_anthropic(req: AnthropicMessagesRequest + ) -> tuple[Path, int, list[dict]]: + msgs = [] + system_text = _anthropic_text_from_content(req.system) if req.system else None + if system_text: + msgs.append({"role": "system", "content": system_text}) + for m in req.messages: + msgs.append({"role": m.role, + "content": _anthropic_text_from_content(m.content)}) + prompt = tokenizer.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=True) + ids = tokenizer.encode(prompt, add_special_tokens=False) + # mkstemp returns (fd, path); discarding fd leaks 1 per request (#15). + fd, path = tempfile.mkstemp(suffix=".bin") + tmp = Path(path) + with os.fdopen(fd, "wb") as f: + for t in ids: + f.write(struct.pack(" tuple[Path, int]: + if not prefill_cfg or not prefill_cfg.enabled: + return prompt_bin, prompt_len + if not prefill_cfg.should_compress(prompt_len) or drafter_tokenizer is None: + return prompt_bin, prompt_len + last_user_idx = next((i for i in range(len(msgs) - 1, -1, -1) + if msgs[i]["role"] == "user"), None) + if last_user_idx is None: + return prompt_bin, prompt_len + long_text = msgs[last_user_idx]["content"] + compressed_text = compress_text_via_daemon( + daemon_stdin=daemon_proc.stdin, + r_pipe=r_pipe, + drafter_tokenizer=drafter_tokenizer, + cfg=prefill_cfg, + prompt_text=long_text, + skip_park=prefill_cfg.skip_park, + ) + new_msgs = list(msgs) + new_msgs[last_user_idx] = {"role": "user", "content": compressed_text} + prompt = tokenizer.apply_chat_template( + new_msgs, tokenize=False, add_generation_prompt=True) + ids = tokenizer.encode(prompt, add_special_tokens=False) + fd, path = tempfile.mkstemp(suffix=".bin") + with os.fdopen(fd, "wb") as f: + for t in ids: + f.write(struct.pack("= n_gen: + hit_stop = True + + @app.post("/v1/messages") + async def anthropic_messages(req: AnthropicMessagesRequest): + prompt_bin, prompt_len, raw_msgs = _tokenize_anthropic(req) + + # Read raw prompt_ids BEFORE compression (for full-cache key). + raw = prompt_bin.read_bytes() + prompt_ids = [struct.unpack_from(" AsyncIterator[str]: + async with daemon_lock: + if full_hit is not None: + slot, cached_cur_bin, _cached_len = full_hit + cmd_line = f"RESTORE {slot} {cached_cur_bin} {gen_len}\n" + snap_prep = None + elif compression_fired: + if full_snap_prep is not None: + fslot, _ = full_snap_prep + cmd_line = f"{cur_bin} {gen_len} snap={len(cur_ids)}:{fslot}\n" + else: + cmd_line = f"{cur_bin} {gen_len}\n" + snap_prep = None + else: + hit = prefix_cache.lookup(prompt_ids) + snap_prep = prefix_cache.prepare_inline_snap(prompt_ids) + if hit: + slot, _prefix_len = hit + cmd_line = f"RESTORE {slot} {cur_bin} {gen_len}" + else: + cmd_line = f"{cur_bin} {gen_len}" + if snap_prep: + cmd_line += f" snap={snap_prep[1]}:{snap_prep[0]}" + cmd_line += "\n" + + message_start = { + "type": "message_start", + "message": { + "id": msg_id, "type": "message", "role": "assistant", + "model": req.model or MODEL_NAME, + "content": [], "stop_reason": None, "stop_sequence": None, + "usage": {"input_tokens": prompt_len, "output_tokens": 0}, + }, + } + yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n" + + cb_start = { + "type": "content_block_start", "index": 0, + "content_block": {"type": "text", "text": ""}, + } + yield f"event: content_block_start\ndata: {json.dumps(cb_start)}\n\n" + + bus.begin_request() + daemon_proc.stdin.write(cmd_line.encode("utf-8")) + daemon_proc.stdin.flush() + + out_tokens = 0 + stop_reason = "end_turn" + matched_stop: str | None = None + holdback = max((len(s) for s in user_stops), default=0) + window = "" + snap_ok = False + try: + async for tok_id in _astream_tokens(r_pipe, gen_len): + out_tokens += 1 + piece = tokenizer.decode([tok_id]) + if user_stops: + window += piece + si = first_stop_match(window, user_stops) + if si != -1: + for s in user_stops: + if window[si:si + len(s)] == s: + matched_stop = s + break + emit = window[:si] + stop_reason = "stop_sequence" + if emit: + delta = { + "type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": emit}, + } + yield f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n" + break + if len(window) > holdback: + emit = window[:-holdback] + window = window[-holdback:] + if emit: + delta = { + "type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": emit}, + } + yield f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n" + else: + delta = { + "type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": piece}, + } + yield f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n" + else: + if user_stops and window: + trimmed, matched_stop = trim_at_stop(window, user_stops) + if matched_stop is not None: + stop_reason = "stop_sequence" + window = trimmed + if window: + delta = { + "type": "content_block_delta", "index": 0, + "delta": {"type": "text_delta", "text": window}, + } + yield f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n" + snap_ok = True + except Exception: + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=cur_bin, + cur_ids=cur_ids, + success=False, + ) + raise + finally: + if full_hit is None: + try: cur_bin.unlink() + except Exception: pass + else: + # On full-cache hit, cur_bin points at the persistent cached file + # (which we MUST keep). The tokenize-stage prompt_bin tempfile, on + # the other hand, was never used (we hit before _maybe_compress) and + # would otherwise leak. + try: prompt_bin.unlink() + except Exception: pass + + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=cur_bin, + cur_ids=cur_ids, + success=snap_ok, + ) + + yield f"event: content_block_stop\ndata: {json.dumps({'type': 'content_block_stop', 'index': 0})}\n\n" + + msg_delta = { + "type": "message_delta", + "delta": { + "stop_reason": stop_reason, + "stop_sequence": matched_stop, + }, + "usage": {"output_tokens": out_tokens}, + } + yield f"event: message_delta\ndata: {json.dumps(msg_delta)}\n\n" + yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") + + # Non-streaming + async with daemon_lock: + if full_hit is not None: + slot, cached_cur_bin, _cached_len = full_hit + cmd_line = f"RESTORE {slot} {cached_cur_bin} {gen_len}\n" + snap_prep = None + elif compression_fired: + if full_snap_prep is not None: + fslot, _ = full_snap_prep + cmd_line = f"{cur_bin} {gen_len} snap={len(cur_ids)}:{fslot}\n" + else: + cmd_line = f"{cur_bin} {gen_len}\n" + snap_prep = None + else: + hit = prefix_cache.lookup(prompt_ids) + snap_prep = prefix_cache.prepare_inline_snap(prompt_ids) + if hit: + slot, _prefix_len = hit + cmd_line = f"RESTORE {slot} {cur_bin} {gen_len}" + else: + cmd_line = f"{cur_bin} {gen_len}" + if snap_prep: + cmd_line += f" snap={snap_prep[1]}:{snap_prep[0]}" + cmd_line += "\n" + bus.begin_request() + daemon_proc.stdin.write(cmd_line.encode("utf-8")) + daemon_proc.stdin.flush() + snap_ok = False + try: + tokens = [t async for t in _astream_tokens(r_pipe, gen_len)] + snap_ok = True + except Exception: + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=cur_bin, + cur_ids=cur_ids, + success=False, + ) + raise + await _finalize_request_snaps( + full_snap_prep=full_snap_prep, + snap_prep=snap_prep, + prompt_ids=prompt_ids, + cur_bin=cur_bin, + cur_ids=cur_ids, + success=snap_ok, + ) + + if full_hit is None: + try: cur_bin.unlink() + except Exception: pass + else: + # On full-cache hit, cur_bin points at the persistent cached file + # (which we MUST keep). The tokenize-stage prompt_bin tempfile, on + # the other hand, was never used (we hit before _maybe_compress) and + # would otherwise leak. + try: prompt_bin.unlink() + except Exception: pass + text = tokenizer.decode(tokens, skip_special_tokens=True) + text, matched_stop = trim_at_stop(text, user_stops) + stop_reason = "stop_sequence" if matched_stop else "end_turn" + return JSONResponse({ + "id": msg_id, + "type": "message", + "role": "assistant", + "model": req.model or MODEL_NAME, + "content": [{"type": "text", "text": text}], + "stop_reason": stop_reason, + "stop_sequence": matched_stop, + "usage": {"input_tokens": prompt_len, + "output_tokens": len(tokens)}, + }) + + return app + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--host", default="0.0.0.0") + ap.add_argument("--port", type=int, default=8000) + ap.add_argument("--target", type=Path, default=DEFAULT_TARGET) + ap.add_argument("--draft", type=Path, default=DEFAULT_DRAFT_ROOT) + ap.add_argument("--bin", type=Path, default=DEFAULT_BIN) + ap.add_argument("--budget", type=int, default=DEFAULT_BUDGET) + # Attention compute currently scales with --max-ctx, not the actual + # prompt+gen length (see issue #10). Default 16384 fits most API + # workloads without the 20×+ slowdown users hit with --max-ctx=131072 + # on short requests. Bump via --max-ctx for long-context serving. + ap.add_argument("--max-ctx", type=int, default=16384, + help="Maximum context length (default: 16384; oversizing " + "this, e.g. 131072 on short prompts, can slow " + "attention 20×+ until issue #10 is fixed)") + ap.add_argument("--kv-f16", action="store_true", + help="Force F16 KV cache. When --max-ctx > 6144 the server " + "auto-enables TQ3_0 KV to fit; pass --kv-f16 to opt out.") + ap.add_argument("--cache-type-k", "--ctk", dest="cache_type_k", default=None, + choices=["f16","bf16","q4_0","q4_1","q5_0","q5_1","q8_0","tq3_0"], + help="K cache element type (overrides --kv-q4/--kv-tq3/--kv-f16 for K). " + "See kv_quant.cpp for supported (K,V) pairs.") + ap.add_argument("--cache-type-v", "--ctv", dest="cache_type_v", default=None, + choices=["f16","bf16","q4_0","q4_1","q5_0","q5_1","q8_0","tq3_0"], + help="V cache element type (overrides --kv-q4/--kv-tq3/--kv-f16 for V).") + ap.add_argument("--fa-window", type=int, default=None, + help="Sliding window for FA layers (KV positions). 0 = full " + "attention. Default 2048 (set in C++); only kicks in " + "once kv_cache > window. Trades attention range for " + "long-context decode speed.") + ap.add_argument("--tokenizer", type=str, default=None, + help="HF tokenizer id; default inferred from target GGUF.") + ap.add_argument("--target-gpu", type=int, default=None, + help="Visible CUDA device id for test_dflash (sets DFLASH_TARGET_GPU)") + ap.add_argument("--draft-gpu", type=int, default=None, + help="Visible CUDA device id for draft (sets DFLASH_DRAFT_GPU)") + ap.add_argument("--target-gpus", type=str, default=None, + help="Comma-separated target GPU ids for target-layer sharding (passes --target-gpus)") + ap.add_argument("--target-layer-split", nargs="?", const="", default=None, + metavar="WEIGHTS", + help="Optional comma-separated layer split weights for --target-gpus " + "(omit WEIGHTS after the flag to use defaults)") + ap.add_argument("--draft-feature-mirror", action="store_true", + help="Pass --draft-feature-mirror to test_dflash (safe cross-GPU feature path)") + ap.add_argument("--peer-access", action="store_true", + help="Pass --peer-access to test_dflash (prefer P2P memcpy when available)") + ap.add_argument("--daemon", action="store_true", + help="No-op: accepted for parity with server.py / Compose; " + "this process always runs test_dflash with --daemon.") + add_cli_flags(ap) + add_tool_split_flags(ap) + ap.add_argument("--prefix-cache-slots", type=int, default=4, + help="Number of prefix-cache snapshot slots (0 to disable)") + ap.add_argument("--prefill-cache-slots", type=int, default=4, + help="Number of full-compress-result cache slots (Option 3). " + "Only active when --prefill-compression is enabled. " + "prefix-cache-slots + prefill-cache-slots must not exceed 8.") + args = ap.parse_args() + prefill_cfg = config_from_args(args) + + tool_split_cfg = tool_split_config_from_args(args) + if getattr(args, "tool_split", None) == "auto": + tool_split_cfg = ToolSplitConfig( + enabled=True, + profile=tool_split_cfg.profile, + plugin_dir=tool_split_cfg.plugin_dir, + pinned_tool_slots=tool_split_cfg.pinned_tool_slots, + compress_conversation=tool_split_cfg.compress_conversation, + ) + elif getattr(args, "tool_split", None) == "on": + tool_split_cfg = ToolSplitConfig( + enabled=True, + profile=tool_split_cfg.profile or "auto", + plugin_dir=tool_split_cfg.plugin_dir, + pinned_tool_slots=tool_split_cfg.pinned_tool_slots, + compress_conversation=tool_split_cfg.compress_conversation, + ) + elif getattr(args, "tool_split", None) == "off": + tool_split_cfg = ToolSplitConfig( + enabled=False, + profile=tool_split_cfg.profile, + plugin_dir=tool_split_cfg.plugin_dir, + pinned_tool_slots=tool_split_cfg.pinned_tool_slots, + compress_conversation=tool_split_cfg.compress_conversation, + ) + + # Auto-enable TQ3_0 KV cache when the requested context exceeds what F16 fits. + # setdefault so an explicit user DFLASH27B_KV_TQ3=0 still wins. + if args.cache_type_k: + os.environ["DFLASH27B_KV_K"] = args.cache_type_k + if args.cache_type_v: + os.environ["DFLASH27B_KV_V"] = args.cache_type_v + if args.max_ctx > 6144 and not args.kv_f16 and not args.cache_type_k and not args.cache_type_v: + os.environ.setdefault("DFLASH27B_KV_TQ3", "1") + + if args.fa_window is not None: + os.environ["DFLASH27B_FA_WINDOW"] = str(args.fa_window) + + if args.target_gpu is not None: + os.environ["DFLASH_TARGET_GPU"] = str(args.target_gpu) + if args.draft_gpu is not None: + os.environ["DFLASH_DRAFT_GPU"] = str(args.draft_gpu) + + # When pflash is on, daemon needs the same env the bench harness uses + # (otherwise post-compress draft graph reserve OOMs at 64K+). + if args.prefill_compression != "off": + os.environ.setdefault("DFLASH27B_LM_HEAD_FIX", "0") + os.environ.setdefault("DFLASH27B_FA_WINDOW", "0") + # FlashPrefill bench-tuned defaults from pflash/README.md headline numbers + # (10x TTFT @ 64K). Without these the drafter falls through to the WMMA + # fallback at the default ALPHA=0.12, which roughly triples cold-start + # TTFT. setdefault so explicit user overrides still win. + os.environ.setdefault("DFLASH_FP_USE_BSA", "1") + os.environ.setdefault("DFLASH_FP_ALPHA", "0.85") + if prefill_cfg.skip_park: + os.environ["DFLASH_COMPRESS_NO_PARK"] = "1" + + if not args.target.is_file(): + raise SystemExit(f"target GGUF not found at {args.target}") + + arch = _arch_from_gguf(args.target) + + if not args.bin.is_file(): + raise SystemExit(f"binary not found at {args.bin} (arch={arch})") + + if arch in _LAGUNA_ARCHES: + draft = None + else: + draft = resolve_draft(args.draft) if args.draft.is_dir() else args.draft + if not draft.is_file(): + raise SystemExit(f"draft safetensors not found at {args.draft}") + + tokenizer_id = args.tokenizer or _tokenizer_id_from_gguf(args.target) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True) + stop_ids = set() + for s in ("<|im_end|>", "<|endoftext|>"): + ids = tokenizer.encode(s, add_special_tokens=False) + if ids: stop_ids.add(ids[0]) + + drafter_tokenizer = None + if prefill_cfg.enabled: + drafter_tokenizer = AutoTokenizer.from_pretrained( + prefill_cfg.drafter_tokenizer_id, trust_remote_code=True) + + tool_split_orchestrator = None + if tool_split_cfg.enabled: + adapter = resolve_tool_split_adapter( + tool_split_cfg, arch=arch, tokenizer_id=tokenizer_id) + if adapter is not None: + # Thick prefix snapshots live in system RAM (CPU snapshot backend), + # so multiple conversation slots coexist with thin tool pins even + # on 24GB cards. Agent stacks (Hermes) interleave 2+ prompt + # families per turn — a single thick slot would thrash on every + # request. No VRAM clamp needed anymore. + if tool_split_cfg.pinned_tool_slots > 0: + total_slots = ( + args.prefix_cache_slots + + args.prefill_cache_slots + + tool_split_cfg.pinned_tool_slots + ) + if total_slots > PrefixCache.DAEMON_MAX_SLOTS: + overflow = total_slots - PrefixCache.DAEMON_MAX_SLOTS + if args.prefill_cache_slots > 0: + cut = min(overflow, args.prefill_cache_slots) + args.prefill_cache_slots -= cut + overflow -= cut + if overflow > 0 and args.prefix_cache_slots > 0: + cut = min(overflow, args.prefix_cache_slots) + args.prefix_cache_slots -= cut + overflow -= cut + if overflow > 0: + tool_split_cfg = ToolSplitConfig( + enabled=tool_split_cfg.enabled, + profile=tool_split_cfg.profile, + plugin_dir=tool_split_cfg.plugin_dir, + pinned_tool_slots=max( + 0, + tool_split_cfg.pinned_tool_slots - overflow, + ), + compress_conversation=tool_split_cfg.compress_conversation, + ) + print( + f" [cfg] tool-split slot budget: prefix={args.prefix_cache_slots} " + f"prefill={args.prefill_cache_slots} " + f"tool_pins={tool_split_cfg.pinned_tool_slots} " + f"(daemon max {PrefixCache.DAEMON_MAX_SLOTS})", + flush=True, + ) + from tool_split.orchestrator import ToolSlotCache + tool_split_orchestrator = ToolSplitOrchestrator( + adapter=adapter, config=tool_split_cfg) + tool_split_orchestrator.tool_slots = ToolSlotCache( + pinned_slots=tool_split_cfg.pinned_tool_slots, + slot_base=args.prefix_cache_slots + args.prefill_cache_slots, + ) + + extra_daemon: list[str] = [] + if args.draft_feature_mirror: + extra_daemon.append("--draft-feature-mirror") + if args.peer_access: + extra_daemon.append("--peer-access") + if args.target_gpus: + extra_daemon.append(f"--target-gpus={args.target_gpus}") + if args.target_layer_split: + extra_daemon.append(f"--target-layer-split={args.target_layer_split}") + extra_daemon.append("--target-split-load-draft") + extra_daemon.append("--target-split-dflash") + if args.prefix_cache_slots > 0 or args.prefill_cache_slots > 0: + print(" [cfg] target-gpus daemon mode disables prefix/full cache slots (snapshot protocol unsupported)") + args.prefix_cache_slots = 0 + args.prefill_cache_slots = 0 + + app = build_app(args.target, draft, args.bin, args.budget, args.max_ctx, + tokenizer, stop_ids, + prefill_cfg=prefill_cfg if prefill_cfg.enabled else None, + drafter_tokenizer=drafter_tokenizer, + prefix_cache_slots=args.prefix_cache_slots, + prefill_cache_slots=args.prefill_cache_slots, + arch=arch, + extra_daemon_args=extra_daemon or None, + tool_split_cfg=tool_split_cfg, + tool_split=tool_split_orchestrator) + + import uvicorn + print(f"Luce DFlash OpenAI server (tool-aware) on http://{args.host}:{args.port}") + print(f" arch = {arch}") + print(f" target = {args.target}") + print(f" draft = {draft}") + print(f" bin = {args.bin}") + print(f" budget = {args.budget}") + print(f" max_ctx = {args.max_ctx}") + print(f" tokenizer = {tokenizer_id}") + if prefill_cfg.enabled: + print(f" pflash = {prefill_cfg.mode} · threshold={prefill_cfg.threshold} " + f"keep={prefill_cfg.keep_ratio} drafter={prefill_cfg.drafter_gguf}") + else: + print(" pflash = off") + if tool_split_orchestrator is not None: + print(f" tool-split = on · profile={tool_split_orchestrator.profile} " + f"pinned_slots={tool_split_cfg.pinned_tool_slots} " + f"compress_conv={tool_split_cfg.compress_conversation}") + if tool_split_cfg.plugin_dir: + print(f" tool-split plugins = {tool_split_cfg.plugin_dir}") + elif tool_split_cfg.enabled: + print(" tool-split = requested but no adapter matched (disabled)") + else: + print(" tool-split = off") + uvicorn.run(app, host=args.host, port=args.port, log_level="info") + + +if __name__ == "__main__": + main() diff --git a/server/scripts/test_prefix_cache_slot_depth.py b/server/scripts/test_prefix_cache_slot_depth.py new file mode 100644 index 000000000..dea55c8d0 --- /dev/null +++ b/server/scripts/test_prefix_cache_slot_depth.py @@ -0,0 +1,118 @@ +"""Unit tests: prefix cache slot depth vs lookup cut (cross-session safety).""" +import unittest +from unittest.mock import MagicMock, patch + +from prefix_cache import PrefixCache, hash_prefix + + +class _FakeTokenizer: + """Minimal Qwen-style markers so PrefixCache enables itself.""" + + _MARKERS = { + "<|im_end|>": [1], + "<|im_start|>": [2], + "system": [3], + } + + def encode(self, text, add_special_tokens=False): + return list(self._MARKERS.get(text, [ord(c) for c in text])) + + +class PrefixCacheSlotDepthTests(unittest.TestCase): + def _make_cache(self) -> PrefixCache: + pc = PrefixCache( + daemon_stdin=MagicMock(), + await_reply=MagicMock(), + daemon_lock=MagicMock(), + tokenizer=_FakeTokenizer(), + kv_k_type="f16", + fa_window=0, + cap=4, + ) + self.assertFalse(pc.disabled) + return pc + + @patch("prefix_cache.find_all_boundaries_markers") + def test_lookup_rejects_stale_shallow_hash(self, mock_bounds): + pc = self._make_cache() + ids = list(range(400)) + mock_bounds.return_value = [50, 376] + + shallow_key = hash_prefix(ids[:376], pc.kv_k_type, pc.fa_window) + pc.entries[shallow_key] = 0 + pc._populated_slots.add(0) + pc._slot_prefix_len[0] = 2411 + + self.assertIsNone(pc.lookup(ids)) + self.assertNotIn(shallow_key, pc.entries) + + @patch("prefix_cache.find_all_boundaries_markers") + def test_lookup_accepts_matching_depth(self, mock_bounds): + pc = self._make_cache() + ids = list(range(400)) + mock_bounds.return_value = [50, 376] + + key = hash_prefix(ids[:376], pc.kv_k_type, pc.fa_window) + pc.entries[key] = 0 + pc._populated_slots.add(0) + pc._slot_prefix_len[0] = 376 + + self.assertEqual(pc.lookup(ids), (0, 376)) + + @patch("prefix_cache.find_all_boundaries_markers") + def test_confirm_evicts_stale_keys_for_slot(self, mock_bounds): + pc = self._make_cache() + ids = list(range(500)) + mock_bounds.return_value = [100, 376, 480] + + shallow = hash_prefix(ids[:376], pc.kv_k_type, pc.fa_window) + deep = hash_prefix(ids[:480], pc.kv_k_type, pc.fa_window) + pc.entries[shallow] = 0 + pc._populated_slots.add(0) + pc._slot_prefix_len[0] = 376 + + pc.confirm_inline_snap(0, 480, ids) + + self.assertNotIn(shallow, pc.entries) + self.assertIn(deep, pc.entries) + self.assertEqual(pc.entries[deep], 0) + self.assertEqual(pc._slot_prefix_len[0], 480) + + @patch("prefix_cache.find_all_boundaries_markers") + def test_abort_inline_snap_purges_reuse_slot_mappings(self, mock_bounds): + pc = self._make_cache() + ids = list(range(400)) + mock_bounds.return_value = [50, 376] + + key = hash_prefix(ids[:376], pc.kv_k_type, pc.fa_window) + pc.entries[key] = 0 + pc._populated_slots.add(0) + pc._slot_prefix_len[0] = 376 + + prep = pc.prepare_inline_snap(ids, reuse_slot=0) + self.assertIsNotNone(prep) + slot, _ = prep + pc.abort_inline_snap(slot) + + self.assertNotIn(key, pc.entries) + self.assertNotIn(0, pc._slot_prefix_len) + + def test_abort_full_snap_purges_stale_lookup(self): + pc = self._make_cache() + pc.init_full_cache(1) + old_ids = list(range(200)) + old_key = hash_prefix(old_ids, pc.kv_k_type, pc.fa_window) + pc.full_entries[old_key] = (pc._full_slot_base, "/tmp/old.bin", 100) + + new_ids = list(range(300)) + prep = pc.prepare_full_snap(new_ids) + self.assertIsNotNone(prep) + slot, _ = prep + self.assertEqual(pc._full_pending_evict_key, old_key) + pc.abort_full_snap(slot) + + self.assertIsNone(pc.lookup_full(old_ids)) + + +if __name__ == "__main__": + unittest.main() diff --git a/server/scripts/test_tool_split_adapters.py b/server/scripts/test_tool_split_adapters.py new file mode 100644 index 000000000..ad9206d83 --- /dev/null +++ b/server/scripts/test_tool_split_adapters.py @@ -0,0 +1,289 @@ +"""Unit tests for tool_split registry and adapters (no GPU).""" +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +SCRIPTS = Path(__file__).resolve().parent +if str(SCRIPTS) not in sys.path: + sys.path.insert(0, str(SCRIPTS)) + +from tool_split.base import PromptSplit, ToolSplitAdapter # noqa: E402 +from tool_split.config import ToolSplitConfig, config_from_env_and_args # noqa: E402 +from tool_split.registry import ( # noqa: E402 + builtin_adapter_names, + get_adapter_class, + register_adapter, + resolve_adapter, +) +from tool_split.qwen3 import Qwen3ToolSplitAdapter # noqa: E402 + + +class _FakeTokenizer: + """Qwen-shaped tokenizer stub for split tests. + + Renders ``<|im_start|>role ... <|im_end|>`` framing (like the real Qwen3 + chat template) and encodes the special markers + role words as single + tokens so ``_tool_prefix_boundary`` can find the first user turn. + """ + + _SPECIALS = ["<|im_start|>", "<|im_end|>", "system", "user", "assistant"] + + def apply_chat_template(self, messages, **kwargs): + parts = [] + first = True + for m in messages: + body = m.get("content", "") or "" + if first and kwargs.get("tools"): + names = ",".join(t["function"]["name"] for t in kwargs["tools"]) + body += f"{names}" + parts.append(f"<|im_start|>{m['role']}\n{body}<|im_end|>\n") + first = False + if kwargs.get("add_generation_prompt", True): + parts.append("<|im_start|>assistant\n") + return "".join(parts) + + def encode(self, text, add_special_tokens=False): + ids = [] + i = 0 + while i < len(text): + for si, s in enumerate(self._SPECIALS): + if text.startswith(s, i): + ids.append(1000 + si) + i += len(s) + break + else: + ids.append(ord(text[i])) + i += 1 + return ids + + +class TestRegistry(unittest.TestCase): + def test_builtins_registered(self): + names = builtin_adapter_names() + self.assertIn("qwen3", names) + self.assertIn("laguna", names) + + def test_get_adapter_class(self): + cls = get_adapter_class("qwen3") + self.assertIs(cls, Qwen3ToolSplitAdapter) + + def test_user_plugin_register(self): + @register_adapter("test_vendor") + class _TestVendor(ToolSplitAdapter): + @classmethod + def detect(cls, *, arch, tokenizer_id): + return False + + def split_prompt(self, *a, **k): + raise NotImplementedError + + self.assertIn("test_vendor", builtin_adapter_names()) + + +class TestQwen3Adapter(unittest.TestCase): + def setUp(self): + self.adapter = Qwen3ToolSplitAdapter() + self.tok = _FakeTokenizer() + + def test_canonical_tools_sorted(self): + tools = [ + {"type": "function", "function": {"name": "z", "description": "", "parameters": {}}}, + {"type": "function", "function": {"name": "a", "description": "", "parameters": {}}}, + ] + canon = self.adapter.canonical_tools(tools) + self.assertEqual([c["function"]["name"] for c in canon], ["a", "z"]) + + def test_fingerprint_stable(self): + tools = [{"type": "function", "function": {"name": "x", "description": "d", "parameters": {}}}] + self.assertEqual( + self.adapter.tools_fingerprint(tools), + self.adapter.tools_fingerprint(tools), + ) + + def test_split_with_tools(self): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + tools = [{"type": "function", "function": {"name": "read_file", "description": "", "parameters": {}}}] + split = self.adapter.split_prompt(self.tok, messages, tools) + self.assertIsInstance(split, PromptSplit) + self.assertGreater(split.tool_prefix_len, 0) + self.assertEqual(split.full_ids, split.tool_prefix_ids + split.conversation_ids) + + def test_split_no_tools(self): + messages = [{"role": "user", "content": "hi"}] + split = self.adapter.split_prompt(self.tok, messages, None) + self.assertEqual(split.tool_prefix_len, 0) + self.assertEqual(split.conversation_ids, split.full_ids) + + def test_tool_boundary_fallback_matches_marker_path(self): + from unittest.mock import patch + + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + tools = [{"type": "function", "function": {"name": "read_file", "description": "", "parameters": {}}}] + with patch("prefix_cache.find_prefix_boundary_markers", return_value=-1): + split = self.adapter.split_prompt(self.tok, messages, tools) + im_start = self.tok.encode("<|im_start|>", add_special_tokens=False)[0] + user_tok = self.tok.encode("user", add_special_tokens=False)[0] + seq = (im_start, user_tok) + idx = next( + i for i in range(len(split.full_ids) - 1) + if tuple(split.full_ids[i:i + 2]) == seq + ) + self.assertEqual(split.tool_prefix_len, idx + 1) + + +class TestConfig(unittest.TestCase): + def test_invalid_env_pinned_falls_back_with_cli_override(self): + import os + from unittest.mock import patch + + class _Args: + tool_split_pinned_slots = 5 + + with patch.dict(os.environ, {"DFLASH_TOOL_SPLIT_PINNED_SLOTS": "not-a-number"}): + cfg = config_from_env_and_args(_Args()) + self.assertEqual(cfg.pinned_tool_slots, 5) + + def test_cli_plugin_dir_expands_user(self): + import os + from unittest.mock import patch + + class _Args: + tool_split_plugin_dir = Path("~/tool_plugins") + + with patch.dict(os.environ, {}, clear=False): + cfg = config_from_env_and_args(_Args()) + self.assertTrue(str(cfg.plugin_dir).endswith("tool_plugins")) + self.assertNotIn("~", str(cfg.plugin_dir)) + + def test_resolve_auto_qwen(self): + cfg = ToolSplitConfig(enabled=True, profile="auto") + adapter = resolve_adapter(cfg, arch="qwen36", tokenizer_id="Qwen/Qwen3.6-27B") + self.assertIsNotNone(adapter) + self.assertEqual(adapter.profile_name, "qwen3") + + +class TestToolSlotCache(unittest.TestCase): + def test_reserve_reuses_evicted_slot_within_range(self): + from tool_split.orchestrator import ToolSlotCache + + cache = ToolSlotCache(pinned_slots=2, slot_base=3) + s0 = cache.reserve("a") + cache.confirm("a", s0) + s1 = cache.reserve("b") + cache.confirm("b", s1) + self.assertEqual({s0, s1}, {3, 4}) + # Reserve "c" — defers eviction of "a" until confirm. + s2 = cache.reserve("c") + self.assertEqual(s2, s0) + self.assertEqual(cache.lookup("a"), s0) + self.assertIsNone(cache.lookup("c")) + cache.confirm("c", s2) + self.assertIsNone(cache.lookup("a")) + self.assertEqual(cache.lookup("c"), s0) + self.assertEqual(cache.lookup("b"), s1) + self.assertTrue(all(3 <= s < 5 for s in cache._confirmed.values())) + + def test_release_reservation_frees_slot(self): + from tool_split.orchestrator import ToolSlotCache + + cache = ToolSlotCache(pinned_slots=2, slot_base=6) + slot = cache.reserve("fp") + cache.release_reservation("fp", slot) + self.assertIsNone(cache.lookup("fp")) + # Freed slot is reusable. + self.assertEqual(cache.reserve("other"), slot) + + def test_release_reservation_restores_evicted_anchor(self): + from tool_split.orchestrator import ToolSlotCache + + cache = ToolSlotCache(pinned_slots=2, slot_base=3) + s0 = cache.reserve("a") + cache.confirm("a", s0) + s1 = cache.reserve("b") + cache.confirm("b", s1) + s2 = cache.reserve("c") + self.assertEqual(cache.lookup("a"), s0) + cache.release_reservation("c", s2) + self.assertEqual(cache.lookup("a"), s0) + self.assertIsNone(cache.lookup("c")) + + +class TestCommitPendingToolSnap(unittest.IsolatedAsyncioTestCase): + async def test_releases_reservation_on_snapshot_failure(self): + import asyncio + + from tool_split.daemon_bridge import commit_pending_tool_snap + from tool_split.orchestrator import ToolSlotCache + + class _Orch: + def __init__(self): + self.tool_slots = ToolSlotCache(pinned_slots=2, slot_base=3) + + orch = _Orch() + slot = orch.tool_slots.reserve("fp") + + async def _fail_reply(prefix, timeout=30.0): + raise asyncio.TimeoutError("daemon timeout") + + class _Stdin: + def write(self, data): + return None + + def flush(self): + return None + + await commit_pending_tool_snap( + orchestrator=orch, + daemon_stdin=_Stdin(), + await_reply=_fail_reply, + fingerprint="fp", + slot=slot, + kv_end=10, + ) + self.assertIsNone(orch.tool_slots.lookup("fp")) + + +class TestOrchestratorDaemonCmd(unittest.TestCase): + def setUp(self): + from tool_split.config import ToolSplitConfig + from tool_split.orchestrator import ToolSplitOrchestrator + + self.orch = ToolSplitOrchestrator( + adapter=Qwen3ToolSplitAdapter(), + config=ToolSplitConfig(enabled=True, pinned_tool_slots=2), + ) + + def test_restore_chain_tool_only(self): + plan = self.orch.build_plan( + split=None, + tools_fingerprint="abc", + prompt_bin=Path("/tmp/p.bin"), + prompt_len=100, + tool_slot_hit=6, + ) + cmd = self.orch.format_daemon_command(plan, 32) + self.assertIn("RESTORE_CHAIN -1 6 /tmp/p.bin 32", cmd) + + def test_restore_chain_tool_and_conv(self): + plan = self.orch.build_plan( + split=None, + tools_fingerprint="abc", + prompt_bin=Path("/tmp/p.bin"), + prompt_len=100, + tool_slot_hit=6, + conv_hit=(2, 80), + ) + cmd = self.orch.format_daemon_command(plan, 16) + self.assertIn("RESTORE_CHAIN 2 6 /tmp/p.bin 16", cmd) + + +if __name__ == "__main__": + unittest.main() diff --git a/server/scripts/tool_split/README.md b/server/scripts/tool_split/README.md new file mode 100644 index 000000000..ddfe804f2 --- /dev/null +++ b/server/scripts/tool_split/README.md @@ -0,0 +1,82 @@ +# Tool-split cache (split tool KV from conversation for PFlash) + +Pins tool-schema KV separately from the growing conversation so PFlash can +compress agent history without re-prefilling static tool definitions. + +## Enable + +```bash +# CLI (server_tools.py) +python3 scripts/server_tools.py --tool-split auto --prefill-compression auto ... + +# Environment (model-runner-v4 / compose) +DFLASH_TOOL_SPLIT_ENABLED=1 +DFLASH_TOOL_SPLIT_PROFILE=auto # or qwen3, laguna, plugin:my_vendor +``` + +## Built-in profiles + +| Profile | Detects | Notes | +|---------|---------|-------| +| `qwen3` | `qwen35` / `qwen36` arch, Qwen HF tokenizer | System + tools prefix split | +| `laguna` | `laguna` arch, Poolside tokenizer | XML boundary fallback | + +`auto` picks the first adapter whose `detect()` matches the loaded GGUF. + +## User plugins + +1. Copy `example_user_plugin.py` to a directory, e.g. `~/.lucebox/tool_split_plugins/my_llm.py`. +2. Implement `ToolSplitAdapter` and decorate with `@register_adapter("my_llm")`. +3. Start the server with: + +```bash +--tool-split on \ +--tool-split-profile plugin:my_llm \ +--tool-split-plugin-dir ~/.lucebox/tool_split_plugins +``` + +Or via env: + +```bash +export DFLASH_TOOL_SPLIT_ENABLED=1 +export DFLASH_TOOL_SPLIT_PROFILE=plugin:my_llm +export DFLASH_TOOL_SPLIT_PLUGIN_DIR=$HOME/.lucebox/tool_split_plugins +``` + +## Settings + +| Flag / env | Default | Purpose | +|------------|---------|---------| +| `--tool-split` / `DFLASH_TOOL_SPLIT_ENABLED` | off | Master switch (`auto` enables when adapter matches) | +| `--tool-split-profile` / `DFLASH_TOOL_SPLIT_PROFILE` | `auto` | Adapter name | +| `--tool-split-plugin-dir` / `DFLASH_TOOL_SPLIT_PLUGIN_DIR` | — | User `.py` plugins | +| `--tool-split-pinned-slots` / `DFLASH_TOOL_SPLIT_PINNED_SLOTS` | `2` | Tool KV slots (daemon) | +| `--tool-split-compress-conv` / `DFLASH_TOOL_SPLIT_COMPRESS_CONV` | on | PFlash on conversation suffix | + +## Daemon flow (wired) + +1. **First request** with a tool fingerprint: full prefill, then + ``SNAPSHOT_THIN 0 `` commits pinned tool KV. +2. **Later requests** with the same tools: ``RESTORE_CHAIN -1 `` + restores tool KV and prefills only the conversation suffix. +3. **Conversation prefix hit** + tool hit: + ``RESTORE_CHAIN ``. + +Slot layout (8 daemon slots max): ``[prefix LRU][full compress][tool pins]``. + +## Requirements + +- `DFLASH_LAYER_SPLIT=0` (or non-sharded daemon): SNAPSHOT / RESTORE_CHAIN are + unsupported in layer-split mode. +- PFlash enabled (`--prefill-compression auto`) for conversation compression. + +## Architecture + +``` +ToolSplitAdapter.split_prompt() → PromptSplit + tool_prefix_ids → pinned ToolSlotCache (RESTORE_CHAIN thin slot / SNAPSHOT_THIN) + conversation_ids → PFlash + PrefixCache inline snap +``` + +Implement `split_prompt()` for each LLM family; the orchestrator and daemon +command formatting are shared. diff --git a/server/scripts/tool_split/__init__.py b/server/scripts/tool_split/__init__.py new file mode 100644 index 000000000..695edad2b --- /dev/null +++ b/server/scripts/tool_split/__init__.py @@ -0,0 +1,37 @@ +"""Split tool KV from conversation for PFlash + pinned tool cache. + +Built-in profiles: ``qwen3``, ``laguna``. User plugins: drop a ``.py`` file +in ``--tool-split-plugin-dir`` that calls ``register_adapter("name")``. + +See ``example_user_plugin.py`` for a template. +""" +from tool_split.base import PromptSplit, ToolSplitAdapter, ToolSplitPlan +from tool_split.config import ToolSplitConfig, add_cli_flags, config_from_env_and_args +from tool_split.orchestrator import ToolSlotCache, ToolSplitOrchestrator +from tool_split.registry import ( + builtin_adapter_names, + get_adapter_class, + load_plugin_dir, + register_adapter, + resolve_adapter, +) + +# Import built-ins so they self-register. +from tool_split import laguna as _laguna # noqa: F401 +from tool_split import qwen3 as _qwen3 # noqa: F401 + +__all__ = [ + "PromptSplit", + "ToolSplitAdapter", + "ToolSplitConfig", + "ToolSplitOrchestrator", + "ToolSplitPlan", + "ToolSlotCache", + "add_cli_flags", + "builtin_adapter_names", + "config_from_env_and_args", + "get_adapter_class", + "load_plugin_dir", + "register_adapter", + "resolve_adapter", +] diff --git a/server/scripts/tool_split/base.py b/server/scripts/tool_split/base.py new file mode 100644 index 000000000..ca6bd87f6 --- /dev/null +++ b/server/scripts/tool_split/base.py @@ -0,0 +1,124 @@ +"""Base types for LLM-specific tool / conversation prompt splitting.""" +from __future__ import annotations + +import hashlib +import json +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Mapping, Sequence + + +@dataclass(frozen=True) +class PromptSplit: + """Tokenized prompt decomposed for split KV + PFlash. + + ``tool_prefix_ids`` is the stable segment pinned in a tool KV slot + (system header + tool schema tokens). ``conversation_ids`` is the + suffix that may grow each turn and is eligible for PFlash compression. + ``full_ids`` is the complete prompt (tool prefix + conversation). + """ + + full_ids: list[int] + tool_prefix_ids: list[int] + conversation_ids: list[int] + tool_prefix_len: int + profile: str + + def __post_init__(self) -> None: + if self.tool_prefix_len != len(self.tool_prefix_ids): + raise ValueError("tool_prefix_len must match len(tool_prefix_ids)") + if self.full_ids != self.tool_prefix_ids + self.conversation_ids: + raise ValueError("full_ids must equal tool_prefix_ids + conversation_ids") + + @property + def conversation_len(self) -> int: + return len(self.conversation_ids) + + +@dataclass +class ToolSplitPlan: + """Daemon command plan produced by :class:`ToolSplitOrchestrator`.""" + + prompt_bin_path: str + prompt_token_count: int + tool_slot: int | None = None + conv_restore_slot: int | None = None + conv_restore_prefix_len: int = 0 + use_restore_chain: bool = False + thin_slot_ids: list[int] = field(default_factory=list) + inline_snap: tuple[int, int] | None = None # (slot_id, cut_pos) + compression_fired: bool = False + started_in_thinking: bool = False + tools_fingerprint: str | None = None + pending_tool_snap: tuple[int, int] | None = None # (slot, kv_end) post-response + tool_prefix_len: int = 0 + + +@dataclass +class ToolRequestContext: + """Per-request tool-split state for daemon command building.""" + + split: PromptSplit | None = None + fingerprint: str | None = None + tool_slot_hit: int | None = None + pending_tool_snap: tuple[int, int] | None = None + + +class ToolSplitAdapter(ABC): + """Per-LLM adapter: tokenize tools separately from conversation.""" + + #: Registry name (e.g. ``qwen3``, ``laguna``). + profile_name: str + + @classmethod + @abstractmethod + def detect(cls, *, arch: str, tokenizer_id: str) -> bool: + """Return True if this adapter should handle *arch* / tokenizer.""" + + @abstractmethod + def split_prompt( + self, + tokenizer: Any, + messages: Sequence[Mapping[str, Any]], + tools: Sequence[Mapping[str, Any]] | None, + *, + chat_template_kwargs: Mapping[str, Any] | None = None, + enable_thinking: bool = False, + ) -> PromptSplit: + """Split a chat request into tool-prefix and conversation token spans.""" + + def canonical_tools( + self, tools: Sequence[Mapping[str, Any]] | None + ) -> list[dict[str, Any]]: + """Normalize tool defs for stable cache keys (override per model).""" + if not tools: + return [] + out: list[dict[str, Any]] = [] + for t in tools: + if isinstance(t, dict): + fn = t.get("function") or {} + name = fn.get("name", "") + out.append({ + "type": t.get("type", "function"), + "function": { + "name": name, + "description": fn.get("description", ""), + "parameters": fn.get("parameters") or {}, + }, + }) + out.sort(key=lambda x: (x.get("function") or {}).get("name", "")) + return out + + def tools_fingerprint(self, tools: Sequence[Mapping[str, Any]] | None) -> str: + """SHA-256 hex digest of canonical tool definitions.""" + payload = json.dumps( + self.canonical_tools(tools), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def supports_pflash_on_conversation(self) -> bool: + """Whether PFlash may run on ``conversation_ids`` only.""" + return True diff --git a/server/scripts/tool_split/config.py b/server/scripts/tool_split/config.py new file mode 100644 index 000000000..55c38579f --- /dev/null +++ b/server/scripts/tool_split/config.py @@ -0,0 +1,106 @@ +"""Configuration for split tool KV + conversation PFlash.""" +from __future__ import annotations + +import argparse +import os +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class ToolSplitConfig: + """Parsed ``--tool-split-*`` flags / ``DFLASH_TOOL_SPLIT_*`` env vars.""" + + enabled: bool = False + #: ``auto`` | built-in name (``qwen3``, ``laguna``) | ``plugin:name`` + profile: str = "auto" + plugin_dir: Path | None = None + #: Dedicated daemon slots pinned for tool KV (not LRU-evicted with conv). + pinned_tool_slots: int = 2 + #: Allow PFlash on the conversation suffix when tools are present. + compress_conversation: bool = True + + def __post_init__(self) -> None: + if self.pinned_tool_slots < 0: + raise ValueError("pinned_tool_slots must be >= 0") + + +def _env_bool(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in ("1", "true", "yes", "on") + + +def config_from_env_and_args(args=None) -> ToolSplitConfig: + """Merge env vars with optional argparse namespace.""" + enabled = _env_bool("DFLASH_TOOL_SPLIT_ENABLED", False) + profile = os.environ.get("DFLASH_TOOL_SPLIT_PROFILE", "auto") + plugin_dir_raw = os.environ.get("DFLASH_TOOL_SPLIT_PLUGIN_DIR") + plugin_dir = Path(plugin_dir_raw).expanduser() if plugin_dir_raw else None + pinned_raw = os.environ.get("DFLASH_TOOL_SPLIT_PINNED_SLOTS", "2") + try: + pinned = int(pinned_raw) + except ValueError: + pinned = 2 + compress_conv = _env_bool("DFLASH_TOOL_SPLIT_COMPRESS_CONV", True) + + if args is not None: + if getattr(args, "tool_split", None) is not None: + enabled = args.tool_split != "off" + if getattr(args, "tool_split_profile", None): + profile = args.tool_split_profile + if getattr(args, "tool_split_plugin_dir", None): + plugin_dir = Path(args.tool_split_plugin_dir).expanduser() + if getattr(args, "tool_split_pinned_slots", None) is not None: + pinned = args.tool_split_pinned_slots + if getattr(args, "tool_split_compress_conv", None) is not None: + compress_conv = args.tool_split_compress_conv + + return ToolSplitConfig( + enabled=enabled, + profile=profile, + plugin_dir=plugin_dir, + pinned_tool_slots=pinned, + compress_conversation=compress_conv, + ) + + +def add_cli_flags(ap) -> None: + """Attach ``--tool-split-*`` flags to an argparse parser.""" + ap.add_argument( + "--tool-split", + choices=["off", "on", "auto"], + default=None, + help="Split tool KV from conversation for PFlash + pinned tool cache. " + "'auto' enables when a matching adapter is found. " + "Env: DFLASH_TOOL_SPLIT_ENABLED=1.", + ) + ap.add_argument( + "--tool-split-profile", + default=None, + metavar="NAME", + help="Adapter profile: auto (default), qwen3, laguna, or plugin:name. " + "Env: DFLASH_TOOL_SPLIT_PROFILE.", + ) + ap.add_argument( + "--tool-split-plugin-dir", + type=Path, + default=None, + help="Directory of user .py plugins calling register_adapter(). " + "Env: DFLASH_TOOL_SPLIT_PLUGIN_DIR.", + ) + ap.add_argument( + "--tool-split-pinned-slots", + type=int, + default=None, + help="Daemon slots reserved for tool KV snapshots (default 2). " + "Env: DFLASH_TOOL_SPLIT_PINNED_SLOTS.", + ) + ap.add_argument( + "--tool-split-compress-conv", + action=argparse.BooleanOptionalAction, + default=None, + help="PFlash-compress conversation suffix when tools are present " + "(default on). Env: DFLASH_TOOL_SPLIT_COMPRESS_CONV.", + ) diff --git a/server/scripts/tool_split/daemon_bridge.py b/server/scripts/tool_split/daemon_bridge.py new file mode 100644 index 000000000..600539f6b --- /dev/null +++ b/server/scripts/tool_split/daemon_bridge.py @@ -0,0 +1,63 @@ +"""Daemon IPC for tool KV thin snapshots.""" +from __future__ import annotations + +import asyncio + + +async def snapshot_thin( + *, + daemon_stdin, + await_reply, + slot: int, + kv_start: int, + kv_end: int, +) -> bool: + """Capture KV range ``[kv_start, kv_end)`` into a thin snapshot slot.""" + line = f"SNAPSHOT_THIN {slot} {kv_start} {kv_end}\n" + daemon_stdin.write(line.encode("utf-8")) + daemon_stdin.flush() + try: + reply = await await_reply("[snap] thin slot=", timeout=30.0) + except (asyncio.TimeoutError, EOFError) as exc: + print(f"[tool-split] SNAPSHOT_THIN slot={slot} failed: {exc}", flush=True) + return False + ok = reply.startswith(f"[snap] thin slot={slot} kv={kv_start},{kv_end}") + if not ok: + print(f"[tool-split] SNAPSHOT_THIN unexpected reply: {reply!r}", flush=True) + return ok + + +async def commit_pending_tool_snap( + *, + orchestrator, + daemon_stdin, + await_reply, + fingerprint: str, + slot: int, + kv_end: int, +) -> None: + if kv_end <= 0: + orchestrator.tool_slots.release_reservation(fingerprint, slot) + return + ok = await snapshot_thin( + daemon_stdin=daemon_stdin, + await_reply=await_reply, + slot=slot, + kv_start=0, + kv_end=kv_end, + ) + if ok: + orchestrator.tool_slots.confirm(fingerprint, slot) + print( + f"[tool-split] tool KV pinned slot={slot} len={kv_end} " + f"fp={fingerprint[:12]}…", + flush=True, + ) + else: + # Reservation must not look like a populated hit on the next request. + orchestrator.tool_slots.release_reservation(fingerprint, slot) + print( + f"[tool-split] SNAPSHOT_THIN failed; released reservation " + f"slot={slot} fp={fingerprint[:12]}…", + flush=True, + ) diff --git a/server/scripts/tool_split/example_user_plugin.py b/server/scripts/tool_split/example_user_plugin.py new file mode 100644 index 000000000..6892388b4 --- /dev/null +++ b/server/scripts/tool_split/example_user_plugin.py @@ -0,0 +1,55 @@ +"""Example user plugin — copy to your ``--tool-split-plugin-dir`` and edit. + +Usage: + mkdir -p ~/.lucebox/tool_split_plugins + cp example_user_plugin.py ~/.lucebox/tool_split_plugins/my_vendor.py + # server: --tool-split on --tool-split-profile plugin:my_vendor \\ + # --tool-split-plugin-dir ~/.lucebox/tool_split_plugins +""" +from __future__ import annotations + +from typing import Any, Mapping, Sequence + +from tool_split.base import PromptSplit, ToolSplitAdapter +from tool_split.registry import register_adapter + + +@register_adapter("my_vendor") +class MyVendorToolSplitAdapter(ToolSplitAdapter): + """Template: replace detection + split logic for your model family.""" + + @classmethod + def detect(cls, *, arch: str, tokenizer_id: str) -> bool: + return "my_vendor" in (arch or "").lower() + + def split_prompt( + self, + tokenizer: Any, + messages: Sequence[Mapping[str, Any]], + tools: Sequence[Mapping[str, Any]] | None, + *, + chat_template_kwargs: Mapping[str, Any] | None = None, + enable_thinking: bool = False, + ) -> PromptSplit: + kwargs: dict[str, Any] = { + "tokenize": False, + "add_generation_prompt": True, + "enable_thinking": enable_thinking, + } + if tools: + kwargs["tools"] = list(tools) + if chat_template_kwargs: + for k, v in chat_template_kwargs.items(): + if k in ("enable_thinking", "add_generation_prompt", "tools"): + kwargs[k] = v + prompt = tokenizer.apply_chat_template(list(messages), **kwargs) + full_ids = tokenizer.encode(prompt, add_special_tokens=False) + # TODO: locate tool-schema token boundary for your template. + split_at = len(full_ids) // 2 # placeholder — do not ship as-is + return PromptSplit( + full_ids=full_ids, + tool_prefix_ids=full_ids[:split_at], + conversation_ids=full_ids[split_at:], + tool_prefix_len=split_at, + profile=self.profile_name, + ) diff --git a/server/scripts/tool_split/laguna.py b/server/scripts/tool_split/laguna.py new file mode 100644 index 000000000..728d48bbd --- /dev/null +++ b/server/scripts/tool_split/laguna.py @@ -0,0 +1,78 @@ +"""Laguna / Poolside XML chat-template splitter (stub + boundary fallback).""" +from __future__ import annotations + +from typing import Any, Mapping, Sequence + +from tool_split.base import PromptSplit, ToolSplitAdapter +from tool_split.registry import register_adapter + +_LAGUNA_ARCHES = frozenset({"laguna"}) +_LAGUNA_TOKENIZER_HINTS = ("laguna", "poolside") + + +@register_adapter("laguna") +class LagunaToolSplitAdapter(ToolSplitAdapter): + """Laguna ```` / ```` XML templates.""" + + @classmethod + def detect(cls, *, arch: str, tokenizer_id: str) -> bool: + a = (arch or "").lower() + t = (tokenizer_id or "").lower() + if a in _LAGUNA_ARCHES: + return True + return any(h in t for h in _LAGUNA_TOKENIZER_HINTS) + + def split_prompt( + self, + tokenizer: Any, + messages: Sequence[Mapping[str, Any]], + tools: Sequence[Mapping[str, Any]] | None, + *, + chat_template_kwargs: Mapping[str, Any] | None = None, + enable_thinking: bool = False, + ) -> PromptSplit: + kwargs: dict[str, Any] = { + "tokenize": False, + "add_generation_prompt": True, + } + if tools: + kwargs["tools"] = list(tools) + if chat_template_kwargs: + kwargs.update(dict(chat_template_kwargs)) + + prompt = tokenizer.apply_chat_template(list(messages), **kwargs) + full_ids = tokenizer.encode(prompt, add_special_tokens=False) + + if not tools: + return PromptSplit( + full_ids=full_ids, + tool_prefix_ids=[], + conversation_ids=full_ids, + tool_prefix_len=0, + profile=self.profile_name, + ) + + from prefix_cache import _resolve_chat_markers, find_prefix_boundary_markers + + markers = _resolve_chat_markers(tokenizer) + boundary = find_prefix_boundary_markers(full_ids, markers) + if boundary < 0: + return PromptSplit( + full_ids=full_ids, + tool_prefix_ids=[], + conversation_ids=full_ids, + tool_prefix_len=0, + profile=self.profile_name, + ) + + return PromptSplit( + full_ids=full_ids, + tool_prefix_ids=full_ids[:boundary], + conversation_ids=full_ids[boundary:], + tool_prefix_len=boundary, + profile=self.profile_name, + ) + + def supports_pflash_on_conversation(self) -> bool: + # Laguna tool format differs; enable when validated on hardware. + return False diff --git a/server/scripts/tool_split/orchestrator.py b/server/scripts/tool_split/orchestrator.py new file mode 100644 index 000000000..a66ec0c54 --- /dev/null +++ b/server/scripts/tool_split/orchestrator.py @@ -0,0 +1,292 @@ +"""Orchestrate split tool KV cache + conversation PFlash.""" +from __future__ import annotations + +import json +import os +import struct +import tempfile +from collections import OrderedDict +from pathlib import Path +from typing import Any, Mapping, Sequence + +from tool_split.base import PromptSplit, ToolSplitAdapter, ToolSplitPlan, ToolRequestContext +from tool_split.config import ToolSplitConfig + + +def _ids_to_bin(ids: list[int]) -> Path: + fd, path = tempfile.mkstemp(suffix=".bin") + with os.fdopen(fd, "wb") as f: + for t in ids: + f.write(struct.pack(" Any: + # Qwen chat templates iterate arguments with `| items`, which requires a + # mapping. OpenAI-shaped requests carry arguments as a JSON string. + if isinstance(args, str): + try: + return json.loads(args) + except (json.JSONDecodeError, ValueError): + return {"_raw": args} + return args + + +def _messages_for_adapter(messages: Sequence[Any]) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for m in messages: + if isinstance(m, dict): + d = dict(m) + if d.get("tool_calls"): + d["tool_calls"] = [ + { + **tc, + "function": { + **(tc.get("function") or {}), + "arguments": _tool_call_args_to_obj( + (tc.get("function") or {}).get("arguments")), + }, + } + for tc in d["tool_calls"] + ] + out.append(d) + continue + d: dict[str, Any] = {"role": m.role} + if m.content is not None: + d["content"] = m.content + if getattr(m, "name", None) is not None: + d["name"] = m.name + if getattr(m, "tool_call_id", None) is not None: + d["tool_call_id"] = m.tool_call_id + if getattr(m, "tool_calls", None): + d["tool_calls"] = [ + { + "id": tc.id, + "type": tc.type, + "function": { + "name": tc.function.name, + "arguments": _tool_call_args_to_obj(tc.function.arguments), + }, + } + for tc in m.tool_calls + ] + out.append(d) + return out + + +def _tools_for_adapter(tools: Sequence[Any] | None) -> list[dict[str, Any]] | None: + if not tools: + return None + return [t.model_dump() if hasattr(t, "model_dump") else dict(t) for t in tools] + + +class ToolSlotCache: + """LRU index: tools fingerprint → daemon thin-snapshot slot (pinned range). + + Slot IDs stay in ``[slot_base, slot_base + pinned_slots)``. Eviction reuses + the LRU entry's slot so RESTORE_CHAIN never references out-of-range IDs. + + Eviction is deferred from ``reserve()`` to ``confirm()`` so a failed + ``SNAPSHOT_THIN`` can call ``release_reservation()`` without losing the + previous anchor (same pattern as ``PrefixCache.prepare_inline_snap``). + """ + + def __init__(self, *, pinned_slots: int, slot_base: int = 0): + self.pinned_slots = max(0, pinned_slots) + self.slot_base = slot_base + self._confirmed: OrderedDict[str, int] = OrderedDict() + self._pending: dict[str, int] = {} + self._pending_evict: tuple[str, int] | None = None + + def lookup(self, fingerprint: str) -> int | None: + if fingerprint in self._pending: + return None + if fingerprint in self._confirmed: + self._confirmed.move_to_end(fingerprint) + return self._confirmed[fingerprint] + return None + + def reserve(self, fingerprint: str) -> int: + if self.pinned_slots <= 0: + raise ValueError("ToolSlotCache.reserve requires pinned_slots > 0") + if fingerprint in self._confirmed: + self._confirmed.move_to_end(fingerprint) + return self._confirmed[fingerprint] + if fingerprint in self._pending: + return self._pending[fingerprint] + if len(self._confirmed) >= self.pinned_slots: + # Peek LRU without evicting — eviction happens in confirm(). + old_fp, slot = next(iter(self._confirmed.items())) + self._pending_evict = (old_fp, slot) + else: + used = set(self._confirmed.values()) | set(self._pending.values()) + slot = next( + s for s in range(self.slot_base, self.slot_base + self.pinned_slots) + if s not in used + ) + self._pending_evict = None + self._pending[fingerprint] = slot + return slot + + def confirm(self, fingerprint: str, slot: int) -> None: + if self._pending_evict is not None: + evict_fp, _ = self._pending_evict + self._confirmed.pop(evict_fp, None) + self._pending_evict = None + self._pending.pop(fingerprint, None) + self._confirmed[fingerprint] = slot + self._confirmed.move_to_end(fingerprint) + + def release_reservation(self, fingerprint: str, slot: int) -> None: + """Drop a fingerprint reserved before SNAPSHOT_THIN succeeded.""" + if self._pending.get(fingerprint) == slot: + del self._pending[fingerprint] + self._pending_evict = None + + +class ToolSplitOrchestrator: + """High-level API used by ``server_tools.py``.""" + + def __init__( + self, + *, + adapter: ToolSplitAdapter, + config: ToolSplitConfig, + tool_slot_cache: ToolSlotCache | None = None, + ): + self.adapter = adapter + self.config = config + self.tool_slots = tool_slot_cache or ToolSlotCache( + pinned_slots=config.pinned_tool_slots, + ) + + @property + def profile(self) -> str: + return self.adapter.profile_name + + def active_for_request(self, tools: Sequence[Any] | None) -> bool: + return bool(tools) + + def split_request( + self, + tokenizer: Any, + messages: Sequence[Any], + tools: Sequence[Any] | None, + *, + chat_template_kwargs: Mapping[str, Any] | None = None, + ) -> PromptSplit: + return self.adapter.split_prompt( + tokenizer, + _messages_for_adapter(messages), + _tools_for_adapter(tools), + chat_template_kwargs=chat_template_kwargs or {}, + enable_thinking=False, + ) + + def tools_fingerprint(self, tools: Sequence[Any] | None) -> str | None: + raw = _tools_for_adapter(tools) + if not raw: + return None + return self.adapter.tools_fingerprint(raw) + + def write_prompt_bin(self, ids: list[int]) -> Path: + return _ids_to_bin(ids) + + def conversation_compressible(self, split: PromptSplit) -> bool: + return ( + self.config.compress_conversation + and self.adapter.supports_pflash_on_conversation() + and split.conversation_len > 0 + ) + + def prepare_request_context( + self, + tokenizer: Any, + messages: Sequence[Any], + tools: Sequence[Any] | None, + *, + chat_template_kwargs: Mapping[str, Any] | None = None, + ) -> ToolRequestContext | None: + if not self.active_for_request(tools): + return None + split = self.split_request( + tokenizer, messages, tools, + chat_template_kwargs=chat_template_kwargs, + ) + fp = self.tools_fingerprint(tools) + if not fp or split.tool_prefix_len <= 0: + return ToolRequestContext(split=split, fingerprint=fp) + + tool_slot_hit = self.tool_slots.lookup(fp) + pending: tuple[int, int] | None = None + if tool_slot_hit is None and self.config.pinned_tool_slots > 0: + slot = self.tool_slots.reserve(fp) + pending = (slot, split.tool_prefix_len) + + return ToolRequestContext( + split=split, + fingerprint=fp, + tool_slot_hit=tool_slot_hit, + pending_tool_snap=pending, + ) + + def build_plan( + self, + *, + split: PromptSplit | None, + tools_fingerprint: str | None, + prompt_bin: Path, + prompt_len: int, + tool_slot_hit: int | None = None, + conv_hit: tuple[int, int] | None = None, + snap_prep: tuple[int, int] | None = None, + compression_fired: bool = False, + started_in_thinking: bool = False, + pending_tool_snap: tuple[int, int] | None = None, + ) -> ToolSplitPlan: + use_chain = tool_slot_hit is not None + thick = conv_hit[0] if conv_hit else None + if use_chain and thick is None: + thick_arg = -1 + elif use_chain: + thick_arg = thick + else: + thick_arg = None + + thin_ids = [tool_slot_hit] if tool_slot_hit is not None else [] + + return ToolSplitPlan( + prompt_bin_path=str(prompt_bin), + prompt_token_count=prompt_len, + tool_slot=tool_slot_hit if use_chain else None, + conv_restore_slot=thick_arg if use_chain else (conv_hit[0] if conv_hit else None), + conv_restore_prefix_len=conv_hit[1] if conv_hit else 0, + use_restore_chain=use_chain, + thin_slot_ids=thin_ids, + inline_snap=snap_prep, + compression_fired=compression_fired, + started_in_thinking=started_in_thinking, + tools_fingerprint=tools_fingerprint, + pending_tool_snap=pending_tool_snap, + tool_prefix_len=split.tool_prefix_len if split else 0, + ) + + def format_daemon_command( + self, + plan: ToolSplitPlan, + gen_len: int, + ) -> str: + """Format stdin line for test_dflash (RESTORE / RESTORE_CHAIN).""" + path = plan.prompt_bin_path + if plan.use_restore_chain and plan.thin_slot_ids: + thick = plan.conv_restore_slot if plan.conv_restore_slot is not None else -1 + thin = ",".join(str(s) for s in plan.thin_slot_ids) + line = f"RESTORE_CHAIN {thick} {thin} {path} {gen_len}" + elif plan.conv_restore_slot is not None and not plan.use_restore_chain: + line = f"RESTORE {plan.conv_restore_slot} {path} {gen_len}" + else: + line = f"{path} {gen_len}" + if plan.inline_snap: + slot, cut = plan.inline_snap + line += f" snap={cut}:{slot}" + return line + "\n" diff --git a/server/scripts/tool_split/qwen3.py b/server/scripts/tool_split/qwen3.py new file mode 100644 index 000000000..67d059c9a --- /dev/null +++ b/server/scripts/tool_split/qwen3.py @@ -0,0 +1,124 @@ +"""Qwen3.x chat-template tool / conversation splitter.""" +from __future__ import annotations + +import re +from typing import Any, Mapping, Sequence + +from tool_split.base import PromptSplit, ToolSplitAdapter +from tool_split.registry import register_adapter + +_QWEN_ARCHES = frozenset({"qwen35", "qwen36", "qwen3"}) +_QWEN_TOKENIZER_HINTS = ("qwen3", "qwen/qwen3") + + +@register_adapter("qwen3") +class Qwen3ToolSplitAdapter(ToolSplitAdapter): + """Qwen ``<|im_start|>`` templates with OpenAI ``tools=`` injection.""" + + @classmethod + def detect(cls, *, arch: str, tokenizer_id: str) -> bool: + a = (arch or "").lower() + t = (tokenizer_id or "").lower() + if a in _QWEN_ARCHES: + return True + return any(h in t for h in _QWEN_TOKENIZER_HINTS) + + def split_prompt( + self, + tokenizer: Any, + messages: Sequence[Mapping[str, Any]], + tools: Sequence[Mapping[str, Any]] | None, + *, + chat_template_kwargs: Mapping[str, Any] | None = None, + enable_thinking: bool = False, + ) -> PromptSplit: + if not tools: + full_ids = self._encode_full( + tokenizer, messages, None, chat_template_kwargs, enable_thinking, + ) + return PromptSplit( + full_ids=full_ids, + tool_prefix_ids=[], + conversation_ids=full_ids, + tool_prefix_len=0, + profile=self.profile_name, + ) + + canon = self.canonical_tools(tools) + full_ids = self._encode_full( + tokenizer, messages, canon, chat_template_kwargs, enable_thinking, + ) + + # Slice the full rendered prompt at the tool/system boundary. + # Do NOT call apply_chat_template on system-only messages: Qwen3.6 + # templates raise "No user query found in messages" when tools= is set + # without a user turn (Hermes webchat always hits this path). + boundary = self._tool_prefix_boundary(tokenizer, full_ids) + if boundary < 0: + tool_prefix_ids: list[int] = [] + conv_ids = full_ids + else: + tool_prefix_ids = full_ids[:boundary] + conv_ids = full_ids[boundary:] + + return PromptSplit( + full_ids=full_ids, + tool_prefix_ids=tool_prefix_ids, + conversation_ids=conv_ids, + tool_prefix_len=len(tool_prefix_ids), + profile=self.profile_name, + ) + + def _encode_full( + self, + tokenizer: Any, + messages: Sequence[Mapping[str, Any]], + tools: list[dict[str, Any]] | None, + chat_template_kwargs: Mapping[str, Any] | None, + enable_thinking: bool, + *, + add_generation_prompt: bool = True, + ) -> list[int]: + kwargs: dict[str, Any] = { + "tokenize": False, + "add_generation_prompt": add_generation_prompt, + "enable_thinking": enable_thinking, + } + if chat_template_kwargs: + for k, v in chat_template_kwargs.items(): + if k in ("enable_thinking", "add_generation_prompt", "tools"): + kwargs[k] = v + if tools: + kwargs["tools"] = tools + kwargs["enable_thinking"] = False + prompt = tokenizer.apply_chat_template(list(messages), **kwargs) + if add_generation_prompt and not enable_thinking: + # Match server_tools: detect thinking prefill at end of template. + pass + return tokenizer.encode(prompt, add_special_tokens=False) + + def thinking_prefill_at_end(self, tokenizer: Any, full_ids: list[int]) -> bool: + prompt = tokenizer.decode(full_ids, skip_special_tokens=False) + return bool(re.search(r"\s*$", prompt)) + + def _tool_prefix_boundary(self, tokenizer: Any, full_ids: list[int]) -> int: + """Index after tool/system header — right before first user turn.""" + from prefix_cache import ( + _find_first_seq, + _resolve_chat_markers, + find_prefix_boundary_markers, + ) + + markers = _resolve_chat_markers(tokenizer) + boundary = find_prefix_boundary_markers(full_ids, markers) + if boundary >= 0: + return boundary + + qs = tokenizer.encode("<|im_start|>", add_special_tokens=False) + user_t = tokenizer.encode("user", add_special_tokens=False) + if len(qs) == 1 and len(user_t) == 1: + user_seq = (qs[0], user_t[0]) + idx = _find_first_seq(full_ids, user_seq) + if idx >= 0: + return idx + len(qs) + return -1 diff --git a/server/scripts/tool_split/registry.py b/server/scripts/tool_split/registry.py new file mode 100644 index 000000000..9d83f1772 --- /dev/null +++ b/server/scripts/tool_split/registry.py @@ -0,0 +1,104 @@ +"""Adapter registry and user plugin discovery.""" +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from typing import TYPE_CHECKING, Type + +from tool_split.base import ToolSplitAdapter + +if TYPE_CHECKING: + from tool_split.config import ToolSplitConfig + +_REGISTRY: dict[str, Type[ToolSplitAdapter]] = {} + + +def register_adapter(name: str): + """Decorator: register a :class:`ToolSplitAdapter` subclass under *name*.""" + + def _wrap(cls: Type[ToolSplitAdapter]) -> Type[ToolSplitAdapter]: + if not issubclass(cls, ToolSplitAdapter): + raise TypeError(f"{cls!r} must subclass ToolSplitAdapter") + key = name.strip().lower() + if not key: + raise ValueError("adapter name must be non-empty") + existing = _REGISTRY.get(key) + if existing is not None and existing is not cls: + raise ValueError( + f"tool-split adapter {key!r} is already registered as {existing.__name__}") + cls.profile_name = key + _REGISTRY[key] = cls + return cls + + return _wrap + + +def builtin_adapter_names() -> list[str]: + return sorted(_REGISTRY.keys()) + + +def get_adapter_class(name: str) -> Type[ToolSplitAdapter]: + key = name.strip().lower() + if key.startswith("plugin:"): + key = key.split(":", 1)[1] + if key not in _REGISTRY: + raise KeyError( + f"unknown tool-split profile {name!r}; " + f"built-ins: {', '.join(builtin_adapter_names()) or '(none loaded)'}" + ) + return _REGISTRY[key] + + +def load_plugin_dir(plugin_dir: Path | None) -> list[str]: + """Import ``*.py`` modules from *plugin_dir* that call ``register_adapter``.""" + if plugin_dir is None or not plugin_dir.is_dir(): + return [] + loaded: list[str] = [] + for path in sorted(plugin_dir.glob("*.py")): + if path.name.startswith("_"): + continue + mod_name = f"_tool_split_plugin_{path.stem}" + spec = importlib.util.spec_from_file_location(mod_name, path) + if spec is None or spec.loader is None: + continue + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + try: + spec.loader.exec_module(mod) + loaded.append(path.name) + except Exception as exc: + print(f"[tool-split] plugin {path.name} failed: {exc}", flush=True) + return loaded + + +def resolve_adapter( + cfg: "ToolSplitConfig", + *, + arch: str, + tokenizer_id: str, +) -> ToolSplitAdapter | None: + """Pick an adapter instance from config (``auto`` uses ``detect()``).""" + if not cfg.enabled: + return None + + if cfg.plugin_dir: + load_plugin_dir(cfg.plugin_dir) + + profile = cfg.profile.strip().lower() + if profile == "auto": + for cls in _REGISTRY.values(): + try: + if cls.detect(arch=arch, tokenizer_id=tokenizer_id): + return cls() + except Exception: + continue + print( + f"[tool-split] auto: no adapter matched arch={arch!r} " + f"tokenizer={tokenizer_id!r}; disabled", + flush=True, + ) + return None + + cls = get_adapter_class(profile) + return cls() diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 420733672..a74babed8 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -234,6 +234,17 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) auto key = hash_prefix(prompt_ids.data(), cut); int idx = find_entry(key); if (idx >= 0) { + const int committed = (int)entries_[idx].ids.size(); + if (committed != cut) { + // Slot was refreshed in-place at a deeper boundary; a shallow + // hash→slot entry would restore the wrong cur_pos. + std::fprintf(stderr, + "[pc] lookup stale slot=%d key_cut=%d committed=%d — evicting\n", + entries_[idx].slot, cut, committed); + entries_.erase(entries_.begin() + idx); + entries_size_count_.fetch_sub(1, std::memory_order_relaxed); + continue; + } if (cut > best_len) { best_slot = entries_[idx].slot; best_len = cut; diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 3ecfc02a7..e1c521624 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -37,6 +37,7 @@ #include "qwen35_layer_split.h" // multi-GPU layer-split daemon args #include "layer_split_daemon_loop.h" // extracted layer-split daemon loop #include "qwen3_daemon.h" // arch dispatch - qwen3 (0.6B standalone) +#include "../src/common/snapshot_backend.h" // CPU-RAM prefix snapshot storage #include "gemma4_daemon.h" // arch dispatch - gemma4 (iSWA + MoE) #include "sampler.h" // shared CPU sampler chain (SamplerCfg / // sample_logits / parse_sampler_token) used by @@ -1131,7 +1132,19 @@ int main(int argc, char ** argv) { // ---- Single-GPU qwen35-family daemon: dispatch to the dedicated daemon ----- // This avoids the duplicated 1800-line inline loop below. The inline // loop remains for one-shot, test-window, and profile-scaling modes. - if (daemon_mode && target_gpus.size() <= 1) { + // + // DFLASH_LEGACY_DAEMON=1 keeps the inline loop for daemon mode too: the + // tool-split orchestrator (server_tools.py) speaks the legacy protocol — + // SNAPSHOT_THIN , RESTORE_CHAIN, and the + // "[snap] inline slot=" ack — which run_qwen35_daemon does not implement. + const char * legacy_env = std::getenv("DFLASH_LEGACY_DAEMON"); + const bool legacy_daemon = legacy_env && std::atoi(legacy_env) != 0; + if (daemon_mode && legacy_daemon) { + std::fprintf(stderr, + "[test_dflash] DFLASH_LEGACY_DAEMON=1 -> using legacy inline daemon loop " + "(tool-split protocol)\n"); + } + if (daemon_mode && !legacy_daemon && target_gpus.size() <= 1) { const int max_ctx_eff = g_max_ctx_override > 0 ? g_max_ctx_override : 4096; dflash::common::Qwen35DaemonArgs qargs; qargs.target_path = target_path; @@ -1182,6 +1195,26 @@ int main(int argc, char ** argv) { } ggml_backend_t backend = target_backend; // legacy target-side alias + // 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; + } + struct SnapBackendGuard { + ggml_backend_t * snap; + ggml_backend_t target; + SnapBackendGuard(ggml_backend_t * s, ggml_backend_t t) : snap(s), target(t) {} + ~SnapBackendGuard() { + if (snap && *snap) { + dflash::common::free_snapshot_backend(*snap, target); + *snap = nullptr; + } + } + } snap_guard(&snap_backend, target_backend); + TargetWeights w; if (!load_target_gguf(target_path, target_backend, w)) { std::fprintf(stderr, "target load: %s\n", dflash27b_last_error()); @@ -2427,7 +2460,7 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[snap] SNAPSHOT_THIN bad args\n"); continue; } - if (!snapshot_target_cache_thin(w, cache, backend, kv_start, kv_end, + if (!snapshot_target_cache_thin(w, cache, snap_backend, kv_start, kv_end, prefix_snapshots[slot])) { std::fprintf(stderr, "[snap] thin failed slot=%d: %s\n", slot, dflash27b_last_error()); @@ -2444,7 +2477,7 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[snap] invalid slot %d\n", slot); continue; } - if (!snapshot_target_cache(w, cache, backend, prefix_snapshots[slot])) { + if (!snapshot_target_cache(w, cache, snap_backend, prefix_snapshots[slot])) { std::fprintf(stderr, "[snap] failed slot=%d: %s\n", slot, dflash27b_last_error()); continue; } @@ -2538,6 +2571,15 @@ int main(int argc, char ** argv) { stream_emit(-1); continue; } + // Optional inline-snap suffix (same as RESTORE / bare prompt): + // snap=: + if (const char * sp = std::strstr(line.c_str(), " snap=")) { + if (std::sscanf(sp + 1, "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; + } + } n_gen = n_gen_local; prompt_file_str = ppath; prompt_path = prompt_file_str.c_str(); @@ -2560,8 +2602,8 @@ int main(int argc, char ** argv) { restore_from_slot = true; restore_slot_id = slot; // Parse optional inline-snap suffix: snap=: - if (const char * sp = std::strstr(line.c_str(), "snap=")) { - if (std::sscanf(sp, "snap=%d:%d", &snap_pos, &snap_slot) != 2 + if (const char * sp = std::strstr(line.c_str(), " snap=")) { + if (std::sscanf(sp + 1, "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; @@ -2576,8 +2618,8 @@ int main(int argc, char ** argv) { prompt_file_str = ppath; prompt_path = prompt_file_str.c_str(); // Parse optional inline-snap suffix: snap=: - if (const char * sp = std::strstr(line.c_str(), "snap=")) { - if (std::sscanf(sp, "snap=%d:%d", &snap_pos, &snap_slot) != 2 + if (const char * sp = std::strstr(line.c_str(), " snap=")) { + if (std::sscanf(sp + 1, "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; @@ -2871,7 +2913,7 @@ int main(int argc, char ** argv) { if (snap_pos >= 0 && snap_pos == start) { cache.cur_pos = start; if (snap_slot >= 0) { - if (snapshot_target_cache(w, cache, backend, prefix_snapshots[snap_slot])) { + if (snapshot_target_cache(w, cache, snap_backend, prefix_snapshots[snap_slot])) { std::printf("[snap] inline slot=%d cur_pos=%d\n", snap_slot, start); std::fflush(stdout); } else { @@ -2968,7 +3010,7 @@ int main(int argc, char ** argv) { cache.cur_pos = committed; cache.last_tok = last_tok; if (snap_slot >= 0) { - if (snapshot_target_cache(w, cache, backend, prefix_snapshots[snap_slot])) { + if (snapshot_target_cache(w, cache, snap_backend, prefix_snapshots[snap_slot])) { std::printf("[snap] inline slot=%d cur_pos=%d\n", snap_slot, committed); std::fflush(stdout); } else {