vllm.model_executor.models.mistral3

Classes:

• Mistral3ForConditionalGeneration –
• Mistral3ImageEmbeddingInputs –

  Projected features with shape (num_images, feature_size, hidden_size).

• Mistral3ImagePixelInputs –

  Dimensions:

• Mistral3PatchMerger –

  Learned merging of spatial_merge_size ** 2 patches.

Mistral3ForConditionalGeneration

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

Methods:

• forward –

  Run forward pass for Mistral3.

• get_mm_mapping –

  Get the module prefix in multimodal models.

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

@MULTIMODAL_REGISTRY.register_processor(
    Mistral3MultiModalProcessor,
    info=Mistral3ProcessingInfo,
    dummy_inputs=Mistral3DummyInputsBuilder,
)
class Mistral3ForConditionalGeneration(
    nn.Module,
    SupportsLoRA,
    SupportsMultiModal,
    SupportsPP,
    SupportsEagle,
    SupportsEagle3,
    PixtralEncoderCudaGraphMixin,
):
    supports_mm_device_do_normalize = True

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

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            # mapping for new names in checkpoint saved after transformers v4.52
            "model.language_model.": "language_model.model.",
            "model.vision_tower.": "vision_tower.",
            "model.multi_modal_projector.": "multi_modal_projector.",
            "lm_head.": "language_model.lm_head.",
            # Some PEFT LoRAs are trained against the text submodule directly
            # and produce names like `base_model.model.model.layers.*`.
            "model.": "language_model.model.",
        },
        orig_to_new_suffix={
            # FP8 quantized HF checkpoints use "activation_scale" and
            # "weight_scale_inv" but vLLM's FP8 linear layers register
            # them as "input_scale" and "weight_scale"
            ".activation_scale": ".input_scale",
            ".weight_scale_inv": ".weight_scale",
        },
    )

    @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 = "") -> None:
        super().__init__()

        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config

        self.config = config
        self.multimodal_config = multimodal_config

        # NOTE: This is a special case for Pixtral-12B in the HF-format
        # https://huggingface.co/mistral-community/pixtral-12b/blob/main/config.json  # noqa
        if (
            config.projector_hidden_act is None
            and config.vision_config.hidden_act == "gelu"
        ):
            config.projector_hidden_act = "gelu"

        with self._mark_tower_model(vllm_config, "image"):
            self.vision_tower = init_vision_tower_for_mistral3(
                config,
                quant_config=quant_config,
                input_norm=build_mm_input_norm(vllm_config.model_config),
                require_post_norm=False,
                prefix=maybe_prefix(prefix, "vision_tower"),
            )
            self.multi_modal_projector = Mistral3MultiModalProjector(
                vision_hidden_size=config.vision_config.hidden_size,
                text_hidden_size=config.text_config.hidden_size,
                projector_hidden_act=config.projector_hidden_act,
                spatial_merge_size=config.spatial_merge_size,
                patch_size=config.vision_config.patch_size,
                multimodal_projector_bias=config.multimodal_projector_bias,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "multi_modal_projector"),
            )

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

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

    def _parse_and_validate_image_input(
        self, **kwargs: object
    ) -> Mistral3ImageInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        image_embeds = kwargs.pop("image_embeds", None)

        if pixel_values is None and image_embeds is None:
            return None

        return Mistral3ImagePixelInputs(
            type="pixel_values_pixtral",
            pixel_values=pixel_values,
        )

    def _process_image_input(
        self,
        image_input: Mistral3ImageInputs,
    ) -> MultiModalEmbeddings:
        if image_input["type"] == "image_embeds":
            return image_input["data"]

        image_sizes = [
            (img.shape[-2], img.shape[-1]) for img in image_input["pixel_values"]
        ]

        image_features = self.vision_tower(image_input["pixel_values"])

        if isinstance(image_features, torch.Tensor):
            return self.multi_modal_projector(image_features, image_sizes)

        feature_sizes = [
            image_feature.shape[0] // self.config.spatial_merge_size**2
            for image_feature in image_features
        ]

        image_embeds = self.multi_modal_projector(
            torch.cat(image_features), image_sizes
        )
        if len(feature_sizes) > 1:
            image_embeds = torch.split(image_embeds, feature_sizes)
        else:
            image_embeds = (image_embeds,)
        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 []

        vision_embeddings = self._process_image_input(image_input)

        return vision_embeddings

    # -- SupportsEncoderCudaGraph --

    encoder_cudagraph_pixel_key = "pixel_values"

    @property
    def encoder_cudagraph_vision_tower(self) -> PixtralHFVisionModel:
        return self.vision_tower

    @property
    def encoder_cudagraph_spatial_merge_size(self) -> int:
        return self.config.spatial_merge_size

    def project_packed_image_features(
        self, features: torch.Tensor, merge_indices: torch.Tensor | None
    ) -> torch.Tensor:
        return self.multi_modal_projector.forward_packed(features, merge_indices)

    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 Mistral3.

        One key thing to understand is the `input_ids` already accounts for the
        positions of the to-be-inserted image embeddings.

        Concretely, consider a text prompt:
        `"USER: <image>\\nWhat's the content of the image?\\nASSISTANT:"`.

        Tokenizer outputs:
        `[1, 3148, 1001, 29901, 29871, 32000, 29871, 13, 5618, 29915, 29879,
        278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566, 29901]`.

        To reserve space in KV cache, we have to insert placeholder tokens
        before they are inputted to the model, so the input processor prepends
        additional image tokens (denoted as `32000`), resulting in:
        `[1, 3148, 1001, 29901, 29871, 32000, ..., 32000, 29871, 13, 5618,
        29915, 29879, 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566,
        29901]`.

        We insert 575 tokens so that including the original image token in the
        input, there are a total of 576 (24 * 24) image tokens, which
        corresponds to the number of image tokens inputted to the language
        model, i.e. the number of image tokens outputted by the visual encoder.

        This way, the `positions` and `attn_metadata` are consistent
        with the `input_ids`.

        Args:
            input_ids: Flattened (concatenated) input_ids corresponding to a
                batch.
            positions: Position indices for the input tokens.
            intermediate_tensors: Intermediate tensors from prior forward pass.
            inputs_embeds: Optional tensor of input embeddings.
            **kwargs: Multimodal inputs for this batch, forwarded to the
                multimodal embedding path.

        Info:
            [`Mistral3ImagePixelInputs`][vllm.model_executor.models.mistral3.Mistral3ImagePixelInputs]

        """
        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 load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    def get_mm_mapping(self) -> MultiModelKeys:
        """Get the module prefix in multimodal models."""
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="multi_modal_projector",
            tower_model="vision_tower",
        )

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

Run forward pass for Mistral3.

One key thing to understand is the input_ids already accounts for the positions of the to-be-inserted image embeddings.

Concretely, consider a text prompt: "USER: <image>\nWhat's the content of the image?\nASSISTANT:".

Tokenizer outputs: [1, 3148, 1001, 29901, 29871, 32000, 29871, 13, 5618, 29915, 29879, 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566, 29901].

To reserve space in KV cache, we have to insert placeholder tokens before they are inputted to the model, so the input processor prepends additional image tokens (denoted as 32000), resulting in: [1, 3148, 1001, 29901, 29871, 32000, ..., 32000, 29871, 13, 5618, 29915, 29879, 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566, 29901].

We insert 575 tokens so that including the original image token in the input, there are a total of 576 (24 * 24) image tokens, which corresponds to the number of image tokens inputted to the language model, i.e. the number of image tokens outputted by the visual encoder.

This way, the positions and attn_metadata are consistent with the input_ids.

Parameters:

• input_ids

  (Tensor | None) –

  Flattened (concatenated) input_ids corresponding to a batch.

• positions

  (Tensor) –

  Position indices for the input tokens.

• intermediate_tensors

  (IntermediateTensors | None, default: None ) –

  Intermediate tensors from prior forward pass.

• inputs_embeds

  (Tensor | None, default: None ) –

  Optional tensor of input embeddings.

• **kwargs

  (object, default: {} ) –

  Multimodal inputs for this batch, forwarded to the multimodal embedding path.

Info

Mistral3ImagePixelInputs

Source code in vllm/model_executor/models/mistral3.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 Mistral3.

    One key thing to understand is the `input_ids` already accounts for the
    positions of the to-be-inserted image embeddings.

    Concretely, consider a text prompt:
    `"USER: <image>\\nWhat's the content of the image?\\nASSISTANT:"`.

    Tokenizer outputs:
    `[1, 3148, 1001, 29901, 29871, 32000, 29871, 13, 5618, 29915, 29879,
    278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566, 29901]`.

    To reserve space in KV cache, we have to insert placeholder tokens
    before they are inputted to the model, so the input processor prepends
    additional image tokens (denoted as `32000`), resulting in:
    `[1, 3148, 1001, 29901, 29871, 32000, ..., 32000, 29871, 13, 5618,
    29915, 29879, 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566,
    29901]`.

    We insert 575 tokens so that including the original image token in the
    input, there are a total of 576 (24 * 24) image tokens, which
    corresponds to the number of image tokens inputted to the language
    model, i.e. the number of image tokens outputted by the visual encoder.

    This way, the `positions` and `attn_metadata` are consistent
    with the `input_ids`.

    Args:
        input_ids: Flattened (concatenated) input_ids corresponding to a
            batch.
        positions: Position indices for the input tokens.
        intermediate_tensors: Intermediate tensors from prior forward pass.
        inputs_embeds: Optional tensor of input embeddings.
        **kwargs: Multimodal inputs for this batch, forwarded to the
            multimodal embedding path.

    Info:
        [`Mistral3ImagePixelInputs`][vllm.model_executor.models.mistral3.Mistral3ImagePixelInputs]

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

get_mm_mapping()

Get the module prefix in multimodal models.

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

def get_mm_mapping(self) -> MultiModelKeys:
    """Get the module prefix in multimodal models."""
    return MultiModelKeys.from_string_field(
        language_model="language_model",
        connector="multi_modal_projector",
        tower_model="vision_tower",
    )

Mistral3ImageEmbeddingInputs

Bases: TensorSchema

Projected features with shape (num_images, feature_size, hidden_size).

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

class Mistral3ImageEmbeddingInputs(TensorSchema):
    """Projected features with shape (num_images, feature_size, hidden_size)."""

    type: Literal["image_embeds"] = "image_embeds"
    data: Annotated[
        torch.Tensor | list[torch.Tensor],
        TensorShape("bn", "ifs", "hs", dynamic_dims={"ifs"}),
    ]

Mistral3ImagePixelInputs

Bases: TensorSchema

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

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

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

    type: Literal["pixel_values_pixtral"] = "pixel_values_pixtral"

    # Note that `height` or `width` may be different per batch and image,
    # in which case the data is passed as a list instead of a batched tensor.
    pixel_values: Annotated[
        torch.Tensor | list[torch.Tensor],
        TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),
    ]

Mistral3PatchMerger

Bases: Module

Learned merging of spatial_merge_size ** 2 patches.

Methods:

• forward_packed –

  Same as forward, gathering blocks with _make_merge_indices

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

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

    def __init__(
        self, vision_hidden_size: int, spatial_merge_size: int, patch_size: int
    ):
        super().__init__()

        self.vision_hidden_size = vision_hidden_size
        self.spatial_merge_size = spatial_merge_size
        self.patch_size = patch_size
        self.merging_layer = nn.Linear(
            vision_hidden_size * self.spatial_merge_size**2,
            vision_hidden_size,
            bias=False,
        )

    def forward(
        self, image_features: torch.Tensor, image_sizes: torch.Tensor
    ) -> torch.Tensor:
        image_sizes = [
            (image_size[0] // self.patch_size, image_size[1] // self.patch_size)
            for image_size in image_sizes
        ]

        tokens_per_image = [h * w for h, w in image_sizes]
        d = image_features.shape[-1]

        permuted_tensor = []
        for image_index, image_tokens in enumerate(
            image_features.split(tokens_per_image)
        ):
            # Reshape image_tokens into a 2D grid
            h, w = image_sizes[image_index]
            image_grid = image_tokens.view(h, w, d).permute(2, 0, 1).unsqueeze(0)
            grid = torch.nn.functional.unfold(
                image_grid,
                kernel_size=self.spatial_merge_size,
                stride=self.spatial_merge_size,
            )
            grid = grid.view(d * self.spatial_merge_size**2, -1).t()
            permuted_tensor.append(grid)

        image_features = torch.cat(permuted_tensor, dim=0)
        image_features = self.merging_layer(image_features)
        return image_features

    def forward_packed(
        self, image_features: torch.Tensor, merge_indices: torch.Tensor | None
    ) -> torch.Tensor:
        """Same as `forward`, gathering blocks with `_make_merge_indices`
        (None when spatial_merge_size is 1)."""
        if merge_indices is not None:
            image_features = _merge_patches_by_index(image_features, merge_indices)
        return self.merging_layer(image_features)

forward_packed(image_features, merge_indices)

Same as forward, gathering blocks with _make_merge_indices (None when spatial_merge_size is 1).

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

def forward_packed(
    self, image_features: torch.Tensor, merge_indices: torch.Tensor | None
) -> torch.Tensor:
    """Same as `forward`, gathering blocks with `_make_merge_indices`
    (None when spatial_merge_size is 1)."""
    if merge_indices is not None:
        image_features = _merge_patches_by_index(image_features, merge_indices)
    return self.merging_layer(image_features)

_get_num_hidden_layers(hf_config)

Determine the number of hidden layers to initialize up to in the visual encoder.

Parameters:

• hf_config

  (Mistral3Config) –

  Model config with vision feature layer(s).

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

def _get_num_hidden_layers(hf_config: Mistral3Config) -> int:
    """Determine the number of hidden layers to initialize up to in the
    visual encoder.

    Args:
        hf_config: Model config with vision feature layer(s).

    """
    feature_layers = hf_config.vision_feature_layer
    num_hidden_layers = hf_config.vision_config.num_hidden_layers
    # If we have one feature layer, initialize up to that layer
    if isinstance(feature_layers, int):
        return get_layer_index(feature_layers, num_hidden_layers)
    # If we have multiple feature layers, initialize up to the deepest one
    elif isinstance(feature_layers, (list, tuple)):
        return max(get_layer_index(idx, num_hidden_layers) for idx in feature_layers)
    raise TypeError(
        f"vision_layer_feature type: {type(feature_layers)} is not supported"
    )