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 9, 06:15 UTC
[ez] Add structured torch.compile logs (#33213)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Jan 31, 2026 at 05:00 UTC
608b556507acbb15c88d4c119414fc8364627472
3 files changed
+227
tests/compile/test_structured_logging.py
new
+121
@@ -0,0 +1,121 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import patch
4
+
5
+import pytest
6
+import regex as re
7
+import torch
8
+from torch import nn
9
+
10
+import tests.compile.silly_attention # noqa
11
+from vllm.compilation.decorators import support_torch_compile
12
+from vllm.config import VllmConfig, set_current_vllm_config
13
+from vllm.config.compilation import (
14
+ CompilationConfig,
15
+ CompilationMode,
16
+ CUDAGraphMode,
17
+)
18
+from vllm.config.scheduler import SchedulerConfig
19
+from vllm.forward_context import set_forward_context
20
+
21
+MLP_SIZE = 64
22
+
23
+
24
+@support_torch_compile
25
+class SimpleModel(nn.Module):
26
+ """A simple model with a splitting op for piecewise compilation."""
27
+
28
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs):
29
+ super().__init__()
30
+
31
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
32
+ x = x + x
33
+ attn_output = torch.empty_like(x)
34
+ torch.ops.silly.attention(x, x, x, attn_output)
35
+ x = attn_output * 2
36
+ return x
37
+
38
+
39
+class TraceStructuredCapture:
40
+ """Captures trace_structured calls for testing."""
41
+
42
+ def __init__(self):
43
+ self.calls: list[dict] = []
44
+
45
+ def __call__(self, event_type: str, metadata_fn=None, payload_fn=None, **kwargs):
46
+ """Capture a trace_structured call."""
47
+ metadata = metadata_fn() if metadata_fn else {}
48
+ self.calls.append(
49
+ {
50
+ "event_type": event_type,
51
+ "metadata": metadata,
52
+ }
53
+ )
54
+
55
+ def get(self, event_type: str, name_pattern: str) -> list[dict]:
56
+ """Get all calls with the given event type and name matching pattern.
57
+
58
+ Args:
59
+ event_type: The event type to filter by (e.g., "artifact", "graph_dump")
60
+ name_pattern: Regex pattern to match against the artifact name
61
+ """
62
+ regex = re.compile(name_pattern)
63
+ return [
64
+ c
65
+ for c in self.calls
66
+ if c["event_type"] == event_type
67
+ and regex.fullmatch(c.get("metadata", {}).get("name", ""))
68
+ ]
69
+
70
+
71
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
72
+def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache):
73
+ """Test that all expected vLLM artifacts are logged during compilation."""
74
+ torch.set_default_device("cuda")
75
+
76
+ capture = TraceStructuredCapture()
77
+
78
+ vllm_config = VllmConfig(
79
+ compilation_config=CompilationConfig(
80
+ mode=CompilationMode.VLLM_COMPILE,
81
+ cudagraph_mode=CUDAGraphMode.PIECEWISE,
82
+ compile_sizes=[8],
83
+ splitting_ops=["silly::attention"],
84
+ ),
85
+ scheduler_config=SchedulerConfig(
86
+ max_num_seqs=8,
87
+ max_model_len=8192,
88
+ is_encoder_decoder=False,
89
+ ),
90
+ )
91
+
92
+ # Patch trace_structured to capture calls
93
+ with (
94
+ patch("vllm.compilation.backends.trace_structured", capture),
95
+ patch("vllm.compilation.piecewise_backend.trace_structured", capture),
96
+ set_current_vllm_config(vllm_config),
97
+ ):
98
+ model = SimpleModel(vllm_config=vllm_config, prefix="test")
99
+ with set_forward_context({}, vllm_config=vllm_config):
100
+ model(torch.randn(8, MLP_SIZE))
101
+
102
+ config_artifacts = capture.get("artifact", "vllm_compilation_config")
103
+ assert len(config_artifacts) == 1, (
104
+ f"Expected 1 vllm_compilation_config, got {len(config_artifacts)}"
105
+ )
106
+ vllm_piecewise_split_graph = capture.get("graph_dump", "vllm_piecewise_split_graph")
107
+ assert len(vllm_piecewise_split_graph) == 1, (
108
+ "Expected 1 toplevel piecewise split graph, "
109
+ f"got {len(vllm_piecewise_split_graph)}"
110
+ )
111
+ compile_start_artifacts = capture.get("artifact", "vllm_piecewise_compile_start")
112
+ assert len(compile_start_artifacts) == 2, (
113
+ "Expected 2 vllm_piecewise_compile_start "
114
+ "(one for dynamic ranges, one for compile size), "
115
+ f"got {len(compile_start_artifacts)}"
116
+ )
117
+ submod_dumps = capture.get("graph_dump", r"vllm_submod_.*")
118
+ assert len(submod_dumps) == 2, (
119
+ "Expected 2 submods (one before attention, one after attention), "
120
+ f"got {len(submod_dumps)}"
121
+ )
vllm/compilation/backends.py
+58
@@ -19,6 +19,7 @@ from typing import Any
19
import torch
20
import torch.fx as fx
21
from torch._dispatch.python import enable_python_dispatcher
22
+from torch._logging._internal import trace_structured
23
24
import vllm.envs as envs
25
from vllm.compilation.inductor_pass import pass_context
@@ -529,6 +530,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
530
sym_shape_indices,
531
self.vllm_backend,
532
graph_returns_tuple(submod),
533
+ submod_name=target,
534
)
535
536
self.module.__dict__[target] = wrap_with_cudagraph_if_needed(
@@ -735,12 +737,61 @@ class VllmBackend:
737
)
738
self.inductor_config[self.pass_key] = self.pass_manager
739
740
+ def _log_compilation_config(self):
741
+ """Log vLLM compilation config for TORCH_TRACE/tlparse."""
742
+ cc = self.compilation_config
743
+ pass_cfg = cc.pass_config
744
+
745
+ # Helper to convert lists to comma-separated strings for tlparse display
746
+ def list_to_str(lst: list | None) -> str:
747
+ if lst is None:
748
+ return ""
749
+ return ", ".join(str(x) for x in lst)
750
+
751
+ # Get enabled passes by introspecting dataclass fields
752
+ enabled_passes = [
753
+ f.name
754
+ for f in dataclasses.fields(pass_cfg)
755
+ if isinstance(getattr(pass_cfg, f.name), bool) and getattr(pass_cfg, f.name)
756
+ ]
757
+
758
+ trace_structured(
759
+ "artifact",
760
+ metadata_fn=lambda: {
761
+ "name": "vllm_compilation_config",
762
+ "encoding": "json",
763
+ },
764
+ payload_fn=lambda: json.dumps(
765
+ {
766
+ "model": self.vllm_config.model_config.model,
767
+ "prefix": self.prefix,
768
+ "mode": str(cc.mode),
769
+ "backend": cc.backend,
770
+ "custom_ops": list_to_str(cc.custom_ops),
771
+ "splitting_ops": list_to_str(cc.splitting_ops),
772
+ "cudagraph_mode": str(cc.cudagraph_mode),
773
+ "compile_sizes": list_to_str(cc.compile_sizes),
774
+ "compile_ranges_split_points": list_to_str(
775
+ cc.compile_ranges_split_points
776
+ ),
777
+ "use_inductor_graph_partition": cc.use_inductor_graph_partition,
778
+ "inductor_passes": list_to_str(list(cc.inductor_passes.keys())),
779
+ "enabled_passes": list_to_str(enabled_passes),
780
+ "dynamic_shapes_type": str(cc.dynamic_shapes_config.type),
781
+ "dynamic_shapes_evaluate_guards": cc.dynamic_shapes_config.evaluate_guards, # noqa: E501
782
+ }
783
+ ),
784
+ )
785
+
786
def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any:
787
from .caching import (
788
VllmSerializableFunction,
789
)
790
791
vllm_config = self.vllm_config
792
+
793
+ self._log_compilation_config()
794
+
795
# Minimal hashing here with existing utilities, reused below.
796
797
env_factors = envs.compile_factors()
@@ -892,6 +943,13 @@ class VllmBackend:
943
lazy_format_graph_code("before split", self.graph)
944
lazy_format_graph_code("after split", self.split_gm)
945
946
+ # Log the piecewise split graph for TORCH_TRACE/tlparse
947
+ trace_structured(
948
+ "graph_dump",
949
+ metadata_fn=lambda: {"name": "vllm_piecewise_split_graph"},
950
+ payload_fn=lambda: self.split_gm.print_readable(print_output=False),
951
+ )
952
+
953
compilation_counter.num_piecewise_graphs_seen += len(self.piecewise_graphs)
954
submod_names_to_compile = [
955
item.submod_name
vllm/compilation/piecewise_backend.py
+48
@@ -3,6 +3,7 @@
3
4
import dataclasses
5
import io
6
+import json
7
import pickle
8
from collections.abc import Callable
9
from pickle import Pickler
@@ -11,6 +12,7 @@ from typing import Any
12
import torch._functorch.config
13
import torch.fx as fx
14
from torch._inductor.runtime.triton_heuristics import CachingAutotuner
15
+from torch._logging._internal import trace_structured
16
17
from vllm.compilation.backends import VllmBackend
18
from vllm.compilation.monitor import end_monitoring_torch_compile
@@ -39,6 +41,7 @@ class PiecewiseBackend:
41
vllm_backend: VllmBackend,
42
returns_tuple: bool,
43
compiled_runnables: dict[str, Callable[..., Any]] | None = None,
44
+ submod_name: str = "",
45
):
46
"""
47
The backend for piecewise compilation.
@@ -70,6 +73,7 @@ class PiecewiseBackend:
73
self.total_piecewise_compiles = total_piecewise_compiles
74
self.vllm_backend = vllm_backend
75
self.compiled_runnables = compiled_runnables
76
+ self.submod_name = submod_name
77
78
self.is_first_graph = piecewise_compile_index == 0
79
self.is_last_graph = piecewise_compile_index == total_piecewise_compiles - 1
@@ -131,6 +135,9 @@ class PiecewiseBackend:
135
compile_range=range,
136
)
137
138
+ # Track whether we've logged the graph for this subgraph (only log once)
139
+ self._graph_logged = False
140
+
141
# get the on_compilation_complete callback from context...
142
# PiecewiseBackend is created during the first call,
143
# which is when the context is set (see compilation/decorators.py)
@@ -221,6 +228,45 @@ class PiecewiseBackend:
228
assert len(fake_example_inputs) == len(args)
229
return fake_example_inputs
230
231
+ def _log_compile_start(self, compile_range: Range):
232
+ """Log compilation event for TORCH_TRACE/tlparse."""
233
+ is_cudagraph_size = (
234
+ self.compile_sizes is not None and compile_range.start in self.compile_sizes
235
+ )
236
+ subgraph_index = self.piecewise_compile_index
237
+ submod_name = self.submod_name
238
+ trace_structured(
239
+ "artifact",
240
+ metadata_fn=lambda: {
241
+ "name": "vllm_piecewise_compile_start",
242
+ "encoding": "json",
243
+ },
244
+ payload_fn=lambda: json.dumps(
245
+ {
246
+ "piecewise_index": subgraph_index,
247
+ "submod_name": submod_name,
248
+ "total_piecewise_compiles": self.total_piecewise_compiles,
249
+ "compile_range_start": compile_range.start,
250
+ "compile_range_end": compile_range.end,
251
+ "is_single_size": compile_range.is_single_size(),
252
+ "is_cudagraph_capture_size": is_cudagraph_size,
253
+ }
254
+ ),
255
+ )
256
+
257
+ # Log the subgraph graph dump only once per subgraph (not per size)
258
+ # to reduce log file size. The graph code is the same for all sizes.
259
+ if not self._graph_logged:
260
+ self._graph_logged = True
261
+ assert self.graph is not None
262
+ trace_structured(
263
+ "graph_dump",
264
+ metadata_fn=lambda: {
265
+ "name": f"vllm_{submod_name}",
266
+ },
267
+ payload_fn=lambda: self.graph.print_readable(print_output=False),
268
+ )
269
+
270
def _maybe_compile_for_range_entry(
271
self, range_entry: RangeEntry, args: tuple[Any, ...]
272
) -> Any:
@@ -230,6 +276,8 @@ class PiecewiseBackend:
276
self.compiled_runnables[str(range_entry.compile_range)]
277
)
278
else:
279
+ self._log_compile_start(range_entry.compile_range)
280
+
281
# args are real arguments
282
# fakify for range, real args for concrete size.
283
# For concrete size, we clear the shape env in