plumbify 0.2.0__py3-none-any.whl

plumbify/core/answers.py

@@ -0,0 +1,75 @@
+"""Option probabilities -> typed answers (choice / noul / score) with a confidence and, optionally, a conformal set."""
+
+from __future__ import annotations
+
+import json
+
+from .row import Row
+
+
+def state_text(state) -> str:
+    """Objects and arrays become labeled text; strings pass through."""
+    if isinstance(state, str):
+        return state
+    if isinstance(state, dict):
+        return "\n".join(
+            f"{k}: {v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)}"
+            for k, v in state.items()
+        )
+    return json.dumps(state, ensure_ascii=False)
+
+
+def choice_confidence(probs: list[float]) -> float:
+    """How far the leading option stands above uniform: (p_max - 1/K) / (1 - 1/K); 1 for a single option."""
+    if len(probs) <= 1:
+        return 1.0
+    u = 1.0 / len(probs)
+    return (max(probs) - u) / (1.0 - u)
+
+
+def score_confidence(probs: list[float]) -> float:
+    """Concentration around the modal level: max(0, 1 - E|level - mode| / D), D = uniform mean absolute deviation."""
+    n = len(probs)
+    if n <= 1:
+        return 1.0
+    mode = max(range(n), key=probs.__getitem__)
+    spread = sum(p * abs(i - mode) for i, p in enumerate(probs))
+    center = (n - 1) / 2
+    d_uniform = sum(abs(i - center) for i in range(n)) / n
+    return max(0.0, 1.0 - spread / d_uniform)
+
+
+def answer(row: Row, probs: list[float], qhat: float | None = None) -> dict:
+    """Probabilities over ``row.options`` (in order) -> the answer object for ``row.qtype``. With a conformal
+    threshold ``qhat``, the answer also carries ``set``: the options that cannot be ruled out."""
+    out = _typed(row, probs)
+    if qhat is not None:
+        from ..calibration import prediction_set
+
+        out["set"] = [row.options[k].name for k in prediction_set(probs, qhat)]
+    return out
+
+
+def _typed(row: Row, probs: list[float]) -> dict:
+    named = {o.name: round(p, 6) for o, p in zip(row.options, probs, strict=True)}
+    if row.qtype == "noul":
+        return {
+            "type": "noul",
+            "noul": named["yes"],
+            "probabilities": named,
+            "confidence": round(choice_confidence(probs), 6),
+        }
+    if row.qtype == "choice":
+        return {
+            "type": "choice",
+            "choice": max(named, key=named.get),
+            "probabilities": named,
+            "confidence": round(choice_confidence(probs), 6),
+        }
+    return {
+        "type": "score",
+        "score": round(sum(i * p for i, p in enumerate(probs)), 4),
+        "legend": {o.name: o.desc for o in row.options},
+        "probabilities": named,
+        "confidence": round(score_confidence(probs), 6),
+    }

plumbify/core/decision_head.py

