vllm.model_executor.models.pixtral

Classes:

• PatchMerger –

  Learned merging of spatial_merge_size ** 2 patches.

• PixtralEncoderCudaGraphMixin –

  Encoder CUDA graph support for models built on a Pixtral vision tower.

• PixtralForConditionalGeneration –
• PixtralHFVisionModel –
• PixtralImagePixelInputs –

  Dimensions:

• VisionTransformer –

Functions:

• apply_rotary_emb_vit –

  Rotate xq/xk of shape (bsz, seq_len, n_heads, head_dim).

• pack_image_patches –

  Pack CHW images into one (num_patches, C * ph * pw) tensor.

• pixtral_patch_embed –

  Normalize and project all CHW images in one matrix multiplication.

• precompute_freqs_cis_2d –

  freqs_cis: 2D complex tensor of shape (height, width, dim // 2)

PatchMerger

Bases: Module

Learned merging of spatial_merge_size ** 2 patches.

Methods:

• forward_packed –

  Same as forward, gathering blocks with _make_merge_indices.

• permute –

  Args:

Source code in vllm/model_executor/models/pixtral.py

class PatchMerger(nn.Module):
    """Learned merging of spatial_merge_size ** 2 patches."""

    def __init__(
        self,
        vision_encoder_dim: int,
        spatial_merge_size: int,
        use_mlp_bias: bool = False,
    ) -> None:
        super().__init__()

        mlp_input_dim = vision_encoder_dim * (spatial_merge_size**2)

        self.spatial_merge_size = spatial_merge_size
        self.mlp_input_dim = mlp_input_dim

        self.merging_layer = ReplicatedLinear(
            mlp_input_dim, vision_encoder_dim, bias=use_mlp_bias, return_bias=False
        )

    def forward(
        self, x: torch.Tensor, image_sizes: list[tuple[int, int]]
    ) -> torch.Tensor:
        # image_sizes specified in tokens
        assert sum([h * w for h, w in image_sizes]) == len(x)

        # x is (N, vision_encoder_dim)
        x = self.permute(x, image_sizes)

        # x is (N / spatial_merge_size ** 2,
        #       vision_encoder_dim * spatial_merge_size ** 2)
        x = self.merging_layer(x)

        # x is (N / spatial_merge_size ** 2, vision_encoder_dim)
        return x

    def forward_packed(
        self, x: torch.Tensor, merge_indices: torch.Tensor
    ) -> torch.Tensor:
        """Same as `forward`, gathering blocks with `_make_merge_indices`."""
        return self.merging_layer(_merge_patches_by_index(x, merge_indices))

    def permute(
        self,
        x: torch.Tensor,
        image_sizes: list[tuple[int, int]],
    ) -> torch.Tensor:
        """Args:
            x: (N, D) where N is flattened and concatenated patch tokens
                for all images
            image_sizes: list of tuple of (height, width) in tokens for
                each image
        Returns:
            image_features: reorders patch tokens so each grid of
                (spatial_merge_size, spatial_merge_size) is contiguous.
                now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2)

        """
        sub_grids = get_sub_grids(
            x=x, image_sizes=image_sizes, spatial_merge_size=self.spatial_merge_size
        )  # list of [d x sub_grid_size x sub_grid_size x n_patches]
        permuted_tensor: list[torch.Tensor] = []
        for grid in sub_grids:
            n_patches = grid.shape[-1]
            permuted_tensor.append(
                grid.view(-1, n_patches).t()
            )  # n_patches x d * sub_grid_size * sub_grid_size
        return torch.cat(
            permuted_tensor, dim=0
        )  # (N / spatial_merge_size ** 2, d * spatial_merge_size ** 2)

forward_packed(x, merge_indices)

Same as forward, gathering blocks with _make_merge_indices.

Source code in vllm/model_executor/models/pixtral.py

def forward_packed(
    self, x: torch.Tensor, merge_indices: torch.Tensor
) -> torch.Tensor:
    """Same as `forward`, gathering blocks with `_make_merge_indices`."""
    return self.merging_layer(_merge_patches_by_index(x, merge_indices))

permute(x, image_sizes)

Parameters:

• x

  (Tensor) –

  (N, D) where N is flattened and concatenated patch tokens for all images

• image_sizes

  (list[tuple[int, int]]) –

  list of tuple of (height, width) in tokens for each image

Returns: image_features: reorders patch tokens so each grid of (spatial_merge_size, spatial_merge_size) is contiguous. now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2)

Source code in vllm/model_executor/models/pixtral.py

def permute(
    self,
    x: torch.Tensor,
    image_sizes: list[tuple[int, int]],
) -> torch.Tensor:
    """Args:
        x: (N, D) where N is flattened and concatenated patch tokens
            for all images
        image_sizes: list of tuple of (height, width) in tokens for
            each image
    Returns:
        image_features: reorders patch tokens so each grid of
            (spatial_merge_size, spatial_merge_size) is contiguous.
            now (N / spatial_merge_size ** 2, D * spatial_merge_size ** 2)

    """
    sub_grids = get_sub_grids(
        x=x, image_sizes=image_sizes, spatial_merge_size=self.spatial_merge_size
    )  # list of [d x sub_grid_size x sub_grid_size x n_patches]
    permuted_tensor: list[torch.Tensor] = []
    for grid in sub_grids:
        n_patches = grid.shape[-1]
        permuted_tensor.append(
            grid.view(-1, n_patches).t()
        )  # n_patches x d * sub_grid_size * sub_grid_size
    return torch.cat(
        permuted_tensor, dim=0
    )  # (N / spatial_merge_size ** 2, d * spatial_merge_size ** 2)

PixtralEncoderCudaGraphMixin

Bases: SupportsEncoderCudaGraph

Encoder CUDA graph support for models built on a Pixtral vision tower.

All images of a batch are packed into one patch sequence. Capture buffers are sized for the token budget, and the batch is padded with zero-length sequences up to max_batch_size, so replays only copy fresh values into fixed-shape buffers.

Methods:

• project_packed_image_features –

  Map (N_patches, D) tower features to language model embeddings.

Attributes:

• encoder_cudagraph_spatial_merge_size (int) –

  Patches merged per side by the projector, 1 when not merging.

Source code in vllm/model_executor/models/pixtral.py

