Metadata-Version: 2.4
Name: interp-engine
Version: 1.0.1
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/hijohnnylin/neuronpedia/tree/main/interp-engine
Project-URL: Repository, https://github.com/hijohnnylin/neuronpedia
Project-URL: Issues, https://github.com/hijohnnylin/neuronpedia/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 having the engine in our monorepo, we can make changes immediately as needed. By adding VLLM support, we can increase speed by orders of magnitude, serving many more researchers.

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`).

## Contents

- [Documentation](#documentation) — which doc answers which question
- [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/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                                          |

## Performance

Both backends capture the same points, so the choice between them 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, ...).
- `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` — additive/orthogonal steering hooks + streaming generation with logprobs.
- `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.
