interp-engine 1.11.2__tar.gz → 1.12.0__tar.gz

{interp_engine-1.11.2 → interp_engine-1.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.11.2
+Version: 1.12.0
 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.11.2 → interp_engine-1.12.0}/docs/USAGE.md

@@ -446,6 +446,13 @@ for step in generate_stream(model, tokens, max_tokens=8, n_logprobs=5):
 accepts either an `Address` or its string form on lookup. `generate_stream` yields a `GenStep` per
 token, which is richer than the text deltas the protocol's streaming method gives you.
 
+`run_with_cache` also takes a padded batch with its `attention_mask`, on both backends. Each row
+counts its positions from 0 at its first real token, as HF `generate` does, so a left-padded row
+gives the same values as the prompt alone. `position_ids_from_mask(mask)` returns those positions.
+Do not read values at masked positions: eager computes them from pad tokens, and vLLM fills them
+with zeros. On vLLM each row is a separate request, and `position_ids` is refused unless it equals
+those positions.
+
 `load_model` on eager still returns a model whose *methods* are async, and calling one from sync code
 is the thing these functions save you from. When you want a method rather than a free function, wrap
 the model once:

{interp_engine-1.11.2 → interp_engine-1.12.0}/interp_engine/__init__.py

@@ -40,6 +40,7 @@ from interp_engine.capture import (
     head_contributions,
     is_rms_norm,
     per_head_value,
+    position_ids_from_mask,
     pre_gain_normalized,
     rms_norm_parts,
     run_with_cache,
@@ -238,6 +239,7 @@ __all__ = [
     "tlens_hook_to_point",
     "tlens_normalized_hook",
     "per_head_value",
+    "position_ids_from_mask",
     "pre_gain_normalized",
     "rms_norm_eps_for_model",
     "rms_norm_parts",

{interp_engine-1.11.2 → interp_engine-1.12.0}/interp_engine/capture.py

@@ -15,6 +15,8 @@ Design invariants (protect future probing/monitoring apps):
 
 from __future__ import annotations
 
+import asyncio
+import inspect
 import math
 from collections.abc import Sequence
 from dataclasses import dataclass, field
@@ -161,6 +163,7 @@ def run_with_cache(
     *,
     detach: bool = True,
     attention_mask: torch.Tensor | None = None,
+    position_ids: torch.Tensor | None = None,
 ) -> Cache:
     """Run a single forward pass, capturing the requested points into a :class:`Cache`.
 
@@ -171,13 +174,85 @@ def run_with_cache(
     ``tokens`` may be the ``[batch, seq]`` tensor ``model.to_tokens`` returns, a bare ``[seq]``
     tensor, or a list of ids.
 
+    ``attention_mask`` marks a padded batch's real tokens with 1. With a mask and no
+    ``position_ids``, each row counts its positions from 0 at its first real token, as HF
+    ``generate`` does (see :func:`position_ids_from_mask`). So a left-padded row gives the values
+    it gives unpadded, on both backends. Values at masked positions have no meaning; do not read them.
+
     ``detach=True`` (the default, for activation endpoints) stores detached clones. Pass
     ``detach=False`` to keep the autograd graph (the lens does this on residuals); it raises
     on vLLM, and on eager unless the model was built with ``requires_grad=True``.
     """
     if not isinstance(model, EagerModel):
-        return _run_with_cache_via_protocol(model, tokens, points, detach=detach, attention_mask=attention_mask)
-    return _run_with_cache_eager(model, tokens, points, detach=detach, attention_mask=attention_mask)
+        return _run_with_cache_via_protocol(
+            model, tokens, points, detach=detach, attention_mask=attention_mask, position_ids=position_ids
+        )
+    return _run_with_cache_eager(
+        model, tokens, points, detach=detach, attention_mask=attention_mask, position_ids=position_ids
+    )
+
+
+def position_ids_from_mask(attention_mask: torch.Tensor) -> torch.Tensor:
+    """The position ids HF ``generate`` builds from a ``[batch, seq]`` attention mask.
+
+    Each row counts from 0 at its first unmasked token, so a left-padded row gets the positions it
+    has unpadded. Masked positions get 0. :func:`run_with_cache` uses this when given a mask alone.
+    """
+    mask = attention_mask.long()
+    return (mask.cumsum(-1) - 1).masked_fill(mask == 0, 0)
+
+
+def _as_mask(attention_mask: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+    """``attention_mask`` on the ids' device, with the ids' ``[batch, seq]`` shape when it is 1-D or 2-D.
+
+    A mask of more dimensions is the caller's own attention pattern and goes to the model unchanged.
+    """
+    mask = attention_mask.to(input_ids.device)
+    if mask.ndim == 1:
+        mask = mask.unsqueeze(0)
+    if mask.ndim == 2 and mask.shape != input_ids.shape:
+        raise ValueError(
+            f"attention_mask has shape {tuple(mask.shape)} and the tokens have shape "
+            f"{tuple(input_ids.shape)}; a padding mask has one entry per token."
+        )
+    return mask
+
+
+def _as_positions(position_ids: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+    """``position_ids`` on the ids' device, with a batch axis when it is 1-D."""
+    given = position_ids.to(input_ids.device)
+    if given.ndim == 1:
+        given = given.unsqueeze(0)
+    if given.shape[-1] != input_ids.shape[-1]:
+        raise ValueError(
+            f"position_ids has shape {tuple(given.shape)} and the tokens have shape "
+            f"{tuple(input_ids.shape)}; it needs one position per token."
+        )
+    return given
+
+
+def _forward_position_ids(
+    hf_model: Any, input_ids: torch.Tensor, mask: torch.Tensor | None, position_ids: torch.Tensor | None
+) -> torch.Tensor | None:
+    """The ``position_ids`` for ``hf_model``, or ``None`` to let the model count from 0.
+
+    The same rule as HF ``generate``: a 2-D mask with no explicit ids gives
+    :func:`position_ids_from_mask`, if the model is decoder-only and its forward takes them.
+    """
+    takes = "position_ids" in inspect.signature(hf_model.forward).parameters
+    if position_ids is not None:
+        if not takes:
+            raise ValueError(
+                f"{type(hf_model).__name__}.forward takes no position_ids, so they cannot be honored. "
+                "Its positions come from the attention mask (ALiBi) or are not an input at all: "
+                "pass attention_mask and omit position_ids."
+            )
+        return _as_positions(position_ids, input_ids)
+    if mask is None or mask.ndim != 2 or not takes:
+        return None
+    if getattr(hf_model.config, "is_encoder_decoder", False):
+        return None
+    return position_ids_from_mask(mask)
 
 
 def _run_with_cache_via_protocol(
@@ -187,32 +262,76 @@ def _run_with_cache_via_protocol(
     *,
     detach: bool,
     attention_mask: torch.Tensor | None,
+    position_ids: torch.Tensor | None,
 ) -> Cache:
-    """The non-eager arm: one prompt through :meth:`InterpModel.capture`, shaped like a Cache.
+    """The non-eager arm: each row through :meth:`InterpModel.capture`, shaped like a Cache.
 
-    ``attention_mask`` is refused rather than ignored -- a single unpadded sequence needs none,
-    and a mask that was going to be dropped is worse than one that raises. The batch refusal
-    lives in :func:`~interp_engine.dispatch.as_token_ids`, for the same reason.
+    The backend takes one unpadded prompt per request and counts its positions from 0. So each
+    row runs as its unmasked tokens alone, which gives the positions :func:`position_ids_from_mask`
+    gives the eager arm. The rows go as concurrent requests, which the backend's scheduler batches.
+    Masked positions hold zeros.
+
+    ``position_ids`` is accepted only where it is those positions, as nothing else can reach
+    the backend.
 
     :attr:`Cache.output` stays ``None``: the forward happened in a worker subprocess and there
     is no HF output object to hand back. Everything reading ``cache.tensors`` is unaffected,
     which is nearly every caller.
     """
-    if attention_mask is not None:
-        raise refuse(model, "run_with_cache", capability="attention_mask")
     steering = active_steering(model)
     if steering is not None and steering.position_mask is not None:
         raise refuse(model, "steer(..., position_mask=...) around a capture", capability="masked_steer_positions")
-    ids = as_token_ids(tokens, model=model, what="run_with_cache")
-    captured = sync_model(model).capture(
-        ids,
-        _normalize_points(points),
-        steering_spec=None if steering is None else steering.spec,
-        detach=detach,
-    )
-    # Restore the batch axis the eager path keeps, so `cache[point][0]` means the same thing on
-    # both backends. A view, so this costs nothing.
-    return Cache(tensors={address: t.unsqueeze(0) for address, t in captured.items()}, output=None)
+    input_ids = as_batched_tokens(tokens).cpu()
+    keep = torch.ones_like(input_ids, dtype=torch.bool)
+    if attention_mask is not None:
+        mask = _as_mask(attention_mask, input_ids)
+        if mask.ndim != 2:
+            raise ValueError(
+                f"attention_mask has shape {tuple(mask.shape)}; this backend takes a [batch, seq] "
+                "padding mask, not an attention pattern."
+            )
+        keep = mask != 0
+    if position_ids is not None:
+        given = _as_positions(position_ids, input_ids)
+        want = position_ids_from_mask(keep)
+        if (
+            given.ndim != 2
+            or given.shape[0] not in (1, want.shape[0])
+            or not torch.equal(given.expand_as(want)[keep], want[keep])
+        ):
+            raise refuse(model, "run_with_cache(position_ids=...)", capability="position_ids")
+    rows = [input_ids[b][keep[b]].tolist() for b in range(input_ids.shape[0])]
+    for b, row in enumerate(rows):
+        if not row:
+            raise ValueError(f"run_with_cache: row {b} has no unmasked token, so it has nothing to run.")
+    addresses = _normalize_points(points)
+    spec = None if steering is None else steering.spec
+
+    async def _every_row() -> list[dict[Address, torch.Tensor]]:
+        return await asyncio.gather(*(model.capture(row, addresses, steering_spec=spec, detach=detach) for row in rows))
+
+    captured = sync_model(model).runner.run(_every_row(), what="run_with_cache()")
+    if len(rows) == 1 and bool(keep.all()):
+        # Restore the batch axis the eager path keeps, so `cache[point][0]` means the same thing on
+        # both backends. A view, so this costs nothing.
+        return Cache(tensors={address: t.unsqueeze(0) for address, t in captured[0].items()}, output=None)
+    return Cache(tensors={a: _pad_rows(a, [c[a] for c in captured], keep) for a in captured[0]}, output=None)
+
+
+def _pad_rows(address: Address, rows: list[torch.Tensor], keep: torch.Tensor) -> torch.Tensor:
+    """One ``[batch, seq, ...]`` tensor from each row's ``[kept, ...]`` capture, zeros where ``keep`` is False."""
+    first = rows[0]
+    out = first.new_zeros((keep.shape[0], keep.shape[1], *first.shape[1:]))
+    for b, row in enumerate(rows):
+        kept = int(keep[b].sum())
+        if row.shape[0] != kept or row.shape[1:] != first.shape[1:]:
+            raise ValueError(
+                f"{address} came back with shape {tuple(row.shape)} for a {kept}-token row, which is "
+                "not one entry per token, so it cannot go back into a padded batch. Capture it one "
+                "prompt at a time."
+            )
+        out[b, keep[b].to(out.device)] = row
+    return out
 
 
 def _run_with_cache_eager(
@@ -222,12 +341,15 @@ def _run_with_cache_eager(
     *,
     detach: bool = True,
     attention_mask: torch.Tensor | None = None,
+    position_ids: torch.Tensor | None = None,
 ) -> Cache:
     """Capture in-process off the live module tree. See :func:`run_with_cache`."""
     # 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)
+    mask = None if attention_mask is None else _as_mask(attention_mask, input_ids)
+    positions = _forward_position_ids(model.hf_model, input_ids, mask, position_ids)
     addresses = _normalize_points(points)
     cache = Cache()
 
@@ -324,12 +446,15 @@ def _run_with_cache_eager(
             hm.read(module, make_reader(keys), point=point)  # pyright: ignore[reportArgumentType]
 
         forward_ctx = torch.no_grad() if detach else torch.enable_grad()
+        # Left out rather than passed as `None`, so a forward with no such argument never sees it.
+        extra = {} if positions is None else {"position_ids": positions}
         with forward_ctx:
             output = model.hf_model(
                 input_ids,
-                attention_mask=attention_mask,
+                attention_mask=mask,
                 output_attentions=wants_attn,
                 use_cache=False,
+                **extra,
             )
         cache.output = output
         for layer, tensor in scores.items():

{interp_engine-1.11.2 → interp_engine-1.12.0}/interp_engine/dispatch.py

@@ -137,10 +137,10 @@ CAPABILITIES: dict[str, Capability] = {
         why="vLLM owns its weights in a worker subprocess and does not expose them to this process",
         instead="running this on an eager model, which holds the live module tree",
     ),
-    "attention_mask": Capability(
-        what="a padding mask",
-        why="this arm forwards one unpadded sequence, so there is nothing for a mask to cover",
-        instead="omitting it, and issuing one call per prompt",
+    "position_ids": Capability(
+        what="position ids other than the ones its attention mask gives",
+        why="each row runs as its own request, and the backend counts that request's positions from 0",
+        instead="omitting position_ids, which counts each row from its first unmasked token on both backends",
     ),
     "masked_steer_positions": Capability(
         what="steering that skips some prompt positions",

{interp_engine-1.11.2 → interp_engine-1.12.0}/interp_engine/vllm_capture/requests.py

@@ -65,7 +65,7 @@ from interp_engine.vllm_capture.capture import worker_addresses
 from interp_engine.vllm_capture.graphs import graph_debug, refuse_writes_reason
 from interp_engine.vllm_capture.lens import _make_lens_modifier
 from interp_engine.vllm_capture.mhc import mhc_taps, require_available, require_steerable
-from interp_engine.vllm_capture.steering import _make_steer_modifier
+from interp_engine.vllm_capture.steering import _compose_modifiers, _make_steer_modifiers
 
 # --- the per-request combined hook body ---------------------------------------
 
@@ -679,9 +679,11 @@ def worker_register_steering(
     demux.registered.add(req_id)
     mods = demux.steer_mods.setdefault(req_id, {})
     skip_set = {int(i) for i in (skip_positions or [])}
+    by_site: dict[Address, list[dict]] = {}
     for s in specs:
-        site = _write_site(worker, s)
-        mods[site] = (_make_steer_modifier(s, demux.dev, demux.dt), skip_set, int(prompt_len))
+        by_site.setdefault(_write_site(worker, s), []).append(s)
+    for site, group in by_site.items():
+        mods[site] = (_make_steer_modifiers(group, demux.dev, demux.dt), skip_set, int(prompt_len))
         _ensure_hook(worker, demux, site)
 
 
@@ -708,10 +710,12 @@ def worker_register_lens(
     demux.registered.add(req_id)
     lm = demux.lens_mods.setdefault(req_id, {})
     skip_set = {int(i) for i in (skip_positions or [])}
+    by_site: dict[Address, list[dict]] = {}
     for s in specs:
-        site = _write_site(worker, s)
+        by_site.setdefault(_write_site(worker, s), []).append(s)
+    for site, group in by_site.items():
         lm[site] = (
-            _make_lens_modifier(s, demux.dev, demux.dt),
+            _compose_modifiers([_make_lens_modifier(s, demux.dev, demux.dt) for s in group]),
             bool(steer_generated),
             skip_set,
             int(prompt_len),

{interp_engine-1.11.2 → interp_engine-1.12.0}/interp_engine/vllm_capture/static.py

@@ -1812,9 +1812,9 @@ def worker_set_static_delta(
 ) -> None:
     """Install static writes from worker specs. Zeros every write site first.
 
-    ``op="additive"`` without a lens scope fills the static ``delta`` buffer. Orthogonal,
-    projection_cap, and lens ops attach a live ``modify`` that reads the residual each
-    forward (breakable ``add_eager``). ``lens_scope`` is the jlens prefill/decode skip.
+    A site of only plain adds, with no lens scope, fills the static ``delta`` buffer with their
+    sum. Any other op attaches one live ``modify`` for the site's ops, in order, that reads the
+    residual each forward (breakable ``add_eager``). ``lens_scope`` is the jlens skip.
 
     A ``stream`` disqualifies the static buffer too. The buffer is added whole to a ``[tokens,
     streams, width]`` activation, so it has no way to say "this stream and not the others"; taking
@@ -1827,59 +1827,78 @@ def worker_set_static_delta(
     if static is None:
         raise RuntimeError("set_static_delta: this worker has no static wraps")
     worker_clear_static_delta(worker)
-    from interp_engine.vllm_capture.lens.intervene import _make_lens_modifier
-    from interp_engine.vllm_capture.steering import _make_steer_modifier
+    for site, group in _group_by_site(static, specs).items():
+        assert site.delta is not None
+        constant = _constant_delta(group, site.delta.device, site.delta.dtype) if lens_scope is None else None
+        if constant is not None:
+            site.delta.copy_(constant.expand_as(site.delta))
+            site.delta_set = True
+            continue
+        site.modify = _group_modifier(group, site.delta.device, site.delta.dtype)
+        site.lens_scope = dict(lens_scope) if lens_scope else None
+
 
+def _group_by_site(static: StaticState, specs: list[dict]) -> dict[_Site, list[dict]]:
+    """``specs`` per write site, in order, each op checked against what a static tap can serve."""
+    out: dict[_Site, list[dict]] = {}
     for spec in specs:
         op = str(spec.get("op", SteerMethod.ADDITIVE))
         if op not in STATIC_WRITE_OPS:
             raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}")
-        site = _write_site(static, Address(str(spec["point"]), int(spec["layer"])))
-        assert site.delta is not None
-        if op == SteerMethod.ADDITIVE and lens_scope is None and spec.get("stream") is None:
-            vec = torch.tensor(spec["vector"], dtype=torch.float32, device=site.delta.device)
-            vec = (vec * float(spec["coeff"])).to(dtype=site.delta.dtype)
-            site.delta.copy_(vec.reshape(1, -1).expand_as(site.delta))
-            site.delta_set = True
-            continue
-        device, dtype = site.delta.device, site.delta.dtype
-        if op in _LENS_OPS:
-            site.modify = _make_lens_modifier(spec, device, dtype)
-        else:
-            site.modify = _make_steer_modifier(spec, device, dtype)
-        site.lens_scope = dict(lens_scope) if lens_scope else None
+        site = _write_site(static, Address(str(spec.get("point") or "resid_post"), int(spec["layer"])))
+        out.setdefault(site, []).append(spec)
+    return out
+
+
+def _constant_delta(group: list[dict], device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
+    """The ``[1, width]`` sum of ``group``'s deltas when every op is a plain add; else ``None``.
+
+    A constant vector broadcasts over a stream axis and so cannot exclude one: a ``stream`` has to
+    go the modifier way.
+    """
+    if any(
+        str(s.get("op", SteerMethod.ADDITIVE)) != SteerMethod.ADDITIVE or s.get("stream") is not None for s in group
+    ):
+        return None
+    total = sum(torch.tensor(s["vector"], dtype=torch.float32, device=device) * float(s["coeff"]) for s in group)
+    return torch.as_tensor(total).to(dtype=dtype).reshape(1, -1)
+
+
+def _group_modifier(group: list[dict], device: torch.device, dtype: torch.dtype):
+    """One ``modify`` for every op in ``group``, steer and lens ops alike, applied in order."""
+    from interp_engine.vllm_capture.lens.intervene import _make_lens_modifier
+    from interp_engine.vllm_capture.steering import _compose_modifiers, _make_steer_modifier
+
+    return _compose_modifiers(
+        [
+            _make_lens_modifier(s, device, dtype)
+            if str(s.get("op", SteerMethod.ADDITIVE)) in _LENS_OPS
+            else _make_steer_modifier(s, device, dtype)
+            for s in group
+        ]
+    )
 
 
 def _compile_write_req(
-    spec: dict,
+    group: list[dict],
     site: _Site,
     *,
     skip_positions: tuple[int, ...],
     prompt_len: int,
     steer_generated: bool,
 ) -> _WriteReq:
-    op = str(spec.get("op", SteerMethod.ADDITIVE))
-    if op not in STATIC_WRITE_OPS:
-        raise ValueError(f"static taps cannot serve op={op!r}; supported ops are {sorted(STATIC_WRITE_OPS)}")
     assert site.delta is not None
     device, dtype = site.delta.device, site.delta.dtype
-    # A constant `[1, width]` vector broadcasts over a stream axis and so cannot exclude one; see
-    # `worker_set_static_delta` for why a `stream` therefore has to go the modifier way.
-    if op == SteerMethod.ADDITIVE and spec.get("stream") is None:
-        vec = torch.tensor(spec["vector"], dtype=torch.float32, device=device)
-        vec = (vec * float(spec["coeff"])).to(dtype=dtype).reshape(1, -1)
+    constant = _constant_delta(group, device, dtype)
+    if constant is not None:
         return _WriteReq(
-            vector=vec,
+            vector=constant,
             skip_positions=skip_positions,
             prompt_len=prompt_len,
             steer_generated=steer_generated,
         )
-    from interp_engine.vllm_capture.lens.intervene import _make_lens_modifier
-    from interp_engine.vllm_capture.steering import _make_steer_modifier
-
-    modify = _make_lens_modifier(spec, device, dtype) if op in _LENS_OPS else _make_steer_modifier(spec, device, dtype)
     return _WriteReq(
-        modify=modify,
+        modify=_group_modifier(group, device, dtype),
         skip_positions=skip_positions,
         prompt_len=prompt_len,
         steer_generated=steer_generated,
@@ -1906,10 +1925,9 @@ def worker_register_static_write(
         length = int(lens_scope.get("prompt_len") or length)
         generated = bool(lens_scope.get("steer_generated", False))
     by_site: dict[str, _WriteReq] = {}
-    for spec in specs:
-        site = _write_site(static, Address(str(spec.get("point") or "resid_post"), int(spec["layer"])))
+    for site, group in _group_by_site(static, specs).items():
         by_site[format_address(site.address)] = _compile_write_req(
-            spec, site, skip_positions=skip, prompt_len=length, steer_generated=generated
+            group, site, skip_positions=skip, prompt_len=length, steer_generated=generated
         )
     static.write_reqs[req_id] = by_site
     static.registered.add(req_id)

{interp_engine-1.11.2 → interp_engine-1.12.0}/interp_engine/vllm_capture/steering.py

@@ -51,6 +51,32 @@ def _one_stream(modify, stream: int):
     return _modify
 
 
+def _compose_modifiers(mods: list):
+    """One ``modify`` for every op at one site, applied in order as eager applies a layer's ops.
+
+    Each op reads the residual the ops before it wrote, and the delta returned is their sum, so a
+    site that holds one write slot still carries all of them.
+    """
+    if len(mods) == 1:
+        return mods[0]
+
+    def _modify_all(full: torch.Tensor) -> torch.Tensor:
+        total = mods[0](full)
+        out = full + total
+        for modify in mods[1:]:
+            delta = modify(out)
+            out = out + delta
+            total = total + delta
+        return total
+
+    return _modify_all
+
+
+def _make_steer_modifiers(specs: list[dict], dev, dt):
+    """:func:`_compose_modifiers` over ``specs``, which all write one site."""
+    return _compose_modifiers([_make_steer_modifier(s, dev, dt) for s in specs])
+
+
 def _make_steer_modifier(spec: dict, dev, dt):
     """Return ``modify(full_resid) -> delta`` (the tensor to ADD to the residual).
 

{interp_engine-1.11.2 → interp_engine-1.12.0}/pyproject.toml

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

{interp_engine-1.11.2 → interp_engine-1.12.0}/tests/test_capability_refusals.py

@@ -130,13 +130,13 @@ def test_layer_logits_refuses_an_explicit_transform_on_a_non_eager_model() -> No
         layer_logits(NotEager(), [1, 2, 3], {"logit_lens": [0]}, softcap=30.0)  # pyright: ignore[reportArgumentType]
 
 
-def test_run_with_cache_refuses_an_attention_mask_on_a_non_eager_model() -> None:
-    with pytest.raises(CapabilityUnsupported, match="mask"):
+def test_run_with_cache_refuses_position_ids_its_mask_does_not_give_on_a_non_eager_model() -> None:
+    with pytest.raises(CapabilityUnsupported, match="position ids"):
         run_with_cache(
             NotEager(),  # pyright: ignore[reportArgumentType]
             torch.tensor([[1, 2, 3]]),
             ["resid_post.0"],
-            attention_mask=torch.ones(1, 3),
+            position_ids=torch.tensor([[5, 6, 7]]),
         )