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 15, 13:30 UTC
[Bugfix] [pytorch] Patch AOTAutogradCache._get_shape_env (#17142)
Signed-off-by: James Wu <jjwu@meta.com>
James Wu committed
Apr 25, 2025 at 23:28 UTC
a6e72e1e4fb450c80f15e09b9f09d5754635724e
1 file changed
+18
-1
vllm/compilation/compiler_interface.py
+18
-1
@@ -195,7 +195,6 @@ class InductorAdaptor(CompilerInterface):
195
hash_str, file_path = None, None
196
from torch._inductor.codecache import (FxGraphCache,
197
compiled_fx_graph_hash)
198
-
198
if torch.__version__.startswith("2.5"):
199
original_load = FxGraphCache.load
200
original_load_name = "torch._inductor.codecache.FxGraphCache.load"
@@ -280,6 +279,16 @@ class InductorAdaptor(CompilerInterface):
279
patch("torch._inductor.codecache.FxGraphCache._get_shape_env",
280
_get_shape_env))
281
282
+ from torch._functorch._aot_autograd.autograd_cache import (
283
+ AOTAutogradCache)
284
+
285
+ # torch 2.8+ on main uses _get_shape_env in AOTAutogradCache
286
+ if hasattr(AOTAutogradCache, "_get_shape_env"):
287
+ stack.enter_context(
288
+ patch(
289
+ "torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env",
290
+ _get_shape_env))
291
+
292
# for forcing the graph to be cached
293
stack.enter_context(
294
patch(
@@ -325,11 +334,19 @@ class InductorAdaptor(CompilerInterface):
334
assert isinstance(handle[1], str)
335
hash_str = handle[0]
336
337
+ from torch._functorch._aot_autograd.autograd_cache import (
338
+ AOTAutogradCache)
339
from torch._inductor.codecache import FxGraphCache
340
with ExitStack() as exit_stack:
341
exit_stack.enter_context(
342
patch("torch._inductor.codecache.FxGraphCache._get_shape_env",
343
lambda *args, **kwargs: AlwaysHitShapeEnv()))
344
+ # torch 2.8+ on main uses _get_shape_env in AOTAutogradCache
345
+ if hasattr(AOTAutogradCache, "_get_shape_env"):
346
+ exit_stack.enter_context(
347
+ patch(
348
+ "torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env",
349
+ lambda *args, **kwargs: AlwaysHitShapeEnv()))
350
351
# Dynamo metrics context, see method for more details.
352
exit_stack.enter_context(self.metrics_context())