Metadata-Version: 2.5
Name: interp-engine
Version: 1.1.0
Summary: Standalone raw-transformers interpretability core: eager PyTorch with its own forward-hook layer, plus a vLLM serving backend.
Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
Project-URL: Repository, https://github.com/decoderesearch/interp-engine
Project-URL: Issues, https://github.com/decoderesearch/interp-engine/issues
License-Expression: Apache-2.0
License-File: LICENSE
Requires-Python: <3.14,>=3.11
Requires-Dist: einops
Requires-Dist: numpy>=1.24
Requires-Dist: torch>=1.10
Requires-Dist: transformers>=4.57.1
Provides-Extra: awq
Requires-Dist: accelerate>=1.0; extra == 'awq'
Requires-Dist: gptqmodel>=5.0; extra == 'awq'
Provides-Extra: dev
Requires-Dist: pyright<1.2,>=1.1.411; extra == 'dev'
Requires-Dist: pytest<9,>=8.3.1; extra == 'dev'
Requires-Dist: pyyaml>=6; extra == 'dev'
Requires-Dist: ruff<0.17,>=0.16.2; extra == 'dev'
Provides-Extra: parity
Requires-Dist: transformer-lens>=3.0; extra == 'parity'
Provides-Extra: quant
Requires-Dist: accelerate>=1.0; extra == 'quant'
Requires-Dist: kernels<0.16.0,>=0.15.2; extra == 'quant'
Provides-Extra: vllm
Requires-Dist: vllm>=0.25.1; (sys_platform == 'linux') and extra == 'vllm'
Description-Content-Type: text/markdown

# interp-engine

An interpretability engine (alternative to TransformerLens/nnsight) that runs both the raw HuggingFace model
in standard eager PyTorch for maximum compatbility, and VLLM for faster inference. `interp-engine` runs all of Neuronpedia's inference work and is checked for accuracy against four other engines.

We built interp-engine so that we can move fast: both for development speed and in serving speed. By adding VLLM support, we can increase speed by orders of magnitude, serving many more researchers.

This repository also holds the two things built around the engine: [`compare-engines/`](compare-engines/), the harness that scores it against TransformerLens, nnsight/nnterp, vLLM and SGLang on real weights, and [`visualizer-web/`](visualizer-web/), a diagram of where each hook point sits in a forward pass and what other stacks call it. The harness runs against the engine in this repo rather than a release, so a change here is scored before it ships.

We also made interp-engine in order to start fresh and standardize.

You are free to use interp-engine for your own projects and contribute back to it. We will keep it maintained with the latest models and improvements. interp-engine is Apache 2 and is very lightweight in dependencies.

We don't reimplement models and we do not adopt a fused inference engine — we let
`transformers` run the forward pass (so every architecture gotcha is applied inside
`forward()`), and we copy only the small per-architecture _knowledge_ (a module-path mapping
derived by inspection + a short known-quirks table).

The canonical model identifier is the **raw HuggingFace repo id** (e.g. `openai-community/gpt2`, `google/gemma-2-2b`).

```bash
pip install interp-engine          # eager backend, runs on CPU/CUDA/MPS
pip install 'interp-engine[vllm]'  # + the vLLM backend (Linux/CUDA only)
```

[docs/USAGE.md](docs/USAGE.md) is the install-to-first-capture walkthrough.

## Contents

