interp-engine 1.3.6__tar.gz → 1.4.0__tar.gz

{interp_engine-1.3.6 → interp_engine-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.3.6
+Version: 1.4.0
 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
@@ -139,6 +139,13 @@ We verify correctness in two main ways:
 1. A test suite that checks results across several models - what each check is designed to catch is in [INTERNALS.md](docs/INTERNALS.md#correctness).
 2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at [`validator/`](validator/).
 
+## Known Issues
+
+- **Gemma 4 requires transformers 5.14.1** because 5.15 moved `head_dim` into `per_layer_config` and vLLM's config read dies before a weight loads ([vllm#51744](https://github.com/vllm-project/vllm/issues/51744)).
+- **Gemma 4 26B refuses `mlp_out`**, because `layer.mlp` there is one branch of a two-branch feed-forward; `mlp_out_post` is the residual contribution and is served ([ARCHITECTURE_QUIRKS.md](docs/ARCHITECTURE_QUIRKS.md#moe-tap-the-block-not-the-experts)).
+- **DeepSeek-V2 on transformers older than 5.15.0** captures a wrong attention temperature — the engine warns at load, and upgrading is the fix ([COMPATIBILITY.md](docs/COMPATIBILITY.md)).
+- **MXFP4 checkpoints (gpt-oss) need `interp-engine[quant]`**, which the `[vllm]` extra does not include; without it transformers dequantizes them to bf16 at roughly 3x the weights, which can turn a model that fits into one that does not.
+
 ## Why use an Interpretability Engine instead of building from scratch?
 
 1. **Speed**: Get performance without sacrificing correctness.

{interp_engine-1.3.6 → interp_engine-1.4.0}/README.md

@@ -108,6 +108,13 @@ We verify correctness in two main ways:
 1. A test suite that checks results across several models - what each check is designed to catch is in [INTERNALS.md](docs/INTERNALS.md#correctness).
 2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at [`validator/`](validator/).
 
+## Known Issues
+
+- **Gemma 4 requires transformers 5.14.1** because 5.15 moved `head_dim` into `per_layer_config` and vLLM's config read dies before a weight loads ([vllm#51744](https://github.com/vllm-project/vllm/issues/51744)).
+- **Gemma 4 26B refuses `mlp_out`**, because `layer.mlp` there is one branch of a two-branch feed-forward; `mlp_out_post` is the residual contribution and is served ([ARCHITECTURE_QUIRKS.md](docs/ARCHITECTURE_QUIRKS.md#moe-tap-the-block-not-the-experts)).
+- **DeepSeek-V2 on transformers older than 5.15.0** captures a wrong attention temperature — the engine warns at load, and upgrading is the fix ([COMPATIBILITY.md](docs/COMPATIBILITY.md)).
+- **MXFP4 checkpoints (gpt-oss) need `interp-engine[quant]`**, which the `[vllm]` extra does not include; without it transformers dequantizes them to bf16 at roughly 3x the weights, which can turn a model that fits into one that does not.
+
 ## Why use an Interpretability Engine instead of building from scratch?
 
 1. **Speed**: Get performance without sacrificing correctness.

{interp_engine-1.3.6 → interp_engine-1.4.0}/docs/ARCHITECTURE_QUIRKS.md

@@ -291,12 +291,49 @@ way for a layer to have no `v_proj`, and it does not overlap with KV sharing: th
 against a key of `rope(k_norm(k_proj(h)))`, so the two are the same projection read through
 different norms, not the same tensor.
 
-**What the engine does with that today is wrong in the quiet direction.** `_first_attr` treats a
-present-but-`None` attribute as absent, so `ArchSpec.v_proj` returns `None` on those layers — which
-is the signal it uses for "fused QKV, the caller must split" — on a family whose `fused_qkv` is
-false. The layer needs a third answer: `value` is capturable, from `k_proj` and the `v_norm` after
-it, and refusing or mis-addressing it is not the same as it not being there. Nothing exercises this
-yet because the 31B has been swept without the value points.
+**The third answer that layer needed is the norm, and it applies to the whole family.**
+`ArchSpec.v_proj` still returns `None` there, which is the signal it uses for "fused QKV, the caller
+must split", on a family whose `fused_qkv` is false — so `value` no longer asks it first.
+`ArchSpec.value_module` does, and it returns the `v_norm` whenever the family has one
+(`facts.ATTN_VALUE_NORM_ATTRS`). Two things that buys, and the first would apply even without
+`attention_k_eq_v`:
+
+- `Gemma4TextAttention.forward` runs `value_states = self.v_norm(value_states)` on **every** layer
+  that projects its own KV, so the projection's output is a norm short of the tensor the attention
+  pattern multiplies. `value` was off by that norm on the sliding layers too, which do have a
+  `v_proj`.
+- vLLM has no value projection to compare against at all — its `Gemma4Attention` splits one fused
+  `qkv_proj`, with the K weights loaded into the V slot on the `k_eq_v` layers — but it runs the same
+  `v_norm` on the V slice. So this is the only boundary at which the two engines hold the same tensor,
+  which is what makes the point checkable rather than merely servable. `_tree._value_module` is the
+  vLLM side, and it prefers the norm in the same order for the same reason.
+
+**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 `run_with_cache`'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,
+because the other 4-D layout in circulation is `[batch, heads, pos, head_dim]` — transformers norms
+after the head transpose on ten families — and flattening that one puts the head count where every
+reader expects the sequence.
+
+**Off Gemma-4, `value` on vLLM is one third of a packed matrix.** Every family with a fused vLLM
+implementation goes through `QKVParallelLinear`, whose output is `[q | k | v]` on the last axis, so an
+output hook there returns all three under the value's name at three times the width.
+`_tree.value_span` measures the q and k widths off the projection's own rank-local geometry
+(`num_heads`, `num_kv_heads`, `head_size`, each divided by the TP size in its `__init__`) and
+`_hooks.value_columns` applies it on both install paths, refusing a geometry it can read only in part —
+every wrong offset into a packed matrix yields a right-shaped tensor of another projection's heads. On
+the per-request path the narrowing happens *before* steering, so a steer on `value` leaves the same
+matrix's q and k alone.
+
+The `k_eq_v` layers stay the interesting case for a *reader*: `value` there is `v_norm(k_proj(h))`
+where the key is `rope(k_norm(k_proj(h)))`, so the two points come off one projection through
+different norms and only one of them has RoPE applied. `ArchSpec.is_k_eq_v_layer` answers which layers
+those are, gated on `not sliding` the way the modeling code gates it — the flag is model-wide and the
+structure is not.
 
 Two more ways the value side can differ from the query side, both of which produce a correctly shaped
 and completely wrong per-head split if you reshape by `head_dim`:
@@ -438,7 +475,7 @@ n_experts]` tensor that is plausible rather than obviously wrong.
   which they are, because the LongCat guard makes the `mlp_pre` refusal require a router **on the MLP
   module**, and Gemma-4's is a sibling. The parameter count of a sparse layer includes the dense MLP,
   so the usual `(n_layers - n_sparse) * dense_mlp` undercounts (the 26B by ~0.5B). And `mlp_out` is
-  **not** the whole feed-forward here, which is the one that has no fix yet — see below.
+  **not** the whole feed-forward here, so it is refused — see below.
 
 Which experts fired _is_ capturable, one level down: `router_logits` / `expert_weights` /
 `expert_indices` are three elements of the router module's own output tuple, resolved by
@@ -484,18 +521,33 @@ module at all. Gemma-4 is the exception noted above, and the refusal already let
 right reason rather than by luck: it is gated on the router being found *on the MLP module*, the same
 guard that keeps LongCat's shortcut MoE from suppressing its two real feed-forwards.
 
-**Open on Gemma-4's 26B**, both from the same fact — the routed branch is a sibling of `layer.mlp`
-rather than a part of it:
-
-- `mlp_out` taps `layer.mlp` and so returns the dense branch alone, missing the experts' contribution
-  entirely. Correctly shaped, wrong tensor, no tell. Either the point has to mean the summed
-  feed-forward here (which has no module boundary — the sum happens in the block's `forward`) or it
-  has to be refused with an explanation, and that choice has not been made.
-- `arch.moe_router` looks for the router under `layer.mlp` and Gemma-4 hangs it on the block, so
-  `router_logits` / `expert_weights` / `expert_indices` do not resolve at all. When they do, note that
-  `Gemma4TextRouter` returns `(router_probabilities, top_k_weights, top_k_index)` — element 0 is
-  already softmaxed over all 128 experts, so it matches the default tuple *order* while not being
-  logits, and `assert_routing_shapes` cannot tell the difference because the width is right.
+**Three things follow on Gemma-4's 26B from that one fact** — the routed branch is a sibling of
+`layer.mlp` rather than a part of it, so the *block* owns both halves of the feed-forward:
+
+- **`mlp_out` is refused**, on both backends. It taps `layer.mlp`, which here returns the dense branch
+  alone: correctly shaped `[tokens, d_model]`, at the right positions, missing the experts entirely,
+  and with nothing about it to notice — both engines build the same tree, so a cross-engine sweep
+  agrees on the same half rather than catching it. The summed tensor has no module boundary to serve
+  instead (the sum is a local of the block's `forward`), and there is already a point that means it:
+  `mlp_out_post` is the `post_feedforward_layernorm`'s output, downstream of the sum, and the one for
+  which `resid_post == resid_mid + mlp_out_post` holds. `ArchSpec.mlp_is_half_the_feed_forward` is the
+  question both refusals ask. Nothing else moves: `mlp_in`, `mlp_pre`, `mlp_pre_linear` and `mlp_act`
+  are the dense branch's own internals and mean what they mean everywhere else.
+- **The router is looked for on the block as well as on the MLP** (`facts.moe_router_owner`). Asking
+  only `layer.mlp` left all three routing points unresolvable on the 26B. Presence alone is not enough
+  now that a decoder layer is searched: `facts.moe_router_attr` requires a callable submodule, because
+  a block can hold a plain flag under one of these names and the caller is about to install a forward
+  hook on whatever comes back.
+- **`router_logits` is read one module deeper than the tuple.** `Gemma4TextRouter.forward` returns
+  `(router_probabilities, top_k_weights, top_k_index)`: element 0 matches the default tuple *order*
+  while being a softmax over all 128 experts, and is exactly as wide as the logits over the same bank,
+  so neither `ROUTER_OUTPUTS` nor a width check can see it — and the block discards it. So the point
+  reads the router's own `proj` output, which is what that softmax consumed and what vLLM's
+  `Gemma4Router.forward` returns outright, leaving the two engines comparable
+  (`facts.ROUTER_LOGITS_SUBMODULE`, plus a `ROUTER_OUTPUTS` row naming slot 0 for what it is).
+  `assert_routing_shapes` also rejects a `router_logits` that is non-negative everywhere and sums to 1
+  per token, which is the family-agnostic version of the same catch. Elements 1 and 2 are the weights
+  and indices the block really routes with, and are read where every other family's are.
 
 In our shipping set only gpt-oss-20b is MoE (32 experts, top-4, no shared expert, every layer
 sparse). Its 3D batched expert weights matter only to code reading MLP weight matrices, and the

{interp_engine-1.3.6 → interp_engine-1.4.0}/interp_engine/arch.py

@@ -86,6 +86,12 @@ class Quirks:
     # no branch on this -- ``mlp_in``/``mlp_out`` tap ``layer.mlp``, which is the whole block
     # including any shared expert -- so it is reported for callers, not consumed here.
     moe_layers: tuple[int, ...] = ()
+    # Whether a sparse layer's routed experts are added *beside* the dense MLP instead of replacing
+    # it (Gemma-4, and so far only Gemma-4). That inverts the sentence above: ``layer.mlp`` is then
+    # one of two branches the block's own forward sums, so tapping it returns half a feed-forward.
+    # Consumed by capture, which refuses ``mlp_out`` on such a layer and points at ``mlp_out_post``.
+    # See :func:`facts.dense_mlp_beside_experts`.
+    dense_mlp_beside_experts: bool = False
     # Each sublayer's output is normalized before being added to the residual (Gemma-2/3/4
     # sandwich norms, OLMo-2/3 post-norms), so the *residual contribution* is a different tensor
     # from the raw submodule output. Enables the ``attn_out_post`` / ``mlp_out_post`` points.
@@ -556,6 +562,10 @@ class ArchSpec:
         # did it while `mlp_in`/`mlp_out` at the same position happily resolved to that same dense
         # module, so the two halves of the MLP vocabulary contradicted each other about one block.
         mlp = self.mlp_projection_holder(layer)
+        # The MLP's *own* router, deliberately, not `facts.moe_router_owner`: a router beside the MLP
+        # rather than inside it means the dense MLP is still there and still has a neuron basis to
+        # capture (Gemma-4 -- see `mlp_is_half_the_feed_forward`), so widening this to the block would
+        # refuse three points that exist. `moe_router` widens; this one must not.
         if self.is_moe_layer(layer) and facts.moe_router_attr(mlp) is not None:
             raise ValueError(
                 f"Layer {layer} of {self.architecture} is a sparse MoE block, so it has no single "
@@ -615,6 +625,17 @@ class ArchSpec:
         """Whether ``layer``'s MLP is a sparse mixture-of-experts block rather than a dense one."""
         return layer in self.quirks.moe_layers
 
+    def mlp_is_half_the_feed_forward(self, layer: int) -> bool:
+        """Whether ``layer.mlp`` is one of two feed-forward branches the block's forward sums.
+
+        True only where a sparse layer keeps its dense MLP *beside* the experts (Gemma-4): both
+        branches read the pre-feedforward residual through norms of their own, and the block adds
+        them. So ``layer.mlp`` is a complete module producing a complete ``d_model`` tensor that is
+        nonetheless not the layer's feed-forward -- which is why this is asked rather than inferred
+        from ``is_moe_layer``, true on families where the MLP *is* the whole block.
+        """
+        return self.is_moe_layer(layer) and self.quirks.dense_mlp_beside_experts
+
     def moe_router(self, layer: int) -> nn.Module:
         """The sparse block's router, whose output is the whole routing decision.
 
@@ -634,12 +655,16 @@ class ArchSpec:
                 "the neuron basis of a dense MLP."
             )
         mlp = self.mlp_module(layer)
-        attr = facts.moe_router_attr(mlp)
-        if attr is None:
+        # Beside the MLP as well as inside it: Gemma-4 hangs the router on the block, and asking only
+        # the MLP reported "no router submodule found" on a checkpoint whose router is one attribute
+        # away. See :func:`facts.moe_router_owner`.
+        owner = facts.moe_router_owner(self.block(layer), mlp)
+        if owner is None:
             raise AttributeError(
-                f"No router submodule found on layer {layer}'s {type(mlp).__name__} "
-                f"({self.architecture}); tried {facts.MOE_ROUTER_ATTRS}"
+                f"No router submodule found on layer {layer}'s {type(mlp).__name__} or on the block "
+                f"itself ({self.architecture}); tried {facts.MOE_ROUTER_ATTRS}"
             )
+        holder, attr = owner
         if "forward" in vars(mlp):
             # A quantizer or kernel loader has replaced the block's forward on the instance, and the
             # replacements route *inline*: transformers' MXFP4 path for gpt-oss calls
@@ -676,7 +701,7 @@ class ArchSpec:
                 "the fused path (for gpt-oss: quantization_config=Mxfp4Config(dequantize=True)) to get the "
                 "eager router, and with it the weights and indices, back."
             )
-        return getattr(mlp, attr)
+        return getattr(holder, attr)
 
     def inline_routing_logits(self, layer: int) -> tuple[nn.Module, str] | None:
         """Address of ``router_logits`` on a block that routes inline, or None if it calls its router.
@@ -696,6 +721,12 @@ class ArchSpec:
 
         Raises on a KV-shared layer, where the value tensor is genuinely produced by a different
         layer rather than merely being hard to find.
+
+        ``None`` also -- and this is why callers must ask :meth:`value_module` *first* rather than
+        reading None as "split a fused QKV" -- on a layer that has no value projection because it uses
+        the key's (:meth:`is_k_eq_v_layer`). That is a third case this return value cannot distinguish,
+        on a family whose :attr:`Quirks.fused_qkv` is false, so a caller reaching here with None on
+        Gemma-4 would go looking for a fused projection that does not exist.
         """
         if self.is_kv_shared_layer(layer):
             raise ValueError(self._kv_shared_refusal(layer, "value", "value projection", "DFA reads it too."))
@@ -953,6 +984,48 @@ class ArchSpec:
         :func:`facts.value_scale`."""
         return facts.value_scale(self.attn_module(layer))
 
+    def is_k_eq_v_layer(self, layer: int) -> bool:
+        """Whether ``layer`` takes its *key* projection's output as the value and has no ``v_proj``.
+
+        Gemma-4's ``attention_k_eq_v``, and only on the layers the modeling code applies it to: the
+        flag is model-wide but ``Gemma4TextAttention`` gates it on ``not is_sliding``, so a sliding
+        layer of the same checkpoint projects its value normally. Reading the flag alone would
+        describe every layer of the 12B, 26B and 31B as having no value projection when half of them
+        do -- and the sliding layers' ``num_key_value_heads`` differs from the global ones', so the
+        two kinds are not interchangeable in any case (:meth:`kv_heads_for_layer`).
+        """
+        if not self.k_eq_v:
+            return False
+        kinds = self.quirks.hybrid_layer_types
+        if not kinds or layer >= len(kinds):
+            # The flag with no layer table is a homogeneous trunk: nothing marks a layer sliding, so
+            # the alternative attention applies throughout.
+            return True
+        return "sliding" not in str(kinds[layer]).lower()
+
+    def value_module(self, layer: int) -> nn.Module | None:
+        """The module whose output is the value attention consumes, when that is not the projection.
+
+        The norm the family runs between projection and attention -- Gemma-4's ``v_norm`` -- or
+        ``None`` on every family where ``value`` is the projection's own output.
+
+        Preferred over :meth:`v_proj` wherever it exists, on *all* of that family's layers rather
+        than only the ones missing a projection. Two reasons, and the first would apply even if the
+        second never came up: the normed tensor is the one the attention pattern multiplies, so it is
+        what ``value`` names and what DFA needs; and vLLM has no separate value projection to compare
+        against (its ``Gemma4Attention`` splits a fused QKV) but runs the same ``v_norm``, so this is
+        the only boundary at which the two engines can be checked against each other at all.
+
+        Raises on a KV-shared layer, where the tensor is genuinely another layer's -- the same
+        refusal :meth:`v_proj` makes, and for the same reason: those layers are built with neither a
+        projection nor a norm.
+        """
+        if self.is_kv_shared_layer(layer):
+            raise ValueError(self._kv_shared_refusal(layer, "value", "value norm", "DFA reads it too."))
+        attn = self.attn_module(layer)
+        attr = facts.value_norm_attr(attn)
+        return getattr(attn, attr) if attr is not None else None
+
     def is_kv_shared_layer(self, layer: int) -> bool:
         """Whether ``layer`` reuses an earlier layer's keys/values and has no k/v projection."""
         return self.first_kv_shared_layer is not None and layer >= self.first_kv_shared_layer
@@ -1127,6 +1200,7 @@ def resolve_arch(model: nn.Module, config: Any) -> ArchSpec:
         qkv_layout=facts.eager_qkv_layout(architecture, config),
         gated_attn_out=gated_attn_out,
         moe_layers=model_facts.moe_layers,
+        dense_mlp_beside_experts=model_facts.dense_mlp_beside_experts,
         sandwich_norms=post_mlp_norm_attr is not None,
         tied_embeddings=tied,
         attn_sinks=bool(hints.get("attn_sinks", False)),

{interp_engine-1.3.6 → interp_engine-1.4.0}/interp_engine/capture.py

@@ -27,7 +27,7 @@ from interp_engine.address import Address, to_address
 from interp_engine.attn_scores import capture_attn_scores
 from interp_engine.dispatch import TokensLike, as_batched_tokens, as_token_ids, refuse
 from interp_engine.facts import text_config
-from interp_engine.hooks import HookManager
+from interp_engine.hooks import HookManager, flat_per_head
 from interp_engine.model import EagerModel
 from interp_engine.points import token_flattened
 from interp_engine.protocol import InterpModel
@@ -386,6 +386,13 @@ def _run_with_cache_eager(
         heads = model.arch.kv_heads_for_layer(key.layer or 0) if key.name.startswith("k_") else model.n_heads
         cache.tensors[key] = _to_token_major(cache.tensors[key], seq, heads)
 
+    for key in {k for k in cache.tensors if k.name == "value"}:
+        # Flat, however the family produced it: a value norm sees the per-head view and a value
+        # projection does not, and `value` is one point. See `hooks.flat_per_head`.
+        layer = key.layer or 0
+        flat, _ = flat_per_head(cache.tensors[key], heads=model.arch.kv_heads_for_layer(layer))
+        cache.tensors[key] = flat
+
     if wants_attn:
         attentions = output.attentions  # one [batch, n_heads, q, k] per softmax-attention layer
         for address in addresses:

{interp_engine-1.3.6 → interp_engine-1.4.0}/interp_engine/facts.py

@@ -1134,9 +1134,46 @@ def moe_router_attr(mlp: Any) -> str | None:
     """Attribute name of the sparse block's router, or ``None`` if this MLP has none.
 
     ``None`` on every dense MLP, including the dense prefix layers of a hybrid MoE model, so a
-    caller must not read it as "this checkpoint is not MoE".
+    caller must not read it as "this checkpoint is not MoE". Also ``None`` on a Gemma-4 sparse layer,
+    whose router is a sibling of the MLP rather than a child of it -- ask :func:`moe_router_owner`,
+    which looks in both places, unless you specifically mean the MLP's own.
     """
-    return next((name for name in MOE_ROUTER_ATTRS if hasattr(mlp, name)), None)
+    return next((name for name in MOE_ROUTER_ATTRS if _is_module(getattr(mlp, name, None))), None)
+
+
+def _is_module(candidate: Any) -> bool:
+    """Whether this attribute is a callable submodule rather than a flag or a tensor.
+
+    Duck-typed because this module holds no torch dependency (see the header). Presence alone is not
+    enough once the *block* is searched for a router as well as the MLP: a decoder layer can hold a
+    plain attribute under one of these names, and ``hasattr`` would hand a bool back to a caller
+    about to install a forward hook on it.
+    """
+    return callable(candidate) and hasattr(candidate, "forward")
+
+
+def moe_router_owner(layer: Any, mlp: Any) -> tuple[Any, str] | None:
+    """The module that holds the sparse block's router, and the attribute name -- or ``None``.
+
+    The MLP first, because that is where every other MoE family puts it: the router is a child of
+    the ``MixtralSparseMoeBlock`` / ``Qwen3MoeSparseMoeBlock`` / ``GptOssMLP`` that consumes it, so
+    on those the block is never searched and nothing about them changes.
+
+    Gemma-4 is the family that needs the second look. Its router is a sibling of ``layer.mlp``
+    (``layer.router``, beside ``layer.experts``), for the same reason its MLP is only half the
+    feed-forward: the routed branch is assembled by the *block's* forward, so the block owns both
+    halves and neither is inside the other. Asking only the MLP left all three routing points
+    unresolvable on the 26B, reported as "no router submodule found" -- a lookup failure that reads
+    like a checkpoint without a router rather than a router one level up.
+
+    ``None`` on a dense layer either way, which a caller must not read as "this checkpoint is not
+    MoE": the dense prefix of a hybrid trunk answers ``None`` too.
+    """
+    if (attr := moe_router_attr(mlp)) is not None:
+        return mlp, attr
+    if layer is not None and (attr := moe_router_attr(layer)) is not None:
+        return layer, attr
+    return None
 
 
 #: Sparse-block ``forward`` replacements that route *inline* -- computing the logits with the router's
@@ -1188,16 +1225,59 @@ def routing_convention(architecture: str) -> str | None:
 #: outside: element 0 is ``[tokens, k]`` where the default reading expects ``[tokens, n_experts]``,
 #: which is a plausible tensor under either name, and it took a cross-engine width mismatch to catch.
 #: :func:`assert_routing_shapes` is the check that makes the next one of these loud instead.
+#: A name for an element that is not one of the three points -- it is a real tensor a router returns
+#: which no canonical point means, and it exists so that a layout can say "the logits are *not* here"
+#: rather than leaving the slot labelled with the point that would then be read out of it.
+_NOT_A_POINT_PROBS = "router_probabilities"
 _DEFAULT_ROUTER_OUTPUT: tuple[str, ...] = ("router_logits", "expert_weights", "expert_indices")
 ROUTER_OUTPUTS: dict[str, tuple[str, ...]] = {
     "GraniteMoeForCausalLM": ("expert_indices", "expert_weights", "router_logits"),
     "GraniteMoeSharedForCausalLM": ("expert_indices", "expert_weights", "router_logits"),
     "GraniteMoeHybridForCausalLM": ("expert_indices", "expert_weights", "router_logits"),
+    # Gemma-4 returns the three tensors in the default *order* while element 0 is not the default
+    # tensor: `Gemma4TextRouter.forward` softmaxes over all 128 experts and returns those
+    # probabilities there. The width check cannot catch it -- probabilities over the bank are exactly
+    # as wide as logits over the bank -- and neither can a caller downstream, since both are
+    # per-token float rows whose entries are plausible either way. The block itself discards element
+    # 0, so it is not even a tensor the model uses. Its logits are one module deeper
+    # (:data:`ROUTER_LOGITS_SUBMODULE`), which is where `router_logits` resolves on this family.
+    "Gemma4ForConditionalGeneration": (_NOT_A_POINT_PROBS, "expert_weights", "expert_indices"),
+    "Gemma4UnifiedForConditionalGeneration": (_NOT_A_POINT_PROBS, "expert_weights", "expert_indices"),
+    "Gemma4ForCausalLM": (_NOT_A_POINT_PROBS, "expert_weights", "expert_indices"),
 }
 
 
+#: Where a family's *pre-softmax* logits are, for the families whose router module does not return
+#: them: the attribute, on the router, of the projection that produces them.
+#:
+#: Gemma-4 is the family (see :data:`ROUTER_OUTPUTS`). ``Gemma4TextRouter`` norms, scales and then
+#: projects, and it is that ``proj`` whose output the softmax consumes -- so this is a *read* of the
+#: tensor the routing decision was made from, not a recomputation of it, and it is bit-identical to
+#: what vLLM's own ``Gemma4Router.forward`` returns, which makes the two engines comparable at this
+#: point rather than only within one.
+#:
+#: Keyed by architecture and deliberately small, like every other table here. A family whose router
+#: returns its logits needs no entry, and guessing an entry for one would address a point at a
+#: submodule whose output has never been checked against what the family routes on.
+ROUTER_LOGITS_SUBMODULE: dict[str, str] = {
+    "Gemma4ForConditionalGeneration": "proj",
+    "Gemma4UnifiedForConditionalGeneration": "proj",
+    "Gemma4ForCausalLM": "proj",
+}
+
+
+def router_logits_submodule(architecture: str) -> str | None:
+    """The router submodule holding this family's pre-softmax logits, or None if the router returns them."""
+    return ROUTER_LOGITS_SUBMODULE.get(architecture)
+
+
 def router_output_index(architecture: str, point: str) -> int:
-    """Which element of this family's router output tuple carries ``point``."""
+    """Which element of this family's router output tuple carries ``point``.
+
+    Raises where the tuple carries no such tensor, which is a real answer and not a gap in the table:
+    a family can return probabilities where the default returns logits, and the caller then has to go
+    somewhere else for them (:func:`router_logits_submodule`) rather than read the slot.
+    """
     layout = ROUTER_OUTPUTS.get(architecture, _DEFAULT_ROUTER_OUTPUT)
     if point not in layout:
         raise ValueError(f"{architecture}'s router output carries {layout}, not {point!r}")
@@ -1221,6 +1301,16 @@ def assert_routing_shapes(point: str, tensor: Any, *, architecture: str, n_exper
             f"Its router's output tuple is not {_DEFAULT_ROUTER_OUTPUT}; register the real order in "
             "`facts.ROUTER_OUTPUTS`."
         )
+    if point == "router_logits" and _looks_like_a_distribution(tensor):
+        raise ValueError(
+            f"{architecture}'s captured 'router_logits' is non-negative everywhere and sums to 1 per "
+            "token, so it is a softmax over the expert bank rather than the logits that went into one. "
+            "A family can return the probabilities where the default tuple returns logits (Gemma-4 "
+            "does, and the width check above cannot see it, because a distribution over the bank is "
+            "exactly as wide as the logits over it). Point 'router_logits' at the router's own "
+            "projection with `facts.ROUTER_LOGITS_SUBMODULE` and mark the slot in "
+            "`facts.ROUTER_OUTPUTS`."
+        )
     if point == "expert_indices" and tensor.dtype.is_floating_point:
         raise ValueError(
             f"{architecture}'s captured 'expert_indices' is {tensor.dtype}, and a selection is integers. "
@@ -1234,6 +1324,28 @@ def assert_routing_shapes(point: str, tensor: Any, *, architecture: str, n_exper
         )
 
 
+def _looks_like_a_distribution(tensor: Any) -> bool:
+    """Whether every row of ``tensor``'s last axis is non-negative and sums to one.
+
+    The tell that separates a softmax's output from its input, and the only property that does: the
+    two are the same shape, the same dtype and the same width, and both are per-token float rows whose
+    entries look like scores. Both halves are needed -- logits happen to sum near 1 sometimes, and a
+    non-negative row need not be normalized -- and together they are a thing real logits over 128
+    experts do not do.
+
+    Written with no torch dependency (see the header) and defensively: anything that does not answer
+    these questions is not a distribution as far as this is concerned, because the caller's job is a
+    shape assertion and it must not be the thing that raises.
+    """
+    try:
+        if tensor.dtype.is_floating_point is False:
+            return False
+        as_float = tensor.detach().float()
+        return bool((as_float >= 0).all()) and bool(((as_float.sum(-1) - 1.0).abs() < 1e-3).all())
+    except (AttributeError, RuntimeError, TypeError):
+        return False
+
+
 def inline_routing_logits_index(mlp: Any) -> int | None:
     """Where this block's own output holds the router logits, when its forward routes inline.
 
@@ -1286,6 +1398,24 @@ def dense_mlp_beside_experts(cfg: Any) -> bool:
     return bool(config_attr(cfg, "enable_moe_block", False))
 
 
+#: The attribute both engines' Gemma-4 decoder layers set from ``enable_moe_block``, marking a block
+#: that hangs the routed branch beside its dense MLP. A *layer* flag rather than a config read, which
+#: is what the vLLM worker needs: it holds a block with no layer index to ask
+#: :func:`is_moe_layer` about, and the two trees set this from the same config field, so asking the
+#: block cannot disagree with what the eager backend concluded from the config.
+DENSE_MLP_BESIDE_EXPERTS_FLAG = "enable_moe_block"
+
+
+def experts_beside_this_layers_mlp(layer: Any) -> bool:
+    """Whether *this* block sums a routed branch with its dense MLP's output.
+
+    False on a Gemma-4 layer that is dense, so this is per block rather than per checkpoint -- and
+    false everywhere outside the family, where a sparse layer has no dense MLP to sum with. See
+    :func:`dense_mlp_beside_experts` for what the arrangement is and why it needs saying.
+    """
+    return bool(getattr(layer, DENSE_MLP_BESIDE_EXPERTS_FLAG, False))
+
+
 def is_moe_layer(cfg: Any, layer: int) -> bool:
     """Whether ``layer``'s ``mlp`` is a sparse MoE block rather than a dense MLP.
 
@@ -1432,6 +1562,41 @@ def value_scale(attn_module: Any) -> float:
     return 1.0
 
 
+#: A norm applied to the value vectors *between* the projection and attention, so the tensor
+#: attention consumes is this module's output rather than the projection's. Gemma-4 is the family:
+#: ``Gemma4TextAttention.forward`` runs ``value_states = self.v_norm(value_states)`` on every layer
+#: that projects its own KV, and vLLM's ``Gemma4Attention`` runs the same line on the V slice of its
+#: fused QKV -- so the two engines agree here in a way they cannot at the projection, which vLLM does
+#: not have separately.
+#:
+#: The same idea as :data:`ATTN_VALUE_SCALE_ATTRS`, one step further: that corrects a scalar the
+#: forward applies after the projection, and this names a module that does. Kept as a module rather
+#: than folded into the scale because an RMS norm is per token, not a constant, so no factor
+#: reproduces it.
+#:
+#: Exact names, and only the ones verified against a family's forward. A ``v_norm`` that a forward
+#: does *not* apply to the value would make this the wrong tensor -- silently, since it is the right
+#: shape -- which is the same trap ``use_qk_norm=True`` beside an ``nn.Identity`` sets.
+ATTN_VALUE_NORM_ATTRS: tuple[str, ...] = ("v_norm",)
+
+
+def value_norm_attr(attn_module: Any) -> str | None:
+    """Attribute name of the norm this layer's value passes through, or ``None`` if there is none.
+
+    ``None`` on every family but Gemma-4, where ``value`` is the projection's own output and nothing
+    about the resolution changes.
+
+    Two things this buys on Gemma-4, and the second is why it is a module question rather than a
+    correction applied afterwards. The value it names is the one attention multiplies the pattern by,
+    where ``v_proj``'s output is a norm short of it -- true on *every* layer of the family, including
+    the sliding ones that do have a ``v_proj``. And on a ``attention_k_eq_v`` layer there is no
+    ``v_proj`` at all: the forward passes the *key* projection's output to ``v_norm``, so this module
+    is the only boundary the value crosses, and asking for the projection there is a question with no
+    answer (see :meth:`interp_engine.arch.ArchSpec.is_k_eq_v_layer`).
+    """
+    return next((name for name in ATTN_VALUE_NORM_ATTRS if _is_module(getattr(attn_module, name, None))), None)
+
+
 def value_head_dim(cfg: Any, head_dim: int) -> int:
     """The width of one *value* head, which is not always the query/key head's width.
 

{interp_engine-1.3.6 → interp_engine-1.4.0}/interp_engine/hooks.py

@@ -46,6 +46,35 @@ def extract_hidden(output: object, index: int = 0) -> torch.Tensor:
     return output  # type: ignore[return-value]
 
 
+def flat_per_head(tensor: torch.Tensor, *, heads: int) -> tuple[torch.Tensor, torch.Size | None]:
+    """A capture that arrived with a head axis, as the flat ``[batch, pos, heads * head_dim]`` a point
+    declares -- with the shape to restore, for a hook that has to write back.
+
+    `value` needs this and `z` does not, for a reason that is about modules rather than about heads.
+    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 `run_with_cache` 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.
+
+    The head count is checked rather than assumed, because the other 4-D layout a norm can produce is
+    `[batch, heads, pos, head_dim]` (transformers norms after the head transpose on ten families), and
+    flattening *that* yields a tensor whose second axis is the head count. Every reader would take it
+    for the sequence, `[batch, seq, width]` would still describe it, and nothing downstream could see
+    it was wrong.
+    """
+    if tensor.ndim != 4:
+        return tensor, None
+    if tensor.shape[-2] != heads:
+        raise ValueError(
+            f"Expected a per-head capture with {heads} heads on its second-to-last axis, got "
+            f"{tuple(tensor.shape)}. A head-major layout ([batch, heads, pos, head_dim]) is the "
+            "likely reason, and flattening it would put the head count where the sequence belongs."
+        )
+    return tensor.flatten(-2, -1), tensor.shape[-2:]
+
+
 def parse_point(point: str) -> tuple[str, int]:
     """Split a hook side into ``(side, tuple index)``: ``"output"`` -> ``("output", 0)``,
     ``"output:2"`` -> ``("output", 2)``.