Coverage for transformer_lens/model_bridge/sources/vllm/overlays/decoder_only.py: 100%

1"""Generic overlay for any decoder-only model vLLM supports. 

2 

3vLLM's decoder-only models all share the same internal structure: 

4``model.embed_tokens`` / ``model.layers.{i}`` (each with ``self_attn`` and 

5``mlp`` submodules) / ``model.norm`` / ``lm_head``. This overlay hooks that 

6shared abstraction, so one file works for Llama, Qwen, Mistral, Gemma, Phi3, 

7Qwen3, Kimi, GLM, and every other model that inherits the standard shape. 

8 

9Non-decoder-only architectures (Mamba SSM, T5 encoder-decoder, BERT, MoE 

10per-expert) break the convention and would need their own overlays. 

11 

12Two hooks capture different points than HF/HookedTransformer: 

13 

14- ``blocks.{i}.hook_out``: vLLM's layer returns ``(mlp_delta, residual)`` 

15 separately (fused-residual). The plugin's hook materializes the sum so the 

16 captured value matches HF's "post-MLP residual stream". 

17- ``ln_final.hook_normalized``: vLLM exposes ``x * rsqrt(var+eps) * weight``; 

18 HF/HT exposes the pre-weight value. The driver un-folds the user-facing 

19 capture (÷ weight, or ÷ (1 + weight) for Gemma) so the cache matches 

20 ``boot_transformers``; logit reconstruction consumes the raw post-weight 

21 value internally. If the norm weight is unreachable the driver warns and 

22 serves the raw post-weight value. 

23""" 

24from __future__ import annotations 

25 

26from typing import Any, Dict, List, Tuple 

27 

28from .base import AdapterOverlay 

29 

30 

31class DecoderOnlyOverlay(AdapterOverlay): 

32 """Default overlay for vLLM decoder-only models.""" 

33 

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

35 d_model = hf_config.hidden_size 

36 n_layers = hf_config.num_hidden_layers 

37 # unembed.hook_out is intentionally NOT captured here — vLLM's sampler 

38 # computes logits via a direct matmul on the final hidden state and 

39 # never invokes lm_head.__call__, so register_forward_hook on lm_head 

40 # would install but never fire. See nonfiring_hooks() below; the 

41 # driver synthesizes the next-token logits from vLLM's sampler output. 

42 specs: Dict[str, Tuple[str, int]] = { 

43 "embed.hook_out": ("model.embed_tokens", d_model), 

44 "ln_final.hook_normalized": ("model.norm", d_model), 

45 } 

46 for i in range(n_layers): 

47 specs[f"blocks.{i}.hook_out"] = (f"model.layers.{i}", d_model) 

48 specs[f"blocks.{i}.attn.hook_out"] = (f"model.layers.{i}.self_attn", d_model) 

49 specs[f"blocks.{i}.mlp.hook_out"] = (f"model.layers.{i}.mlp", d_model) 

50 return specs 

51 

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

53 # vLLM's universal fused-kernel limitations — PagedAttention fuses the 

54 # QK^T → softmax → attn-weight path; QKVParallelLinear fuses RoPE inside 

55 # the projection. Same restriction on every decoder-only model. 

56 return [ 

57 "blocks.{i}.attn.hook_pattern", 

58 "blocks.{i}.attn.hook_attn_scores", 

59 "blocks.{i}.attn.hook_rot_q", 

60 "blocks.{i}.attn.hook_rot_k", 

61 # Head-split projections live inside the fused QKVParallelLinear / 

62 # RowParallelLinear kernels — unobservable per-head. 

63 "blocks.{i}.attn.hook_q", 

64 "blocks.{i}.attn.hook_k", 

65 "blocks.{i}.attn.hook_v", 

66 "blocks.{i}.attn.hook_z", 

67 # Block/MLP inputs: vLLM passes (hidden, residual) fused between layers, 

68 # so the HF-convention pre-block/pre-MLP stream is never materialized. 

69 "blocks.{i}.hook_in", 

70 "blocks.{i}.mlp.hook_in", 

71 "blocks.{i}.attn.hook_in", 

72 # vLLM's sampler bypasses lm_head.__call__ — capture-via-forward-hook 

73 # never fires; driver synthesizes argmax-matching logits from the 

74 # sampler's returned token. 

75 "unembed.hook_out", 

76 ]