[Perf][Multimodal] Sort multimodal encoder inputs by modality (#60196)

Signed-off-by: Jaebok Lee <jaebok9541@naver.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/v1/worker/test_encoder_runner.py

@@ -78,7 +78,7 @@ def test_encoder_graph_unsupported_modality_uses_eager_output():
         "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs",
         return_value=[("audio", 1, {"audio_values": torch.zeros(2)})],
     ):
-        result = runner.execute_mm_encoder([])
+        result = runner.execute_mm_encoder([("audio", MagicMock())])
     assert len(result) == 1
     assert result[0] is output
     manager.execute.assert_not_called()
@@ -290,6 +290,21 @@ def test_execute_mm_encoder_caches_outputs_without_gathering():
     state.encoder_runner.gather_mm_embeddings.assert_not_called()
 
 
+def test_execute_mm_encoder_sorts_by_modality():
+    """Items are encoded sorted by modality; outputs keep the input order."""
+    mm_kwargs = [("video", 0), ("image", 1), ("video", 2), ("image", 3)]
+    runner = _make_runner([], [])
+    runner.model = MagicMock()
+    runner.model.embed_multimodal.side_effect = lambda x: [torch.full((1, 1), x)]
+    with patch(
+        "vllm.v1.worker.gpu.mm.encoder_runner.group_and_batch_mm_kwargs",
+        side_effect=lambda items, **_: [(m, 1, {"x": x}) for m, x in items],
+    ) as group:
+        outputs = runner.execute_mm_encoder(mm_kwargs)
+    assert group.call_args.args[0] == [mm_kwargs[i] for i in (1, 3, 0, 2)]
+    assert [int(o) for o in outputs] == [0, 1, 2, 3]
+
+
 def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
     """A step that schedules no encoder input must not touch the encoder."""
     cache = EncoderCache()

tests/v1/worker/test_gpu_model_runner.py

@@ -624,6 +624,70 @@ def test_set_active_mm_loras_builds_tower_and_connector_mappings():
     assert connector_mapping.index_mapping == ((7,) * 14 + (7,) * 13 + (0,) * 12)
 
 
+def test_set_active_mm_loras_follows_modality_sorted_encoder_order():
+    model = Mock()
+    model.get_mm_lora_token_counts.side_effect = (
+        lambda *, modality, mm_kwargs, num_mm_embeds: (num_mm_embeds, None)
+    )
+    lora_manager = Mock()
+    lora_manager.supports_tower_connector_lora.return_value = True
+
+    encoder_cache = EncoderCache()
+    for req_id, modality, length in [("req-a", "video", 1), ("req-b", "image", 2)]:
+        encoder_cache.mm_features[req_id] = [
+            MultiModalFeatureSpec(
+                data=None,
+                modality=modality,
+                identifier=f"{req_id}-0",
+                mm_position=PlaceholderRange(offset=0, length=length),
+            )
+        ]
+    lora_state = LoraState(max_num_reqs=2)
+    lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora")
+    lora_state.add_request("req-a", 0, lora_request)
+    lora_state.add_request("req-b", 1, None)
+
+    set_active_mm_loras(
+        model=model,
+        lora_manager=lora_manager,
+        encoder_cache=encoder_cache,
+        req_id_to_index={"req-a": 0, "req-b": 1},
+        lora_state=lora_state,
+        scheduled_encoder_inputs={"req-a": [0], "req-b": [0]},
+    )
+
+    # The image is encoded before the video, so its tokens come first.
+    _, tower_mapping = lora_manager.set_active_adapters.call_args_list[0].args
+    assert tower_mapping.index_mapping == (0, 0, 7)
+
+
+def test_batch_mm_inputs_from_scheduler_sorts_by_modality():
+    features = {
+        f"req{i}": MultiModalFeatureSpec(
+            data=Mock(),
+            modality=m,
+            identifier=f"hash{i}",
+            mm_position=PlaceholderRange(offset=i, length=1),
+        )
+        for i, m in enumerate(["video", "image", "video"])
+    }
+    runner = SimpleNamespace(
+        requests={r: SimpleNamespace(mm_features=[f]) for r, f in features.items()}
+    )
+    scheduler_output = SimpleNamespace(
+        scheduled_encoder_inputs={r: [0] for r in features}
+    )
+
+    mm_hashes, mm_kwargs, mm_lora_refs = GPUModelRunner._batch_mm_inputs_from_scheduler(
+        runner, scheduler_output
+    )
+
+    order = ["req1", "req0", "req2"]
+    assert mm_hashes == [features[r].identifier for r in order]
+    assert mm_kwargs == [(features[r].modality, features[r].data) for r in order]
+    assert mm_lora_refs == [(r, features[r].mm_position) for r in order]
+
+
 def test_update_states_new_request(model_runner, dist_init):
     req_id = "req_0"
 

vllm/v1/worker/gpu/mm/encoder_runner.py

