[MRV2][Metrics] Support `CUDAGraphStat` in MRV2 (#52358)
Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Yizhou committed
Sep 4, 2026 at 15:19 UTC
8f816a3f665489d7f0d222115d4f72ebab01076b
3 files changed
+22
-1
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1
@@ -185,6 +185,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
185
finished_req_ids=set(),
186
ec_connector_output=None,
187
routed_experts=None,
188
+ cudagraph_stats=None,
189
)
190
runner.req_states = SimpleNamespace()
191
vllm/v1/worker/gpu/cudagraph_utils.py
+12
-1
@@ -17,7 +17,7 @@ from vllm.compilation.breakable_cudagraph import (
17
is_breakable_cudagraph_enabled,
18
)
19
from vllm.compilation.counter import compilation_counter
20
-from vllm.compilation.cuda_graph import CUDAGraphWrapper
20
+from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper
21
from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
22
from vllm.config import VllmConfig, set_current_vllm_config
23
from vllm.config.compilation import CUDAGraphMode
@@ -70,6 +70,17 @@ class BatchExecutionDescriptor:
70
num_active_loras: int = 0
71
72
73
+def make_cudagraph_stats(
74
+ batch_desc: BatchExecutionDescriptor, num_tokens: int
75
+) -> CUDAGraphStat:
76
+ return CUDAGraphStat(
77
+ num_unpadded_tokens=num_tokens,
78
+ num_padded_tokens=batch_desc.num_tokens,
79
+ num_paddings=batch_desc.num_tokens - num_tokens,
80
+ runtime_mode=str(batch_desc.cg_mode),
81
+ )
82
+
83
+
84
class CreateForwardFn(Protocol):
85
"""Factory that prepares inputs (OUTSIDE the graph) and returns a
86
forward_fn. Called with warmup=True for the warmup pass and warmup=False
vllm/v1/worker/gpu/model_runner.py
+9
@@ -30,6 +30,7 @@ import torch.nn as nn
30
31
import vllm.envs as envs
32
from vllm.compilation.counter import compilation_counter
33
+from vllm.compilation.cuda_graph import CUDAGraphStat
34
from vllm.config import VllmConfig
35
from vllm.config.compilation import CUDAGraphMode
36
from vllm.distributed.parallel_state import (
@@ -100,6 +101,7 @@ from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
101
from vllm.v1.worker.gpu.cudagraph_utils import (
102
BatchExecutionDescriptor,
103
ModelCudaGraphManager,
104
+ make_cudagraph_stats,
105
)
106
from vllm.v1.worker.gpu.cudagraph_utils import (
107
profile_cudagraph_memory as _profile_cudagraph_memory,
@@ -1602,9 +1604,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1604
empty_output = self.kv_connector.no_forward(scheduler_output)
1605
return self._merge_ec_connector_no_forward(scheduler_output, empty_output)
1606
1607
+ cudagraph_stats = None
1608
if not dummy_run:
1609
# Common case.
1610
# Prepare all the inputs and copy to the input buffers.
1611
+ if self.observability_config.cudagraph_metrics:
1612
+ cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks)
1613
assert batch_req_state is not None
1614
input_batch = self.prepare_inputs(
1615
scheduler_output, batch_req_state, batch_desc
@@ -1826,6 +1831,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1831
finished_req_ids=finished_req_ids,
1832
ec_connector_output=ec_connector_output,
1833
routed_experts=routed_experts,
1834
+ cudagraph_stats=cudagraph_stats,
1835
)
1836
1837
if not self.is_last_pp_rank:
@@ -1851,6 +1857,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1857
finished_req_ids = self.execute_model_state.finished_req_ids
1858
ec_connector_output = self.execute_model_state.ec_connector_output
1859
routed_experts = self.execute_model_state.routed_experts
1860
+ cudagraph_stats = self.execute_model_state.cudagraph_stats
1861
self.execute_model_state = None
1862
1863
if not self.is_last_pp_rank:
@@ -1909,6 +1916,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1916
req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
1917
sampled_token_ids=None, # type: ignore
1918
prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]
1919
+ cudagraph_stats=cudagraph_stats,
1920
)
1921
# Start async output copy here so that it can overlap with speculator proposal.
1922
async_output = AsyncOutput(
@@ -2124,6 +2132,7 @@ class ExecuteModelState(NamedTuple):
2132
finished_req_ids: set[str]
2133
ec_connector_output: ECConnectorOutput | None
2134
routed_experts: RoutedExpertsTensors | None
2135
+ cudagraph_stats: CUDAGraphStat | None
2136
2137
2138
class BatchReqState(NamedTuple):