class PixtralEncoderCudaGraphMixin(SupportsEncoderCudaGraph):
    """Encoder CUDA graph support for models built on a Pixtral vision tower.

    All images of a batch are packed into one patch sequence. Capture buffers
    are sized for the token budget, and the batch is padded with zero-length
    sequences up to `max_batch_size`, so replays only copy fresh values into
    fixed-shape buffers.
    """

    # mm_kwargs key holding the list of CHW images.
    encoder_cudagraph_pixel_key: ClassVar[str]

    config: Any
    multimodal_config: MultiModalConfig | None

    @property
    def encoder_cudagraph_vision_tower(
        self,
    ) -> "VisionTransformer | PixtralHFVisionModel":
        raise NotImplementedError

    @property
    def encoder_cudagraph_spatial_merge_size(self) -> int:
        """Patches merged per side by the projector, 1 when not merging."""
        raise NotImplementedError

    def project_packed_image_features(
        self, features: torch.Tensor, merge_indices: torch.Tensor | None
    ) -> torch.Tensor:
        """Map (N_patches, D) tower features to language model embeddings."""
        raise NotImplementedError

    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
        tower = self.encoder_cudagraph_vision_tower
        attn = tower.encoder_attention
        # SDPA splits the packed batch with Python sequence lengths, and
        # dynamic FP8 scales update their state on the host.
        graph_safe = (
            tower.attn_backend != AttentionBackendEnum.TORCH_SDPA
            and not attn._fp8_dynamic_scale
        )
        # Graph replays only handle pixel inputs, not precomputed embeddings.
        mm_config = self.multimodal_config
        pixels_only = mm_config is None or not mm_config.enable_mm_embeds
        return EncoderCudaGraphConfig(
            modalities=["image"] if graph_safe and pixels_only else [],
            buffer_keys=[
                "patches",
                "positions",
                "cu_seqlens",
                "max_seqlen",
                "sequence_lengths",
                "merge_indices",
            ],
            out_hidden_size=self.config.text_config.hidden_size,
        )

    def get_encoder_cudagraph_budget_range(
        self, vllm_config: VllmConfig
    ) -> tuple[int, int]:
        tower = self.encoder_cudagraph_vision_tower
        merged_patch_size = tower.patch_size * self.encoder_cudagraph_spatial_merge_size
        # Smallest budget: a 224x224 image.
        min_budget = math.ceil(224 / merged_patch_size) ** 2
        max_budget = min(
            vllm_config.scheduler_config.max_num_batched_tokens,
            vllm_config.model_config.max_model_len,
        )
        return min_budget, max_budget

    def get_encoder_cudagraph_item_specs(
        self, mm_kwargs: dict[str, Any]
    ) -> list[EncoderItemSpec]:
        m = self.encoder_cudagraph_spatial_merge_size
        grid_sizes = _patch_grid_sizes(
            mm_kwargs[self.encoder_cudagraph_pixel_key],
            self.encoder_cudagraph_vision_tower.patch_size,
        )
        specs = []
        for h, w in grid_sizes:
            # The processors resize images to a multiple of the merged patch
            # size; the capture buffers rely on it too.
            if h % m or w % m:
                raise ValueError(
                    f"Image patch grid {h}x{w} is not divisible by the "
                    f"spatial merge size {m}."
                )
            specs.append(
                EncoderItemSpec(input_size=h * w, output_tokens=(h // m) * (w // m))
            )
        return specs

    def select_encoder_cudagraph_items(
        self, mm_kwargs: dict[str, Any], indices: list[int]
    ) -> dict[str, Any]:
        images = mm_kwargs[self.encoder_cudagraph_pixel_key]
        return {self.encoder_cudagraph_pixel_key: [images[i] for i in indices]}

    def prepare_encoder_cudagraph_capture_inputs(
        self,
        token_budget: int,
        max_batch_size: int,
        max_frames_per_batch: int,
        device: torch.device,
        dtype: torch.dtype,
        path: str = "default",
        axis_keys: tuple[Hashable, ...] | None = None,
    ) -> EncoderCudaGraphCaptureInputs:
        tower = self.encoder_cudagraph_vision_tower
        m = self.encoder_cudagraph_spatial_merge_size
        num_patches = token_budget * m * m
        input_norm = getattr(tower, "input_norm", None)
        # max_seqlen is baked into the graph, so capture a single sequence
        # spanning the whole buffer to cover any replay.
        metadata = tower.packed_attention_metadata(
            [num_patches], device, num_sequences=max_batch_size
        )
        values: dict[str, torch.Tensor | None] = {
            "patches": torch.zeros(
                num_patches,
                tower.patch_conv.weight[0].numel(),
                device=device,
                dtype=getattr(input_norm, "input_dtype", None) or dtype,
            ),
            # Replays zero-pad positions; (0, 0) is a valid rotary index.
            "positions": torch.zeros(num_patches, 2, dtype=torch.long, device=device),
            "cu_seqlens": metadata["cu_seqlens"],
            "max_seqlen": metadata["max_seqlen"],
            "sequence_lengths": metadata["sequence_lengths"],
        }
        if m > 1:
            values["merge_indices"] = torch.arange(num_patches, device=device).view(
                token_budget, m * m
            )
        return EncoderCudaGraphCaptureInputs(
            values={key: value for key, value in values.items() if value is not None}
        )

    def prepare_encoder_cudagraph_replay_buffers(
        self,
        mm_kwargs: dict[str, Any],
        max_batch_size: int,
        max_frames_per_batch: int,
        path: str = "default",
    ) -> EncoderCudaGraphReplayBuffers:
        tower = self.encoder_cudagraph_vision_tower
        m = self.encoder_cudagraph_spatial_merge_size
        images = [
            image.to(device=tower.device)
            for image in mm_kwargs[self.encoder_cudagraph_pixel_key]
        ]
        patches, grid_sizes = pack_image_patches(
            images, (tower.patch_size, tower.patch_size)
        )
        metadata = tower.packed_attention_metadata(
            [h * w for h, w in grid_sizes],
            patches.device,
            num_sequences=max_batch_size,
        )
        # max_seqlen is not replayed, keeping the captured upper bound.
        values = {
            "patches": patches,
            "positions": _patch_grid_positions(grid_sizes, patches.device),
            "cu_seqlens": metadata["cu_seqlens"],
            "sequence_lengths": metadata["sequence_lengths"],
        }
        if m > 1:
            values["merge_indices"] = _make_merge_indices(grid_sizes, m, patches.device)
        return EncoderCudaGraphReplayBuffers(values=values)

    def encoder_cudagraph_forward(
        self,
        inputs: dict[str, torch.Tensor],
        path: str = "default",
    ) -> torch.Tensor:
        features = self.encoder_cudagraph_vision_tower.forward_packed(
            inputs["patches"],
            inputs["positions"],
            inputs["cu_seqlens"],
            inputs["max_seqlen"],
            inputs.get("sequence_lengths"),
        )
        return self.project_packed_image_features(features, inputs.get("merge_indices"))

    def encoder_eager_forward(
        self,
        mm_kwargs: dict[str, Any],
        path: str = "default",
    ) -> torch.Tensor:
        embeddings = cast(SupportsMultiModal, self).embed_multimodal(**mm_kwargs)
        return torch.cat(tuple(embeddings))

encoder_cudagraph_spatial_merge_size property

Patches merged per side by the projector, 1 when not merging.

project_packed_image_features(features, merge_indices)

Map (N_patches, D) tower features to language model embeddings.

Source code in vllm/model_executor/models/pixtral.py

def project_packed_image_features(
    self, features: torch.Tensor, merge_indices: torch.Tensor | None
) -> torch.Tensor:
    """Map (N_patches, D) tower features to language model embeddings."""
    raise NotImplementedError

PixtralForConditionalGeneration

Bases: Module, SupportsLoRA, SupportsEagle3, SupportsMultiModal, SupportsPP, PixtralEncoderCudaGraphMixin

Methods:

• forward –

  Run forward pass for pixtral.

Source code in vllm/model_executor/models/pixtral.py

@MULTIMODAL_REGISTRY.register_processor(
    PixtralMultiModalProcessor,
    info=PixtralProcessingInfo,
    dummy_inputs=PixtralDummyInputsBuilder,
)
class PixtralForConditionalGeneration(
    nn.Module,
    SupportsLoRA,
    SupportsEagle3,
    SupportsMultiModal,
    SupportsPP,
    PixtralEncoderCudaGraphMixin,
):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "model.language_model.": "language_model.model.",
            "model.vision_tower.": "vision_encoder.",
            "model.multi_modal_projector.": "vision_language_adapter.",
        },
        orig_to_new_substr={
            ".linear_1.": ".w_in.",
            ".linear_2.": ".w_out.",
        },
    )

    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }

    supports_tower_connector_lora = True

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return None

        raise ValueError("Only image modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config
        self.config = config
        self.multimodal_config = multimodal_config

        dataclass_fields = {field.name for field in fields(VisionEncoderArgs)}
        vision_args = {
            key: value
            for key, value in self.config.vision_config.to_dict().items()
            if key in dataclass_fields
        }

        self.vision_args = VisionEncoderArgs(**vision_args)

        # init MistralForCausalLM
        with self._mark_language_model(vllm_config):
            self.language_model = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
            )

        with self._mark_tower_model(vllm_config, "image"):
            self.vision_encoder = VisionTransformer(
                self.vision_args,
                prefix=maybe_prefix(prefix, "vision_encoder"),
            )
            self.pre_mm_projector_norm = (
                RMSNorm(self.vision_args.hidden_size, eps=1e-5)
                if self.vision_args.add_pre_mm_projector_layer_norm
                else None
            )
            self.patch_merger = (
                PatchMerger(
                    vision_encoder_dim=self.vision_args.hidden_size,
                    spatial_merge_size=self.vision_args.spatial_merge_size,
                    use_mlp_bias=False,
                )
                if self.vision_args.mm_projector_id == PATCH_MERGE
                else None
            )
            self.vision_language_adapter = VisionLanguageAdapter(
                self.vision_args, dim=config.text_config.hidden_size
            )

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

    def _parse_and_validate_image_input(
        self, **kwargs: object
    ) -> PixtralImagePixelInputs | None:
        images = kwargs.pop("images", None)
        if images is None:
            return None

        return PixtralImagePixelInputs(
            type="pixel_values",
            images=images,
        )

    def _process_image_input(
        self,
        image_input: PixtralImagePixelInputs,
    ) -> tuple[torch.Tensor, ...]:
        images = image_input["images"]
        image_features = self.vision_encoder(images)
        feature_sizes = [image_feature.shape[0] for image_feature in image_features]
        image_features = torch.cat(image_features)
        if self.pre_mm_projector_norm is not None:
            image_features = self.pre_mm_projector_norm(image_features)
        if self.patch_merger is not None:
            patch_size = self.vision_args.patch_size
            spatial_merge_size_square = self.vision_args.spatial_merge_size**2
            img_patch_dims = [
                (img.shape[1] // patch_size, img.shape[2] // patch_size)
                for img in images
            ]
            feature_sizes = [
                feature_size // spatial_merge_size_square
                for feature_size in feature_sizes
            ]
            image_features = self.patch_merger(
                image_features, image_sizes=img_patch_dims
            )
        image_embeds = self.vision_language_adapter(image_features)
        image_embeds = torch.split(image_embeds, feature_sizes)
        return image_embeds

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
            return []

        return self._process_image_input(image_input)

    # -- SupportsEncoderCudaGraph --

    encoder_cudagraph_pixel_key = "images"

    @property
    def encoder_cudagraph_vision_tower(self) -> "VisionTransformer":
        return self.vision_encoder

    @property
    def encoder_cudagraph_spatial_merge_size(self) -> int:
        if self.patch_merger is None:
            return 1
        return self.vision_args.spatial_merge_size

    def project_packed_image_features(
        self, features: torch.Tensor, merge_indices: torch.Tensor | None
    ) -> torch.Tensor:
        if self.pre_mm_projector_norm is not None:
            features = self.pre_mm_projector_norm(features)
        if self.patch_merger is not None:
            assert merge_indices is not None
            features = self.patch_merger.forward_packed(features, merge_indices)
        return self.vision_language_adapter(features)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: object,
    ) -> torch.Tensor | IntermediateTensors:
        """Run forward pass for pixtral."""
        if intermediate_tensors is not None:
            inputs_embeds = None

        hidden_states = self.language_model.model(
            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
        )

        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        return self.language_model.compute_logits(hidden_states)

    def _require_language_model_eagle3(self) -> None:
        if not supports_eagle3(self.language_model):
            raise RuntimeError(
                f"EAGLE-3 speculative decoding requires the language model to "
                f"support EAGLE-3, but {type(self.language_model).__name__} does not."
            )

    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
        self._require_language_model_eagle3()
        self.language_model.set_aux_hidden_state_layers(layers)

    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
        self._require_language_model_eagle3()
        return self.language_model.get_eagle3_aux_hidden_state_layers()

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
        _vision_encoder_stacked_params = [
            # (param_name, shard_name, shard_id)
            # HF format
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
            # Mistral native (consolidated) format
            (".qkv_proj", ".wq", "q"),
            (".qkv_proj", ".wk", "k"),
            (".qkv_proj", ".wv", "v"),
            (".gate_up_proj", ".w1", 0),
            (".gate_up_proj", ".w3", 1),
        ]

        # Remap Mistral native names to HF-style names
        # used by the vLLM vision encoder modules.
        _vision_encoder_name_remap = {
            ".wo.": ".o_proj.",
            ".w2.": ".down_proj.",
        }

        def is_vision_encoder_weights(weight: tuple[str, torch.Tensor]):
            return weight[0].startswith(("vision_encoder", "vision_tower"))

        def is_vision_lang_adapter_weights(weight: tuple[str, torch.Tensor]):
            return weight[0].startswith(
                ("vision_language_adapter", "multi_modal_projector")
            )

        def is_patch_merger(weight: tuple[str, torch.Tensor]):
            return weight[0].startswith("patch_merger")

        def is_pre_mm_projector_norm(weight: tuple[str, torch.Tensor]):
            return weight[0].startswith("pre_mm_projector_norm")

        vision_encoder_dict = (
            dict(self.vision_encoder.named_parameters())
            if self.vision_encoder is not None
            else {}
        )
        patch_merger_dict = (
            dict(self.patch_merger.named_parameters())
            if self.patch_merger is not None
            else {}
        )
        pre_mm_projector_norm_dict = (
            dict(self.pre_mm_projector_norm.named_parameters())
            if self.pre_mm_projector_norm is not None
            else {}
        )
        vision_lang_adapter_dict = (
            dict(self.vision_language_adapter.named_parameters())
            if self.vision_language_adapter is not None
            else {}
        )

        def llm_weights_generator():
            for name, w in weights:
                if is_vision_encoder_weights((name, w)):
                    if _is_layer_none_or_staged(self.vision_encoder):
                        continue
                    trimmed_name = ".".join(name.split(".")[1:])
                    for (
                        param_name,
                        weight_name,
                        shard_id,
                    ) in _vision_encoder_stacked_params:
                        if weight_name in trimmed_name:
                            trimmed_name = trimmed_name.replace(weight_name, param_name)
                            param = vision_encoder_dict[trimmed_name]
                            weight_loader = param.weight_loader
                            weight_loader(param, w, shard_id)
                            break
                    else:
                        for old, new in _vision_encoder_name_remap.items():
                            if old in trimmed_name:
                                trimmed_name = trimmed_name.replace(old, new)
                                break

                        param = vision_encoder_dict.get(trimmed_name)
                        if param is not None:
                            weight_loader = getattr(
                                param, "weight_loader", default_weight_loader
                            )
                            weight_loader(param, w)
                elif is_patch_merger((name, w)):
                    if _is_layer_none_or_staged(self.patch_merger):
                        continue
                    trimmed_name = ".".join(name.split(".")[1:])
                    param = patch_merger_dict[trimmed_name]
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, w)
                elif is_pre_mm_projector_norm((name, w)):
                    if _is_layer_none_or_staged(self.pre_mm_projector_norm):
                        continue
                    trimmed_name = ".".join(name.split(".")[1:])
                    param = pre_mm_projector_norm_dict[trimmed_name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
                elif is_vision_lang_adapter_weights((name, w)):
                    if _is_layer_none_or_staged(self.vision_language_adapter):
                        continue
                    trimmed_name = ".".join(name.split(".")[1:])
                    param = vision_lang_adapter_dict.get(trimmed_name)
                    if param is not None:
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, w)
                else:
                    name = name.removeprefix("language_model.")
                    yield (name, w)

        self.language_model.load_weights(llm_weights_generator())

    def get_mm_mapping(self) -> MultiModelKeys:
        return MultiModelKeys.from_string_field(
            language_model="language_model.",
            connector="vision_language_adapter.",
            tower_model="vision_encoder",
        )

    def get_mm_lora_token_counts(
        self,
        *,
        modality: str,
        mm_kwargs: MultiModalKwargsItem | None,
        num_mm_embeds: int,
    ) -> tuple[int, int | None]:
        del modality, mm_kwargs
        if getattr(self, "patch_merger", None) is None:
            return num_mm_embeds, num_mm_embeds
        merge_size = self.vision_args.spatial_merge_size
        return num_mm_embeds * (merge_size**2), num_mm_embeds

