How a vLLM model file is written (llama.py)

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

The idea#

A decoder model is a stack of layers: an embedding, N decoder layers and an LM head. Each decoder layer has a norm, attention, a norm and an MLP (Inside a transformer↗). vLLM does not run the Transformers class of the checkpoint. It runs its own class for the architecture name, for example LlamaForCausalLM in vllm/model_executor/models/llama.py (Out-of-tree models and the model registry).

The vLLM class has the same math and the same parameter tree as the Transformers class. It differs in four ways:

• Constructor contract. Every class takes vllm_config (the whole VllmConfig) and a prefix (of a layer), the dotted name of the layer in the tree.
• Tensor-parallel layers. Linear layers and the embedding are vLLM classes that hold only the slice of this tensor parallelism rank, and do the communication themselves (The tensor-parallel layers (linear.py and vocab_parallel_embedding.py)).
• An attention layer object. The model gives Q, K and V to an Attention layer. That layer finds its KV cache and the metadata of the step itself, through the forward context (The forward context: how layers find per-step data).
• Pipeline stages. Layers that another pipeline parallelism stage owns are placeholders, and a stage that is not the last returns IntermediateTensors instead of hidden states.

Other work is not in the model file. A decorator adds torch.compile; the compile backend and the model runner add CUDA graphs. A loader streams the weights in, and each parameter cuts its own slice ("Weight loading: from checkpoint files to GPU parameters"). Mixin classes tell the registry what the model supports (Model capability protocols (SupportsPP, SupportsLoRA, ...)).

An everyday comparison. Think of a new office building. The builders put a room number on every door, such as "3.12.B" (floor, corridor, room). Later, the furniture company, the phone list and the cleaning plan all use these numbers. A door with a wrong number gets the wrong desk or no phone. The comparison stops here. In vLLM, nobody checks the number on the door. Two systems read the number, and two systems walk the real module tree. When the number and the path are different, the systems disagree.

The picture#

Read the tree from the top. Each parent passes a longer prefix to its children, so the prefix of a row is the dotted path to it. Blue and violet rows are the tensor-parallel linear layers. The amber row is the attention layer that registers itself under its prefix.

In the code#

Step 1: the constructor contract#

The worker builds the model in initialize_model(). It checks by name that the constructor has the parameters vllm_config and prefix. Then it calls the class with both, inside set_current_vllm_config (vllm/model_executor/model_loader/utils.py:96-103). A class without them gets a DeprecationWarning and an argument guess; Out-of-tree models and the model registry explains that path. The root gets prefix="": Model Runner V2 calls load_model() without a prefix, so the default "" goes to create_model() (vllm/v1/worker/gpu/model_runner.py:403-406, vllm/model_executor/model_loader/base_loader.py:63-80).

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config

        self.model = self._init_model(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "model"),
            layer_type=layer_type,
        )

        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
            if config.tie_word_embeddings:
                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

            logit_scale = getattr(config, "logit_scale", 1.0)
            self.logits_processor = LogitsProcessor(
                config.vocab_size, scale=logit_scale
            )
        else:
            self.lm_head = PPMissingLayer()

        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors
        )

1. Lines 487-488. Keyword-only vllm_config and prefix. The registry also checks the keyword vllm_config when it inspects the class (Model capability protocols (SupportsPP, SupportsLoRA, ...)).
2. Lines 492-493. The model reads what it needs from vllm_config: the Hugging Face config and the quantization config. Inner classes do the same.
3. Lines 496-500. The inner model gets the prefix maybe_prefix(prefix, "model"), which is "model" for the root (vllm/model_executor/models/utils.py:971-982).
4. Lines 502-508. Only the last pipeline stage builds the LM head.
5. Lines 519-521. The model exposes a factory for empty intermediate tensors; pipeline parallelism needs it.

Step 2: how each layer gets its prefix, and why it matters#

