Coverage for transformer_lens/model_bridge/sources/vllm/overlays/base.py: 71%

1"""Base class for vLLM overlays.""" 

2from __future__ import annotations 

3 

4from typing import Any, Dict, List, Tuple 

5 

6 

7class AdapterOverlay: 

8 """A vLLM overlay for a single architecture family. 

9 

10 :meth:`capture_specs` is called BEFORE ``LLM(...)`` to register the dot-paths 

11 and output widths the plugin should pre-allocate GPU buffers for. The plugin 

12 reads these during ``Worker.load_model`` so capture hooks are present when 

13 ``torch.compile`` traces the model. 

14 

15 :meth:`nonfiring_hooks` enumerates hooks that vLLM's fused kernels prevent 

16 from firing; surfaced as a single boot-time warning. 

17 """ 

18 

19 def capture_specs(self, hf_config: Any) -> Dict[str, Tuple[str, int]]: 

20 """Return ``{canonical_hook_name: (dot_path_in_vllm_model, output_width)}``.""" 

21 raise NotImplementedError 

22 

23 def nonfiring_hooks(self) -> List[str]: 

24 """Canonical hook names that vLLM's fused kernels cannot expose.""" 

25 return []