forward(input_ids, positions, intermediate_tensors=None, inputs_embeds=None, **kwargs)

Run forward pass for pixtral.

Source code in vllm/model_executor/models/pixtral.py

def forward(
    self,
    input_ids: torch.Tensor | None,
    positions: torch.Tensor,
    intermediate_tensors: IntermediateTensors | None = None,
    inputs_embeds: torch.Tensor | None = None,
    **kwargs: object,
) -> torch.Tensor | IntermediateTensors:
    """Run forward pass for pixtral."""
    if intermediate_tensors is not None:
        inputs_embeds = None

    hidden_states = self.language_model.model(
        input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
    )

    return hidden_states

PixtralHFVisionModel

Bases: Module

Methods:

• forward –

  Args:

• forward_packed –

  Encode images packed by pack_image_patches.

Attributes:

• encoder_attention (MMEncoderAttention) –

  Attention of the first layer; all layers share its settings.

Source code in vllm/model_executor/models/pixtral.py

class PixtralHFVisionModel(nn.Module):
    def __init__(
        self,
        config: PixtralVisionConfig,
        quant_config: QuantizationConfig | None = None,
        *,
        input_norm: nn.Module | None = None,
        num_hidden_layers_override: int | None = None,
        require_post_norm: bool | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()

        self.config = config
        self.input_norm = input_norm

        self.patch_conv = Conv2dLayer(
            in_channels=config.num_channels,
            out_channels=config.hidden_size,
            kernel_size=config.patch_size,
            stride=config.patch_size,
            bias=False,
        )
        self.ln_pre = RMSNorm(config.hidden_size, eps=1e-5)
        self.transformer = PixtralHFTransformer(
            config,
            quant_config=quant_config,
            num_hidden_layers_override=num_hidden_layers_override,
            prefix=f"{prefix}.transformer",
        )

        num_hidden_layers = config.num_hidden_layers
        if len(self.transformer.layers) > config.num_hidden_layers:
            raise ValueError(
                f"The original encoder only has {num_hidden_layers} "
                f"layers, but you requested {len(self.transformer.layers)} "
                "layers."
            )

        if require_post_norm is True:
            msg = "PixtralHFVisionModel does not have post-layernorm"
            raise ValueError(msg)

        self.dtype = next(self.parameters()).dtype
        self.device = next(self.parameters()).device
        self.patch_positional_embedding = PixtralVisionRotaryEmbedding(config).to(
            self.device
        )
        self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(
            config.num_attention_heads, config.hidden_size // config.num_attention_heads
        )

    @property
    def patch_size(self) -> int:
        return cast(int, self.config.patch_size)

    @property
    def encoder_attention(self) -> MMEncoderAttention:
        """Attention of the first layer; all layers share its settings."""
        return self._first_attention.attn

    @property
    def _first_attention(self) -> "PixtralHFAttention":
        return cast(PixtralHFTransformerBlock, self.transformer.layers[0]).attention

    @property
    def attn_backend(self) -> AttentionBackendEnum:
        return self.encoder_attention.attn_backend

    def packed_attention_metadata(
        self,
        sequence_lengths: list[int],
        device: torch.device,
        *,
        num_sequences: int | None = None,
    ) -> PackedAttentionMetadata:
        cu_seqlens, max_seqlen, packed_sequence_lengths = (
            _make_packed_sequence_metadata(
                sequence_lengths,
                self.attn_backend,
                self.config.hidden_size,
                self._first_attention.tp_size,
                device,
                num_sequences=num_sequences,
                fp8_padded_hidden_size=self.fp8_padded_hidden_size,
            )
        )
        return {
            "cu_seqlens": cu_seqlens,
            "max_seqlen": max_seqlen,
            "sequence_lengths": packed_sequence_lengths,
        }

    def forward_packed(
        self,
        patches: torch.Tensor,
        positions: torch.Tensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: torch.Tensor,
        sequence_lengths: torch.Tensor | None,
        *,
        select_layers: list[int] | None = None,
        feature_select_strategy: VisionFeatureSelectStrategy | None = None,
    ) -> torch.Tensor:
        """Encode images packed by `pack_image_patches`.

        `patches` holds (N_toks, C * patch_size**2) pixel patches (raw pixels
        when `input_norm` normalizes on device) and `positions` their
        (N_toks, 2) (row, col) coordinates; the attention arguments come from
        `packed_attention_metadata`. Returns (N_toks, D) token features.
        """
        weight = self.patch_conv.weight
        if self.input_norm is not None:
            patches = self.input_norm(patches, weight.dtype)
        else:
            patches = patches.to(weight.dtype)
        patch_embeds = F.linear(patches, weight.reshape(weight.shape[0], -1))
        patch_embeds = self.ln_pre(patch_embeds.unsqueeze(0))

        if not TRANSFORMERS_WITH_AXIAL_ROPE:
            max_width = self.config.image_size // self.config.patch_size
            positions = positions[:, 0] * max_width + positions[:, 1]
        position_embedding = self.patch_positional_embedding(patch_embeds, positions)

        out = self.transformer(
            patch_embeds,
            position_embedding,
            cu_seqlens,
            max_seqlen,
            sequence_lengths,
            return_all_hidden_states=select_layers is not None,
        )

        out = resolve_visual_encoder_outputs(
            out,
            None,
            select_layers=select_layers,
            max_possible_layers=self.config.num_hidden_layers,
            feature_select_strategy=feature_select_strategy,
        )
        return out.squeeze(0)

    def forward(
        self,
        pixel_values: list[torch.Tensor],
        *,
        select_layers: list[int] | None = None,
        feature_select_strategy: VisionFeatureSelectStrategy | None = None,
    ) -> tuple[torch.Tensor, ...]:
        """Args:
            pixel_values: Each image to be processed will be a separate tensor
                in pixel_values. This means it will be a list of tensors
                because multiple requests batched can have multiple images,
                each with their own shape potentially
            select_layers: Layer indices whose features should be
                concatenated and used as the visual encoder output. If none
                are provided, the last layer is used.

        Returns:
            image_features: tensor of token features for
                all tokens of all images of shape (N_toks, D)

        """
        images = [image.to(device=self.device) for image in pixel_values]
        patches, grid_sizes = pack_image_patches(
            images, (self.patch_size, self.patch_size)
        )
        embed_sizes = [h * w for h, w in grid_sizes]
        out = self.forward_packed(
            patches,
            _patch_grid_positions(grid_sizes, self.device),
            **self.packed_attention_metadata(embed_sizes, self.device),
            select_layers=select_layers,
            feature_select_strategy=feature_select_strategy,
        )

        # split into separate tensors for each image
        return torch.split(out, embed_sizes)

    # (TODO) Add prefix argument for filtering out weights to be loaded
    #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        layer_count = len(self.transformer.layers)

        for name, loaded_weight in weights:
            # omit layers when num_hidden_layers_override is set
            if name.startswith("transformer.layers"):
                layer_idx = int(name.split(".")[2])
                if layer_idx >= layer_count:
                    continue

            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

encoder_attention property

Attention of the first layer; all layers share its settings.

forward(pixel_values, *, select_layers=None, feature_select_strategy=None)

Parameters:

• pixel_values

  (list[Tensor]) –

  Each image to be processed will be a separate tensor in pixel_values. This means it will be a list of tensors because multiple requests batched can have multiple images, each with their own shape potentially

• select_layers

  (list[int] | None, default: None ) –

  Layer indices whose features should be concatenated and used as the visual encoder output. If none are provided, the last layer is used.

Returns:

• image_features ( tuple[Tensor, ...] ) –

  tensor of token features for all tokens of all images of shape (N_toks, D)

Source code in vllm/model_executor/models/pixtral.py

def forward(
    self,
    pixel_values: list[torch.Tensor],
    *,
    select_layers: list[int] | None = None,
    feature_select_strategy: VisionFeatureSelectStrategy | None = None,
) -> tuple[torch.Tensor, ...]:
    """Args:
        pixel_values: Each image to be processed will be a separate tensor
            in pixel_values. This means it will be a list of tensors
            because multiple requests batched can have multiple images,
            each with their own shape potentially
        select_layers: Layer indices whose features should be
            concatenated and used as the visual encoder output. If none
            are provided, the last layer is used.

    Returns:
        image_features: tensor of token features for
            all tokens of all images of shape (N_toks, D)

    """
    images = [image.to(device=self.device) for image in pixel_values]
    patches, grid_sizes = pack_image_patches(
        images, (self.patch_size, self.patch_size)
    )
    embed_sizes = [h * w for h, w in grid_sizes]
    out = self.forward_packed(
        patches,
        _patch_grid_positions(grid_sizes, self.device),
        **self.packed_attention_metadata(embed_sizes, self.device),
        select_layers=select_layers,
        feature_select_strategy=feature_select_strategy,
    )

    # split into separate tensors for each image
    return torch.split(out, embed_sizes)

forward_packed(patches, positions, cu_seqlens, max_seqlen, sequence_lengths, *, select_layers=None, feature_select_strategy=None)

Encode images packed by pack_image_patches.

patches holds (N_toks, C * patch_size**2) pixel patches (raw pixels when input_norm normalizes on device) and positions their (N_toks, 2) (row, col) coordinates; the attention arguments come from packed_attention_metadata. Returns (N_toks, D) token features.

Source code in vllm/model_executor/models/pixtral.py

def forward_packed(
    self,
    patches: torch.Tensor,
    positions: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: torch.Tensor,
    sequence_lengths: torch.Tensor | None,
    *,
    select_layers: list[int] | None = None,
    feature_select_strategy: VisionFeatureSelectStrategy | None = None,
) -> torch.Tensor:
    """Encode images packed by `pack_image_patches`.

    `patches` holds (N_toks, C * patch_size**2) pixel patches (raw pixels
    when `input_norm` normalizes on device) and `positions` their
    (N_toks, 2) (row, col) coordinates; the attention arguments come from
    `packed_attention_metadata`. Returns (N_toks, D) token features.
    """
    weight = self.patch_conv.weight
    if self.input_norm is not None:
        patches = self.input_norm(patches, weight.dtype)
    else:
        patches = patches.to(weight.dtype)
    patch_embeds = F.linear(patches, weight.reshape(weight.shape[0], -1))
    patch_embeds = self.ln_pre(patch_embeds.unsqueeze(0))

    if not TRANSFORMERS_WITH_AXIAL_ROPE:
        max_width = self.config.image_size // self.config.patch_size
        positions = positions[:, 0] * max_width + positions[:, 1]
    position_embedding = self.patch_positional_embedding(patch_embeds, positions)

    out = self.transformer(
        patch_embeds,
        position_embedding,
        cu_seqlens,
        max_seqlen,
        sequence_lengths,
        return_all_hidden_states=select_layers is not None,
    )

    out = resolve_visual_encoder_outputs(
        out,
        None,
        select_layers=select_layers,
        max_possible_layers=self.config.num_hidden_layers,
        feature_select_strategy=feature_select_strategy,
    )
    return out.squeeze(0)

PixtralImagePixelInputs

Bases: TensorSchema

Dimensions

• bn: Batch size * number of images
• c: Number of channels (3)
• h: Height of each image
• w: Width of each image

The result of stacking ImageEncoding.tokens from each prompt.

Source code in vllm/model_executor/models/pixtral.py

class PixtralImagePixelInputs(TensorSchema):
    """Dimensions:
        - bn: Batch size * number of images
        - c: Number of channels (3)
        - h: Height of each image
        - w: Width of each image

    The result of stacking `ImageEncoding.tokens` from each prompt.
    """

    type: Literal["pixel_values"] = "pixel_values"

    images: Annotated[
        torch.Tensor | list[torch.Tensor],
        TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),
    ]

