interp-engine 1.13.0__tar.gz → 1.13.2__tar.gz

{interp_engine-1.13.0 → interp_engine-1.13.2}/PKG-INFO

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

{interp_engine-1.13.0 → interp_engine-1.13.2}/docs/INTERNALS.md

@@ -119,7 +119,9 @@ index names the layer HF's does, which nothing checked before and which would fa
 than raise. It needs `interp-engine[vllm]`, so it self-skips elsewhere; note that running it via
 `.venv-vllm/bin/python` needs that directory on `PATH` too, because vLLM shells out to `ninja` to
 build a sampler kernel at startup. `tests/test_vllm_wire_grammar.py` covers the same process
-boundary on CPU, over a synthetic demux. `tests/test_multigpu.py` (`-m multigpu`, two CUDA cards)
+boundary on CPU, over a synthetic demux. vLLM's Transformers backend hands each hook `[1, tokens, d]`;
+`tests/test_vllm_batch_axis.py` pins how the demux reads that on CPU, and
+`tests/test_vllm_transformers_backend_gpu.py` runs OLMo-2 on that backend against eager. `tests/test_multigpu.py` (`-m multigpu`, two CUDA cards)
 repeats the comparison at `num_gpus=2`: eager under accelerate's layer placement and vLLM under
 tensor parallelism, where `vllm_capture/_tp.py` gathers the head- and neuron-sharded points across
 ranks, against the one-card eager reference — captures, attention, steering, decode and the lens.
@@ -129,6 +131,16 @@ while the others' replies stay queued, and the next RPC consumes those — so a
 took every point of a TP=8 cell with it. `resolvable_points` and `resolvable_attn` are the ask-first
 calls; `tests/test_vllm_new_points.py` pins that they install nothing.
 
+**Static writes inside a recorded graph.** vLLM's V2 runner records FULL decode graphs with plain
+`torch.cuda.graph`, so a static write that Python decides at record time is replayed as "no write"
+on every decode step. `vllm_capture/static_program.py` makes the write data instead: one Triton
+kernel per write site reads which writer owns each row, and that writer's ops, from fixed-address
+tables that the host refills between steps. `tests/test_static_program.py` checks the tables on CPU,
+`tests/test_steer_math_parity.py` holds the kernel's single op form to the worker modifiers,
+`tests/test_static_program_gpu.py` replays a graph recorded with nothing registered, and
+`tests/test_static_decode_writes_gpu.py` requires static decode rows to match hooked vLLM for every
+method, with and without `generated`.
+
 **The docs are parsed, not maintained.** The point table in
 [SUPPORTED_POINTS.md](SUPPORTED_POINTS.md) and the footnote markers in
 [ENGINE_HOOK_MAPPINGS.md](ENGINE_HOOK_MAPPINGS.md) are read back and compared against `points.py` by

{interp_engine-1.13.0 → interp_engine-1.13.2}/interp_engine/vllm_capture/_demux.py

@@ -9,6 +9,7 @@ refcounting -- state manipulation only, no hook bodies.
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from typing import Any, cast
 
 import torch
@@ -64,6 +65,10 @@ class _Demux:
         self.runner_api: str | None = None
         # Snapshot of the in-flight forward's row layout: (req_ids, seq_lens) or None.
         self.current_meta: tuple[list[str], list[int]] | None = None
+        # Per request in ``current_meta`` order: the absolute position of its first row, or None.
+        self.current_starts: list[int] | None = None
+        # Called after each snapshot, before the forward. Static writes fill their row map here.
+        self.on_meta: list[Callable[[], None]] = []
         # Diagnostics (see worker_demux_debug).
         self.dbg_exec_calls = 0
         self.dbg_last_meta: tuple[list[str], list[int]] | None = None
@@ -89,6 +94,9 @@ class _Demux:
         # Installed hooks: hook site -> [handle, refcount]. Keyed by site and not by address, so
         # two requests reading different streams of one point share a handle -- see `hook_site`.
         self.hooks: dict[Address, list] = {}
+        # Hook site -> whether its tensor arrives with a leading batch axis of one, as seen on the
+        # last forward of more than one token (see requests._batch_axis).
+        self.batch_axis: dict[Address, bool] = {}
         self.dev: Any = None
         self.dt: Any = None
 
