Coverage for transformer_lens/model_bridge/sources/vllm/driver.py: 92%

1"""vLLM Driver: forward dispatches via ``llm.generate``; captures via ``collective_rpc``.""" 

2from __future__ import annotations 

3 

4import gc 

5import logging 

6import threading 

7import warnings 

8from typing import Any, Mapping, Optional, Sequence, Union 

9 

10import torch 

11 

12from transformer_lens.model_bridge.driver_protocol import ( 

13 ForwardResult, 

14 Intervention, 

15 TensorLike, 

16) 

17from transformer_lens.model_bridge.sources._driver_base import DriverBase 

18 

19from .intervention_specs import validate_spec 

20from .worker_extension import _TL_TENSOR_KEY, decode_tensor 

21 

22# vLLM's distributed teardown operates on process-wide globals, so close() may only run 

23# it when no other VLLMDriver-owned engine is alive (notebook re-binding boots B before 

24# dropping A; A's __del__ must not destroy the process groups B is using). 

25_LIVE_DRIVERS = 0 

26_LIVE_DRIVERS_LOCK = threading.Lock() 

27 

28 

29class VLLMDriver(DriverBase): 

30 """Driver wrapping a vLLM ``LLM``; captures via ``collective_rpc``.""" 

31 

32 # vLLM owns the model in a worker — no torch module surface (parameters/state_dict/ 

33 # grads). Named-weight reads ARE served: get_param() returns CPU clones via the 

34 # tl_get_param RPC (what logit reconstruction and direct-logit-attribution use). 

35 _supported_features = frozenset({"weight_access"}) 

36 # Full-sequence logits reconstructed host-side (ln_final @ lm_head.weight.T); vLLM's 

37 # sampler only hands back the final position, so the driver rebuilds the rest. 

38 provides_sequence_logits = True 

39 

40 # Post-weight final-norm capture that lm_head consumes — reconstruction reads it. 

41 _LN_FINAL = "ln_final.hook_normalized" 

42 

43 def __init__( 

44 self, 

45 llm: Any, 

46 adapter: Any, 

47 tokenizer: Any, 

48 overlay: Any, 

49 hf_config: Any, 

50 max_num_batched_tokens: int, 

51 enable_batching: bool = False, 

52 enable_position_interventions: bool = False, 

53 tensor_parallel_size: int = 1, 

54 pipeline_parallel_size: int = 1, 

55 ) -> None: 

56 super().__init__(adapter.cfg, tokenizer) 

57 self._llm = llm 

58 self._max_num_batched_tokens = max_num_batched_tokens 

59 self._enable_batching = enable_batching 

60 # Rank layout: PP stages own disjoint layer (hence hook) subsets; the TP 

61 # ranks within a stage hold replicas of that stage's post-all-reduce hook 

62 # points. Reads merge across stages and take the first replica within one; 

63 # the first capture-bearing forward cross-checks the whole layout 

64 # (see _verify_rank_layout) and then trusts it. 

65 self._tp_size = tensor_parallel_size 

66 self._layout_verified = tensor_parallel_size == 1 and pipeline_parallel_size == 1 

67 # Position-scoped 'pos' interventions need (max_n, width) affine buffers, 

68 # allocated at boot only when this is set (see plugin.patched_load_model). 

69 self._enable_position_interventions = enable_position_interventions 

70 # Logprobs per forward = real vocab (boot's max_logprobs). d_vocab can be 

71 # padded larger, which vLLM would reject; the logits tensor stays d_vocab. 

72 self._n_logprobs = int(getattr(hf_config, "vocab_size", self.bridge_config.d_vocab)) 

73 # Unembedding cache: (weight_fp32, bias_fp32|None, scale, soft_cap) once the probe 

74 # runs — a per-forward re-fetch clones a d_vocab×d_model tensor to CPU every call. 

75 self._unembed: tuple[torch.Tensor, torch.Tensor | None, float, float | None] | None = None 

76 self._unembed_probed = False 

77 # fp32 reciprocal of the guarded final-norm weight; None after a failed probe. 

78 self._lnf_inv_denom: torch.Tensor | None = None 

79 self._lnf_probed = False 

80 

81 capture_specs = overlay.capture_specs(hf_config) 

82 self._hook_widths = {name: width for name, (_path, width) in capture_specs.items()} 

83 self._capture_paths = {name: path for name, (path, _width) in capture_specs.items()} 

84 self.supported_hook_points = frozenset(capture_specs.keys()) 

85 

86 global _LIVE_DRIVERS 

87 with _LIVE_DRIVERS_LOCK: 

88 _LIVE_DRIVERS += 1 

89 

90 n_layers = getattr(hf_config, "num_hidden_layers", 0) 

91 if not isinstance(n_layers, int) or n_layers <= 0: 

92 # A raw "{i}" template would land unexpanded in non_fireable_hook_points — 

93 # a broken config should fail at boot, not surface as a garbled hook name. 

94 raise ValueError( 

95 f"VLLMDriver: hf_config.num_hidden_layers={n_layers!r} — expected a " 

96 "positive int; the config is missing or malformed." 

97 ) 

98 nonfiring: list[str] = [] 

99 for tmpl in overlay.nonfiring_hooks(): 

100 if "{i}" in tmpl: 

101 nonfiring.extend(tmpl.replace("{i}", str(i)) for i in range(n_layers)) 

102 else: 

103 nonfiring.append(tmpl) 

104 self.non_fireable_hook_points = frozenset(nonfiring) 

105 

106 def forward( 

107 self, 

108 # Wider than the protocol's TensorLike: the batched path documents plain 

109 # (ragged) list[int] / list[list[int]] prompts, which have no shape/dtype. 

110 input_ids: Optional[Union[TensorLike, Sequence[Any]]] = None, 

111 *, 

112 capture: tuple[str, ...] = (), 

113 intervene: Mapping[str, Intervention] | None = None, 

114 max_new_tokens: int = 1, 

115 return_logits: bool = True, 

116 **kwargs: Any, 

117 ) -> ForwardResult: 

118 if input_ids is None: 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true

119 raise ValueError("VLLMDriver requires input_ids") 

120 if int(max_new_tokens) != 1: 

121 raise NotImplementedError( 

122 "VLLMDriver supports max_new_tokens=1 only — decode-step writes " 

123 "overwrite the prefill buffer; multi-step capture is multi-buffer work." 

124 ) 

125 intervene_specs = self._validate_interventions(intervene or {}) 

126 # vLLM treats pad IDs as content, so send only the mask's real-token span. 

127 attention_mask = kwargs.pop("attention_mask", None) 

128 

129 # capture is authoritative — the bridge sends exactly the hooked names, so () 

130 # means "capture nothing" and a plain forward(tokens) skips the GPU→CPU copy 

131 # entirely. (The worker's None-means-all convention is never triggered from here; 

132 # an empty tuple used to collapse to None and silently copy every buffer.) 

133 names = list(capture) 

134 

135 if self._enable_batching: 

136 return self._forward_batched( 

137 input_ids, intervene_specs, return_logits, names, attention_mask 

138 ) 

139 

140 ids_list = self._normalize_input_ids(input_ids) 

141 original_length = len(ids_list) 

142 start = 0 

143 if attention_mask is not None: 

144 mask = torch.as_tensor(attention_mask) 

145 if mask.ndim == 2 and mask.shape[0] == 1: 145 ↛ 147line 145 didn't jump to line 147 because the condition on line 145 was always true

146 mask = mask[0] 

147 ids_list, start = self._trim_masked_prompt(ids_list, mask) 

148 intervene_specs = self._remap_intervention_positions( 

149 intervene_specs, start, len(ids_list) 

150 ) 

151 if len(ids_list) > self._max_num_batched_tokens: 

152 # Worker buffers silently clamp on overflow — fail loud here instead. 

153 raise ValueError( 

154 f"Prompt length {len(ids_list)} exceeds max_num_batched_tokens=" 

155 f"{self._max_num_batched_tokens}; raise the boot_vllm kwarg or " 

156 "shorten the prompt." 

157 ) 

158 # 'pos'-scoped edits target affine-buffer rows, but the compiled hook only reads 

159 # rows [0, len(ids_list)); a pos past the prompt length would be a silent no-op. 

160 self._reject_pos_beyond_seq(intervene_specs, len(ids_list)) 

161 

162 from vllm import SamplingParams 

163 from vllm.inputs import TokensPrompt 

164 

165 # Push intervention state (possibly empty) before generate — this also 

166 # resets stale interventions from prior forwards. 

167 self._llm.collective_rpc("tl_set_interventions", args=(intervene_specs,)) 

168 # Open per-hook capture gates; first-write-wins means a fresh prefill writes and 

169 # subsequent forwards self-copy. Without this, repeated bridge.forward calls would 

170 # see the gate closed from the prior call and read stale buffers. 

171 self._llm.collective_rpc("tl_reset_capture_flags") 

172 outputs = self._llm.generate( 

173 prompts=[TokensPrompt(prompt_token_ids=ids_list)], 

174 sampling_params=SamplingParams( 

175 max_tokens=int(max_new_tokens), 

176 temperature=0.0, 

177 logprobs=self._sampler_logprobs(return_logits), 

178 ), 

179 ) 

180 

181 n_tokens = len(ids_list) 

182 read_names = self._read_names(names, return_logits) 

183 # collective_rpc returns one result per worker; each returns only the hooks 

184 # it owns (PP shards) with TP replicas within a stage — merge across ranks. 

185 # Nothing to read (no captures, logits off) → skip the crossing altogether. 

186 if read_names: 

187 per_rank_raw = self._llm.collective_rpc( 

188 "tl_read_captures", args=([n_tokens], read_names) 

189 ) 

