[BugFix] Revert incorrect MM keep_on_cpu=True changes (#52881)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Aug 19, 2026 at 00:07 UTC
ee117307512f7ab0fe5ce70f8a2963bb83dfc442
7 files changed
+17
-12
vllm/model_executor/models/ernie45_vl.py
+8
-4
@@ -1726,7 +1726,9 @@ class Ernie4_5_VLMoeForConditionalGeneration(
1726
# Eager fallback: run the full pipeline (ViT + resampler). The result
1727
# is scattered directly, so it must be the post-merge embeddings.
1728
pixel_values = mm_kwargs["pixel_values"].type(self.vision_model.dtype)
1729
- grid_thw = mm_kwargs["image_grid_thw"].to(self.vision_model.device)
1729
+ grid_thw = mm_kwargs["image_grid_thw"].to(
1730
+ self.vision_model.device, non_blocking=True
1731
+ )
1732
image_features = self.vision_model(pixel_values, grid_thw)
1733
return self.resampler_model(image_features, grid_thw)
1734
@@ -1744,11 +1746,13 @@ class Ernie4_5_VLMoeForConditionalGeneration(
1746
# the actual batch grid_thw, then scatter the post-merge embeddings.
1747
# Ernie only uses the single "default" encoder path.
1748
output = outputs["default"]
1747
- grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device)
1749
+ grid_thw_cpu = batch_mm_kwargs["image_grid_thw"]
1750
+ grid_thw = grid_thw_cpu.to(output.device, non_blocking=True)
1751
# The valid token count slices the graph output for the eager
1752
# resampler call, so it has to come back to the host.
1750
- with gpu_sync_allowed():
1751
- num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum())
1753
+ num_valid = int(
1754
+ (grid_thw_cpu[:, 0] * grid_thw_cpu[:, 1] * grid_thw_cpu[:, 2]).sum()
1755
+ )
1756
image_embeds = self.resampler_model(output[:num_valid], grid_thw)
1757
scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone)
1758
vllm/model_executor/models/fireredlid.py
+1
-1
@@ -482,7 +482,7 @@ class FireRedLIDMultiModalProcessor(
482
) -> Mapping[str, MultiModalFieldConfig]:
483
return dict(
484
input_features=MultiModalFieldConfig.batched("audio"),
485
- speech_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
485
+ speech_lengths=MultiModalFieldConfig.batched("audio"),
486
fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
487
)
488
vllm/model_executor/models/isaac.py
+3
-1
@@ -961,7 +961,9 @@ class IsaacForConditionalGeneration(
961
device = next(self.language_model.parameters()).device
962
dtype = self.vision_embedding.linear_fc1.weight.dtype
963
pixel_values = pixel_values.to(device=device, dtype=dtype)
964
- spatial_grids = image_grid_thw[:, 1:3].to(device, dtype=torch.int32)
964
+ spatial_grids = image_grid_thw[:, 1:3].to(
965
+ device, dtype=torch.int32, non_blocking=True
966
+ )
967
968
vision_embeddings = self.vision_embedding((pixel_values, spatial_grids))
969
merge_size = self.config.vision_config.pixel_shuffle_scale_factor
vllm/model_executor/models/llava_onevision2.py
+2
-2
@@ -435,7 +435,7 @@ def _create_field_factory(
435
# OV2 first-class MM kwarg: per-patch (t,h,w)
436
# positions required by the 3-D vision RoPE.
437
patch_positions=MultiModalFieldConfig.flat_from_sizes(
438
- "image", image_pixel_grid_sizes, keep_on_cpu=True
438
+ "image", image_pixel_grid_sizes
439
),
440
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
441
"video", video_patch_sizes
@@ -444,7 +444,7 @@ def _create_field_factory(
444
"video", video_num_frames, keep_on_cpu=True
445
),
446
patch_positions_videos=MultiModalFieldConfig.flat_from_sizes(
447
- "video", video_patch_sizes, keep_on_cpu=True
447
+ "video", video_patch_sizes
448
),
449
video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
450
frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
vllm/model_executor/models/minicpmv4_6.py
+2
-2
@@ -75,11 +75,11 @@ from .vision import is_vit_use_data_parallel
75
def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]):
76
fields = dict(
77
pixel_values=MultiModalFieldConfig.batched("image"),
78
- tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
78
+ tgt_sizes=MultiModalFieldConfig.batched("image"),
79
image_embeds=MultiModalFieldConfig.batched("image"),
80
video_pixel_values=MultiModalFieldConfig.batched("video"),
81
video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
82
- video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
82
+ video_tgt_sizes=MultiModalFieldConfig.batched("video"),
83
video_embeds=MultiModalFieldConfig.batched("video"),
84
)
85
if "use_vit_merger" in hf_inputs:
vllm/model_executor/models/moss_transcribe_diarize.py
-1
@@ -352,7 +352,6 @@ def _mtd_field_config(
352
audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes(
353
"audio",
354
audio_chunk_counts,
355
- keep_on_cpu=True,
355
),
356
audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
357
audio_token_lengths=MultiModalFieldConfig.batched(
vllm/model_executor/models/qwen2_audio.py
+1
-1
@@ -130,7 +130,7 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
130
return dict(
131
audio_embeds=MultiModalFieldConfig.batched("audio"),
132
input_features=MultiModalFieldConfig.batched("audio"),
133
- feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
133
+ feature_attention_mask=MultiModalFieldConfig.batched("audio"),
134
)
135
136