@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 51d7c6a2b2
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 16:00 UTC

[Model] Support Mistral3 in the HF Transformers format (#15505)

Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Michael Goin committed Apr 1, 2025 at 07:10 UTC 51d7c6a2b23e100cd9e7d85b8e7c0eea656b331e
9 files changed +723 -4
docs/source/models/supported_models.md
+7
@@ -865,6 +865,13 @@ See [this page](#generative-models) for more information on how to use generativ
865 * ✅︎
866 * ✅︎
867 * ✅︎
868 +- * `Mistral3ForConditionalGeneration`
869 + * Mistral3
870 + * T + I<sup>+</sup>
871 + * `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc.
872 + *
873 + * ✅︎
874 + *
875 - * `MllamaForConditionalGeneration`
876 * Llama 3.2
877 * T + I<sup>+</sup>
examples/offline_inference/vision_language.py
+24
@@ -498,6 +498,29 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
498 return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")
499
500
501 +# Mistral-3 HF-format
502 +def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
503 + assert modality == "image"
504 +
505 + model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
506 +
507 + # NOTE: Need L40 (or equivalent) to avoid OOM
508 + engine_args = EngineArgs(
509 + model=model_name,
510 + max_model_len=8192,
511 + max_num_seqs=2,
512 + tensor_parallel_size=2,
513 + disable_mm_preprocessor_cache=args.disable_mm_preprocessor_cache,
514 + )
515 +
516 + prompts = [f"<s>[INST]{question}\n[IMG][/INST]" for question in questions]
517 +
518 + return ModelRequestData(
519 + engine_args=engine_args,
520 + prompts=prompts,
521 + )
522 +
523 +
524 # LLama 3.2
525 def run_mllama(questions: list[str], modality: str) -> ModelRequestData:
526 assert modality == "image"
@@ -859,6 +882,7 @@ model_example_map = {
882 "mantis": run_mantis,
883 "minicpmo": run_minicpmo,
884 "minicpmv": run_minicpmv,
885 + "mistral3": run_mistral3,
886 "mllama": run_mllama,
887 "molmo": run_molmo,
888 "NVLM_D": run_nvlm_d,
examples/offline_inference/vision_language_multi_image.py
+23
@@ -218,6 +218,28 @@ def load_internvl(question: str, image_urls: list[str]) -> ModelRequestData:
218 )
219
220
221 +def load_mistral3(question: str, image_urls: list[str]) -> ModelRequestData:
222 + model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
223 +
224 + # Adjust this as necessary to fit in GPU
225 + engine_args = EngineArgs(
226 + model=model_name,
227 + max_model_len=8192,
228 + max_num_seqs=2,
229 + tensor_parallel_size=2,
230 + limit_mm_per_prompt={"image": len(image_urls)},
231 + )
232 +
233 + placeholders = "[IMG]" * len(image_urls)
234 + prompt = f"<s>[INST]{question}\n{placeholders}[/INST]"
235 +
236 + return ModelRequestData(
237 + engine_args=engine_args,
238 + prompt=prompt,
239 + image_data=[fetch_image(url) for url in image_urls],
240 + )
241 +
242 +
243 def load_mllama(question: str, image_urls: list[str]) -> ModelRequestData:
244 model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct"
245
@@ -509,6 +531,7 @@ model_example_map = {
531 "h2ovl_chat": load_h2ovl,
532 "idefics3": load_idefics3,
533 "internvl_chat": load_internvl,
534 + "mistral3": load_mistral3,
535 "mllama": load_mllama,
536 "NVLM_D": load_nvlm_d,
537 "phi3_v": load_phi3v,
tests/models/registry.py
+3
@@ -297,6 +297,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
297 "MiniCPMV": _HfExamplesInfo("openbmb/MiniCPM-Llama3-V-2_5",
298 extras={"2.6": "openbmb/MiniCPM-V-2_6"}, # noqa: E501
299 trust_remote_code=True),
300 + "Mistral3ForConditionalGeneration": _HfExamplesInfo("mistralai/Mistral-Small-3.1-24B-Instruct-2503", # noqa: E501
301 + min_transformers_version="4.50", # noqa: E501
302 + extras={"fp8": "nm-testing/Mistral-Small-3.1-24B-Instruct-2503-FP8-dynamic"}), # noqa: E501
303 "MolmoForCausalLM": _HfExamplesInfo("allenai/Molmo-7B-D-0924",
304 max_transformers_version="4.48",
305 transformers_version_reason="Use of private method which no longer exists.", # noqa: E501
vllm/entrypoints/chat_utils.py
+2 -1
@@ -487,7 +487,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
487 return "<|endoftext10|>" # 200010 (see vocab.json in hf model)
488 if model_type in ("minicpmo", "minicpmv"):
489 return "(<image>./</image>)"
490 - if model_type in ("blip-2", "fuyu", "paligemma", "pixtral"):
490 + if model_type in ("blip-2", "fuyu", "paligemma", "pixtral",
491 + "mistral3"):
492 # These models do not use image tokens in the prompt
493 return None
494 if model_type == "qwen":
vllm/model_executor/models/mistral3.py new
+656
@@ -0,0 +1,656 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from abc import abstractmethod
4 +from collections.abc import Iterable, Mapping, Sequence
5 +from functools import cached_property
6 +from typing import (Final, Literal, Optional, Protocol, Set, Tuple, TypedDict,
7 + TypeVar, Union)
8 +
9 +import torch
10 +import torch.nn as nn
11 +from transformers import (BatchFeature, Mistral3Config, PixtralVisionConfig,
12 + PretrainedConfig)
13 +from transformers.models.pixtral import PixtralProcessor
14 +
15 +from vllm.config import VllmConfig
16 +from vllm.inputs import InputProcessingContext
17 +from vllm.model_executor.layers.activation import get_act_fn
18 +from vllm.model_executor.layers.layernorm import RMSNorm
19 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
20 + RowParallelLinear)
21 +from vllm.model_executor.layers.quantization import QuantizationConfig
22 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
23 +from vllm.model_executor.sampling_metadata import SamplingMetadata
24 +from vllm.multimodal import MULTIMODAL_REGISTRY
25 +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargs
26 +from vllm.multimodal.parse import (ImageProcessorItems, ImageSize,
27 + MultiModalDataItems)
28 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
29 + BaseProcessingInfo, ProcessingCache,
30 + PromptReplacement, PromptUpdate)
31 +from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
32 +from vllm.sequence import IntermediateTensors
33 +
34 +from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,
35 + SupportsV0Only)
36 +from .pixtral import PixtralHFEncoderInfo, PixtralHFVisionModel
37 +from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
38 + maybe_prefix, merge_multimodal_embeddings)
39 +from .vision import get_vision_encoder_info, select_patch_features
40 +
41 +
42 +class Mistral3ImagePixelInputs(TypedDict):
43 + type: Literal["pixel_values_pixtral"]
44 + pixel_values: Union[torch.Tensor, list[torch.Tensor]]
45 + """
46 + Shape: `(batch_size * num_images, num_channels, height, width)`
47 +
48 + Note that `height` or `width` may be different per batch and image,
49 + in which case the data is passed as a list instead of a batched tensor.
50 + """
51 +
52 + embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]
53 + """
54 + A boolean mask indicating which image embeddings correspond
55 + to patch tokens.
56 +
57 + Shape: `(batch_size, num_images, num_embeds)`
58 + """
59 +
60 +
61 +class Mistral3PatchMerger(nn.Module):
62 + """
63 + Learned merging of spatial_merge_size ** 2 patches
64 + """
65 +
66 + def __init__(self, vision_hidden_size: int, spatial_merge_size: int,
67 + patch_size: int):
68 + super().__init__()
69 +
70 + self.vision_hidden_size = vision_hidden_size
71 + self.spatial_merge_size = spatial_merge_size
72 + self.patch_size = patch_size
73 + self.merging_layer = nn.Linear(vision_hidden_size *
74 + self.spatial_merge_size**2,
75 + vision_hidden_size,
76 + bias=False)
77 +
78 + def forward(self, image_features: torch.Tensor,
79 + image_sizes: torch.Tensor) -> torch.Tensor:
80 + image_sizes = [(image_size[0] // self.patch_size,
81 + image_size[1] // self.patch_size)
82 + for image_size in image_sizes]
83 +
84 + tokens_per_image = [h * w for h, w in image_sizes]
85 + d = image_features.shape[-1]
86 +
87 + permuted_tensor = []
88 + for image_index, image_tokens in enumerate(
89 + image_features.split(tokens_per_image)):
90 + # Reshape image_tokens into a 2D grid
91 + h, w = image_sizes[image_index]
92 + image_grid = image_tokens.view(h, w, d).permute(2, 0,
93 + 1).unsqueeze(0)
94 + grid = torch.nn.functional.unfold(
95 + image_grid,
96 + kernel_size=self.spatial_merge_size,
97 + stride=self.spatial_merge_size)
98 + grid = grid.view(d * self.spatial_merge_size**2, -1).t()
99 + permuted_tensor.append(grid)
100 +
101 + image_features = torch.cat(permuted_tensor, dim=0)
102 + image_features = self.merging_layer(image_features)
103 + return image_features
104 +
105 +
106 +class Mistral3MultiModalProjector(nn.Module):
107 +
108 + def __init__(self,
109 + vision_hidden_size: int,
110 + text_hidden_size: int,
111 + spatial_merge_size: int,
112 + patch_size: int,
113 + projector_hidden_act: str,
114 + multimodal_projector_bias: bool,
115 + quant_config: Optional[QuantizationConfig] = None,
116 + prefix: str = ""):
117 + super().__init__()
118 +
119 + self.norm = RMSNorm(vision_hidden_size, eps=1e-5)
120 + self.patch_merger = Mistral3PatchMerger(
121 + vision_hidden_size=vision_hidden_size,
122 + spatial_merge_size=spatial_merge_size,
123 + patch_size=patch_size)
124 +
125 + self.linear_1 = ColumnParallelLinear(vision_hidden_size,
126 + text_hidden_size,
127 + bias=multimodal_projector_bias,
128 + quant_config=quant_config,
129 + prefix=f"{prefix}.linear_1")
130 + self.act = get_act_fn(projector_hidden_act)
131 + self.linear_2 = RowParallelLinear(text_hidden_size,
132 + text_hidden_size,
133 + bias=multimodal_projector_bias,
134 + quant_config=quant_config,
135 + prefix=f"{prefix}.linear_2")
136 +
137 + def forward(self, image_features: torch.Tensor,
138 + image_sizes: torch.Tensor) -> torch.Tensor:
139 + image_features = self.norm(image_features)
140 + image_features = self.patch_merger(image_features, image_sizes)
141 + hidden_states, _ = self.linear_1(image_features)
142 + hidden_states = self.act(hidden_states)
143 + hidden_states, _ = self.linear_2(hidden_states)
144 + return hidden_states
145 +
146 +
147 +class LlavaLikeConfig(Protocol):
148 + vision_config: Final[PretrainedConfig]
149 + image_token_index: Final[int]
150 + vision_feature_select_strategy: Final[str]
151 + vision_feature_layer: Final[Union[int, list[int]]]
152 +
153 +
154 +class LlavaLikeProcessor(Protocol):
155 + image_token: Final[str]
156 +
157 +
158 +class BaseLlavaProcessingInfo(BaseProcessingInfo):
159 +
160 + def get_hf_config(self) -> LlavaLikeConfig:
161 + return self.ctx.get_hf_config(Mistral3Config)
162 +
163 + def get_vision_encoder_info(self):
164 + return get_vision_encoder_info(self.get_hf_config())
165 +
166 + @abstractmethod
167 + def get_hf_processor(self, **kwargs: object) -> LlavaLikeProcessor:
168 + raise NotImplementedError
169 +
170 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
171 + return {"image": None}
172 +
173 + def get_mm_max_tokens_per_item(
174 + self,
175 + seq_len: int,
176 + mm_counts: Mapping[str, int],
177 + ) -> Mapping[str, int]:
178 + return {"image": self.get_max_image_tokens()}
179 +
180 + def get_num_image_tokens(
181 + self,
182 + *,
183 + image_width: int,
184 + image_height: int,
185 + ) -> int:
186 + vision_encoder_info = self.get_vision_encoder_info()
187 + return vision_encoder_info.get_num_image_tokens(
188 + image_width=image_width,
189 + image_height=image_height,
190 + )
191 +
192 + def get_image_size_with_most_features(self) -> ImageSize:
193 + vision_encoder_info = self.get_vision_encoder_info()
194 + width = height = vision_encoder_info.get_image_size()
195 + return ImageSize(width=width, height=height)
196 +
197 + def get_max_image_tokens(self) -> int:
198 + target_width, target_height = self.get_image_size_with_most_features()
199 +
200 + return self.get_num_image_tokens(
201 + image_width=target_width,
202 + image_height=target_height,
203 + )
204 +
205 +
206 +_I = TypeVar("_I", bound=BaseLlavaProcessingInfo)
207 +
208 +
209 +class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder[_I]):
210 +
211 + def get_dummy_processor_inputs(
212 + self,
213 + seq_len: int,
214 + mm_counts: Mapping[str, int],
215 + ) -> ProcessorInputs:
216 + num_images = mm_counts.get("image", 0)
217 +
218 + processor = self.info.get_hf_processor()
219 + image_token = processor.image_token
220 + target_width, target_height = \
221 + self.info.get_image_size_with_most_features()
222 +
223 + mm_data = {
224 + "image":
225 + self._get_dummy_images(width=target_width,
226 + height=target_height,
227 + num_images=num_images)
228 + }
229 +
230 + return ProcessorInputs(
231 + prompt_text=image_token * num_images,
232 + mm_data=mm_data,
233 + )
234 +
235 +
236 +class Mistral3ProcessingInfo(BaseLlavaProcessingInfo):
237 +
238 + def get_hf_processor(self, **kwargs: object):
239 + return self.ctx.get_hf_processor(PixtralProcessor, **kwargs)
240 +
241 +
242 +class Mistral3MultiModalProcessor(
243 + BaseMultiModalProcessor[Mistral3ProcessingInfo]):
244 +
245 + def _call_hf_processor(
246 + self,
247 + prompt: str,
248 + mm_data: Mapping[str, object],
249 + mm_kwargs: Mapping[str, object],
250 + ) -> BatchFeature:
251 + processed_outputs = super()._call_hf_processor(
252 + prompt=prompt,
253 + mm_data=mm_data,
254 + mm_kwargs=mm_kwargs,
255 + )
256 +
257 + pixel_values = processed_outputs.get("pixel_values")
258 + if pixel_values is not None:
259 +
260 + # Avoid padding since we need the output for each image to be
261 + # independent of other images for the cache to work correctly
262 + image_sizes = processed_outputs["image_sizes"]
263 + assert len(pixel_values) == len(image_sizes)
264 +
265 + processed_outputs["pixel_values"] = [
266 + p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)
267 + ]
268 +
269 + hf_config = self.info.get_hf_config()
270 + vision_config = hf_config.vision_config
271 + assert isinstance(vision_config, PixtralVisionConfig)
272 + encoder_info = PixtralHFEncoderInfo(vision_config)
273 +
274 + tile_sizes = [
275 + encoder_info.get_patch_grid_size(
276 + image_width=pixel_value.shape[-1],
277 + image_height=pixel_value.shape[-2],
278 + ) for pixel_value in processed_outputs["pixel_values"]
279 + ]
280 + embed_is_patch = [
281 + torch.tensor(([True] * ncols + [False]) * nrows)
282 + for ncols, nrows in tile_sizes
283 + ]
284 + processed_outputs["embed_is_patch"] = embed_is_patch
285 +
286 + return processed_outputs
287 +
288 + def _get_mm_fields_config(
289 + self,
290 + hf_inputs: BatchFeature,
291 + hf_processor_mm_kwargs: Mapping[str, object],
292 + ) -> Mapping[str, MultiModalFieldConfig]:
293 + return dict(
294 + pixel_values=MultiModalFieldConfig.batched("image"),
295 + embed_is_patch=MultiModalFieldConfig.batched("image"),
296 + image_embeds=MultiModalFieldConfig.batched("image"),
297 + )
298 +
299 + def _get_prompt_updates(
300 + self,
301 + mm_items: MultiModalDataItems,
302 + hf_processor_mm_kwargs: Mapping[str, object],
303 + out_mm_kwargs: MultiModalKwargs,
304 + ) -> Sequence[PromptUpdate]:
305 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
306 + hf_config = self.info.get_hf_config()
307 + tokenizer = self.info.get_tokenizer()
308 + vocab = tokenizer.get_vocab()
309 +
310 + image_break_id = vocab[processor.image_break_token]
311 + image_token_id = hf_config.image_token_index
312 + image_end_id = vocab[processor.image_end_token]
313 +
314 + vision_config = hf_config.vision_config
315 + assert isinstance(vision_config, PixtralVisionConfig)
316 + encoder_info = PixtralHFEncoderInfo(vision_config)
317 +
318 + def get_replacement(item_idx: int):
319 + images = mm_items.get_items("image", ImageProcessorItems)
320 + image_size = images.get_image_size(item_idx)
321 +
322 + ncols, nrows = encoder_info.get_patch_grid_size(
323 + image_width=image_size.width,
324 + image_height=image_size.height,
325 + )
326 +
327 + tokens = ([image_token_id] * ncols + [image_break_id]) * nrows
328 + tokens[-1] = image_end_id
329 +
330 + return tokens
331 +
332 + return [
333 + PromptReplacement(
334 + modality="image",
335 + target=[image_token_id],
336 + replacement=get_replacement,
337 + ),
338 + ]
339 +
340 +
341 +def _build_mistral3_info(
342 + ctx: InputProcessingContext, ) -> BaseLlavaProcessingInfo:
343 + hf_config = ctx.get_hf_config(Mistral3Config)
344 + assert isinstance(hf_config.vision_config, PixtralVisionConfig)
345 + return Mistral3ProcessingInfo(ctx)
346 +
347 +
348 +def _build_mistral3_processor(
349 + info: _I,
350 + dummy_inputs: BaseDummyInputsBuilder[_I],
351 + *,
352 + cache: Optional[ProcessingCache] = None,
353 + enable_sanity_checks: bool = True,
354 +) -> BaseMultiModalProcessor:
355 + assert isinstance(info, Mistral3ProcessingInfo)
356 + return Mistral3MultiModalProcessor(
357 + info,
358 + dummy_inputs, # type: ignore
359 + cache=cache,
360 + enable_sanity_checks=enable_sanity_checks,
361 + )
362 +
363 +
364 +def _get_num_hidden_layers(hf_config: LlavaLikeConfig) -> int:
365 + """Determine the number of hidden layers to initialize up to in the
366 + visual encoder.
367 +
368 + Args:
369 + hf_config: Model config with vision feature layer(s).
370 + """
371 + feature_layers = hf_config.vision_feature_layer
372 + num_hidden_layers = hf_config.vision_config.num_hidden_layers
373 + # If we have one feature layer, initialize up to that layer
374 + if isinstance(feature_layers, int):
375 + return _get_layer_index(feature_layers, num_hidden_layers)
376 + # If we have multiple feature layers, initialize up to the deepest one
377 + elif isinstance(feature_layers, (list, tuple)):
378 + return max(
379 + _get_layer_index(idx, num_hidden_layers) for idx in feature_layers)
380 + raise TypeError(f"vision_layer_feature type: {type(feature_layers)}"
381 + " is not supported")
382 +
383 +
384 +def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:
385 + """Given a signed vision feature layer, get the number of hidden layers
386 + needed to leverage it.
387 +
388 + Args:
389 + feature_layer_index: Index of a required layer in the visual encoder.
390 + num_hidden_layers: The total number of hidden layers in the visual
391 + encoder.
392 + """
393 + if feature_layer_index < 0:
394 + return num_hidden_layers + feature_layer_index + 1
395 + return feature_layer_index
396 +
397 +
398 +def init_vision_tower_for_llava(
399 + hf_config: LlavaLikeConfig,
400 + quant_config: Optional[QuantizationConfig],
401 + *,
402 + require_post_norm: Optional[bool] = None,
403 + prefix: str = "",
404 +) -> PixtralHFVisionModel:
405 + vision_config = hf_config.vision_config
406 +
407 + # Initialize the vision tower only up to the deepest required feature layer
408 + num_hidden_layers = _get_num_hidden_layers(hf_config)
409 +
410 + assert isinstance(vision_config, PixtralVisionConfig)
411 +
412 + return PixtralHFVisionModel(
413 + vision_config,
414 + quant_config=quant_config,
415 + num_hidden_layers_override=num_hidden_layers,
416 + require_post_norm=require_post_norm,
417 + prefix=prefix,
418 + )
419 +
420 +
421 +# TODO(mgoin): Support V1, there are issues with image batching/chunking
422 +# that need to be resolved first.
423 +@MULTIMODAL_REGISTRY.register_processor(
424 + _build_mistral3_processor,
425 + info=_build_mistral3_info,
426 + dummy_inputs=Mistral3DummyInputsBuilder)
427 +class Mistral3ForConditionalGeneration(nn.Module, SupportsMultiModal,
428 + SupportsPP, SupportsV0Only):
429 +
430 + packed_modules_mapping = {
431 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
432 + "gate_up_proj": ["gate_proj", "up_proj"]
433 + }
434 +
435 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
436 + super().__init__()
437 +
438 + config = vllm_config.model_config.hf_config
439 + quant_config = vllm_config.quant_config
440 + multimodal_config = vllm_config.model_config.multimodal_config
441 +
442 + self.config = config
443 + self.multimodal_config = multimodal_config
444 +
445 + # NOTE: These are special cases for Pixtral-12B in the HF-format
446 + # https://huggingface.co/mistral-community/pixtral-12b/blob/main/config.json # noqa
447 + if (config.text_config.architectures is None
448 + and config.text_config.model_type == "mistral"):
449 + config.text_config.architectures = ["MistralForCausalLM"]
450 + if (config.projector_hidden_act is None
451 + and config.vision_config.hidden_act == "gelu"):
452 + config.projector_hidden_act = "gelu"
453 +
454 + # TODO: Optionally initializes this for supporting embeddings.
455 + self.vision_tower = init_vision_tower_for_llava(
456 + config,
457 + quant_config,
458 + require_post_norm=False,
459 + prefix=maybe_prefix(prefix, "vision_tower"))
460 + self.multi_modal_projector = Mistral3MultiModalProjector(
461 + vision_hidden_size=config.vision_config.hidden_size,
462 + text_hidden_size=config.text_config.hidden_size,
463 + projector_hidden_act=config.projector_hidden_act,
464 + spatial_merge_size=config.spatial_merge_size,
465 + patch_size=config.vision_config.patch_size,
466 + multimodal_projector_bias=config.multimodal_projector_bias,
467 + quant_config=quant_config,
468 + prefix=maybe_prefix(prefix, "multi_modal_projector"))
469 +
470 + self.language_model = init_vllm_registered_model(
471 + vllm_config=vllm_config,
472 + hf_config=config.text_config,
473 + prefix=maybe_prefix(prefix, "language_model"),
474 + )
475 +
476 + self.make_empty_intermediate_tensors = (
477 + self.language_model.make_empty_intermediate_tensors)
478 +
479 + @cached_property
480 + def sampler(self):
481 + if hasattr(self.language_model, "sampler"):
482 + return self.language_model.sampler
483 +
484 + return get_sampler()
485 +
486 + def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
487 + h = w = self.config.vision_config.image_size
488 + expected_dims = (3, h, w)
489 + actual_dims = tuple(data.shape[1:])
490 +
491 + if actual_dims != expected_dims:
492 + expected_expr = ("batch_size", *map(str, expected_dims))
493 + raise ValueError(
494 + f"The expected shape of pixel values is {expected_expr}. "
495 + f"You supplied {tuple(data.shape)}.")
496 +
497 + return data
498 +
499 + def _parse_and_validate_image_input(
500 + self, **kwargs: object) -> Optional[Mistral3ImagePixelInputs]:
501 + pixel_values = kwargs.pop("pixel_values", None)
502 + image_embeds = kwargs.pop("image_embeds", None)
503 +
504 + if pixel_values is None and image_embeds is None:
505 + return None
506 +
507 + assert pixel_values is not None
508 + if not isinstance(pixel_values, (torch.Tensor, list)):
509 + raise ValueError("Incorrect type of pixel values. "
510 + f"Got type: {type(pixel_values)}")
511 +
512 + assert self.config.vision_config.model_type == "pixtral"
513 + embed_is_patch = kwargs.pop("embed_is_patch")
514 + if not isinstance(embed_is_patch, (torch.Tensor, list)):
515 + raise ValueError("Incorrect type of embed_is_patch. "
516 + f"Got type: {type(embed_is_patch)}")
517 +
518 + return Mistral3ImagePixelInputs(
519 + type="pixel_values_pixtral",
520 + pixel_values=flatten_bn(pixel_values),
521 + embed_is_patch=embed_is_patch,
522 + )
523 +
524 + def _process_image_input(
525 + self,
526 + image_input: Mistral3ImagePixelInputs,
527 + ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
528 + if image_input["type"] == "image_embeds":
529 + return image_input["data"]
530 +
531 + image_sizes = [(img.shape[-2], img.shape[-1])
532 + for img in image_input["pixel_values"]]
533 +
534 + image_features = self.vision_tower(image_input["pixel_values"])
535 +
536 + if isinstance(image_features, torch.Tensor):
537 + return self.multi_modal_projector(image_features, image_sizes)
538 +
539 + feature_sizes = [
540 + image_feature.shape[0] // self.config.spatial_merge_size**2
541 + for image_feature in image_features
542 + ]
543 +
544 + image_embeds = self.multi_modal_projector(torch.cat(image_features),
545 + image_sizes)
546 + if len(feature_sizes) > 1:
547 + image_embeds = torch.split(image_embeds, feature_sizes)
548 + else:
549 + image_embeds = (image_embeds, )
550 + return image_embeds
551 +
552 + def get_multimodal_embeddings(
553 + self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
554 + image_input = self._parse_and_validate_image_input(**kwargs)
555 + if image_input is None:
556 + return None
557 +
558 + vision_embeddings = self._process_image_input(image_input)
559 +
560 + return vision_embeddings
561 +
562 + def get_input_embeddings(
563 + self,
564 + input_ids: torch.Tensor,
565 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
566 + ) -> torch.Tensor:
567 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
568 + if multimodal_embeddings is not None:
569 + inputs_embeds = merge_multimodal_embeddings(
570 + input_ids,
571 + inputs_embeds,
572 + select_patch_features(multimodal_embeddings),
573 + self.config.image_token_index,
574 + )
575 + return inputs_embeds
576 +
577 + def forward(
578 + self,
579 + input_ids: torch.Tensor,
580 + positions: torch.Tensor,
581 + intermediate_tensors: Optional[IntermediateTensors] = None,
582 + inputs_embeds: Optional[torch.Tensor] = None,
583 + **kwargs: object,
584 + ) -> Union[torch.Tensor, IntermediateTensors]:
585 + """Run forward pass for Mistral3.
586 +
587 + One key thing to understand is the `input_ids` already accounts for the
588 + positions of the to-be-inserted image embeddings.
589 +
590 + Concretely, consider a text prompt:
591 + `"USER: <image>\\nWhat's the content of the image?\\nASSISTANT:"`.
592 +
593 + Tokenizer outputs:
594 + `[1, 3148, 1001, 29901, 29871, 32000, 29871, 13, 5618, 29915, 29879,
595 + 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566, 29901]`.
596 +
597 + To reserve space in KV cache, we have to insert placeholder tokens
598 + before they are inputted to the model, so the input processor prepends
599 + additional image tokens (denoted as `32000`), resulting in:
600 + `[1, 3148, 1001, 29901, 29871, 32000, ..., 32000, 29871, 13, 5618,
601 + 29915, 29879, 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566,
602 + 29901]`.
603 +
604 + We insert 575 tokens so that including the original image token in the
605 + input, there are a total of 576 (24 * 24) image tokens, which
606 + corresponds to the number of image tokens inputted to the language
607 + model, i.e. the number of image tokens outputted by the visual encoder.
608 +
609 + This way, the `positions` and `attn_metadata` are consistent
610 + with the `input_ids`.
611 +
612 + Args:
613 + input_ids: Flattened (concatenated) input_ids corresponding to a
614 + batch.
615 + pixel_values: The pixels in each input image.
616 +
617 + See also:
618 + :class:`Mistral3ImagePixelInputs`
619 + """
620 + if intermediate_tensors is not None:
621 + inputs_embeds = None
622 +
623 + # NOTE: In v1, inputs_embeds is always generated at model runner, this
624 + # condition is for v0 compatibility.
625 + elif inputs_embeds is None:
626 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)
627 + inputs_embeds = self.get_input_embeddings(input_ids,
628 + vision_embeddings)
629 + input_ids = None
630 +
631 + hidden_states = self.language_model.model(input_ids,
632 + positions,
633 + intermediate_tensors,
634 + inputs_embeds=inputs_embeds)
635 +
636 + return hidden_states
637 +
638 + def compute_logits(
639 + self,
640 + hidden_states: torch.Tensor,
641 + sampling_metadata: SamplingMetadata,
642 + ) -> Optional[torch.Tensor]:
643 + return self.language_model.compute_logits(hidden_states,
644 + sampling_metadata)
645 +
646 + def sample(
647 + self,
648 + logits: torch.Tensor,
649 + sampling_metadata: SamplingMetadata,
650 + ) -> Optional[SamplerOutput]:
651 + return self.language_model.sample(logits, sampling_metadata)
652 +
653 + def load_weights(self, weights: Iterable[Tuple[str,
654 + torch.Tensor]]) -> Set[str]:
655 + loader = AutoWeightsLoader(self)
656 + return loader.load_weights(weights)
vllm/model_executor/models/pixtral.py
+4 -3
@@ -979,7 +979,8 @@ class PixtralHFEncoderInfo(VisionEncoderInfo[PixtralVisionConfig]):
979 return self.vision_config.image_size
980
981 def get_patch_size(self) -> int:
982 - return self.vision_config.patch_size
982 + return (self.vision_config.patch_size *
983 + self.vision_config.spatial_merge_size)
984
985 def get_patch_grid_length(self) -> int:
986 image_size, patch_size = self.get_image_size(), self.get_patch_size()
@@ -1001,8 +1002,8 @@ class PixtralHFEncoderInfo(VisionEncoderInfo[PixtralVisionConfig]):
1002 ratio = max(image_width / max_width, image_height / max_height)
1003
1004 if ratio > 1:
1004 - image_width = int(math.ceil(image_width / ratio))
1005 - image_height = int(math.ceil(image_height / ratio))
1005 + image_width = int(math.floor(image_width / ratio))
1006 + image_height = int(math.floor(image_height / ratio))
1007
1008 nrows, ncols = _get_pixtral_hf_num_image_tokens(
1009 (image_height, image_width),
vllm/model_executor/models/registry.py
+1
@@ -177,6 +177,7 @@ _MULTIMODAL_MODELS = {
177 "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"), # noqa: E501
178 "MiniCPMO": ("minicpmo", "MiniCPMO"),
179 "MiniCPMV": ("minicpmv", "MiniCPMV"),
180 + "Mistral3ForConditionalGeneration": ("mistral3", "Mistral3ForConditionalGeneration"), # noqa: E501
181 "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
182 "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
183 "PaliGemmaForConditionalGeneration": ("paligemma", "PaliGemmaForConditionalGeneration"), # noqa: E501
vllm/model_executor/models/vision.py
+3
@@ -69,6 +69,9 @@ def get_vision_encoder_info(
69 if isinstance(vision_config, CLIPVisionConfig):
70 return CLIPEncoderInfo(vision_config)
71 if isinstance(vision_config, PixtralVisionConfig):
72 + # Need to sneak in spatial_merge_size for Mistral3
73 + vision_config.spatial_merge_size = getattr(hf_config,
74 + "spatial_merge_size", 1)
75 return PixtralHFEncoderInfo(vision_config)
76 if isinstance(vision_config, SiglipVisionConfig):
77 return SiglipEncoderInfo(vision_config)