transformer_lens.model_bridge.sources.vllm.overlays.decoder_only module

Generic overlay for any decoder-only model vLLM supports.

vLLM’s decoder-only models all share the same internal structure: model.embed_tokens / model.layers.{i} (each with self_attn and mlp submodules) / model.norm / lm_head. This overlay hooks that shared abstraction, so one file works for Llama, Qwen, Mistral, Gemma, Phi3, Qwen3, Kimi, GLM, and every other model that inherits the standard shape.

Non-decoder-only architectures (Mamba SSM, T5 encoder-decoder, BERT, MoE per-expert) break the convention and would need their own overlays.

Two hooks capture different points than HF/HookedTransformer:

  • blocks.{i}.hook_out: vLLM’s layer returns (mlp_delta, residual) separately (fused-residual). The plugin’s hook materializes the sum so the captured value matches HF’s “post-MLP residual stream”.

  • ln_final.hook_normalized: vLLM exposes x * rsqrt(var+eps) * weight; HF/HT exposes the pre-weight value. The driver un-folds the user-facing capture (÷ weight, or ÷ (1 + weight) for Gemma) so the cache matches boot_transformers; logit reconstruction consumes the raw post-weight value internally. If the norm weight is unreachable the driver warns and serves the raw post-weight value.

class transformer_lens.model_bridge.sources.vllm.overlays.decoder_only.DecoderOnlyOverlay

Bases: AdapterOverlay

Default overlay for vLLM decoder-only models.

capture_specs(hf_config: Any) → Dict[str, Tuple[str, int]]

Return {canonical_hook_name: (dot_path_in_vllm_model, output_width)}.

nonfiring_hooks() → List[str]

Canonical hook names that vLLM’s fused kernels cannot expose.