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 10, 12:45 UTC
fix(compile): apply partition wrapper when loading AOT cached functions (#31536)
Signed-off-by: Devbyteai <abud6673@gmail.com> Signed-off-by: DevByteAI <161969603+devbyteai@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
DevByteAI committed
Jan 8, 2026 at 11:27 UTC
1f214290d65a6a69b898e058f1408f2d929f8fa7
2 files changed
+103
-3
tests/compile/test_aot_compile.py
+95
@@ -5,6 +5,7 @@ import functools
5
import multiprocessing
6
import tempfile
7
from contextlib import contextmanager
8
+from pathlib import Path
9
10
import pytest
11
import torch
@@ -24,6 +25,13 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
25
from ..utils import create_new_process_for_each_test
26
27
28
+@pytest.fixture
29
+def vllm_tmp_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
30
+ """Fixture that sets VLLM_CACHE_ROOT to a temporary directory."""
31
+ monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm_cache"))
32
+ return tmp_path
33
+
34
+
35
def reference_fn(x: torch.Tensor):
36
assert x.shape[0] <= 42
37
assert x.shape[0] % 2 == 0
@@ -148,6 +156,93 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
156
assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
157
158
159
+@pytest.mark.skipif(
160
+ not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
161
+)
162
+def test_partition_wrapper_applied_on_aot_load(
163
+ monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker
164
+):
165
+ """
166
+ Test that partition wrappers are applied when loading AOT cached functions.
167
+
168
+ This test verifies the fix for GitHub issue #31439 where AOT compile
169
+ caused 2x latency regression when use_inductor_graph_partition=True.
170
+ The root cause was that partition wrapper context was bypassed when
171
+ loading from AOT cache.
172
+ """
173
+ from vllm.config import CUDAGraphMode
174
+
175
+ args = (torch.randn(10, 10),)
176
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
177
+
178
+ # Create config with partition enabled
179
+ vllm_config = VllmConfig(
180
+ compilation_config=CompilationConfig(
181
+ mode=CompilationMode.VLLM_COMPILE,
182
+ use_inductor_graph_partition=True,
183
+ cudagraph_mode=CUDAGraphMode.PIECEWISE,
184
+ )
185
+ )
186
+
187
+ # First compilation - save to cache
188
+ with use_vllm_config(vllm_config):
189
+ compiled_mod = CompiledMod(vllm_config=vllm_config)
190
+ compiled_mod(*args)
191
+ disable_envs_cache()
192
+
193
+ # Second run - load from cache, verify partition wrapper applied
194
+ monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", "1")
195
+ vllm_config = VllmConfig(
196
+ compilation_config=CompilationConfig(
197
+ mode=CompilationMode.VLLM_COMPILE,
198
+ use_inductor_graph_partition=True,
199
+ cudagraph_mode=CUDAGraphMode.PIECEWISE,
200
+ )
201
+ )
202
+
203
+ # Use mocker to spy on set_customized_partition_wrappers
204
+ spy = mocker.spy(torch._inductor.utils, "set_customized_partition_wrappers")
205
+
206
+ with use_vllm_config(vllm_config):
207
+ compiled_mod = CompiledMod(vllm_config=vllm_config)
208
+
209
+ # First call after restart: loads from AOT cache.
210
+ # This tests the fix for the first call after a restart.
211
+ compiled_mod(*args)
212
+
213
+ # Verify partition wrapper was called on AOT load.
214
+ assert spy.call_count >= 2, (
215
+ "Expected partition wrapper to be set and cleared on AOT load, "
216
+ f"got {spy.call_count} calls"
217
+ )
218
+ # First call should set a wrapper, last call should clear it
219
+ assert spy.call_args_list[0][0][0] is not None, (
220
+ "First call on AOT load should set a wrapper function"
221
+ )
222
+ assert spy.call_args_list[-1][0][0] is None, (
223
+ "Last call on AOT load should clear the wrapper"
224
+ )
225
+
226
+ # Reset for the next check.
227
+ spy.reset_mock()
228
+
229
+ # Subsequent call: uses the cached `aot_compiled_fn`.
230
+ # This tests the fix for subsequent calls.
231
+ compiled_mod(*args)
232
+
233
+ # Verify partition wrapper was called on the subsequent call.
234
+ assert spy.call_count >= 2, (
235
+ "Expected partition wrapper set and cleared on subsequent "
236
+ f"call, got {spy.call_count} calls"
237
+ )
238
+ assert spy.call_args_list[0][0][0] is not None, (
239
+ "First call on subsequent call should set a wrapper function"
240
+ )
241
+ assert spy.call_args_list[-1][0][0] is None, (
242
+ "Last call on subsequent call should clear the wrapper"
243
+ )
244
+
245
+
246
@pytest.mark.skipif(
247
not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
248
)
vllm/compilation/decorators.py
+8
-3
@@ -371,9 +371,12 @@ def _support_torch_compile(
371
if self.do_not_compile or torch.compiler.is_compiling():
372
return self.forward(*args, **kwargs)
373
374
- # if aot_compiled_fn is set, just call it.
374
+ # if aot_compiled_fn is set, call it with partition wrapper context.
375
+ # The partition wrapper must be active at runtime for CUDA graph
376
+ # capture to work correctly with inductor graph partitioning.
377
if getattr(self, "aot_compiled_fn", None) is not None:
376
- return self.aot_compiled_fn(self, *args, **kwargs)
378
+ with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
379
+ return self.aot_compiled_fn(self, *args, **kwargs)
380
381
ds_type = self.compilation_config.dynamic_shapes_config.type
382
cache_dir = None
@@ -432,7 +435,9 @@ def _support_torch_compile(
435
logger.info(
436
"Directly load AOT compilation from path %s", aot_compilation_path
437
)
435
- return self.aot_compiled_fn(self, *args, **kwargs)
438
+ # Apply partition wrapper context for proper CUDA graph capture
439
+ with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
440
+ return self.aot_compiled_fn(self, *args, **kwargs)
441
442
if self.compiled:
443
assert (