transformer_lens.model_bridge.sources.vllm.worker_extension module¶

Worker extension exposed to collective_rpc for capture reads and intervention writes.

Hook installation lives in plugin (must happen pre-compile). This class only exposes the read/write surface. State is per-Worker so concurrent boot_vllm calls don’t collide. All methods prefixed tl_ to avoid colliding with vLLM Worker attributes.

Two capture modes, selected at boot:

Compiled (default): per-hook GPU buffers + affine scale/bias swap. Single prompt. tl_read_captures / tl_set_interventions.
Batched (enable_batching=True, eager): per-(req_id, hook) CPU accumulators filled in the hook via query_start_loc segmentation; arbitrary batch + chunked prefill. tl_read_batched_captures / tl_set_batched_interventions / tl_reset_accumulators.

class transformer_lens.model_bridge.sources.vllm.worker_extension.TLWorkerExtension¶

Bases: object

Mixed into vLLM’s Worker via worker_extension_cls.

tl_absent_hooks() → List[str] | None¶: Spec names that installed no hook on this rank — the boot site verifies their union across ranks covers every spec. None means installation never ran at all (plugin patch absent or spec channel empty), which the coverage check must treat as fatal rather than vacuously complete.

tl_get_param(dotted_name: str) → Dict[str, Any] | None¶

Read a named model tensor (e.g. model.norm.weight) as a wire-encoded CPU clone.

None if the path doesn’t resolve to a tensor. The bridge has no general weight surface, so this is how callers reach e.g. the ln_final weight.

tl_read_batched_captures(names: List[str] | None = None) → Dict[str, Dict[str, Tensor]]¶

Cat per-request chunks into {req_id: {hook: (seq, width)}} (token-order).

names restricts to those hooks (None = all). Note the per-chunk GPU→CPU copy already happened in the hook, so this only saves the cat.

tl_read_captures(prompt_lens: List[int], names: List[str] | None = None) → Dict[str, Dict[str, Any]]¶

Slice each capture buffer to sum(prompt_lens) rows; wire-encoded CPU copies.

names restricts the read (None = all) — this is the only GPU→CPU crossing, so it’s where a names_filtered run saves bandwidth. Caller gates sum(prompt_lens) <= max_num_batched_tokens, so total is in bounds.

tl_read_counter() → int¶: Total hook fires since the last reset.

tl_remove_hooks() → None¶: Detach all capture hooks and drop buffer references. Idempotent.

tl_reset_accumulators() → None¶: Clear capture chunks before each generate, else prior chunks leak into the cat.

tl_reset_capture_flags() → None¶

Open every per-hook capture gate so the next forward writes to the buffers.

First-write-wins gating means decode-step forwards self-copy and never overwrite prefill activations — the driver calls this once before any capture-needing generate (single-forward or multi-token eval) and tl_read_captures afterward.

tl_reset_counter() → None¶: Zero the shared hook-fire counter before a forward.

tl_set_batched_interventions(specs: Dict[str, Dict[str, Any]]) → None¶: Store the global spec dict the eager hook reads; {} clears.

tl_set_interventions(specs: Dict[str, Dict[str, Any]]) → None¶

Reset all affine buffers to identity, then apply each spec.

Driver pushes the full spec set every forward (or {} to reset). Spec format: {hook_name: {"op": <op>, ...op-specific params>}}. Supported ops: suppress, scale (factor: float), add and set (value: scalar broadcast across width, or 1-D shape (width,)).

transformer_lens.model_bridge.sources.vllm.worker_extension.decode_tensor(payload: Dict[str, Any]) → Tensor¶: Inverse of encode_tensor(); used driver-side.

transformer_lens.model_bridge.sources.vllm.worker_extension.dtype_name(dtype: dtype) → str¶: torch.float32 → "float32" — the one spelling for every dtype-string site.

transformer_lens.model_bridge.sources.vllm.worker_extension.encode_tensor(t: Tensor) → Dict[str, Any]¶: Encode a CPU tensor for the RPC wire. bf16 rides as fp32 bytes (exact).

transformer_lens.model_bridge.sources.vllm.worker_extension.is_pp_missing_layer(module: Any) → bool¶: vLLM PP keeps full-length module lists on every rank, filling non-owned slots (layers, embed_tokens, norm, lm_head) with PPMissingLayer identity stubs that the forward never calls — a hook installed there would serve its dead zero buffer as a real capture. Matched by name so non-vllm test doubles work; the pinned band’s class is named exactly this, and a rename fails loud via the rank-layout check.

transformer_lens.model_bridge.sources.vllm.worker_extension.resolve_dot_path(root: Any, dot_path: str) → Any¶: Walk a dot-path; None when any segment is missing or the target is a PPMissingLayer stub. Per-rank absence is legal under pipeline parallelism (each rank owns a layer subset) — the boot site verifies every hook landed on at least one rank.