Each constructor adds the attribute name to the prefix it got. The decoder layer gives f"{prefix}.self_attn" and f"{prefix}.mlp" (vllm/model_executor/models/llama.py:285-307). The attention class gives .qkv_proj, .o_proj and .attn (vllm/model_executor/models/llama.py:165-181, vllm/model_executor/models/llama.py:212-222). make_layers() builds the list of decoder layers with prefix=f"{prefix}.{idx}" (vllm/model_executor/models/utils.py:911-917). So the QKV layer of layer 0 gets the prefix model.layers.0.self_attn.qkv_proj.

Four parts of vLLM use the name of a layer:

No code checks that the prefix equals the attribute path of the layer. But the quantization config compares the prefix with checkpoint names, and the weight loader and LoRA use the module tree. So use the attribute path as the prefix; then all four parts use the same name (derived from the four rows above). The section What can go wrong gives the failures.

Step 3: tensor-parallel layers#

LlamaAttention divides the heads by the TP size and builds a QKVParallelLinear and a RowParallelLinear:

        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

        head_dim = getattr(config, "head_dim", None)
        self.head_dim = head_dim or self.hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.max_position_embeddings = max_position_embeddings

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )

        self.o_proj = RowParallelLinear(
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
            bias=bias_o_proj,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

1. Lines 145-146. The query heads must divide by the TP size. Start-up checks this earlier, with a clear error (vllm/config/model.py:1557-1564).
2. Lines 148-156. KV heads are split when there are enough of them, and replicated when TP is larger (Tensor and pipeline parallelism in practice has the table).
3. Lines 160-161. q_size and kv_size are per rank. The forward pass splits the QKV output with them.
4. Lines 165-181. The model gives the total head counts to QKVParallelLinear and the full input size to o_proj. The layers divide by the TP size themselves.

LlamaMLP uses a MergedColumnParallelLinear for gate and up, and a RowParallelLinear for down (vllm/model_executor/models/llama.py:95-111); Tensor and pipeline parallelism in practice shows the shapes. The embedding is a VocabParallelEmbedding and the head is a ParallelLMHead (vllm/model_executor/models/llama.py:382-386, vllm/model_executor/models/llama.py:503-508). The model code calls no collective itself. The page The tensor-parallel layers (linear.py and vocab_parallel_embedding.py) explains each layer class.

Step 4: the attention layer and the CustomOps#

        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",
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output

1. Lines 212-222. Attention gets the per-rank head counts, the scale, the cache and quantization configs, an optional sliding window, and its prefix. When it is built, it selects an attention backend, unless the caller gives one (vllm/model_executor/layers/attention/attention.py:349-350). It also registers itself in the forward context.
2. Line 229. One GEMM computes Q, K and V for the heads of this rank: [T, q_size + 2 × kv_size].
3. Lines 230-231. The split, then RoPE on Q and K.
4. Line 232. self.attn(q, k, v) passes no KV cache and no metadata. The attention op uses the layer name to find the metadata and the layer object (with its KV cache) in the forward context (vllm/model_executor/layers/attention/attention.py:705-717). The page The forward context: how layers find per-step data explains how.
5. Line 233. The row-parallel o_proj adds the partial sums of the ranks.

Three small layers are CustomOps: RMSNorm (vllm/model_executor/layers/layernorm.py:38-39), the rotary embedding from get_rope() (vllm/model_executor/layers/rotary_embedding/base.py:14-15) and SiluAndMul (vllm/model_executor/layers/activation.py:115-116). Each one picks its forward method when it is built (CustomOp, out-of-tree ops and IR op providers). get_rope() keeps a cache keyed by its arguments, so all 32 layers share one rotary object (vllm/model_executor/layers/rotary_embedding/__init__.py:84-85, vllm/model_executor/layers/rotary_embedding/__init__.py:351).

Step 5: LlamaModel, compilation and pipeline stages#

@support_torch_compile(
    # TODO[#32068]: Investigate recompilation
    # mark_unbacked_dims={"input_ids": 0},
    dynamic_arg_dims={
        "input_ids": {0: "b"},
        "positions": {0: "b"},
        "intermediate_tensors": {0: "b"},
        "inputs_embeds": {0: "b"},
    },
)

The decorator wraps LlamaModel.forward for torch.compile and CUDA graphs (torch.compile and CUDA graphs in vLLM). It sits on LlamaModel, not on LlamaForCausalLM, so compute_logits stays outside the compiled graph.

The constructor builds the parts of this pipeline stage only:

        if (
            get_pp_group().is_first_rank
            or (config.tie_word_embeddings and get_pp_group().is_last_rank)
            or spec_decode_needs_target_embed(vllm_config)
        ):
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
            )
        else:
            self.embed_tokens = PPMissingLayer()
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix),
            prefix=f"{prefix}.layers",
        )
        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()

        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )

