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 8, 12:30 UTC
[torch.compile] Disable recursive pre_grad_passes (#34092)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Feb 10, 2026 at 18:02 UTC
341eed3d30b7579b730e9959213d83b5dbd4731c
2 files changed
+24
-1
vllm/compilation/compiler_interface.py
+14
-1
@@ -257,7 +257,20 @@ class InductorStandaloneAdaptor(CompilerInterface):
257
if use_aot:
258
compile_kwargs["aot"] = True # type: ignore[assignment]
259
260
- compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
260
+ # Inductor's pre-grad passes don't do anything for vLLM.
261
+ # The pre-grad passes get run even on cache-hit and negatively impact
262
+ # vllm cold compile times by O(1s)
263
+ # Can remove this after the following issue gets fixed
264
+ # https://github.com/pytorch/pytorch/issues/174502
265
+ if envs.VLLM_ENABLE_PREGRAD_PASSES:
266
+ ctx: Any = contextlib.nullcontext()
267
+ else:
268
+ ctx = patch(
269
+ "torch._inductor.compile_fx._recursive_pre_grad_passes",
270
+ lambda gm, _: gm,
271
+ )
272
+ with ctx:
273
+ compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
274
275
if use_aot:
276
from torch._inductor.standalone_compile import AOTCompiledArtifact
vllm/envs.py
+10
@@ -132,6 +132,7 @@ if TYPE_CHECKING:
132
VLLM_DP_RANK_LOCAL: int = -1
133
VLLM_DP_SIZE: int = 1
134
VLLM_USE_STANDALONE_COMPILE: bool = True
135
+ VLLM_ENABLE_PREGRAD_PASSES: bool = False
136
VLLM_DP_MASTER_IP: str = ""
137
VLLM_DP_MASTER_PORT: int = 0
138
VLLM_MOE_DP_CHUNK_SIZE: int = 256
@@ -568,6 +569,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
569
"VLLM_USE_STANDALONE_COMPILE", "1"
570
)
571
== "1",
572
+ # Inductor's pre-grad passes don't do anything for vLLM.
573
+ # The pre-grad passes get run even on cache-hit and negatively impact
574
+ # vllm cold compile times by O(1s)
575
+ # Can remove this after the following issue gets fixed
576
+ # https://github.com/pytorch/pytorch/issues/174502
577
+ "VLLM_ENABLE_PREGRAD_PASSES": lambda: os.environ.get(
578
+ "VLLM_ENABLE_PREGRAD_PASSES", "0"
579
+ )
580
+ == "1",
581
# Debug pattern matching inside custom passes.
582
# Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').
583
"VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get(