diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index eeb222ff9..33583cf1f 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -371,6 +371,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--enable-single-batch-overlap` | Let computation and communication overlap within one micro batch. | `False` | bool flag (set to enable) | | `--tbo-token-distribution-threshold` | The threshold of token distribution between two batches in micro-batch-overlap, determines whether to two-batch-overlap or two-chunk-overlap. Set to 0 denote disable two-chunk-overlap. | `0.48` | Type: float | | `--enable-torch-compile` | Optimize the model with torch.compile. Experimental feature. | `False` | bool flag (set to enable) | +| `--enable-torch-compile-debug-mode` | Enable debug mode for torch compile. | `False` | bool flag (set to enable) | | `--enable-piecewise-cuda-graph` | Optimize the model with piecewise cuda graph for extend/prefill only. Experimental feature. | `False` | bool flag (set to enable) | | `--piecewise-cuda-graph-tokens` | Set the list of tokens when using piecewise cuda graph. | `None` | Type: JSON list | | `--torch-compile-max-bs` | Set the maximum batch size when using torch compile. | `32` | Type: int | diff --git a/python/sglang/srt/compilation/compilation_config.py b/python/sglang/srt/compilation/compilation_config.py index fbf1493e1..5ddafe8af 100644 --- a/python/sglang/srt/compilation/compilation_config.py +++ b/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 diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index 6d4f514c1..2e45d34d3 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/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) diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 95f486abe..3ccf613dc 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 04cc868ab..4d546d79b 100644 --- a/python/sglang/srt/server_args.py +++ b/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",