diff --git a/benchmarks/bench_accuracy_baseline.py b/benchmarks/bench_accuracy_baseline.py new file mode 100644 index 0000000..356f0d7 --- /dev/null +++ b/benchmarks/bench_accuracy_baseline.py @@ -0,0 +1,648 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Accuracy baseline: chunk-mode KV reuse vs full prefill on LongBench QA. + +Drives the production DaseR HTTP service (``/documents`` + ``/infer``) for +multi-document QA samples. Each sample's passages are uploaded as separate +documents (independently prefilled chunks), then the same assembled prompt is +answered twice with identical greedy decoding: + +- ``full``: ``use_kv_cache=false`` -> full prefill (accuracy upper bound). +- ``reuse``: ``use_kv_cache=true`` -> concatenated chunk KV, RoPE-relocated, + no cross-attention repair (the lossy path under measurement). + +Reported per dataset: LongBench-style QA F1 for both modes, the F1 gap, +extracted-answer F1 (bold span / first sentence, separating factual errors +from verbosity-induced format drift), gold-answer recall, normalized answer +agreement, chunk-hit token coverage, and TTFT. Accuracy means are computed +over samples whose reuse request verifiably loaded chunk KV; samples that +silently fell back to full recompute are excluded and reported separately. + +Usage: + python benchmarks/bench_accuracy_baseline.py \\ + --daser-url http://127.0.0.1:2046 \\ + --tokenizer /path/to/model \\ + --data-dir /path/to/daser-baseline/data \\ + --datasets hotpotqa,2wikimqa,musique \\ + --num-samples 50 \\ + --out /path/to/daser-baseline/results/baseline.json + +The DaseR server must run with ``--cache-reuse-mode chunk`` and vLLM must be +started with prefix caching disabled (the standard DaseR deployment). +""" + +from __future__ import annotations + +# Standard +import argparse +from collections import Counter +from dataclasses import asdict, dataclass, field +import json +from pathlib import Path +import re +import string +import sys +import time +from typing import Any + +# Third Party +import requests + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +# First Party +from daser.logging import init_logger + +logger = init_logger(__name__) + +PASSAGE_SPLIT_RE = re.compile(r"(?=Passage \d+:\n)") +BOLD_SPAN_RE = re.compile(r"\*\*(.+?)\*\*") +TASK_INSTRUCTION = ( + "Answer the question based on the given documents. Only give me the " + "answer and do not output any other words.\n\nQuestion: {question}" +) +HTTP_TIMEOUT_S = 600.0 +DEFAULT_MAX_NEW_TOKENS = 256 +DEFAULT_NUM_SAMPLES = 50 +DEFAULT_MAX_CONTEXT_TOKENS = 24000 + + +@dataclass +class SampleResult: + """Accuracy comparison result for one QA sample. + + Attributes: + sample_id: LongBench record ``_id``. + question: QA question text. + gold_answers: reference answers. + num_docs: number of uploaded passage documents. + prompt_tokens: assembled prompt length reported by ``/infer``. + full_text: answer text from the full-prefill mode. + reuse_text: answer text from the chunk-reuse mode. + full_f1: QA F1 of the full-prefill answer. + reuse_f1: QA F1 of the chunk-reuse answer. + full_f1_extracted: QA F1 of the extracted full-prefill answer span. + reuse_f1_extracted: QA F1 of the extracted chunk-reuse answer span. + full_recall: gold-answer substring recall of full prefill (0/1). + reuse_recall: gold-answer substring recall of chunk reuse (0/1). + agreement: normalized exact match between the two mode answers. + reuse_coverage: fraction of prompt tokens served from cached chunks. + full_ttft_ms: TTFT of the full-prefill request. + reuse_ttft_ms: TTFT of the chunk-reuse request. + reuse_loaded: whether the reuse path actually loaded cached KV, + inferred from reuse TTFT being meaningfully below full TTFT. + error: failure description when the sample was skipped. + """ + + sample_id: str + question: str + gold_answers: list[str] + num_docs: int = 0 + prompt_tokens: int = 0 + full_text: str = "" + reuse_text: str = "" + full_f1: float = 0.0 + reuse_f1: float = 0.0 + full_f1_extracted: float = 0.0 + reuse_f1_extracted: float = 0.0 + full_recall: float = 0.0 + reuse_recall: float = 0.0 + agreement: bool = False + reuse_coverage: float = 0.0 + full_ttft_ms: float = 0.0 + reuse_ttft_ms: float = 0.0 + reuse_loaded: bool = False + error: str = "" + + +@dataclass +class DatasetSummary: + """Aggregated accuracy metrics for one dataset. + + Accuracy means (`full_f1` .. `agreement_rate`) are computed only over + samples whose reuse request verifiably loaded chunk KV (``reuse_loaded``), + so silent full-recompute fallbacks cannot dilute the comparison. + + Attributes: + dataset: dataset name. + samples: number of scored samples (no error). + samples_loaded: scored samples whose reuse path actually loaded KV; + accuracy means below are over this subset. + skipped: number of skipped samples (errors / over-length). + full_f1: mean QA F1 of full prefill. + reuse_f1: mean QA F1 of chunk reuse. + f1_gap: ``full_f1 - reuse_f1``. + full_f1_extracted: mean QA F1 of the extracted full-prefill answer. + reuse_f1_extracted: mean QA F1 of the extracted chunk-reuse answer. + f1_gap_extracted: ``full_f1_extracted - reuse_f1_extracted``. + full_recall: mean gold-answer substring recall of full prefill. + reuse_recall: mean gold-answer substring recall of chunk reuse. + recall_gap: ``full_recall - reuse_recall``. + agreement_rate: fraction of samples with matching normalized answers. + mean_coverage: mean chunk-hit token coverage over scored samples. + loaded_rate: fraction of scored samples with a verified KV load. + mean_full_ttft_ms: mean full-prefill TTFT over scored samples. + mean_reuse_ttft_ms: mean chunk-reuse TTFT over scored samples. + results: per-sample records. + """ + + dataset: str + samples: int = 0 + samples_loaded: int = 0 + skipped: int = 0 + full_f1: float = 0.0 + reuse_f1: float = 0.0 + f1_gap: float = 0.0 + full_f1_extracted: float = 0.0 + reuse_f1_extracted: float = 0.0 + f1_gap_extracted: float = 0.0 + full_recall: float = 0.0 + reuse_recall: float = 0.0 + recall_gap: float = 0.0 + agreement_rate: float = 0.0 + mean_coverage: float = 0.0 + loaded_rate: float = 0.0 + mean_full_ttft_ms: float = 0.0 + mean_reuse_ttft_ms: float = 0.0 + results: list[SampleResult] = field(default_factory=list) + + +def normalize_answer(text: str) -> str: + """Normalize a QA answer for comparison (LongBench convention). + + Args: + text: raw answer text. + + Returns: + Lower-cased text without articles, punctuation, or extra whitespace. + """ + text = text.lower() + text = "".join(ch for ch in text if ch not in set(string.punctuation)) + text = re.sub(r"\b(a|an|the)\b", " ", text) + return " ".join(text.split()) + + +def qa_f1_score(prediction: str, gold_answers: list[str]) -> float: + """Token-level QA F1 against the best-matching gold answer. + + Args: + prediction: model answer text. + gold_answers: list of acceptable reference answers. + + Returns: + Maximum F1 over the gold answers, in [0, 1]. + """ + pred_tokens = normalize_answer(prediction).split() + best = 0.0 + for gold in gold_answers: + gold_tokens = normalize_answer(gold).split() + if not pred_tokens or not gold_tokens: + best = max(best, float(pred_tokens == gold_tokens)) + continue + common = Counter(pred_tokens) & Counter(gold_tokens) + num_same = sum(common.values()) + if num_same == 0: + continue + precision = num_same / len(pred_tokens) + recall = num_same / len(gold_tokens) + best = max(best, 2 * precision * recall / (precision + recall)) + return best + + +def answer_recall(prediction: str, gold_answers: list[str]) -> float: + """Return 1.0 when any gold answer appears in the prediction. + + Substring containment after normalization. Robust to answer verbosity, + so it isolates factual correctness from instruction-following format + (terse vs verbose) drift that token-level F1 conflates. + + Args: + prediction: model answer text. + gold_answers: list of acceptable reference answers. + + Returns: + 1.0 if a normalized gold answer is contained in the normalized + prediction, else 0.0. + """ + pred_norm = normalize_answer(prediction) + for gold in gold_answers: + gold_norm = normalize_answer(gold) + if gold_norm and gold_norm in pred_norm: + return 1.0 + return 0.0 + + +def extract_short_answer(text: str) -> str: + """Extract the concise answer span from a possibly verbose response. + + Degraded chunk-reuse answers tend to wrap the correct answer in + explanatory prose (often bolding the answer itself). Scoring the + extracted span separates factual errors from verbosity-induced + token-F1 dilution. + + Args: + text: raw model answer text. + + Returns: + The first bold ``**span**`` when present, otherwise the first + sentence of the first line. + """ + match = BOLD_SPAN_RE.search(text) + if match: + return match.group(1) + first_line = text.strip().split("\n", 1)[0] + return first_line.split(". ", 1)[0] + + +def split_passages(context: str) -> list[str]: + """Split a LongBench context into passage documents. + + Args: + context: raw context text, typically ``Passage N:`` delimited. + + Returns: + Non-empty passage texts; the whole context when no markers exist. + """ + parts = [part.strip() for part in PASSAGE_SPLIT_RE.split(context)] + parts = [part for part in parts if part] + if len(parts) <= 1: + return [context.strip()] if context.strip() else [] + return parts + + +def load_samples(path: Path, limit: int) -> list[dict[str, Any]]: + """Load LongBench JSONL records. + + Args: + path: dataset JSONL path. + limit: maximum number of records (0 = all). + + Returns: + Parsed records in file order. + """ + records: list[dict[str, Any]] = [] + with path.open(encoding="utf-8") as fp: + for line in fp: + line = line.strip() + if not line: + continue + records.append(json.loads(line)) + if limit and len(records) >= limit: + break + return records + + +class DaserClient: + """Thin synchronous client for the DaseR HTTP service. + + Args: + base_url: DaseR HTTP server base URL. + + Async/thread-safety: + Uses one ``requests.Session`` from a single benchmark thread. + """ + + def __init__(self, base_url: str) -> None: + self._base_url = base_url.rstrip("/") + self._session = requests.Session() + + def health(self) -> bool: + """Return whether DaseR and vLLM report healthy.""" + resp = self._session.get(f"{self._base_url}/health", timeout=30) + resp.raise_for_status() + return resp.json().get("status") == "ok" + + def upload_document(self, title: str, text: str) -> str: + """Upload one document and return its doc_id.""" + resp = self._session.post( + f"{self._base_url}/documents", + json={"title": title, "text": text}, + timeout=HTTP_TIMEOUT_S, + ) + resp.raise_for_status() + return str(resp.json()["doc_id"]) + + def delete_document(self, doc_id: str) -> None: + """Delete one document, ignoring missing IDs.""" + try: + self._session.delete( + f"{self._base_url}/documents/{doc_id}", + timeout=HTTP_TIMEOUT_S, + ) + except requests.RequestException as exc: # pragma: no cover - cleanup + logger.warning("delete_document %s failed: %s", doc_id[:8], exc) + + def infer( + self, + doc_ids: list[str], + task: str, + use_kv_cache: bool, + max_new_tokens: int, + ) -> dict[str, Any]: + """Run one ``/infer`` request with greedy decoding. + + Args: + doc_ids: documents to include in the prompt. + task: task text appended after the documents. + use_kv_cache: enable DaseR chunk KV loading. + max_new_tokens: generation cap. + + Returns: + ``/infer`` response payload. + """ + resp = self._session.post( + f"{self._base_url}/infer", + json={ + "doc_ids": doc_ids, + "task": task, + "use_kv_cache": use_kv_cache, + "trace_cache": use_kv_cache, + "gen_params": { + "max_tokens": max_new_tokens, + "temperature": 0.0, + "top_p": 1.0, + }, + }, + timeout=HTTP_TIMEOUT_S, + ) + resp.raise_for_status() + return resp.json() + + +def hit_token_coverage(response: dict[str, Any]) -> float: + """Return the fraction of prompt tokens covered by cache hits. + + Args: + response: ``/infer`` response with ``cache_hits`` trace entries. + + Returns: + Covered tokens / prompt tokens, in [0, 1]. + """ + prompt_tokens = int(response.get("prompt_tokens", 0)) + if prompt_tokens <= 0: + return 0.0 + covered = sum( + int(hit.get("token_count", 0)) for hit in response.get("cache_hits", []) + ) + return min(1.0, covered / prompt_tokens) + + +def run_sample( + client: DaserClient, + record: dict[str, Any], + max_new_tokens: int, + delete_docs: bool, +) -> SampleResult: + """Measure one sample in both modes. + + Args: + client: DaseR HTTP client. + record: LongBench QA record. + max_new_tokens: generation cap per request. + delete_docs: delete uploaded documents after the sample. Keeping + them avoids delete/ring-eviction interplay and lets identical + passages dedup across samples. + + Returns: + Populated SampleResult; ``error`` is set on failure. + """ + result = SampleResult( + sample_id=str(record.get("_id", "")), + question=str(record.get("input", "")), + gold_answers=[str(ans) for ans in record.get("answers", [])], + ) + passages = split_passages(str(record.get("context", ""))) + if not passages: + result.error = "empty context" + return result + + doc_ids: list[str] = [] + try: + for i, passage in enumerate(passages): + doc_ids.append( + client.upload_document(f"{result.sample_id}-p{i:02d}", passage) + ) + result.num_docs = len(doc_ids) + task = TASK_INSTRUCTION.format(question=result.question) + + full = client.infer( + doc_ids, task, use_kv_cache=False, max_new_tokens=max_new_tokens + ) + reuse = client.infer( + doc_ids, task, use_kv_cache=True, max_new_tokens=max_new_tokens + ) + except requests.RequestException as exc: + result.error = f"http error: {exc}" + return result + finally: + if delete_docs: + for doc_id in doc_ids: + client.delete_document(doc_id) + + result.prompt_tokens = int(full.get("prompt_tokens", 0)) + result.full_text = str(full.get("text", "")).strip() + result.reuse_text = str(reuse.get("text", "")).strip() + result.full_f1 = qa_f1_score(result.full_text, result.gold_answers) + result.reuse_f1 = qa_f1_score(result.reuse_text, result.gold_answers) + result.full_f1_extracted = qa_f1_score( + extract_short_answer(result.full_text), result.gold_answers + ) + result.reuse_f1_extracted = qa_f1_score( + extract_short_answer(result.reuse_text), result.gold_answers + ) + result.full_recall = answer_recall(result.full_text, result.gold_answers) + result.reuse_recall = answer_recall(result.reuse_text, result.gold_answers) + result.agreement = normalize_answer(result.full_text) == normalize_answer( + result.reuse_text + ) + result.reuse_coverage = hit_token_coverage(reuse) + result.full_ttft_ms = float(full.get("ttft_ms") or 0.0) + result.reuse_ttft_ms = float(reuse.get("ttft_ms") or 0.0) + # A genuine chunk-reuse load skips most prefill, so its TTFT is well + # below full prefill. When reuse TTFT is not clearly lower, the reuse + # path silently fell back to full recompute (e.g. evicted prefix), which + # makes the accuracy comparison meaningless for that sample. + result.reuse_loaded = result.reuse_ttft_ms < 0.85 * result.full_ttft_ms + return result + + +def summarize(dataset: str, results: list[SampleResult]) -> DatasetSummary: + """Aggregate per-sample results for one dataset. + + Accuracy means are computed only over samples with a verified chunk-KV + load (``reuse_loaded``); coverage, loaded rate, and TTFT means are over + all scored samples. + + Args: + dataset: dataset name. + results: all per-sample records including failures. + + Returns: + DatasetSummary with means as described above. + """ + summary = DatasetSummary(dataset=dataset, results=results) + scored = [r for r in results if not r.error] + loaded = [r for r in scored if r.reuse_loaded] + summary.samples = len(scored) + summary.samples_loaded = len(loaded) + summary.skipped = len(results) - len(scored) + if not scored: + return summary + n = float(len(scored)) + summary.mean_coverage = sum(r.reuse_coverage for r in scored) / n + summary.loaded_rate = len(loaded) / n + summary.mean_full_ttft_ms = sum(r.full_ttft_ms for r in scored) / n + summary.mean_reuse_ttft_ms = sum(r.reuse_ttft_ms for r in scored) / n + if not loaded: + return summary + m = float(len(loaded)) + summary.full_f1 = sum(r.full_f1 for r in loaded) / m + summary.reuse_f1 = sum(r.reuse_f1 for r in loaded) / m + summary.f1_gap = summary.full_f1 - summary.reuse_f1 + summary.full_f1_extracted = sum(r.full_f1_extracted for r in loaded) / m + summary.reuse_f1_extracted = sum(r.reuse_f1_extracted for r in loaded) / m + summary.f1_gap_extracted = summary.full_f1_extracted - summary.reuse_f1_extracted + summary.full_recall = sum(r.full_recall for r in loaded) / m + summary.reuse_recall = sum(r.reuse_recall for r in loaded) / m + summary.recall_gap = summary.full_recall - summary.reuse_recall + summary.agreement_rate = sum(1 for r in loaded if r.agreement) / m + return summary + + +def print_report(summaries: list[DatasetSummary]) -> None: + """Print the aggregate accuracy comparison table. + + Args: + summaries: per-dataset aggregates. + """ + print("\n" + "=" * 118) + print("ACCURACY BASELINE — chunk-mode KV reuse vs full prefill (greedy)") + print("=" * 118) + header = ( + f"{'Dataset':<12} {'N':>4} {'Full F1':>8} {'Reuse F1':>9} {'F1 gap':>8} " + f"{'ExtF1 gap':>10} {'Full Rec':>9} {'Reuse Rec':>10} {'Rec gap':>8} " + f"{'Agree':>7} {'Loaded':>7}" + ) + print(header) + print("-" * 118) + for s in summaries: + print( + f"{s.dataset:<12} {s.samples_loaded:>4} {s.full_f1:>8.4f} " + f"{s.reuse_f1:>9.4f} {s.f1_gap:>+8.4f} " + f"{s.f1_gap_extracted:>+10.4f} {s.full_recall:>9.4f} " + f"{s.reuse_recall:>10.4f} {s.recall_gap:>+8.4f} " + f"{s.agreement_rate:>6.1%} {s.loaded_rate:>6.1%}" + ) + print("-" * 118) + total = sum(s.samples_loaded for s in summaries) + if total: + full = sum(s.full_f1 * s.samples_loaded for s in summaries) / total + reuse = sum(s.reuse_f1 * s.samples_loaded for s in summaries) / total + full_r = sum(s.full_recall * s.samples_loaded for s in summaries) / total + reuse_r = sum(s.reuse_recall * s.samples_loaded for s in summaries) / total + ext_gap = sum(s.f1_gap_extracted * s.samples_loaded for s in summaries) / total + print( + f"{'OVERALL':<12} {total:>4} {full:>8.4f} {reuse:>9.4f} " + f"{full - reuse:>+8.4f} {ext_gap:>+10.4f} {full_r:>9.4f} " + f"{reuse_r:>10.4f} {full_r - reuse_r:>+8.4f}" + ) + print("=" * 118) + print( + "Note: N counts samples with a verified chunk-KV load; accuracy " + "means exclude silent full-recompute fallbacks.\n'ExtF1 gap' scores " + "the extracted answer span (bold/first sentence), separating factual " + "errors from verbosity drift." + ) + + +def main() -> None: + """Entry point.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--daser-url", default="http://127.0.0.1:2046") + parser.add_argument( + "--tokenizer", + required=True, + help="HF tokenizer path for context-length filtering", + ) + parser.add_argument("--data-dir", required=True) + parser.add_argument("--datasets", default="hotpotqa,2wikimqa,musique") + parser.add_argument("--num-samples", type=int, default=DEFAULT_NUM_SAMPLES) + parser.add_argument( + "--max-context-tokens", type=int, default=DEFAULT_MAX_CONTEXT_TOKENS + ) + parser.add_argument("--max-new-tokens", type=int, default=DEFAULT_MAX_NEW_TOKENS) + parser.add_argument( + "--delete-docs", + action="store_true", + help="Delete uploaded documents after each sample (default: keep)", + ) + parser.add_argument("--out", default=None) + args = parser.parse_args() + + # Third Party + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + client = DaserClient(args.daser_url) + if not client.health(): + raise SystemExit("DaseR /health is not ok; start the services first") + + summaries: list[DatasetSummary] = [] + for dataset in [d.strip() for d in args.datasets.split(",") if d.strip()]: + path = Path(args.data_dir) / f"{dataset}.jsonl" + records = load_samples(path, 0) + results: list[SampleResult] = [] + scored = 0 + for record in records: + if scored >= args.num_samples: + break + context_tokens = len( + tokenizer.encode( + str(record.get("context", "")), add_special_tokens=False + ) + ) + if context_tokens > args.max_context_tokens: + continue + t0 = time.time() + result = run_sample(client, record, args.max_new_tokens, args.delete_docs) + results.append(result) + if not result.error: + scored += 1 + logger.info( + "[%s %d/%d] docs=%d tokens=%d full_f1=%.3f reuse_f1=%.3f " + "cov=%.0f%% %.1fs %s", + dataset, + scored, + args.num_samples, + result.num_docs, + result.prompt_tokens, + result.full_f1, + result.reuse_f1, + result.reuse_coverage * 100, + time.time() - t0, + result.error or "", + ) + summaries.append(summarize(dataset, results)) + + print_report(summaries) + if args.out: + payload = { + "config": { + "daser_url": args.daser_url, + "datasets": args.datasets, + "num_samples": args.num_samples, + "max_context_tokens": args.max_context_tokens, + "max_new_tokens": args.max_new_tokens, + "tokenizer": args.tokenizer, + "decoding": "greedy", + "delete_docs": args.delete_docs, + }, + "summaries": [asdict(s) for s in summaries], + } + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False)) + print(f"\nJSON results written to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/docs/optimizations/5_chunk_cross_attention_repair.md b/docs/optimizations/5_chunk_cross_attention_repair.md new file mode 100644 index 0000000..6d102ad --- /dev/null +++ b/docs/optimizations/5_chunk_cross_attention_repair.md @@ -0,0 +1,400 @@ +# Chunk-Mode Cross-Attention Repair: Analysis and Optimization Directions + +**Date:** 2026-06-10 +**Target:** accuracy of multi-chunk KV reuse in `--cache-reuse-mode chunk` +**Scope:** `daser/server/`, `daser/connector/`, `daser/ops/`, `benchmarks/` +**Status:** analysis and staged proposal; no code change in this document + +Chunks are prefilled independently, so their KV tensors carry no +cross-attention against the system prompt or other documents. This document +decomposes the resulting accuracy error, states the hard constraints the +current architecture imposes on any repair, and lays out optimization +directions ranked by leverage. It supersedes the Stage C/D portions of the +earlier multi-document production plan where the two overlap. + +## 1. Current State + +The only load-time transform on master is positional: RoPE delta relocation +plus an optional global K/V scale in +`daser/connector/staging.py::_transform_loaded_staging_batch`. No +cross-attention repair exists. + +The chunk-mode pipeline: + +1. **Upload** (`POST /documents`): the document is padded to a block boundary + (`Chunker.pad_to_block_boundary`, default `block_tokens=16`) and prefilled + as a standalone prompt through vLLM (`max_tokens=1`). The chunk's KV is + computed in isolation starting at position 0 and stored with + `pos_offset=0`. Fixed segments (`chat_prefix`, `doc_separator`) are + prewarmed the same way. +2. **Lookup** (`/infer`): `ChunkReuseIndex` scans block-aligned prompt windows + and returns full-chunk hits at arbitrary offsets. +3. **Supply**: `SchedulerConnectorMixin.get_num_new_matched_tokens` accepts + only the contiguous chunk coverage starting at `num_computed_tokens` + (`_contiguous_prefix_tokens`), capped at `available - 1` so at least the + final prompt token is always recomputed. +4. **Load**: staging bytes are repositioned by a TileLang RoPE delta kernel + (`daser/ops/rope_apply.py`) and copied into the vLLM KV cache. + +## 2. Error Decomposition + +Four independent error sources, in decreasing order of impact: + +1. **Fake attention sinks.** Each chunk is prefilled from position 0, so the + model treats its first tokens as an attention sink (abnormally large K + norms). A concatenated prompt with N chunks contains N fake sinks that + compete for attention mass. Published ablations (EPIC/LegoLink) find this + — not the missing cross-attention itself — is the dominant cause of + accuracy loss in position-independent caching. +2. **Missing cross-chunk attention.** Document B's KV never attended to the + system prompt or document A. +3. **Pad-token noise.** Block-boundary padding KV is stored, loaded, and + attended to by every subsequent token (up to `block_tokens - 1` pad + tokens per segment). +4. **RoPE relocation numerics.** bf16 rotation error at large position + deltas. Minor relative to the above. + +The query/task suffix is always recomputed with full attention over all +loaded KV, so the tail seam is healthy by construction; the damage is +concentrated inside and between the reused chunks. + +## 3. Hard Constraints + +Any repair plan must respect these; the first two are properties of the vLLM +`KVConnectorBase_V1` interface (read-only third-party dependency, CLAUDE.md +rule 6), the third is a property of our supply chain. + +1. **Contiguous-only external supply.** `get_num_new_matched_tokens` returns + a single int. "Skip the first k tokens of each chunk and let vLLM + recompute them" (LegoLink-style selective recompute) is not expressible + without vLLM changes. Holes can only be ceded at the head or tail of the + whole reused prefix. +2. **No attention introspection.** The connector cannot read QK^T scores or + attention weights, and cannot change the attention computation + (temperature, masking) in-kernel. +3. **Block granularity.** `_trim_chunk_to_external_window` and the whole + supply path align to `block_tokens` (16). Token-grain trim parameters + round up to whole blocks; plan trims in block units. + +## 4. Optimization Directions + +### 4.1 Shared-prefix prefill — eliminate fake sinks at store time + +**Highest leverage.** Instead of correcting sink drift numerically at load +time, prevent the fake sink from forming: prefill `[chat_prefix + doc]` at +upload time and store only the KV of the doc's block range. + +- This is the shared-prefix mechanism from APE (ICLR'25): every chunk + computes its KV facing the same real sink, so K-norm distributions stay + normal and chunk heads are never "sinkified". +- Cost is marginal: chunk-mode lookup already hits the prewarmed + `chat_prefix` chunk, so the prefix portion of the upload prefill loads + from cache and only doc-attends-prefix compute is added. +- The missing piece is **sub-range store**: `ChunkReuseStrategy` only + supports whole-prefix stores keyed by the full aligned prefix hash. The + store spec needs a "store only the trailing N blocks, keyed by the doc + segment hash" form. The load side already does the symmetric metadata + surgery in `_trim_chunk_to_external_window`, so this is a contained + connector + HTTP-upload change, entirely inside `daser/`. +- Orthogonal to load-time sink correction; once landed, the numeric + correction term shrinks and may be disabled. + +### 4.2 Composed-prefix partial hits + canonical chunk ordering — the lossless path + +The contiguous-supply constraint is also an opportunity: "reuse a composed +prefix `[SYS, A, B]` + fully recompute the tail `[C, task]`" is natively +expressible and **strictly lossless** — the composed prefix carries full +internal cross-attention and the tail is recomputed. + +- Requires a composed-KV index keyed by the ordered chunk sequence with + longest-common-prefix matching, alongside the existing per-chunk index. + Full-combination hits are a special case; partial-prefix hits are where + the storage reuse rate comes from. +- **Canonical ordering** multiplies the hit rate: when assembling the + `/infer` prompt, normalize `doc_ids` order (by access frequency, then a + stable tiebreak) so `[A, C, B]` and `[A, B, C]` collapse into one + combination. RAG document order is usually semantically exchangeable; + gate it behind a flag for callers that require explicit order. + Frequency-first ordering also concentrates heat near the prefix-tree + root (RAGCache's knowledge-tree insight), which composes well with + ring-buffer eviction. + +### 4.3 Seam patches — residuals instead of seam chunks + +Storing whole doc-pair seam KV costs O(pair length). The difference between +full-prefill KV and independently concatenated KV is concentrated in a small +number of boundary tokens (the HKVD observation behind CacheBlend's ~15% +recompute budget). Store only the **KV residual** of the k seam tokens and +add it back in `_transform_loaded_staging_batch` after RoPE relocation. + +- Storage drops from O(pair) to O(k_seam) and the residual compresses well. +- Residual generation needs no HF dependency: a background full prefill of + the combination (`daser_skip_load=true` + store) diffed against the + independent chunk KV produces it, and the same pass yields the online + calibration signal of §4.6. + +### 4.4 Per-chunk load-time correction — upgrade the existing hooks + +`load_key_scale` / `load_value_scale` already exist as connector-level +constants (default 1.0, plumbed from `kv_connector_extra_config` through +`copy_staging_to_kv_cache`). APE shows per-chunk K/V scaling factors close +most of the parallel-vs-sequential encoding gap when combined with a shared +prefix. + +- Promote the two scalars to per-chunk metadata: `ChunkMeta` → + IPC `lookup` payload → `ReqLoadSpec` → staging transform. +- Watch the copy-run fragmentation: `build_load_copy_runs` merges adjacent + loads only when `pos_offset` matches; per-chunk parameters would split + runs and multiply kernel launches. Extend the TileLang kernels to accept + per-slot parameter tables (same pattern as the existing cos/sin tables) + so one launch handles heterogeneous chunks. +- Sink-drift correction coefficients, if kept, should be calibrated + per-layer rather than as a global constant — error concentrates in + retrieval heads and shallow layers, and the stat tensor is already + shaped `[layers, kv_heads, head_dim]`. +- APE's third knob (attention temperature) requires kernel changes and + stays out of scope; uniform K scaling is its first-order approximation. + +### 4.5 Low-cost hygiene fixes + +- **Pad strategy.** Replace `pad_token` block padding with semantically + neutral natural tokens (newline / separator text), or zero the V rows of + pad positions at load time. Removes error source 3 outright. +- **Tail awareness.** `extra_tokens = available - 1` already guarantees one + recomputed token; head-side block-aligned trim + (`_trim_chunk_to_external_window`) is the only other recompute lever and + rounds to 16-token blocks — size any head-trim parameter accordingly. + +### 4.6 Traffic-driven composed upgrade + KV-space calibration + +The cost-based decision between "independent chunks + correction" and +"composed KV" needs an execution mechanism that does not depend on users +registering doc-sets: + +- First sighting of a combination: serve via independent chunks (fast, + lossy), then schedule a low-priority background full prefill + (`daser_skip_load` + store) that writes the composed KV. Second sighting + onward hits losslessly. Combination-level access counters plus a bounded + background scheduler are the only new pieces; GPU contention is managed + by concurrency limits and off-peak scheduling. +- The same shadow prefill produces an **online accuracy proxy**: per-layer + distance between composed and concatenated KV, measured in KV space on + the serving model. This replaces offline HF + `output_attentions` calibration — no separate pipeline, no + calibration-vs-production model drift. + +## 5. Priority and Sequencing + +1. **Accuracy baseline first — done, see §7.** + `benchmarks/bench_accuracy_baseline.py` measures raw concatenation vs + full prefill on a LongBench QA subset. The measured gap (extracted-answer + F1 +0.21 to +0.45, recall +0.05 to +0.26) is large enough to justify the + full sequence below; multi-hop QA (HotpotQA) is confirmed as the + highest-impact, highest-signal regression set. +2. **Shared-prefix prefill (§4.1)** — root-cause fix for the dominant error + term, contained change, no vLLM involvement. +3. **Pad hygiene (§4.5)** — trivial cost, pure win. +4. **Composed partial hits + canonical ordering (§4.2)** — the lossless + path; mostly control-plane work in `daser/server/`. +5. **Per-chunk correction (§4.4), seam patches (§4.3), background upgrade + (§4.6)** — invest based on the residual gap the baseline shows after + steps 2-4. + +The overall stance: load-time numeric correction is the last safety net, not +the primary repair. The biggest levers are **store-time conditions** (§4.1, +which removes the dominant error source) and **combination reuse rate** +(§4.2/§4.6, which routes ever more traffic off the lossy path entirely). + +## 6. Out of Scope (requires vLLM or kernel changes) + +Tracked for upstream watching, not actionable inside `daser/` today: + +- Token-grain selective recompute (CacheBlend-style HKVD) — needs + non-contiguous external supply and attention-weight export. +- Per-chunk head-token recompute (LegoLink) — needs non-contiguous supply. +- In-kernel attention temperature / sink masking (APE's full recipe). +- NoPE-format KV storage with RoPE fused at attention time — removes + relocation numerics entirely. + +## 7. Accuracy Baseline Results (2026-06-10) + +### 7.1 Experiment Setup + +| Item | Value | +|------|-------| +| GPU | NVIDIA GeForce RTX 4090 (24 GiB) | +| Model | Qwen3-8B, bf16, max_model_len=16384 | +| DaseR transfer | iouring, L1=4 GiB, L2=30 GiB | +| Dataset | LongBench QA subset: hotpotqa, 2wikimqa, musique | +| Decoding | greedy (temperature=0) | +| Max new tokens | 256 (an initial run used 64 and truncated 35–69% of reuse answers; see §7.7) | +| Max context tokens | 14,500 (filtered before chat-template expansion) | +| Script | `benchmarks/bench_accuracy_baseline.py` | + +### 7.2 Service Startup + +```bash +# vLLM (GPU 1, port 8021) +CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 \ +vllm serve /data/zwt/model/models/Qwen/Qwen3-8B \ + --port 8021 \ + --no-enable-prefix-caching \ + --max-model-len 16384 \ + --max-num-seqs 8 \ + --gpu-memory-utilization 0.92 \ + --kv-transfer-config '{"kv_connector":"DaserConnector",...}' + +# DaseR HTTP server (port 2046) +CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 \ +python -m daser.server \ + --vllm-base-url http://127.0.0.1:8021 \ + --model-path /data/zwt/model/models/Qwen/Qwen3-8B \ + --store-dir /data/$USER/daser-baseline/store \ + --l2-size 30gb --l1-size 4gb \ + --transfer-mode iouring \ + --cache-reuse-mode chunk \ + --socket-path /tmp/daser_baseline_$USER.sock \ + --host 127.0.0.1 --port 2046 +``` + +### 7.3 Experiment Script + +```bash +# IMPORTANT: restart DaseR (and clean daser.store) before each dataset +# to prevent ring-buffer exhaustion from cross-dataset chunk accumulation. +for dataset in hotpotqa 2wikimqa musique; do + python benchmarks/bench_accuracy_baseline.py \ + --daser-url http://127.0.0.1:2046 \ + --tokenizer /data/zwt/model/models/Qwen/Qwen3-8B \ + --data-dir /data/$USER/daser-baseline/data \ + --datasets $dataset \ + --num-samples 20 \ + --max-context-tokens 14500 \ + --max-new-tokens 256 \ + --out /data/$USER/daser-baseline/results/${dataset}_v2.json +done +``` + +Each sample uploads every passage in the LongBench context as a separate +document (5–20 docs per sample) via `POST /documents`, then runs two +`POST /infer` requests with identical greedy parameters: + +- `full`: `use_kv_cache=false` — full prefill (accuracy upper bound). +- `reuse`: `use_kv_cache=true` — concatenated chunk KV with RoPE relocation, + no cross-attention repair (the lossy path). + +Documents are kept after each sample (script default, recorded as +`delete_docs=false` in the result JSON). Keeping them avoids delete/ring- +eviction interplay within a dataset run; cross-dataset accumulation is +handled by the per-dataset DaseR restart above. + +### 7.4 Results + +Accuracy means are computed only over samples with a verified chunk-KV load +(N below); the one HotpotQA sample whose reuse request silently fell back to +full recompute is excluded. + +| Dataset | N | Full F1 | Reuse F1 | **F1 Gap** | **Ext-F1 Gap** | Full Rec | Reuse Rec | **Rec Gap** | Agree | +|---------|---|---------|----------|------------|----------------|----------|-----------|-------------|-------| +| HotpotQA | 19/20 | **0.7338** | 0.1714 | **+0.5624** | +0.4486 | 0.6842 | 0.4211 | **+0.2632** | 5% | +| 2WikimQA | 20/20 | 0.2951 | 0.0665 | **+0.2286** | +0.2060 | 0.4500 | 0.4000 | **+0.0500** | 0% | +| MuSiQue | 16/16 | 0.2784 | 0.0138 | **+0.2646** | +0.2073 | 0.1875 | 0.0625 | **+0.1250** | 0% | + +Columns: + +- **F1 / Rec**: token-level QA F1 and gold-answer substring recall against + LongBench gold answers, on the raw answer text. +- **Ext-F1 Gap**: F1 gap after extracting the concise answer span (bold + ``**span**`` or first sentence) from both answers. The difference between + F1 Gap and Ext-F1 Gap is the portion attributable to verbosity-induced + token-F1 dilution rather than factual error. +- **Agree**: fraction of samples where the normalized full answer equals the + normalized reuse answer. Near-zero is expected given the verbosity contrast + (terse vs rambling); treat as a diagnostic, not an accuracy measure. +- **N**: samples with a TTFT-verified chunk-KV load / scored samples. + +### 7.5 Key Findings + +1. **Chunk reuse causes significant accuracy loss on every dataset, on every + metric.** The raw F1 gap is +0.23 to +0.56; the extracted-answer F1 gap — + which removes verbosity dilution — remains +0.21 to +0.45; gold-answer + recall drops on all three datasets (+0.05 to +0.26). The loss is real + factual degradation, not merely formatting drift. + +2. **The degradation has two distinguishable components.** + - *Factual errors*: wrong entities hallucinated from unrelated context + ("Hafsa Hatun", "Sven Nijdam"-style substitutions), captured by the + recall gap and the extracted-F1 gap. + - *Instruction-following collapse*: reuse answers average 320–850 + characters versus 29–58 for full prefill, frequently rambling, + self-contradicting ("Wait, I cannot determine…"), or failing to stop — + 14/16 MuSiQue reuse answers still hit the 256-token cap. This unbounded- + generation behavior is itself a symptom of attention-sink damage (§2.1) + and makes answer-length ratio a cheap regression indicator for repair + work. + +3. **HotpotQA suffers the worst (+0.56 F1, +0.26 Rec).** HotpotQA's + multi-hop reasoning requires aggregating evidence across passages, which + directly exercises the cross-chunk attention paths that chunk-mode reuse + breaks. This aligns with the §2 prediction that multi-doc QA is the + highest-impact scenario. With full-prefill F1 at 0.73, it also has the + best signal-to-noise ratio and should be the primary regression set for + §4 repair work. + +4. **Even tasks where the model already struggles are made worse.** + 2WikimQA and MuSiQue full-prefill F1 is only ~0.28–0.30 (Qwen3-8B at + 16k context is insufficient for these tasks), but chunk reuse pushes F1 + down to 0.01–0.07 and MuSiQue recall from 0.19 to 0.06. + +5. **Loaded ≈ 100% with coverage ≈ 99% confirms the gap is caused by KV + content differences, not cache misses.** The chunks are found, loaded + (reuse TTFT 112–208 ms vs full 800–1464 ms), and copied into the vLLM KV + cache. The accuracy loss comes from what is **inside** those chunks: + independently-prefilled KV with fake sinks at every chunk boundary. + +### 7.6 Measurement Pitfall: Ring-Buffer Exhaustion + +An initial run that processed all three datasets in a single process (without +restarting DaseR between them) reported zero F1 gap for 2WikimQA and MuSiQue. +Investigation revealed the cause: + +- The ring buffer (12,715 slots, L2=30 GiB) accumulated chunks across all + datasets. After the first dataset (HotpotQA, ~20 samples × ~10 docs × ~75 + chunks/doc ≈ 15,000 chunks), the buffer was full. +- By the time 2WikimQA and MuSiQue ran, older chunks had been evicted and + the vLLM external prefix cache hit rate had decayed from ~30% to ~9%. +- The reuse path fell back to full prefill, producing byte-identical output + and masking the true accuracy gap. + +**Remedy:** restart DaseR (or use `--skip-l2` for volatile L1-only runs) +between datasets so every sample in the reuse path exercises a real chunk-KV +load. The `Loaded` column in the report detects this condition: a dataset +with Loaded ≪ 100% should be re-run. Since the v2 run, `summarize()` also +excludes non-loaded samples from accuracy means automatically, so isolated +fallbacks no longer dilute the gap. + +### 7.7 Measurement Pitfall: Generation-Cap Truncation + +The initial run used `max_new_tokens=64` and reported a HotpotQA F1 gap of ++0.53. Per-sample inspection showed that degraded reuse answers are not just +wrong but **verbose** (4–5× the character length of full-prefill answers), +and 35–69% of reuse answers were cut mid-sentence at the 64-token cap. The +truncation depressed both reuse F1 (long answers dilute token precision) and +reuse recall (the gold span can be cut off before it appears), conflating +generation-cap artifacts with genuine KV damage. + +The v2 measurement therefore made three corrections to +`benchmarks/bench_accuracy_baseline.py`: + +1. `max_new_tokens` raised 64 → 256 (script default), removing most + truncation; only answers that genuinely fail to stop still hit the cap. +2. Accuracy means computed only over samples with a verified chunk-KV load + (`reuse_loaded`), reported as `samples_loaded`. +3. Added extracted-answer F1 (bold span / first sentence) to decompose the + raw F1 gap into factual error versus verbosity dilution, and recorded + `delete_docs` in the result JSON for reproducibility. + +The corrected numbers (§7.4) confirm the original conclusion — the v1 and v2 +gaps agree within a few points — while making the metric decomposition and +the per-sample validity explicit. Result artifacts are written to +`/data/$USER/daser-baseline/results/{hotpotqa,2wikimqa,musique}_v2.json`.