From 73a3b6f49c79c1f3e68d0cbe2c2a86619053387b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 22 Jun 2026 02:59:48 +0000 Subject: [PATCH] feat(early-exit): CALM-style adaptive per-token threshold + arch gating (v0.27.0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EarlyExitDecoder gains an opt-in CALM-style (Schuster et al. 2022) adaptive confidence threshold that decays geometrically over decoding steps from confidence_threshold toward min_confidence: early tokens need high confidence (an early mistake propagates), later tokens exit more readily. effective_ confidence(step) exposes the schedule; the static path is unchanged. The constructor now validates its inputs; generate() stats gain adaptive + final_confidence_threshold. auto_profile gains _early_exit_suitable() arch gating — encoder-style families (BERT/RoBERTa/T5/...) and sub-6-layer models fall back to vanilla with an explanatory rationale, for both early_exit and layered_early_exit. wrapper exposes adaptive_early_exit (parallel to adaptive_gamma) so the feature is reachable from wrap_model. New tests/test_early_exit.py takes the module 21% -> 100%; 3 arch-gate tests added. Suite: 732 passed, 4 HF-gated skipped. Version 0.26.0 -> 0.27.0. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TmSqgYqcYfzPnkpDVR5vdw --- CHANGELOG.md | 40 +++++++++ PLAN.md | 17 +++- kairu/__init__.py | 2 +- kairu/auto_profile.py | 45 ++++++++++- kairu/early_exit.py | 104 ++++++++++++++++++------ kairu/wrapper.py | 2 + pyproject.toml | 2 +- tests/test_auto_profile.py | 25 ++++++ tests/test_early_exit.py | 162 +++++++++++++++++++++++++++++++++++++ 9 files changed, 371 insertions(+), 28 deletions(-) create mode 100644 tests/test_early_exit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 806f40b..abc908c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,46 @@ All notable changes to Kairu follow [Conventional Commits](https://www.conventio --- +## [0.27.0] — 2026-06-22 + +### Added — Adaptive per-token early exit + early-exit architecture gating + +**`kairu/early_exit.py`** — `EarlyExitDecoder` gains a CALM-style (Schuster et +al. 2022, "Confident Adaptive Language Modeling") per-token *adaptive* threshold, +opt-in via `adaptive=True` (the static path is unchanged and bit-for-bit +identical): + +- The effective confidence bar **decays geometrically** over decoding steps: + `threshold(t) = min_confidence + (confidence_threshold - min_confidence) * + exp(-adapt_decay * t)`. Early tokens must clear a high bar (an early mistake + propagates through every later token); the bar relaxes toward `min_confidence` + as decoding proceeds. New `effective_confidence(step)` exposes the schedule. +- `min_confidence` is clamped to `confidence_threshold` so the schedule is always + monotone non-increasing. The constructor now **validates** its inputs (the + bars, the entropy floor, the decay rate) instead of accepting nonsense. +- `generate()` stats gain `adaptive` and `final_confidence_threshold` (the + effective bar at the step generation stopped). + +**`kairu/auto_profile.py`** — early exit is now **architecture-gated**. The new +`_early_exit_suitable(model, name)` rules out encoder-style families +(BERT/RoBERTa/ELECTRA/DeBERTa/T5/encoder — no left-to-right confidence +trajectory) and models shallower than `_MIN_EARLY_EXIT_LAYERS` (6) layers, which +have too little depth for an exit head to save compute. Both the `early_exit` +and `layered_early_exit` recommendations fall back to `vanilla` (with an +explanatory rationale) when the gate fails. + +**`kairu/wrapper.py`** — `ModelWrapper` / `wrap_model` gain an +`adaptive_early_exit: bool = False` flag (parallel to `adaptive_gamma`) that +forwards to the decoder, making the feature reachable from the public entry +point. + +**Tests:** new `tests/test_early_exit.py` (14 tests — validation, the decay +schedule, every exit reason, and an adaptive-vs-static divergence case) takes +the module from 21% → **100%** coverage; 3 new arch-gate tests in +`tests/test_auto_profile.py`. Suite: **732 passed**, 4 HF-gated skipped. + +--- + ## [0.26.0] — 2026-06-21 ### Added — Attention-weighted KV eviction + INT8/INT4 quantised storage tier diff --git a/PLAN.md b/PLAN.md index 68019cf..ec1178a 100644 --- a/PLAN.md +++ b/PLAN.md @@ -2,7 +2,7 @@ > 流 · *to flow, to stream* -Current version: **v0.26.0** +Current version: **v0.27.0** --- @@ -81,6 +81,21 @@ 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.27.0 — Adaptive per-token early exit + arch-suitability gating *(DONE)* + +- **`kairu/early_exit.py`** — CALM-style (Schuster et al. 2022) per-token + adaptive confidence threshold, opt-in via `adaptive=True` (static path + unchanged). The bar decays geometrically from `confidence_threshold` toward + `min_confidence` over decoding steps via `effective_confidence(step)`; early + tokens need high confidence, later tokens exit more readily. Constructor now + validates inputs; stats gain `adaptive` + `final_confidence_threshold`. +- **`kairu/auto_profile.py`** — `_early_exit_suitable` gates early exit out of + encoder-style architectures (BERT/RoBERTa/T5/…) and sub-6-layer models, + falling back to `vanilla` with an explanatory rationale. +- **`kairu/wrapper.py`** — `adaptive_early_exit` flag (parallel to + `adaptive_gamma`) forwards to the decoder from the public entry point. +- 17 new tests; new `tests/test_early_exit.py` takes the module 21% → 100%. + ### ✅ v0.26.0 — Attention-weighted KV eviction + INT8/INT4 quant tier *(DONE)* - **`kairu/kv_cache.py`** — upgrades `LogitsCache` from plain-recency LRU with diff --git a/kairu/__init__.py b/kairu/__init__.py index c1252df..d64357f 100644 --- a/kairu/__init__.py +++ b/kairu/__init__.py @@ -28,7 +28,7 @@ from __future__ import annotations -__version__ = "0.26.0" +__version__ = "0.27.0" from kairu.adversarial import ( AdversarialMatch, diff --git a/kairu/auto_profile.py b/kairu/auto_profile.py index cd32873..4cc7276 100644 --- a/kairu/auto_profile.py +++ b/kairu/auto_profile.py @@ -29,6 +29,34 @@ re.IGNORECASE, ) _DRAFT_HINTS = re.compile(r"\b(tiny|small|mini|draft|125m|350m)\b", re.IGNORECASE) +# Encoder-style / non-causal architectures do not benefit from confidence +# early exit — there is no left-to-right confidence trajectory to ride. +_EARLY_EXIT_UNSUITABLE = re.compile( + r"\b(bert|roberta|electra|distilbert|deberta|t5|encoder)\b", + re.IGNORECASE, +) +# 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 + + +def _early_exit_suitable(model: ModelInterface, name: str) -> tuple[bool, str]: + """Gate whether early exit is architecturally sensible for ``model``. + + Returns ``(suitable, reason)``; ``reason`` is empty when suitable and + otherwise explains why early exit was ruled out (for the profile rationale). + """ + if name and _EARLY_EXIT_UNSUITABLE.search(name): + return False, f"'{name}' is an encoder-style architecture (no exit trajectory)" + depth = getattr(model, "num_layers", None) + if callable(depth): # LayeredModelInterface exposes num_layers() as a method + depth = depth() + if isinstance(depth, int) and depth < _MIN_EARLY_EXIT_LAYERS: + return ( + False, + f"depth {depth} < {_MIN_EARLY_EXIT_LAYERS} layers — too shallow to exit", + ) + return True, "" @dataclass(frozen=True) @@ -86,7 +114,9 @@ def recommend( ), ) - if is_layered: + ee_suitable, ee_reason = _early_exit_suitable(model, name) + + if is_layered and ee_suitable: return DecoderProfile( strategy="layered_early_exit", gamma=1, @@ -97,7 +127,7 @@ def recommend( rationale="model exposes per-layer logits → architecture-aware early exit", ) - if vocab >= 5_000: + if vocab >= 5_000 and ee_suitable: return DecoderProfile( strategy="early_exit", gamma=1, @@ -108,6 +138,17 @@ def recommend( rationale=f"mid-size single-backbone model (vocab={vocab}) → confidence early exit", ) + if vocab >= 5_000 and not ee_suitable: + return DecoderProfile( + strategy="vanilla", + gamma=1, + early_exit_threshold=1.0, + temperature=1.0, + use_cache=True, + cache_capacity=128, + rationale=f"mid-size model (vocab={vocab}) but {ee_reason} → vanilla decoding", + ) + return DecoderProfile( strategy="vanilla", gamma=1, diff --git a/kairu/early_exit.py b/kairu/early_exit.py index 0a4827b..7c1a8f7 100644 --- a/kairu/early_exit.py +++ b/kairu/early_exit.py @@ -1,15 +1,38 @@ +"""Dynamic early-exit decoding — halt generation once the model is confident. + +Two threshold regimes are available: + +* **Static** (default) — a fixed ``confidence_threshold`` and ``entropy_floor``, + identical to prior releases. +* **Adaptive** (``adaptive=True``) — a CALM-style per-token *decaying* confidence + threshold (Schuster et al. 2022, "Confident Adaptive Language Modeling"). + Early tokens must clear a high bar to trigger an exit (an early mistake + propagates through every later token), while the bar relaxes geometrically + toward ``min_confidence`` as decoding proceeds:: + + threshold(t) = min_confidence + + (confidence_threshold - min_confidence) * exp(-decay * t) + + At ``t = 0`` the threshold equals ``confidence_threshold``; as ``t → ∞`` it + approaches ``min_confidence``. The entropy floor is unchanged — confidence is + the adaptive signal, entropy remains the complementary static guard. +""" + +from __future__ import annotations + +import math + import numpy as np + from kairu.base import ModelInterface class EarlyExitDecoder: - """ - Dynamic early exit: stop generating once the model is sufficiently - confident — either the top-token probability exceeds confidence_threshold - OR the distribution entropy drops below entropy_floor. + """Stop generating once the model is sufficiently confident. - Both conditions measure the same signal (certainty) from complementary - angles: high max-prob → peaked distribution; low entropy → peaked distribution. + An exit fires when the top-token probability clears the (optionally + adaptive) confidence threshold OR the distribution entropy drops below + ``entropy_floor`` — both measure certainty from complementary angles. """ def __init__( @@ -18,11 +41,28 @@ def __init__( confidence_threshold: float = 0.9, entropy_floor: float = 0.5, temperature: float = 1.0, + *, + adaptive: bool = False, + min_confidence: float = 0.5, + adapt_decay: float = 0.2, ): + if not 0.0 < confidence_threshold <= 1.0: + raise ValueError("confidence_threshold must be in (0, 1]") + if not 0.0 <= min_confidence <= 1.0: + raise ValueError("min_confidence must be in [0, 1]") + if entropy_floor < 0.0: + raise ValueError("entropy_floor must be >= 0") + if adapt_decay < 0.0: + raise ValueError("adapt_decay must be >= 0") self.model = model self.confidence_threshold = confidence_threshold self.entropy_floor = entropy_floor self.temperature = temperature + self.adaptive = adaptive + # The relaxed floor can never exceed the base bar, or the schedule would + # rise instead of decay — clamp so it stays monotone non-increasing. + self.min_confidence = min(min_confidence, confidence_threshold) + self.adapt_decay = adapt_decay self._rng = np.random.default_rng(42) def _softmax(self, x: np.ndarray) -> np.ndarray: @@ -35,41 +75,54 @@ def _entropy(self, probs: np.ndarray) -> float: p = np.clip(probs, 1e-10, 1.0) return float(-np.sum(p * np.log(p))) - def _should_exit(self, probs: np.ndarray) -> bool: - top_prob = float(probs.max()) - entropy = self._entropy(probs) - return top_prob >= self.confidence_threshold or entropy <= self.entropy_floor + def effective_confidence(self, step: int) -> float: + """The confidence threshold in force at generation ``step`` (0-based). + + Constant in static mode; a decaying geometric schedule from + ``confidence_threshold`` down toward ``min_confidence`` when adaptive. + """ + if not self.adaptive: + return self.confidence_threshold + span = self.confidence_threshold - self.min_confidence + return self.min_confidence + span * math.exp(-self.adapt_decay * step) + + def _should_exit(self, probs: np.ndarray, step: int) -> tuple[bool, str]: + """Return ``(exit_now, reason)`` for the distribution at ``step``.""" + if float(probs.max()) >= self.effective_confidence(step): + return True, "confidence" + if self._entropy(probs) <= self.entropy_floor: + return True, "entropy" + return False, "" def generate( self, prompt_ids: list[int], max_new_tokens: int = 50, ) -> tuple[list[int], dict]: - """ - Generate tokens, stopping early when the model is sufficiently confident. + """Generate tokens, halting early once the model is confident enough. - Returns: - generated_ids: list of generated token ids - stats: {tokens_generated, exit_reason, max_new_tokens, early_exit} + Returns ``(generated_ids, stats)`` where ``stats`` carries + ``tokens_generated``, ``exit_reason``, ``max_new_tokens``, + ``early_exit``, ``adaptive``, and ``final_confidence_threshold`` (the + effective threshold at the step where generation stopped). """ tokens = list(prompt_ids) generated: list[int] = [] exit_reason = "max_tokens" + threshold = self.confidence_threshold - for _ in range(max_new_tokens): + for step in range(max_new_tokens): logits = self.model.next_token_logits(tokens) probs = self._softmax(logits / max(self.temperature, 1e-8)) + threshold = self.effective_confidence(step) - if self._should_exit(probs): - # Emit the argmax token (most confident choice), then halt + exit_now, reason = self._should_exit(probs, step) + if exit_now: + # Emit the argmax token (most confident choice), then halt. tok = int(np.argmax(probs)) generated.append(tok) tokens.append(tok) - exit_reason = ( - "confidence" - if float(probs.max()) >= self.confidence_threshold - else "entropy" - ) + exit_reason = reason break tok = int(self._rng.choice(len(probs), p=probs)) @@ -81,5 +134,10 @@ def generate( "exit_reason": exit_reason, "max_new_tokens": max_new_tokens, "early_exit": exit_reason != "max_tokens", + "adaptive": self.adaptive, + "final_confidence_threshold": threshold, } return generated, stats + + +__all__ = ["EarlyExitDecoder"] diff --git a/kairu/wrapper.py b/kairu/wrapper.py index d37d1c1..a1bb40b 100644 --- a/kairu/wrapper.py +++ b/kairu/wrapper.py @@ -51,6 +51,7 @@ def __init__( temperature: float = 1.0, cache_capacity: int = 0, adaptive_gamma: bool = False, + adaptive_early_exit: bool = False, ): if cache_capacity > 0: model = CachedModel(model, cache_capacity=cache_capacity) @@ -78,6 +79,7 @@ def __init__( model=model, confidence_threshold=early_exit_threshold, temperature=temperature, + adaptive=adaptive_early_exit, ) @property diff --git a/pyproject.toml b/pyproject.toml index 629e072..fa13771 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "kairu" -version = "0.26.0" +version = "0.27.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 78a4345..96a1660 100644 --- a/tests/test_auto_profile.py +++ b/tests/test_auto_profile.py @@ -85,3 +85,28 @@ def test_recommendation_is_deterministic(): a = AutoProfile.recommend(MockModel(), name_hint="qwen", has_draft=False) b = AutoProfile.recommend(MockModel(), name_hint="qwen", has_draft=False) assert a == b + + +def test_early_exit_gated_for_encoder_architecture(): + """A mid-size model is normally early_exit, but an encoder-style name + (no left-to-right confidence trajectory) must fall back to vanilla.""" + + class Mid(MockModel): + @property + def vocab_size(self) -> int: + return 32_000 + + p = AutoProfile.recommend(Mid(), name_hint="bert-large") + assert p.strategy == "vanilla" + assert "encoder" in p.rationale.lower() + + +def test_layered_gated_when_too_shallow(): + """A layered model below the minimum depth is too shallow to exit.""" + p = AutoProfile.recommend(MockLayeredModel(num_layers=4)) + assert p.strategy == "vanilla" + + +def test_layered_deep_enough_still_early_exits(): + p = AutoProfile.recommend(MockLayeredModel(num_layers=12)) + assert p.strategy == "layered_early_exit" diff --git a/tests/test_early_exit.py b/tests/test_early_exit.py new file mode 100644 index 0000000..07fad45 --- /dev/null +++ b/tests/test_early_exit.py @@ -0,0 +1,162 @@ +"""Tests for kairu.early_exit — static + CALM-style adaptive early exit.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from kairu.base import ModelInterface +from kairu.early_exit import EarlyExitDecoder + + +class FixedLogitsModel(ModelInterface): + """Returns the same logits at every step — makes exit decisions deterministic.""" + + def __init__(self, logits) -> None: + self._logits = np.asarray(logits, dtype=np.float32) + + @property + def vocab_size(self) -> int: + return int(self._logits.size) + + def next_token_logits(self, token_ids: list[int]) -> np.ndarray: + return self._logits + + def max_seq_len(self) -> int: + return 4096 + + +# --------------------------------------------------------------------------- # +# Construction validation +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "kwargs", + [ + {"confidence_threshold": 0.0}, + {"confidence_threshold": 1.5}, + {"min_confidence": -0.1}, + {"min_confidence": 1.1}, + {"entropy_floor": -1.0}, + {"adapt_decay": -0.5}, + ], +) +def test_rejects_invalid_params(kwargs): + with pytest.raises(ValueError): + EarlyExitDecoder(FixedLogitsModel([1.0, 1.0]), **kwargs) + + +def test_min_confidence_clamped_to_base(): + dec = EarlyExitDecoder( + FixedLogitsModel([1.0, 1.0]), + confidence_threshold=0.6, + adaptive=True, + min_confidence=0.9, # above the base → must clamp down to 0.6 + ) + assert dec.min_confidence == pytest.approx(0.6) + # span collapses to zero → schedule is constant at the base. + assert dec.effective_confidence(0) == pytest.approx(0.6) + assert dec.effective_confidence(50) == pytest.approx(0.6) + + +# --------------------------------------------------------------------------- # +# Effective threshold schedule +# --------------------------------------------------------------------------- # + + +def test_static_threshold_is_constant(): + dec = EarlyExitDecoder(FixedLogitsModel([1.0, 1.0]), confidence_threshold=0.8) + assert dec.effective_confidence(0) == 0.8 + assert dec.effective_confidence(100) == 0.8 + + +def test_adaptive_threshold_decays_monotonically_toward_floor(): + dec = EarlyExitDecoder( + FixedLogitsModel([1.0, 1.0]), + confidence_threshold=0.9, + adaptive=True, + min_confidence=0.5, + adapt_decay=0.5, + ) + assert dec.effective_confidence(0) == pytest.approx(0.9) # base at t=0 + assert dec.effective_confidence(0) > dec.effective_confidence(1) + assert dec.effective_confidence(1) > dec.effective_confidence(5) + # As t → ∞ the schedule approaches the floor. + assert dec.effective_confidence(1000) == pytest.approx(0.5, abs=1e-3) + + +# --------------------------------------------------------------------------- # +# Generation / exit reasons +# --------------------------------------------------------------------------- # + + +def test_confidence_exit_on_peaked_distribution(): + dec = EarlyExitDecoder( + FixedLogitsModel([10.0, 0.0, 0.0, 0.0]), confidence_threshold=0.9 + ) + gen, stats = dec.generate([1], max_new_tokens=5) + assert stats["exit_reason"] == "confidence" + assert stats["early_exit"] is True + assert stats["tokens_generated"] == 1 + assert gen[0] == 0 # emits the argmax token before halting + + +def test_entropy_exit_when_confidence_bar_unreached(): + # top prob 0.85 < 0.99 (no confidence exit) but entropy ≈ 0.59 < floor 1.0. + soft = FixedLogitsModel(np.log([0.85, 0.05, 0.05, 0.05])) + dec = EarlyExitDecoder(soft, confidence_threshold=0.99, entropy_floor=1.0) + _, stats = dec.generate([1], max_new_tokens=5) + assert stats["exit_reason"] == "entropy" + assert stats["early_exit"] is True + + +def test_runs_to_max_tokens_when_uncertain(): + uniform = FixedLogitsModel([1.0, 1.0, 1.0, 1.0]) # top 0.25, entropy ln4 ≈ 1.39 + dec = EarlyExitDecoder(uniform, confidence_threshold=0.9, entropy_floor=0.5) + gen, stats = dec.generate([1], max_new_tokens=4) + assert stats["exit_reason"] == "max_tokens" + assert stats["early_exit"] is False + assert stats["tokens_generated"] == 4 + assert len(gen) == 4 + + +def test_stats_carry_adaptive_metadata(): + dec = EarlyExitDecoder(FixedLogitsModel([1.0, 1.0, 1.0, 1.0]), entropy_floor=0.5) + _, stats = dec.generate([1], max_new_tokens=2) + assert set(stats) >= { + "tokens_generated", + "exit_reason", + "max_new_tokens", + "early_exit", + "adaptive", + "final_confidence_threshold", + } + assert stats["adaptive"] is False + assert stats["final_confidence_threshold"] == pytest.approx(0.9) + + +def test_adaptive_exits_where_static_would_not(): + """A mid-confidence (0.7) model never clears a static 0.9 bar, but the + adaptive schedule relaxes below 0.7 and exits at step 2.""" + model = FixedLogitsModel(np.log([0.7, 0.1, 0.1, 0.1])) + + static = EarlyExitDecoder(model, confidence_threshold=0.9, entropy_floor=0.0) + _, s_stats = static.generate([1], max_new_tokens=5) + assert s_stats["exit_reason"] == "max_tokens" + + adaptive = EarlyExitDecoder( + model, + confidence_threshold=0.9, + entropy_floor=0.0, + adaptive=True, + min_confidence=0.5, + adapt_decay=0.5, + ) + gen, a_stats = adaptive.generate([1], max_new_tokens=5) + assert a_stats["adaptive"] is True + assert a_stats["exit_reason"] == "confidence" + # threshold(2) = 0.5 + 0.4·e^-1 ≈ 0.647 ≤ 0.7 → exits at step index 2. + assert a_stats["tokens_generated"] == 3 + assert a_stats["final_confidence_threshold"] == pytest.approx(0.6472, abs=1e-3) + assert gen[-1] == 0 # final token is the argmax at the exit step