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 14, 05:30 UTC
[torch.compile] Support conditional torch.compile per module (#22269)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Aug 20, 2025 at 09:52 UTC
dfd2382039c38be80d6c2c9b56e441b5bd7cd0ad
4 files changed
+308
-103
.buildkite/test-pipeline.yaml
+2
@@ -328,6 +328,7 @@ steps:
328
- pytest -v -s compile/test_sequence_parallelism.py
329
- pytest -v -s compile/test_async_tp.py
330
- pytest -v -s compile/test_fusion_all_reduce.py
331
+ - pytest -v -s compile/test_decorator.py
332
333
- label: PyTorch Fullgraph Smoke Test # 9min
334
mirror_hardwares: [amdexperimental]
@@ -341,6 +342,7 @@ steps:
342
- pytest -v -s compile/piecewise/test_simple.py
343
- pytest -v -s compile/piecewise/test_toy_llama.py
344
- pytest -v -s compile/piecewise/test_full_cudagraph.py
345
+ - pytest -v -s compile/piecewise/test_multiple_graphs.py
346
347
- label: PyTorch Fullgraph Test # 18min
348
mirror_hardwares: [amdexperimental]
tests/compile/piecewise/test_multiple_graphs.py
+36
-101
@@ -12,10 +12,9 @@ from vllm.compilation.backends import set_model_tag
12
from vllm.compilation.counter import compilation_counter
13
from vllm.compilation.decorators import (ignore_torch_compile,
14
support_torch_compile)
15
-from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
16
- set_current_vllm_config)
17
-from vllm.envs import VLLM_USE_V1
18
-from vllm.forward_context import set_forward_context
15
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
16
+ VllmConfig, set_current_vllm_config)
17
+from vllm.forward_context import BatchDescriptor, set_forward_context
18
from vllm.utils import direct_register_custom_op
19
20
# create a library to hold the custom op
@@ -164,104 +163,34 @@ class SimpleModelWithTwoGraphs(ParentModel):
163
return x
164
165
167
-def test_ignore_torch_compile_decorator():
168
- assert VLLM_USE_V1
169
-
170
- # piecewise
171
- vllm_config = VllmConfig(compilation_config=CompilationConfig(
172
- level=CompilationLevel.PIECEWISE,
173
- use_cudagraph=True,
174
- splitting_ops=["silly.attention"],
175
- cudagraph_capture_sizes=[1, 2],
176
- ))
177
-
178
- @support_torch_compile
179
- class A(nn.Module):
180
-
181
- def __init__(self,
182
- *,
183
- vllm_config: VllmConfig,
184
- prefix: str = '',
185
- **kwargs) -> None:
186
- super().__init__()
187
-
188
- def forward(self, x: torch.Tensor) -> torch.Tensor:
189
- x = x + x
190
- attn_output = torch.empty_like(x)
191
- torch.ops.silly.attention(x, x, x, attn_output)
192
- x = attn_output
193
- x = x * 3
194
- return x
195
-
196
- @ignore_torch_compile
197
- class B(A):
198
- ...
199
-
200
- @support_torch_compile
201
- class C(B):
202
- ...
203
-
204
- with set_current_vllm_config(vllm_config):
205
- mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
206
-
207
- # A has support_torch_compile
208
- with compilation_counter.expect(
209
- num_graphs_seen=1,
210
- num_piecewise_graphs_seen=3,
211
- num_piecewise_capturable_graphs_seen=2,
212
- num_backend_compilations=2,
213
- num_cudagraph_captured=4,
214
- # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
215
- ), set_forward_context({}, vllm_config=vllm_config):
216
- # first run is for compile
217
- mod_A(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
218
- # run cudagraph captured sizes
219
- mod_A(torch.randn(2, MLP_SIZE).cuda())
220
- mod_A(torch.randn(1, MLP_SIZE).cuda())
221
-
222
- with set_current_vllm_config(vllm_config):
223
- mod_B = B(vllm_config=vllm_config, prefix='').eval().cuda()
224
-
225
- # B's ignore_torch_compile should override A's support_torch_compile
226
- with compilation_counter.expect(
227
- num_graphs_seen=0,
228
- num_piecewise_graphs_seen=0,
229
- num_piecewise_capturable_graphs_seen=0,
230
- num_backend_compilations=0,
231
- num_cudagraph_captured=0,
232
- ), set_forward_context({}, vllm_config=vllm_config):
233
- mod_B(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
234
- mod_B(torch.randn(2, MLP_SIZE).cuda())
235
- mod_B(torch.randn(1, MLP_SIZE).cuda())
236
-
237
- with set_current_vllm_config(vllm_config):
238
- mod_C = C(vllm_config=vllm_config, prefix='').eval().cuda()
239
-
240
- # C's support_torch_compile should override B's ignore_torch_compile
241
- with compilation_counter.expect(
242
- num_graphs_seen=1,
243
- num_piecewise_graphs_seen=3,
244
- num_piecewise_capturable_graphs_seen=2,
245
- num_backend_compilations=2,
246
- num_cudagraph_captured=4,
247
- # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
248
- ), set_forward_context({}, vllm_config=vllm_config):
249
- mod_C(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
250
- mod_C(torch.randn(2, MLP_SIZE).cuda())
251
- mod_C(torch.randn(1, MLP_SIZE).cuda())
252
-
253
-
166
@torch.inference_mode
255
-def run_model(vllm_config, model: nn.Module, inputs: torch.Tensor):
167
+def run_model(vllm_config: VllmConfig, model: nn.Module, inputs: torch.Tensor,
168
+ cudagraph_runtime_mode: CUDAGraphMode):
169
with set_forward_context({}, vllm_config=vllm_config):
257
- # First run is for compile
170
+ # warmup for the model with cudagraph_mode NONE
171
model(inputs)
172
260
- # Run CUDAGraph captured sizes
261
- model(inputs[:2])
262
- model(inputs[:1])
263
-
264
- output = model(inputs[:2])
173
+ # simulate cudagraphs capturing
174
+ with set_forward_context({},
175
+ vllm_config=vllm_config,
176
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
177
+ batch_descriptor=BatchDescriptor(
178
+ num_tokens=2, )):
179
+ model(inputs[:2])
180
+ with set_forward_context({},
181
+ vllm_config=vllm_config,
182
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
183
+ batch_descriptor=BatchDescriptor(
184
+ num_tokens=1, )):
185
+ model(inputs[:1])
186
+
187
+ # simulate cudagraphs replay
188
+ with set_forward_context({},
189
+ vllm_config=vllm_config,
190
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
191
+ batch_descriptor=BatchDescriptor(
192
+ num_tokens=2, )):
193
+ output = model(inputs[:2])
194
195
output = output.cpu()
196
return output.cpu()
@@ -277,6 +206,7 @@ def test_multi_graph_piecewise_compile_outputs_equal():
206
splitting_ops=["silly.attention"],
207
cudagraph_capture_sizes=[1, 2],
208
))
209
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
210
211
with set_current_vllm_config(vllm_config):
212
model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
@@ -299,11 +229,13 @@ def test_multi_graph_piecewise_compile_outputs_equal():
229
num_cudagraph_captured=8,
230
# num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
231
):
302
- outputs.append(run_model(vllm_config, model, inputs))
232
+ outputs.append(
233
+ run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
234
235
# no compile or cudagraph
236
vllm_config = VllmConfig(compilation_config=CompilationConfig(
237
level=CompilationLevel.NO_COMPILATION, ))
238
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
239
240
with set_current_vllm_config(vllm_config):
241
model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
@@ -318,7 +250,8 @@ def test_multi_graph_piecewise_compile_outputs_equal():
250
num_backend_compilations=0,
251
num_cudagraph_captured=0,
252
):
321
- outputs.append(run_model(vllm_config, model, inputs))
253
+ outputs.append(
254
+ run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
255
256
# piecewise compile without CUDA graph
257
vllm_config = VllmConfig(compilation_config=CompilationConfig(
@@ -326,6 +259,7 @@ def test_multi_graph_piecewise_compile_outputs_equal():
259
use_cudagraph=False,
260
splitting_ops=["silly.attention"],
261
))
262
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
263
264
with set_current_vllm_config(vllm_config):
265
model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
@@ -340,7 +274,8 @@ def test_multi_graph_piecewise_compile_outputs_equal():
274
num_backend_compilations=4,
275
num_cudagraph_captured=0, # no cudagraph captured
276
):
343
- outputs.append(run_model(vllm_config, model, inputs))
277
+ outputs.append(
278
+ run_model(vllm_config, model, inputs, cudagraph_runtime_mode))
279
280
# Generally don't expect outputs with and without inductor
281
# to be bitwise equivalent
tests/compile/test_decorator.py
new
+251
@@ -0,0 +1,251 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+from torch import nn
5
+from torch.library import Library
6
+
7
+from vllm.compilation.counter import compilation_counter
8
+from vllm.compilation.decorators import (ignore_torch_compile,
9
+ support_torch_compile)
10
+from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
11
+ CUDAGraphMode, VllmConfig, set_current_vllm_config)
12
+from vllm.forward_context import BatchDescriptor, set_forward_context
13
+from vllm.utils import direct_register_custom_op
14
+
15
+# create a library to hold the custom op
16
+silly_lib = Library("silly", "FRAGMENT") # noqa
17
+
18
+BATCH_SIZE = 32
19
+MLP_SIZE = 128
20
+
21
+
22
+def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
23
+ out: torch.Tensor) -> None:
24
+ out.copy_(q)
25
+ out += k
26
+ out += v
27
+
28
+
29
+def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
30
+ out: torch.Tensor) -> None:
31
+ return
32
+
33
+
34
+direct_register_custom_op(
35
+ op_name="attention",
36
+ op_func=silly_attention,
37
+ mutates_args=["out"],
38
+ fake_impl=silly_attention_fake,
39
+ target_lib=silly_lib,
40
+)
41
+
42
+
43
+@torch.inference_mode
44
+def run_model(vllm_config: VllmConfig, model: nn.Module,
45
+ cudagraph_runtime_mode: CUDAGraphMode):
46
+ with set_forward_context({}, vllm_config=vllm_config):
47
+ # warmup for the model with cudagraph_mode NONE
48
+ model(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
49
+
50
+ # simulate cudagraphs capturing
51
+ with set_forward_context({},
52
+ vllm_config=vllm_config,
53
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
54
+ batch_descriptor=BatchDescriptor(
55
+ num_tokens=2, )):
56
+ model(torch.randn(2, MLP_SIZE).cuda())
57
+ with set_forward_context({},
58
+ vllm_config=vllm_config,
59
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
60
+ batch_descriptor=BatchDescriptor(
61
+ num_tokens=1, )):
62
+ model(torch.randn(1, MLP_SIZE).cuda())
63
+
64
+ # simulate cudagraphs replay
65
+ with set_forward_context({},
66
+ vllm_config=vllm_config,
67
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
68
+ batch_descriptor=BatchDescriptor(
69
+ num_tokens=2, )):
70
+ output = model(torch.randn(2, MLP_SIZE).cuda())
71
+
72
+ output = output.cpu()
73
+ return output.cpu()
74
+
75
+
76
+def test_ignore_torch_compile_decorator():
77
+ # piecewise
78
+ vllm_config = VllmConfig(compilation_config=CompilationConfig(
79
+ level=CompilationLevel.PIECEWISE,
80
+ use_cudagraph=True,
81
+ splitting_ops=["silly.attention"],
82
+ cudagraph_capture_sizes=[1, 2],
83
+ ))
84
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
85
+
86
+ @support_torch_compile
87
+ class A(nn.Module):
88
+
89
+ def __init__(self,
90
+ *,
91
+ vllm_config: VllmConfig,
92
+ prefix: str = '',
93
+ **kwargs) -> None:
94
+ super().__init__()
95
+
96
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
97
+ x = x + x
98
+ attn_output = torch.empty_like(x)
99
+ torch.ops.silly.attention(x, x, x, attn_output)
100
+ x = attn_output
101
+ x = x * 3
102
+ return x
103
+
104
+ @ignore_torch_compile
105
+ class B(A):
106
+ ...
107
+
108
+ @support_torch_compile
109
+ class C(B):
110
+ ...
111
+
112
+ with set_current_vllm_config(vllm_config):
113
+ mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
114
+
115
+ # A has support_torch_compile
116
+ with compilation_counter.expect(
117
+ num_graphs_seen=1,
118
+ num_piecewise_graphs_seen=3,
119
+ num_piecewise_capturable_graphs_seen=2,
120
+ num_backend_compilations=2,
121
+ num_cudagraph_captured=4,
122
+ # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
123
+ ):
124
+ run_model(vllm_config, mod_A, cudagraph_runtime_mode)
125
+
126
+ with set_current_vllm_config(vllm_config):
127
+ mod_B = B(vllm_config=vllm_config, prefix='').eval().cuda()
128
+
129
+ # B's ignore_torch_compile should override A's support_torch_compile
130
+ with compilation_counter.expect(
131
+ num_graphs_seen=0,
132
+ num_piecewise_graphs_seen=0,
133
+ num_piecewise_capturable_graphs_seen=0,
134
+ num_backend_compilations=0,
135
+ num_cudagraph_captured=0,
136
+ ):
137
+ run_model(vllm_config, mod_B, cudagraph_runtime_mode)
138
+
139
+ with set_current_vllm_config(vllm_config):
140
+ mod_C = C(vllm_config=vllm_config, prefix='').eval().cuda()
141
+
142
+ # C's support_torch_compile should override B's ignore_torch_compile
143
+ with compilation_counter.expect(
144
+ num_graphs_seen=1,
145
+ num_piecewise_graphs_seen=3,
146
+ num_piecewise_capturable_graphs_seen=2,
147
+ num_backend_compilations=2,
148
+ num_cudagraph_captured=4,
149
+ # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
150
+ ):
151
+ run_model(vllm_config, mod_C, cudagraph_runtime_mode)
152
+
153
+
154
+# Only enable torch.compile if
155
+# vllm_config.cache_config.kv_sharing_fast_prefill=True
156
+@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.
157
+ kv_sharing_fast_prefill)
158
+class B(nn.Module):
159
+
160
+ def __init__(self,
161
+ *,
162
+ vllm_config: VllmConfig,
163
+ prefix: str = '',
164
+ **kwargs) -> None:
165
+ super().__init__()
166
+
167
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
168
+ x = x + x
169
+ attn_output = torch.empty_like(x)
170
+ torch.ops.silly.attention(x, x, x, attn_output)
171
+ x = attn_output
172
+ x = x + x
173
+ return x
174
+
175
+
176
+# Only enable torch.compile if
177
+# vllm_config.cache_config.kv_sharing_fast_prefill=False
178
+@support_torch_compile(enable_if=lambda vllm_config: not vllm_config.
179
+ cache_config.kv_sharing_fast_prefill)
180
+class A(nn.Module):
181
+
182
+ def __init__(self,
183
+ *,
184
+ vllm_config: VllmConfig,
185
+ prefix: str = '',
186
+ **kwargs) -> None:
187
+ super().__init__()
188
+ self.mod1 = B(vllm_config=vllm_config, prefix=prefix, **kwargs)
189
+ self.mod2 = B(vllm_config=vllm_config, prefix=prefix, **kwargs)
190
+
191
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
192
+ x = self.mod1(x)
193
+ attn_output = torch.empty_like(x)
194
+ torch.ops.silly.attention(x, x, x, attn_output)
195
+ x = attn_output
196
+ x = self.mod2(x)
197
+ return x
198
+
199
+
200
+def test_conditional_compile_enable_if():
201
+ vllm_config = VllmConfig(cache_config=CacheConfig(
202
+ kv_sharing_fast_prefill=True, ),
203
+ compilation_config=CompilationConfig(
204
+ level=CompilationLevel.PIECEWISE,
205
+ use_cudagraph=True,
206
+ splitting_ops=["silly.attention"],
207
+ cudagraph_capture_sizes=[1, 2],
208
+ ))
209
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
210
+
211
+ with set_current_vllm_config(vllm_config):
212
+ mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
213
+
214
+ # A has support_torch_compile but enable_if fn returns False
215
+ # enalbe_if will be True for B, so we expect mod1 and mod2
216
+ # to be compiled
217
+ with compilation_counter.expect(
218
+ num_graphs_seen=2,
219
+ num_piecewise_graphs_seen=6,
220
+ # 3 piecewise graphs per instance of B()
221
+ num_piecewise_capturable_graphs_seen=4,
222
+ num_backend_compilations=4,
223
+ num_cudagraph_captured=8,
224
+ # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
225
+ ):
226
+ run_model(vllm_config, mod_A, cudagraph_runtime_mode)
227
+
228
+ # Set kv_sharing_fast_prefill=False
229
+ # which will cause A to be compiled and B to not be compiled
230
+ vllm_config = VllmConfig(cache_config=CacheConfig(
231
+ kv_sharing_fast_prefill=False, ),
232
+ compilation_config=CompilationConfig(
233
+ level=CompilationLevel.PIECEWISE,
234
+ use_cudagraph=True,
235
+ splitting_ops=["silly.attention"],
236
+ cudagraph_capture_sizes=[1, 2],
237
+ ))
238
+
239
+ with set_current_vllm_config(vllm_config):
240
+ mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
241
+
242
+ with compilation_counter.expect(
243
+ num_graphs_seen=1,
244
+ num_piecewise_graphs_seen=7,
245
+ # 3 attn ops and 4 non-attn ops
246
+ num_piecewise_capturable_graphs_seen=4,
247
+ num_backend_compilations=4,
248
+ num_cudagraph_captured=8,
249
+ # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
250
+ ):
251
+ run_model(vllm_config, mod_A, cudagraph_runtime_mode)
vllm/compilation/decorators.py
+19
-2
@@ -52,6 +52,14 @@ def _should_ignore_torch_compile(cls) -> bool:
52
return getattr(cls, IGNORE_COMPILE_KEY, False)
53
54
55
+@overload
56
+def support_torch_compile(
57
+ *,
58
+ enable_if: Optional[Callable[[VllmConfig], bool]] = None,
59
+) -> Callable[[_T], _T]:
60
+ ...
61
+
62
+
63
@overload
64
def support_torch_compile(
65
*,
@@ -69,6 +77,7 @@ def support_torch_compile(
77
cls: Optional[_T] = None,
78
*,
79
dynamic_arg_dims: Optional[dict[str, Union[int, list[int]]]] = None,
80
+ enable_if: Optional[Callable[[VllmConfig], bool]] = None,
81
) -> Union[Callable[[_T], _T], _T]:
82
"""
83
A decorator to add support for compiling the forward method of a class.
@@ -118,6 +127,11 @@ def support_torch_compile(
127
NOTE: if an argument is `None`, it should always be passed as `None` during
128
the lifetime of the model, otherwise, it cannot be captured as a single
129
computation graph.
130
+
131
+ `enable_if` is a function that takes a `VllmConfig` object as input and
132
+ returns a boolean value indicating whether to compile the model or not.
133
+ This is useful if you want to compile the model only when certain
134
+ conditions are met.
135
"""
136
137
def cls_decorator_helper(cls: _T) -> _T:
@@ -149,7 +163,8 @@ def support_torch_compile(
163
if k not in sig.parameters:
164
raise ValueError(
165
f"Argument {k} not found in the forward method of {cls}")
152
- return _support_torch_compile(cls, inferred_dynamic_arg_dims)
166
+ return _support_torch_compile(cls, inferred_dynamic_arg_dims,
167
+ enable_if)
168
169
if cls is not None:
170
# use `support_torch_compile` as a decorator without arguments
@@ -162,6 +177,7 @@ def support_torch_compile(
177
def _support_torch_compile(
178
cls: _T,
179
dynamic_arg_dims: dict[str, Union[int, list[int]]],
180
+ enable_if: Optional[Callable[[VllmConfig], bool]] = None,
181
) -> _T:
182
"""
183
A decorator to add support for compiling the forward method of a class.
@@ -182,13 +198,14 @@ def _support_torch_compile(
198
def __init__(self, *, vllm_config: VllmConfig, prefix: str = '', **kwargs):
199
old_init(self, vllm_config=vllm_config, prefix=prefix, **kwargs)
200
self.vllm_config = vllm_config
201
+ enable_compile = enable_if is None or enable_if(vllm_config)
202
# for CompilationLevel.DYNAMO_AS_IS , the upper level model runner
203
# will handle the compilation, so we don't need to do anything here.
204
self.do_not_compile = \
205
vllm_config.compilation_config.level in [
206
CompilationLevel.NO_COMPILATION, CompilationLevel.DYNAMO_AS_IS
207
] or not supports_dynamo() or _should_ignore_torch_compile(
191
- self.__class__)
208
+ self.__class__) or not enable_compile
209
if self.do_not_compile:
210
return
211