From ea0e9c8f2e4b037e6fa4d914c82bf178c9f65f55 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Tue, 21 Jul 2026 11:18:14 +0800 Subject: [PATCH] [MRV2] Add encoder cache profiling implementation (#47985) Signed-off-by: Isotr0py --- vllm/multimodal/encoder_budget.py | 29 ++++++++++++++++- vllm/v1/worker/gpu/mm/encoder_cache.py | 2 +- vllm/v1/worker/gpu/mm/encoder_runner.py | 43 +++++++++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 21 ++++++++++++ 4 files changed, 93 insertions(+), 2 deletions(-) diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py index c1ff600869..3042091611 100644 --- a/vllm/multimodal/encoder_budget.py +++ b/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 diff --git a/vllm/v1/worker/gpu/mm/encoder_cache.py b/vllm/v1/worker/gpu/mm/encoder_cache.py index 065df2975c..21d9218182 100644 --- a/vllm/v1/worker/gpu/mm/encoder_cache.py +++ b/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: diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index c337ddd4aa..f0d9b7431f 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/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]] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 293b06ff97..6392d150d4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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: