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 13, 04:30 UTC
Avoid bytecode hook and simplify TorchCompileWrapperWithCustomDipatch (#25110)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Nov 14, 2025 at 14:11 UTC
2e0ad629b0422358d424e1fcfddeb22d102936e8
10 files changed
+422
-236
tests/compile/piecewise/test_multiple_graphs.py
+10
-1
@@ -22,6 +22,8 @@ from vllm.config import (
22
from vllm.forward_context import BatchDescriptor, set_forward_context
23
from vllm.utils.torch_utils import is_torch_equal_or_newer
24
25
+from ...utils import create_new_process_for_each_test
26
+
27
# This import automatically registers `torch.ops.silly.attention`
28
from .. import silly_attention # noqa: F401
29
@@ -193,7 +195,14 @@ def run_model(
195
196
197
@pytest.mark.parametrize("use_inductor_graph_partition", [False, True])
196
-def test_multi_graph_piecewise_compile(use_inductor_graph_partition: bool):
198
+@pytest.mark.parametrize("use_bytecode_hook", [True, False])
199
+@create_new_process_for_each_test("spawn")
200
+def test_multi_graph_piecewise_compile(
201
+ use_inductor_graph_partition: bool, use_bytecode_hook: bool, monkeypatch
202
+):
203
+ # Set the environment variable for this test
204
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
205
+
206
if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
207
pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
208
tests/compile/piecewise/test_simple.py
+3
@@ -21,6 +21,8 @@ from vllm.config import (
21
from vllm.forward_context import BatchDescriptor, set_forward_context
22
from vllm.utils.torch_utils import is_torch_equal_or_newer
23
24
+from ...utils import create_new_process_for_each_test
25
+
26
# This import automatically registers `torch.ops.silly.attention`
27
from ..silly_attention import get_global_counter, reset_global_counter
28
@@ -124,6 +126,7 @@ def _run_simple_model(
126
127
@pytest.mark.parametrize("use_inductor", [True, False])
128
@torch.inference_mode()
129
+@create_new_process_for_each_test("spawn")
130
def test_simple_piecewise_compile(use_inductor):
131
_run_simple_model(
132
splitting_ops=["silly::attention"],
tests/compile/piecewise/test_toy_llama.py
+8
-1
@@ -29,6 +29,8 @@ from vllm.config import (
29
from vllm.forward_context import BatchDescriptor, set_forward_context
30
from vllm.utils.torch_utils import is_torch_equal_or_newer
31
32
+from ...utils import create_new_process_for_each_test
33
+
34
# This import automatically registers `torch.ops.silly.attention`
35
from .. import silly_attention # noqa: F401
36
@@ -334,6 +336,7 @@ def run_model(llama_config, compile_config: CompilationConfig) -> torch.Tensor:
336
("inductor", True), # Inductor, Inductor partition
337
],
338
)
339
+@create_new_process_for_each_test("spawn")
340
def test_toy_llama(
341
backend: str, use_inductor_graph_partition: bool, monkeypatch, tmp_path
342
):
@@ -513,4 +516,8 @@ def benchmark():
516
517
518
if __name__ == "__main__":
516
- benchmark()
519
+ # Protect against subprocess reimport when using spawn_new_process_for_each_test
520
+ import os
521
+
522
+ if os.environ.get("RUNNING_IN_SUBPROCESS") != "1":
523
+ benchmark()
tests/compile/test_wrapper.py
+115
-40
@@ -2,59 +2,134 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
+import os
6
+
7
+import pytest
8
import torch
9
7
-from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
8
-from vllm.config import CompilationMode
10
+from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
11
+from vllm.config import (
12
+ CompilationConfig,
13
+ CompilationMode,
14
+ VllmConfig,
15
+ set_current_vllm_config,
16
+)
17
18
19
class MyMod(torch.nn.Module):
20
def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None):
13
- if cache is not None:
14
- return x + cache
15
- return x * 2
21
+ if x.size()[0] >= 4:
22
+ return x * 2
23
+ else:
24
+ return x * 100
25
26
18
-class MyWrapper(TorchCompileWrapperWithCustomDispatcher):
27
+class MyWrapper(TorchCompileWithNoGuardsWrapper):
28
def __init__(self, model):
29
self.model = model
21
- compiled_callable = torch.compile(self.forward, backend="eager")
22
- super().__init__(
23
- compiled_callable, compilation_mode=CompilationMode.DYNAMO_TRACE_ONCE
24
- )
30
+ super().__init__()
31
26
- def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None):
32
+ def forward(self, x: torch.Tensor): # type: ignore[override]
33
# this is the function to be compiled
28
- return self.model(x, cache)
29
-
30
- def __call__(self, x: torch.Tensor, cache: torch.Tensor | None = None):
31
- # let torch.compile compile twice
32
- if len(self.compiled_codes) == 2:
33
- dispatch_id = 0 if cache is None else 1
34
- with self.dispatch_to_code(dispatch_id):
35
- return self.forward(x, cache)
36
- else:
37
- return self.compiled_callable(x, cache)
34
+ return self.model(x)
35
+
36
37
+@pytest.mark.parametrize("use_bytecode_hook", [True, False])
38
+def test_torch_compile_wrapper(use_bytecode_hook, monkeypatch):
39
+ """Test basic functionality of TorchCompileWithNoGuardsWrapper."""
40
+ # Set the environment variable for this test
41
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
42
40
-def test_torch_compile_wrapper():
41
- mod = MyMod()
42
- wrappers = []
43
- for i in range(3):
43
+ # Create a proper vLLM config instead of mocking
44
+ vllm_config = VllmConfig()
45
+ vllm_config.compilation_config = CompilationConfig()
46
+ vllm_config.compilation_config.mode = CompilationMode.DYNAMO_TRACE_ONCE
47
+ vllm_config.compilation_config.backend = "inductor"
48
+
49
+ # Test DYNAMO_TRACE_ONCE
50
+ with set_current_vllm_config(vllm_config):
51
torch._dynamo.reset()
52
+ mod = MyMod()
53
+ wrapper = MyWrapper(mod)
54
+
55
+ # First call should trigger compilation
56
+ x = torch.tensor([1, 2, 3, 4])
57
+ torch._dynamo.mark_dynamic(x, 0)
58
+
59
+ result1 = wrapper(x)
60
+ expected1 = torch.tensor([2, 4, 6, 8])
61
+ assert torch.allclose(result1, expected1), (
62
+ f"Expected {expected1}, got {result1}"
63
+ )
64
+
65
+ # Second call should use compiled code
66
+ x2 = torch.tensor([1, 2, 3])
67
+ result2 = wrapper(x2)
68
+ expected2 = torch.tensor([2, 4, 6])
69
+ assert torch.allclose(result2, expected2), (
70
+ f"Expected {expected2}, got {result2}"
71
+ )
72
+
73
+ # without the wrapper result would be different.
74
+ result3 = mod(x2)
75
+ expected3 = torch.tensor([100, 200, 300])
76
+
77
+ assert torch.allclose(result3, expected3), (
78
+ f"Expected {result3}, got {expected3}"
79
+ )
80
+
81
+ # with STOCK_TORCH_COMPILE we do not remove guards.
82
+ vllm_config.compilation_config.mode = CompilationMode.STOCK_TORCH_COMPILE
83
+ torch._dynamo.reset()
84
+ with set_current_vllm_config(vllm_config):
85
+ mod = MyMod()
86
wrapper = MyWrapper(mod)
46
- wrappers.append(wrapper)
47
- x = torch.tensor([1])
48
- wrapper(x, None) # profile run, compile
49
- # create a cache tensor
50
- cache = torch.tensor([2])
51
- wrapper(x, cache) # warm up with cache, recompile
52
-
53
- # for new input, dispatch to the compiled code directly
54
- new_x = torch.tensor([3])
55
- assert wrapper(new_x, None).item() == 6 # dispatch to the first compiled code
56
- assert wrapper(new_x, cache).item() == 5 # dispatch to the second compiled code
57
-
58
- for wrapper in wrappers:
59
- # make sure they have independent compiled codes
60
- assert len(wrapper.compiled_codes) == 2
87
+
88
+ # First call should trigger compilation
89
+ x = torch.tensor([1, 2, 3, 4])
90
+ torch._dynamo.mark_dynamic(x, 0)
91
+
92
+ result1 = wrapper(x)
93
+ expected1 = torch.tensor([2, 4, 6, 8])
94
+ assert torch.allclose(result1, expected1), (
95
+ f"Expected {expected1}, got {result1}"
96
+ )
97
+
98
+ # Second call should triger another compilation
99
+ x2 = torch.tensor([1, 2, 3])
100
+ result2 = wrapper(x2)
101
+ expected2 = torch.tensor([100, 200, 300])
102
+ assert torch.allclose(result2, expected2), (
103
+ f"Expected {expected2}, got {result2}"
104
+ )
105
+
106
+ # NO_COMPILATION level not supported.
107
+ vllm_config.compilation_config.mode = None
108
+ torch._dynamo.reset()
109
+ with set_current_vllm_config(vllm_config):
110
+ torch._dynamo.reset()
111
+ mod = MyMod()
112
+
113
+ try:
114
+ wrapper = MyWrapper(mod)
115
+ except Exception:
116
+ return
117
+ raise AssertionError("expected an exception to be raised")
118
+
119
+
120
+if __name__ == "__main__":
121
+ # Run with both parameter values
122
+
123
+ class MockMonkeypatch:
124
+ def setenv(self, name, value):
125
+ os.environ[name] = value
126
+
127
+ mp = MockMonkeypatch()
128
+
129
+ print("Testing with VLLM_USE_BYTECODE_HOOK=False")
130
+ test_torch_compile_wrapper(False, mp)
131
+
132
+ print("Testing with VLLM_USE_BYTECODE_HOOK=True")
133
+ test_torch_compile_wrapper(True, mp)
134
+
135
+ print("All tests passed!")
tests/models/multimodal/generation/test_qwen2_5_vl.py
+10
@@ -34,6 +34,7 @@ VIDEO_PROMPTS = VIDEO_ASSETS.prompts(
34
@pytest.mark.parametrize("num_frames", [16])
35
@pytest.mark.parametrize("dtype", [target_dtype])
36
@pytest.mark.parametrize("max_tokens", [128])
37
+@pytest.mark.parametrize("use_bytecode_hook", [True, False])
38
def test_qwen2_5_vl_evs_functionality(
39
vllm_runner,
40
video_assets,
@@ -42,10 +43,14 @@ def test_qwen2_5_vl_evs_functionality(
43
num_frames: int,
44
dtype: str,
45
max_tokens: int,
46
+ use_bytecode_hook: bool,
47
+ monkeypatch,
48
) -> None:
49
"""Test EVS (Efficient Video Sampling) functionality with different
50
pruning rates.
51
"""
52
+ # Set the environment variable for this test
53
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
54
55
# Sample frames from video assets
56
sampled_vids = [
@@ -86,6 +91,7 @@ def test_qwen2_5_vl_evs_functionality(
91
@pytest.mark.parametrize("num_frames", [16])
92
@pytest.mark.parametrize("dtype", [target_dtype])
93
@pytest.mark.parametrize("max_tokens", [128])
94
+@pytest.mark.parametrize("use_bytecode_hook", [True, False])
95
def test_qwen2_5_vl_evs_batched_videos(
96
vllm_runner,
97
video_assets,
@@ -94,6 +100,8 @@ def test_qwen2_5_vl_evs_batched_videos(
100
num_frames: int,
101
dtype: str,
102
max_tokens: int,
103
+ use_bytecode_hook: bool,
104
+ monkeypatch,
105
) -> None:
106
"""Test EVS functionality with batched videos.
107
@@ -102,6 +110,8 @@ def test_qwen2_5_vl_evs_batched_videos(
110
2. Both pruning configurations work with multiple videos
111
3. The model doesn't crash when processing multiple videos simultaneously
112
"""
113
+ # Set the environment variable for this test
114
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1" if use_bytecode_hook else "0")
115
# Sample frames from video assets
116
sampled_vids = [
117
sample_frames_from_video(asset.np_ndarrays, num_frames)
tests/v1/e2e/test_spec_decode.py
+8
@@ -75,6 +75,14 @@ def model_name():
75
return "meta-llama/Llama-3.1-8B-Instruct"
76
77
78
+@pytest.fixture(autouse=True)
79
+def reset_torch_dynamo():
80
+ """Reset torch dynamo cache before each test"""
81
+ yield
82
+ # Cleanup after test
83
+ torch._dynamo.reset()
84
+
85
+
86
@pytest.mark.parametrize(
87
"speculative_config",
88
[
vllm/compilation/decorators.py
+115
-119
@@ -17,7 +17,7 @@ from torch._dynamo.symbolic_convert import InliningInstructionTranslator
17
18
import vllm.envs as envs
19
from vllm.compilation.counter import compilation_counter
20
-from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
20
+from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
21
from vllm.config import (
22
CompilationMode,
23
VllmConfig,
@@ -246,14 +246,14 @@ def _support_torch_compile(
246
"""
247
A decorator to add support for compiling the forward method of a class.
248
"""
249
- if TorchCompileWrapperWithCustomDispatcher in cls.__bases__:
249
+ if TorchCompileWithNoGuardsWrapper in cls.__bases__:
250
# support decorating multiple times
251
return cls
252
253
# take care of method resolution order
254
# make sure super().__init__ is called on the base class
255
- # other than TorchCompileWrapperWithCustomDispatcher
256
- cls.__bases__ = cls.__bases__ + (TorchCompileWrapperWithCustomDispatcher,)
255
+ # other than TorchCompileWithNoGuardsWrapper
256
+ cls.__bases__ = cls.__bases__ + (TorchCompileWithNoGuardsWrapper,)
257
258
old_init = cls.__init__
259
@@ -290,12 +290,43 @@ def _support_torch_compile(
290
return
291
292
compilation_counter.num_models_seen += 1
293
- TorchCompileWrapperWithCustomDispatcher.__init__(
294
- self, compilation_mode=vllm_config.compilation_config.mode
295
- )
293
+ self.compiled = False
294
+ TorchCompileWithNoGuardsWrapper.__init__(self)
295
296
cls.__init__ = __init__
297
298
+ def _mark_dynamic_inputs(mod, *args, **kwargs):
299
+ sig = inspect.signature(mod.__class__.forward)
300
+ bound_args = sig.bind(mod, *args, **kwargs)
301
+ bound_args.apply_defaults()
302
+ for k, dims in dynamic_arg_dims.items():
303
+ arg = bound_args.arguments.get(k)
304
+ if arg is not None:
305
+ dims = [dims] if isinstance(dims, int) else dims
306
+ if isinstance(arg, torch.Tensor):
307
+ # In case dims is specified with negative indexing
308
+ dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
309
+ torch._dynamo.mark_dynamic(arg, dims)
310
+ elif isinstance(arg, IntermediateTensors):
311
+ for tensor in arg.tensors.values():
312
+ # In case dims is specified with negative indexing
313
+ dims = [tensor.ndim + dim if dim < 0 else dim for dim in dims]
314
+ torch._dynamo.mark_dynamic(tensor, dims)
315
+ else:
316
+ raise ValueError(
317
+ "Unsupported dynamic dimensions"
318
+ f" {dims} for argument {k} with type {type(arg)}."
319
+ )
320
+ if mark_unbacked_dims:
321
+ for k, dims in mark_unbacked_dims.items():
322
+ arg = bound_args.arguments.get(k)
323
+ if arg is not None:
324
+ dims = [dims] if isinstance(dims, int) else dims
325
+ if isinstance(arg, torch.Tensor):
326
+ # In case dims is specified with negative indexing
327
+ dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
328
+ torch._dynamo.decorators.mark_unbacked(arg, dims)
329
+
330
def __call__(self, *args, **kwargs):
331
# torch.compiler.is_compiling() means we are inside the compilation
332
# e.g. TPU has the compilation logic in model runner, so we don't
@@ -303,6 +334,7 @@ def _support_torch_compile(
334
if self.do_not_compile or torch.compiler.is_compiling():
335
return self.forward(*args, **kwargs)
336
337
+ # if aot_compiled_fn is set, just call it.
338
if getattr(self, "aot_compiled_fn", None) is not None:
339
return self.aot_compiled_fn(self, *args, **kwargs)
340
@@ -362,120 +394,84 @@ def _support_torch_compile(
394
)
395
return self.aot_compiled_fn(self, *args, **kwargs)
396
397
+ if self.compiled:
398
+ assert not envs.VLLM_USE_AOT_COMPILE
399
+ return TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)
400
+
401
+ # This is the path for the first compilation.
402
+
403
# the first compilation needs to have dynamic shapes marked
366
- if len(self.compiled_codes) < 1:
367
- sig = inspect.signature(self.__class__.forward)
368
- bound_args = sig.bind(self, *args, **kwargs)
369
- bound_args.apply_defaults()
370
- for k, dims in dynamic_arg_dims.items():
371
- arg = bound_args.arguments.get(k)
372
- if arg is not None:
373
- dims = [dims] if isinstance(dims, int) else dims
374
- if isinstance(arg, torch.Tensor):
375
- # In case dims is specified with negative indexing
376
- dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
377
- torch._dynamo.mark_dynamic(arg, dims)
378
- elif isinstance(arg, IntermediateTensors):
379
- for tensor in arg.tensors.values():
380
- # In case dims is specified with negative indexing
381
- dims = [
382
- tensor.ndim + dim if dim < 0 else dim for dim in dims
383
- ]
384
- torch._dynamo.mark_dynamic(tensor, dims)
385
- else:
386
- raise ValueError(
387
- "Unsupported dynamic dimensions"
388
- f" {dims} for argument {k} with type {type(arg)}."
389
- )
390
- if mark_unbacked_dims:
391
- for k, dims in mark_unbacked_dims.items():
392
- arg = bound_args.arguments.get(k)
393
- if arg is not None:
394
- dims = [dims] if isinstance(dims, int) else dims
395
- if isinstance(arg, torch.Tensor):
396
- # In case dims is specified with negative indexing
397
- dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
398
- torch._dynamo.decorators.mark_unbacked(arg, dims)
399
- # here, it is the starting point of the `torch.compile` process
400
- start_monitoring_torch_compile(self.vllm_config)
401
- logger.debug("Start compiling function %s", self.original_code_object)
402
-
403
- # if we don't use custom dispatcher, we can directly call the
404
- # compiled function and let torch.compile handle the dispatching,
405
- # with the overhead of guard evaluation and recompilation.
406
- if len(self.compiled_codes) < 1 or not self.use_custom_dispatcher:
407
- # it seems Dynamo reuse the compilation across instances,
408
- # while we need to make sure the compiled code is not reused.
409
- # we need to control all the compilation of the model.
410
- torch._dynamo.eval_frame.remove_from_cache(self.original_code_object)
411
-
412
- # collect all relevant files traced by Dynamo,
413
- # so that the compilation cache can trigger re-compilation
414
- # properly when any of these files change.
415
-
416
- # 1. the file containing the top-level forward function
417
- self.vllm_config.compilation_config.traced_files.add(
418
- self.original_code_object.co_filename
419
- )
404
+ _mark_dynamic_inputs(self, *args, **kwargs)
405
421
- # 2. every time Dynamo sees a function call, it will inline
422
- # the function by calling InliningInstructionTranslator.inline_call_
423
- # we hijack this function to know all the functions called
424
- # during Dynamo tracing, and their corresponding files
425
- inline_call = InliningInstructionTranslator.inline_call_
426
-
427
- def patched_inline_call(self_):
428
- code = self_.f_code
429
- self.vllm_config.compilation_config.traced_files.add(code.co_filename)
430
- return inline_call(self_)
431
-
432
- # Disable the C++ compilation of symbolic shape guards. C++-fication
433
- # of symbolic shape guards can improve guard overhead. But, since
434
- # vllm skip guards anyways, setting this flag to False can improve
435
- # compile time.
436
- dynamo_config_patches = {}
437
- try:
438
- _ = torch._dynamo.config.enable_cpp_symbolic_shape_guards
439
- dynamo_config_patches["enable_cpp_symbolic_shape_guards"] = False
440
- except AttributeError:
441
- # Note: this config is not available in torch 2.6, we can skip
442
- # if the config doesn't exist
443
- logger.debug("enable_cpp_symbolic_shape_guards config not available")
444
-
445
- with (
446
- patch.object(
447
- InliningInstructionTranslator, "inline_call_", patched_inline_call
448
- ),
449
- torch._dynamo.config.patch(**dynamo_config_patches),
450
- maybe_use_cudagraph_partition_wrapper(self.vllm_config),
451
- _torch27_patch_tensor_subclasses(),
452
- ):
453
- if envs.VLLM_USE_AOT_COMPILE:
454
- self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
455
- output = self.aot_compiled_fn(self, *args, **kwargs)
456
- assert aot_compilation_path is not None
457
- assert cache_dir is not None
458
- try:
459
- os.makedirs(cache_dir, exist_ok=True)
460
- self.aot_compiled_fn.save_compiled_function(
461
- aot_compilation_path
462
- )
463
- except Exception as e:
464
- logger.warning(
465
- "Cannot save aot compilation to path %s, error: %s",
466
- aot_compilation_path,
467
- str(e),
468
- )
469
- else:
470
- output = self.compiled_callable(*args, **kwargs)
471
- return output
472
-
473
- # usually, capturing the model once is enough, and then we can
474
- # dispatch to the compiled code directly, without going through
475
- # the Dynamo guard mechanism.
476
- with self.dispatch_to_code(0):
477
- model_output = self.forward(*args, **kwargs)
478
- return model_output
406
+ # here, it is the starting point of the `torch.compile` process
407
+ start_monitoring_torch_compile(self.vllm_config)
408
+ original_code_object = self.original_code_object()
409
+ logger.debug("Start compiling function %s", original_code_object)
410
+
411
+ # we do not want tp delete the original code object entries since
412
+ # we depend on them now to look up cached compiled functions.
413
+ # torch._dynamo.eval_frame.remove_from_cache(original_code_object)
414
+
415
+ # collect all relevant files traced by Dynamo,
416
+ # so that the compilation cache can trigger re-compilation
417
+ # properly when any of these files change.
418
+
419
+ # 1. the file containing the top-level forward function
420
+ self.vllm_config.compilation_config.traced_files.add(
421
+ original_code_object.co_filename
422
+ )
423
+
424
+ # 2. every time Dynamo sees a function call, it will inline
425
+ # the function by calling InliningInstructionTranslator.inline_call_
426
+ # we hijack this function to know all the functions called
427
+ # during Dynamo tracing, and their corresponding files
428
+ inline_call = InliningInstructionTranslator.inline_call_
429
+
430
+ def patched_inline_call(self_):
431
+ code = self_.f_code
432
+ self.vllm_config.compilation_config.traced_files.add(code.co_filename)
433
+ return inline_call(self_)
434
+
435
+ # Disable the C++ compilation of symbolic shape guards. C++-fication
436
+ # of symbolic shape guards can improve guard overhead. But, since
437
+ # vllm skip guards anyways, setting this flag to False can improve
438
+ # compile time.
439
+ dynamo_config_patches = {}
440
+ try:
441
+ _ = torch._dynamo.config.enable_cpp_symbolic_shape_guards
442
+ dynamo_config_patches["enable_cpp_symbolic_shape_guards"] = False
443
+ except AttributeError:
444
+ # Note: this config is not available in torch 2.6, we can skip
445
+ # if the config doesn't exist
446
+ logger.debug("enable_cpp_symbolic_shape_guards config not available")
447
+
448
+ with (
449
+ patch.object(
450
+ InliningInstructionTranslator, "inline_call_", patched_inline_call
451
+ ),
452
+ torch._dynamo.config.patch(**dynamo_config_patches),
453
+ maybe_use_cudagraph_partition_wrapper(self.vllm_config),
454
+ _torch27_patch_tensor_subclasses(),
455
+ ):
456
+ if envs.VLLM_USE_AOT_COMPILE:
457
+ self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
458
+ output = self.aot_compiled_fn(self, *args, **kwargs)
459
+ assert aot_compilation_path is not None
460
+ assert cache_dir is not None
461
+ try:
462
+ os.makedirs(cache_dir, exist_ok=True)
463
+ self.aot_compiled_fn.save_compiled_function(aot_compilation_path)
464
+ except Exception as e:
465
+ logger.warning(
466
+ "Cannot save aot compilation to path %s, error: %s",
467
+ aot_compilation_path,
468
+ str(e),
469
+ )
470
+ else:
471
+ output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)
472
+
473
+ self.compiled = True
474
+ return output
475
476
cls.__call__ = __call__
477
return cls
vllm/compilation/wrapper.py
+141
-71
@@ -4,11 +4,11 @@
4
import os
5
import sys
6
from abc import abstractmethod
7
-from collections.abc import Callable
7
from contextlib import contextmanager
8
from types import CodeType
9
10
import torch
11
+import torch._C._dynamo.guards
12
13
import vllm.envs as envs
14
from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
@@ -17,88 +17,153 @@ from vllm.logger import init_logger
17
logger = init_logger(__name__)
18
19
20
-class TorchCompileWrapperWithCustomDispatcher:
20
+def _noop_add_global_state_guard(self, *args, **kwargs):
21
+ """No-op to skip the GLOBAL_STATE guard entirely"""
22
+ pass
23
+
24
+
25
+def _noop_add_torch_function_mode_stack_guard(self, *args, **kwargs):
26
+ """No-op to skip the TORCH_FUNCTION_MODE_STACK guard entirely"""
27
+ pass
28
+
29
+
30
+@contextmanager
31
+def _compilation_context():
32
+ """Context manager for compilation settings and patches.
33
+
34
+ This manager:
35
+ 1. Sets higher dynamo cache limits for compilation. (Needed for
36
+ qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
37
+ Generally a recompilation can happen whenever we use a new
38
+ backend instance in torch.compile.
39
+ 2. Patches out add_global_state_guard to skip GLOBAL_STATE guards
40
+ 3. Patches out add_torch_function_mode_stack_guard to skip
41
+ TORCH_FUNCTION_MODE_STACK guards.
42
+ 4. Restores everything when compilation completes
43
"""
22
- A wrapper class for torch.compile, with a custom dispatch logic.
23
- Subclasses should:
24
- 1. Implement the forward method
25
- 2. Implement the dispatch logic in the __call__ method
26
- It can use `self.compiled_codes` to access the compiled bytecode,
27
- and `with self.dispatch_to_code(index):` to dispatch to
28
- the compiled code.
29
- 3. Implement the `__init__` method to determine how to call
30
- `torch.compile` over the forward method.
44
+ # Save original values
45
+ original_global_state_guard = (
46
+ torch._C._dynamo.guards.GuardManager.add_global_state_guard
47
+ )
48
+ original_torch_function_mode_stack_guard = (
49
+ torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard
50
+ )
51
+ original_cache_size = torch._dynamo.config.cache_size_limit
52
+ original_accumulated_cache = torch._dynamo.config.accumulated_cache_size_limit
53
+
54
+ try:
55
+ # Set higher cache limits for compilation
56
+ torch._dynamo.config.cache_size_limit = 2048
57
+ torch._dynamo.config.accumulated_cache_size_limit = 8192
58
+
59
+ # Patch guard manager
60
+ torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
61
+ _noop_add_global_state_guard
62
+ )
63
+ torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
64
+ _noop_add_torch_function_mode_stack_guard
65
+ )
66
+ yield
67
+ finally:
68
+ # Restore original values
69
+ torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
70
+ original_global_state_guard
71
+ )
72
+ torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
73
+ original_torch_function_mode_stack_guard
74
+ )
75
+ torch._dynamo.config.cache_size_limit = original_cache_size
76
+ torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache
77
+
78
+
79
+class TorchCompileWithNoGuardsWrapper:
80
"""
81
+ A wrapper class for torch.compile, it ensures that all guards are dropped
82
+ when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE.
83
+ When guards are dropped, the first time __call__ is invoked, a single
84
+ compilation is triggered. Dynamo should never be traced again after that
85
+ since we drop all guards.
86
+ """
87
+
88
+ def __init__(self):
89
+ self.compiled = False
90
33
- def __init__(
34
- self,
35
- compiled_callable: Callable | None = None,
36
- compilation_mode: CompilationMode = CompilationMode.NONE,
37
- ):
91
vllm_config = get_current_vllm_config()
92
self.vllm_config = vllm_config
40
- if compiled_callable is None:
41
- # default compilation settings
42
- # compiling the forward method
43
-
44
- backend = vllm_config.compilation_config.init_backend(vllm_config)
45
- options = None
46
- if isinstance(backend, str) and backend == "inductor":
47
- options = (
48
- get_current_vllm_config().compilation_config.inductor_compile_config
49
- )
50
- if envs.VLLM_USE_AOT_COMPILE:
51
- options = options or {}
52
- # This effectively drop all the guards.
53
- # We need this because bytecode hook is not used any more to
54
- # drop guards in the AOT compile mode.
55
- options["guard_filter_fn"] = lambda guards: [False for _ in guards]
56
- if hasattr(torch._dynamo.config, "enable_aot_compile"):
57
- torch._dynamo.config.enable_aot_compile = True
58
- else:
59
- msg = "torch._dynamo.config.enable_aot_compile is not "
60
- msg += "available. AOT compile is disabled and please "
61
- msg += "upgrade PyTorch version to use AOT compile."
62
- logger.warning(msg)
63
-
64
- compiled_callable = torch.compile(
65
- self.forward, fullgraph=True, backend=backend, options=options
66
- )
67
-
68
- self.compiled_callable = compiled_callable
69
- self.original_code_object = self.__class__.forward.__code__
70
- self.compiled_codes: list[CodeType] = []
71
- torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)
72
-
73
- # read the env var to determine whether to use the custom dispatcher
74
- # subclasses can use this to switch between the custom dispatcher
75
- # and the default Dynamo guard mechanism.
76
- self.use_custom_dispatcher: bool = (
77
- compilation_mode >= CompilationMode.DYNAMO_TRACE_ONCE
93
+ mode = vllm_config.compilation_config.mode
94
+ if mode is None:
95
+ raise RuntimeError("Compilation mode cannot be NO_COMPILATION")
96
+
97
+ backend = vllm_config.compilation_config.init_backend(vllm_config)
98
+ options = {}
99
+
100
+ if isinstance(backend, str) and backend == "inductor":
101
+ options = vllm_config.compilation_config.inductor_compile_config
102
+
103
+ if mode != CompilationMode.STOCK_TORCH_COMPILE:
104
+ # Drop all the guards.
105
+ options["guard_filter_fn"] = lambda x: [False for _ in x]
106
+
107
+ if envs.VLLM_USE_AOT_COMPILE:
108
+ if hasattr(torch._dynamo.config, "enable_aot_compile"):
109
+ torch._dynamo.config.enable_aot_compile = True
110
+ else:
111
+ msg = "torch._dynamo.config.enable_aot_compile is not "
112
+ msg += "available. AOT compile is disabled and please "
113
+ msg += "upgrade PyTorch version to use AOT compile."
114
+ logger.warning(msg)
115
+
116
+ self._compiled_callable = torch.compile(
117
+ self.forward,
118
+ fullgraph=True,
119
+ dynamic=False,
120
+ backend=backend,
121
+ options=options,
122
)
123
124
+ if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE:
125
+ torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)
126
+ self._compiled_bytecode = None
127
+
128
def aot_compile(self, *args, **kwargs):
81
- if not hasattr(self.compiled_callable, "aot_compile"):
129
+ if not hasattr(self._compiled_callable, "aot_compile"):
130
raise RuntimeError(
131
"aot_compile is not supported by the current configuration. "
132
+ "Please make sure torch.compile is enabled with the latest "
133
+ f"version of PyTorch (current using torch: {torch.__version__})"
134
)
87
- return self.compiled_callable.aot_compile((args, kwargs))
135
+ return self._compiled_callable.aot_compile((args, kwargs))
136
137
def __call__(self, *args, **kwargs):
90
- """Implement the dispatch logic here, beyond the torch.compile mode.
91
- NOTE: this function can have additional arguments beyond the forward
92
- method, for directly dispatching to the compiled code.
93
- """
94
- return self.compiled_callable(*args, **kwargs)
138
+ if envs.VLLM_USE_BYTECODE_HOOK:
139
+ if (
140
+ self.vllm_config.compilation_config.mode
141
+ == CompilationMode.STOCK_TORCH_COMPILE
142
+ ):
143
+ return self._compiled_callable(*args, **kwargs)
144
+
145
+ if not self._compiled_bytecode:
146
+ # Make sure a compilation is triggered by clearing dynamo
147
+ # cache.
148
+ torch._dynamo.eval_frame.remove_from_cache(self.original_code_object())
149
+ return self._compiled_callable(*args, **kwargs)
150
+ else:
151
+ with self._dispatch_to_compiled_code():
152
+ return self.forward(*args, **kwargs)
153
+ else:
154
+ with _compilation_context():
155
+ return self._compiled_callable(*args, **kwargs)
156
157
@abstractmethod
158
def forward(self, *args, **kwargs): ...
159
160
+ def original_code_object(self) -> CodeType:
161
+ """Return the original code object of the forward method."""
162
+ return self.__class__.forward.__code__
163
+
164
def bytecode_hook(self, old_code: CodeType, new_code: CodeType):
165
"""Hook to save the compiled bytecode for direct execution."""
101
- if old_code is not self.original_code_object:
166
+ if old_code is not self.original_code_object():
167
return
168
# code borrowed from https://github.com/thuml/depyf/blob/f4ad79fadee27ea113b4c75202db1eb1a11c0dbc/depyf/explain/enable_debugging.py#L25
169
frame = sys._getframe()
@@ -114,7 +179,7 @@ class TorchCompileWrapperWithCustomDispatcher:
179
if frame.f_locals["self"] is not self:
180
return
181
117
- self.compiled_codes.append(new_code)
182
+ self._compiled_bytecode = new_code
183
184
path = self.vllm_config.compile_debug_dump_path()
185
if path:
@@ -153,16 +218,21 @@ class TorchCompileWrapperWithCustomDispatcher:
218
raise RuntimeError(msg)
219
220
@contextmanager
156
- def dispatch_to_code(self, index: int):
157
- """Context manager to dispatch to the compiled code.
221
+ def _dispatch_to_compiled_code(self):
222
+ # noqa: E501
223
+ """
224
+ Context manager to dispatch to internally compiled code for torch<2.8.
225
Why does this work? Because Dynamo guarantees that the compiled
226
bytecode has exactly the same arguments, cell variables, and free
227
variables as the original code. Therefore we can directly switch
228
the code object in the function and call it.
229
163
- See https://dev-discuss.pytorch.org/t/what-is-the-relationship-requirement-among-original-bytecode-transformed-bytecode-and-bytecode-returned-by-hooks-in-dynamo/1693/7
164
- for more details.
165
- """
166
- self.__class__.forward.__code__ = self.compiled_codes[index]
167
- yield
168
- self.__class__.forward.__code__ = self.original_code_object
230
+ See https://dev-discuss.pytorch.org/t/what-is-the-relationship-requirement-among-original-bytecode-transformed-bytecode-and-bytecode-returned-by-hooks-in-dynamo/1693/7 for more details.
231
+ """ # noqa: E501 line too long
232
+ original = self.original_code_object()
233
+ assert self._compiled_bytecode is not None
234
+ self.__class__.forward.__code__ = self._compiled_bytecode
235
+ try:
236
+ yield
237
+ finally:
238
+ self.__class__.forward.__code__ = original
vllm/envs.py
+6
@@ -92,6 +92,7 @@ if TYPE_CHECKING:
92
VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False
93
VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
94
VLLM_USE_AOT_COMPILE: bool = False
95
+ VLLM_USE_BYTECODE_HOOK: bool = False
96
VLLM_FORCE_AOT_LOAD: bool = False
97
VLLM_TORCH_PROFILER_WITH_STACK: bool = True
98
VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
@@ -556,6 +557,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
557
# compilation is done in warmup phase and the compilation will be
558
# reused in subsequent calls.
559
"VLLM_USE_AOT_COMPILE": use_aot_compile,
560
+ # Feature flag to enable/disable bytecode in
561
+ # TorchCompileWithNoGuardsWrapper.
562
+ "VLLM_USE_BYTECODE_HOOK": lambda: bool(
563
+ int(os.environ.get("VLLM_USE_BYTECODE_HOOK", "1"))
564
+ ),
565
# Force vllm to always load AOT compiled models from disk. Failure
566
# to load will result in a hard error when this is enabled.
567
# Will be ignored when VLLM_USE_AOT_COMPILE is disabled.
vllm/v1/worker/tpu_model_runner.py
+6
-4
@@ -21,7 +21,7 @@ from vllm.attention import Attention
21
from vllm.attention.backends.abstract import AttentionType
22
from vllm.attention.layer import MLAAttention
23
from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
24
-from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
24
+from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
25
from vllm.config import (
26
ParallelConfig,
27
VllmConfig,
@@ -1895,12 +1895,14 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1895
compiled_model = self.model.get_language_model().model
1896
else:
1897
compiled_model = self.model.model
1898
- if isinstance(compiled_model, TorchCompileWrapperWithCustomDispatcher):
1898
+ if isinstance(compiled_model, TorchCompileWithNoGuardsWrapper):
1899
logger.info("Clear dynamo cache and cached dynamo bytecode.")
1900
torch._dynamo.eval_frame.remove_from_cache(
1901
- compiled_model.original_code_object
1901
+ compiled_model.original_code_object()
1902
)
1903
- compiled_model.compiled_codes.clear()
1903
+ # Reset the wrapper to re-initialize.
1904
+ compiled_model.compiled = False
1905
+ TorchCompileWithNoGuardsWrapper.__init__(compiled_model)
1906
1907
@torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1908
def select_hidden_states(self, hidden_states, indices_do_sample):