@@ -183,6 +191,22 @@ def _meta_from_v1_inputs(input_batch: object, num_scheduled_tokens: object) -> t
     return _pair_meta(getattr(input_batch, "req_ids", None), num_scheduled_tokens)
 
 
+def _starts(computed: object, meta: tuple[list[str], list[int]] | None) -> list[int] | None:
+    """Each request's first absolute position this step, from the runner's per-request computed-token count."""
+    if meta is None or computed is None:
+        return None
+    starts = [int(x) for x in cast(Any, computed)[: len(meta[0])]]
+    return starts if len(starts) == len(meta[0]) else None
+
+
+def _run_on_meta(demux: _Demux) -> None:
+    for callback in demux.on_meta:
+        try:
+            callback()
+        except Exception as exc:  # noqa: BLE001 - a callback must fail closed, not stop the forward
+            demux.dbg_last_error = f"on_meta {type(exc).__name__}: {exc}"
+
+
 def _ensure_patched(worker: object, demux: _Demux) -> None:
     """Patch the model runner once to snapshot each forward's row layout.
 
@@ -212,10 +236,12 @@ def _ensure_patched(worker: object, demux: _Demux) -> None:
             demux.dbg_exec_calls += 1
             try:
                 demux.current_meta = _meta_from_input_batch(input_batch)
+                demux.current_starts = _starts(getattr(input_batch, "num_computed_tokens_np", None), demux.current_meta)
                 demux.dbg_last_meta = demux.current_meta
             except Exception as exc:  # noqa: BLE001 - never break the forward on metadata errors
-                demux.current_meta = None
+                demux.current_meta = demux.current_starts = None
                 demux.dbg_last_error = f"{type(exc).__name__}: {exc}"
+            _run_on_meta(demux)
             return input_batch
 
         model_runner.prepare_inputs = _patched
@@ -229,10 +255,14 @@ def _ensure_patched(worker: object, demux: _Demux) -> None:
             try:
                 nst = args[0] if args else kwargs.get("num_scheduled_tokens")
                 demux.current_meta = _meta_from_v1_inputs(model_runner.input_batch, nst)
+                demux.current_starts = _starts(
+                    getattr(model_runner.input_batch, "num_computed_tokens_cpu", None), demux.current_meta
+                )
                 demux.dbg_last_meta = demux.current_meta
             except Exception as exc:  # noqa: BLE001 - never break the forward on metadata errors
-                demux.current_meta = None
+                demux.current_meta = demux.current_starts = None
                 demux.dbg_last_error = f"{type(exc).__name__}: {exc}"
+            _run_on_meta(demux)
             return out
 
         model_runner._prepare_inputs = _patched

{interp_engine-1.13.0 → interp_engine-1.13.2}/interp_engine/vllm_capture/mhc.py

@@ -648,9 +648,13 @@ class _MhcTaps:
         detectably different from no steer at all. So the rows the recorder actually changed take the
         re-run and every other row keeps the kernel's own output, which is also why a zero delta needs
         no re-run: no row changed.
+
+        Inside a graph capture the host cannot read ``changed``, so the re-run is always recorded and
+        the row mask alone decides at replay.
         """
         changed = (edited != stack).flatten(start_dim=1).any(dim=1)
-        if not bool(changed.any()):
+        capturing = stack.is_cuda and torch.cuda.is_current_stream_capturing()
+        if not capturing and not bool(changed.any()):
             return fused
         bound = signature.bind(*args, **kwargs)
         bound.apply_defaults()

{interp_engine-1.13.0 → interp_engine-1.13.2}/interp_engine/vllm_capture/requests.py

@@ -86,6 +86,32 @@ def _process_point(demux: _Demux, site: Address, full: torch.Tensor) -> torch.Te
     if meta is None:
         return full
     req_ids, seq_lens = meta
