Eager Compiler for Torch Compile (#11803)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>

python/sglang/srt/compilation/backend.py

@@ -17,24 +17,30 @@ from torch._dispatch.python import enable_python_dispatcher
 
 from sglang.srt.compilation.compilation_config import CompilationConfig
 from sglang.srt.compilation.compilation_counter import compilation_counter
-from sglang.srt.compilation.compiler_interface import InductorAdaptor
+from sglang.srt.compilation.compiler_interface import EagerAdapter, InductorAdaptor
 from sglang.srt.compilation.cuda_piecewise_backend import CUDAPiecewiseBackend
 from sglang.srt.compilation.pass_manager import PostGradPassManager
 
 logger = logging.getLogger(__name__)
 
 
-def make_compiler():
-    return InductorAdaptor()
+def make_compiler(config: CompilationConfig):
+    if config.compiler == "eager":
+        return EagerAdapter()
+    elif config.compiler == "inductor":
+        return InductorAdaptor()
+    else:
+        raise ValueError(f"Unknown compiler: {config.compiler}")
 
 
 class CompilerManager:
     def __init__(
         self,
+        config: CompilationConfig,
     ):
         self.cache = dict()
         self.is_cache_updated = False
-        self.compiler = make_compiler()
+        self.compiler = make_compiler(config)
 
     def compute_hash(self):
         return self.compiler.compute_hash()
@@ -348,7 +354,7 @@ class SGLangBackend:
         self.sym_tensor_indices = []
         self.input_buffers = []
 
-        self.compiler_manager = CompilerManager()
+        self.compiler_manager = CompilerManager(config)
         self.inductor_config = {
             "enable_auto_functionalized_v2": False,
         }

python/sglang/srt/compilation/compilation_config.py

@@ -5,9 +5,10 @@ from typing import List
 
 # TODO(Yuwei): support better compile config support
 class CompilationConfig:
-    def __init__(self, capture_sizes: List[int]):
+    def __init__(self, capture_sizes: List[int], compiler: str = "eager"):
         self.traced_files = set()
         self.capture_sizes = capture_sizes
+        self.compiler = compiler
 
     def add_traced_file(self, file_path: str):
         self.traced_files.add(file_path)

python/sglang/srt/compilation/compiler_interface.py

@@ -475,3 +475,29 @@ def set_inductor_config(config, runtime_shape):
         # can be beneficial
         config["max_autotune"] = True
         config["coordinate_descent_tuning"] = True
+
+
+class EagerAdapter(CompilerInterface):
+    name = "eager"
+
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        compiler_config: dict[str, Any],
+        runtime_shape: Optional[int] = None,
+        key: Optional[str] = None,
+        num_graphs: int = 1,
+    ) -> tuple[Optional[Callable], Optional[Any]]:
+        return graph, None
+
+    def load(
+        self,
+        handle: Any,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        graph_index: int,
+        runtime_shape: Optional[int] = None,
+        num_graphs: int = 1,
+    ) -> Callable:
+        raise NotImplementedError("eager compilation is not supported")

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -9,6 +9,7 @@ from unittest.mock import patch
 import torch
 import torch.fx as fx
 
+import sglang.srt.compilation.weak_ref_tensor_jit  # noqa: F401
 from sglang.srt.compilation.compilation_config import CompilationConfig
 from sglang.srt.compilation.compilation_counter import compilation_counter