190 if not self._layout_verified: 190 ↛ 191line 190 didn't jump to line 191 because the condition on line 190 was never true

191 self._verify_rank_layout([self._rpc_captures(caps) for caps in per_rank_raw]) 

192 self._layout_verified = True 

193 # Merge on the wire dicts, decode only the survivors — decoding every TP 

194 # replica just to discard all but the first is wasted CPU per forward. 

195 worker_captures = self._rpc_captures( 

196 self._merge_rank_captures(per_rank_raw, read_names) 

197 ) 

198 else: 

199 worker_captures = {} 

200 

201 logits: torch.Tensor | None = None 

202 if return_logits: 

203 recon = self._reconstruct_logits(worker_captures.get(self._LN_FINAL)) 

204 # Fall back to final-position log-probs if the unembedding isn't fetchable. 

205 logits = ( 

206 recon.unsqueeze(0) 

207 if recon is not None 

208 else self._synthesize_logits(outputs[0], n_tokens, self.bridge_config.d_vocab) 

209 ) 

210 

211 captured = self._expose_captured( 

212 {name: t.unsqueeze(0) for name, t in worker_captures.items()}, names 

213 ) 

214 if attention_mask is not None: 

215 if logits is not None: 

216 logits = self._restore_masked_positions( 

217 logits, [start], [original_length], [n_tokens], float("-inf") 

218 ) 

219 captured = { 

220 name: self._restore_masked_positions(t, [start], [original_length], [n_tokens], 0.0) 

221 for name, t in captured.items() 

222 } 

223 return ForwardResult(logits=logits, captured=captured, raw_output=outputs[0]) 

224 

225 def _sampler_logprobs(self, return_logits: bool) -> int | None: 

226 # Sampler logprobs are the reconstruction fallback only; skipping them avoids 

227 # marshaling a d_vocab-entry Logprob dict host-side per request. 

228 return self._n_logprobs if return_logits and self._unembed is None else None 

229 

230 def _read_names(self, names: list[str], return_logits: bool) -> list[str]: 

231 # Reconstruction needs ln_final even when uncaptured; skip once probed-unavailable. 

232 recon_possible = self._unembed is not None or not self._unembed_probed 

233 if return_logits and recon_possible and self._LN_FINAL not in names: 

234 return names + [self._LN_FINAL] 

235 return names 

236 

237 def _expose_captured( 

238 self, captured: Mapping[str, torch.Tensor], names: list[str] 

239 ) -> dict[str, torch.Tensor]: 

240 """Filter to the caller's requested hooks (dropping any forced ln_final) and 

241 convert the exposed ln_final to the pre-weight convention its name promises — 

242 reconstruction consumed the raw post-weight value already.""" 

243 out = {name: t for name, t in captured.items() if name in names} 

244 if self._LN_FINAL in out: 244 ↛ 245line 244 didn't jump to line 245 because the condition on line 244 was never true

245 out[self._LN_FINAL] = self._unfold_ln_final(out[self._LN_FINAL]) 

246 return out 

247 

248 def _forward_batched( 

249 self, 

250 input_ids: Union[TensorLike, Sequence[Any]], 

251 intervene_specs: dict, 

252 return_logits: bool, 

253 names: list[str], 

254 attention_mask: Any = None, 

255 ) -> ForwardResult: 

256 """Eager batched path: per-request capture, right-padded to (B, S, W). 

257 

258 No per-prompt length gate — chunked prefill accumulates long prompts 

259 across forwards. Interventions are global across the batch. ``names`` is 

260 authoritative: exactly the hooks to return (empty = none). 

261 """ 

262 prompts_ids = self._normalize_input_ids_batched(input_ids) 

263 original_lengths = [len(ids) for ids in prompts_ids] 

264 starts = [0] * len(prompts_ids) 

265 if attention_mask is not None: 

266 mask = torch.as_tensor(attention_mask) 

267 if mask.dim() == 1: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true

268 mask = mask.unsqueeze(0) 

269 if mask.dim() != 2: 269 ↛ 270line 269 didn't jump to line 270 because the condition on line 269 was never true

270 raise ValueError("attention_mask must have shape [batch, seq].") 

271 if mask.shape[0] != len(prompts_ids): 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true

272 raise ValueError( 

273 f"attention_mask batch dim {mask.shape[0]} != number of prompts " 

274 f"{len(prompts_ids)}." 

275 ) 

276 trimmed = [self._trim_masked_prompt(ids, row) for ids, row in zip(prompts_ids, mask)] 

277 prompts_ids = [ids for ids, _start in trimmed] 

278 starts = [start for _ids, start in trimmed] 

279 prompt_lens = [len(ids) for ids in prompts_ids] 

280 

281 from vllm import SamplingParams 

282 from vllm.inputs import TokensPrompt 

283 

284 # Reset accumulators so prior-forward chunks don't leak into the cat. 

285 self._llm.collective_rpc("tl_reset_accumulators") 

