[MM] Move Nemotron VL image normalization to device

docs/design/mm_processing.md

@@ -90,6 +90,7 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:
 | `llama-nemotron-vl-embed`  | `LlamaNemotronVLForEmbedding`              | `nvidia/llama-nemotron-embed-vl-1b-v2`                |
 | `llama-nemotron-vl-rerank` | `LlamaNemotronVLForSequenceClassification` | `nvidia/llama-nemotron-rerank-vl-1b-v2`               |
 | `mistral3`                 | `Mistral3ForConditionalGeneration`         | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. |
+| `nano-nemotron-vl`         | `NemotronH_Nano_VL_V2`                     | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4` |
 | `nemotron-parse`           | `NemotronParseForConditionalGeneration`    | `nvidia/NVIDIA-Nemotron-Parse-2.0`, etc.              |
 | `qwen2-vl`                 | `Qwen2VLForConditionalGeneration`          | `Qwen/Qwen2-VL-2B-Instruct`, etc.                     |
 | `qwen2.5-vl`               | `Qwen2_5_VLForConditionalGeneration`       | `Qwen/Qwen2.5-VL-3B-Instruct`, etc.                   |

tests/models/multimodal/test_nano_nemotron_vl.py

@@ -1,13 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import numpy as np
 import pytest
 import torch
+import torch.nn as nn
+from PIL import Image
 
 from vllm import envs
+from vllm.config import MultiModalConfig
+from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm
 from vllm.model_executor.models.nano_nemotron_vl import (
     NanoNemotronVLMultiModalProcessor,
     NemotronH_Nano_VL_V2,
@@ -17,6 +22,183 @@
     MultiModalDataParser,
     VideoProcessorItems,
 )
+from vllm.platforms import current_platform
+from vllm.transformers_utils.processors.nano_nemotron_vl import (
+    DynamicResolutionImageTiler,
+    NanoNemotronVLProcessor,
+    _bicubic_resize_and_normalize,
+)
+
+# FusedMMInputNorm runs a Triton kernel on GPU platforms, so its inputs must
+# live on the platform's device.
+_DEVICE = torch.device(current_platform.device_type)
+
+
+@pytest.mark.parametrize("do_cpu_normalize", [True, False])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_nemotron_resized_normalization_matches_reference(
+    dtype: torch.dtype, do_cpu_normalize: bool, default_vllm_config
+):
+    pixels = torch.randint(
+        0,
+        256,
+        (1, 19, 23, 3),
+        dtype=torch.uint8,
+        generator=torch.Generator().manual_seed(0),
+    )
+    mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)
+    std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)
+    output = _bicubic_resize_and_normalize(
+        pixels,
+        size=(16, 16),
+        norm_mean=mean,
+        norm_std=std,
+        do_cpu_normalize=do_cpu_normalize,
+        dtype=dtype,
+    )
+
+    if not do_cpu_normalize:
+        assert output.dtype == torch.uint8
+        model = object.__new__(NemotronH_Nano_VL_V2)
+        nn.Module.__init__(model)
+        model.input_norm = FusedMMInputNorm(
+            mean.flatten().tolist(), std.flatten().tolist(), 1.0 / 255.0
+        ).to(_DEVICE)
+        model.llm_dtype = dtype
+        output = model._normalize_pixel_values(output.to(_DEVICE)).cpu()
+
+    resized = torch.nn.functional.interpolate(
+        pixels.permute(0, 3, 1, 2).float(),
+        size=(16, 16),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    )
+    expected = ((resized.float() / 255.0 - mean) / std).to(dtype)
+    assert output.dtype == dtype
+    torch.testing.assert_close(output, expected, rtol=0.02, atol=0.06)
+    assert (output.float() - expected.float()).abs().mean() < 0.01
+
+
+def test_nemotron_dynamic_device_normalization_matches_cpu_reference(
+    default_vllm_config,
+):
+    image_mean = [0.485, 0.456, 0.406]
+    image_std = [0.229, 0.224, 0.225]
+    shape = (1, 7, 3 * 16 * 16)
+    pixels = torch.randint(0, 256, shape, dtype=torch.uint8)
+
+    model = object.__new__(NemotronH_Nano_VL_V2)
+    nn.Module.__init__(model)
+    model.input_norm = FusedMMInputNorm(image_mean, image_std, 1.0 / 255.0).to(_DEVICE)
+    model.llm_dtype = torch.bfloat16
+
+    output = model._normalize_pixel_values(pixels.to(_DEVICE)).cpu()
+    batch, patches, _ = pixels.shape
+    expected = pixels.to(torch.float32).view(batch, patches, 3, -1)
+    expected = (
+        expected / 255.0 - torch.tensor(image_mean).view(1, 1, 3, 1)
+    ) / torch.tensor(image_std).view(1, 1, 3, 1)
+    expected = expected.view(shape)
+
+    torch.testing.assert_close(output, expected.to(torch.bfloat16))
+
+
+def test_nemotron_cpu_normalized_pixels_bypass_input_norm():
+    pixels = torch.randn(1, 3, 16, 16)
+    model = object.__new__(NemotronH_Nano_VL_V2)
+    nn.Module.__init__(model)
+    model.input_norm = None
+
+    assert model._normalize_pixel_values(pixels) is pixels
+
+
+@pytest.mark.parametrize("modality", ["image", "video", "dynamic_image"])
+def test_nemotron_processor_defers_normalization_to_device(
+    modality: str, default_vllm_config
+):
+    dynamic_image = modality == "dynamic_image"
+    config = SimpleNamespace(
+        force_image_size=16,
+        patch_size=4,
+        downsample_ratio=0.5,
+        use_thumbnail=False,
+        norm_mean=[0.485, 0.456, 0.406],
+        norm_std=[0.229, 0.224, 0.225],
+        dtype=torch.bfloat16,
+        vision_config=SimpleNamespace(
+            args={"min_num_patches": 4, "max_num_patches": 16} if dynamic_image else {}
+        ),
+        sound_config=None,
+    )
+    tokenizer = SimpleNamespace(encode=lambda *args, **kwargs: [1])
+    common_kwargs = {
+        "config": config,
+        "tokenizer": tokenizer,
+        "max_model_len": 1024,
+    }
+    cpu_processor = NanoNemotronVLProcessor(**common_kwargs)
+    device_kwargs = MultiModalConfig(
+        mm_device_do_normalize=True
+    ).merge_mm_processor_kwargs({})
+    device_processor = NanoNemotronVLProcessor(**common_kwargs, **device_kwargs)
+
+    pixels = np.random.default_rng(0).integers(0, 256, (4, 19, 23, 3), dtype=np.uint8)
+    if modality == "image":
+        image = Image.fromarray(pixels[0])
+        cpu_values = cpu_processor._images_to_pixel_values_lst([image], 1)[0]
+        raw_values = device_processor._images_to_pixel_values_lst([image], 1)[0]
+    elif modality == "video":
+        cpu_values = cpu_processor._videos_to_pixel_values_lst(
+            [pixels], dtype=torch.bfloat16
+        )[0]
+        raw_values = device_processor._videos_to_pixel_values_lst([pixels])[0]
+    else:
+        params = DynamicResolutionImageTiler.DynamicResolutionParams(
+            media=Image.fromarray(pixels[0]),
+            num_tiles=1,
+            num_embeddings=4,
+            patch_size=(4, 4),
+        )
+        assert cpu_processor.dynamic_tiler is not None
+        assert device_processor.dynamic_tiler is not None
+        cpu_values = cpu_processor.dynamic_tiler.apply_params(
+            params, dtype=torch.bfloat16
+        )[0]
+        raw_values = device_processor.dynamic_tiler.apply_params(params)[0]
+
+    assert raw_values.dtype == torch.uint8
+    assert cpu_values.dtype == torch.bfloat16
+    assert raw_values.shape == cpu_values.shape
+
+    # Normalizing on the device hands the encoder the same dtype as the CPU
+    # path. Values differ only where the uint8 resize rounds or clips the
+    # bicubic overshoot: a few elements by several levels, tiny on average.
+    model = object.__new__(NemotronH_Nano_VL_V2)
+    nn.Module.__init__(model)
+    model.input_norm = FusedMMInputNorm(
+        config.norm_mean, config.norm_std, 1.0 / 255.0
+    ).to(_DEVICE)
+    model.llm_dtype = config.dtype
+    if dynamic_image:
+        cpu_values = DynamicResolutionImageTiler.stack([cpu_values], config.patch_size)
+        raw_values = DynamicResolutionImageTiler.stack([raw_values], config.patch_size)
+    device_values = model._normalize_pixel_values(raw_values.to(_DEVICE)).cpu()
+    assert device_values.dtype == cpu_values.dtype
+    assert (device_values.float() - cpu_values.float()).abs().mean() < 0.01
+
+
+def test_nemotron_processor_rejects_partial_normalization():
+    with pytest.raises(
+        ValueError, match="requires do_rescale and do_normalize to have the same value"
+    ):
+        NanoNemotronVLProcessor(
+            config=None,
+            tokenizer=None,
+            max_model_len=1024,
+            do_rescale=False,
+            do_normalize=True,
+        )
 
 
 @pytest.mark.parametrize("input_key", ["image_embeds", "video_embeds"])
@@ -163,6 +345,7 @@ def _vision_norm_model(use_layernorm=True):
     model.language_model = _LanguageModel()
     model.vision_model = _VisionModel()
     model.sound_encoder = None
+    model.input_norm = None
     model.mlp1 = torch.nn.Identity()
     model.vision_final_layernorm = (
         torch.nn.LayerNorm(4, eps=1e-6, dtype=torch.bfloat16) if use_layernorm else None

vllm/model_executor/models/nano_nemotron_vl.py

@@ -28,6 +28,7 @@
 from vllm.logger import init_logger
 from vllm.lora.layers.base import BaseLayerWithLoRA
 from vllm.model_executor.layers.activation import ReLUSquaredActivation
+from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (
@@ -934,6 +935,7 @@ class NemotronH_Nano_VL_V2(
     SupportsMultiModalPruning,
     SupportsLoRA,
 ):
+    supports_mm_device_do_normalize = True
     requires_sequential_video_encoding = True
     """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet"""
 
@@ -992,6 +994,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         assert isinstance(llm_dtype, torch.dtype)
         self.llm_dtype = llm_dtype
         with self._mark_tower_model(vllm_config, {"image", "video", "audio"}):
+            self.input_norm: FusedMMInputNorm | None = (
+                FusedMMInputNorm(
+                    image_mean=config.norm_mean,
+                    image_std=config.norm_std,
+                    rescale_factor=1.0 / 255.0,
+                )
+                if multimodal_config.mm_device_do_normalize
+                else None
+            )
             self.vision_model = self.get_vit_model_from_radio_config(config).to(
                 llm_dtype
             )
@@ -1099,10 +1110,23 @@ def pixel_shuffle_dynamic_res(
 
         return x
 
+    def _normalize_pixel_values(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        if self.input_norm is None:
+            return pixel_values
+        shape = pixel_values.shape
+        if pixel_values.ndim == 4:
+            flattened = pixel_values.reshape(shape[0], -1)
+        else:
+            assert pixel_values.ndim == 3
+            flattened = pixel_values.reshape(-1, shape[-1])
+        normalized = self.input_norm(flattened, self.llm_dtype)
+        return normalized.reshape(shape)
+
     def extract_feature_dynamic(
         self, pixel_values: torch.Tensor, imgs_sizes: list[tuple[int, int]]
     ):
         """Dynamic resolution extract_feature for images."""
+        pixel_values = self._normalize_pixel_values(pixel_values)
         _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes)
         vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
         if self.vision_final_layernorm is not None:
@@ -1134,6 +1158,7 @@ def extract_feature(
         vit_embeds_list = []
         for i in range(0, N, micro_batch_size):
             chunk = pixel_values[i : i + micro_batch_size]
+            chunk = self._normalize_pixel_values(chunk)
             if num_frames is not None and T > 1:
                 _, vit_embeds = self.vision_model(chunk, num_frames=chunk.shape[0])
             else: