diff --git a/docs/INTERNALS.md b/docs/INTERNALS.md index 61d7642..1b6af4e 100644 --- a/docs/INTERNALS.md +++ b/docs/INTERNALS.md @@ -54,7 +54,9 @@ files and why they are where they are. - `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. + take, registering per-request on vLLM rather than installing a global hook. The methods are named + once, in `steer_specs.SteerMethod`: the eager `SteerSpec.method` and the worker dict's `op` both + read from it, so a method has one spelling on either backend. - `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](PORTING.md#porting-from-transformerlens-nnsight-or-nnterp). diff --git a/interp_engine/__init__.py b/interp_engine/__init__.py index 1f0235c..86b7da7 100644 --- a/interp_engine/__init__.py +++ b/interp_engine/__init__.py @@ -120,6 +120,8 @@ ProjectionCapSpec, SteeringOp, SteeringSpec, + SteerMethod, + steer_method, steering_spec_to_worker_specs, ) from interp_engine.sync import SyncModel, sync_model @@ -245,6 +247,7 @@ "PositionMask", "ProjectionCapSpec", "SteerMask", + "SteerMethod", "SteerSpec", "SteeringOp", "SteeringSpec", @@ -252,6 +255,7 @@ "resolve_masked_positions", "steer", "steer_delta", + "steer_method", "steering_spec_to_eager_specs", "steering_spec_to_worker_specs", "unit_vector", diff --git a/interp_engine/steer.py b/interp_engine/steer.py index 524631e..92dc9c8 100644 --- a/interp_engine/steer.py +++ b/interp_engine/steer.py @@ -32,6 +32,8 @@ OrthogonalDecompSpec, ProjectionCapSpec, SteeringSpec, + SteerMethod, + steer_method, ) from interp_engine.sync import sync_model @@ -156,10 +158,10 @@ def projection_cap_delta( return (capped - projection) * unit -#: The steering methods :class:`SteerSpec` accepts, each with the same arithmetic as the -#: worker-side op of the same name (``vllm_capture/steering.py``). Kept as a tuple so the -#: refusal for an unknown method can list them. -STEER_METHODS = ("additive", "orthogonal", "projection_cap") +#: The steering methods :class:`SteerSpec` accepts, as their spellings. Derived from +#: :class:`~interp_engine.steer_specs.SteerMethod`, which the worker-side op reads from too, and +#: kept under this name so callers that listed the methods keep working. +STEER_METHODS: tuple[str, ...] = tuple(m.value for m in SteerMethod) @dataclass @@ -169,7 +171,7 @@ class SteerSpec: vector: torch.Tensor layer: int coeff: float = 1.0 - method: str = "additive" # one of STEER_METHODS + method: SteerMethod = SteerMethod.ADDITIVE point: str = "resid_post" normalize: bool = False stream: int | None = None @@ -195,6 +197,11 @@ class SteerSpec: max: float | None = None """Upper bound for ``method="projection_cap"``. See :attr:`min`.""" + def __post_init__(self) -> None: + # A spec built from a string -- a config file, a wire payload -- is checked here, at + # construction, rather than at the first forward it reaches. An unknown name lists the members. + self.method = steer_method(self.method) + def _prepared_vector(spec: SteerSpec, ref: torch.Tensor) -> torch.Tensor: vec = spec.vector.to(device=ref.device, dtype=ref.dtype) @@ -214,13 +221,13 @@ def steer_delta(spec: SteerSpec, activations: torch.Tensor, vector: torch.Tensor orthogonal case subtly wrong if either half is forgotten. It is also the shape the vLLM worker's modifiers already have, so the two backends run the same expressions. """ - if spec.method == "additive": - return spec.coeff * vector - if spec.method == "orthogonal": - return OrthogonalProjector(vector).delta(activations, spec.coeff) - if spec.method == "projection_cap": - return projection_cap_delta(activations, vector, minimum=spec.min, maximum=spec.max) - raise ValueError(f"Unknown steering method {spec.method!r}; expected one of {', '.join(STEER_METHODS)}") + match steer_method(spec.method): + case SteerMethod.ADDITIVE: + return spec.coeff * vector + case SteerMethod.ORTHOGONAL: + return OrthogonalProjector(vector).delta(activations, spec.coeff) + case SteerMethod.PROJECTION_CAP: + return projection_cap_delta(activations, vector, minimum=spec.min, maximum=spec.max) @dataclass(frozen=True) @@ -420,12 +427,16 @@ def steering_spec_to_eager_specs(spec: SteeringSpec, *, point: str | None = None if isinstance(op, AddSpec): vector = op.vector if isinstance(op.vector, torch.Tensor) else torch.tensor(op.vector) out.append( - SteerSpec(vector=vector, layer=int(layer), coeff=float(op.scale), method="additive", **where) + SteerSpec( + vector=vector, layer=int(layer), coeff=float(op.scale), method=SteerMethod.ADDITIVE, **where + ) ) elif isinstance(op, OrthogonalDecompSpec): vector = op.vector if isinstance(op.vector, torch.Tensor) else torch.tensor(op.vector) out.append( - SteerSpec(vector=vector, layer=int(layer), coeff=float(op.coeff), method="orthogonal", **where) + SteerSpec( + vector=vector, layer=int(layer), coeff=float(op.coeff), method=SteerMethod.ORTHOGONAL, **where + ) ) elif isinstance(op, ProjectionCapSpec): vector = op.vector if isinstance(op.vector, torch.Tensor) else torch.tensor(op.vector) @@ -433,7 +444,7 @@ def steering_spec_to_eager_specs(spec: SteeringSpec, *, point: str | None = None SteerSpec( vector=vector, layer=int(layer), - method="projection_cap", + method=SteerMethod.PROJECTION_CAP, min=op.min, max=op.max, **where, diff --git a/interp_engine/steer_specs.py b/interp_engine/steer_specs.py index a8a0f62..d83670b 100644 --- a/interp_engine/steer_specs.py +++ b/interp_engine/steer_specs.py @@ -9,10 +9,43 @@ from __future__ import annotations from dataclasses import dataclass, field +from enum import StrEnum import torch +class SteerMethod(StrEnum): + """The steering arithmetics, named once for both backends. + + The eager ``SteerSpec.method``, the worker dict's ``op`` and the static tap's op set all read from + here, so a method has one spelling on every path. It used to have two -- ``additive`` on the + eager spec and ``add`` on the worker dict -- and nothing but a runtime set-membership check told + a caller which one it was speaking to. A ``StrEnum`` so a member compares equal to its own + spelling: a spec loaded from JSON with ``"orthogonal"`` in it needs no translation. + + A new method is a new member here, plus a delta in ``steer.py`` and a modifier in + ``vllm_capture/steering.py`` -- see ``AGENTS.md``, "Steering arithmetic is one function per + method, shared with the worker". + """ + + ADDITIVE = "additive" + ORTHOGONAL = "orthogonal" + PROJECTION_CAP = "projection_cap" + + +def steer_method(value: SteerMethod | str) -> SteerMethod: + """``value`` as a :class:`SteerMethod`, refusing anything else with the members listed. + + The enum's own ``ValueError`` says only that the value is invalid; a caller who wrote ``"add"`` + is better served by seeing ``additive`` in the message. + """ + try: + return SteerMethod(value) + except ValueError: + expected = ", ".join(m.value for m in SteerMethod) + raise ValueError(f"Unknown steering method {value!r}; expected one of {expected}") from None + + def _as_list(vector: torch.Tensor | list[float]) -> list[float]: if isinstance(vector, torch.Tensor): return vector.detach().float().flatten().tolist() @@ -102,7 +135,7 @@ def steering_spec_to_worker_specs(spec: SteeringSpec, *, point: str | None = Non { "layer": int(layer), **where, - "op": "add", + "op": SteerMethod.ADDITIVE.value, "vector": _as_list(op.vector), "coeff": float(op.scale), } @@ -112,7 +145,7 @@ def steering_spec_to_worker_specs(spec: SteeringSpec, *, point: str | None = Non { "layer": int(layer), **where, - "op": "projection_cap", + "op": SteerMethod.PROJECTION_CAP.value, "vector": _as_list(op.vector), "min": op.min, "max": op.max, @@ -123,7 +156,7 @@ def steering_spec_to_worker_specs(spec: SteeringSpec, *, point: str | None = Non { "layer": int(layer), **where, - "op": "orthogonal", + "op": SteerMethod.ORTHOGONAL.value, "vector": _as_list(op.vector), "coeff": float(op.coeff), } diff --git a/interp_engine/vllm_backend.py b/interp_engine/vllm_backend.py index 61f3377..ff00ff8 100644 --- a/interp_engine/vllm_backend.py +++ b/interp_engine/vllm_backend.py @@ -34,6 +34,7 @@ from interp_engine.points import d_model_wide, hyper_connection_names, refusal_reasons from interp_engine.points import steer_refusal_reason as points_steer_refusal from interp_engine.residual_basis import ResidualBasis, vllm_residual_basis +from interp_engine.steer_specs import SteerMethod from interp_engine.vllm_capture import ( _GLOBAL_POINTS, DEFAULT_HS_STORAGE_PATH, @@ -52,6 +53,7 @@ DECODE_ONLY_GRAPHS, STATIC_ENV, STATIC_SKIP_ABSENT_ENV, + STATIC_WRITE_OPS, apply_breakable_env, decode_only_graphs_reason, encode_static_env, @@ -1504,20 +1506,20 @@ def _require_capture_points(self, pts: Sequence[str], what: str) -> None: def _require_static_writes(self, specs: Sequence[dict], what: str) -> None: """Allow static writes when hooks run, or when every write site is declared. - Additive ``add_`` and live-read ops (orthogonal, projection_cap, lens steer/ablate/swap) - all ride the same static wrap. A site miss is a 400, not a silent no-op. + The additive static ``add_`` and live-read ops (orthogonal, projection_cap, lens + steer/ablate/swap) all ride the same static wrap. A site miss is a 400, not a silent no-op. """ if self.hooks_available: return writes = getattr(self, "_static_writes", frozenset()) if self.graph_replay and writes: for spec in specs: - op = spec.get("op", "add") - if op not in {"add", "orthogonal", "projection_cap", "steer", "ablate", "swap"}: + op = spec.get("op", SteerMethod.ADDITIVE) + if op not in STATIC_WRITE_OPS: raise RuntimeError( f"{what} op {op!r} is not one a static write tap can apply; " - f"backend='vllm-static' serves add, orthogonal, projection_cap, steer, " - f"ablate and swap. Use backend='vllm' for this one." + f"backend='vllm-static' serves {', '.join(sorted(STATIC_WRITE_OPS))}. " + f"Use backend='vllm' for this one." ) site = Address(str(spec.get("point") or "resid_post"), int(spec["layer"])) if not any(alias in writes for alias in resid_stream_aliases(site)): diff --git a/interp_engine/vllm_capture/static.py b/interp_engine/vllm_capture/static.py index 6829918..65ebe30 100644 --- a/interp_engine/vllm_capture/static.py +++ b/interp_engine/vllm_capture/static.py @@ -25,6 +25,7 @@ from interp_engine.facts import is_linear_attention_layer, unclassified_layer_kinds from interp_engine.hooks import hidden_arg_index, hidden_from_call from interp_engine.points import steer_refusal_reason +from interp_engine.steer_specs import SteerMethod from interp_engine.vllm_capture._demux import _ensure_patched, _get_demux, _resolve_rid from interp_engine.vllm_capture._hooks import ( flat_value, @@ -675,10 +676,10 @@ def fit_max_num_batched_tokens( # --- worker state ------------------------------------------------------------ -#: Ops static writes can serve. ``add`` uses a static ``add_``; the rest read the live residual. -_STEER_OPS = frozenset({"add", "orthogonal", "projection_cap"}) +#: Ops static writes can serve. ``additive`` uses a static ``add_``; the rest read the live residual. +_STEER_OPS = frozenset(m.value for m in SteerMethod) _LENS_OPS = frozenset({"steer", "ablate", "swap"}) -_STATIC_WRITE_OPS = _STEER_OPS | _LENS_OPS +STATIC_WRITE_OPS = _STEER_OPS | _LENS_OPS @dataclass @@ -1811,7 +1812,7 @@ def worker_set_static_delta( ) -> None: """Install static writes from worker specs. Zeros every write site first. - Additive ``op="add"`` without a lens scope fills the static ``delta`` buffer. Orthogonal, + ``op="additive"`` 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. @@ -1830,12 +1831,12 @@ def worker_set_static_delta( from interp_engine.vllm_capture.steering import _make_steer_modifier for spec in specs: - op = str(spec.get("op", "add")) - if op not in _STATIC_WRITE_OPS: - raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(_STATIC_WRITE_OPS)}") + op = str(spec.get("op", SteerMethod.ADDITIVE)) + if op not in STATIC_WRITE_OPS: + raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}") site = _write_site(static, Address(str(spec["point"]), int(spec["layer"]))) assert site.delta is not None - if op == "add" and lens_scope is None and spec.get("stream") is None: + if op == SteerMethod.ADDITIVE 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)) @@ -1857,14 +1858,14 @@ def _compile_write_req( prompt_len: int, steer_generated: bool, ) -> _WriteReq: - op = str(spec.get("op", "add")) - if op not in _STATIC_WRITE_OPS: - raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(_STATIC_WRITE_OPS)}") + op = str(spec.get("op", SteerMethod.ADDITIVE)) + if op not in STATIC_WRITE_OPS: + raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}") assert site.delta is not None device, dtype = site.delta.device, site.delta.dtype # A constant `[1, width]` vector broadcasts over a stream axis and so cannot exclude one; see # `worker_set_static_delta` for why a `stream` therefore has to go the modifier way. - if op == "add" and spec.get("stream") is None: + if op == SteerMethod.ADDITIVE 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( diff --git a/interp_engine/vllm_capture/steering.py b/interp_engine/vllm_capture/steering.py index 0cf3d4d..ff73ff6 100644 --- a/interp_engine/vllm_capture/steering.py +++ b/interp_engine/vllm_capture/steering.py @@ -11,6 +11,7 @@ import torch +from interp_engine.steer_specs import SteerMethod, steer_method from interp_engine.vllm_capture._tree import _attn_out_proj, _get_layers, _worker_model, absent_point_reason # --- worker-side steering (additive write-hooks) ----------------------------- @@ -53,10 +54,10 @@ def _modify(full: torch.Tensor) -> torch.Tensor: def _make_steer_modifier(spec: dict, dev, dt): """Return ``modify(full_resid) -> delta`` (the tensor to ADD to the residual). - ``op="add"`` -> constant ``coeff*vector`` (broadcast). ``op="projection_cap"`` - -> clamp the residual's projection onto ``vector`` into ``[min, max]`` by adding - ``(clamp(proj)-proj)*unit_vector``. ``op="orthogonal"`` -> rescale the projection - onto ``vector`` by ``coeff`` (``h -> (I-P)h + coeff*P h``) by adding + ``op`` is a :class:`~interp_engine.steer_specs.SteerMethod`. ``additive`` -> constant + ``coeff*vector`` (broadcast). ``projection_cap`` -> clamp the residual's projection onto + ``vector`` into ``[min, max]`` by adding ``(clamp(proj)-proj)*unit_vector``. ``orthogonal`` -> + rescale the projection onto ``vector`` by ``coeff`` (``h -> (I-P)h + coeff*P h``) by adding ``(coeff-1)*proj*unit_vector``, matching the eager ``OrthogonalProjector``. ``spec["stream"]`` restricts the delta to one residual stream of a hyper-connection trunk; see @@ -64,55 +65,54 @@ def _make_steer_modifier(spec: dict, dev, dt): right default for a point that has no such axis and the only sane one for a point that does -- the alternative would be to pick a stream on the caller's behalf. """ - op = spec.get("op", "add") + op = steer_method(spec.get("op", SteerMethod.ADDITIVE)) vec = torch.tensor(spec["vector"], dtype=torch.float32).to(dev, dt) stream = spec.get("stream") if stream is not None: return _one_stream(_make_steer_modifier({**spec, "stream": None}, dev, dt), int(stream)) - if op == "add": - delta = vec * float(spec["coeff"]) + match op: + case SteerMethod.ADDITIVE: + delta = vec * float(spec["coeff"]) - def _modify_add(_full: torch.Tensor) -> torch.Tensor: - return delta + def _modify_add(_full: torch.Tensor) -> torch.Tensor: + return delta - return _modify_add + return _modify_add - if op == "orthogonal": - unit = vec / vec.norm().clamp_min(1e-12) - coeff = float(spec["coeff"]) + case SteerMethod.ORTHOGONAL: + unit = vec / vec.norm().clamp_min(1e-12) + coeff = float(spec["coeff"]) - def _modify_orthogonal(full: torch.Tensor) -> torch.Tensor: - proj = (full.to(unit.dtype) * unit).sum(dim=-1, keepdim=True) # [T,1] - return (coeff - 1.0) * proj * unit + def _modify_orthogonal(full: torch.Tensor) -> torch.Tensor: + proj = (full.to(unit.dtype) * unit).sum(dim=-1, keepdim=True) # [T,1] + return (coeff - 1.0) * proj * unit - return _modify_orthogonal + return _modify_orthogonal - if op == "projection_cap": - unit = vec / vec.norm().clamp_min(1e-12) - lo = spec.get("min") - hi = spec.get("max") + case SteerMethod.PROJECTION_CAP: + unit = vec / vec.norm().clamp_min(1e-12) + lo = spec.get("min") + hi = spec.get("max") - def _modify_projection_cap(full: torch.Tensor) -> torch.Tensor: - proj = (full.to(unit.dtype) * unit).sum(dim=-1, keepdim=True) # [T,1] - capped = proj - if lo is not None: - capped = torch.clamp(capped, min=float(lo)) - if hi is not None: - capped = torch.clamp(capped, max=float(hi)) - return (capped - proj) * unit + def _modify_projection_cap(full: torch.Tensor) -> torch.Tensor: + proj = (full.to(unit.dtype) * unit).sum(dim=-1, keepdim=True) # [T,1] + capped = proj + if lo is not None: + capped = torch.clamp(capped, min=float(lo)) + if hi is not None: + capped = torch.clamp(capped, max=float(hi)) + return (capped - proj) * unit - return _modify_projection_cap - - raise ValueError(f"Unsupported steering op {op!r}") + return _modify_projection_cap def worker_install_steering(worker: object, specs: list[dict]) -> None: """Install steering write-hooks on the worker model. ``specs`` items: ``{"layer": int, "point": "resid_post"|"resid_pre"|"z", - "op": "add"|"projection_cap", "vector": list[float], and op args ("coeff" for - add; "min"/"max" for projection_cap)}``. Applied to all token positions. + "op": a SteerMethod value, "vector": list[float], and op args ("coeff" for + additive and orthogonal; "min"/"max" for projection_cap)}``. Applied to all token positions. Single request-locked use (no per-request demux). Requires enforce_eager. """ model = _worker_model(worker) diff --git a/tests/test_static_set.py b/tests/test_static_set.py index 90bf35a..ab3209d 100644 --- a/tests/test_static_set.py +++ b/tests/test_static_set.py @@ -1184,7 +1184,7 @@ def forward(self, hidden, residual): # noqa: ANN001 worker_register_static_write( worker, "steer-me", - [{"point": "resid_post", "layer": 0, "op": "add", "vector": [1.0, 0.0], "coeff": 1.0}], + [{"point": "resid_post", "layer": 0, "op": "additive", "vector": [1.0, 0.0], "coeff": 1.0}], ) demux.current_meta = (["steer-me", "leave-me"], [2, 2]) _wrap_module(layer, [(site, "write")], worker) @@ -1211,7 +1211,7 @@ def forward(self, hidden, residual): # noqa: ANN001 worker_register_static_write( worker, "r", - [{"point": "resid_post", "layer": 0, "op": "add", "vector": [1.0, 0.0], "coeff": 1.0}], + [{"point": "resid_post", "layer": 0, "op": "additive", "vector": [1.0, 0.0], "coeff": 1.0}], skip_positions=[0], prompt_len=2, ) @@ -1240,12 +1240,12 @@ def forward(self, hidden, residual): # noqa: ANN001 worker_register_static_write( worker, "a", - [{"point": "resid_post", "layer": 0, "op": "add", "vector": [1.0, 0.0], "coeff": 1.0}], + [{"point": "resid_post", "layer": 0, "op": "additive", "vector": [1.0, 0.0], "coeff": 1.0}], ) worker_register_static_write( worker, "b", - [{"point": "resid_post", "layer": 0, "op": "add", "vector": [0.0, 1.0], "coeff": 1.0}], + [{"point": "resid_post", "layer": 0, "op": "additive", "vector": [0.0, 1.0], "coeff": 1.0}], ) demux.current_meta = (["a", "b"], [1, 1]) _wrap_module(layer, [(site, "write")], worker) diff --git a/tests/test_steer_math_parity.py b/tests/test_steer_math_parity.py index 735b21e..2cd436d 100644 --- a/tests/test_steer_math_parity.py +++ b/tests/test_steer_math_parity.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import cast + import pytest import torch @@ -31,6 +33,8 @@ OrthogonalDecompSpec, ProjectionCapSpec, SteeringSpec, + SteerMethod, + steering_spec_to_worker_specs, ) from interp_engine.vllm_capture.steering import _make_steer_modifier @@ -61,14 +65,14 @@ def test_additive_matches_the_worker() -> None: """Both sides return a ``[d_model]`` delta that broadcasts across positions.""" residual, vector, coeff = _residual(), _vector(), 2.5 eager = steer_delta(SteerSpec(vector=vector, layer=0, coeff=coeff), residual, vector) - worker = _worker_delta({"op": "add", "vector": vector.tolist(), "coeff": coeff}, residual) + worker = _worker_delta({"op": "additive", "vector": vector.tolist(), "coeff": coeff}, residual) torch.testing.assert_close(eager, worker) @pytest.mark.parametrize("coeff", [0.0, 0.5, 1.0, -1.0, 3.0]) def test_orthogonal_matches_the_worker(coeff: float) -> None: residual, vector = _residual(), _vector() - eager = steer_delta(SteerSpec(vector=vector, layer=0, coeff=coeff, method="orthogonal"), residual, vector) + eager = steer_delta(SteerSpec(vector=vector, layer=0, coeff=coeff, method=SteerMethod.ORTHOGONAL), residual, vector) worker = _worker_delta({"op": "orthogonal", "vector": vector.tolist(), "coeff": coeff}, residual) torch.testing.assert_close(eager, worker) @@ -77,7 +81,9 @@ def test_orthogonal_matches_the_worker(coeff: float) -> None: def test_projection_cap_matches_the_worker(lo: float | None, hi: float | None) -> None: """The method that had no eager implementation at all until this change.""" residual, vector = _residual(), _vector() - eager = steer_delta(SteerSpec(vector=vector, layer=0, method="projection_cap", min=lo, max=hi), residual, vector) + eager = steer_delta( + SteerSpec(vector=vector, layer=0, method=SteerMethod.PROJECTION_CAP, min=lo, max=hi), residual, vector + ) worker = _worker_delta({"op": "projection_cap", "vector": vector.tolist(), "min": lo, "max": hi}, residual) torch.testing.assert_close(eager, worker) @@ -175,7 +181,7 @@ def test_every_backend_agnostic_op_converts_to_an_eager_spec() -> None: got = steering_spec_to_eager_specs(spec) - assert [s.method for s in got] == ["additive", "orthogonal", "projection_cap"] + assert [s.method for s in got] == [SteerMethod.ADDITIVE, SteerMethod.ORTHOGONAL, SteerMethod.PROJECTION_CAP] assert all(s.layer == 3 and s.point == "resid_post" for s in got) assert (got[2].min, got[2].max) == (-1.0, 1.0) @@ -267,12 +273,50 @@ def test_a_stream_on_a_point_without_one_is_refused_rather_than_ignored() -> Non from interp_engine.vllm_capture.steering import _make_steer_modifier modify = _make_steer_modifier( - {"op": "add", "vector": _vector().tolist(), "coeff": 1.0, "stream": 0}, torch.device("cpu"), torch.float32 + {"op": "additive", "vector": _vector().tolist(), "coeff": 1.0, "stream": 0}, torch.device("cpu"), torch.float32 ) with pytest.raises(ValueError, match="no stream axis"): modify(torch.zeros(4, D_MODEL)) def test_an_unknown_method_names_the_ones_that_exist() -> None: - with pytest.raises(ValueError, match="projection_cap"): - steer_delta(SteerSpec(vector=_vector(), layer=0, method="nope"), _residual(), _vector()) + # Refused at construction, before any forward, and the message lists every member. + with pytest.raises(ValueError, match="additive, orthogonal, projection_cap"): + SteerSpec(vector=_vector(), layer=0, method=cast(SteerMethod, "nope")) + + +# ── one vocabulary on both sides ──────────────────────────────────────────────────────────── + + +def test_a_spec_built_from_a_string_holds_the_enum() -> None: + """A spec loaded from JSON says ``"orthogonal"``; it must become the member, not stay a str.""" + spec = SteerSpec(vector=_vector(), layer=0, method=cast(SteerMethod, "orthogonal")) + assert spec.method is SteerMethod.ORTHOGONAL + + +def test_worker_specs_carry_steer_method_values() -> None: + """The worker dict's ``op`` is the same spelling as the eager ``method``, member for member.""" + vector = _vector() + spec = SteeringSpec( + layers={ + 0: LayerSteeringSpec( + operations=[ + AddSpec(vector=vector, scale=1.0), + OrthogonalDecompSpec(vector=vector, coeff=0.0), + ProjectionCapSpec(vector=vector, max=1.0), + ] + ) + } + ) + ops = [s["op"] for s in steering_spec_to_worker_specs(spec)] + methods = [s.method for s in steering_spec_to_eager_specs(spec)] + assert ops == ["additive", "orthogonal", "projection_cap"] + assert [SteerMethod(op) for op in ops] == methods + + +def test_the_worker_refuses_the_old_spelling() -> None: + """``add`` was the worker's own name for ``additive``; it is gone rather than aliased.""" + with pytest.raises(ValueError, match="additive"): + _make_steer_modifier( + {"op": "add", "vector": _vector().tolist(), "coeff": 1.0}, torch.device("cpu"), torch.float32 + ) diff --git a/tests/test_vllm_hyper_connections.py b/tests/test_vllm_hyper_connections.py index 6f7bf1f..cc72e9c 100644 --- a/tests/test_vllm_hyper_connections.py +++ b/tests/test_vllm_hyper_connections.py @@ -1200,7 +1200,7 @@ def test_a_stack_steer_can_be_confined_to_one_stream(): trunk = _HyperTrunk(n_layers=2) site = Address("resid_streams", 0) modifier = _make_steer_modifier( - {"op": "add", "vector": [1.0] * D_MODEL, "coeff": 1.0, "stream": 2}, torch.device("cpu"), torch.float32 + {"op": "additive", "vector": [1.0] * D_MODEL, "coeff": 1.0, "stream": 2}, torch.device("cpu"), torch.float32 ) demux, handles = _steering_demux(trunk, {site: modifier}, capture={site}) with torch.no_grad(): @@ -1384,7 +1384,7 @@ def test_a_static_additive_steer_confined_to_one_stream_does_not_take_the_consta from interp_engine.vllm_capture.static import _compile_write_req, _Site site = _Site(Address("resid_streams", 0), delta=torch.zeros(8, HC_MULT, D_MODEL)) - spec = {"op": "add", "vector": [1.0] * D_MODEL, "coeff": 1.0} + spec = {"op": "additive", "vector": [1.0] * D_MODEL, "coeff": 1.0} plain = _compile_write_req(spec, site, skip_positions=(), prompt_len=0, steer_generated=True) assert plain.vector is not None and plain.modify is None, "no stream is still the constant path" @@ -1402,7 +1402,7 @@ def test_a_global_additive_steer_confined_to_one_stream_leaves_the_static_buffer site = _Site(Address("resid_streams", 0), delta=torch.zeros(8, HC_MULT, D_MODEL)) worker = SimpleNamespace(_ie_static=StaticState(writes={"resid_streams.0": site})) - spec = {"op": "add", "point": "resid_streams", "layer": 0, "vector": [1.0] * D_MODEL, "coeff": 1.0} + spec = {"op": "additive", "point": "resid_streams", "layer": 0, "vector": [1.0] * D_MODEL, "coeff": 1.0} worker_set_static_delta(worker, [spec]) assert site.modify is None and site.delta is not None diff --git a/visualizer-web/knowledge/bundle.generated.ts b/visualizer-web/knowledge/bundle.generated.ts index 2977ee7..a779316 100644 --- a/visualizer-web/knowledge/bundle.generated.ts +++ b/visualizer-web/knowledge/bundle.generated.ts @@ -29,8 +29,8 @@ export const KNOWLEDGE_DOCUMENTS: readonly KnowledgeDocument[] = [ {"path":"docs/GRADIENTS.md","title":"Gradients","why":"What is differentiable on which backend, and why gradient support never gates loading.","kind":"markdown","text":"# Gradients\n\nTwo rules, and they are the whole design.\n\n**Gradient support never gates loading.** Someone switching from vLLM to `VLLMModel` loads\nsuccessfully on any kernel, any quantization scheme, any cudagraph mode, and finds out about\ngradients only if they ask. Nothing in `autograd_support.py` is consulted during construction: the\nverdict is computed lazily on first read of `model.grad_support`, from configuration alone — it\nbuilds no engine, loads no kernel, and runs no forward.\n\n**And nothing degrades silently.** Asking for gradients where they are unavailable raises\n`GradientsUnsupported` naming the specific blockers. It never quietly hands back detached tensors,\nand it never quietly flips to a slower kernel to make gradients possible.\n\n### The two kinds of support are not the same question\n\n| | what it means | eager | vLLM |\n| --- | --- | --- | --- |\n| `downstream` | the captured tensor works in an autograd graph *you* build — fit a probe on it, backprop into a decoder | yes | **yes** |\n| `through_forward` | gradients flow back through the model's own forward, to its weights or inputs | with `requires_grad=True` | **never** |\n\nSo a linear probe on captured activations trains fine on either backend. Attribution that needs a\ngradient *of the model* is eager-only.\n\n### Getting gradients (eager)\n\n```python\nmodel = EagerModel(\"openai-community/gpt2\", requires_grad=True) # default is False\ncache = run_with_cache(model, input_ids, [(\"resid_post\", 5)], detach=False)\ncache.get(\"resid_post\", 5).sum().backward()\nmodel.arch.embed.weight.grad # populated: the graph reached the embedding\n```\n\n`requires_grad` defaults to `False` because serving is forward-only and a live graph costs\nactivation memory on every capture. The async `capture()` takes the same `detach=False`, and returns\ndevice tensors rather than CPU ones in that case (moving them is a graph node nobody asked for).\n\n`detach=False` on a **frozen** model raises rather than returning graph-free tensors. It cannot do\nanything useful: the inputs are token ids, so with frozen parameters there is nothing to\ndifferentiate with respect to, and the silent version of this is someone debugging all-zero\ngradients for an afternoon.\n\n### float16, and why one architecture is refused outright\n\nA backward only amplifies a marginal forward, so `float16` is where a differentiable model quietly\nstops being useful. The verdict gives two different answers on purpose:\n\n- **On an architecture in `facts.FP16_EAGER_OVERFLOW_ARCHS` it is a blocker**, and\n `EagerModel(..., requires_grad=True)` refuses at construction. `pythia-70m-deduped` in fp16 already\n NaNs from layer 3 in transformers' GPT-NeoX eager attention kernel, so its gradients would be NaN\n too — and a NaN gradient is worse than a missing one, because it looks like a result. The error\n names `dtype='float32'` as the fix. A *plain* fp16 load of the same model is untouched: gradient\n support never gates a load that did not ask for gradients.\n- **Anywhere else it is a caveat, not a refusal** — logged once at construction and reported in\n `grad_support.caveats`. fp16 gradients are legitimate on a healthy model, and refusing them would be\n the tail wagging the dog.\n\nWorth knowing that `EagerModel.__init__` defaults to `float32` while `from_pretrained` defaults to\n`auto` (the checkpoint's native precision), so the entry point you pick decides this for you. The\nverdict reads the *loaded* dtype, so `auto` is judged by what the checkpoint turned out to be.\n\n### Quantization: capture does not care, the backward does\n\nHooks read activations, which `transformers` dequantizes to a compute dtype at module boundaries, so\nevery point the engine serves is quantization-agnostic down to 4-bit. The backward is a different\nquestion, and it splits by **implementation rather than bit width** — the same two-answer shape as\nfloat16, from the two tables in `facts.py`:\n\n- **A fused forward-only kernel is a blocker** (`facts.FORWARD_ONLY_QUANT_METHODS`: AWQ, GPTQ, FP8 in\n its spellings, native MXFP4). Their dequantize-matmul is registered without an autograd formula, so\n a backward through it raises from inside the op. Refusing up front turns that into a sentence naming\n the scheme, and the remedy says the part people get wrong: capture is unaffected, only the backward\n is.\n- **bitsandbytes is a caveat** (`facts.DIFFERENTIABLE_QUANT_METHODS`, 4-bit and 8-bit). Its matmuls\n route through autograd Functions with a real backward, which is why an offline lens fit works on\n these. Qualified twice: the gradient is with respect to activations only — the quantized weights are\n frozen and receive nothing — and dequantization adds noise, so it is not the gradient the\n unquantized model would give.\n- **An unrecognised scheme is also a caveat**, not a blocker, and this asymmetry is deliberate. A\n wrongly-listed scheme refuses a request that would have worked and the caller cannot override it;\n an unlisted one fails inside the kernel with a caveat already saying why. `torchao` and `quanto` are\n the reason the tables are short: both support quantization-aware training on some paths, so neither\n belongs in either table without being measured.\n\nThe scheme is read from the *loaded* config (`EagerModel.quant_method`), so a pre-quantized checkpoint\nand a quantize-on-load are judged the same way. See\n[PERFORMANCE.md](PERFORMANCE.md#quantization-support) for which install extra each scheme needs.\n\n### Three more things that will bite you, which are documented rather than guarded\n\nThe gradient path itself is architecture-agnostic — one generic hook and one context-manager flip, no\nper-family code — but these are not:\n\n- **MoE changes what the gradient means.** Top-k expert routing is a discrete choice: gradient flows\n through the selected experts and the router's gate weights, never through the selection. Which\n layers are affected is reported in `quirks.moe_layers`.\n- **On sandwich-norm models the capture point silently picks your attribution target**, because\n `mlp_out` and `mlp_out_post` are different graph nodes there (Gemma-2/3/4, OLMo-2/3).\n\n### Generation is not differentiable, and that is deliberate\n\n`steer.py`'s generation loop is unconditionally `no_grad` with no `detach` escape hatch. A tape over\n`max_tokens` sequential forwards retains every step's activations at once, so memory grows with\nthe generation length and a few hundred tokens will OOM a card that generates the same text fine.\nDifferentiating a rollout is a real thing to want, but it needs a purpose-built path (a fixed short\nrollout, or gradient checkpointing) rather than a flag here.\n\nThe lens read-out *is* differentiable, via `decode_residuals(..., detach=False)`:\n\n```python\nresid = torch.randn(n, model.d_model, requires_grad=True)\ndecode_residuals(model, resid, detach=False).sum().backward()\nresid.grad # populated\n```\n\nThat works on a **frozen** model and deliberately does not consult `grad_support`, because the useful\ngradient here is with respect to the residual you passed in — optimizing a steering vector against a\nlogit objective, say — and it never needs to reach a parameter. `layer_logits` has no such flag on\npurpose: it is the serving read-out path, where a graph is pure overhead.\n\n### Why vLLM can never differentiate through its forward\n\nNot a kernel problem, and not fixable by configuration. vLLM decorates `GPUModelRunner.execute_model`\n(and `Worker.execute_model`) with `@torch.inference_mode()`, which is strictly stronger than\n`no_grad`: tensors created inside are *inference tensors*, and autograd refuses them outright —\n`RuntimeError: Inference tensors cannot be saved for backward` — rather than merely not tracking\nthem. So no attention backend, cudagraph mode or quantization scheme buys gradients through a vLLM\nforward. `through_forward` is `False` on every vLLM configuration, and the verdict says so.\n\nThe probe still reports the other blockers it can see (`enforce_eager=False`, `cudagraph_mode=FULL`,\n`quantization=...`, per-layer attention kernels with no exposed backward), because a caller\ndebugging this deserves all the reasons at once, and because those are what the verdict would turn\non if the `inference_mode` decorator were ever lifted upstream. `BACKWARD_CAPABLE_ATTENTION_BACKENDS`\nis deliberately tiny — `TORCH_SDPA` and `FLEX_ATTENTION`, the two built from differentiable torch\nops. Everything else in vLLM's roster is a hand-written kernel, and **anything unrecognised counts\nas unsupported**, so a backend added upstream is a blocker until someone checks it.\n\nGiven all that, do not try to add a \"capture with gradients\" worker RPC. Its precondition — a vLLM\nconfig the probe calls differentiable — is unsatisfiable, so it would be untestable dead code.\nReaching gradients on vLLM would mean stripping `inference_mode` off the hot path *and* replacing\nthe paged attention kernel *and* dealing with a KV cache written in place: at that point you have\nwritten the eager backend, which is already here.\n\n### Why `downstream` is still true on vLLM\n\nBecause captures cross the process boundary as raw bytes. `encode_tensor_payload` /\n`decode_tensor_payload` rebuild the tensor on the client, which launders vLLM's inference tensors\ninto ordinary ones. Verified both ways in `tests/test_autograd_support.py`, and end to end against a\nreal vLLM capture — the raw worker tensor refuses `backward()`, the decoded one does not. This is\nload-bearing for anyone fitting probes against a serving pod, so if you ever change the payload\nencoding to pass tensors by reference, that guarantee goes with it.\n\n### Asking, rather than guessing from the backend name\n\n`model.grad_support` is on the protocol, so a caller gates on fact:\n\n```python\nif model.grad_support.through_forward:\n ...\nelse:\n print(model.grad_support.blockers) # why not\n```\n\n`apps/inference` surfaces `grad_support` in `/capabilities` (via `.describe()`) for the same reason —\nthe webapp should not infer gradient availability from the string `\"vllm\"`, since an eager pod is\nalso non-differentiable unless it was loaded with `requires_grad=True`, which serving does not do.\n\n\n[← back to the interp-engine README](../README.md)"}, {"path":"docs/PERFORMANCE.md","title":"Performance","why":"vLLM-side speed and feature tradeoffs, and which quantized checkpoints load.","kind":"markdown","text":"# Performance and quantization\n\nBoth knobs are vLLM-side: what you trade away for speed, and which quantized\ncheckpoints load. Neither affects the eager backend.\n\n## Performance configuration (vLLM)\n\nThe vLLM backend exists to get vLLM's throughput — continuous batching, paged attention, fused\nkernels — and it keeps all of that. But **one engine flag is defaulted for correctness, not speed,\nand it is paid on every request even by a pod that never captures anything.** If you are comparing\nagainst stock vLLM and the numbers look disappointing, that is why; it is not the interp machinery.\n\n| flag | default | why | cost when left on |\n| --- | --- | --- | --- |\n| `enforce_eager` | `True` | CUDA-graph replay does **not** re-execute the Python forward, so `register_forward_hook` never fires and capture returns nothing. | vLLM bundles CUDA graphs *and* inductor compile under non-eager, so you lose both — **4x to 11x** of decode throughput, measured below. It is the only correctness default left, and `backend=\"vllm-static\"` recovers most of it without giving up capture. |\n\n`enable_prefix_caching` **used** to be the second row, and is now on. The hazard was real — a cache\nhit means the cached positions are never forwarded, so a capture comes back short and a steered\nrequest inherits un-steered KV — but it is a property of individual requests, not of the engine.\n`VLLMModel._prompt` now opts each intervening request out with a unique vLLM `cache_salt`, which\nmakes its block hashes unique and so isolates it in both directions, while plain generation shares\nthe cache normally. Worth ~1.75x on time-to-first-token for a repeated long prefix (48ms → 27ms on a\n2862-token shared prefix, gemma-3-1b) — i.e. the chat/system-prompt case. `tests/test_vllm_kv_isolation.py`\npins which requests opt out; `tests/test_vllm_capture_gpu.py` pins the effect on a real engine,\nincluding a control that reproduces the short capture when the salt is removed. Pass\n`extra_vllm_kwargs={\"enable_prefix_caching\": False}` to get the old engine-wide behaviour back.\n\nOne further flag is set for you, and only on the architectures that cannot boot without it.\n`kv_cache_dtype` is left to vLLM everywhere else, but a few families serve attention through a KV\nlayout that exists in one dtype only — DeepSeek-V4 stores its compressed KV in `fp8_ds_mla`, which\nhas no 16-bit form — and vLLM's own `auto` means \"match the model dtype\", so it resolves to something\nthe layout does not implement and the model class asserts: *DeepseekV4 fp8_ds_mla layout only supports\nfp8 kv-cache, got auto*. The engine therefore derives the dtype from the architecture\n(`facts.mandatory_kv_cache_dtype`) rather than expecting every caller to know. Note it is a numerics\nchoice as well as a boot requirement, since an FP8 cache quantizes what decode reads back; the\nalternative on these architectures is not a bf16 cache but no engine.\n\nAnything else you pass as `extra_vllm_kwargs` is applied **last**, over the engine's own defaults, so\na caller can override any of this — including the correctness defaults above. That is deliberate:\nthere are legitimate reasons to want stock vLLM behaviour, and the alternative (an allowlist) turns\nevery new upstream flag into an engine change.\n\n### The vLLM version is an exact pin\n\n`interp-engine[vllm]` declares `vllm==0.28.0` on Linux. One engine release runs on one vLLM:\n\n- **The engine patches vLLM internals.** `vllm_capture/` walks worker attributes and replaces\n kernels by name (`mhc.py`), so a vLLM the engine has not run on is a risk, not a feature. A drift\n now fails at install, from the resolver, rather than on a pod after the weights load.\n- **A vLLM release fixes the whole CUDA stack.** Each one requires one exact `flashinfer-python`,\n one `torch` and so one CUDA build. With vLLM pinned, those are known constants per engine release.\n- **Why 0.28.0.** 0.27.x cannot run a head_dim=256 model on Blackwell (sm_100+): the FA4 kernel rejects\n the `seqused_q/k` vLLM's decoder attention supplies, fixed by vllm-project/vllm#52050 in 0.28.0 and\n confirmed on a B200 with google/gemma-3-12b-pt. `vllm_capture/mhc.py` needs\n `mhc_pre_broadcast_tilelang`, present from 0.26.0. The performance sweep in\n [`benchmarks/results-latest.md`](../benchmarks/results-latest.md) ran on 0.26.0 and the validator's\n registry snapshot (`tests/vllm_supported_archs.json`) tracks 0.26.0; neither is a bug.\n- **To bump:** change the pin, run a capture on a GPU (Blackwell if available), relock the root and\n `validator/`, and say in the commit what was run. The usual breakage on a new vLLM is a\n `LogitsProcessor` change (see\n [vLLM `compute_logits` is not a bare unembed](ARCHITECTURE_QUIRKS.md#vllm-compute_logits-is-not-a-bare-unembed))\n or a worker-attribute rename that `_walk_trunk` no longer finds; both surface as a loud refusal\n rather than a wrong number.\n\n### FlashInfer needs a compiler on PATH, or the prebuilt kernels\n\nOn Blackwell GPUs (compute capability 10.0 and up) vLLM selects FlashInfer as the attention backend,\nand FlashInfer compiles its kernels on first use. vLLM enables it only when `nvcc` is on PATH or the\n`flashinfer-cubin` package is installed; with neither, the engine fails after the weights load with\n`RuntimeError: FlashInfer backend is not available`, from a stub that names no cause.\n\nThe usual cause is not a missing compiler but a missing PATH entry. The compiler sits at\n`/usr/local/cuda/bin/nvcc`, the image puts that directory on PATH, and the operator's interactive\nshell has it, but the process that launches the server does not: a non-interactive `ssh host cmd`\nnever runs `.bashrc` (where RunPod images source `/etc/rp_environment`), a tmux server keeps the\nPATH it was started with, and a service manager gives its children its own. Everything under that\nlauncher, vLLM's engine core included, inherits the stripped PATH.\n\n`load_model` runs `check_flashinfer` before the engine is built, so this is refused in milliseconds\nrather than after a 27B load. The message prints the PATH the process has and both fixes:\n\n```bash\nexport PATH=/usr/local/cuda/bin:$PATH # in whatever launches the server, not only in your shell\n```\n\nor, for a box with no compiler at all, the prebuilt kernels (both at the exact `flashinfer-python`\nversion vLLM pins, from FlashInfer's index rather than PyPI; the `cu` suffix follows the CUDA build\nof the torch wheel, `cu130` for torch 2.13):\n\n```bash\npip install flashinfer-cubin==0.6.16.post3 flashinfer-jit-cache==0.6.16.post3 \\\n --extra-index-url https://flashinfer.ai/whl --extra-index-url https://flashinfer.ai/whl/cu130\n```\n\nA prebuilt package at another version is refused on every GPU, because FlashInfer itself raises on\nthat pair at import. On pre-Blackwell GPUs FlashInfer is opt-in, so a missing kernel source is a\nwarning. `VLLM_HAS_FLASHINFER_CUBIN=1` skips the check, as it does in vLLM.\n\nWhat is genuinely **free** — do not attribute slowness to these:\n\n- **Attention pattern / DFA recompute** runs only on the attention endpoint, never on generation.\n- **Capture hooks** are installed per request and reference-counted (`_ensure_hook` /\n `_release_hook`), then removed. No capture request means no hooks are attached at all.\n- **Native hidden-state extraction** (`enable_extraction`) is **off** by default. It adds a\n speculative draft plus a KV connector, i.e. extra forwards; worker-hook capture serves every\n point without it.\n\n### Trading interp features for speed\n\n`enforce_eager` is required *while capturing*, not permanently, so a generation-only instance\nshould turn it off — preferably as `backend=\"vllm-generate\"`, which is the same engine configuration\nreached through a name that says what it gives up. There is already a precedent in-tree:\n`apps/nla/vllm_verbalizer.py` runs its own `VLLMModel` with `enforce_eager=False` (CUDA graphs +\ninductor on) precisely because the verbalizer installs no hooks, with\n`NLA_VERBALIZER_ENFORCE_EAGER=1` as a debug escape hatch. It also passes\n`compilation_config.cudagraph_capture_sizes` to match graph capture to its expected fan-out.\n\nWhat it costs, measured (`benchmarks/results-latest.md`, and gpt2 separately for the graph modes).\nThree configurations: the hooked default, `backend=\"vllm-static\"`, and vanilla graphs with no\ncapture at all — which is `backend=\"vllm-generate\"`. Multipliers are against the default in the first column.\n\nSingle-stream decode (tok/s):\n\n| | `backend=\"vllm\"` | `backend=\"vllm-static\"` | `backend=\"vllm-generate\"` |\n| --- | --- | --- | --- |\n| `gemma-2-2b` | 37.5 | 212 (5.7x) | 399 (10.6x) |\n| `qwen3-4b` | 51.0 | 317 (6.2x) | 362 (7.1x) |\n| `llama-3.1-8b` | 64.7 | 262 (4.0x) | 278 (4.3x) |\n\nAt concurrency 8 (tok/s):\n\n| | `backend=\"vllm\"` | `backend=\"vllm-static\"` | `backend=\"vllm-generate\"` |\n| --- | --- | --- | --- |\n| `gemma-2-2b` | 283 | 1,202 (4.2x) | 2,131 (7.5x) |\n| `qwen3-4b` | 391 | 1,831 (4.7x) | 2,035 (5.2x) |\n| `llama-3.1-8b` | 433 | 1,520 (3.5x) | 1,751 (4.0x) |\n\n**The default costs 4x to 11x, not the couple of percent an earlier revision of this file claimed\nfrom a smaller card — worth re-checking any capacity estimate made on that number.** Graph replay,\nnot compilation, is where nearly all of it sits; see the gpt2 modes below.\n\nThe important column is the middle one. Static recovers most of the graph win *while still\ncapturing*: at 4B and 8B it lands within 5-8% of vanilla (317 against 362, 262 against 278), so a\ncapture-serving pod gives up single digits rather than a factor of four. On `gemma-2-2b` it recovers\nabout half. A dedicated generation-only pod is still the fastest thing available, but it is no longer\nthe only way to get graph replay.\n\n### The middle ground: `backend=\"vllm-static\"`\n\nThere *is* a middle ground, which was an open question in this file and is now a shipped backend: a\npreallocated buffer plus a `copy_` that graph capture records, so replay re-executes the tap instead\nof skipping it. That is `backend=\"vllm-static\"`, and it is opt-in:\n\n```python\nmodel = load_model(\"Qwen/Qwen3-8B\", backend=\"vllm-static\") # resid_post at every layer, read and write, under graphs\n```\n\nThe three vLLM backends are three engines, not three settings, and the tap set is what separates\nthem:\n\n- **`backend=\"vllm\"`** — hooked vLLM, the default. Keeps the Python forward, chooses its points per\n request, and is the only configuration that serves every point.\n- **`backend=\"vllm-static\"`** — CUDA-graph replay over the taps named by `static_points`, which\n defaults to `\"auto\"`: `resid_post` at every layer, to read *and* to write, via static `copy_`\n taps. Turns Dynamo off (`VLLM_USE_BREAKABLE_CUDAGRAPH=1`), so the win is replay without compile.\n Because breakable `add_eager` keeps the wrap as ordinary PyTorch on the live tensor, the write path\n is not limited to additive — `orthogonal`, `projection_cap` and the j-lens `steer`/`ablate`/`swap`\n ops all ride the same wrap, per-request and with `position_mask` honoured\n (`register_static_write`). An op outside that set is refused rather than silently skipped. Auto\n covers the write because the two halves are one decision: a read tap alone serves the lens\n read-out and refuses every steer, ablation and swap derived from it, at an address already tapped.\n Pass `static_writes=[]` for the reads without the write sites, and an explicit `static_points`\n list to name both halves yourself — neither is filled in for a caller who said something.\n `static_points=[]` is refused, because an engine with no taps is the next one under another name.\n- **`backend=\"vllm-generate\"`** — no taps, graphs and inductor both on. Generation only: every\n capture, steer and lens entry point refuses rather than returning plausible unsteered text. This\n is the configuration a dedicated completion pod wants, reached through a name that says so rather\n than through `enforce_eager=False`.\n\n`static_points` and `static_writes` are only accepted on `backend=\"vllm-static\"`. Passing either\nanywhere else is an error naming the backend, rather than a flag that quietly turns graphs on.\n\nWhat the static backend does not do, which is why `backend=\"vllm\"` is still the default:\n\n- **Not every point.** `embeddings` and `final_norm` are refused by name\n (`static_unsupported_reason`) because they hang off the trunk rather than a decoder layer, and\n `attn_scores` / `attn_probs` because they are rebuilt off-kernel — static `Address(\"attn\", layer)`\n instead and let `capture_attention` recompute the matrix. The hyper-connection points are NOT in\n that set: `resid_streams` and the stream collapse / write / mix taps static through\n `_install_mhc_static`, and `\"auto\"` resolves to `resid_streams` at every layer on such a trunk\n (`tests/test_static_dsv4_gpu.py`) rather than refusing it. The static set is fixed at engine\n build — a request's `points=` can only filter it, and a miss is refused on the client rather than\n served short.\n- **A narrower batch.** Static *read* buffers are `max_num_batched_tokens` rows tall and have to fit\n alongside the graphs, so the batch is stepped down (16384 → 1024) to make room, and refuses rather\n than OOM-ing in graph capture. Only the reads: a write is the same constant vector on every token,\n so since 1.6.0 its delta is one row and broadcasts. That halves what a read-and-write set costs and\n is why `static_writes=[]` is now a capability switch rather than a way to buy batch width back.\n- **Four times the buffer on a hyper-connection trunk.** `deepseek-v4-flash-0731`'s block carries\n four parallel residual streams, so `\"auto\"` declares `resid_streams` — the whole stack per layer,\n four times the width and so four times the buffer. That is what its static row in the sweep prices.\n- **Graphs for decode only on a hybrid trunk.** `qwen3.8-27b`, and any other linear-attention or\n unclassified trunk, runs prefill eagerly and keeps the decode graphs\n (`static.decode_only_graphs_reason`). Static pins that mode rather than offering it: breakable\n graphs turn `torch.compile` off, and vLLM's mixed prefill-decode graph capture then miscomputes\n prefill on a gated-delta trunk — a whole wrong forward, not a bad tap, and reproducible on plain\n `vllm.LLM`. Eager prefill costs little here, since the wraps are ordinary PyTorch and replay's win\n is in decode.\n- **One engine per process.** `VLLM_USE_BREAKABLE_CUDAGRAPH` is process-global, so a static engine\n and a compiled engine cannot share a process.\n- **Measured for throughput, not yet for capture latency.** The sweep runs `generate` and\n `generate_x8` on the static variant, and the tables above compare only the three conventional\n trunks; its capture, steering and lens cells are unmeasured. Capture *correctness* under static is covered separately —\n `tests/test_static_parity_gpu.py` holds the harvest to cosine ≥ 0.999 and greedy token-id parity\n against hooked capture, across concurrent requests, chunked prefill and decode rows — and the\n backend self-tests each tap at startup, refusing to serve if a `copy_` produced a non-finite or\n all-zero harvest rather than returning a plausible wrong tensor.\n\nThe gpt2 graph modes are what ruled out the *other* candidate middle ground, and they still hold:\n\n| mode | decode | capture |\n| --- | --- | --- |\n| `enforce_eager=True` | 608 tok/s | works |\n| vLLM default graphs | 2,224 tok/s | returns nothing |\n| `cudagraph_mode=NONE` (inductor compile, no replay) | 630 tok/s | returns nothing |\n| `cudagraph_mode=PIECEWISE` | 1,147 tok/s | returns nothing |\n\nPiecewise does **not** let hooks fire on the tapped modules; Dynamo traces the module forwards away\nwhether or not attention is split out — so piecewise is no help, and static deliberately takes the\nbreakable-CUDA-graph path rather than a Dynamo one. And compiling without graph replay buys 1.04x\nwhile still breaking capture, so the entire win is graph replay rather than compilation. That is the\nfinding static is built on: keeping replay and dropping compile costs almost nothing, which is why\n`backend=\"vllm-static\"` turns inductor off and still lands near vanilla above.\n\nWhat a Python hook cannot do is be traced — it appends to a list — which is why the tap had to become\na preallocated buffer plus a `copy_` that graph capture records, roughly what vLLM's own\n`enable_extraction` does. That was the redesign `backend=\"vllm-static\"` required, and it is why static is a\ndifferent mechanism from the hook path rather than the same hooks under graphs.\n\n#### What it actually turns off, and how that is enforced\n\n\"Generation-only\" is not a figure of speech. CUDA graph replay skips the Python `forward` that\n*every* hook is attached to, so the same switch removes capture, steering, DFA/attention and the lens\ninterventions in one go. Steering is the dangerous one: a capture that returns nothing is at least\nnoticeable, but a steering hook that never fires writes nothing and reports nothing, so the request\nreturns fluent, **unsteered** text under a steered label. Nothing downstream can tell.\n\nSo `VLLMModel.hooks_available` is False on such an engine, and `_require_hooks` refuses every\nhook-dependent entry point — `capture`, `capture_generation`(`_stream`), `capture_attention`,\n`set_steering`, `set_lens_intervention`, and `generate_steered` *when* it was given a spec or capture\npoints — before any work is done. Plain generation, `clear_steering` and `capture_resid_post` are\ndeliberately not gated: the first is the point of the mode, the second must stay callable from a\n`finally`, and the third rides vLLM's native extraction rather than hooks.\n\n`_assert_points_captured` remains as a backstop for anything that gets past the front door, but it is\nno longer the first line of defence — a serving pod should read `hooks_available` at startup and\nadvertise the endpoints it can actually serve. `apps/inference` does exactly that behind\n`GENERATION_ONLY=true`: it refuses to start if SAE sets are configured (an SAE read *is* a capture)\nor if the pod resolved to the eager backend (where the flag would cost the endpoints and buy nothing),\nreports `hooks_available: false` plus an empty `capture_points` from `/capabilities`, and returns a\n400 naming the flag for the steered completion types.\n\n## Quantization support\n\nBecause hooks read/write *activations* (which `transformers` dequantizes to a compute dtype at\nmodule boundaries), everything the engine serves is quantization-agnostic down to 4-bit. Serving\ntakes no gradients at all — `EagerModel` calls `requires_grad_(False)`, and the Jacobian lens applies\na *precomputed* matrix rather than differentiating — so a forward-only quantized load\n(GPTQ/AWQ/FP8/native-MXFP4) costs no endpoint. Gradients only matter for *fitting* a lens offline,\nwhich wants fp16/bf16 or bitsandbytes (4/8-bit).\n\nInstall the backend that matches the checkpoint's quantization (all optional + CUDA-only, so the\nbase install stays CPU-friendly):\n\n- **MXFP4** (e.g. gpt-oss): `uv pip install -e '.[quant]'` (the `kernels` Triton loader).\n Included in `apps/inference`. Watch for a silent downgrade here: `transformers` needs `kernels`,\n Triton >= 3.4, and compute capability >= 7.5, and if any is missing it *warns* and dequantizes to\n bf16 instead of failing — gpt-oss-20b goes from 12.9 GiB of weights to ~42 GiB, which is the\n difference between fitting on a 48 GiB card and needing an 80 GiB one. The native path also\n fetches `kernels-community/gpt-oss-triton-kernels` from the Hub on first load, so an offline pod\n needs that pre-cached (a failed fetch raises rather than falling back).\n- **Block-FP8** (e.g. DeepSeek-V4-Flash): the same `'.[quant]'`, and the same Hub fetch — the kernel is\n `kernels-community/finegrained-fp8`, and a missing or out-of-window `kernels` raises from the first\n attention matmul rather than downgrading. Note that transformers moves that window between minor\n releases (5.14 wants `kernels>=0.15.2,<0.16.0`, 5.15 wants `>=0.16.0,<0.17.0`) and the windows do not\n overlap, so the extra's range spans both and the venv's transformers picks the end.\n\n These weights are also the one case where the engine chooses *where* to load rather than loading and\n moving: a quantized checkpoint is placed on the target device by `device_map` at load time, because a\n quantizer with no kernels for the initial device dequantizes to the compute dtype to have something\n runnable — and CPU is such a device for every FP8 scheme. Loading DeepSeek-V4-Flash to CPU first\n materializes ~285 GiB of bf16 on the way to a card that holds its 156 GiB of FP8 with room to spare.\n- **AWQ / GPTQ** (e.g. `casperhansen/*-awq`, quantized Llama-70B): `uv pip install -e '.[awq]'`.\n On transformers v5 these load via the `gptqmodel` backend (it replaced the deprecated `autoawq`,\n which pinned `transformers<=4.47.1`); `accelerate` (included) also enables `device_map` multi-GPU\n sharding needed to fit a 70B AWQ. AWQ CUDA kernels run fp16 (bf16 is auto-downcast).\n- **bitsandbytes** 4/8-bit: `uv pip install bitsandbytes` (differentiable, so the Jacobian lens fit\n works on these).\n\nThat last distinction is not just prose: `grad_support` encodes it, so a forward-only scheme refuses a\ndifferentiable load with a sentence naming the scheme rather than raising from inside a quantized\nmatmul. See [GRADIENTS.md](GRADIENTS.md#quantization-capture-does-not-care-the-backward-does).\n\nNew quantized model? Validate `eager` against an independent implementation before wiring it in — a\nquantized load that returns plausible logits is not evidence that the captured activations are right.\n\n[← back to the interp-engine README](../README.md)"}, {"path":"docs/COMPATIBILITY.md","title":"Compatibility","why":"Which transformers versions are exercised, and which are known to compute a model wrongly.","kind":"markdown","text":"# transformers versions\n\nThis engine has no forward pass of its own. It attaches hooks to `transformers` modules and reports\nwhat they compute, which is the property that makes its captures trustworthy — and it means the\nversion of `transformers` you install is part of your numerical result, the same way the checkpoint\nis. `pyproject.toml` asks for `transformers>=4.57.1` with no ceiling. This page is what that floor\ndoes and does not promise.\n\n## What is actually tested\n\nTwo versions, and they are tested for different things:\n\n| version | what runs | what it proves |\n| --- | --- | --- |\n| 4.57.1, the declared floor | the config-facing tests, no weights, no network (`floor` job in `.github/workflows/engine-tests.yml`) | the engine still *reads* an old transformers correctly |\n| whatever CI resolves today | the full CPU and GPU suites, including the gpt-2 golden parity gate | the engine is right on the version people install |\n\nEverything between the two is inference from those endpoints. Raising the floor is therefore a\none-line edit plus a CI run, and it should be done whenever the old version stops being read\ncorrectly rather than kept out of sentiment.\n\n## What \"supported\" means, and what it does not\n\nSupported means the engine resolves this version's configs into the right facts and hooks the right\nmodules. It does **not** mean two supported versions produce the same numbers. They do not, and the\ndifference is not always visible:\n\n- **The forward pass changes.** transformers 5.15.0 fixed DeepSeek-V2's YaRN `mscale`: before it, the\n attention softmax ran at the wrong temperature, so every attention-derived activation on that\n family was wrong by 1–5% cosine. A capture on 5.14.1 is a faithful capture of a wrong forward pass.\n- **Configs move.** transformers 5.15.0 made Gemma-4's `head_dim` and `num_key_value_heads` per-layer\n attributes. Reading them the old way now raises rather than returning a plausible-looking global\n number, which is the good case; the bad case is a field that silently keeps a stale value.\n- **New architectures arrive.** A checkpoint added in 5.15 does not load on 4.57 at all — a loud,\n obvious failure, and the least interesting of the three.\n\nThe first two are the reason this page exists. Neither shows up as an error, a shape mismatch, or a\nNaN. If you have one engine, you have no second opinion, so the engine has to say something itself.\n\n## Combinations known to be wrong\n\n`facts.TRANSFORMERS_CAVEATS` is the list, and `resolve_facts` warns once per process when the loaded\narchitecture matches one on the installed version:\n\n```\nRuntimeWarning: DeepseekV2ForCausalLM on transformers 5.14.1: attention runs at the wrong\ntemperature -- the YaRN `mscale_all_dim` factor is missing from the softmax scale ... Fixed in\ntransformers 5.15.0 (https://github.com/huggingface/transformers/pull/47435) -- upgrade before\ntrusting these activations.\n```\n\nThe bar for adding a row is deliberately high: a named upstream fix, the release it landed in, and an\neffect on *captured activations*. A deprecation, a speed regression or a bug in a code path the\nengine does not read is not a row here — a warning people learn to ignore is worse than no warning.\nRows narrow themselves further where they can: the DeepSeek entry checks for YaRN scaling in the\nconfig, so the same architecture without it stays quiet.\n\nRemoving a row is fine once the floor rises above its `fixed_in`, and at that point it is dead code\nrather than history worth keeping.\n\n## Practical advice\n\n- Record the transformers version alongside any activations you publish or cache. It belongs with the\n checkpoint revision, not in a footnote.\n- Pin it in an environment whose captures are compared over time. Two runs of the same prompt on the\n same checkpoint can differ across a transformers upgrade, and nothing in the output says so.\n- When cross-checking against another implementation (vLLM, TransformerLens, someone's paper), rule\n the version out first. On DeepSeek-V2 it was the whole answer, and it took a cross-engine sweep to\n see it — see `interp-engine-validator`'s `docs/ENGINE_DIFFERENCES.md`, \"When the reference is the\n one that is wrong\"."}, - {"path":"docs/INTERNALS.md","title":"Modules and correctness","why":"What each file in interp_engine/ owns, and what the test suite checks. Reach for it when the question is about how the engine is built or how it is verified.","kind":"markdown","text":"# Internals\n\nWhat each file in `interp_engine/` owns, and what the engine checks about itself. Read this before\nediting the engine; [AGENTS.md](../AGENTS.md) is the companion piece on the boundaries between those\nfiles and why they are where they are.\n\n- [Modules](#modules)\n- [Correctness](#correctness)\n\n## Modules\n\n- `model.py` — `EagerModel`: wraps `AutoModelForCausalLM` (eager, `no_processing` semantics),\n holds the tokenizer + config-derived dims, canonical hook-point resolution, and an optional\n `quantization_config` passthrough to `from_pretrained`.\n- `facts.py` — the single source of truth for model facts, shared by both backends: structural\n attribute-name vocabularies, config-derived dims, per-layer window/linear-attention predicates,\n and the per-backend tables (fused-QKV layout, parallel-block architectures). Config arithmetic\n and string tables only — no torch, no live model — so the vLLM client can answer dims for a model\n it never builds.\n- `arch.py` — the **eager adapter**: binds a live HF module tree to the structural roles in\n `facts.py`, plus the machine-readable known-quirks table (attention sinks, softcapping, hybrid\n attention, ...).\n- `protocol.py` — `InterpModel`: the surface both backends implement, and the contract a sync free\n function dispatches against. Adding a method here means adding it to both backends and to\n `sync.py`; `tests/test_sync_parity.py` fails on a missing twin.\n- `sync.py` / `_loop.py` — `sync_model(model)`: the protocol without an event loop, one explicit\n wrapper per method over a background loop thread that is created lazily, reused per model, and\n refuses rather than deadlocks when called from inside a running loop.\n- `dispatch.py` — the shared plumbing every free function's two arms sit on: token coercion\n (`TokensLike`, batch refusals) and `CAPABILITIES`, the table each `CapabilityUnsupported` message\n is built from.\n- `points.py` — the canonical point registry: every point's scope, width and vLLM support, and the\n engine's own sentence for each refusal. [SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) is that table\n as prose, and is checked against this module by `tests/test_points_registry.py`.\n- `hooks.py` — the low-level read/write forward-hook substrate.\n- `capture.py` — capture context manager returning a cache keyed by canonical names\n (`resid_post`, `resid_mid`, `mlp_in`, `mlp_act`, `attn_probs`, `value`, `router_logits`,\n `embeddings`, ...), plus the post-processing a captured tensor needs to be usable (fused-QKV\n splits, the attention gate, a norm's scale and gain, per-head residual contributions, dense expert\n assignments).\n- `attn_scores.py` — the pre-softmax attention scores, which no module boundary carries: it registers\n a wrapping attention implementation for the duration of a capture and delegates to the\n checkpoint's own eager function, so the forward is unchanged.\n- `tokenize.py` — `to_tokens`/`to_str_tokens`/`to_string` (TransformerLens-parity), chat\n templating, and per-token span metadata (the single source of truth for message boundaries).\n- `chat_conventions.py` — the only per-model chat table: harmony markers, reasoning delimiters,\n turn-end tokens. Selected by tokenizer capability, never by model name (see\n [Where model-specific config lives](ARCHITECTURE_QUIRKS.md#where-model-specific-config-lives)).\n- `chat_compose.py` — rebuilds assistant messages from a generation (`compose_assistant_turns`),\n reading the generated text only; callers pair it with the prompt messages they already have.\n- `lens.py` — logit + Jacobian lens by calling the real `final_norm` + `lm_head`. Returns **raw**\n logits unless the caller passes a `softcap`. The vLLM path never returns raw logits (see [vLLM\n `compute_logits` is not a bare unembed](ARCHITECTURE_QUIRKS.md#vllm-compute_logits-is-not-a-bare-unembed)).\n- `steer.py` — steering, and the per-token generation stream. Each method's arithmetic is one\n `steer_delta` branch, which is what the vLLM worker's modifier computes too and what\n `tests/test_steer_math_parity.py` runs against it on CPU; `steer()` is the context both backends\n take, registering per-request on vLLM rather than installing a global hook.\n- `mappers.py` — translation between canonical points and other frameworks' names:\n TransformerLens hook strings and nnsight/nnterp accessors, both directions. See [Porting from\n TransformerLens, nnsight or nnterp](PORTING.md#porting-from-transformerlens-nnsight-or-nnterp).\n- `autograd_support.py` — the `GradSupport` verdict: whether a model can give you gradients, and\n which specific thing is blocking it. Pure config arithmetic, so it is safe to call before\n `warmup()`. See [Gradients](GRADIENTS.md#gradients).\n- `cuda_preflight.py` — `check_cuda_driver`: compares the host CUDA driver against the CUDA\n version torch was built for and raises with the forward-compat fix (`cuda-compat--`\n - `LD_LIBRARY_PATH`) before the first CUDA call, instead of failing ten frames deep in\n `torch.cuda._lazy_init`. Lives here because every app on the engine inherits the same CUDA\n floor — the `[vllm]` wheels link `libcudart.so.13` directly. `check_flashinfer`: refuses, before\n the vLLM engine is built, a box where vLLM would disable FlashInfer (no `nvcc` on PATH, no\n `flashinfer-cubin` + `flashinfer-jit-cache`) or where a prebuilt package is not\n `flashinfer-python`'s version; an error on Blackwell, a warning elsewhere. `load_model` calls it\n on the vLLM path. Reads the GPU through NVML so it creates no CUDA context in the parent.\n- `notebook_stdout.py` — `ensure_stdout_descriptor`: gives a notebook kernel's `sys.stdout` the file\n descriptor vLLM's engine start needs. vLLM silences C-level output by dup'ing over one and forks\n its EngineCore child, so under ipykernel — whose stream writes to a socket and has none — the\n child dies before it loads anything, and the caller is told to see a root cause that is in another\n process. Called from `_ensure_engine`; a no-op in a script, a server, or a kernel that captured\n the real descriptors itself.\n\n## Correctness\n\nThe engine's job is to hand back the tensor a module actually produced, so most of what can go wrong\nis quiet: a point resolves to a plausible neighbour, the shapes agree, and the numbers are wrong.\nThe test suite is built around checks that a shape-correct guess cannot pass. The cross-engine half\nof the story — the same points scored against TransformerLens, nnsight/nnterp, vLLM and SGLang on\n50+ real checkpoints — is [`validator/`](../validator/README.md).\n\n**Golden parity.** `tests/test_parity_gpt2.py` pins every capture point on gpt2 against\nTransformerLens, from a committed golden file. It is the one place another framework is loaded, and\nCI treats a skip as a failure (`IE_REQUIRE_PARITY=1`) so a missing dependency or a cold cache cannot\nquietly retire the gate.\n\n**Invariants over attribute names.** Three identities hold on any model of a family, so they catch a\nmisresolved point without needing a reference implementation: `probs @ value == z` for the per-head\nvalue and DFA (`tests/test_qkv_layout.py`, which also asserts the _wrong_ layout fails — otherwise\nthe test would pass on a single-head model), `resid_pre + attn_out_post + mlp_out_post == resid_post`\nfor sandwich norms and residual multipliers, and `down_proj(mlp_act) == mlp_out` for the neuron\nbasis. Where a point genuinely does not exist — a Mamba block's attention, a latent-attention model's\n`value`, the residual between the sublayers of a parallel block — it is refused with an explanation\nrather than returned as a plausible tensor.\n\n**Self-consistency on real weights.** `tests/test_new_models_gpu.py` decodes the last layer's\nresidual through the real `final_norm` + `lm_head` and requires the model's true next-token argmax\nback, which validates the whole arch map end-to-end without a second framework.\n`tests/test_sliding_window_attn.py` pins the vLLM attention recompute's band and sink terms, and\n`tests/test_vllm_only_families.py` checks the spellings for families `transformers` has no class for\nagainst a synthetic tree in the shape their own modeling file describes.\n\n**The two backends against each other.** `tests/test_vllm_capture_gpu.py` runs a real vLLM engine and\nrequires each captured point to match the eager backend's — including that vLLM's positional layer\nindex names the layer HF's does, which nothing checked before and which would fail silently rather\nthan raise. It needs `interp-engine[vllm]`, so it self-skips elsewhere; note that running it via\n`.venv-vllm/bin/python` needs that directory on `PATH` too, because vLLM shells out to `ninja` to\nbuild a sampler kernel at startup. `tests/test_vllm_wire_grammar.py` covers the same process\nboundary on CPU, over a synthetic demux. `tests/test_multigpu.py` (`-m multigpu`, two CUDA cards)\nrepeats the comparison at `num_gpus=2`: eager under accelerate's layer placement and vLLM under\ntensor parallelism, where `vllm_capture/_tp.py` gathers the head- and neuron-sharded points across\nranks, against the one-card eager reference — captures, attention, steering, decode and the lens.\nOne rule follows from tensor parallelism: a worker RPC refuses by *returning* a reason, never by\nraising. Under vLLM's multiprocess executor a raise inside `collective_rpc` is read from one rank\nwhile the others' replies stay queued, and the next RPC consumes those — so a single refused layer\ntook every point of a TP=8 cell with it. `resolvable_points` and `resolvable_attn` are the ask-first\ncalls; `tests/test_vllm_new_points.py` pins that they install nothing.\n\n**The docs are parsed, not maintained.** The point table in\n[SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) and the footnote markers in\n[ENGINE_HOOK_MAPPINGS.md](ENGINE_HOOK_MAPPINGS.md) are read back and compared against `points.py` by\n`tests/test_points_registry.py`, and every `interp_engine` name in a doc code fence has to still\nimport and still be in `__all__` (`tests/test_doc_code_fences.py`). A point that quietly became\nservable while the tick stayed put is a claim the engine does not make.\n\n**How CI is split.** Two jobs, both on every non-Markdown change: a CPU job\n(`-m \"not gpu and not xl\"`) that owns the golden gate, the lint/format/type gates and the small\nmodels eagerly, and a managed-L4 GPU job (`-m \"gpu and not xl\"`) running the same models on\nCUDA/bf16. The `xl` models are tens of GB and run nowhere automatically — `pytest -m xl` on a big\nbox. Locally the GPU tests self-skip without CUDA and model loads skip when weights aren't cached, so\na plain `pytest tests` on a laptop runs the fast suite.\n\n[← back to the interp-engine README](../README.md)"}, - {"path":"interp_engine/__init__.py","title":"Public API surface","why":"`__all__`, grouped by task. A name absent from here is not public API and must not appear in a snippet.","kind":"python","text":"\"\"\"InterpEngine: standalone raw-transformers interpretability core.\n\nRuns the plain HuggingFace model in eager PyTorch and captures/steers via our own\nforward-hook layer. Eager (unfused) execution is what lets us read per-head attention\npatterns, tap arbitrary hook points, and match reference numerics; serving is forward-only.\nKeyed by the raw HuggingFace model id.\n\"\"\"\n\nfrom importlib.metadata import PackageNotFoundError, version\n\nfrom interp_engine.address import (\n Address,\n AddressError,\n UnknownCoordinate,\n format_address,\n parse_address,\n to_address,\n)\nfrom interp_engine.arch import (\n ArchSpec,\n Quirks,\n resolve_arch,\n special_token_ids,\n special_token_positions,\n)\nfrom interp_engine.attn_config import unsupported_attn_config\nfrom interp_engine.autograd_support import (\n BACKWARD_CAPABLE_ATTENTION_BACKENDS,\n GradientsUnsupported,\n GradSupport,\n eager_grad_support,\n vllm_grad_support,\n)\nfrom interp_engine.capture import (\n Cache,\n attn_out_gate,\n capture_attention,\n capture_generation,\n expert_assignment,\n head_contributions,\n is_rms_norm,\n per_head_value,\n pre_gain_normalized,\n rms_norm_parts,\n run_with_cache,\n split_fused_qkv,\n)\nfrom interp_engine.chat_compose import (\n ChatTurn,\n compose_assistant_turns,\n strip_wire_reasoning,\n)\nfrom interp_engine.chat_conventions import (\n REASONING_TAGS,\n TURN_END_TOKENS,\n ReasoningTags,\n added_vocab,\n detect_reasoning_tags,\n is_harmony,\n)\nfrom interp_engine.chat_formatters import (\n CODE_CHAT_FORMATS,\n ChatFormatter,\n ChatFormatterUnavailable,\n RenderedChat,\n resolve_chat_formatter,\n)\nfrom interp_engine.cuda_preflight import check_cuda_driver, check_flashinfer\nfrom interp_engine.dispatch import CAPABILITIES, Capability, CapabilityUnsupported, TokensLike\nfrom interp_engine.facts import rms_norm_eps_for_model\nfrom interp_engine.hooks import HookManager\nfrom interp_engine.lens import (\n apply_final_logit_softcap,\n apply_logit_transform,\n capture_residuals,\n decode_residuals,\n layer_logits,\n)\nfrom interp_engine.load import BACKENDS, VLLM_BACKENDS, load_model, vllm_installed\nfrom interp_engine.mappers import (\n UnmappedHook,\n nnsight_accessor_to_point,\n point_to_nnsight_accessor,\n point_to_tlens_hook,\n tlens_hook_to_point,\n tlens_normalized_hook,\n)\nfrom interp_engine.model import EagerModel, HubKernelUnsupported, deepgemm_fallback_kwargs\nfrom interp_engine.protocol import Completion, InterpModel, Point\nfrom interp_engine.residual_basis import (\n RESIDUAL_POINTS,\n STREAM_REDUCTIONS,\n ResidualBasis,\n ResidualBasisUnsupported,\n eager_residual_basis,\n reduce_streams,\n vllm_residual_basis,\n)\nfrom interp_engine.select import BackendSelection, select_backend\nfrom interp_engine.steer import (\n STEER_METHODS,\n GenStep,\n OrthogonalProjector,\n PositionMask,\n SteerMask,\n SteerSpec,\n generate_stream,\n projection_cap_delta,\n resolve_masked_positions,\n steer,\n steer_delta,\n steering_spec_to_eager_specs,\n top_logprobs,\n unit_vector,\n)\nfrom interp_engine.steer_specs import (\n AddSpec,\n LayerSteeringSpec,\n OrthogonalDecompSpec,\n ProjectionCapSpec,\n SteeringOp,\n SteeringSpec,\n steering_spec_to_worker_specs,\n)\nfrom interp_engine.sync import SyncModel, sync_model\nfrom interp_engine.tokenize import (\n GeneratedTurnSpans,\n NoChatTemplateError,\n Tokenize,\n TokenSpan,\n)\nfrom interp_engine.vllm_backend import (\n VLLMModel,\n attn_capture_layers,\n is_linear_attention_layer,\n read_attn_dims,\n recompute_attn_from_payloads,\n sliding_window_for_layer,\n)\n\n# Only the pieces a caller driving their own vLLM engine needs: the point names, and the\n# payload codecs for reading collective_rpc results back. The ``worker_*`` functions those\n# RPCs run are NOT re-exported -- the supported way to reach them is by name through\n# ``interp_engine.vllm_plugin``, and code that really wants the callables (to hand one to\n# ``collective_rpc`` directly, which needs VLLM_ALLOW_INSECURE_SERIALIZATION) can import\n# them from ``interp_engine.vllm_capture``. Same for the recompute internals.\nfrom interp_engine.vllm_capture import (\n HOOK_CAPTURE_POINTS,\n decode_capture_payload,\n decode_tensor_payload,\n encode_tensor_payload,\n)\nfrom interp_engine.vllm_plugin import (\n WORKER_EXTENSION_CLS,\n InterpWorkerExtension,\n capture_engine_kwargs,\n native_extraction_engine_kwargs,\n)\n\n# Read from the installed distribution rather than hardcoded, so it cannot drift from\n# pyproject.toml. Falls back for a source tree that was never installed.\ntry:\n __version__ = version(\"interp-engine\")\nexcept PackageNotFoundError: # pragma: no cover -- source checkout, not installed\n __version__ = \"0.0.0.dev0\"\n\n# Grouped by what you reach for, since alphabetical order buries the entry points among the\n# helpers. Anything not listed here is reachable from its own module but is not API: it can\n# change without a major version.\n__all__ = [\n # The installed version, as `vllm` and `transformers` expose it.\n \"__version__\",\n # Loading a model -- start here.\n \"BACKENDS\",\n \"VLLM_BACKENDS\",\n \"BackendSelection\",\n \"load_model\",\n \"select_backend\",\n \"vllm_installed\",\n # The models, and the surface they share.\n \"Completion\",\n \"EagerModel\",\n \"HubKernelUnsupported\",\n \"deepgemm_fallback_kwargs\",\n \"InterpModel\",\n \"VLLMModel\",\n # That same surface without an event loop, for scripts and notebooks. The sync free\n # functions below dispatch through this, so most callers never name it.\n \"SyncModel\",\n \"sync_model\",\n # What one signature over two backends is allowed to do when only one of them can: refuse,\n # from a table, naming the capability. Never a warning and never a silent no-op.\n \"CAPABILITIES\",\n \"Capability\",\n \"CapabilityUnsupported\",\n \"TokensLike\",\n # Addressing a tensor: the type, its string form, and the refusals.\n \"Address\",\n \"AddressError\",\n \"Point\",\n \"UnknownCoordinate\",\n \"format_address\",\n \"parse_address\",\n \"to_address\",\n # Capture: one signature over either backend, dispatched on the model you pass.\n \"Cache\",\n \"UnmappedHook\",\n \"attn_out_gate\",\n \"capture_attention\",\n \"capture_generation\",\n \"expert_assignment\",\n \"head_contributions\",\n \"is_rms_norm\",\n \"nnsight_accessor_to_point\",\n \"point_to_nnsight_accessor\",\n \"point_to_tlens_hook\",\n \"tlens_hook_to_point\",\n \"tlens_normalized_hook\",\n \"per_head_value\",\n \"pre_gain_normalized\",\n \"rms_norm_eps_for_model\",\n \"rms_norm_parts\",\n \"run_with_cache\",\n \"split_fused_qkv\",\n # Gradients: whether they are available, and the refusal when they are not.\n \"BACKWARD_CAPABLE_ATTENTION_BACKENDS\",\n \"GradSupport\",\n \"GradientsUnsupported\",\n \"eager_grad_support\",\n \"vllm_grad_support\",\n # The residual stream's structure: how many, additive, sequenced -- and the refusals when not.\n \"RESIDUAL_POINTS\",\n \"STREAM_REDUCTIONS\",\n \"ResidualBasis\",\n \"ResidualBasisUnsupported\",\n \"eager_residual_basis\",\n \"reduce_streams\",\n \"vllm_residual_basis\",\n # Steering. Every method's arithmetic is one delta function, shared with the vLLM worker.\n \"STEER_METHODS\",\n \"AddSpec\",\n \"LayerSteeringSpec\",\n \"OrthogonalDecompSpec\",\n \"OrthogonalProjector\",\n \"PositionMask\",\n \"ProjectionCapSpec\",\n \"SteerMask\",\n \"SteerSpec\",\n \"SteeringOp\",\n \"SteeringSpec\",\n \"projection_cap_delta\",\n \"resolve_masked_positions\",\n \"steer\",\n \"steer_delta\",\n \"steering_spec_to_eager_specs\",\n \"steering_spec_to_worker_specs\",\n \"unit_vector\",\n # Generation (sync, eager).\n \"GenStep\",\n \"generate_stream\",\n \"top_logprobs\",\n # Lens read-out.\n \"apply_final_logit_softcap\",\n \"apply_logit_transform\",\n \"capture_residuals\",\n \"decode_residuals\",\n \"layer_logits\",\n # Architecture introspection.\n \"ArchSpec\",\n \"HookManager\",\n \"Quirks\",\n \"attn_capture_layers\",\n \"read_attn_dims\",\n \"recompute_attn_from_payloads\",\n \"resolve_arch\",\n \"unsupported_attn_config\",\n # Tokenization and chat templates.\n \"CODE_CHAT_FORMATS\",\n \"ChatFormatter\",\n \"ChatFormatterUnavailable\",\n \"ChatTurn\",\n \"GeneratedTurnSpans\",\n \"NoChatTemplateError\",\n \"REASONING_TAGS\",\n \"ReasoningTags\",\n \"RenderedChat\",\n \"TURN_END_TOKENS\",\n \"TokenSpan\",\n \"Tokenize\",\n \"added_vocab\",\n \"compose_assistant_turns\",\n \"detect_reasoning_tags\",\n \"is_harmony\",\n \"resolve_chat_formatter\",\n \"special_token_ids\",\n \"special_token_positions\",\n \"strip_wire_reasoning\",\n # Driving your own vLLM engine (see interp_engine.vllm_plugin).\n \"HOOK_CAPTURE_POINTS\",\n \"InterpWorkerExtension\",\n \"WORKER_EXTENSION_CLS\",\n \"capture_engine_kwargs\",\n \"decode_capture_payload\",\n \"decode_tensor_payload\",\n \"encode_tensor_payload\",\n \"is_linear_attention_layer\",\n \"native_extraction_engine_kwargs\",\n \"sliding_window_for_layer\",\n # Environment checks.\n \"check_cuda_driver\",\n \"check_flashinfer\",\n]"}, + {"path":"docs/INTERNALS.md","title":"Modules and correctness","why":"What each file in interp_engine/ owns, and what the test suite checks. Reach for it when the question is about how the engine is built or how it is verified.","kind":"markdown","text":"# Internals\n\nWhat each file in `interp_engine/` owns, and what the engine checks about itself. Read this before\nediting the engine; [AGENTS.md](../AGENTS.md) is the companion piece on the boundaries between those\nfiles and why they are where they are.\n\n- [Modules](#modules)\n- [Correctness](#correctness)\n\n## Modules\n\n- `model.py` — `EagerModel`: wraps `AutoModelForCausalLM` (eager, `no_processing` semantics),\n holds the tokenizer + config-derived dims, canonical hook-point resolution, and an optional\n `quantization_config` passthrough to `from_pretrained`.\n- `facts.py` — the single source of truth for model facts, shared by both backends: structural\n attribute-name vocabularies, config-derived dims, per-layer window/linear-attention predicates,\n and the per-backend tables (fused-QKV layout, parallel-block architectures). Config arithmetic\n and string tables only — no torch, no live model — so the vLLM client can answer dims for a model\n it never builds.\n- `arch.py` — the **eager adapter**: binds a live HF module tree to the structural roles in\n `facts.py`, plus the machine-readable known-quirks table (attention sinks, softcapping, hybrid\n attention, ...).\n- `protocol.py` — `InterpModel`: the surface both backends implement, and the contract a sync free\n function dispatches against. Adding a method here means adding it to both backends and to\n `sync.py`; `tests/test_sync_parity.py` fails on a missing twin.\n- `sync.py` / `_loop.py` — `sync_model(model)`: the protocol without an event loop, one explicit\n wrapper per method over a background loop thread that is created lazily, reused per model, and\n refuses rather than deadlocks when called from inside a running loop.\n- `dispatch.py` — the shared plumbing every free function's two arms sit on: token coercion\n (`TokensLike`, batch refusals) and `CAPABILITIES`, the table each `CapabilityUnsupported` message\n is built from.\n- `points.py` — the canonical point registry: every point's scope, width and vLLM support, and the\n engine's own sentence for each refusal. [SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) is that table\n as prose, and is checked against this module by `tests/test_points_registry.py`.\n- `hooks.py` — the low-level read/write forward-hook substrate.\n- `capture.py` — capture context manager returning a cache keyed by canonical names\n (`resid_post`, `resid_mid`, `mlp_in`, `mlp_act`, `attn_probs`, `value`, `router_logits`,\n `embeddings`, ...), plus the post-processing a captured tensor needs to be usable (fused-QKV\n splits, the attention gate, a norm's scale and gain, per-head residual contributions, dense expert\n assignments).\n- `attn_scores.py` — the pre-softmax attention scores, which no module boundary carries: it registers\n a wrapping attention implementation for the duration of a capture and delegates to the\n checkpoint's own eager function, so the forward is unchanged.\n- `tokenize.py` — `to_tokens`/`to_str_tokens`/`to_string` (TransformerLens-parity), chat\n templating, and per-token span metadata (the single source of truth for message boundaries).\n- `chat_conventions.py` — the only per-model chat table: harmony markers, reasoning delimiters,\n turn-end tokens. Selected by tokenizer capability, never by model name (see\n [Where model-specific config lives](ARCHITECTURE_QUIRKS.md#where-model-specific-config-lives)).\n- `chat_compose.py` — rebuilds assistant messages from a generation (`compose_assistant_turns`),\n reading the generated text only; callers pair it with the prompt messages they already have.\n- `lens.py` — logit + Jacobian lens by calling the real `final_norm` + `lm_head`. Returns **raw**\n logits unless the caller passes a `softcap`. The vLLM path never returns raw logits (see [vLLM\n `compute_logits` is not a bare unembed](ARCHITECTURE_QUIRKS.md#vllm-compute_logits-is-not-a-bare-unembed)).\n- `steer.py` — steering, and the per-token generation stream. Each method's arithmetic is one\n `steer_delta` branch, which is what the vLLM worker's modifier computes too and what\n `tests/test_steer_math_parity.py` runs against it on CPU; `steer()` is the context both backends\n take, registering per-request on vLLM rather than installing a global hook. The methods are named\n once, in `steer_specs.SteerMethod`: the eager `SteerSpec.method` and the worker dict's `op` both\n read from it, so a method has one spelling on either backend.\n- `mappers.py` — translation between canonical points and other frameworks' names:\n TransformerLens hook strings and nnsight/nnterp accessors, both directions. See [Porting from\n TransformerLens, nnsight or nnterp](PORTING.md#porting-from-transformerlens-nnsight-or-nnterp).\n- `autograd_support.py` — the `GradSupport` verdict: whether a model can give you gradients, and\n which specific thing is blocking it. Pure config arithmetic, so it is safe to call before\n `warmup()`. See [Gradients](GRADIENTS.md#gradients).\n- `cuda_preflight.py` — `check_cuda_driver`: compares the host CUDA driver against the CUDA\n version torch was built for and raises with the forward-compat fix (`cuda-compat--`\n - `LD_LIBRARY_PATH`) before the first CUDA call, instead of failing ten frames deep in\n `torch.cuda._lazy_init`. Lives here because every app on the engine inherits the same CUDA\n floor — the `[vllm]` wheels link `libcudart.so.13` directly. `check_flashinfer`: refuses, before\n the vLLM engine is built, a box where vLLM would disable FlashInfer (no `nvcc` on PATH, no\n `flashinfer-cubin` + `flashinfer-jit-cache`) or where a prebuilt package is not\n `flashinfer-python`'s version; an error on Blackwell, a warning elsewhere. `load_model` calls it\n on the vLLM path. Reads the GPU through NVML so it creates no CUDA context in the parent.\n- `notebook_stdout.py` — `ensure_stdout_descriptor`: gives a notebook kernel's `sys.stdout` the file\n descriptor vLLM's engine start needs. vLLM silences C-level output by dup'ing over one and forks\n its EngineCore child, so under ipykernel — whose stream writes to a socket and has none — the\n child dies before it loads anything, and the caller is told to see a root cause that is in another\n process. Called from `_ensure_engine`; a no-op in a script, a server, or a kernel that captured\n the real descriptors itself.\n\n## Correctness\n\nThe engine's job is to hand back the tensor a module actually produced, so most of what can go wrong\nis quiet: a point resolves to a plausible neighbour, the shapes agree, and the numbers are wrong.\nThe test suite is built around checks that a shape-correct guess cannot pass. The cross-engine half\nof the story — the same points scored against TransformerLens, nnsight/nnterp, vLLM and SGLang on\n50+ real checkpoints — is [`validator/`](../validator/README.md).\n\n**Golden parity.** `tests/test_parity_gpt2.py` pins every capture point on gpt2 against\nTransformerLens, from a committed golden file. It is the one place another framework is loaded, and\nCI treats a skip as a failure (`IE_REQUIRE_PARITY=1`) so a missing dependency or a cold cache cannot\nquietly retire the gate.\n\n**Invariants over attribute names.** Three identities hold on any model of a family, so they catch a\nmisresolved point without needing a reference implementation: `probs @ value == z` for the per-head\nvalue and DFA (`tests/test_qkv_layout.py`, which also asserts the _wrong_ layout fails — otherwise\nthe test would pass on a single-head model), `resid_pre + attn_out_post + mlp_out_post == resid_post`\nfor sandwich norms and residual multipliers, and `down_proj(mlp_act) == mlp_out` for the neuron\nbasis. Where a point genuinely does not exist — a Mamba block's attention, a latent-attention model's\n`value`, the residual between the sublayers of a parallel block — it is refused with an explanation\nrather than returned as a plausible tensor.\n\n**Self-consistency on real weights.** `tests/test_new_models_gpu.py` decodes the last layer's\nresidual through the real `final_norm` + `lm_head` and requires the model's true next-token argmax\nback, which validates the whole arch map end-to-end without a second framework.\n`tests/test_sliding_window_attn.py` pins the vLLM attention recompute's band and sink terms, and\n`tests/test_vllm_only_families.py` checks the spellings for families `transformers` has no class for\nagainst a synthetic tree in the shape their own modeling file describes.\n\n**The two backends against each other.** `tests/test_vllm_capture_gpu.py` runs a real vLLM engine and\nrequires each captured point to match the eager backend's — including that vLLM's positional layer\nindex names the layer HF's does, which nothing checked before and which would fail silently rather\nthan raise. It needs `interp-engine[vllm]`, so it self-skips elsewhere; note that running it via\n`.venv-vllm/bin/python` needs that directory on `PATH` too, because vLLM shells out to `ninja` to\nbuild a sampler kernel at startup. `tests/test_vllm_wire_grammar.py` covers the same process\nboundary on CPU, over a synthetic demux. `tests/test_multigpu.py` (`-m multigpu`, two CUDA cards)\nrepeats the comparison at `num_gpus=2`: eager under accelerate's layer placement and vLLM under\ntensor parallelism, where `vllm_capture/_tp.py` gathers the head- and neuron-sharded points across\nranks, against the one-card eager reference — captures, attention, steering, decode and the lens.\nOne rule follows from tensor parallelism: a worker RPC refuses by *returning* a reason, never by\nraising. Under vLLM's multiprocess executor a raise inside `collective_rpc` is read from one rank\nwhile the others' replies stay queued, and the next RPC consumes those — so a single refused layer\ntook every point of a TP=8 cell with it. `resolvable_points` and `resolvable_attn` are the ask-first\ncalls; `tests/test_vllm_new_points.py` pins that they install nothing.\n\n**The docs are parsed, not maintained.** The point table in\n[SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) and the footnote markers in\n[ENGINE_HOOK_MAPPINGS.md](ENGINE_HOOK_MAPPINGS.md) are read back and compared against `points.py` by\n`tests/test_points_registry.py`, and every `interp_engine` name in a doc code fence has to still\nimport and still be in `__all__` (`tests/test_doc_code_fences.py`). A point that quietly became\nservable while the tick stayed put is a claim the engine does not make.\n\n**How CI is split.** Two jobs, both on every non-Markdown change: a CPU job\n(`-m \"not gpu and not xl\"`) that owns the golden gate, the lint/format/type gates and the small\nmodels eagerly, and a managed-L4 GPU job (`-m \"gpu and not xl\"`) running the same models on\nCUDA/bf16. The `xl` models are tens of GB and run nowhere automatically — `pytest -m xl` on a big\nbox. Locally the GPU tests self-skip without CUDA and model loads skip when weights aren't cached, so\na plain `pytest tests` on a laptop runs the fast suite.\n\n[← back to the interp-engine README](../README.md)"}, + {"path":"interp_engine/__init__.py","title":"Public API surface","why":"`__all__`, grouped by task. A name absent from here is not public API and must not appear in a snippet.","kind":"python","text":"\"\"\"InterpEngine: standalone raw-transformers interpretability core.\n\nRuns the plain HuggingFace model in eager PyTorch and captures/steers via our own\nforward-hook layer. Eager (unfused) execution is what lets us read per-head attention\npatterns, tap arbitrary hook points, and match reference numerics; serving is forward-only.\nKeyed by the raw HuggingFace model id.\n\"\"\"\n\nfrom importlib.metadata import PackageNotFoundError, version\n\nfrom interp_engine.address import (\n Address,\n AddressError,\n UnknownCoordinate,\n format_address,\n parse_address,\n to_address,\n)\nfrom interp_engine.arch import (\n ArchSpec,\n Quirks,\n resolve_arch,\n special_token_ids,\n special_token_positions,\n)\nfrom interp_engine.attn_config import unsupported_attn_config\nfrom interp_engine.autograd_support import (\n BACKWARD_CAPABLE_ATTENTION_BACKENDS,\n GradientsUnsupported,\n GradSupport,\n eager_grad_support,\n vllm_grad_support,\n)\nfrom interp_engine.capture import (\n Cache,\n attn_out_gate,\n capture_attention,\n capture_generation,\n expert_assignment,\n head_contributions,\n is_rms_norm,\n per_head_value,\n pre_gain_normalized,\n rms_norm_parts,\n run_with_cache,\n split_fused_qkv,\n)\nfrom interp_engine.chat_compose import (\n ChatTurn,\n compose_assistant_turns,\n strip_wire_reasoning,\n)\nfrom interp_engine.chat_conventions import (\n REASONING_TAGS,\n TURN_END_TOKENS,\n ReasoningTags,\n added_vocab,\n detect_reasoning_tags,\n is_harmony,\n)\nfrom interp_engine.chat_formatters import (\n CODE_CHAT_FORMATS,\n ChatFormatter,\n ChatFormatterUnavailable,\n RenderedChat,\n resolve_chat_formatter,\n)\nfrom interp_engine.cuda_preflight import check_cuda_driver, check_flashinfer\nfrom interp_engine.dispatch import CAPABILITIES, Capability, CapabilityUnsupported, TokensLike\nfrom interp_engine.facts import rms_norm_eps_for_model\nfrom interp_engine.hooks import HookManager\nfrom interp_engine.lens import (\n apply_final_logit_softcap,\n apply_logit_transform,\n capture_residuals,\n decode_residuals,\n layer_logits,\n)\nfrom interp_engine.load import BACKENDS, VLLM_BACKENDS, load_model, vllm_installed\nfrom interp_engine.mappers import (\n UnmappedHook,\n nnsight_accessor_to_point,\n point_to_nnsight_accessor,\n point_to_tlens_hook,\n tlens_hook_to_point,\n tlens_normalized_hook,\n)\nfrom interp_engine.model import EagerModel, HubKernelUnsupported, deepgemm_fallback_kwargs\nfrom interp_engine.protocol import Completion, InterpModel, Point\nfrom interp_engine.residual_basis import (\n RESIDUAL_POINTS,\n STREAM_REDUCTIONS,\n ResidualBasis,\n ResidualBasisUnsupported,\n eager_residual_basis,\n reduce_streams,\n vllm_residual_basis,\n)\nfrom interp_engine.select import BackendSelection, select_backend\nfrom interp_engine.steer import (\n STEER_METHODS,\n GenStep,\n OrthogonalProjector,\n PositionMask,\n SteerMask,\n SteerSpec,\n generate_stream,\n projection_cap_delta,\n resolve_masked_positions,\n steer,\n steer_delta,\n steering_spec_to_eager_specs,\n top_logprobs,\n unit_vector,\n)\nfrom interp_engine.steer_specs import (\n AddSpec,\n LayerSteeringSpec,\n OrthogonalDecompSpec,\n ProjectionCapSpec,\n SteeringOp,\n SteeringSpec,\n SteerMethod,\n steer_method,\n steering_spec_to_worker_specs,\n)\nfrom interp_engine.sync import SyncModel, sync_model\nfrom interp_engine.tokenize import (\n GeneratedTurnSpans,\n NoChatTemplateError,\n Tokenize,\n TokenSpan,\n)\nfrom interp_engine.vllm_backend import (\n VLLMModel,\n attn_capture_layers,\n is_linear_attention_layer,\n read_attn_dims,\n recompute_attn_from_payloads,\n sliding_window_for_layer,\n)\n\n# Only the pieces a caller driving their own vLLM engine needs: the point names, and the\n# payload codecs for reading collective_rpc results back. The ``worker_*`` functions those\n# RPCs run are NOT re-exported -- the supported way to reach them is by name through\n# ``interp_engine.vllm_plugin``, and code that really wants the callables (to hand one to\n# ``collective_rpc`` directly, which needs VLLM_ALLOW_INSECURE_SERIALIZATION) can import\n# them from ``interp_engine.vllm_capture``. Same for the recompute internals.\nfrom interp_engine.vllm_capture import (\n HOOK_CAPTURE_POINTS,\n decode_capture_payload,\n decode_tensor_payload,\n encode_tensor_payload,\n)\nfrom interp_engine.vllm_plugin import (\n WORKER_EXTENSION_CLS,\n InterpWorkerExtension,\n capture_engine_kwargs,\n native_extraction_engine_kwargs,\n)\n\n# Read from the installed distribution rather than hardcoded, so it cannot drift from\n# pyproject.toml. Falls back for a source tree that was never installed.\ntry:\n __version__ = version(\"interp-engine\")\nexcept PackageNotFoundError: # pragma: no cover -- source checkout, not installed\n __version__ = \"0.0.0.dev0\"\n\n# Grouped by what you reach for, since alphabetical order buries the entry points among the\n# helpers. Anything not listed here is reachable from its own module but is not API: it can\n# change without a major version.\n__all__ = [\n # The installed version, as `vllm` and `transformers` expose it.\n \"__version__\",\n # Loading a model -- start here.\n \"BACKENDS\",\n \"VLLM_BACKENDS\",\n \"BackendSelection\",\n \"load_model\",\n \"select_backend\",\n \"vllm_installed\",\n # The models, and the surface they share.\n \"Completion\",\n \"EagerModel\",\n \"HubKernelUnsupported\",\n \"deepgemm_fallback_kwargs\",\n \"InterpModel\",\n \"VLLMModel\",\n # That same surface without an event loop, for scripts and notebooks. The sync free\n # functions below dispatch through this, so most callers never name it.\n \"SyncModel\",\n \"sync_model\",\n # What one signature over two backends is allowed to do when only one of them can: refuse,\n # from a table, naming the capability. Never a warning and never a silent no-op.\n \"CAPABILITIES\",\n \"Capability\",\n \"CapabilityUnsupported\",\n \"TokensLike\",\n # Addressing a tensor: the type, its string form, and the refusals.\n \"Address\",\n \"AddressError\",\n \"Point\",\n \"UnknownCoordinate\",\n \"format_address\",\n \"parse_address\",\n \"to_address\",\n # Capture: one signature over either backend, dispatched on the model you pass.\n \"Cache\",\n \"UnmappedHook\",\n \"attn_out_gate\",\n \"capture_attention\",\n \"capture_generation\",\n \"expert_assignment\",\n \"head_contributions\",\n \"is_rms_norm\",\n \"nnsight_accessor_to_point\",\n \"point_to_nnsight_accessor\",\n \"point_to_tlens_hook\",\n \"tlens_hook_to_point\",\n \"tlens_normalized_hook\",\n \"per_head_value\",\n \"pre_gain_normalized\",\n \"rms_norm_eps_for_model\",\n \"rms_norm_parts\",\n \"run_with_cache\",\n \"split_fused_qkv\",\n # Gradients: whether they are available, and the refusal when they are not.\n \"BACKWARD_CAPABLE_ATTENTION_BACKENDS\",\n \"GradSupport\",\n \"GradientsUnsupported\",\n \"eager_grad_support\",\n \"vllm_grad_support\",\n # The residual stream's structure: how many, additive, sequenced -- and the refusals when not.\n \"RESIDUAL_POINTS\",\n \"STREAM_REDUCTIONS\",\n \"ResidualBasis\",\n \"ResidualBasisUnsupported\",\n \"eager_residual_basis\",\n \"reduce_streams\",\n \"vllm_residual_basis\",\n # Steering. Every method's arithmetic is one delta function, shared with the vLLM worker.\n \"STEER_METHODS\",\n \"AddSpec\",\n \"LayerSteeringSpec\",\n \"OrthogonalDecompSpec\",\n \"OrthogonalProjector\",\n \"PositionMask\",\n \"ProjectionCapSpec\",\n \"SteerMask\",\n \"SteerMethod\",\n \"SteerSpec\",\n \"SteeringOp\",\n \"SteeringSpec\",\n \"projection_cap_delta\",\n \"resolve_masked_positions\",\n \"steer\",\n \"steer_delta\",\n \"steer_method\",\n \"steering_spec_to_eager_specs\",\n \"steering_spec_to_worker_specs\",\n \"unit_vector\",\n # Generation (sync, eager).\n \"GenStep\",\n \"generate_stream\",\n \"top_logprobs\",\n # Lens read-out.\n \"apply_final_logit_softcap\",\n \"apply_logit_transform\",\n \"capture_residuals\",\n \"decode_residuals\",\n \"layer_logits\",\n # Architecture introspection.\n \"ArchSpec\",\n \"HookManager\",\n \"Quirks\",\n \"attn_capture_layers\",\n \"read_attn_dims\",\n \"recompute_attn_from_payloads\",\n \"resolve_arch\",\n \"unsupported_attn_config\",\n # Tokenization and chat templates.\n \"CODE_CHAT_FORMATS\",\n \"ChatFormatter\",\n \"ChatFormatterUnavailable\",\n \"ChatTurn\",\n \"GeneratedTurnSpans\",\n \"NoChatTemplateError\",\n \"REASONING_TAGS\",\n \"ReasoningTags\",\n \"RenderedChat\",\n \"TURN_END_TOKENS\",\n \"TokenSpan\",\n \"Tokenize\",\n \"added_vocab\",\n \"compose_assistant_turns\",\n \"detect_reasoning_tags\",\n \"is_harmony\",\n \"resolve_chat_formatter\",\n \"special_token_ids\",\n \"special_token_positions\",\n \"strip_wire_reasoning\",\n # Driving your own vLLM engine (see interp_engine.vllm_plugin).\n \"HOOK_CAPTURE_POINTS\",\n \"InterpWorkerExtension\",\n \"WORKER_EXTENSION_CLS\",\n \"capture_engine_kwargs\",\n \"decode_capture_payload\",\n \"decode_tensor_payload\",\n \"encode_tensor_payload\",\n \"is_linear_attention_layer\",\n \"native_extraction_engine_kwargs\",\n \"sliding_window_for_layer\",\n # Environment checks.\n \"check_cuda_driver\",\n \"check_flashinfer\",\n]"}, {"path":"interp_engine/points.py","title":"Canonical point registry","why":"The point table with each point's scope, width and vLLM support, plus the engine's own refusal reasons.","kind":"python","text":"\"\"\"The canonical hook points, as data: one row per point, and the properties consumers derive.\n\nEverything that needs to know something about a point reads it from this table: the resolver, both\nmappers, the vLLM client and worker, the reshape in :mod:`interp_engine.capture`, the protocol\ndocstring, the comparison spec, and the doc tables with their footnote markers. One table rather than\na description per consumer, because a point missing from any single consumer fails *silently* -- one\nabsent from ``HOOK_CAPTURE_POINTS`` is merely refused under vLLM, one absent from the width guard is\nmerely unchecked, one absent from the reshape set merely comes back with its batch and sequence axes\nflattened together. A row here is what makes a point either fully added or plainly absent.\n\n**What belongs here: properties of a point that hold across architectures.** Where the tensor comes\nfrom does *not* -- that stays in :meth:`interp_engine.model.EagerModel.resolve_point`, because a\nresolution is a module plus a side plus, for several points, a refusal that has to explain itself\n(a parallel block has no ``resid_mid``; a sparse layer has no neuron basis; a fused ``gate_up_proj``\nhas no separable branches). Flattening those into a table would either lose the explanations or\ngrow a callable column per row, which is an if-chain with extra steps. So the table declares *what\neach point is*, the resolver decides *where it is*, and :func:`known_names` is what lets a test\nassert the two agree.\n\n**What deliberately does not belong here: other frameworks' names.** ``mappers.py`` imports this\nmodule's point set and is tested for coverage against it; the dependency runs that way and not the\nother, so the engine never carries a vocabulary it does not own. See ``AGENTS.md``.\n\nThe set of points stays **open** -- an unrecognized name with no layer falls through to a dotted\nmodule path, which is how a consumer taps something the core does not enumerate. This table is the\nclosed set of names with *declared semantics*, not a whitelist of what may be captured.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass\nfrom enum import Enum\n\n__all__ = [\n \"HYPER_CONNECTION_POINTS\",\n \"POINTS\",\n \"PointSpec\",\n \"Scope\",\n \"VllmSupport\",\n \"Width\",\n \"conditional_points\",\n \"d_model_wide\",\n \"eager_only\",\n \"hyper_connection_names\",\n \"known_names\",\n \"point_spec\",\n \"points_for\",\n \"reason\",\n \"refusal_reasons\",\n \"token_flattened\",\n \"tp_sharded\",\n \"vllm_hookable\",\n]\n\n\nclass Scope(Enum):\n \"\"\"Whether a point is addressed with a layer index.\"\"\"\n\n GLOBAL = \"global\"\n LAYER = \"layer\"\n\n\nclass Width(Enum):\n \"\"\"What the last axis of the captured tensor counts -- the axis a shard would narrow.\n\n This is the tensor-parallel question, which is why it is an enum rather than a\n ``d_model: bool``. ``HEADS``, ``NEURONS`` and ``VOCAB`` are exactly the axes vLLM shards, so\n under TP rank 0 holds a slice of them -- and for the first two there is no width to check it\n against either (``n_heads * head_dim`` coincides with ``hidden_size`` on Llama and not on\n Gemma-3, so a width check would pass on one and fail on the other for reasons unrelated to\n sharding). :func:`tp_sharded` is that set.\n\n The rest reach a hook at full width, for three different reasons, which is why \"is it\n ``D_MODEL``\" is *not* the same question: ``D_MODEL`` is all-reduced before the hook sees it,\n ``ROUTING`` comes off a replicated gate (vLLM routes with a ``ReplicatedLinear``, so every rank\n computes the same logits), and ``STREAMS`` is replicated because each block mixes all the\n streams together and a per-rank slice of that mixture is not a partial sum of anything.\n Only ``D_MODEL`` can be *checked* against ``hidden_size``, though, which is what\n :func:`d_model_wide` is for -- a narrower one there is a shard.\n \"\"\"\n\n D_MODEL = \"d_model\"\n HEADS = \"n_heads * head_dim\"\n NEURONS = \"d_mlp\"\n ROUTING = \"experts\"\n SCORES = \"n_heads * query * key\"\n VOCAB = \"vocab_size\"\n # A hyper-connection trunk's residual streams (DeepSeek-V4's `hc_mult`). Answering this enum's\n # actual question -- which axis a shard narrows -- the answer is *none*: the streams are\n # replicated on every rank, because each block mixes all of them together with a learned\n # doubly-stochastic matrix and a per-rank slice of that mixture is not a partial sum of\n # anything. So a width check here is meaningful under TP in a way `HEADS` and `NEURONS` are not.\n STREAMS = \"n_residual_streams\"\n\n @classmethod\n def sharded(cls) -> frozenset[Width]:\n \"\"\"The widths a tensor-parallel rank holds only a slice of. See the class docstring.\"\"\"\n return frozenset({cls.HEADS, cls.NEURONS, cls.VOCAB})\n\n\nclass VllmSupport(Enum):\n \"\"\"How -- or whether -- the vLLM backend can serve a point.\n\n Three states rather than a bool because the reasons differ in kind, and a caller reading a\n refusal deserves to know which: ``NONE`` splits into \"unimplemented\" and \"unreachable\", and the\n per-point ``note`` says which one applies.\n\n The line between the first two is *where the tensor is assembled*, not how hard it was to get.\n ``HOOKS`` means the worker returns the point itself, whether that comes off a module boundary\n (almost every row) or off a wrapped kernel call with a little arithmetic behind it (the mHC rows,\n whose tensors are locals of a decoder layer's forward -- see\n :mod:`interp_engine.vllm_capture.mhc`). ``RECOMPUTE`` means the worker returns something *else*\n and the client rebuilds the point from it, which is a different contract: the client needs the\n architecture's own softmax rules to do it, and the helper tensors travel under their own keys.\n \"\"\"\n\n HOOKS = \"worker-side capture\"\n RECOMPUTE = \"rebuilt on the client from captured q/k/v\"\n NONE = \"eager only\"\n\n\n@dataclass(frozen=True)\nclass PointSpec:\n name: str\n scope: Scope\n width: Width\n vllm: VllmSupport\n #: Set where the module reports the whole batch as one token axis, so the capture arrives\n #: ``[batch * pos, ...]`` and has to be reshaped before it indexes like every other point.\n token_flattened: bool = False\n #: Set where no module boundary carries the tensor, so ``resolve_point`` does not answer for it\n #: and the capture path special-cases it instead.\n module_resolved: bool = True\n #: Why this point is limited the way it is, in the caller's terms. Not decoration: for the\n #: eager-only points this is the difference between \"nobody wrote it yet\" and \"there is no such\n #: tensor in a fused engine\", which is what a reader needs to know whether to file a bug.\n note: str = \"\"\n\n\ndef _p(name: str, scope: Scope, width: Width, vllm: VllmSupport, **kw: object) -> PointSpec:\n return PointSpec(name=name, scope=scope, width=width, vllm=vllm, **kw) # type: ignore[arg-type]\n\n\n_G, _L = Scope.GLOBAL, Scope.LAYER\n_HOOKS, _RECOMPUTE, _NONE = VllmSupport.HOOKS, VllmSupport.RECOMPUTE, VllmSupport.NONE\n\n#: Every canonical point, in forward order -- which is also the order the docs table reads in.\nPOINTS: tuple[PointSpec, ...] = (\n _p(\"embeddings\", _G, Width.D_MODEL, _HOOKS),\n _p(\"resid_pre\", _L, Width.D_MODEL, _HOOKS),\n _p(\"attn_in\", _L, Width.D_MODEL, _HOOKS),\n _p(\"q_norm_in\", _L, Width.HEADS, _HOOKS),\n _p(\"q_norm_out\", _L, Width.HEADS, _HOOKS),\n _p(\"k_norm_in\", _L, Width.HEADS, _HOOKS),\n _p(\"k_norm_out\", _L, Width.HEADS, _HOOKS),\n _p(\"value\", _L, Width.HEADS, _HOOKS),\n _p(\n \"attn_scores\",\n _L,\n Width.SCORES,\n _RECOMPUTE,\n module_resolved=False,\n note=\"no module boundary holds the pre-softmax scores on either backend, and the paged \"\n \"kernel never forms the matrix at all; rebuilt alongside attn_probs from the same captured \"\n \"post-RoPE q/k, which is the tensor the softmax is taken over\",\n ),\n _p(\n \"attn_probs\",\n _L,\n Width.SCORES,\n _RECOMPUTE,\n module_resolved=False,\n note=\"fused paged attention never materializes the probabilities; rebuilt from captured \"\n \"post-RoPE q/k with the checkpoint's own window, softcap and sinks reapplied\",\n ),\n _p(\"z\", _L, Width.HEADS, _HOOKS),\n _p(\n \"attn_gate\",\n _L,\n Width.HEADS,\n _NONE,\n note=\"unimplemented: q_proj is a real module on both trees; double width, so rank-sliced under TP\",\n ),\n _p(\"attn_out\", _L, Width.D_MODEL, _HOOKS),\n _p(\"attn_out_post\", _L, Width.D_MODEL, _HOOKS),\n _p(\"resid_mid\", _L, Width.D_MODEL, _HOOKS),\n _p(\"mlp_in\", _L, Width.D_MODEL, _HOOKS),\n _p(\n \"mlp_pre\",\n _L,\n Width.NEURONS,\n _NONE,\n note=\"unreachable as a module boundary on vLLM: it fuses gate_proj and up_proj into one \"\n \"gate_up_proj, so neither branch is a module output -- and being d_mlp wide it is also \"\n \"interleaved per rank under TP, so rank 0's slice cannot be reassembled on its own\",\n ),\n _p(\"mlp_pre_linear\", _L, Width.NEURONS, _NONE, note=\"see mlp_pre; gated MLPs only\"),\n _p(\"mlp_act\", _L, Width.NEURONS, _HOOKS),\n _p(\"router_logits\", _L, Width.ROUTING, _HOOKS, token_flattened=True),\n _p(\n \"expert_weights\",\n _L,\n Width.ROUTING,\n _NONE,\n token_flattened=True,\n note=\"unreachable: the top-k happens inside the FusedMoE kernel, which takes the logits and \"\n \"returns the combined output with the selection never materialized\",\n ),\n _p(\n \"expert_indices\",\n _L,\n Width.ROUTING,\n _NONE,\n token_flattened=True,\n note=\"see expert_weights; integer-valued, so it is the one point that is not a differentiable activation\",\n ),\n _p(\"mlp_out\", _L, Width.D_MODEL, _HOOKS),\n _p(\"mlp_out_post\", _L, Width.D_MODEL, _HOOKS),\n _p(\"resid_post\", _L, Width.D_MODEL, _HOOKS),\n _p(\"final_norm\", _G, Width.D_MODEL, _HOOKS),\n _p(\n \"lm_head\",\n _G,\n Width.VOCAB,\n _NONE,\n note=\"unreachable as a bare unembed: vLLM's compute_logits can fold scaling and softcapping, \"\n \"so what it returns is not this point\",\n ),\n)\n\n_BY_NAME: dict[str, PointSpec] = {p.name: p for p in POINTS}\nif len(_BY_NAME) != len(POINTS): # pragma: no cover - a duplicate row would shadow silently\n raise RuntimeError(\"duplicate canonical point name in POINTS\")\n\n\n# --- the rows a hyper-connection trunk adds -----------------------------------------------------\n#\n# Points that exist only where the trunk carries several residual streams, in the same dataclass as\n# the global rows so every consumer keeps reading one table.\n#\n# The bar for a conditional row is high, and manifold-constrained hyper-connections (mHC) are the\n# case that clears it: the quantities below are not variations on a shared point (that is what the\n# ``stream`` coordinate is for) but tensors with no counterpart on a conventional trunk at all -- a\n# per-token collapse of several streams into one, and the doubly-stochastic matrix that remixes them\n# afterwards. Giving them global names would put seven rows in front of every reader of every other\n# family, and giving them none would leave the tensor that a steering vector or an SAE actually wants\n# reachable only by dotted module path.\n#\n# Gated on the trunk a model was found to have and deliberately NOT on its architecture name, which\n# for this trunk carries no information: `MotifForCausalLM` is Motif 2.6B, Motif 2-12.7B and Motif 3\n# alike, and only the last has hyper-connections. A name-keyed table would offer these rows on the\n# two that do not and refuse them at resolution time, one layer too late to be honest, while a\n# checkpoint that switches the trunk off (`mhc_enabled=false`) would be mis-answered either way.\n# :func:`interp_engine.facts.residual_streams` is where the count comes from; the module layouts each\n# family puts these tensors in are :data:`interp_engine.facts.HYPER_CONNECTION_LAYOUTS`.\n#\n# A conditional name may not shadow a global one (checked below): one name, one meaning, everywhere.\n\n#: What the two collapse rows share: the sublayer's argument is the collapse *after* the block's norm,\n#: so it is `attn_in`/`mlp_in` and not this point, and the point is rebuilt rather than read. Quoted by\n#: both, because a reader who takes the argument for the collapse gets a correctly-shaped tensor one\n#: norm away -- and `facts.HyperConnectionLayout` already warns that this is the one place where being\n#: approximately right is indistinguishable from being right.\n_COLLAPSE_FUSED_NORM = (\n \"the obvious address for it is WRONG on vllm's NVIDIA tree. The argument passed to the sublayer \"\n \"is not this tensor: attn_norm/ffn_norm are fused INTO the mHC pre kernel (nvidia/model.py passes \"\n \"norm_weight to mhc_pre_broadcast / mhc_fused_post_pre and its own comment says the norm is fused \"\n \"there), so that argument is the collapse already normed -- the engine's attn_in/mlp_in. Measured \"\n \"on DeepSeek-V4-Flash at layers 0/21/42: it matches RMSNorm(eager collapse) * norm.weight to 5e-3 \"\n \"and differs from eager's collapse itself by up to 24x, and the norm dropped a per-token scalar, \"\n \"so it cannot be inverted back. So this point is rebuilt from the stream stack the kernel was \"\n \"handed and the layer's flat hc_{site}_{fn,base,scale} -- the collapse half of vllm's own \"\n \"mhc_pre_torch, re-derived to 5e-3 against transformers' module. It is the one mHC row whose \"\n \"value is arithmetic rather than a tensor the engine computed. On the amd/ and xpu/ trees none of \"\n \"that is needed: they instantiate MHCPreOp/MHCFusedPostPreOp (CustomOp, so nn.Module) and apply \"\n \"attn_norm as a SEPARATE call afterwards, so there the unnormed collapse is a module output. \"\n \"Steerable, and for the same reason by an unobvious route: since the recompute agrees with the \"\n \"kernel only to 5e-3, the edit is applied as the DIFFERENCE it makes to the fused norm, \"\n \"norm(c + delta) - norm(c), which is exactly zero when delta is -- substituting norm(c + delta) \"\n \"outright would impose that 5e-3 on an unsteered request's sublayer\"\n)\n\n#: What the four write/mix rows share, and the one property of the mixing matrix that the loose\n#: version of gets wrong on the real checkpoint.\n_SINKHORN_AXIS = (\n \"the mixing matrix is COLUMN stochastic to 1e-6 and only roughly row stochastic (up to 7e-2 at \"\n \"hidden_size 4096). Sinkhorn ends on a column normalization, and the columns are also the axis \"\n \"the post phase contracts, so that is both the exact one and the meaningful one; the row error \"\n \"shrinks with more iterations and grows with width, so a fixture narrow enough to run cheaply \"\n \"will not show it\"\n)\n\n#: The mHC rows -- every one eager-served on both families that have the trunk, and every one now\n#: served under vLLM too. The notes below are about vLLM's DeepSeek-V4, which is the only mHC\n#: implementation upstream vLLM has at all -- Motif 3 is served through its authors' fork, so on this\n#: backend the question does not arise for it yet.\n#:\n#: The notes state what was MEASURED, on deepseek-ai/DeepSeek-V4-Flash under vLLM 0.26.0 on one B200\n#: with enforce_eager, at layers 0/21/42 of 43 (`plans/scripts/verify_dsv4_mhc_vllm.py`), and against\n#: transformers' own `DeepseekV4HyperConnection` on the same checkpoint's weights\n#: (`plans/scripts/compare_dsv4_mhc_eager.py`).\n#:\n#: One fact of vllm's implementation shapes all seven rows, and it is worth stating once here:\n#: **each sublayer's post phase is deferred into the next sublayer's pre phase kernel.** So within a\n#: layer's forward the attention coefficients are computed and then overwritten, and the stream stack\n#: that crosses the layer boundary is the one the MLP read rather than the one the block produced.\n#: Only the MLP's pair survives to a module boundary; the other five are locals, and are served by\n#: wrapping the kernel calls themselves (`vllm_capture.mhc`, and `vllm_capture._tree` for which\n#: mechanism serves which row). That is a real distinction and not an implementation detail to\n#: paper over: a wrapped kernel is NVIDIA-tree-specific, where a module hook would not be.\nHYPER_CONNECTION_POINTS: tuple[PointSpec, ...] = (\n _p(\n \"resid_streams\",\n _L,\n Width.STREAMS,\n _HOOKS,\n module_resolved=False,\n note=\"the block's whole (num_tokens, hc_mult, hidden_size) output stack, and the row it took a \"\n \"measurement to address, because the obvious answer looks right and is not. The layer returns \"\n \"(x, residual, post_mix, res_mix) and `residual` has exactly this shape -- (13, 4, 4096) at \"\n \"layers 0/21/42, stream axis second-from-last as select_stream expects -- and is a different \"\n \"tensor: because the post phase is deferred it is the stack the MLP's pre phase READ, so it is \"\n \"resid_mid in stream form. Measured rather than argued: collapsing it at the ffn site \"\n \"reproduces the argument to self.ffn to 3e-3 and at the attn site misses the argument to \"\n \"self.attn by 0.31-0.74, which places it strictly between the sublayers. This point is served \"\n \"one layer downstream instead, off the NEXT layer's first fused kernel, whose post half is \"\n \"where the MLP's scatter actually happens -- and for the last layer off the model's own \"\n \"standalone mhc_post call after the loop. Refused at the last layer when a draft/EAGLE model \"\n \"is attached, which reconstructs aux hidden states with that same function and makes the call \"\n \"ambiguous. Steerable, by re-running that fused call's second half (mhc_pre) on the edited \"\n \"stack so the collapse the same call computes sees the edit -- writing the call's OUTPUT \"\n \"would reach every later layer and miss its own first reader. The re-run's numbers go back \"\n \"only to the rows that were steered, since mhc_pre alone differs from the fused kernel by up \"\n \"to 2e-2 in bf16 and a request that steers nothing must not move\",\n ),\n _p(\n \"attn_stream_collapse\",\n _L,\n Width.D_MODEL,\n _HOOKS,\n module_resolved=False,\n note=\"the one d_model tensor on such a trunk that a steering vector or an SAE wants, since it \"\n \"is what attention actually reads. Served, and \" + _COLLAPSE_FUSED_NORM,\n ),\n _p(\n \"attn_stream_write\",\n _L,\n Width.STREAMS,\n _HOOKS,\n module_resolved=False,\n note=\"the attention half of the write/mix pair, and the half the deferral hides: the layer's \"\n \"first fused call computes it and the second overwrites it, so unlike the mlp pair it reaches \"\n \"no return tuple and is read off that first call instead. Confirmed by measurement rather than \"\n \"by reading -- driving transformers' own mHC module with the stack the layer returns \"\n \"reproduces the RETURNED coefficients at the ffn site to 4e-4 and disagrees at the attention \"\n \"site, which is how the two halves were told apart. \" + _SINKHORN_AXIS,\n ),\n _p(\"attn_stream_mix\", _L, Width.STREAMS, _HOOKS, module_resolved=False, note=\"see attn_stream_write\"),\n _p(\n \"mlp_stream_collapse\",\n _L,\n Width.D_MODEL,\n _HOOKS,\n module_resolved=False,\n note=\"the MLP's counterpart to attn_stream_collapse: what the FFN actually reads, one norm \"\n \"before mlp_in. Served, and \" + _COLLAPSE_FUSED_NORM,\n ),\n # The two rows the deferral leaves *intact*: the pair that survives to the layer's return is the\n # MLP's, so for these the obvious address is the right one, and both are verified against the\n # reference implementation on identical weights. Served on vLLM by index into the layer's return\n # (`vllm_capture._tree.LAYER_RETURN_INDEX`), which is also the only place the index is written.\n _p(\n \"mlp_stream_write\",\n _L,\n Width.STREAMS,\n _HOOKS,\n note=\"the layer's output:2, verified: matches transformers' DeepseekV4HyperConnection `post` \"\n \"to 6e-4 on the same checkpoint's weights at layers 0/21/42. vllm returns it as \"\n \"(num_tokens, hc_mult, 1) where eager returns (batch, seq, hc_mult); the capture squeezes \"\n \"that trailing axis, so the width is the stream count on both backends rather than 1 on one \"\n \"of them\",\n ),\n _p(\n \"mlp_stream_mix\",\n _L,\n Width.STREAMS,\n _HOOKS,\n note=\"the layer's output:3, (num_tokens, hc_mult, hc_mult), verified: matches transformers' \"\n \"`comb` to 4e-4 on the same weights. One property to state exactly, because the loose version \"\n \"of it fails on the real checkpoint: \" + _SINKHORN_AXIS,\n ),\n)\n\n_BY_CONDITIONAL_NAME: dict[str, PointSpec] = {p.name: p for p in HYPER_CONNECTION_POINTS}\n_shadowed = sorted(set(_BY_CONDITIONAL_NAME) & set(_BY_NAME))\nif _shadowed: # pragma: no cover - a shadowing row would make a name mean two things\n raise RuntimeError(f\"conditional point name(s) shadow a global point: {_shadowed}\")\n\n\ndef known_names() -> frozenset[str]:\n \"\"\"Every *global* name with declared semantics. Not a whitelist -- the point set is open.\n\n Global only, which is a narrower set than it looks and should not be reached for as \"every\n point\". ``mappers.py`` used to be tested for coverage against this alone, on the reasoning that\n an mHC point had no foreign name to map to; TransformerLens 3 then grew a DeepSeek-V4 bridge\n that names all seven, and the guard could not see that they had gone unmapped. It now covers\n this set together with :func:`hyper_connection_names`. Use :func:`points_for` when the question\n is \"what can I ask this model for\".\n \"\"\"\n return frozenset(_BY_NAME)\n\n\ndef hyper_connection_names() -> frozenset[str]:\n \"\"\"The conditional rows' names -- points that exist only on a hyper-connection trunk.\n\n For a consumer that has a name in hand and needs to know whether the *model* must have streams\n for it to mean anything, which is a different question from whether the backend can serve it.\n \"\"\"\n return frozenset(_BY_CONDITIONAL_NAME)\n\n\ndef mhc_coefficient_names() -> frozenset[str]:\n \"\"\"The hyper-connection rows that are coefficients rather than activations.\n\n The per-stream write weights and the Sinkhorn-normalized mixing matrix, at both sites. Split out\n because the distinction decides what may be *written*: the other three rows are tensors on the\n residual trunk and an additive edit to one is a steer, while an additive edit to a doubly\n stochastic matrix leaves it neither stochastic nor a mixture of anything. Both backends refuse a\n steer at these, and the vocabulary belongs here rather than in either of them.\n\n By ``Width.STREAMS`` on a conditional row minus the stack itself, so a row added later is\n classified by what it is rather than by whether someone remembered to add it to a list.\n \"\"\"\n return frozenset(\n spec.name for spec in HYPER_CONNECTION_POINTS if spec.width is Width.STREAMS and spec.name != \"resid_streams\"\n )\n\n\ndef stream_stack_points() -> frozenset[str]:\n \"\"\"The points whose capture carries the residual stream axis, so a ``d_model`` read must reduce it.\n\n The distinction this answers is a claim about the *point*, and it does not follow from the trunk:\n on a hyper-connection model ``attn_out`` and ``mlp_stream_collapse`` are ``d_model``-wide -- a\n sublayer's own output, and the collapsed vector a sublayer reads -- while ``resid_streams`` is the\n whole stack. So \"this model has four streams\" does not tell a lens read-out whether the tensor in\n front of it has an axis to average, and guessing from the shape cannot either, since both end in\n ``d_model``.\n\n Complement of :func:`mhc_coefficient_names` within ``Width.STREAMS``, and derived rather than\n listed for the same reason: the coefficient rows do carry a stream axis, but they are not\n ``d_model``-wide, so reducing one to a residual is not a thing to ask for.\n \"\"\"\n coefficients = mhc_coefficient_names()\n return frozenset(\n spec.name for spec in HYPER_CONNECTION_POINTS if spec.width is Width.STREAMS and spec.name not in coefficients\n )\n\n\ndef steer_refusal_reason(name: str) -> str | None:\n \"\"\"Why nothing can steer ``name`` on any backend, or None when the point itself does not forbid it.\n\n A claim about the *point* rather than about a backend or a model, which is why it lives here and\n is quoted by both the client gate and the worker's registration check rather than written twice:\n the two refuse in different places and a caller must not get different reasons depending on which\n one they reached first.\n\n Returning None is not a promise that a steer will be accepted. Whether a point can be written\n depends on the family's block (``resid_mid``), on the kernels (``resid_streams``) and on the wire\n (the trunk-level points), and each of those is refused where it becomes knowable.\n \"\"\"\n if name not in mhc_coefficient_names():\n return None\n return (\n f\"{name} is a hyper-connection coefficient -- the per-stream write weights, or the \"\n \"Sinkhorn-normalized mixing matrix -- and not an activation on the residual stream. An \"\n \"additive edit leaves it neither stochastic nor a mixture of anything, so there is no \"\n \"intervention here that means what a steer means. Capture is supported, and so is steering \"\n \"the hyper-connection points that ARE activations: resid_streams, attn_stream_collapse, \"\n \"mlp_stream_collapse.\"\n )\n\n\ndef conditional_points(n_streams: int) -> tuple[PointSpec, ...]:\n \"\"\"The rows a trunk carrying ``n_streams`` residual streams adds. Empty on a conventional one.\"\"\"\n return HYPER_CONNECTION_POINTS if n_streams > 1 else ()\n\n\ndef points_for(n_streams: int = 1) -> tuple[PointSpec, ...]:\n \"\"\"Every point addressable on a model whose trunk carries ``n_streams`` residual streams.\n\n Takes the stream count rather than an architecture name because the name does not answer the\n question -- see the comment above :data:`HYPER_CONNECTION_POINTS`. Callers holding a model should\n prefer ``model.points()``, which reads the count off the model rather than being told it.\n \"\"\"\n return POINTS + conditional_points(n_streams)\n\n\ndef point_spec(name: str, n_streams: int = 1) -> PointSpec | None:\n \"\"\"The row for ``name``, or ``None`` for a name the core does not enumerate (a dotted path).\n\n At the default single stream this answers about the global table only. Passing a real count also\n considers the rows a hyper-connection trunk adds -- and only then, so a caller cannot resolve a\n stream point against a Llama and be told it exists.\n \"\"\"\n found = _BY_NAME.get(name)\n if found is not None:\n return found\n return next((p for p in conditional_points(n_streams) if p.name == name), None)\n\n\n#: Every row, global and conditional. The derived sets below are membership guards -- the reshape\n#: in `capture`, the vLLM width check -- and a conditional point missing from one of those fails\n#: exactly as silently as a global one would, which is what this module exists to prevent.\n_ALL_POINTS: tuple[PointSpec, ...] = POINTS + HYPER_CONNECTION_POINTS\n\n\ndef _names(**match: object) -> frozenset[str]:\n return frozenset(p.name for p in _ALL_POINTS if all(getattr(p, k) == v for k, v in match.items()))\n\n\ndef vllm_hookable() -> frozenset[str]:\n \"\"\"Points the vLLM worker's forward hooks can serve, as module I/O.\"\"\"\n return _names(vllm=VllmSupport.HOOKS)\n\n\ndef eager_only() -> frozenset[str]:\n \"\"\"Points with no vLLM path at all -- the ``*`` footnote in the mapping table.\"\"\"\n return _names(vllm=VllmSupport.NONE)\n\n\ndef d_model_wide() -> frozenset[str]:\n \"\"\"Points whose last axis is ``hidden_size`` on every architecture, so a narrower one is a shard.\"\"\"\n return _names(width=Width.D_MODEL)\n\n\ndef tp_sharded() -> frozenset[str]:\n \"\"\"Points a tensor-parallel rank holds only a slice of, so the worker gathers them at collect\n (:mod:`interp_engine.vllm_capture._tp`) before rank 0's payload is read.\n\n **Not the complement of** :func:`d_model_wide`, which is what a serving pod used to narrow its\n served set by. That proxy holds for the head- and neuron-wide points and fails for the other\n two full-width cases: ``router_logits`` comes off a *replicated* gate and a hyper-connection\n trunk replicates its streams, so both reach rank 0 whole and neither should be refused for a\n sharding that does not apply to it.\n \"\"\"\n return frozenset(p.name for p in _ALL_POINTS if p.width in Width.sharded())\n\n\ndef token_flattened() -> frozenset[str]:\n \"\"\"Points whose module reports ``[batch * pos, ...]`` and need reshaping to ``[batch, pos, ...]``.\"\"\"\n return _names(token_flattened=True)\n\n\ndef reason(name: str) -> str:\n \"\"\"The reason for a point's vLLM limit, with any ``see `` indirection followed.\n\n The table uses ``see mlp_pre`` so that four QK-norm rows do not restate one paragraph, but an\n *error message* is the wrong place to send someone to another table row -- so the indirection is\n resolved here rather than printed. A test asserts every reference resolves and that none chain.\n \"\"\"\n found = _BY_NAME.get(name) or _BY_CONDITIONAL_NAME.get(name)\n if found is None:\n return \"not a canonical point name\"\n if not found.note:\n return found.vllm.value\n if not found.note.startswith(\"see \"):\n return found.note\n referent, _, extra = found.note.removeprefix(\"see \").partition(\";\")\n target = _BY_NAME.get(referent.strip()) or _BY_CONDITIONAL_NAME.get(referent.strip())\n if target is None: # pragma: no cover - pinned by test_every_reference_resolves\n return found.note\n return f\"as {referent.strip()}: {target.note}\" + (f\"; {extra.strip()}\" if extra.strip() else \"\")\n\n\ndef refusal_reasons(names: object) -> str:\n \"\"\"One line per refused point, quoting its own reason -- for a backend's error message.\n\n Built from the table so the message cannot go stale the way a hand-written list of point\n families does. Unknown names are reported as unknown rather than skipped.\n \"\"\"\n return \"\\n\".join(f\" {name}: {reason(name)}\" for name in sorted(str(n) for n in names)) # type: ignore[union-attr]"}, {"path":"interp_engine/protocol.py","title":"The InterpModel contract","why":"What both backends implement, and what is deliberately excluded from the protocol.","kind":"python","text":"\"\"\"The surface both backends share, so callers can hold a model without knowing which.\n\n:class:`InterpModel` is what :func:`interp_engine.load_model` returns, and the contract\n:class:`~interp_engine.model.EagerModel` and :class:`~interp_engine.vllm_backend.VLLMModel`\nboth satisfy. Code written against it runs unchanged on either.\n\nEverything that touches the model is ``async``, including on the eager backend where the\nwork is synchronous underneath. The alternative -- a sync protocol with an async vLLM\nescape hatch -- pushes the difference back onto every caller, which is the thing this is\nhere to remove. Eager's wrappers are thin (see :meth:`EagerModel.capture`), so a caller\nwith no event loop can drive them through ``asyncio.run`` or reach past the protocol to\nthe free functions (``run_with_cache``, ``steer``, ``generate_stream``), which stay sync\nand are the better fit for notebook use.\n\n``asyncio.run`` is eager-only advice. A vLLM model is bound to the loop that built its\nengine, and ``asyncio.run`` closes its loop on the way out, so a second call would reach\nan engine nothing is driving -- :func:`interp_engine._loop.refuse_foreign_loop` raises\nthere rather than letting it hang. The free functions and ``sync_model`` are loop-free on\nboth backends, and are what a sync caller should use when the backend is not known.\n\nDeliberately NOT in the protocol:\n\n- **Per-head points** (``value``, ``attn_probs``) and attention patterns. vLLM's paged\n attention kernel never materializes a probability matrix, so it needs the off-kernel\n recompute in ``capture_attention``, which does not exist eagerly (eager just reads the\n real softmax). Ask for these behind an ``isinstance`` check or a capability query.\n- **Prompt embeddings** (``generate_from_embeds``), which requires the vLLM engine to have\n been built with ``enable_prompt_embeds``.\n- **Weight and module access** (``hf_model``, ``resolve_point``, and gradients *through the\n forward*). vLLM owns its weights in a worker subprocess; anything reaching for a module is\n eager-only by nature. The gradient *verdict* is in the protocol (``grad_support``) even though\n the capability is not, so a caller can gate on fact rather than on backend name.\n\nSo a protocol-typed caller gets capture at the residual/MLP points, generation, and the\nlens read-out -- which covers the serving paths -- and everything else stays an explicit\nbackend choice rather than a method that raises on one of them.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import AsyncIterator, Sequence\nfrom dataclasses import dataclass, field\nfrom typing import Any, Protocol, runtime_checkable\n\nimport torch\n\nfrom interp_engine.address import Address\nfrom interp_engine.autograd_support import GradSupport\nfrom interp_engine.residual_basis import ResidualBasis\n\n# Re-exported so a protocol-typed caller needs one import. The type itself lives in\n# `interp_engine.address`, which owns the grammar; this module defines a Protocol and should not.\n__all__ = [\"Address\", \"Completion\", \"InterpModel\", \"Point\"]\n\n#: Deprecated alias for the two-tuple an address used to be. Kept because ``Point`` is a public\n#: export and downstream repos import this package unpinned, so deleting it would break them at\n#: import time rather than at the call site that needs updating. Accepted wherever an address is\n#: taken (see :func:`interp_engine.address.to_address`); never returned.\nPoint = tuple[str, int]\n\n\n@dataclass\nclass Completion:\n \"\"\"One generated completion, in the shape vLLM's ``CompletionOutput`` exposes.\n\n The eager backend returns this so that a caller reading ``.text`` / ``.token_ids`` off\n :meth:`InterpModel.capture_generation` does not have to care which backend produced it.\n vLLM returns its own richer object (also carrying ``.logprobs`` and ``.finish_reason``)\n rather than being narrowed to this, since callers that know they are on vLLM use those.\n \"\"\"\n\n text: str\n token_ids: list[int] = field(default_factory=list)\n\n\n@runtime_checkable\nclass InterpModel(Protocol):\n \"\"\"A language model you can capture from, steer, and generate with.\n\n ``runtime_checkable`` so ``isinstance(model, InterpModel)`` works as a coarse guard.\n Note that this only checks method *presence*, not signatures -- it will not catch a\n third-party model whose ``capture`` means something else.\n \"\"\"\n\n # --- identity and shape -------------------------------------------------\n hf_model_id: str\n \"\"\"The HuggingFace repo id this was loaded from.\"\"\"\n\n @property\n def n_layers(self) -> int:\n \"\"\"Number of decoder layers, so ``range(n_layers)`` enumerates valid layers.\"\"\"\n ...\n\n @property\n def d_model(self) -> int:\n \"\"\"Residual stream width. Note that ``z`` is ``n_heads * head_dim``, which is\n NOT ``d_model`` on every family (Gemma 3), so do not use this to size a ``z``.\"\"\"\n ...\n\n @property\n def grad_support(self) -> GradSupport:\n \"\"\"Whether this model can provide gradients, and what is blocking the rest.\n\n Cheap and side-effect-free on both backends -- it consults configuration only, never a\n forward pass or a worker, so it is safe to call before ``warmup()``. ``downstream`` is True\n everywhere; ``through_forward`` is eager-only, and only with ``requires_grad=True``. See\n :mod:`interp_engine.autograd_support`.\n \"\"\"\n ...\n\n @property\n def hooks_available(self) -> bool:\n \"\"\"Whether capture and steering can work at all on this instance.\n\n Always True on eager, which holds the module tree and hooks it in-process. On vLLM it is\n False when the engine was built with ``enforce_eager=False``, because CUDA graph replay never\n calls the Python ``forward`` a hook is attached to. That is **dynamic** hooks only: graph\n static taps are advertised separately as :attr:`static_points` / :attr:`static_writes`.\n Answerable without a forward, like the two verdicts above, so a server can advertise its\n endpoint set at startup -- the hook-dependent methods gate on it themselves either way.\n \"\"\"\n ...\n\n @property\n def graph_replay(self) -> bool:\n \"\"\"Whether this instance replays CUDA graphs instead of running Python ``forward``.\n\n False on eager. On vLLM, True iff ``enforce_eager`` is False (including static-mode\n engines). Unsteered generate still works; capture and steer need static sites or hooks.\n \"\"\"\n ...\n\n @property\n def static_points(self) -> tuple[Any, ...]:\n \"\"\"Sites baked into CUDA graphs as ``copy_`` taps. Empty on eager / hooked vLLM.\"\"\"\n ...\n\n @property\n def static_writes(self) -> tuple[Any, ...]:\n \"\"\"Static additive write sites. Empty when this instance declared no static writes.\"\"\"\n ...\n\n @property\n def residual_basis(self) -> ResidualBasis:\n \"\"\"How this model's residual stream is structured, and what that rules out.\n\n The same shape as :attr:`grad_support` and for the same reason: a capability that must not\n gate loading, must be answerable without a forward, and must produce one error text wherever\n the request came from. See :mod:`interp_engine.residual_basis`.\n \"\"\"\n ...\n\n # --- tokenization -------------------------------------------------------\n tokenizer: Any\n \"\"\"The HF tokenizer (or processor on multimodal archs), for chat templating and\n decoding. Untyped because those two have no common base class.\"\"\"\n\n def to_tokens(self, text: str | list[str], **kwargs: Any) -> torch.Tensor: ...\n\n def to_str_tokens(self, text: str | torch.Tensor, **kwargs: Any) -> list[str]: ...\n\n def to_string(self, tokens: Any) -> str | list[str]: ...\n\n # --- lifecycle ----------------------------------------------------------\n async def warmup(self) -> None:\n \"\"\"Pay any deferred load cost now rather than on the first request.\n\n On a graph-static engine this also runs a sentinel capture and write. A dead\n ``copy_`` or ``add_`` raises rather than serving unsteered text.\n \"\"\"\n ...\n\n async def shutdown(self) -> None:\n \"\"\"Release the model's device memory. Idempotent, and required before loading\n another model in the same process on vLLM, whose KV cache lives in a child\n process that a dropped Python reference does not reap.\"\"\"\n ...\n\n # --- capture ------------------------------------------------------------\n async def capture(\n self,\n prompt_token_ids: Sequence[int],\n points: Sequence[Address | str | Point],\n *,\n steering_spec: Any = None,\n detach: bool = True,\n ) -> dict[Address, torch.Tensor]:\n \"\"\"Capture ``points`` over one prompt's forward pass.\n\n Returns ``{Address: [n_prompt_tokens, width]}`` on CPU -- one row per prompt\n token, in order, with no batch dimension on either backend. ``width`` is\n :attr:`d_model` for the residual and MLP points and ``n_heads * head_dim`` for\n ``z``. Requests may be ``Address``es, their canonical string forms, or the\n ``(name, layer)`` tuples this used to take; the keys coming back are always\n ``Address``es, since that is the only form that can carry every coordinate.\n\n ``points`` names are the canonical ones (``resid_pre``, ``resid_mid``, ``resid_post``,\n ``mlp_in``, ``mlp_out``, ``attn_out``, ``mlp_out_post``, ``attn_out_post``, ``z``); ``value``,\n ``attn_probs``, ``attn_scores``, the QK-norm points, the MLP-internal points (``mlp_act``,\n ``mlp_pre``, ``mlp_pre_linear``) and the MoE routing points (``router_logits``,\n ``expert_weights``, ``expert_indices``) are eager-only, see the module\n docstring. The ``*_post`` pair is the\n sublayer's residual *contribution*, which differs from the raw output only on post-norm\n architectures (Gemma-2/3/4, OLMo-2/3) and aliases it everywhere else. With ``steering_spec`` (an engine ``SteeringSpec``) the\n activations are captured from the *steered* forward, not a separate one.\n\n ``detach=False`` keeps the autograd graph and returns device tensors instead of CPU ones.\n It raises :class:`~interp_engine.autograd_support.GradientsUnsupported` wherever\n :attr:`grad_support` says gradients cannot flow through the forward -- which is always on\n vLLM, and on eager unless the model was built with ``requires_grad=True``. It never\n silently returns detached tensors instead.\n \"\"\"\n ...\n\n async def capture_generation(\n self,\n prompt_token_ids: Sequence[int],\n points: Sequence[Address | str | Point],\n *,\n max_tokens: int = 8,\n temperature: float = 0.0,\n seed: int | None = None,\n steering_spec: Any = None,\n ) -> tuple[Any, dict[Address, torch.Tensor]]:\n \"\"\"Generate, capturing ``points`` at prompt AND generated positions.\n\n Returns ``(completion, {Address: [prompt_len + generated_len - 1, width]})``.\n The captured length is one short of prompt plus generated because the final sampled\n token is never fed back through the model -- autoregressive behavior, not a backend\n quirk. ``completion`` exposes ``.text`` and ``.token_ids``.\n \"\"\"\n ...\n\n async def capture_attention(\n self, prompt_token_ids: Sequence[int], layers: Sequence[int]\n ) -> dict[int, dict[str, torch.Tensor]]:\n \"\"\"Attention scores, probabilities and per-head values for ``layers``, one prompt.\n\n Returns ``{layer: {\"scores\": [heads, q, k], \"probs\": [heads, q, k],\n \"value\": [pos, kv_heads, v_head_dim]}}``. ``probs`` is the softmax of ``scores`` and both\n come from one pass; ``value`` is the per-head, family-scaled tensor that satisfies\n ``probs @ value == z``, not the raw projection output.\n\n Neither backend has these as module boundaries -- a fused kernel never forms the score\n matrix -- so both reconstruct them: eager from ``output_attentions``, which requires the\n model to have been loaded with eager attention, and vLLM off-kernel from captured\n post-RoPE q/k, gathered across ranks under tensor parallelism. Each names its own refusal.\n \"\"\"\n ...\n\n # --- generation ---------------------------------------------------------\n async def generate_text(\n self,\n prompt_token_ids: Sequence[int],\n *,\n max_tokens: int = 200,\n temperature: float = 1.0,\n seed: int | None = None,\n ) -> str:\n \"\"\"Generate and return the completion text (no prompt echo).\"\"\"\n ...\n\n def generate_stream(\n self,\n prompt_token_ids: Sequence[int],\n *,\n max_tokens: int = 200,\n temperature: float = 1.0,\n seed: int | None = None,\n ) -> AsyncIterator[str]:\n \"\"\"Yield decoded text deltas as they are produced.\n\n Deltas concatenate to what :meth:`generate_text` would have returned. Not declared\n ``async def`` here because an async generator's type is the iterator it returns, so\n an implementation may be either an ``async def`` generator or a method returning\n one. Eager streaming yields per-token; use ``interp_engine.generate_stream`` for the\n richer per-step form with logits and logprobs.\n \"\"\"\n ...\n\n # --- lens ---------------------------------------------------------------\n async def decode_residuals(self, residuals: torch.Tensor, *, detach: bool = True) -> torch.Tensor:\n \"\"\"Decode ``[n_rows, d_model]`` residuals to ``[n_rows, vocab]`` logits.\n\n Both backends apply the model's configured ``final_logit_softcapping`` when it has\n one, so the two are comparable; do not apply it again. (The sync free function\n ``interp_engine.decode_residuals`` returns RAW logits and takes ``softcap``\n explicitly -- this method is the one that normalizes across backends.)\n\n ``detach=False`` keeps the graph so the read-out can be differentiated with respect to the\n residuals you passed in. Eager honors it on a frozen model, because that gradient never needs\n to reach a parameter; vLLM raises, because the unembed happens in another process.\n \"\"\"\n ..."}, {"path":"interp_engine/dispatch.py","title":"Capability refusals","why":"The CAPABILITIES table: everything exactly one backend can do, why the other cannot, and what to call instead.","kind":"python","text":"\"\"\"What the unified free functions share to serve either backend.\n\n``run_with_cache``, ``generate_stream``, ``capture_residuals`` and the rest each keep their\neager body and gain a second arm that goes through :class:`~interp_engine.sync.SyncModel`. This\nmodule holds the two things all of those arms need: coercing whatever the caller passed as\ntokens into the shape each backend wants, and the capability table that decides what an arm\nmust refuse.\n\n**The dispatch shape to copy** when adding a function here -- a two-line front door, so the\neager path stays exactly what it was and the branch is visible at the top rather than threaded\nthrough the body:\n\n def thing(model: InterpModel, tokens: TokensLike, ...) -> Result:\n if not isinstance(model, EagerModel):\n return _thing_via_protocol(model, tokens, ...)\n return _thing_eager(model, tokens, ...)\n\nEager first in the sense that its body is untouched: ``engine_adapter.py`` calls\n``run_with_cache`` rather than ``await model.capture`` specifically because eager can hand back\na tensor on the device the SAE encode is about to run on, where the protocol method must return\nCPU tensors. Routing eager through the facade would quietly add a device round-trip to a serving\npath, so it does not.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Sequence\nfrom dataclasses import dataclass\nfrom typing import TYPE_CHECKING, TypeGuard\n\nimport torch\n\nfrom interp_engine.model import EagerModel\n\nif TYPE_CHECKING:\n from interp_engine.protocol import InterpModel\n\n#: What any unified entry point accepts as a token sequence: the ``[batch, seq]`` tensor the\n#: eager functions have always taken, a bare ``[seq]`` tensor, or a plain list of ids. Accepting\n#: all three because ``model.to_tokens(...)`` gives the first and the protocol methods take the\n#: third, and a caller moving between them should not have to convert.\nTokensLike = torch.Tensor | Sequence[int]\n\n\ndef as_token_ids(tokens: TokensLike, *, model: InterpModel, what: str) -> list[int]:\n \"\"\"One prompt's token ids, for the protocol arm.\n\n Refuses a batch rather than capturing row 0, which is the silent-wrong-answer version of\n the same call. Batching on vLLM belongs to its scheduler: issue concurrent requests, which\n it will batch better than stacking rows here would.\n \"\"\"\n if not isinstance(tokens, torch.Tensor):\n return [int(t) for t in tokens]\n if tokens.ndim == 2 and tokens.shape[0] != 1:\n raise ValueError(\n f\"{what} takes one prompt at a time on the {type(model).__name__} backend, got a batch \"\n f\"of {tokens.shape[0]}. Loop over the rows, or issue concurrent requests -- how to batch \"\n \"them into a forward pass is the engine's scheduling decision, not one this call can \"\n \"make on your behalf.\"\n )\n return [int(t) for t in tokens.reshape(-1).tolist()]\n\n\ndef as_batched_tokens(tokens: TokensLike, *, device: torch.device | str | None = None) -> torch.Tensor:\n \"\"\"A ``[batch, seq]`` int tensor, for the eager arm, which feeds it to ``hf_model``.\"\"\"\n ids = tokens if isinstance(tokens, torch.Tensor) else torch.tensor([int(t) for t in tokens], dtype=torch.long)\n if ids.ndim == 1:\n ids = ids.unsqueeze(0)\n return ids if device is None else ids.to(device)\n\n\n# ── Refusals ────────────────────────────────────────────────────────────────────────────────\n#\n# One signature over two backends creates exactly one new failure mode, and it is the dangerous\n# kind: an argument that means something on one backend and nothing on the other returns a\n# plausible, *unaffected* result with nothing anywhere to say so. That is the same shape of bug\n# as a steering hook that never fires under CUDA graph replay, which is what `hooks_available`\n# exists to prevent.\n#\n# So: every argument the chosen backend cannot honor raises, before any work is done, naming the\n# capability and what to do instead. Not a warning, not a no-op, not a best-effort\n# approximation. The messages are built from the table below rather than written at each call\n# site, for the same reason `points.refusal_reasons()` is a table -- so the text cannot go stale\n# in one arm while staying right in another, and so `tests/test_capability_refusals.py` can\n# enumerate them.\n\n\nclass CapabilityUnsupported(ValueError):\n \"\"\"A capability was requested on a backend that cannot provide it.\n\n A ``ValueError``, matching :class:`~interp_engine.residual_basis.ResidualBasisUnsupported`,\n since these refusals sit on the same call paths and callers already catch that.\n \"\"\"\n\n\n@dataclass(frozen=True)\nclass Capability:\n \"\"\"One thing a backend may or may not be able to do, and how to say so.\n\n ``what`` is the capability in the caller's terms, ``why`` is the mechanical reason the\n backend cannot, and ``instead`` is the call that does work. All three are needed for the\n message to be actionable: a refusal that says only \"unsupported\" sends the reader into the\n source to find out whether they are holding it wrong or asking for the impossible.\n \"\"\"\n\n what: str\n why: str\n instead: str\n\n\n#: Every capability one backend has and the other does not. Keyed by a short stable string so a\n#: refusal and its test name the same thing. Add a row here rather than a string literal in an\n#: arm; ``tests/test_capability_refusals.py`` walks this table.\nCAPABILITIES: dict[str, Capability] = {\n \"raw_logits\": Capability(\n what=\"logits with the family's post-unembed arithmetic left to the caller\",\n why=(\n \"the unembed runs in a worker through vLLM's own compute_logits, which applies the \"\n \"model's final-logit softcap inside itself, so raw logits never cross the process \"\n \"boundary\"\n ),\n instead=(\n \"sync_model(model).decode_residuals(residuals) (or `await model.decode_residuals`), \"\n \"which returns the model-normalized read-out on both backends\"\n ),\n ),\n \"explicit_logit_transform\": Capability(\n what=\"a caller-supplied softcap or logit multiplier\",\n why=(\n \"the worker's unembed has already applied the model's own, so applying one here would \"\n \"double it rather than override it\"\n ),\n instead=\"omitting them, which yields the model's own transform on both backends\",\n ),\n \"module_weights\": Capability(\n what=\"a computation over the model's own weight matrices\",\n why=\"vLLM owns its weights in a worker subprocess and does not expose them to this process\",\n instead=\"running this on an eager model, which holds the live module tree\",\n ),\n \"attention_mask\": Capability(\n what=\"a padding mask\",\n why=\"this arm forwards one unpadded sequence, so there is nothing for a mask to cover\",\n instead=\"omitting it, and issuing one call per prompt\",\n ),\n \"masked_steer_positions\": Capability(\n what=\"steering that skips some prompt positions\",\n why=(\n \"the per-request steer hook applies to every row of the request it is registered for, \"\n \"and the wire format carries no position selection\"\n ),\n instead=(\n \"steering every position, or generating with model.generate_steered, which does resolve a position_mask\"\n ),\n ),\n \"full_logits_per_step\": Capability(\n what=\"the full vocab-wide logit vector at each generated position\",\n why=(\n \"vLLM's sampler returns the top-n logprobs it was asked for rather than the logit \"\n \"tensor, which is never shipped out of the worker\"\n ),\n instead=(\"n_logprobs=N for the top N at each step, which is what GenStep.logprobs carries on both backends\"),\n ),\n}\n\n\ndef is_eager(model: InterpModel) -> TypeGuard[EagerModel]:\n \"\"\"Whether this is the in-process eager backend, narrowed for the type checker.\"\"\"\n return isinstance(model, EagerModel)\n\n\ndef refuse(model: InterpModel, what: str, *, capability: str) -> CapabilityUnsupported:\n \"\"\"The exception for ``capability`` being unavailable on ``model``. Returned, not raised.\n\n Returned so the call site reads ``raise refuse(...)`` and the traceback starts there rather\n than in here.\n \"\"\"\n cap = CAPABILITIES[capability]\n return CapabilityUnsupported(\n f\"{what} needs {cap.what}, which the {type(model).__name__} backend cannot provide: \"\n f\"{cap.why}. Instead: {cap.instead}.\"\n )\n\n\ndef require_eager(model: InterpModel, what: str, *, capability: str) -> None:\n \"\"\"Refuse unless ``model`` is the eager backend. Call before doing any work.\"\"\"\n if not is_eager(model):\n raise refuse(model, what, capability=capability)\n\n\ndef refuse_arguments(model: InterpModel, what: str, *, capability: str, given: dict[str, object]) -> None:\n \"\"\"Refuse if any of ``given`` was actually passed (is not ``None``), naming which.\n\n For the arguments that exist on the unified signature because one backend honors them. The\n check is \"was it passed\", not \"is it non-default\", so a caller who explicitly asked for\n something is told no rather than served a result that ignored them.\n \"\"\"\n passed = sorted(name for name, value in given.items() if value is not None)\n if not passed:\n return\n cap = CAPABILITIES[capability]\n raise CapabilityUnsupported(\n f\"{what} cannot honor {', '.join(passed)} on the {type(model).__name__} backend: \"\n f\"{cap.why}. Instead: {cap.instead}.\"\n )"},