286 self._llm.collective_rpc("tl_set_batched_interventions", args=(intervene_specs,)) 

287 

288 d_vocab = self.bridge_config.d_vocab 

289 outputs = self._llm.generate( 

290 prompts=[TokensPrompt(prompt_token_ids=ids) for ids in prompts_ids], 

291 sampling_params=SamplingParams( 

292 max_tokens=1, 

293 temperature=0.0, 

294 logprobs=self._sampler_logprobs(return_logits), 

295 ), 

296 ) 

297 

298 read_names = self._read_names(names, return_logits) 

299 # Keyed by req_id (no guaranteed order) — _assemble_padded joins to slot 

300 # k via outputs[k].request_id, not by position. Empty read → skip both the 

301 # crossing AND the join (_assemble_padded requires one worker key per request, 

302 # so it can't run on an empty dict — mirror the single path's empty captured). 

303 if read_names: 

304 worker_captures = self._llm.collective_rpc( 

305 "tl_read_batched_captures", args=(read_names,) 

306 )[0] 

307 # Batched runs single-rank today, but coerce anyway — multiproc RPC 

308 # serializes tensors to lists (see _rpc_tensor). 

309 worker_captures = { 

310 req_id: self._rpc_captures(caps) for req_id, caps in worker_captures.items() 

311 } 

312 captured = self._assemble_padded(outputs, worker_captures, prompt_lens) 

313 else: 

314 captured = {} 

315 

316 logits: torch.Tensor | None = None 

317 if return_logits: 

318 recon = self._reconstruct_logits( 

319 captured.get(self._LN_FINAL) 

320 ) # (batch, max_seq, d_vocab) 

321 if recon is not None: 

322 # Pad rows reconstruct from zero-filled ln_final into finite garbage 

323 # (0 @ W = plausible uniform logits) — mask them to the -inf convention 

324 # the fallback and per-row consumers rely on. 

325 for k, n in enumerate(prompt_lens): 

326 recon[k, n:] = float("-inf") 

327 logits = ( 

328 recon 

329 if recon is not None 

330 else self._synthesize_logits_batched(outputs, prompt_lens, d_vocab) 

331 ) 

332 captured = self._expose_captured(captured, names) 

333 if attention_mask is not None: 

334 if logits is not None: 

335 logits = self._restore_masked_positions( 

336 logits, starts, original_lengths, prompt_lens, float("-inf") 

337 ) 

338 captured = { 

339 name: self._restore_masked_positions(t, starts, original_lengths, prompt_lens, 0.0) 

340 for name, t in captured.items() 

341 } 

342 

343 return ForwardResult(logits=logits, captured=captured, raw_output=outputs) 

344 

345 @staticmethod 

346 def _assemble_padded( 

347 outputs: list, 

348 worker_captures: Mapping[str, Mapping[str, torch.Tensor]], 

349 prompt_lens: list[int], 

350 ) -> dict[str, torch.Tensor]: 

351 """Stack per-request captures into right-padded ``(batch, max_seq, width)``. 

352 

353 Pad is a cache-assembly artifact only: vLLM computes each request 

354 independently, so real-token activations don't depend on the padding. 

355 """ 

356 batch = len(outputs) 

357 max_seq = max(prompt_lens) if prompt_lens else 0 

358 # Worker keys are engine-internal req_ids ("10-83c3532c"); RequestOutput 

359 # carries only the public id ("10"). Join exact-or-prefix; the "-" keeps 

360 # "1" from matching "10-...". 

361 worker_keys = list(worker_captures.keys()) 

362 

363 def _captures_for(public_rid: str) -> Mapping[str, torch.Tensor]: 

364 if public_rid in worker_captures: 

365 return worker_captures[public_rid] 

366 matches = [k for k in worker_keys if k.startswith(f"{public_rid}-")] 

367 # Raise, never silently zero-fill the row: a missing or ambiguous join 

368 # is indistinguishable from a genuine zero activation, which is silent 

369 # data loss on the collection path. 

370 if len(matches) != 1: 370 ↛ 376line 370 didn't jump to line 376 because the condition on line 370 was always true

371 raise RuntimeError( 

372 f"Cannot join request {public_rid!r} to worker captures: found " 

373 f"{len(matches)} key(s) in {worker_keys}. Expected exactly one " 

374 f"(exact or '{public_rid}-<hash>')." 

375 ) 

376 return worker_captures[matches[0]] 

377 

378 per_slot = [_captures_for(o.request_id) for o in outputs] 

379 

380 hook_names: set[str] = set() 

381 for caps in per_slot: 

382 hook_names |= set(caps.keys()) 

383 

384 assembled: dict[str, torch.Tensor] = {} 

385 for name in hook_names: 

386 sample = next(caps[name] for caps in per_slot if name in caps) 

387 buf = torch.zeros(batch, max_seq, sample.shape[-1], dtype=sample.dtype) 

