Coverage for transformer_lens/model_bridge/sources/vllm/worker_extension.py: 94%
161 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Worker extension exposed to collective_rpc for capture reads and
2intervention writes.
4Hook *installation* lives in :mod:`plugin` (must happen pre-compile). This class
5only exposes the read/write surface. State is per-Worker so concurrent
6``boot_vllm`` calls don't collide. All methods prefixed ``tl_`` to avoid
7colliding with vLLM ``Worker`` attributes.
9Two capture modes, selected at boot:
10 * Compiled (default): per-hook GPU buffers + affine scale/bias swap. Single
11 prompt. ``tl_read_captures`` / ``tl_set_interventions``.
12 * Batched (``enable_batching=True``, eager): per-(req_id, hook) CPU
13 accumulators filled in the hook via query_start_loc segmentation; arbitrary
14 batch + chunked prefill. ``tl_read_batched_captures`` /
15 ``tl_set_batched_interventions`` / ``tl_reset_accumulators``.
16"""
17from __future__ import annotations
19from typing import Any, Dict, List, Optional, Set
21import torch
23from .intervention_specs import SUPPORTED_OPS, validate_spec
25# Explicit tensor wire format for collective_rpc returns. vLLM's multiproc RPC
26# cannot round-trip raw tensors from extension methods — its tensor-aware msgpack
27# encoder's header (['float32', [rows, cols], buf_idx]) leaks through undecoded on
28# the response path (GPU-verified on vllm 0.20.2). dtype/shape/bytes are all
29# msgpack-native, so this survives any executor topology; the in-process executor
30# pays one extra CPU copy.
31_TL_TENSOR_KEY = "__tl_tensor__"
34def dtype_name(dtype: torch.dtype) -> str:
35 """``torch.float32`` → ``"float32"`` — the one spelling for every dtype-string site."""
36 return str(dtype).removeprefix("torch.")
39def is_pp_missing_layer(module: Any) -> bool:
40 """vLLM PP keeps full-length module lists on every rank, filling non-owned slots
41 (layers, embed_tokens, norm, lm_head) with PPMissingLayer identity stubs that the
42 forward never calls — a hook installed there would serve its dead zero buffer as
43 a real capture. Matched by name so non-vllm test doubles work; the pinned band's
44 class is named exactly this, and a rename fails loud via the rank-layout check."""
45 return type(module).__name__ == "PPMissingLayer"
48def resolve_dot_path(root: Any, dot_path: str) -> Any:
49 """Walk a dot-path; ``None`` when any segment is missing or the target is a
50 PPMissingLayer stub. Per-rank absence is legal under pipeline parallelism (each
51 rank owns a layer subset) — the boot site verifies every hook landed on at least
52 one rank."""
53 target = root
54 for seg in dot_path.split("."):
55 if seg.isdigit():
56 try:
57 target = target[int(seg)]
58 except (IndexError, KeyError, TypeError):
59 return None
60 else:
61 target = getattr(target, seg, None)
62 if target is None:
63 return None
64 return None if is_pp_missing_layer(target) else target
67def encode_tensor(t: torch.Tensor) -> Dict[str, Any]:
68 """Encode a CPU tensor for the RPC wire. bf16 rides as fp32 bytes (exact)."""
69 t = t.detach().cpu().contiguous()
70 name = dtype_name(t.dtype)
71 if t.dtype == torch.bfloat16:
72 t = t.to(torch.float32)
73 return {
74 _TL_TENSOR_KEY: True,
75 "dtype": name,
76 "shape": list(t.shape),
77 "data": t.numpy().tobytes(),
78 }
81def decode_tensor(payload: Dict[str, Any]) -> torch.Tensor:
82 """Inverse of :func:`encode_tensor`; used driver-side."""
83 import numpy as np
85 dtype = getattr(torch, payload["dtype"])
86 wire_dtype = torch.float32 if dtype == torch.bfloat16 else dtype
87 np_dtype = torch.empty(0, dtype=wire_dtype).numpy().dtype
88 array = np.frombuffer(payload["data"], dtype=np_dtype).copy()
89 tensor = torch.from_numpy(array).reshape(payload["shape"])
90 return tensor.to(torch.bfloat16) if dtype == torch.bfloat16 else tensor
93class TLWorkerExtension:
94 """Mixed into vLLM's ``Worker`` via ``worker_extension_cls``."""
96 _tl_hook_handles: list
97 _tl_buffers: Dict[str, torch.Tensor]
98 _tl_scale_buffers: Dict[str, torch.Tensor]
99 _tl_bias_buffers: Dict[str, torch.Tensor]
100 # Per-hook first-write-wins gates (compiled mode); see plugin._gated_capture.
101 _tl_capture_flags: Dict[str, torch.Tensor]
102 _tl_fire_counter: torch.Tensor
103 # Batched-mode state (eager).
104 _tl_accum: Dict[tuple, List[torch.Tensor]]
105 _tl_intervention_specs: Dict[str, Dict[str, Any]]
106 # Specs whose module doesn't exist on this rank (PP layer shards).
107 _tl_absent_hooks: set
109 def tl_absent_hooks(self) -> Optional[List[str]]:
110 """Spec names that installed no hook on this rank — the boot site verifies
111 their union across ranks covers every spec. ``None`` means installation
112 never ran at all (plugin patch absent or spec channel empty), which the
113 coverage check must treat as fatal rather than vacuously complete."""
114 if not hasattr(self, "_tl_absent_hooks"):
115 return None
116 return sorted(self._tl_absent_hooks)
118 def tl_read_captures(
119 self, prompt_lens: List[int], names: Optional[List[str]] = None
120 ) -> Dict[str, Dict[str, Any]]:
121 """Slice each capture buffer to ``sum(prompt_lens)`` rows; wire-encoded CPU copies.
123 ``names`` restricts the read (``None`` = all) — this is the only GPU→CPU
124 crossing, so it's where a names_filtered run saves bandwidth. Caller gates
125 ``sum(prompt_lens) <= max_num_batched_tokens``, so ``total`` is in bounds.
126 """
127 total = sum(prompt_lens)
128 buffers: Dict[str, torch.Tensor] = getattr(self, "_tl_buffers", {})
129 # Ownership is dynamic, not structural: a hook can be installed on a module
130 # this rank holds but never runs (PP stages share tied-embedding aliases and
131 # some archs instantiate norm on every rank), so a still-open first-write
132 # flag means the buffer holds no data from this forward — don't serve it.
133 flags: Dict[str, torch.Tensor] = getattr(self, "_tl_capture_flags", {})
134 wanted = buffers.keys() if names is None else [n for n in names if n in buffers]
135 return {
136 name: encode_tensor(buffers[name][:total])
137 for name in wanted
138 if name not in flags or bool(flags[name].item())
139 }
141 def tl_set_interventions(self, specs: Dict[str, Dict[str, Any]]) -> None:
142 """Reset all affine buffers to identity, then apply each spec.
144 Driver pushes the full spec set every forward (or ``{}`` to reset).
145 Spec format: ``{hook_name: {"op": <op>, ...op-specific params>}}``.
146 Supported ops: suppress, scale (``factor``: float), add and set
147 (``value``: scalar broadcast across width, or 1-D shape ``(width,)``).
148 """
149 scale_bufs: Dict[str, torch.Tensor] = getattr(self, "_tl_scale_buffers", {})
150 bias_bufs: Dict[str, torch.Tensor] = getattr(self, "_tl_bias_buffers", {})
151 # Reset every hook to identity first — clears any stale state from
152 # the previous forward.
153 for name, sb in scale_bufs.items():
154 sb.fill_(1.0)
155 bias_bufs[name].zero_()
156 absent: Set[str] = getattr(self, "_tl_absent_hooks", set())
157 for hook_name, spec in specs.items():
158 if hook_name in absent:
159 # Another PP stage owns this hook; its rank applies the spec. The
160 # driver's supported_hook_points check already caught real typos.
161 continue
162 if hook_name not in scale_bufs: 162 ↛ 166line 162 didn't jump to line 166 because the condition on line 162 was always true
163 raise KeyError(f"Unknown hook for intervention: {hook_name!r}")
164 # Authoritative validation: producers that bypass VLLMDriver (the Inspect
165 # vLLM provider pushes specs over this same RPC) get identical rejection.
166 validate_spec(hook_name, spec, width=scale_bufs[hook_name].shape[-1])
167 _apply_intervention(scale_bufs[hook_name], bias_bufs[hook_name], spec)
169 def tl_get_param(self, dotted_name: str) -> Optional[Dict[str, Any]]:
170 """Read a named model tensor (e.g. ``model.norm.weight``) as a wire-encoded
171 CPU clone.
173 ``None`` if the path doesn't resolve to a tensor. The bridge has no general
174 weight surface, so this is how callers reach e.g. the ln_final weight.
175 """
176 target = resolve_dot_path(getattr(self, "model_runner").model, dotted_name)
177 return encode_tensor(target) if isinstance(target, torch.Tensor) else None
179 def tl_reset_counter(self) -> None:
180 """Zero the shared hook-fire counter before a forward."""
181 counter = getattr(self, "_tl_fire_counter", None)
182 if counter is not None:
183 counter.zero_()
185 def tl_reset_capture_flags(self) -> None:
186 """Open every per-hook capture gate so the next forward writes to the buffers.
188 First-write-wins gating means decode-step forwards self-copy and never overwrite
189 prefill activations — the driver calls this once before any capture-needing
190 generate (single-forward or multi-token eval) and ``tl_read_captures`` afterward.
191 """
192 flags: Dict[str, torch.Tensor] = getattr(self, "_tl_capture_flags", {})
193 for flag in flags.values():
194 flag.zero_()
196 def tl_read_counter(self) -> int:
197 """Total hook fires since the last reset."""
198 counter = getattr(self, "_tl_fire_counter", None)
199 return int(counter.item()) if counter is not None else 0
201 def tl_reset_accumulators(self) -> None:
202 """Clear capture chunks before each generate, else prior chunks leak into the cat."""
203 self._tl_accum = {}
205 def tl_read_batched_captures(
206 self, names: Optional[List[str]] = None
207 ) -> Dict[str, Dict[str, torch.Tensor]]:
208 """Cat per-request chunks into ``{req_id: {hook: (seq, width)}}`` (token-order).
210 ``names`` restricts to those hooks (``None`` = all). Note the per-chunk
211 GPU→CPU copy already happened in the hook, so this only saves the cat.
212 """
213 accum: Dict[tuple, List[torch.Tensor]] = getattr(self, "_tl_accum", {})
214 nameset = None if names is None else set(names)
215 out: Dict[str, Dict[str, Any]] = {}
216 for (req_id, name), chunks in accum.items():
217 if nameset is not None and name not in nameset:
218 continue
219 out.setdefault(req_id, {})[name] = encode_tensor(torch.cat(chunks, dim=0))
220 return out
222 def tl_set_batched_interventions(self, specs: Dict[str, Dict[str, Any]]) -> None:
223 """Store the global spec dict the eager hook reads; ``{}`` clears."""
224 for hook_name, spec in specs.items():
225 validate_spec(hook_name, spec)
226 self._tl_intervention_specs = dict(specs)
228 def tl_remove_hooks(self) -> None:
229 """Detach all capture hooks and drop buffer references. Idempotent."""
230 for handle in getattr(self, "_tl_hook_handles", []): 230 ↛ 231line 230 didn't jump to line 231 because the loop on line 230 never started
231 handle.remove()
232 self._tl_hook_handles = []
233 self._tl_buffers = {}
234 self._tl_scale_buffers = {}
235 self._tl_bias_buffers = {}
236 self._tl_capture_flags = {}
237 self._tl_accum = {}
238 self._tl_intervention_specs = {}
239 self._tl_absent_hooks = set()
242def _apply_op(t: torch.Tensor, spec: Dict[str, Any]) -> torch.Tensor:
243 """Apply a spec to a tensor in-line (eager path; no GPU buffer to swap)."""
244 op = spec.get("op")
245 if op not in SUPPORTED_OPS:
246 raise ValueError(f"Unsupported intervention op: {op!r}. Supported: {sorted(SUPPORTED_OPS)}")
247 if op == "suppress":
248 return torch.zeros_like(t)
249 if op == "scale":
250 return t * float(spec["factor"])
251 value = torch.as_tensor(spec["value"], device=t.device, dtype=t.dtype)
252 if value.ndim != 0 and value.shape != (t.shape[-1],):
253 raise ValueError(
254 f"Intervention 'value' must be a scalar or shape {(t.shape[-1],)}; "
255 f"got shape {tuple(value.shape)}"
256 )
257 if op == "add":
258 return t + value
259 return torch.zeros_like(t) + value # set
262def _write_rows(buf: torch.Tensor, idx: Optional[List[int]], value: Any) -> None:
263 """In-place write of ``value`` (Python scalar or ``(width,)`` tensor) into ``buf``.
265 ``idx is None`` writes the whole buffer — for a 2-D ``(max_n, width)`` buffer that
266 broadcasts a ``(width,)`` value across every row. ``idx`` (a list of row indices)
267 writes only those rows of a 2-D buffer.
268 """
269 if idx is None:
270 if isinstance(value, torch.Tensor):
271 buf[...] = value # (width,) broadcasts across rows for a 2-D buffer
272 else:
273 buf.fill_(value)
274 else:
275 buf[idx] = value # advanced-index rows; scalar or (width,) broadcasts
278def _apply_intervention(
279 scale_buf: torch.Tensor, bias_buf: torch.Tensor, spec: Dict[str, Any]
280) -> None:
281 """Translate a spec dict to in-place buffer writes.
283 ``spec['pos']`` (int or list[int]) scopes the edit to those sequence rows and
284 requires 2-D ``(max_n, width)`` affine buffers (position interventions enabled at
285 boot). Absent ``pos`` writes the whole buffer, applying to every position. The
286 caller resets both buffers to identity first, so a ``pos`` edit leaves other rows
287 untouched.
288 """
289 op = spec.get("op")
290 if op not in SUPPORTED_OPS: 290 ↛ 291line 290 didn't jump to line 291 because the condition on line 290 was never true
291 raise ValueError(f"Unsupported intervention op: {op!r}. Supported: {sorted(SUPPORTED_OPS)}")
292 pos = spec.get("pos")
293 idx: Optional[List[int]] = None
294 if pos is not None:
295 if scale_buf.ndim != 2:
296 raise ValueError(
297 "Per-position 'pos' requires 2-D affine buffers; boot with "
298 "enable_position_interventions=True."
299 )
300 idx = [pos] if isinstance(pos, int) else list(pos)
301 max_n = scale_buf.shape[0]
302 bad = [p for p in idx if p < 0 or p >= max_n]
303 if bad:
304 raise ValueError(f"Intervention 'pos' {bad} out of range [0, {max_n}).")
306 if op == "suppress":
307 _write_rows(scale_buf, idx, 0.0)
308 _write_rows(bias_buf, idx, 0.0)
309 return
310 if op == "scale":
311 _write_rows(scale_buf, idx, float(spec["factor"]))
312 _write_rows(bias_buf, idx, 0.0)
313 return
314 value = torch.as_tensor(spec["value"], device=bias_buf.device, dtype=bias_buf.dtype)
315 width = bias_buf.shape[-1]
316 # 0-d broadcasts across width (e.g. "shift all dims by 0.5"); width-shaped
317 # writes element-wise (e.g. SAE steering vector). Anything else is an error.
318 if value.ndim != 0 and value.shape != (width,):
319 raise ValueError(
320 f"Intervention 'value' must be a scalar or shape {(width,)}; "
321 f"got shape {tuple(value.shape)}"
322 )
323 _write_rows(bias_buf, idx, value.item() if value.ndim == 0 else value)
324 if op == "add":
325 _write_rows(scale_buf, idx, 1.0)
326 elif op == "set": 326 ↛ 329line 326 didn't jump to line 329 because the condition on line 326 was always true
327 _write_rows(scale_buf, idx, 0.0)
328 else:
329 raise RuntimeError(
330 f"op {op!r} is in SUPPORTED_OPS but _apply_intervention has no branch for it."
331 )