+    if _batch_axis(demux, site, full, sum(seq_lens)):
+        rows = full[0]
+        new = _process_rows(demux, site, rows, req_ids, seq_lens)
+        return full if new is rows else new.unsqueeze(0)
+    return _process_rows(demux, site, full, req_ids, seq_lens)
+
+
+def _batch_axis(demux: _Demux, site: Address, full: torch.Tensor, tokens: int) -> bool:
+    """Whether ``full`` is ``[1, tokens, ...]`` rather than tokens first.
+
+    vLLM's Transformers backend calls the HF model with a batch axis, and so do some native families
+    at attention. For a forward of more than one token the shape decides it. For one token it does
+    not: ``[1, 1, d_model]`` with a batch axis has the shape of ``[1, 1, head_dim]`` on a model with one
+    KV head, or of one token on a hyper-connection trunk. So the site keeps the answer from its last
+    forward of more than one token, and a request's prefill always comes before its decode.
+    """
+    if full.dim() < 3 or full.shape[0] != 1:
+        return False
+    if tokens > 1:
+        demux.batch_axis[site] = full.shape[1] >= tokens
+    return demux.batch_axis.get(site, False)
+
+
+def _process_rows(
+    demux: _Demux, site: Address, full: torch.Tensor, req_ids: list[str], seq_lens: list[int]
+) -> torch.Tensor:
     if full.dim() not in (2, 3) or full.shape[0] < sum(seq_lens):
         # Fewer rows than the layout expects (unexpected shape); skip rather than corrupt.
         # More rows is fine: trailing padding after sum(seq_lens) is ignored by the loop.
@@ -266,10 +292,8 @@ def _mk_value_hook(demux: _Demux, site: Address):
     (``_hooks.flat_value``), because a norm over ``head_dim`` is handed the per-head view.
 
     The narrowed slice goes through ``flat_value`` too, and not only for symmetry: where vLLM leaves a
-    batch axis on the hidden state it hands the projection, the slice is ``[1, tokens, width]``, whose
-    first axis is one rather than the token count. :func:`_process_point` reads that as a trunk with
-    one row and returns the tensor untouched rather than steering rows it cannot line up -- so on those
-    families this point was quietly serving nothing at all.
+    batch axis on the hidden state it hands the projection, the slice is ``[1, tokens, width]``, and
+    the packed projection's geometry is what tells that apart from a per-head view.
 
     Either way the adjustment happens *before* steering and is undone after it, so a steer on this
     point writes the value in the shape a capture of it reports, leaves the queries and keys of the same

{interp_engine-1.13.0 → interp_engine-1.13.2}/interp_engine/vllm_capture/static.py

@@ -50,6 +50,7 @@ from interp_engine.vllm_capture._tree import (
     resolve_capture_module,
     scale_capture,
 )
+from interp_engine.vllm_capture.static_program import StaticWriteProgram
 
 #: Post-RoPE q/k/v at ``self_attn.attn``. Not a public capture point; ``capture_attention``
 #: harvests the three roles and recomputes scores/probs.
@@ -749,6 +750,9 @@ class StaticState:
     #: read-out or collect raises this rather than return a sequence with a hole in it.
     failed: dict[str, str] = field(default_factory=dict)
     patched_execute: bool = False
+    #: The write sites' device tables on a CUDA worker. When set, every write goes through it, so a
+    #: FULL graph replays each request's own write. None on CPU, where the Python path serves.
+    program: StaticWriteProgram | None = None
 
 
 def _state(worker: object) -> StaticState | None:
@@ -1156,7 +1160,11 @@ def worker_install_static(worker: object) -> None:
     if mhc_sites:
         _install_mhc_static(worker, mhc_sites)
 
-    _ensure_patched(worker, _get_demux(worker))
+    demux = _get_demux(worker)
+    _ensure_patched(worker, demux)
+    if state.writes and device.type == "cuda":
+        state.program = StaticWriteProgram(sites=list(state.writes.values()), device=device, max_n=max_n)
+        demux.on_meta.append(lambda: _fill_program_rows(worker))
     _patch_execute_model(worker)
     logger.info(
         "interp-engine static: %d read site(s), %d write site(s), max_n=%d",
@@ -1166,6 +1174,25 @@ def worker_install_static(worker: object) -> None:
     )
 
 
+def _fill_program_rows(worker: object) -> None:
+    """Point the next forward's rows at their writers. On any failure, no row is written."""
+    static = _state(worker)
+    program = static.program if static is not None else None
+    if program is None:
+        return
+    demux = _get_demux(worker)
+    meta = demux.current_meta
+    try:
+        if meta is None:
+            program.fill_rows([], [], None, lambda rid: rid)
+        else:
+            program.fill_rows(meta[0], meta[1], demux.current_starts, lambda rid: _resolve_rid(demux, rid))
+    except Exception:
+        logger.exception("static writes: could not map this step's rows to their writers; writing none")
+        program.fill_rows([], [], None, lambda rid: rid)
+        raise
+
+
 def _install_mhc_static(worker: object, mhc_sites: Sequence[_Site]) -> None:
     """Wrap the mHC kernels for every stacked static site, refusing one this model cannot serve.
 
@@ -1198,9 +1225,14 @@ def _static_mhc_recorder(worker: object, site: _Site):
         if n == 0:
             return tensor
         static = _state(worker)
-        reqs = _write_reqs_for(static, site) if static is not None else []
+        program = static.program if static is not None else None
         live = tensor
-        if reqs or site.modify is not None or site.delta_set:
+        if program is not None and program.has_site(site):
+            # A recorded graph must hold the edit whether or not anything is registered now.
+            if torch.cuda.is_current_stream_capturing() or not program.idle(site):
+                live = tensor.clone()
+                program.launch(site, live, None, n)
+        elif (static is not None and _write_reqs_for(static, site)) or site.modify is not None or site.delta_set:
             live = tensor.clone()
             _apply_write(live, None, site, n, fused=False, worker=worker)
         if site.buf is not None:
@@ -1637,6 +1669,12 @@ def _apply_write(
     if n <= 0:
         return
     static = _state(worker) if worker is not None else None
+    program = static.program if static is not None else None
+    if program is not None and program.has_site(site):
+        # A recorded graph launches the kernel on every replay; an eager forward may skip an idle site.
+        if torch.cuda.is_current_stream_capturing() or not program.idle(site):
+            program.launch(site, hidden, residual if fused else None, n)
+        return
     reqs = _write_reqs_for(static, site) if static is not None else []
     if reqs:
         _apply_demuxed_writes(hidden, residual, n, fused, worker, reqs)
@@ -1868,9 +1906,11 @@ def worker_set_static_delta(
 ) -> None:
     """Install static writes from worker specs. Zeros every write site first.
 
-    Additive ``op="additive"`` without a scope fills the static ``delta`` buffer. Every other op attaches
-    a live ``modify`` that reads the residual each forward (breakable ``add_eager``). ``lens_scope``
-    is the write's position scope: ``skip_positions``, ``prompt_len`` and ``steer_generated``.
+    On a CUDA worker the ops go to :class:`~interp_engine.vllm_capture.static_program.StaticWriteProgram`
+    as its global write. Without one, additive ``op="additive"`` without a scope fills the static
+    ``delta`` buffer, and every other op attaches a live ``modify`` that reads the residual each
+    forward. ``lens_scope`` is the write's position scope: ``skip_positions``, ``prompt_len`` and
+    ``steer_generated``.
 
     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
@@ -1885,7 +1925,10 @@ def worker_set_static_delta(
     worker_clear_static_delta(worker)
     from interp_engine.vllm_capture.steering import _make_steer_modifiers
 
-    for site, group in _group_by_site(static, specs).items():
+    groups = _group_by_site(static, specs)
+    if static.program is not None:
+        static.program.set_global(groups, lens_scope)
+    for site, group in groups.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:
@@ -1969,8 +2012,11 @@ def worker_register_static_write(
         skip = tuple(int(i) for i in (lens_scope.get("skip_positions") or skip))
         length = int(lens_scope.get("prompt_len") or length)
         generated = bool(lens_scope.get("steer_generated", False))
+    groups = _group_by_site(static, specs)
+    if static.program is not None:
+        static.program.register(req_id, groups, skip=skip, prompt_len=length, generated=generated)
     by_site: dict[str, _WriteReq] = {}
-    for site, group in _group_by_site(static, specs).items():
+    for site, group in groups.items():
         by_site[format_address(site.address)] = _compile_write_req(
             group, site, skip_positions=skip, prompt_len=length, steer_generated=generated
         )
@@ -1984,6 +2030,8 @@ def worker_unregister_static_write(worker: object, req_id: str) -> None:
     if static is None:
         return
     static.write_reqs.pop(req_id, None)
+    if static.program is not None:
+        static.program.unregister(req_id)
     if req_id not in static.cap_points:
         static.registered.discard(req_id)
 
@@ -1992,6 +2040,8 @@ def worker_clear_static_delta(worker: object) -> None:
     static = _state(worker)
     if static is None:
         return
+    if static.program is not None:
+        static.program.clear_global()
     for site in static.writes.values():
         site.modify = None
         site.lens_scope = None