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, 07:00 UTC
[compile] Initialize passes at VllmBackend init (#35216)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Mar 20, 2026 at 11:40 UTC
12fd17eb5198708523008dda6809143d0f7234ed
3 files changed
+19
-5
tests/test_config.py
+2
-2
@@ -32,9 +32,9 @@ from vllm.platforms import current_platform
32
33
def test_compile_config_repr_succeeds():
34
# setup: VllmBackend mutates the config object
35
+ # Note: VllmBackend.__init__ already calls configure_post_pass()
36
config = VllmConfig()
36
- backend = VllmBackend(config)
37
- backend.configure_post_pass()
37
+ _ = VllmBackend(config)
38
39
# test that repr(config) succeeds
40
val = repr(config)
vllm/compilation/backends.py
+12
-3
@@ -836,8 +836,18 @@ class VllmBackend:
836
# in future we need PostGradPassManager.uuid() to be executed
837
# only at compile time.
838
self.inductor_config = deepcopy(self.compilation_config.inductor_compile_config)
839
- # `torch.compile` is JIT compiled, so we don't need to
840
- # do anything here
839
+
840
+ # Configure post-grad passes (including AllReduceFusionPass) during
841
+ # backend init rather than at torch.compile time, so that expensive
842
+ # one-time setup (e.g. FlashInfer workspace allocation) is not
843
+ # attributed to compilation latency.
844
+ start = time.time()
845
+ self.configure_post_pass()
846
+ logger.info_once(
847
+ "Post-grad pass configuration time: %.2f s",
848
+ time.time() - start,
849
+ scope="local",
850
+ )
851
852
def collect_standalone_compile_artifacts(
853
self,
@@ -1118,7 +1128,6 @@ class VllmBackend:
1128
assert not self._called, "VllmBackend can only be called once"
1129
1130
self.graph = graph
1121
- self.configure_post_pass()
1131
1132
if self.compilation_config.use_inductor_graph_partition:
1133
# Let Inductor decide partitioning; avoid FX-level pre-splitting.
vllm/compilation/decorators.py
+5
@@ -380,6 +380,11 @@ def _support_torch_compile(
380
compilation_counter.num_models_seen += 1
381
self.compiled = False
382
383
+ # Skip if a parent class's @support_torch_compile already
384
+ # initialized the compile wrapper
385
+ if hasattr(self, "_compiled_callable"):
386
+ return
387
+
388
# Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class
389
TorchCompileWithNoGuardsWrapper.__init__(
390
self,