[Bugfix][MRV2] Support encoder timing stats in model runner V2 (#50020)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Guan-Ming Chiu committed
Aug 12, 2026 at 00:00 UTC
457a5f34db04b082239c11ee34601e5ee5754a7c
7 files changed
+99
-22
tests/v1/worker/test_encoder_runner.py
+16
@@ -224,3 +224,19 @@ def test_execute_mm_encoder_is_a_noop_without_scheduled_items():
224
225
assert not cache.encoder_outputs
226
state.encoder_runner.execute_mm_encoder.assert_not_called()
227
+
228
+
229
+def test_encoder_timing_stats_registry():
230
+ runner = _make_runner([], [])
231
+ runner.enable_timing = True
232
+
233
+ with runner.timed_encoder_operation({"r1"}):
234
+ pass
235
+ with runner.timed_encoder_operation({"r1"}):
236
+ pass
237
+
238
+ stats = runner.get_encoder_timing_stats()
239
+ assert set(stats) == {"r1"}
240
+ assert stats["r1"]["num_encoder_calls"] == 2
241
+ assert stats["r1"]["encoder_forward_secs"] >= 0
242
+ assert runner.get_encoder_timing_stats() == {}
vllm/v1/worker/gpu/mm/encoder_runner.py
+43
-1
@@ -1,5 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import threading
4
+import time
5
+from collections.abc import Collection
6
+from contextlib import contextmanager
7
+
8
import numpy as np
9
import torch
10
@@ -13,7 +18,10 @@ from vllm.multimodal.utils import (
18
set_mm_embedding_modality,
19
)
20
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
16
-from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
21
+from vllm.v1.worker.utils import (
22
+ EncoderTimingStats,
23
+ sanity_check_mm_encoder_outputs,
24
+)
25
26
logger = init_logger(__name__)
27
@@ -27,6 +35,7 @@ class EncoderRunner:
35
encoder_cache: EncoderCache,
36
dtype: torch.dtype,
37
device: torch.device,
38
+ enable_timing: bool = False,
39
):
40
self.model = model
41
self.max_num_tokens = max_num_tokens
@@ -35,6 +44,9 @@ class EncoderRunner:
44
self.dtype = dtype
45
self.device = device
46
self.is_realtime = supports_realtime(model)
47
+ self.enable_timing = enable_timing
48
+ self.encoder_timing_registry: dict[str, EncoderTimingStats] = {}
49
+ self._timing_lock = threading.Lock()
50
51
self.inputs_embeds = torch.zeros(
52
max_num_tokens, hidden_size, dtype=dtype, device=device
@@ -110,6 +122,36 @@ class EncoderRunner:
122
encoder_outputs.extend(batch_outputs)
123
return encoder_outputs
124
125
+ @contextmanager
126
+ def timed_encoder_operation(self, request_ids: Collection[str]):
127
+ if not (self.enable_timing and request_ids):
128
+ yield
129
+ return
130
+
131
+ torch.accelerator.synchronize()
132
+ start_time = time.perf_counter()
133
+ try:
134
+ yield
135
+ finally:
136
+ torch.accelerator.synchronize()
137
+ per_request_time = (time.perf_counter() - start_time) / len(request_ids)
138
+ with self._timing_lock:
139
+ for req_id in request_ids:
140
+ stats = self.encoder_timing_registry.setdefault(
141
+ req_id, EncoderTimingStats()
142
+ )
143
+ stats.encoder_forward_secs += per_request_time
144
+ stats.num_encoder_calls += 1
145
+
146
+ def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
147
+ with self._timing_lock:
148
+ stats = {
149
+ req_id: stats_obj.to_dict()
150
+ for req_id, stats_obj in self.encoder_timing_registry.items()
151
+ }
152
+ self.encoder_timing_registry.clear()
153
+ return stats
154
+
155
def gather_mm_embeddings(
156
self,
157
req_ids: list[str],
vllm/v1/worker/gpu/model_runner.py
+6
@@ -453,6 +453,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
453
# Cache the default CUDA stream to avoid lookup overhead.
454
return torch.cuda.current_stream(self.device)
455
456
+ def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
457
+ encoder_runner = getattr(self.model_state, "encoder_runner", None)
458
+ if encoder_runner is None:
459
+ return {}
460
+ return encoder_runner.get_encoder_timing_stats()
461
+
462
def get_kv_cache_spec(self):
463
if self.is_encoder_only:
464
return {}
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+2
-1
@@ -86,7 +86,8 @@ class EncoderDecoderModelState(ModelState):
86
# so execute_mm_encoder preserves request order; use its return value
87
# directly. No need to store in encoder_cache: cross-attention K/V are
88
# written to the KV cache on the first step; decode steps use the cache.
89
- self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
89
+ with self.encoder_runner.timed_encoder_operation(encoder_inputs.keys()):
90
+ self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
91
else:
92
# Decode steps: encoder K/V are in cross-attention KV cache.
93
self.encoder_outputs = []
vllm/v1/worker/gpu/model_states/interface.py
+9
-1
@@ -60,6 +60,7 @@ class ModelState(ABC):
60
self.supports_mm_inputs = encoder_cache is not None
61
if encoder_cache is not None:
62
self.encoder_cache = encoder_cache
63
+ observability_config = vllm_config.observability_config
64
self.encoder_runner = EncoderRunner(
65
model=self.model,
66
max_num_tokens=self.max_num_tokens,
@@ -67,6 +68,10 @@ class ModelState(ABC):
68
encoder_cache=encoder_cache,
69
dtype=self.dtype,
70
device=self.device,
71
+ enable_timing=bool(
72
+ observability_config
73
+ and observability_config.enable_mm_processor_stats
74
+ ),
75
)
76
77
def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
@@ -150,7 +155,10 @@ class ModelState(ABC):
155
scheduled_encoder_inputs
156
)
157
if mm_kwargs:
153
- encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
158
+ with self.encoder_runner.timed_encoder_operation(
159
+ scheduled_encoder_inputs.keys()
160
+ ):
161
+ encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
162
self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))
163
164
def gather_mm_embeddings(
vllm/v1/worker/gpu_model_runner.py
+6
-19
@@ -10,7 +10,7 @@ from collections import defaultdict
10
from collections.abc import Callable, Iterable, Iterator, Sequence
11
from contextlib import AbstractContextManager, contextmanager, nullcontext
12
from copy import copy, deepcopy
13
-from dataclasses import dataclass, replace
13
+from dataclasses import replace
14
from functools import reduce
15
from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
16
@@ -230,7 +230,11 @@ from vllm.v1.worker.ubatch_utils import (
230
maybe_create_ubatch_slices,
231
split_attn_metadata,
232
)
233
-from vllm.v1.worker.utils import is_residual_scattered_for_sp, raise_if_nan_logits
233
+from vllm.v1.worker.utils import (
234
+ EncoderTimingStats,
235
+ is_residual_scattered_for_sp,
236
+ raise_if_nan_logits,
237
+)
238
from vllm.v1.worker.workspace import lock_workspace
239
240
from .utils import (
@@ -7948,20 +7952,3 @@ class GPUModelRunner(
7952
stats = self.encoder_timing_registry[req_id]
7953
stats.encoder_forward_secs += per_request_time
7954
stats.num_encoder_calls += 1
7951
-
7952
-
7953
-@dataclass
7954
-class EncoderTimingStats:
7955
- """Per-request timing statistics for encoder forward pass."""
7956
-
7957
- encoder_forward_secs: float = 0.0
7958
- """Time spent in vision encoder forward pass (seconds)."""
7959
-
7960
- num_encoder_calls: int = 0
7961
- """Number of times encoder was called for this request."""
7962
-
7963
- def to_dict(self) -> dict[str, float | int]:
7964
- return {
7965
- "encoder_forward_secs": self.encoder_forward_secs,
7966
- "num_encoder_calls": self.num_encoder_calls,
7967
- }
vllm/v1/worker/utils.py
+17
@@ -639,3 +639,20 @@ def is_residual_scattered_for_sp(
639
assert num_input_tokens % tp == 0
640
641
return True
642
+
643
+
644
+@dataclass
645
+class EncoderTimingStats:
646
+ """Per-request timing statistics for encoder forward pass."""
647
+
648
+ encoder_forward_secs: float = 0.0
649
+ """Time spent in vision encoder forward pass (seconds)."""
650
+
651
+ num_encoder_calls: int = 0
652
+ """Number of times encoder was called for this request."""
653
+
654
+ def to_dict(self) -> dict[str, float | int]:
655
+ return {
656
+ "encoder_forward_secs": self.encoder_forward_secs,
657
+ "num_encoder_calls": self.num_encoder_calls,
658
+ }