1. Lines 377-388. The embedding is real on the first stage. It is also real on the last stage when the head shares its weight, or when a speculative drafter needs it. Elsewhere it is a PPMissingLayer.
2. Lines 389-393. make_layers() builds real layers only for indices start_layer to end_layer - 1. All other slots are PPMissingLayer, so the indices and names stay the same on every stage (vllm/model_executor/models/utils.py:886-919).
3. Lines 394-397. The final norm exists on the last stage only.
4. Lines 399-401. Empty intermediate tensors have the keys hidden_states and residual, each [batch_size, hidden_size] of zeros (vllm/model_executor/models/utils.py:955-968).

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
        **extra_layer_kwargs,
    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.embed_input_ids(input_ids)
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

        remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)

        aux_hidden_states: list[torch.Tensor] = []
        if get_pp_group().is_first_rank:
            self._maybe_add_hidden_state(
                aux_hidden_states, self.start_layer, hidden_states, residual
            )
        for idx, layer in enumerate(
            islice(self.layers, self.start_layer, self.end_layer),
            start=self.start_layer,
        ):
            hidden_states, residual = layer(
                positions, hidden_states, residual, **extra_layer_kwargs
            )
            self._maybe_add_hidden_state(
                aux_hidden_states, idx + 1, hidden_states, residual
            )

        if not get_pp_group().is_last_rank:
            return IntermediateTensors(
                {
                    "hidden_states": hidden_states,
                    "residual": residual,
                    **self.pack_local_aux_hidden_states(aux_hidden_states),
                }
            )

        hidden_states, _ = self.norm(hidden_states, residual)

        aux_hidden_states = remote_aux + aux_hidden_states
        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states
        return hidden_states

1. Lines 414-423. The first stage starts from the token embeddings, or from inputs_embeds when the runner gives them. Other stages read hidden_states and residual from the stage before.
2. Lines 432-438. The loop runs only the layers of this stage. Each layer returns the new hidden states and the residual.
3. Lines 443-450. A stage that is not the last returns IntermediateTensors.
4. Line 452. The last stage applies the final norm. It returns hidden states of shape [T, hidden_size], not logits.

The decoder layer keeps the residual outside the norm. input_layernorm(hidden_states, residual) adds the residual and normalizes in one call, and returns both tensors (vllm/model_executor/models/llama.py:320-328, vllm/model_executor/layers/layernorm.py:85-99).

Step 6: LlamaForCausalLM: forward, logits and mixins#

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
        model_output = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
        return model_output

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        logits = self.logits_processor(self.lm_head, hidden_states)
        return logits

    def compute_logits_local(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)

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

1. Lines 541-543. forward(input_ids, positions, intermediate_tensors, inputs_embeds) only calls the inner model. The model runner builds these four arguments (vllm/v1/worker/gpu/model_runner.py:1984-2009).
2. Line 550. compute_logits applies the head through LogitsProcessor. To sample, the runner calls it only on the rows that sample a token, hidden_states[logits_indices] (vllm/v1/worker/gpu/model_runner.py:1637-1638). Prompt logprobs call it on prompt rows too (vllm/v1/worker/gpu/model_runner.py:2177-2178).
3. Line 557. compute_logits_local skips the gather. The runner uses it for batch-sharded sampling.
4. Lines 560-561. load_weights hands the checkpoint stream to AutoWeightsLoader ("Weight loading: from checkpoint files to GPU parameters").

