[MM] Add Triton kernel for mm_input_normal. (#56798)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
Signed-off-by: wang.yuqi <noooop@126.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>

docs/design/mm_processing.md

@@ -42,34 +42,39 @@ When new data is passed in, we first check which items are in the cache, and whi
 
 ### Fused Normalisation on the Device
 
-To accelerate the multi‑modal data pipeline (decoding, resizing, normalisation, and rescaling), we offload the heavy numerical preprocessing from the CPU to the GPU and optimise data movement.
+To accelerate the multi‑modal data pipeline (decoding, resizing, normalisation, and rescaling), we offload the **normalisation and rescaling** steps from the CPU to the device and optimises data movement.
 
 #### Fusing Normalisation and Rescaling on the GPU
 
-Traditionally, the CPU would divide pixel values by 255, then subtract the mean and divide by the standard deviation. We fuse these steps into one operation and run it entirely on the GPU.
+Traditionally, the CPU would divide pixel values by 255, then subtracts the mean and divides by the standard deviation. We fuse these steps into a single per-channel affine transform on the device.
 
-**How it works**: We use a dedicated `FusedInputNorm` module that bakes the rescale factor (1/255) directly into the layer's `weight` and `bias` parameters. Instead of performing three separate steps (scale, subtract, divide), the module does everything in a single affine transformation: `y = x * weight + bias`.
+`FusedMMInputNorm` dispatches to `fused_mm_input_norm_triton` on CUDA, which applies:
 
-The parameters are set as follows:
+    y = x * weight[c] + bias[c]
+      = (x * rescale_factor - image_mean[c]) / image_std[c]
 
-- `weight` controls both the standard deviation and the rescale factor
-- `bias` centers the data using the mean and the same rescale factor
+with:
 
-This means the module takes raw pixel values (0–255) and outputs properly normalised values without ever explicitly dividing by 255 as a separate step.
+    weight[c] = rescale_factor / image_std[c]
+    bias[c]   = -image_mean[c] / image_std[c]
 
-#### Optimized Data Path for Fused Normalisation
+The kernel accepts `uint8`, `float16`, `bfloat16` and `float32` inputs. The `uint8` path is the primary fast path: raw bytes travel to the device unprocessed, and the rescale factor is folded into `weight` — no separate divide-by-255 step. Compute is always done in fp32 inside the kernel.
 
-Performing fused normalisation directly on the device allows us to keep the entire transfer path—from **Entrypoint** through **Engine Core** to **GPU memory**—in **`uint8`**. This halves PCIe bandwidth and reduces CPU memory footprint.
+#### Optimized Data Path
 
-Only after data reaches GPU memory do we cast to `fp32` for the `FusedInputNorm` layer (to ensure numerical accuracy), then cast to `bf16` for subsequent layers—all within the GPU, avoiding any host‑side conversions.
+    Before: CPU decode → CPU resize → CPU rescale (÷255) → CPU normalize
+            → cast to bf16 → PCIe (2 B/elem) → GPU
 
-Overall path: **`Entrypoint (uint8) → Engine Core (uint8) → GPU Memory (uint8)`** → GPU‑local `fp32` `FusedInputNorm` → `bf16` output.
+    After:  CPU decode → CPU resize → uint8 pixel_values
+            → PCIe (1 B/elem) → GPU fused affine (fp32 compute) → bf16
+
+The transfer path **Entrypoint (API server / offline LLM) → Engine Core(scheduler + workers) → Device Memory** stays in `uint8`. On the device,`fused_mm_input_norm_triton` computes in fp32 internally and writes the requested `visual_dtype` (a per-call argument to `forward_*`, commonly `bf16`) directly, without materialising a global fp32 intermediate.
 
 #### Toggle: `mm_device_do_normalize`
 
-This GPU‑side fusion is controlled by a config flag called **`mm_device_do_normalize`**.
+GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:
 
-- When `True`, normalisation and rescaling are done on the GPU using the `FusedInputNorm` layer; when `False`, we fall back to the old CPU‑side path.
+- When `True`, normalisation and rescaling are done on the GPU using the `FusedMMInputNorm` layer; when `False`, we fall back to the old CPU‑side path.
 - The flag is **enabled by default** for all models that support it.
 - Currently, it’s on by default for these architectures:
 
@@ -78,8 +83,9 @@ This GPU‑side fusion is controlled by a config flag called **`mm_device_do_nor
 | `qwen2-vl`   | `Qwen2VLForConditionalGeneration`    | `Qwen/Qwen2-VL-2B-Instruct`, etc.   |
 | `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
 
-#### What We Gain Overall
+#### Key Properties and Gains
 
-- **CPU offload**: The arithmetic for normalisation and rescaling is completely gone from the CPU.
-- **PCIe savings**: Sending `uint8` (1 byte) instead of `bf16` (2 bytes) slashes data transfer volume by **50%** .
-- **GPU overhead**: The fused kernel is very lightweight and can often be merged with subsequent CUDA operations, so it hardly adds any extra cost.
+- **CPU offload & 50% PCIe savings** — normalisation/rescaling leaves the CPU entirely; sending `uint8` (1 byte) instead of `bf16` (2 bytes) halves transfer volume.
+- **Single-pass fusion** — `y = x * weight[c] + bias[c]` in one kernel launch via `fused_mm_input_norm_triton`; input is read in its native dtype (`uint8`).
+- **float32 compute for free** — arithmetic runs in fp32 inside the kernel regardless of I/O dtypes. Accuracy matches fp32, with no extra bandwidth: intermediates stay in registers, and no global fp32 tensor is materialised.
+- **Preallocated output with batch-dim padding** — callers can pass a larger buffer; only the leading `N` rows are written, so buffers are reusable across calls.

tests/model_executor/layers/test_mm_input_norm.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for FusedMMInputNorm and the fused affine transform kernel."""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fusion.mm_input_norm import (
+    FusedMMInputNorm,
+    IdentityInputNorm,
+    NormParams,
+    fused_mm_input_norm_triton,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
+from vllm.utils.torch_utils import set_random_seed
+
+# Module-level tests run on any platform: the CPU platform exercises
+# forward_native, accelerators exercise the fused kernel. Only the raw
+# kernel tests require an accelerator.
+_DEVICE = torch.device(current_platform.device_type)
+
+requires_accelerator = pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason="fused Triton kernel requires a CUDA/XPU accelerator",
+)
+requires_triton = pytest.mark.skipif(not HAS_TRITON, reason="requires Triton")
+requires_vllm_config = pytest.mark.usefixtures("default_vllm_config")
+
+
+def _reference_input_norm(
+    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: (x * rescale - mean) / std."""
+    patches, size = pixel_values.shape
+    patch_size = size // channel
+    mean = torch.tensor(
+        image_mean, dtype=torch.float32, device=pixel_values.device
+    ).view(1, channel, 1)
+    std = torch.tensor(image_std, dtype=torch.float32, device=pixel_values.device).view(
+        1, channel, 1
+    )
+    x = pixel_values.to(torch.float32).reshape(patches, channel, patch_size)
+    x = (x * rescale_factor - mean) / std
+    return x.view(patches, size).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]
+_RGB_RESCALE = 1.0 / 255.0
+
+
+def _check_against_reference(
+    pixel_values: torch.Tensor,
+    out_dtype: torch.dtype,
+    channel: int = 3,
+    image_mean: list[float] | None = None,
+    image_std: list[float] | None = None,
+    rescale_factor: float = _RGB_RESCALE,
+    **close_kwargs,
+):
+    """Run ``FusedMMInputNorm`` on ``pixel_values`` and compare against the
+    eager per-channel affine reference. Uses the RGB constants by default."""
+    if image_mean is None:
+        image_mean = list(_RGB_MEAN)
+    if image_std is None:
+        image_std = list(_RGB_STD)
+    norm = FusedMMInputNorm(
+        image_mean=image_mean,
+        image_std=image_std,
+        rescale_factor=rescale_factor,
+        channel=channel,
+    ).to(_DEVICE)
+    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
+    )
+
+
+# ===========================================================================
+# Module-level behavior
+# ===========================================================================
+@requires_vllm_config
+class TestFusedMMInputNormModule:
+    def test_matches_reference_above_cudnn_grid_limit(self):
+        """num_patches above the old cuDNN batch-norm grid limit (~65535)
+        used to raise CUDNN_STATUS_INTERNAL_ERROR."""
+        set_random_seed(0)
+        pixel_values = torch.randint(
+            0, 256, (70000, 3 * 196), dtype=torch.float32, device=_DEVICE
+        )
+        _check_against_reference(pixel_values, torch.float32)
+
+    def test_identity_passthrough(self):
+        norm = IdentityInputNorm()
+
+        pixel_values = torch.randn(8, 3 * 196, dtype=torch.float32, device=_DEVICE)
+        out = norm(pixel_values, visual_dtype=torch.bfloat16)
+        torch.testing.assert_close(out, pixel_values.to(torch.bfloat16))
+
+
+# ===========================================================================
+# dtype coverage
+# ===========================================================================
+@requires_vllm_config
+class TestFusedMMInputNormDtypes:
+    @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(self, in_dtype: torch.dtype, out_dtype: torch.dtype):
+        set_random_seed(0)
+        if in_dtype == torch.uint8:
+            pixel_values = torch.randint(
+                0, 256, (32, 3 * 16), dtype=torch.uint8, device=_DEVICE
+            )
+        else:
+            pixel_values = torch.rand(
+                (32, 3 * 16), dtype=torch.float32, device=_DEVICE
+            ).to(in_dtype)
+
+        _check_against_reference(
+            pixel_values,
+            out_dtype,
+            image_mean=[0.5, 0.5, 0.5],
+            image_std=[0.5, 0.5, 0.5],
+            atol=1e-2,
+            rtol=1e-2,
+        )
+
+
+# ===========================================================================
+# Shape / channel coverage
+# ===========================================================================
+@requires_vllm_config
+class TestFusedMMInputNormShapes:
+    @pytest.mark.parametrize("channel", [1, 3, 4])
+    def test_channel_variants(self, channel: int):
+        set_random_seed(0)
+        pixel_values = torch.randint(
+            0, 256, (16, channel * 64), dtype=torch.float32, device=_DEVICE
+        )
+        _check_against_reference(
+            pixel_values,
+            torch.float32,
+            channel=channel,
+            image_mean=[0.5] * channel,
+            image_std=[0.25] * channel,
+        )
+
+    @pytest.mark.parametrize("patch_size", [1, 2048])
+    def test_block_boundaries(self, patch_size: int):
+        """Masked-tail (12 elements) and exact-multiple (24 blocks) paths."""
+        set_random_seed(0)
+        pixel_values = torch.randint(
+            0, 256, (4, 3 * patch_size), dtype=torch.float32, device=_DEVICE
+        )
+        _check_against_reference(pixel_values, torch.float32)
+
+
+# ===========================================================================
+# Input handling: non-contiguous inputs
+# ===========================================================================
+@requires_vllm_config
+class TestFusedMMInputNormInputHandling:
+    def test_non_contiguous_input_matches_reference(self):
+        set_random_seed(0)
+        base = torch.randint(
+            0, 256, (8, 3 * 32, 2), dtype=torch.float32, device=_DEVICE
+        )
+        non_contig = base[..., 0]
+        assert not non_contig.is_contiguous()
+
+        _check_against_reference(non_contig, torch.float32)
+
+
+# ===========================================================================
+# Raw kernel entry point
+# ===========================================================================
+@requires_accelerator
+@requires_triton
+class TestFusedMMInputNormKernel:
+    @pytest.mark.parametrize(
+        "N, C, L",
+        [
+            (1, 1, 1),
+            (1, 3, 1),
+            (3, 3, 7),  # tile straddles channel boundaries
+            (2, 8, 63),
+            (4, 3, 1000),
+        ],
+    )
+    def test_channel_boundary_crossing(self, N: 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)
+        w = torch.randn(C, dtype=torch.float32, device=_DEVICE)
+        b = torch.randn(C, dtype=torch.float32, device=_DEVICE)
+
+        out = torch.empty_like(x)
+        fused_mm_input_norm_triton(x, out, w, b)
+
+        expected = x * w.view(1, C, 1) + b.view(1, C, 1)
+        torch.testing.assert_close(out, expected)
+
+    def test_non_contiguous_inputs_materialized(self):
+        """Read-only tensors are made contiguous internally."""
+        N, C, L = 3, 3, 100
+        set_random_seed(0)
+        base = torch.randn(N, 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)
+        fused_mm_input_norm_triton(x, out, w, b)
+
+        expected = x * w.view(1, C, 1) + b.view(1, C, 1)
+        torch.testing.assert_close(out, expected)
+
+    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)
+        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)):
+            with pytest.raises(AssertionError):
+                fused_mm_input_norm_triton(
+                    x,
+                    torch.empty(*out_shape, dtype=torch.float32, device=_DEVICE),
+                    w,
+                    b,
+                )
+
+
+# ===========================================================================
+# Construction / configuration
+# ===========================================================================
+@requires_vllm_config
+class TestFusedMMInputNormConstruction:
+    """Weight/bias buffer semantics at init time."""
+
+    def test_identity_config_buffers(self):
+        """Numerically identity parameters still build plain weight/bias
+        buffers; the identity shortcut lives in ``from_model_config``."""
+        norm = FusedMMInputNorm(
+            image_mean=[0.0, 0.0, 0.0],
+            image_std=[1.0, 1.0, 1.0],
+            rescale_factor=1.0,
+        )
+        torch.testing.assert_close(norm.weight, torch.ones(3))
+        torch.testing.assert_close(norm.bias, torch.zeros(3))
+
+    def test_non_identity_buffers(self):
+        channel = 3
+        norm = FusedMMInputNorm(
+            image_mean=_RGB_MEAN,
+            image_std=_RGB_STD,
+            rescale_factor=_RGB_RESCALE,
+            channel=channel,
+        )
+        assert norm.weight.shape == (channel,)
+        assert norm.bias.shape == (channel,)
+
+        mean = torch.tensor(_RGB_MEAN, dtype=torch.float32)
+        std = torch.tensor(_RGB_STD, dtype=torch.float32)
+        torch.testing.assert_close(norm.weight, _RGB_RESCALE / std)
+        torch.testing.assert_close(norm.bias, -mean / std)
+
+    def test_is_identity_params(self):
+        """The numeric-identity check mirrors the old allclose detection."""
+        assert NormParams([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0).is_identity
+        assert not NormParams(_RGB_MEAN, _RGB_STD, _RGB_RESCALE).is_identity
+
+    @pytest.mark.parametrize(
+        ("image_mean", "image_std", "rescale_factor"),
+        [
+            ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0),
+            ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25], 1 / 255),
+        ],
+    )
+    def test_buffers_initialized_on_default_device(
+        self,
+        image_mean: list[float],
+        image_std: list[float],
+        rescale_factor: float,
+    ):
+        """Buffers are built on CPU first, then moved to the caller's
+        default device, so init never touches the accelerator."""
+        # The meta device gives the CPU-only test shard the same non-CPU
+        # default-device semantics without requiring a CUDA build.
+        default_device = "cuda" if torch.cuda.is_available() else "meta"
+        with torch.device(default_device):
+            norm = FusedMMInputNorm(image_mean, image_std, rescale_factor)
+
+        assert norm.weight.device.type == default_device
+        assert norm.bias.device.type == default_device

tests/models/multimodal/processing/test_qwen2_vl.py

@@ -8,10 +8,11 @@ import pytest
 import torch
 from PIL import Image
 
-from vllm.model_executor.models.vision import FusedInputNorm
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
 from vllm.multimodal.inputs import batched_tensors_equal
+from vllm.platforms import current_platform
 
 from ....conftest import ImageTestAssets
 from ...utils import build_model_context
@@ -65,45 +66,6 @@ def test_jina_vl_processing_order() -> None:
     ]
 
 
-@pytest.mark.parametrize(
-    ("image_mean", "image_std", "rescale_factor", "is_identity"),
-    [
-        ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0, True),
-        ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25], 1 / 255, False),
-    ],
-)
-def test_fused_input_norm_initialization_on_device(
-    monkeypatch: pytest.MonkeyPatch,
-    image_mean: list[float],
-    image_std: list[float],
-    rescale_factor: float,
-    is_identity: bool,
-):
-    """Identity detection must not synchronize the default device."""
-    original_allclose = torch.allclose
-
-    def cpu_allclose(input: torch.Tensor, other: torch.Tensor, *args, **kwargs):
-        assert input.device.type == "cpu"
-        assert other.device.type == "cpu"
-        return original_allclose(input, other, *args, **kwargs)
-
-    monkeypatch.setattr(torch, "allclose", cpu_allclose)
-    # Exercise the real accelerator when available. The meta device gives the
-    # CPU-only test shard the same non-CPU default-device semantics without
-    # requiring a CUDA-enabled PyTorch build.
-    default_device = "cuda" if torch.cuda.is_available() else "meta"
-    with torch.device(default_device):
-        input_norm = FusedInputNorm(image_mean, image_std, rescale_factor)
-
-    assert input_norm.is_identity is is_identity
-    if is_identity:
-        assert input_norm.weight is None
-        assert input_norm.bias is None
-    else:
-        assert input_norm.weight.device.type == default_device
-        assert input_norm.bias.device.type == default_device
-
-
 @pytest.mark.parametrize("model_id", ["Qwen/Qwen2-VL-2B-Instruct"])
 @pytest.mark.parametrize(
     ("mm_processor_kwargs", "expected_toks_per_img", "expected_pixels_shape"),
@@ -240,16 +202,20 @@ def test_qwen2_5_vl_video_fps_to_second_per_grid_ts(model_id: str) -> None:
         assert float(spg) == pytest.approx(temporal_patch_size / fps, rel=2e-2)
 
 
+@pytest.mark.usefixtures("default_vllm_config")
 @pytest.mark.parametrize(
     "model_id", ["Qwen/Qwen2-VL-2B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct"]
 )
 @pytest.mark.parametrize("num_imgs", [1, 2])
 def test_mm_device_do_normalize(
-    image_assets: ImageTestAssets,
-    model_id: str,
-    num_imgs: int,
-):
-    """Ensure that enable mm_device_do_normalize yields the correct result."""
+    image_assets: ImageTestAssets, model_id: str, num_imgs: int
+) -> None:
+    """Device-side normalisation must reproduce the on-CPU processor result.
+
+    Runs on any platform: the CPU platform exercises the ``forward_native``
+    fallback, accelerators exercise the fused kernel.
+    """
+    device = current_platform.device_type
     ctx = build_model_context(
         model_id,
         limit_mm_per_prompt={"image": num_imgs},
@@ -280,9 +246,15 @@ def test_mm_device_do_normalize(
     ].get_data()["pixel_values"]
 
     ctx.model_config.multimodal_config.mm_device_do_normalize = True
-    input_norm = FusedInputNorm.from_model_config(ctx.model_config)
+    input_norm = build_mm_input_norm(ctx.model_config).to(device)
+
+    # With normalisation disabled, the processor emits raw uint8 pixels,
+    # matching the production mm_device_do_normalize path.
+    assert pixel_values_without_normalize.dtype == torch.uint8
     pixel_values_do_input_norm = input_norm(
-        pixel_values_without_normalize.to(dtype), dtype
+        pixel_values_without_normalize.to(device), dtype
     )
 
-    torch.testing.assert_close(pixel_values_with_normalize, pixel_values_do_input_norm)
+    torch.testing.assert_close(
+        pixel_values_with_normalize.to(device), pixel_values_do_input_norm
+    )

tests/models/test_vision.py

@@ -13,7 +13,6 @@ from vllm.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm.model_executor.models.vision import (
-    FusedInputNorm,
     get_load_balance_assignment,
     resolve_visual_encoder_outputs,
     run_dp_sharded_mrope_vision_model,
@@ -488,62 +487,3 @@ def test_simple_mrope_vision_model_spatial_merge(spatial_merge_size: int):
 
     assert output.shape[0] == expected_output_patches
     assert output.shape[1] == vision_model.out_hidden_size
-
-
-def _reference_input_norm(
-    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: (x * rescale - mean) / std."""
-    patches, size = pixel_values.shape
-    patch_size = size // channel
-    mean = torch.tensor(image_mean, dtype=torch.float32).view(1, channel, 1)
-    std = torch.tensor(image_std, dtype=torch.float32).view(1, channel, 1)
-    x = pixel_values.to(torch.float32).view(patches, channel, patch_size)
-    x = (x * rescale_factor - mean) / std
-    return x.view(patches, size).to(out_dtype)
-
-
-@pytest.mark.parametrize("num_patches", [1, 37, 70000])
-def test_fused_input_norm_matches_reference(num_patches: int):
-    """FusedInputNorm must equal the plain affine, including for num_patches
-    above the cuDNN batch-norm grid limit (~65535) that previously raised
-    CUDNN_STATUS_INTERNAL_ERROR (issue #51717)."""
-    channel = 3
-    patch_size = 14 * 14
-    image_mean = [0.48145466, 0.4578275, 0.40821073]
-    image_std = [0.26862954, 0.26130258, 0.27577711]
-    rescale_factor = 1.0 / 255.0
-
-    set_random_seed(0)
-    pixel_values = torch.randint(
-        0, 256, (num_patches, channel * patch_size), dtype=torch.float32
-    )
-
-    norm = FusedInputNorm(
-        image_mean=image_mean,
-        image_std=image_std,
-        rescale_factor=rescale_factor,
-        channel=channel,
-    )
-    assert not norm.is_identity
-
-    out = norm(pixel_values, visual_dtype=torch.float32)
-    expected = _reference_input_norm(
-        pixel_values, image_mean, image_std, rescale_factor, channel, torch.float32
-    )
-    torch.testing.assert_close(out, expected)
-
-
-def test_fused_input_norm_identity_passthrough():
-    """The identity configuration returns the input unchanged (cast only)."""
-    norm = FusedInputNorm.identity()
-    assert norm.is_identity
-
-    pixel_values = torch.randn(8, 3 * 196, dtype=torch.float32)
-    out = norm(pixel_values, visual_dtype=torch.bfloat16)
-    torch.testing.assert_close(out, pixel_values.to(torch.bfloat16))

vllm/config/multimodal.py

@@ -534,11 +534,19 @@ class MultiModalConfig:
         """Get the keyword arguments to pass to the multi-modal processor
         according to the extra arguments passed during inference.
         """
-        kwargs = self.mm_processor_kwargs or {}
+        merged = dict(self.mm_processor_kwargs or {}) | dict(inference_kwargs)
         if self.mm_device_do_normalize:
-            kwargs["do_normalize"] = False
-            kwargs["do_rescale"] = False
-        return kwargs | dict(inference_kwargs)
+            # The model expects raw pixels and normalises on device;
+            # overriding these per request would silently double-normalise.
+            for key in ("do_normalize", "do_rescale"):
+                if key in inference_kwargs:
+                    logger.warning_once(
+                        "mm_device_do_normalize is enabled; ignoring "
+                        "per-request mm_processor_kwargs[%r].",
+                        key,
+                    )
+                merged[key] = False
+        return merged
 
     def use_gpu_video_backend(self) -> bool:
         """Return whether the configured video loader or codec uses the GPU."""

vllm/model_executor/layers/fusion/mm_input_norm.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Fused Normalisation on the Device.
+
+Equivalent to::
+
+    output = (input * rescale_factor - image_mean) / image_std
+
+This is implemented as a single per-channel affine transform::
+
+    output = input * weight[c] + bias[c]
+
+where::
+
+    weight = rescale_factor / image_std
+    bias = -image_mean / image_std
+"""
+
+import math
+from typing import Any, NamedTuple
+
+import torch
+from torch import nn
+
+from vllm.config import ModelConfig
+from vllm.model_executor.custom_op import CustomOp
+from vllm.transformers_utils.processor import cached_get_processor, get_processor_config
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _fused_mm_input_norm_kernel(
+    x_ptr,
+    y_ptr,
+    w_ptr,
+    b_ptr,
+    numel,
+    L,
+    C: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # 1D grid over the flattened (N, 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
+    # numel < 2**31 so the flat index never overflows.
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    c = (offs // L) % C
+
+    x = tl.load(
+        x_ptr + offs,
+        mask=mask,
+        other=0,
+        eviction_policy="evict_first",
+    ).to(tl.float32)
+    w = tl.load(w_ptr + c, mask=mask, other=0).to(tl.float32)
+    b = tl.load(b_ptr + c, mask=mask, other=0).to(tl.float32)
+    tl.store(
+        y_ptr + offs,
+        x * w + b,
+        mask=mask,
+        eviction_policy="evict_first",
+    )
+
+
+def fused_mm_input_norm_triton(
+    inputs: torch.Tensor,
+    outputs: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+):
+    """Fused per-channel affine transform for normalisation.
+
+    Args:
+        inputs: Input tensor, shape ``(N, C, L)``. A contiguous copy is
+            materialized internally if needed.
+        outputs: Output tensor, contiguous and shaped exactly ``(N, C, L)``.
+        weight: Per-channel scale, shape ``(C,)``.
+        bias: Per-channel shift, shape ``(C,)``.
+
+    Returns:
+        ``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)}"
+    )
+    assert weight.numel() == C and bias.numel() == C, (
+        f"weight/bias must have {C} elements, got {weight.numel()} / {bias.numel()}"
+    )
+    # The flat 1D kernel writes outputs by raw offset; a non-contiguous
+    # buffer would be written at the wrong locations and cannot be
+    # auto-materialized without breaking the caller's aliasing.
+    assert outputs.is_contiguous(), "outputs must be contiguous"
+
+    # Tile size along the flattened element axis. 1024 with 4 warps gives
+    # 8 elements per lane and benchmarks fastest across realistic image
+    # sizes; larger tiles starve the SMs on small inputs.
+    BLOCK_SIZE = 1024
+
+    numel = N * 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"
+
+    grid = (triton.cdiv(numel, BLOCK_SIZE),)
+
+    _fused_mm_input_norm_kernel[grid](
+        inputs.contiguous(),
+        outputs,
+        weight.contiguous(),
+        bias.contiguous(),
+        numel,
+        L,
+        C=C,
+        BLOCK=BLOCK_SIZE,
+        num_warps=4,  # 8 elements per lane at BLOCK_SIZE=1024
+    )
+    return outputs
+
+
+class NormParams(NamedTuple):
+    """Resolved per-channel affine parameters (flags already folded in)."""
+
+    image_mean: list[float]
+    image_std: list[float]
+    rescale_factor: float
+
+    @property
+    def is_identity(self) -> bool:
+        """Whether ``weight = rescale/std`` and ``bias = -mean/std`` are
+        numerically the identity transform (mirrors ``torch.allclose``'s
+        default fp32 tolerances)."""
+        return all(
+            math.isclose(self.rescale_factor / s, 1.0, rel_tol=1e-5, abs_tol=1e-8)
+            and math.isclose(m / s, 0.0, abs_tol=1e-8)
+            for m, s in zip(self.image_mean, self.image_std)
+        )
+
+
+def _load_norm_params(model_config: ModelConfig) -> NormParams:
+    """Resolve the per-channel affine parameters ``(image_mean, image_std,
+    rescale_factor)`` from the processor config, falling back to the image
+    processor object."""
+    model = model_config.model
+    revision = model_config.revision
+
+    config = get_processor_config(model, revision=revision)
+    # NOTE: do not use cached_image_processor_from_config here — it merges
+    # mm_processor_kwargs, which mm_device_do_normalize poisons with
+    # do_normalize=False.
+    image_processor = cached_get_processor(
+        model, revision=revision, trust_remote_code=model_config.trust_remote_code
+    ).image_processor
+
+    def resolve(key: str) -> Any:
+        """Processor config value, falling back to the image_processor."""
+        if (value := config.get(key)) is not None:
+            return value
+        return getattr(image_processor, key, None)
+
+    do_rescale = bool(resolve("do_rescale"))
+    do_normalize = bool(resolve("do_normalize"))
+
+    # Parameters whose flag is off are unused; default them to no-ops
+    # without resolving them.
+    rescale_factor = resolve("rescale_factor") if do_rescale else 1.0
+    image_mean = resolve("image_mean") if do_normalize else [0.0] * 3
+    image_std = resolve("image_std") if do_normalize else [1.0] * 3
+
+    assert rescale_factor is not None, "rescale_factor is still None after resolution."
+    assert image_mean is not None, "image_mean is still None after resolution."
+    assert image_std is not None, "image_std is still None after resolution."
+    assert len(image_mean) == len(image_std), (
+        f"image_mean and image_std have different lengths: "
+        f"{len(image_mean)} vs {len(image_std)}"
+    )
+
+    return NormParams(
+        image_mean=[float(v) for v in image_mean],
+        image_std=[float(v) for v in image_std],
+        rescale_factor=float(rescale_factor),
+    )
+
+
+class IdentityInputNorm(nn.Module):
+    """Stand-in used when the processor requires no rescale/normalise.
+
+    Not a no-op: with ``mm_device_do_normalize`` enabled, raw ``uint8``
+    pixels travel to the device unprocessed and must be cast to
+    ``visual_dtype`` here.
+    """
+
+    def forward(
+        self, pixel_values: torch.Tensor, visual_dtype: torch.dtype
+    ) -> torch.Tensor:
+        return pixel_values.to(visual_dtype, copy=False)
+
+
+@CustomOp.register("fused_mm_input_norm")
+class FusedMMInputNorm(CustomOp):
+    """Module that applies rescaling and normalisation to input images.
+    Equivalent to: output = (input * rescale_factor - mean) / std
+
+    Dtype semantics:
+
+    * Input dtype — the dtype of the ``pixel_values`` argument to ``forward_*``.
+      It is ``uint8`` when ``mm_device_do_normalize`` is enabled (raw bytes
+      travel to the device unprocessed) and equals ``visual_dtype``
+      otherwise.
+    * Output dtype — the ``visual_dtype`` argument of ``forward_*``. The
+      computation itself is always fp32, independent of the output dtype
+      (e.g. compute fp32, emit bf16).
+
+    Platform dispatch:
+
+    * ``forward_native`` — pure PyTorch eager path, used as the semantic
+      reference and default fallback on all platforms.
+    * ``forward_cuda`` — Triton kernel path for CUDA devices.
+    * ``forward_xpu`` — custom XPU kernel path.
+    * ``forward_oot`` — out-of-tree platform override entry point; falls back
+      to ``forward_native`` unless a plugin overrides it.
+    """
+
+    def __init__(
+        self,
+        image_mean: list[float],
+        image_std: list[float],
+        rescale_factor: float,
+        channel: int = 3,
+    ):
+        super().__init__()
+
+        assert len(image_mean) == len(image_std) == channel, (
+            f"image_mean/image_std must have {channel} entries, "
+            f"got {len(image_mean)} / {len(image_std)}"
+        )
+        assert rescale_factor != 0.0, "rescale_factor must be non-zero"
+
+        self.channel = channel
+
+        # Model construction can set the accelerator as PyTorch's default
+        # device; build the buffers on CPU first, then move them over.
+        mean = torch.tensor(image_mean, dtype=torch.float32, device="cpu")
+        std = torch.tensor(image_std, dtype=torch.float32, device="cpu")
+        device = torch.get_default_device()
+        self.register_buffer("weight", (rescale_factor / std).to(device))
+        self.register_buffer("bias", (-mean / std).to(device))
+
+    # ------------------------------------------------------------------
+    # Internal helpers shared by the platform-specific forward_* methods
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _unpack_2d(pixel_values: torch.Tensor) -> tuple[int, int]:
+        assert pixel_values.ndim == 2, (
+            f"pixel_values must be 2D (patches, size), got {pixel_values.dim()}D "
+            f"with shape {tuple(pixel_values.shape)}"
+        )
+        patches, size = pixel_values.shape
+        return patches, size
+
+    def _patch_size(self, size: int) -> int:
+        assert size % self.channel == 0, (
+            f"size={size} is not divisible by channel={self.channel}"
+        )
+        return size // self.channel
+
+    # ------------------------------------------------------------------
+    # Platform-specific implementations
+    # ------------------------------------------------------------------
+
+    def forward_native(
+        self, pixel_values: torch.Tensor, visual_dtype: torch.dtype
+    ) -> torch.Tensor:
+        """Pure PyTorch eager implementation.
+
+        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)
+
+        # 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(
+            1, self.channel, 1
+        )
+        return x.view(patches, size).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)
+
+        fused_mm_input_norm_triton(x3, y3, self.weight, self.bias)
+        return y
+
+    def forward_xpu(
+        self, pixel_values: torch.Tensor, visual_dtype: torch.dtype
+    ) -> torch.Tensor:
+        """XPU fused custom kernel path.
+
+        On XPU, fuse the whole rescale + normalise into a single custom
+        kernel. The eager path materializes an fp32 intermediate and then
+        casts back, which adds device-side compute that cancels the
+        bandwidth saving of transferring uint8 pixel_values. The fused
+        kernel reads uint8 directly and writes ``visual_dtype`` in one pass.
+        """
+        # The out-of-tree XPU kernel only supports the uint8 input that
+        # device-side normalisation guarantees; fail loudly otherwise.
+        assert pixel_values.dtype == torch.uint8, (
+            f"xpu_fused_input_norm requires uint8 input, got {pixel_values.dtype}"
+        )
+        return torch.ops.vllm.xpu_fused_input_norm(
+            pixel_values, self.weight, self.bias, visual_dtype
+        )
+
+    def forward_oot(
+        self, pixel_values: torch.Tensor, visual_dtype: torch.dtype
+    ) -> torch.Tensor:
+        """Out-of-tree platform override entrypoint."""
+        return self.forward_native(pixel_values, visual_dtype)
+
+
+def build_mm_input_norm(model_config: ModelConfig) -> 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.
+    """
+    mm_config = getattr(model_config, "multimodal_config", None)
+    if not getattr(mm_config, "mm_device_do_normalize", False):
+        return IdentityInputNorm()
+
+    params = _load_norm_params(model_config)
+
+    # Flags off are already defaulted to identity values by
+    # ``_load_norm_params``; this also catches configs that request
+    # rescale/normalise but are numerically the identity transform.
+    if params.is_identity:
+        return IdentityInputNorm()
+
+    return FusedMMInputNorm(
+        image_mean=params.image_mean,
+        image_std=params.image_std,
+        rescale_factor=params.rescale_factor,
+        channel=len(params.image_mean),
+    )

vllm/model_executor/models/qwen2_5_vl.py

@@ -55,6 +55,10 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_and_mul_fn
 from vllm.model_executor.layers.attention import MMEncoderAttention
 from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.fusion.mm_input_norm import (
+    IdentityInputNorm,
+    build_mm_input_norm,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -120,7 +124,6 @@ from .utils import (
     maybe_prefix,
 )
 from .vision import (
-    FusedInputNorm,
     get_fp8_padded_hidden_size,
     get_vit_attn_backend,
     is_vit_use_data_parallel,
@@ -664,9 +667,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
             in_channels=in_channels,
             hidden_size=self.hidden_size,
         )
-        self.input_norm = (
-            input_norm if input_norm is not None else FusedInputNorm.identity()
-        )
+        self.input_norm = input_norm if input_norm is not None else IdentityInputNorm()
 
         norm_layer = partial(RMSNorm, eps=norm_eps)
         head_dim = self.hidden_size // self.num_heads
@@ -1425,7 +1426,7 @@ class Qwen2_5_VLForConditionalGeneration(
                 vision_config=config.vision_config,
                 norm_eps=getattr(config, "rms_norm_eps", 1e-6),
                 quant_config=self.quant_config,
-                input_norm=FusedInputNorm.from_model_config(self.model_config),
+                input_norm=build_mm_input_norm(self.model_config),
                 prefix=maybe_prefix(prefix, "visual"),
             )
 

vllm/model_executor/models/qwen2_vl.py

@@ -51,6 +51,10 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import QuickGELU
 from vllm.model_executor.layers.attention import MMEncoderAttention
 from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.fusion.mm_input_norm import (
+    IdentityInputNorm,
+    build_mm_input_norm,
+)
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     RowParallelLinear,
@@ -105,7 +109,6 @@ from .utils import (
     maybe_prefix,
 )
 from .vision import (
-    FusedInputNorm,
     get_vit_attn_backend,
     is_vit_use_data_parallel,
     run_dp_sharded_mrope_vision_model,
@@ -562,9 +565,7 @@ class Qwen2VisionTransformer(nn.Module):
             in_channels=in_channels,
             embed_dim=embed_dim,
         )
-        self.input_norm = (
-            input_norm if input_norm is not None else FusedInputNorm.identity()
-        )
+        self.input_norm = input_norm if input_norm is not None else IdentityInputNorm()
 
         norm_layer = partial(nn.LayerNorm, eps=norm_eps)
         head_dim = embed_dim // num_heads
@@ -1295,7 +1296,7 @@ class Qwen2VLForConditionalGeneration(
                 config.vision_config,
                 norm_eps=getattr(config, "rms_norm_eps", 1e-6),
                 quant_config=quant_config,
-                input_norm=FusedInputNorm.from_model_config(self.model_config),
+                input_norm=build_mm_input_norm(self.model_config),
                 prefix=maybe_prefix(prefix, "visual"),
             )
 

vllm/model_executor/models/vision.py

@@ -8,10 +8,9 @@ from collections.abc import Callable
 from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar
 
 import torch
-import torch.nn as nn
 from transformers import PreTrainedConfig
 
-from vllm.config import ModelConfig, MultiModalConfig, get_current_vllm_config_or_none
+from vllm.config import MultiModalConfig, get_current_vllm_config_or_none
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
@@ -19,7 +18,6 @@ from vllm.distributed import (
 )
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.transformers_utils.processor import get_processor, get_processor_config
 from vllm.utils.math_utils import round_up
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -578,165 +576,3 @@ def run_dp_sharded_mrope_vision_model(
         "Found unassigned embeddings"
     )
     return out_embeddings
-
-
-class FusedInputNorm(nn.Module):
-    """Module that applies rescaling and normalization to input images.
-    Equivalent to: output = (input * rescale_factor - mean) / std
-    """
-
-    def __init__(
-        self,
-        image_mean: list[float],
-        image_std: list[float],
-        rescale_factor: float,
-        channel: int = 3,
-        dtype: torch.dtype = torch.float32,
-    ):
-        super().__init__()
-
-        self.channel = channel
-
-        # Model construction can set the accelerator as PyTorch's default
-        # device. Determine whether the normalization is an identity on CPU
-        # so torch.allclose does not introduce a device synchronization while
-        # the model is being initialized. The actual buffers below still use
-        # the caller's default device.
-        image_mean_cpu = torch.tensor(image_mean, dtype=dtype, device="cpu") * (
-            1.0 / rescale_factor
-        )
-        image_std_cpu = torch.tensor(image_std, dtype=dtype, device="cpu") * (
-            1.0 / rescale_factor
-        )
-        weight_cpu = 1.0 / image_std_cpu
-        bias_cpu = -image_mean_cpu / image_std_cpu
-        self.is_identity = bool(
-            torch.allclose(weight_cpu, torch.ones_like(weight_cpu))
-            and torch.allclose(bias_cpu, torch.zeros_like(bias_cpu))
-        )
-
-        if not self.is_identity:
-            image_mean_tensor = torch.tensor(image_mean, dtype=dtype) * (
-                1.0 / rescale_factor
-            )
-            image_std_tensor = torch.tensor(image_std, dtype=dtype) * (
-                1.0 / rescale_factor
-            )
-            weight = 1.0 / image_std_tensor
-            bias = -image_mean_tensor / image_std_tensor
-            self.register_buffer("weight", weight)
-            self.register_buffer("bias", bias)
-        else:
-            self.register_buffer("weight", None)
-            self.register_buffer("bias", None)
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.weight.dtype
-
-    @classmethod
-    def identity(
-        cls, channel: int = 3, dtype: torch.dtype = torch.float32
-    ) -> "FusedInputNorm":
-        return cls(
-            image_mean=[0.0, 0.0, 0.0],
-            image_std=[1.0, 1.0, 1.0],
-            rescale_factor=1.0,
-            channel=channel,
-            dtype=dtype,
-        )
-
-    @classmethod
-    def from_model_config(cls, model_config: "ModelConfig") -> nn.Module:
-        if not model_config.multimodal_config.mm_device_do_normalize:
-            return cls.identity()
-
-        model = model_config.model
-        revision = model_config.revision
-
-        # Try to read parameters from the processor config
-        config = get_processor_config(model, revision=revision)
-        do_rescale = config.get("do_rescale", None)
-        do_normalize = config.get("do_normalize", None)
-        image_mean = config.get("image_mean", None)
-        image_std = config.get("image_std", None)
-        rescale_factor = config.get("rescale_factor", None)
-
-        # Fallback to the image_processor object if any parameter is missing
-        if None in [do_rescale, do_normalize, image_mean, image_std, rescale_factor]:
-            image_processor = get_processor(model, revision=revision).image_processor
-
-            if do_rescale is None:
-                do_rescale = getattr(image_processor, "do_rescale", None)
-            if do_normalize is None:
-                do_normalize = getattr(image_processor, "do_normalize", None)
-            if image_mean is None:
-                image_mean = getattr(image_processor, "image_mean", None)
-            if image_std is None:
-                image_std = getattr(image_processor, "image_std", None)
-            if rescale_factor is None:
-                rescale_factor = getattr(image_processor, "rescale_factor", None)
-
-        # Apply defaults based on flags
-        if not do_rescale:
-            rescale_factor = 1.0
-        if not do_normalize:
-            image_mean = [0.0, 0.0, 0.0]
-            image_std = [1.0, 1.0, 1.0]
-
-        # Ensure all required parameters are resolved
-        assert None not in [
-            do_rescale,
-            do_normalize,
-            image_mean,
-            image_std,
-            rescale_factor,
-        ], "Some normalization parameters are still None after resolution."
-
-        # If no processing is needed, return an identity module
-        if not do_rescale and not do_normalize:
-            return cls.identity()
-
-        return cls(
-            image_mean=image_mean, image_std=image_std, rescale_factor=rescale_factor
-        )
-
-    def forward(
-        self,
-        grid_thw: torch.Tensor,
-        visual_dtype: torch.dtype,
-    ) -> torch.Tensor:
-        if self.is_identity:
-            return grid_thw.to(visual_dtype)
-
-        assert grid_thw.ndim == 2
-        patches, size = grid_thw.shape
-        patch_size = size // self.channel
-
-        # On XPU, fuse the whole rescale + normalize into a single custom
-        # kernel. The eager path below materializes an fp32 intermediate and
-        # then casts back, which adds device-side compute that cancels the
-        # bandwidth saving of transferring uint8 pixel_values. The fused
-        # kernel reads uint8 directly and writes ``visual_dtype`` in one pass.
-        if (
-            current_platform.is_xpu()
-            and grid_thw.dtype == torch.uint8
-            and self.weight.dtype == torch.float32
-        ):
-            return torch.ops.vllm.xpu_fused_input_norm(
-                grid_thw, self.weight, self.bias, visual_dtype
-            )
-
-        # Apply the per-channel affine transform directly instead of via
-        # F.batch_norm. batch_norm dispatches to cuDNN, whose batch-norm
-        # kernels cap the batch dimension near the CUDA grid limit (~65535);
-        # here that dimension is the number of patches, which grows unbounded
-        # with image resolution and batch size and overflows the cap on large
-        # image-heavy requests (CUDNN_STATUS_INTERNAL_ERROR). The plain
-        # broadcasted multiply-add is numerically identical and has no such
-        # limit.
-        x = grid_thw.to(self.dtype).view(patches, self.channel, patch_size)
-        x = x * self.weight.view(1, self.channel, 1) + self.bias.view(
-            1, self.channel, 1
-        )
-        return x.view(patches, size).to(visual_dtype)