interp-engine 1.3.3__tar.gz → 1.3.5__tar.gz

{interp_engine-1.3.3 → interp_engine-1.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.3.3
+Version: 1.3.5
 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.3.3 → interp_engine-1.3.5}/docs/ARCHITECTURE_QUIRKS.md

@@ -277,9 +277,26 @@ is therefore asked per side (`facts.has_qk_norm(attn, "q")`), because the pair-w
 "no" for such a layer and refusing the query norm on the key norm's absence declines a tensor the model
 plainly computes.
 
-Related config fields on this family that nothing reads yet: `attention_k_eq_v` (with
-`num_global_key_value_heads`) makes `v_proj` **None** on non-sliding layers and sets `value = key`;
-E2B and E4B both ship with it off.
+- **`kv_heads_for_layer(layer)`** — the kv-head count moves with the head width, and in the opposite
+  direction: 4 against 16 on the 31B, 2 against 8 on the 26B, 1 against 8 on the 12B. transformers
+  5.15 states this per layer (`per_layer_config`); below that the family spells it
+  `num_global_key_value_heads`, which `Gemma4TextAttention` applies to a layer only when
+  `attention_k_eq_v` is set **and** the layer is not sliding — so `ArchSpec` carries the flag beside
+  the count rather than reading the count on its own.
+
+`attention_k_eq_v` is on for the 26B, 31B and 12B, and off for E2B and E4B. It is a second, separate
+way for a layer to have no `v_proj`, and it does not overlap with KV sharing: those three SKUs set
+`num_kv_shared_layers: 0`, so **every** one of their `full_attention` layers is built with
+`v_proj = None` and takes the key projection's output as the value instead — `v_norm(k_proj(h))`
+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.
 
 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`:
@@ -408,9 +425,20 @@ n_experts]` tensor that is plausible rather than obviously wrong.
   Mistral-4, dots1, GLM-4.5) or **every k-th layer sparse with a dense opt-out list** (`(layer + 1)
 % decoder_sparse_step == 0 and layer not in mlp_only_layers`; Qwen2/3-MoE, Qwen3-Next,
   Qwen3-VL/Omni). Expert counts have four live spellings (`num_local_experts`, `num_experts`,
-  `n_routed_experts`, `moe_num_experts`), all read. A third idiom, an explicit `mlp_layer_types`
-  pattern, wins over both — see the spelling warning in [block
+  `n_routed_experts`, `moe_num_experts`), all read, and the top-k has four of its own
+  (`num_experts_per_tok`, `experts_per_token`, `moe_topk`, `top_k_experts`). A third idiom, an
+  explicit `mlp_layer_types` pattern, wins over both — see the spelling warning in [block
   types](#block-types-classify-them-do-not-pattern-match-them).
+- **Gemma-4 breaks the first sentence of this section.** Its sparse layers do not swap the MLP for an
+  expert bank: `Gemma4TextDecoderLayer` builds `self.mlp` on every layer and, where
+  `enable_moe_block` is set, hangs `self.router` and `self.experts` **beside** it on the block,
+  summing the two branches (`post_ffn_norm_1(mlp(x)) + post_ffn_norm_2(experts(x))`, both reading the
+  pre-feedforward residual). Three consequences, and `facts.dense_mlp_beside_experts` is the fact they
+  hang off. A sparse layer keeps a real neuron basis, so `mlp_pre`/`mlp_act` must stay served there —
+  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.
 
 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
@@ -452,7 +480,22 @@ router_logits)`. `facts.ROUTER_OUTPUTS` holds the exceptions; `facts.assert_rout
 
 The MLP's neuron basis (`mlp_pre`, `mlp_pre_linear`, `mlp_act`) is refused on a sparse layer, since
 the projections live on the experts — often as one fused 3-D parameter per bank, with no per-expert
-module at all.
+module at all. Gemma-4 is the exception noted above, and the refusal already lets it through for the
+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.
 
 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.3 → interp_engine-1.3.5}/docs/INTERNALS.md

@@ -66,6 +66,12 @@ files and why they are where they are.
   - `LD_LIBRARY_PATH`) before the first CUDA call, instead of failing ten frames deep in
     `torch.cuda._lazy_init`. Lives here because every app on the engine inherits the same CUDA
     floor — the `[vllm]` wheels link `libcudart.so.13` directly.
+- `notebook_stdout.py` — `ensure_stdout_descriptor`: gives a notebook kernel's `sys.stdout` the file
+  descriptor vLLM's engine start needs. vLLM silences C-level output by dup'ing over one and forks
+  its EngineCore child, so under ipykernel — whose stream writes to a socket and has none — the
+  child dies before it loads anything, and the caller is told to see a root cause that is in another
+  process. Called from `_ensure_engine`; a no-op in a script, a server, or a kernel that captured
+  the real descriptors itself.
 
 ## Correctness
 

{interp_engine-1.3.3 → interp_engine-1.3.5}/interp_engine/arch.py

@@ -324,6 +324,11 @@ class ArchSpec:
     # elsewhere; use :meth:`head_dim_for_layer` / :meth:`kv_heads_for_layer`.
     per_layer_head_dim: tuple[int, ...] = ()
     per_layer_kv_heads: tuple[int, ...] = ()
+    # Gemma-4's ``num_global_key_value_heads`` and ``attention_k_eq_v``: the older spelling of the
+    # per-layer kv-head count, and the flag that decides whether it applies. See
+    # :func:`facts.kv_heads_for_layer`.
+    global_kv_heads: int | None = None
+    k_eq_v: bool = False
     # The width of one *value* head, where the family makes it differ from the q/k head (MiMo-V2,
     # DeepSeek). 0 means "same as the q/k head"; use :meth:`value_head_dim_for_layer`.
     v_head_dim: int = 0
@@ -914,12 +919,20 @@ class ArchSpec:
     def kv_heads_for_layer(self, layer: int) -> int:
         """How many key/value heads ``layer`` attends with.
 
-        :attr:`n_kv_heads` for every family but Gemma-4-31B, whose full-attention layers carry 4 where
-        its sliding ones carry 16. Prefer this anywhere ``k``, ``value`` or a k-side norm is reshaped
-        per head: the model-wide number divides cleanly into the other layers' widths, so getting it
-        wrong scrambles heads rather than raising. See :func:`facts.kv_heads_for_layer`.
+        :attr:`n_kv_heads` for every family but Gemma-4, whose full-attention layers carry 4 where its
+        sliding ones carry 16 on the 31B, and 2 against 8 on the 26B. Prefer this anywhere ``k``,
+        ``value`` or a k-side norm is reshaped per head: the model-wide number divides cleanly into
+        the other layers' widths, so getting it wrong scrambles heads rather than raising. See
+        :func:`facts.kv_heads_for_layer`.
         """
