Passing only necessary compilation config to inductor pass config

Summary: we should not pass the weakref to compilation_config, which  include static_forward_context that will holds the pointers to the model layers (e.g. moe, attention), which is dangerous, as this will be passed as config to torch.compile

Differential Revision: D84790018

Author: Lucia (Lu) Fang

vllm/compilation/vllm_inductor_pass.py

@@ -3,21 +3,34 @@
 import functools
 import operator
 import time
-import weakref
 from typing import ClassVar
 
 import regex as re
 import torch
+from dataclasses import dataclass
 from torch._dynamo.utils import lazy_format_graph_code
 from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter
 
-from vllm.config import VllmConfig
+from vllm.config import CompilationConfig, VllmConfig
 from vllm.logger import init_logger
 
 from .inductor_pass import InductorPass
 
 logger = init_logger(__name__)
 
+@dataclass
+class SimplifiedCompilationConfig:
+    splitting_ops: list[str] | None = None
+    use_inductor_graph_partition: bool = False
+    compile_sizes: bool = False
+
+def copy_necessary_config_for_pass(config: CompilationConfig) -> SimplifiedCompilationConfig:
+    """Get the necessary CompilationConfig for the current pass."""
+    return SimplifiedCompilationConfig(
+        splitting_ops=config.splitting_ops,
+        use_inductor_graph_partition=config.use_inductor_graph_partition,
+        compile_sizes=config.compile_sizes,
+    )
 
 class VllmInductorPass(InductorPass):
     """
@@ -29,7 +42,7 @@ class VllmInductorPass(InductorPass):
     """Keep track of pass index for debug dump ordering."""
 
     def __init__(self, config: VllmConfig):
-        self.compilation_config = weakref.proxy(config.compilation_config)
+        self.compilation_config = copy_necessary_config_for_pass(config.compilation_config)
         self.pass_config = config.compilation_config.pass_config
         self.model_dtype = config.model_config.dtype if config.model_config else None
         self.device = config.device_config.device if config.device_config else None