Coverage for transformer_lens/model_bridge/sources/vllm/internals.py: 60%

1"""Single chokepoint for vLLM internal API access. 

2 

3vLLM rearranges its internal class paths every 4-6 weeks. Centralize every 

4``llm.llm_engine.…`` walk here so version drift is patched in one place. 

5 

6**Validated against ``vllm==0.20.2``** (also the version pinned in 

7``demos/vLLM_Bridge_Integration_Test.ipynb``). The patched-load-model path in 

8``plugin.py`` and the ``hf_config`` walk below have been confirmed on that 

9release; newer releases may move attributes — re-validate before bumping. 

10""" 

11from __future__ import annotations 

12 

13from typing import Any 

14 

15 

16def extract_hf_config(llm: Any) -> Any: 

17 """Return the HF config that vLLM loaded the model from.""" 

18 try: 

19 return llm.llm_engine.model_config.hf_config 

20 except AttributeError as e: 

21 raise RuntimeError( 

22 "Could not locate hf_config under llm.llm_engine.model_config. " 

23 "vLLM may have moved it; update extract_hf_config() to match." 

24 ) from e 

25 

26 

27def verify_hook_coverage(llm: Any) -> None: 

28 """Raise if any configured capture hook installed on NO rank. 

29 

30 Per-rank absence is legal (pipeline-parallel ranks own layer subsets), so 

31 hook installation skips missing modules instead of raising — this boot-time 

32 check restores the fail-loud contract: a hook absent everywhere is a broken 

33 overlay dot-path and would otherwise read back as silent zeros. 

34 """ 

35 absent_per_rank = llm.collective_rpc("tl_absent_hooks") 

36 if not absent_per_rank: 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true

37 return 

38 never_ran = [rank for rank, absent in enumerate(absent_per_rank) if absent is None] 

39 if never_ran: 

40 raise RuntimeError( 

41 f"Capture-hook installation never ran on rank(s) {never_ran}: the TL vLLM " 

42 "plugin did not execute in those worker processes. The vllm.general_plugins " 

43 "entry point registers it — if this install predates that, rerun `uv sync`." 

44 ) 

45 absent_everywhere = set(absent_per_rank[0]) 

46 for rank_absent in absent_per_rank[1:]: 

47 absent_everywhere &= set(rank_absent) 

48 if absent_everywhere: 

49 raise RuntimeError( 

50 f"Capture hooks failed to install on every worker: {sorted(absent_everywhere)}. " 

51 "The overlay's dot-paths don't match this model's module tree." 

52 ) 

53 

54 

55# Cumulative per-request query offsets: FlashAttention/Triton name them 

56# query_start_loc, FlashInfer names them qo_indptr. Request i = rows i:i+1. 

57_QUERY_OFFSET_ATTRS = ("query_start_loc", "qo_indptr") 

58 

59 

60def _to_cpu_offsets(buf: Any) -> Any: 

61 """Coerce a CpuGpuBuffer (``.cpu``/``.np``/``.gpu``) or tensor to a CPU tensor.""" 

62 import torch 

63 

64 for accessor in ("cpu", "np", "gpu"): 

65 data = getattr(buf, accessor, None) 

66 if data is not None and hasattr(data, "__len__"): 

67 return torch.as_tensor(data) 

68 return torch.as_tensor(buf) if hasattr(buf, "shape") else None 

69 

70 

71def segment_by_request(model_runner: Any) -> Any: 

72 """Return ``(query_offsets_cpu, req_ids)``; request i = rows offsets[i]:offsets[i+1]. 

73 

74 Only valid inside a forward. ``req_ids`` is row order, NOT submission order — 

75 join on it. Reads ``model_runner.query_start_loc`` (backend-agnostic; the 

76 runner builds it before any attention backend, whereas FlashInfer buries its 

77 offsets in an opaque C++ wrapper), falling back to attn metadata for backends 

78 that surface them directly. ``(None, req_ids)`` ⇒ caller single-slices. 

79 """ 

80 req_ids = list(model_runner.input_batch.req_ids) 

81 n = len(req_ids) 

82 

83 qsl = getattr(model_runner, "query_start_loc", None) 

84 if qsl is not None: 84 ↛ 89line 84 didn't jump to line 89 because the condition on line 84 was always true

85 offsets = _to_cpu_offsets(qsl) 

86 if offsets is not None and len(offsets) >= n + 1: # buffer is padded to max batch 86 ↛ 89line 86 didn't jump to line 89 because the condition on line 86 was always true

87 return offsets[: n + 1].detach().cpu(), req_ids 

88 

89 from vllm.forward_context import get_forward_context 

90 

91 attn_metadata = get_forward_context().attn_metadata 

92 if isinstance(attn_metadata, list): # dual-batch-overlap returns a list of dicts 

93 attn_metadata = attn_metadata[0] 

94 if isinstance(attn_metadata, dict): 

95 for meta in attn_metadata.values(): 

96 for attr in _QUERY_OFFSET_ATTRS: 

97 off = getattr(meta, attr, None) 

98 if off is not None: 

99 return off.detach().cpu(), req_ids 

100 return None, req_ids