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, 23:45 UTC
Add test case for compiling multiple graphs (#21044)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Jul 23, 2025 at 11:00 UTC
4ac7713e32f3d372a3177d02ff19f67cadfa2a4b
3 files changed
+390
-1
tests/compile/piecewise/test_multiple_graphs.py
new
+350
@@ -0,0 +1,350 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Test (piecewise) compilation with a simple model where multiple submodules
5
+are compiled and graph captured separately.
6
+"""
7
+import torch
8
+from torch import nn
9
+from torch.library import Library
10
+
11
+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
19
+from vllm.utils import direct_register_custom_op
20
+
21
+# create a library to hold the custom op
22
+silly_lib = Library("silly", "FRAGMENT") # noqa
23
+
24
+BATCH_SIZE = 32
25
+MLP_SIZE = 128
26
+HIDDEN_SIZE = 1024
27
+RANDOM_SEED = 0
28
+
29
+
30
+def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
31
+ out: torch.Tensor) -> None:
32
+ out.copy_(q)
33
+ out += k
34
+ out += v
35
+
36
+
37
+def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
38
+ out: torch.Tensor) -> None:
39
+ return
40
+
41
+
42
+direct_register_custom_op(
43
+ op_name="attention",
44
+ op_func=silly_attention,
45
+ mutates_args=["out"],
46
+ fake_impl=silly_attention_fake,
47
+ target_lib=silly_lib,
48
+)
49
+
50
+
51
+@support_torch_compile
52
+class ParentModel(nn.Module):
53
+
54
+ def __init__(self,
55
+ *,
56
+ vllm_config: VllmConfig,
57
+ prefix: str = '',
58
+ **kwargs) -> None:
59
+ super().__init__()
60
+
61
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
62
+ return x
63
+
64
+
65
+class Attention(nn.Module):
66
+
67
+ def __init__(self, mlp_size: int, hidden_size: int) -> None:
68
+ super().__init__()
69
+ self.pre_attn = nn.Linear(mlp_size, hidden_size, bias=False)
70
+ self.post_attn = nn.Linear(hidden_size, mlp_size, bias=False)
71
+ self.rms_norm_weight = nn.Parameter(torch.ones(hidden_size))
72
+
73
+ # Initialize to same weights for testing
74
+ nn.init.xavier_normal_(
75
+ self.pre_attn.weight.data,
76
+ generator=torch.Generator().manual_seed(RANDOM_SEED),
77
+ gain=0.001)
78
+ nn.init.xavier_normal_(
79
+ self.post_attn.weight.data,
80
+ generator=torch.Generator().manual_seed(RANDOM_SEED),
81
+ gain=0.001)
82
+
83
+ def rms_norm_ref(self, x: torch.Tensor) -> torch.Tensor:
84
+ x_f32 = x.float()
85
+ return (x_f32 * torch.rsqrt(
86
+ torch.mean(x_f32.square(), dim=-1, keepdim=True) + 1e-6) *
87
+ self.rms_norm_weight).to(x.dtype)
88
+
89
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
90
+ x = self.pre_attn(x)
91
+ x = self.rms_norm_ref(x)
92
+ attn_output = torch.empty_like(x)
93
+ torch.ops.silly.attention(x, x, x, attn_output)
94
+ x = attn_output
95
+ x = self.rms_norm_ref(x)
96
+ x = self.post_attn(x)
97
+ return x
98
+
99
+
100
+@support_torch_compile
101
+class CompiledAttention(nn.Module):
102
+
103
+ def __init__(self,
104
+ *,
105
+ mlp_size: int,
106
+ hidden_size: int,
107
+ vllm_config: VllmConfig,
108
+ prefix: str = '',
109
+ **kwargs) -> None:
110
+ super().__init__()
111
+ self.attn = Attention(mlp_size, hidden_size)
112
+
113
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
114
+ return self.attn(x)
115
+
116
+
117
+@support_torch_compile
118
+class CompiledAttentionTwo(CompiledAttention):
119
+
120
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
121
+ return self.attn(x) + x
122
+
123
+
124
+@ignore_torch_compile
125
+class SimpleModelWithTwoGraphs(ParentModel):
126
+
127
+ def __init__(self,
128
+ *,
129
+ mlp_size: int,
130
+ hidden_size: int,
131
+ vllm_config: VllmConfig,
132
+ prefix: str = '',
133
+ **kwargs) -> None:
134
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
135
+ # Test will fail without set_model_tag here with error:
136
+ # "ValueError: too many values to unpack (expected 3)"
137
+ # This is because CompiledAttention and CompiledAttentionTwo
138
+ # have different implmentations but the same torch.compile
139
+ # cache dir will be used as default prefix is 'model_tag'
140
+ with set_model_tag("attn_one"):
141
+ self.attn_one = CompiledAttention(
142
+ mlp_size=mlp_size,
143
+ hidden_size=hidden_size,
144
+ vllm_config=vllm_config,
145
+ prefix=f"{prefix}.attn_one",
146
+ )
147
+ with set_model_tag("attn_two"):
148
+ self.attn_two = CompiledAttentionTwo(
149
+ mlp_size=mlp_size,
150
+ hidden_size=hidden_size,
151
+ vllm_config=vllm_config,
152
+ prefix=f"{prefix}.attn_two",
153
+ )
154
+
155
+ self.hidden_states = torch.zeros((BATCH_SIZE, MLP_SIZE)).cuda()
156
+
157
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
158
+ bsz = x.shape[0]
159
+ # CUDAGraph expects same tensor addresses for each run
160
+ self.hidden_states[:bsz].copy_(x)
161
+ x = self.attn_one(self.hidden_states[:bsz])
162
+ self.hidden_states[:bsz].copy_(x)
163
+ x = self.attn_two(self.hidden_states[:bsz])
164
+ return x
165
+
166
+
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
+
254
+@torch.inference_mode
255
+def run_model(vllm_config, model: nn.Module, inputs: torch.Tensor):
256
+ with set_forward_context({}, vllm_config=vllm_config):
257
+ # First run is for compile
258
+ model(inputs)
259
+
260
+ # Run CUDAGraph captured sizes
261
+ model(inputs[:2])
262
+ model(inputs[:1])
263
+
264
+ output = model(inputs[:2])
265
+
266
+ output = output.cpu()
267
+ return output.cpu()
268
+
269
+
270
+def test_multi_graph_piecewise_compile_outputs_equal():
271
+ outputs = []
272
+
273
+ # piecewise compile
274
+ vllm_config = VllmConfig(compilation_config=CompilationConfig(
275
+ level=CompilationLevel.PIECEWISE,
276
+ use_cudagraph=True,
277
+ splitting_ops=["silly.attention"],
278
+ cudagraph_capture_sizes=[1, 2],
279
+ ))
280
+
281
+ with set_current_vllm_config(vllm_config):
282
+ model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
283
+ hidden_size=HIDDEN_SIZE,
284
+ vllm_config=vllm_config,
285
+ prefix='').eval().cuda()
286
+
287
+ # Pre-allocate memory for CUDAGraph which expects
288
+ # static tensor addresses
289
+ inputs = torch.randn(BATCH_SIZE, MLP_SIZE).cuda()
290
+
291
+ with compilation_counter.expect(
292
+ num_graphs_seen=2, # two graphs for the model
293
+ num_piecewise_graphs_seen=6,
294
+ # attn_one, attn_two each has 3 piecewise graphs
295
+ # (pre attn, post attn, silly_attention) each
296
+ num_piecewise_capturable_graphs_seen=4,
297
+ # attn_one, attn_two has pre attn and post attn each, total=4
298
+ num_backend_compilations=4, # num_piecewise_capturable_graphs_seen
299
+ num_cudagraph_captured=8,
300
+ # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
301
+ ):
302
+ outputs.append(run_model(vllm_config, model, inputs))
303
+
304
+ # no compile or cudagraph
305
+ vllm_config = VllmConfig(compilation_config=CompilationConfig(
306
+ level=CompilationLevel.NO_COMPILATION, ))
307
+
308
+ with set_current_vllm_config(vllm_config):
309
+ model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
310
+ hidden_size=HIDDEN_SIZE,
311
+ vllm_config=vllm_config,
312
+ prefix='').eval().cuda()
313
+
314
+ with compilation_counter.expect(
315
+ num_graphs_seen=0,
316
+ num_piecewise_graphs_seen=0,
317
+ num_piecewise_capturable_graphs_seen=0,
318
+ num_backend_compilations=0,
319
+ num_cudagraph_captured=0,
320
+ ):
321
+ outputs.append(run_model(vllm_config, model, inputs))
322
+
323
+ # piecewise compile without CUDA graph
324
+ vllm_config = VllmConfig(compilation_config=CompilationConfig(
325
+ level=CompilationLevel.PIECEWISE,
326
+ use_cudagraph=False,
327
+ splitting_ops=["silly.attention"],
328
+ ))
329
+
330
+ with set_current_vllm_config(vllm_config):
331
+ model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
332
+ hidden_size=HIDDEN_SIZE,
333
+ vllm_config=vllm_config,
334
+ prefix='').eval().cuda()
335
+
336
+ with compilation_counter.expect(
337
+ num_graphs_seen=2,
338
+ num_piecewise_graphs_seen=6,
339
+ num_piecewise_capturable_graphs_seen=4,
340
+ num_backend_compilations=4,
341
+ num_cudagraph_captured=0, # no cudagraph captured
342
+ ):
343
+ outputs.append(run_model(vllm_config, model, inputs))
344
+
345
+ # Generally don't expect outputs with and without inductor
346
+ # to be bitwise equivalent
347
+ assert torch.allclose(outputs[0], outputs[1])
348
+
349
+ # Expect bitwise equivalence using inductor w/ and w/o cudagraph
350
+ assert torch.equal(outputs[0], outputs[2])
vllm/compilation/compiler_interface.py
+6
@@ -423,6 +423,12 @@ class InductorAdaptor(CompilerInterface):
423
if is_torch_equal_or_newer("2.6"):
424
stack.enter_context(
425
torch._inductor.config.patch(fx_graph_remote_cache=False))
426
+ # InductorAdaptor (unfortunately) requires AOTAutogradCache
427
+ # to be turned off to run. It will fail to acquire the hash_str
428
+ # and error if not.
429
+ # StandaloneInductorAdaptor (PyTorch 2.8+) fixes this problem.
430
+ stack.enter_context(
431
+ torch._functorch.config.patch(enable_autograd_cache=False))
432
stack.enter_context(
433
torch._functorch.config.patch(
434
enable_remote_autograd_cache=False))
vllm/compilation/decorators.py
+34
-1
@@ -20,9 +20,38 @@ from .monitor import start_monitoring_torch_compile
20
21
logger = init_logger(__name__)
22
23
+IGNORE_COMPILE_KEY = "_ignore_compile_vllm"
24
+
25
_T = TypeVar("_T", bound=type[nn.Module])
26
27
28
+def ignore_torch_compile(cls: _T) -> _T:
29
+ """
30
+ A decorator to ignore support_torch_compile decorator
31
+ on the class. This is useful when a parent class has
32
+ a support_torch_compile decorator, but we don't want to
33
+ compile the class `cls` that inherits the parent class.
34
+ This only ignores compiling the forward of the class the
35
+ decorator is applied to.
36
+
37
+ If the parent has ignore_torch_compile but the child has
38
+ support_torch_compile, the child will still be compiled.
39
+
40
+ If the class has one or more submodules
41
+ that have support_torch_compile decorator applied, compile will
42
+ not be ignored for those submodules.
43
+ """
44
+ setattr(cls, IGNORE_COMPILE_KEY, True)
45
+ return cls
46
+
47
+
48
+def _should_ignore_torch_compile(cls) -> bool:
49
+ """
50
+ Check if the class should be ignored for torch.compile.
51
+ """
52
+ return getattr(cls, IGNORE_COMPILE_KEY, False)
53
+
54
+
55
@overload
56
def support_torch_compile(
57
*,
@@ -148,6 +177,8 @@ def _support_torch_compile(
177
178
old_init = cls.__init__
179
180
+ setattr(cls, IGNORE_COMPILE_KEY, False)
181
+
182
def __init__(self, *, vllm_config: VllmConfig, prefix: str = '', **kwargs):
183
old_init(self, vllm_config=vllm_config, prefix=prefix, **kwargs)
184
self.vllm_config = vllm_config
@@ -156,9 +187,11 @@ def _support_torch_compile(
187
self.do_not_compile = \
188
vllm_config.compilation_config.level in [
189
CompilationLevel.NO_COMPILATION, CompilationLevel.DYNAMO_AS_IS
159
- ] or not supports_dynamo()
190
+ ] or not supports_dynamo() or _should_ignore_torch_compile(
191
+ self.__class__)
192
if self.do_not_compile:
193
return
194
+
195
compilation_counter.num_models_seen += 1
196
TorchCompileWrapperWithCustomDispatcher.__init__(
197
self, compilation_level=vllm_config.compilation_config.level)