[MRV2] Add encoder cache profiling implementation (#47985)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Isotr0py committed
Jul 21, 2026 at 11:18 UTC
ea0e9c8f2e4b037e6fa4d914c82bf178c9f65f55
4 files changed
+93
-2
vllm/multimodal/encoder_budget.py
+28
-1
@@ -4,6 +4,7 @@ from collections.abc import Mapping
4
5
from vllm.config import ModelConfig, VllmConfig
6
from vllm.logger import init_logger
7
+from vllm.multimodal.inputs import MultiModalKwargsItem
8
from vllm.multimodal.processing import BaseMultiModalProcessor
9
from vllm.multimodal.registry import MultiModalRegistry
10
from vllm.utils.torch_utils import set_default_torch_num_threads
@@ -48,6 +49,8 @@ class MultiModalBudget:
49
self,
50
vllm_config: VllmConfig,
51
mm_registry: MultiModalRegistry,
52
+ *,
53
+ enable_cache: bool = True,
54
) -> None:
55
super().__init__()
56
@@ -58,7 +61,11 @@ class MultiModalBudget:
61
self.max_num_reqs = scheduler_config.max_num_seqs
62
63
with set_default_torch_num_threads(): # Avoid hang during startup
61
- cache = mm_registry.processor_only_cache_from_config(vllm_config)
64
+ cache = (
65
+ mm_registry.processor_only_cache_from_config(vllm_config)
66
+ if enable_cache
67
+ else None
68
+ )
69
processor = mm_registry.create_processor(model_config, cache=cache)
70
71
self.cache = cache
@@ -191,3 +198,23 @@ class MultiModalBudget:
198
def reset_cache(self) -> None:
199
if self.cache is not None:
200
self.cache.clear_cache()
201
+
202
+
203
+def get_dummy_encoder_profile_inputs(
204
+ mm_registry: MultiModalRegistry,
205
+ budget: MultiModalBudget,
206
+) -> list[tuple[str, MultiModalKwargsItem]]:
207
+ if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item:
208
+ return []
209
+
210
+ modality = budget.get_modality_with_max_tokens()
211
+ max_items_per_batch = budget.mm_max_items_per_batch[modality]
212
+ dummy_mm_inputs = mm_registry.get_dummy_mm_inputs(
213
+ budget.model_config,
214
+ mm_counts={modality: 1},
215
+ processor=budget.processor,
216
+ )
217
+ dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
218
+ assert dummy_mm_item is not None, "Dummy item should be generated"
219
+
220
+ return [(modality, dummy_mm_item)] * max_items_per_batch
vllm/v1/worker/gpu/mm/encoder_cache.py
+1
-1
@@ -28,7 +28,7 @@ class EncoderCache:
28
Clear the multi-modal cache that was used during profiling,
29
but no longer needed during inference.
30
"""
31
- # TODO: Implement MM budget for encoder dummy run
31
+ # NOTE: v2 encoder cache profiling skips the multi-modal cache
32
pass
33
34
def reset_encoder_cache(self) -> None:
vllm/v1/worker/gpu/mm/encoder_runner.py
+43
@@ -3,7 +3,9 @@
3
import numpy as np
4
import torch
5
6
+from vllm.logger import init_logger
7
from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
8
+from vllm.multimodal.encoder_budget import MultiModalBudget
9
from vllm.multimodal.inputs import MultiModalKwargsItem
10
from vllm.multimodal.utils import (
11
get_mm_features_in_window,
@@ -13,6 +15,8 @@ from vllm.multimodal.utils import (
15
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
16
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
17
18
+logger = init_logger(__name__)
19
+
20
21
class EncoderRunner:
22
def __init__(
@@ -52,6 +56,45 @@ class EncoderRunner:
56
57
return mm_hashes, mm_kwargs
58
59
+ @torch.inference_mode()
60
+ def profile_encoder_cache(
61
+ self,
62
+ dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]],
63
+ budget: MultiModalBudget,
64
+ ) -> None:
65
+ """Profile multimodal encoder and temporary encoder cache memory."""
66
+ if (encoder_budget := budget.get_encoder_budget()) <= 0:
67
+ return
68
+
69
+ if not budget.mm_max_toks_per_item:
70
+ logger.info(
71
+ "Skipping encoder profiling for embedding-only mode "
72
+ "(all modality limits=0 with enable_mm_embeds=True).",
73
+ )
74
+ return
75
+
76
+ assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling"
77
+ dummy_modality = dummy_mm_inputs[0][0]
78
+ max_mm_items_per_batch = len(dummy_mm_inputs)
79
+
80
+ logger.info_once(
81
+ "Encoder cache will be initialized with a budget of %s tokens, "
82
+ "and profiled with %s %s items of the maximum feature size.",
83
+ encoder_budget,
84
+ max_mm_items_per_batch,
85
+ dummy_modality,
86
+ )
87
+
88
+ dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs)
89
+
90
+ sanity_check_mm_encoder_outputs(
91
+ dummy_encoder_outputs,
92
+ expected_num_items=max_mm_items_per_batch,
93
+ )
94
+ self.encoder_cache.encoder_outputs.update(
95
+ (f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs)
96
+ )
97
+
98
@torch.inference_mode()
99
def execute_mm_encoder(
100
self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
vllm/v1/worker/gpu/model_runner.py
+21
@@ -43,6 +43,10 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
43
)
44
from vllm.model_executor.model_loader import get_model_loader
45
from vllm.multimodal import MULTIMODAL_REGISTRY
46
+from vllm.multimodal.encoder_budget import (
47
+ MultiModalBudget,
48
+ get_dummy_encoder_profile_inputs,
49
+)
50
from vllm.sequence import IntermediateTensors
51
from vllm.tasks import SupportedTask
52
from vllm.utils.math_utils import cdiv
@@ -671,6 +675,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
675
676
@torch.inference_mode()
677
def profile_run(self) -> None:
678
+ if self.supports_mm_inputs and self.is_first_pp_rank:
679
+ mm_config = self.model_config.multimodal_config
680
+ if mm_config is not None and not mm_config.skip_mm_profiling:
681
+ mm_budget = MultiModalBudget(
682
+ self.vllm_config,
683
+ self.mm_registry,
684
+ enable_cache=False,
685
+ )
686
+ dummy_mm_inputs = get_dummy_encoder_profile_inputs(
687
+ self.mm_registry,
688
+ mm_budget,
689
+ )
690
+ self.model_state.encoder_runner.profile_encoder_cache(
691
+ dummy_mm_inputs, mm_budget
692
+ )
693
+
694
hidden_states, sample_hidden_states = self._dummy_run(
695
self.max_num_tokens, skip_attn=True, is_profile=True
696
)
@@ -685,6 +705,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
705
706
torch.accelerator.synchronize()
707
del hidden_states, sample_hidden_states
708
+ self.reset_encoder_cache()
709
gc.collect()
710
711
def post_kv_cache_wake_up(self) -> None: