feat(early-exit): CALM-style adaptive per-token threshold + arch gating (v0.27.0)

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

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

kairu/__init__.py

@@ -28,7 +28,7 @@
 
 from __future__ import annotations
 
-__version__ = "0.26.0"
+__version__ = "0.27.0"
 
 from kairu.adversarial import (
     AdversarialMatch,

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,

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"]

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

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"

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"