interp-engine 1.7.0__tar.gz → 1.7.2__tar.gz

{interp_engine-1.7.0 → interp_engine-1.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.7.0
+Version: 1.7.2
 Summary: A fast, standardized, and easy to use interpretability engine.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine
@@ -162,6 +162,7 @@ Never OOM again - `interp-engine` includes gpu-sizer, an intuitive UI which tell
 - **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)).
 - **DeepSeek-V4 on the eager backend needs two flags outside Hopper.** Its FP8 paths reach for `kernels-community/deep-gemm`, which declares a `9.0a` build, and with `kernels` >= 0.16.1 the arch refusal escapes transformers' Triton fallback and kills the first forward. Load with `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` and `experts_implementation="grouped_mm"`, which transformers [already recommends on B200](https://github.com/huggingface/transformers/blob/93c8b7b485963a10800c91f55304db6be211c2bd/src/transformers/integrations/finegrained_fp8.py#L252-L254) for an unrelated DeepGEMM accuracy problem. The engine raises `HubKernelUnsupported` naming both; vLLM is unaffected.
 - **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.
+- **`gemma-4-26B-A4B-it` on the static vLLM backend needs `VLLM_BATCH_INVARIANT=1` on B200/B300.** On compute capability 10.x, vLLM's own greedy output moves by up to 6.8 nats under the CUDA-graph replay a static set requires, so the taps would report a forward pass vLLM itself gets wrong ([vllm#55238](https://github.com/vllm-project/vllm/issues/55238)). The engine refuses this combination at load and names the flag; capability 9.x and 12.x are bit-identical and unaffected.
 
 Per-architecture structural quirks — which points a family serves and why — are not caveats but facts about the architecture, and live in [ARCHITECTURE_QUIRKS.md](docs/ARCHITECTURE_QUIRKS.md).
 

{interp_engine-1.7.0 → interp_engine-1.7.2}/README.md

@@ -131,6 +131,7 @@ Never OOM again - `interp-engine` includes gpu-sizer, an intuitive UI which tell
 - **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)).
 - **DeepSeek-V4 on the eager backend needs two flags outside Hopper.** Its FP8 paths reach for `kernels-community/deep-gemm`, which declares a `9.0a` build, and with `kernels` >= 0.16.1 the arch refusal escapes transformers' Triton fallback and kills the first forward. Load with `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` and `experts_implementation="grouped_mm"`, which transformers [already recommends on B200](https://github.com/huggingface/transformers/blob/93c8b7b485963a10800c91f55304db6be211c2bd/src/transformers/integrations/finegrained_fp8.py#L252-L254) for an unrelated DeepGEMM accuracy problem. The engine raises `HubKernelUnsupported` naming both; vLLM is unaffected.
 - **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.
+- **`gemma-4-26B-A4B-it` on the static vLLM backend needs `VLLM_BATCH_INVARIANT=1` on B200/B300.** On compute capability 10.x, vLLM's own greedy output moves by up to 6.8 nats under the CUDA-graph replay a static set requires, so the taps would report a forward pass vLLM itself gets wrong ([vllm#55238](https://github.com/vllm-project/vllm/issues/55238)). The engine refuses this combination at load and names the flag; capability 9.x and 12.x are bit-identical and unaffected.
 
 Per-architecture structural quirks — which points a family serves and why — are not caveats but facts about the architecture, and live in [ARCHITECTURE_QUIRKS.md](docs/ARCHITECTURE_QUIRKS.md).
 

{interp_engine-1.7.0 → interp_engine-1.7.2}/interp_engine/capture.py

@@ -224,7 +224,10 @@ def _run_with_cache_eager(
     attention_mask: torch.Tensor | None = None,
 ) -> Cache:
     """Capture in-process off the live module tree. See :func:`run_with_cache`."""
-    input_ids = as_batched_tokens(tokens)
+    # Placed here rather than by each caller: these ids go straight into `hf_model`, so a list or a
+    # host tensor -- both documented inputs -- fails on every accelerator otherwise, through either
+    # entry point that shares this body. A tensor already on the device is unmoved.
+    input_ids = as_batched_tokens(tokens, device=model.device)
     addresses = _normalize_points(points)
     cache = Cache()
 