388 for k, caps in enumerate(per_slot): 

389 t = caps.get(name) 

390 if t is not None: 390 ↛ 388line 390 didn't jump to line 388 because the condition on line 390 was always true

391 buf[k, : t.shape[0]] = t 

392 assembled[name] = buf 

393 return assembled 

394 

395 @staticmethod 

396 def _synthesize_logits_batched( 

397 outputs: list, prompt_lens: list[int], d_vocab: int 

398 ) -> torch.Tensor: 

399 """Build ``(batch, max_seq, d_vocab)`` logits; next-token dist at each row's 

400 ``prompt_lens[k] - 1``, never ``-1`` (a pad position for shorter prompts).""" 

401 batch = len(outputs) 

402 max_seq = max(prompt_lens) if prompt_lens else 0 

403 logits = torch.full((batch, max_seq, d_vocab), float("-inf"), dtype=torch.float16) 

404 for k, request_output in enumerate(outputs): 

405 gen = request_output.outputs[0] if request_output.outputs else None 

406 if gen is None: 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true

407 continue 

408 pos = prompt_lens[k] - 1 

409 if gen.logprobs: 409 ↛ 412line 409 didn't jump to line 412 because the condition on line 409 was always true

410 for token_id, lp_obj in gen.logprobs[0].items(): 

411 logits[k, pos, int(token_id)] = float(lp_obj.logprob) 

412 elif gen.token_ids: 

413 logits[k, pos, int(gen.token_ids[0])] = 0.0 

414 return logits 

415 

416 def get_param(self, dotted_name: str) -> torch.Tensor | None: 

417 """Fetch a named model tensor (e.g. ``model.norm.weight``) for conversions 

418 the bridge can't otherwise do (ln_final post→pre-weight; see the overlay). 

419 Gathered across ranks: replicated params return one copy, vocab-sharded 

420 weights concatenate, stage-local params (PP) come from whichever rank owns 

421 them. None if closed or the path resolves nowhere.""" 

422 return self._gather_param(dotted_name) 

423 

424 @staticmethod 

425 def _rpc_tensor(value: Any) -> torch.Tensor: 

426 """Decode a non-None collective_rpc payload back to a tensor. 

427 

428 Worker methods return the explicit wire format (see 

429 ``worker_extension.encode_tensor``) because vLLM's multiproc RPC can't 

430 round-trip raw tensors. Raw tensors still pass through for mocks and any 

431 legacy in-process payloads; anything else is best-effort coerced.""" 

432 if isinstance(value, torch.Tensor): 

433 return value 

434 if isinstance(value, Mapping) and value.get(_TL_TENSOR_KEY): 

435 return decode_tensor(dict(value)) 

436 return torch.as_tensor(value) 

437 

438 @classmethod 

439 def _rpc_captures(cls, captures: Mapping[str, Any]) -> dict[str, torch.Tensor]: 

440 return {name: cls._rpc_tensor(value) for name, value in captures.items()} 

441 

442 def _gather_param(self, dotted_name: str, dim: int = 0) -> torch.Tensor | None: 

443 """Fetch a param across all TP ranks: replicated → rank 0; sharded → concat. 

444 

445 Vocab-parallel weights (``lm_head.weight``, ``embed_tokens.weight``) hold 

446 contiguous per-rank slices in rank order, which is also collective_rpc's 

447 result order — concatenation along ``dim`` reassembles the full tensor 

448 (vocab padding lands in the tail rows, sliced off by reconstruction). 

449 """ 

450 if self._llm is None: 

451 return None 

452 shards = [ 

453 self._rpc_tensor(s) 

454 for s in self._llm.collective_rpc("tl_get_param", args=(dotted_name,)) 

455 if s is not None 

456 ] 

457 if not shards: 

458 return None 

459 if len(shards) == 1: 

460 return shards[0] 

461 if all(s.shape == shards[0].shape and torch.equal(s, shards[0]) for s in shards[1:]): 

462 return shards[0] # replicated (norm weights, biases on some archs) 

463 return torch.cat(shards, dim=dim) 

464 

465 def _merge_rank_captures(self, per_rank_captures: list, requested: list[str]) -> dict[str, Any]: 

466 """Merge per-rank capture dicts into one (values stay opaque — callers decode): 

467 PP stages own disjoint hook subsets; TP replicas within a stage agree (verified 

468 once by _verify_rank_layout), so the first copy wins. A requested, supported 

469 hook served by NO rank fails loud — zero-filling would be silent data loss 

470 (same policy as the batched join).""" 

471 merged: dict[str, Any] = {} 

472 for captures in per_rank_captures: 

473 for name, tensor in captures.items(): 

474 if name not in merged: 

475 merged[name] = tensor 

476 missing = [ 

477 name for name in requested if name not in merged and name in self.supported_hook_points 

478 ] 

479 if missing: 