- [Documentation](#documentation) — which doc answers which question
- [Switching backend](#switching-backend) — the same code on eager and vLLM, and where it stops
- [Hook points](#hook-points) — every point, and whether each backend can serve it
- [Performance](#performance) — what the vLLM backend buys, and what capture costs it
- [Modules](#modules) — what each file in `interp_engine/` owns
- [Correctness](#correctness) — what the engine checks about itself, and how

## Documentation

| doc                                                          | when you need it                                                                               |
| ------------------------------------------------------------ | ---------------------------------------------------------------------------------------------- |
| [docs/USAGE.md](docs/USAGE.md)                               | install, load a model, capture, generate, steer, lens — start here                             |
| [docs/AGENT_INTEGRATION.md](docs/AGENT_INTEGRATION.md)       | porting code onto the engine: recipes, hard rules, error-to-fix (written for a coding agent)   |
| [docs/ARCHITECTURE_QUIRKS.md](docs/ARCHITECTURE_QUIRKS.md)   | every architecture quirk the engine knows about, and where a per-model fact is allowed to live |
| [docs/GRADIENTS.md](docs/GRADIENTS.md)                       | what is differentiable, on which backend, and what is silently not                             |
| [docs/ENGINE_HOOK_MAPPINGS.md](docs/ENGINE_HOOK_MAPPINGS.md) | every hook point mapped across interp-engine, TransformerLens and nnsight                      |
| [docs/PORTING.md](docs/PORTING.md)                           | translating code from TransformerLens, nnsight or nnterp                                       |
| [docs/PERFORMANCE.md](docs/PERFORMANCE.md)                   | vLLM speed/feature tradeoffs and quantization support                                          |
| [docs/COMPATIBILITY.md](docs/COMPATIBILITY.md)               | which transformers versions are tested, and the ones known to compute a model wrongly          |

## Switching backend

`backend=` is the only line that changes. The sync free functions dispatch on the model you hand
them, and the async methods are one protocol both backends implement, so the same script runs on the
eager PyTorch model and on a served vLLM engine:

```python
from interp_engine import capture_attention, capture_generation, generate_stream, load_model, run_with_cache, steer

model = load_model("google/gemma-2-2b-it", backend="eager")  # or backend="vllm"
tokens = model.to_tokens("The capital of France is")

cache = run_with_cache(model, tokens, ["resid_post.10"])
attn = capture_attention(model, tokens, [10])
with steer(model, spec):
    completion, acts = capture_generation(model, tokens, ["resid_post.10"], max_tokens=8)
    for step in generate_stream(model, tokens, max_tokens=8, n_logprobs=5):
        print(step.token_str, step.logprobs)
```

There is no `vllm=` flag on any function and no `**kwargs` forwarded to a backend, because both turn
"can this backend do it" into a silent behavior difference. Where only one backend can serve
something, the call raises `CapabilityUnsupported` naming the capability, why this backend cannot,
and what to call instead — `interp_engine.CAPABILITIES` is that table, and the [hook
points](#hook-points) below are the per-point half of it. Two asymmetries are worth knowing up front:
`GenStep.logits` is `None` on vLLM (ask for `n_logprobs` instead), and the free `decode_residuals`
returns raw logits and is eager-only, while `sync_model(model).decode_residuals` normalizes across
both. [docs/USAGE.md](docs/USAGE.md#without-an-event-loop) has both in full.

`sync_model(model)` is the same trick applied to the methods: it mirrors the whole protocol
synchronously on one background event loop, for a notebook that does not want `await`.

## Hook points

The eager backend can capture every point below; the vLLM backend serves most of them. Where it
does not, the reason is one of two kinds, and the difference is the difference between filing a bug
and switching backend:

- **unimplemented** — the module is right there on vLLM's tree and nobody has wired the point up.
- **unreachable** — a fused kernel ate the tensor, or no module boundary holds it. Wiring is not the
  problem; the quantity does not exist at a hookable point.

`interp_engine.points` is the table this is generated from, and
[docs/ENGINE_HOOK_MAPPINGS.md](docs/ENGINE_HOOK_MAPPINGS.md) carries each point's definition and its
TransformerLens/nnsight translation. Ask the code rather than this table if you are branching on it:
`points.vllm_hookable()` is the set, and `points.reason(name)` is the sentence explaining a refusal.

| point                      | width                   | eager | vLLM | vLLM notes                                                                                                                                                  |
| -------------------------- | ----------------------- | :---: | :--: | ----------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `embeddings`               | `d_model`               |  ✅   |  ✅  | trunk-level, so addressed with no layer index; distinct from `resid_pre` at layer 0 only where the trunk adds positional embeddings or scales the embedding |
| `resid_pre`                | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `attn_in`                  | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `q_norm_in` / `q_norm_out` | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                            |
| `k_norm_in` / `k_norm_out` | `n_kv_heads * head_dim` |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                            |
| `value`                    | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                            |
| `attn_scores`              | `n_heads * query * key` |  ✅   |  ♻️  | no module boundary holds the pre-softmax matrix on **either** backend; vLLM rebuilds it from captured post-RoPE q/k                                         |
| `attn_probs`               | `n_heads * query * key` |  ✅   |  ♻️  | fused paged attention never materializes the probabilities; same recompute                                                                                  |
| `z`                        | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded, so single-GPU only                                                                                                                            |
| `attn_gate`                | `n_heads * head_dim`    |  ✅   |  ❌  | unimplemented — a real module on both trees                                                                                                                 |
| `attn_out`                 | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `attn_out_post`            | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `resid_mid`                | `d_model`               |  ✅   |  ✅  | capture works everywhere; _steering_ it is refused on families where vLLM adds the residual before the norm                                                 |
| `mlp_in`                   | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `mlp_pre`                  | `d_mlp`                 |  ✅   |  ❌  | unreachable — vLLM fuses `gate_proj` and `up_proj` into one `gate_up_proj`, so neither branch is a module output                                            |
| `mlp_pre_linear`           | `d_mlp`                 |  ✅   |  ❌  | as `mlp_pre`; gated MLPs only                                                                                                                               |
| `mlp_act`                  | `d_mlp`                 |  ✅   |  ✅  | neuron-sharded, so single-GPU only                                                                                                                          |
| `router_logits`            | `n_experts`             |  ✅   |  ✅  | replicated gate, so it survives tensor parallelism                                                                                                          |
| `expert_weights`           | `n_experts`             |  ✅   |  ❌  | unreachable — the top-k happens inside the FusedMoE kernel, which returns the combined output with the selection never materialized                         |
| `expert_indices`           | `n_experts`             |  ✅   |  ❌  | as `expert_weights`                                                                                                                                         |
| `mlp_out`                  | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `mlp_out_post`             | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `resid_post`               | `d_model`               |  ✅   |  ✅  |                                                                                                                                                             |
| `final_norm`               | `d_model`               |  ✅   |  ✅  | trunk-level, so addressed with no layer index; runs over every position, not just the ones being decoded                                                    |
| `lm_head`                  | `vocab_size`            |  ✅   |  ❌  | unreachable as a bare unembed — vLLM's `compute_logits` folds scaling and softcapping in, so hooking it returns something other than `W_U @ x`              |

✅ served · ♻️ served by recompute rather than a hook · ❌ not served

**DeepSeek-V4 only** — the multi-head-channel (mHC) points, which exist on a hyper-connection trunk
and nowhere else. `interp_engine.points.points_for(architecture)` adds them for that family, and an
`Address` carries the `stream` coordinate that selects one trunk out of the stack
(`resid_post.5.stream-2`). All seven are eager-only for now, and **unimplemented rather than
unreachable**: most are one `output:N` hook away on vLLM, but the smallest V4 is 284B, so nobody has
been able to check the captures against eager. [`docs/ENGINE_HOOK_MAPPINGS.md`](docs/ENGINE_HOOK_MAPPINGS.md)
has the per-point detail.

| point                                          | width                | eager | vLLM |
| ---------------------------------------------- | -------------------- | :---: | :--: |
| `resid_streams`                                | `n_residual_streams` |  ✅   |  ❌  |
| `attn_stream_collapse` / `mlp_stream_collapse` | `d_model`            |  ✅   |  ❌  |
| `attn_stream_write` / `mlp_stream_write`       | `n_residual_streams` |  ✅   |  ❌  |
| `attn_stream_mix` / `mlp_stream_mix`           | `n_residual_streams` |  ✅   |  ❌  |

**Tensor parallelism narrows the vLLM column further.** The capture path reads rank 0's payload
alone, so a point whose last axis vLLM shards comes back as a slice: `z`, `value`, `mlp_act` and the
four QK-norm points are refused on a multi-GPU pod rather than returned short, and so is the
attention recompute (q/k/v are head-sharded). Everything `d_model` wide is all-reduced before the
hook sees it, and `router_logits` comes off a replicated gate, so those are unaffected.

## Performance

The two backends do **not** capture quite the same set of points — see the table above — so between
two that both serve what you need, the choice is a speed choice.
**The reason to pick vLLM is concurrency.** On a single stream it is modestly faster than raw
eager PyTorch; served several requests at once it is roughly an order of magnitude faster, because
it batches them into shared forwards while the eager backend's generation loop is synchronous
underneath and serializes them.

One RTX 5090, bf16, 512-token prompt, 128 new tokens, greedy — decode throughput, vLLM against the
eager backend:

| model          | eager     | vLLM, one stream | vLLM, 8 concurrent |
| -------------- | --------- | ---------------- | ------------------ |
| `gemma-3-1b`   | 101 tok/s | 112 tok/s (+11%) | 845 tok/s (8.4x)   |
| `gemma-2-2b`   | 110 tok/s | 140 tok/s (+27%) | 999 tok/s (9.5x)   |
| `qwen3-4b`     | 96 tok/s  | 165 tok/s (+71%) | 879 tok/s (9.3x)   |
| `llama-3.1-8b` | 82 tok/s  | 101 tok/s (+22%) | 585 tok/s (7.3x)   |

### Our VLLM against stock vLLM

The engine runs vLLM with `enforce_eager=True`, because CUDA-graph replay does not re-execute the
Python forward and so a `register_forward_hook` never fires — with graphs on, a capture returns
nothing. That is the one place the interp machinery is slower than stock vLLM on generation, and it
is a **small-model** tax: graph replay removes per-kernel launch overhead, which is most of a 1B
model's decode step and noise for an 8B one.

The comparison below is the same engine with `enforce_eager=False`, which is vLLM's own default
configuration, so it is the tax and nothing else:

| model          | single-stream decode, stock vLLM | capture-capable (our default) | throughput lost |
| -------------- | -------------------------------- | ----------------------------- | --------------- |
| `gemma-3-1b`   | 391 tok/s                        | 112 tok/s                     | -71%            |
| `gemma-2-2b`   | 228 tok/s                        | 140 tok/s                     | -39%            |
| `qwen3-4b`     | 170 tok/s                        | 165 tok/s                     | -3%             |
| `llama-3.1-8b` | 102 tok/s                        | 101 tok/s                     | -1%             |

So at 4B and up the engine is stock vLLM's speed; below ~2B a generation-only pod is worth running
with `enforce_eager=False`, which is supported and already done in-tree. Everything else — the
attention recompute, the capture hooks, native extraction — is off unless requested and costs a
generation request nothing. [docs/PERFORMANCE.md](docs/PERFORMANCE.md) has the reasoning and the
graph-mode measurements that rule out a middle ground; the full report, including capture, lens and
steering latencies and peak VRAM, is at
[benchmarks/results-latest.md](benchmarks/results-latest.md) with the suite in
[benchmarks/](benchmarks/README.md).

## Modules

- `model.py` — `EagerModel`: wraps `AutoModelForCausalLM` (eager, `no_processing` semantics),
  holds the tokenizer + config-derived dims, canonical hook-point resolution, and an optional
  `quantization_config` passthrough to `from_pretrained`.
- `facts.py` — the single source of truth for model facts, shared by both backends: structural
  attribute-name vocabularies, config-derived dims, per-layer window/linear-attention predicates,
  and the per-backend tables (fused-QKV layout, parallel-block architectures). Config arithmetic
  and string tables only — no torch, no live model — so the vLLM client can answer dims for a model
  it never builds.
- `arch.py` — the **eager adapter**: binds a live HF module tree to the structural roles in
  `facts.py`, plus the machine-readable known-quirks table (attention sinks, softcapping, hybrid
  attention, ...).
- `protocol.py` — `InterpModel`: the surface both backends implement, and the contract a sync free
  function dispatches against. Adding a method here means adding it to both backends and to
  `sync.py`; `tests/test_sync_parity.py` fails on a missing twin.
- `sync.py` / `_loop.py` — `sync_model(model)`: the protocol without an event loop, one explicit
  wrapper per method over a background loop thread that is created lazily, reused per model, and
  refuses rather than deadlocks when called from inside a running loop.
- `dispatch.py` — the shared plumbing every free function's two arms sit on: token coercion
  (`TokensLike`, batch refusals) and `CAPABILITIES`, the table each `CapabilityUnsupported` message
  is built from.
- `hooks.py` — the low-level read/write forward-hook substrate.
- `capture.py` — capture context manager returning a cache keyed by canonical names
  (`resid_post`, `resid_mid`, `mlp_in`, `mlp_act`, `attn_probs`, `value`, `router_logits`,
  `embeddings`, ...), plus the post-processing a captured tensor needs to be usable (fused-QKV
  splits, the attention gate, a norm's scale and gain, per-head residual contributions, dense expert
  assignments).
- `attn_scores.py` — the pre-softmax attention scores, which no module boundary carries: it registers
  a wrapping attention implementation for the duration of a capture and delegates to the
  checkpoint's own eager function, so the forward is unchanged.
- `tokenize.py` — `to_tokens`/`to_str_tokens`/`to_string` (TransformerLens-parity), chat
  templating, and per-token span metadata (the single source of truth for message boundaries).
- `chat_conventions.py` — the only per-model chat table: harmony markers, reasoning delimiters,
  turn-end tokens. Selected by tokenizer capability, never by model name (see
  [Where model-specific config lives](docs/ARCHITECTURE_QUIRKS.md#where-model-specific-config-lives)).
- `chat_compose.py` — rebuilds assistant messages from a generation (`compose_assistant_turns`),
  reading the generated text only; callers pair it with the prompt messages they already have.
- `lens.py` — logit + Jacobian lens by calling the real `final_norm` + `lm_head`. Returns **raw**
  logits unless the caller passes a `softcap`. The vLLM path never returns raw logits (see [vLLM
  `compute_logits` is not a bare unembed](docs/ARCHITECTURE_QUIRKS.md#vllm-compute_logits-is-not-a-bare-unembed)).
- `steer.py` — steering, and the per-token generation stream. Each method's arithmetic is one
  `steer_delta` branch, which is what the vLLM worker's modifier computes too and what
  `tests/test_steer_math_parity.py` runs against it on CPU; `steer()` is the context both backends
  take, registering per-request on vLLM rather than installing a global hook.
- `mappers.py` — translation between canonical points and other frameworks' names:
  TransformerLens hook strings and nnsight/nnterp accessors, both directions. See [Porting from
  TransformerLens, nnsight or nnterp](docs/PORTING.md#porting-from-transformerlens-nnsight-or-nnterp).
- `autograd_support.py` — the `GradSupport` verdict: whether a model can give you gradients, and
  which specific thing is blocking it. Pure config arithmetic, so it is safe to call before
  `warmup()`. See [Gradients](docs/GRADIENTS.md#gradients).
- `cuda_preflight.py` — `check_cuda_driver`: compares the host CUDA driver against the CUDA
  version torch was built for and raises with the forward-compat fix (`cuda-compat-<major>-<minor>`
  - `LD_LIBRARY_PATH`) before the first CUDA call, instead of failing ten frames deep in
    `torch.cuda._lazy_init`. Lives here because every app on the engine inherits the same CUDA
    floor — the `[vllm]` wheels link `libcudart.so.13` directly.

## Correctness

The engine's job is to hand back the tensor a module actually produced, so most of what can go wrong
is quiet: a point resolves to a plausible neighbour, the shapes agree, and the numbers are wrong.
The test suite is built around checks that a shape-correct guess cannot pass.

**Golden parity.** `tests/test_parity_gpt2.py` pins every capture point on gpt2 against
TransformerLens, from a committed golden file. It is the one place another framework is loaded, and
CI treats a skip as a failure (`IE_REQUIRE_PARITY=1`) so a missing dependency or a cold cache cannot
quietly retire the gate.

**Invariants over attribute names.** Three identities hold on any model of a family, so they catch a
misresolved point without needing a reference implementation: `probs @ value == z` for the per-head
value and DFA (`tests/test_qkv_layout.py`, which also asserts the _wrong_ layout fails — otherwise
the test would pass on a single-head model), `resid_pre + attn_out_post + mlp_out_post == resid_post`
for sandwich norms and residual multipliers, and `down_proj(mlp_act) == mlp_out` for the neuron
basis. Where a point genuinely does not exist — a Mamba block's attention, a latent-attention model's
`value`, the residual between the sublayers of a parallel block — it is refused with an explanation
rather than returned as a plausible tensor.

**Self-consistency on real weights.** `tests/test_new_models_gpu.py` decodes the last layer's
residual through the real `final_norm` + `lm_head` and requires the model's true next-token argmax
back, which validates the whole arch map end-to-end without a second framework.
`tests/test_sliding_window_attn.py` pins the vLLM attention recompute's band and sink terms, and
`tests/test_vllm_only_families.py` checks the spellings for families `transformers` has no class for
against a synthetic tree in the shape their own modeling file describes.

**The two backends against each other.** `tests/test_vllm_capture_gpu.py` runs a real vLLM engine and
requires each captured point to match the eager backend's — including that vLLM's positional layer
index names the layer HF's does, which nothing checked before and which would fail silently rather
than raise. It needs `interp-engine[vllm]`, so it self-skips elsewhere; note that running it via
`.venv-vllm/bin/python` needs that directory on `PATH` too, because vLLM shells out to `ninja` to
build a sampler kernel at startup. `tests/test_vllm_wire_grammar.py` covers the same process
boundary on CPU, over a synthetic demux.

**How CI is split.** Two jobs, both on every non-Markdown change: a CPU job
(`-m "not gpu and not xl"`) that owns the golden gate, the lint/format/type gates and the small
models eagerly, and a managed-L4 GPU job (`-m "gpu and not xl"`) running the same models on
CUDA/bf16. The `xl` models are tens of GB and run nowhere automatically — `pytest -m xl` on a big
box. Locally the GPU tests self-skip without CUDA and model loads skip when weights aren't cached, so
a plain `pytest tests` on a laptop runs the fast suite.

### A hook point's name is not its definition

The engine has its own canonical point names, and translating them is a real hazard rather than a
formality: `blocks.5.hook_mlp_out` (TransformerLens) and `mlps_output[5]` (nnsight) are the same
tensor on Llama and _different_ tensors on Gemma, because TransformerLens' block-level hook fires
after the post-sublayer norm. Here that distinction is two separate points — `mlp_out` is the raw
module output and `mlp_out_post` is the residual contribution.

`interp_engine.mappers` translates names in both directions.
[docs/ENGINE_HOOK_MAPPINGS.md](docs/ENGINE_HOOK_MAPPINGS.md) maps every point across the three
hookable stacks, including the ones TransformerLens has and we do not, and
[docs/PORTING.md](docs/PORTING.md) is the migration guide.
