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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.13.2
+Version: 1.14.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.13.2 → interp_engine-1.14.0}/docs/INTERNALS.md

@@ -83,6 +83,12 @@ files and why they are where they are.
   child dies before it loads anything, and the caller is told to see a root cause that is in another
   process. Called from `_ensure_engine`; a no-op in a script, a server, or a kernel that captured
   the real descriptors itself.
+- `vllm_seq_fit.py` — fits `max_num_seqs` to the Mamba state cache on a hybrid trunk. With full
+  decode CUDA graphs, vLLM needs one Mamba block per sequence, and it knows the block count only
+  inside its EngineCore child, after it measures memory. The worker plugin wraps vLLM's check to
+  write that count to a temporary file; `VLLMModel._start_engine` reads it after the failed start
+  and starts once more with `max_num_seqs` set to it. A caller's explicit `max_num_seqs` gets no
+  retry. `tests/test_vllm_seq_fit.py` covers both halves on CPU.
 
 ## Correctness
 

{interp_engine-1.13.2 → interp_engine-1.14.0}/interp_engine/vllm_backend.py

@@ -18,9 +18,11 @@ macOS/CPU.
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import importlib.util
 import logging
 import os
+import tempfile
 import warnings
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence
 from typing import Any, TypeVar
@@ -82,6 +84,7 @@ from interp_engine.vllm_capture.static import (
     static_unsupported_reason,
 )
 from interp_engine.vllm_plugin import WORKER_EXTENSION_CLS
+from interp_engine.vllm_seq_fit import SEQ_FIT_ENV, read_block_limit
 
 logger = logging.getLogger(__name__)
 
@@ -1128,6 +1131,7 @@ class VLLMModel:
             extra_vllm_kwargs=extra_vllm_kwargs,
             max_lora_rank=max_lora_rank,
         )
