transformer_lens.model_bridge.sources.vllm.plugin module

vLLM plugin entry point.

Monkey-patches Worker.load_model to install capture hooks after weights load and before compile_or_warm_up_model — the only window where hooks make it into the compiled FX graph (PyTorch #117758).

Two hook flavors, selected by configure(enable_batching=...):

  • Compiled (default): in-place writes to a pre-allocated GPU tensor; no .cpu() (illegal during CUDA-graph capture); SymInt-indexed slicing only (Python .shape access forces specialization). Single prompt. Interventions ride the same hook as an affine transform output = output * scale_buf + bias_buf (defaults identity). The driver swaps buffer contents between forwards via tl_set_interventions — the FX graph references the buffers, so swaps take effect without recompiling. Memory cost: the affine allocates a transient output-shape tensor per hook per forward, even at identity — peak forward memory ~1.5× capture-only. Branching to skip the affine would defeat the swap trick and break the graph.

  • Batched (enable_batching=True, runs enforce_eager): the hook reads per-request token boundaries via segment_by_request (only valid inside a forward, untraceable under compile — hence eager), slices each request’s rows to CPU, and appends to per-(req_id, hook) accumulators across chunked-prefill forwards. Interventions apply directly to the tensor via _apply_op.

transformer_lens.model_bridge.sources.vllm.plugin.clear_config() → None

Unset the spec channel. Boot sites call this in a finally after LLM(...) so a later non-TL engine can’t inherit the specs.

transformer_lens.model_bridge.sources.vllm.plugin.configure(capture_specs: Dict[str, Tuple[str, int]], max_num_batched_tokens: int, dtype: dtype, enable_batching: bool = False, enable_position_interventions: bool = False) → None

Set capture specs, buffer length, dtype, and hook flavor before LLM(...).

transformer_lens.model_bridge.sources.vllm.plugin.register() → None

Idempotent monkey-patch of Worker.load_model.

vLLM calls register() once per process at entry-points discovery. Idempotent so re-imports (notebook restarts, repeated boot_vllm in the same process) don’t double-wrap.

No unregister() symmetry: the patch stays for process lifetime. Benign because patched_load_model no-ops when no spec channel is populated — and boot sites call clear_config() after each LLM(...), so any subsequent non-TL LLM(...) in the same process hits the no-op path.