Coverage for transformer_lens/model_bridge/sources/vllm/plugin.py: 36%
148 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""vLLM plugin entry point.
3Monkey-patches ``Worker.load_model`` to install capture hooks after weights
4load and before ``compile_or_warm_up_model`` — the only window where hooks
5make it into the compiled FX graph (PyTorch #117758).
7Two hook flavors, selected by ``configure(enable_batching=...)``:
9* Compiled (default): in-place writes to a pre-allocated GPU tensor; no
10 ``.cpu()`` (illegal during CUDA-graph capture); SymInt-indexed slicing only
11 (Python ``.shape`` access forces specialization). Single prompt.
12 Interventions ride the same hook as an affine transform
13 ``output = output * scale_buf + bias_buf`` (defaults identity). The driver
14 swaps buffer contents between forwards via ``tl_set_interventions`` — the FX
15 graph references the buffers, so swaps take effect without recompiling.
16 Memory cost: the affine allocates a transient output-shape tensor per hook
17 per forward, even at identity — peak forward memory ~1.5× capture-only.
18 Branching to skip the affine would defeat the swap trick and break the graph.
20* Batched (``enable_batching=True``, runs ``enforce_eager``): the hook reads
21 per-request token boundaries via ``segment_by_request`` (only valid inside a
22 forward, untraceable under compile — hence eager), slices each request's rows
23 to CPU, and appends to per-(req_id, hook) accumulators across chunked-prefill
24 forwards. Interventions apply directly to the tensor via ``_apply_op``.
25"""
26from __future__ import annotations
28import json
29import os
30import re
31from typing import Any, Dict, Optional, Tuple
33import torch
35from .internals import segment_by_request
36from .worker_extension import _apply_op, dtype_name, resolve_dot_path
38# Decoder layers return the fused-residual (mlp_delta, residual) 2-tuple, so
39# their hooks materialize the sum (see _make_capture_hook); other modules don't.
40_DECODER_LAYER_PATH = re.compile(r"^model\.layers\.\d+$")
42# Transient signal driver → worker during LLM construction: an env var, because
43# spawned worker processes (TP/PP > 1) re-import this module and would see any
44# module global empty — vLLM runs ``register()`` in every worker, so the patch
45# itself propagates; only this payload needs a cross-process channel. Per-Worker
46# buffers live on Worker instances so concurrent boot_vllm calls don't collide.
47# Single-node only: env vars don't reach Ray remote workers.
48_ENV_CONFIG_KEY = "TL_VLLM_PLUGIN_CONFIG"
49_install_patched = False
50_orig_load_model = None
53def configure(
54 capture_specs: Dict[str, Tuple[str, int]],
55 max_num_batched_tokens: int,
56 dtype: torch.dtype,
57 enable_batching: bool = False,
58 enable_position_interventions: bool = False,
59) -> None:
60 """Set capture specs, buffer length, dtype, and hook flavor before ``LLM(...)``."""
61 os.environ[_ENV_CONFIG_KEY] = _serialize_config(
62 {
63 "capture_specs": capture_specs,
64 "max_num_batched_tokens": max_num_batched_tokens,
65 "dtype": dtype,
66 "enable_batching": enable_batching,
67 "enable_position_interventions": enable_position_interventions,
68 }
69 )
72def clear_config() -> None:
73 """Unset the spec channel. Boot sites call this in a ``finally`` after
74 ``LLM(...)`` so a later non-TL engine can't inherit the specs."""
75 os.environ.pop(_ENV_CONFIG_KEY, None)
78def _serialize_config(config: Dict[str, Any]) -> str:
79 return json.dumps(
80 {
81 "capture_specs": {
82 name: [path, width] for name, (path, width) in config["capture_specs"].items()
83 },
84 "max_num_batched_tokens": config["max_num_batched_tokens"],
85 "dtype": dtype_name(config["dtype"]),
86 "enable_batching": config["enable_batching"],
87 "enable_position_interventions": config["enable_position_interventions"],
88 }
89 )
92def _deserialize_config(raw: str) -> Dict[str, Any]:
93 data = json.loads(raw)
94 dtype = getattr(torch, data["dtype"], None)
95 if not isinstance(dtype, torch.dtype):
96 raise ValueError(f"{_ENV_CONFIG_KEY}: {data['dtype']!r} is not a torch dtype.")
97 return {
98 "capture_specs": {
99 name: (path, int(width)) for name, (path, width) in data["capture_specs"].items()
100 },
101 "max_num_batched_tokens": int(data["max_num_batched_tokens"]),
102 "dtype": dtype,
103 "enable_batching": bool(data["enable_batching"]),
104 "enable_position_interventions": bool(data["enable_position_interventions"]),
105 }
108def _active_config() -> Optional[Dict[str, Any]]:
109 raw = os.environ.get(_ENV_CONFIG_KEY)
110 return _deserialize_config(raw) if raw else None
113def register() -> None:
114 """Idempotent monkey-patch of ``Worker.load_model``.
116 vLLM calls ``register()`` once per process at entry-points discovery. Idempotent
117 so re-imports (notebook restarts, repeated ``boot_vllm`` in the same process)
118 don't double-wrap.
120 No ``unregister()`` symmetry: the patch stays for process lifetime. Benign
121 because ``patched_load_model`` no-ops when no spec channel is populated — and
122 boot sites call ``clear_config()`` after each ``LLM(...)``, so any subsequent
123 non-TL ``LLM(...)`` in the same process hits the no-op path.
124 """
125 global _install_patched, _orig_load_model
126 if _install_patched:
127 return
128 from vllm.v1.worker.gpu_worker import Worker
130 _orig_load_model = Worker.load_model
132 def patched_load_model(self):
133 _orig_load_model(self)
134 config = _active_config()
135 if config is None:
136 return # not a TL-driven LLM; no hooks to install
137 specs = config["capture_specs"]
138 max_n = config["max_num_batched_tokens"]
139 dtype = config["dtype"]
140 enable_batching = config.get("enable_batching", False)
141 # Per-position affine buffers are (max_n, width) instead of (width,), so each
142 # sequence row can carry a distinct scale/bias (position-scoped patching).
143 per_position = config.get("enable_position_interventions", False)
144 device = next(self.model_runner.model.parameters()).device
146 # Detach prior handles before reassigning — vLLM doesn't double-load
147 # today, but unconditional reassignment would orphan hooks if it ever did.
148 for handle in getattr(self, "_tl_hook_handles", []):
149 handle.remove()
150 self._tl_buffers = {}
151 self._tl_scale_buffers = {}
152 self._tl_bias_buffers = {}
153 # Per-hook first-write-wins flag (compiled mode). 0 = open (next forward captures),
154 # 1 = closed (subsequent forwards self-copy and don't overwrite). Driver opens via
155 # tl_reset_capture_flags before any capture-needing call so a multi-token generate's
156 # prefill captures cleanly and decode steps don't overwrite row 0.
157 self._tl_capture_flags = {}
158 self._tl_hook_handles = []
159 # Batched-mode per-(req_id, hook) accumulators + global spec dict.
160 self._tl_accum = {}
161 self._tl_intervention_specs = {}
162 # Specs whose module isn't on this rank (PP layer shards); the boot site
163 # verifies via tl_absent_hooks that every spec landed somewhere.
164 self._tl_absent_hooks = set()
165 # Shared counter — surfaces hook double-fire under compile via tl_read_counter.
166 self._tl_fire_counter = torch.zeros(1, device=device, dtype=torch.int64)
167 for canonical_name, (dot_path, width) in specs.items():
168 target = resolve_dot_path(self.model_runner.model, dot_path)
169 if target is None:
170 self._tl_absent_hooks.add(canonical_name)
171 continue
172 # Decoder layers return vLLM's (mlp_delta, residual) tuple; the
173 # hook materializes their sum so the capture semantically matches
174 # HF's full residual stream. Other modules use the default path.
175 materialize = bool(_DECODER_LAYER_PATH.match(dot_path))
176 if enable_batching:
177 handle = target.register_forward_hook(
178 _make_batched_hook(
179 self,
180 canonical_name,
181 self._tl_fire_counter,
182 materialize=materialize,
183 )
184 )
185 else:
186 capture_buf = torch.zeros(max_n, width, device=device, dtype=dtype)
187 # Affine identity at install. Driver swaps via tl_set_interventions
188 # to enable suppress/scale/add/set ops between forwards. Shape is
189 # (max_n, width) when position interventions are enabled so each row
190 # can differ; (width,) otherwise (broadcast across all positions).
191 affine_shape = (max_n, width) if per_position else (width,)
192 scale_buf = torch.ones(affine_shape, device=device, dtype=dtype)
193 bias_buf = torch.zeros(affine_shape, device=device, dtype=dtype)
194 # Default closed — opened explicitly by tl_reset_capture_flags for the
195 # next forward(s) that need to capture.
196 capture_flag = torch.ones(1, device=device, dtype=torch.int64)
197 self._tl_buffers[canonical_name] = capture_buf
198 self._tl_scale_buffers[canonical_name] = scale_buf
199 self._tl_bias_buffers[canonical_name] = bias_buf
200 self._tl_capture_flags[canonical_name] = capture_flag
201 handle = target.register_forward_hook(
202 _make_capture_hook(
203 capture_buf,
204 scale_buf,
205 bias_buf,
206 self._tl_fire_counter,
207 capture_flag,
208 materialize=materialize,
209 per_position=per_position,
210 )
211 )
212 self._tl_hook_handles.append(handle)
214 Worker.load_model = patched_load_model
215 _install_patched = True
218def _make_capture_hook(
219 capture_buf: torch.Tensor,
220 scale_buf: torch.Tensor,
221 bias_buf: torch.Tensor,
222 fire_counter: torch.Tensor,
223 capture_flag: torch.Tensor,
224 *,
225 materialize: bool = False,
226 per_position: bool = False,
227):
228 """GPU-only, dynamic-shape-safe affine + first-write-wins capture into pre-allocated buffers.
230 When ``materialize=True`` (decoder layers), treat the module's output as
231 vLLM's fused-residual ``(mlp_delta, residual)`` tuple: capture
232 ``mlp_delta + residual`` (the full residual stream, matching HF's
233 blocks.{i}.hook_out semantics) and return ``(modified - residual, residual)``
234 so the next layer's input_layernorm sees the same fused sum. Mutations
235 propagate through both the capture and the downstream graph.
237 ``capture_flag`` (0 = open, 1 = closed) gates the buffer write — driver opens it
238 via ``tl_reset_capture_flags`` before each capture-needing forward, the hook closes
239 it on first fire, so a multi-token generate's prefill captures cleanly and decode
240 steps self-copy (no overwrite). Interventions still apply on every forward
241 regardless of the flag — the gate only affects the capture write.
243 ``per_position`` (compile-time constant): when set, ``scale_buf``/``bias_buf`` are
244 ``(max_n, width)`` and the affine is row-scoped (``buf.narrow(0, 0, n)``) so a
245 ``pos``-scoped intervention edits only its rows; otherwise they are ``(width,)``
246 and broadcast across every position.
248 ``fire_counter`` is incremented per call for the fire-once check.
249 """
251 # NO type annotations on _affine: the pytest jaxtyping hook wraps annotated
252 # transformer_lens functions, and dynamo tracing that wrapper into the compiled
253 # graph corrupts jaxtyping's thread-local memo stack for the whole process
254 # (GPU-verified: unrelated outer calls die with 'pop from empty list').
255 def _affine(t, n):
256 # per_position is a closure constant → torch.compile specializes this branch.
257 # narrow(0, 0, n) keeps the SymInt dynamic shape (same trick as _gated_capture);
258 # the (width,) path broadcasts across all rows.
259 if per_position:
260 return t * scale_buf.narrow(0, 0, n) + bias_buf.narrow(0, 0, n)
261 return t * scale_buf + bias_buf
263 @torch.no_grad()
264 def hook(_module, _inputs, output):
265 fire_counter.add_(1)
266 if materialize and isinstance(output, tuple) and len(output) == 2:
267 hidden, residual = output
268 if isinstance(hidden, torch.Tensor) and isinstance(residual, torch.Tensor):
269 n = hidden.shape[0]
270 modified = _affine(hidden + residual, n)
271 _gated_capture(capture_buf, n, modified, capture_flag)
272 # Reconstructs ``modified`` in the next layer's fused norm: exact at
273 # identity, bounded fp16 error under intervention.
274 return (modified - residual, residual)
276 tuple_tail: tuple = ()
277 if isinstance(output, tuple):
278 t = output[0]
279 tuple_tail = output[1:]
280 else:
281 t = output
282 if not isinstance(t, torch.Tensor):
283 return None
284 # Affine transform; default scale=1 / bias=0 means identity. Driver
285 # swaps buffer contents to enable interventions.
286 n = t.shape[0]
287 modified = _affine(t, n)
288 _gated_capture(capture_buf, n, modified, capture_flag)
289 # Gate on isinstance, not truthy tuple_tail — a 1-tuple has an empty tail.
290 if isinstance(output, tuple):
291 return (modified,) + tuple_tail
292 return modified
294 return hook
297# NO type annotations: this is traced into the compiled graph — see _affine's note.
298def _gated_capture(capture_buf, n, modified, capture_flag):
299 """First-write-wins via torch.where, compile-safe (no Python branching).
301 When ``capture_flag == 0`` (open), writes ``modified`` to ``capture_buf[:n]``.
302 When ``capture_flag == 1`` (closed), self-copies ``capture_buf[:n]`` (no-op).
303 Always closes the flag — driver explicitly opens it before each capture forward.
304 """
305 existing = capture_buf.narrow(0, 0, n)
306 to_write = torch.where(capture_flag.bool(), existing, modified)
307 capture_buf.narrow(0, 0, n).copy_(to_write)
308 capture_flag.fill_(1)
311def _make_batched_hook(
312 worker: Any,
313 canonical_name: str,
314 fire_counter: torch.Tensor,
315 *,
316 materialize: bool = False,
317):
318 """Eager hook: intervene, then append each request's rows to ``worker._tl_accum``.
320 Chunked prefill fires this once per chunk; appends are token-order for the
321 later cat. ``materialize`` mirrors the compiled hook (fused-residual sum).
322 """
324 @torch.no_grad()
325 def hook(_module, _inputs, output):
326 fire_counter.add_(1)
327 residual = None
328 if materialize and isinstance(output, tuple) and len(output) == 2: 328 ↛ 329line 328 didn't jump to line 329 because the condition on line 328 was never true
329 hidden, residual = output
330 if isinstance(hidden, torch.Tensor) and isinstance(residual, torch.Tensor):
331 t = hidden + residual
332 else:
333 return None
334 tuple_tail: tuple = ()
335 elif isinstance(output, tuple): 335 ↛ 336line 335 didn't jump to line 336 because the condition on line 335 was never true
336 t = output[0]
337 tuple_tail = output[1:]
338 else:
339 t = output
340 tuple_tail = ()
341 if not isinstance(t, torch.Tensor): 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true
342 return None
344 # Only this hook's spec — keyed by name like the compiled per-hook buffers.
345 # Iterating all specs would apply every intervention to every hook.
346 modified = t
347 spec = getattr(worker, "_tl_intervention_specs", {}).get(canonical_name)
348 if spec is not None:
349 modified = _apply_op(modified, spec)
351 qsl, req_ids = segment_by_request(worker.model_runner)
352 accum: Dict[tuple, list] = worker._tl_accum
353 if qsl is None: 353 ↛ 355line 353 didn't jump to line 355 because the condition on line 353 was never true
354 # No per-request boundaries available — treat the batch as one request.
355 req_id = req_ids[0] if req_ids else "0"
356 accum.setdefault((req_id, canonical_name), []).append(modified.detach().cpu())
357 else:
358 for i, req_id in enumerate(req_ids):
359 start, end = int(qsl[i]), int(qsl[i + 1])
360 if end <= start: 360 ↛ 361line 360 didn't jump to line 361 because the condition on line 360 was never true
361 continue
362 chunk = modified[start:end].detach().cpu()
363 accum.setdefault((req_id, canonical_name), []).append(chunk)
365 if residual is not None: 365 ↛ 366line 365 didn't jump to line 366 because the condition on line 365 was never true
366 return (modified - residual, residual)
367 if isinstance(output, tuple): 367 ↛ 368line 367 didn't jump to line 368 because the condition on line 367 was never true
368 return (modified,) + tuple_tail
369 return modified
371 return hook