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 6, 17:00 UTC
use skip_all_guards_unsafe to drop global_state and torch_function_mode_stack guards instead of previous hacks (#36204)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Mar 16, 2026 at 01:52 UTC
52131f88d9f8d3257530ac492d9db40ca81b4872
1 file changed
+7
-50
vllm/compilation/wrapper.py
+7
-50
@@ -10,7 +10,6 @@ from types import CodeType
10
from typing import Any, ParamSpec, TypeVar
11
12
import torch
13
-import torch._C._dynamo.guards
13
14
import vllm.envs as envs
15
from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
@@ -24,65 +23,23 @@ R = TypeVar("R")
23
P = ParamSpec("P")
24
25
27
-def _noop_add_global_state_guard(
28
- self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any
29
-) -> None:
30
- """No-op to skip the GLOBAL_STATE guard entirely"""
31
- pass
32
-
33
-
34
-def _noop_add_torch_function_mode_stack_guard(
35
- self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any
36
-) -> None:
37
- """No-op to skip the TORCH_FUNCTION_MODE_STACK guard entirely"""
38
- pass
39
-
40
-
26
@contextmanager
27
def _compilation_context() -> Generator[None, None, None]:
43
- """Context manager for compilation settings and patches.
44
-
45
- This manager:
46
- 1. Sets higher dynamo cache limits for compilation. (Needed for
47
- qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
48
- Generally a recompilation can happen whenever we use a new
49
- backend instance in torch.compile.
50
- 2. Patches out add_global_state_guard to skip GLOBAL_STATE guards
51
- 3. Patches out add_torch_function_mode_stack_guard to skip
52
- TORCH_FUNCTION_MODE_STACK guards.
53
- 4. Restores everything when compilation completes
28
+ """Context manager for compilation settings.
29
+
30
+ This manager sets higher dynamo cache limits for compilation.
31
+ (Needed for qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
32
+ Generally a recompilation can happen whenever we use a new
33
+ backend instance in torch.compile.
34
"""
55
- # Save original values
56
- original_global_state_guard = (
57
- torch._C._dynamo.guards.GuardManager.add_global_state_guard
58
- )
59
- original_torch_function_mode_stack_guard = (
60
- torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard
61
- )
35
original_cache_size = torch._dynamo.config.cache_size_limit
36
original_accumulated_cache = torch._dynamo.config.accumulated_cache_size_limit
37
38
try:
66
- # Set higher cache limits for compilation
39
torch._dynamo.config.cache_size_limit = 2048
40
torch._dynamo.config.accumulated_cache_size_limit = 8192
69
-
70
- # Patch guard manager
71
- torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
72
- _noop_add_global_state_guard
73
- )
74
- torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
75
- _noop_add_torch_function_mode_stack_guard
76
- )
41
yield
42
finally:
79
- # Restore original values
80
- torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
81
- original_global_state_guard
82
- )
83
- torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
84
- original_torch_function_mode_stack_guard
85
- )
43
torch._dynamo.config.cache_size_limit = original_cache_size
44
torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache
45
@@ -155,7 +112,7 @@ class TorchCompileWithNoGuardsWrapper:
112
entry.guard_type == "SHAPE_ENV" for entry in x
113
]
114
else:
158
- options["guard_filter_fn"] = lambda x: [False for _ in x]
115
+ options["guard_filter_fn"] = torch.compiler.skip_all_guards_unsafe
116
117
compiled_ptr: Any = self.forward
118
# Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False