[MM] Keep device input normalization fused when encoder compilation is enabled (#59195)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Co-authored-by: Codex <codex@openai.com>
Jakub Zakrzewski committed
Sep 30, 2026 at 13:08 UTC
5a4351ab096268944bf29a9b804800066309baa9
2 files changed
+23
tests/model_executor/layers/test_mm_input_norm.py
+19
@@ -6,6 +6,8 @@
6
import pytest
7
import torch
8
9
+from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
10
+from vllm.config.compilation import CompilationMode
11
from vllm.model_executor.layers.fusion.mm_input_norm import (
12
FusedMMInputNorm,
13
IdentityInputNorm,
@@ -263,6 +265,23 @@ class TestFusedMMInputNormKernel:
265
class TestFusedMMInputNormConstruction:
266
"""Weight/bias buffer semantics at init time."""
267
268
+ def test_input_norm_is_always_enabled(self):
269
+ config = VllmConfig(
270
+ compilation_config=CompilationConfig(
271
+ mode=CompilationMode.VLLM_COMPILE,
272
+ compile_mm_encoder=True,
273
+ custom_ops=["none"],
274
+ )
275
+ )
276
+ with set_current_vllm_config(config):
277
+ FusedMMInputNorm(_RGB_MEAN, _RGB_STD, _RGB_RESCALE)
278
+ assert (
279
+ config.compilation_config.enabled_custom_ops["fused_mm_input_norm"] == 1
280
+ )
281
+ assert not config.compilation_config.disabled_custom_ops[
282
+ "fused_mm_input_norm"
283
+ ]
284
+
285
def test_identity_config_buffers(self):
286
"""Numerically identity parameters still build plain weight/bias
287
buffers; the identity shortcut lives in ``from_model_config``."""
vllm/model_executor/layers/fusion/mm_input_norm.py
+4
@@ -246,6 +246,10 @@ class FusedMMInputNorm(CustomOp):
246
self.register_buffer("weight", (rescale_factor / std).to(device))
247
self.register_buffer("bias", (-mean / std).to(device))
248
249
+ @classmethod
250
+ def enabled(cls) -> bool:
251
+ return True
252
+
253
@property
254
def input_dtype(self) -> torch.dtype | None:
255
return torch.uint8