-        return facts.kv_heads_for_layer(self.n_kv_heads, layer, self.per_layer_kv_heads)
+        return facts.kv_heads_for_layer(
+            self.n_kv_heads,
+            layer,
+            self.per_layer_kv_heads,
+            self.global_kv_heads,
+            self.quirks.hybrid_layer_types,
+            self.k_eq_v,
+        )
 
     def value_head_dim_for_layer(self, layer: int) -> int:
         """``layer``'s *value* head width -- what ``value`` and ``z`` are per head.
@@ -1160,6 +1173,8 @@ def resolve_arch(model: nn.Module, config: Any) -> ArchSpec:
         global_head_dim=model_facts.global_head_dim,
         per_layer_head_dim=model_facts.per_layer_head_dim,
         per_layer_kv_heads=model_facts.per_layer_kv_heads,
+        global_kv_heads=model_facts.global_kv_heads,
+        k_eq_v=model_facts.k_eq_v,
         v_head_dim=model_facts.v_head_dim,
         first_kv_shared_layer=model_facts.first_kv_shared_layer,
         layer_slots=layer_slots,

{interp_engine-1.3.3 → interp_engine-1.3.5}/interp_engine/facts.py

@@ -1072,7 +1072,15 @@ def logit_multiplier(cfg: Any) -> tuple[float | None, str]:
 
 # Four spellings of the routed-expert count, all in current use.
 _N_EXPERTS_FIELDS: tuple[str, ...] = ("num_local_experts", "num_experts", "n_routed_experts", "moe_num_experts")
-_EXPERTS_PER_TOKEN_FIELDS: tuple[str, ...] = ("num_experts_per_tok", "experts_per_token", "moe_topk")
+# ``top_k_experts`` is Gemma-4's. Missing it did not merely lose a number: :func:`assert_routing_shapes`
+# guards its width check with ``and top_k``, so a zero here switched off the check that catches a
+# router tuple read in the wrong order -- on the newest MoE family, and silently.
+_EXPERTS_PER_TOKEN_FIELDS: tuple[str, ...] = (
+    "num_experts_per_tok",
+    "experts_per_token",
+    "moe_topk",
+    "top_k_experts",
+)
 _N_SHARED_EXPERTS_FIELDS: tuple[str, ...] = ("n_shared_experts", "num_shared_experts")
 
 
@@ -1228,6 +1236,27 @@ def block_types_name_the_feed_forward(cfg: Any) -> bool:
     return any(str(kind).lower() in SPARSE_MLP_LAYER_KINDS | DENSE_MLP_LAYER_KINDS for kind in kinds)
 
 
+def dense_mlp_beside_experts(cfg: Any) -> bool:
+    """Whether a sparse layer *also* runs a dense MLP, whose output is added to the experts'.
+
+    Every other MoE family replaces the feed-forward with an expert bank, so a sparse layer has no
+    dense MLP and no neuron basis. Gemma-4 does not: ``Gemma4TextDecoderLayer`` builds ``self.mlp`` on
+    every layer, and where ``enable_moe_block`` is set it adds a routed branch *beside* it, combining
+    them as ``hidden_states_1 + hidden_states_2`` -- two separately normed branches that both read the
+    pre-feedforward residual.
+
+    Two consequences, and they pull in opposite directions, which is why this is its own fact. The
+    parameter count of a sparse layer includes the dense MLP, so subtracting it (as the arithmetic
+    does everywhere else) undercounts. And the dense neuron basis is real on a sparse layer here, so
+    refusing ``mlp_act`` on one would be wrong.
+
+    Not the same thing as a shared expert: that lives inside the MoE block and is sized by
+    ``moe_intermediate_size``, while this branch is the ordinary ``intermediate_size`` MLP that the
+    checkpoint's dense siblings also carry.
+    """
+    return bool(config_attr(cfg, "enable_moe_block", False))
+
+
 def is_moe_layer(cfg: Any, layer: int) -> bool:
     """Whether ``layer``'s ``mlp`` is a sparse MoE block rather than a dense MLP.
 
