From df1630c3c7aa0ea0fd56c362f4089060f74be006 Mon Sep 17 00:00:00 2001 From: RhizoNymph Date: Thu, 18 Jun 2026 17:07:08 -0700 Subject: [PATCH 1/6] fix(config): fall back to v1 model runner when steering or capture is configured --- tests/test_config.py | 35 +++++++++++++++++++++++++++++++++++ vllm/config/vllm.py | 23 +++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index 5c01d652a17a..df40453a8f4a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -148,6 +148,41 @@ def test_is_default_v2_model_runner_model(model_config, expected): assert VllmConfig._is_default_v2_model_runner_model(config) is expected +# A non-None sentinel stands in for a configured SteeringConfig / +# CaptureConsumersConfig; the predicate only checks ``is not None``. +_ENABLED = object() + + +@pytest.mark.parametrize( + ("steering_config", "capture_config", "expected"), + [ + (None, None, []), + (_ENABLED, None, ["activation steering"]), + (None, _ENABLED, ["activation capture"]), + (_ENABLED, _ENABLED, ["activation steering", "activation capture"]), + ], +) +def test_v2_silently_broken_features(steering_config, capture_config, expected): + config = SimpleNamespace( + steering_config=steering_config, + capture_consumers_config=capture_config, + ) + + assert VllmConfig._v2_model_runner_silently_broken_features(config) == expected + + +def test_validate_v2_model_runner_fails_closed_on_silently_broken(monkeypatch): + # An explicit VLLM_USE_V2_MODEL_RUNNER opt-in is enforced here: a config the + # v2 runner would silently no-op must raise rather than be honored. + monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) + config = SimpleNamespace( + _get_v2_model_runner_unsupported_features=lambda: ["activation steering"], + ) + + with pytest.raises(ValueError, match="activation steering"): + VllmConfig._validate_v2_model_runner(config) + + @pytest.mark.skip_global_cleanup def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping( monkeypatch, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 53f8cf147f88..3c95e6ecf7a1 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -509,6 +509,10 @@ def num_speculative_tokens(self) -> int: def use_v2_model_runner(self) -> bool: use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER if use_v2_model_runner is not None: + # An explicit opt-in still runs through _validate_v2_model_runner(), + # which fails closed on _get_v2_model_runner_unsupported_features() + # (steering/capture included), so the override can't silently land + # on a runner that would no-op those features. return use_v2_model_runner if not self._is_default_v2_model_runner_model(): @@ -2068,8 +2072,27 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: # Will be added by https://github.com/vllm-project/vllm/pull/38390 unsupported.append("EC transfer") + unsupported.extend(self._v2_model_runner_silently_broken_features()) + return unsupported + def _v2_model_runner_silently_broken_features(self) -> list[str]: + """Features the v2 model runner would silently no-op rather than fail on. + + The activation steering and capture control planes are only wired into + the v1 model runner. The v2 runner loads the same model (so the layer + hooks exist) but never populates the steering tables or installs a + capture manager, so these features would silently do nothing instead of + raising. Callers must treat them as unsupported until the control plane + is ported to v2. + """ + features: list[str] = [] + if self.steering_config is not None: + features.append("activation steering") + if self.capture_consumers_config is not None: + features.append("activation capture") + return features + def _validate_v2_model_runner(self) -> None: """Check for features not yet supported by the V2 model runner.""" if not HAS_TRITON: From 2e680d3ddc59015a707621b76216290ae78b4786 Mon Sep 17 00:00:00 2001 From: RhizoNymph Date: Thu, 18 Jun 2026 17:07:08 -0700 Subject: [PATCH 2/6] feat(v2-runner): port activation steering and capture control planes to the v2 model runner --- docs/design/v2_runner_steering_capture.md | 135 ++++++ tests/v1/worker/test_gpu_v2_capture_glue.py | 122 ++++++ tests/v1/worker/test_gpu_v2_steering_glue.py | 198 +++++++++ vllm/v1/worker/gpu/capture_runner_mixin.py | 438 +++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 57 ++- vllm/v1/worker/gpu/steering_runner_mixin.py | 309 +++++++++++++ 6 files changed, 1257 insertions(+), 2 deletions(-) create mode 100644 docs/design/v2_runner_steering_capture.md create mode 100644 tests/v1/worker/test_gpu_v2_capture_glue.py create mode 100644 tests/v1/worker/test_gpu_v2_steering_glue.py create mode 100644 vllm/v1/worker/gpu/capture_runner_mixin.py create mode 100644 vllm/v1/worker/gpu/steering_runner_mixin.py diff --git a/docs/design/v2_runner_steering_capture.md b/docs/design/v2_runner_steering_capture.md new file mode 100644 index 000000000000..8bc92a5fa0c8 --- /dev/null +++ b/docs/design/v2_runner_steering_capture.md @@ -0,0 +1,135 @@ +# Steering + Capture on the V2 Model Runner + +Design/implementation contract for porting the **activation steering** and +**activation capture** control planes from the v1 GPU model runner +(`vllm/v1/worker/gpu_model_runner.py`) to the experimental v2 runner +(`vllm/v1/worker/gpu/model_runner.py`). + +## Scope + +In scope: + +- Wiring the runner-agnostic steering/capture subsystems into the v2 runner's + lifecycle so both features behave identically to v1. +- A v2-native control plane (new modules under `vllm/v1/worker/gpu/`) that keeps + its own per-request state, since v2 does not retain a `CachedRequestState` + dict the way v1 does. + +Out of scope: + +- The data plane (model-side custom ops, layer buffers, Triton kernels). These + live in `vllm/model_executor/` and are **already shared** by both runners. +- Refactoring the v1 mixins. The v1 path is validated/production; we leave it + untouched and write v2-native modules instead. +- Dynamic-steering (steer-from-capture feedback) and routed-experts capture + (tracked separately). + +## Key architectural fact + +The activation read/write mechanism is **not** in either model runner. Decoder +layers in `vllm/model_executor/models/*.py` call `apply_layer_steering()` and +`maybe_capture_residual()`; `register_steering_buffers()` runs in the model's +`__init__`. Both runners load the *same* model object via +`model_loader.load_model()`, so in v2 the buffers already exist, the custom ops +(`torch.ops.vllm.apply_steering`, `torch.ops.vllm.capture_residual`) already +fire, and they safely no-op when no control plane drives them (steering tables +stay zero → `any_active=False`; `get_active_capture_manager()` is `None` → +constant-folds under `torch.compile`). + +The split is therefore: + +| Plane | Location | Status in v2 | +| --- | --- | --- | +| Data plane (ops, buffers, kernels, store, managers, gate, types) | `model_executor/`, `v1/capture/`, `v1/worker/steering_manager.py` | shared, reused unchanged | +| Scheduler handoff (`NewRequestData.{prefill,decode}_steering_config_hash`, `capture_block_hashes`, `sampling_params.capture`; `ModelRunnerOutput.capture_results`) | `v1/core/sched/output.py`, `v1/outputs.py` | shared, already present | +| Control plane (init, per-step buffer fill / plan build, force-eager, request lifecycle, output drain) | runner | **absent — this port** | + +## V2 runner seams + +The v2 runner splits the monolithic v1 `execute_model` into discrete methods. +The port attaches to these (all in `gpu/model_runner.py`): + +- `load_model` (266): construct managers/gate/store; init steerable-layer + discovery. Buffers are already registered model-side. +- `add_requests` (691): per `new_req_data` in `scheduled_new_reqs` — register + steering config + track phase; `gate.register` (all ranks) + capture + `register_request` (TP0). Note `add_requests` calls `_remove_request` first + for streaming re-adds, so refresh state accordingly. +- `update_requests` (736): prefill→decode transition / resumption bookkeeping. +- `finish_requests` (678): use `scheduler_output.finished_req_ids` for steering + release + capture finalize + `gate.drop`; `preempted_req_ids` → steering + reset, **not** capture finalize. +- `execute_model` (1009): + - Force-eager seam at the `dispatch_cg_and_sync_dp(..., need_eager=...)` call + (1042–1050): OR in `capture_pending` (client-spec captures only; global + specs ride the cudagraph-safe persistent-buffer path). **Steering needs no + force-eager** — its tables/index are persistent buffers written before the + forward, so graph replay reads them correctly. + - After `prepare_inputs` (1060) and before the model forward (1167): build the + per-step view, `_update_steering_buffers(view)`, and + `capture_manager.build_step_plan(view)` (TP0). + - After the forward (after 1210, before non-last-PP return at 1220): + `_finalize_capture_step()` (consume plan, async dispatch). +- `sample_tokens` (1229): attach drained `_pending_capture_results` to the + `ModelRunnerOutput` (1276); `_finalize_capture_for_request_async` results land + here, same as v1's `get_output`. + +### Per-request state ownership + +v2's `RequestState` (`gpu/states.py`) holds only tokens/lengths — not +`sampling_params` or steering hashes. The control plane therefore keeps its own +dicts (`req_id → (prefill_hash, decode_hash, phase)` for steering; +gate selectors + manager registration for capture), populated from +`NewRequestData` in `add_requests`. The per-step view is built from v2's +`InputBatch` (`req_ids` ordering + `idx_mapping_np` + `num_scheduled_tokens`) +plus `req_states` (`num_computed_tokens_np`, `prefill_len`, `prompt_len`). + +## Rank-replication invariant + +Preserved exactly as in v1: the force-eager decision (`CaptureStepGate`) and +steering-manager row allocation are rank-local and deterministic, fed by the +broadcast `scheduler_output`. No hot-path collectives. Every new seam must read +only rank-identical inputs. + +## CUDA-graph interaction + +- Steering: persistent buffers (`steering_table_*`, `steering_index`) written + in-place before the forward → FULL-graph replay reads them. Safe. +- Capture global specs: fixed-shape full-residual copy into persistent + `_global_buffers`, baked at warmup. Safe (no force-eager). +- Capture client specs: dynamic `index_select` → not graph-capturable → gate + forces eager for that step only. + +## Workstreams + +1. **Capture control plane** — DONE (CPU-tested, GPU pending). + `gpu/capture_runner_mixin.py` (`CaptureRunnerMixin`): init, gate, + force-eager seam, `_build_capture_{gate,batch}_view`, + `_register_capture_request`, `_finalize_capture_step`, + `_finalize_capture_for_request_async`, output drain, activation store. + Tests: `tests/v1/worker/test_gpu_v2_capture_glue.py`. +2. **Steering control plane** — DONE (CPU-tested, GPU pending). + `gpu/steering_runner_mixin.py` (`SteeringRunnerMixin`, a subclass of + `SteeringModelRunnerMixin` that reuses init / discovery / validation / the + public RPC API / `_resolve_request_steering` and overrides only the three + v1-state-coupled paths). Keeps its own `_steering_reqs` per-request state; + `_steering_add_request` (register + streaming re-add), `_steering_finish_requests` + (release on finish/preempt), `_update_steering_buffers_v2` (transition + + per-token index). No force-eager seam (persistent buffers). `gpu_worker.py` + already forwards the RPCs to `self.model_runner.*`. + Tests: `tests/v1/worker/test_gpu_v2_steering_glue.py`. + +### Known gaps to confirm on GPU + +- Capture: streaming re-add (manager re-register may error) and preemption resume. +- Steering: preemption resume assumes the request re-enters via `add_requests` + (not `update_requests`); spec-decode + steering token layout not yet exercised. + +## Validation + +- CPU unit tests for the v2 control-plane glue (state bookkeeping, view + construction, gate decisions) where the data plane can be exercised via the + python fns (the CUDA ops are dispatch-only on GPU). +- GPU end-to-end on node2 (gemma + qwen3): steering eager-vs-cudagraph parity, + capture client-spec (`all_prompt` → recapture) and `all_generated` → reuse, + TP/PP rank agreement. Mirrors the v1 validation matrix. diff --git a/tests/v1/worker/test_gpu_v2_capture_glue.py b/tests/v1/worker/test_gpu_v2_capture_glue.py new file mode 100644 index 000000000000..624234a04a34 --- /dev/null +++ b/tests/v1/worker/test_gpu_v2_capture_glue.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU unit tests for the v2 runner's capture control-plane glue. + +These cover the v2-specific projection logic (``CaptureBatchView`` builders and +the finalized-result drain) without a CUDA device or a real model. The data +plane and managers are exercised separately in ``tests/v1/capture``. +""" + +from __future__ import annotations + +import threading +from types import SimpleNamespace + +import numpy as np + +from vllm.v1.worker.gpu.capture_runner_mixin import CaptureRunnerMixin + + +class _Glue(CaptureRunnerMixin): + """Minimal host exposing only what the view builders / drain read.""" + + def __init__(self, req_states): + self.req_states = req_states + self._capture_feature_enabled = True + self._capture_step_gate = None + self._capture_manager = None + self._pending_capture_results = {} + self._pending_capture_results_lock = threading.Lock() + + +def _req_states(prompt_len, num_computed, req_id_to_index): + return SimpleNamespace( + prompt_len=SimpleNamespace(np=np.asarray(prompt_len, dtype=np.int32)), + num_computed_tokens_np=np.asarray(num_computed, dtype=np.int32), + req_id_to_index=req_id_to_index, + ) + + +def test_gate_view_from_scheduler_output(): + # req_state slots: a->1, b->0 (deliberately not in batch order). + rs = _req_states( + prompt_len=[7, 5], + num_computed=[5, 0], + req_id_to_index={"a": 1, "b": 0}, + ) + glue = _Glue(rs) + # Scheduler dict order is the iteration order used for the gate view. + sched = SimpleNamespace(num_scheduled_tokens={"a": 2, "b": 5}) + + view = glue._build_capture_gate_view(sched) + + assert view.req_ids == ["a", "b"] + assert view.num_prompt_tokens == [5, 7] # a->idx1=5, b->idx0=7 + assert view.num_computed_tokens == [0, 5] + assert view.num_scheduled_tokens == [2, 5] + assert view.token_offsets == [0, 2] # cumulative scheduled tokens + + +def test_gate_view_unknown_request_defaults_zero(): + rs = _req_states([3], [0], {"a": 0}) + glue = _Glue(rs) + sched = SimpleNamespace(num_scheduled_tokens={"ghost": 4}) + + view = glue._build_capture_gate_view(sched) + + assert view.req_ids == ["ghost"] + assert view.num_prompt_tokens == [0] + assert view.num_computed_tokens == [0] + assert view.num_scheduled_tokens == [4] + + +def test_batch_view_uses_input_batch_offsets(): + # req_state slots: d->0, p->1; batch is decode-first so idx_mapping=[0, 1]. + rs = _req_states( + prompt_len=[10, 20], + num_computed=[10, 0], + req_id_to_index={"d": 0, "p": 1}, + ) + glue = _Glue(rs) + input_batch = SimpleNamespace( + num_reqs=2, + req_ids=["d", "p"], + idx_mapping_np=np.asarray([0, 1], dtype=np.int32), + num_scheduled_tokens=np.asarray([1, 20], dtype=np.int32), + # query_start_loc_np carries one extra (cumulative) entry; the builder + # slices [:num_reqs]. + query_start_loc_np=np.asarray([0, 1, 21], dtype=np.int32), + ) + + view = glue._build_capture_batch_view(input_batch) + + assert view.req_ids == ["d", "p"] + assert view.num_prompt_tokens == [10, 20] + assert view.num_computed_tokens == [10, 0] + assert view.num_scheduled_tokens == [1, 20] + assert view.token_offsets == [0, 1] # from query_start_loc_np[:2] + + +def test_gate_decision_false_without_gate(): + glue = _Glue(_req_states([1], [0], {"a": 0})) + sched = SimpleNamespace(num_scheduled_tokens={"a": 1}) + + assert glue._capture_gate_decision(sched) is False + + +def test_drain_capture_results_empties_buffer(): + glue = _Glue(_req_states([1], [0], {"a": 0})) + glue._pending_capture_results = {"a": {"c": object()}} + + drained = glue._drain_capture_results() + + assert set(drained) == {"a"} + assert glue._drain_capture_results() == {} # buffer cleared + + +def test_drain_disabled_returns_empty(): + glue = _Glue(_req_states([1], [0], {"a": 0})) + glue._capture_feature_enabled = False + glue._pending_capture_results = {"a": {"c": object()}} + + assert glue._drain_capture_results() == {} diff --git a/tests/v1/worker/test_gpu_v2_steering_glue.py b/tests/v1/worker/test_gpu_v2_steering_glue.py new file mode 100644 index 000000000000..b5528d080f1e --- /dev/null +++ b/tests/v1/worker/test_gpu_v2_steering_glue.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU unit tests for the v2 runner's steering control-plane glue. + +Covers the v2-specific lifecycle (register on add, release on finish, +prefill->decode transition) and the per-step index build, using a fake +``SteeringManager`` and CPU tensors. The fused kernel / real manager are +exercised separately in ``tests/v1/worker/test_steering_manager*.py``. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import torch + +from vllm.model_executor.layers.steering import ( + HOOK_POINT_ANY_ACTIVE_ATTR, + SteeringHookPoint, +) +from vllm.v1.worker.gpu.steering_runner_mixin import SteeringRunnerMixin + + +class _FakeManager: + def __init__(self): + self.config_to_row: dict = {} + self.global_base_vectors: dict = {} + self.global_prefill_vectors: dict = {} + self.global_decode_vectors: dict = {} + self._tables_dirty = False + self.registered: list[tuple[int, str]] = [] + self.released: list[tuple[int, str]] = [] + self.populated = 0 + + def register_config(self, h, effective, phase, locally_owned_layers): + self.registered.append((h, phase)) + self.config_to_row[(h, phase)] = len(self.config_to_row) + 3 + + def release_config(self, h, phase): + self.released.append((h, phase)) + + def get_row_for_config(self, h, is_prefill): + return int(h) # row == hash keeps the expected index readable + + def populate_steering_tables(self, layers): + self.populated += 1 + self._tables_dirty = False + + +def _layer(num_tokens=16): + layer = SimpleNamespace(steering_index=torch.zeros(num_tokens, dtype=torch.long)) + for hp in SteeringHookPoint: + setattr(layer, HOOK_POINT_ANY_ACTIVE_ATTR[hp], torch.ones(1, dtype=torch.bool)) + return layer + + +def _make_glue(num_computed, max_tokens=16, max_seqs=8): + glue = SteeringRunnerMixin.__new__(SteeringRunnerMixin) + glue._steering_manager = _FakeManager() + glue._steerable_layers_cache = {0: _layer(max_tokens)} + glue._steering_reqs = {} + glue._steering_index_dirty = False + glue._locally_owned_layers = frozenset({0}) + glue._steering_rows_scratch = np.zeros(max_seqs, dtype=np.int64) + glue._steering_n_tokens_scratch = np.zeros(max_seqs, dtype=np.int64) + glue._steering_index_pinned = torch.zeros(max_tokens, dtype=torch.long) + glue.req_states = SimpleNamespace( + num_computed_tokens_np=np.asarray(num_computed, dtype=np.int32) + ) + # Avoid building real SamplingParams: the resolve step just needs to be truthy. + glue._resolve_request_steering = lambda sp, phase: {"pre_attn": {0: [1.0]}} + return glue + + +def _new_req(req_id, prefill_hash, decode_hash, prompt_len, num_computed=0): + return SimpleNamespace( + req_id=req_id, + sampling_params=object(), + prefill_steering_config_hash=prefill_hash, + decode_steering_config_hash=decode_hash, + prompt_token_ids=list(range(prompt_len)), + prompt_embeds=None, + num_computed_tokens=num_computed, + ) + + +def test_add_request_registers_prefill(): + glue = _make_glue(num_computed=[0]) + glue._steering_add_request( + _new_req("a", prefill_hash=7, decode_hash=9, prompt_len=10) + ) + + rs = glue._steering_reqs["a"] + assert rs.phase == "prefill" + assert rs.num_prompt_tokens == 10 + assert glue._steering_manager.registered == [(7, "prefill")] + + +def test_add_request_direct_to_decode_on_full_prefix_hit(): + glue = _make_glue(num_computed=[10]) + glue._steering_add_request( + _new_req("a", prefill_hash=0, decode_hash=5, prompt_len=10, num_computed=10) + ) + + rs = glue._steering_reqs["a"] + assert rs.phase == "decode" + assert glue._steering_manager.registered == [(5, "decode")] + + +def test_add_request_no_hashes_is_untracked(): + glue = _make_glue(num_computed=[0]) + glue._steering_add_request( + _new_req("a", prefill_hash=0, decode_hash=0, prompt_len=4) + ) + + assert "a" not in glue._steering_reqs + assert glue._steering_manager.registered == [] + + +def test_finish_request_releases_current_phase(): + glue = _make_glue(num_computed=[0]) + glue._steering_add_request( + _new_req("a", prefill_hash=7, decode_hash=9, prompt_len=10) + ) + + glue._steering_finish_requests(["a"]) + + assert "a" not in glue._steering_reqs + assert glue._steering_manager.released == [(7, "prefill")] + + +def test_streaming_readd_releases_old_then_registers_new(): + glue = _make_glue(num_computed=[0]) + glue._steering_add_request( + _new_req("a", prefill_hash=7, decode_hash=9, prompt_len=10) + ) + # Re-add same id with a different config (streaming update). + glue._steering_add_request( + _new_req("a", prefill_hash=11, decode_hash=9, prompt_len=10) + ) + + assert glue._steering_manager.released == [(7, "prefill")] + assert glue._steering_manager.registered == [(7, "prefill"), (11, "prefill")] + + +def test_update_buffers_builds_per_token_index_and_transition(): + # Two requests, batch order [decode "d", prefill "p"]. + glue = _make_glue(num_computed=[10, 8]) + # d: direct-to-decode (computed 10 >= prompt 10), decode_hash 5. + glue._steering_add_request( + _new_req("d", prefill_hash=0, decode_hash=5, prompt_len=10, num_computed=10) + ) + # p: prefilling, computed 8 of 10, prefill_hash 7 / decode_hash 9; this + # step schedules 3 tokens -> crosses the boundary -> transition fires. + glue._steering_add_request( + _new_req("p", prefill_hash=7, decode_hash=9, prompt_len=10) + ) + glue._steering_reqs["p"].num_prompt_tokens = 10 # ensure boundary at 10 + + input_batch = SimpleNamespace( + num_reqs=2, + req_ids=["d", "p"], + idx_mapping_np=np.asarray([0, 1], dtype=np.int32), + ) + sched = SimpleNamespace(num_scheduled_tokens={"d": 1, "p": 3}) + + glue._update_steering_buffers_v2(sched, input_batch) + + steering_index = glue._steerable_layers_cache[0].steering_index + # d -> row 5 (1 token); p -> row 7 (3 tokens); tail zeroed. + assert steering_index[:4].tolist() == [5, 7, 7, 7] + assert steering_index[4:].sum().item() == 0 + # Boundary crossed for p (8 + 3 >= 10): prefill 7 released, decode 9 added. + assert (7, "prefill") in glue._steering_manager.released + assert (9, "decode") in glue._steering_manager.registered + assert glue._steering_reqs["p"].phase == "decode" + + +def test_update_buffers_short_circuit_zeroes_dirty_index(): + glue = _make_glue(num_computed=[0]) + # No tracked requests and no globals -> nothing active. + layer = glue._steerable_layers_cache[0] + layer.steering_index[:3] = torch.tensor([1, 2, 3]) + glue._steering_index_dirty = True + + input_batch = SimpleNamespace( + num_reqs=0, req_ids=[], idx_mapping_np=np.asarray([], dtype=np.int32) + ) + sched = SimpleNamespace(num_scheduled_tokens={}) + + glue._update_steering_buffers_v2(sched, input_batch) + + assert layer.steering_index.sum().item() == 0 + assert glue._steering_index_dirty is False + # any_active flags cleared so apply_steering short-circuits. + for hp in SteeringHookPoint: + assert getattr(layer, HOOK_POINT_ANY_ACTIVE_ATTR[hp]).item() is False diff --git a/vllm/v1/worker/gpu/capture_runner_mixin.py b/vllm/v1/worker/gpu/capture_runner_mixin.py new file mode 100644 index 000000000000..09a68a9c6be9 --- /dev/null +++ b/vllm/v1/worker/gpu/capture_runner_mixin.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Activation-capture control plane for the v2 GPU model runner. + +The capture *data plane* (the ``capture_residual`` custom op, per-layer taps, +persistent global buffers, kernels) lives in ``vllm/model_executor`` and +``vllm/v1/capture`` and is shared with the v1 runner unchanged. This mixin is +only the v2 runner-side glue: it builds the managers/gate/store, registers and +finalizes requests, drives the per-step force-eager decision and gather plan, +and drains finalized results onto ``ModelRunnerOutput``. + +It deliberately keeps its own per-request bookkeeping rather than piggybacking +on runner state, because v2's ``RequestState`` does not retain ``sampling_params`` +or capture specs past ``add_requests``. + +See ``docs/design/v2_runner_steering_capture.md`` for the full contract. +""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any + +from vllm.distributed.parallel_state import get_tp_group +from vllm.logger import init_logger +from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.torch_utils import get_dtype_size + +if TYPE_CHECKING: + from vllm.v1.capture.plan import CaptureBatchView + from vllm.v1.capture.types import CaptureResult + from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput + from vllm.v1.worker.gpu.input_batch import InputBatch + +logger = init_logger(__name__) + + +class CaptureRunnerMixin: + """Mixin adding the activation-capture control plane to the v2 runner. + + Assumes the concrete runner provides: ``vllm_config``, ``model_config``, + ``parallel_config``, ``device``, ``max_num_tokens`` and ``req_states``. + """ + + # ---- state (set by _init_capture_state) ------------------------------- + _capture_feature_enabled: bool = False + _capture_manager: Any = None + _capture_step_gate: Any = None + _capture_validators: list[Any] + _capture_name_to_index: dict[str, int] + _capture_index_to_name: dict[int, str] + _pending_capture_results: dict[str, dict[str, CaptureResult]] + _pending_capture_results_lock: threading.Lock + + def _init_capture_state(self) -> None: + """Construct the capture managers/gate/store. Idempotent no-op when + capture is not configured. + + Mirrors the v1 runner's ``__init__`` capture block: the rank-replicated + ``CaptureStepGate`` is built on every rank so the eager-vs-cudagraph + choice agrees across the TP/PP topology without a per-step collective, + while the ``CaptureManager`` and process-global active manager are + installed only on TP rank 0 (the residual stream is byte-identical + across the TP group after the all-reduce, so exactly one rank captures). + """ + self._capture_manager = None + self._capture_step_gate = None + self._capture_validators = [] + self._capture_name_to_index = {} + self._capture_index_to_name = {} + self._pending_capture_results = {} + self._pending_capture_results_lock = threading.Lock() + + cc_config = self.vllm_config.capture_consumers_config + self._capture_feature_enabled = cc_config is not None + if cc_config is None: + return + + from vllm.model_executor.layers.activation_capture import ( + set_active_capture_manager, + ) + from vllm.v1.capture.step_gate import CaptureStepGate + + # Global capture specs ride a CUDA-graph-safe persistent-buffer path, so + # the gate forces eager only when a per-request *client* spec captures. + self._capture_step_gate = CaptureStepGate() + + if get_tp_group().rank_in_group != 0: + # Non-capturer rank: no manager, cold-path custom op. + set_active_capture_manager(None) + else: + from vllm.v1.capture import registry as _capture_registry + from vllm.v1.capture.manager import CaptureManager + + instances = list(cc_config.instances) + sinks, validators, name_to_index = _capture_registry.build_consumers( + self.vllm_config, consumer_instances=instances + ) + + global_specs: list[Any] = [] + for validator in validators: + spec = None + try: + if hasattr(validator, "global_capture_spec"): + spec = validator.global_capture_spec() + except Exception: + spec = None + global_specs.append(spec) + + self._capture_manager = CaptureManager( + consumers=sinks, + consumer_specs=tuple(global_specs), + num_hidden_layers=self.model_config.get_total_num_hidden_layers(), + local_layer_range=self.model_config.get_layers_start_end_indices( + self.parallel_config + ), + hidden_size=self.model_config.get_hidden_size(), + model_dtype=self.model_config.dtype, + device=self.device, + max_num_tokens=self.max_num_tokens, + dispatch_queue_size=getattr(cc_config, "dispatch_queue_size", 256), + overload_policy=getattr(cc_config, "overload_policy", "spill"), + spill_dir=getattr(cc_config, "spill_dir", None), + spill_max_bytes=getattr(cc_config, "spill_max_bytes", 4 << 30), + ) + self._capture_validators = validators + self._capture_name_to_index = dict(name_to_index) + self._capture_index_to_name = { + idx: name for name, idx in name_to_index.items() + } + set_active_capture_manager(self._capture_manager) + + budget = cc_config.activation_cache_bytes + if budget > 0: + from vllm.v1.capture.activation_store import ( + ActivationStore, + set_active_activation_store, + ) + + set_active_activation_store(ActivationStore(max_bytes=budget)) + logger.info( + "Capture activation store enabled: budget=%.3f GB", + budget / 1_000_000_000, + ) + + # ---- request lifecycle ------------------------------------------------- + + def _capture_add_request(self, new_req_data: NewRequestData) -> None: + """Hook a newly admitted request into capture tracking. + + ``register`` runs on every rank (the gate is rank-replicated); + manager registration runs only on the capturer rank. + """ + if not self._capture_feature_enabled: + return + sp = new_req_data.sampling_params + if self._capture_step_gate is not None: + self._capture_step_gate.register( + new_req_data.req_id, getattr(sp, "capture", None) + ) + if self._capture_manager is not None: + self._register_capture_request(new_req_data) + + def _capture_finish_request(self, req_id: str) -> None: + """Drop a finished request from the gate and finalize its capture.""" + if not self._capture_feature_enabled: + return + if self._capture_step_gate is not None: + self._capture_step_gate.drop(req_id) + if self._capture_manager is not None: + self._finalize_capture_for_request_async(req_id) + + def _register_capture_request(self, new_req_data: NewRequestData) -> None: + """Admit a request into the capture framework (capturer rank only). + + Resolves ``sampling_params.capture`` consumer names against the registry, + validates raw specs, and registers with the manager. Admission errors + never abort generation — they surface as ``CaptureResult(status="error")`` + on finalize. + """ + assert self._capture_manager is not None + mgr = self._capture_manager + + sp = new_req_data.sampling_params + if sp is None: + return + + prompt_len = length_from_prompt_token_ids_or_embeds( + new_req_data.prompt_token_ids, + new_req_data.prompt_embeds, + ) + + try: + element_size_bytes = get_dtype_size(self.model_config.dtype) + except Exception: + element_size_bytes = 2 + + from vllm.v1.capture.activation_store import pop_pending_serve + from vllm.v1.capture.errors import CaptureValidationError + from vllm.v1.capture.types import ( + CaptureContext, + CaptureSpec, + VllmInternalRequestId, + capture_expert_parallel_size, + ) + + # Step A serve: whole-prefix store serve means the prompt prefix was + # reused from KV cache; validate against num_computed=0 so the validator + # accepts those positions, then inject served rows after registration. + served_rows = pop_pending_serve(new_req_data.req_id) + ctx_num_computed = ( + 0 if served_rows is not None else new_req_data.num_computed_tokens + ) + + parallel_config = self.parallel_config + ctx = CaptureContext( + vllm_internal_request_id=VllmInternalRequestId(new_req_data.req_id), + num_prompt_tokens=prompt_len, + num_computed_tokens=ctx_num_computed, + num_hidden_layers=self.model_config.get_total_num_hidden_layers(), + hidden_size=self.model_config.get_hidden_size(), + element_size_bytes=element_size_bytes, + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + expert_parallel_size=capture_expert_parallel_size(parallel_config), + data_parallel_size=parallel_config.data_parallel_size, + ) + + raw_client = getattr(sp, "capture", None) + client_specs: dict[int, CaptureSpec] = {} + + if raw_client: + if not isinstance(raw_client, dict): + mgr.record_request_error( + new_req_data.req_id, + "SamplingParams.capture must be a dict keyed by consumer " + f"name, got {type(raw_client).__name__}", + ) + return + + for name, raw in raw_client.items(): + idx = self._capture_name_to_index.get(name) + if idx is None: + mgr.record_request_error( + new_req_data.req_id, + f"capture consumer {name!r} is not registered; known " + f"consumers: {sorted(self._capture_name_to_index)}", + ) + logger.warning( + "capture admission rejected req=%s: unknown consumer %s", + new_req_data.req_id, + name, + ) + return + + if isinstance(raw, CaptureSpec): + client_specs[idx] = raw + continue + + validator = self._capture_validators[idx] + try: + resolved = validator.validate_client_spec(raw, ctx) + except CaptureValidationError as exc: + mgr.record_request_error(new_req_data.req_id, str(exc)) + logger.warning( + "capture admission rejected req=%s consumer=%s: %s", + new_req_data.req_id, + name, + exc, + ) + return + except Exception as exc: # noqa: BLE001 + mgr.record_request_error( + new_req_data.req_id, + f"consumer {name!r} validator raised: {exc}", + ) + logger.warning( + "capture admission rejected req=%s consumer=%s: %s", + new_req_data.req_id, + name, + exc, + ) + return + + client_specs[idx] = resolved + + sidecar_fields: dict[str, Any] = { + "vllm_internal_request_id": new_req_data.req_id, + "prompt_token_ids": ( + list(new_req_data.prompt_token_ids) + if new_req_data.prompt_token_ids is not None + else [] + ), + } + + try: + mgr.register_request( + new_req_data.req_id, + client_specs=client_specs, + num_prompt_tokens=prompt_len, + sidecar_fields=sidecar_fields, + block_hashes=new_req_data.capture_block_hashes, + hash_block_size=new_req_data.capture_hash_block_size, + ) + if served_rows is not None: + mgr.serve_from_store(new_req_data.req_id, served_rows) + except ValueError as exc: + mgr.record_request_error(new_req_data.req_id, str(exc)) + logger.warning( + "capture register rejected req=%s: %s", new_req_data.req_id, exc + ) + + # ---- per-step ---------------------------------------------------------- + + def _build_capture_gate_view( + self, scheduler_output: SchedulerOutput + ) -> CaptureBatchView: + """Build an (unordered) view for the force-eager gate decision. + + v2 resolves the cudagraph-vs-eager batch descriptor *before* + ``prepare_inputs`` builds the ``InputBatch``, so the gate view is built + here from ``scheduler_output`` + ``req_states``. The gate only inspects + per-request prompt/computed/scheduled token counts (``token_offsets`` is + unused for the boolean), so request ordering does not matter. + """ + from vllm.v1.capture.plan import CaptureBatchView + + req_states = self.req_states + req_ids: list[str] = [] + num_prompt_tokens: list[int] = [] + num_computed_tokens: list[int] = [] + num_scheduled_tokens: list[int] = [] + token_offsets: list[int] = [] + + offset = 0 + for req_id, n_tokens in scheduler_output.num_scheduled_tokens.items(): + req_ids.append(req_id) + req_idx = req_states.req_id_to_index.get(req_id) + if req_idx is None: + num_prompt_tokens.append(0) + num_computed_tokens.append(0) + else: + num_prompt_tokens.append(int(req_states.prompt_len.np[req_idx])) + num_computed_tokens.append( + int(req_states.num_computed_tokens_np[req_idx]) + ) + num_scheduled_tokens.append(int(n_tokens)) + token_offsets.append(offset) + offset += int(n_tokens) + + return CaptureBatchView( + req_ids=req_ids, + num_prompt_tokens=num_prompt_tokens, + num_computed_tokens=num_computed_tokens, + num_scheduled_tokens=num_scheduled_tokens, + token_offsets=token_offsets, + ) + + def _build_capture_batch_view(self, input_batch: InputBatch) -> CaptureBatchView: + """Project v2's ``InputBatch`` into a :class:`CaptureBatchView`. + + Used for the gather plan, which needs token offsets that match the + actual forward batch layout. v2's batch is sorted decode-first; + ``query_start_loc_np`` gives the per-request token offset and + ``idx_mapping_np`` maps batch index to the request-state slot holding + prompt/computed lengths. + """ + from vllm.v1.capture.plan import CaptureBatchView + + num_reqs = input_batch.num_reqs + idx = input_batch.idx_mapping_np[:num_reqs] + req_states = self.req_states + return CaptureBatchView( + req_ids=list(input_batch.req_ids), + num_prompt_tokens=req_states.prompt_len.np[idx].tolist(), + num_computed_tokens=req_states.num_computed_tokens_np[idx].tolist(), + num_scheduled_tokens=input_batch.num_scheduled_tokens[:num_reqs].tolist(), + token_offsets=input_batch.query_start_loc_np[:num_reqs].tolist(), + ) + + def _capture_gate_decision(self, scheduler_output: SchedulerOutput) -> bool: + """Rank-replicated force-eager decision for this step. + + Returns ``True`` iff a per-request client spec captures this step, so the + runner must run eager (the dynamic ``index_select`` gather cannot be + recorded into a CUDA graph). Global specs ride the persistent-buffer path + and never force eager. + """ + if self._capture_step_gate is None: + return False + view = self._build_capture_gate_view(scheduler_output) + return self._capture_step_gate.step_captures(view) + + def _capture_build_plan(self, input_batch: InputBatch) -> None: + """Build the per-step gather plan on the capturer rank (pre-forward).""" + if self._capture_manager is not None and self._capture_manager.is_active(): + view = self._build_capture_batch_view(input_batch) + self._capture_manager.build_step_plan(view) + + def _finalize_capture_step(self) -> None: + """Dispatch captured rows to consumer sinks (after the forward).""" + if self._capture_manager is None: + return + plan = self._capture_manager.consume_step_plan() + if plan is None: + return + self._capture_manager.dispatch_step_captures(plan) + + def _finalize_capture_for_request_async(self, req_id: str) -> None: + """Finalize *req_id* off the step thread; stash results for draining.""" + mgr = self._capture_manager + if mgr is None: + return + + index_to_name = self._capture_index_to_name + pending = self._pending_capture_results + lock = self._pending_capture_results_lock + + def _on_complete(indexed: dict[int, CaptureResult]) -> None: + if not indexed: + return + named = { + index_to_name.get(idx, f"consumer_{idx}"): result + for idx, result in indexed.items() + } + with lock: + pending.setdefault(req_id, {}).update(named) + + mgr.finalize_request_async(req_id, _on_complete) + + def _drain_capture_results(self) -> dict[str, dict[str, CaptureResult]]: + """Atomically take the finalized results buffered since the last step.""" + if not self._capture_feature_enabled: + return {} + with self._pending_capture_results_lock: + results = self._pending_capture_results + self._pending_capture_results = {} + return results diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c0a95618c111..0f06c14172a4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -61,6 +61,7 @@ ) from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.v1.worker.gpu.capture_runner_mixin import CaptureRunnerMixin from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -101,13 +102,14 @@ from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.gpu.steering_runner_mixin import SteeringRunnerMixin from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin logger = init_logger(__name__) -class GPUModelRunner(LoRAModelRunnerMixin): +class GPUModelRunner(LoRAModelRunnerMixin, CaptureRunnerMixin, SteeringRunnerMixin): def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -327,6 +329,13 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: device=self.device, ) + # Activation-capture control plane. The data-plane hooks already live in + # the loaded model; this installs the managers/gate/store that drive them. + self._init_capture_state() + # Activation-steering control plane. Discovers steerable layers on the + # loaded model and builds the SteeringManager (no-op when disabled). + self._init_steering_state() + def get_model(self) -> nn.Module: return self.model @@ -677,9 +686,18 @@ def _remove_request(self, req_id: str) -> bool: def finish_requests(self, scheduler_output: SchedulerOutput) -> None: finished_req_ids = scheduler_output.finished_req_ids + # Finalize capture only for genuinely finished requests; preempted + # requests resume later and must not have their capture closed. + if self._capture_feature_enabled: + for req_id in finished_req_ids: + self._capture_finish_request(req_id) preempted_req_ids = scheduler_output.preempted_req_ids if preempted_req_ids: finished_req_ids = finished_req_ids.union(preempted_req_ids) + # Release steering configs for finished AND preempted requests; a + # resumed request re-registers a fresh prefill config via add_requests. + if self._steering_manager is not None: + self._steering_finish_requests(finished_req_ids) for req_id in finished_req_ids: self._remove_request(req_id) @@ -727,6 +745,15 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: req_id, req_index, new_req_data.sampling_params ) + # Register the request with the capture control plane (every rank + # registers with the gate; the capturer rank also registers with + # the manager). Runs on all PP ranks since capturer layers may live + # on any stage. + self._capture_add_request(new_req_data) + # Register the request's steering config (rank-local; deterministic + # across the TP/PP topology from the broadcast scheduler output). + self._steering_add_request(new_req_data) + if scheduler_output.scheduled_new_reqs: self.req_states.apply_staged_writes() self.model_state.apply_staged_writes() @@ -1014,6 +1041,7 @@ def execute_model( skip_attn_for_dummy_run: bool = False, is_profile: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: + capture_pending = False if not dummy_run: # Update the request states. self.finish_requests(scheduler_output) @@ -1025,6 +1053,12 @@ def execute_model( # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) return empty_output + # Rank-replicated force-eager decision: a per-request client capture + # spec uses a dynamic gather that cannot be recorded into a CUDA + # graph, so this step must run eager. Computed from scheduler_output + # (not InputBatch) because the batch descriptor is resolved below, + # before prepare_inputs runs. + capture_pending = self._capture_gate_decision(scheduler_output) # Get batch descriptor and sync across DP ranks. num_reqs = len(scheduler_output.num_scheduled_tokens) @@ -1046,7 +1080,7 @@ def execute_model( uniform_tok_count, self.dp_size, self.dp_rank, - need_eager=is_profile or skip_compiled, + need_eager=is_profile or skip_compiled or capture_pending, ) if batch_desc.num_tokens == 0: @@ -1060,6 +1094,17 @@ def execute_model( input_batch = self.prepare_inputs(scheduler_output, batch_desc) block_tables, slot_mappings = self.prepare_attn(input_batch) + # Build the capture gather plan (capturer rank only) from the final + # batch layout; the in-forward capture_residual op populates it. + if self._capture_manager is not None: + self._capture_build_plan(input_batch) + + # Populate steering tables + per-token index before the forward. + # The buffers are persistent, so a FULL cudagraph replay reads this + # step's values — no force-eager needed. + if self._steering_manager is not None: + self._update_steering_buffers_v2(scheduler_output, input_batch) + if self.lora_config: # Activate LoRA adapters. lora_inputs = self.lora_state.make_lora_inputs( @@ -1192,6 +1237,12 @@ def execute_model( self.kv_connector.pre_forward(scheduler_output) model_output = self.model(**model_inputs) + # Dispatch the rows the in-forward capture op gathered this step. Runs + # on every PP stage (each capturer rank owns its stage's layers) and + # before non-last stages return their intermediate tensors below. + if self._capture_manager is not None: + self._finalize_capture_step() + if self.is_last_pp_rank: if self.use_aux_hidden_state_outputs: assert isinstance(model_output, tuple) @@ -1281,6 +1332,8 @@ def sample_tokens( sampled_token_ids=None, # type: ignore prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type] kv_connector_output=kv_connector_output, + # Capture results finalized (off-thread) since the last step. + capture_results=self._drain_capture_results(), ) async_output = AsyncOutput( model_runner_output=model_runner_output, diff --git a/vllm/v1/worker/gpu/steering_runner_mixin.py b/vllm/v1/worker/gpu/steering_runner_mixin.py new file mode 100644 index 000000000000..27461583acb1 --- /dev/null +++ b/vllm/v1/worker/gpu/steering_runner_mixin.py @@ -0,0 +1,309 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Activation-steering control plane for the v2 GPU model runner. + +The steering *data plane* (the ``apply_steering`` custom op, per-layer table +buffers, the fused Triton kernel, ``SteeringManager``) is shared with the v1 +runner unchanged. So is the entire runner-agnostic half of +:class:`SteeringModelRunnerMixin`: state init, steerable-layer discovery, spec +validation, the public RPC API (``set_steering_vectors`` etc., which +``gpu_worker`` already forwards to ``self.model_runner``), and +``_resolve_request_steering``. + +Only three v1 methods read v1-runner state (``self.input_batch`` / +``self.requests``): the per-step ``_update_steering_buffers`` hot path, the +prefill->decode transition, and finished-config release. v2 retains no +``CachedRequestState`` dict, so this subclass keeps its own per-request steering +state (populated from ``NewRequestData`` in ``add_requests``) and reimplements +those three against v2's ``InputBatch`` + ``RequestState``. + +Steering needs no force-eager seam: the per-layer tables and ``steering_index`` +are persistent buffers written in place before the forward, so a FULL cudagraph +replay reads the current step's values. + +See ``docs/design/v2_runner_steering_capture.md`` for the full contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +import numpy as np +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.steering import ( + HOOK_POINT_ANY_ACTIVE_ATTR, + SteeringHookPoint, +) +from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.v1.worker.steering_model_runner_mixin import SteeringModelRunnerMixin + +if TYPE_CHECKING: + from vllm.sampling_params import SamplingParams + from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput + from vllm.v1.worker.gpu.input_batch import InputBatch + +logger = init_logger(__name__) + + +@dataclass +class _SteeringReqState: + """Per-request steering state the v2 runner must retain itself. + + v2 does not keep ``sampling_params`` or the steering hashes past + ``add_requests``, so we capture what the transition / release / resolve + paths need: the params (for re-resolving the decode tier lazily), both + config hashes, the prompt length (for the prefill->decode boundary), and + the currently-registered phase. + """ + + sampling_params: SamplingParams + prefill_hash: int + decode_hash: int + num_prompt_tokens: int + phase: str # "prefill" | "decode" + + +class SteeringRunnerMixin(SteeringModelRunnerMixin): + """v2 steering control plane. Mixes the shared steering logic in via the + v1 mixin and overrides only the v2-runner-coupled paths. + + Assumes the concrete runner provides ``req_states`` (v2 ``RequestState``) + in addition to everything :class:`SteeringModelRunnerMixin` needs. + """ + + _steering_reqs: dict[str, _SteeringReqState] + + def _init_steering_state(self) -> None: + super()._init_steering_state() + self._steering_reqs = {} + + # ---- request lifecycle ------------------------------------------------- + + def _steering_add_request(self, new_req_data: NewRequestData) -> None: + """Track a newly admitted request and register its initial config. + + Also covers streaming re-adds: ``add_requests`` removes the prior + instance first, so any state we already held for this id is released + before the fresh prefill config is registered. + """ + mgr = self._steering_manager + if mgr is None: + return + + # Streaming re-add / stale state: release whatever we held before. + old = self._steering_reqs.pop(new_req_data.req_id, None) + if old is not None: + self._steering_release_state(old) + + sp = new_req_data.sampling_params + prefill_hash = new_req_data.prefill_steering_config_hash + decode_hash = new_req_data.decode_steering_config_hash + if sp is None or (prefill_hash == 0 and decode_hash == 0): + return + + num_prompt = length_from_prompt_token_ids_or_embeds( + new_req_data.prompt_token_ids, + new_req_data.prompt_embeds, + ) + rs = _SteeringReqState( + sampling_params=sp, + prefill_hash=prefill_hash, + decode_hash=decode_hash, + num_prompt_tokens=num_prompt, + phase="prefill", + ) + self._steering_reqs[new_req_data.req_id] = rs + + # A full prefix-cache hit admits the request directly into decode; the + # scheduler reserves the matching row, so register_config is expected to + # succeed (a RuntimeError indicates a scheduler accounting bug). + if new_req_data.num_computed_tokens >= num_prompt: + effective_decode = self._resolve_request_steering(sp, "decode") + if decode_hash != 0 and effective_decode: + mgr.register_config( + decode_hash, + effective_decode, + phase="decode", + locally_owned_layers=self._locally_owned_layers, + ) + rs.phase = "decode" + else: + # Normal: start in prefill; the decode config is registered lazily + # at the prefill->decode boundary in _update_steering_buffers_v2. + effective_prefill = self._resolve_request_steering(sp, "prefill") + if prefill_hash != 0 and effective_prefill: + mgr.register_config( + prefill_hash, + effective_prefill, + phase="prefill", + locally_owned_layers=self._locally_owned_layers, + ) + rs.phase = "prefill" + + def _steering_finish_requests(self, req_ids: set[str] | list[str]) -> None: + """Release configs for finished (or preempted) requests. + + Preempted requests are released too: they re-enter through + ``add_requests`` on resume, which re-registers a fresh prefill config. + """ + if self._steering_manager is None: + return + for req_id in req_ids: + rs = self._steering_reqs.pop(req_id, None) + if rs is not None: + self._steering_release_state(rs) + + def _steering_release_state(self, rs: _SteeringReqState) -> None: + """Release the config for whichever phase ``rs`` is currently in.""" + mgr = self._steering_manager + if mgr is None: + return + if rs.phase == "prefill" and rs.prefill_hash != 0: + mgr.release_config(rs.prefill_hash, "prefill") + elif rs.phase == "decode" and rs.decode_hash != 0: + mgr.release_config(rs.decode_hash, "decode") + + def _steering_transition(self, rs: _SteeringReqState) -> None: + """Handle a request crossing the prefill->decode boundary this step. + + Releases the prefill config and registers the decode config so it is + ready for the next step's table population. The scheduler reserves the + decode row at the step prefill completes, so register_config succeeds. + """ + mgr = self._steering_manager + assert mgr is not None + if rs.prefill_hash != 0: + mgr.release_config(rs.prefill_hash, "prefill") + if rs.decode_hash != 0: + effective_decode = self._resolve_request_steering( + rs.sampling_params, "decode" + ) + if effective_decode: + mgr.register_config( + rs.decode_hash, + effective_decode, + phase="decode", + locally_owned_layers=self._locally_owned_layers, + ) + rs.phase = "decode" + + # ---- per-step buffer / index maintenance ------------------------------- + + def _update_steering_buffers_v2( + self, scheduler_output: SchedulerOutput, input_batch: InputBatch + ) -> None: + """Populate per-layer steering tables and the shared steering index. + + v2 port of ``SteeringModelRunnerMixin._update_steering_buffers``: the + per-request hashes come from ``self._steering_reqs`` (not the input + batch), and token counts / phase come from ``input_batch`` + + ``req_states``. + """ + mgr = self._steering_manager + if mgr is None or not self._steerable_layers_cache: + return + + reqs = self._steering_reqs + num_reqs = input_batch.num_reqs + req_ids = input_batch.req_ids + idx_np = input_batch.idx_mapping_np + + # Short-circuit when nothing is active (no per-request config in this + # batch and no globals). A decode-only request (prefill_hash == 0, + # decode_hash != 0) registers its config lazily at the transition below, + # so the batch scan must not let the short-circuit swallow it. + batch_has_per_request_steering = any( + (rs := reqs.get(req_ids[i])) is not None + and (rs.prefill_hash != 0 or rs.decode_hash != 0) + for i in range(num_reqs) + ) + if ( + not batch_has_per_request_steering + and not mgr.config_to_row + and not mgr.global_base_vectors + and not mgr.global_prefill_vectors + and not mgr.global_decode_vectors + ): + if self._steering_index_dirty: + any_layer = next(iter(self._steerable_layers_cache.values())) + steering_index = cast(torch.Tensor, any_layer.steering_index) + steering_index.zero_() + for mod in self._steerable_layers_cache.values(): + for hp in SteeringHookPoint: + flag_buf = getattr(mod, HOOK_POINT_ANY_ACTIVE_ATTR[hp], None) + if flag_buf is not None: + flag_buf.zero_() + self._steering_index_dirty = False + return + + # 1. Populate tables only when state changed since the last populate. + if mgr._tables_dirty: + mgr.populate_steering_tables(self._steerable_layers_cache) + + # 2. Build the per-token steering index. + any_layer = next(iter(self._steerable_layers_cache.values())) + steering_index = cast(torch.Tensor, any_layer.steering_index) + + rows_scratch = self._steering_rows_scratch + n_tokens_scratch = self._steering_n_tokens_scratch + index_pinned = self._steering_index_pinned + assert rows_scratch is not None + assert n_tokens_scratch is not None + assert index_pinned is not None + if rows_scratch.shape[0] < num_reqs: + rows_scratch = np.zeros(num_reqs, dtype=np.int64) + n_tokens_scratch = np.zeros(num_reqs, dtype=np.int64) + self._steering_rows_scratch = rows_scratch + self._steering_n_tokens_scratch = n_tokens_scratch + + num_computed_np = self.req_states.num_computed_tokens_np + active_count = 0 + for i in range(num_reqs): + req_id = req_ids[i] + n_tokens = scheduler_output.num_scheduled_tokens.get(req_id, 0) + if n_tokens == 0: + continue + + rs = reqs.get(req_id) + if rs is None: + # No steering for this request — row 0 is the no-steer sentinel. + rows_scratch[active_count] = 0 + n_tokens_scratch[active_count] = n_tokens + active_count += 1 + continue + + num_computed = int(num_computed_np[int(idx_np[i])]) + num_prompt = rs.num_prompt_tokens + if num_computed < num_prompt: + row = mgr.get_row_for_config(rs.prefill_hash, is_prefill=True) + rows_scratch[active_count] = row + n_tokens_scratch[active_count] = n_tokens + if num_computed + n_tokens >= num_prompt: + self._steering_transition(rs) + else: + row = mgr.get_row_for_config(rs.decode_hash, is_prefill=False) + rows_scratch[active_count] = row + n_tokens_scratch[active_count] = n_tokens + active_count += 1 + + if active_count > 0: + expanded = np.repeat( + rows_scratch[:active_count], + n_tokens_scratch[:active_count], + ) + n_expanded = int(expanded.shape[0]) + n_expanded = min(n_expanded, index_pinned.shape[0], steering_index.shape[0]) + index_pinned[:n_expanded].copy_(torch.from_numpy(expanded[:n_expanded])) + steering_index[:n_expanded].copy_( + index_pinned[:n_expanded], non_blocking=True + ) + else: + n_expanded = 0 + + if n_expanded < steering_index.shape[0]: + steering_index[n_expanded:].zero_() + + self._steering_index_dirty = True From 40becc7c3fe0a455e93b4c9a836b02ff139de13b Mon Sep 17 00:00:00 2001 From: RhizoNymph Date: Thu, 18 Jun 2026 17:28:32 -0700 Subject: [PATCH 3/6] feat(v2-runner): enable steering and capture on v2 (remove interim v1 fallback guard); fix global-steering row for untracked requests --- docs/design/v2_runner_steering_capture.md | 20 ++++++++-- tests/test_config.py | 35 ----------------- tests/v1/worker/test_gpu_v2_steering_glue.py | 41 +++++++++++++++++--- vllm/config/vllm.py | 23 ----------- vllm/v1/worker/gpu/steering_runner_mixin.py | 35 ++++++++++------- 5 files changed, 71 insertions(+), 83 deletions(-) diff --git a/docs/design/v2_runner_steering_capture.md b/docs/design/v2_runner_steering_capture.md index 8bc92a5fa0c8..c784416ef372 100644 --- a/docs/design/v2_runner_steering_capture.md +++ b/docs/design/v2_runner_steering_capture.md @@ -119,11 +119,23 @@ only rank-identical inputs. already forwards the RPCs to `self.model_runner.*`. Tests: `tests/v1/worker/test_gpu_v2_steering_glue.py`. -### Known gaps to confirm on GPU +### Validation -- Capture: streaming re-add (manager re-register may error) and preemption resume. -- Steering: preemption resume assumes the request re-enters via `add_requests` - (not `update_requests`); spec-decode + steering token layout not yet exercised. +GPU-validated on Qwen3-0.6B (RTX 3090, TP1/PP1), forcing +`VLLM_USE_V2_MODEL_RUNNER=1`: + +- Steering (eager **and** cudagraph): global `set_steering_vectors` shifts the + output and `clear_steering_vectors` restores the exact baseline — confirming + the persistent-buffer path is cudagraph-safe (no force-eager). +- Capture (eager): a client-spec request (`post_attn`, layer 5, `last_prompt`) + delivers one `(1, hidden)` bf16 row to a driver consumer's `on_capture`. + +Once validated, the interim Phase-1 fallback guard was removed so v2 actually +runs these features (auto-selected for Qwen3, or via the env override). + +Not yet exercised on GPU (mirrors v1, but unverified here): TP>1 / PP>1, +per-request inline steering and named modules, capture prefix-cache reuse, +preemption resume, and spec-decode token layout. ## Validation diff --git a/tests/test_config.py b/tests/test_config.py index df40453a8f4a..5c01d652a17a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -148,41 +148,6 @@ def test_is_default_v2_model_runner_model(model_config, expected): assert VllmConfig._is_default_v2_model_runner_model(config) is expected -# A non-None sentinel stands in for a configured SteeringConfig / -# CaptureConsumersConfig; the predicate only checks ``is not None``. -_ENABLED = object() - - -@pytest.mark.parametrize( - ("steering_config", "capture_config", "expected"), - [ - (None, None, []), - (_ENABLED, None, ["activation steering"]), - (None, _ENABLED, ["activation capture"]), - (_ENABLED, _ENABLED, ["activation steering", "activation capture"]), - ], -) -def test_v2_silently_broken_features(steering_config, capture_config, expected): - config = SimpleNamespace( - steering_config=steering_config, - capture_consumers_config=capture_config, - ) - - assert VllmConfig._v2_model_runner_silently_broken_features(config) == expected - - -def test_validate_v2_model_runner_fails_closed_on_silently_broken(monkeypatch): - # An explicit VLLM_USE_V2_MODEL_RUNNER opt-in is enforced here: a config the - # v2 runner would silently no-op must raise rather than be honored. - monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True) - config = SimpleNamespace( - _get_v2_model_runner_unsupported_features=lambda: ["activation steering"], - ) - - with pytest.raises(ValueError, match="activation steering"): - VllmConfig._validate_v2_model_runner(config) - - @pytest.mark.skip_global_cleanup def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping( monkeypatch, diff --git a/tests/v1/worker/test_gpu_v2_steering_glue.py b/tests/v1/worker/test_gpu_v2_steering_glue.py index b5528d080f1e..d47beb1a0c8c 100644 --- a/tests/v1/worker/test_gpu_v2_steering_glue.py +++ b/tests/v1/worker/test_gpu_v2_steering_glue.py @@ -23,9 +23,9 @@ class _FakeManager: - def __init__(self): + def __init__(self, has_globals=False): self.config_to_row: dict = {} - self.global_base_vectors: dict = {} + self.global_base_vectors: dict = {"pre_attn": {0: [1.0]}} if has_globals else {} self.global_prefill_vectors: dict = {} self.global_decode_vectors: dict = {} self._tables_dirty = False @@ -41,7 +41,14 @@ def release_config(self, h, phase): self.released.append((h, phase)) def get_row_for_config(self, h, is_prefill): - return int(h) # row == hash keeps the expected index readable + # Mirror SteeringManager: a per-request hash maps to its own row; hash 0 + # maps to the global prefill (1) / decode (2) row when globals are set, + # else to the row-0 no-steer sentinel. + if h != 0: + return int(h) + if self.global_base_vectors: + return 1 if is_prefill else 2 + return 0 def populate_steering_tables(self, layers): self.populated += 1 @@ -55,9 +62,11 @@ def _layer(num_tokens=16): return layer -def _make_glue(num_computed, max_tokens=16, max_seqs=8): +def _make_glue( + num_computed, prompt_len=None, has_globals=False, max_tokens=16, max_seqs=8 +): glue = SteeringRunnerMixin.__new__(SteeringRunnerMixin) - glue._steering_manager = _FakeManager() + glue._steering_manager = _FakeManager(has_globals=has_globals) glue._steerable_layers_cache = {0: _layer(max_tokens)} glue._steering_reqs = {} glue._steering_index_dirty = False @@ -65,8 +74,11 @@ def _make_glue(num_computed, max_tokens=16, max_seqs=8): glue._steering_rows_scratch = np.zeros(max_seqs, dtype=np.int64) glue._steering_n_tokens_scratch = np.zeros(max_seqs, dtype=np.int64) glue._steering_index_pinned = torch.zeros(max_tokens, dtype=torch.long) + if prompt_len is None: + prompt_len = [0] * len(num_computed) glue.req_states = SimpleNamespace( - num_computed_tokens_np=np.asarray(num_computed, dtype=np.int32) + num_computed_tokens_np=np.asarray(num_computed, dtype=np.int32), + prompt_len=SimpleNamespace(np=np.asarray(prompt_len, dtype=np.int32)), ) # Avoid building real SamplingParams: the resolve step just needs to be truthy. glue._resolve_request_steering = lambda sp, phase: {"pre_attn": {0: [1.0]}} @@ -177,6 +189,23 @@ def test_update_buffers_builds_per_token_index_and_transition(): assert glue._steering_reqs["p"].phase == "decode" +def test_global_steering_applies_to_untracked_request(): + # Globals are set but the request carries no per-request config (untracked); + # it must still pick up the global row, not the no-steer sentinel. + glue = _make_glue(num_computed=[0], prompt_len=[5], has_globals=True) + input_batch = SimpleNamespace( + num_reqs=1, req_ids=["g"], idx_mapping_np=np.asarray([0], dtype=np.int32) + ) + sched = SimpleNamespace(num_scheduled_tokens={"g": 5}) + + glue._update_steering_buffers_v2(sched, input_batch) + + steering_index = glue._steerable_layers_cache[0].steering_index + # Prefilling (computed 0 < prompt 5) with globals -> global prefill row 1. + assert steering_index[:5].tolist() == [1, 1, 1, 1, 1] + assert steering_index[5:].sum().item() == 0 + + def test_update_buffers_short_circuit_zeroes_dirty_index(): glue = _make_glue(num_computed=[0]) # No tracked requests and no globals -> nothing active. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3c95e6ecf7a1..53f8cf147f88 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -509,10 +509,6 @@ def num_speculative_tokens(self) -> int: def use_v2_model_runner(self) -> bool: use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER if use_v2_model_runner is not None: - # An explicit opt-in still runs through _validate_v2_model_runner(), - # which fails closed on _get_v2_model_runner_unsupported_features() - # (steering/capture included), so the override can't silently land - # on a runner that would no-op those features. return use_v2_model_runner if not self._is_default_v2_model_runner_model(): @@ -2072,27 +2068,8 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: # Will be added by https://github.com/vllm-project/vllm/pull/38390 unsupported.append("EC transfer") - unsupported.extend(self._v2_model_runner_silently_broken_features()) - return unsupported - def _v2_model_runner_silently_broken_features(self) -> list[str]: - """Features the v2 model runner would silently no-op rather than fail on. - - The activation steering and capture control planes are only wired into - the v1 model runner. The v2 runner loads the same model (so the layer - hooks exist) but never populates the steering tables or installs a - capture manager, so these features would silently do nothing instead of - raising. Callers must treat them as unsupported until the control plane - is ported to v2. - """ - features: list[str] = [] - if self.steering_config is not None: - features.append("activation steering") - if self.capture_consumers_config is not None: - features.append("activation capture") - return features - def _validate_v2_model_runner(self) -> None: """Check for features not yet supported by the V2 model runner.""" if not HAS_TRITON: diff --git a/vllm/v1/worker/gpu/steering_runner_mixin.py b/vllm/v1/worker/gpu/steering_runner_mixin.py index 27461583acb1..cd3b639803c4 100644 --- a/vllm/v1/worker/gpu/steering_runner_mixin.py +++ b/vllm/v1/worker/gpu/steering_runner_mixin.py @@ -260,6 +260,7 @@ def _update_steering_buffers_v2( self._steering_n_tokens_scratch = n_tokens_scratch num_computed_np = self.req_states.num_computed_tokens_np + prompt_len_np = self.req_states.prompt_len.np active_count = 0 for i in range(num_reqs): req_id = req_ids[i] @@ -267,26 +268,30 @@ def _update_steering_buffers_v2( if n_tokens == 0: continue + req_idx = int(idx_np[i]) + num_computed = int(num_computed_np[req_idx]) + # Untracked requests (no per-request config) still pass through + # get_row_for_config with hash 0 so any global vectors apply — the + # manager maps hash 0 to the global prefill/decode row (or to the + # row-0 no-steer sentinel when no globals are set). rs = reqs.get(req_id) - if rs is None: - # No steering for this request — row 0 is the no-steer sentinel. - rows_scratch[active_count] = 0 - n_tokens_scratch[active_count] = n_tokens - active_count += 1 - continue + if rs is not None: + num_prompt = rs.num_prompt_tokens + prefill_hash = rs.prefill_hash + decode_hash = rs.decode_hash + else: + num_prompt = int(prompt_len_np[req_idx]) + prefill_hash = 0 + decode_hash = 0 - num_computed = int(num_computed_np[int(idx_np[i])]) - num_prompt = rs.num_prompt_tokens if num_computed < num_prompt: - row = mgr.get_row_for_config(rs.prefill_hash, is_prefill=True) - rows_scratch[active_count] = row - n_tokens_scratch[active_count] = n_tokens - if num_computed + n_tokens >= num_prompt: + row = mgr.get_row_for_config(prefill_hash, is_prefill=True) + if rs is not None and num_computed + n_tokens >= num_prompt: self._steering_transition(rs) else: - row = mgr.get_row_for_config(rs.decode_hash, is_prefill=False) - rows_scratch[active_count] = row - n_tokens_scratch[active_count] = n_tokens + row = mgr.get_row_for_config(decode_hash, is_prefill=False) + rows_scratch[active_count] = row + n_tokens_scratch[active_count] = n_tokens active_count += 1 if active_count > 0: From edec90554f0659e6c5c0144cd5746ea6861cd3a1 Mon Sep 17 00:00:00 2001 From: RhizoNymph Date: Thu, 18 Jun 2026 23:18:34 -0700 Subject: [PATCH 4/6] docs(v2-runner): record full steering/capture validation matrix (single-node, TP/PP, preemption, streaming) --- docs/design/v2_runner_steering_capture.md | 64 +++++++++++++++++++++-- 1 file changed, 61 insertions(+), 3 deletions(-) diff --git a/docs/design/v2_runner_steering_capture.md b/docs/design/v2_runner_steering_capture.md index c784416ef372..63ca7a83468a 100644 --- a/docs/design/v2_runner_steering_capture.md +++ b/docs/design/v2_runner_steering_capture.md @@ -133,9 +133,67 @@ GPU-validated on Qwen3-0.6B (RTX 3090, TP1/PP1), forcing Once validated, the interim Phase-1 fallback guard was removed so v2 actually runs these features (auto-selected for Qwen3, or via the env override). -Not yet exercised on GPU (mirrors v1, but unverified here): TP>1 / PP>1, -per-request inline steering and named modules, capture prefix-cache reuse, -preemption resume, and spec-decode token layout. +Expanded GPU matrix (Qwen3-0.6B unless noted): + +- Steering: global (eager + cudagraph); per-request inline; **mixed batch** (a + steered and an unsteered request together — the unsteered output is + byte-identical to baseline, so per-request rows don't cross-contaminate); + decode-only (lazy decode-config registration at the prefill→decode boundary); + per-request under cudagraph; chunked prefill (multi-step prefill). +- Capture: client-spec eager; client-spec **under cudagraph** (the force-eager + gate fires for that step); `all_generated` positions (multi-step decode); + global-spec under cudagraph (persistent-buffer path, no force-eager). +- Cross-node (2×3090, Ray): steering under TP=2 and PP=2 (rank-replication and + `locally_owned_layers` confirmed). +- Model coverage: gemma-3-4b-it runs on v2 with steering (hidden 2560 / 34 + layers) — the port is not Qwen3-specific. + +Additional GPU coverage: + +- Steering: named-module (`register_steering_modules` + `steering_module_ref`) + and per-request scale (scale 0 → baseline, scale 1 → steered); async + scheduling. +- Capture: filesystem consumer (worker-location, files read back); multiple + consumers (filesystem + global logging); activation-store **write** path + (64 prompt rows written with prefix caching on — block-hash wiring works); + async scheduling. +- Capture under **TP=2** (exactly one rank — TP rank 0 — writes; the other + writes nothing) and **PP=2** (stage 0 captures its layer 5, stage 1 captures + its layer 20 — per-stage `local_layer_range` filtering correct), via the + worker-location filesystem consumer. + +- Capture + prefix caching via the **OpenAI server** (v2, activation store on): + a repeated prefix reuses under `all_generated` (32-token prefix-cache hit on + the 2nd request) and recaptures under `all_prompt` (0 hits, full re-forward) — + identical to the documented v1 behavior. Store write validated separately + (64 rows). The Step-A store *serve* path (`pop_pending_serve` / + `serve_from_store`) was not observed to trigger, consistent with v1's + "all_prompt → full recapture", so those two lines stay formally unexercised. + +- **Preemption resume**: under a tiny KV cache, the worker observed 248 + preemption events; all 16 steered requests still produced the correct steered + output (no config leak). Capture under preemption: 72 preemption events, all + 24/24 capturing requests still delivered — preempted capturing requests resume + and capture cleanly (no lost/double captures). +- **Steering hook points**: pre_attn and post_mlp both shift/clear correctly + (post_attn was already covered); the prefill-only tier + (`prefill_steering_vectors`) steers. +- **Capture positions**: `all` (prompt+generated rows) and an explicit index + list (`[0, 2]` → exactly 2 rows), in addition to `last_prompt`/`all_generated`. + +- **Streaming re-add**: an async streaming-input session (prompt fed in chunks + via `AsyncLLM.generate(prompt=)`) with + steering produced steered output, and the port's re-add branch fired + (`_steering_add_request` saw an already-tracked req_id → released the old + config + registered the new one). No crash. + +Still unverified: spec-decode; DP; the async-dispatch overload policies +(`spill`/`drop`/`block`); the store *serve* path (doesn't trigger for +`all_prompt` even on v1 — full recapture by design); and capture under streaming +re-add specifically (steering's re-add is validated; capture's +`register_request` on a re-added id is the narrow analog). (The capture-consumer +entry points were missing from one prebuilt install's metadata — a stale +dist-info issue fixed by reinstalling; pyproject already declares them.) ## Validation From 40133ec11ac250558a86c337dfa1d22308c0e945 Mon Sep 17 00:00:00 2001 From: RhizoNymph Date: Fri, 19 Jun 2026 19:40:23 -0700 Subject: [PATCH 5/6] fix(v2-runner): refresh capture registration on streaming re-add and preemption resume --- docs/design/v2_runner_steering_capture.md | 28 ++++- tests/v1/worker/test_gpu_v2_capture_glue.py | 124 ++++++++++++++++++++ vllm/v1/worker/gpu/capture_runner_mixin.py | 36 +++++- vllm/v1/worker/gpu/model_runner.py | 9 +- 4 files changed, 185 insertions(+), 12 deletions(-) diff --git a/docs/design/v2_runner_steering_capture.md b/docs/design/v2_runner_steering_capture.md index 63ca7a83468a..74ea932715d0 100644 --- a/docs/design/v2_runner_steering_capture.md +++ b/docs/design/v2_runner_steering_capture.md @@ -187,11 +187,31 @@ Additional GPU coverage: (`_steering_add_request` saw an already-tracked req_id → released the old config + registered the new one). No crash. +- **Capture re-add / preemption resume** (the asymmetry steering already + handled): `add_requests` calls `_remove_request` first, which does *not* touch + capture state, so a re-admitted request would re-register an already-registered + id. The capture manager raises on duplicate ids (`already registered`), caught + as a request error. Two paths reach this: + - *Streaming re-add* — the request is still live (`_remove_request` returns + `True`); the grown prompt makes the prior chunk's registration stale, so + `_capture_add_request` now discards it (gate `drop` + `unregister_request`, + no finalize) and re-registers against the new prompt. + - *Preemption resume* — on v2 the scheduler folds `scheduled_resumed_reqs` + into `scheduled_new_reqs`, so a resumed request flows through + `_capture_add_request` with `was_present=False` while its registration + survived (capture is intentionally not finalized on preempt). It is kept + as-is (skip re-registration), preserving rows captured before preemption. + GPU-validated on Qwen3-0.6B with a clean before/after: pre-fix, a streaming + session logged `capture request '...' is already registered` on each re-add, + and the preemption scenario (24 capturing requests, 64-block KV cache) logged + 20 such rejections; post-fix, both paths log zero rejections and deliver all + captures (24/24 under preemption). CPU glue tests cover the three branches + (fresh / streaming re-add / preemption resume) plus the non-capturer rank. + Still unverified: spec-decode; DP; the async-dispatch overload policies -(`spill`/`drop`/`block`); the store *serve* path (doesn't trigger for -`all_prompt` even on v1 — full recapture by design); and capture under streaming -re-add specifically (steering's re-add is validated; capture's -`register_request` on a re-added id is the narrow analog). (The capture-consumer +(`spill`/`drop`/`block` — runner-agnostic transport code shared with v1, not +touched by the port); and the store *serve* path (doesn't trigger for +`all_prompt` even on v1 — full recapture by design). (The capture-consumer entry points were missing from one prebuilt install's metadata — a stale dist-info issue fixed by reinstalling; pyproject already declares them.) diff --git a/tests/v1/worker/test_gpu_v2_capture_glue.py b/tests/v1/worker/test_gpu_v2_capture_glue.py index 624234a04a34..18a0a1147677 100644 --- a/tests/v1/worker/test_gpu_v2_capture_glue.py +++ b/tests/v1/worker/test_gpu_v2_capture_glue.py @@ -120,3 +120,127 @@ def test_drain_disabled_returns_empty(): glue._pending_capture_results = {"a": {"c": object()}} assert glue._drain_capture_results() == {} + + +# ---- re-add / preemption-resume admission --------------------------------- + + +class _FakeGate: + def __init__(self): + self.registered = {} + self.dropped = [] + + def register(self, req_id, raw): + self.registered[req_id] = raw + + def drop(self, req_id): + self.dropped.append(req_id) + self.registered.pop(req_id, None) + + +class _FakeManager: + def __init__(self): + self._reqs = set() + self.unregistered = [] + + def has_request(self, req_id): + return req_id in self._reqs + + def unregister_request(self, req_id): + self.unregistered.append(req_id) + self._reqs.discard(req_id) + + +class _AddGlue(CaptureRunnerMixin): + """Host that records admissions and isolates the re-add branching.""" + + def __init__(self, gate, mgr): + self._capture_feature_enabled = True + self._capture_step_gate = gate + self._capture_manager = mgr + self.registered_calls = [] + + # Stub the full registration machinery; only the branching is under test. + def _register_capture_request(self, new_req_data): + self.registered_calls.append(new_req_data.req_id) + if self._capture_manager is not None: + self._capture_manager._reqs.add(new_req_data.req_id) + + +def _new_req(req_id, capture): + return SimpleNamespace( + req_id=req_id, + sampling_params=SimpleNamespace(capture=capture), + ) + + +def test_capture_add_fresh_request_registers(): + gate, mgr = _FakeGate(), _FakeManager() + glue = _AddGlue(gate, mgr) + + glue._capture_add_request(_new_req("a", {"c": {}}), was_present=False) + + assert glue.registered_calls == ["a"] + assert gate.registered == {"a": {"c": {}}} + assert gate.dropped == [] + assert mgr.unregistered == [] + + +def test_capture_add_streaming_readd_discards_and_reregisters(): + gate, mgr = _FakeGate(), _FakeManager() + # Prior chunk already admitted. + mgr._reqs.add("a") + gate.registered["a"] = {"c": {"positions": "last_prompt"}} + glue = _AddGlue(gate, mgr) + + # Re-add (still live) with a different capture spec. + glue._capture_add_request( + _new_req("a", {"c": {"positions": "all_generated"}}), was_present=True + ) + + # Stale registration dropped, then re-registered against the new prompt. + assert gate.dropped == ["a"] + assert mgr.unregistered == ["a"] + assert glue.registered_calls == ["a"] + assert gate.registered["a"] == {"c": {"positions": "all_generated"}} + + +def test_capture_add_preemption_resume_keeps_registration(): + gate, mgr = _FakeGate(), _FakeManager() + # Registration survived preemption (finish_requests did not finalize it). + mgr._reqs.add("a") + gate.registered["a"] = {"c": {}} + glue = _AddGlue(gate, mgr) + + # Resumed req is folded into scheduled_new_reqs on v2, but was_present is + # False because finish_requests removed it from req_states on preempt. + glue._capture_add_request(_new_req("a", {"c": {}}), was_present=False) + + # No discard, no re-registration — the open registration is reused. + assert gate.dropped == [] + assert mgr.unregistered == [] + assert glue.registered_calls == [] + + +def test_capture_add_readd_on_non_capturer_rank_no_manager(): + gate = _FakeGate() + gate.registered["a"] = {"c": {}} + glue = _AddGlue(gate, None) # non-capturer rank: no manager + + # Streaming re-add still refreshes the rank-replicated gate, no crash. + glue._capture_add_request(_new_req("a", {"c": {"positions": "all"}}), True) + + assert gate.dropped == ["a"] + assert gate.registered["a"] == {"c": {"positions": "all"}} + assert glue.registered_calls == [] + + +def test_capture_add_disabled_is_noop(): + gate, mgr = _FakeGate(), _FakeManager() + glue = _AddGlue(gate, mgr) + glue._capture_feature_enabled = False + + glue._capture_add_request(_new_req("a", {"c": {}}), was_present=True) + + assert gate.registered == {} and gate.dropped == [] + assert glue.registered_calls == [] diff --git a/vllm/v1/worker/gpu/capture_runner_mixin.py b/vllm/v1/worker/gpu/capture_runner_mixin.py index 09a68a9c6be9..0923df5b6388 100644 --- a/vllm/v1/worker/gpu/capture_runner_mixin.py +++ b/vllm/v1/worker/gpu/capture_runner_mixin.py @@ -145,20 +145,46 @@ def _init_capture_state(self) -> None: # ---- request lifecycle ------------------------------------------------- - def _capture_add_request(self, new_req_data: NewRequestData) -> None: + def _capture_add_request( + self, new_req_data: NewRequestData, was_present: bool + ) -> None: """Hook a newly admitted request into capture tracking. ``register`` runs on every rank (the gate is rank-replicated); manager registration runs only on the capturer rank. + + ``was_present`` is ``True`` only for a **streaming re-add** — a still + live request re-admitted with a grown prompt. In that case the prior + chunk's gate selector and manager registration are stale and must be + discarded (without finalizing — a partial first-chunk capture is + dropped, not emitted) before re-registering against the new prompt. + This mirrors the steering control plane's re-add handling. + + A **preemption resume** reaches this path too (the v2 scheduler folds + ``scheduled_resumed_reqs`` into ``scheduled_new_reqs``), but with + ``was_present`` ``False`` because ``finish_requests`` removed the + request on preemption while intentionally leaving its capture + registration open. Such a request is re-prefilled (recompute) into the + existing registration, so we keep it and skip re-registration — both + avoiding the manager's duplicate-register error and preserving any + rows already captured before preemption. """ if not self._capture_feature_enabled: return + req_id = new_req_data.req_id + mgr = self._capture_manager + if was_present: + # Streaming re-add: prior chunk's capture state is stale. + if self._capture_step_gate is not None: + self._capture_step_gate.drop(req_id) + if mgr is not None: + mgr.unregister_request(req_id) sp = new_req_data.sampling_params if self._capture_step_gate is not None: - self._capture_step_gate.register( - new_req_data.req_id, getattr(sp, "capture", None) - ) - if self._capture_manager is not None: + self._capture_step_gate.register(req_id, getattr(sp, "capture", None)) + # Skip re-registration if the request is already registered (a + # preemption resume whose registration survived); otherwise admit it. + if mgr is not None and not mgr.has_request(req_id): self._register_capture_request(new_req_data) def _capture_finish_request(self, req_id: str) -> None: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0f06c14172a4..5ae950722ae9 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -714,8 +714,11 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: # Streaming input update: request already exists from a prior # chunk. Remove old state so it can be cleanly re-added below - # with the updated prompt_token_ids and mm_features. - self._remove_request(req_id) + # with the updated prompt_token_ids and mm_features. The return + # value distinguishes a streaming re-add (request was still live) + # from a fresh admit / preemption resume (already gone) so the + # capture control plane can refresh stale registrations correctly. + was_present = self._remove_request(req_id) prompt_len = len(new_req_data.prompt_token_ids) self.req_states.add_request( @@ -749,7 +752,7 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: # registers with the gate; the capturer rank also registers with # the manager). Runs on all PP ranks since capturer layers may live # on any stage. - self._capture_add_request(new_req_data) + self._capture_add_request(new_req_data, was_present) # Register the request's steering config (rank-local; deterministic # across the TP/PP topology from the broadcast scheduler output). self._steering_add_request(new_req_data) From 4646f0c7557d7488b5825a0d12fcf4d6bcb01a99 Mon Sep 17 00:00:00 2001 From: RhizoNymph Date: Sat, 20 Jun 2026 22:39:13 -0700 Subject: [PATCH 6/6] docs(v2-runner): record steering-under-APC and capture TP/PP-cudagraph validation --- docs/design/v2_runner_steering_capture.md | 38 +++++++++++++++++++---- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/docs/design/v2_runner_steering_capture.md b/docs/design/v2_runner_steering_capture.md index 74ea932715d0..87d45402b938 100644 --- a/docs/design/v2_runner_steering_capture.md +++ b/docs/design/v2_runner_steering_capture.md @@ -161,6 +161,28 @@ Additional GPU coverage: writes nothing) and **PP=2** (stage 0 captures its layer 5, stage 1 captures its layer 20 — per-stage `local_layer_range` filtering correct), via the worker-location filesystem consumer. +- Capture under **TP=2 and PP=2 *with cudagraph*** (cross-node, 2×3090 Ray; + `enforce_eager=False`, FULL_AND_PIECEWISE graphs compiled): mixing plain + (cudagraph) and client-spec capturing (force-eager gate) requests in the same + run does not hang — the force-eager decision stays rank-replicated across + ranks/stages so every rank toggles eager↔graph in lockstep (including the PP + P2P send/recv between stages). TP rank 0 / each PP stage wrote exactly its own + 30720-byte (`15 gen tokens × 1024 × bf16`) capture and no other rank did. +- Steering **under prefix caching**: a steered request whose prompt is largely + served from the KV cache (partial hit — the scheduler always reserves the last + block to recompute logits) still steers correctly (degenerate `οοο` output vs + the unsteered baseline of the same tokens). KV block hashes are steering-aware, + so a steered request does not reuse an unsteered cache. The narrower + *admit-straight-to-decode* branch (`_steering_add_request`, `num_computed >= + num_prompt`) is **not reachable under single-engine APC** — the scheduler caps + `num_computed` at `num_prompt − block_size`, never `>= num_prompt`; that branch + is only reachable when an external mechanism (KV connector / disaggregated + prefill) sets `num_computed_tokens` to the full prompt at admission. It mirrors + the v1 mixin's identical defensive branch and stays formally unexercised + without a connector. (Reminder: inline `SamplingParams.steering_vectors` + requires `enable_steering=True` at engine init — otherwise the worker steering + manager is `None` and steering silently no-ops; the per-request hash is still + packed host-side.) - Capture + prefix caching via the **OpenAI server** (v2, activation store on): a repeated prefix reuses under `all_generated` (32-token prefix-cache hit on @@ -208,12 +230,16 @@ Additional GPU coverage: captures (24/24 under preemption). CPU glue tests cover the three branches (fresh / streaming re-add / preemption resume) plus the non-capturer rank. -Still unverified: spec-decode; DP; the async-dispatch overload policies -(`spill`/`drop`/`block` — runner-agnostic transport code shared with v1, not -touched by the port); and the store *serve* path (doesn't trigger for -`all_prompt` even on v1 — full recapture by design). (The capture-consumer -entry points were missing from one prebuilt install's metadata — a stale -dist-info issue fixed by reinstalling; pyproject already declares them.) +Still unverified: spec-decode; DP; combined 2-D parallelism (TP *and* PP on the +same request, and TP/PP > 2 — each validated independently, intersection needs +≥4 GPUs); the async-dispatch overload policies (`spill`/`drop`/`block` — +runner-agnostic transport code shared with v1, not touched by the port); the +store *serve* path (doesn't trigger for `all_prompt` even on v1 — full recapture +by design); and the steering *admit-straight-to-decode* branch (needs a KV +connector / disaggregated prefill — unreachable under single-engine APC, see +above). (The capture-consumer entry points were missing from one prebuilt +install's metadata — a stale dist-info issue fixed by reinstalling; pyproject +already declares them.) ## Validation