480 raise RuntimeError( 

481 f"Capture(s) {missing} returned by no rank — a hook fired nowhere " 

482 "despite passing the boot-time coverage check. Report this." 

483 ) 

484 return merged 

485 

486 def _verify_rank_layout(self, per_rank_captures: list) -> None: 

487 """One-time cross-rank check of the capture layout under TP/PP. 

488 

489 Invariant: each hook lives on exactly one PP stage and is replicated across 

490 that stage's ``tp_size`` ranks (every overlay hook point is post-all-reduce). 

491 A wrong copy count means installation drift; divergent replicas mean a hook 

492 point moved pre-all-reduce — either way rank-merged reads would be silently 

493 wrong, so fail loud on the first capture-bearing forward. Ranks serving the 

494 same hook set are TP replicas of one stage and must report identical fire 

495 counters; counters are NOT comparable across stages — per-rank installed 

496 counts overcount modules that exist but never run (tied-embedding aliases, 

497 archs that instantiate norm on every rank). 

498 """ 

499 all_names: set[str] = set() 

500 for captures in per_rank_captures: 

501 all_names |= set(captures.keys()) 

502 for name in sorted(all_names): 

503 copies = [ 

504 (rank, caps[name]) for rank, caps in enumerate(per_rank_captures) if name in caps 

505 ] 

506 if len(copies) != self._tp_size: 

507 raise RuntimeError( 

508 f"Rank-layout check failed for {name!r}: found on {len(copies)} rank(s) " 

509 f"but expected exactly tp_size={self._tp_size} (one PP stage × its TP " 

510 "replicas). Per-rank installation drifted — report this." 

511 ) 

512 rank0, t0 = copies[0] 

513 for rank, t_r in copies[1:]: 

514 if t0.shape != t_r.shape or not torch.allclose( 

515 t0.float(), t_r.float(), atol=1e-5, rtol=1e-5 

516 ): 

517 diff = ( 

518 (t0.float() - t_r.float()).abs().max().item() 

519 if t0.shape == t_r.shape 

520 else float("inf") 

521 ) 

522 raise RuntimeError( 

523 f"TP replication check failed for {name!r}: rank {rank0} vs rank " 

524 f"{rank} max abs diff {diff:.3e}. This hook point is no longer " 

525 "replicated across ranks (vLLM may have moved it pre-all-reduce) — " 

526 "merged capture reads would be silently wrong. Report this." 

527 ) 

528 counters = [int(c) for c in self._llm.collective_rpc("tl_read_counter")] 

529 stages: dict = {} 

530 for rank, captures in enumerate(per_rank_captures): 

531 stages.setdefault(frozenset(captures.keys()), []).append(rank) 

532 for ranks in stages.values(): 

533 group = [counters[r] for r in ranks] 

534 if len(set(group)) > 1: 

535 raise RuntimeError( 

536 f"Fire-counter mismatch within a TP replica group: ranks {ranks} " 

537 f"serve the same hooks but fired {group} times. Replicated forwards " 

538 "diverged — merged capture reads would be silently wrong. Report this." 

539 ) 

540 

541 def probe_logit_reconstruction(self) -> bool: 

542 """One-time unembedding fetch + cache; downgrades ``provides_sequence_logits`` 

543 honestly when no unembedding or logits processor is reachable (the fallback 

544 path is final-position log-probs, which cannot back a loss). Idempotent — later 

545 calls return the cached availability.""" 

546 if self._unembed_probed: 

547 return self._unembed is not None 

548 # Gathered reads: under TP the unembedding is vocab-sharded per rank. 

549 weight = self._gather_param("lm_head.weight") 

550 if weight is None: # tied embeddings expose no separate lm_head 

551 weight = self._gather_param("model.embed_tokens.weight") 

552 if weight is None: 

553 self._unembed_probed = True 

554 self.provides_sequence_logits = False 

555 return False 

556 # Config fields alone miss architecture composition and runtime overrides. 

557 processors = [ 

558 p for p in self._llm.collective_rpc("tl_get_logits_processor") if p is not None 

559 ] 

560 if not processors: 

561 self._unembed_probed = True 

562 self.provides_sequence_logits = False 

563 warnings.warn( 

564 "vLLM logits processor is unavailable; disabling full-sequence logit reconstruction.", 

565 UserWarning, 

566 stacklevel=2, 

567 ) 

568 return False 

569 if any(p != processors[0] for p in processors[1:]): 

570 raise RuntimeError("vLLM logits processor metadata disagrees across ranks.") 

571 scale, cap = processors[0]["scale"], processors[0]["soft_cap"] 

572 bias = self._gather_param("lm_head.bias") 

573 d_vocab = int(self.bridge_config.d_vocab) 

574 # Slice vLLM's vocab-pad rows at cache time; fp32 residency (~2× checkpoint 

575 # dtype on CPU) trades memory for skipping a full-matrix upcast per forward. 

