interp-engine 1.6.0__tar.gz → 1.7.1__tar.gz

{interp_engine-1.6.0 → interp_engine-1.7.1}/PKG-INFO

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

@@ -478,6 +478,7 @@ class GateUpLayout(StrEnum):
 
     #: ``[all_gate | all_up]``, two contiguous halves, gate first. Phi-3's ``gate_up_proj``, whose
     #: forward is ``gate, up = self.gate_up_proj(x).chunk(2, dim=-1); down(up * act(gate))``.
+    #: GLM-4's ``Glm4MLP`` performs the identical chunk.
     GATE_FIRST = "gate_first"
     #: ``[gate_0 up_0 | gate_1 up_1 | ...]``, the two branches interleaved per neuron. gpt-oss's
     #: MXFP4 expert weights, where the kernel reads ``[..., ::2]`` and ``[..., 1::2]``.
@@ -488,8 +489,10 @@ class GateUpLayout(StrEnum):
 #: packed. Absent means the refusal in ``ArchSpec.mlp_projection`` stands: a family may fuse and pack
 #: either way, and "probably gate first" is exactly the guess that produces silent garbage.
 #: Verified by the identity the branches exist to satisfy -- ``act(mlp_pre) * mlp_pre_linear`` is the
-#: down projection's input, which is captured independently as ``mlp_act`` (tests/test_fused_mlp.py).
+#: down projection's input, which is captured independently as ``mlp_act``
+#: (tests/test_mlp_internals.py).
 FUSED_GATE_UP_LAYOUTS: dict[str, GateUpLayout] = {
+    "Glm4ForCausalLM": GateUpLayout.GATE_FIRST,
     "Phi3ForCausalLM": GateUpLayout.GATE_FIRST,
 }
 

{interp_engine-1.6.0 → interp_engine-1.7.1}/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.6.0 → interp_engine-1.7.1}/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.6.0 → interp_engine-1.7.1}/pyproject.toml

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

{interp_engine-1.6.0 → interp_engine-1.7.1}/tests/synthetic_families.py

@@ -437,6 +437,44 @@ def shrunk_phi3(*, seed: int = 0) -> Any:
         return EagerModel("Phi3ForCausalLM", hf_model=model, tokenizer=NoTokenizer(), device=None, dtype="float32")
 
 
+#: Dims for :func:`shrunk_glm4`. Same subject as :data:`_PHI3_SHRUNK` -- the MLP's single
+#: ``gate_up_proj`` -- so the only constraint is an even ``intermediate_size``. GLM-4 rotates half
+#: its head dimension, so ``head_dim`` stays a multiple of four.
+_GLM4_SHRUNK: dict[str, Any] = {
+    "hidden_size": 32,
+    "intermediate_size": 64,
+    "num_hidden_layers": 2,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 2,
+    "head_dim": 8,
+    "vocab_size": 128,
+    "max_position_embeddings": 64,
+    "pad_token_id": 0,
+    "eos_token_id": 1,
+}
+
+
+def shrunk_glm4(*, seed: int = 0) -> Any:
+    """A tiny real GLM-4 with float32 weights on CPU, as an ``EagerModel``.
+
+    The second family to fuse its two pre-activation projections into one ``gate_up_proj``, and the
+    reason the packing is a table rather than a default: ``Glm4MLP`` chunks gate-first exactly as
+    Phi-3 does, but that is a fact about the family, established by running it, not one the shared
+    attribute name implies.
+    """
+    from transformers import Glm4Config, Glm4ForCausalLM
+
+    from interp_engine import EagerModel
+
+    config = Glm4Config(**_GLM4_SHRUNK)
+    config.architectures = ["Glm4ForCausalLM"]
+    torch.manual_seed(seed)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        model = Glm4ForCausalLM(config).to(torch.float32).eval()
+        return EagerModel("Glm4ForCausalLM", hf_model=model, tokenizer=NoTokenizer(), device=None, dtype="float32")
+
+
 def shrunk_lfm2_moe(*, seed: int = 0) -> Any:
     """A tiny LFM2-MoE with real float32 weights on CPU, as an ``EagerModel``.
 

{interp_engine-1.6.0 → interp_engine-1.7.1}/tests/test_mlp_internals.py

@@ -29,7 +29,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 from harness import GEMMA_IT, GPT2, QWEN_THINKING, ModelSpec, load_model, require_hf_token
-from synthetic_families import shrunk_opt, shrunk_phi3
+from synthetic_families import shrunk_glm4, shrunk_opt, shrunk_phi3
 
 from interp_engine import run_with_cache
 from interp_engine.capture import AddressLike
@@ -176,8 +176,9 @@ def test_a_plain_mlp_refuses_the_multiplied_branch():
         model.resolve_point("mlp_pre_linear", 0)
 
 
-def test_a_fused_gate_up_is_sliced_into_its_two_branches():
-    """Phi-3's shape: one projection holding both branches, so neither is a module output.
+@pytest.mark.parametrize("build", [shrunk_phi3, shrunk_glm4], ids=["phi3", "glm4"])
+def test_a_fused_gate_up_is_sliced_into_its_two_branches(build):
+    """Phi-3's and GLM-4's shape: one projection holding both branches, so neither is a module output.
 
     Served rather than refused, because a dense MLP's neuron basis exists whether or not the
     checkpoint stores the two matrices concatenated -- and this is a read plus a last-axis slice, the
@@ -186,9 +187,10 @@ def test_a_fused_gate_up_is_sliced_into_its_two_branches():
     Checked by the identity the branches exist to satisfy: `act(mlp_pre) * mlp_pre_linear` is the
     down projection's input, which arrives independently as `mlp_act`. That is what says the halves
     were cut the right way round, and it is exact -- swapping them changes the answer on any
-    non-symmetric activation, which SiLU is.
+    non-symmetric activation, which SiLU is. Run per family rather than once, because the packing is
+    a property of the family: both of these chunk gate-first, and nothing but running them says so.
     """
-    model = shrunk_phi3()
+    model = build()
     ids = torch.arange(7).unsqueeze(0) % 128
     cache = run_with_cache(model, ids, [("mlp_pre", 1), ("mlp_pre_linear", 1), ("mlp_act", 1)])
     pre, linear, act = (cache.get(name, 1) for name in ("mlp_pre", "mlp_pre_linear", "mlp_act"))

{interp_engine-1.6.0 → interp_engine-1.7.1}/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``."""