[MRV2][Multimodal] Enable encoder cuda graph for model runner v2 (#49852)
Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Isotr0py committed
Aug 14, 2026 at 17:33 UTC
66728feb1fbe0c6d32dcce2d4ce6e827712118c8
4 files changed
+101
-26
vllm/v1/worker/encoder_cudagraph.py
+4
@@ -203,6 +203,10 @@ class EncoderCudaGraphManager:
203
"""Check if a modality is supported by this manager."""
204
return modality in self.config.modalities
205
206
+ def is_captured(self) -> bool:
207
+ """Return whether a CUDA graph pool is active."""
208
+ return self.graph_pool is not None
209
+
210
def clear(self) -> None:
211
"""Release captured encoder CUDA graphs and the manager-local pool."""
212
for graph_set in self.budget_graphs.values():
vllm/v1/worker/gpu/mm/encoder_runner.py
+38
-1
@@ -4,6 +4,7 @@ import threading
4
import time
5
from collections.abc import Collection
6
from contextlib import contextmanager
7
+from typing import TYPE_CHECKING
8
9
import numpy as np
10
import torch
@@ -24,6 +25,9 @@ from vllm.v1.worker.utils import (
25
sanity_check_mm_encoder_outputs,
26
)
27
28
+if TYPE_CHECKING:
29
+ from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
30
+
31
logger = init_logger(__name__)
32
33
@@ -36,6 +40,7 @@ class EncoderRunner:
40
encoder_cache: EncoderCache,
41
dtype: torch.dtype,
42
device: torch.device,
43
+ cudagraph_manager: "EncoderCudaGraphManager | None" = None,
44
enable_timing: bool = False,
45
):
46
self.model = model
@@ -45,6 +50,7 @@ class EncoderRunner:
50
self.dtype = dtype
51
self.device = device
52
self.is_realtime = supports_realtime(model)
53
+ self.cudagraph_manager = cudagraph_manager
54
self.enable_timing = enable_timing
55
self.encoder_timing_registry: dict[str, EncoderTimingStats] = {}
56
self._timing_lock = threading.Lock()
@@ -53,6 +59,25 @@ class EncoderRunner:
59
max_num_tokens, hidden_size, dtype=dtype, device=device
60
)
61
62
+ def has_cudagraph(self) -> bool:
63
+ return self.cudagraph_manager is not None
64
+
65
+ @torch.inference_mode()
66
+ def capture(self) -> None:
67
+ manager = self.cudagraph_manager
68
+ assert manager is not None
69
+
70
+ from vllm.distributed.parallel_state import graph_capture
71
+ from vllm.platforms import current_platform
72
+
73
+ with graph_capture(device=self.device):
74
+ manager.capture(graph_pool=current_platform.graph_pool_handle())
75
+ torch.accelerator.synchronize()
76
+
77
+ def clear(self) -> None:
78
+ if self.cudagraph_manager is not None:
79
+ self.cudagraph_manager.clear()
80
+
81
def prepare_mm_inputs(
82
self, scheduled_encoder_inputs: dict[str, list[int]]
83
) -> tuple[list[str], list[tuple[str, MultiModalKwargsItem]]]:
@@ -118,7 +143,19 @@ class EncoderRunner:
143
for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(
144
mm_kwargs, device=self.device, pin_memory=PIN_MEMORY
145
):
121
- batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch)
146
+ cg_manager = self.cudagraph_manager
147
+ cudagraph_output = (
148
+ cg_manager.execute(mm_kwargs_batch)
149
+ if cg_manager is not None
150
+ and cg_manager.is_captured()
151
+ and cg_manager.supports_modality(modality)
152
+ else None
153
+ )
154
+ batch_outputs = (
155
+ cudagraph_output
156
+ if cudagraph_output is not None
157
+ else self.model.embed_multimodal(**mm_kwargs_batch)
158
+ )
159
sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)
160
encoder_outputs.extend(batch_outputs)
161
return encoder_outputs
vllm/v1/worker/gpu/model_runner.py
+36
-24
@@ -842,10 +842,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
842
return 0
843
844
assert self.cudagraph_manager is not None
845
- if not self.cudagraph_manager.needs_capture():
845
+ capture_encoder = (
846
+ self.model_state.supports_mm_inputs
847
+ and self.model_state.encoder_runner.has_cudagraph()
848
+ )
849
+ capture_decoder = self.cudagraph_manager.needs_capture()
850
+ if not capture_encoder and not capture_decoder:
851
logger.warning(
847
- "Skipping CUDA graph capture. To turn on CUDA graph capture, "
848
- "ensure `cudagraph_mode` was not manually set to `NONE`"
852
+ "Skipping encoder and decoder CUDA graph capture. To enable "
853
+ "encoder capture, ensure `cudagraph_mm_encoder` is enabled; "
854
+ "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`."
855
)
856
return 0
857
@@ -857,27 +863,31 @@ class GPUModelRunner(LoRAModelRunnerMixin):
863
start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
864
865
with self.maybe_setup_dummy_loras(self.lora_config):
860
- self.cudagraph_manager.capture(
861
- self.model,
862
- self.model_state,
863
- self.input_buffers,
864
- self.intermediate_tensors,
865
- self.block_tables,
866
- self.attn_groups,
867
- self.kv_cache_config,
868
- has_lora=self.lora_config is not None,
869
- use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
870
- lora_capture_hook=create_lora_capture_hook(self.lora_config, self),
871
- )
872
- if self.speculator is not None:
873
- self.speculator.capture()
874
- if self.adaptive_verification is not None:
875
- with self.step_timing.collect() as timings:
876
- for batch in self.adaptive_verification.batches_to_profile(
877
- self.cudagraph_manager.captured_token_counts()
878
- ):
879
- self._dummy_run(**batch)
880
- self.adaptive_verification.set_initial_cost_curves(timings)
866
+ if capture_encoder:
867
+ self.model_state.encoder_runner.capture()
868
+
869
+ if capture_decoder:
870
+ self.cudagraph_manager.capture(
871
+ self.model,
872
+ self.model_state,
873
+ self.input_buffers,
874
+ self.intermediate_tensors,
875
+ self.block_tables,
876
+ self.attn_groups,
877
+ self.kv_cache_config,
878
+ has_lora=self.lora_config is not None,
879
+ use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
880
+ lora_capture_hook=create_lora_capture_hook(self.lora_config, self),
881
+ )
882
+ if self.speculator is not None:
883
+ self.speculator.capture()
884
+ if self.adaptive_verification is not None:
885
+ with self.step_timing.collect() as timings:
886
+ for batch in self.adaptive_verification.batches_to_profile(
887
+ self.cudagraph_manager.captured_token_counts()
888
+ ):
889
+ self._dummy_run(**batch)
890
+ self.adaptive_verification.set_initial_cost_curves(timings)
891
892
end_time = time.perf_counter()
893
end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
@@ -1885,6 +1895,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1895
self.attn_groups.clear()
1896
if hasattr(self, "kv_cache_config"):
1897
del self.kv_cache_config
1898
+ if hasattr(self, "model_state") and self.model_state.supports_mm_inputs:
1899
+ self.model_state.encoder_runner.clear()
1900
free_before_shutdown(self.vllm_config)
1901
if hasattr(self, "model_state"):
1902
del self.model_state
vllm/v1/worker/gpu/model_states/interface.py
+23
-1
@@ -1,17 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from abc import ABC, abstractmethod
4
-from typing import Any
4
+from typing import Any, cast
5
6
import torch
7
import torch.nn as nn
8
9
from vllm.config import VllmConfig
10
from vllm.config.compilation import CUDAGraphMode
11
+from vllm.model_executor.models.interfaces import (
12
+ SupportsEncoderCudaGraph,
13
+ supports_encoder_cudagraph,
14
+)
15
from vllm.tasks import GenerationTask
16
from vllm.v1.attention.backend import AttentionCGSupport
17
from vllm.v1.core.sched.output import NewRequestData
18
from vllm.v1.kv_cache_interface import KVCacheConfig
19
+from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
20
from vllm.v1.worker.gpu.input_batch import InputBatch
21
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
22
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
@@ -59,6 +64,22 @@ class ModelState(ABC):
64
65
self.supports_mm_inputs = encoder_cache is not None
66
if encoder_cache is not None:
67
+ enable_encoder_cuda_graph = (
68
+ not self.model_config.enforce_eager
69
+ and vllm_config.compilation_config.cudagraph_mm_encoder
70
+ and supports_encoder_cudagraph(model)
71
+ )
72
+ cudagraph_manager = (
73
+ EncoderCudaGraphManager(
74
+ vllm_config=vllm_config,
75
+ device=device,
76
+ dtype=self.dtype,
77
+ model=cast(SupportsEncoderCudaGraph, model),
78
+ )
79
+ if enable_encoder_cuda_graph
80
+ else None
81
+ )
82
+
83
self.encoder_cache = encoder_cache
84
observability_config = vllm_config.observability_config
85
self.encoder_runner = EncoderRunner(
@@ -68,6 +89,7 @@ class ModelState(ABC):
89
encoder_cache=encoder_cache,
90
dtype=self.dtype,
91
device=self.device,
92
+ cudagraph_manager=cudagraph_manager,
93
enable_timing=bool(
94
observability_config
95
and observability_config.enable_mm_processor_stats