+        self._max_num_seqs_pinned = "max_num_seqs" in (extra_vllm_kwargs or {})
         facts = read_residual_facts(hf_model_id, trust_remote_code)
         self._attn_dims = read_attn_dims(hf_model_id, trust_remote_code)
         reads, writes, graph = resolve_static_points(
@@ -1389,10 +1393,43 @@ class VLLMModel:
                     # The line below forks a child that suppresses stdout by descriptor, which a
                     # notebook kernel's stdout does not have. See `notebook_stdout`.
                     ensure_stdout_descriptor()
-                    self.engine = AsyncLLM.from_engine_args(AsyncEngineArgs(**self._engine_kwargs))
+                    self.engine = self._start_engine(AsyncLLM, AsyncEngineArgs)
                     self._engine_loop = asyncio.get_running_loop()
         return self.engine
 
+    def _start_engine(self, async_llm: Any, engine_args: Any) -> Any:
+        """Build the engine. On a hybrid trunk, start again once if ``max_num_seqs`` does not fit.
+
+        See :mod:`interp_engine.vllm_seq_fit`. The retry happens only when a worker recorded the
+        Mamba block count and the caller did not set ``max_num_seqs``.
+        """
+        if getattr(self, "_max_num_seqs_pinned", False):
+            return async_llm.from_engine_args(engine_args(**self._engine_kwargs))
+        fd, path = tempfile.mkstemp(prefix="interp-engine-seq-fit-")
+        os.close(fd)
+        os.environ[SEQ_FIT_ENV] = path
+        try:
+            try:
+                return async_llm.from_engine_args(engine_args(**self._engine_kwargs))
+            except Exception:
+                blocks = read_block_limit(path)
+                if blocks is None:
+                    raise
+            logger.warning(
+                "%s: max_num_seqs %s exceeds the %s Mamba cache blocks left after memory profiling; "
+                "starting the engine again with max_num_seqs=%s. Set max_num_seqs to skip the retry.",
+                self.hf_model_id,
+                self._engine_kwargs.get("max_num_seqs", "(vLLM default)"),
+                blocks,
+                blocks,
+            )
+            self._engine_kwargs["max_num_seqs"] = blocks
+            return async_llm.from_engine_args(engine_args(**self._engine_kwargs))
+        finally:
+            os.environ.pop(SEQ_FIT_ENV, None)
+            with contextlib.suppress(OSError):
+                os.unlink(path)
+
     async def warmup(self) -> None:
         """Build the engine now instead of on the first request, and run one throwaway decode.
 

{interp_engine-1.13.2 → interp_engine-1.14.0}/interp_engine/vllm_plugin.py

@@ -80,11 +80,13 @@ from interp_engine.vllm_capture import (
     worker_unregister_steering,
 )
 from interp_engine.vllm_capture.static import patch_worker_for_static
+from interp_engine.vllm_seq_fit import patch_mamba_seq_check
 
 # Worker process imports this module for worker_extension_cls. Patching load_model here
 # (not as an extension method named load_model) is what lets static wraps run before
 # CUDA-graph capture. vLLM refuses an extension attribute that collides with Worker.
 patch_worker_for_static()
+patch_mamba_seq_check()
 
 # Pass as ``worker_extension_cls=`` when constructing your engine. vLLM resolves this
 # qualified name inside each worker process, so it must stay importable under this path.

interp_engine-1.14.0/interp_engine/vllm_seq_fit.py

@@ -0,0 +1,85 @@
+"""Fit ``max_num_seqs`` to the Mamba state cache, so a hybrid trunk starts without a hand-set value.
+
+With full decode CUDA graphs, vLLM needs one Mamba cache block per sequence, up to
+``max_num_seqs`` (128 by default). It knows the block count only after it measures memory,
+inside the engine-core child, and refuses to start when there are fewer blocks than sequences.
+Static attention buffers take memory from that cache, so a config that starts without them can
+fail with them. The parent sees only "Engine core initialization failed".
+
+The worker half wraps vLLM's check: when it refuses, the worker writes the block count to the
+file named by :data:`SEQ_FIT_ENV`, then raises as before. The parent half reads that file after
+the failed start and starts again with ``max_num_seqs`` set to the count. A caller who sets
+``max_num_seqs`` gets no retry: an explicit value wins, and fails the same way it did before.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+from typing import Any
+
+# Path of the file the worker writes the Mamba block count to. The parent sets it before it
+# starts the engine, and the engine-core and worker processes inherit it.
+SEQ_FIT_ENV = "INTERP_ENGINE_SEQ_FIT_FILE"
+
+
+def mamba_block_limit(kv_cache_config: Any, max_num_reqs: int | None) -> int | None:
+    """Return the Mamba block count when it is below ``max_num_reqs``, else ``None``."""
+    if kv_cache_config is None or max_num_reqs is None:
+        return None
+    if not getattr(kv_cache_config, "has_mamba_layers", False):
+        return None
+    blocks = int(getattr(kv_cache_config, "num_blocks", 0))
+    if blocks <= 0 or max_num_reqs <= blocks:
+        return None
+    return blocks
+
+
+def record_block_limit(blocks: int) -> None:
+    """Write ``blocks`` to the file named by :data:`SEQ_FIT_ENV`. Do nothing when it is unset."""
+    path = os.environ.get(SEQ_FIT_ENV)
+    if not path:
+        return
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(str(blocks))
+
+
+def read_block_limit(path: str) -> int | None:
+    """Return the block count a worker wrote to ``path``, or ``None`` when none did."""
+    try:
+        with open(path, encoding="utf-8") as f:
+            text = f.read().strip()
+    except OSError:
+        return None
+    return int(text) if text.isdigit() and int(text) > 0 else None
+
+
+def patch_mamba_seq_check() -> None:
+    """Wrap ``CompilationConfig.resolve_cudagraph_mode_and_sizes`` to record the block count.
+
+    The wrap changes no result. It only writes the count before vLLM's own error propagates.
+    """
+    try:
+        from vllm.config.compilation import CompilationConfig  # pyright: ignore[reportMissingImports]
+    except Exception:  # noqa: BLE001
+        return
+    orig = CompilationConfig.resolve_cudagraph_mode_and_sizes
+    if getattr(orig, "_ie_seq_fit", False):
+        return
+    signature = inspect.signature(orig)
+
+    def resolve_cudagraph_mode_and_sizes(self: Any, *args: Any, **kwargs: Any) -> Any:
+        try:
+            return orig(self, *args, **kwargs)
+        except ValueError:
+            try:
+                bound = dict(signature.bind(self, *args, **kwargs).arguments)
+            except TypeError:
+                bound = {}
+            blocks = mamba_block_limit(bound.get("kv_cache_config"), bound.get("max_num_reqs"))
+            if blocks is not None:
+                record_block_limit(blocks)
+            raise
+
+    resolve_cudagraph_mode_and_sizes._ie_seq_fit = True  # pyright: ignore[reportFunctionMemberAccess]
+    CompilationConfig.resolve_cudagraph_mode_and_sizes = resolve_cudagraph_mode_and_sizes

{interp_engine-1.13.2 → interp_engine-1.14.0}/pyproject.toml

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

{interp_engine-1.13.2 → interp_engine-1.14.0}/tests/test_static_set.py

@@ -776,8 +776,9 @@ def test_ensure_engine_sets_breakable_beside_static_env():
     assert "apply_breakable_env" in src
     static_at = src.index("STATIC_ENV")
     breakable_at = src.index("apply_breakable_env")
-    engine_at = src.index("from_engine_args")
+    engine_at = src.index("_start_engine")
     assert static_at < breakable_at < engine_at
+    assert "from_engine_args" in inspect.getsource(VLLMModel._start_engine)
 
 
 def test_a_read_buffer_is_not_used_as_a_boolean_when_sizing_the_write():

interp_engine-1.14.0/tests/test_vllm_seq_fit.py

@@ -0,0 +1,191 @@
+"""A hybrid trunk whose Mamba cache holds fewer blocks than ``max_num_seqs`` starts on the retry.
+
+vLLM measures the block count inside its engine-core child and refuses to start when there are
+fewer blocks than sequences. The worker records the count, and the parent starts again with
+``max_num_seqs`` set to it. Stand-ins replace vLLM: what is under test is the handoff between the
+two halves, which needs no GPU.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import types
+from collections.abc import Iterator
+from typing import Any
+
+import pytest
+
+from interp_engine import vllm_seq_fit
+from interp_engine.vllm_backend import VLLMModel
+from interp_engine.vllm_seq_fit import (
+    SEQ_FIT_ENV,
+    mamba_block_limit,
+    patch_mamba_seq_check,
+    read_block_limit,
+    record_block_limit,
+)
+
+
+class _KVConfig:
+    def __init__(self, num_blocks: int, has_mamba_layers: bool = True) -> None:
+        self.num_blocks = num_blocks
+        self.has_mamba_layers = has_mamba_layers
+
+
+def test_block_limit_only_when_mamba_blocks_are_short() -> None:
+    assert mamba_block_limit(_KVConfig(88), 128) == 88
+    assert mamba_block_limit(_KVConfig(128), 128) is None, "enough blocks"
+    assert mamba_block_limit(_KVConfig(88, has_mamba_layers=False), 128) is None, "no Mamba layers"
+    assert mamba_block_limit(None, 128) is None
+    assert mamba_block_limit(_KVConfig(88), None) is None
+
+
+def test_record_and_read_round_trip(tmp_path: Any, monkeypatch: pytest.MonkeyPatch) -> None:
+    path = str(tmp_path / "fit")
+    monkeypatch.setenv(SEQ_FIT_ENV, path)
+    record_block_limit(88)
+    assert read_block_limit(path) == 88
+
+
+def test_read_ignores_a_missing_or_empty_file(tmp_path: Any) -> None:
+    empty = tmp_path / "empty"
+    empty.write_text("")
+    assert read_block_limit(str(empty)) is None
+    assert read_block_limit(str(tmp_path / "absent")) is None
+
+
+def test_record_without_the_env_writes_nothing(tmp_path: Any, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv(SEQ_FIT_ENV, raising=False)
+    monkeypatch.chdir(tmp_path)
+    record_block_limit(88)
+    assert list(tmp_path.iterdir()) == []
+
+
+@pytest.fixture
+def fake_compilation(monkeypatch: pytest.MonkeyPatch) -> Iterator[type]:
+    """Install a stand-in ``vllm.config.compilation`` whose check matches vLLM's signature."""
+
+    class CompilationConfig:
+        def resolve_cudagraph_mode_and_sizes(
+            self,
+            min_cg_support: Any,
+            min_cg_attn_backend: str | None,
+            uniform_decode_query_len: int = 1,
+            use_v2_model_runner: bool = False,
+            tensor_parallel_size: int = 1,
+            kv_cache_config: Any = None,
+            max_num_reqs: int | None = None,
+            is_profiling: bool = False,
+        ) -> str:
+            blocks = mamba_block_limit(kv_cache_config, max_num_reqs)
+            if blocks is not None:
+                raise ValueError(f"max_num_seqs ({max_num_reqs}) exceeds available Mamba cache blocks ({blocks}).")
+            if min_cg_support == "bad":
+                raise ValueError("unrelated")
+            return "FULL"
+
+    module = types.ModuleType("vllm.config.compilation")
+    module.CompilationConfig = CompilationConfig  # pyright: ignore[reportAttributeAccessIssue]
+    for name in ("vllm", "vllm.config"):
+        monkeypatch.setitem(sys.modules, name, sys.modules.get(name) or types.ModuleType(name))
+    monkeypatch.setitem(sys.modules, "vllm.config.compilation", module)
+    yield CompilationConfig
+
+
+def test_patched_check_records_the_block_count_and_still_raises(
+    fake_compilation: type, tmp_path: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    path = str(tmp_path / "fit")
+    monkeypatch.setenv(SEQ_FIT_ENV, path)
+    patch_mamba_seq_check()
+    patch_mamba_seq_check()
+
+    config = fake_compilation()
+    with pytest.raises(ValueError, match="exceeds available Mamba cache blocks"):
+        config.resolve_cudagraph_mode_and_sizes("ok", None, 1, kv_cache_config=_KVConfig(88), max_num_reqs=128)
+    assert read_block_limit(path) == 88
+
+
+def test_patched_check_changes_nothing_else(
+    fake_compilation: type, tmp_path: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    path = str(tmp_path / "fit")
+    monkeypatch.setenv(SEQ_FIT_ENV, path)
+    patch_mamba_seq_check()
+
+    config = fake_compilation()
+    assert (
+        config.resolve_cudagraph_mode_and_sizes("ok", None, kv_cache_config=_KVConfig(128), max_num_reqs=128) == "FULL"
+    )
+    with pytest.raises(ValueError, match="unrelated"):
+        config.resolve_cudagraph_mode_and_sizes("bad", None)
+    assert not os.path.exists(path), "no block count for an unrelated error"
+
+
+class _FakeAsyncLLM:
+    """Fails like vLLM while ``max_num_seqs`` exceeds ``blocks``, recording the count as the worker does."""
+
+    def __init__(self, blocks: int | None) -> None:
+        self.blocks = blocks
+        self.calls: list[dict[str, Any]] = []
+
+    def from_engine_args(self, args: dict[str, Any]) -> str:
+        self.calls.append(dict(args))
+        if self.blocks is not None and args.get("max_num_seqs", 128) > self.blocks:
+            record_block_limit(self.blocks)
+            raise RuntimeError("Engine core initialization failed. See root cause above.")
+        return "engine"
+
+
+def _model(extra: dict[str, Any] | None = None) -> VLLMModel:
+    model = object.__new__(VLLMModel)
+    model.hf_model_id = "stand-in/hybrid"
+    model._engine_kwargs = {"model": "stand-in/hybrid", **(extra or {})}
+    model._max_num_seqs_pinned = "max_num_seqs" in (extra or {})
+    return model
+
+
+def test_start_engine_retries_with_the_recorded_block_count() -> None:
+    llm = _FakeAsyncLLM(blocks=88)
+    model = _model()
+
+    assert model._start_engine(llm, dict) == "engine"
+    assert [c.get("max_num_seqs") for c in llm.calls] == [None, 88]
+    assert model._engine_kwargs["max_num_seqs"] == 88
+    assert SEQ_FIT_ENV not in os.environ
+
+
+def test_start_engine_starts_once_when_it_fits() -> None:
+    llm = _FakeAsyncLLM(blocks=None)
+
+    assert _model()._start_engine(llm, dict) == "engine"
+    assert len(llm.calls) == 1
+
+
+def test_start_engine_reraises_a_failure_with_no_block_count() -> None:
+    class Broken(_FakeAsyncLLM):
+        def from_engine_args(self, args: dict[str, Any]) -> str:
+            self.calls.append(dict(args))
+            raise RuntimeError("out of memory")
+
+    llm = Broken(blocks=None)
+    with pytest.raises(RuntimeError, match="out of memory"):
+        _model()._start_engine(llm, dict)
+    assert len(llm.calls) == 1
+
+
+def test_an_explicit_max_num_seqs_wins_and_is_not_retried() -> None:
+    llm = _FakeAsyncLLM(blocks=88)
+    model = _model({"max_num_seqs": 128})
+
+    with pytest.raises(RuntimeError, match="initialization failed"):
+        model._start_engine(llm, dict)
+    assert len(llm.calls) == 1
+    assert model._engine_kwargs["max_num_seqs"] == 128
+
+
+def test_the_worker_plugin_installs_the_patch() -> None:
+    import interp_engine.vllm_plugin as plugin
+
+    assert plugin.patch_mamba_seq_check is vllm_seq_fit.patch_mamba_seq_check