diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md index 530e1f903284..2198fb107f1a 100644 --- a/docs/design/mm_processing.md +++ b/docs/design/mm_processing.md @@ -80,6 +80,7 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`: | name | Architecture | Example HF Models | | -------------------------- | ------------------------------------------ | ----------------------------------------------------- | +| `gemma-3` | `Gemma3ForConditionalGeneration` | `google/gemma-3-4b-it`, etc. | | `glm-4v` | `Glm4vForConditionalGeneration` | `zai-org/GLM-4.1V-9B-Thinking`, etc. | | `glm-5-next` | `Glm5NextForConditionalGeneration` | `zai-org/GLM-5.3-Flash` | | `glm-ocr` | `GlmOcrForConditionalGeneration` | `zai-org/GLM-OCR` | diff --git a/tests/model_executor/layers/test_mm_input_norm.py b/tests/model_executor/layers/test_mm_input_norm.py index 97a715728048..ebedb78f6c3d 100644 --- a/tests/model_executor/layers/test_mm_input_norm.py +++ b/tests/model_executor/layers/test_mm_input_norm.py @@ -3,6 +3,8 @@ """Tests for FusedMMInputNorm and the fused affine transform kernel.""" +from types import SimpleNamespace +from typing import get_args from unittest.mock import Mock import pytest @@ -19,7 +21,9 @@ from vllm.model_executor.layers.fusion.mm_input_norm import ( FusedMMInputNorm, IdentityInputNorm, + InputLayout, NormParams, + _validate_layout, fused_mm_input_norm_triton, ) from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader @@ -48,7 +52,7 @@ def _reference_input_norm( channel: int, out_dtype: torch.dtype, ) -> torch.Tensor: - """Straightforward per-channel affine: (x * rescale - mean) / std.""" + """Straightforward per-channel affine for (patches, C*L).""" patches, size = pixel_values.shape patch_size = size // channel mean = torch.tensor( @@ -62,6 +66,28 @@ def _reference_input_norm( return x.view(patches, size).to(out_dtype) +def _reference_input_norm_nchw( + pixel_values: torch.Tensor, + image_mean: list[float], + image_std: list[float], + rescale_factor: float, + channel: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Straightforward per-channel affine for (B, C, H, W).""" + B, C, H, W = pixel_values.shape + assert channel == C, f"reference got C={C}, expected channel={channel}" + mean = torch.tensor( + image_mean, dtype=torch.float32, device=pixel_values.device + ).view(1, C, 1, 1) + std = torch.tensor(image_std, dtype=torch.float32, device=pixel_values.device).view( + 1, C, 1, 1 + ) + x = pixel_values.to(torch.float32) + x = (x * rescale_factor - mean) / std + return x.to(out_dtype) + + # CLIP-style RGB normalization constants reused across tests. _RGB_MEAN = [0.48145466, 0.4578275, 0.40821073] _RGB_STD = [0.26862954, 0.26130258, 0.27577711] @@ -75,6 +101,7 @@ def _check_against_reference( image_mean: list[float] | None = None, image_std: list[float] | None = None, rescale_factor: float = _RGB_RESCALE, + layout: InputLayout = "pcl", **close_kwargs, ): """Run ``FusedMMInputNorm`` on ``pixel_values`` and compare against the @@ -83,15 +110,25 @@ def _check_against_reference( image_mean = list(_RGB_MEAN) if image_std is None: image_std = list(_RGB_STD) + layout = _validate_layout(layout) + norm = FusedMMInputNorm( image_mean=image_mean, image_std=image_std, rescale_factor=rescale_factor, channel=channel, + layout=layout, ).to(_DEVICE) - expected = _reference_input_norm( - pixel_values, image_mean, image_std, rescale_factor, channel, out_dtype - ) + + if layout == "bchw": + expected = _reference_input_norm_nchw( + pixel_values, image_mean, image_std, rescale_factor, channel, out_dtype + ) + else: + expected = _reference_input_norm( + pixel_values, image_mean, image_std, rescale_factor, channel, out_dtype + ) + torch.testing.assert_close( norm(pixel_values, visual_dtype=out_dtype), expected, **close_kwargs ) @@ -204,6 +241,12 @@ def test_block_boundaries(self, patch_size: int): ) _check_against_reference(pixel_values, torch.float32) + def test_rejects_indivisible_pcl_size(self): + """Size % channel != 0 must fail loudly, not silently truncate.""" + norm = FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE).to(_DEVICE) + with pytest.raises(AssertionError): + norm(torch.randn(2, 3 * 10 + 1, device=_DEVICE), torch.float32) + # =========================================================================== # Input handling: non-contiguous inputs @@ -221,6 +264,201 @@ def test_non_contiguous_input_matches_reference(self): _check_against_reference(non_contig, torch.float32) +# =========================================================================== +# BCHW layout +# =========================================================================== +@requires_vllm_config +class TestFusedMMInputNormBCHWLayout: + @pytest.mark.parametrize( + "in_dtype,out_dtype", + [ + (torch.float32, torch.float32), + (torch.float32, torch.bfloat16), + (torch.float32, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.bfloat16, torch.float32), + (torch.float16, torch.float16), + (torch.uint8, torch.bfloat16), + (torch.uint8, torch.float32), + ], + ) + def test_dtype_combinations_bchw( + self, in_dtype: torch.dtype, out_dtype: torch.dtype + ): + set_random_seed(0) + shape = (4, 3, 8, 16) + if in_dtype == torch.uint8: + pixel_values = torch.randint( + 0, 256, shape, dtype=torch.uint8, device=_DEVICE + ) + else: + pixel_values = torch.rand(shape, dtype=torch.float32, device=_DEVICE).to( + in_dtype + ) + + _check_against_reference( + pixel_values, + out_dtype, + layout="bchw", + image_mean=[0.5, 0.5, 0.5], + image_std=[0.5, 0.5, 0.5], + atol=1e-2, + rtol=1e-2, + ) + + @pytest.mark.parametrize("channel", [1, 3, 4]) + def test_channel_variants_bchw(self, channel: int): + set_random_seed(0) + pixel_values = torch.randint( + 0, 256, (4, channel, 16, 16), dtype=torch.float32, device=_DEVICE + ) + _check_against_reference( + pixel_values, + torch.float32, + channel=channel, + layout="bchw", + image_mean=[0.5] * channel, + image_std=[0.25] * channel, + ) + + @pytest.mark.parametrize("hw", [(1, 1), (2, 8), (32, 64)]) + def test_spatial_shapes_bchw(self, hw: tuple[int, int]): + """H*W=1 exercises L=1; 32*64=2048 is exactly two blocks.""" + H, W = hw + set_random_seed(0) + pixel_values = torch.randint( + 0, 256, (4, 3, H, W), dtype=torch.float32, device=_DEVICE + ) + _check_against_reference(pixel_values, torch.float32, layout="bchw") + + def test_bchw_matches_pcl_equivalent(self): + """The same bytes viewed as (P, C, H*W) and (B, C, H, W) must + produce numerically identical outputs.""" + set_random_seed(0) + B, C, H, W = 4, 3, 16, 32 + x_bchw = torch.randint( + 0, 256, (B, C, H, W), dtype=torch.float32, device=_DEVICE + ) + x_pcl = x_bchw.reshape(B, C * H * W) + + norm_bchw = FusedMMInputNorm( + _RGB_MEAN, _RGB_STD, _RGB_RESCALE, layout="bchw" + ).to(_DEVICE) + norm_pcl = FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE, layout="pcl").to( + _DEVICE + ) + + y_bchw = norm_bchw(x_bchw, torch.float32) + y_pcl = norm_pcl(x_pcl, torch.float32) + + assert y_bchw.shape == (B, C, H, W) + torch.testing.assert_close(y_bchw.reshape(B, C * H * W), y_pcl) + + def test_large_batch_bchw(self): + """Regression guard for the cuDNN grid-limit size, now in BCHW.""" + set_random_seed(0) + pixel_values = torch.randint( + 0, 256, (70000, 3, 1, 1), dtype=torch.float32, device=_DEVICE + ) + _check_against_reference(pixel_values, torch.float32, layout="bchw") + + +# =========================================================================== +# BCHW validation / non-contiguous +# =========================================================================== +@requires_vllm_config +class TestFusedMMInputNormBCHWValidation: + def _norm(self, channel: int = 3) -> FusedMMInputNorm: + return FusedMMInputNorm( + _RGB_MEAN, _RGB_STD, _RGB_RESCALE, channel=channel, layout="bchw" + ).to(_DEVICE) + + def test_rejects_2d_input(self): + norm = self._norm() + with pytest.raises(AssertionError): + norm(torch.randn(4, 3 * 64, device=_DEVICE), torch.float32) + + def test_rejects_5d_input(self): + norm = self._norm() + with pytest.raises(AssertionError): + norm(torch.randn(2, 3, 4, 4, 4, device=_DEVICE), torch.float32) + + def test_rejects_channel_mismatch(self): + norm = self._norm(channel=3) + with pytest.raises(AssertionError): + norm(torch.randn(2, 4, 8, 8, device=_DEVICE), torch.float32) + + def test_non_contiguous_bchw_matches_reference(self): + """A permuted (still 4D) tensor must be materialized correctly.""" + set_random_seed(0) + base = torch.randint( + 0, 256, (4, 3, 16, 32), dtype=torch.float32, device=_DEVICE + ) + # Transpose to (4, 3, 32, 16), make contiguous, transpose back so + # the resulting (4, 3, 16, 32) view is non-contiguous. + non_contig = base.transpose(2, 3).contiguous().transpose(2, 3) + assert not non_contig.is_contiguous() + + _check_against_reference(non_contig, torch.float32, layout="bchw") + + +# =========================================================================== +# layout constructor semantics +# =========================================================================== +@requires_vllm_config +class TestFusedMMInputInputLayoutConfig: + def test_default_is_pcl(self): + norm = FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE) + assert norm.layout == "pcl" + + @pytest.mark.parametrize("layout", ["pcl", "bchw"]) + def test_accepts_valid_string(self, layout): + norm = FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE, layout=layout) + assert norm.layout == layout + + def test_invalid_layout_raises(self): + with pytest.raises(ValueError): + FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE, layout="nchw") + + def test_layout_literal_members(self): + # The Literal alias must remain exactly the two supported layouts so + # config plumbing that passes strings keeps working. + assert set(get_args(InputLayout)) == {"pcl", "bchw"} + + def test_input_dtype_and_enabled(self): + # Sanity check for the CustomOp contract used by the processor. + norm = FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE) + assert norm.input_dtype is torch.uint8 + assert FusedMMInputNorm.enabled() is True + + +# =========================================================================== +# forward_oot / IdentityInputNorm 4D passthrough +# =========================================================================== +@requires_vllm_config +class TestFusedMMInputNormOOTAndIdentity: + @pytest.mark.parametrize("layout", ["pcl", "bchw"]) + def test_forward_oot_matches_native(self, layout): + set_random_seed(0) + shape = (4, 3, 8, 8) if layout == "bchw" else (4, 3 * 64) + pixel_values = torch.randn(shape, dtype=torch.float32, device=_DEVICE) + + norm = FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE, layout=layout).to( + _DEVICE + ) + torch.testing.assert_close( + norm.forward_oot(pixel_values, torch.bfloat16), + norm.forward_native(pixel_values, torch.bfloat16), + ) + + def test_identity_passthrough_bchw(self): + norm = IdentityInputNorm() + x = torch.randn(2, 3, 8, 8, dtype=torch.float32, device=_DEVICE) + out = norm(x, visual_dtype=torch.bfloat16) + assert out.shape == x.shape + torch.testing.assert_close(out, x.to(torch.bfloat16)) + + # =========================================================================== # Raw kernel entry point # =========================================================================== @@ -228,7 +466,7 @@ def test_non_contiguous_input_matches_reference(self): @requires_triton class TestFusedMMInputNormKernel: @pytest.mark.parametrize( - "N, C, L", + "P, C, L", [ (1, 1, 1), (1, 3, 1), @@ -237,12 +475,12 @@ class TestFusedMMInputNormKernel: (4, 3, 1000), ], ) - def test_channel_boundary_crossing(self, N: int, C: int, L: int): + def test_channel_boundary_crossing(self, P: int, C: int, L: int): """The 1D kernel recovers ``c = (offs // L) % C`` per lane, so a tile can straddle a channel boundary and gather weight/bias per lane. These shapes force that path by keeping ``L`` well below the block.""" set_random_seed(0) - x = torch.randn(N, C, L, dtype=torch.float32, device=_DEVICE) + x = torch.randn(P, C, L, dtype=torch.float32, device=_DEVICE) w = torch.randn(C, dtype=torch.float32, device=_DEVICE) b = torch.randn(C, dtype=torch.float32, device=_DEVICE) @@ -254,15 +492,15 @@ def test_channel_boundary_crossing(self, N: int, C: int, L: int): def test_non_contiguous_inputs_materialized(self): """Read-only tensors are made contiguous internally.""" - N, C, L = 3, 3, 100 + P, C, L = 3, 3, 100 set_random_seed(0) - base = torch.randn(N, C, L, 2, dtype=torch.float32, device=_DEVICE) + base = torch.randn(P, C, L, 2, dtype=torch.float32, device=_DEVICE) x = base[..., 0] assert not x.is_contiguous() w = torch.randn(C, dtype=torch.float32, device=_DEVICE) b = torch.randn(C, dtype=torch.float32, device=_DEVICE) - out = torch.empty(N, C, L, dtype=torch.float32, device=_DEVICE) + out = torch.empty(P, C, L, dtype=torch.float32, device=_DEVICE) fused_mm_input_norm_triton(x, out, w, b) expected = x * w.view(1, C, 1) + b.view(1, C, 1) @@ -270,12 +508,12 @@ def test_non_contiguous_inputs_materialized(self): def test_rejects_mismatched_output_shape(self): """``outputs`` must be shaped exactly like ``inputs``.""" - N, C, L = 3, 3, 100 - x = torch.randn(N, C, L, dtype=torch.float32, device=_DEVICE) + P, C, L = 3, 3, 100 + x = torch.randn(P, C, L, dtype=torch.float32, device=_DEVICE) w = torch.randn(C, dtype=torch.float32, device=_DEVICE) b = torch.randn(C, dtype=torch.float32, device=_DEVICE) - for out_shape in ((N, C + 1, L), (N, C, L + 1), (N + 2, C, L)): + for out_shape in ((P, C + 1, L), (P, C, L + 1), (P + 2, C, L)): with pytest.raises(AssertionError): fused_mm_input_norm_triton( x, @@ -364,3 +602,66 @@ def test_buffers_initialized_on_default_device( assert norm.weight.device.type == default_device assert norm.bias.device.type == default_device + + +# =========================================================================== +# Factory: build_mm_input_norm +# =========================================================================== +@requires_vllm_config +class TestBuildMMInputNorm: + """``build_mm_input_norm`` picks IdentityInputNorm vs FusedMMInputNorm + and forwards ``layout`` to the constructed module.""" + + @staticmethod + def _patch_processor(monkeypatch, cfg: dict): + from vllm.model_executor.layers.fusion import mm_input_norm as mod + + monkeypatch.setattr(mod, "get_processor_config", lambda *a, **k: dict(cfg)) + return mod + + @staticmethod + def _model_config(do_normalize: bool = True) -> SimpleNamespace: + return SimpleNamespace( + multimodal_config=SimpleNamespace(mm_device_do_normalize=do_normalize), + model="dummy", + revision=None, + ) + + _NON_IDENTITY_CFG = { + "do_normalize": True, + "do_rescale": True, + "image_mean": _RGB_MEAN, + "image_std": _RGB_STD, + "rescale_factor": _RGB_RESCALE, + } + + _IDENTITY_CFG = { + "do_normalize": True, + "do_rescale": True, + "image_mean": [0.0, 0.0, 0.0], + "image_std": [1.0, 1.0, 1.0], + "rescale_factor": 1.0, + } + + @pytest.mark.parametrize("layout", ["pcl", "bchw"]) + def test_layout_passthrough(self, monkeypatch, layout): + mod = self._patch_processor(monkeypatch, self._NON_IDENTITY_CFG) + norm = mod.build_mm_input_norm(self._model_config(), layout=layout) + assert isinstance(norm, FusedMMInputNorm) + assert norm.layout == layout + + def test_default_layout_is_pcl(self, monkeypatch): + mod = self._patch_processor(monkeypatch, self._NON_IDENTITY_CFG) + norm = mod.build_mm_input_norm(self._model_config()) + assert isinstance(norm, FusedMMInputNorm) + assert norm.layout == "pcl" + + def test_identity_params_return_identity_norm(self, monkeypatch): + mod = self._patch_processor(monkeypatch, self._IDENTITY_CFG) + norm = mod.build_mm_input_norm(self._model_config(), layout="bchw") + assert isinstance(norm, IdentityInputNorm) + + def test_normalize_disabled_returns_identity_norm(self, monkeypatch): + mod = self._patch_processor(monkeypatch, {}) + norm = mod.build_mm_input_norm(self._model_config(do_normalize=False)) + assert isinstance(norm, IdentityInputNorm) diff --git a/tests/models/multimodal/generation_ppl_test/gemma.py b/tests/models/multimodal/generation_ppl_test/gemma.py new file mode 100644 index 000000000000..aeff52e0c470 --- /dev/null +++ b/tests/models/multimodal/generation_ppl_test/gemma.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from tests.models.utils import GenerateModelInfo + +from .ppl_utils import vqa_ppl_test + +GEMMA_PPL_TOL = 0.02 + +MODELS = [ + GenerateModelInfo( + "google/gemma-3-4b-it", image_token_id=262144, hf_ppl=119.61958312988281 + ), +] + + +@pytest.mark.parametrize("model_info", MODELS) +@pytest.mark.parametrize("mm_device_do_normalize", [False, True]) +def test_ppl( + hf_runner, + vllm_runner, + model_info: GenerateModelInfo, + mm_device_do_normalize: bool, +): + vqa_ppl_test( + hf_runner, + vllm_runner, + model_info, + vllm_extra_kwargs={ + "mm_device_do_normalize": mm_device_do_normalize, + }, + tol=GEMMA_PPL_TOL, + ) + + +@pytest.mark.parametrize("model_info", MODELS) +@pytest.mark.parametrize("mm_device_do_normalize", [False, True]) +def test_ppl_cg( + hf_runner, + vllm_runner, + model_info: GenerateModelInfo, + mm_device_do_normalize: bool, +): + vqa_ppl_test( + hf_runner, + vllm_runner, + model_info, + vllm_extra_kwargs={ + "mm_device_do_normalize": mm_device_do_normalize, + "compilation_config": {"cudagraph_mm_encoder": True}, + }, + tol=GEMMA_PPL_TOL, + ) diff --git a/tests/models/multimodal/generation_ppl_test/ppl_utils.py b/tests/models/multimodal/generation_ppl_test/ppl_utils.py index d19764740995..edb5f0415bf1 100644 --- a/tests/models/multimodal/generation_ppl_test/ppl_utils.py +++ b/tests/models/multimodal/generation_ppl_test/ppl_utils.py @@ -34,12 +34,12 @@ def vqa_ppl_test( mm_processor_kwargs = {} processor = AutoProcessor.from_pretrained(model_info.name) - image_token_id = processor.image_token_id + image_token_id = model_info.image_token_id or processor.image_token_id images = [] prompts = [] for row in dataset: - image = row["image"] + image = _ensure_rgb(row["image"]) images.append(image) question = row["question"] @@ -101,6 +101,7 @@ def vqa_ppl_test( assert token_data is not None assert len(token_data) == 1 token_id, token_log_prob = next(iter(token_data.items())) + if token_id == image_token_id: continue token_log_probs.append(token_log_prob.logprob) @@ -114,7 +115,8 @@ def vqa_ppl_test( vllm_dtype = model_config.dtype head_dtype = model_config.head_dtype - # Accelerate ppl test by setting Transformers ppl score to a constant + # Accelerate the PPL test by setting the Transformers PPL score to a + # constant, if provided in model_info. if model_info.hf_ppl is None: with hf_runner( model_info.name, @@ -132,6 +134,7 @@ def vqa_ppl_test( return_tensors="pt", **mm_processor_kwargs, ) + inputs.pop("labels", None) inputs = inputs.to("cuda") input_ids = inputs["input_ids"] @@ -159,3 +162,9 @@ def vqa_ppl_test( assert abs(differ) < tol return vllm_ppl + + +def _ensure_rgb(image): + if image.mode != "RGB": + image = image.convert("RGB") + return image diff --git a/tests/models/utils.py b/tests/models/utils.py index f5c0cf97b36f..edd54bf99937 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -414,6 +414,9 @@ class RerankModelInfo(ModelInfo): class GenerateModelInfo(ModelInfo): hf_dtype: str = "auto" hf_ppl: float | None = None + pixel_factor: int = 0 + extra_kwargs: dict | None = None + image_token_id: int | None = None def get_vllm_extra_kwargs(model_info: ModelInfo, vllm_extra_kwargs): diff --git a/vllm/model_executor/layers/fusion/mm_input_norm.py b/vllm/model_executor/layers/fusion/mm_input_norm.py index 4f8da46a6607..e5c1d74c7715 100644 --- a/vllm/model_executor/layers/fusion/mm_input_norm.py +++ b/vllm/model_executor/layers/fusion/mm_input_norm.py @@ -18,7 +18,7 @@ """ import math -from typing import NamedTuple +from typing import Literal, NamedTuple, get_args import torch from torch import nn @@ -40,7 +40,7 @@ def _fused_mm_input_norm_kernel( C: tl.constexpr, BLOCK: tl.constexpr, ): - # 1D grid over the flattened (N, C, L) tensor. Each program processes + # 1D grid over the flattened (P, C, L) tensor. Each program processes # BLOCK contiguous elements; the channel index is recovered from the # flat offset via `(offs // L) % C`. int32 offsets: int64 division is # ~2x slower here (measured on GB200); the wrapper asserts @@ -75,9 +75,9 @@ def fused_mm_input_norm_triton( """Fused per-channel affine transform for normalisation. Args: - inputs: Input tensor, shape ``(N, C, L)``. A contiguous copy is + inputs: Input tensor, shape ``(P, C, L)``. A contiguous copy is materialized internally if needed. - outputs: Output tensor, contiguous and shaped exactly ``(N, C, L)``. + outputs: Output tensor, contiguous and shaped exactly ``(P, C, L)``. weight: Per-channel scale, shape ``(C,)``. bias: Per-channel shift, shape ``(C,)``. @@ -85,9 +85,9 @@ def fused_mm_input_norm_triton( ``outputs``, for chaining. """ - N, C, L = inputs.shape - assert outputs.shape == (N, C, L), ( - f"outputs shape {tuple(outputs.shape)} != inputs shape {(N, C, L)}" + P, C, L = inputs.shape + assert outputs.shape == (P, C, L), ( + f"outputs shape {tuple(outputs.shape)} != inputs shape {(P, C, L)}" ) assert weight.numel() == C and bias.numel() == C, ( f"weight/bias must have {C} elements, got {weight.numel()} / {bias.numel()}" @@ -102,7 +102,7 @@ def fused_mm_input_norm_triton( # sizes; larger tiles starve the SMs on small inputs. BLOCK_SIZE = 1024 - numel = N * C * L + numel = P * C * L # The kernel indexes with int32 offsets (int64 division costs ~2x # throughput); guard the flat index range here. assert numel < 2**31, f"numel={numel} exceeds the int32 index range" @@ -179,6 +179,26 @@ def _load_norm_params(model_config: ModelConfig) -> NormParams: ) +# Layout of the tensors (a.k.a. ``pixel_values``) passed to ``FusedMMInputNorm``. +# +# * ``"pcl"`` — ``(patches, C * L)``: flattened per-patch pixels. +# * ``"bchw"`` — ``(B, C, H, W)``: NCHW image tensors. +InputLayout = Literal["pcl", "bchw"] +_INPUT_LAYOUTS: tuple[str, ...] = get_args(InputLayout) + + +def _validate_layout(layout: str) -> InputLayout: + """Runtime validation for layouts supplied via config/plumbing. + + A ``Literal`` is only checked statically; validating here preserves the + fail-fast behaviour of the previous ``InputLayout(layout)`` call so that + typos do not silently fall through to the ``PCL`` branch. + """ + if layout not in _INPUT_LAYOUTS: + raise ValueError(f"invalid layout {layout!r}; expected one of {_INPUT_LAYOUTS}") + return layout # type: ignore[return-value] + + class IdentityInputNorm(nn.Module): """Stand-in used when the processor requires no rescale/normalise. @@ -212,6 +232,13 @@ class FusedMMInputNorm(CustomOp): computation itself is always fp32, independent of the output dtype (e.g. compute fp32, emit bf16). + Input Layout semantics: + + * ``layout="pcl"`` (default) — ``pixel_values`` is ``(patches, C*L)``, + output has the same shape. + * ``layout="bchw"`` — ``pixel_values`` is ``(B, C, H, W)``, output has + the same shape. + Platform dispatch: * ``forward_native`` — pure PyTorch eager path, used as the semantic @@ -228,6 +255,7 @@ def __init__( image_std: list[float], rescale_factor: float, channel: int = 3, + layout: InputLayout = "pcl", ): super().__init__() @@ -238,6 +266,7 @@ def __init__( assert rescale_factor != 0.0, "rescale_factor must be non-zero" self.channel = channel + self.layout = _validate_layout(layout) # Model construction can set the accelerator as PyTorch's default # device; build the buffers on CPU first, then move them over. @@ -277,6 +306,37 @@ def _patch_size(self, size: int) -> int: ) return size // self.channel + def _unpack( + self, pixel_values: torch.Tensor + ) -> tuple[torch.Tensor, tuple[int, ...]]: + """Normalise either layout to ``(P, C, L)`` plus the output shape. + + Returns a ``(x3, out_shape)`` pair where ``x3`` is a + ``(P, C, L)`` view of the input and ``out_shape`` is the shape + the final output must have. The ``reshape`` below is a pure view + as long as the caller's tensor is contiguous; the platform + forward_* methods call ``.contiguous()`` inside + ``fused_mm_input_norm_triton`` when needed. + """ + if self.layout == "bchw": + assert pixel_values.ndim == 4, ( + f"pixel_values must be 4D (B, C, H, W) for layout='bchw', " + f"got {pixel_values.dim()}D with shape " + f"{tuple(pixel_values.shape)}" + ) + B, C, H, W = pixel_values.shape + assert self.channel == C, ( + f"input has {C} channels, module expects {self.channel}" + ) + return pixel_values.reshape(B, C, H * W), (B, C, H, W) + + patches, size = self._unpack_2d(pixel_values) + patch_size = self._patch_size(size) + return ( + pixel_values.reshape(patches, self.channel, patch_size), + (patches, size), + ) + # ------------------------------------------------------------------ # Platform-specific implementations # ------------------------------------------------------------------ @@ -289,26 +349,21 @@ def forward_native( This is the semantic reference implementation and the fallback used on any platform without a specialised kernel. """ - patches, size = self._unpack_2d(pixel_values) - patch_size = self._patch_size(size) + x3, out_shape = self._unpack(pixel_values) - # weight/bias are fp32, so type promotion makes the arithmetic fp32 - x = pixel_values.reshape(patches, self.channel, patch_size) - x = x * self.weight.view(1, self.channel, 1) + self.bias.view( + # weight/bias are fp32, so type promotion makes the arithmetic fp32. + x = x3 * self.weight.view(1, self.channel, 1) + self.bias.view( 1, self.channel, 1 ) - return x.view(patches, size).to(visual_dtype) + return x.reshape(out_shape).to(visual_dtype) def forward_cuda( self, pixel_values: torch.Tensor, visual_dtype: torch.dtype ) -> torch.Tensor: """Triton kernel path for CUDA devices.""" - patches, size = self._unpack_2d(pixel_values) - patch_size = self._patch_size(size) - - x3 = pixel_values.reshape(patches, self.channel, patch_size) - y = torch.empty((patches, size), dtype=visual_dtype, device=pixel_values.device) - y3 = y.view(patches, self.channel, patch_size) + x3, out_shape = self._unpack(pixel_values) + y = torch.empty(out_shape, dtype=visual_dtype, device=pixel_values.device) + y3 = y.view(x3.shape) fused_mm_input_norm_triton(x3, y3, self.weight, self.bias) return y @@ -319,7 +374,7 @@ def forward_xpu( """XPU fused custom kernel path for uint8 and Triton kernel path for others. """ - if pixel_values.dtype != torch.uint8: + if pixel_values.dtype != torch.uint8 or self.layout == "bchw": return self.forward_cuda(pixel_values, visual_dtype) # For uint8 input, falls to xpu_fused_input_norm from xpu kernels @@ -336,14 +391,16 @@ def forward_oot( return self.forward_native(pixel_values, visual_dtype) -def build_mm_input_norm(model_config: ModelConfig) -> nn.Module: +def build_mm_input_norm( + model_config: ModelConfig, + layout: InputLayout = "pcl", +) -> nn.Module: """Build the input normalisation module for a model. Returns an ``IdentityInputNorm`` when device-side normalisation is disabled or the processor's rescale/normalise is numerically the identity transform; otherwise a ``FusedMMInputNorm`` built from the processor's - parameters. - """ + parameters.""" mm_config = getattr(model_config, "multimodal_config", None) if not getattr(mm_config, "mm_device_do_normalize", False): return IdentityInputNorm() @@ -361,4 +418,5 @@ def build_mm_input_norm(model_config: ModelConfig) -> nn.Module: image_std=params.image_std, rescale_factor=params.rescale_factor, channel=len(params.image_mean), + layout=layout, ) diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index b5b0c0d48bc4..72a402c874b5 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -39,6 +39,7 @@ from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape +from ..layers.fusion.mm_input_norm import build_mm_input_norm from .interfaces import ( MultiModalEmbeddings, SupportsEncoderCudaGraph, @@ -523,6 +524,7 @@ class Gemma3ForConditionalGeneration( ) supports_tower_connector_lora = True + supports_mm_device_do_normalize = True @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: @@ -554,6 +556,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vision_config, quant_config, prefix=maybe_prefix(prefix, "vision_tower"), + input_norm=build_mm_input_norm(self.model_config, layout="bchw"), ) self.multi_modal_projector = Gemma3MultiModalProjector(config) @@ -782,7 +785,8 @@ def prepare_encoder_cudagraph_capture_inputs( ) image_size = self.config.vision_config.image_size - dummy_pixel_values = torch.randn( + dtype = self.vision_tower.vision_model.input_norm.input_dtype or dtype + dummy_pixel_values = torch.zeros( num_images, 3, image_size, diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index bcaaf11f0d07..c4afa9822e8f 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -59,6 +59,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import async_tensor_h2d +from ..layers.fusion.mm_input_norm import IdentityInputNorm from .clip import dual_encoder_has_text_tokens, merge_dual_encoder_text_and_vision from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant from .interfaces_base import default_pooling_type @@ -319,7 +320,7 @@ def forward( target_dtype = self.patch_embedding.weight.dtype patch_embeds = self.patch_embedding( pixel_values.to(dtype=target_dtype) - ) # shape = [*, width, grid, grid] + ) # shape = [*, embed_dim, grid, grid] embeddings = patch_embeds.flatten(2).transpose(1, 2) if interpolate_pos_encoding: @@ -659,6 +660,7 @@ def __init__( require_post_norm: bool | None = None, prefix: str = "", use_head: bool | None = False, + input_norm: nn.Module | None = None, ) -> None: super().__init__() @@ -726,6 +728,8 @@ def __init__( orig_to_new_prefix=drops ) + self.input_norm = input_norm if input_norm is not None else IdentityInputNorm() + @property def dtype(self): return next(self.parameters()).dtype @@ -742,6 +746,7 @@ def forward( select_layers: list[int] | None = None, feature_select_strategy: VisionFeatureSelectStrategy | None = None, ) -> torch.Tensor: + pixel_values = self.input_norm(pixel_values.to(device=self.device), self.dtype) hidden_states = self.embeddings( pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, @@ -812,6 +817,7 @@ def __init__( require_post_norm: bool | None = None, prefix: str = "", use_head: bool | None = False, + input_norm: nn.Module | None = None, ) -> None: super().__init__() @@ -823,6 +829,7 @@ def __init__( require_post_norm=require_post_norm, prefix=maybe_prefix(prefix, "vision_model"), use_head=use_head, + input_norm=input_norm, ) def get_input_embeddings(self) -> nn.Module: