From 940f54793f3343dbb7a3bb4bbaf803cc23a60d92 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 22 Jun 2026 10:52:19 +0000 Subject: [PATCH] feat(bench): SPEED-Bench task splits + speculative spec/quant warnings (v0.28.0) New kairu/speed_bench.py runs the benchmark across named semantic task splits (SPEED-Bench methodology) because speculative speedup and early-exit savings are strongly task-dependent. run_speed_bench -> SpeedBenchReport reports per-split throughput plus a throughput coefficient of variation that quantifies task-sensitivity (0 -> uniform). Six DEFAULT_SPLITS (translation, summarization, qa, code, dialogue, math), each with a deterministic prompt so MockModel runs offline. BenchmarkRunner gains an optional prompt arg to drive the splits. auto_profile: DecoderProfile.warnings + recommend(quant=, draft_kind=) flag speculative configs that erode the expected speedup -- a 4-bit draft (tanks draft-target acceptance) and a tree-structured draft (the sequential verifier cannot exploit it). Empty unless the hints are supplied, so existing recommendations are unchanged. New tests/test_speed_bench.py (100% module) + 5 warning tests. Suite: 746 passed, 4 HF-gated skipped. Version 0.27.0 -> 0.28.0. Closes the session's Discovery sweep. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TmSqgYqcYfzPnkpDVR5vdw --- CHANGELOG.md | 42 +++++++++++ PLAN.md | 18 ++++- kairu/__init__.py | 14 +++- kairu/auto_profile.py | 31 ++++++++ kairu/bench.py | 12 ++- kairu/speed_bench.py | 149 +++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- tests/test_auto_profile.py | 39 ++++++++++ tests/test_speed_bench.py | 94 +++++++++++++++++++++++ 9 files changed, 396 insertions(+), 5 deletions(-) create mode 100644 kairu/speed_bench.py create mode 100644 tests/test_speed_bench.py diff --git a/CHANGELOG.md b/CHANGELOG.md index abc908c..bc12bff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,48 @@ All notable changes to Kairu follow [Conventional Commits](https://www.conventio --- +## [0.28.0] — 2026-06-22 + +### Added — SPEED-Bench semantic task splits + speculative spec/quant warnings + +**New module `kairu/speed_bench.py`** (pure stdlib + `kairu.bench`): +- `run_speed_bench(model, splits=None, ...)` → `SpeedBenchReport` — runs the + benchmark across named semantic task splits (SPEED-Bench methodology) because + speculative speedup and early-exit savings are strongly *task-dependent*: a + single aggregate latency hides the variance that matters at deploy time. +- `TaskSplit` (name + representative prompt + description) and six + `DEFAULT_SPLITS` (translation, summarization, qa, code, dialogue, math), each + with a deterministic integer prompt so `MockModel` runs fully offline. +- `SpeedBenchReport` carries per-split `SplitResult`s, the fastest/slowest split + by tokens/s, the cross-split mean throughput, and a **throughput coefficient + of variation** — one number quantifying how task-sensitive the configuration + is (`0.0` → uniform). All result types expose `as_dict()` for JSON transport. + +**`kairu/bench.py`** — `BenchmarkRunner` gains an optional `prompt` argument +(defaults to the fixed class prompt) so SPEED-Bench can drive distinct task +splits through the same runner. Existing behaviour is unchanged. + +**`kairu/auto_profile.py` — speculative configuration warnings:** +- `DecoderProfile` gains a `warnings: tuple[str, ...] = ()` field. +- `AutoProfile.recommend(..., quant=None, draft_kind=None)` now flags + speculative configs known to erode the expected speedup: a **4-bit draft** + (`int4`/`4bit`/`nf4`/`fp4` — tanks draft–target acceptance) and a + **tree-structured draft** (the sequential verifier here cannot exploit it). + Warnings are empty unless the relevant hints are supplied, so existing + recommendations are unchanged. + +**`kairu/__init__.py`** exports `TaskSplit`, `SplitResult`, `SpeedBenchReport`, +`SPEED_BENCH_SPLITS`, `run_speed_bench`; version `0.27.0 → 0.28.0`. + +**Tests:** new `tests/test_speed_bench.py` (9 tests, module 100%) + 5 warning +tests in `tests/test_auto_profile.py`. Suite: **746 passed**, 4 HF-gated skipped. + +> This closes the last item in this session's researched Discovery sweep +> (CyclicJudge → reliability → KV eviction/quant → adaptive early exit → +> SPEED-Bench), completing both the eval and inference-optimizer tracks. + +--- + ## [0.27.0] — 2026-06-22 ### Added — Adaptive per-token early exit + early-exit architecture gating diff --git a/PLAN.md b/PLAN.md index ec1178a..d096106 100644 --- a/PLAN.md +++ b/PLAN.md @@ -2,7 +2,7 @@ > 流 · *to flow, to stream* -Current version: **v0.27.0** +Current version: **v0.28.0** --- @@ -81,6 +81,22 @@ has shipped in v0.15.0 — the four `🔴` rows below are now **DONE**. error recovery, goal progress/completion, efficiency (steps taken vs. optimal). Returns a per-step breakdown plus an overall trajectory grade. +### ✅ v0.28.0 — SPEED-Bench task splits + speculative spec/quant warnings *(DONE)* + +- **`kairu/speed_bench.py`** — `run_speed_bench` → `SpeedBenchReport` runs the + benchmark across six semantic task splits (`TaskSplit`/`DEFAULT_SPLITS`: + translation, summarization, qa, code, dialogue, math) and reports per-split + throughput + a throughput coefficient of variation quantifying + task-dependence (SPEED-Bench methodology). `BenchmarkRunner` gained an + optional `prompt` arg to drive distinct splits. +- **`kairu/auto_profile.py`** — `DecoderProfile.warnings` + `recommend(quant=, + draft_kind=)` flag 4-bit draft and tree-draft speculative configs that erode + acceptance/speedup (empty unless the hints are supplied). +- 14 new tests; new `tests/test_speed_bench.py` at 100% module coverage. +- **Closes this session's Discovery sweep** — both the eval track (CyclicJudge, + reliability) and the inference-optimizer track (KV eviction/quant, adaptive + early exit, SPEED-Bench) are complete. + ### ✅ v0.27.0 — Adaptive per-token early exit + arch-suitability gating *(DONE)* - **`kairu/early_exit.py`** — CALM-style (Schuster et al. 2022) per-token diff --git a/kairu/__init__.py b/kairu/__init__.py index d64357f..a8345fa 100644 --- a/kairu/__init__.py +++ b/kairu/__init__.py @@ -28,7 +28,7 @@ from __future__ import annotations -__version__ = "0.27.0" +__version__ = "0.28.0" from kairu.adversarial import ( AdversarialMatch, @@ -60,6 +60,13 @@ from kairu.significance import SignificanceResult, paired_t_test, per_criterion_diffs from kairu.auto_profile import AutoProfile, DecoderProfile from kairu.bench import BenchmarkResult, BenchmarkRunner +from kairu.speed_bench import ( + DEFAULT_SPLITS as SPEED_BENCH_SPLITS, + SpeedBenchReport, + SplitResult, + TaskSplit, + run_speed_bench, +) from kairu.benchmarks import BENCHMARKS, BenchmarkStats, percentile_rank from kairu.budget import TokenBudget from kairu.ci_regression import ( @@ -254,6 +261,11 @@ "HFTokenizer", "BenchmarkRunner", "BenchmarkResult", + "TaskSplit", + "SplitResult", + "SpeedBenchReport", + "SPEED_BENCH_SPLITS", + "run_speed_bench", "create_app", "ServerConfig", "RateLimiter", diff --git a/kairu/auto_profile.py b/kairu/auto_profile.py index 4cc7276..86ad733 100644 --- a/kairu/auto_profile.py +++ b/kairu/auto_profile.py @@ -38,6 +38,30 @@ # Below this depth there are too few layers for an early-exit head to skip # enough compute to be worth the confidence bookkeeping. _MIN_EARLY_EXIT_LAYERS = 6 +# Quantisation tiers at which draft-model quality erodes enough to threaten +# speculative acceptance. +_FOUR_BIT_QUANTS = frozenset({"int4", "4bit", "nf4", "fp4"}) + + +def _speculative_warnings(quant: str | None, draft_kind: str | None) -> tuple[str, ...]: + """Flag speculative-decoding configs known to erode the expected speedup. + + Returns a (possibly empty) tuple of human-readable caveats — a 4-bit draft + that tanks draft–target agreement, or a tree-structured draft the sequential + verifier here cannot exploit. + """ + warnings: list[str] = [] + if quant is not None and quant.lower() in _FOUR_BIT_QUANTS: + warnings.append( + "4-bit draft quantisation sharply lowers draft–target agreement; " + "speculative speedup can vanish — benchmark acceptance before deploying" + ) + if draft_kind is not None and draft_kind.lower() == "tree": + warnings.append( + "tree-structured drafts need tree-aware verification; the sequential " + "verifier here realises none of their speedup" + ) + return tuple(warnings) def _early_exit_suitable(model: ModelInterface, name: str) -> tuple[bool, str]: @@ -70,6 +94,9 @@ class DecoderProfile: use_cache: bool cache_capacity: int rationale: str + # Non-fatal caveats about the chosen configuration (e.g. a 4-bit draft that + # may erode speculative acceptance). Empty for unconfigured recommendations. + warnings: tuple[str, ...] = () class AutoProfile: @@ -80,6 +107,9 @@ def recommend( model: ModelInterface, name_hint: str | None = None, has_draft: bool = False, + *, + quant: str | None = None, + draft_kind: str | None = None, ) -> DecoderProfile: from kairu.layered import LayeredModelInterface # local — avoid cycle @@ -112,6 +142,7 @@ def recommend( f"{' / frontier family' if _FRONTIER_FAMILIES.search(name) else ''}" f" with draft → speculative γ={gamma}" ), + warnings=_speculative_warnings(quant, draft_kind), ) ee_suitable, ee_reason = _early_exit_suitable(model, name) diff --git a/kairu/bench.py b/kairu/bench.py index 751fbf3..f2801e6 100644 --- a/kairu/bench.py +++ b/kairu/bench.py @@ -244,9 +244,17 @@ class BenchmarkRunner: # A fixed prompt so every run starts from the same context. _PROMPT: list[int] = [1, 2, 3, 4, 5] - def __init__(self, model: ModelInterface, name: str = "benchmark") -> None: + def __init__( + self, + model: ModelInterface, + name: str = "benchmark", + prompt: list[int] | None = None, + ) -> None: self._model = model self.name = name + # Per-runner prompt lets SPEED-Bench drive distinct task splits through + # the same runner; defaults to the fixed class prompt for back-compat. + self._prompt = list(prompt) if prompt else list(self._PROMPT) def run( self, @@ -274,7 +282,7 @@ def run( t0 = time.perf_counter() # Consume the full stream — each generated token is timed as part of # the per-run total latency. - for _ in decoder.stream(self._PROMPT, max_new_tokens=num_tokens): + for _ in decoder.stream(self._prompt, max_new_tokens=num_tokens): pass elapsed = time.perf_counter() - t0 diff --git a/kairu/speed_bench.py b/kairu/speed_bench.py new file mode 100644 index 0000000..7c2da79 --- /dev/null +++ b/kairu/speed_bench.py @@ -0,0 +1,149 @@ +"""SPEED-Bench-style semantic task splits for benchmarking. + +Speculative-decoding speedup and early-exit savings are strongly +*task-dependent*: a model drafts structured code or repetitive text far more +predictably than open-ended dialogue, so a single aggregate latency number +hides the variance that actually matters at deploy time. Following the +SPEED-Bench methodology, this module runs the benchmark across a set of named +semantic task splits and reports per-split throughput plus a dispersion measure +(coefficient of variation) that quantifies exactly how task-sensitive the +configuration is. + +Pure stdlib + :mod:`kairu.bench` — no ML frameworks. Each split carries a +deterministic integer prompt so :class:`~kairu.mock_model.MockModel` runs fully +offline; real tokenised prompts plug in behind the same :class:`TaskSplit` +contract. +""" + +from __future__ import annotations + +import statistics +from dataclasses import dataclass + +from kairu.base import ModelInterface +from kairu.bench import BenchmarkRunner + + +@dataclass(frozen=True) +class TaskSplit: + """A named semantic task paired with a representative prompt.""" + + name: str + prompt: list[int] + description: str = "" + + +# Distinct deterministic prompts — the split *labels* carry the semantics; the +# token ids merely need to differ so each split exercises a different prefix. +DEFAULT_SPLITS: tuple[TaskSplit, ...] = ( + TaskSplit("translation", [2, 4, 8, 16, 32], "short-context sequence transduction"), + TaskSplit("summarization", [5, 10, 15, 20, 25, 30], "long-context compression"), + TaskSplit("qa", [7, 14, 21], "short factual question answering"), + TaskSplit("code", [1, 1, 2, 3, 5, 8, 13], "structured, low-entropy generation"), + TaskSplit("dialogue", [9, 18, 27, 36], "open-ended conversational turns"), + TaskSplit("math", [3, 9, 27, 81], "step-by-step numeric reasoning"), +) + + +@dataclass(frozen=True) +class SplitResult: + """Per-split latency / throughput summary.""" + + name: str + p50_s: float + mean_s: float + tokens_per_s: float + + def as_dict(self) -> dict: + """JSON-friendly mapping of the split's metrics.""" + return { + "name": self.name, + "p50_s": self.p50_s, + "mean_s": self.mean_s, + "tokens_per_s": self.tokens_per_s, + } + + +@dataclass(frozen=True) +class SpeedBenchReport: + """Cross-split benchmark report quantifying task-dependence.""" + + splits: tuple[SplitResult, ...] + fastest: str + slowest: str + mean_tokens_per_s: float + # Coefficient of variation of throughput across splits — 0.0 means the + # configuration performs uniformly regardless of task; larger means more + # task-sensitive (the speedup you measure depends heavily on the workload). + throughput_cv: float + + def as_dict(self) -> dict: + """JSON-friendly mapping of the full report.""" + return { + "splits": [s.as_dict() for s in self.splits], + "fastest": self.fastest, + "slowest": self.slowest, + "mean_tokens_per_s": self.mean_tokens_per_s, + "throughput_cv": self.throughput_cv, + } + + +def run_speed_bench( + model: ModelInterface, + splits: tuple[TaskSplit, ...] | None = None, + *, + num_tokens: int = 64, + num_runs: int = 20, + warmup: int = 3, +) -> SpeedBenchReport: + """Benchmark ``model`` across semantic task splits. + + Drives one :class:`~kairu.bench.BenchmarkRunner` per split and returns a + :class:`SpeedBenchReport` carrying per-split throughput, the fastest/slowest + split by tokens/s, the cross-split mean throughput, and the throughput + coefficient of variation — a single number capturing how task-sensitive the + configuration is (``0.0`` → uniform across tasks). + """ + chosen = tuple(splits) if splits is not None else DEFAULT_SPLITS + if not chosen: + raise ValueError("at least one task split is required") + + results: list[SplitResult] = [] + for split in chosen: + runner = BenchmarkRunner(model, name=split.name, prompt=split.prompt) + r = runner.run(num_tokens=num_tokens, num_runs=num_runs, warmup=warmup) + results.append( + SplitResult( + name=split.name, + p50_s=r.p50, + mean_s=r.mean, + tokens_per_s=r.tokens_per_s_mean, + ) + ) + + throughputs = [s.tokens_per_s for s in results] + mean_tps = statistics.mean(throughputs) + cv = ( + statistics.stdev(throughputs) / mean_tps + if len(throughputs) >= 2 and mean_tps > 0 + else 0.0 + ) + fastest = max(results, key=lambda s: s.tokens_per_s).name + slowest = min(results, key=lambda s: s.tokens_per_s).name + + return SpeedBenchReport( + splits=tuple(results), + fastest=fastest, + slowest=slowest, + mean_tokens_per_s=mean_tps, + throughput_cv=cv, + ) + + +__all__ = [ + "TaskSplit", + "SplitResult", + "SpeedBenchReport", + "DEFAULT_SPLITS", + "run_speed_bench", +] diff --git a/pyproject.toml b/pyproject.toml index fa13771..51f755d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "kairu" -version = "0.27.0" +version = "0.28.0" description = "Real-time LLM inference optimizer: speculative decoding, layered early exit, KV-cache recycling, adaptive γ, SSE/JSONL streaming, Prometheus metrics, Redis-backed rate limiting, OpenTelemetry tracing, cluster token budgets, token watermarking, squish quantization-quality eval, rubric-based response evaluation + A/B comparison HTTP API, eight named rubrics + prism UI, prompt shield content policy, judge ensemble + CI regression + log-to-eval pipeline, constitutional rubric generation, agentic trajectory scoring, judge bias calibration + A-BB bias bounds" license = { text = "BUSL-1.1" } requires-python = ">=3.10" diff --git a/tests/test_auto_profile.py b/tests/test_auto_profile.py index 96a1660..315fa4c 100644 --- a/tests/test_auto_profile.py +++ b/tests/test_auto_profile.py @@ -110,3 +110,42 @@ def test_layered_gated_when_too_shallow(): def test_layered_deep_enough_still_early_exits(): p = AutoProfile.recommend(MockLayeredModel(num_layers=12)) assert p.strategy == "layered_early_exit" + + +def _big_with_draft(**kwargs): + class Big(MockModel): + @property + def vocab_size(self) -> int: + return 50_000 + + return AutoProfile.recommend( + Big(), name_hint="llama-3-8b", has_draft=True, **kwargs + ) + + +def test_speculative_has_no_warnings_by_default(): + p = _big_with_draft() + assert p.strategy == "speculative" + assert p.warnings == () + + +def test_four_bit_draft_emits_warning(): + p = _big_with_draft(quant="int4") + assert p.strategy == "speculative" + assert any("4-bit" in w for w in p.warnings) + + +def test_tree_draft_emits_warning(): + p = _big_with_draft(draft_kind="tree") + assert any("tree" in w.lower() for w in p.warnings) + + +def test_both_caveats_stack(): + p = _big_with_draft(quant="NF4", draft_kind="Tree") # case-insensitive + assert len(p.warnings) == 2 + + +def test_non_speculative_profile_has_no_warnings(): + # quant/draft_kind only matter on the speculative path. + p = AutoProfile.recommend(MockModel(), quant="int4", draft_kind="tree") + assert p.warnings == () diff --git a/tests/test_speed_bench.py b/tests/test_speed_bench.py new file mode 100644 index 0000000..e7ddde9 --- /dev/null +++ b/tests/test_speed_bench.py @@ -0,0 +1,94 @@ +"""Tests for kairu.speed_bench — SPEED-Bench semantic task splits.""" + +from __future__ import annotations + +import pytest + +from kairu.mock_model import MockModel +from kairu.speed_bench import ( + DEFAULT_SPLITS, + SpeedBenchReport, + SplitResult, + TaskSplit, + run_speed_bench, +) + +# Small, fast splits for the assertions that don't need the full default set. +_SPLITS = ( + TaskSplit("alpha", [1, 2, 3], "first"), + TaskSplit("beta", [4, 5, 6, 7], "second"), +) + + +def _run(splits=_SPLITS): + return run_speed_bench( + MockModel(), splits=splits, num_tokens=5, num_runs=2, warmup=1 + ) + + +def test_default_splits_cover_six_semantic_tasks(): + names = {s.name for s in DEFAULT_SPLITS} + assert names == {"translation", "summarization", "qa", "code", "dialogue", "math"} + assert all(isinstance(s, TaskSplit) and s.prompt for s in DEFAULT_SPLITS) + + +def test_report_has_one_result_per_split(): + report = _run() + assert isinstance(report, SpeedBenchReport) + assert [s.name for s in report.splits] == ["alpha", "beta"] + assert all(s.tokens_per_s > 0 for s in report.splits) + assert all(s.mean_s > 0 for s in report.splits) + + +def test_fastest_and_slowest_are_actual_splits(): + report = _run() + names = {s.name for s in report.splits} + assert report.fastest in names + assert report.slowest in names + # With ≥2 splits the CV is a non-negative dispersion measure. + assert report.throughput_cv >= 0.0 + + +def test_single_split_has_zero_dispersion(): + report = _run(splits=(TaskSplit("solo", [1, 2]),)) + assert len(report.splits) == 1 + assert report.throughput_cv == 0.0 + assert report.fastest == "solo" + assert report.slowest == "solo" + + +def test_default_splits_used_when_none_passed(): + report = run_speed_bench(MockModel(), num_tokens=4, num_runs=2, warmup=1) + assert len(report.splits) == len(DEFAULT_SPLITS) + + +def test_empty_splits_rejected(): + with pytest.raises(ValueError): + run_speed_bench(MockModel(), splits=()) + + +def test_report_round_trips_to_dict(): + report = _run() + d = report.as_dict() + assert set(d) == { + "splits", + "fastest", + "slowest", + "mean_tokens_per_s", + "throughput_cv", + } + assert d["splits"][0].keys() == {"name", "p50_s", "mean_s", "tokens_per_s"} + + +def test_split_result_as_dict(): + sr = SplitResult(name="x", p50_s=1.0, mean_s=2.0, tokens_per_s=3.0) + assert sr.as_dict() == { + "name": "x", + "p50_s": 1.0, + "mean_s": 2.0, + "tokens_per_s": 3.0, + } + + +def test_task_split_description_defaults_empty(): + assert TaskSplit("t", [1]).description == ""