[Model] Extend device-side mm normalization to Qwen3VL/Qwen3.5/Qwen4Next (#51289)
Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
cjackal committed
Sep 29, 2026 at 12:44 UTC
491f44adfa4220ad9b958e9b61e6d49df30fd307
11 files changed
+87
-25
docs/design/mm_processing.md
+11
-4
@@ -78,10 +78,17 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:
78
- The flag is **enabled by default** for all models that support it.
79
- Currently, it’s on by default for these architectures:
80
81
-| name | Architecture | Example HF Models |
82
-|--------------|--------------------------------------|-------------------------------------|
83
-| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |
84
-| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
81
+| name | Architecture | Example HF Models |
82
+|--------------------|------------------------------------------|----------------------------------------|
83
+| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` |
84
+| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` |
85
+| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |
86
+| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
87
+| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. |
88
+| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. |
89
+| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. |
90
+| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. |
91
+| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` |
92
93
#### Key Properties and Gains
94
tests/models/multimodal/generation_ppl_test/ppl_utils.py
+8
-4
@@ -34,6 +34,7 @@ def vqa_ppl_test(
34
mm_processor_kwargs = {}
35
36
processor = AutoProcessor.from_pretrained(model_info.name)
37
+ image_token_id = processor.image_token_id
38
39
images = []
40
prompts = []
@@ -99,8 +100,10 @@ def vqa_ppl_test(
100
for token_data in token_datas[1:]:
101
assert token_data is not None
102
assert len(token_data) == 1
102
- token_log_prob = list(token_data.values())[0].logprob
103
- token_log_probs.append(token_log_prob)
103
+ token_id, token_log_prob = next(iter(token_data.items()))
104
+ if token_id == image_token_id:
105
+ continue
106
+ token_log_probs.append(token_log_prob.logprob)
107
108
neg_log_likelihood = -torch.tensor(
109
token_log_probs, dtype=torch.float32, device="cpu"
@@ -132,11 +135,12 @@ def vqa_ppl_test(
135
inputs = inputs.to("cuda")
136
input_ids = inputs["input_ids"]
137
135
- outputs = hf_model.model(**inputs, labels=input_ids)
138
+ labels = input_ids.masked_fill(input_ids == image_token_id, -100)
139
+ outputs = hf_model.model(**inputs, labels=labels)
140
neg_log_likelihood = outputs.loss
141
neg_log_likelihood = neg_log_likelihood.to(torch.float32).cpu()
142
139
- num_loss_tokens = input_ids.shape[1] - 1
143
+ num_loss_tokens = int((labels[:, 1:] != -100).sum())
144
nll_sum += neg_log_likelihood * num_loss_tokens
145
n_tokens += num_loss_tokens
146
tests/models/multimodal/generation_ppl_test/test_qwen.py
+35
-10
@@ -8,26 +8,51 @@ from tests.models.utils import GenerateModelInfo
8
from .ppl_utils import vqa_ppl_test
9
10
MODELS = [
11
- GenerateModelInfo("Qwen/Qwen2-VL-2B-Instruct", hf_ppl=41081356.0),
12
- GenerateModelInfo("Qwen/Qwen2.5-VL-3B-Instruct", hf_ppl=18330016.0),
11
+ (
12
+ GenerateModelInfo("Qwen/Qwen2-VL-2B-Instruct", hf_ppl=9.810060501098633),
13
+ 28,
14
+ {},
15
+ ),
16
+ (
17
+ GenerateModelInfo("Qwen/Qwen2.5-VL-3B-Instruct", hf_ppl=22.17719078063965),
18
+ 28,
19
+ {},
20
+ ),
21
+ (
22
+ GenerateModelInfo("Qwen/Qwen3-VL-4B-Instruct", hf_ppl=11.916606903076172),
23
+ 32,
24
+ {},
25
+ ),
26
+ (
27
+ GenerateModelInfo("Qwen/Qwen3.5-4B", hf_ppl=6.478141784667969),
28
+ 32,
29
+ {"enable_prefix_caching": False},
30
+ ),
31
]
32
33
16
-mm_processor_kwargs = {
17
- "min_pixels": 28 * 28,
18
- "max_pixels": 1280 * 28 * 28,
19
-}
34
+def _pixel_factor(factor: int):
35
+ """`factor = patch_size * merge_size` for Qwen models"""
36
+ return {"min_pixels": factor**2, "max_pixels": 1280 * factor**2}
37
38
22
-@pytest.mark.parametrize("model_info", MODELS)
39
+@pytest.mark.parametrize("model_info,factor,extra_kwargs", MODELS)
40
@pytest.mark.parametrize("mm_device_do_normalize", [True, False])
41
def test_ppl(
25
- hf_runner, vllm_runner, model_info: GenerateModelInfo, mm_device_do_normalize: bool
42
+ hf_runner,
43
+ vllm_runner,
44
+ model_info: GenerateModelInfo,
45
+ factor: int,
46
+ extra_kwargs: dict,
47
+ mm_device_do_normalize: bool,
48
):
49
vqa_ppl_test(
50
hf_runner,
51
vllm_runner,
52
model_info,
31
- vllm_extra_kwargs={"mm_device_do_normalize": mm_device_do_normalize},
32
- mm_processor_kwargs=mm_processor_kwargs,
53
+ vllm_extra_kwargs={
54
+ "mm_device_do_normalize": mm_device_do_normalize,
55
+ **extra_kwargs,
56
+ },
57
+ mm_processor_kwargs=_pixel_factor(factor),
58
)
vllm/model_executor/models/colqwen3.py
+2
@@ -134,6 +134,8 @@ class ColQwen3Model(Qwen3VLForConditionalGeneration, SupportsLateInteraction):
134
135
# Mark this as a pooling model so vLLM routes to pooler path
136
is_pooling_model = True
137
+ # Mark as false to avoid collision with the checkpoint's custom code path
138
+ supports_mm_device_do_normalize = False
139
140
# Override hf_to_vllm_mapper to handle ColQwen3 weight naming.
141
# NOTE: WeightsMapper applies ALL matching prefix rules sequentially
vllm/model_executor/models/interns1_pro.py
+2
@@ -45,6 +45,7 @@ from vllm.model_executor.layers.fused_moe import (
45
FusedMoEFactory,
46
GateLinear,
47
)
48
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
49
from vllm.model_executor.layers.layernorm import RMSNorm
50
from vllm.model_executor.layers.linear import (
51
MergedColumnParallelLinear,
@@ -575,6 +576,7 @@ class InternS1ProForConditionalGeneration(
576
self.visual = Qwen3_VisionTransformer(
577
config.vision_config,
578
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
579
+ input_norm=build_mm_input_norm(self.model_config),
580
prefix=maybe_prefix(prefix, "visual"),
581
)
582
vllm/model_executor/models/interns2_mobius.py
+2
@@ -21,6 +21,7 @@ from vllm.model_executor.layers.fused_moe import (
21
FusedMoEFactory,
22
GateLinear,
23
)
24
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
25
from vllm.model_executor.layers.linear import ReplicatedLinear
26
from vllm.model_executor.layers.logits_processor import LogitsProcessor
27
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
@@ -496,6 +497,7 @@ class InternS2MobiusForConditionalGeneration(InternS2PreviewForConditionalGenera
497
config.vision_config,
498
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
499
quant_config=quant_config,
500
+ input_norm=build_mm_input_norm(self.model_config),
501
prefix=maybe_prefix(prefix, "visual"),
502
)
503
vllm/model_executor/models/qwen3_5.py
+3
@@ -39,6 +39,7 @@ from vllm.logger import init_logger
39
from vllm.model_executor.layers.fused_moe.utils import (
40
is_model_fused_shared_expert_compatible,
41
)
42
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
43
from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm
44
from vllm.model_executor.layers.logits_processor import LogitsProcessor
45
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
@@ -517,6 +518,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
518
config.vision_config,
519
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
520
quant_config=quant_config,
521
+ input_norm=build_mm_input_norm(self.model_config),
522
prefix=maybe_prefix(prefix, "visual"),
523
)
524
@@ -739,6 +741,7 @@ class Qwen3_5MoeForConditionalGeneration(
741
config.vision_config,
742
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
743
quant_config=quant_config,
744
+ input_norm=build_mm_input_norm(self.model_config),
745
prefix=maybe_prefix(prefix, "visual"),
746
)
747
vllm/model_executor/models/qwen3_vl.py
+18
-7
@@ -67,6 +67,10 @@ from vllm.model_executor.layers.attention.mm_encoder_attention import (
67
MMEncoderAttention,
68
)
69
from vllm.model_executor.layers.conv import Conv3dLayer
70
+from vllm.model_executor.layers.fusion.mm_input_norm import (
71
+ IdentityInputNorm,
72
+ build_mm_input_norm,
73
+)
74
from vllm.model_executor.layers.linear import (
75
ColumnParallelLinear,
76
RowParallelLinear,
@@ -566,6 +570,7 @@ class Qwen3_VisionTransformer(nn.Module):
570
vision_config: Qwen3VLVisionConfig,
571
norm_eps: float = 1e-6,
572
quant_config: QuantizationConfig | None = None,
573
+ input_norm: nn.Module | None = None,
574
prefix: str = "",
575
) -> None:
576
super().__init__()
@@ -605,6 +610,7 @@ class Qwen3_VisionTransformer(nn.Module):
610
in_channels=vision_config.in_channels,
611
hidden_size=self.hidden_size,
612
)
613
+ self.input_norm = input_norm if input_norm is not None else IdentityInputNorm()
614
615
self.pos_embed = nn.Embedding(self.num_position_embeddings, self.hidden_size)
616
@@ -858,7 +864,9 @@ class Qwen3_VisionTransformer(nn.Module):
864
*,
865
encoder_metadata: dict[str, torch.Tensor] | None = None,
866
) -> torch.Tensor:
861
- hidden_states = x.to(device=self.device, dtype=self.dtype, non_blocking=True)
867
+ hidden_states = self.input_norm(
868
+ x.to(device=self.device, non_blocking=True), self.dtype
869
+ )
870
hidden_states = self.patch_embed(hidden_states)
871
872
if encoder_metadata is None:
@@ -1813,6 +1821,7 @@ class Qwen3VLForConditionalGeneration(
1821
1822
supports_encoder_tp_data = True
1823
supports_tower_connector_lora = True
1824
+ supports_mm_device_do_normalize = True
1825
1826
supported_video_pruning_methods = ("evs", "vidcom2")
1827
@@ -1863,6 +1872,7 @@ class Qwen3VLForConditionalGeneration(
1872
config.vision_config,
1873
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1874
quant_config=quant_config,
1875
+ input_norm=build_mm_input_norm(self.model_config),
1876
prefix=maybe_prefix(prefix, "visual"),
1877
)
1878
@@ -2172,8 +2182,11 @@ class Qwen3VLForConditionalGeneration(
2182
flattened_patch_size = (
2183
in_channels * temporal_patch_size * patch_size * patch_size
2184
)
2175
- dummy_pixel_values = torch.randn(
2176
- total_patches, flattened_patch_size, device=device, dtype=dtype
2185
+ dummy_pixel_values = torch.zeros(
2186
+ total_patches,
2187
+ flattened_patch_size,
2188
+ device=device,
2189
+ dtype=self.visual.input_norm.input_dtype or dtype,
2190
)
2191
2192
# Override max_seqlen with a safe upper bound for capture.
@@ -2307,7 +2320,7 @@ class Qwen3VLForConditionalGeneration(
2320
if image_input["type"] == "image_embeds":
2321
image_embeds = image_input["image_embeds"].type(self.visual.dtype)
2322
else:
2310
- pixel_values = image_input["pixel_values"].type(self.visual.dtype)
2323
+ pixel_values = image_input["pixel_values"]
2324
if self.use_data_parallel:
2325
return run_dp_sharded_mrope_vision_model(
2326
self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"
@@ -2329,9 +2342,7 @@ class Qwen3VLForConditionalGeneration(
2342
if video_input["type"] == "video_embeds":
2343
video_embeds = video_input["video_embeds"].type(self.visual.dtype)
2344
else:
2332
- pixel_values_videos = video_input["pixel_values_videos"].type(
2333
- self.visual.dtype
2334
- )
2345
+ pixel_values_videos = video_input["pixel_values_videos"]
2346
if self.use_data_parallel:
2347
grid_thw_list = grid_thw.tolist()
2348
return run_dp_sharded_mrope_vision_model(
vllm/model_executor/models/qwen3_vl_moe.py
+2
@@ -36,6 +36,7 @@ from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import VllmConfig
37
from vllm.distributed import get_pp_group
38
from vllm.logger import init_logger
39
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
41
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
42
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -254,6 +255,7 @@ class Qwen3VLMoeForConditionalGeneration(
255
config.vision_config,
256
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
257
quant_config=quant_config,
258
+ input_norm=build_mm_input_norm(self.model_config),
259
prefix=maybe_prefix(prefix, "visual"),
260
)
261
vllm/models/qwen4_exp/amd/model.py
+2
@@ -14,6 +14,7 @@ from vllm.distributed import get_pp_group
14
from vllm.model_executor.layers.fused_moe.utils import (
15
is_model_fused_shared_expert_compatible,
16
)
17
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
18
from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
20
QwenGatedDeltaNetAttention,
@@ -897,6 +898,7 @@ class Qwen4ExpForConditionalGeneration(
898
config.vision_config,
899
norm_eps=config.text_config.rms_norm_eps,
900
quant_config=quant_config,
901
+ input_norm=build_mm_input_norm(self.model_config),
902
prefix=maybe_prefix(prefix, "visual"),
903
)
904
vllm/models/qwen4_exp/nvidia/model.py
+2
@@ -13,6 +13,7 @@ from vllm.distributed import get_pp_group
13
from vllm.model_executor.layers.fused_moe.utils import (
14
is_model_fused_shared_expert_compatible,
15
)
16
+from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm
17
from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
19
QwenGatedDeltaNetAttention,
@@ -939,6 +940,7 @@ class Qwen4ExpForConditionalGeneration(
940
config.vision_config,
941
norm_eps=config.text_config.rms_norm_eps,
942
quant_config=quant_config,
943
+ input_norm=build_mm_input_norm(self.model_config),
944
prefix=maybe_prefix(prefix, "visual"),
945
)
946