open-interpretability 0.0.20__py3-none-any.whl

openinterp/harness/layer_cache.py

@@ -0,0 +1,301 @@
+"""On-disk cache of a scan's whole-layer activations.
+
+One forward pass over the corpus captures every neuron of a layer, so a scan stores
+all of them. Later audits of any neuron in that layer, and per-token analysis
+(``open-interpretability activations``), replay the stored tokens instead of re-running the model.
+
+An entry is keyed by everything that sets its numbers. That is the model, its revision,
+the weights' dtype and the device type (``coreml-ane`` for a scan on the ANE). It is also
+the torch and transformers versions and the layer. Last, it is the token budget per
+batch (which sets padding), and the exact pieces scanned (texts and windows). Replaying
+an entry gives the same scan result as running the model, bit for bit.
+
+Each entry is a directory ``<root>/<model>/L<layer>/<key>/``. It holds ``meta.json``,
+``pieces.json`` and ``acts.bin`` (``[rows, width]`` in the activation dtype). It also
+holds the per-row ``ids.npy``, ``positions.npy``, ``piece.npy`` and ``counted.npy``.
+The rows are every real (non-padding) token of every piece, in scan order. ``meta.json``
+lists how many pieces each forward pass took. Entries are written to a temporary
+directory and renamed into place, so concurrent scans never see a partial one. Once the
+cache outgrows its size cap, the least recently used entries go.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import secrets
+import shutil
+from collections.abc import Iterator, Sequence
+from contextlib import suppress
+from dataclasses import asdict, dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from openinterp.core.cache_dirs import cache_root
+from openinterp.harness.hooks import NeuronHookRunner
+from openinterp.harness.scan import Chunk, Piece
+from openinterp.harness.scan_backend import scan_device
+
+FORMAT = 2  # 2: batches by token budget, recorded in meta.json
+ENV = "OPENINTERP_SCAN_CACHE"  # cache directory, or "off"
+SIZE_ENV = "OPENINTERP_SCAN_CACHE_GB"
+DEFAULT_GB = 40.0
+
+# numpy has no bfloat16: store its bits as int16 and view them back.
+_BITS = {torch.bfloat16: torch.int16}
+
+
+def default_root() -> Path | None:
+    """``$OPENINTERP_SCAN_CACHE``, else ``$XDG_CACHE_HOME/openinterp/scans``; None if off."""
+    return cache_root(ENV, "scans")
+
+
+@dataclass(frozen=True)
+class Entry:
+    """One cached layer scan."""
+
+    path: Path
+    meta: dict
+
+    @property
+    def size(self) -> int:
+        """The entry's size on disk, in bytes."""
+        return sum(f.stat().st_size for f in self.path.iterdir())
+
+    @property
+    def last_used(self) -> float:
+        """When the entry was last read or written, as a timestamp."""
+        return (self.path / "meta.json").stat().st_mtime
+
+    def pieces(self) -> list[Piece]:
+        """The pieces that the entry scanned, in scan order."""
+        return [Piece(**p) for p in json.loads((self.path / "pieces.json").read_text())]
+
+    def rows(self, name: str) -> np.ndarray:
+        """A per-row array: ``ids``, ``positions``, ``piece`` or ``counted``."""
+        return np.load(self.path / f"{name}.npy", mmap_mode="r")
+
+    def acts(self) -> np.memmap:
+        """The whole layer's activations, ``[rows, width]``, memory-mapped."""
+        shape = (self.meta["rows"], self.meta["width"])
+        return np.memmap(self.path / "acts.bin", dtype=self.meta["storage"], mode="r", shape=shape)
+
+    def column(self, neuron_idx: int) -> torch.Tensor:
+        """One neuron's activation on every row, as float32 (as the scan reads it)."""
+        return self._tensor(np.ascontiguousarray(self.acts()[:, neuron_idx])).float()
+
+    def _tensor(self, array: np.ndarray) -> torch.Tensor:
+        tensor = torch.from_numpy(np.array(array))
+        dtype = getattr(torch, self.meta["dtype"])
+        return tensor.view(dtype) if dtype in _BITS else tensor
+
+    def chunks(self) -> Iterator[Chunk]:
+        """The scan's batches again, as the model produced them."""
+        pieces = self.pieces()
+        acts = self.acts()
+        ids, positions = self.rows("ids"), self.rows("positions")
+        counted, piece = self.rows("counted"), self.rows("piece")
+        lengths = np.bincount(piece, minlength=len(pieces)).tolist()
+        start = row = 0
+        for size in self.meta["batches"]:
+            batch_lengths = lengths[start : start + size]
+            end = row + sum(batch_lengths)
+            yield Chunk(
+                pieces=pieces[start : start + size],
+                lengths=batch_lengths,
+                ids=torch.from_numpy(np.array(ids[row:end])).long(),
+                positions=torch.from_numpy(np.array(positions[row:end])).long(),
+                counted=torch.from_numpy(np.array(counted[row:end])),
+                acts=self._tensor(acts[row:end]),
+            )
+            start, row = start + size, end
+
+
+class LayerCache:
+    """The cache of whole-layer scans under ``root``, capped at ``max_bytes``."""
+
+    def __init__(self, root: Path, max_bytes: int | None = None) -> None:
+        self.root = Path(root)
+        if max_bytes is None:
+            max_bytes = int(float(os.environ.get(SIZE_ENV, DEFAULT_GB)) * 1e9)
+        self.max_bytes = max_bytes
+        self.last: Path | None = None  # the entry the last scan replayed or stored
+        self.replayed = False
+
+    @classmethod
+    def default(cls) -> LayerCache | None:
+        """The cache at the default location, or ``None`` when caching is off."""
+        root = default_root()
+        return None if root is None else cls(root)
+
+    # -- keys ----------------------------------------------------------------
+
+    @staticmethod
+    def fields(
+        runner: NeuronHookRunner, layer: int, pieces: Sequence[Piece], batch_tokens: int
+    ) -> dict | None:
+        """What an entry is keyed by; None for a model with no name to key it by."""
+        config = getattr(runner.model, "config", None)
+        model = getattr(runner, "model_name", None) or getattr(config, "_name_or_path", None)
+        if not model:
+            return None
+        try:
+            import transformers
+
+            transformers_version = transformers.__version__
+        except ImportError:
+            transformers_version = None
+        digest = hashlib.sha256()
+        for piece in pieces:
+            digest.update(json.dumps(asdict(piece), ensure_ascii=False).encode() + b"\n")
+        return {
+            "format": FORMAT,
+            "model": model,
+            "revision": getattr(runner, "model_revision", None),
+            "dtype": str(next(runner.model.parameters()).dtype).removeprefix("torch."),
+            "device": scan_device(
+                runner, pieces, layer, batch_tokens
+            ),  # the ANE's activations differ from torch's
+            "torch": torch.__version__,
+            "transformers": transformers_version,
+            "layer": layer,
+            "batch_tokens": batch_tokens,
+            "pieces_sha256": digest.hexdigest(),
+        }
+
+    def path(self, fields: dict) -> Path:
+        """The entry directory for ``fields``."""
+        key = hashlib.sha256(json.dumps(fields, sort_keys=True).encode()).hexdigest()[:20]
+        return self.root / fields["model"].replace("/", "--") / f"L{fields['layer']}" / key
+
+    # -- reading -------------------------------------------------------------
+
+    def lookup(
+        self, runner: NeuronHookRunner, layer: int, pieces: Sequence[Piece], batch_tokens: int
+    ) -> Entry | None:
+        """The entry for this scan, or ``None`` when there is none."""
+        fields = self.fields(runner, layer, pieces, batch_tokens)
+        if fields is None:
+            return None
+        path = self.path(fields)
+        try:
+            meta = json.loads((path / "meta.json").read_text())
+        except (OSError, ValueError):
+            return None
+        os.utime(path / "meta.json")  # least recently used goes first
+        return Entry(path, meta)
+
+    def entries(self) -> list[Entry]:
+        """Every complete entry in the cache."""
+        found = []
+        for meta in self.root.glob("*/L*/*/meta.json"):
+            if meta.parent.name.startswith("."):
+                continue  # still being written
+            try:
+                found.append(Entry(meta.parent, json.loads(meta.read_text())))
+            except (OSError, ValueError):
+                continue
+        return sorted(found, key=lambda e: e.last_used)
+
+    def cached_layers(self, model: str) -> set[int]:
+        """The layers of ``model`` that the cache holds."""
+        entries = self.entries()
+        return {e.meta["fields"]["layer"] for e in entries if e.meta["fields"]["model"] == model}
+
+    # -- writing -------------------------------------------------------------
+
+    def through(
+        self,
+        runner: NeuronHookRunner,
+        layer: int,
+        pieces: Sequence[Piece],
+        batch_tokens: int,
+        chunks: Iterator[Chunk],
+    ) -> Iterator[Chunk]:
+        """The cached chunks if this scan is cached, else ``chunks``, stored as they pass."""
+        entry = self.lookup(runner, layer, pieces, batch_tokens)
+        self.replayed, self.last = entry is not None, entry and entry.path
+        if entry is not None:
+            return entry.chunks()
+        fields = self.fields(runner, layer, pieces, batch_tokens)
+        if fields is None:
+            return chunks
+        self.last = self.path(fields)
+        return self._record(fields, pieces, chunks)
+
+    def _record(self, fields: dict, pieces: Sequence[Piece], chunks: Iterator[Chunk]):
+        final = self.path(fields)
+        tmp = final.with_name(f".{final.name}.{os.getpid()}.{secrets.token_hex(4)}")
+        tmp.mkdir(parents=True)
+        rows: dict[str, list[np.ndarray]] = {"ids": [], "positions": [], "piece": [], "counted": []}
+        sizes: list[int] = []
+        index = 0
+        width = storage = None
+        try:
+            with open(tmp / "acts.bin", "wb") as fh:
+                for chunk in chunks:
+                    acts = chunk.acts.contiguous()
+                    bits = acts.view(_BITS[acts.dtype]) if acts.dtype in _BITS else acts
+                    fh.write(bits.numpy().tobytes())
+                    width, storage = acts.shape[1], str(bits.numpy().dtype)
+                    rows["ids"].append(chunk.ids.numpy().astype(np.int32))
+                    rows["positions"].append(chunk.positions.numpy().astype(np.int32))
+                    rows["counted"].append(chunk.counted.numpy())
+                    sizes.append(len(chunk.pieces))
+                    for n in chunk.lengths:
+                        rows["piece"].append(np.full(n, index, dtype=np.int32))
+                        index += 1
+                    yield chunk
+            _write_index(tmp, fields, pieces, rows, width, storage, sizes)
+            with suppress(OSError):  # another scan stored it first
+                tmp.rename(final)
+            self.evict(keep=final)
+        finally:
+            shutil.rmtree(tmp, ignore_errors=True)
+
+    def evict(self, keep: Path | None = None) -> list[Entry]:
+        """Delete least recently used entries until the cache fits its cap."""
+        entries = self.entries()
+        total = sum(e.size for e in entries)
+        removed = []
+        for entry in entries:
+            if total <= self.max_bytes:
+                break
+            if entry.path == keep:
+                continue
+            total -= entry.size
+            shutil.rmtree(entry.path, ignore_errors=True)
+            removed.append(entry)
+        return removed
+
+    def clear(self) -> int:
+        """Delete every entry; returns how many there were."""
+        entries = self.entries()
+        for entry in entries:
+            shutil.rmtree(entry.path, ignore_errors=True)
+        return len(entries)
+
+
+def _write_index(tmp: Path, fields: dict, pieces, rows: dict, width, storage, sizes) -> None:
+    """Write everything that an entry holds besides ``acts.bin``.
+
+    That is the per-row arrays, the pieces, and ``meta.json``.
+    """
+    for name, parts in rows.items():
+        np.save(tmp / f"{name}.npy", np.concatenate(parts) if parts else np.zeros(0))
+    (tmp / "pieces.json").write_text(json.dumps([asdict(p) for p in pieces]))
+    total = int(sum(len(p) for p in rows["ids"]))
+    meta = {
+        "fields": fields,
+        "rows": total,
+        "width": width or 0,
+        "dtype": fields["dtype"],
+        "storage": storage or "float32",
+        "batches": sizes,
+        "created": datetime.now(UTC).isoformat(timespec="seconds"),
+    }
+    (tmp / "meta.json").write_text(json.dumps(meta, indent=2))

openinterp/harness/layer_major.py

@@ -0,0 +1,274 @@
+"""Run independent forward passes layer by layer over a streamed model.
+
+A streamed model (``--memory stream``, :mod:`openinterp.harness.memory`) reads each
+decoder layer from disk on every pass. The review makes hundreds of passes, one after
+another, so it reads every layer hundreds of times. :func:`map_layer_major` runs the
+items of a loop at once, each in its own thread. It runs each layer once for all the
+passes that wait for it:
+
+1. Each thread runs the model's own forward, so each pass computes what it computes
+   alone. Only the order of the passes changes.
+2. Every thread waits at a layer or has finished. Then the scheduler loads the next
+   layer that a thread waits at, in the order of the layers, and lets those threads run
+   it. It then frees the layer and loads the next one. So the layers are read in sweeps
+   from the first to the last, and a sweep serves every waiting pass.
+3. While the threads run, a hook is applied only in the thread that registered it.
+   The harness's hooks stay where they are, on modules that every thread shares.
+
+A model that isn't streamed runs the loop as before, one item after another.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from typing import Any, TypeVar
+
+import torch
+from torch import nn
+
+from openinterp.harness.memory import is_streamed, shared_by, stream_rows
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+# The most passes that run at once. Each holds its own activations.
+MAX_WORKERS = 32
+
+# The prompt length that :func:`workers_for` plans each pass's memory with.
+PLANNED_TOKENS = 256
+
+_active = False  # a layer-major loop is running: a loop inside it runs in order
+
+
+class LayerScheduler:
+    """Lets threads through the streamed layers in sweeps, one loaded layer at a time."""
+
+    def __init__(self, layers: Sequence[nn.Module]) -> None:
+        self.layers = list(layers)
+        self.streamed = [
+            [m for m in layer.modules() if getattr(getattr(m, "_hf_hook", None), "offload", False)]
+            for layer in self.layers
+        ]
+        self.cond = threading.Condition()
+        self.running = 0  # threads outside the scheduler's waits
+        self.inside = 0  # threads inside the loaded layer
+        self.waiting: dict[int, int] = {}  # thread id -> the layer it waits at
+        self.loaded: int | None = None
+        self.loads = 0  # layers read from disk, for tests and reports
+
+    def start(self) -> None:
+        """A worker thread starts."""
+        with self.cond:
+            self.running += 1
+
+    def finish(self) -> None:
+        """A worker thread ends, normally or with an exception."""
+        with self.cond:
+            self.running -= 1
+            self._advance_if_idle()
+            self.cond.notify_all()
+
+    def enter(self, layer: int) -> None:
+        """Block until ``layer`` is loaded for this thread's pass."""
+        me = threading.get_ident()
+        with self.cond:
+            self.running -= 1
+            self.waiting[me] = layer
+            self._advance_if_idle()
+            self.cond.notify_all()
+            while self.loaded != layer or me not in self.waiting:
+                self.cond.wait()
+            del self.waiting[me]
+            self.running += 1
+            self.inside += 1
+
+    def leave(self, layer: int) -> None:
+        """This thread's pass has run ``layer``, or stopped in it."""
+        with self.cond:
+            self.inside -= 1
+
+    def _advance_if_idle(self) -> None:
+        # Called with the lock held. Only when no thread can still reach the loaded layer.
+        if self.running or self.inside or not self.waiting:
+            return
+        targets = set(self.waiting.values())
+        if self.loaded in targets:
+            return  # its waiters are being let through
+        after = [layer for layer in targets if self.loaded is None or layer > self.loaded]
+        self._load(min(after) if after else min(targets))
+
+    def _load(self, layer: int) -> None:
+        if self.loaded is not None:
+            self._release(self.loaded)
+        for module in self.streamed[layer]:
+            module._hf_hook.pre_forward(module)  # reads the weights, as each pass did
+            module._hf_hook.offload = False  # so a pass doesn't read or free them again
+        self.loaded = layer
+        self.loads += 1
+
+    def _release(self, layer: int) -> None:
+        for module in self.streamed[layer]:
+            module._hf_hook.offload = True
+            module._hf_hook.post_forward(module, None)  # frees the weights
+
+    def close(self) -> None:
+        """Free the layer that is still loaded."""
+        if self.loaded is not None:
+            self._release(self.loaded)
+            self.loaded = None
+
+    @contextmanager
+    def installed(self) -> Iterator[None]:
+        """Route every pass through :meth:`enter` and :meth:`leave` at each layer."""
+        originals = [layer.forward for layer in self.layers]
+
+        def through(index: int, forward: Callable) -> Callable:
+            def run(*args, **kwargs):
+                self.enter(index)
+                try:
+                    return forward(*args, **kwargs)
+                finally:
+                    self.leave(index)
+
+            return run
+
+        for index, (layer, forward) in enumerate(zip(self.layers, originals, strict=True)):
+            layer.forward = through(index, forward)
+        try:
+            with _hooks_per_thread():
+                yield
+        finally:
+            for layer, forward in zip(self.layers, originals, strict=True):
+                layer.forward = forward
+            self.close()
+
+
+@contextmanager
+def _hooks_per_thread() -> Iterator[None]:
+    """While it is open, a forward hook registered in a thread is applied only in it."""
+    register_hook, register_pre_hook = (
+        nn.Module.register_forward_hook,
+        nn.Module.register_forward_pre_hook,
+    )
+
+    def owned(fn: Callable) -> Callable:
+        owner = threading.get_ident()
+
+        def hook(*args, **kwargs):
+            return fn(*args, **kwargs) if threading.get_ident() == owner else None
+
+        return hook
+
+    def forward_hook(self, hook, *args, **kwargs):
+        return register_hook(self, owned(hook), *args, **kwargs)
+
+    def forward_pre_hook(self, hook, *args, **kwargs):
+        return register_pre_hook(self, owned(hook), *args, **kwargs)
+
+    nn.Module.register_forward_hook = forward_hook
+    nn.Module.register_forward_pre_hook = forward_pre_hook
+    try:
+        yield
+    finally:
+        nn.Module.register_forward_hook = register_hook
+        nn.Module.register_forward_pre_hook = register_pre_hook
+
+
+@contextmanager
+def _locked(obj: Any, name: str) -> Iterator[None]:
+    """Let one thread at a time call ``obj.name``: a fast tokenizer isn't thread-safe."""
+    own = vars(obj).get(name)  # an attribute set on the instance, to put back after
+    method, lock = getattr(obj, name), threading.Lock()
+
+    def call(*args, **kwargs):
+        with lock:
+            return method(*args, **kwargs)
+
+    setattr(obj, name, call)
+    try:
+        yield
+    finally:
+        if own is None:
+            delattr(obj, name)
+        else:
+            setattr(obj, name, own)
+
+
+def workers_for(runner, n_items: int, tokens: int, logits: bool) -> int:
+    """How many passes run at once: as many as a quarter of the planned memory holds.
+
+    Each pass is planned at ``tokens`` tokens, with its logits if ``logits``.
+    """
+    rows = stream_rows(runner, tokens, 1, logits=logits)
+    return max(1, min(n_items, MAX_WORKERS, rows))
+
+
+def prompt_tokens(runner, texts: Sequence[str]) -> int:
+    """The token count of the longest of ``texts``: what a pass over one of them holds."""
+    return max((runner.encode([t])["input_ids"].shape[1] for t in texts), default=1)
+
+
+def map_layer_major(
+    runner,
+    fn: Callable[[T], R],
+    items: Sequence[T],
+    *,
+    tokens: int = PLANNED_TOKENS,
+    logits: bool = True,
+) -> list[R]:
+    """``[fn(item) for item in items]``, run layer-major when ``runner``'s model streams.
+
+    ``tokens`` and ``logits`` describe the largest pass an item makes, so that the
+    passes that run at once fit in the planned memory (:func:`workers_for`). ``fn`` must
+    not change state that another item reads. The results are in the order of
+    ``items``. The first exception raised by an item is raised again.
+    """
+    global _active
+    items = list(items)
+    if len(items) < 2 or _active or not is_streamed(runner.model):
+        return [fn(item) for item in items]  # a loop inside a running one runs in order
+    n_workers = workers_for(runner, len(items), tokens, logits)
+    scheduler = LayerScheduler(runner.layers)
+    _active = True
+    try:
+        with scheduler.installed(), _locked(runner, "encode"), shared_by(n_workers):
+            results = _run_workers(scheduler, fn, items, n_workers)
+    finally:
+        _active = False
+    runner.layer_loads = getattr(runner, "layer_loads", 0) + scheduler.loads
+    return results
+
+
+def _run_workers(scheduler: LayerScheduler, fn: Callable, items: list, n_workers: int) -> list:
+    """Run ``fn`` over ``items`` in ``n_workers`` threads; raise the first exception again."""
+    results: list[Any] = [None] * len(items)
+    errors: list[BaseException] = []
+    queue = iter(range(len(items)))
+    lock = threading.Lock()
+    grad = torch.is_grad_enabled()  # grad mode is per thread: carry the caller's over
+
+    def work() -> None:
+        torch.set_grad_enabled(grad)
+        try:
+            while not errors:
+                with lock:
+                    index = next(queue, None)
+                if index is None:
+                    return
+                results[index] = fn(items[index])
+        except BaseException as exc:  # raised again in the calling thread
+            errors.append(exc)
+        finally:
+            scheduler.finish()
+
+    threads = [threading.Thread(target=work) for _ in range(n_workers)]
+    for thread in threads:
+        scheduler.start()
+        thread.start()
+    for thread in threads:
+        thread.join()
+    if errors:
+        raise errors[0]
+    return results

