From 702ce313f4f85f5f2e76b7796aa5cdc80f3595a1 Mon Sep 17 00:00:00 2001 From: zzaebok <44357534+zzaebok@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:55:36 +0800 Subject: [PATCH] [Perf][Multimodal] Sort multimodal encoder inputs by modality (#60196) Signed-off-by: Jaebok Lee Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_encoder_runner.py | 17 ++++++- tests/v1/worker/test_gpu_model_runner.py | 64 ++++++++++++++++++++++++ vllm/v1/worker/gpu/mm/encoder_runner.py | 7 ++- vllm/v1/worker/gpu/mm/lora.py | 35 +++++++------ vllm/v1/worker/gpu_model_runner.py | 16 +++--- 5 files changed, 112 insertions(+), 27 deletions(-) diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 98aaa568f0..69afc6ea93 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/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() diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index deb4f7b421..3851d5dad6 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/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" diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index 16e42d140a..4d7ab92013 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/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]): diff --git a/vllm/v1/worker/gpu/mm/lora.py b/vllm/v1/worker/gpu/mm/lora.py index 819f3b1cb3..e73e028f30 100644 --- a/vllm/v1/worker/gpu/mm/lora.py +++ b/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 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 21b0f67e2f..c02bdc74e2 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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():