Piecewise Cuda Graph set default (#16331)

python/sglang/srt/compilation/compile.py

@@ -1,31 +1,19 @@
-import contextvars
 import inspect
 import logging
 import os
 import sys
 import types
-from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any, Callable, Optional, Union
 
 import torch
 
 from sglang.srt.compilation.compilation_config import CompilationConfig
+from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.utils.common import rank0_log
 
 logger = logging.getLogger(__name__)
 
-_COMPILE_ENABLED = contextvars.ContextVar("_COMPILE_ENABLED", default=False)
-
-
-@contextmanager
-def set_compiled(enabled: bool = True):
-    token = _COMPILE_ENABLED.set(enabled)
-    try:
-        yield
-    finally:
-        _COMPILE_ENABLED.reset(token)
-
 
 @dataclass
 class IntermediateTensors:
@@ -200,7 +188,7 @@ def install_torch_compiled(
         state["compiled_callable"] = compiled_callable
 
     def trampoline(self, *args, **kwargs):
-        use_compiled = _COMPILE_ENABLED.get()
+        use_compiled = is_in_piecewise_cuda_graph()
         if use_compiled:
             if not state["compiled"]:
                 _ensure_compiled(self, *args, **kwargs)

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
+import logging
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, List, Optional
 
 import torch
 
+logger = logging.getLogger(__name__)
+
+
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
@@ -38,10 +42,17 @@ def enable_piecewise_cuda_graph_compile():
 def enable_piecewise_cuda_graph():
     global _in_piecewise_cuda_graph
     _in_piecewise_cuda_graph = True
-
-    yield
-
-    _in_piecewise_cuda_graph = False
+    try:
+        yield
+    except Exception as e:
+        logger.error(
+            "Piecewise CUDA Graph failed with error: %s\n%s",
+            e,
+            PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+        )
+        raise
+    finally:
+        _in_piecewise_cuda_graph = False
 
 
 @contextmanager
@@ -56,7 +67,7 @@ def set_pcg_capture_stream(stream: torch.cuda.Stream):
 class ForwardContext:
     def __init__(self):
         self.forward_batch = None
-        self.attention_layer = None
+        self.attention_layers = None
         self.quant_config = None
         self.moe_layers = None
         self.moe_fusions = None
@@ -105,3 +116,10 @@ def set_forward_context(
         yield
     finally:
         _forward_context = None
+
+
+PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "Piecewise CUDA Graph is enabled by default as an experimental feature.\n"
+    "To work around this error, add --disable-piecewise-cuda-graph to your launch command.\n"
+    "Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose"
+)