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