feat(capture): graph-safe global capture for the filesystem consumer

docs/design/capture_consumers.md

@@ -595,12 +595,26 @@ without buffering the whole tensor. Owns a private `ActivationWriter`
 thread pool (`writer.py`).
 
 - `location = "worker"`, `reads_client_spec = True`.
-- `global_capture_spec()` returns `None` — captures are always
-  per-request via `SamplingParams.capture["filesystem"]`.
+- `global_capture_spec()` returns `None` by default (captures are
+  per-request via `SamplingParams.capture["filesystem"]`), **or** a
+  `CaptureSpec` built from the consumer-level `global_hooks` /
+  `global_positions` params when those are configured. A global spec
+  captures every request uniformly and rides the manager's
+  CUDA-graph-safe persistent-buffer path, so filesystem capture no
+  longer forces eager every step (the dominant cost at
+  `positions="all_generated"` under cudagraph). Global-driven requests
+  have no per-request `tag`/`request_id`, so files are named
+  `{root}/{default_tag}/{engine_request_id}/{layer}_{hook}.bin` — the
+  engine request id is already threaded through the manager dispatch
+  path (`CapturePositionEntry.request_id` → `CaptureChunk.key[0]` →
+  the finalize key), so no extra plumbing is needed. `default_tag`
+  defaults to `"default"` (the legacy fallback directory name).
 - `validate_client_spec` accepts `FilesystemCaptureRequest` or a
   matching dict, then lazily delegates to
   `validation.validate_filesystem_request` (lazy to avoid pulling
-  pydantic in at module import).
+  pydantic in at module import). Per-request client specs and the
+  global spec coexist: a client spec for a request overrides the
+  global spec for that request (manager merge rule).
 
 Per-chunk flow:
 

tests/v1/capture/consumers/filesystem/test_consumer.py

@@ -55,15 +55,20 @@ def _make_vllm_config(
     *,
     root_path: str | None = "/tmp/activations",
     max_bytes: int = 0,
+    num_hidden_layers: int = 32,
 ) -> MagicMock:
     """Build a minimal mock ``VllmConfig`` for the filesystem consumer.
 
     ``root_path`` / ``max_bytes`` are accepted for backwards compatibility
     with older call sites but are no longer consulted by the validator;
     the consumer receives its ``root`` via its constructor ``params``.
+    ``num_hidden_layers`` backs ``model_config.get_total_num_hidden_layers()``,
+    which the consumer reads to resolve ``global_hooks`` (e.g. ``"all"``).
     """
     del root_path, max_bytes
-    return MagicMock()
+    cfg = MagicMock()
+    cfg.model_config.get_total_num_hidden_layers.return_value = num_hidden_layers
+    return cfg
 
 
 def _make_context(
@@ -666,6 +671,166 @@ def test_global_capture_spec_is_none(self, tmp_path: pathlib.Path) -> None:
             consumer.shutdown(timeout=5.0)
 
 
+class TestGlobalCaptureSpec:
+    """Consumer-level global capture spec (CUDA-graph-safe path)."""
+
+    def test_none_without_global_hooks(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(tmp_path)
+        try:
+            assert consumer.global_capture_spec() is None
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_returns_configured_spec(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(
+            tmp_path,
+            global_hooks={"post_mlp": [0, 2]},
+            global_positions="all_generated",
+        )
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.hooks == {"post_mlp": [0, 2]}
+            assert spec.positions == "all_generated"
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_default_positions_all_prompt(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(tmp_path, global_hooks={"post_mlp": [1]})
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.positions == "all_prompt"
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_explicit_positions_list(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(
+            tmp_path,
+            global_hooks={"pre_attn": [0]},
+            global_positions=[0, 1, 2],
+        )
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.positions == [0, 1, 2]
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_layers_sorted_deduped(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(tmp_path, global_hooks={"post_attn": [3, 1, 1, 0]})
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.hooks == {"post_attn": [0, 1, 3]}
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_empty_global_hooks_disables(self, tmp_path: pathlib.Path) -> None:
+        # An empty dict is treated as "not configured" → no global spec.
+        consumer = _make_consumer(tmp_path, global_hooks={})
+        try:
+            assert consumer.global_capture_spec() is None
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_invalid_hook_name_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="not a valid hook point"):
+            _make_consumer(tmp_path, global_hooks={"bogus": [0]})
+
+    def test_non_dict_global_hooks_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="must be a dict"):
+            _make_consumer(tmp_path, global_hooks=[0, 1])
+
+    def test_empty_layer_list_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="resolved to no layers"):
+            _make_consumer(tmp_path, global_hooks={"post_mlp": []})
+
+    def test_negative_layer_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="non-negative"):
+            _make_consumer(tmp_path, global_hooks={"post_mlp": [-1]})
+
+    def test_non_int_layer_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="must be ints"):
+            _make_consumer(tmp_path, global_hooks={"post_mlp": ["x"]})
+
+    def test_default_tag_configurable(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(
+            tmp_path,
+            global_hooks={"post_mlp": [0]},
+            default_tag="run42",
+        )
+        try:
+            assert consumer._params.default_tag == "run42"
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_default_tag_defaults_to_default(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(tmp_path)
+        try:
+            assert consumer._params.default_tag == "default"
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    # ---- string DSL + range / all / dot-list (CLI-safe forms) ----
+
+    def test_string_dsl_multiple_hooks(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(
+            tmp_path, global_hooks="pre_attn:0-2;post_mlp:20"
+        )
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.hooks == {"pre_attn": [0, 1, 2], "post_mlp": [20]}
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_string_dsl_all(self, tmp_path: pathlib.Path) -> None:
+        # num_hidden_layers defaults to 32 in the fixture.
+        consumer = _make_consumer(tmp_path, global_hooks="post_mlp:all")
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.hooks == {"post_mlp": list(range(32))}
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_dict_value_range_string(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(tmp_path, global_hooks={"post_mlp": "0-3"})
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.hooks == {"post_mlp": [0, 1, 2, 3]}
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_dict_value_dot_list_and_all(self, tmp_path: pathlib.Path) -> None:
+        consumer = _make_consumer(
+            tmp_path, global_hooks={"post_mlp": "1.5.9", "pre_attn": "all"}
+        )
+        try:
+            spec = consumer.global_capture_spec()
+            assert spec is not None
+            assert spec.hooks["post_mlp"] == [1, 5, 9]
+            assert spec.hooks["pre_attn"] == list(range(32))
+        finally:
+            consumer.shutdown(timeout=5.0)
+
+    def test_layer_out_of_range_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="out of range"):
+            _make_consumer(tmp_path, global_hooks={"post_mlp": [99]})
+
+    def test_string_dsl_missing_colon_rejected(
+        self, tmp_path: pathlib.Path
+    ) -> None:
+        with pytest.raises(ValueError, match="must be '<hook>:<layers>'"):
+            _make_consumer(tmp_path, global_hooks="post_mlp")
+
+    def test_bad_layer_spec_string_rejected(self, tmp_path: pathlib.Path) -> None:
+        with pytest.raises(ValueError, match="layer spec"):
+            _make_consumer(tmp_path, global_hooks={"post_mlp": "1,2,3"})
+
+
 class TestClassVars:
     """Verify the ClassVar metadata on FilesystemConsumer."""
 

tests/v1/capture/test_runner_integration.py

@@ -63,12 +63,24 @@ def _wait_for_status(
     pytest.fail(f"timeout waiting for {key} to finalize")
 
 
+class _FakeModelConfig:
+    """Stand-in exposing the layer count the consumer reads to resolve
+    ``global_hooks`` (e.g. ``"all"``)."""
+
+    def __init__(self, num_hidden_layers: int = 32) -> None:
+        self._num_hidden_layers = num_hidden_layers
+
+    def get_total_num_hidden_layers(self) -> int:
+        return self._num_hidden_layers
+
+
 class _FakeVllmConfig:
     """Minimal stand-in for ``VllmConfig`` — enough for the filesystem
     consumer's constructor to run without pulling in pydantic."""
 
-    def __init__(self) -> None:
+    def __init__(self, num_hidden_layers: int = 32) -> None:
         self.capture_consumers_config = None
+        self.model_config = _FakeModelConfig(num_hidden_layers)
 
 
 # ---------------------------------------------------------------------------
@@ -155,6 +167,150 @@ def test_filesystem_consumer_end_to_end_via_manager(tmp_path: pathlib.Path) -> N
     assert sidecar["hook"] == "post_mlp"
 
 
+# ---------------------------------------------------------------------------
+# 1b. Global-spec driven capture — the CUDA-graph-safe persistent-buffer
+#     path. No per-request client spec; the consumer advertises a global
+#     spec and files are named by the engine request id under default_tag.
+# ---------------------------------------------------------------------------
+
+
+def test_global_spec_drives_per_request_files(tmp_path: pathlib.Path) -> None:
+    """A consumer-level global spec captures every request via the
+    persistent-buffer path and still writes one file set per request,
+    keyed by the engine request id under the configured ``default_tag``.
+    """
+    vllm_config = _FakeVllmConfig()
+    consumer = FilesystemConsumer(
+        vllm_config=vllm_config,
+        params={
+            "root": str(tmp_path),
+            "writer_threads": 2,
+            "global_hooks": {"post_mlp": [1]},
+            "global_positions": "last_prompt",
+            "default_tag": "run-global",
+        },
+    )
+
+    # The runner installs the consumer's global spec into the manager.
+    global_spec = consumer.global_capture_spec()
+    assert global_spec is not None
+
+    mgr = CaptureManager(
+        consumers=(consumer,),
+        consumer_specs=(global_spec,),
+        num_hidden_layers=4,
+        hidden_size=8,
+        model_dtype=torch.float32,
+        device="cpu",
+        # Engage the CUDA-graph-safe persistent-buffer path: with
+        # max_num_tokens>0 the global key is served from a persistent buffer
+        # instead of the eager dynamic gather.
+        max_num_tokens=16,
+    )
+
+    # The global key gets a persistent buffer (the graph-safe path), and
+    # routes to the global gather (not the dynamic index_select).
+    assert mgr._global_keys == frozenset({(1, "post_mlp")})
+
+    req_id = "req-global-1"
+    # No client_specs and no admission slugs — purely global-driven.
+    mgr.register_request(
+        req_id,
+        client_specs=None,
+        num_prompt_tokens=3,
+        sidecar_fields={"vllm_internal_request_id": req_id},
+    )
+
+    batch_view = CaptureBatchView(
+        req_ids=[req_id],
+        num_prompt_tokens=[3],
+        num_computed_tokens=[0],
+        num_scheduled_tokens=[3],
+        token_offsets=[0],
+    )
+    plan = mgr.build_step_plan(batch_view)
+    # Global key routed to the buffer path, not the dynamic gather.
+    assert (1, "post_mlp") in plan.global_gather_indices
+    assert (1, "post_mlp") not in plan.gather_indices
+
+    # Simulate the graph-recorded full-residual copy into the persistent
+    # buffer (what on_hook does for a global key).
+    hidden = torch.arange(24, dtype=torch.float32).reshape(3, 8)
+    mgr.on_hook(1, "post_mlp", hidden)
+
+    mgr.dispatch_step_captures(plan)
+    results = mgr.finalize_request(req_id)
+    assert list(results.keys()) == [0]
+
+    _wait_for_status(consumer, (req_id, 1, "post_mlp"))
+    consumer.shutdown()
+
+    # File named by the engine request id under the configured tag.
+    bin_path = tmp_path / "run-global" / req_id / "1_post_mlp.bin"
+    sidecar_path = bin_path.with_suffix(".json")
+    assert bin_path.exists(), f"missing bin file {bin_path}"
+    assert sidecar_path.exists(), f"missing sidecar {sidecar_path}"
+
+    sidecar = json.loads(sidecar_path.read_text())
+    assert sidecar["request_id"] == req_id
+    assert sidecar["layer"] == 1
+    assert sidecar["hook"] == "post_mlp"
+
+    # The captured row is the last prompt position of the residual.
+    captured = torch.frombuffer(bytearray(bin_path.read_bytes()), dtype=torch.float32)
+    assert torch.equal(captured, hidden[-1])
+
+
+def test_global_spec_two_requests_distinct_dirs(tmp_path: pathlib.Path) -> None:
+    """Two global-driven requests land in separate request directories."""
+    consumer = FilesystemConsumer(
+        vllm_config=_FakeVllmConfig(),
+        params={
+            "root": str(tmp_path),
+            "writer_threads": 2,
+            "global_hooks": {"post_mlp": [0]},
+            "global_positions": "last_prompt",
+        },
+    )
+    mgr = CaptureManager(
+        consumers=(consumer,),
+        consumer_specs=(consumer.global_capture_spec(),),
+        num_hidden_layers=2,
+        hidden_size=4,
+        model_dtype=torch.float32,
+        device="cpu",
+        max_num_tokens=16,
+    )
+
+    for req_id, base in (("req-A", 0), ("req-B", 100)):
+        mgr.register_request(
+            req_id,
+            client_specs=None,
+            num_prompt_tokens=2,
+            sidecar_fields={"vllm_internal_request_id": req_id},
+        )
+        batch_view = CaptureBatchView(
+            req_ids=[req_id],
+            num_prompt_tokens=[2],
+            num_computed_tokens=[0],
+            num_scheduled_tokens=[2],
+            token_offsets=[0],
+        )
+        plan = mgr.build_step_plan(batch_view)
+        hidden = torch.arange(base, base + 8, dtype=torch.float32).reshape(2, 4)
+        mgr.on_hook(0, "post_mlp", hidden)
+        mgr.dispatch_step_captures(plan)
+        mgr.finalize_request(req_id)
+        _wait_for_status(consumer, (req_id, 0, "post_mlp"))
+
+    consumer.shutdown()
+
+    # Default tag is "default" (legacy fallback name); each request gets its
+    # own directory keyed by the engine request id.
+    assert (tmp_path / "default" / "req-A" / "0_post_mlp.bin").exists()
+    assert (tmp_path / "default" / "req-B" / "0_post_mlp.bin").exists()
+
+
 # ---------------------------------------------------------------------------
 # 2. Plan-level admission error is surfaced as a terminal result.
 # ---------------------------------------------------------------------------