"model_executor/models/llama.py: the Llama model"

Code version: vLLM @ 3e44c0d · 2026-10-22

Purpose#

vllm/model_executor/models/llama.py implements the Llama architecture with vLLM layers, so that the worker can build it, load a Hugging Face checkpoint into it, and run it in the forward pass of each step.

The page How a vLLM model file is written (llama.py) explains the patterns of this file. This page reads the file class by class and follows the load_weights path.

Where it sits#

Who calls it#

• The registry. Nine names map to this file, seven of them to LlamaForCausalLM: LlamaForCausalLM, LLaMAForCausalLM, InternLM3ForCausalLM, IQuestCoderForCausalLM, TeleChat3ForCausalLM, and the embedding names LlamaModel and MistralModel (vllm/model_executor/models/registry.py:140-155, vllm/model_executor/models/registry.py:212, vllm/model_executor/models/registry.py:231-238). Two more map to the bidirectional wrappers (vllm/model_executor/models/registry.py:230, vllm/model_executor/models/registry.py:326-329). (Counts from an evaluation of the tables.)
• The loader. initialize_model() builds the class; the default loader calls load_weights() (vllm/model_executor/model_loader/default_loader.py:503).
• The model runner. Each step calls the model with input_ids, positions, inputs_embeds and intermediate_tensors (vllm/v1/worker/gpu/model_runner.py:1984-2009), then compute_logits() on the sampled rows.
• Other model files. For example MistralForCausalLM, Phi3ForCausalLM and TeleChat2ForCausalLM subclass LlamaForCausalLM (vllm/model_executor/models/mistral.py:236, vllm/model_executor/models/phi3.py:10, vllm/model_executor/models/telechat2.py:107).

What it calls#

• QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear (vllm/model_executor/models/llama.py:45-49), VocabParallelEmbedding and ParallelLMHead (vllm/model_executor/models/llama.py:53-56).
• Attention or EncoderOnlyAttention (vllm/model_executor/models/llama.py:39-42).
• CustomOps: RMSNorm, SiluAndMul, and the rotary embedding from get_rope (vllm/model_executor/models/llama.py:38, vllm/model_executor/models/llama.py:44, vllm/model_executor/models/llama.py:52).
• maybe_fused_act_quant, which can fuse the activation with the quantization of the next layer's input (vllm/model_executor/models/llama.py:43, vllm/model_executor/models/llama.py:120).
• Helpers from models/utils.py and mixins from models/interfaces.py (vllm/model_executor/models/llama.py:61-79).

Key classes and functions#

Walk through the main methods#

LlamaAttention.__init__: sliding windows and the attention class#

The head arithmetic and the two linear layers are on How a vLLM model file is written (llama.py) (lines 143-181). The second half of the constructor:

        self._init_rotary_emb(config, quant_config=quant_config)

        sliding_window = None
        if layer_types := getattr(config, "layer_types", None):
            # Fix for Eagle3 compatibility:
            # for draft models, subtract target layer count
            # to get draft-relative layer index starting from 0
            if hasattr(config, "target_layer_count"):
                # This is a draft model,
                # adjust layer_idx to be relative to draft layers
                effective_layer_idx = layer_idx - config.target_layer_count
            else:
                # This is a target model, use layer_idx directly
                effective_layer_idx = layer_idx
            assert effective_layer_idx < len(layer_types), (
                f"effective_layer_idx: {effective_layer_idx} "
                f"is out of bounds for layer_types: {layer_types}"
            )

            is_sliding = layer_types[effective_layer_idx] == "sliding_attention"
            if is_sliding:
                sliding_window = config.sliding_window

        attn_cls = (
            EncoderOnlyAttention
            if attn_type == AttentionType.ENCODER_ONLY
            else Attention
        )

        self.attn = attn_cls(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            per_layer_sliding_window=sliding_window,
            attn_type=attn_type,
            prefix=f"{prefix}.attn",
        )

1. Line 183. _init_rotary_emb calls get_rope with the head size, max_position_embeddings and the config's rope_parameters, NeoX style (lines 236-248).
2. Lines 186-204. If the config has layer_types, a layer marked sliding_attention gets config.sliding_window. The layer index comes from the prefix (line 141). For an EAGLE-3 draft, the index is shifted by target_layer_count.
3. Lines 206-210. ENCODER_ONLY selects EncoderOnlyAttention; everything else uses Attention.
4. Lines 212-222. The attention layer gets per-rank heads (self.num_heads, self.num_kv_heads), the scale head_dim ** -0.5 (line 162) and its prefix.

LlamaDecoderLayer: config knobs and the residual#

        config = config or vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = self.get_quant_config(vllm_config)

        self.hidden_size = config.hidden_size
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
        # Support abacusai/Smaug-72B-v0.1 with attention_bias
        attention_bias = getattr(config, "attention_bias", False) or getattr(
            config, "bias", False
        )
        bias_o_proj = attention_bias
        # support internlm/internlm3-8b with qkv_bias
        if hasattr(config, "qkv_bias"):
            attention_bias = config.qkv_bias

        # By default, Llama uses causal attention as it is a decoder-only model.
        # You can override the HF config with `is_causal=False` to enable
        # bidirectional attention, which is used in some embedding models
        # (e.g. nvidia/llama-nemotron-embed-1b-v2)
        if getattr(config, "is_causal", True):
            attn_type = AttentionType.DECODER
        else:
            attn_type = AttentionType.ENCODER_ONLY