VisionTransformer

Bases: Module

Methods:

• forward –

  Args:

• forward_packed –

  Encode images packed by pack_image_patches.

Attributes:

• encoder_attention (MMEncoderAttention) –

  Attention of the first layer; all layers share its settings.

Source code in vllm/model_executor/models/pixtral.py

class VisionTransformer(nn.Module):
    def __init__(
        self,
        args: VisionEncoderArgs,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.args = args
        disable_tp = is_vit_use_data_parallel()
        self.patch_conv = Conv2dLayer(
            in_channels=args.num_channels,
            out_channels=args.hidden_size,
            kernel_size=args.patch_size,
            stride=args.patch_size,
            bias=False,
        )
        self.ln_pre = RMSNorm(args.hidden_size, eps=1e-5)
        self.transformer = Transformer(
            args,
            quant_config=quant_config,
            prefix=f"{prefix}.transformer",
            disable_tp=disable_tp,
        )

        head_dim = self.args.hidden_size // self.args.num_attention_heads
        assert head_dim % 2 == 0, "ROPE requires even head_dim"
        self._freqs_cis: torch.Tensor | None = None
        self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(
            args.num_attention_heads, head_dim
        )

    @property
    def patch_size(self) -> int:
        return self.args.patch_size

    @property
    def max_patches_per_side(self) -> int:
        return self.args.image_size // self.args.patch_size

    @property
    def device(self) -> torch.types.Device:
        return next(self.parameters()).device

    @property
    def dtype(self) -> torch.dtype:
        return next(self.parameters()).dtype

    @property
    def freqs_cis(self) -> torch.Tensor:
        if self._freqs_cis is None:
            self._freqs_cis = torch.view_as_real(
                precompute_freqs_cis_2d(
                    dim=self.args.hidden_size // self.args.num_attention_heads,
                    height=self.max_patches_per_side,
                    width=self.max_patches_per_side,
                    theta=self.args.rope_theta,
                )
            )

        if self._freqs_cis.device != self.device:
            self._freqs_cis = self._freqs_cis.to(device=self.device)

        return self._freqs_cis

    @property
    def encoder_attention(self) -> MMEncoderAttention:
        """Attention of the first layer; all layers share its settings."""
        return cast(TransformerBlock, self.transformer.layers[0]).attention.attn

    @property
    def attn_backend(self) -> AttentionBackendEnum:
        return self.encoder_attention.attn_backend

    def packed_attention_metadata(
        self,
        sequence_lengths: list[int],
        device: torch.device,
        *,
        num_sequences: int | None = None,
    ) -> PackedAttentionMetadata:
        cu_seqlens, max_seqlen, packed_sequence_lengths = (
            _make_packed_sequence_metadata(
                sequence_lengths,
                self.attn_backend,
                self.args.hidden_size,
                1
                if is_vit_use_data_parallel()
                else get_tensor_model_parallel_world_size(),
                device,
                num_sequences=num_sequences,
                fp8_padded_hidden_size=self.fp8_padded_hidden_size,
            )
        )
        return {
            "cu_seqlens": cu_seqlens,
            "max_seqlen": max_seqlen,
            "sequence_lengths": packed_sequence_lengths,
        }

    def forward_packed(
        self,
        patches: torch.Tensor,
        positions: torch.Tensor,
        cu_seqlens: torch.Tensor,
        max_seqlen: torch.Tensor,
        sequence_lengths: torch.Tensor | None,
    ) -> torch.Tensor:
        """Encode images packed by `pack_image_patches`.

        `patches` holds (N_toks, C * patch_size**2) pixel patches and
        `positions` their (N_toks, 2) (row, col) coordinates; the attention
        arguments come from `packed_attention_metadata`. Returns (N_toks, D)
        token features.
        """
        weight = self.patch_conv.weight
        patch_embeds = F.linear(
            patches.to(weight.dtype), weight.reshape(weight.shape[0], -1)
        )
        patch_embeds = self.ln_pre(patch_embeds.unsqueeze(0))
        freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]
        out = self.transformer(
            patch_embeds,
            freqs_cis=freqs_cis,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
            sequence_lengths=sequence_lengths,
        )
        return out.squeeze(0)

    def forward(
        self,
        images: list[torch.Tensor],
    ) -> tuple[torch.Tensor, ...]:
        """Args:
            images: list of N_img images of variable sizes,
                each of shape (C, H, W)

        Returns:
            image_features: per-image token features, each of shape
                (N_toks_i, D)

        """
        patch_size = self.args.patch_size
        patches, grid_sizes = pack_image_patches(images, (patch_size, patch_size))
        embed_sizes = [h * w for h, w in grid_sizes]
        out = self.forward_packed(
            patches,
            _patch_grid_positions(grid_sizes, patches.device),
            **self.packed_attention_metadata(embed_sizes, patches.device),
        )
        return torch.split(out, embed_sizes)

