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 14, 05:30 UTC
[BugFix] Work around graph partition x torch.compile cache issue (#26956)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Oct 15, 2025 at 23:06 UTC
9b6504c30751ab2cd73a1fb43ba82969e478b8dd
2 files changed
+31
-12
tests/compile/piecewise/test_toy_llama.py
+1
-11
@@ -337,9 +337,8 @@ def run_model(llama_config, compile_config: CompilationConfig) -> torch.Tensor:
337
def test_toy_llama(
338
backend: str, use_inductor_graph_partition: bool, monkeypatch, tmp_path
339
):
340
- # We disable the vLLM compile cache into a new tmp dir for 2 reasons:
340
+ # We disable the vLLM compile cache into a new tmp dir for 1 reason:
341
# 1. To make sure we can properly track the number of Inductor compilations.
342
- # 2. Inductor partitioning does not play nicely with Autograd cache (below)
342
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
343
344
if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
@@ -369,15 +368,6 @@ def test_toy_llama(
368
cudagraph_capture_sizes=[1, 2],
369
)
370
372
- # FIXME(luka/boyuan): the graph from the previous test case
373
- # (no inductor partition) gets cached by AotAutograd so then the
374
- # compilation with inductor partitioning incorrectly loads an unpartitioned
375
- # graph and never partitions. I think this is a bug with custom inductor
376
- # partitioning but does not affect vLLM more generally as vLLM uses its own
377
- # cache (which takes inductor partitioning into account).
378
- if use_inductor_graph_partition:
379
- compile_config_no_split.inductor_compile_config["force_disable_caches"] = True
380
-
371
compile_config_split = deepcopy(compile_config_no_split)
372
compile_config_split.splitting_ops = ["silly::attention"]
373
vllm/compilation/pass_manager.py
+30
-1
@@ -110,6 +110,27 @@ class PostGradPassManager(CustomGraphPass):
110
self.post_cleanup = PostCleanupPass(config)
111
self.fix_functionalization = FixFunctionalizationPass(config)
112
113
+ # [HACK: Bug with Inductor graph partition and torch.compile cache]
114
+ # In PyTorch 2.9, torch.compile has a bug where the graph
115
+ # partition is not taken into account during caching.
116
+ # Because vLLM's Mode.VLLM_COMPILE is the only mode that uses
117
+ # Inductor graph partition, and VLLM_COMPILE implies there
118
+ # is a PostGradPassManager, we put the list of operators to graph
119
+ # partition into the PostGradPassManager's uuid (which
120
+ # then gets incorporated into Inductor's FX graph cache key).
121
+ # Remove this hack whenever torch.compile fixes it.
122
+
123
+ # This is the list of operators that vLLM asks Inductor to split.
124
+ self.inductor_splitting_ops = []
125
+ if (
126
+ config.compilation_config.use_inductor_graph_partition
127
+ and config.compilation_config.splitting_ops is not None
128
+ ):
129
+ # Sort them so we're not dependent on the ordering.
130
+ self.inductor_splitting_ops = sorted(
131
+ config.compilation_config.splitting_ops
132
+ )
133
+
134
def add(self, pass_: InductorPass):
135
assert isinstance(pass_, InductorPass)
136
self.passes.append(pass_)
@@ -120,8 +141,16 @@ class PostGradPassManager(CustomGraphPass):
141
affects compilation caching. Its uuid depends on the UUIDs of all
142
dependent passes and the pass config. See InductorPass for more info.
143
"""
123
- state = {"pass_config": self.pass_config.uuid(), "passes": []}
144
+ state = {
145
+ "pass_config": self.pass_config.uuid(),
146
+ "passes": [],
147
+ "inductor_splitting_ops": [],
148
+ }
149
for pass_ in self.passes:
150
state["passes"].append(pass_.uuid())
151
state["passes"].append(self.fix_functionalization.uuid())
152
+
153
+ # See [HACK: Bug with Inductor graph partition and torch.compile cache]
154
+ state["inductor_splitting_ops"].extend(self.inductor_splitting_ops)
155
+
156
return InductorPass.hash_dict(state)