[MM][Mistral3] Image preprocessing optimization (#57531)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jakub Zakrzewski committed
Sep 29, 2026 at 08:53 UTC
d86cf91240988310d04f75217a2ea7e76f1db3d6
6 files changed
+228
-25
docs/design/mm_processing.md
+12
-11
@@ -78,17 +78,18 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:
78
- The flag is **enabled by default** for all models that support it.
79
- Currently, it’s on by default for these architectures:
80
81
-| name | Architecture | Example HF Models |
82
-|--------------------|------------------------------------------|----------------------------------------|
83
-| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` |
84
-| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` |
85
-| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |
86
-| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
87
-| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. |
88
-| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. |
89
-| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. |
90
-| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. |
91
-| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` |
81
+| name | Architecture | Example HF Models |
82
+|--------------------|------------------------------------------|-------------------------------------------------------|
83
+| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` |
84
+| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` |
85
+| `mistral3` | `Mistral3ForConditionalGeneration` | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. |
86
+| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |
87
+| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
88
+| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. |
89
+| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. |
90
+| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. |
91
+| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. |
92
+| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` |
93
94
#### Key Properties and Gains
95
tests/models/multimodal/processing/test_mistral3.py
+128
-2
@@ -4,14 +4,29 @@
4
5
import pytest
6
import torch
7
+import torch.nn.functional as F
8
from PIL import Image
9
from transformers import AutoProcessor, BatchFeature
10
+from transformers.models.pixtral import PixtralProcessor
11
10
-from vllm.model_executor.models.lightonocr import LightOnOCRProcessingInfo
12
+from vllm.model_executor.layers.fusion.mm_input_norm import (
13
+ FusedMMInputNorm,
14
+ IdentityInputNorm,
15
+ build_mm_input_norm,
16
+)
17
+from vllm.model_executor.models.lightonocr import (
18
+ LightOnOCRForConditionalGeneration,
19
+ LightOnOCRProcessingInfo,
20
+)
21
from vllm.model_executor.models.mistral3 import Mistral3HFEncoderInfo
12
-from vllm.model_executor.models.pixtral import PixtralHFEncoderInfo
22
+from vllm.model_executor.models.pixtral import (
23
+ PixtralHFEncoderInfo,
24
+ pixtral_patch_embed,
25
+)
26
from vllm.multimodal import MULTIMODAL_REGISTRY
27
from vllm.multimodal.inputs import MultiModalKwargsItems
28
+from vllm.platforms import current_platform
29
+from vllm.triton_utils import HAS_TRITON
30
31
from ...utils import build_model_context
32
@@ -21,6 +36,65 @@ _MODEL_CONFIG_KWARGS = {"config_format": "hf"}
36
_MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
37
_LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025"
38
39
+_RGB_MEAN = [0.48145466, 0.4578275, 0.40821073]
40
+_RGB_STD = [0.26862954, 0.26130258, 0.27577711]
41
+_RGB_RESCALE = 1.0 / 255.0
42
+
43
+
44
+@pytest.mark.usefixtures("default_vllm_config")
45
+@pytest.mark.skipif(
46
+ not HAS_TRITON or not current_platform.is_cuda(), reason="CUDA Triton kernel"
47
+)
48
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
49
+@pytest.mark.parametrize("normalize", [False, True])
50
+@pytest.mark.parametrize(("channels", "patch_shape"), [(3, (4, 4)), (2, (2, 4))])
51
+def test_pixtral_patch_embed_preserves_image_order_and_strides(
52
+ dtype: torch.dtype,
53
+ normalize: bool,
54
+ channels: int,
55
+ patch_shape: tuple[int, int],
56
+):
57
+ """Packed projection matches CHW normalization and patch convolution."""
58
+ device = torch.device(current_platform.device_type)
59
+ norm = (
60
+ FusedMMInputNorm(
61
+ _RGB_MEAN[:channels],
62
+ _RGB_STD[:channels],
63
+ _RGB_RESCALE,
64
+ channel=channels,
65
+ ).to(device)
66
+ if normalize
67
+ else IdentityInputNorm()
68
+ )
69
+ weight = torch.randn(32, channels, *patch_shape, device=device, dtype=dtype)
70
+ images = [
71
+ torch.randint(0, 256, (channels, 17, 18), device=device, dtype=torch.uint8),
72
+ torch.randint(0, 256, (channels, 25, 17), device=device, dtype=torch.uint8)[
73
+ :, 1:17, 1:17
74
+ ],
75
+ torch.randint(0, 256, (channels, 19, 20), device=device, dtype=torch.uint8)[
76
+ :, 1:17, 2:18
77
+ ],
78
+ ]
79
+ reference = []
80
+ for image in images:
81
+ if normalize:
82
+ normalized = (
83
+ image.float() * norm.weight[:, None, None] + norm.bias[:, None, None]
84
+ )
85
+ else:
86
+ normalized = image
87
+ reference.append(
88
+ F.conv2d(normalized.to(dtype).unsqueeze(0), weight, stride=patch_shape)
89
+ )
90
+ packed, actual = pixtral_patch_embed(images, weight, norm)
91
+ for expected, result in zip(reference, actual):
92
+ torch.testing.assert_close(result, expected, rtol=0.01, atol=0.02)
93
+ expected_packed = torch.cat(
94
+ [result.flatten(2).transpose(1, 2) for result in reference], dim=1
95
+ )
96
+ torch.testing.assert_close(packed, expected_packed, rtol=0.01, atol=0.02)
97
+
98
99
def _process_images_with_hf(
100
hf_processor,
@@ -128,6 +202,9 @@ def test_processor_size_override(
202
)
203
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
204
hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs
205
+ vllm_hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
206
+ assert isinstance(vllm_hf_processor, PixtralProcessor)
207
+
208
hf_processor = AutoProcessor.from_pretrained(
209
model_id,
210
fix_mistral_regex=True,
@@ -153,6 +230,53 @@ def test_processor_size_override(
230
assert prompt_update_tokens == hf_placeholder_tokens
231
232
233
+@pytest.mark.usefixtures("default_vllm_config")
234
+def test_mm_device_do_normalize():
235
+ ctx = build_model_context(
236
+ _MODEL_ID,
237
+ limit_mm_per_prompt={"image": 2},
238
+ model_config_kwargs=_MODEL_CONFIG_KWARGS,
239
+ )
240
+ assert ctx.model_config.multimodal_config.mm_device_do_normalize
241
+
242
+ ctx.model_config.multimodal_config.mm_device_do_normalize = False
243
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
244
+ images = [
245
+ Image.new("RGB", (31, 47), color=(17, 89, 231)),
246
+ Image.new("RGB", (48, 32), color=(201, 13, 127)),
247
+ ]
248
+ prompt = [processor.info.get_hf_config().image_token_index] * len(images)
249
+ mm_items = processor.info.parse_mm_data({"image": images})
250
+
251
+ normalized_inputs = processor(prompt, mm_items=mm_items)
252
+ normalized_values = normalized_inputs["mm_kwargs"].get_data()["pixel_values"]
253
+
254
+ ctx.model_config.multimodal_config.mm_device_do_normalize = True
255
+ raw_inputs = processor(prompt, mm_items=mm_items)
256
+ raw_values = raw_inputs["mm_kwargs"].get_data()["pixel_values"]
257
+ assert all(value.dtype == torch.uint8 for value in raw_values)
258
+
259
+ input_norm = build_mm_input_norm(ctx.model_config)
260
+ patch_size = processor.info.get_hf_config().vision_config.patch_size
261
+
262
+ def pack_patches(image: torch.Tensor) -> torch.Tensor:
263
+ channels, height, width = image.shape
264
+ rows = height // patch_size
265
+ cols = width // patch_size
266
+ return (
267
+ image[:, : rows * patch_size, : cols * patch_size]
268
+ .reshape(channels, rows, patch_size, cols, patch_size)
269
+ .permute(1, 3, 0, 2, 4)
270
+ .reshape(-1, channels * patch_size**2)
271
+ )
272
+
273
+ for raw, normalized in zip(raw_values, normalized_values):
274
+ output = input_norm(pack_patches(raw), normalized.dtype)
275
+ torch.testing.assert_close(
276
+ output, pack_patches(normalized), rtol=1e-5, atol=1e-6
277
+ )
278
+
279
+
280
def test_scoped_request_size_overrides_configured_flat_size():
281
ctx = build_model_context(
282
_MODEL_ID,
@@ -188,6 +312,8 @@ def test_scoped_request_size_overrides_configured_flat_size():
312
313
314
def test_lightonocr_keeps_vision_config_image_size():
315
+ assert not LightOnOCRForConditionalGeneration.supports_mm_device_do_normalize
316
+
317
ctx = build_model_context(
318
_LIGHTON_MODEL_ID,
319
mm_processor_kwargs={"size": {"longest_edge": 1008}},
vllm/model_executor/layers/fusion/mm_input_norm.py
+5
@@ -147,6 +147,11 @@ def _load_norm_params(model_config: ModelConfig) -> NormParams:
147
rescale_factor)`` from the processor config."""
148
config = get_processor_config(model_config.model, revision=model_config.revision)
149
config = config.get("image_processor", config)
150
+ if not any(
151
+ key in config
152
+ for key in ("do_normalize", "do_rescale", "image_mean", "image_std")
153
+ ):
154
+ config = model_config.hf_image_processor_config
155
156
has_norm_params = "image_mean" in config and "image_std" in config
157
do_normalize = bool(config.get("do_normalize", has_norm_params))
vllm/model_executor/models/lightonocr.py
+2
@@ -115,6 +115,8 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessing
115
dummy_inputs=Mistral3DummyInputsBuilder,
116
)
117
class LightOnOCRForConditionalGeneration(Mistral3ForConditionalGeneration):
118
+ supports_mm_device_do_normalize = False
119
+
120
hf_to_vllm_mapper = WeightsMapper(
121
orig_to_new_prefix={
122
"model.vision_encoder.": "vision_tower.",
vllm/model_executor/models/mistral3.py
+19
-3
@@ -7,6 +7,8 @@ from typing import Annotated, Any, Literal, cast
7
8
import torch
9
import torch.nn as nn
10
+import transformers
11
+from packaging.version import Version
12
from transformers import BatchFeature, Mistral3Config, PixtralVisionConfig
13
from transformers.models.pixtral import PixtralProcessor
14
@@ -14,6 +16,7 @@ from vllm.config import VllmConfig
16
from vllm.config.multimodal import MultiModalDummyOptions
17
from vllm.inputs import MultiModalDataDict
18
from vllm.model_executor.layers.activation import get_act_fn
19
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
22
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -56,6 +59,10 @@ from .utils import (
59
maybe_prefix,
60
)
61
62
+TRANSFORMERS_SUPPORTS_PIXTRAL_IMAGE_ONLY = Version(transformers.__version__) >= Version(
63
+ "5.15.0"
64
+)
65
+
66
67
class Mistral3ImagePixelInputs(TensorSchema):
68
"""Dimensions:
@@ -225,7 +232,7 @@ class Mistral3ProcessingInfo(BaseProcessingInfo):
232
image_size = size["longest_edge"]
233
return Mistral3HFEncoderInfo(self.get_hf_config(), image_size)
234
228
- def get_hf_processor(self, **kwargs: object):
235
+ def get_hf_processor(self, **kwargs: object) -> PixtralProcessor:
236
return self.ctx.get_hf_processor(PixtralProcessor, **kwargs)
237
238
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
@@ -277,8 +284,12 @@ class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder):
284
285
286
class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]):
280
- def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
281
- return self.dummy_inputs.get_dummy_text(mm_counts)
287
+ def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str | None:
288
+ # PixtralProcessor supports image-only calls starting in transformers 5.15.
289
+ # Older releases need dummy text and take the slower tokenization path.
290
+ if not TRANSFORMERS_SUPPORTS_PIXTRAL_IMAGE_ONLY:
291
+ return self.dummy_inputs.get_dummy_text(mm_counts)
292
+ return None
293
294
def _postprocess_hf_mm_data(
295
self,
@@ -378,6 +389,7 @@ def init_vision_tower_for_mistral3(
389
hf_config: Mistral3Config,
390
quant_config: QuantizationConfig | None,
391
*,
392
+ input_norm: nn.Module | None = None,
393
require_post_norm: bool | None = None,
394
prefix: str = "",
395
) -> PixtralHFVisionModel:
@@ -391,6 +403,7 @@ def init_vision_tower_for_mistral3(
403
return PixtralHFVisionModel(
404
vision_config,
405
quant_config=quant_config,
406
+ input_norm=input_norm,
407
num_hidden_layers_override=num_hidden_layers,
408
require_post_norm=require_post_norm,
409
prefix=prefix,
@@ -410,6 +423,8 @@ class Mistral3ForConditionalGeneration(
423
SupportsEagle,
424
SupportsEagle3,
425
):
426
+ supports_mm_device_do_normalize = True
427
+
428
packed_modules_mapping = {
429
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
430
"gate_up_proj": ["gate_proj", "up_proj"],
@@ -464,6 +479,7 @@ class Mistral3ForConditionalGeneration(
479
self.vision_tower = init_vision_tower_for_mistral3(
480
config,
481
quant_config=quant_config,
482
+ input_norm=build_mm_input_norm(vllm_config.model_config),
483
require_post_norm=False,
484
prefix=maybe_prefix(prefix, "vision_tower"),
485
)
vllm/model_executor/models/pixtral.py
+62
-9
@@ -10,6 +10,7 @@ from typing import Annotated, Literal
10
import numpy as np
11
import torch
12
import torch.nn as nn
13
+import torch.nn.functional as F
14
from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk
15
from mistral_common.protocol.instruct.messages import UserMessage
16
from mistral_common.protocol.instruct.request import ChatCompletionRequest
@@ -1427,12 +1428,55 @@ class PixtralHFTransformer(nn.Module):
1428
return x
1429
1430
1431
+def pixtral_patch_embed(
1432
+ images: list[torch.Tensor],
1433
+ weight: torch.Tensor,
1434
+ input_norm: nn.Module,
1435
+) -> tuple[torch.Tensor, list[torch.Tensor]]:
1436
+ """Normalize and project all CHW images in one matrix multiplication."""
1437
+ out_channels, in_channels, patch_height, patch_width = weight.shape
1438
+ grid_sizes = []
1439
+ for image in images:
1440
+ assert image.ndim == 3 and image.shape[0] == in_channels
1441
+ _, height, width = image.shape
1442
+ patch_rows = height // patch_height
1443
+ patch_cols = width // patch_width
1444
+ grid_sizes.append((patch_rows, patch_cols))
1445
+
1446
+ patch_counts = [rows * cols for rows, cols in grid_sizes]
1447
+ packed = torch.empty(
1448
+ (sum(patch_counts), in_channels * patch_height * patch_width),
1449
+ dtype=images[0].dtype,
1450
+ device=images[0].device,
1451
+ )
1452
+ offset = 0
1453
+ for image, (rows, cols), count in zip(images, grid_sizes, patch_counts):
1454
+ source = (
1455
+ image[:, : rows * patch_height, : cols * patch_width]
1456
+ .reshape(in_channels, rows, patch_height, cols, patch_width)
1457
+ .permute(1, 3, 0, 2, 4)
1458
+ )
1459
+ packed[offset : offset + count].view(
1460
+ rows, cols, in_channels, patch_height, patch_width
1461
+ ).copy_(source)
1462
+ offset += count
1463
+
1464
+ normalized = input_norm(packed, weight.dtype)
1465
+ projected = F.linear(normalized, weight.reshape(out_channels, -1))
1466
+ embeddings = [
1467
+ tokens.transpose(0, 1).reshape(1, out_channels, rows, cols)
1468
+ for tokens, (rows, cols) in zip(projected.split(patch_counts), grid_sizes)
1469
+ ]
1470
+ return projected.unsqueeze(0), embeddings
1471
+
1472
+
1473
class PixtralHFVisionModel(nn.Module):
1474
def __init__(
1475
self,
1476
config: PixtralVisionConfig,
1477
quant_config: QuantizationConfig | None = None,
1478
*,
1479
+ input_norm: nn.Module | None = None,
1480
num_hidden_layers_override: int | None = None,
1481
require_post_norm: bool | None = None,
1482
prefix: str = "",
@@ -1440,6 +1484,7 @@ class PixtralHFVisionModel(nn.Module):
1484
super().__init__()
1485
1486
self.config = config
1487
+ self.input_norm = input_norm
1488
1489
self.patch_conv = Conv2dLayer(
1490
in_channels=config.num_channels,
@@ -1495,16 +1540,24 @@ class PixtralHFVisionModel(nn.Module):
1540
all tokens of all images of shape (N_toks, D)
1541
1542
"""
1498
- # pass images through initial convolution independently
1499
- patch_embeds_list = [
1500
- self.patch_conv(img.unsqueeze(0).to(self.dtype)) for img in pixel_values
1501
- ]
1502
-
1503
- patch_embeds = [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list]
1504
- embed_sizes = [p.shape[1] for p in patch_embeds]
1543
+ images = [image.to(device=self.device) for image in pixel_values]
1544
+ if self.input_norm is not None:
1545
+ patch_embeds, patch_embeds_list = pixtral_patch_embed(
1546
+ images,
1547
+ self.patch_conv.weight,
1548
+ self.input_norm,
1549
+ )
1550
+ else:
1551
+ patch_embeds_list = []
1552
+ for image in images:
1553
+ patch_embeds_list.append(
1554
+ self.patch_conv(image.unsqueeze(0).to(self.dtype))
1555
+ )
1556
+ patch_embeds = torch.cat(
1557
+ [p.flatten(2).permute(0, 2, 1) for p in patch_embeds_list], dim=1
1558
+ )
1559
1506
- # flatten to a single sequence
1507
- patch_embeds = torch.cat(patch_embeds, dim=1)
1560
+ embed_sizes = [p.shape[-2] * p.shape[-1] for p in patch_embeds_list]
1561
patch_embeds = self.ln_pre(patch_embeds)
1562
1563
# positional embeddings