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 12, 13:30 UTC
[Compile] Refactor. Move PostGradPassManager out of Compilation config (#29340)
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Ilya Markov committed
Nov 25, 2025 at 20:58 UTC
e7d776273de379bb6c9fc11ce070c57e0fcd84f9
2 files changed
+21
-14
vllm/compilation/backends.py
+20
-13
@@ -11,6 +11,7 @@ import pprint
11
import time
12
from collections.abc import Callable, Sequence
13
from contextlib import contextmanager
14
+from copy import deepcopy
15
from functools import partial
16
from typing import Any
17
@@ -429,7 +430,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
430
self.vllm_backend.compiler_manager.compile(
431
submod,
432
args,
432
- self.compilation_config.inductor_compile_config,
433
+ self.vllm_backend.inductor_config,
434
self.compilation_config,
435
graph_index=index,
436
num_graphs=len(self.compile_submod_names),
@@ -531,6 +532,9 @@ class VllmBackend:
532
sym_tensor_indices: list[int]
533
input_buffers: list[torch.Tensor]
534
compiler_manager: CompilerManager
535
+ # Copy of CompilationConfig.inductor_compile_config +
536
+ # an entry for PostGradPassManager
537
+ inductor_config: dict[str, Any]
538
539
def __init__(
540
self,
@@ -561,25 +565,30 @@ class VllmBackend:
565
self.compilation_config
566
)
567
568
+ # Deepcopy the inductor config to detach the post-grad custom pass
569
+ # from CompilationConfig.
570
+ # We want to avoid PostGradPassManager in CompilationConfig because
571
+ # in future we need PostGradPassManager.uuid() to be executed
572
+ # only at compile time.
573
+ self.inductor_config = deepcopy(self.compilation_config.inductor_compile_config)
574
# `torch.compile` is JIT compiled, so we don't need to
575
# do anything here
576
577
def configure_post_pass(self):
568
- config = self.compilation_config
578
self.pass_manager.configure(self.vllm_config)
579
580
# Post-grad custom passes are run using the post_grad_custom_post_pass
581
# hook. If a pass for that hook exists, add it to the pass manager.
573
- inductor_config = config.inductor_compile_config
574
- if self.pass_key in inductor_config:
575
- if isinstance(inductor_config[self.pass_key], PostGradPassManager):
576
- # PassManager already added to config, make sure it's correct
577
- assert inductor_config[self.pass_key].uuid() == self.pass_manager.uuid()
582
+ if self.pass_key in self.inductor_config:
583
+ if isinstance(self.inductor_config[self.pass_key], PostGradPassManager):
584
+ raise ValueError(
585
+ "PostGradPassManager can not be kept in CompilationConfig."
586
+ )
587
else:
588
# Config should automatically wrap all inductor passes
580
- assert isinstance(inductor_config[self.pass_key], InductorPass)
581
- self.pass_manager.add(inductor_config[self.pass_key])
582
- inductor_config[self.pass_key] = self.pass_manager
589
+ assert isinstance(self.inductor_config[self.pass_key], InductorPass)
590
+ self.pass_manager.add(self.inductor_config[self.pass_key])
591
+ self.inductor_config[self.pass_key] = self.pass_manager
592
593
def __call__(
594
self, graph: fx.GraphModule, example_inputs
@@ -638,9 +647,7 @@ class VllmBackend:
647
self.compilation_config.local_cache_dir = local_cache_dir
648
649
# Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE.
641
- disable_cache = not is_compile_cache_enabled(
642
- self.compilation_config.inductor_compile_config
643
- )
650
+ disable_cache = not is_compile_cache_enabled(self.inductor_config)
651
652
if disable_cache:
653
logger.info_once("vLLM's torch.compile cache is disabled.", scope="local")
vllm/compilation/piecewise_backend.py
+1
-1
@@ -107,7 +107,7 @@ class PiecewiseBackend:
107
entry.runnable = self.vllm_backend.compiler_manager.compile(
108
self.graph,
109
args,
110
- self.compilation_config.inductor_compile_config,
110
+ self.vllm_backend.inductor_config,
111
self.compilation_config,
112
graph_index=self.piecewise_compile_index,
113
num_graphs=self.total_piecewise_compiles,