@@ -0,0 +1,178 @@
+"""DecisionHead: a small transformer that reads a FROZEN trunk and returns a distribution over the supplied options.
+
+Input per question: the trunk's hidden states at the decision suffix only (question, option and DECIDE tokens), taken
+from several layers ("taps"). The state itself is never re-read here: the trunk already attended over it, and not
+touching state positions is what lets a decision reuse a generator's KV cache.
+
+    taps [S, T, d] --per-tap LayerNorm, concat, Linear--> x [S, d_h] (+ segment embedding: question/option/decide)
+      --token blocks (bidirectional attention over the suffix)-->
+      --pool each option: Linear([last ; mean])--> o_k,   q = x[DECIDE]
+      --set blocks over {q, o_1..o_K} (no positions: order-invariant)-->
+      z_k = <W_q q, W_o o_k> / sqrt(d_proj) + u . o_k
+
+Dense on purpose: the head sees a few hundred tokens and costs ~1% of a trunk forward, so capacity is cheap and exact
+attention beats linear-attention or MoE variants at this length.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+SEG_QUESTION, SEG_OPTION, SEG_DECIDE, SEG_OTHER = 0, 1, 2, 3
+
+
+@dataclass
+class HeadConfig:
+    d: int  # trunk hidden size
+    n_taps: int  # number of trunk layers read
+    d_h: int = 1024
+    n_heads: int = 16
+    token_blocks: int = 1
+    set_blocks: int = 2
+    d_proj: int = 512
+    dropout: float = 0.1
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+class Block(nn.Module):
+    """Pre-norm transformer block with a key-padding mask (bidirectional)."""
+
+    def __init__(self, d: int, n_heads: int, dropout: float):
+        super().__init__()
+        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.qkv = nn.Linear(d, 3 * d)
+        self.out = nn.Linear(d, d)
+        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
+        self.h, self.p = n_heads, dropout
+
+    def forward(self, x: torch.Tensor, keep: torch.Tensor) -> torch.Tensor:
+        """x [B, S, d]; keep [B, S] bool (False = padding)."""
+        B, S, d = x.shape
+        q, k, v = self.qkv(self.n1(x)).view(B, S, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
+        mask = keep[:, None, None, :]  # attend only to real positions
+        a = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, dropout_p=self.p if self.training else 0.0
+        )
+        x = x + self.out(a.transpose(1, 2).reshape(B, S, d))
+        return x + self.mlp(self.n2(x))
+
+
+@dataclass
+class SuffixBatch:
+    """Padded decision suffixes. feats [B, S, T, d]; seg [B, S] segment ids; keep [B, S];
+    opt_spans per row [(start, end)] within the suffix; decide per row (index within the suffix)."""
+
+    feats: torch.Tensor
+    seg: torch.Tensor
+    keep: torch.Tensor
+    opt_spans: list[list[tuple[int, int]]]
+    decide: list[int]
+
+    @staticmethod
+    def collate(items: list[dict], device=None) -> SuffixBatch:
+        """items: {"feats": [S, T, d] tensor, "opt_spans": [...], "decide": int} with spans relative to the suffix."""
+        B, S = len(items), max(it["feats"].shape[0] for it in items)
+        T, d = items[0]["feats"].shape[1:]
+        dev = device or items[0]["feats"].device
+        feats = torch.zeros(B, S, T, d, dtype=items[0]["feats"].dtype, device=dev)
+        seg = torch.full((B, S), SEG_OTHER, dtype=torch.long, device=dev)
+        keep = torch.zeros(B, S, dtype=torch.bool, device=dev)
+        for b, it in enumerate(items):
+            n = it["feats"].shape[0]
+            feats[b, :n] = it["feats"].to(dev)
+            keep[b, :n] = True
+            seg[b, :n] = SEG_QUESTION
+            for s, e in it["opt_spans"]:
+                seg[b, s:e] = SEG_OPTION
+            seg[b, it["decide"]] = SEG_DECIDE
+        return SuffixBatch(
+            feats, seg, keep, [it["opt_spans"] for it in items], [it["decide"] for it in items]
+        )
+
+
+class DecisionHead(nn.Module):
+    def __init__(self, cfg: HeadConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.tap_norm = nn.ModuleList(nn.LayerNorm(cfg.d) for _ in range(cfg.n_taps))
+        self.inp = nn.Linear(cfg.n_taps * cfg.d, cfg.d_h)
+        self.seg = nn.Embedding(4, cfg.d_h)
+        self.tok_blocks = nn.ModuleList(
+            Block(cfg.d_h, cfg.n_heads, cfg.dropout) for _ in range(cfg.token_blocks)
+        )
+        self.pool = nn.Linear(2 * cfg.d_h, cfg.d_h)
+        self.set_type = nn.Embedding(2, cfg.d_h)  # 0 = decide, 1 = option
+        self.set_blocks = nn.ModuleList(
+            Block(cfg.d_h, cfg.n_heads, cfg.dropout) for _ in range(cfg.set_blocks)
+        )
+        self.norm = nn.LayerNorm(cfg.d_h)
+        self.wq = nn.Linear(cfg.d_h, cfg.d_proj, bias=False)
+        self.wo = nn.Linear(cfg.d_h, cfg.d_proj, bias=False)
+        self.u = nn.Linear(cfg.d_h, 1, bias=False)
+        nn.init.zeros_(self.u.weight)
+        self.scale = 1.0 / math.sqrt(cfg.d_proj)
+        self.register_buffer("temperature", torch.ones(()))
+
+    def forward(self, b: SuffixBatch) -> torch.Tensor:
+        """Returns logits [B, K_max] in fp32, -inf where a row has fewer options."""
+        with torch.autocast(b.feats.device.type, enabled=False):
+            f = b.feats.float()
+            x = self.inp(
+                torch.cat([n(f[:, :, t]) for t, n in enumerate(self.tap_norm)], -1)
+            ) + self.seg(b.seg)
+            for blk in self.tok_blocks:
+                x = blk(x, b.keep)
+            B, K = x.shape[0], max(len(s) for s in b.opt_spans)
+            # option pooling: [last ; mean] over each span, gathered into a [B, 1 + K] set (decide first)
+            last = x.new_zeros(B, K, x.shape[-1])
+            mean = x.new_zeros(B, K, x.shape[-1])
+            okeep = torch.zeros(B, 1 + K, dtype=torch.bool, device=x.device)
+            okeep[:, 0] = True
+            for i, spans in enumerate(b.opt_spans):  # K is small; spans are ragged
+                for k, (s, e) in enumerate(spans):
+                    last[i, k] = x[i, e - 1]
+                    mean[i, k] = x[i, s:e].mean(0)
+                    okeep[i, 1 + k] = True
+            o = self.pool(torch.cat([last, mean], -1)) + self.set_type.weight[1]
+            q = (
+                x[torch.arange(B, device=x.device), torch.tensor(b.decide, device=x.device)]
+                + self.set_type.weight[0]
+            )
+            s = torch.cat([q[:, None], o], 1)
+            for blk in self.set_blocks:
+                s = blk(s, okeep)
+            s = self.norm(s)
+            q, o = s[:, 0], s[:, 1:]
+            z = (self.wo(o) * self.wq(q)[:, None]).sum(-1) * self.scale + self.u(o).squeeze(-1)
+            return (z / self.temperature).masked_fill(~okeep[:, 1:], float("-inf"))
+
+
+def decision_loss(logits: torch.Tensor, gold, teacher, lam: float = 0.5) -> torch.Tensor:
+    """CE on gold when present; + lam * KL(teacher || p) when a teacher distribution is present.
+    logits [B,K] (-inf padded). gold: list[int|None]. teacher: list[list[float]|None]."""
+    logp = F.log_softmax(logits, -1)
+    B, K = logits.shape
+    terms = []
+    has_gold = [b for b in range(B) if gold[b] is not None]
+    if has_gold:
+        idx = torch.tensor(has_gold, device=logits.device)
+        g = torch.tensor([gold[b] for b in has_gold], device=logits.device)
+        terms.append(-logp[idx, g])
+    has_t = [b for b in range(B) if teacher[b] is not None]
+    if has_t:
+        t = logits.new_zeros((len(has_t), K))
+        for i, b in enumerate(has_t):
+            t[i, : len(teacher[b])] = torch.tensor(teacher[b], dtype=t.dtype)
+        t = t / t.sum(-1, keepdim=True).clamp_min(1e-9)
+        lp = logp[torch.tensor(has_t, device=logits.device)].masked_fill(t == 0, 0.0)
+        terms.append(lam * (t * (t.clamp_min(1e-9).log() - lp)).sum(-1))
+    if not terms:
+        return logits.sum() * 0.0
+    return torch.cat(terms).sum() / len(torch.cat(terms))

plumbify/core/render.py

@@ -0,0 +1,107 @@
+"""Render v2: a decision is a user turn in the model's own chat format, appended after the context.
+
+    <context: a conversation, or a state>  <user turn: question, options, instruction>  <assistant-open>
+    '------------ prefix (cacheable) -----''---------------- decision suffix ---------------------------'
+
+The frozen trunk reads the question in the format it was trained on, and the prefix is exactly what a generator has
+already cached for the same conversation, so a decision only prefills its suffix. The head reads the suffix only:
+option spans, and DECIDE = the last token of the assistant-open (the position where the model would start answering).
+
+Pieces are tokenised separately (prefix, question, each option, closing) so option spans are exact; training and
+serving both go through ``render_decision``, so the token ids always agree.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .row import Row
+
+SENTINEL = "\x00PLUMB_DECISION\x00"
+LEADS = {
+    "choice": ("Options:", "Answer with exactly one option."),
+    "noul": ("Options:", "Answer yes or no."),
+    "score": ("Levels, lowest to highest:", "Answer with exactly one level."),
+}
+
+
+@dataclass
+class Rendered:
+    input_ids: list[int]
+    suffix_start: int  # first token of the decision suffix; everything before is shared context
+    opt_spans: list[tuple[int, int]]  # absolute positions, presented order
+    decide_pos: int
+    perm: list[int]  # presented position -> original option index
+    gold: int | None
+    teacher: list[float] | None
+    qtype: str
+    id: str
+
+    def suffix(self) -> dict:
+        """Spans relative to the suffix, as the head consumes them."""
+        s = self.suffix_start
+        return {
+            "opt_spans": [(a - s, b - s) for a, b in self.opt_spans],
+            "decide": self.decide_pos - s,
+        }
+
+
+def _ids(tok, text: str) -> list[int]:
+    return tok(text, add_special_tokens=False).input_ids
+
+
+def _template(tok, messages: list[dict], **kw) -> tuple[str, str]:
+    """Chat-template text before and after the decision turn's content."""
+    text = tok.apply_chat_template(
+        [*messages, {"role": "user", "content": SENTINEL}],
+        tokenize=False,
+        add_generation_prompt=True,
+        **kw,
+    )
+    head, tail = text.split(SENTINEL)
+    return head, tail
+
+
+def render_decision(
+    tok,
+    row: Row,
+    context: list[dict] | None = None,
+    perm: list[int] | None = None,
+    template_kwargs: dict | None = None,
+) -> Rendered:
+    """Token ids + positions for one decision.
+
+    ``context``: prior chat messages to decide over (the state is then ``row.state`` if non-empty, placed at the top
+    of the decision turn). Without ``context`` the state opens the decision turn. ``perm`` presents options in a
+    given order (training shuffles; never for score).
+    """
+    kw = {"enable_thinking": False, **(template_kwargs or {})}
+    head, tail = _template(tok, list(context or []), **kw)
+    K = len(row.options)
+    perm = list(range(K)) if perm is None or row.qtype == "score" else perm
+    lead, instruction = LEADS[row.qtype]
+    from ..core.answers import state_text
+
+    state = f"{state_text(row.state)}\n\n" if row.state else ""
+    ids = _ids(tok, head + state)  # prefix: template head (+ state), tokenised as one run
+    suffix_start = len(ids)
+    ids += _ids(tok, f"Decision: {state_text(row.question)}\n{lead}\n")
+    spans = []
+    for j in perm:
+        o = row.options[j]
+        s = len(ids)
+        ids += _ids(tok, f"- {o.name}: {o.desc}\n" if o.desc else f"- {o.name}\n")
+        spans.append((s, len(ids)))
+    ids += _ids(tok, instruction + tail)
+    inv = {orig: pos for pos, orig in enumerate(perm)}
+    return Rendered(
+        ids,
+        suffix_start,
+        spans,
+        len(ids) - 1,
+        perm,
+        None if row.gold is None else inv[row.gold],
+        None if row.teacher is None else [row.teacher[j] for j in perm],
+        row.qtype,
+        row.id,
+    )

