vllm.model_executor.models.pixtral

Classes:

• PatchMerger –

  Learned merging of spatial_merge_size ** 2 patches.

• PixtralForConditionalGeneration –
• PixtralHFVisionModel –
• PixtralImagePixelInputs –

  Dimensions:

• VisionTransformer –

Functions:

• 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:

• 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 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)

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)

PixtralForConditionalGeneration

Bases: Module, SupportsLoRA, SupportsEagle3, SupportsMultiModal, SupportsPP

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
):
    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)

    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:

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
        )

    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]
        if self.input_norm is not None:
            patch_embeds, patch_embeds_list = pixtral_patch_embed(
                images,
                self.patch_conv.weight,
                self.input_norm,
            )
        else:
            patch_embeds_list = []
            for image in images:
                patch_embeds_list.append(
                    self.patch_conv(image.unsqueeze(0).to(self.dtype))
                )
            patch_embeds = torch.cat(
                [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1
            )

        embed_sizes = [p.shape[-2] * p.shape[-1] for p in patch_embeds_list]
        patch_embeds = self.ln_pre(patch_embeds)

        # positional embeddings
        position_ids = position_ids_in_meshgrid(
            patch_embeds_list,
            max_width=self.config.image_size // self.config.patch_size,
        ).to(self.device)
        position_embedding = self.patch_positional_embedding(patch_embeds, position_ids)

        attention = self.transformer.layers[0].attention.attn
        cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(
            embed_sizes,
            attention.attn_backend,
            self.config.hidden_size,
            self.transformer.layers[0].attention.tp_size,
            patch_embeds.device,
        )

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

        # squeeze dim 0 and split into separate tensors for each image
        return torch.split(out.squeeze(0), 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

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]
    if self.input_norm is not None:
        patch_embeds, patch_embeds_list = pixtral_patch_embed(
            images,
            self.patch_conv.weight,
            self.input_norm,
        )
    else:
        patch_embeds_list = []
        for image in images:
            patch_embeds_list.append(
                self.patch_conv(image.unsqueeze(0).to(self.dtype))
            )
        patch_embeds = torch.cat(
            [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1
        )

    embed_sizes = [p.shape[-2] * p.shape[-1] for p in patch_embeds_list]
    patch_embeds = self.ln_pre(patch_embeds)

    # positional embeddings
    position_ids = position_ids_in_meshgrid(
        patch_embeds_list,
        max_width=self.config.image_size // self.config.patch_size,
    ).to(self.device)
    position_embedding = self.patch_positional_embedding(patch_embeds, position_ids)

    attention = self.transformer.layers[0].attention.attn
    cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(
        embed_sizes,
        attention.attn_backend,
        self.config.hidden_size,
        self.transformer.layers[0].attention.tp_size,
        patch_embeds.device,
    )

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

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

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:

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

    @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 = 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

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

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

        """
        # pass images through initial convolution independently
        patch_embeds_list = [
            self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in images
        ]

        patch_embeds = [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list]
        embed_sizes = [p.shape[1] for p in patch_embeds]

        # flatten to a single sequence
        patch_embeds = torch.cat(patch_embeds, dim=1)
        patch_embeds = self.ln_pre(patch_embeds)

        # positional embeddings
        positions = position_meshgrid(patch_embeds_list).to(self.device)
        freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]

        attention = self.transformer.layers[0].attention.attn
        cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(
            embed_sizes,
            attention.attn_backend,
            self.args.hidden_size,
            1 if is_vit_use_data_parallel() else get_tensor_model_parallel_world_size(),
            patch_embeds.device,
        )
        out = self.transformer(
            patch_embeds,
            freqs_cis=freqs_cis,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
            sequence_lengths=sequence_lengths,
        )

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

forward(images)

Parameters:

• images

  (list[Tensor]) –

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

Returns:

• image_features ( 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,
    images: list[torch.Tensor],
) -> torch.Tensor:
    """Args:
        images: list of N_img images of variable sizes,
            each of shape (C, H, W)

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

    """
    # pass images through initial convolution independently
    patch_embeds_list = [
        self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in images
    ]

    patch_embeds = [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list]
    embed_sizes = [p.shape[1] for p in patch_embeds]

    # flatten to a single sequence
    patch_embeds = torch.cat(patch_embeds, dim=1)
    patch_embeds = self.ln_pre(patch_embeds)

    # positional embeddings
    positions = position_meshgrid(patch_embeds_list).to(self.device)
    freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]

    attention = self.transformer.layers[0].attention.attn
    cu_seqlens, max_seqlen, sequence_lengths = _make_packed_sequence_metadata(
        embed_sizes,
        attention.attn_backend,
        self.args.hidden_size,
        1 if is_vit_use_data_parallel() else get_tensor_model_parallel_world_size(),
        patch_embeds.device,
    )
    out = self.transformer(
        patch_embeds,
        freqs_cis=freqs_cis,
        cu_seqlens=cu_seqlens,
        max_seqlen=max_seqlen,
        sequence_lengths=sequence_lengths,
    )

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

_reshape_for_broadcast(freqs_cis, x)

freqs_cis: complex - (seq_len, head_dim / 2) x: complex - (bsz, seq_len, head_dim / 2)

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

def _reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    """freqs_cis: complex - (seq_len, head_dim / 2)
    x: complex - (bsz, seq_len, head_dim / 2)
    """
    ndim = x.ndim
    assert ndim > 1
    assert freqs_cis.shape == (x.shape[1], x.shape[-1]), (
        freqs_cis.shape,
        (x.shape[1], x.shape[-1]),
    )
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

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
    grid_sizes = []
    for image in images:
        assert image.ndim == 3 and image.shape[0] == in_channels
        _, height, width = image.shape
        patch_rows = height // patch_height
        patch_cols = width // patch_width
        grid_sizes.append((patch_rows, patch_cols))

    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

    normalized = input_norm(packed, weight.dtype)
    projected = F.linear(normalized, weight.reshape(out_channels, -1))
    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)