576 self._unembed = ( 

577 weight.to(torch.float32)[:d_vocab], 

578 bias.to(torch.float32)[:d_vocab] if bias is not None else None, 

579 scale, 

580 cap, 

581 ) 

582 self.provides_sequence_logits = True 

583 self._unembed_probed = True 

584 return True 

585 

586 def _reconstruct_logits(self, ln_final: Any) -> torch.Tensor | None: 

587 """Rebuild real logits from the captured post-weight ln_final: 

588 ``ln_final @ lm_head.weight.T`` (+ bias), then the engine's softcap and scale. 

589 

590 vLLM's ``ln_final.hook_normalized`` is the POST-weight RMSNorm value lm_head 

591 consumes (verified empirically: it equals HF's pre-weight value times the norm 

592 weight), so no un-fold is needed here — the raw worker value feeds this directly. 

593 Accepts any ``(..., d_model)`` tensor and returns ``(..., d_vocab)`` on CPU. 

594 ``None`` if ln_final wasn't captured or no unembedding weight is fetchable — 

595 the caller then falls back to the sampler's log-probs. 

596 """ 

597 if ln_final is None or not self.probe_logit_reconstruction(): 

598 return None 

599 assert self._unembed is not None # probe returned True 

600 weight32, bias32, scale, cap = self._unembed 

601 lf = ln_final.to(device=weight32.device, dtype=torch.float32) 

602 logits = lf @ weight32.T 

603 if bias32 is not None: 

604 logits = logits + bias32.to(device=logits.device) 

605 d_vocab = int(self.bridge_config.d_vocab) 

606 if logits.shape[-1] > d_vocab: 

607 # vLLM pads vocab embeddings to a multiple of 64; its own sampler slices to 

608 # org_vocab_size before sampling — mirror that, or pad columns (zero-filled 

609 # at load) become phantom argmax candidates and bias softmax denominators. 

610 logits = logits[..., :d_vocab] 

611 if cap is not None: 

612 logits = cap * torch.tanh(logits / cap) 

613 if scale != 1.0: 

614 logits = logits * scale 

615 if logits.shape[-1] < d_vocab: # pad the padded-vocab tail (never predicted) 615 ↛ 616line 615 didn't jump to line 616 because the condition on line 615 was never true

616 pad = logits.new_full((*logits.shape[:-1], d_vocab - logits.shape[-1]), float("-inf")) 

617 logits = torch.cat([logits, pad], dim=-1) 

618 return logits.cpu() 

619 

620 def _unfold_ln_final(self, t: torch.Tensor) -> torch.Tensor: 

621 """Convert vLLM's post-weight RMSNorm capture to the pre-weight value the 

622 canonical hook name promises (÷ weight; Gemma folds ``1 + weight``). Warns 

623 once and serves the raw value when the norm weight is unreachable — loud 

624 beats silent cross-backend mismatch.""" 

625 if not self._lnf_probed: 

626 self._lnf_probed = True 

627 # The overlay's capture spec owns the module path; derive the weight from 

628 # it. Gathered read: under PP the final norm lives only on the last stage, 

629 # so a rank-0 get_param would miss it. 

630 path = self._capture_paths.get(self._LN_FINAL, "model.norm") 

631 weight = self._gather_param(f"{path}.weight") 

632 if weight is None: 

633 warnings.warn( 

634 "ln_final.hook_normalized: norm weight unreachable — the captured " 

635 "value stays POST-weight and will not match boot_transformers.", 

636 UserWarning, 

637 stacklevel=3, 

638 ) 

639 else: 

640 w = weight.detach().to(torch.float32) 

641 denom = (1.0 + w) if "gemma" in self.architecture.lower() else w 

642 # Near-zero weight entries would blow up the division; identity beats inf. 

643 denom = torch.where(denom.abs() < 1e-6, torch.ones_like(denom), denom) 

644 self._lnf_inv_denom = denom.reciprocal() 

645 if self._lnf_inv_denom is None: 

646 return t 

647 inv = self._lnf_inv_denom.to(device=t.device) 

648 return (t.to(torch.float32) * inv).to(t.dtype) 

649 

650 def close(self) -> None: 

651 if self._llm is None: # already closed — keep the refcount single-shot 

652 return 

653 # Detach hooks before dropping the LLM so they don't stay registered on 

654 # worker modules for the life of the process (long-running notebooks). 

655 log = logging.getLogger("transformer_lens.vllm") 

656 try: 

657 self._llm.collective_rpc("tl_remove_hooks") 

658 except Exception as e: 

659 # Best-effort: engine may already be torn down or the RPC surface 

660 # gone. Log so hook-leak debugging has a thread to pull. 

661 log.debug("tl_remove_hooks failed during close(): %s", e) 

662 self._llm = None 

663 self._unembed = None 

664 self._lnf_inv_denom = None 

665 

666 global _LIVE_DRIVERS 

667 with _LIVE_DRIVERS_LOCK: 

668 _LIVE_DRIVERS -= 1 

