transformer_lens.model_bridge.sources.vllm.overlays package

Submodules

• transformer_lens.model_bridge.sources.vllm.overlays.base module
  • AdapterOverlay
    • AdapterOverlay.capture_specs()
    • AdapterOverlay.nonfiring_hooks()
• transformer_lens.model_bridge.sources.vllm.overlays.decoder_only module
  • DecoderOnlyOverlay
    • DecoderOnlyOverlay.capture_specs()
    • DecoderOnlyOverlay.nonfiring_hooks()

Module contents

vLLM overlay registry.

One DecoderOnlyOverlay handles every vLLM decoder-only model (Llama / Qwen / Mistral / Gemma / Phi3 / Qwen3 / Kimi / GLM / …). It’s the default for any architecture; per-architecture overlays would only land if a model breaks vLLM’s conventional decoder-only shape.

class transformer_lens.model_bridge.sources.vllm.overlays.AdapterOverlay

Bases: object

A vLLM overlay for a single architecture family.

capture_specs() is called BEFORE LLM(...) to register the dot-paths and output widths the plugin should pre-allocate GPU buffers for. The plugin reads these during Worker.load_model so capture hooks are present when torch.compile traces the model.

nonfiring_hooks() enumerates hooks that vLLM’s fused kernels prevent from firing; surfaced as a single boot-time warning.

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.

class transformer_lens.model_bridge.sources.vllm.overlays.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.

transformer_lens.model_bridge.sources.vllm.overlays.get_overlay(architecture: str) → AdapterOverlay

Return the overlay for an architecture; falls back to the decoder-only default.