The class header lists the mixins (vllm/model_executor/models/llama.py:464-472):

LlamaModel also inherits EagleModelMixin, which collects the hidden states of chosen layers for EAGLE-3 (vllm/model_executor/models/interfaces.py:1608-1681). The page Model capability protocols (SupportsPP, SupportsLoRA, ...) explains how the registry detects each protocol.

Step 7: how other files reuse llama.py#

Nine registry names point to llama.py. Seven of them use LlamaForCausalLM directly, for example InternLM3ForCausalLM and TeleChat3ForCausalLM (vllm/model_executor/models/registry.py:140-155, vllm/model_executor/models/registry.py:212). The embedding table also copies every name that maps to LlamaForCausalLM (vllm/model_executor/models/registry.py:231-238). (Counts from an evaluation of the tables at this commit.)

Other model files subclass the Llama classes. MistralForCausalLM subclasses LlamaForCausalLM (vllm/model_executor/models/mistral.py:236), and so does Phi3ForCausalLM (vllm/model_executor/models/phi3.py:10). TeleChat2ForCausalLM adds its own name mapper and builds TeleChat2Model, a subclass of LlamaModel (vllm/model_executor/models/telechat2.py:107-127, vllm/model_executor/models/telechat2.py:35). 21 files in vllm/ import from llama.py (git grep at this commit). The file has two hooks for this reuse:

• A layer_type argument, so a subclass can give its own decoder layer class (vllm/model_executor/models/llama.py:360-366, vllm/model_executor/models/llama.py:523-529).
• LlamaDecoderLayer.get_quant_config(), which the docstring says to override in subclasses (vllm/model_executor/models/llama.py:332-334).

A second model tree: vllm/models/#

Some new model families do not live in vllm/model_executor/models/. They have one package per family in vllm/models/, with sub-packages per hardware vendor. The package __init__.py imports the implementation for current_platform, for example the NVIDIA, AMD, XPU or CPU variant of DeepSeek V4 (vllm/models/deepseek_v4/__init__.py:1-48). For the registry they are normal lazy entries (Out-of-tree models and the model registry). This site does not explain these packages further; llama.py is the one model example.

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.

What can go wrong#

• Two attention layers with the same prefix. The second one raises ValueError: Duplicate layer name while the model is built (vllm/model_executor/layers/attention/attention.py:443-444). A model that builds its layers in a loop must put the index in the prefix.
• A prefix that is not the real path. No error occurs at once. The quantization config compares the wrong name with its lists, so a layer can get the wrong method. For example, it can get FP8 although the checkpoint lists it in ignored_layers (vllm/model_executor/layers/quantization/fp8.py:180-190). An attention prefix without exactly one layer number fails the assertion of extract_layer_index() (vllm/model_executor/models/llama.py:141, vllm/model_executor/models/utils.py:1027-1030).
• Heads that do not divide. 32 query heads with TP = 3 stop start-up with "must be divisible by tensor parallel size" (vllm/config/model.py:1559-1564).
• Another activation. LlamaMLP accepts only silu and raises ValueError for other values of hidden_act (vllm/model_executor/models/llama.py:112-115).
• PP on a model without the mixin. Start-up raises NotImplementedError when the class does not implement SupportsPP (vllm/config/model.py:1569-1576).

Why it matters. A new architecture, a fork of a model, or an out-of-tree model plugin starts with a file like llama.py. Learn the contract, the prefixes and the TP layer classes. Then you can read any model file in vLLM. You can also write a new one that works with quantization, LoRA, PP and torch.compile.