[Bugfix] Make `mm_device_do_normalize` encoder-cudagraph safe (#55370)
Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
cjackal committed
Sep 10, 2026 at 02:43 UTC
22d95d1adc24d5834ec45bb697170a76cf4c9a95
3 files changed
+28
-18
tests/models/multimodal/generation/test_qwen2_vl.py
+6
-2
@@ -126,8 +126,10 @@ def batch_make_image_embeddings(
126
# image to pixel values
127
image_processor = processor.image_processor
128
129
+ mm_config = llm.get_llm().model_config.multimodal_config
130
+ mm_kwargs = mm_config.merge_mm_processor_kwargs({})
131
preprocess_result = image_processor.preprocess(
130
- images=images, return_tensors="pt"
132
+ images=images, return_tensors="pt", **mm_kwargs
133
).data
134
pixel_values = preprocess_result["pixel_values"]
135
image_grid_thw = preprocess_result["image_grid_thw"]
@@ -210,8 +212,10 @@ def batch_make_video_embeddings(
212
# video to pixel values
213
video_processor = processor.video_processor
214
215
+ mm_config = llm.get_llm().model_config.multimodal_config
216
+ mm_kwargs = mm_config.merge_mm_processor_kwargs({})
217
preprocess_result = video_processor.preprocess(
214
- videos=videos, return_tensors="pt"
218
+ videos=videos, return_tensors="pt", **mm_kwargs
219
).data
220
pixel_values = preprocess_result["pixel_values_videos"]
221
video_grid_thw = preprocess_result["video_grid_thw"]
vllm/model_executor/models/qwen2_5_vl.py
+11
-8
@@ -630,6 +630,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
630
vision_config: Qwen2_5_VLVisionConfig,
631
norm_eps: float = 1e-6,
632
quant_config: QuantizationConfig | None = None,
633
+ input_norm: nn.Module | None = None,
634
prefix: str = "",
635
) -> None:
636
super().__init__()
@@ -661,6 +662,9 @@ class Qwen2_5_VisionTransformer(nn.Module):
662
in_channels=in_channels,
663
hidden_size=self.hidden_size,
664
)
665
+ self.input_norm = (
666
+ input_norm if input_norm is not None else FusedInputNorm.identity()
667
+ )
668
669
norm_layer = partial(RMSNorm, eps=norm_eps)
670
head_dim = self.hidden_size // self.num_heads
@@ -1076,7 +1080,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
1080
*,
1081
encoder_metadata: dict[str, torch.Tensor] | None = None,
1082
) -> torch.Tensor:
1079
- hidden_states = x.to(device=self.device, dtype=self.dtype)
1083
+ hidden_states = self.input_norm(x.to(device=self.device), self.dtype)
1084
hidden_states = self.patch_embed(hidden_states)
1085
1086
seq_len = hidden_states.shape[0]
@@ -1359,9 +1363,9 @@ class Qwen2_5_VLForConditionalGeneration(
1363
vision_config=config.vision_config,
1364
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1365
quant_config=self.quant_config,
1366
+ input_norm=FusedInputNorm.from_model_config(self.model_config),
1367
prefix=maybe_prefix(prefix, "visual"),
1368
)
1364
- self.input_norm = FusedInputNorm.from_model_config(self.model_config)
1369
1370
with self._mark_language_model(vllm_config):
1371
self.language_model = init_vllm_registered_model(
@@ -1436,11 +1440,13 @@ class Qwen2_5_VLForConditionalGeneration(
1440
image_embeds = image_input["image_embeds"].type(self.visual.dtype)
1441
else:
1442
pixel_values = image_input["pixel_values"]
1439
- pixel_values = self.input_norm(pixel_values, self.visual.dtype)
1443
1444
if self.use_data_parallel:
1445
return run_dp_sharded_mrope_vision_model(
1443
- self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
1446
+ self.visual,
1447
+ pixel_values.type(self.visual.dtype),
1448
+ grid_thw_list,
1449
+ rope_type="rope_3d",
1450
)
1451
else:
1452
image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
@@ -1494,14 +1500,11 @@ class Qwen2_5_VLForConditionalGeneration(
1500
video_embeds = video_input["video_embeds"].type(self.visual.dtype)
1501
else:
1502
pixel_values_videos = video_input["pixel_values_videos"]
1497
- pixel_values_videos = self.input_norm(
1498
- pixel_values_videos, self.visual.dtype
1499
- )
1503
1504
if self.use_data_parallel:
1505
return run_dp_sharded_mrope_vision_model(
1506
self.visual,
1504
- pixel_values_videos,
1507
+ pixel_values_videos.type(self.visual.dtype),
1508
grid_thw_list,
1509
rope_type="rope_3d",
1510
)
vllm/model_executor/models/qwen2_vl.py
+11
-8
@@ -537,6 +537,7 @@ class Qwen2VisionTransformer(nn.Module):
537
vision_config: Qwen2VLVisionConfig,
538
norm_eps: float = 1e-6,
539
quant_config: QuantizationConfig | None = None,
540
+ input_norm: nn.Module | None = None,
541
prefix: str = "",
542
) -> None:
543
super().__init__()
@@ -564,6 +565,9 @@ class Qwen2VisionTransformer(nn.Module):
565
in_channels=in_channels,
566
embed_dim=embed_dim,
567
)
568
+ self.input_norm = (
569
+ input_norm if input_norm is not None else FusedInputNorm.identity()
570
+ )
571
572
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
573
head_dim = embed_dim // num_heads
@@ -720,7 +724,7 @@ class Qwen2VisionTransformer(nn.Module):
724
encoder_metadata: dict[str, torch.Tensor] | None = None,
725
) -> torch.Tensor:
726
# patchify
723
- x = x.to(device=self.device, dtype=self.dtype)
727
+ x = self.input_norm(x.to(device=self.device), self.dtype)
728
x = self.patch_embed(x)
729
730
if encoder_metadata is None:
@@ -1301,9 +1305,9 @@ class Qwen2VLForConditionalGeneration(
1305
config.vision_config,
1306
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1307
quant_config=quant_config,
1308
+ input_norm=FusedInputNorm.from_model_config(self.model_config),
1309
prefix=maybe_prefix(prefix, "visual"),
1310
)
1306
- self.input_norm = FusedInputNorm.from_model_config(self.model_config)
1311
1312
with self._mark_language_model(vllm_config):
1313
self.language_model = init_vllm_registered_model(
@@ -1374,11 +1378,13 @@ class Qwen2VLForConditionalGeneration(
1378
image_embeds = image_input["image_embeds"].type(self.visual.dtype)
1379
else:
1380
pixel_values = image_input["pixel_values"]
1377
- pixel_values = self.input_norm(pixel_values, self.visual.dtype)
1381
1382
if self.use_data_parallel:
1383
return run_dp_sharded_mrope_vision_model(
1381
- self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"
1384
+ self.visual,
1385
+ pixel_values.type(self.visual.dtype),
1386
+ grid_thw.tolist(),
1387
+ rope_type="rope_3d",
1388
)
1389
else:
1390
image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
@@ -1398,13 +1404,10 @@ class Qwen2VLForConditionalGeneration(
1404
video_embeds = video_input["video_embeds"].type(self.visual.dtype)
1405
else:
1406
pixel_values_videos = video_input["pixel_values_videos"]
1401
- pixel_values_videos = self.input_norm(
1402
- pixel_values_videos, self.visual.dtype
1403
- )
1407
if self.use_data_parallel:
1408
return run_dp_sharded_mrope_vision_model(
1409
self.visual,
1407
- pixel_values_videos,
1410
+ pixel_values_videos.type(self.visual.dtype),
1411
grid_thw.tolist(),
1412
rope_type="rope_3d",
1413
)