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 11, 23:45 UTC
[Feature] Add Layer-wise NVTX Support (#29990)
Signed-off-by: Max Hu <hyoung2991@gmail.com> Signed-off-by: Max Hu <maxhu@nvidia.com> Co-authored-by: Max Hu <maxhu@nvidia.com>
Max Hu committed
Dec 5, 2025 at 06:20 UTC
c2894d3883c1b21299822c685e6be73ec3be046b
5 files changed
+375
-3
vllm/compilation/wrapper.py
+27
-3
@@ -14,6 +14,7 @@ import torch._C._dynamo.guards
14
import vllm.envs as envs
15
from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
16
from vllm.logger import init_logger
17
+from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context
18
19
logger = init_logger(__name__)
20
@@ -92,12 +93,29 @@ class TorchCompileWithNoGuardsWrapper:
93
94
return self.forward(*args, **kwargs)
95
96
+ def _call_with_optional_nvtx_range(self, callable_fn, *args, **kwargs):
97
+ if self.layerwise_nvtx_tracing_enabled:
98
+ args_list = list(args)
99
+ kwargs_dict = dict(kwargs)
100
+ with layerwise_nvtx_marker_context(
101
+ "Torch Compiled Module (input):{}".format(self.__class__.__name__),
102
+ self,
103
+ in_tensor=args_list,
104
+ kwargs=kwargs_dict,
105
+ ) as ctx:
106
+ ctx.result = callable_fn(*args, **kwargs)
107
+ return ctx.result
108
+ return callable_fn(*args, **kwargs)
109
+
110
def __init__(self):
111
self.compiled = False
112
113
vllm_config = get_current_vllm_config()
114
self.vllm_config = vllm_config
115
mode = vllm_config.compilation_config.mode
116
+ self.layerwise_nvtx_tracing_enabled = (
117
+ vllm_config.observability_config.enable_layerwise_nvtx_tracing
118
+ )
119
if mode is None:
120
raise RuntimeError("Compilation mode cannot be NO_COMPILATION")
121
@@ -168,13 +186,19 @@ class TorchCompileWithNoGuardsWrapper:
186
# Make sure a compilation is triggered by clearing dynamo
187
# cache.
188
torch._dynamo.eval_frame.remove_from_cache(self.original_code_object())
171
- return self._compiled_callable(*args, **kwargs)
189
+ return self._call_with_optional_nvtx_range(
190
+ self._compiled_callable, *args, **kwargs
191
+ )
192
else:
193
with self._dispatch_to_compiled_code():
174
- return self.forward(*args, **kwargs)
194
+ return self._call_with_optional_nvtx_range(
195
+ self.forward, *args, **kwargs
196
+ )
197
else:
198
with _compilation_context():
177
- return self._compiled_callable(*args, **kwargs)
199
+ return self._call_with_optional_nvtx_range(
200
+ self._compiled_callable, *args, **kwargs
201
+ )
202
203
@abstractmethod
204
def forward(self, *args, **kwargs): ...
vllm/config/observability.py
+5
@@ -59,6 +59,11 @@ class ObservabilityConfig:
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
+ enable_layerwise_nvtx_tracing: bool = False
63
+ """Enable layerwise NVTX tracing. This traces the execution of each layer or
64
+ module in the model and attach informations such as input/output shapes to
65
+ nvtx range markers. Noted that this doesn't work with CUDA graphs enabled."""
66
+
67
@cached_property
68
def collect_model_forward_time(self) -> bool:
69
"""Whether to collect model forward time for the request."""
vllm/engine/arg_utils.py
+8
@@ -519,6 +519,9 @@ class EngineArgs:
519
ObservabilityConfig, "kv_cache_metrics_sample"
520
)
521
cudagraph_metrics: bool = ObservabilityConfig.cudagraph_metrics
522
+ enable_layerwise_nvtx_tracing: bool = (
523
+ ObservabilityConfig.enable_layerwise_nvtx_tracing
524
+ )
525
scheduling_policy: SchedulerPolicy = SchedulerConfig.policy
526
scheduler_cls: str | type[object] | None = SchedulerConfig.scheduler_cls
527
@@ -1026,6 +1029,10 @@ class EngineArgs:
1029
"--cudagraph-metrics",
1030
**observability_kwargs["cudagraph_metrics"],
1031
)
1032
+ observability_group.add_argument(
1033
+ "--enable-layerwise-nvtx-tracing",
1034
+ **observability_kwargs["enable_layerwise_nvtx_tracing"],
1035
+ )
1036
1037
# Scheduler arguments
1038
scheduler_kwargs = get_kwargs(SchedulerConfig)
@@ -1704,6 +1711,7 @@ class EngineArgs:
1711
kv_cache_metrics=self.kv_cache_metrics,
1712
kv_cache_metrics_sample=self.kv_cache_metrics_sample,
1713
cudagraph_metrics=self.cudagraph_metrics,
1714
+ enable_layerwise_nvtx_tracing=self.enable_layerwise_nvtx_tracing,
1715
)
1716
1717
# Compilation config overrides
vllm/utils/nvtx_pytorch_hooks.py
new
+286
@@ -0,0 +1,286 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from contextlib import contextmanager
5
+
6
+import torch
7
+import torch.cuda.nvtx as nvtx
8
+
9
+
10
+def print_tensor(tensor_obj, prefix, tensor_list=None):
11
+ """Descends iterators that contains Tensors and prints the Tensor.
12
+ Recursive function that descends iterator type arguments until
13
+ it finds a Tensor object.
14
+ """
15
+ if tensor_list is None:
16
+ tensor_list = []
17
+
18
+ if isinstance(tensor_obj, (list, tuple)):
19
+ for ten in tensor_obj:
20
+ tensor_list = print_tensor(ten, prefix, tensor_list)
21
+ elif isinstance(tensor_obj, torch.Tensor):
22
+ tensor_dims = list(tensor_obj.size())
23
+ tensor_list.append(tensor_dims)
24
+ return tensor_list
25
+
26
+
27
+def process_layer_params(module_obj):
28
+ """Extract the static parameters from LLM and VLM relevant layer types"""
29
+ param_info = {}
30
+ # Extract parameters for layers commonly used in LLMs and VLMs
31
+ if isinstance(module_obj, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)):
32
+ conv_params = {}
33
+ conv_params["in_chan"] = module_obj.in_channels
34
+ conv_params["out_chan"] = module_obj.out_channels
35
+ conv_params["filter_dim"] = module_obj.kernel_size
36
+ conv_params["stride"] = module_obj.stride
37
+ conv_params["padding"] = module_obj.padding
38
+ conv_params["dilation"] = module_obj.dilation
39
+ conv_params["transposed"] = module_obj.transposed
40
+ conv_params["output_padding"] = module_obj.output_padding
41
+ conv_params["groups"] = module_obj.groups
42
+ conv_params["padding_mode"] = module_obj.padding_mode
43
+ param_info = conv_params
44
+ elif isinstance(
45
+ module_obj,
46
+ (
47
+ torch.nn.ConvTranspose1d,
48
+ torch.nn.ConvTranspose2d,
49
+ torch.nn.ConvTranspose3d,
50
+ ),
51
+ ):
52
+ convtranspose_params = {}
53
+ convtranspose_params["in_chan"] = module_obj.in_channels
54
+ convtranspose_params["out_chan"] = module_obj.out_channels
55
+ convtranspose_params["filter_dim"] = module_obj.kernel_size
56
+ convtranspose_params["stride"] = module_obj.stride
57
+ convtranspose_params["padding"] = module_obj.padding
58
+ convtranspose_params["dilation"] = module_obj.dilation
59
+ convtranspose_params["transposed"] = module_obj.transposed
60
+ convtranspose_params["output_padding"] = module_obj.output_padding
61
+ convtranspose_params["groups"] = module_obj.groups
62
+ convtranspose_params["padding_mode"] = module_obj.padding_mode
63
+ param_info = convtranspose_params
64
+ elif isinstance(
65
+ module_obj, (torch.nn.MaxPool1d, torch.nn.MaxPool2d, torch.nn.MaxPool3d)
66
+ ):
67
+
68
+ def _handle_int_or_tuple(parameter):
69
+ if isinstance(parameter, tuple):
70
+ return list(parameter)
71
+ elif isinstance(parameter, int):
72
+ return [parameter, parameter]
73
+
74
+ pooling_params = {}
75
+ pooling_params["filter_dim"] = _handle_int_or_tuple(module_obj.kernel_size)
76
+ pooling_params["stride"] = _handle_int_or_tuple(module_obj.stride)
77
+ pooling_params["padding"] = _handle_int_or_tuple(module_obj.padding)
78
+ pooling_params["dilation"] = _handle_int_or_tuple(module_obj.dilation)
79
+ param_info = pooling_params
80
+ elif isinstance(
81
+ module_obj, (torch.nn.AvgPool1d, torch.nn.AvgPool2d, torch.nn.AvgPool3d)
82
+ ):
83
+ pooling_params = {}
84
+ pooling_params["filter_dim"] = [
85
+ module_obj.kernel_size,
86
+ module_obj.kernel_size,
87
+ ]
88
+ pooling_params["stride"] = [module_obj.stride, module_obj.stride]
89
+ pooling_params["padding"] = [module_obj.padding, module_obj.padding]
90
+ pooling_params["ceil_mode"] = module_obj.ceil_mode
91
+ pooling_params["count_include_pad"] = module_obj.count_include_pad
92
+ param_info = pooling_params
93
+ elif isinstance(
94
+ module_obj,
95
+ (
96
+ torch.nn.AdaptiveAvgPool1d,
97
+ torch.nn.AdaptiveAvgPool2d,
98
+ torch.nn.AdaptiveAvgPool3d,
99
+ ),
100
+ ):
101
+ pooling_params = {}
102
+ pooling_params["output_size"] = [
103
+ module_obj.output_size,
104
+ module_obj.output_size,
105
+ ]
106
+ param_info = pooling_params
107
+ elif isinstance(module_obj, torch.nn.Linear):
108
+ param_info["in_features"] = module_obj.in_features
109
+ param_info["out_features"] = module_obj.out_features
110
+ elif isinstance(
111
+ module_obj,
112
+ (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d),
113
+ ):
114
+ param_info["num_features"] = module_obj.num_features
115
+ param_info["epsilon"] = module_obj.eps
116
+ param_info["momentum"] = module_obj.momentum
117
+ elif isinstance(module_obj, torch.nn.ReLU):
118
+ param_info["in_place"] = module_obj.inplace
119
+ elif isinstance(module_obj, torch.nn.Dropout):
120
+ param_info["p"] = module_obj.p
121
+ param_info["in_place"] = module_obj.inplace
122
+ elif isinstance(module_obj, torch.nn.Embedding):
123
+ param_info["num_embeddings"] = module_obj.num_embeddings
124
+ param_info["embedding_dim"] = module_obj.embedding_dim
125
+ elif isinstance(
126
+ module_obj,
127
+ (
128
+ torch.nn.Upsample,
129
+ torch.nn.UpsamplingNearest2d,
130
+ torch.nn.UpsamplingBilinear2d,
131
+ ),
132
+ ):
133
+ param_info["scale_factor"] = module_obj.scale_factor
134
+
135
+ return param_info
136
+
137
+
138
+def construct_marker_dict_and_push(
139
+ module_name, module_obj, in_tensor, kwargs=None, out_tensor=None
140
+):
141
+ marker_dict = {}
142
+ marker_dict["Module"] = module_name
143
+
144
+ ## Get trainable parameters like weights and bias
145
+ module_params = module_obj.named_parameters(recurse=False)
146
+ for idx, (param_name, param_obj) in enumerate(module_params):
147
+ if idx == 0:
148
+ marker_dict["TrainableParams"] = {}
149
+ marker_dict["TrainableParams"][param_name] = list(param_obj.size())
150
+
151
+ in_tensor_list = print_tensor(in_tensor, "Input")
152
+ if in_tensor_list:
153
+ marker_dict["Inputs"] = in_tensor_list
154
+
155
+ out_tensor_list = print_tensor(out_tensor, "Output")
156
+ if out_tensor_list:
157
+ marker_dict["Outputs"] = out_tensor_list
158
+
159
+ ## Get Kwargs like input_ids and positions for the top module
160
+ if kwargs:
161
+ for key, value in kwargs.items():
162
+ if isinstance(value, (torch.Tensor, list, tuple)):
163
+ tensor_list = print_tensor(value, key)
164
+ if tensor_list:
165
+ marker_dict[key] = tensor_list
166
+
167
+ param_info = process_layer_params(module_obj)
168
+ if param_info:
169
+ marker_dict["StaticParams"] = param_info
170
+ nvtx.range_push("{}".format(marker_dict))
171
+
172
+
173
+class ResultHolder:
174
+ """Holder for storing results from within a context manager."""
175
+
176
+ result = None
177
+
178
+
179
+@contextmanager
180
+def layerwise_nvtx_marker_context(module_name, module_obj, in_tensor=None, kwargs=None):
181
+ """Context manager for NVTX markers that automatically pushes on enter
182
+ and pops on exit.
183
+
184
+ Example:
185
+ with nvtx_marker_context("Module:MyModule", module, in_tensor=args,
186
+ kwargs=kwargs) as ctx:
187
+ ctx.result = module(*args, **kwargs)
188
+ return ctx.result
189
+ """
190
+ holder = ResultHolder()
191
+
192
+ # Push input marker
193
+ construct_marker_dict_and_push(
194
+ module_name,
195
+ module_obj,
196
+ in_tensor=in_tensor,
197
+ kwargs=kwargs,
198
+ )
199
+ try:
200
+ yield holder
201
+ finally:
202
+ # Pop input marker
203
+ nvtx.range_pop()
204
+ # Push and pop output marker
205
+ output_name = module_name.replace("(input)", "(output)")
206
+ construct_marker_dict_and_push(
207
+ output_name,
208
+ module_obj,
209
+ in_tensor=None,
210
+ kwargs=None,
211
+ out_tensor=holder.result,
212
+ )
213
+ nvtx.range_pop()
214
+
215
+
216
+class PytHooks:
217
+ """This module contains all the code needed to enable forward hooks
218
+ in a pytorch network.
219
+
220
+ To register the hooks for a given network, the user needs to instantiate
221
+ a PytHook object. Then call the register_hooks method.
222
+
223
+ Example:
224
+
225
+ my_hook = PytHook()
226
+ my_hook.register_hooks(my_network_model)
227
+ """
228
+
229
+ def __init__(self):
230
+ """Initialize module variables."""
231
+ super().__init__()
232
+ self.module_to_name_map = {}
233
+
234
+ def _process_layer_params(self, module_obj):
235
+ return process_layer_params(module_obj)
236
+
237
+ def module_fwd_hook(self, module_obj, in_tensor, out_tensor):
238
+ """Callback function that ends the NVTX marker.
239
+ Records the module name and tensor information.
240
+ Called after the module executes the forward method.
241
+ """
242
+ nvtx.range_pop()
243
+ module_name = self.module_to_name_map.get(module_obj, "unknown")
244
+ construct_marker_dict_and_push(
245
+ module_name, module_obj, in_tensor=None, kwargs=None, out_tensor=out_tensor
246
+ )
247
+ nvtx.range_pop()
248
+ return
249
+
250
+ def module_fwd_pre_hook(self, module_obj, in_tensor, kwargs):
251
+ """Creates an NVTX marker with the module name in it.
252
+ This function is called before the module executes.
253
+ """
254
+ module_name = self.module_to_name_map.get(module_obj, "unknown")
255
+ construct_marker_dict_and_push(
256
+ module_name, module_obj, in_tensor=in_tensor, kwargs=kwargs, out_tensor=None
257
+ )
258
+ return
259
+
260
+ def register_hooks(self, network_model, module_prefix="top"):
261
+ """User level function that activates all the hooks.
262
+ The user needs to call this method from the network source code.
263
+ The code descends all the modules in the network and registers their
264
+ respective hooks.
265
+ """
266
+ # Module types to skip (simple operations that don't need detailed profiling)
267
+ skip_types = (
268
+ torch.nn.Identity,
269
+ torch.nn.Dropout,
270
+ torch.nn.Dropout1d,
271
+ torch.nn.Dropout2d,
272
+ torch.nn.Dropout3d,
273
+ )
274
+
275
+ for name, module in network_model.named_modules(prefix=module_prefix):
276
+ # Skip certain module types to reduce profiling overhead
277
+ if isinstance(module, skip_types):
278
+ continue
279
+
280
+ module.register_forward_pre_hook(self.module_fwd_pre_hook, with_kwargs=True)
281
+ module.register_forward_hook(self.module_fwd_hook)
282
+ if module not in self.module_to_name_map:
283
+ self.module_to_name_map[module] = name
284
+ else:
285
+ raise ValueError("Module instance {} is not unique ".format(module))
286
+ return
vllm/v1/worker/gpu_model_runner.py
+49
@@ -88,6 +88,7 @@ from vllm.utils.jsontree import json_map_leaves
88
from vllm.utils.math_utils import cdiv, round_up
89
from vllm.utils.mem_constants import GiB_bytes
90
from vllm.utils.mem_utils import DeviceMemoryProfiler
91
+from vllm.utils.nvtx_pytorch_hooks import PytHooks
92
from vllm.utils.platform_utils import is_pin_memory_available
93
from vllm.utils.torch_utils import (
94
get_dtype_size,
@@ -599,6 +600,7 @@ class GPUModelRunner(
600
# Ephemeral state transferred between execute_model() and sample_tokens().
601
self.execute_model_state: ExecuteModelState | None = None
602
self.kv_connector_output: KVConnectorOutput | None = None
603
+ self.layerwise_nvtx_hooks_registered = False
604
605
def reset_mm_cache(self) -> None:
606
if self.mm_budget:
@@ -2828,6 +2830,42 @@ class GPUModelRunner(
2830
cudagraph_stats,
2831
)
2832
2833
+ def _register_layerwise_nvtx_hooks(self) -> None:
2834
+ """
2835
+ Register layerwise NVTX hooks if --enable-layerwise-nvtx-tracing is enabled
2836
+ to trace detailed information of each layer or module in the model.
2837
+ """
2838
+
2839
+ if (
2840
+ self.vllm_config.observability_config.enable_layerwise_nvtx_tracing
2841
+ and not self.layerwise_nvtx_hooks_registered
2842
+ ):
2843
+ if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
2844
+ logger.debug_once(
2845
+ "layerwise NVTX tracing is not supported when CUDA graph is "
2846
+ "turned off; you may observe part or all of the model "
2847
+ "missing NVTX markers"
2848
+ )
2849
+
2850
+ # In STOCK_TORCH_COMPILE mode, after registering hooks here,
2851
+ # the __call__ function of nn.module will be recompiled with
2852
+ # fullgraph=True. Since nvtx.range_push/pop are not traceable
2853
+ # by torch dynamo, we can't register hook functions here
2854
+ # because hook functions will also be traced by torch dynamo.
2855
+ if (
2856
+ self.vllm_config.compilation_config.mode
2857
+ == CompilationMode.STOCK_TORCH_COMPILE
2858
+ ):
2859
+ logger.debug_once(
2860
+ "layerwise NVTX tracing is not supported when "
2861
+ "CompilationMode is STOCK_TORCH_COMPILE, skipping "
2862
+ "function hooks registration"
2863
+ )
2864
+ else:
2865
+ pyt_hooks = PytHooks()
2866
+ pyt_hooks.register_hooks(self.model, self.model.__class__.__name__)
2867
+ self.layerwise_nvtx_hooks_registered = True
2868
+
2869
@torch.inference_mode()
2870
def execute_model(
2871
self,
@@ -4122,6 +4160,17 @@ class GPUModelRunner(
4160
is_graph_capturing=is_graph_capturing,
4161
)
4162
4163
+ # We register layerwise NVTX hooks here after the first dynamo tracing is
4164
+ # done to avoid nvtx operations in hook functions being traced by
4165
+ # torch dynamo and causing graph breaks.
4166
+ # Note that for DYNAMO_ONCE and VLLM_COMPILE mode,
4167
+ # compiled model's dynamo tracing is only done once and the compiled model's
4168
+ # __call__ function is replaced by calling the compiled function.
4169
+ # So it's safe to register hooks here. Hooks will be registered to
4170
+ # both compiled and uncompiled models but they will never
4171
+ # be called on the compiled model execution path.
4172
+ self._register_layerwise_nvtx_hooks()
4173
+
4174
# This is necessary to avoid blocking DP.
4175
# For dummy runs, we typically skip EPLB since we don't have any real
4176
# requests to process.