[MRV2] Add encoder cache profiling implementation (#47985)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

vllm/multimodal/encoder_budget.py

@@ -4,6 +4,7 @@ from collections.abc import Mapping
 
 from vllm.config import ModelConfig, VllmConfig
 from vllm.logger import init_logger
+from vllm.multimodal.inputs import MultiModalKwargsItem
 from vllm.multimodal.processing import BaseMultiModalProcessor
 from vllm.multimodal.registry import MultiModalRegistry
 from vllm.utils.torch_utils import set_default_torch_num_threads
@@ -48,6 +49,8 @@ class MultiModalBudget:
         self,
         vllm_config: VllmConfig,
         mm_registry: MultiModalRegistry,
+        *,
+        enable_cache: bool = True,
     ) -> None:
         super().__init__()
 
@@ -58,7 +61,11 @@ class MultiModalBudget:
         self.max_num_reqs = scheduler_config.max_num_seqs
 
         with set_default_torch_num_threads():  # Avoid hang during startup
-            cache = mm_registry.processor_only_cache_from_config(vllm_config)
+            cache = (
+                mm_registry.processor_only_cache_from_config(vllm_config)
+                if enable_cache
+                else None
+            )
             processor = mm_registry.create_processor(model_config, cache=cache)
 
             self.cache = cache
@@ -191,3 +198,23 @@ class MultiModalBudget:
     def reset_cache(self) -> None:
         if self.cache is not None:
             self.cache.clear_cache()
+
+
+def get_dummy_encoder_profile_inputs(
+    mm_registry: MultiModalRegistry,
+    budget: MultiModalBudget,
+) -> list[tuple[str, MultiModalKwargsItem]]:
+    if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item:
+        return []
+
+    modality = budget.get_modality_with_max_tokens()
+    max_items_per_batch = budget.mm_max_items_per_batch[modality]
+    dummy_mm_inputs = mm_registry.get_dummy_mm_inputs(
+        budget.model_config,
+        mm_counts={modality: 1},
+        processor=budget.processor,
+    )
+    dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
+    assert dummy_mm_item is not None, "Dummy item should be generated"
+
+    return [(modality, dummy_mm_item)] * max_items_per_batch

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

@@ -28,7 +28,7 @@ class EncoderCache:
         Clear the multi-modal cache that was used during profiling,
         but no longer needed during inference.
         """
-        # TODO: Implement MM budget for encoder dummy run
+        # NOTE: v2 encoder cache profiling skips the multi-modal cache
         pass
 
     def reset_encoder_cache(self) -> None:

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

@@ -3,7 +3,9 @@
 import numpy as np
 import torch
 
+from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
+from vllm.multimodal.encoder_budget import MultiModalBudget
 from vllm.multimodal.inputs import MultiModalKwargsItem
 from vllm.multimodal.utils import (
     get_mm_features_in_window,
@@ -13,6 +15,8 @@ from vllm.multimodal.utils import (
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
+logger = init_logger(__name__)
+
 
 class EncoderRunner:
     def __init__(
@@ -52,6 +56,45 @@ class EncoderRunner:
 
         return mm_hashes, mm_kwargs
 
+    @torch.inference_mode()
+    def profile_encoder_cache(
+        self,
+        dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]],
+        budget: MultiModalBudget,
+    ) -> None:
+        """Profile multimodal encoder and temporary encoder cache memory."""
+        if (encoder_budget := budget.get_encoder_budget()) <= 0:
+            return
+
+        if not budget.mm_max_toks_per_item:
+            logger.info(
+                "Skipping encoder profiling for embedding-only mode "
+                "(all modality limits=0 with enable_mm_embeds=True).",
+            )
+            return
+
+        assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling"
+        dummy_modality = dummy_mm_inputs[0][0]
+        max_mm_items_per_batch = len(dummy_mm_inputs)
+
+        logger.info_once(
+            "Encoder cache will be initialized with a budget of %s tokens, "
+            "and profiled with %s %s items of the maximum feature size.",
+            encoder_budget,
+            max_mm_items_per_batch,
+            dummy_modality,
+        )
+
+        dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs)
+
+        sanity_check_mm_encoder_outputs(
+            dummy_encoder_outputs,
+            expected_num_items=max_mm_items_per_batch,
+        )
+        self.encoder_cache.encoder_outputs.update(
+            (f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs)
+        )
+
     @torch.inference_mode()
     def execute_mm_encoder(
         self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]

vllm/v1/worker/gpu/model_runner.py

@@ -43,6 +43,10 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
 )
 from vllm.model_executor.model_loader import get_model_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.encoder_budget import (
+    MultiModalBudget,
+    get_dummy_encoder_profile_inputs,
+)
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
 from vllm.utils.math_utils import cdiv
@@ -671,6 +675,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
     @torch.inference_mode()
     def profile_run(self) -> None:
+        if self.supports_mm_inputs and self.is_first_pp_rank:
+            mm_config = self.model_config.multimodal_config
+            if mm_config is not None and not mm_config.skip_mm_profiling:
+                mm_budget = MultiModalBudget(
+                    self.vllm_config,
+                    self.mm_registry,
+                    enable_cache=False,
+                )
+                dummy_mm_inputs = get_dummy_encoder_profile_inputs(
+                    self.mm_registry,
+                    mm_budget,
+                )
+                self.model_state.encoder_runner.profile_encoder_cache(
+                    dummy_mm_inputs, mm_budget
+                )
+
         hidden_states, sample_hidden_states = self._dummy_run(
             self.max_num_tokens, skip_attn=True, is_profile=True
         )
@@ -685,6 +705,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         torch.accelerator.synchronize()
         del hidden_states, sample_hidden_states
+        self.reset_encoder_cache()
         gc.collect()
 
     def post_kv_cache_wake_up(self) -> None: