Patch AOTAutogradCache._get_shape_env

Author: jamesjwu

vllm/compilation/compiler_interface.py

@@ -280,6 +280,10 @@ def _get_shape_env() -> AlwaysHitShapeEnv:
                 patch("torch._inductor.codecache.FxGraphCache._get_shape_env",
                       _get_shape_env))
 
+            stack.enter_context(
+                patch("torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env",
+                      _get_shape_env))
+
             # for forcing the graph to be cached
             stack.enter_context(
                 patch(
@@ -330,6 +334,9 @@ def load(self,
             exit_stack.enter_context(
                 patch("torch._inductor.codecache.FxGraphCache._get_shape_env",
                       lambda *args, **kwargs: AlwaysHitShapeEnv()))
+            exit_stack.enter_context(
+                patch("torch._inductor.codecache.GuardedCache._get_shape_env",
+                      lambda *args, **kwargs: AlwaysHitShapeEnv()))
 
             # Dynamo metrics context, see method for more details.
             exit_stack.enter_context(self.metrics_context())