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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.3.5
+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.5 → 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

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

@@ -1134,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.5 → interp_engine-1.3.6}/interp_engine/vllm_capture/static.py

@@ -544,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.
 
@@ -561,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.5 → interp_engine-1.3.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "interp-engine"
-version = "1.3.5"
+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.5 → 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(