[BugFix] Revert incorrect MM keep_on_cpu=True changes (#52881)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

vllm/model_executor/models/ernie45_vl.py

@@ -1726,7 +1726,9 @@ class Ernie4_5_VLMoeForConditionalGeneration(
         # Eager fallback: run the full pipeline (ViT + resampler). The result
         # is scattered directly, so it must be the post-merge embeddings.
         pixel_values = mm_kwargs["pixel_values"].type(self.vision_model.dtype)
-        grid_thw = mm_kwargs["image_grid_thw"].to(self.vision_model.device)
+        grid_thw = mm_kwargs["image_grid_thw"].to(
+            self.vision_model.device, non_blocking=True
+        )
         image_features = self.vision_model(pixel_values, grid_thw)
         return self.resampler_model(image_features, grid_thw)
 
@@ -1744,11 +1746,13 @@ class Ernie4_5_VLMoeForConditionalGeneration(
         # the actual batch grid_thw, then scatter the post-merge embeddings.
         # Ernie only uses the single "default" encoder path.
         output = outputs["default"]
-        grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device)
+        grid_thw_cpu = batch_mm_kwargs["image_grid_thw"]
+        grid_thw = grid_thw_cpu.to(output.device, non_blocking=True)
         # The valid token count slices the graph output for the eager
         # resampler call, so it has to come back to the host.
-        with gpu_sync_allowed():
-            num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum())
+        num_valid = int(
+            (grid_thw_cpu[:, 0] * grid_thw_cpu[:, 1] * grid_thw_cpu[:, 2]).sum()
+        )
         image_embeds = self.resampler_model(output[:num_valid], grid_thw)
         scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone)
 

vllm/model_executor/models/fireredlid.py

@@ -482,7 +482,7 @@ class FireRedLIDMultiModalProcessor(
     ) -> Mapping[str, MultiModalFieldConfig]:
         return dict(
             input_features=MultiModalFieldConfig.batched("audio"),
-            speech_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
+            speech_lengths=MultiModalFieldConfig.batched("audio"),
             fake_token_lengths=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
         )
 

vllm/model_executor/models/isaac.py

@@ -961,7 +961,9 @@ class IsaacForConditionalGeneration(
         device = next(self.language_model.parameters()).device
         dtype = self.vision_embedding.linear_fc1.weight.dtype
         pixel_values = pixel_values.to(device=device, dtype=dtype)
-        spatial_grids = image_grid_thw[:, 1:3].to(device, dtype=torch.int32)
+        spatial_grids = image_grid_thw[:, 1:3].to(
+            device, dtype=torch.int32, non_blocking=True
+        )
 
         vision_embeddings = self.vision_embedding((pixel_values, spatial_grids))
         merge_size = self.config.vision_config.pixel_shuffle_scale_factor

vllm/model_executor/models/llava_onevision2.py

@@ -435,7 +435,7 @@ def _create_field_factory(
             # OV2 first-class MM kwarg: per-patch (t,h,w)
             # positions required by the 3-D vision RoPE.
             patch_positions=MultiModalFieldConfig.flat_from_sizes(
-                "image", image_pixel_grid_sizes, keep_on_cpu=True
+                "image", image_pixel_grid_sizes
             ),
             pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
                 "video", video_patch_sizes
@@ -444,7 +444,7 @@ def _create_field_factory(
                 "video", video_num_frames, keep_on_cpu=True
             ),
             patch_positions_videos=MultiModalFieldConfig.flat_from_sizes(
-                "video", video_patch_sizes, keep_on_cpu=True
+                "video", video_patch_sizes
             ),
             video_num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
             frame_timestamps=MultiModalFieldConfig.batched("video", keep_on_cpu=True),

vllm/model_executor/models/minicpmv4_6.py

@@ -75,11 +75,11 @@ from .vision import is_vit_use_data_parallel
 def _minicpmv4_6_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     fields = dict(
         pixel_values=MultiModalFieldConfig.batched("image"),
-        tgt_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
+        tgt_sizes=MultiModalFieldConfig.batched("image"),
         image_embeds=MultiModalFieldConfig.batched("image"),
         video_pixel_values=MultiModalFieldConfig.batched("video"),
         video_image_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
-        video_tgt_sizes=MultiModalFieldConfig.batched("video", keep_on_cpu=True),
+        video_tgt_sizes=MultiModalFieldConfig.batched("video"),
         video_embeds=MultiModalFieldConfig.batched("video"),
     )
     if "use_vit_merger" in hf_inputs:

vllm/model_executor/models/moss_transcribe_diarize.py

@@ -352,7 +352,6 @@ def _mtd_field_config(
             audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes(
                 "audio",
                 audio_chunk_counts,
-                keep_on_cpu=True,
             ),
             audio_chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
             audio_token_lengths=MultiModalFieldConfig.batched(

vllm/model_executor/models/qwen2_audio.py

@@ -130,7 +130,7 @@ def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
     return dict(
         audio_embeds=MultiModalFieldConfig.batched("audio"),
         input_features=MultiModalFieldConfig.batched("audio"),
-        feature_attention_mask=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
+        feature_attention_mask=MultiModalFieldConfig.batched("audio"),
     )