@@ -1322,18 +1351,36 @@ def head_dim_for_layer(
     return head_dim if "sliding" in str(layer_types[layer]).lower() else global_head_dim
 
 
-def kv_heads_for_layer(n_kv_heads: int, layer: int, per_layer: tuple[int, ...] = ()) -> int:
+def kv_heads_for_layer(
+    n_kv_heads: int,
+    layer: int,
+    per_layer: tuple[int, ...] = (),
+    global_kv_heads: int | None = None,
+    layer_types: tuple[str, ...] | None = None,
+    k_eq_v: bool = False,
+) -> int:
     """How many key/value heads ``layer`` attends with.
 
-    One number for the whole model on every family but Gemma-4-31B, whose full-attention layers carry
-    4 where its sliding ones carry 16 -- and whose config says so only through ``per_layer_config``,
-    so there is no older spelling to fall back to and ``per_layer`` empty means the model-wide value
-    is the answer. Getting it wrong is not a shape error (see :func:`effective_kv_heads`): the
-    reshape succeeds into a head count the layer does not have.
+    One number for the whole model on every family but Gemma-4, whose full-attention layers carry 4
+    where its sliding ones carry 16 on the 31B (2 against 8 on the 26B, 1 against 8 on the 12B).
+    Getting it wrong is not a shape error (see :func:`effective_kv_heads`): the reshape succeeds into
+    a head count the layer does not have.
+
+    Two spellings, as with :func:`head_dim_for_layer`. ``per_layer`` is the table a heterogeneous
+    config states outright (transformers >= 5.15) and wins where it exists. ``global_kv_heads`` is the
+    older ``num_global_key_value_heads``, and it applies only where the modeling code applies it:
+    ``Gemma4TextAttention`` takes it when ``attention_k_eq_v and not is_sliding``, and the model-wide
+    count otherwise. The ``k_eq_v`` gate is not decoration -- E2B and E4B set that flag false and
+    carry ``num_global_key_value_heads: null``, so reading the field unconditionally would be wrong
+    the moment a checkpoint states one without switching the flag on.
     """
     if per_layer and layer < len(per_layer):
         return per_layer[layer]
-    return n_kv_heads
+    if not (k_eq_v and global_kv_heads):
+        return n_kv_heads
+    if not layer_types or layer >= len(layer_types):
+        return n_kv_heads
+    return n_kv_heads if "sliding" in str(layer_types[layer]).lower() else global_kv_heads
 
 
 # A scalar the attention module multiplies its value vectors by *after* the projection, so the tensor
@@ -1747,6 +1794,15 @@ class ModelFacts:
     # prefer these; see :func:`per_layer_ints`.
     per_layer_head_dim: tuple[int, ...] = ()
     per_layer_kv_heads: tuple[int, ...] = ()
+    # Gemma-4's kv-head count for its ``full_attention`` layers, which it states as
+    # ``num_global_key_value_heads`` and applies only when ``attention_k_eq_v`` is on. The older
+    # spelling of what ``per_layer_kv_heads`` carries on transformers >= 5.15; ask
+    # :meth:`kv_heads_for_layer`.
+    global_kv_heads: int | None = None
+    # Gemma-4's ``attention_k_eq_v``: on its ``full_attention`` layers the value tensor *is* the key
+    # projection's output (differently normed and un-RoPE'd), and those layers are built with
+    # ``v_proj = None``. True on 26B/31B/12B, false on E2B/E4B.
+    k_eq_v: bool = False
     # The width of one value head, which differs from ``head_dim`` on MiMo-V2 and the DeepSeek MLA
     # families. See :func:`value_head_dim`; ``value`` and ``z`` are this wide per head, not ``head_dim``.
     v_head_dim: int = 0
@@ -1759,6 +1815,9 @@ class ModelFacts:
     n_experts: int = 0
     experts_per_token: int = 0
     n_shared_experts: int = 0
+    # Gemma-4: a sparse layer keeps its dense MLP and adds the routed branch beside it, so the two
+    # coexist rather than the experts replacing the MLP. See :func:`dense_mlp_beside_experts`.
+    dense_mlp_beside_experts: bool = False
     # Layers whose ``mlp`` is a sparse block. Precomputed here because the branch needs config
     # fields that the vLLM client does not carry across the process boundary.
     moe_layers: tuple[int, ...] = ()
@@ -1842,8 +1901,10 @@ class ModelFacts:
         return head_dim_for_layer(self.head_dim, self.global_head_dim, self.layer_types, layer, self.per_layer_head_dim)
 
     def kv_heads_for_layer(self, layer: int) -> int:
-        """``layer``'s kv-head count. Prefer this to :attr:`n_kv_heads`, which is wrong on Gemma-4-31B."""
-        return kv_heads_for_layer(self.n_kv_heads, layer, self.per_layer_kv_heads)
+        """``layer``'s kv-head count. Prefer this to :attr:`n_kv_heads`, which is wrong on Gemma-4."""
+        return kv_heads_for_layer(
+            self.n_kv_heads, layer, self.per_layer_kv_heads, self.global_kv_heads, self.layer_types, self.k_eq_v
+        )
 
     def value_head_dim_for_layer(self, layer: int) -> int:
         """``layer``'s *value* head width, for reshaping ``value`` and ``z``. See :func:`value_head_dim`.
@@ -2198,11 +2259,14 @@ def resolve_facts(config: Any, *, n_layers_fallback: int | None = None) -> Model
         global_head_dim=_first_int(cfg, ("global_head_dim",)) or None,
         per_layer_head_dim=per_layer_ints(cfg, "head_dim", n_layers),
         per_layer_kv_heads=per_layer_ints(cfg, "num_key_value_heads", n_layers),
+        global_kv_heads=_first_int(cfg, ("num_global_key_value_heads",)) or None,
+        k_eq_v=bool(config_attr(cfg, "attention_k_eq_v", False)),
         v_head_dim=value_head_dim(cfg, head_dim),
         first_kv_shared_layer=first_kv_shared_layer(cfg, n_layers),
         n_experts=n_experts(cfg),
         experts_per_token=_first_int(cfg, _EXPERTS_PER_TOKEN_FIELDS) or 0,
         n_shared_experts=_first_int(cfg, _N_SHARED_EXPERTS_FIELDS) or 0,
+        dense_mlp_beside_experts=dense_mlp_beside_experts(cfg),
         moe_layers=tuple(layer for layer in range(n_layers) if is_moe_layer(cfg, layer)),
         logit_multiplier=multiplier,
         logit_multiplier_source=multiplier_source,

interp_engine-1.3.5/interp_engine/notebook_stdout.py

@@ -0,0 +1,90 @@
+"""Give a notebook kernel's stdout the file descriptor vLLM's engine start needs.
+
+vLLM silences C-level output around ``torch.distributed.new_group`` by dup'ing over a
+descriptor (``suppress_stdout`` in ``vllm/utils/system_utils.py``), and its EngineCore
+child is forked, so the child inherits whichever ``sys.stdout`` the parent had. Under
+Jupyter that is ``ipykernel.iostream.OutStream``, which writes through a ZMQ socket and
+raises ``io.UnsupportedOperation`` from ``fileno()``. The child then dies before it loads
+anything, and what reaches the caller is vLLM's ``Engine core initialization failed. See
+root cause above.`` -- above being a traceback in another process, about a descriptor.
+
+A descriptor is what that stream is supposed to hand back: ipykernel answers ``fileno()``
+with ``_original_stdstream_copy``, a dup of the real stdout it took over, whenever the
+kernel was started capturing the low-level ones. This supplies the one a kernel that was
+not started that way -- Colab's -- did not keep. Colab is the case that matters, because
+this repo's own notebook templates run there.
+
+The descriptor stays on the stream rather than being restored after the engine is built.
+The fork is not the only caller: vLLM suppresses stdout the same way around its stateless
+process groups, from whichever process reaches one.
+
+Forcing ``spawn`` would also give the child a stdout of its own, and is what
+``VLLM_WORKER_MULTIPROC_METHOD`` is for. It is not what this does, because a spawned
+child writes to the kernel's real stdout rather than to the cell -- which takes the engine
+logs, and the root cause behind that message, out of the notebook with it.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+
+logger = logging.getLogger(__name__)
+
+
+def _answers_fileno(stream: object) -> bool:
+    """Whether ``stream.fileno()`` returns rather than raises.
+
+    ``io.UnsupportedOperation`` is both an ``OSError`` and a ``ValueError``, which is what
+    a kernel stream raises; a detached or closed one raises ``ValueError`` on its own.
+    """
+    try:
+        stream.fileno()  # type: ignore[attr-defined]
+    except (AttributeError, OSError, ValueError):
+        return False
+    return True
+
+
+def _process_stdout_fd() -> int | None:
+    """The descriptor this process's stdout is, or None when it has none.
+
+    ``sys.__stdout__`` is the stream Python opened at startup, which Jupyter leaves alone
+    when it replaces ``sys.stdout``, so it is the first place to ask. Fd 1 covers an
+    embedding that replaced both. Either is confirmed with ``fstat`` before being handed
+    out, since a daemonized process may have closed it.
+    """
+    original = sys.__stdout__
+    fd = original.fileno() if original is not None and _answers_fileno(original) else 1
+    try:
+        os.fstat(fd)
+    except OSError:
+        return None
+    return fd
+
+
+def ensure_stdout_descriptor() -> bool:
+    """Give ``sys.stdout`` a ``fileno()`` if it has none, and say whether it needed one.
+
+    Idempotent, and a no-op outside a notebook. A stream that already answers is left
+    alone, and so is a process whose stdout is closed: there is no descriptor to offer
+    there, and vLLM's refusal is then the honest outcome rather than one to work around.
+    """
+    stream = sys.stdout
+    if stream is None or _answers_fileno(stream):
+        return False
+    fd = _process_stdout_fd()
+    if fd is None:
+        return False
+    try:
+        stream.fileno = lambda: fd  # type: ignore[method-assign]
+    except (AttributeError, TypeError):
+        # A stream that takes no new attribute, which a C-level or slotted one does not.
+        # Nothing is lost: this is the refusal vLLM was about to raise anyway.
+        return False
+    logger.info(
+        "%s has no file descriptor, which vLLM's engine start requires; answering fileno() with fd %d.",
+        type(stream).__name__,
+        fd,
+    )
+    return True

{interp_engine-1.3.3 → interp_engine-1.3.5}/interp_engine/vllm_backend.py

@@ -30,6 +30,7 @@ from interp_engine import facts
 from interp_engine._loop import refuse_foreign_loop
 from interp_engine.address import Address, format_address, to_address
 from interp_engine.autograd_support import GradSupport, vllm_grad_support
+from interp_engine.notebook_stdout import ensure_stdout_descriptor
 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
@@ -314,6 +315,11 @@ def read_attn_dims(hf_model_id: str, trust_remote_code: bool = True) -> dict[str
         # Empty on a config that describes one shape for the whole model.
         "per_layer_head_dim": model_facts.per_layer_head_dim,
         "per_layer_kv_heads": model_facts.per_layer_kv_heads,
+        # The older spelling of that kv-head count (`num_global_key_value_heads`), and the flag
+        # Gemma-4 gates it on. Carried so a transformers below 5.15 -- which states no per-layer
+        # table -- still gets the wide layers' count right rather than the sliding one's.
+        "global_kv_heads": model_facts.global_kv_heads,
+        "k_eq_v": model_facts.k_eq_v,
         # From here on a layer reuses an earlier layer's keys/values and has no v_proj to hook, so
         # `value`/DFA is unavailable there (Gemma-4). None when every layer projects its own.
         "first_kv_shared_layer": model_facts.first_kv_shared_layer,
@@ -406,7 +412,12 @@ def kv_heads_for_layer(dims: dict[str, Any], layer: int) -> int:
     disagreement between the two is worth saying out loud.
     """
     return facts.kv_heads_for_layer(
-        int(dims.get("n_kv_heads") or 0), layer, tuple(dims.get("per_layer_kv_heads") or ())
+        int(dims.get("n_kv_heads") or 0),
+        layer,
+        tuple(dims.get("per_layer_kv_heads") or ()),
+        dims.get("global_kv_heads"),
+        tuple(dims.get("layer_types") or ()),
+        bool(dims.get("k_eq_v")),
     )
 
 
@@ -550,7 +561,10 @@ def recompute_attn_from_payloads(payloads, layers, dims, tensor_parallel_size: i
         # one is not a claim about this layer -- it disagrees with Gemma-4's wide layers by design, and
         # on an MLA or tensor-parallel capture it describes something other than the width in hand --
         # so passing it as `expected` would turn a working recompute into a raise.
-        stated_kv = kv_heads_for_layer(dims, int(layer)) if dims.get("per_layer_kv_heads") else None
+        states_per_layer_kv = bool(dims.get("per_layer_kv_heads")) or bool(
+            dims.get("k_eq_v") and dims.get("global_kv_heads")
+        )
+        stated_kv = kv_heads_for_layer(dims, int(layer)) if states_per_layer_kv else None
         n_kv_heads = _heads_in(k, head_dim, "k", int(layer), expected=stated_kv)
         scores = recompute_attn_scores(
             q,
@@ -1213,6 +1227,9 @@ class VLLMModel:
                         tuple(getattr(self, "_static_reads", ())),
                         tuple(getattr(self, "_static_writes", ())),
                     )
+                    # The line below forks a child that suppresses stdout by descriptor, which a
+                    # notebook kernel's stdout does not have. See `notebook_stdout`.
+                    ensure_stdout_descriptor()
                     self.engine = AsyncLLM.from_engine_args(AsyncEngineArgs(**self._engine_kwargs))
                     self._engine_loop = asyncio.get_running_loop()
         return self.engine

{interp_engine-1.3.3 → interp_engine-1.3.5}/interp_engine/vllm_capture/static.py

@@ -446,9 +446,12 @@ def _config_weight_bytes(config: Any) -> int:
     routed_mlp = f.n_experts * 3 * f.d_model * moe_inter
     shared_mlp = f.n_shared_experts * 3 * f.d_model * moe_inter
     n_sparse = len(f.moe_layers) if f.n_experts else 0
+    # Gemma-4's sparse layers keep their dense MLP beside the experts, so it is paid on every layer.
+    # It stays at the model dtype either way -- it is not part of the quantized routed bank.
+    n_dense = f.n_layers if f.dense_mlp_beside_experts else f.n_layers - n_sparse
     embeddings = f.vocab_size * f.d_model * (1 if f.tied_embeddings else 2)
     stored = _storage_dtype_bytes(config)
-    at_model_dtype = f.n_layers * attn + (f.n_layers - n_sparse) * dense_mlp + n_sparse * shared_mlp + embeddings
+    at_model_dtype = f.n_layers * attn + n_dense * dense_mlp + n_sparse * shared_mlp + embeddings
     return int(at_model_dtype * stored + n_sparse * routed_mlp * _expert_dtype_bytes(config, stored))
 
 

{interp_engine-1.3.3 → interp_engine-1.3.5}/pyproject.toml

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

interp_engine-1.3.5/tests/test_notebook_stdout.py

@@ -0,0 +1,132 @@
+"""Unit tests for the stdout descriptor a notebook kernel does not provide.
+
+No vLLM and no GPU: what is under test is a property of ``sys.stdout``, and the stream
+that has that property in the field -- ipykernel's ``OutStream`` -- is stood in for by a
+class which raises from ``fileno()`` the same way. The engine-shaped half is a vLLM engine
+started under such a stream, which fails in the EngineCore child rather than here.
+
+``sys.stdout`` is replaced **inside each test body** rather than by a fixture. pytest
+captures at the descriptor level and reassigns ``sys.stdout`` every time it resumes
+capture, including on the way from setup into the call phase -- so a patch applied in a
+fixture is gone by the time the test runs, and the test passes or fails on the runner's
+stream instead of the one it meant to use.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import sys
+from typing import IO, Any
+
+import pytest
+
+from interp_engine.notebook_stdout import ensure_stdout_descriptor
+
+
+class KernelStdout:
+    """``ipykernel.iostream.OutStream``, as far as this matters.
+
+    Writes go somewhere that is not a file -- a ZMQ socket, in the real one -- so there is
+    no descriptor to answer with, which is the whole of the failure being reproduced.
+    """
+
+    def __init__(self) -> None:
+        self.written: list[str] = []
+
+    def write(self, text: str) -> int:
+        self.written.append(text)
+        return len(text)
+
+    def flush(self) -> None:
+        pass
+
+    def fileno(self) -> int:
+        raise io.UnsupportedOperation("fileno")
+
+
+class SlottedStdout:
+    """A stream that takes no new attribute, as a C-level one does not."""
+
+    __slots__ = ()
+
+    def fileno(self) -> int:
+        raise io.UnsupportedOperation("fileno")
+
+
+def _as_stdout(monkeypatch: pytest.MonkeyPatch, stream: object) -> None:
+    """Make ``stream`` this test's ``sys.stdout``. Call from the test body; see above."""
+    monkeypatch.setattr(sys, "stdout", stream)
+
+
+def test_a_kernel_stream_gains_a_descriptor_the_os_accepts(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The point of the exercise: ``fileno()`` answers, and with a live descriptor.
+
+    ``fstat`` is the assertion rather than a particular number, because what vLLM does
+    with the answer is ``os.dup``, which fails on anything ``fstat`` would reject too.
+    """
+    _as_stdout(monkeypatch, KernelStdout())
+    assert ensure_stdout_descriptor() is True
+    os.fstat(sys.stdout.fileno())
+
+
+def test_writes_still_go_where_they_did(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A descriptor is added; nothing is redirected. The cell keeps its output."""
+    stream = KernelStdout()
+    _as_stdout(monkeypatch, stream)
+    ensure_stdout_descriptor()
+    print("to the notebook")
+    assert "to the notebook" in "".join(stream.written)
+
+
+def test_it_is_idempotent(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Every engine build calls this, and only the first one has anything to do."""
+    _as_stdout(monkeypatch, KernelStdout())
+    assert ensure_stdout_descriptor() is True
+    assert ensure_stdout_descriptor() is False
+
+
+def test_a_real_stream_is_left_alone(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None:
+    """A script's stdout already answers, so this must not touch it.
+
+    The second assertion is what says *untouched* rather than merely still working: a
+    descriptor of our own, put on the instance, would satisfy the first one too.
+    """
+    with open(tmp_path / "out", "w") as real:
+        _as_stdout(monkeypatch, real)
+        assert ensure_stdout_descriptor() is False
+        assert "fileno" not in getattr(real, "__dict__", {})
+
+
+def test_a_stream_that_takes_no_attribute_is_refused(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Refused rather than raised: vLLM's own error is a better one than ours here."""
+    _as_stdout(monkeypatch, SlottedStdout())
+    assert ensure_stdout_descriptor() is False
+
+
+def test_a_closed_process_stdout_offers_nothing(monkeypatch: pytest.MonkeyPatch) -> None:
+    """With none to lend, the answer is no rather than a number that fails later.
+
+    This is the daemonized process, whose stdout was closed before Python started. Both
+    sources are taken away: ``sys.__stdout__``, and the fd 1 fallback that stands in for
+    it, which is open in a test runner and would otherwise be handed out.
+    """
+    _as_stdout(monkeypatch, KernelStdout())
+    monkeypatch.setattr(sys, "__stdout__", None)
+    monkeypatch.setattr("interp_engine.notebook_stdout.os.fstat", _raise_ebadf)
+    assert ensure_stdout_descriptor() is False
+
+
+def _raise_ebadf(fd: int) -> os.stat_result:
+    raise OSError(9, "Bad file descriptor")
+
+
+def test_the_stand_in_matches_ipykernel() -> None:
+    """The premise, asserted where it can be: ``fileno()`` raises what is caught for.
+
+    ``io.UnsupportedOperation`` is both an ``OSError`` and a ``ValueError``, which is why
+    the module catches those two rather than importing ipykernel to name the class.
+    """
+    stream: IO[Any] = KernelStdout()  # type: ignore[assignment]
+    with pytest.raises((OSError, ValueError)):
+        stream.fileno()