669 last_driver = _LIVE_DRIVERS <= 0 

670 # vLLM 0.20.2 has no LLM.shutdown(); the distributed teardown below hits 

671 # process-global state, so it runs only for the last live driver (see 

672 # _LIVE_DRIVERS). Both calls are best-effort no-ops otherwise. 

673 if last_driver: 

674 try: 

675 from vllm.distributed.parallel_state import ( 

676 destroy_distributed_environment, 

677 destroy_model_parallel, 

678 ) 

679 

680 destroy_model_parallel() 

681 destroy_distributed_environment() 

682 except Exception as e: 

683 log.debug("vLLM distributed teardown failed during close(): %s", e) 

684 # Free the caching allocator's blocks here so the caller doesn't have to. 

685 gc.collect() 

686 try: 

687 if torch.cuda.is_available(): 687 ↛ 688line 687 didn't jump to line 688 because the condition on line 687 was never true

688 torch.cuda.empty_cache() 

689 except Exception as e: 

690 log.debug("torch.cuda.empty_cache failed during close(): %s", e) 

691 

692 @staticmethod 

693 def _synthesize_logits(request_output: Any, n_tokens: int, d_vocab: int) -> torch.Tensor: 

694 """Build a (1, n_tokens, d_vocab) logits-like tensor from vLLM's sampler output. 

695 

696 Values are **log-probs**, not raw logits (vLLM returns log_softmax): fine 

697 for argmax/next-token, wrong for absolute scale (temperature, logit-lens). 

698 lm_head is bypassed so only position -1 is filled (next-token); earlier 

699 positions stay -inf — populating them needs prompt_logprobs per call. 

700 """ 

701 logits = torch.full((1, n_tokens, d_vocab), float("-inf"), dtype=torch.float16) 

702 gen = request_output.outputs[0] if request_output.outputs else None 

703 if gen is None: 703 ↛ 704line 703 didn't jump to line 704 because the condition on line 703 was never true

704 return logits 

705 # Prefer real logprobs; fall back to the generated token id (one-hot-ish) 

706 # if logprobs weren't requested (e.g. return_logits=False elsewhere). 

707 if gen.logprobs: 707 ↛ 710line 707 didn't jump to line 710 because the condition on line 707 was always true

708 for token_id, lp_obj in gen.logprobs[0].items(): 

709 logits[0, -1, int(token_id)] = float(lp_obj.logprob) 

710 elif gen.token_ids: 

711 logits[0, -1, int(gen.token_ids[0])] = 0.0 

712 return logits 

713 

714 def _validate_interventions(self, intervene: Mapping[str, Any]) -> dict: 

715 """Shared spec validation plus driver-side gating (hook membership, pos support).""" 

716 out: dict = {} 

717 for hook_name, spec in intervene.items(): 

718 validated = validate_spec(hook_name, spec, width=self._hook_widths.get(hook_name)) 

719 if hook_name not in self.supported_hook_points: 

720 raise ValueError( 

721 f"Cannot intervene on {hook_name!r}: not in supported_hook_points." 

722 ) 

723 if validated.get("pos") is not None: 

724 if self._enable_batching: 

725 # The batched/eager path applies ops to the raw tensor, not the 

726 # (max_n, width) affine buffers, so it has no position surface. 

727 raise NotImplementedError( 

728 f"Intervention {hook_name!r}: per-position 'pos' is not supported on the " 

729 "batched/eager path. Boot the compiled path (enable_batching=False) with " 

730 "enable_position_interventions=True." 

731 ) 

732 if not self._enable_position_interventions: 

733 # Default affine buffers are (width,) and broadcast across every position; 

734 # honoring 'pos' needs the (max_n, width) buffers allocated at boot. 

735 raise NotImplementedError( 

736 f"Intervention {hook_name!r}: per-position 'pos' requires " 

737 "boot_vllm(enable_position_interventions=True) (its default affine " 

738 "buffers broadcast across all positions). Use the Inspect/HF backend for " 

739 "position-scoped patching, or drop 'pos' for a whole-sequence edit." 

740 ) 

741 out[hook_name] = validated 

742 return out 

743 

744 @staticmethod 

745 def _reject_pos_beyond_seq(specs: Mapping[str, Any], seq_len: int) -> None: 

746 """Fail loud if a spec's 'pos' targets a row past the actual prompt length. 

747 

748 The compiled hook applies the affine over rows ``[0, seq_len)`` only, so a ``pos`` 

749 in ``[seq_len, max_num_batched_tokens)`` clears the driver's non-negativity check 

750 and the worker's buffer-capacity check yet is never read — a silent no-op. Bound it 

751 against the real sequence length (known once ``ids_list`` exists) and raise instead. 

752 """ 

753 for hook_name, spec in specs.items(): 

754 pos = spec.get("pos") 

755 if pos is None: 755 ↛ 756line 755 didn't jump to line 756 because the condition on line 755 was never true

756 continue