encoder_attention property

Attention of the first layer; all layers share its settings.

forward(images)

Parameters:

• images

  (list[Tensor]) –

  list of N_img images of variable sizes, each of shape (C, H, W)

Returns:

• image_features ( tuple[Tensor, ...] ) –

  per-image token features, each of shape (N_toks_i, D)

Source code in vllm/model_executor/models/pixtral.py

def forward(
    self,
    images: list[torch.Tensor],
) -> tuple[torch.Tensor, ...]:
    """Args:
        images: list of N_img images of variable sizes,
            each of shape (C, H, W)

    Returns:
        image_features: per-image token features, each of shape
            (N_toks_i, D)

    """
    patch_size = self.args.patch_size
    patches, grid_sizes = pack_image_patches(images, (patch_size, patch_size))
    embed_sizes = [h * w for h, w in grid_sizes]
    out = self.forward_packed(
        patches,
        _patch_grid_positions(grid_sizes, patches.device),
        **self.packed_attention_metadata(embed_sizes, patches.device),
    )
    return torch.split(out, embed_sizes)

forward_packed(patches, positions, cu_seqlens, max_seqlen, sequence_lengths)

Encode images packed by pack_image_patches.

patches holds (N_toks, C * patch_size**2) pixel patches and positions their (N_toks, 2) (row, col) coordinates; the attention arguments come from packed_attention_metadata. Returns (N_toks, D) token features.