@@ -150,9 +150,12 @@ class EncoderRunner:
     def execute_mm_encoder(
         self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
     ) -> list[torch.Tensor]:
+        # Stable-sort items by modality so each modality is encoded in as few
+        # batches as possible. The returned outputs are restored to input order.
+        order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0])
         encoder_outputs: list[torch.Tensor] = []
         for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(
-            mm_kwargs, device=self.device, pin_memory=PIN_MEMORY
+            [mm_kwargs[i] for i in order], device=self.device, pin_memory=PIN_MEMORY
         ):
             cg_manager = self.cudagraph_manager
             cudagraph_output = (
@@ -169,7 +172,7 @@ class EncoderRunner:
             )
             sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)
             encoder_outputs.extend(batch_outputs)
-        return encoder_outputs
+        return [encoder_outputs[i] for i in np.argsort(order)]
 
     @contextmanager
     def timed_encoder_operation(self, request_ids: Collection[str]):

vllm/v1/worker/gpu/mm/lora.py

@@ -6,6 +6,7 @@ import numpy as np
 
 from vllm.lora.layers import LoRAMapping, LoRAMappingType
 from vllm.lora.worker_manager import WorkerLoRAManager
+from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.v1.worker.gpu.lora_utils import LoraState
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 
@@ -31,7 +32,7 @@ def set_active_mm_loras(
     encoder_token_counts: list[int] = []
     connector_token_counts: list[int] = []
 
-    # iterate through images
+    mm_items: list[tuple[int, MultiModalFeatureSpec]] = []
     for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
         req_idx = req_id_to_index.get(req_id)
         if req_idx is None:
@@ -39,27 +40,29 @@ def set_active_mm_loras(
 
         lora_id = int(lora_state.lora_ids[req_idx])
         mm_features = encoder_cache.mm_features[req_id]
-
-        # iterate through visual tokens
-        for mm_input_id in encoder_input_ids:
-            pos_info = mm_features[mm_input_id].mm_position
-
-            tower_tokens, connector_tokens = model.get_mm_lora_token_counts(
-                modality=mm_features[mm_input_id].modality,
-                mm_kwargs=mm_features[mm_input_id].data,
-                num_mm_embeds=pos_info.get_num_embeds(),
-            )
-
-            prompt_lora_mapping.append(lora_id)
-            token_lora_mapping.extend([lora_id] * tower_tokens)
-            encoder_token_counts.append(tower_tokens)
-            connector_token_counts.append(connector_tokens)
+        mm_items.extend((lora_id, mm_features[i]) for i in encoder_input_ids)
 
         if lora_id > 0:
             lora_request = lora_state.lora_requests.get(req_id)
             if lora_request is not None:
                 lora_requests.add(lora_request)
 
+    # Runs items sorted by modality.
+    mm_items.sort(key=lambda x: x[1].modality)
+    for lora_id, mm_feature in mm_items:
+        pos_info = mm_feature.mm_position
+
+        tower_tokens, connector_tokens = model.get_mm_lora_token_counts(
+            modality=mm_feature.modality,
+            mm_kwargs=mm_feature.data,
+            num_mm_embeds=pos_info.get_num_embeds(),
+        )
+
+        prompt_lora_mapping.append(lora_id)
+        token_lora_mapping.extend([lora_id] * tower_tokens)
+        encoder_token_counts.append(tower_tokens)
+        connector_token_counts.append(connector_tokens)
+
     if not prompt_lora_mapping:
         return
 

vllm/v1/worker/gpu_model_runner.py

@@ -2929,7 +2929,14 @@ class GPUModelRunner(
                 mm_kwargs.append((mm_feature.modality, mm_feature.data))
                 mm_lora_refs.append((req_id, mm_feature.mm_position))
 
-        return mm_hashes, mm_kwargs, mm_lora_refs
+        # Stable-sort by modality so each modality is encoded in as few batches
+        # as possible. Encoder outputs are cached by mm_hash, so order is free.
+        order = sorted(range(len(mm_kwargs)), key=lambda i: mm_kwargs[i][0])
+        return (
+            [mm_hashes[i] for i in order],
+            [mm_kwargs[i] for i in order],
+            [mm_lora_refs[i] for i in order],
+        )
 
     def _cache_encoder_output(
         self,
@@ -2990,13 +2997,6 @@ class GPUModelRunner(
             and scheduler_output.scheduled_encoder_inputs
         )
 
-        # Batch mm inputs as much as we can: if a request in the batch has
-        # multiple modalities or a different modality than the previous one,
-        # we process it separately to preserve item order.
-        # FIXME(ywang96): This is a hacky way to deal with multiple modalities
-        # in the same batch while still being able to benefit from batching
-        # multimodal inputs. The proper solution should be reordering the
-        # encoder outputs.
         model = cast(SupportsMultiModal, self.model)
 
         if self.lora_config and self.lora_manager.supports_tower_connector_lora():