diff --git a/README.md b/README.md index 42f91ad..eb69793 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,8 @@ Since prefill is compute-bound, FLOP is a good proxy for latency. **Per token pe > Block hybrid shown for CUDA (B=16). Note, that full-attention cache and ssm cache can be combined; and for long realistic context size adding logarithmic hybrid cache to kvcache introduce deminishing cache size overhead while reducing computations by a margin. +> **TODO:** although default for vllm steps on cuda, `blocksize=16` seem to be suboptimal for prefix caching. Sweep this. + ### Huggingface Transformers prototype > Note: the code is a mess and neurosloppy, I will (hopefully) update it once @@ -80,22 +82,16 @@ The code can be found in `benchmark_baselines.py`. To evaluate caching strategies under realistic workloads, `benchmark_e2e.py` replays multi-turn conversations from the [tucnguyen/ShareGPT](https://huggingface.co/datasets/tucnguyen/ShareChat) dataset through a single Qwen3.5-0.8B layer group with random weights. -> **IMPORTANT:** a number of simplifications is applied in this tests. See `benchmark_e2e.py` for detailed discussion and real-world applicability of these numbers. +> **IMPORTANT:** a number of simplifications is applied in this tests. See `benchmark_e2e.py` for detailed discussion and real-world applicability of these numbers. -> **TODO:** these are results over 100 sequences (~500 requests) and 1G cache size. TODO: scale 10 times (should fit kaggle limits) -> **TODO:** add relative speedup boxplot +![ese_speedup](assets/e2e_time_vs_length.png) + +> Left panel — Prefill time vs context length. Here i have 40 dialog histories and send them in random order preserving in-conversation order. I measure prefill wall time, cache is stored at CPU. In total, there are 852 requests and 3 GB prefix cache budget (5M tokens overall). Right panel — per-request relative speedup. -``` +Speedup is possible because logarithmic checkpoints use $\sim O(\log L)$ memory per entry (vs $O(L/B)$ for block), so more conversations fit in the 3 GB budget simultaneously, yielding higher hit rates (80% cache hits vs 10% for block-boundary). + + +> **TODO:** add relative speedup boxplot +> **TODO:** there seem to be some memory overhead in transformers -- OOM happen much earlier that expected. Study this -Strategy Time (s) Speedup Hit rate GDN saved ------------------------------------------------------------- -no_cache 42.4 1.00x -attn_only 35.1 1.21x -block 40.0 1.06x 15.0% 18.6% -block_and_attn 39.6 1.07x 15.0% 18.6% -log 33.0 1.28x 57.6% 45.2% -log_and_attn 31.1 1.36x 54.2% 42.9% ---- -``` -> **TODO:** most entries in this dataset have short context lengths (<1K toks for the whole conversation) hence it does not highlight the asymptotic improvements. Filter to larger lengths, and plot acceleration VS length. diff --git a/assets/.DS_Store b/assets/.DS_Store index 8d33405..a6f9f82 100644 Binary files a/assets/.DS_Store and b/assets/.DS_Store differ diff --git a/assets/e2e_time_vs_length.png b/assets/e2e_time_vs_length.png new file mode 100644 index 0000000..4f97962 Binary files /dev/null and b/assets/e2e_time_vs_length.png differ diff --git a/simulation/README.md b/simulation/README.md new file mode 100644 index 0000000..2760acc --- /dev/null +++ b/simulation/README.md @@ -0,0 +1,152 @@ +# Simulation Code for Sparse Prefix Caching + +This directory contains the distributional and trace-driven simulators that +validate the theoretical results in the paper *Sparse Prefix Caching for +Hybrid and Recurrent LLM Serving*. + +## Directory structure + +``` +simulation/ +├── sparse_prefix_caching_sim.py # Distributional simulator (main entry point) +├── trace_cache_sim/ # Trace-driven simulator package +│ ├── __init__.py +│ ├── dp_opt.py # O(NM) DP solver with convex hull trick +│ ├── policies.py # Cache policies (uniform, geometric, DP, branch-aware, online) +│ ├── simulator.py # Trace-driven simulation engine +│ ├── traces.py # Trace generation & prefix trie utilities +│ ├── experiments.py # Experiment runner & figure generation +│ └── toy_recurrence.py # Toy recurrence for checkpoint-resume exactness check +├── tests/ +│ ├── __init__.py +│ └── test_trace_cache_sim.py # Unit tests +└── README.md # This file +``` + +## Requirements + +``` +numpy +matplotlib +seaborn +``` + +Install with: +```bash +pip install numpy matplotlib seaborn +``` + +## Distributional simulator + +The distributional simulator evaluates checkpoint placement strategies under +four synthetic overlap distributions (uniform, Zipf, bimodal, realistic proxy) +and validates theoretical bounds from the paper. + +**Run all experiments and generate figures:** + +```bash +python sparse_prefix_caching_sim.py +``` + +Options: +- `--prefix-length N` — maximum cached prefix length (default: 1024) +- `--budget M` — checkpoint budget for validation experiments (default: 32) +- `--samples S` — Monte Carlo samples per strategy/distribution (default: 100,000) + +**Experiments included:** +1. Strategy comparison on synthetic distributions (Table 1 in paper) +2. Pareto frontier: memory vs. compute trade-off +3. Savings heatmap across strategies and distributions +4. Realistic workload TTFT-proxy reduction at varying budgets +5. Online adaptation convergence (flat vs. exponential histogram) +6. Empirical histogram oracle convergence (validates Theorem 6) +7. Drift tracking under piecewise-stationary overlap shifts (validates Theorem 7) +8. Dynamic regret under drift (validates Theorem 8) +9. Full non-stationary sweep over γ × H × drift regime (validates Corollary 10) + +Output figures are written to a `figures/` directory next to the script. + +## Trace-driven simulator + +The trace-driven simulator evaluates caching policies on five synthetic request +families that cover qualitatively different sharing patterns: + +| Family | Description | +|--------|-------------| +| `exact-hot-prefix` | Many requests share one long prefix, diverge after it | +| `append-only-chat` | Each request extends a previously completed chat trace | +| `diffuse-cutpoints` | Requests share a long document but branch at varying positions | +| `agent-tree` | Branching workload with multiple depths | +| `adversarial-uniform-overlap` | Overlaps spread uniformly over a wide depth range | + +**Run the trace-driven experiments:** + +```bash +python -m trace_cache_sim.experiments +``` + +This runs all five trace families against all policies and generates: +- Per-family comparison tables +- Memory vs. token-hit-ratio, recompute, and TTFT figures +- Eviction sweep and SSM fraction sweep analyses +- Checkpoint-resume exactness verification (zero numerical error) + +**Policies implemented:** +- `NoCachePolicy` — baseline +- `DenseEveryKPolicy` — dense checkpoints every k positions +- `UniformBudgetMPolicy` — evenly spaced within budget +- `GeometricBaseRPolicy` / `GeometricEpsPolicy` — geometric schedules +- `BranchOnlyPolicy` / `BranchPlusGeometricPolicy` — branch-aware +- `OfflineDPOptimalPolicy` — offline O(NM) DP on empirical histogram +- `OnlineHistogramPolicy` — periodic replanning on flat histogram +- `ExponentialHistogramPolicy` — periodic replanning with exponential decay +- `BoundedCachePolicy` — LRU-bounded wrapper for finite cache capacity + +## Offline DP solver + +`trace_cache_sim/dp_opt.py` implements the O(NM) offline checkpoint placement +solver from Theorem 5 (the "monotone convex hull trick" DP). It can be used +standalone: + +```python +from trace_cache_sim.dp_opt import offline_optimal_checkpoints +import numpy as np + +# Example: bimodal distribution on N=1024 +N = 1024 +dist = np.zeros(N + 1) +dist[200:220] = 1.0 +dist[830:850] = 1.0 +dist /= dist.sum() + +checkpoints, cost = offline_optimal_checkpoints(N, budget=32, distribution=dist) +print(f"Optimal checkpoints: {checkpoints}") +print(f"Expected recompute cost: {cost:.4f}") +``` + +## Tests + +```bash +cd simulation +python -m pytest tests/ -v +# or +python -m unittest tests.test_trace_cache_sim -v +``` + +## Key theoretical results validated + +| Result | Simulator | +|--------|-----------| +| Balanced spacing optimal & minimax under uniform (Thm 1–2) | Distributional | +| Geometric schedule ≥ 1/r savings (Thm 3) | Distributional | +| Lower bounds Ω(N/M) (Thm 4) | Distributional | +| Stability / Lipschitz (Thm 5) | Both | +| Empirical oracle n^{-1/2} convergence (Thm 6) | Distributional | +| Exponential histogram drift tracking (Thm 7) | Distributional | +| Dynamic regret decomposition (Thm 8) | Distributional | +| Optimal γ* under bounded drift (Cor 10) | Distributional | +| Exact DP O(NM) solver (Thm 9) | Both | +| Checkpoint-resume exactness (Thm 10) | Trace-driven | +| Branch-aware policies for tree workloads | Trace-driven | +| LRU eviction robustness | Trace-driven | +| SSM fraction scaling | Trace-driven | diff --git a/simulation/sparse_prefix_caching_sim.py b/simulation/sparse_prefix_caching_sim.py new file mode 100644 index 0000000..9ed6f11 --- /dev/null +++ b/simulation/sparse_prefix_caching_sim.py @@ -0,0 +1,1082 @@ +#!/usr/bin/env python3 +"""Sparse prefix caching simulation and theoretical validation for SSM checkpoints.""" + +from __future__ import annotations + +import argparse +import math +import os +from collections import deque +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +import numpy as np + + +REPO_ROOT = Path(__file__).resolve().parent +os.environ.setdefault("MPLCONFIGDIR", str(REPO_ROOT / ".mplconfig")) +(REPO_ROOT / ".mplconfig").mkdir(exist_ok=True) + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import seaborn as sns + + +SEED = 42 +EPS = 1e-12 + + +@dataclass(frozen=True) +class StrategyResult: + name: str + checkpoints: list[int] + expected_cost: float + monte_carlo_cost: float + savings: float + relative_error: float + + +def dedupe_sorted_positions(positions: Iterable[int], prefix_length: int, budget: int) -> list[int]: + clipped = [min(prefix_length, max(1, int(p))) for p in positions] + unique = sorted(set(clipped)) + if budget <= 0: + return [] + if len(unique) <= budget: + return unique + # Preserve overall coverage when rounding produces too many candidates. + indices = np.linspace(0, len(unique) - 1, num=budget) + return sorted({unique[int(round(i))] for i in indices}) + + +class CheckpointStrategy: + """Base class for checkpoint placement strategies.""" + + name = "base" + + def place_checkpoints(self, prefix_length: int, budget: int) -> list[int]: + raise NotImplementedError + + +class UniformStrategy(CheckpointStrategy): + """Place checkpoints at evenly spaced positions.""" + + name = "uniform" + + def place_checkpoints(self, prefix_length: int, budget: int) -> list[int]: + if budget <= 0: + return [] + if budget >= prefix_length: + return list(range(1, prefix_length + 1)) + positions = np.linspace(1, prefix_length, num=budget + 2)[1:-1] + return dedupe_sorted_positions(np.rint(positions), prefix_length, budget) + + +class PowerOfTwoStrategy(CheckpointStrategy): + """Place checkpoints on a geometric grid; powers-of-two are recovered near log2(N) budget.""" + + name = "power_of_two" + + def place_checkpoints(self, prefix_length: int, budget: int) -> list[int]: + if budget <= 0: + return [] + if budget >= prefix_length: + return list(range(1, prefix_length + 1)) + if budget == 1: + return [1] + positions = np.geomspace(1, prefix_length, num=budget) + return dedupe_sorted_positions(np.rint(positions), prefix_length, budget) + + +class SqrtStrategy(CheckpointStrategy): + """Place checkpoints on the canonical sqrt(N)-spaced grid, capped or densified to match the budget.""" + + name = "sqrt" + + def place_checkpoints(self, prefix_length: int, budget: int) -> list[int]: + if budget <= 0: + return [] + if budget >= prefix_length: + return list(range(1, prefix_length + 1)) + step = max(1, int(round(math.sqrt(prefix_length)))) + natural = np.arange(step, prefix_length + 1, step) + if len(natural) >= budget: + indices = np.linspace(0, len(natural) - 1, num=budget) + return dedupe_sorted_positions(np.rint(natural[indices.astype(int)]), prefix_length, budget) + extra = np.linspace(1, prefix_length, num=budget) + positions = np.concatenate([natural, extra]) + return dedupe_sorted_positions(np.rint(positions), prefix_length, budget) + + +class ConvexHullTrick: + """Monotone lower hull for lines queried at nondecreasing x.""" + + def __init__(self) -> None: + self.lines: deque[tuple[float, float, int]] = deque() + + @staticmethod + def _value(line: tuple[float, float, int], x: float) -> float: + m, b, _ = line + return m * x + b + + @staticmethod + def _is_redundant( + l1: tuple[float, float, int], + l2: tuple[float, float, int], + l3: tuple[float, float, int], + ) -> bool: + m1, b1, _ = l1 + m2, b2, _ = l2 + m3, b3, _ = l3 + return (b3 - b1) * (m1 - m2) <= (b2 - b1) * (m1 - m3) + EPS + + def add_line(self, slope: float, intercept: float, arg: int) -> None: + line = (slope, intercept, arg) + while len(self.lines) >= 2 and self._is_redundant(self.lines[-2], self.lines[-1], line): + self.lines.pop() + self.lines.append(line) + + def query(self, x: float) -> tuple[float, int]: + while len(self.lines) >= 2 and self._value(self.lines[0], x) >= self._value(self.lines[1], x) - EPS: + self.lines.popleft() + best = self.lines[0] + return self._value(best, x), best[2] + + +class AdaptiveStrategy(CheckpointStrategy): + """Find optimal checkpoint positions for an empirical divergence distribution via exact DP.""" + + name = "adaptive" + + def __init__(self, divergence_distribution: np.ndarray): + distribution = np.asarray(divergence_distribution, dtype=float) + if distribution.ndim != 1 or distribution.size < 2: + raise ValueError("distribution must be a 1D array indexed from 0..N") + total = distribution[1:].sum() + if total <= 0: + raise ValueError("distribution mass over positions 1..N must be positive") + self.p = distribution.copy() + self.p[0] = 0.0 + self.p[1:] /= total + + def place_checkpoints(self, prefix_length: int, budget: int) -> list[int]: + if budget <= 0: + return [] + budget = min(budget, prefix_length) + if prefix_length != len(self.p) - 1: + raise ValueError("prefix_length must match the adaptive distribution length") + if budget >= prefix_length: + return list(range(1, prefix_length + 1)) + + prefix_prob = np.zeros(prefix_length + 1, dtype=float) + prefix_tp = np.zeros(prefix_length + 1, dtype=float) + positions = np.arange(prefix_length + 1, dtype=float) + prefix_prob[1:] = np.cumsum(self.p[1:]) + prefix_tp[1:] = np.cumsum(self.p[1:] * positions[1:]) + + dp_prev = prefix_tp.copy() + backpointers = np.zeros((budget + 1, prefix_length + 1), dtype=np.int32) + + for m in range(1, budget + 1): + dp_curr = np.zeros(prefix_length + 1, dtype=float) + dp_curr[0] = 0.0 + hull = ConvexHullTrick() + hull.add_line(-1.0, dp_prev[0], 1) + for j in range(1, prefix_length + 1): + s = j + intercept = dp_prev[s - 1] - prefix_tp[s - 1] + s * prefix_prob[s - 1] + hull.add_line(-float(s), float(intercept), s) + line_value, argmin_s = hull.query(prefix_prob[j]) + dp_curr[j] = prefix_tp[j] + line_value + backpointers[m, j] = argmin_s + dp_prev = dp_curr + + checkpoints: list[int] = [] + j = prefix_length + for m in range(budget, 0, -1): + s = int(backpointers[m, j]) + if s <= 0: + break + checkpoints.append(s) + j = s - 1 + return sorted(set(checkpoints)) + + +def normalize_distribution(probabilities: np.ndarray) -> np.ndarray: + distribution = np.asarray(probabilities, dtype=float) + distribution[0] = 0.0 + mass = distribution[1:].sum() + if mass <= 0: + raise ValueError("distribution has zero probability mass") + distribution[1:] /= mass + return distribution + + +def uniform_distribution(prefix_length: int) -> np.ndarray: + distribution = np.zeros(prefix_length + 1, dtype=float) + distribution[1:] = 1.0 / prefix_length + return distribution + + +def zipf_distribution(prefix_length: int, alpha: float = 1.2) -> np.ndarray: + positions = np.arange(1, prefix_length + 1, dtype=float) + distribution = np.zeros(prefix_length + 1, dtype=float) + distribution[1:] = 1.0 / np.power(positions, alpha) + return normalize_distribution(distribution) + + +def gaussian_bump(prefix_length: int, center: float, width: float) -> np.ndarray: + positions = np.arange(1, prefix_length + 1, dtype=float) + return np.exp(-0.5 * np.square((positions - center) / max(width, 1.0))) + + +def bimodal_distribution(prefix_length: int) -> np.ndarray: + early = gaussian_bump(prefix_length, 0.2 * prefix_length, 0.06 * prefix_length) + late = gaussian_bump(prefix_length, 0.82 * prefix_length, 0.05 * prefix_length) + distribution = np.zeros(prefix_length + 1, dtype=float) + distribution[1:] = 0.55 * early + 0.45 * late + return normalize_distribution(distribution) + + +def realistic_workload_distribution(prefix_length: int, shared_system_prompt: int = 500) -> np.ndarray: + positions = np.arange(1, prefix_length + 1, dtype=float) + shoulder = gaussian_bump(prefix_length, min(shared_system_prompt + 120, prefix_length), 90) + long_context = gaussian_bump(prefix_length, 0.7 * prefix_length, 0.12 * prefix_length) + tail = np.exp(-(prefix_length - positions) / max(prefix_length / 8.0, 1.0)) + distribution = np.zeros(prefix_length + 1, dtype=float) + distribution[1:] = 0.5 * shoulder + 0.3 * long_context + 0.2 * tail + return normalize_distribution(distribution) + + +def late_peak_distribution(prefix_length: int) -> np.ndarray: + distribution = np.zeros(prefix_length + 1, dtype=float) + distribution[1:] = gaussian_bump(prefix_length, 0.88 * prefix_length, 0.05 * prefix_length) + return normalize_distribution(distribution) + + +def expected_recomputation_cost(distribution: np.ndarray, checkpoints: list[int]) -> float: + prefix_length = len(distribution) - 1 + latest = np.zeros(prefix_length + 1, dtype=int) + if checkpoints: + checkpoints_array = np.asarray(sorted(set(checkpoints)), dtype=int) + latest[checkpoints_array] = checkpoints_array + latest = np.maximum.accumulate(latest) + positions = np.arange(prefix_length + 1, dtype=float) + costs = positions - latest + return float(np.dot(distribution[1:], costs[1:])) + + +def monte_carlo_cost( + distribution: np.ndarray, + checkpoints: list[int], + rng: np.random.Generator, + samples: int, +) -> float: + prefix_length = len(distribution) - 1 + positions = np.arange(prefix_length + 1) + draws = rng.choice(positions[1:], size=samples, p=distribution[1:]) + latest = np.zeros(prefix_length + 1, dtype=int) + if checkpoints: + checkpoints_array = np.asarray(sorted(set(checkpoints)), dtype=int) + latest[checkpoints_array] = checkpoints_array + latest = np.maximum.accumulate(latest) + costs = draws - latest[draws] + return float(np.mean(costs)) + + +def baseline_cost(distribution: np.ndarray) -> float: + positions = np.arange(len(distribution), dtype=float) + return float(np.dot(distribution[1:], positions[1:])) + + +def empirical_distribution_from_draws(draws: np.ndarray, prefix_length: int) -> np.ndarray: + counts = np.bincount(draws, minlength=prefix_length + 1).astype(float) + return normalize_distribution(counts) + + +def compute_savings(distribution: np.ndarray, checkpoints: list[int]) -> float: + base = baseline_cost(distribution) + return 1.0 - expected_recomputation_cost(distribution, checkpoints) / max(base, EPS) + + +def theoretical_uniform_savings(memory: np.ndarray) -> np.ndarray: + return memory / (memory + 1.0) + + +def evaluate_strategy( + strategy: CheckpointStrategy, + distribution: np.ndarray, + prefix_length: int, + budget: int, + rng: np.random.Generator, + samples: int, +) -> StrategyResult: + checkpoints = strategy.place_checkpoints(prefix_length, budget) + expected_cost = expected_recomputation_cost(distribution, checkpoints) + monte_carlo = monte_carlo_cost(distribution, checkpoints, rng, samples=samples) + base = baseline_cost(distribution) + savings = 1.0 - expected_cost / max(base, EPS) + relative_error = abs(monte_carlo - expected_cost) / max(expected_cost, EPS) + return StrategyResult( + name=strategy.name, + checkpoints=checkpoints, + expected_cost=expected_cost, + monte_carlo_cost=monte_carlo, + savings=savings, + relative_error=relative_error, + ) + + +def format_checkpoints(checkpoints: list[int], limit: int = 8) -> str: + if len(checkpoints) <= limit: + return str(checkpoints) + head = ", ".join(str(x) for x in checkpoints[: limit // 2]) + tail = ", ".join(str(x) for x in checkpoints[-(limit // 2) :]) + return f"[{head}, ..., {tail}]" + + +def latex_results_table(rows: list[dict[str, object]]) -> str: + lines = [ + r"\begin{tabular}{llrrrr}", + r"\toprule", + r"Distribution & Strategy & Budget & Used & Savings (\%) & RelErr (\%) \\", + r"\midrule", + ] + for row in rows: + lines.append( + f"{row['distribution']} & {row['strategy']} & {row['budget']} & {row['used']} & " + f"{100.0 * row['savings']:.2f} & {100.0 * row['relative_error']:.2f} \\\\" + ) + lines.extend([r"\bottomrule", r"\end{tabular}"]) + return "\n".join(lines) + + +def exact_match_savings(distribution: np.ndarray) -> float: + return float(distribution[-1]) + + +def plot_pareto_frontier( + figure_dir: Path, + prefix_length: int, + budgets: np.ndarray, + strategies: list[CheckpointStrategy], + distribution: np.ndarray, +) -> list[dict[str, float | str]]: + records: list[dict[str, float | str]] = [] + for strategy in strategies: + for budget in budgets: + checkpoints = strategy.place_checkpoints(prefix_length, int(budget)) + records.append( + { + "strategy": strategy.name, + "budget": int(budget), + "used_memory": len(checkpoints), + "savings": 100.0 * compute_savings(distribution, checkpoints), + } + ) + + plt.figure(figsize=(8.5, 5.4)) + palette = sns.color_palette("Set2", n_colors=len(strategies)) + for color, strategy in zip(palette, strategies): + xs = [r["used_memory"] for r in records if r["strategy"] == strategy.name] + ys = [r["savings"] for r in records if r["strategy"] == strategy.name] + plt.plot(xs, ys, marker="o", linewidth=2, markersize=4, label=strategy.name, color=color) + + theoretical_x = np.unique(np.array([r["used_memory"] for r in records if r["used_memory"] > 0], dtype=float)) + plt.plot( + theoretical_x, + 100.0 * theoretical_uniform_savings(theoretical_x), + linestyle="--", + linewidth=2, + color="#333333", + label="uniform theory", + ) + plt.axhline(50.0, linestyle="--", linewidth=1.5, color="#9b2226", label="power-of-two 50% bound") + plt.scatter([prefix_length], [100.0], color="#000000", marker="*", s=100, label="full cache") + plt.xscale("log") + plt.xlabel("Checkpoint slots used") + plt.ylabel("Compute savings (%)") + plt.title("Pareto frontier under uniform divergence") + plt.ylim(0, 103) + plt.legend(frameon=False, ncol=2) + plt.tight_layout() + plt.savefig(figure_dir / "pareto_frontier_uniform.pdf") + plt.close() + return records + + +def plot_strategy_heatmap( + figure_dir: Path, + prefix_length: int, + budget: int, + strategies: list[CheckpointStrategy], + distributions: dict[str, np.ndarray], +) -> dict[str, dict[str, float]]: + matrix: dict[str, dict[str, float]] = {strategy.name: {} for strategy in strategies} + for strategy in strategies: + for dist_name, distribution in distributions.items(): + checkpoints = strategy.place_checkpoints(prefix_length, budget) + score = compute_savings(distribution, checkpoints) / max(len(checkpoints), 1) + matrix[strategy.name][dist_name] = 100.0 * score + + heatmap_data = np.array( + [[matrix[strategy.name][dist_name] for dist_name in distributions] for strategy in strategies], + dtype=float, + ) + plt.figure(figsize=(7.4, 4.2)) + sns.heatmap( + heatmap_data, + annot=True, + fmt=".2f", + cmap="YlGnBu", + xticklabels=list(distributions.keys()), + yticklabels=[strategy.name for strategy in strategies], + cbar_kws={"label": "Savings per checkpoint slot (%)"}, + ) + plt.title(f"Strategy efficiency at budget M={budget}") + plt.tight_layout() + plt.savefig(figure_dir / "strategy_heatmap.pdf") + plt.close() + return matrix + + +def simulate_online_adaptation( + true_distribution: np.ndarray, + prefix_length: int, + budget: int, + horizon: int, + update_every: int, + trials: int, + rng_seed: int, + decay: float = 1.0, +) -> tuple[np.ndarray, np.ndarray]: + oracle_checkpoints = AdaptiveStrategy(true_distribution).place_checkpoints(prefix_length, budget) + oracle_cost = np.arange(prefix_length + 1) - np.maximum.accumulate( + np.bincount(oracle_checkpoints, weights=oracle_checkpoints, minlength=prefix_length + 1).astype(int) + ) + all_regrets = np.zeros((trials, horizon), dtype=float) + + for trial in range(trials): + rng = np.random.default_rng(rng_seed + trial) + empirical_counts = np.ones(prefix_length + 1, dtype=float) + current_strategy: CheckpointStrategy = PowerOfTwoStrategy() + current_checkpoints = current_strategy.place_checkpoints(prefix_length, budget) + latest_current = np.zeros(prefix_length + 1, dtype=int) + latest_current[current_checkpoints] = current_checkpoints + latest_current = np.maximum.accumulate(latest_current) + + latest_oracle = np.zeros(prefix_length + 1, dtype=int) + latest_oracle[oracle_checkpoints] = oracle_checkpoints + latest_oracle = np.maximum.accumulate(latest_oracle) + + cumulative = 0.0 + for t in range(horizon): + draw = int(rng.choice(np.arange(1, prefix_length + 1), p=true_distribution[1:])) + observed_cost = draw - latest_current[draw] + best_cost = draw - latest_oracle[draw] + cumulative += observed_cost - best_cost + all_regrets[trial, t] = cumulative + if decay < 1.0: + empirical_counts *= decay + empirical_counts[draw] += 1.0 + + if (t + 1) % update_every == 0: + empirical_distribution = normalize_distribution(empirical_counts.copy()) + current_checkpoints = AdaptiveStrategy(empirical_distribution).place_checkpoints(prefix_length, budget) + latest_current = np.zeros(prefix_length + 1, dtype=int) + latest_current[current_checkpoints] = current_checkpoints + latest_current = np.maximum.accumulate(latest_current) + + return all_regrets.mean(axis=0), all_regrets.std(axis=0) + + +def plot_online_adaptation( + figure_dir: Path, + distribution: np.ndarray, + prefix_length: int, + budget: int, +) -> tuple[float, float, float, float]: + common = dict( + true_distribution=distribution, + prefix_length=prefix_length, + budget=budget, + horizon=2000, + update_every=50, + trials=24, + rng_seed=SEED, + ) + mean_flat, std_flat = simulate_online_adaptation(**common, decay=1.0) + mean_exp, std_exp = simulate_online_adaptation(**common, decay=0.99) + + xs = np.arange(1, len(mean_flat) + 1) + plt.figure(figsize=(8.2, 4.8)) + plt.plot(xs, mean_flat, linewidth=2.2, color="#005f73", label="flat histogram") + plt.fill_between(xs, mean_flat - std_flat, mean_flat + std_flat, color="#94d2bd", alpha=0.25) + plt.plot(xs, mean_exp, linewidth=2.2, color="#ae2012", label="exponential (\u03b3=0.99)") + plt.fill_between(xs, mean_exp - std_exp, mean_exp + std_exp, color="#ee9b00", alpha=0.25) + plt.xlabel("Requests processed") + plt.ylabel("Cumulative regret") + plt.title("Online adaptation: exponential weighting reduces regret") + plt.legend(frameon=False) + plt.tight_layout() + plt.savefig(figure_dir / "online_adaptation_convergence.pdf") + plt.close() + return float(mean_flat[-1]), float(std_flat[-1]), float(mean_exp[-1]), float(std_exp[-1]) + + +def plot_empirical_oracle_convergence( + figure_dir: Path, + distribution: np.ndarray, + prefix_length: int, + budget: int, + sample_sizes: list[int] | None = None, + trials: int = 40, + rng_seed: int = SEED, +) -> dict[str, object]: + if sample_sizes is None: + sample_sizes = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096] + + oracle_checkpoints = AdaptiveStrategy(distribution).place_checkpoints(prefix_length, budget) + oracle_cost = expected_recomputation_cost(distribution, oracle_checkpoints) + support = np.arange(1, prefix_length + 1) + + mean_l1: list[float] = [] + std_l1: list[float] = [] + mean_gap: list[float] = [] + std_gap: list[float] = [] + + for size_index, sample_size in enumerate(sample_sizes): + l1_errors = np.zeros(trials, dtype=float) + suboptimality_gaps = np.zeros(trials, dtype=float) + for trial in range(trials): + rng = np.random.default_rng(rng_seed + 10_000 * size_index + trial) + draws = rng.choice(support, size=sample_size, p=distribution[1:]) + empirical_distribution = empirical_distribution_from_draws(draws, prefix_length) + l1_errors[trial] = float(np.abs(empirical_distribution[1:] - distribution[1:]).sum()) + empirical_checkpoints = AdaptiveStrategy(empirical_distribution).place_checkpoints(prefix_length, budget) + plugin_cost = expected_recomputation_cost(distribution, empirical_checkpoints) + suboptimality_gaps[trial] = max(0.0, plugin_cost - oracle_cost) + + mean_l1.append(float(l1_errors.mean())) + std_l1.append(float(l1_errors.std())) + mean_gap.append(float(suboptimality_gaps.mean())) + std_gap.append(float(suboptimality_gaps.std())) + + xs = np.asarray(sample_sizes, dtype=float) + mean_l1_arr = np.asarray(mean_l1, dtype=float) + std_l1_arr = np.asarray(std_l1, dtype=float) + mean_gap_arr = np.asarray(mean_gap, dtype=float) + std_gap_arr = np.asarray(std_gap, dtype=float) + reference_l1 = mean_l1_arr[0] * np.sqrt(xs[0] / xs) + reference_gap = max(mean_gap_arr[0], EPS) * np.sqrt(xs[0] / xs) + + fig, axes = plt.subplots(1, 2, figsize=(10.2, 4.4)) + + axes[0].plot(xs, mean_l1_arr, marker="o", linewidth=2.2, color="#005f73", label=r"mean $\|\hat p-p\|_1$") + axes[0].fill_between( + xs, + np.maximum(mean_l1_arr - std_l1_arr, EPS), + mean_l1_arr + std_l1_arr, + color="#94d2bd", + alpha=0.25, + ) + axes[0].plot(xs, reference_l1, linestyle="--", linewidth=1.8, color="#9b2226", label=r"$n^{-1/2}$ reference") + axes[0].set_xscale("log", base=2) + axes[0].set_yscale("log") + axes[0].set_xlabel("Observed overlaps") + axes[0].set_ylabel(r"$\ell_1$ estimation error") + axes[0].set_title("Empirical law estimation") + axes[0].legend(frameon=False) + + axes[1].plot(xs, mean_gap_arr, marker="o", linewidth=2.2, color="#bb3e03", label="mean plug-in gap") + axes[1].fill_between( + xs, + np.maximum(mean_gap_arr - std_gap_arr, EPS), + mean_gap_arr + std_gap_arr, + color="#ee9b00", + alpha=0.25, + ) + axes[1].plot(xs, reference_gap, linestyle="--", linewidth=1.8, color="#0a9396", label=r"$n^{-1/2}$ reference") + axes[1].set_xscale("log", base=2) + axes[1].set_yscale("log") + axes[1].set_xlabel("Observed overlaps") + axes[1].set_ylabel("Suboptimality gap") + axes[1].set_title("Empirical oracle under true law") + axes[1].legend(frameon=False) + + fig.suptitle("Stationary sample complexity of the empirical histogram oracle", y=1.02) + fig.tight_layout() + fig.savefig(figure_dir / "empirical_oracle_convergence.pdf") + plt.close(fig) + + return { + "sample_sizes": sample_sizes, + "mean_l1": mean_l1, + "std_l1": std_l1, + "mean_gap": mean_gap, + "std_gap": std_gap, + } + + +def simulate_drift_tracking( + schedule: list[np.ndarray], + prefix_length: int, + budget: int, + update_every: int, + trials: int, + rng_seed: int, + decay: float, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + horizon = len(schedule) + support = np.arange(1, prefix_length + 1) + + oracle_cache: dict[int, float] = {} + for distribution in schedule: + key = id(distribution) + if key not in oracle_cache: + oracle_checkpoints = AdaptiveStrategy(distribution).place_checkpoints(prefix_length, budget) + oracle_cache[key] = expected_recomputation_cost(distribution, oracle_checkpoints) + + l1_curves = np.zeros((trials, horizon), dtype=float) + gap_curves = np.zeros((trials, horizon), dtype=float) + + for trial in range(trials): + rng = np.random.default_rng(rng_seed + trial) + counts = np.ones(prefix_length + 1, dtype=float) + current_l1 = 0.0 + current_gap = 0.0 + + for t, distribution in enumerate(schedule): + draw = int(rng.choice(support, p=distribution[1:])) + if decay < 1.0: + counts *= decay + counts[draw] += 1.0 + + if t == 0 or (t + 1) % update_every == 0: + empirical_distribution = normalize_distribution(counts.copy()) + checkpoints = AdaptiveStrategy(empirical_distribution).place_checkpoints(prefix_length, budget) + oracle_cost = oracle_cache[id(distribution)] + current_l1 = float(np.abs(empirical_distribution[1:] - distribution[1:]).sum()) + current_gap = max(0.0, expected_recomputation_cost(distribution, checkpoints) - oracle_cost) + + l1_curves[trial, t] = current_l1 + gap_curves[trial, t] = current_gap + + return ( + l1_curves.mean(axis=0), + l1_curves.std(axis=0), + gap_curves.mean(axis=0), + gap_curves.std(axis=0), + ) + + +def plot_drift_tracking( + figure_dir: Path, + prefix_length: int, + budget: int, + update_every: int = 25, + segment_length: int = 800, + trials: int = 20, +) -> dict[str, float | list[int]]: + phases = [ + ("zipf", zipf_distribution(prefix_length, alpha=1.25)), + ("bimodal", bimodal_distribution(prefix_length)), + ("late", late_peak_distribution(prefix_length)), + ] + schedule = [distribution for _, distribution in phases for _ in range(segment_length)] + horizon = len(schedule) + change_points = [segment_length, 2 * segment_length] + + flat_mean_l1, flat_std_l1, flat_mean_gap, flat_std_gap = simulate_drift_tracking( + schedule=schedule, + prefix_length=prefix_length, + budget=budget, + update_every=update_every, + trials=trials, + rng_seed=SEED + 2000, + decay=1.0, + ) + exp_mean_l1, exp_std_l1, exp_mean_gap, exp_std_gap = simulate_drift_tracking( + schedule=schedule, + prefix_length=prefix_length, + budget=budget, + update_every=update_every, + trials=trials, + rng_seed=SEED + 4000, + decay=0.99, + ) + + xs = np.arange(1, horizon + 1) + fig, axes = plt.subplots(2, 1, figsize=(10.2, 6.8), sharex=True) + + axes[0].plot(xs, flat_mean_l1, linewidth=2.0, color="#005f73", label="flat histogram") + axes[0].fill_between(xs, flat_mean_l1 - flat_std_l1, flat_mean_l1 + flat_std_l1, color="#94d2bd", alpha=0.22) + axes[0].plot(xs, exp_mean_l1, linewidth=2.0, color="#ae2012", label="exponential ($\\gamma=0.99$)") + axes[0].fill_between(xs, exp_mean_l1 - exp_std_l1, exp_mean_l1 + exp_std_l1, color="#ee9b00", alpha=0.18) + axes[0].set_ylabel(r"Current-law $\ell_1$ error") + axes[0].set_title("Tracking a drifting overlap law") + axes[0].legend(frameon=False, ncol=2) + + axes[1].plot(xs, flat_mean_gap, linewidth=2.0, color="#005f73", label="flat histogram") + axes[1].fill_between(xs, np.maximum(flat_mean_gap - flat_std_gap, 0.0), flat_mean_gap + flat_std_gap, color="#94d2bd", alpha=0.22) + axes[1].plot(xs, exp_mean_gap, linewidth=2.0, color="#ae2012", label="exponential ($\\gamma=0.99$)") + axes[1].fill_between(xs, np.maximum(exp_mean_gap - exp_std_gap, 0.0), exp_mean_gap + exp_std_gap, color="#ee9b00", alpha=0.18) + axes[1].set_xlabel("Requests processed") + axes[1].set_ylabel("Current-law plug-in gap") + + for ax in axes: + for change_point in change_points: + ax.axvline(change_point, linestyle="--", linewidth=1.2, color="#6c757d", alpha=0.9) + + fig.tight_layout() + fig.savefig(figure_dir / "drift_tracking_tradeoff.pdf") + plt.close(fig) + + transition_window = 150 + post_shift_mask = np.zeros(horizon, dtype=bool) + for change_point in change_points: + start = change_point + stop = min(horizon, change_point + transition_window) + post_shift_mask[start:stop] = True + + return { + "flat_mean_l1": float(flat_mean_l1.mean()), + "exp_mean_l1": float(exp_mean_l1.mean()), + "flat_mean_gap": float(flat_mean_gap.mean()), + "exp_mean_gap": float(exp_mean_gap.mean()), + "flat_post_shift_gap": float(flat_mean_gap[post_shift_mask].mean()), + "exp_post_shift_gap": float(exp_mean_gap[post_shift_mask].mean()), + "change_points": change_points, + } + + +def simulate_nonstationary_dynamic_regret( + schedule: list[np.ndarray], + prefix_length: int, + budget: int, + update_every: int, + trials: int, + rng_seed: int, + decay: float, +) -> tuple[np.ndarray, np.ndarray]: + horizon = len(schedule) + support = np.arange(1, prefix_length + 1) + oracle_costs = np.zeros(horizon, dtype=float) + oracle_cache: dict[int, float] = {} + + for t, distribution in enumerate(schedule): + key = id(distribution) + if key not in oracle_cache: + oracle_checkpoints = AdaptiveStrategy(distribution).place_checkpoints(prefix_length, budget) + oracle_cache[key] = expected_recomputation_cost(distribution, oracle_checkpoints) + oracle_costs[t] = oracle_cache[key] + + regret_curves = np.zeros((trials, horizon), dtype=float) + for trial in range(trials): + rng = np.random.default_rng(rng_seed + trial) + counts = np.ones(prefix_length + 1, dtype=float) + current_checkpoints = PowerOfTwoStrategy().place_checkpoints(prefix_length, budget) + cumulative = 0.0 + + for t, distribution in enumerate(schedule): + draw = int(rng.choice(support, p=distribution[1:])) + if decay < 1.0: + counts *= decay + counts[draw] += 1.0 + + if t == 0 or (t + 1) % update_every == 0: + empirical_distribution = normalize_distribution(counts.copy()) + current_checkpoints = AdaptiveStrategy(empirical_distribution).place_checkpoints(prefix_length, budget) + + cumulative += expected_recomputation_cost(distribution, current_checkpoints) - oracle_costs[t] + regret_curves[trial, t] = cumulative + + return regret_curves.mean(axis=0), regret_curves.std(axis=0) + + +def plot_nonstationary_dynamic_regret( + figure_dir: Path, + prefix_length: int, + budget: int, + update_every: int = 25, + segment_length: int = 600, + trials: int = 12, +) -> dict[str, float | list[int]]: + phases = [ + ("zipf", zipf_distribution(prefix_length, alpha=1.25)), + ("bimodal", bimodal_distribution(prefix_length)), + ("late", late_peak_distribution(prefix_length)), + ] + schedule = [distribution for _, distribution in phases for _ in range(segment_length)] + horizon = len(schedule) + change_points = [segment_length, 2 * segment_length] + + flat_mean, flat_std = simulate_nonstationary_dynamic_regret( + schedule=schedule, + prefix_length=prefix_length, + budget=budget, + update_every=update_every, + trials=trials, + rng_seed=SEED + 6000, + decay=1.0, + ) + exp_mean, exp_std = simulate_nonstationary_dynamic_regret( + schedule=schedule, + prefix_length=prefix_length, + budget=budget, + update_every=update_every, + trials=trials, + rng_seed=SEED + 8000, + decay=0.99, + ) + + xs = np.arange(1, horizon + 1) + plt.figure(figsize=(10.0, 4.8)) + plt.plot(xs, flat_mean, linewidth=2.2, color="#005f73", label="flat histogram") + plt.fill_between(xs, np.maximum(flat_mean - flat_std, 0.0), flat_mean + flat_std, color="#94d2bd", alpha=0.22) + plt.plot(xs, exp_mean, linewidth=2.2, color="#ae2012", label="exponential ($\\gamma=0.99$)") + plt.fill_between(xs, np.maximum(exp_mean - exp_std, 0.0), exp_mean + exp_std, color="#ee9b00", alpha=0.18) + for change_point in change_points: + plt.axvline(change_point, linestyle="--", linewidth=1.2, color="#6c757d", alpha=0.9) + plt.xlabel("Requests processed") + plt.ylabel("Cumulative dynamic regret") + plt.title("Dynamic regret under piecewise-stationary overlap drift") + plt.legend(frameon=False, ncol=2) + plt.tight_layout() + plt.savefig(figure_dir / "drift_dynamic_regret.pdf") + plt.close() + + return { + "flat_final_regret": float(flat_mean[-1]), + "exp_final_regret": float(exp_mean[-1]), + "improvement": float(1.0 - exp_mean[-1] / max(flat_mean[-1], EPS)), + "change_points": change_points, + } + + +def plot_realistic_workload_bars( + figure_dir: Path, + prefix_length: int, + budgets: list[int], + strategies: list[CheckpointStrategy], + distribution: np.ndarray, +) -> dict[int, dict[str, float]]: + data: dict[int, dict[str, float]] = {} + for budget in budgets: + data[budget] = {} + for strategy in strategies: + checkpoints = strategy.place_checkpoints(prefix_length, budget) + data[budget][strategy.name] = 100.0 * compute_savings(distribution, checkpoints) + data[budget]["full_cache"] = 100.0 + data[budget]["exact_match"] = 100.0 * exact_match_savings(distribution) + + labels = [strategy.name for strategy in strategies] + ["full_cache", "exact_match"] + x = np.arange(len(labels)) + width = 0.22 + plt.figure(figsize=(9.2, 5.2)) + colors = ["#0a9396", "#ee9b00", "#ca6702"] + for idx, budget in enumerate(budgets): + heights = [data[budget][label] for label in labels] + plt.bar(x + (idx - 1) * width, heights, width=width, label=f"M={budget}", color=colors[idx]) + plt.xticks(x, labels) + plt.ylabel("TTFT reduction proxy (%)") + plt.title("Realistic workload comparison at fixed memory budgets") + plt.ylim(0, 103) + plt.legend(frameon=False) + plt.tight_layout() + plt.savefig(figure_dir / "realistic_workload_bars.pdf") + plt.close() + return data + + +def print_experiment_summary( + results: list[dict[str, object]], + realistic_data: dict[int, dict[str, float]], + final_regret: float, + final_regret_std: float, + final_regret_exp: float | None = None, + final_regret_exp_std: float | None = None, + empirical_oracle_stats: dict[str, object] | None = None, + drift_tracking_stats: dict[str, float | list[int]] | None = None, + dynamic_regret_stats: dict[str, float | list[int]] | None = None, +) -> None: + print("Sparse Prefix Caching for SSM/Hybrid Models") + print("=" * 48) + for row in results: + print( + f"[{row['distribution']}] {row['strategy']:>12} | budget={row['budget']:>3} | " + f"used={row['used']:>3} | savings={100.0 * row['savings']:.2f}% | " + f"expected={row['expected_cost']:.3f} | mc={row['monte_carlo_cost']:.3f} | " + f"relerr={100.0 * row['relative_error']:.2f}% | checkpoints={row['checkpoints']}" + ) + print() + print("Realistic workload TTFT proxy reduction:") + for budget, metrics in realistic_data.items(): + ordered = ", ".join(f"{name}={value:.2f}%" for name, value in metrics.items()) + print(f" M={budget}: {ordered}") + print() + print(f"Online adaptation final cumulative regret (flat): {final_regret:.2f} +/- {final_regret_std:.2f}") + if final_regret_exp is not None: + print(f"Online adaptation final cumulative regret (exp): {final_regret_exp:.2f} +/- {final_regret_exp_std:.2f}") + if empirical_oracle_stats is not None: + sample_sizes = empirical_oracle_stats["sample_sizes"] + mean_l1 = empirical_oracle_stats["mean_l1"] + mean_gap = empirical_oracle_stats["mean_gap"] + print( + "Empirical oracle convergence:" + f" n={sample_sizes[0]} -> l1={mean_l1[0]:.4f}, gap={mean_gap[0]:.4f};" + f" n={sample_sizes[-1]} -> l1={mean_l1[-1]:.4f}, gap={mean_gap[-1]:.4f}" + ) + if drift_tracking_stats is not None: + print( + "Drift tracking:" + f" flat gap={drift_tracking_stats['flat_mean_gap']:.3f}," + f" exp gap={drift_tracking_stats['exp_mean_gap']:.3f};" + f" post-shift flat={drift_tracking_stats['flat_post_shift_gap']:.3f}," + f" post-shift exp={drift_tracking_stats['exp_post_shift_gap']:.3f}" + ) + if dynamic_regret_stats is not None: + print( + "Dynamic regret under drift:" + f" flat={dynamic_regret_stats['flat_final_regret']:.1f}," + f" exp={dynamic_regret_stats['exp_final_regret']:.1f}," + f" improvement={100.0 * dynamic_regret_stats['improvement']:.1f}%" + ) + print() + print("LaTeX table:") + print(latex_results_table(results)) + + +def run_experiments(prefix_length: int, validation_budget: int, monte_carlo_samples: int) -> None: + sns.set_theme(style="whitegrid", context="talk") + figure_dir = REPO_ROOT / "figures" + figure_dir.mkdir(exist_ok=True) + + rng = np.random.default_rng(SEED) + distributions = { + "uniform": uniform_distribution(prefix_length), + "zipf": zipf_distribution(prefix_length), + "bimodal": bimodal_distribution(prefix_length), + } + realistic_distribution = realistic_workload_distribution(prefix_length) + + fixed_strategies: list[CheckpointStrategy] = [ + UniformStrategy(), + PowerOfTwoStrategy(), + SqrtStrategy(), + ] + results: list[dict[str, object]] = [] + + for dist_name, distribution in distributions.items(): + strategy_set = fixed_strategies + [AdaptiveStrategy(distribution)] + for strategy in strategy_set: + outcome = evaluate_strategy( + strategy=strategy, + distribution=distribution, + prefix_length=prefix_length, + budget=validation_budget, + rng=rng, + samples=monte_carlo_samples, + ) + results.append( + { + "distribution": dist_name, + "strategy": strategy.name, + "budget": validation_budget, + "used": len(outcome.checkpoints), + "savings": outcome.savings, + "expected_cost": outcome.expected_cost, + "monte_carlo_cost": outcome.monte_carlo_cost, + "relative_error": outcome.relative_error, + "checkpoints": format_checkpoints(outcome.checkpoints), + } + ) + + pareto_budgets = np.unique(np.rint(np.geomspace(1, prefix_length, num=18)).astype(int)) + plot_pareto_frontier( + figure_dir=figure_dir, + prefix_length=prefix_length, + budgets=pareto_budgets, + strategies=fixed_strategies + [AdaptiveStrategy(distributions["uniform"])], + distribution=distributions["uniform"], + ) + plot_strategy_heatmap( + figure_dir=figure_dir, + prefix_length=prefix_length, + budget=validation_budget, + strategies=fixed_strategies + [AdaptiveStrategy(distributions["bimodal"])], + distributions=distributions, + ) + final_regret, final_regret_std, final_regret_exp, final_regret_exp_std = plot_online_adaptation( + figure_dir=figure_dir, + distribution=distributions["bimodal"], + prefix_length=prefix_length, + budget=validation_budget, + ) + empirical_oracle_stats = plot_empirical_oracle_convergence( + figure_dir=figure_dir, + distribution=distributions["bimodal"], + prefix_length=prefix_length, + budget=validation_budget, + ) + drift_tracking_stats = plot_drift_tracking( + figure_dir=figure_dir, + prefix_length=prefix_length, + budget=validation_budget, + ) + dynamic_regret_stats = plot_nonstationary_dynamic_regret( + figure_dir=figure_dir, + prefix_length=prefix_length, + budget=validation_budget, + ) + realistic_data = plot_realistic_workload_bars( + figure_dir=figure_dir, + prefix_length=prefix_length, + budgets=[10, 50, 100], + strategies=fixed_strategies + [AdaptiveStrategy(realistic_distribution)], + distribution=realistic_distribution, + ) + print_experiment_summary( + results, + realistic_data, + final_regret, + final_regret_std, + final_regret_exp, + final_regret_exp_std, + empirical_oracle_stats, + drift_tracking_stats, + dynamic_regret_stats, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--prefix-length", type=int, default=1024, help="Maximum cached prefix length N.") + parser.add_argument("--budget", type=int, default=32, help="Checkpoint budget M for validation experiments.") + parser.add_argument( + "--samples", + type=int, + default=100_000, + help="Monte Carlo samples per strategy/distribution pair.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.prefix_length < 8: + raise ValueError("prefix length must be at least 8") + if args.budget < 1: + raise ValueError("budget must be positive") + run_experiments( + prefix_length=args.prefix_length, + validation_budget=args.budget, + monte_carlo_samples=args.samples, + ) + + +if __name__ == "__main__": + main() diff --git a/simulation/tests/__init__.py b/simulation/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simulation/tests/test_trace_cache_sim.py b/simulation/tests/test_trace_cache_sim.py new file mode 100644 index 0000000..b14402c --- /dev/null +++ b/simulation/tests/test_trace_cache_sim.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import unittest + +from trace_cache_sim.policies import deepest_cached_depth, geometric_base_lengths, uniform_lengths +from trace_cache_sim.toy_recurrence import checkpoint_resume_error +from trace_cache_sim.traces import longest_common_prefix + + +class TraceCacheSimTests(unittest.TestCase): + def test_longest_common_prefix(self) -> None: + self.assertEqual(longest_common_prefix((1, 2, 3), (1, 2, 4)), 2) + self.assertEqual(longest_common_prefix((1, 2), (3, 4)), 0) + self.assertEqual(longest_common_prefix((1, 2, 3), (1, 2, 3)), 3) + + def test_geometric_guarantee_sanity(self) -> None: + checkpoints = geometric_base_lengths(128, base=2.0) + prefixes = frozenset(tuple(range(length)) for length in checkpoints) + for length in range(1, 129): + tokens = tuple(range(length)) + cached = deepest_cached_depth(tokens, prefixes) + gap = length - cached + self.assertLessEqual(gap, length / 2 + 1) + + def test_uniform_additive_gap_sanity(self) -> None: + schedule = uniform_lengths(120, 12) + prefixes = frozenset(tuple(range(length)) for length in schedule) + max_gap = 0 + for length in range(1, 121): + tokens = tuple(range(length)) + cached = deepest_cached_depth(tokens, prefixes) + max_gap = max(max_gap, length - cached) + self.assertLessEqual(max_gap, 12) + + def test_toy_recurrence_exactness(self) -> None: + tokens = tuple(range(1, 40)) + error = checkpoint_resume_error(tokens, checkpoint_depth=17) + self.assertLess(error, 1e-10) + + +if __name__ == "__main__": + unittest.main() diff --git a/simulation/trace_cache_sim/__init__.py b/simulation/trace_cache_sim/__init__.py new file mode 100644 index 0000000..7b6a9da --- /dev/null +++ b/simulation/trace_cache_sim/__init__.py @@ -0,0 +1,2 @@ +"""Trace-driven sparse prefix caching simulator.""" + diff --git a/simulation/trace_cache_sim/dp_opt.py b/simulation/trace_cache_sim/dp_opt.py new file mode 100644 index 0000000..cd54bde --- /dev/null +++ b/simulation/trace_cache_sim/dp_opt.py @@ -0,0 +1,111 @@ +"""Offline dynamic programming for sparse checkpoint placement.""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass + +import numpy as np + + +EPS = 1e-12 + + +def normalize_distribution(distribution: np.ndarray) -> np.ndarray: + dist = np.asarray(distribution, dtype=float).copy() + dist[0] = 0.0 + mass = dist[1:].sum() + if mass <= 0: + return dist + dist[1:] /= mass + return dist + + +def histogram_from_lengths(lengths: list[int], max_length: int) -> np.ndarray: + histogram = np.zeros(max_length + 1, dtype=float) + for length in lengths: + if 0 <= length <= max_length: + histogram[length] += 1.0 + return normalize_distribution(histogram) + + +def expected_gap_from_lengths(lengths: list[int], checkpoints: list[int]) -> float: + if not lengths: + return 0.0 + ordered = sorted(set(int(c) for c in checkpoints if c > 0)) + total = 0.0 + for length in lengths: + eligible = [c for c in ordered if c <= length] + previous = eligible[-1] if eligible else 0 + total += length - previous + return total / len(lengths) + + +@dataclass +class Line: + slope: float + intercept: float + arg: int + + +class MonotoneHull: + def __init__(self) -> None: + self.lines: deque[Line] = deque() + + @staticmethod + def _value(line: Line, x: float) -> float: + return line.slope * x + line.intercept + + @staticmethod + def _redundant(a: Line, b: Line, c: Line) -> bool: + return (c.intercept - a.intercept) * (a.slope - b.slope) <= (b.intercept - a.intercept) * (a.slope - c.slope) + EPS + + def add_line(self, line: Line) -> None: + while len(self.lines) >= 2 and self._redundant(self.lines[-2], self.lines[-1], line): + self.lines.pop() + self.lines.append(line) + + def query(self, x: float) -> tuple[float, int]: + while len(self.lines) >= 2 and self._value(self.lines[0], x) >= self._value(self.lines[1], x) - EPS: + self.lines.popleft() + best = self.lines[0] + return self._value(best, x), best.arg + + +def offline_optimal_checkpoints(max_length: int, budget: int, distribution: np.ndarray) -> tuple[list[int], float]: + if budget <= 0 or max_length <= 0: + return [], float(np.dot(np.arange(len(distribution)), distribution)) + budget = min(budget, max_length) + distribution = normalize_distribution(distribution) + prefix_prob = np.zeros(max_length + 1, dtype=float) + prefix_tp = np.zeros(max_length + 1, dtype=float) + positions = np.arange(max_length + 1, dtype=float) + prefix_prob[1:] = np.cumsum(distribution[1:]) + prefix_tp[1:] = np.cumsum(distribution[1:] * positions[1:]) + + dp_prev = prefix_tp.copy() + back = np.zeros((budget + 1, max_length + 1), dtype=np.int32) + + for used in range(1, budget + 1): + dp_curr = np.zeros(max_length + 1, dtype=float) + hull = MonotoneHull() + hull.add_line(Line(slope=-1.0, intercept=float(dp_prev[0]), arg=1)) + for end in range(1, max_length + 1): + start = end + intercept = float(dp_prev[start - 1] - prefix_tp[start - 1] + start * prefix_prob[start - 1]) + hull.add_line(Line(slope=-float(start), intercept=intercept, arg=start)) + value, argmin = hull.query(float(prefix_prob[end])) + dp_curr[end] = prefix_tp[end] + value + back[used, end] = argmin + dp_prev = dp_curr + + checkpoints: list[int] = [] + end = max_length + for used in range(budget, 0, -1): + start = int(back[used, end]) + if start <= 0: + break + checkpoints.append(start) + end = start - 1 + return sorted(set(checkpoints)), float(dp_prev[max_length]) + diff --git a/simulation/trace_cache_sim/experiments.py b/simulation/trace_cache_sim/experiments.py new file mode 100644 index 0000000..c0f0b35 --- /dev/null +++ b/simulation/trace_cache_sim/experiments.py @@ -0,0 +1,535 @@ +"""Experiment runner for the trace-driven sparse prefix caching simulator.""" + +from __future__ import annotations + +import os +from pathlib import Path + +from .policies import ( + BoundedCachePolicy, + BranchOnlyPolicy, + BranchPlusGeometricPolicy, + DenseEveryKPolicy, + ExponentialHistogramPolicy, + GeometricBaseRPolicy, + GeometricEpsPolicy, + NoCachePolicy, + OfflineDPOptimalPolicy, + OnlineHistogramPolicy, + UniformBudgetMPolicy, +) +from .simulator import SimulationConfig, SimulationResult, simulate_trace +from .toy_recurrence import checkpoint_resume_error +from .traces import Request, build_trace_families + + +REPO_ROOT = Path(__file__).resolve().parents[1] +os.environ.setdefault("MPLCONFIGDIR", str(REPO_ROOT / ".mplconfig")) +(REPO_ROOT / ".mplconfig").mkdir(exist_ok=True) + +import matplotlib +import numpy as np +import seaborn as sns + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib import ticker # noqa: E402 + + +POLICY_ORDER = [ + "no_cache", + "dense_every_k", + "uniform_budget_m", + "geometric_base_r", + "geometric_eps", + "branch_only", + "branch_plus_geometric", + "offline_dp_optimal", + "online_histogram_policy", + "exp_histogram_policy", +] + + +def build_policies(chunk_size: int) -> list: + return [ + NoCachePolicy(chunk_size=chunk_size), + DenseEveryKPolicy(every_k=32, chunk_size=chunk_size), + UniformBudgetMPolicy(budget=16, chunk_size=chunk_size), + GeometricBaseRPolicy(base=2.0, chunk_size=chunk_size), + GeometricEpsPolicy(eps=0.35, chunk_size=chunk_size), + BranchOnlyPolicy(chunk_size=chunk_size), + BranchPlusGeometricPolicy(base=2.0, chunk_size=chunk_size), + OfflineDPOptimalPolicy(budget=16, chunk_size=chunk_size), + OnlineHistogramPolicy(budget=16, update_every=8, chunk_size=chunk_size), + ExponentialHistogramPolicy(budget=16, update_every=8, decay=0.95, chunk_size=chunk_size), + ] + + +def policy_label(name: str) -> str: + return { + "no_cache": "none", + "dense_every_k": "dense-k", + "uniform_budget_m": "uniform", + "geometric_base_r": "geom-r", + "geometric_eps": "geom-eps", + "branch_only": "branch", + "branch_plus_geometric": "branch+geom", + "offline_dp_optimal": "offline-dp", + "online_histogram_policy": "online-hist", + "exp_histogram_policy": "exp-hist", + }[name] + + +def policy_color_map() -> dict[str, tuple[float, float, float]]: + colors = sns.color_palette("tab10", n_colors=len(POLICY_ORDER)) + return {policy: color for policy, color in zip(POLICY_ORDER, colors)} + + +def short_bytes_label(value: float) -> str: + if value >= 1_000_000: + return f"{value / 1_000_000:.1f}M" + if value >= 1_000: + rounded = value / 1_000 + if rounded >= 100: + return f"{rounded:.0f}k" + if rounded >= 10: + return f"{rounded:.0f}k" + return f"{rounded:.1f}k" + return f"{value:.0f}" + + +def byte_tick_formatter(value: float, _: float) -> str: + return short_bytes_label(value) + + +def exactness_suite(traces: dict[str, list[Request]]) -> dict[str, float]: + errors: dict[str, float] = {} + for family, requests in traces.items(): + sample = requests[min(2, len(requests) - 1)].tokens + checkpoint = max(1, len(sample) // 2) + errors[family] = checkpoint_resume_error(sample, checkpoint) + return errors + + +def eviction_sweep( + traces: dict[str, list[Request]], + config: SimulationConfig, + figure_dir: Path, +) -> list[dict[str, float | str | int]]: + """Run cache eviction sweep: vary max cache entries for key policies.""" + base_specs: list[tuple[str, type, dict]] = [ + ("branch_only", BranchOnlyPolicy, {"chunk_size": config.chunk_size}), + ("offline_dp", OfflineDPOptimalPolicy, {"budget": 16, "chunk_size": config.chunk_size}), + ("dense_every_k", DenseEveryKPolicy, {"every_k": 32, "chunk_size": config.chunk_size}), + ] + cache_budgets = [4, 8, 16, 32, 64, 128] + records: list[dict[str, float | str | int]] = [] + + for family, requests in traces.items(): + for label, cls, kwargs in base_specs: + base = cls(**kwargs) + result = simulate_trace(requests, base, config) + records.append({ + "family": family, + "policy": label, + "max_entries": 9999, + "avg_token_hit": result.summary.avg_token_hit_ratio, + "avg_hybrid_cost": result.summary.avg_hybrid_cost, + }) + for max_ent in cache_budgets: + base = cls(**kwargs) + bounded = BoundedCachePolicy(base, max_entries=max_ent, chunk_size=config.chunk_size) + result = simulate_trace(requests, bounded, config) + records.append({ + "family": family, + "policy": label, + "max_entries": max_ent, + "avg_token_hit": result.summary.avg_token_hit_ratio, + "avg_hybrid_cost": result.summary.avg_hybrid_cost, + }) + + plot_eviction_sweep(records, list(traces.keys()), figure_dir / "trace_eviction_sweep.pdf") + return records + + +def plot_eviction_sweep( + records: list[dict[str, float | str | int]], + families: list[str], + filename: Path, +) -> None: + policy_colors = {"branch_only": "#0a9396", "offline_dp": "#ee9b00", "dense_every_k": "#ca6702"} + families_plus = families + ["all_families"] + fig, axes = plt.subplots(2, 3, figsize=(14.2, 7.6)) + axes_flat = axes.flatten() + legend_handles = [] + + for ax, family in zip(axes_flat, families_plus): + family_rows = records if family == "all_families" else [r for r in records if r["family"] == family] + for policy_name in ["branch_only", "offline_dp", "dense_every_k"]: + rows = sorted( + [r for r in family_rows if r["policy"] == policy_name], + key=lambda r: int(r["max_entries"]), + ) + if not rows: + continue + xs = [int(r["max_entries"]) for r in rows] + ys = [float(r["avg_token_hit"]) for r in rows] + handle, = ax.plot( + xs, ys, marker="o", markersize=5, linewidth=2, + label=policy_name.replace("_", " "), color=policy_colors[policy_name], + ) + if family == families_plus[0]: + legend_handles.append(handle) + title = "all families" if family == "all_families" else family.replace("_", "\n") + ax.set_title(title, fontsize=13) + ax.set_xscale("log") + ax.set_ylim(-0.05, 1.05) + ax.grid(alpha=0.3) + if family == "all_families": + ax.set_facecolor("#f7f7f7") + + for ax in axes[0, :]: + ax.tick_params(labelbottom=False) + if len(families_plus) < len(axes_flat): + for ax in axes_flat[len(families_plus):]: + ax.set_visible(False) + + fig.legend( + legend_handles, + ["branch only", "offline dp", "dense every-k"], + loc="lower center", ncol=3, frameon=False, fontsize=10.5, bbox_to_anchor=(0.5, 0.01), + ) + fig.supxlabel("Max cache entries (log scale)", fontsize=13, y=0.07) + fig.supylabel("Average token-hit ratio", fontsize=13, x=0.04) + fig.subplots_adjust(left=0.08, right=0.98, top=0.91, bottom=0.17, hspace=0.16, wspace=0.18) + plt.savefig(filename) + plt.close(fig) + + +def ssm_fraction_sweep( + traces: dict[str, list[Request]], + figure_dir: Path, +) -> list[dict[str, float | str]]: + """Sweep SSM layer fraction and measure sparse caching benefit.""" + fractions = np.linspace(0.1, 1.0, 10) + records: list[dict[str, float | str]] = [] + + for ssm_frac in fractions: + config = SimulationConfig( + a_rec=1.0 * ssm_frac, + a_attn=0.002 * (1.0 - ssm_frac), + ) + for family, requests in traces.items(): + no_cache_result = simulate_trace( + requests, NoCachePolicy(chunk_size=config.chunk_size), config, + ) + dp_result = simulate_trace( + requests, OfflineDPOptimalPolicy(budget=16, chunk_size=config.chunk_size), config, + ) + branch_result = simulate_trace( + requests, BranchOnlyPolicy(chunk_size=config.chunk_size), config, + ) + no_cost = no_cache_result.summary.avg_hybrid_cost + best_cost = min(dp_result.summary.avg_hybrid_cost, branch_result.summary.avg_hybrid_cost) + reduction = 1.0 - best_cost / max(no_cost, 1e-12) if no_cost > 1e-12 else 0.0 + records.append({ + "family": family, + "ssm_fraction": float(ssm_frac), + "no_cache_cost": no_cost, + "best_sparse_cost": best_cost, + "cost_reduction": reduction, + }) + + plot_ssm_fraction_sweep(records, list(traces.keys()), figure_dir / "trace_ssm_fraction_sweep.pdf") + return records + + +def plot_ssm_fraction_sweep( + records: list[dict[str, float | str]], + families: list[str], + filename: Path, +) -> None: + plt.figure(figsize=(8.5, 5.4)) + colors = sns.color_palette("Set2", n_colors=len(families)) + for color, family in zip(colors, families): + rows = sorted( + [r for r in records if r["family"] == family], + key=lambda r: float(r["ssm_fraction"]), + ) + xs = [float(r["ssm_fraction"]) for r in rows] + ys = [100.0 * float(r["cost_reduction"]) for r in rows] + plt.plot(xs, ys, marker="o", label=family.replace("_", " "), color=color, linewidth=2, markersize=5) + + plt.xlabel("SSM layer fraction") + plt.ylabel("Cost reduction from sparse caching (%)") + plt.title("Sparse checkpoint benefit scales with SSM fraction") + plt.legend(frameon=False, fontsize=10) + plt.grid(alpha=0.3) + plt.tight_layout() + plt.savefig(filename) + plt.close() + + +def run_all_experiments() -> tuple[list[dict[str, float | str]], dict[str, float]]: + sns.set_theme(style="whitegrid", context="talk") + figure_dir = REPO_ROOT / "figures" + figure_dir.mkdir(exist_ok=True) + summary_path = REPO_ROOT / "trace_driven_summary.md" + + config = SimulationConfig() + traces = build_trace_families(seed=42) + exactness = exactness_suite(traces) + + records: list[dict[str, float | str]] = [] + all_results: dict[tuple[str, str], SimulationResult] = {} + for family, requests in traces.items(): + for policy in build_policies(config.chunk_size): + result = simulate_trace(requests, policy, config) + all_results[(family, policy.name)] = result + summary = result.summary + records.append( + { + "family": family, + "policy": policy.name, + "policy_label": policy_label(policy.name), + "avg_recompute_tokens": summary.avg_recompute_tokens, + "max_recompute_tokens": summary.max_recompute_tokens, + "avg_token_hit_ratio": summary.avg_token_hit_ratio, + "avg_alignment_error": summary.avg_alignment_error, + "avg_hybrid_cost": summary.avg_hybrid_cost, + "avg_recurrent_cost": summary.avg_recurrent_cost, + "avg_checkpoint_count": summary.avg_checkpoint_count, + "max_checkpoint_count": summary.max_checkpoint_count, + "max_recurrent_bytes": summary.max_recurrent_bytes, + "max_total_bytes": summary.max_total_bytes, + "avg_bytes_written_total": summary.avg_bytes_written_total, + "avg_bytes_read_total": summary.avg_bytes_read_total, + "exactness_error": exactness[family], + } + ) + + no_cache_map = {family: all_results[(family, "no_cache")] for family in traces} + branch_map = {family: all_results[(family, "branch_only")] for family in traces} + for record in records: + family = str(record["family"]) + policy = str(record["policy"]) + metrics = all_results[(family, policy)].request_metrics + no_metrics = no_cache_map[family].request_metrics + branch_metrics = branch_map[family].request_metrics + wins_vs_no = sum(metric.hybrid_cost < base.hybrid_cost for metric, base in zip(metrics, no_metrics)) + wins_vs_branch = sum(metric.hybrid_cost < base.hybrid_cost for metric, base in zip(metrics, branch_metrics)) + record["win_rate_vs_no_cache"] = wins_vs_no / len(metrics) + record["win_rate_vs_branch_only"] = wins_vs_branch / len(metrics) + record["hybrid_cost_ratio_vs_no_cache"] = record["avg_hybrid_cost"] / max(no_cache_map[family].summary.avg_hybrid_cost, 1e-12) + + plot_memory_vs_metric( + records, + families=list(traces.keys()), + x_key="max_recurrent_bytes", + y_key="avg_token_hit_ratio", + y_label="Average token-hit ratio", + filename=figure_dir / "trace_memory_vs_token_hit.pdf", + y_limits=(0.5, 1.0), + ) + plot_memory_vs_metric( + records, + families=list(traces.keys()), + x_key="max_recurrent_bytes", + y_key="avg_recompute_tokens", + y_label="Average recompute tokens", + filename=figure_dir / "trace_memory_vs_expected_recompute.pdf", + y_limits=(0.0, 100.0), + ) + plot_memory_vs_metric( + records, + families=list(traces.keys()), + x_key="max_total_bytes", + y_key="avg_hybrid_cost", + y_label="Average hybrid proxy TTFT cost", + filename=figure_dir / "trace_memory_vs_estimated_ttft.pdf", + y_limits=(0.0, 250.0), + per_family_y_limits={"agent_tree": (0.0, 60.0)}, + ) + plot_trace_family_heatmap(records, list(traces.keys()), figure_dir / "trace_family_comparisons.pdf") + + eviction_records = eviction_sweep(traces, config, figure_dir) + ssm_records = ssm_fraction_sweep(traces, figure_dir) + + write_markdown_summary(records, exactness, summary_path) + print_summary(records, exactness) + print() + print("Eviction sweep (sample):") + for row in eviction_records: + if row["family"] == "exact_hot_prefix": + print(f" {row['policy']:>15} max_entries={int(row['max_entries']):>5} hit={float(row['avg_token_hit']):.3f}") + print() + print("SSM fraction sweep (sample):") + for row in ssm_records: + if row["family"] == "diffuse_cutpoints": + print(f" ssm_frac={float(row['ssm_fraction']):.2f} reduction={100*float(row['cost_reduction']):.1f}%") + return records, exactness + + +def plot_memory_vs_metric( + records: list[dict[str, float | str]], + families: list[str], + x_key: str, + y_key: str, + y_label: str, + filename: Path, + y_limits: tuple[float, float] | None = None, + per_family_y_limits: dict[str, tuple[float, float]] | None = None, +) -> None: + families_with_all = families + ["all_families"] + fig, axes = plt.subplots(2, 3, figsize=(14.2, 7.6)) + axes_flat = axes.flatten() + colors = policy_color_map() + legend_handles = [] + per_family_y_limits = per_family_y_limits or {} + + for axis, family in zip(axes_flat, families_with_all): + family_rows = records if family == "all_families" else [row for row in records if row["family"] == family] + for policy in POLICY_ORDER: + matching_rows = [row for row in family_rows if row["policy"] == policy] + if not matching_rows: + continue + color = colors[policy] + if family == "all_families": + xs = [float(row[x_key]) for row in matching_rows] + ys = [float(row[y_key]) for row in matching_rows] + handle = axis.scatter(xs, ys, s=55, color=color, alpha=0.75, edgecolors="white", linewidths=0.4) + else: + row = matching_rows[0] + x = float(row[x_key]) + y = float(row[y_key]) + handle = axis.scatter(x, y, s=80, color=color, edgecolors="white", linewidths=0.5) + if family == families_with_all[0]: + legend_handles.append(handle) + + title = "all families" if family == "all_families" else family.replace("_", "\n") + axis.set_title(title, fontsize=13) + axis.tick_params(labelsize=11) + axis.set_xscale("log") + axis.xaxis.set_major_locator(ticker.LogLocator(base=10, subs=(1.0, 2.0, 5.0), numticks=4)) + axis.xaxis.set_major_formatter(ticker.FuncFormatter(byte_tick_formatter)) + axis.xaxis.set_minor_locator(ticker.NullLocator()) + axis.ticklabel_format(axis="y", style="plain") + axis.grid(alpha=0.3) + if family in per_family_y_limits: + axis.set_ylim(*per_family_y_limits[family]) + elif y_limits is not None: + axis.set_ylim(*y_limits) + if family == "all_families": + axis.set_facecolor("#f7f7f7") + + for axis in axes[0, :]: + axis.tick_params(labelbottom=False) + + fig.legend( + legend_handles, + [policy_label(policy) for policy in POLICY_ORDER], + loc="lower center", + ncol=5, + frameon=False, + fontsize=10.5, + bbox_to_anchor=(0.5, 0.01), + ) + fig.supxlabel("Max recurrent bytes" if x_key == "max_recurrent_bytes" else "Max total bytes", fontsize=13, y=0.07) + fig.supylabel(y_label, fontsize=13, x=0.04) + fig.subplots_adjust(left=0.08, right=0.98, top=0.91, bottom=0.17, hspace=0.16, wspace=0.18) + plt.savefig(filename) + plt.close(fig) + + +def plot_trace_family_heatmap(records: list[dict[str, float | str]], families: list[str], filename: Path) -> None: + policies = POLICY_ORDER + matrix = np.zeros((len(policies), len(families))) + for i, policy in enumerate(policies): + for j, family in enumerate(families): + row = next(item for item in records if item["family"] == family and item["policy"] == policy) + matrix[i, j] = 100.0 * float(row["win_rate_vs_branch_only"]) + with sns.plotting_context("paper", font_scale=1.05): + plt.figure(figsize=(7.7, 4.2)) + ax = sns.heatmap( + matrix, + annot=True, + fmt=".1f", + cmap="YlOrBr", + xticklabels=[family.replace("_", "\n") for family in families], + yticklabels=[policy_label(name) for name in policies], + annot_kws={"size": 7.5}, + cbar_kws={"label": "Request win-rate vs branch-only (%)"}, + ) + ax.set_title("Per-family win-rate against the branch-only baseline", fontsize=10.5, pad=6) + ax.set_xlabel("") + ax.set_ylabel("") + ax.tick_params(axis="x", labelrotation=25, labelsize=7.5) + ax.tick_params(axis="y", labelrotation=0, labelsize=7.5) + plt.tight_layout() + plt.savefig(filename) + plt.close() + + +def write_markdown_summary(records: list[dict[str, float | str]], exactness: dict[str, float], path: Path) -> None: + best_by_family: list[str] = [] + for family in sorted(exactness): + family_rows = [row for row in records if row["family"] == family and row["policy"] != "no_cache"] + best = min(family_rows, key=lambda row: float(row["avg_hybrid_cost"])) + best_by_family.append( + f"- `{family}`: best hybrid proxy cost from `{best['policy']}` with " + f"token-hit {100.0 * float(best['avg_token_hit_ratio']):.1f}% and " + f"alignment error {float(best['avg_alignment_error']):.2f}." + ) + + lines = [ + "# Trace-Driven Sparse Prefix Caching Summary", + "", + "## Assumptions", + "", + "- Requests are explicit token sequences; overlap is exact longest common prefix.", + "- Checkpoints are exact recurrent states; resume always recomputes the missing suffix exactly.", + "- Memory is tracked separately for recurrent checkpoints and dense attention KV storage induced by cached prefix depth.", + "", + "## Main Findings", + "", + *best_by_family, + "", + "## Exactness", + "", + ] + for family, error in exactness.items(): + lines.append(f"- `{family}`: max checkpoint-resume error `{error:.3e}`.") + lines.extend( + [ + "", + "## Publication signal", + "", + "- The trace-driven simulator suggests sparse checkpoint placement is not just a tuning knob.", + "- Branch-aware caching matters on tree-like workloads.", + "- Total hybrid memory can remain large even when recurrent checkpoint memory is sparse.", + ] + ) + path.write_text("\n".join(lines), encoding="utf-8") + + +def print_summary(records: list[dict[str, float | str]], exactness: dict[str, float]) -> None: + print("Trace-driven sparse prefix caching experiments") + print("=" * 52) + for row in records: + print( + f"[{row['family']}] {str(row['policy']).rjust(22)} | " + f"hit={100.0 * float(row['avg_token_hit_ratio']):6.2f}% | " + f"recompute={float(row['avg_recompute_tokens']):7.2f} | " + f"hybrid={float(row['avg_hybrid_cost']):9.2f} | " + f"recurrent_mem={int(row['max_recurrent_bytes']):8d} | " + f"total_mem={int(row['max_total_bytes']):9d} | " + f"align={float(row['avg_alignment_error']):5.2f}" + ) + print() + print("Toy exactness:") + for family, error in exactness.items(): + print(f" {family}: max error={error:.3e}") + + +if __name__ == "__main__": + run_all_experiments() diff --git a/simulation/trace_cache_sim/policies.py b/simulation/trace_cache_sim/policies.py new file mode 100644 index 0000000..bf793a3 --- /dev/null +++ b/simulation/trace_cache_sim/policies.py @@ -0,0 +1,298 @@ +"""Cache policy definitions for trace-driven sparse checkpointing.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .dp_opt import histogram_from_lengths, normalize_distribution, offline_optimal_checkpoints +from .traces import PrefixTrie, Request + + +Prefix = tuple[int, ...] + + +@dataclass(frozen=True) +class CacheSnapshot: + raw_prefixes: frozenset[Prefix] + aligned_prefixes: frozenset[Prefix] + + +def align_length(length: int, chunk_size: int) -> int: + if chunk_size <= 1: + return max(0, length) + return max(0, (length // chunk_size) * chunk_size) + + +def uniform_lengths(max_length: int, budget: int) -> list[int]: + if budget <= 0 or max_length <= 0: + return [] + if budget >= max_length: + return list(range(1, max_length + 1)) + return sorted({max(1, min(max_length, round(i * (max_length + 1) / (budget + 1)))) for i in range(1, budget + 1)}) + + +def geometric_base_lengths(max_length: int, base: float) -> list[int]: + if max_length <= 0: + return [] + points = [1] + while points[-1] < max_length: + next_point = max(points[-1] + 1, int(points[-1] * base)) + points.append(min(max_length, next_point)) + if points[-1] == max_length: + break + return sorted(set(points)) + + +def geometric_eps_lengths(max_length: int, eps: float) -> list[int]: + if max_length <= 0: + return [] + current = 1 + points = [current] + factor = 1.0 + eps + while current < max_length: + current = max(current + 1, int(current * factor)) + points.append(min(max_length, current)) + if current >= max_length: + break + return sorted(set(points)) + + +def collect_prefixes_for_lengths(requests: list[Request], lengths: list[int]) -> set[Prefix]: + prefixes: set[Prefix] = set() + positive_lengths = [length for length in lengths if length > 0] + for request in requests: + for length in positive_lengths: + if length <= len(request.tokens): + prefixes.add(request.tokens[:length]) + return prefixes + + +def apply_alignment(prefixes: set[Prefix], chunk_size: int) -> set[Prefix]: + aligned: set[Prefix] = set() + for prefix in prefixes: + length = align_length(len(prefix), chunk_size) + if length > 0: + aligned.add(prefix[:length]) + return aligned + + +def deepest_cached_depth(tokens: tuple[int, ...], prefixes: frozenset[Prefix]) -> int: + for length in range(len(tokens), 0, -1): + if tokens[:length] in prefixes: + return length + return 0 + + +class CachePolicy: + name = "base" + + def __init__(self, chunk_size: int = 1) -> None: + self.chunk_size = chunk_size + + def reset(self, full_trace: list[Request] | None = None) -> None: + return None + + def observe(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> None: + return None + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + return set() + + def snapshot(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> CacheSnapshot: + raw = self.raw_prefixes(seen_requests, trie, overlap_history) + return CacheSnapshot(raw_prefixes=frozenset(raw), aligned_prefixes=frozenset(apply_alignment(raw, self.chunk_size))) + + +class NoCachePolicy(CachePolicy): + name = "no_cache" + + +class DenseEveryKPolicy(CachePolicy): + name = "dense_every_k" + + def __init__(self, every_k: int, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.every_k = every_k + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + if self.every_k <= 0: + return set() + max_length = max((len(request.tokens) for request in seen_requests), default=0) + lengths = list(range(self.every_k, max_length + 1, self.every_k)) + return collect_prefixes_for_lengths(seen_requests, lengths) + + +class UniformBudgetMPolicy(CachePolicy): + name = "uniform_budget_m" + + def __init__(self, budget: int, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.budget = budget + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + max_length = max((len(request.tokens) for request in seen_requests), default=0) + return collect_prefixes_for_lengths(seen_requests, uniform_lengths(max_length, self.budget)) + + +class GeometricBaseRPolicy(CachePolicy): + name = "geometric_base_r" + + def __init__(self, base: float, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.base = base + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + max_length = max((len(request.tokens) for request in seen_requests), default=0) + return collect_prefixes_for_lengths(seen_requests, geometric_base_lengths(max_length, self.base)) + + +class GeometricEpsPolicy(CachePolicy): + name = "geometric_eps" + + def __init__(self, eps: float, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.eps = eps + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + max_length = max((len(request.tokens) for request in seen_requests), default=0) + return collect_prefixes_for_lengths(seen_requests, geometric_eps_lengths(max_length, self.eps)) + + +class BranchOnlyPolicy(CachePolicy): + name = "branch_only" + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + prefixes = set(trie.branch_prefixes()) + prefixes.update(request.tokens for request in seen_requests) + return prefixes + + +class BranchPlusGeometricPolicy(CachePolicy): + name = "branch_plus_geometric" + + def __init__(self, base: float, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.base = base + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + prefixes = BranchOnlyPolicy().raw_prefixes(seen_requests, trie, overlap_history) + max_length = max((len(request.tokens) for request in seen_requests), default=0) + prefixes.update(collect_prefixes_for_lengths(seen_requests, geometric_base_lengths(max_length, self.base))) + return prefixes + + +class OfflineDPOptimalPolicy(CachePolicy): + name = "offline_dp_optimal" + + def __init__(self, budget: int, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.budget = budget + self.schedule: list[int] = [] + + def reset(self, full_trace: list[Request] | None = None) -> None: + if not full_trace: + self.schedule = [] + return + from .traces import sequential_overlap_lengths + + overlaps = sequential_overlap_lengths(full_trace) + max_length = max(overlaps, default=0) + histogram = histogram_from_lengths(overlaps, max_length=max_length) + self.schedule, _ = offline_optimal_checkpoints(max_length, self.budget, histogram) + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + return collect_prefixes_for_lengths(seen_requests, self.schedule) + + +class OnlineHistogramPolicy(CachePolicy): + name = "online_histogram_policy" + + def __init__(self, budget: int, update_every: int, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.budget = budget + self.update_every = update_every + self.schedule: list[int] = [] + + def reset(self, full_trace: list[Request] | None = None) -> None: + self.schedule = [] + + def observe(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> None: + if not overlap_history or len(overlap_history) % self.update_every != 0: + return + max_length = max(overlap_history) + histogram = histogram_from_lengths(overlap_history, max_length=max_length) + self.schedule, _ = offline_optimal_checkpoints(max_length, self.budget, histogram) + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + return collect_prefixes_for_lengths(seen_requests, self.schedule) + + +class ExponentialHistogramPolicy(CachePolicy): + """Online policy with exponentially weighted overlap histogram. + + Recent observations receive more weight, allowing faster adaptation + to non-stationary workloads. + """ + + name = "exp_histogram_policy" + + def __init__(self, budget: int, update_every: int, decay: float = 0.95, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.budget = budget + self.update_every = update_every + self.decay = decay + self.schedule: list[int] = [] + + def reset(self, full_trace: list[Request] | None = None) -> None: + self.schedule = [] + + def observe(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> None: + if not overlap_history or len(overlap_history) % self.update_every != 0: + return + max_length = max(overlap_history) + if max_length <= 0: + return + n = len(overlap_history) + histogram = np.zeros(max_length + 1, dtype=float) + for i, length in enumerate(overlap_history): + weight = self.decay ** (n - 1 - i) + if 0 <= length <= max_length: + histogram[length] += weight + histogram = normalize_distribution(histogram) + self.schedule, _ = offline_optimal_checkpoints(max_length, self.budget, histogram) + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + return collect_prefixes_for_lengths(seen_requests, self.schedule) + + +class BoundedCachePolicy(CachePolicy): + """Wraps any base policy and enforces a maximum cache size with LRU eviction.""" + + def __init__(self, base_policy: CachePolicy, max_entries: int, chunk_size: int = 1) -> None: + super().__init__(chunk_size=chunk_size) + self.base = base_policy + self.max_entries = max_entries + self.name = f"{base_policy.name}+lru{max_entries}" + self._cache: dict[tuple[int, ...], int] = {} + self._step: int = 0 + + def reset(self, full_trace: list[Request] | None = None) -> None: + self.base.reset(full_trace) + self._cache = {} + self._step = 0 + + def observe(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> None: + self.base.observe(seen_requests, trie, overlap_history) + self._step = len(overlap_history) + desired = self.base.raw_prefixes(seen_requests, trie, overlap_history) + for prefix in desired: + self._cache[prefix] = self._step + while len(self._cache) > self.max_entries: + oldest_key = min(self._cache, key=self._cache.get) + del self._cache[oldest_key] + + def raw_prefixes(self, seen_requests: list[Request], trie: PrefixTrie, overlap_history: list[int]) -> set[Prefix]: + return set(self._cache.keys()) + diff --git a/simulation/trace_cache_sim/simulator.py b/simulation/trace_cache_sim/simulator.py new file mode 100644 index 0000000..6cb9ead --- /dev/null +++ b/simulation/trace_cache_sim/simulator.py @@ -0,0 +1,162 @@ +"""Trace-driven simulation of sparse recurrent-state prefix caching.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .policies import CachePolicy, CacheSnapshot, align_length, deepest_cached_depth +from .traces import PrefixTrie, Request + + +@dataclass(frozen=True) +class SimulationConfig: + recurrent_state_bytes: int = 2048 + attention_bytes_per_token: int = 256 + a_rec: float = 1.0 + a_attn: float = 0.002 + a_load: float = 0.0001 + a_write: float = 0.00002 + chunk_size: int = 16 + + +@dataclass +class RequestMetrics: + overlap_length: int + raw_reuse_depth: int + reuse_depth: int + recompute_tokens: int + alignment_error: int + token_hit_ratio: float + recurrent_cost: float + hybrid_cost: float + bytes_read_recurrent: int + bytes_read_total: int + bytes_written_recurrent: int + bytes_written_total: int + checkpoint_count: int + recurrent_bytes: int + total_bytes: int + + +@dataclass +class SimulationSummary: + family: str + policy: str + requests: int + avg_recompute_tokens: float + max_recompute_tokens: int + avg_token_hit_ratio: float + avg_alignment_error: float + avg_recurrent_cost: float + avg_hybrid_cost: float + avg_bytes_read_total: float + avg_bytes_written_total: float + avg_checkpoint_count: float + max_checkpoint_count: int + avg_recurrent_bytes: float + max_recurrent_bytes: int + avg_total_bytes: float + max_total_bytes: int + + +@dataclass +class SimulationResult: + summary: SimulationSummary + request_metrics: list[RequestMetrics] + + +def recurrent_bytes(snapshot: CacheSnapshot, config: SimulationConfig) -> int: + return len(snapshot.aligned_prefixes) * config.recurrent_state_bytes + + +def attention_bytes(snapshot: CacheSnapshot, config: SimulationConfig) -> int: + return sum(len(prefix) * config.attention_bytes_per_token for prefix in snapshot.aligned_prefixes) + + +def total_bytes(snapshot: CacheSnapshot, config: SimulationConfig) -> int: + return recurrent_bytes(snapshot, config) + attention_bytes(snapshot, config) + + +def bytes_for_new_prefixes(prefixes: set[tuple[int, ...]], config: SimulationConfig) -> tuple[int, int]: + rec = len(prefixes) * config.recurrent_state_bytes + attn = sum(len(prefix) * config.attention_bytes_per_token for prefix in prefixes) + return rec, rec + attn + + +def simulate_trace(requests: list[Request], policy: CachePolicy, config: SimulationConfig) -> SimulationResult: + trie = PrefixTrie() + seen_requests: list[Request] = [] + overlap_history: list[int] = [] + policy.reset(full_trace=requests) + + metrics: list[RequestMetrics] = [] + family = requests[0].family if requests else "unknown" + + for request in requests: + before = policy.snapshot(seen_requests, trie, overlap_history) + overlap = trie.longest_prefix(request.tokens) + raw_reuse = min(overlap, deepest_cached_depth(request.tokens, before.raw_prefixes)) + reuse = min(overlap, deepest_cached_depth(request.tokens, before.aligned_prefixes)) + alignment_error = raw_reuse - align_length(raw_reuse, config.chunk_size) + recompute = overlap - reuse + token_hit = reuse / overlap if overlap > 0 else 0.0 + + read_recurrent = config.recurrent_state_bytes if reuse > 0 else 0 + read_total = read_recurrent + reuse * config.attention_bytes_per_token if reuse > 0 else 0 + + recurrent_cost = config.a_rec * recompute + config.a_load * read_recurrent + hybrid_cost = config.a_rec * recompute + config.a_attn * (overlap * overlap - reuse * reuse) + config.a_load * read_recurrent + + trie.insert(request.tokens) + seen_requests.append(request) + overlap_history.append(overlap) + policy.observe(seen_requests, trie, overlap_history) + after = policy.snapshot(seen_requests, trie, overlap_history) + + new_prefixes = set(after.aligned_prefixes) - set(before.aligned_prefixes) + written_recurrent, written_total = bytes_for_new_prefixes(new_prefixes, config) + recurrent_cost += config.a_write * written_recurrent + hybrid_cost += config.a_write * written_total + + metrics.append( + RequestMetrics( + overlap_length=overlap, + raw_reuse_depth=raw_reuse, + reuse_depth=reuse, + recompute_tokens=recompute, + alignment_error=alignment_error, + token_hit_ratio=token_hit, + recurrent_cost=recurrent_cost, + hybrid_cost=hybrid_cost, + bytes_read_recurrent=read_recurrent, + bytes_read_total=read_total, + bytes_written_recurrent=written_recurrent, + bytes_written_total=written_total, + checkpoint_count=len(after.aligned_prefixes), + recurrent_bytes=recurrent_bytes(after, config), + total_bytes=total_bytes(after, config), + ) + ) + + summary = SimulationSummary( + family=family, + policy=policy.name, + requests=len(metrics), + avg_recompute_tokens=float(np.mean([item.recompute_tokens for item in metrics])) if metrics else 0.0, + max_recompute_tokens=max((item.recompute_tokens for item in metrics), default=0), + avg_token_hit_ratio=float(np.mean([item.token_hit_ratio for item in metrics])) if metrics else 0.0, + avg_alignment_error=float(np.mean([item.alignment_error for item in metrics])) if metrics else 0.0, + avg_recurrent_cost=float(np.mean([item.recurrent_cost for item in metrics])) if metrics else 0.0, + avg_hybrid_cost=float(np.mean([item.hybrid_cost for item in metrics])) if metrics else 0.0, + avg_bytes_read_total=float(np.mean([item.bytes_read_total for item in metrics])) if metrics else 0.0, + avg_bytes_written_total=float(np.mean([item.bytes_written_total for item in metrics])) if metrics else 0.0, + avg_checkpoint_count=float(np.mean([item.checkpoint_count for item in metrics])) if metrics else 0.0, + max_checkpoint_count=max((item.checkpoint_count for item in metrics), default=0), + avg_recurrent_bytes=float(np.mean([item.recurrent_bytes for item in metrics])) if metrics else 0.0, + max_recurrent_bytes=max((item.recurrent_bytes for item in metrics), default=0), + avg_total_bytes=float(np.mean([item.total_bytes for item in metrics])) if metrics else 0.0, + max_total_bytes=max((item.total_bytes for item in metrics), default=0), + ) + return SimulationResult(summary=summary, request_metrics=metrics) diff --git a/simulation/trace_cache_sim/toy_recurrence.py b/simulation/trace_cache_sim/toy_recurrence.py new file mode 100644 index 0000000..d215f78 --- /dev/null +++ b/simulation/trace_cache_sim/toy_recurrence.py @@ -0,0 +1,39 @@ +"""Tiny exact recurrent model for checkpoint-resume validation.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass +class ToyRecurrentModel: + dim: int = 4 + seed: int = 42 + + def token_features(self, token: int) -> tuple[float, np.ndarray, np.ndarray]: + rng = np.random.default_rng(self.seed + 97 * token) + lam = float(rng.uniform(0.1, 0.9)) + v = rng.normal(size=self.dim) + k = rng.normal(size=self.dim) + return lam, v, k + + def step(self, state: np.ndarray, token: int) -> np.ndarray: + lam, v, k = self.token_features(token) + return lam * state + np.outer(v, k) + + def run(self, tokens: tuple[int, ...], start_state: np.ndarray | None = None) -> np.ndarray: + state = np.zeros((self.dim, self.dim), dtype=float) if start_state is None else start_state.copy() + for token in tokens: + state = self.step(state, token) + return state + + +def checkpoint_resume_error(tokens: tuple[int, ...], checkpoint_depth: int, dim: int = 4, seed: int = 42) -> float: + model = ToyRecurrentModel(dim=dim, seed=seed) + full = model.run(tokens) + checkpoint_state = model.run(tokens[:checkpoint_depth]) + resumed = model.run(tokens[checkpoint_depth:], start_state=checkpoint_state) + return float(np.max(np.abs(full - resumed))) + diff --git a/simulation/trace_cache_sim/traces.py b/simulation/trace_cache_sim/traces.py new file mode 100644 index 0000000..1d4fcb3 --- /dev/null +++ b/simulation/trace_cache_sim/traces.py @@ -0,0 +1,169 @@ +"""Trace generation and prefix trie utilities.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Iterable + +import numpy as np + + +@dataclass(frozen=True) +class Request: + request_id: int + family: str + tokens: tuple[int, ...] + + +def longest_common_prefix(a: tuple[int, ...], b: tuple[int, ...]) -> int: + limit = min(len(a), len(b)) + index = 0 + while index < limit and a[index] == b[index]: + index += 1 + return index + + +@dataclass +class TrieNode: + depth: int + children: dict[int, "TrieNode"] = field(default_factory=dict) + pass_count: int = 0 + terminal_count: int = 0 + + +class PrefixTrie: + """Simple prefix trie used for overlap queries and branch-point extraction.""" + + def __init__(self) -> None: + self.root = TrieNode(depth=0) + + def insert(self, tokens: tuple[int, ...]) -> None: + node = self.root + node.pass_count += 1 + for token in tokens: + child = node.children.get(token) + if child is None: + child = TrieNode(depth=node.depth + 1) + node.children[token] = child + node = child + node.pass_count += 1 + node.terminal_count += 1 + + def longest_prefix(self, tokens: tuple[int, ...]) -> int: + node = self.root + matched = 0 + for token in tokens: + child = node.children.get(token) + if child is None: + break + matched += 1 + node = child + return matched + + def branch_prefixes(self) -> set[tuple[int, ...]]: + results: set[tuple[int, ...]] = set() + + def dfs(node: TrieNode, prefix: list[int]) -> None: + continuation_count = len(node.children) + (1 if node.terminal_count > 0 else 0) + if node.depth > 0 and node.pass_count >= 2 and continuation_count >= 2: + results.add(tuple(prefix)) + for token, child in node.children.items(): + prefix.append(token) + dfs(child, prefix) + prefix.pop() + + dfs(self.root, []) + return results + + +def _random_tokens(rng: np.random.Generator, length: int, low: int = 1_000, high: int = 1_000_000) -> tuple[int, ...]: + if length <= 0: + return () + return tuple(int(x) for x in rng.integers(low, high, size=length)) + + +def exact_hot_prefix(num_requests: int = 72, prefix_length: int = 320, suffix_length: int = 48, seed: int = 42) -> list[Request]: + rng = np.random.default_rng(seed) + shared_prefix = _random_tokens(rng, prefix_length) + requests: list[Request] = [] + for index in range(num_requests): + suffix = _random_tokens(rng, suffix_length) + requests.append(Request(index, "exact_hot_prefix", shared_prefix + suffix)) + return requests + + +def append_only_chat(num_requests: int = 64, system_length: int = 48, turn_length: int = 10, seed: int = 42) -> list[Request]: + rng = np.random.default_rng(seed + 1) + prefix = list(_random_tokens(rng, system_length)) + requests: list[Request] = [] + for index in range(num_requests): + prefix.extend(_random_tokens(rng, turn_length)) + requests.append(Request(index, "append_only_chat", tuple(prefix))) + return requests + + +def diffuse_cutpoints(num_requests: int = 72, document_length: int = 512, tail_length: int = 24, seed: int = 42) -> list[Request]: + rng = np.random.default_rng(seed + 2) + document = _random_tokens(rng, document_length) + requests: list[Request] = [] + for index in range(num_requests): + cut = int(rng.integers(document_length // 6, document_length + 1)) + suffix = _random_tokens(rng, tail_length) + requests.append(Request(index, "diffuse_cutpoints", document[:cut] + suffix)) + return requests + + +def agent_tree(branching_factor: int = 3, depth: int = 4, segment_length: int = 20, seed: int = 42) -> list[Request]: + rng = np.random.default_rng(seed + 3) + segment_bank: dict[tuple[int, ...], tuple[int, ...]] = {} + requests: list[Request] = [] + + def make_segment(path: tuple[int, ...]) -> tuple[int, ...]: + if path not in segment_bank: + segment_bank[path] = _random_tokens(rng, segment_length) + return segment_bank[path] + + def walk(path: tuple[int, ...], level: int) -> None: + if level == depth: + tokens: tuple[int, ...] = () + for stop in range(1, len(path) + 1): + tokens += make_segment(path[:stop]) + tokens += _random_tokens(rng, 8) + requests.append(Request(len(requests), "agent_tree", tokens)) + return + for child in range(branching_factor): + walk(path + (child,), level + 1) + + walk((), 0) + return requests + + +def adversarial_uniform_overlap(num_requests: int = 72, max_overlap: int = 384, suffix_length: int = 32, seed: int = 42) -> list[Request]: + rng = np.random.default_rng(seed + 4) + backbone = _random_tokens(rng, max_overlap) + requests: list[Request] = [Request(0, "adversarial_uniform_overlap", backbone + _random_tokens(rng, suffix_length))] + for index in range(1, num_requests): + overlap = int(rng.integers(1, max_overlap + 1)) + suffix = _random_tokens(rng, suffix_length) + requests.append(Request(index, "adversarial_uniform_overlap", backbone[:overlap] + suffix)) + return requests + + +def build_trace_families(seed: int = 42) -> dict[str, list[Request]]: + return { + "exact_hot_prefix": exact_hot_prefix(seed=seed), + "append_only_chat": append_only_chat(seed=seed), + "diffuse_cutpoints": diffuse_cutpoints(seed=seed), + "agent_tree": agent_tree(seed=seed), + "adversarial_uniform_overlap": adversarial_uniform_overlap(seed=seed), + } + + +def sequential_overlap_lengths(requests: Iterable[Request]) -> list[int]: + trie = PrefixTrie() + overlaps: list[int] = [] + for request in requests: + overlaps.append(trie.longest_prefix(request.tokens)) + trie.insert(request.tokens) + return overlaps +