Source code in vllm/model_executor/models/pixtral.py

def forward_packed(
    self,
    patches: torch.Tensor,
    positions: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: torch.Tensor,
    sequence_lengths: torch.Tensor | None,
) -> torch.Tensor:
    """Encode images packed by `pack_image_patches`.

    `patches` holds (N_toks, C * patch_size**2) pixel patches and
    `positions` their (N_toks, 2) (row, col) coordinates; the attention
    arguments come from `packed_attention_metadata`. Returns (N_toks, D)
    token features.
    """
    weight = self.patch_conv.weight
    patch_embeds = F.linear(
        patches.to(weight.dtype), weight.reshape(weight.shape[0], -1)
    )
    patch_embeds = self.ln_pre(patch_embeds.unsqueeze(0))
    freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]
    out = self.transformer(
        patch_embeds,
        freqs_cis=freqs_cis,
        cu_seqlens=cu_seqlens,
        max_seqlen=max_seqlen,
        sequence_lengths=sequence_lengths,
    )
    return out.squeeze(0)

_make_merge_indices(grid_sizes, spatial_merge_size, device)

Indices gathering each spatial_merge_size^2 patch block of the packed sequence into one row, matching the patch mergers' unfold order.

Source code in vllm/model_executor/models/pixtral.py

