Steering + capture on the v2 model runner

[RhizoNymph]

What

Ports the activation steering and capture control planes to the experimental v2 GPU model runner (vllm/v1/worker/gpu/model_runner.py), so they work there the same way they do on the v1 runner.

Why

The steering/capture data plane (the apply_steering / capture_residual custom ops, per-layer buffers, kernels, SteeringManager/CaptureManager/CaptureStepGate/ActivationStore) already lives in model_executor/ + v1/capture/ and is shared by both runners — both load the same model, so the in-forward hooks already fire on v2 and safely no-op when nothing drives them. Only the runner-side control plane was v1-only. Without it, turning on the v2 runner with steering/capture configured silently did nothing.

How

Two v2-native modules wire the existing subsystems into v2's lifecycle (add_requests / finish_requests / execute_model / sample_tokens); the v1 runner is untouched.

• gpu/capture_runner_mixin.py (CaptureRunnerMixin): manager/gate/store init, request register/finalize, the per-step force-eager decision (client-spec captures only — global specs ride the cudagraph-safe persistent-buffer path), gather-plan build, and draining results onto ModelRunnerOutput.capture_results.
• gpu/steering_runner_mixin.py (SteeringRunnerMixin): subclasses the v1 SteeringModelRunnerMixin, reusing init / layer discovery / validation / the public RPC API / _resolve_request_steering unchanged, and overrides only the three methods that touched v1-runner state. It keeps its own per-request state (v2 retains no CachedRequestState), drives prefill→decode transitions and the per-token steering index, and needs no force-eager (persistent buffers are cudagraph-safe). gpu_worker.py already forwards the steering RPCs.

Two v2-specific correctness fixes are included:

• Global steering for untracked requests: requests with no per-request config are routed through get_row_for_config(0, …) so any global steering vectors still apply (instead of the row-0 no-steer sentinel).
• Capture registration on re-add (_capture_add_request now takes was_present): 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 manager raises already registered). A streaming re-add (still-live request, grown prompt) discards the stale registration and re-registers; a preemption resume (v2 folds scheduled_resumed_reqs into scheduled_new_reqs) keeps the surviving registration. This mirrors the symmetric handling steering already had.

Design notes: docs/design/v2_runner_steering_capture.md.

Validation

CPU glue tests — tests/v1/worker/test_gpu_v2_{capture,steering}_glue.py (19 tests: state bookkeeping, view construction, gate decisions, and the re-add/preemption/non-capturer branches).

GPU — Qwen3-0.6B unless noted, VLLM_USE_V2_MODEL_RUNNER=1, on 1×/2× RTX 3090. Mirrors the v1 validation matrix.

Steering:

