interp-engine 1.2.7__tar.gz → 1.2.8__tar.gz

{interp_engine-1.2.7 → interp_engine-1.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.2.7
+Version: 1.2.8
 Summary: A fast, standardized interpretability engine that supports most modern models and architectures. Powers Neuronpedia.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine
@@ -31,33 +31,22 @@ Description-Content-Type: text/markdown
 
 # interp-engine
 
-<p align="center">
-  <img src="https://neuronpedia.s3.amazonaws.com/site-assets/ielogo.png" alt="interp-engine logo, a magnifying glass where the handle is a rocket" width="160">
-</p>
-<p align="center">
-  🔗 <a href="https://interp-engine.org"><strong>interp-engine.org</strong></a>
-</p>
-<p align="center">
-  <a href="https://github.com/decoderesearch/interp-engine/actions/workflows/engine-tests.yml"><img src="https://github.com/decoderesearch/interp-engine/actions/workflows/engine-tests.yml/badge.svg?branch=main" alt="CI status"></a>
-  <a href="https://pypi.org/project/interp-engine/"><img src="https://img.shields.io/pypi/v/interp-engine.svg" alt="PyPI version"></a>
-  <a href="LICENSE"><img src="https://img.shields.io/pypi/l/interp-engine.svg" alt="Apache-2.0 license"></a>
-  <a href="https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-3z9o0hxjl-MDX9pbATO2qESOazNDLpdQ"><img src="https://img.shields.io/badge/Slack-Open%20Source%20Mechanistic%20Interpretability-4A154B?logo=slack&logoColor=white" alt="Join the Slack"></a>
-</p>
+
+
+🔗 **[interp-engine.org](https://interp-engine.org)**
+
 
 
 `interp-engine` is an interpretability engine that is fast (>40x tok/s vs HF eager), standardized (34 'points'/addresses across architectures), and easy to use and debug. It powers all of [Neuronpedia](https://neuronpedia.org)'s inference and is checked for accuracy against HF Transformers and other engines.
 
-<p align="center">
-  <img src="https://neuronpedia.s3.amazonaws.com/site-assets/ie-perf.png" alt="Tokens per second while capturing and generating: eager against IE-vLLM and IE-vLLM-freeze on qwen3.8-27b and deepseek-v4-flash-0731, 8 requests in flight" width="100%">
-</p>
-<p align="center">
-  <img src="https://neuronpedia.s3.amazonaws.com/site-assets/interp-engine-demo.gif" alt="interp-engine demo gif" width="100%">
-</p>
+
+
+
 
 This repo contains:
 
-1. [`validator/`](validator/), which compares/validates it against TransformerLens, and nnsight/nnterp on real architectures.
-2. [`visualizer-web/`](visualizer-web/), a "cheat sheet" hosted at [interp-engine.org](https://interp-engine.org) of each 'point' (eg `resid_post.16`), standardized across architectures.
+1. `[validator/](validator/)`, which compares/validates it against TransformerLens, and nnsight/nnterp on real architectures.
+2. `[visualizer-web/](visualizer-web/)`, a "cheat sheet" hosted at [interp-engine.org](https://interp-engine.org) of each 'point' (eg `resid_post.16`), standardized across architectures.
 
 ## Installation
 
@@ -71,7 +60,15 @@ pip install interp-engine          # eager backend only
 ```python
 from interp_engine import Address, load_model, run_with_cache
 
-model = load_model("Qwen/Qwen3-8B")    # vLLM default, use backend='eager' to override
+# VLLM MODE (default): low VRAM, medium speed
+model = load_model("Qwen/Qwen3-8B") 
+
+# VLLM-FREEZE MODE: high VRAM, high speed, only frozen points (default resid_post)
+# model = load_model("Qwen/Qwen3-8B", freeze_points="auto")
+
+# EAGER MODE: low VRAM, low speed
+# model = load_model("Qwen/Qwen3-8B", backend="eager")
+
 point = Address("resid_post", 10)      # or string: "resid_post.10"
 cache = run_with_cache(model, model.to_tokens("Hello, world"), [point])
 cache[point]  # [batch, pos, ...]
@@ -87,7 +84,7 @@ Add "use interp-engine" to your prompt and let your agent figure it out - everyt
 
 ## Performance / Speed
 
-vLLM gives `interp-engine` high throughput via concurrency, and **graph freeze** adds CUDA-graph replay on top of that _without giving up capture or steering_. Every column below is capture-capable.
+vLLM gives `interp-engine` high throughput via concurrency, and **graph freeze** adds CUDA-graph replay on top of that *without giving up capture or steering*. Every column below is capture-capable.
 
 <!-- THROUGHPUT:START -->
 
@@ -117,22 +114,14 @@ One stream (tok/s):
 
 <!-- THROUGHPUT:END -->
 
-Every multiplier is against eager on the same workload, and is the ratio of the _unrounded_ figures — tok/s is printed whole at 10 and above, so dividing two cells by hand can differ in the last place. Eager's generation loop is synchronous underneath, so eight requests serialize rather than batch — which is why it earns almost nothing from the second table and why the multipliers there are so much larger. `deepseek-v4-flash-0731` is block-quantized FP8 on both backends rather than bf16, so read its row as two backends serving the same quantized weights; it also keeps a tenth, because at 3 tok/s that digit is worth several percent. [`benchmarks/results-latest.md`](benchmarks/results-latest.md) has every figure at full precision.
-
-**Graph freeze** is opt-in, via `freeze_points`. Ordinary hooks cannot survive CUDA graphs — replay never calls the Python forward a hook is attached to — so the tap becomes a preallocated buffer plus a `copy_` that graph capture records and replay re-executes. Steering rides the same wrap, so additive, orthogonal, projection-cap and j-lens steer/ablate/swap all keep working under graphs, position masks included.
-
-```python
-model = load_model("Qwen/Qwen3-8B", freeze_points="auto")  # resid_post at every layer, under graphs
-```
-
-The trade is that a frozen engine serves _the set it froze_ rather than any point on request: `"auto"` covers `resid_post` at every layer on a conventional trunk, an explicit list covers whatever you name, and anything outside the set is refused rather than quietly returned empty. Omit `freeze_points` for today's hooked vLLM, which still serves every point. `qwen3.8-27b` is a hybrid trunk, so its freeze row runs prefill eagerly and keeps the decode graphs — freeze pins that mode, because breakable graphs turn `torch.compile` off and vLLM's mixed prefill-decode capture then miscomputes prefill on a gated-delta trunk. On a hyper-connection trunk such as `deepseek-v4-flash-0731`, whose block carries four parallel residual streams, `"auto"` freezes `resid_streams` instead — the whole stack per layer, which is four times the width and so four times the buffer, and is what that row's freeze cell prices. See [PERFORMANCE.md](docs/PERFORMANCE.md) for the full trade-off, and [`benchmarks/results-latest.md`](benchmarks/results-latest.md) for the sweep, including capture, steering and lens latencies.
+**Graph freeze** is opt-in via `freeze_points`, and a frozen engine serves only the set it froze. [PERFORMANCE.md](docs/PERFORMANCE.md) has how it works and what it trades; `[benchmarks/results-latest.md](benchmarks/results-latest.md)` has every figure at full precision, including capture, steering and lens latencies; `[benchmarks/README.md](benchmarks/README.md)` has how the tables above are rounded.
 
 ## Correctness
 
 We verify correctness in two main ways:
 
 1. A test suite that checks results across several models - what each check is designed to catch is in [INTERNALS.md](docs/INTERNALS.md#correctness).
-2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at [`validator/`](validator/).
+2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at `[validator/](validator/)`.
 
 ## Why use an Interpretability Engine, instead of just having my AI code whatever it needs on the fly?
 
@@ -156,4 +145,4 @@ Bugs and feature requests belong in [issues](https://github.com/decoderesearch/i
 
 ## License
 
-Apache 2.0
+Apache 2.0

interp_engine-1.2.8/README.md

@@ -0,0 +1,117 @@
+# interp-engine
+
+
+
+🔗 **[interp-engine.org](https://interp-engine.org)**
+
+
+
+`interp-engine` is an interpretability engine that is fast (>40x tok/s vs HF eager), standardized (34 'points'/addresses across architectures), and easy to use and debug. It powers all of [Neuronpedia](https://neuronpedia.org)'s inference and is checked for accuracy against HF Transformers and other engines.
+
+
+
+
+
+This repo contains:
+
+1. `[validator/](validator/)`, which compares/validates it against TransformerLens, and nnsight/nnterp on real architectures.
+2. `[visualizer-web/](visualizer-web/)`, a "cheat sheet" hosted at [interp-engine.org](https://interp-engine.org) of each 'point' (eg `resid_post.16`), standardized across architectures.
+
+## Installation
+
+```bash
+pip install 'interp-engine[vllm]'  # preferred install: includes vLLM support (CUDA required)
+pip install interp-engine          # eager backend only
+```
+
+## Simple Usage
+
+```python
+from interp_engine import Address, load_model, run_with_cache
+
+# VLLM MODE (default): low VRAM, medium speed
+model = load_model("Qwen/Qwen3-8B") 
+
+# VLLM-FREEZE MODE: high VRAM, high speed, only frozen points (default resid_post)
+# model = load_model("Qwen/Qwen3-8B", freeze_points="auto")
+
+# EAGER MODE: low VRAM, low speed
+# model = load_model("Qwen/Qwen3-8B", backend="eager")
+
+point = Address("resid_post", 10)      # or string: "resid_post.10"
+cache = run_with_cache(model, model.to_tokens("Hello, world"), [point])
+cache[point]  # [batch, pos, ...]
+```
+
+### AI Agents
+
+Add "use interp-engine" to your prompt and let your agent figure it out - everything is fully documented in this repo and open source.
+
+## Supported Points ("Addresses")
+
+`interp-engine` supports 34 standardized points ("Addresses") across architectures: every one of them on the eager backend, 28 of them on vLLM. Check [interp-engine.org](https://interp-engine.org) for the "cheat sheet", or [SUPPORTED_POINTS.md](docs/SUPPORTED_POINTS.md) for a markdown version with the per-backend detail.
+
+## Performance / Speed
+
+vLLM gives `interp-engine` high throughput via concurrency, and **graph freeze** adds CUDA-graph replay on top of that *without giving up capture or steering*. Every column below is capture-capable.
+
+<!-- THROUGHPUT:START -->
+
+<!-- Generated by `python -m benchmarks.report_bench`. Do not edit: rerun the sweep. -->
+
+Measured on NVIDIA B200, bf16, 512-token prompt, 128 new tokens.
+
+One stream (tok/s):
+
+| model                    | eager | vLLM       | vLLM + graph freeze |
+| ------------------------ | ----- | ---------- | ------------------- |
+| `gemma-2-2b`             | 31    | 31 (1.0x)  | **214 (6.9x)**      |
+| `qwen3-4b`               | 24    | 47 (2.0x)  | **296 (12.3x)**     |
+| `llama-3.1-8b`           | 33    | 57 (1.7x)  | **256 (7.9x)**      |
+| `qwen3.8-27b`            | 9.9   | 12 (1.2x)  | **63 (6.4x)**       |
+| `deepseek-v4-flash-0731` | 3.3   | 2.9 (0.9x) | **119 (36x)**       |
+
+8 concurrent requests (aggregate tok/s):
+
+| model                    | eager | vLLM        | vLLM + graph freeze |
+| ------------------------ | ----- | ----------- | ------------------- |
+| `gemma-2-2b`             | 30    | 226 (7.5x)  | **1,238 (41x)**     |
+| `qwen3-4b`               | 24    | 333 (14.0x) | **1,018 (43x)**     |
+| `llama-3.1-8b`           | 32    | 419 (13.0x) | **1,536 (48x)**     |
+| `qwen3.8-27b`            | 9.5   | 87 (9.2x)   | **386 (41x)**       |
+| `deepseek-v4-flash-0731` | 3.2   | 23 (7.2x)   | **402 (127x)**      |
+
+<!-- THROUGHPUT:END -->
+
+**Graph freeze** is opt-in via `freeze_points`, and a frozen engine serves only the set it froze. [PERFORMANCE.md](docs/PERFORMANCE.md) has how it works and what it trades; `[benchmarks/results-latest.md](benchmarks/results-latest.md)` has every figure at full precision, including capture, steering and lens latencies; `[benchmarks/README.md](benchmarks/README.md)` has how the tables above are rounded.
+
+## Correctness
+
+We verify correctness in two main ways:
+
+1. A test suite that checks results across several models - what each check is designed to catch is in [INTERNALS.md](docs/INTERNALS.md#correctness).
+2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at `[validator/](validator/)`.
+
+## Why use an Interpretability Engine, instead of just having my AI code whatever it needs on the fly?
+
+Software is easy to make in 2026, here's why an engine makes sense:
+
+1. **Speed**: Get performance without sacrificing correctness.
+2. **Standardization + Verification**: Eliminate ambiguity when referring to points, plus a full test suite included.
+3. **Faster Dev / Fewer Tokens Used**: You could spend ten million tokens and have your AI write, test, and make production-ready an interpretability engine. Or you could just `pip install interp-engine[vllm]`.
+
+## Development / Contributing
+
+Activate the shared git hooks once per clone — they format staged Python, rebuild the generated files, and run CI's static checks before a push. Details in [CONTRIBUTING.md](CONTRIBUTING.md).
+
+```bash
+make hooks    # or: git config core.hooksPath .githooks
+```
+
+## Contact
+
+Bugs and feature requests belong in [issues](https://github.com/decoderesearch/interp-engine/issues). For anything else: [johnny@neuronpedia.org](mailto:johnny@neuronpedia.org).
+
+## License
+
+Apache 2.0

{interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/README.md

@@ -104,6 +104,21 @@ python -m benchmarks.publish          # rewrite both
 python -m benchmarks.publish --check  # exit 1 if either has drifted, and name it
 ```
 
+Both print the same cells in a display form of their own, applied by `publish.py` and pinned by
+`tests/test_published_benchmarks.py`:
+
+- **tok/s is whole at 10 and above, one decimal below.** A tenth beside a four-digit figure in the
+  next column claims a resolution the reader cannot use. At 3 tok/s that same tenth is worth several
+  percent, so the small rows keep it.
+- **every comparison is a multiplier, never a percent**, with one decimal below 20x, where it is
+  still checkable against the two printed figures. `+20%` beside `27x` makes the reader convert one
+  of them.
+- **multipliers are ratios of the unrounded metrics**, each against eager on the same workload, so
+  dividing two printed figures by hand can differ in the last place. Ratios of the rounded figures
+  would make the published win depend on the rounding, which is worse.
+
+`results-latest.md` keeps every figure at full precision, so nothing is lost to those rules.
+
 `tests/test_published_benchmarks.py` runs that check over the committed cells, so a stale copy is a
 red suite rather than a claim nobody re-read. The visualizer's chatbot answers out of a bundle holding
 the README verbatim, so a publish that changed the README also wants `make viz-knowledge` -- the
@@ -225,12 +240,15 @@ already refers to, so a replay serves capture and steering without a Python forw
 `freeze_points="auto"` takes in production, and the row exists to price it against the
 `enforce_eager=True` column capture would otherwise have to use.
 
-`"auto"` installs **reads** only, and a steering op needs a write tap to land in, so this row priced
-half the feature and reported the other half as `n/a` -- with a message that blamed graph replay for
-it, which is the thing freeze exists to work around. It now also passes `freeze_writes`, whose value
-is the sentinel `run_bench.STEER_WRITES` rather than a site: the `steer` workload writes mid-stack,
-that layer differs per model, and a freeze write is a `load_model` argument, so it has to be resolved
-from the config before a model exists to ask.
+`"auto"` once installed **reads** only, and a steering op needs a write tap to land in, so this row
+priced half the feature and reported the other half as `n/a` -- with a message that blamed graph
+replay for it, which is the thing freeze exists to work around. Auto now covers both halves, so the
+cell would be a number either way, and this row still passes `freeze_writes` on purpose: an explicit
+list *narrows* what auto would install, to the one mid-stack site the `steer` workload actually
+writes. A row that priced a write buffer at every layer would not be comparable with the ones beside
+it, which is the whole job of the column. Its value is the sentinel `run_bench.STEER_WRITES` rather
+than a site, because that layer differs per model and a freeze write is a `load_model` argument, so
+it has to be resolved from the config before a model exists to ask.
 
 `VariantSpec.models` restricts the row to the checkpoints freeze has been shown correct on, so a model
 missing from it renders `--` rather than a number nobody checked.

{interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/bench_spec.py

@@ -264,12 +264,14 @@ VARIANTS: tuple[VariantSpec, ...] = (
     # Its write tap is `mlp_out`, not `resid_post`: the steer workload addresses the model's
     # `capture_point`, and a hyper-connection trunk refuses the default name (`run_bench._steer_site`).
     #
-    # `freeze_writes` is what makes the `steer` cell a number instead of `n/a`. `"auto"` installs
+    # `freeze_writes` was once what made the `steer` cell a number instead of `n/a`: `"auto"` installed
     # reads, and a steering op needs a write tap to land in, so this row priced capture under replay
     # and left the other half of the feature unmeasured -- with a message that blamed graph replay for
-    # it. The value is the sentinel `run_bench.STEER_WRITES`, resolved there to the mid-stack
-    # `resid_post` the workload steers, because the layer differs per model and a freeze write has to
-    # be named before the model exists.
+    # it. Auto covers writes now, so the cell stands either way, and naming them here has become a
+    # *narrowing*: one write buffer at the site the workload steers rather than one per layer, which
+    # is what keeps this row's memory comparable with the columns beside it. The value is the sentinel
+    # `run_bench.STEER_WRITES`, resolved there to the mid-stack `resid_post` the workload steers,
+    # because the layer differs per model and a freeze write has to be named before the model exists.
     VariantSpec(
         "vllm-freeze",
         "vllm",

{interp_engine-1.2.7 → interp_engine-1.2.8}/docs/PERFORMANCE.md

@@ -124,15 +124,24 @@ There *is* a middle ground, which was an open question in this file and is now a
 preallocated buffer plus a `copy_` that graph capture records, so replay re-executes the tap instead
 of skipping it. That is `freeze_points` on `VLLMModel`, and it is opt-in:
 
+```python
+model = load_model("Qwen/Qwen3-8B", freeze_points="auto")  # resid_post at every layer, read and write, under graphs
+```
+
 - **`freeze_points=None`** (omit it) — hooked vLLM, `enforce_eager=True`. The default, and the only
   configuration that serves every point.
-- **`freeze_points="auto"`** — `resid_post` at every layer, via static `copy_` taps. Keeps graph
-  replay and turns Dynamo off (`VLLM_USE_BREAKABLE_CUDAGRAPH=1`), so the win is replay without
-  compile. Reads only: name write sites with `freeze_writes`, which wraps them for steering. Because
-  breakable `add_eager` keeps the wrap as ordinary PyTorch on the live tensor, the write path is not
-  limited to additive — `orthogonal`, `projection_cap` and the j-lens `steer`/`ablate`/`swap` ops all
-  ride the same wrap, per-request and with `position_mask` honoured (`register_freeze_write`). An op
-  outside that set is refused rather than silently skipped.
+- **`freeze_points="auto"`** — `resid_post` at every layer, to read *and* to write, via static
+  `copy_` taps. Keeps graph replay and turns Dynamo off (`VLLM_USE_BREAKABLE_CUDAGRAPH=1`), so the
+  win is replay without compile. Because breakable `add_eager` keeps the wrap as ordinary PyTorch on
+  the live tensor, the write path is not limited to additive — `orthogonal`, `projection_cap` and the
+  j-lens `steer`/`ablate`/`swap` ops all ride the same wrap, per-request and with `position_mask`
+  honoured (`register_freeze_write`). An op outside that set is refused rather than silently skipped.
+  Auto covers the write because the two halves are one decision: a read tap alone serves the lens
+  read-out and refuses every steer, ablation and swap derived from it, at an address already tapped.
+  Pass `freeze_writes=[]` for the reads without the write buffers, and an explicit `freeze_points`
+  list to name both halves yourself — neither is filled in for a caller who said something.
+- **`freeze_writes=[...]`** — write sites, and on their own (no `freeze_points`) an engine that
+  writes and does not read.
 - **`freeze_points=[]`** — no taps, graphs and inductor both on. Generation-only, the same thing as
   `enforce_eager=False` but reached through the flag that knows what it costs. Prefer it to setting
   `enforce_eager` yourself.
@@ -150,11 +159,23 @@ What freeze does not do, which is why `enforce_eager=True` is still the default:
   served short.
 - **A narrower batch.** Freeze buffers have to fit alongside the graphs, so `max_num_batched_tokens`
   is stepped down (16384 → 1024) to make room, and refuses rather than OOM-ing in graph capture.
+  `freeze_writes=[]` buys some of that width back when the write buffers are what stands between you
+  and a wide enough batch.
+- **Four times the buffer on a hyper-connection trunk.** `deepseek-v4-flash-0731`'s block carries
+  four parallel residual streams, so `"auto"` freezes `resid_streams` — the whole stack per layer,
+  four times the width and so four times the buffer. That is what its freeze row in the sweep prices.
+- **Graphs for decode only on a hybrid trunk.** `qwen3.8-27b`, and any other linear-attention or
+  unclassified trunk, runs prefill eagerly and keeps the decode graphs
+  (`freeze.decode_only_graphs_reason`). Freeze pins that mode rather than offering it: breakable
+  graphs turn `torch.compile` off, and vLLM's mixed prefill-decode graph capture then miscomputes
+  prefill on a gated-delta trunk — a whole wrong forward, not a bad tap, and reproducible on plain
+  `vllm.LLM`. Eager prefill costs little here, since the wraps are ordinary PyTorch and replay's win
+  is in decode.
 - **One engine per process.** `VLLM_USE_BREAKABLE_CUDAGRAPH` is process-global, so a freeze engine
   and a compiled engine cannot share a process.
 - **Measured for throughput, not yet for capture latency.** The sweep runs `generate` and
-  `generate_x8` on the freeze variant and only on the three conventional trunks above; its capture,
-  steering and lens cells are unmeasured. Capture *correctness* under freeze is covered separately —
+  `generate_x8` on the freeze variant, and the tables above compare only the three conventional
+  trunks; its capture, steering and lens cells are unmeasured. Capture *correctness* under freeze is covered separately —
   `tests/test_freeze_parity_gpu.py` holds the harvest to cosine ≥ 0.999 and greedy token-id parity
   against hooked capture, across concurrent requests, chunked prefill and decode rows — and the
   backend self-tests each tap at startup, refusing to serve if a `copy_` produced a non-finite or

{interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_hooks.py

@@ -8,6 +8,8 @@ hooks (steering then capture) in :mod:`~interp_engine.vllm_capture.requests` ins
 
 from __future__ import annotations
 
+from collections.abc import Iterable
+
 import torch
 
 from interp_engine.hooks import hidden_arg_index
@@ -15,6 +17,29 @@ from interp_engine.hooks import hidden_from_call as _hidden_from_call
 from interp_engine.vllm_capture._payload import select_stream
 from interp_engine.vllm_capture._tree import LAYER_RETURN_INDEX
 
+
+def position_mask(positions: Iterable[int], num_tokens: int, like: torch.Tensor) -> torch.Tensor:
+    """A boolean row mask over ``positions``, shaped to broadcast against ``like``.
+
+    One trailing singleton axis per axis ``like`` has after the token one, rather than exactly one.
+    On a hyper-connection trunk a point is ``[tokens, streams, width]``, and a ``[tokens, 1]`` mask
+    would line its token axis up against the *stream* axis: a shape error for most prompts, and --
+    worse -- a silent masking of the wrong thing for a prompt whose length happens to equal the
+    stream count.
+
+    Lives here, in the leaf both write paths already import, because there is one right answer and
+    two callers: the hooked path in :mod:`~interp_engine.vllm_capture.requests` and the frozen one
+    in :mod:`~interp_engine.vllm_capture.freeze`. It was fixed in the first and not the second while
+    they were separate copies, which is how a jlens intervention on a stream stack came to crash on
+    a shape under CUDA graphs and work under hooks.
+    """
+    mask = torch.zeros(num_tokens, *([1] * (like.dim() - 1)), dtype=torch.bool, device=like.device)
+    for position in positions:
+        if 0 <= position < num_tokens:
+            mask[position] = True
+    return mask
+
+
 # vLLM's decoder-layer convention is `forward(positions, hidden, residual) -> (hidden, residual)`
 # where the returned `hidden` is the part NOT yet added to the residual stream: the fused add+norm at
 # the top of the *next* block does that add, so the stream is the sum of the pair.

{interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/freeze.py

@@ -26,7 +26,7 @@ from interp_engine.facts import is_linear_attention_layer, unclassified_layer_ki
 from interp_engine.hooks import hidden_arg_index, hidden_from_call
 from interp_engine.points import steer_refusal_reason
 from interp_engine.vllm_capture._demux import _ensure_patched, _get_demux, _resolve_rid
-from interp_engine.vllm_capture._hooks import layer_return_tensor, returns_full_residual
+from interp_engine.vllm_capture._hooks import layer_return_tensor, position_mask, returns_full_residual
 from interp_engine.vllm_capture._payload import attn_payload_key, decode_capture_payload, encode_tensor_payload
 from interp_engine.vllm_capture._tree import (
     _INPUT_POINTS,
@@ -1394,7 +1394,15 @@ def _is_fused_resid(module: torch.nn.Module, name: str, residual: torch.Tensor |
 
 
 def _apply_lens_scope(delta: torch.Tensor, n: int, scope: dict[str, Any] | None) -> torch.Tensor | None:
-    """Prefill-vs-decode skip used by jlens. None means leave the live tensor alone."""
+    """Prefill-vs-decode skip used by jlens. None means leave the live tensor alone.
+
+    The mask comes from :func:`~interp_engine.vllm_capture._hooks.position_mask`, the one the hooked
+    path uses, rather than being built here: a ``[tokens, 1]`` mask is right for every point with one
+    width axis and wrong for a hyper-connection trunk, where the delta is ``[tokens, streams,
+    width]`` and broadcasting -- which pads on the LEFT -- lines the token axis up against the stream
+    axis. There is nothing about the answer that differs between the two paths, so there is no
+    longer a second construction of it here.
+    """
     if not scope:
         return delta
     is_prefill = n > 1
@@ -1403,10 +1411,7 @@ def _apply_lens_scope(delta: torch.Tensor, n: int, scope: dict[str, Any] | None)
     skip = scope.get("skip_positions") or []
     prompt_len = int(scope.get("prompt_len") or 0)
     if is_prefill and skip and n == prompt_len:
-        mask = torch.zeros(n, 1, dtype=torch.bool, device=delta.device)
-        for i in skip:
-            if 0 <= int(i) < n:
-                mask[int(i)] = True
+        mask = position_mask((int(i) for i in skip), n, delta)
         delta = torch.where(mask, torch.zeros_like(delta), delta)
     return delta
 
@@ -1607,6 +1612,13 @@ def worker_set_freeze_delta(
     Additive ``op="add"`` without a lens scope fills the static ``delta`` buffer. Orthogonal,
     projection_cap, and lens ops attach a live ``modify`` that reads the residual each
     forward (breakable ``add_eager``). ``lens_scope`` is the jlens prefill/decode skip.
+
+    A ``stream`` disqualifies the static buffer too. The buffer is added whole to a ``[tokens,
+    streams, width]`` activation, so it has no way to say "this stream and not the others"; taking
+    the fast path anyway would steer all four streams of a hyper-connection trunk on a request that
+    asked for one, and report success. :func:`~interp_engine.vllm_capture.steering._one_stream`
+    knows how, so a stream is served by the modifier path like every other op that is not a plain
+    constant.
     """
     freeze = _state(worker)
     if freeze is None:
@@ -1621,7 +1633,7 @@ def worker_set_freeze_delta(
             raise ValueError(f"graph freeze cannot serve op={op!r}; supported ops are {sorted(_FREEZE_WRITE_OPS)}")
         site = _write_site(freeze, Address(str(spec["point"]), int(spec["layer"])))
         assert site.delta is not None
-        if op == "add" and lens_scope is None:
+        if op == "add" and lens_scope is None and spec.get("stream") is None:
             vec = torch.tensor(spec["vector"], dtype=torch.float32, device=site.delta.device)
             vec = (vec * float(spec["coeff"])).to(dtype=site.delta.dtype)
             site.delta.copy_(vec.reshape(1, -1).expand_as(site.delta))
@@ -1647,7 +1659,9 @@ def _compile_write_req(
         raise ValueError(f"graph freeze cannot serve op={op!r}; supported ops are {sorted(_FREEZE_WRITE_OPS)}")
     assert site.delta is not None
     device, dtype = site.delta.device, site.delta.dtype
-    if op == "add":
+    # A constant `[1, width]` vector broadcasts over a stream axis and so cannot exclude one; see
+    # `worker_set_freeze_delta` for why a `stream` therefore has to go the modifier way.
+    if op == "add" and spec.get("stream") is None:
         vec = torch.tensor(spec["vector"], dtype=torch.float32, device=device)
         vec = (vec * float(spec["coeff"])).to(dtype=dtype).reshape(1, -1)
         return _WriteReq(

{interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/requests.py

@@ -14,7 +14,6 @@ cycle around :mod:`~interp_engine.vllm_capture._demux`.
 from __future__ import annotations
 
 import inspect
-from collections.abc import Iterable
 from typing import cast
 
 import torch
@@ -37,6 +36,7 @@ from interp_engine.vllm_capture._hooks import (
     layer_return_tensor,
     returns_full_residual,
 )
+from interp_engine.vllm_capture._hooks import position_mask as _position_mask
 from interp_engine.vllm_capture._payload import (
     attn_payload_key,
     encode_tensor_payload,
@@ -64,22 +64,6 @@ from interp_engine.vllm_capture.steering import _make_steer_modifier
 # --- the per-request combined hook body ---------------------------------------
 
 
-def _position_mask(positions: Iterable[int], num_tokens: int, like: torch.Tensor) -> torch.Tensor:
-    """A boolean row mask over ``positions``, shaped to broadcast against ``like``.
-
-    One trailing singleton axis per axis ``like`` has after the token one, rather than exactly one.
-    On a hyper-connection trunk a point is ``[tokens, streams, width]``, and a ``[tokens, 1]`` mask
-    would line its token axis up against the *stream* axis: a shape error for most prompts, and --
-    worse -- a silent masking of the wrong thing for a prompt whose length happens to equal the
-    stream count.
-    """
-    mask = torch.zeros(num_tokens, *([1] * (like.dim() - 1)), dtype=torch.bool, device=like.device)
-    for position in positions:
-        if 0 <= position < num_tokens:
-            mask[position] = True
-    return mask
-
-
 def _process_point(demux: _Demux, site: Address, full: torch.Tensor) -> torch.Tensor:
     """Apply per-request steering/lens then capture at one hook ``site``.
 

{interp_engine-1.2.7 → interp_engine-1.2.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "interp-engine"
-version = "1.2.7"
+version = "1.2.8"
 description = "A fast, standardized interpretability engine that supports most modern models and architectures. Powers Neuronpedia."
 readme = "README.md"
 license = "Apache-2.0"