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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.3.4
+Version: 1.3.6
 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.4 → interp_engine-1.3.6}/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.4 → interp_engine-1.3.6}/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.4 → interp_engine-1.3.6}/interp_engine/facts.py

@@ -777,6 +777,35 @@ def mandatory_kv_cache_dtype(architectures: Sequence[str] | None) -> str | None:
     return None
 
 
+# A multimodal model whose image tokens attend bidirectionally (Gemma 3, Gemma 4) has chunked
+# multimodal input forced OFF by vLLM, and one whole item must then fit in a single batch or the
+# engine refuses to start: "Chunked MM input disabled but max_tokens_per_mm_item (2496) is larger
+# than max_num_batched_tokens (512)". vLLM raises its own floor for this, but only while defaulting
+# `max_num_batched_tokens`, and clamps the result against `max_model_len` immediately after -- so a
+# caller who sizes the window to its prompt, which capture does, defeats the fix and gets the raise.
+#
+# A floor rather than the number, because the number is not in the HF config: the per-item token
+# count comes from vLLM's processing info for that family (2496 on Gemma 4, 256 on Gemma 3), and
+# deriving it here would mean reimplementing each family's image-token math against a private API.
+# Generous enough for every multimodal family served so far; if one ever needs more, vLLM's refusal
+# names the value it wanted, which is the only honest way to revise this.
+MM_MIN_BATCHED_TOKENS = 8192
+
+
+def min_batched_tokens(config: Any) -> int | None:
+    """The ``max_num_batched_tokens`` this checkpoint cannot boot below, or ``None`` if unconstrained.
+
+    Companion to :func:`mandatory_kv_cache_dtype`, and here for the same reason: a boot requirement
+    every harness has to know is one the library should answer, not one each caller rediscovers from
+    a stack trace.
+
+    Keyed on the config being a multimodal wrapper rather than on the family, since the constraint
+    follows from having a non-text modality at all. Text-only checkpoints get ``None`` and keep
+    vLLM's own defaults, which is what makes this safe to ask about unconditionally.
+    """
+    return MM_MIN_BATCHED_TOKENS if config is not None and text_config(config) is not config else None
+
+
 # --- quantization and the backward pass ---------------------------------------
 #
 # Quantization is invisible to *capture*: hooks read activations, which transformers dequantizes to a
@@ -1072,7 +1101,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 +1265,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 +1380,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 +1823,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 +1844,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 +1930,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 +2288,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.4 → interp_engine-1.3.6}/interp_engine/vllm_backend.py

@@ -315,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,
@@ -407,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")),
     )
 
 
@@ -551,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,
@@ -1121,6 +1134,7 @@ class VLLMModel:
             ),
             n_layers=self.num_hidden_layers,
             tensor_parallel_size=self.tensor_parallel_size,
+            min_n=facts.min_batched_tokens(cfg) or 0,
         )
         if fitted != max_n:
             logger.warning("lowering max_num_batched_tokens %s -> %s so static buffers fit", max_n, fitted)

{interp_engine-1.3.4 → interp_engine-1.3.6}/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))
 
 
@@ -541,6 +544,7 @@ def fit_max_num_batched_tokens(
     kv_width: int,
     n_layers: int,
     tensor_parallel_size: int = 1,
+    min_n: int = 0,
 ) -> int:
     """Largest capture size at or below ``max_n`` whose static buffers fit, or raise.
 
@@ -558,14 +562,20 @@ def fit_max_num_batched_tokens(
     The 1024 floor is for ``"auto"`` / default ``max_n``: vLLM will not usefully serve with a
     tiny batch. A caller who already passed a smaller ``max_n`` (chunked-prefill tests) keeps it
     when it fits.
+
+    ``min_n`` raises that floor to a size the *engine* will not start below, which is a different
+    kind of limit and has to outrank the fit: shrinking under it trades an out-of-memory error for a
+    refusal to boot, and the refusal comes from vLLM's scheduler talking about multimodal items,
+    which reads as anything but a static-buffer problem. Above the floor it changes nothing --
+    ``fitted == asked`` stays the common case. See :func:`interp_engine.facts.min_batched_tokens`.
     """
     graph_fudge = 3 * 1024**3
     tp = max(int(tensor_parallel_size), 1)
     min_kv = max(int(n_layers), 1) * max(int(max_model_len), 1) * max(int(kv_width), 1) * 2
     budget = int(gpu_memory_utilization * device_memory) - int(weight_bytes) // tp - graph_fudge - min_kv
     asked = int(max_n)
-    floor = min(1024, asked)
-    candidates = [asked] + [s for s in _CAPTURE_SIZES if s < asked]
+    floor = max(min(1024, asked), int(min_n))
+    candidates = [max(asked, floor)] + [s for s in _CAPTURE_SIZES if s < asked]
     for n in candidates:
         if n < floor:
             continue

{interp_engine-1.3.4 → interp_engine-1.3.6}/pyproject.toml

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

{interp_engine-1.3.4 → interp_engine-1.3.6}/tests/test_static_set.py

@@ -221,6 +221,33 @@ def test_vram_check_lowers_max_n_instead_of_ooming():
     assert fitted >= 1024
 
 
+def test_vram_check_will_not_shrink_below_an_engine_boot_floor():
+    """A multimodal prefix-LM will not start below one whole image, so shrinking under it is no fit.
+
+    Better to refuse here, naming the buffers, than to hand back a size vLLM rejects later with a
+    scheduler error about multimodal items that reads as anything but a static-buffer problem.
+    """
+
+    def fit(device_gib: int, **extra: int) -> int:
+        return fit_max_num_batched_tokens(
+            n_sites=80,
+            width=5120,
+            max_n=16384,
+            device_memory=device_gib * 1024**3,
+            gpu_memory_utilization=0.95,
+            weight_bytes=10 * 1024**3,
+            max_model_len=4096,
+            kv_width=kv_cache_width(d_model=5120),
+            n_layers=32,
+            **extra,
+        )
+
+    assert fit(20) == 4096  # left alone, this card shrinks under the floor
+    with pytest.raises(ValueError, match="do not fit even at max_num_batched_tokens=8192"):
+        fit(20, min_n=8192)
+    assert fit(24, min_n=8192) == 8192  # a floor it can meet changes nothing
+
+
 def test_vram_check_keeps_a_caller_max_n_below_the_1024_floor():
     """Chunked-prefill tests pass max_num_batched_tokens=32; skipping that candidate used to raise."""
     fitted = fit_max_num_batched_tokens(