Move unnecessary input_addr capture under debug mode flag for speed-up (#13690)

python/sglang/srt/compilation/compilation_config.py

@@ -5,10 +5,16 @@ from typing import List
 
 # TODO(Yuwei): support better compile config support
 class CompilationConfig:
-    def __init__(self, capture_sizes: List[int], compiler: str = "eager"):
+    def __init__(
+        self,
+        capture_sizes: List[int],
+        compiler: str = "eager",
+        enable_debug_mode: bool = False,
+    ):
         self.traced_files = set()
         self.capture_sizes = capture_sizes
         self.compiler = compiler
+        self.enable_debug_mode = enable_debug_mode
 
     def add_traced_file(self, file_path: str):
         self.traced_files.add(file_path)
@@ -18,3 +24,6 @@ class CompilationConfig:
 
     def get_capture_sizes(self):
         return self.capture_sizes
+
+    def get_enable_debug_mode(self):
+        return self.enable_debug_mode

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -96,8 +96,6 @@ class CUDAPiecewiseBackend:
 
         self.sym_shape_indices = sym_shape_indices
 
-        self.is_debugging_mode = True
-
         # the entries for different shapes that we need to either
         # compile or capture cudagraph
         self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}
@@ -161,10 +159,11 @@ class CUDAPiecewiseBackend:
                 entry.num_finished_warmup += 1
                 return entry.runnable(*args)
 
-            input_addresses = [
-                x.data_ptr() for x in args if isinstance(x, torch.Tensor)
-            ]
-            entry.input_addresses = input_addresses
+            if self.compile_config.get_enable_debug_mode():
+                input_addresses = [
+                    x.data_ptr() for x in args if isinstance(x, torch.Tensor)
+                ]
+                entry.input_addresses = input_addresses
             cudagraph = torch.cuda.CUDAGraph()
 
             with ExitStack() as stack:
@@ -202,7 +201,7 @@ class CUDAPiecewiseBackend:
             # manage the memory during cuda graph capture
             return output
 
-        if self.is_debugging_mode:
+        if self.compile_config.get_enable_debug_mode():
             # check if the input addresses are the same
             new_input_addresses = [
                 x.data_ptr() for x in args if isinstance(x, torch.Tensor)

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -170,6 +170,7 @@ class PiecewiseCudaGraphRunner:
         self.compile_config = CompilationConfig(
             self.model_runner.server_args.piecewise_cuda_graph_tokens,
             self.model_runner.server_args.piecewise_cuda_graph_compiler,
+            self.model_runner.server_args.enable_torch_compile_debug_mode,
         )
         self.quant_config = getattr(self.model_runner.model, "quant_config", None)
 

python/sglang/srt/server_args.py

@@ -501,6 +501,7 @@ class ServerArgs:
     tbo_token_distribution_threshold: float = 0.48
     enable_torch_compile: bool = False
     enable_piecewise_cuda_graph: bool = False
+    enable_torch_compile_debug_mode: bool = False
     torch_compile_max_bs: int = 32
     piecewise_cuda_graph_max_tokens: int = 4096
     piecewise_cuda_graph_tokens: Optional[List[int]] = None
@@ -3413,6 +3414,11 @@ class ServerArgs:
             action="store_true",
             help="Optimize the model with torch.compile. Experimental feature.",
         )
+        parser.add_argument(
+            "--enable-torch-compile-debug-mode",
+            action="store_true",
+            help="Enable debug mode for torch compile",
+        )
         parser.add_argument(
             "--enable-piecewise-cuda-graph",
             action="store_true",