interp-engine 1.14.1__tar.gz → 1.14.2__tar.gz

{interp_engine-1.14.1 → interp_engine-1.14.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.14.1
+Version: 1.14.2
 Summary: A fast, standardized, and easy to use interpretability engine.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine

{interp_engine-1.14.1 → interp_engine-1.14.2}/docs/AGENT_INTEGRATION.md

@@ -134,7 +134,7 @@ engine_kwargs = {"worker_extension_cls": WORKER_EXTENSION_CLS, **capture_engine_
 
 The decision rule, in order:
 
-1. Do you need eager-only points (MLP internals, MoE selection), gradients, or module/weight access?
+1. Do you need eager-only points (MLP input branches, MoE selection), gradients, or module/weight access?
    → **eager**, `load_model(..., backend="eager")`. Neither tier of vLLM serves these; see [the
    vocabulary table](#the-point-vocabulary). Attention is *not* on that list any more —
    `capture_attention` serves it on both.
@@ -189,14 +189,16 @@ Canonical names, with the layer after a dot: `resid_post.10`. Extra coordinates
 | `attn_out` | attention's raw output | both |
 | `mlp_out_post`, `attn_out_post` | that sublayer's **residual contribution** — differs from the raw output only on post-norm architectures | both |
 | `z` | per-head attention output, `n_heads * head_dim` wide (**not** `d_model` on every family) | both |
-| `value`, `attn_probs`, `attn_scores` | per-head values and the attention matrix | both, but through `capture_attention` rather than as points on vLLM |
-| `mlp_act`, `mlp_pre`, `mlp_pre_linear` | MLP internals, `d_mlp` wide | **eager only** |
+| `value` | per-head values, the value projection's output on most families | both (head-sharded; gathered across TP ranks) |
+| `attn_probs`, `attn_scores` | the attention matrix | both, but through `capture_attention` rather than as points on vLLM |
+| `mlp_act` | post-activation neurons, the down projection's input, `d_mlp` wide | both (neuron-sharded; gathered across TP ranks) |
+| `mlp_pre`, `mlp_pre_linear` | the MLP's input branches, `d_mlp` wide | **eager only** |
 | `router_logits` | MoE routing scores, every expert | both |
 | `expert_weights`, `expert_indices` | the top-k the router selected, and its weights | **eager only** |
 | the QK-norm points | inside the attention module | both (head-sharded; gathered across TP ranks) |
 
-Eager-only is not an omission: vLLM's fused MLP and MoE kernels compute those tensors inline, so
-there is no module boundary to hook. This table is the working subset;
+Eager-only is not an omission: vLLM fuses `gate_proj` and `up_proj` into one `gate_up_proj`, and its
+MoE kernel forms the selection inline, so there is no module boundary to hook. This table is the working subset;
 [SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) is all 34 points with the per-backend verdict on each,
 and it is checked against the registry rather than maintained by hand.
 

{interp_engine-1.14.1 → interp_engine-1.14.2}/docs/ARCHITECTURE_QUIRKS.md

@@ -311,7 +311,7 @@ must split", on a family whose `fused_qkv` is false — so `value` no longer ask
 **Reading a norm costs a rank, and the point does not pay it.** A norm over `head_dim` can only be
 given the per-head view, so both engines hand `v_norm` a `(…, n_kv_heads, head_dim)` tensor where a
 `v_proj` would have produced the flat one — and `value` would then be 4-D on this family and 3-D on the
-next, against a `Width.HEADS` declaration and `capture`'s `[batch, seq, width]`. Flattened back
+next, against a `Width.KV_HEADS` declaration and `capture`'s `[batch, seq, width]`. Flattened back
 on both sides (`hooks.flat_per_head`, `vllm_capture._hooks.flat_value`), including on the way *into* a
 steer, so a vector measured on a capture of `value` is the shape the write expects; the module gets its
 own rank back before the attention reshapes it. Eagerly the head count is checked rather than assumed,

{interp_engine-1.14.1 → interp_engine-1.14.2}/docs/SUPPORTED_POINTS.md

@@ -16,7 +16,7 @@ nnsight and nnterp.
 | [`attn_in`][attn_in]                                  | `d_model`               |  ✅   |  ✅  |                                                                                                                                                                                                                                           |
 | [`q_norm_in`][q_norm_in] / [`q_norm_out`][q_norm_out] | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
 | [`k_norm_in`][k_norm_in] / [`k_norm_out`][k_norm_out] | `n_kv_heads * head_dim` |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
-| [`value`][value]                                      | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
+| [`value`][value]                                      | `n_kv_heads * head_dim` |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |
 | [`attn_scores`][attn_scores]                          | `n_heads * query * key` |  ✅   |  ♻️  | no module boundary holds the pre-softmax matrix on **either** backend; vLLM rebuilds it from captured post-RoPE q/k                                                                                                                       |
 | [`attn_probs`][attn_probs]                            | `n_heads * query * key` |  ✅   |  ♻️  | fused paged attention never materializes the probabilities; same recompute                                                                                                                                                                |
 | [`z`][z]                                              | `n_heads * head_dim`    |  ✅   |  ✅  | head-sharded; gathered across TP ranks at collect                                                                                                                                                                                         |

{interp_engine-1.14.1 → interp_engine-1.14.2}/docs/USAGE.md

@@ -148,8 +148,8 @@ Three things to know about what comes back:
   one up, or `format_address(key)` to print one. (The eager `Cache` from `capture` is the
   exception — it accepts either form on lookup.)
 - **Shape is `[n_prompt_tokens, width]`, on CPU, with no batch dimension** — on either backend.
-  `width` is `d_model` for the residual and MLP points, and `n_heads * head_dim` for `z`, which is
-  *not* `d_model` on every family.
+  `width` is `d_model` for the residual and sublayer-boundary points, `d_mlp` for `mlp_act`, and
+  `n_heads * head_dim` for `z`, which is *not* `d_model` on every family.
 - **Rows are prompt positions in order**, so `acts[address][3]` is the activation at token 3.
 
 ## Naming a point
@@ -167,16 +167,16 @@ Points every backend serves: `resid_pre`, `resid_mid`, `resid_post`, `mlp_in`, `
 `attn_out`, `mlp_out_post`, `attn_out_post`, `z`. [SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) is the
 full list, with the width and the per-backend verdict for each.
 
-Eager-only, because they are formed inside a fused kernel vLLM never unfolds: the MLP internals
-(`mlp_act`, `mlp_pre`, `mlp_pre_linear`) and the MoE *selection* (`expert_weights`,
-`expert_indices`) — its `router_logits` are served on both, since the gate produces them before any
-kernel takes over.
+Eager-only, because they are formed inside a fused kernel vLLM never unfolds: the MLP input branches
+(`mlp_pre`, `mlp_pre_linear`, which vLLM fuses into one `gate_up_proj`) and the MoE *selection*
+(`expert_weights`, `expert_indices`). The neuron basis `mlp_act` is the down projection's input, and
+`router_logits` comes from the gate before any kernel takes over, so both are served on both.
 
-The attention trio (`attn_scores`, `attn_probs`, `value`) is served on both, but not as a point on
-vLLM: no module boundary holds a score matrix the paged kernel never forms, so both backends
-reconstruct it. Ask for it with `capture_attention`, which returns the same
-`{layer: {"scores", "probs", "value"}}` either way — see [Without an event
-loop](#without-an-event-loop).
+The attention matrix (`attn_scores`, `attn_probs`) is served on both, but not as a point on vLLM:
+no module boundary holds a score matrix the paged kernel never forms, so both backends reconstruct
+it. Ask for it with `capture_attention`, which returns the same `{layer: {"scores", "probs",
+"value"}}` either way — see [Without an event loop](#without-an-event-loop). The raw `value` point
+is an ordinary capture point on both.
 
 The `*_post` pair deserves attention: `mlp_out` is the raw module output and `mlp_out_post` is that
 sublayer's *contribution to the residual stream*. They differ only on post-norm architectures

{interp_engine-1.14.1 → interp_engine-1.14.2}/interp_engine/capture.py

@@ -554,7 +554,7 @@ def _capture_eager(
         raw = cache.tensors[key]
         if model.arch.quirks.fused_qkv:
             # The value alone, not the `[q | k | v]` slab the fused projection emits. This point
-            # declares `Width.HEADS` and both vLLM backends return exactly that, so leaving the slab
+            # declares `Width.KV_HEADS` and both vLLM backends return exactly that, so leaving the slab
             # here made the reference three times too wide on gpt2, GPT-NeoX and Falcon -- and the
             # engine every other one is scored against was the one disagreeing.
             raw = split_fused_qkv(model, raw)["v"]
@@ -699,7 +699,7 @@ def per_head_value(model: EagerModel, cache: Cache, layer: int) -> torch.Tensor:
     point, is what satisfies ``probs @ value == z``.
 
     A fused qkv is already split by the time the point lands in the cache, so there is nothing to
-    unpack here: the raw point is the value at ``Width.HEADS`` on every family.
+    unpack here: the raw point is the value at ``Width.KV_HEADS`` on every family.
     """
     raw = cache.get("value", layer)  # [batch, seq, n_kv_heads * v_head_dim]
     batch, seq, _ = raw.shape

{interp_engine-1.14.1 → interp_engine-1.14.2}/interp_engine/hooks.py

@@ -54,7 +54,7 @@ def flat_per_head(tensor: torch.Tensor, *, heads: int) -> tuple[torch.Tensor, to
     The module `value` resolves to is a *norm* wherever the family has one, and a norm over `head_dim`
     has to be given the per-head view: Gemma-4 runs `v_norm` on `v_proj(h).view(..., n_kv, head_dim)`
     and vLLM runs the same norm on `v.unflatten(-1, (n_kv, head_dim))`. So the point arrives per head
-    on one family and flat on the next, while `Width.HEADS` and `capture` both promise one
+    on one family and flat on the next, while `Width.KV_HEADS` and `capture` both promise one
     rank -- and TransformerLens' `hook_v` is the per-head spelling of this point, which is a
     distinction the mapping table would lose if ours moved family to family.
 

{interp_engine-1.14.1 → interp_engine-1.14.2}/interp_engine/points.py

@@ -64,11 +64,13 @@ class Width(Enum):
     """What the last axis of the captured tensor counts -- the axis a shard would narrow.
 
     This is the tensor-parallel question, which is why it is an enum rather than a
-    ``d_model: bool``. ``HEADS``, ``NEURONS`` and ``VOCAB`` are exactly the axes vLLM shards, so
-    under TP rank 0 holds a slice of them -- and for the first two there is no width to check it
-    against either (``n_heads * head_dim`` coincides with ``hidden_size`` on Llama and not on
-    Gemma-3, so a width check would pass on one and fail on the other for reasons unrelated to
-    sharding). :func:`tp_sharded` is that set.
+    ``d_model: bool``. ``HEADS``, ``KV_HEADS``, ``NEURONS`` and ``VOCAB`` are exactly the axes vLLM
+    shards, so under TP rank 0 holds a slice of them -- and for the first three there is no width to
+    check it against either (``n_heads * head_dim`` coincides with ``hidden_size`` on Llama and not
+    on Gemma-3, so a width check would pass on one and fail on the other for reasons unrelated to
+    sharding). :func:`tp_sharded` is that set. ``KV_HEADS`` is its own member because a GQA model
+    has fewer KV heads than query heads, and because vLLM replicates KV heads, not shards them,
+    when there are fewer of them than ranks.
 
     The rest reach a hook at full width, for three different reasons, which is why "is it
     ``D_MODEL``" is *not* the same question: ``D_MODEL`` is all-reduced before the hook sees it,
@@ -81,6 +83,7 @@ class Width(Enum):
 
     D_MODEL = "d_model"
     HEADS = "n_heads * head_dim"
+    KV_HEADS = "n_kv_heads * head_dim"
     NEURONS = "d_mlp"
     ROUTING = "experts"
     SCORES = "n_heads * query * key"
@@ -95,7 +98,7 @@ class Width(Enum):
     @classmethod
     def sharded(cls) -> frozenset[Width]:
         """The widths a tensor-parallel rank holds only a slice of. See the class docstring."""
-        return frozenset({cls.HEADS, cls.NEURONS, cls.VOCAB})
+        return frozenset({cls.HEADS, cls.KV_HEADS, cls.NEURONS, cls.VOCAB})
 
 
 class VllmSupport(Enum):
@@ -151,9 +154,9 @@ POINTS: tuple[PointSpec, ...] = (
     _p("attn_in", _L, Width.D_MODEL, _HOOKS),
     _p("q_norm_in", _L, Width.HEADS, _HOOKS),
     _p("q_norm_out", _L, Width.HEADS, _HOOKS),
-    _p("k_norm_in", _L, Width.HEADS, _HOOKS),
-    _p("k_norm_out", _L, Width.HEADS, _HOOKS),
-    _p("value", _L, Width.HEADS, _HOOKS),
+    _p("k_norm_in", _L, Width.KV_HEADS, _HOOKS),
+    _p("k_norm_out", _L, Width.KV_HEADS, _HOOKS),
+    _p("value", _L, Width.KV_HEADS, _HOOKS),
     _p(
         "attn_scores",
         _L,

{interp_engine-1.14.1 → interp_engine-1.14.2}/interp_engine/protocol.py

@@ -20,8 +20,8 @@ both backends, and are what a sync caller should use when the backend is not kno
 
 Deliberately NOT in the protocol:
 
-- **Per-head points through ``capture``** (``value``, ``attn_probs``). No backend holds these
-  at a module boundary, so each reconstructs them, and the reconstruction is
+- **The attention matrix through ``capture``** (``attn_scores``, ``attn_probs``). No backend
+  holds these at a module boundary, so each reconstructs them, and the reconstruction is
   ``capture_attention`` -- which *is* in the protocol. What is not is asking for them as
   ordinary capture points. Gate on ``refuses(point, layer)`` rather than on the backend name.
 - **vLLM ``SamplingParams`` over embeddings** (``VLLMModel.generate_from_embeds``, deprecated).
@@ -398,17 +398,17 @@ class InterpModel(Protocol):
         token, in order, with no batch dimension on either backend. With ``rows``, only those
         positions come back, in that order (``[len(rows), width]``); on vLLM the worker drops
         the others before they cross the process boundary. ``width`` is
-        :attr:`d_model` for the residual and MLP points and ``n_heads * head_dim`` for
-        ``z``. Requests may be ``Address``es, their canonical string forms, or the
-        ``(name, layer)`` tuples this used to take; the keys coming back are always
-        ``Address``es, since that is the only form that can carry every coordinate.
+        :attr:`d_model` for the residual and sublayer-boundary points, ``d_mlp`` for
+        ``mlp_act``, and ``n_heads * head_dim`` for ``z``. Requests may be ``Address``es, their
+        canonical string forms, or the ``(name, layer)`` tuples this used to take; the keys coming
+        back are always ``Address``es, since that is the only form that can carry every coordinate.
 
         ``points`` names are the canonical ones (``resid_pre``, ``resid_mid``, ``resid_post``,
-        ``mlp_in``, ``mlp_out``, ``attn_out``, ``mlp_out_post``, ``attn_out_post``, ``z``); ``value``,
-        ``attn_probs``, ``attn_scores``, the QK-norm points, the MLP-internal points (``mlp_act``,
-        ``mlp_pre``, ``mlp_pre_linear``) and the MoE routing points (``router_logits``,
-        ``expert_weights``, ``expert_indices``) are eager-only, see the module
-        docstring. The ``*_post`` pair is the
+        ``mlp_in``, ``mlp_out``, ``attn_out``, ``mlp_out_post``, ``attn_out_post``, ``z``,
+        ``value``, ``mlp_act``, ``router_logits``, the QK-norm points); ``mlp_pre``,
+        ``mlp_pre_linear``, ``expert_weights`` and ``expert_indices`` are eager-only, and
+        ``attn_probs`` / ``attn_scores`` come from ``capture_attention``, see the module
+        docstring. :func:`interp_engine.points.eager_only` is the full list. The ``*_post`` pair is the
         sublayer's residual *contribution*, which differs from the raw output only on post-norm
         architectures (Gemma-2/3/4, OLMo-2/3) and aliases it everywhere else. With ``steering_spec`` (a ``SteeringSpec``, or a list for several points) the
         activations are captured from the *steered* forward, not a separate one.

{interp_engine-1.14.1 → interp_engine-1.14.2}/interp_engine/vllm_capture/_tp.py

@@ -29,8 +29,6 @@ from interp_engine.address import parse_address
 from interp_engine.points import Width, point_spec, tp_sharded
 from interp_engine.vllm_capture._tree import _attn_module, _get_layers, _worker_tp_world_size
 
-#: Points whose heads are the *KV* heads, and so may be replicated across ranks rather than sharded.
-_KV_HEAD_POINTS = frozenset({"k_norm_in", "k_norm_out", "value"})
 #: Attention payload roles by which head count they carry.
 _ATTN_ROLE_KV = frozenset({"k", "v"})
 _SHARDED_POINTS = tp_sharded()
@@ -138,7 +136,7 @@ def gather_capture(model: torch.nn.Module, key: str, tensor: torch.Tensor) -> to
     if counts is None:
         return _all_gather(tensor, -1)
     n_heads, n_kv, total_kv, head_dim = counts
-    if address.name in _KV_HEAD_POINTS:
+    if spec is not None and spec.width is Width.KV_HEADS:
         return gather_heads(tensor, per_rank=n_kv, total=total_kv, head_dim=head_dim)
     return gather_heads(tensor, per_rank=n_heads, total=None, head_dim=head_dim)
 

{interp_engine-1.14.1 → interp_engine-1.14.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "interp-engine"
-version = "1.14.1"
+version = "1.14.2"
 description = "A fast, standardized, and easy to use interpretability engine."
 readme = "README.md"
 license = "Apache-2.0"

{interp_engine-1.14.1 → interp_engine-1.14.2}/tests/test_points_registry.py

@@ -251,6 +251,13 @@ def test_the_sharded_set_is_the_narrowed_axes_and_not_just_the_non_d_model_ones(
     assert not ({"router_logits", "resid_post", "attn_in", "resid_streams"} & sharded)
 
 
+def test_the_kv_head_points_are_declared_kv_head_wide():
+    """On a GQA model these are ``n_kv_heads * head_dim``, and the TP gather reads this width to
+    pick the KV head count and drop replicated KV heads."""
+    kv = {p.name for p in points.POINTS if p.width is points.Width.KV_HEADS}
+    assert kv == {"k_norm_in", "k_norm_out", "value"}
+
+
 def test_the_mapping_doc_marks_exactly_the_eager_only_points():
     """The footnote markers in the docs table, checked against the table in code.
 

{interp_engine-1.14.1 → interp_engine-1.14.2}/tests/test_qkv_layout.py

@@ -136,7 +136,7 @@ def test_per_head_value_agrees_with_the_split(spec: ModelSpec, expected: QKVLayo
 
 @pytest.mark.parametrize(("spec", "expected"), FUSED_MODELS)
 def test_the_raw_value_point_is_the_value_and_not_the_whole_fused_slab(spec: ModelSpec, expected: QKVLayout):
-    """`value` declares ``Width.HEADS``, so a fused family has to be split before the point is served.
+    """`value` declares ``Width.KV_HEADS``, so a fused family has to be split before the point is served.
 
     It was not: the raw point handed back ``[q | k | v]``, three times too wide on gpt2 -- and on the
     one engine every other one is scored against. Both vLLM backends already returned the value alone,