openinterp/harness/lookahead.py

@@ -0,0 +1,93 @@
+"""Distinguish retrospective neurons from lookahead (next-token preparing) neurons.
+
+Two complementary measurements on native weights:
+
+* :func:`output_token_projection` reads the neuron's output direction through the
+  unembedding, to list the tokens that the pole directly promotes. This is a logit-lens
+  approximation that skips the final norm.
+* :func:`measure_lookahead` zero-ablates the pole on real prompts and compares the
+  next-token distributions. A lookahead neuron's ablation moves probability mass
+  in a concentrated way onto/off a small set of upcoming tokens. :func:`ablation_profile`
+  does the same for any unit's ablation (an attention head is mean-ablated).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+import torch.nn.functional as F
+
+from openinterp.core.schemas import LookaheadProfile, Pole
+from openinterp.harness.hooks import Inputs, Intervention, NeuronHookRunner, ablate_pole
+
+
+def output_token_projection(
+    runner: NeuronHookRunner, layer: int, neuron_idx: int, pole: Pole, k: int = 10
+) -> list[tuple[int, float]]:
+    """Top-``k`` ``(token_id, score)`` pairs promoted by the pole through the direct path."""
+    unembed = runner.model.get_output_embeddings().weight.detach()  # [vocab, d_model]
+    direction = runner.neuron_output_direction(layer, neuron_idx).to(unembed.dtype)
+    scores = unembed @ direction
+    if pole == "negative":
+        scores = -scores
+    top = scores.topk(min(k, scores.numel()))
+    return list(zip(top.indices.tolist(), top.values.tolist(), strict=True))
+
+
+def measure_lookahead(
+    runner: NeuronHookRunner,
+    prompts: Inputs,
+    layer: int,
+    neuron_idx: int,
+    pole: Pole,
+    **kwargs,
+) -> LookaheadProfile:
+    """Zero-ablate the pole on ``prompts``; see :func:`ablation_profile`."""
+    return ablation_profile(
+        runner, prompts, {"interventions": {layer: ablate_pole(neuron_idx, pole)}}, **kwargs
+    )
+
+
+@torch.no_grad()
+def ablation_profile(
+    runner: NeuronHookRunner,
+    prompts: Inputs,
+    ablation: Mapping[str, Mapping[int, Intervention]],
+    *,
+    k: int = 10,
+    kl_threshold: float = 1e-3,
+    concentration_threshold: float = 0.5,
+) -> LookaheadProfile:
+    """Ablate a unit on ``prompts`` and profile the effect on next-token predictions.
+
+    ``ablation`` holds :meth:`NeuronHookRunner.run_with_interventions` keyword arguments
+    (a unit's ``ablation()``). ``kl_divergence`` is KL(clean || ablated) averaged over
+    non-padding positions. The unit is flagged as lookahead when two things are true.
+    That KL exceeds ``kl_threshold``. And the ``k`` tokens that lose the most probability
+    hold at least ``concentration_threshold`` of the total probability mass moved.
+    """
+    encoded = runner.encode(prompts)
+    mask = encoded.get("attention_mask", torch.ones_like(encoded["input_ids"])).bool()
+    clean = F.log_softmax(runner.run_with_interventions(encoded).float(), dim=-1)[mask]
+    ablated = F.log_softmax(runner.run_with_interventions(encoded, **ablation).float(), dim=-1)[
+        mask
+    ]
+
+    kl = F.kl_div(ablated, clean, log_target=True, reduction="batchmean").clamp(min=0.0)
+    prob_shift = (clean.exp() - ablated.exp()).mean(dim=0)  # positive = unit supports token
+    total_moved = prob_shift.abs().sum()
+    top = prob_shift.topk(min(k, prob_shift.numel()))
+    concentration = float(top.values.clamp(min=0).sum() / total_moved) if total_moved > 0 else 0.0
+
+    tokenizer = runner.tokenizer
+    tokens = [
+        tokenizer.decode([i]) if tokenizer is not None else str(i)
+        for i, v in zip(top.indices.tolist(), top.values.tolist(), strict=True)
+        if v > 0
+    ]
+    return LookaheadProfile(
+        is_lookahead=float(kl) >= kl_threshold and concentration >= concentration_threshold,
+        target_upcoming_tokens=tokens,
+        kl_divergence=float(kl),
+    )