1. Line 263. The quant config comes from get_quant_config(), which subclasses can override (lines 332-334).
2. Lines 268-274. attention_bias or bias turns on biases in QKV and o_proj; qkv_bias sets the QKV bias alone. The comments name the models that need this.
3. Lines 280-283. is_causal=False in the config makes the layer bidirectional, for embedding models.

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual

1. Lines 320-322. The first layer has no residual yet: it keeps the input as the residual and normalizes.
2. Line 324. Later layers call the norm with the residual. RMSNorm adds it and normalizes in one op, and returns the new residual too.
3. Line 328. The same fused add-and-norm before the MLP. The layer returns (hidden_states, residual); the final norm of LlamaModel adds the last residual (line 452).

LlamaModel.load_weights and LlamaForCausalLM.load_weights: the load path#

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)

Follow one tensor, model.layers.0.mlp.up_proj.weight, in a TP = 2 worker:

1. The default loader calls LlamaForCausalLM.load_weights(iterator) (vllm/model_executor/model_loader/default_loader.py:503).
2. The outer AutoWeightsLoader has no model mapper. It adds the quant config mappers, if any, and the rule that drops old rotary tables (vllm/model_executor/models/utils.py:460-473).
3. The walk groups by model and finds self.model, a LlamaModel with its own load_weights. It passes the rest of the name: layers.0.mlp.up_proj.weight (vllm/model_executor/models/utils.py:387-400).
4. LlamaModel.load_weights makes a second AutoWeightsLoader with hf_to_vllm_mapper. .up_proj becomes .gate_up_proj with shard id 1 (vllm/model_executor/models/llama.py:354-355).
5. The walk goes layers → 0 → mlp → gate_up_proj. MergedColumnParallelLinear has a load_weights; it reads shard_id = 1 from the tensor and calls the parameter's loader (vllm/model_executor/layers/linear.py:985-1000).
6. The v2 loader computes offset 14336 / 2 = 7168 and size 7168 in the parameter. It copies checkpoint rows rank × 7168 to rank × 7168 + 7167 (vllm/model_executor/layers/linear.py:965-983, vllm/model_executor/parameter.py:153-173).
7. The loaded name goes back up the tree with its prefixes: model.layers.0.mlp.gate_up_proj.weight. The default loader uses the set of these names for its completeness check (vllm/model_executor/model_loader/default_loader.py:510-523).

The tensor lm_head.weight takes another path. In the outer walk, lm_head is a ParallelLMHead without load_weights, so its weight parameter uses VocabParallelEmbedding.weight_loader (vllm/model_executor/models/utils.py:416-424). If tie_word_embeddings is set, lm_head.weight is the same object as model.embed_tokens.weight (lines 509-510), and the loader skips the second name (vllm/model_executor/models/utils.py:255-259). One exception: if the config says tied but the checkpoint has its own lm_head, VllmConfig turns the flag off, so the head loads (vllm/config/vllm.py:1594-1599). After loading, vLLM ties the two again if they are equal (vllm/config/model.py:1255-1281, vllm/model_executor/model_loader/utils.py:155).

LlamaForCausalLM: class attributes#

class LlamaForCausalLM(
    LocalArgmaxMixin,
    nn.Module,
    SupportsLoRA,
    SupportsPP,
    SupportsEagle,
    SupportsEagle3,
    SupportsQuant,
):
    hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper
    # LoRA specific attributes
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }

1. Lines 464-472. The mixins; How a vLLM model file is written (llama.py) has a table of what each one adds.
2. Line 473. The outer class shares the mapper of LlamaModel. SupportsQuant.__new__ gives its rename-only part to the quant config, so names in ignored_layers can match vLLM names (vllm/model_executor/models/interfaces.py:1292-1299). For Llama this part is empty: get_rename_mapper() drops the stacked rules, which are the only rules of this mapper (vllm/model_executor/models/utils.py:180-199).
3. Lines 475-478. packed_modules_mapping tells LoRA and the quant config which checkpoint layers form each stacked layer.
4. Lines 479-482. embedding_modules lists the layers that get LoRA embedding buffers; Mistral clears it (vllm/model_executor/models/mistral.py:237-238).

Deep dive#

The sections below give the full detail. Open them one at a time, or use Open all deep dives at the top of the page.

Related concepts#

• How a vLLM model file is written (llama.py): the patterns of this file, with the prefix rules and PP.
• "Weight loading: from checkpoint files to GPU parameters": the loader, the mapper and the weight loaders.
• The tensor-parallel layers (linear.py and vocab_parallel_embedding.py): the layer classes that this file builds.
• Model capability protocols (SupportsPP, SupportsLoRA, ...): what the mixins mean to the registry.
• CustomOp, out-of-tree ops and IR op providers: RMSNorm, SiluAndMul and the rotary embedding.