{interp_engine-1.7.0 → interp_engine-1.7.2}/interp_engine/vllm_backend.py

@@ -60,6 +60,7 @@ from interp_engine.vllm_capture.static import (
     kv_cache_width,
     resid_stream_aliases,
     resolve_static_points,
+    sm100_cudagraph_refusal_reason,
     static_read_width,
     static_unsupported_reason,
 )
@@ -67,6 +68,22 @@ from interp_engine.vllm_plugin import WORKER_EXTENSION_CLS
 
 logger = logging.getLogger(__name__)
 
+
+def _device_capability() -> tuple[int, int] | None:
+    """This GPU's compute capability, or None when there is no CUDA device to ask.
+
+    None rather than a default, because callers gate refusals on it: a guess would either refuse a
+    machine nobody measured or wave through the one that is known wrong.
+    """
+    try:
+        if not torch.cuda.is_available():
+            return None
+        major, minor = torch.cuda.get_device_capability(0)
+        return int(major), int(minor)
+    except Exception:  # noqa: BLE001 - no driver, or a torch without the call
+        return None
+
+
 # Residual-width sites the warmup sentinel can prove. ``mlp_act`` / ``z`` need a different
 # vector length than ``d_model``, which this process does not know until a worker wrap exists.
 _STATIC_SENTINEL_WRITE_POINTS = frozenset(
@@ -1099,6 +1116,7 @@ class VLLMModel:
         # Same condition as `apply_breakable_env`: only a non-empty static set turns torch.compile
         # off, and it is that combination a linear-attention trunk cannot survive.
         if reads or writes:
+            self._refuse_static_where_vllm_reads_wrong()
             self._pin_decode_only_graphs_on_hybrid_trunk()
         # Reads only. A write allocates a `[1, width]` delta (see `static._alloc_site`), so it does
         # not scale with `max_num_batched_tokens` and has no business in a budget whose whole job is
@@ -1144,6 +1162,23 @@ class VLLMModel:
             logger.warning("lowering max_num_batched_tokens %s -> %s so static buffers fit", max_n, fitted)
             self._engine_kwargs["max_num_batched_tokens"] = fitted
 
+    def _refuse_static_where_vllm_reads_wrong(self) -> None:
+        """Refuse a static set that would report a forward pass vLLM itself gets wrong.
+
+        Raised here rather than after the engine exists, because it is decided by the checkpoint and
+        the device: building 26B of weights first only delays the same answer. See
+        :func:`~interp_engine.vllm_capture.static.sm100_cudagraph_refusal_reason`.
+        """
+        compilation = self._engine_kwargs.get("compilation_config")
+        reason = sm100_cudagraph_refusal_reason(
+            getattr(self, "hf_model_id", None),
+            _device_capability(),
+            cudagraph_mode=str(compilation.get("cudagraph_mode") or "") if isinstance(compilation, dict) else "",
+            batch_invariant=os.environ.get("VLLM_BATCH_INVARIANT") == "1",
+        )
+        if reason:
+            raise ValueError(reason)
+
     def _pin_decode_only_graphs_on_hybrid_trunk(self) -> None:
         """Capture graphs for decode only when the trunk is linear attention.
 

{interp_engine-1.7.0 → interp_engine-1.7.2}/interp_engine/vllm_capture/static.py

@@ -142,6 +142,54 @@ def decode_only_graphs_reason(layer_types: Sequence[str] | None, n_layers: int)
     )
 
 
+#: The upstream issue behind :func:`sm100_cudagraph_refusal_reason`.
+SM100_CUDAGRAPH_ISSUE = "https://github.com/vllm-project/vllm/issues/55238"
+
+#: The one setting measured to make that capture correct again.
+SM100_CUDAGRAPH_WORKAROUND = "VLLM_BATCH_INVARIANT=1"
+
+#: Checkpoints measured wrong under a static set on compute capability 10.x. Exact ids rather than a
+#: glob over MoE Gemma 4: this is the one that was measured, and refusing a checkpoint nobody has
+#: tested trades a silent wrong answer for a wrong refusal.
+SM100_CUDAGRAPH_MODELS = ("google/gemma-4-26B-A4B-it",)
+
+
+def sm100_cudagraph_refusal_reason(
+    hf_id: str | None,
+    capability: tuple[int, int] | None,
+    *,
+    cudagraph_mode: str = "",
+    batch_invariant: bool = False,
+) -> str | None:
+    """Why this checkpoint cannot be captured under a static set on this GPU, or None.
+
+    A static set needs graph replay with torch.compile off, and on compute capability 10.x (B200,
+    B300) vLLM's own greedy output moves under exactly that: up to 6.8 nats of logprob and a different
+    token, deterministic, and bit-identical on 9.x and 12.x. It follows CUDA-graph padding rather than
+    the prompt, so no choice of prompt avoids it, and the taps sit downstream -- they would report a
+    forward pass vLLM itself gets wrong, at every point of the affected layers.
+
+    Refused rather than worked around quietly. ``VLLM_BATCH_INVARIANT=1`` fixes it but changes kernel
+    selection for the whole run, which is a decision about what the captured numbers mean, and that
+    belongs to the caller. See :data:`SM100_CUDAGRAPH_ISSUE` and validator/comparison/engine_bugs.py.
+    """
+    if not hf_id or hf_id not in SM100_CUDAGRAPH_MODELS:
+        return None
+    if capability is None or capability[0] != 10:
+        return None
+    # Either of these removes the defect: the flag was measured bit-identical at every prompt length,
+    # and without replay there is no padding to be wrong about.
+    if batch_invariant or str(cudagraph_mode).upper() == "NONE":
+        return None
+    return (
+        f"{hf_id} on compute capability {capability[0]}.{capability[1]}: vLLM's own output moves under "
+        "CUDA-graph replay with torch.compile off, which is what a static set requires -- up to 6.8 nats "
+        "of logprob and a different greedy token -- so these taps would read a forward pass vLLM itself "
+        f"gets wrong ({SM100_CUDAGRAPH_ISSUE}). Set {SM100_CUDAGRAPH_WORKAROUND} to capture anyway, or "
+        "use the hooked `vllm` backend, or capture on capability 9.x/12.x."
+    )
+
+
 def multi_stream_refusal_reason(name: str, n_streams: int) -> str | None:
     """Why ``name`` cannot be declared as a static tap on a trunk carrying ``n_streams`` residual streams, or None.
 

{interp_engine-1.7.0 → interp_engine-1.7.2}/pyproject.toml

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

{interp_engine-1.7.0 → interp_engine-1.7.2}/tests/test_static_set.py

@@ -11,6 +11,8 @@ from interp_engine.address import Address
 from interp_engine.points import steer_refusal_reason
 from interp_engine.vllm_capture.static import (
     BREAKABLE_ENV,
+    SM100_CUDAGRAPH_ISSUE,
+    SM100_CUDAGRAPH_WORKAROUND,
     STATIC_SKIP_ABSENT_ENV,
     StaticState,
     _activation_width,
@@ -41,6 +43,7 @@ from interp_engine.vllm_capture.static import (
     multi_stream_refusal_reason,
     resid_stream_aliases,
     resolve_static_points,
+    sm100_cudagraph_refusal_reason,
     static_buffer_bytes,
     static_unsupported_reason,
     steer_write_for_sae_point,
@@ -603,6 +606,38 @@ def test_unknown_layer_types_are_not_a_hybrid_claim():
     assert decode_only_graphs_reason([], 12) is None
 
 
+def test_the_checkpoint_vllm_reads_wrong_on_sm100_is_refused_rather_than_captured():
+    """A static set on this pair would report vLLM's own wrong forward as an interp result.
+
+    The refusal names the workaround and the issue, because a caller who hits this needs to decide
+    between `VLLM_BATCH_INVARIANT=1` and different hardware, and cannot do that from "unsupported".
+    """
+    reason = sm100_cudagraph_refusal_reason("google/gemma-4-26B-A4B-it", (10, 0))
+    assert reason is not None
+    assert SM100_CUDAGRAPH_WORKAROUND in reason
+    assert SM100_CUDAGRAPH_ISSUE in reason
+    assert sm100_cudagraph_refusal_reason("google/gemma-4-26B-A4B-it", (10, 3)) is not None
+
+
+def test_the_sm100_refusal_lifts_for_each_setting_measured_to_remove_the_defect():
+    """`VLLM_BATCH_INVARIANT=1` was bit-identical at every prompt length, and without replay there is
+    no padding to be wrong about. Both are the caller's call to make, so neither is refused."""
+    args = ("google/gemma-4-26B-A4B-it", (10, 0))
+    assert sm100_cudagraph_refusal_reason(*args, batch_invariant=True) is None
+    assert sm100_cudagraph_refusal_reason(*args, cudagraph_mode="NONE") is None
+    assert sm100_cudagraph_refusal_reason(*args, cudagraph_mode="none") is None
+
+
+def test_the_sm100_refusal_does_not_reach_hardware_or_checkpoints_it_was_not_measured_on():
+    """Capability 9.0 and 12.0 were measured bit-identical across all 60 points, and no other
+    checkpoint was measured at all -- refusing those would trade a wrong answer for a wrong refusal."""
+    assert sm100_cudagraph_refusal_reason("google/gemma-4-26B-A4B-it", (9, 0)) is None
+    assert sm100_cudagraph_refusal_reason("google/gemma-4-26B-A4B-it", (12, 0)) is None
+    assert sm100_cudagraph_refusal_reason("google/gemma-4-26B-A4B-it", None) is None
+    assert sm100_cudagraph_refusal_reason("google/gemma-4-12B-it", (10, 0)) is None
+    assert sm100_cudagraph_refusal_reason(None, (10, 0)) is None
+
+
 def test_every_recurrent_kind_upstream_leaves_unbroken_is_pinned():
     """The shared recurrent layers carry no `eager_break_during_capture`, so each of these families is
     exposed the same way Qwen3.5 was: mamba_mixer (Jamba), mamba_mixer2 (Nemotron-H, Bamba,
@@ -1239,6 +1274,76 @@ def test_capturing_attn_wrap_weak_refs_qkv_before_add_eager(monkeypatch):
     assert torch.equal(v_site.buf[:3], v)
 
 
+class _ReplayableBreakableCapture(_FakeBreakableCapture):
+    """:class:`_FakeBreakableCapture`, plus the replays a real graph runs after recording.
+
+    The base class stops at capture time. That is what the weak-ref tests above need -- and it is
+    also why they cannot see this: a recorded eager segment runs again on *every* replay, reading
+    whatever the weak-reffed addresses hold then, and :func:`_harvest` reads the site buffer after
+    the replay rather than after the recording.
+    """
+
+    def replay(self) -> None:
+        for fn in self.fns:
+            fn()
+
+
+class _KwargAttn(torch.nn.Module):
+    """An attention module called the way vLLM calls it on a Gemma-shaped block."""
+
+    def forward(self, positions=None, hidden_states=None):  # noqa: ANN001
+        return hidden_states * 0.5
+
+
+def _tap_attn_in(monkeypatch) -> tuple[_ReplayableBreakableCapture, _Site, torch.Tensor]:
+    """An ``attn_in`` read tapped on a recording graph, with the trunk tensor it was called with."""
+    import interp_engine.vllm_capture.static as static_mod
+
+    cap = _ReplayableBreakableCapture()
+    monkeypatch.setattr(static_mod, "_breakable_capture", lambda: cap)
+    attn = _KwargAttn()
+    site = _Site(address=Address("attn_in", 22), buf=torch.zeros(4, 2), module=attn)
+    _wrap_module(attn, [(site, "read")])
+    hidden = torch.full((3, 2), 22.0)
+    attn(positions=torch.arange(3), hidden_states=hidden)
+    return cap, site, hidden
+
+
+def test_replay_recopies_the_same_rows_when_nothing_overwrites_the_tapped_buffer(monkeypatch):
+    """Replaying the segment is harmless by itself -- it recopies what it copied at capture.
+
+    Read with the test below: deferral alone is safe, and only a reused block moves the rows. That
+    separates the two candidate explanations whenever a static column disagrees with a hooked one.
+    """
+    cap, site, _ = _tap_attn_in(monkeypatch)
+    assert torch.equal(site.buf[:3], torch.full((3, 2), 22.0))
+    cap.replay()
+    assert torch.equal(site.buf[:3], torch.full((3, 2), 22.0))
+
+
+def test_a_replayed_read_follows_the_address_so_vllms_callback_order_is_load_bearing(monkeypatch):
+    """The recorded read holds an address, so its correctness rests on vLLM's callback order.
+
+    The weak ref handed to ``add_eager`` is deliberate: a strong Python ref pins the capture-time
+    cudagraph-pool slot, and replay then recopies the profile run into the site buffer instead of the
+    live request (:func:`_cuda_weak_ref`). Materialising a copy at registration would instead cost a
+    copy and its VRAM per site per step, on every model. So the tap reads an address and depends on
+    vLLM running the callback where it was registered, ahead of anything that writes that block.
+
+    Overwrite the block first and the read follows the address. That is the shape to watch for if a
+    future vLLM stops ordering eager segments this way, and this test is here to catch that.
+
+    It is not the explanation for ``attn_in.22`` on capability 10.x. Plain vLLM moves by 6.8 nats on
+    that hardware with no tap in the process, and a fresh engine and a reversed run order return the
+    same wrong numbers, which a reused block would not
+    (validator/docs/VLLM_SM100_CUDAGRAPH.md).
+    """
+    cap, site, hidden = _tap_attn_in(monkeypatch)
+    hidden.copy_(torch.full((3, 2), 23.0))  # a later layer reusing the block
+    cap.replay()
+    assert torch.equal(site.buf[:3], torch.full((3, 2), 23.0))
+
+
 class _FusedQKV(torch.nn.Module):
     """The three numbers :func:`value_span` reads off vLLM's ``QKVParallelLinear``."""
 

{interp_engine-1.7.0 → interp_engine-1.7.2}/tests/test_unified_free_functions.py

@@ -8,6 +8,7 @@ checked against these same shapes by ``scripts/vllm_capture_generation_check.py`
 
 import asyncio
 
+import pytest
 import torch
 
 from interp_engine import (
@@ -75,6 +76,29 @@ def test_capture_attention_returns_scores_probs_and_per_head_value(gpt2: EagerMo
         assert got[layer]["value"].shape == (seq, gpt2.n_kv_heads, gpt2.head_dim)
 
 
+def test_every_eager_entry_point_takes_a_plain_list_of_ids(gpt2: EagerModel, prompt: str):
+    """A list of ints is a documented input, and it is the one that used to reach the device wrong.
+
+    The tests around this one all pass ``to_tokens`` output, a tensor already on the model's device,
+    so the host tensor a list builds was never exercised and both functions below raised on any
+    accelerator. ``EagerModel.capture`` hid it by placing its own tensor. On a CPU box both arms are
+    the same tensor and this only pins the signature; the device half is the gpu test below.
+    """
+    ids = [int(t) for t in gpt2.to_tokens(prompt)[0]]
+    assert sorted(capture_attention(gpt2, ids, [0])) == [0]
+    assert run_with_cache(gpt2, ids, [("resid_post", 0)]).get("resid_post", 0).shape[1] == len(ids)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="a host tensor only misplaces on an accelerator")
+def test_a_list_of_ids_reaches_the_accelerator(prompt: str):
+    """The regression itself: ids built off-device must be moved before they reach ``hf_model``."""
+    model = EagerModel("openai-community/gpt2", device="cuda", attn_implementation="eager")
+    ids = [int(t) for t in model.to_tokens(prompt)[0]]
+    assert capture_attention(model, ids, [0])[0]["scores"].device.type == "cuda"
+    assert run_with_cache(model, ids, [("resid_post", 0)]).get("resid_post", 0).device.type == "cuda"
+
+
 def test_capture_attention_agrees_with_the_points_it_is_built_from(gpt2: EagerModel, prompt: str):
     ids = gpt2.to_tokens(prompt)
     got = capture_attention(gpt2, ids, [4])