plumbify/core/row.py

@@ -0,0 +1,45 @@
+"""One decision: a state, a question, typed options, and (for training) a gold index or a teacher distribution."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+QTYPES = ("choice", "noul", "score")
+
+
+@dataclass
+class Option:
+    name: str
+    desc: str = ""
+
+
+@dataclass
+class Row:
+    id: str
+    state: str
+    question: str
+    qtype: str  # choice | noul | score
+    options: list[Option]  # noul: exactly [yes, no]; score: ordered levels low -> high
+    gold: int | None = None  # index into options, or None (teacher-only row)
+    teacher: list[float] | None = None  # distribution over options in the ORIGINAL order, or None
+    meta: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        assert self.qtype in QTYPES, self.qtype
+        if self.qtype == "noul":
+            assert len(self.options) == 2
+        if self.gold is not None:
+            assert 0 <= self.gold < len(self.options)
+        if self.teacher is not None:
+            assert len(self.teacher) == len(self.options)
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), ensure_ascii=False)
+
+    @staticmethod
+    def from_json(s: str) -> Row:
+        d = json.loads(s)
+        d["options"] = [Option(**o) for o in d["options"]]
+        return Row(**d)

plumbify/core/suffix_lora.py

@@ -0,0 +1,124 @@
+"""Suffix-only LoRA: the adapter acts on the decision suffix tokens and nothing else.
+
+Every adapted linear layer computes ``base(x) + mask * B(A(x))`` with ``mask`` = 1 on decision-suffix tokens (question,
+options, DECIDE) and 0 on context tokens (conversation / state). So:
+
+- context tokens are processed by the untouched base weights: their KV and recurrent state are exactly the base model's,
+  shareable with a generator that has already read the same conversation;
+- generation (no suffix) is bit-identical to the base model; the adapter is never merged, and with no mask set it
+  contributes nothing (System1 clears the mask after every inference forward);
+- decision tokens get a trained trunk, reading the base model's context through adapted projections: the computation a
+  frozen trunk plus head cannot learn.
+
+The mask is applied by a forward hook on every ``lora_B`` module, so it holds through gradient checkpointing (the hook
+reruns on recompute with the same mask). Adapters are injected in place (no PeftModel wrapper), so layer paths and tap
+hooks are unchanged. Files: ``suffix_adapter.json`` (LoRA config) + ``suffix_adapter.safetensors`` (weights).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+
+import torch
+from torch import nn
+
+from .targets import coverage, select_targets
+
+CONFIG, WEIGHTS = "suffix_adapter.json", "suffix_adapter.safetensors"
+
+
+class SuffixLoRA:
+    def __init__(
+        self,
+        trunk: nn.Module,
+        r: int = 32,
+        alpha: int = 64,
+        dropout: float = 0.05,
+        targets: list[str] | None = None,
+    ):
+        from peft import LoraConfig, inject_adapter_in_model
+
+        mtype = getattr(trunk.config, "model_type", None)
+        self.targets = list(targets or select_targets(trunk, mtype))
+        cov = coverage(trunk, self.targets)
+        if cov["_layers_total"] and cov["_layers_adapted"] < 0.9 * cov["_layers_total"]:
+            raise ValueError(
+                f"LoRA targets touch {cov['_layers_adapted']}/{cov['_layers_total']} layers"
+            )
+        self.cfg = {
+            "r": r,
+            "lora_alpha": alpha,
+            "lora_dropout": dropout,
+            "target_modules": self.targets,
+        }
+        inject_adapter_in_model(LoraConfig(bias="none", **self.cfg), trunk)
+        self.trunk = trunk
+        self.mask: torch.Tensor | None = (
+            None  # [B, L] for a decision forward; None = adapter off (base model)
+        )
+        self.lora_B = [m for n, m in trunk.named_modules() if n.endswith("lora_B.default")]
+        if not self.lora_B:
+            raise ValueError(f"no LoRA modules were injected (targets={self.targets})")
+        self._hooks = [m.register_forward_hook(self._gate) for m in self.lora_B]
+        dev = next(trunk.parameters()).device
+        for p in self.parameters():  # injected in fp32 on CPU by default; live next to the trunk
+            p.data = p.data.to(dev)
+
+    def _gate(self, _mod, _args, out):
+        m = self.mask
+        if m is None:  # no decision in flight: the base model, e.g. for generation
+            return out * 0
+        if out.dim() == 3 and out.shape[:2] == m.shape:
+            return out * m.unsqueeze(-1).to(out.dtype)
+        if out.dim() == 2 and out.shape[0] == m.numel():
+            return out * m.reshape(-1, 1).to(out.dtype)
+        raise RuntimeError(
+            f"suffix mask {tuple(m.shape)} does not fit adapter output {tuple(out.shape)}"
+        )
+
+    def parameters(self) -> list[nn.Parameter]:
+        return [p for n, p in self.trunk.named_parameters() if ".lora_" in n]
+
+    def trainable(self, on: bool = True) -> SuffixLoRA:
+        for p in self.parameters():
+            p.requires_grad_(on)
+        return self
+
+    @staticmethod
+    def suffix_mask(suffix_starts: list[int], length: int, device) -> torch.Tensor:
+        pos = torch.arange(length, device=device)
+        return (pos[None, :] >= torch.tensor(suffix_starts, device=device)[:, None]).float()
+
+    def save(self, out: str) -> None:
+        from safetensors.torch import save_file
+
+        sd = {
+            n: p.detach().contiguous().cpu()
+            for n, p in self.trunk.named_parameters()
+            if ".lora_" in n
+        }
+        save_file(sd, os.path.join(out, WEIGHTS))
+        with open(os.path.join(out, CONFIG), "w") as f:
+            json.dump({**self.cfg, "scope": "suffix"}, f, indent=1)
+
+    @classmethod
+    def load(cls, trunk: nn.Module, plumb: str) -> SuffixLoRA:
+        from safetensors.torch import load_file
+
+        from ..artifact import fetch
+
+        cfg = json.load(open(fetch(plumb, CONFIG)))
+        lora = cls(trunk, cfg["r"], cfg["lora_alpha"], cfg["lora_dropout"], cfg["target_modules"])
+        sd = load_file(fetch(plumb, WEIGHTS), device=str(next(trunk.parameters()).device))
+        params = dict(trunk.named_parameters())
+        missing = [n for n in params if ".lora_" in n and n not in sd]
+        if missing or any(n not in params for n in sd):
+            raise RuntimeError(
+                f"suffix adapter mismatch: {len(missing)} missing, "
+                f"{sum(n not in params for n in sd)} unexpected"
+            )
+        with torch.no_grad():
+            for n, t in sd.items():
+                params[n].copy_(t)
+        return lora.trainable(False)

