@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 69520bc695

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 03:30 UTC

Add logging for cudagraph related info (#29825)

Signed-off-by: Yong Hoon Shin <yhshin@meta.com>

Yong Hoon Shin committed Dec 2, 2025 at 23:01 UTC 69520bc695ff8fa7fda66ef7c1a16761824ad354

9 files changed +161 -6

vllm/compilation/cuda_graph.py

+94

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import dataclasses

5 +from collections import Counter

6 from collections.abc import Callable

7 from contextlib import ExitStack

8 from typing import Any

@@ -22,6 +23,99 @@ from vllm.utils.torch_utils import weak_ref_tensors

23 logger = init_logger(__name__)

26 +@dataclasses.dataclass(frozen=True)

27 +class CUDAGraphStat:

28 + num_unpadded_tokens: int

29 + num_padded_tokens: int

30 + num_paddings: int

31 + runtime_mode: str

32 +

33 +

34 +class CUDAGraphLogging:

35 + """Aggregate and log cudagraph metrics"""

36 +

37 + COLUMN_HEADERS = [

38 + "Unpadded Tokens",

39 + "Padded Tokens",

40 + "Num Paddings",

41 + "Runtime Mode",

42 + "Count",

43 + ]

44 +

45 + def __init__(self, cg_mode: CUDAGraphMode, cg_capture_sizes: list[int] | None):

46 + self.reset()

47 + self.cg_mode = str(cg_mode)

48 + self.cg_capture_sizes = str(cg_capture_sizes or [])

49 +

50 + self.settings_header = (

51 + "**CUDAGraph Config Settings:**\n\n"

52 + f"- Mode: {self.cg_mode}\n"

53 + f"- Capture sizes: {self.cg_capture_sizes}\n\n"

54 + "**CUDAGraph Stats:**\n\n"

55 + )

56 +

57 + def reset(self):

58 + self.stats = []

59 +

60 + def observe(self, cudagraph_stat: CUDAGraphStat):

61 + self.stats.append(cudagraph_stat)

62 +

63 + def generate_metric_table(self) -> str:

64 + stats_counts = Counter(self.stats)

65 +

66 + # Convert stats to rows of strings, in descending order of observed frequencies

67 + rows = []

68 + for stat, count in sorted(

69 + stats_counts.items(), key=lambda item: item[1], reverse=True

70 + ):

71 + rows.append(

72 + [

73 + str(stat.num_unpadded_tokens),

74 + str(stat.num_padded_tokens),

75 + str(stat.num_paddings),

76 + stat.runtime_mode,

77 + str(count),

78 + ]

79 + )

80 +

81 + # Calculate column widths (max of header and data)

82 + col_widths = []

83 + for i, header_text in enumerate(self.COLUMN_HEADERS):

84 + max_width = len(header_text)

85 + for row in rows:

86 + max_width = max(max_width, len(row[i]))

87 + col_widths.append(max_width)

88 +

89 + table_header_list = [

90 + h.ljust(w) for h, w in zip(self.COLUMN_HEADERS, col_widths)

91 + ]

92 + table_header = "| " + " | ".join(table_header_list) + " |\n"

93 +

94 + table_separator = "|" + "|".join("-" * (w + 2) for w in col_widths) + "|\n"

95 +

96 + # Create data rows with proper alignment

97 + data_rows = []

98 + for row in rows:

99 + formatted_row = [

100 + str(val).ljust(width) for val, width in zip(row, col_widths)

101 + ]

102 + data_rows.append("| " + " | ".join(formatted_row) + " |")

103 +

104 + return (

105 + self.settings_header

106 + + table_header

107 + + table_separator

108 + + "\n".join(data_rows)

109 + + "\n"

110 + )

111 +

112 + def log(self, log_fn=logger.info):

113 + if not self.stats:

114 + return

115 + log_fn(self.generate_metric_table())

116 + self.reset()

117 +

118 +

119 @dataclasses.dataclass

120 class CUDAGraphEntry:

121 batch_descriptor: BatchDescriptor

vllm/config/observability.py

@@ -55,6 +55,10 @@ class ObservabilityConfig:

55 kv_cache_metrics_sample: float = Field(default=0.01, gt=0, le=1)

56 """Sampling rate for KV cache metrics (0.0, 1.0]. Default 0.01 = 1% of blocks."""

58 + cudagraph_metrics: bool = False

59 + """Enable CUDA graph metrics (number of padded/unpadded tokens, runtime cudagraph

60 + dispatch modes, and their observed frequencies at every logging interval)."""

61 +

62 @cached_property

63 def collect_model_forward_time(self) -> bool:

64 """Whether to collect model forward time for the request."""

vllm/engine/arg_utils.py

@@ -518,6 +518,7 @@ class EngineArgs:

518 kv_cache_metrics_sample: float = get_field(

519 ObservabilityConfig, "kv_cache_metrics_sample"

520 )

521 + cudagraph_metrics: bool = ObservabilityConfig.cudagraph_metrics

522 scheduling_policy: SchedulerPolicy = SchedulerConfig.policy

523 scheduler_cls: str | type[object] | None = SchedulerConfig.scheduler_cls

524

@@ -1021,6 +1022,10 @@ class EngineArgs:

1022 "--kv-cache-metrics-sample",

1023 **observability_kwargs["kv_cache_metrics_sample"],

1024 )

1025 + observability_group.add_argument(

1026 + "--cudagraph-metrics",

1027 + **observability_kwargs["cudagraph_metrics"],

1028 + )

1029

1030 # Scheduler arguments

1031 scheduler_kwargs = get_kwargs(SchedulerConfig)

@@ -1698,6 +1703,7 @@ class EngineArgs:

1703 collect_detailed_traces=self.collect_detailed_traces,

1704 kv_cache_metrics=self.kv_cache_metrics,

1705 kv_cache_metrics_sample=self.kv_cache_metrics_sample,

1706 + cudagraph_metrics=self.cudagraph_metrics,

1707 )

1708

1709 # Compilation config overrides

vllm/v1/core/sched/scheduler.py

+7 -1

@@ -7,6 +7,7 @@ from collections.abc import Iterable

7 from typing import Any

9 from vllm import envs

10 +from vllm.compilation.cuda_graph import CUDAGraphStat

11 from vllm.config import VllmConfig

12 from vllm.distributed.ec_transfer.ec_connector.base import (

13 ECConnectorMetadata,

@@ -1037,6 +1038,7 @@ class Scheduler(SchedulerInterface):

1038 pooler_outputs = model_runner_output.pooler_output

1039 num_nans_in_logits = model_runner_output.num_nans_in_logits

1040 kv_connector_output = model_runner_output.kv_connector_output

1041 + cudagraph_stats = model_runner_output.cudagraph_stats

1042

1043 outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)

1044 spec_decoding_stats: SpecDecodingStats | None = None

@@ -1219,7 +1221,9 @@ class Scheduler(SchedulerInterface):

1221 finished_req_ids.clear()

1222

1223 if (

1222 - stats := self.make_stats(spec_decoding_stats, kv_connector_stats)

1224 + stats := self.make_stats(

1225 + spec_decoding_stats, kv_connector_stats, cudagraph_stats

1226 + )

1227 ) is not None:

1228 # Return stats to only one of the front-ends.

1229 if (eco := next(iter(engine_core_outputs.values()), None)) is None:

@@ -1420,6 +1424,7 @@ class Scheduler(SchedulerInterface):

1424 self,

1425 spec_decoding_stats: SpecDecodingStats | None = None,

1426 kv_connector_stats: KVConnectorStats | None = None,

1427 + cudagraph_stats: CUDAGraphStat | None = None,

1428 ) -> SchedulerStats | None:

1429 if not self.log_stats:

1430 return None

@@ -1444,6 +1449,7 @@ class Scheduler(SchedulerInterface):

1449 kv_cache_eviction_events=eviction_events,

1450 spec_decoding_stats=spec_stats,

1451 kv_connector_stats=connector_stats_payload,

1452 + cudagraph_stats=cudagraph_stats,

1453 )

1454

1455 def make_spec_decoding_stats(

vllm/v1/metrics/loggers.py

+14

@@ -10,6 +10,7 @@ from typing import TypeAlias

10 from prometheus_client import Counter, Gauge, Histogram

12 import vllm.envs as envs

13 +from vllm.compilation.cuda_graph import CUDAGraphLogging

14 from vllm.config import SupportsMetricsInfo, VllmConfig

15 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (

16 KVConnectorLogging,

@@ -106,6 +107,12 @@ class LoggingStatLogger(StatLoggerBase):

107 self.spec_decoding_logging = SpecDecodingLogging()

108 kv_transfer_config = self.vllm_config.kv_transfer_config

109 self.kv_connector_logging = KVConnectorLogging(kv_transfer_config)

110 + self.cudagraph_logging = None

111 + if self.vllm_config.observability_config.cudagraph_metrics:

112 + self.cudagraph_logging = CUDAGraphLogging(

113 + self.vllm_config.compilation_config.cudagraph_mode,

114 + self.vllm_config.compilation_config.cudagraph_capture_sizes,

115 + )

116 self.last_prompt_throughput: float = 0.0

117 self.last_generation_throughput: float = 0.0

118 self.engine_is_idle = False

@@ -161,6 +168,11 @@ class LoggingStatLogger(StatLoggerBase):

168 self.spec_decoding_logging.observe(scheduler_stats.spec_decoding_stats)

169 if kv_connector_stats := scheduler_stats.kv_connector_stats:

170 self.kv_connector_logging.observe(kv_connector_stats)

171 + if (

172 + self.cudagraph_logging is not None

173 + and scheduler_stats.cudagraph_stats is not None

174 + ):

175 + self.cudagraph_logging.observe(scheduler_stats.cudagraph_stats)

176 if not self.aggregated:

177 self.last_scheduler_stats = scheduler_stats

178 if mm_cache_stats:

@@ -240,6 +252,8 @@ class LoggingStatLogger(StatLoggerBase):

252

253 self.spec_decoding_logging.log(log_fn=log_fn)

254 self.kv_connector_logging.log(log_fn=log_fn)

255 + if self.cudagraph_logging is not None:

256 + self.cudagraph_logging.log(log_fn=log_fn)

257

258 def log_engine_initialized(self):

259 if self.vllm_config.cache_config.num_gpu_blocks:

vllm/v1/metrics/stats.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass, field

7 from typing import TYPE_CHECKING, Any

9 import vllm.envs as envs

10 +from vllm.compilation.cuda_graph import CUDAGraphStat

11 from vllm.v1.spec_decode.metrics import SpecDecodingStats

13 if TYPE_CHECKING:

@@ -183,6 +184,8 @@ class SchedulerStats:

184 waiting_lora_adapters: dict[str, int] = field(default_factory=dict)

185 running_lora_adapters: dict[str, int] = field(default_factory=dict)

186

187 + cudagraph_stats: CUDAGraphStat | None = None

188 +

189

190 @dataclass

191 class RequestStateStats:

vllm/v1/outputs.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, NamedTuple

8 import numpy as np

9 import torch

11 +from vllm.compilation.cuda_graph import CUDAGraphStat

12 from vllm.v1.core.sched.output import SchedulerOutput

14 if TYPE_CHECKING:

@@ -169,6 +170,9 @@ class ModelRunnerOutput:

170 # req_id -> num_nans_in_logits

171 num_nans_in_logits: dict[str, int] | None = None

172

173 + # information related to cudagraph execution

174 + cudagraph_stats: CUDAGraphStat | None = None

175 +

176

177 # ModelRunnerOutput wrapper for async scheduling.

178 class AsyncModelRunnerOutput(ABC):

vllm/v1/worker/gpu_model_runner.py

+28 -4

@@ -27,7 +27,7 @@ from vllm.attention.backends.abstract import (

27 )

28 from vllm.attention.layer import Attention, MLAAttention

29 from vllm.compilation.counter import compilation_counter

30 -from vllm.compilation.cuda_graph import CUDAGraphWrapper

30 +from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper

31 from vllm.compilation.monitor import set_cudagraph_capturing_enabled

32 from vllm.config import (

33 CompilationMode,

@@ -257,6 +257,7 @@ class ExecuteModelState(NamedTuple):

257 sample_hidden_states: torch.Tensor

258 aux_hidden_states: list[torch.Tensor] | None

259 ec_connector_output: ECConnectorOutput | None

260 + cudagraph_stats: CUDAGraphStat | None

261

262

263 class GPUModelRunner(

@@ -2755,7 +2756,11 @@ class GPUModelRunner(

2756 force_uniform_decode: bool | None = None,

2757 force_has_lora: bool | None = None,

2758 ) -> tuple[

2758 - CUDAGraphMode, BatchDescriptor, UBatchSlices | None, torch.Tensor | None

2759 + CUDAGraphMode,

2760 + BatchDescriptor,

2761 + UBatchSlices | None,

2762 + torch.Tensor | None,

2763 + CUDAGraphStat | None,

2764 ]:

2765 num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens)

2766 uniform_decode = (

@@ -2820,7 +2825,22 @@ class GPUModelRunner(

2825 # num_tokens_across_dp will no-longer be valid

2826 assert batch_descriptor.num_tokens == num_tokens_padded

2827

2823 - return cudagraph_mode, batch_descriptor, ubatch_slices, num_tokens_across_dp

2828 + cudagraph_stats = None

2829 + if self.vllm_config.observability_config.cudagraph_metrics:

2830 + cudagraph_stats = CUDAGraphStat(

2831 + num_unpadded_tokens=num_tokens,

2832 + num_padded_tokens=batch_descriptor.num_tokens,

2833 + num_paddings=batch_descriptor.num_tokens - num_tokens,

2834 + runtime_mode=str(cudagraph_mode),

2835 + )

2836 +

2837 + return (

2838 + cudagraph_mode,

2839 + batch_descriptor,

2840 + ubatch_slices,

2841 + num_tokens_across_dp,

2842 + cudagraph_stats,

2843 + )

2844

2845 @torch.inference_mode()

2846 def execute_model(

@@ -2918,6 +2938,7 @@ class GPUModelRunner(

2938 batch_desc,

2939 ubatch_slices,

2940 num_tokens_across_dp,

2941 + cudagraph_stats,

2942 ) = self._determine_batch_execution_and_padding(

2943 num_tokens=num_tokens_unpadded,

2944 num_reqs=num_reqs,

@@ -3067,6 +3088,7 @@ class GPUModelRunner(

3088 sample_hidden_states,

3089 aux_hidden_states,

3090 ec_connector_output,

3091 + cudagraph_stats,

3092 )

3093 self.kv_connector_output = kv_connector_output

3094 return None

@@ -3102,6 +3124,7 @@ class GPUModelRunner(

3124 sample_hidden_states,

3125 aux_hidden_states,

3126 ec_connector_output,

3127 + cudagraph_stats,

3128 ) = self.execute_model_state

3129 # Clear ephemeral state.

3130 self.execute_model_state = None

@@ -3217,6 +3240,7 @@ class GPUModelRunner(

3240 if self.supports_mm_inputs

3241 else None,

3242 num_nans_in_logits=num_nans_in_logits,

3243 + cudagraph_stats=cudagraph_stats,

3244 )

3245

3246 if not self.use_async_scheduling:

@@ -3937,7 +3961,7 @@ class GPUModelRunner(

3961

3962 num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)

3963

3940 - _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp = (

3964 + _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp, _ = (

3965 self._determine_batch_execution_and_padding(

3966 num_tokens=num_tokens_unpadded,

3967 num_reqs=num_reqs,

vllm/v1/worker/gpu_worker.py

+1 -1

@@ -564,7 +564,7 @@ class Worker(WorkerBase):

564 # TODO(lucas): This is pretty gross; ideally we should only ever call

565 # `_determine_batch_execution_and_padding` once (will get called again

566 # in `execute_model`) but this requires a larger refactor of PP.

567 - _, batch_desc, _, _ = (

567 + _, batch_desc, _, _, _ = (

568 self.model_runner._determine_batch_execution_and_padding(

569 num_tokens=num_scheduled_tokens,

570 num_reqs=len(num_scheduled_tokens_np),