• Global vectors eager and cudagraph: output shifts; clearing restores the exact baseline (persistent-buffer path is cudagraph-safe — no force-eager).
• Per-request inline vectors and named modules (register_steering_modules + steering_module_ref); per-request scale (0 → baseline, 1 → steered); all three hook points (pre_attn/post_attn/post_mlp); the prefill-only tier.
• Mixed batch — a steered and an unsteered request together: the unsteered output is byte-identical to baseline (per-request rows don't cross-contaminate).
• Decode-only (lazy decode-config registration at the prefill→decode boundary); per-request under cudagraph; chunked prefill.
• Under prefix caching: a steered request whose prompt is largely served from the KV cache still steers correctly (KV block hashes are steering-aware, so a steered request never reuses an unsteered cache).
• Streaming re-add: an async streaming-input session re-feeds the prompt in chunks; the re-add branch releases the old config and registers the new one. No crash.
• Preemption resume: under a tiny KV cache (248 preemption events observed), all 16 steered requests still produced correct steered output (no config leak).
• Cross-node (2×3090, Ray): steering under TP=2 and PP=2 (rank-replication and locally_owned_layers filtering confirmed; no NCCL hang).
• Model coverage: gemma-3-4b-it runs on v2 with steering (hidden 2560 / 34 layers) — the port is not Qwen3-specific.
• Async scheduling.

Capture:

• Client-spec eager (delivers a (1, hidden) bf16 row to a consumer's on_capture) and under cudagraph (the force-eager gate fires for that step); global-spec under cudagraph (persistent-buffer path, no force-eager).
• Position modes: last_prompt, all_generated, all, and an explicit index list.
• Consumers: filesystem (worker-location, files read back) and global logging; multiple consumers at once; the activation-store write path (64 prompt rows with prefix caching on — block-hash wiring works).
• Prefix-cache reuse via the OpenAI server: a repeated prefix reuses under all_generated (cache hit on the 2nd request) and recaptures under all_prompt (full re-forward) — identical to the v1 behavior.
• Under preemption (72 events): all 24 capturing requests still delivered — preempted capturers resume and capture cleanly (no lost/double captures).
• Cross-node (2×3090, Ray) under TP=2 (exactly TP rank 0 writes; the other writes nothing) and PP=2 (stage 0 captures its layer, stage 1 captures its own — per-stage local_layer_range filtering correct).
• TP=2 and PP=2 with cudagraph (cross-node): mixing plain (cudagraph) and client-spec capturing (force-eager) requests does not hang — the force-eager decision stays rank-replicated, so every rank toggles eager↔graph in lockstep (including the PP P2P send/recv). Each owning rank/stage wrote exactly its own 15 gen × 1024 × bf16 capture and no other rank did.
• Async scheduling.

Performance — v1-vs-v2 latency (Qwen3-0.6B, steering-bench, cudagraphs on): every point within ±1% (measurement noise). The v2 port adds no latency penalty, and the steering-overhead profile matches v1 (~4% fixed for steering-on-idle, ~8–10% for real per-request steering).

Not exercised (out of scope / unreachable on this setup, all shared with or matching v1): spec-decode; DP; combined 2-D TP+PP and TP/PP > 2 (each validated independently; intersection needs ≥4 GPUs); the async-dispatch overload policies (spill/drop/block — runner-agnostic transport shared with v1); 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 (num_computed >= num_prompt), which the single-engine scheduler never produces — it always reserves the last block to recompute logits — so it's reachable only via a KV connector / disaggregated prefill, and mirrors v1's identical defensive branch.

Notes

• Includes an interim commit that added a fail-closed v1 fallback when steering/capture is configured, and a later commit that removes it now that the port exists — the two net out, so the runtime/test diff is purely additive (+1311 / −4 across the runner, two mixins, and glue tests).
• AI assistance was used for this change.

[RhizoNymph]

Also validated cross-node on 2×RTX 3090 (Ray, NCCL over bond0): steering passes under both TP=2 and PP=2 — global set_steering_vectors shifts output and clear restores baseline; per-worker results confirm the rank-replication invariant (TP: both ranks register; PP: stage 1 correctly reports no owned steered layers via locally_owned_layers).

[RhizoNymph]

Full validation matrix (GPU, 2×RTX 3090)

Extended the validation well beyond the initial pass. All on the v2 runner (VLLM_USE_V2_MODEL_RUNNER=1); see docs/design/v2_runner_steering_capture.md for details.

Steering

• Global (set_steering_vectors) and per-request inline (SamplingParams.steering_vectors), named-module (register_steering_modules + steering_module_ref), and per-request scale (0 → baseline, 1 → steered).
• Mixed batch: a steered and an unsteered request together — the unsteered output is byte-identical to baseline (per-request rows don't cross-contaminate).
• Decode-only / prefill→decode transition (lazy decode-config registration at the boundary); prefill-only tier.
• All three hook points (pre_attn / post_attn / post_mlp); eager and cudagraph (persistent buffers → no force-eager); chunked prefill.

Capture

• Client-spec eager and under cudagraph (the force-eager gate fires for the capturing step); global-spec persistent-buffer path.
• Position selectors: last_prompt, all_generated, all, explicit index list.
• Filesystem + logging consumers, multiple consumers together; activation-store write path (block-hash wiring); server-side prefix-cache reuse (all_generated reuses, all_prompt recaptures — matches v1).

Distributed (Ray, NCCL/bond0)

• Steering under TP=2 and PP=2 (rank-replication and per-stage locally_owned_layers confirmed).
• Capture under TP=2 (only TP rank 0 writes) and PP=2 (stage 0 → its layer, stage 1 → its layer).

Hard edges

• Preemption resume: 248 preemption events, all steered requests stayed steered; capture under preemption: 72 events, 24/24 requests delivered.
• Streaming re-add (async streaming-input session): re-add branch fires, output steered.

Coverage: Qwen3-0.6B and gemma-3-4b-it (v2 supports gemma-3; the port is not Qwen3-specific). Plus CPU glue tests for the v2-specific projection logic.

Out of scope / not exercised: spec-decode, DP, async-dispatch overload policies, and the store serve path (dormant by design — all_prompt recaptures on v1 too).

[RhizoNymph]

Follow-up fix: capture registration on re-add paths

Discovered while auditing whether the "untested" items actually touch the port's diff. Two paths re-enter _capture_add_request for an id the capture manager already holds, and the manager raises on duplicate ids (already registered) — caught as a request error. Steering already handled this symmetric case; capture didn't.

• Streaming re-add — add_requests → _remove_request (which does not touch capture state) → re-add of a still-live request with a grown prompt. The stale prior-chunk registration is now discarded (gate drop + unregister_request, no finalize) and re-registered 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 intentionally survived preemption. It's now kept as-is (skip re-registration), preserving rows captured before preemption.

The runtime change threads was_present (the _remove_request return) into _capture_add_request to distinguish a live re-add from a fresh admit / resume.

GPU-validated on Qwen3-0.6B (RTX 3090, v2 runner), clean before/after:

• Streaming session: pre-fix logged capture request '...' is already registered on each re-add; post-fix zero, capture delivered.
• Preemption (24 capturing requests, 64-block KV cache, forced eviction): pre-fix logged 20 already registered rejections; post-fix 0, 24/24 captures delivered.

CPU glue tests added for the three branches (fresh / streaming re-add / preemption resume) + non-capturer rank; full v2 glue suite 19 passed.

Also confirmed the async-dispatch overload policies (spill/drop/block) are not touched by the port — they live in the runner-agnostic transport shared with v1, so they stay out of scope here.

Commit: 40133ec

[RhizoNymph]

Follow-up tracked in #188: the capture_residual data-plane op returns its input tensor, which torch.compile now flags as a deprecated alias (hard error in PyTorch 2.12). Pre-existing in the shared data plane — not introduced here — but newly visible because this PR runs capture on the v2 torch.compile/cudagraph path. Benign on the current torch; not a blocker for this PR.