plumbify/core/taps.py

@@ -0,0 +1,83 @@
+"""Read a frozen trunk at several layers, at the decision suffix only.
+
+A tap is the residual stream after decoder layer ``i`` (0-based); tap ``-1`` is the trunk's final, normed output. That
+is the same quantity vLLM exposes as auxiliary hidden states (EAGLE3) plus the pooler input, so the head sees the same
+features in training (transformers) and in serving (vLLM).
+
+Hooks slice positions >= ``suffix_start`` inside the forward, so a 32k-token context never materialises 32k x taps
+hidden states.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def decoder_layers(trunk: nn.Module) -> nn.ModuleList:
+    for name in ("layers", "h", "blocks"):
+        layers = getattr(trunk, name, None)
+        if isinstance(layers, nn.ModuleList):
+            return layers
+    raise ValueError(f"cannot find the decoder layer list on {type(trunk).__name__}")
+
+
+def default_taps(n_layers: int, fractions=(0.5, 0.625, 0.75, 0.875)) -> list[int]:
+    """Four mid-to-upper layers plus the final output: decision-relevant features peak mid-depth in frozen LMs,
+    and the final layer is specialised for next-token prediction."""
+    return [*sorted({min(n_layers - 2, round(f * n_layers) - 1) for f in fractions}), -1]
+
+
+class TapReader:
+    """``reader.run(input_ids, mask, suffix_starts, lengths)`` -> per-row suffix features ``[S_b, T, d]``.
+
+    ``suffix_starts[b]`` is the first suffix position of row b *in this forward's input* (0 when the input is the
+    suffix itself, e.g. a continuation over a cached prefix); ``lengths[b]`` excludes right padding.
+    """
+
+    def __init__(self, trunk: nn.Module, taps: list[int]):
+        self.trunk, self.taps = trunk, list(taps)
+        self.layers = decoder_layers(trunk)
+        n = len(self.layers)
+        bad = [t for t in self.taps if t != -1 and not 0 <= t < n]
+        if bad:
+            raise ValueError(f"taps {bad} out of range for {n} layers")
+
+    def run(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        suffix_starts: list[int],
+        lengths: list[int],
+        detach: bool = True,
+        **trunk_kw,
+    ) -> list[torch.Tensor]:
+        lo = min(suffix_starts)  # keep only positions any row needs, inside the hook
+        grabbed: dict[int, torch.Tensor] = {}
+
+        def hook(t):
+            def f(_m, _args, out):
+                h = (out[0] if isinstance(out, tuple) else out)[:, lo:]
+                grabbed[t] = h.detach() if detach else h
+
+            return f
+
+        handles = [self.layers[t].register_forward_hook(hook(t)) for t in self.taps if t != -1]
+        try:
+            final = self.trunk(
+                input_ids=input_ids, attention_mask=attention_mask, **trunk_kw
+            ).last_hidden_state
+        finally:
+            for h in handles:
+                h.remove()
+        missing = [t for t in self.taps if t != -1 and t not in grabbed]
+        if missing:
+            raise RuntimeError(f"taps {missing} were not reached in the forward")
+        dev = final.device
+        stack = [
+            grabbed[t].to(dev) if t != -1 else final[:, lo:] for t in self.taps
+        ]  # T x [B, L - lo, d]
+        return [
+            torch.stack([x[b, suffix_starts[b] - lo : lengths[b] - lo] for x in stack], 1)
+            for b in range(input_ids.shape[0])
+        ]