iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[MM][Mistral3] Image preprocessing optimization (#57531)

Browse Source 
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>
This commit is contained in:
Jakub Zakrzewski
2026-09-29 06:53:38 +00:00
committed by GitHub
co-authored by OpenAI Codex OpenAI Codex mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 65ed0a2642
commit d86cf91240
6 changed files with 228 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/mm_processing.md
+12 -11
View File
@@ -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
tests/models/multimodal/processing/test_mistral3.py
+128 -2
View File
@@ -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}},
vllm/model_executor/layers/fusion/mm_input_norm.py
+5
View File
@@ -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))
vllm/model_executor/models/lightonocr.py
+2
View File
@@ -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.",
vllm/model_executor/models/mistral3.py
+19 -3
View File
@@ -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"),
)
vllm/model_executor/models/pixtral.py
+62 -9
View File
@@ -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