def _make_merge_indices(
    grid_sizes: Sequence[tuple[int, int]],
    spatial_merge_size: int,
    device: torch.device,
) -> torch.Tensor:
    """Indices gathering each `spatial_merge_size`^2 patch block of the packed
    sequence into one row, matching the patch mergers' unfold order."""
    m = spatial_merge_size
    indices = []
    offset = 0
    for h, w in grid_sizes:
        grid = np.arange(offset, offset + h * w).reshape(h // m, m, w // m, m)
        indices.append(grid.transpose(0, 2, 1, 3).reshape(-1, m * m))
        offset += h * w
    return async_tensor_h2d(np.concatenate(indices), device=device)

_make_packed_sequence_metadata(sequence_lengths, attn_backend, hidden_size, tp_size, device, *, num_sequences=None, fp8_padded_hidden_size=None)

Build varlen attention metadata for images packed in one sequence.

num_sequences pads the batch with zero-length sequences, which keeps the metadata shapes fixed for encoder CUDA graphs while attention does no work for the padding rows.

Source code in vllm/model_executor/models/pixtral.py

def _make_packed_sequence_metadata(
    sequence_lengths: list[int],
    attn_backend: AttentionBackendEnum,
    hidden_size: int,
    tp_size: int,
    device: torch.device,
    *,
    num_sequences: int | None = None,
    fp8_padded_hidden_size: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Build varlen attention metadata for images packed in one sequence.

    `num_sequences` pads the batch with zero-length sequences, which keeps the
    metadata shapes fixed for encoder CUDA graphs while attention does no work
    for the padding rows.
    """
    lengths = np.array(sequence_lengths, dtype=np.int32)
    if num_sequences is not None:
        assert len(lengths) <= num_sequences
        lengths = np.pad(lengths, (0, num_sequences - len(lengths)))
    cu_seqlens = np.concatenate(
        [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)]
    )
    sequence_lengths_tensor = MMEncoderAttention.maybe_compute_seq_lens(
        attn_backend, cu_seqlens, device
    )
    # Attention wrappers call .item() on max_seqlen; keeping it on the CPU
    # avoids a device sync (and a D2H copy recorded in encoder CUDA graphs).
    max_seqlen = torch.tensor(
        MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens),
        dtype=torch.int32,
    )
    cu_seqlens_tensor = MMEncoderAttention.maybe_recompute_cu_seqlens(
        attn_backend,
        cu_seqlens,
        hidden_size,
        tp_size,
        device,
        fp8_padded_hidden_size=fp8_padded_hidden_size,
    )
    return cu_seqlens_tensor, max_seqlen, sequence_lengths_tensor

_patch_grid_positions(grid_sizes, device)

(row, col) coordinates of every patch, packed as (num_patches, 2).

Source code in vllm/model_executor/models/pixtral.py

def _patch_grid_positions(
    grid_sizes: Sequence[tuple[int, int]], device: torch.device
) -> torch.Tensor:
    """(row, col) coordinates of every patch, packed as (num_patches, 2)."""
    positions = np.concatenate(
        [np.indices((h, w)).reshape(2, -1).T for h, w in grid_sizes]
    )
    return async_tensor_h2d(positions, device=device)

apply_rotary_emb_vit(xq, xk, freqs_cis)

Rotate xq/xk of shape (bsz, seq_len, n_heads, head_dim).

freqs_cis is the real view (seq_len, head_dim / 2, 2) of the complex frequencies.

Source code in vllm/model_executor/models/pixtral.py

def apply_rotary_emb_vit(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Rotate `xq`/`xk` of shape (bsz, seq_len, n_heads, head_dim).

    `freqs_cis` is the real view (seq_len, head_dim / 2, 2) of the complex
    frequencies.
    """
    assert freqs_cis.dtype == torch.float32
    assert freqs_cis.shape == (xq.shape[1], xq.shape[-1] // 2, 2), (
        freqs_cis.shape,
        xq.shape,
    )
    if torch.compiler.is_compiling():
        # Inductor does not generate code for complex operators and would
        # fall back to eager around RoPE.
        return _apply_rotary_emb_vit_real(xq, xk, freqs_cis)
    # In eager mode the complex multiply is a single kernel per tensor, while
    # the real formulation launches several elementwise kernels.
    return _apply_rotary_emb_vit_complex(xq, xk, freqs_cis)

pack_image_patches(images, patch_shape)

Pack CHW images into one (num_patches, C * ph * pw) tensor.

Patches are laid out row-major per image, with each row flattened in the (C, ph, pw) order of a patch convolution weight. Returns the packed tensor and the (rows, cols) patch grid of every image.

Source code in vllm/model_executor/models/pixtral.py

def pack_image_patches(
    images: Sequence[torch.Tensor],
    patch_shape: tuple[int, int],
) -> tuple[torch.Tensor, list[tuple[int, int]]]:
    """Pack CHW images into one (num_patches, C * ph * pw) tensor.

    Patches are laid out row-major per image, with each row flattened in the
    (C, ph, pw) order of a patch convolution weight. Returns the packed tensor
    and the (rows, cols) patch grid of every image.
    """
    patch_height, patch_width = patch_shape
    in_channels = images[0].shape[0]
    grid_sizes = []
    for image in images:
        assert image.ndim == 3 and image.shape[0] == in_channels
        _, height, width = image.shape
        grid_sizes.append((height // patch_height, width // patch_width))

    patch_counts = [rows * cols for rows, cols in grid_sizes]
    packed = torch.empty(
        (sum(patch_counts), in_channels * patch_height * patch_width),
        dtype=images[0].dtype,
        device=images[0].device,
    )
    offset = 0
    for image, (rows, cols), count in zip(images, grid_sizes, patch_counts):
        source = (
            image[:, : rows * patch_height, : cols * patch_width]
            .reshape(in_channels, rows, patch_height, cols, patch_width)
            .permute(1, 3, 0, 2, 4)
        )
        packed[offset : offset + count].view(
            rows, cols, in_channels, patch_height, patch_width
        ).copy_(source)
        offset += count
    return packed, grid_sizes

pixtral_patch_embed(images, weight, input_norm)

Normalize and project all CHW images in one matrix multiplication.

Source code in vllm/model_executor/models/pixtral.py

def pixtral_patch_embed(
    images: list[torch.Tensor],
    weight: torch.Tensor,
    input_norm: nn.Module,
) -> tuple[torch.Tensor, list[torch.Tensor]]:
    """Normalize and project all CHW images in one matrix multiplication."""
    out_channels, in_channels, patch_height, patch_width = weight.shape
    assert images[0].shape[0] == in_channels
    packed, grid_sizes = pack_image_patches(images, (patch_height, patch_width))
    normalized = input_norm(packed, weight.dtype)
    projected = F.linear(normalized, weight.reshape(out_channels, -1))
    patch_counts = [rows * cols for rows, cols in grid_sizes]
    embeddings = [
        tokens.transpose(0, 1).reshape(1, out_channels, rows, cols)
        for tokens, (rows, cols) in zip(projected.split(patch_counts), grid_sizes)
    ]
    return projected.unsqueeze(0), embeddings

precompute_freqs_cis_2d(dim, height, width, theta)

freqs_cis: 2D complex tensor of shape (height, width, dim // 2) to be indexed by (height, width) position tuples

Source code in vllm/model_executor/models/pixtral.py

def precompute_freqs_cis_2d(
    dim: int,
    height: int,
    width: int,
    theta: float,
) -> torch.Tensor:
    """freqs_cis: 2D complex tensor of shape (height, width, dim // 2)
    to be indexed by (height, width) position tuples
    """
    # (dim / 2) frequency bases
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))

    h = torch.arange(height, device=freqs.device)
    w = torch.arange(width, device=freqs.device)

    freqs_h = torch.outer(h, freqs[::2]).float()
    freqs_w = torch.outer(w, freqs[1::2]).float()
    freqs_2d = torch.cat(
        [
            freqs_h[:, None, :].repeat(1, width, 1),
            freqs_w[None, :, :].repeat(height, 1, 1),
        ],
        dim=-1,
    )
    return torch.polar(torch.ones_like(freqs_2d), freqs_2d)