From d86cf91240988310d04f75217a2ea7e76f1db3d6 Mon Sep 17 00:00:00 2001 From: Jakub Zakrzewski Date: Tue, 29 Sep 2026 08:53:38 +0200 Subject: [PATCH] [MM][Mistral3] Image preprocessing optimization (#57531) Signed-off-by: Jakub Zakrzewski Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/design/mm_processing.md | 23 ++-- .../multimodal/processing/test_mistral3.py | 130 +++++++++++++++++- .../layers/fusion/mm_input_norm.py | 5 + vllm/model_executor/models/lightonocr.py | 2 + vllm/model_executor/models/mistral3.py | 22 ++- vllm/model_executor/models/pixtral.py | 71 ++++++++-- 6 files changed, 228 insertions(+), 25 deletions(-) diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md index d350836ca8..f53f7d2dd1 100644 --- a/docs/design/mm_processing.md +++ b/docs/design/mm_processing.md @@ -78,17 +78,18 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`: - The flag is **enabled by default** for all models that support it. - Currently, it’s on by default for these architectures: -| name | Architecture | Example HF Models | -|--------------------|------------------------------------------|----------------------------------------| -| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` | -| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` | -| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. | -| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. | -| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. | -| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | -| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. | -| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. | -| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` | +| name | Architecture | Example HF Models | +|--------------------|------------------------------------------|-------------------------------------------------------| +| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` | +| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` | +| `mistral3` | `Mistral3ForConditionalGeneration` | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. | +| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. | +| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. | +| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. | +| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | +| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. | +| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. | +| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` | #### Key Properties and Gains diff --git a/tests/models/multimodal/processing/test_mistral3.py b/tests/models/multimodal/processing/test_mistral3.py index 3445dd86fe..87ca8c842b 100644 --- a/tests/models/multimodal/processing/test_mistral3.py +++ b/tests/models/multimodal/processing/test_mistral3.py @@ -4,14 +4,29 @@ import pytest import torch +import torch.nn.functional as F from PIL import Image from transformers import AutoProcessor, BatchFeature +from transformers.models.pixtral import PixtralProcessor -from vllm.model_executor.models.lightonocr import LightOnOCRProcessingInfo +from vllm.model_executor.layers.fusion.mm_input_norm import ( + FusedMMInputNorm, + IdentityInputNorm, + build_mm_input_norm, +) +from vllm.model_executor.models.lightonocr import ( + LightOnOCRForConditionalGeneration, + LightOnOCRProcessingInfo, +) from vllm.model_executor.models.mistral3 import Mistral3HFEncoderInfo -from vllm.model_executor.models.pixtral import PixtralHFEncoderInfo +from vllm.model_executor.models.pixtral import ( + PixtralHFEncoderInfo, + pixtral_patch_embed, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalKwargsItems +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON from ...utils import build_model_context @@ -21,6 +36,65 @@ _MODEL_CONFIG_KWARGS = {"config_format": "hf"} _MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" _LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025" +_RGB_MEAN = [0.48145466, 0.4578275, 0.40821073] +_RGB_STD = [0.26862954, 0.26130258, 0.27577711] +_RGB_RESCALE = 1.0 / 255.0 + + +@pytest.mark.usefixtures("default_vllm_config") +@pytest.mark.skipif( + not HAS_TRITON or not current_platform.is_cuda(), reason="CUDA Triton kernel" +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("normalize", [False, True]) +@pytest.mark.parametrize(("channels", "patch_shape"), [(3, (4, 4)), (2, (2, 4))]) +def test_pixtral_patch_embed_preserves_image_order_and_strides( + dtype: torch.dtype, + normalize: bool, + channels: int, + patch_shape: tuple[int, int], +): + """Packed projection matches CHW normalization and patch convolution.""" + device = torch.device(current_platform.device_type) + norm = ( + FusedMMInputNorm( + _RGB_MEAN[:channels], + _RGB_STD[:channels], + _RGB_RESCALE, + channel=channels, + ).to(device) + if normalize + else IdentityInputNorm() + ) + weight = torch.randn(32, channels, *patch_shape, device=device, dtype=dtype) + images = [ + torch.randint(0, 256, (channels, 17, 18), device=device, dtype=torch.uint8), + torch.randint(0, 256, (channels, 25, 17), device=device, dtype=torch.uint8)[ + :, 1:17, 1:17 + ], + torch.randint(0, 256, (channels, 19, 20), device=device, dtype=torch.uint8)[ + :, 1:17, 2:18 + ], + ] + reference = [] + for image in images: + if normalize: + normalized = ( + image.float() * norm.weight[:, None, None] + norm.bias[:, None, None] + ) + else: + normalized = image + reference.append( + F.conv2d(normalized.to(dtype).unsqueeze(0), weight, stride=patch_shape) + ) + packed, actual = pixtral_patch_embed(images, weight, norm) + for expected, result in zip(reference, actual): + torch.testing.assert_close(result, expected, rtol=0.01, atol=0.02) + expected_packed = torch.cat( + [result.flatten(2).transpose(1, 2) for result in reference], dim=1 + ) + torch.testing.assert_close(packed, expected_packed, rtol=0.01, atol=0.02) + def _process_images_with_hf( hf_processor, @@ -128,6 +202,9 @@ def test_processor_size_override( ) processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs + vllm_hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs) + assert isinstance(vllm_hf_processor, PixtralProcessor) + hf_processor = AutoProcessor.from_pretrained( model_id, fix_mistral_regex=True, @@ -153,6 +230,53 @@ def test_processor_size_override( assert prompt_update_tokens == hf_placeholder_tokens +@pytest.mark.usefixtures("default_vllm_config") +def test_mm_device_do_normalize(): + ctx = build_model_context( + _MODEL_ID, + limit_mm_per_prompt={"image": 2}, + model_config_kwargs=_MODEL_CONFIG_KWARGS, + ) + assert ctx.model_config.multimodal_config.mm_device_do_normalize + + ctx.model_config.multimodal_config.mm_device_do_normalize = False + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + images = [ + Image.new("RGB", (31, 47), color=(17, 89, 231)), + Image.new("RGB", (48, 32), color=(201, 13, 127)), + ] + prompt = [processor.info.get_hf_config().image_token_index] * len(images) + mm_items = processor.info.parse_mm_data({"image": images}) + + normalized_inputs = processor(prompt, mm_items=mm_items) + normalized_values = normalized_inputs["mm_kwargs"].get_data()["pixel_values"] + + ctx.model_config.multimodal_config.mm_device_do_normalize = True + raw_inputs = processor(prompt, mm_items=mm_items) + raw_values = raw_inputs["mm_kwargs"].get_data()["pixel_values"] + assert all(value.dtype == torch.uint8 for value in raw_values) + + input_norm = build_mm_input_norm(ctx.model_config) + patch_size = processor.info.get_hf_config().vision_config.patch_size + + def pack_patches(image: torch.Tensor) -> torch.Tensor: + channels, height, width = image.shape + rows = height // patch_size + cols = width // patch_size + return ( + image[:, : rows * patch_size, : cols * patch_size] + .reshape(channels, rows, patch_size, cols, patch_size) + .permute(1, 3, 0, 2, 4) + .reshape(-1, channels * patch_size**2) + ) + + for raw, normalized in zip(raw_values, normalized_values): + output = input_norm(pack_patches(raw), normalized.dtype) + torch.testing.assert_close( + output, pack_patches(normalized), rtol=1e-5, atol=1e-6 + ) + + def test_scoped_request_size_overrides_configured_flat_size(): ctx = build_model_context( _MODEL_ID, @@ -188,6 +312,8 @@ def test_scoped_request_size_overrides_configured_flat_size(): def test_lightonocr_keeps_vision_config_image_size(): + assert not LightOnOCRForConditionalGeneration.supports_mm_device_do_normalize + ctx = build_model_context( _LIGHTON_MODEL_ID, mm_processor_kwargs={"size": {"longest_edge": 1008}}, diff --git a/vllm/model_executor/layers/fusion/mm_input_norm.py b/vllm/model_executor/layers/fusion/mm_input_norm.py index f03cfa4c36..5127f1007e 100644 --- a/vllm/model_executor/layers/fusion/mm_input_norm.py +++ b/vllm/model_executor/layers/fusion/mm_input_norm.py @@ -147,6 +147,11 @@ def _load_norm_params(model_config: ModelConfig) -> NormParams: rescale_factor)`` from the processor config.""" config = get_processor_config(model_config.model, revision=model_config.revision) config = config.get("image_processor", config) + if not any( + key in config + for key in ("do_normalize", "do_rescale", "image_mean", "image_std") + ): + config = model_config.hf_image_processor_config has_norm_params = "image_mean" in config and "image_std" in config do_normalize = bool(config.get("do_normalize", has_norm_params)) diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py index d8a1240242..bc0292ed55 100644 --- a/vllm/model_executor/models/lightonocr.py +++ b/vllm/model_executor/models/lightonocr.py @@ -115,6 +115,8 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessing dummy_inputs=Mistral3DummyInputsBuilder, ) class LightOnOCRForConditionalGeneration(Mistral3ForConditionalGeneration): + supports_mm_device_do_normalize = False + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "model.vision_encoder.": "vision_tower.", diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 593ba90ec0..40af706e3a 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -7,6 +7,8 @@ from typing import Annotated, Any, Literal, cast import torch import torch.nn as nn +import transformers +from packaging.version import Version from transformers import BatchFeature, Mistral3Config, PixtralVisionConfig from transformers.models.pixtral import PixtralProcessor @@ -14,6 +16,7 @@ from vllm.config import VllmConfig from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -56,6 +59,10 @@ from .utils import ( maybe_prefix, ) +TRANSFORMERS_SUPPORTS_PIXTRAL_IMAGE_ONLY = Version(transformers.__version__) >= Version( + "5.15.0" +) + class Mistral3ImagePixelInputs(TensorSchema): """Dimensions: @@ -225,7 +232,7 @@ class Mistral3ProcessingInfo(BaseProcessingInfo): image_size = size["longest_edge"] return Mistral3HFEncoderInfo(self.get_hf_config(), image_size) - def get_hf_processor(self, **kwargs: object): + def get_hf_processor(self, **kwargs: object) -> PixtralProcessor: return self.ctx.get_hf_processor(PixtralProcessor, **kwargs) def get_supported_mm_limits(self) -> Mapping[str, int | None]: @@ -277,8 +284,12 @@ class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder): class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]): - def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: - return self.dummy_inputs.get_dummy_text(mm_counts) + def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str | None: + # PixtralProcessor supports image-only calls starting in transformers 5.15. + # Older releases need dummy text and take the slower tokenization path. + if not TRANSFORMERS_SUPPORTS_PIXTRAL_IMAGE_ONLY: + return self.dummy_inputs.get_dummy_text(mm_counts) + return None def _postprocess_hf_mm_data( self, @@ -378,6 +389,7 @@ def init_vision_tower_for_mistral3( hf_config: Mistral3Config, quant_config: QuantizationConfig | None, *, + input_norm: nn.Module | None = None, require_post_norm: bool | None = None, prefix: str = "", ) -> PixtralHFVisionModel: @@ -391,6 +403,7 @@ def init_vision_tower_for_mistral3( return PixtralHFVisionModel( vision_config, quant_config=quant_config, + input_norm=input_norm, num_hidden_layers_override=num_hidden_layers, require_post_norm=require_post_norm, prefix=prefix, @@ -410,6 +423,8 @@ class Mistral3ForConditionalGeneration( SupportsEagle, SupportsEagle3, ): + supports_mm_device_do_normalize = True + packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -464,6 +479,7 @@ class Mistral3ForConditionalGeneration( 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"), ) diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index 91ac38e110..521df62a92 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -10,6 +10,7 @@ from typing import Annotated, Literal import numpy as np import torch import torch.nn as nn +import torch.nn.functional as F from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk from mistral_common.protocol.instruct.messages import UserMessage from mistral_common.protocol.instruct.request import ChatCompletionRequest @@ -1427,12 +1428,55 @@ class PixtralHFTransformer(nn.Module): return x +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 + + 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 = "", @@ -1440,6 +1484,7 @@ class PixtralHFVisionModel(nn.Module): super().__init__() self.config = config + self.input_norm = input_norm self.patch_conv = Conv2dLayer( in_channels=config.num_channels, @@ -1495,16 +1540,24 @@ class PixtralHFVisionModel(nn.Module): 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 pixel_values - ] + 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 + ) - 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) + embed_sizes = [p.shape[-2] * p.shape[-1] for p in patch_embeds_list] patch_embeds = self.ln_pre(patch_embeds) # positional embeddings