feat(plugins): dynamic hidden-states layer swap for extract_hidden_states

[reed-meyerson]

Summary

Adds a self-contained vLLM plugin (vllm_plugins/dynamic_hidden_states) that
lets you swap which target-model hidden layers are captured during
speculator training-data generation (method="extract_hidden_states") at
runtime over HTTP, without restarting the engine. Only the set of layer
indices changes; the count is fixed at launch (proposer buffers, dummy-proposer
KV-cache shape and on-disk safetensors layout depend on it), so count changes are
rejected (fail closed).

This is opt-in and touches no vLLM core files — it wires in entirely through
public plugin entry points and --worker-extension-cls.

How it works

• graph_patch.install (vllm.general_plugins) — a pre-compile
  monkeypatch of EagleModelMixin._maybe_add_hidden_state /
  _set_aux_hidden_state_layers. The stock capture uses a Python membership test
  (if layer_idx in self.aux_hidden_state_layers) that torch.compile
  constant-folds into the compiled forward, freezing the selection — so a naive
  attribute swap is a silent no-op unless the engine runs --enforce-eager.
  Instead we keep N fixed accumulator buffers and fold every candidate layer's
  residual into them weighted by a one-hot selection mask held in a registered
  buffer. The graph becomes selection-agnostic, so a swap is an in-place
  mask.copy_(...) — same storage address, read live by both the compiled graph
  and captured CUDA graphs. No recompile, no --enforce-eager. Memory stays
  O(N); the cost is a small multiply-add over every candidate layer's residual
  (~0.006% of forward FLOPs).
• AuxLayerWorkerExtension (--worker-extension-cls) — engine-side RPC
  (get/set_aux_hidden_state_layers_rpc) callable through collective_rpc.
• AuxLayerEndpoint (vllm.endpoint_plugins) — registers
  GET/POST /aux_hidden_state_layers on the OpenAI-compatible server.

Testing

Verified on Qwen3-8B (1×H100) under the default torch.compile +
FULL_AND_PIECEWISE CUDA graphs (no --enforce-eager):

• Swapping only the first captured layer (2→6) changes exactly column 0
  (max|Δ|=21.125) while columns for unchanged layers 18/34 stay
  bit-identical (Δ=0).
• The compiled masked path reproduces the eager baked path's values exactly.
• No recompilation is triggered by a swap.
• Count-mismatch and empty-layer requests are rejected with 400.

An end-to-end integration script is included at
vllm_plugins/tests/test_swap_e2e.py (run manually against a live server; see
its docstring and vllm_plugins/README.md).

Scope / limitations

• Covers models that capture through EagleModelMixin._maybe_add_hidden_state
  (Qwen2/Qwen3, Llama, generic dense/MoE). Models that inline the membership test
  in their own forward (e.g. deepseek_v2, qwen3_next) are not covered by
  the compile patch and still need --enforce-eager for a live swap.
• The hidden-states connector's safetensors carry no layer-set metadata, so a
  controller that swaps mid-run must record its own {time/req-range → layer set}
  mapping.

Notes

Opening as a draft for discussion — feedback welcome on whether this belongs
here as a plugin vs. upstreaming the masked-accumulate capture into vLLM core.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 8c2db222-0097-4302-979b-4b3c7fdc0c4e

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz