EP Support for Piecewise Cuda Graph (#14164)

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

python/sglang/srt/compilation/backend.py

@@ -26,16 +26,6 @@ from sglang.srt.utils.common import is_npu, rank0_log
 logger = logging.getLogger(__name__)
 
 
-SPLIT_OPS = [
-    "sglang.unified_attention_with_output",
-    "sglang.gdn_with_output",
-]
-
-
-def add_split_ops(ops):
-    SPLIT_OPS.extend(ops)
-
-
 def make_compiler(config: CompilationConfig):
     if config.compiler == "eager":
         return EagerAdapter()
@@ -433,7 +423,7 @@ class SGLangBackend:
 
         self.split_gm, self.piecewise_graphs = split_graph(
             graph,
-            SPLIT_OPS,
+            self.compile_config.split_ops,
         )
         from torch._dynamo.utils import lazy_format_graph_code
 

python/sglang/srt/compilation/compilation_config.py

@@ -15,6 +15,13 @@ class CompilationConfig:
         self.capture_sizes = capture_sizes
         self.compiler = compiler
         self.enable_debug_mode = enable_debug_mode
+        self.split_ops = [
+            "sglang.unified_attention_with_output",
+            "sglang.gdn_with_output",
+        ]
+
+    def add_split_op(self, op: str):
+        self.split_ops.append(op)
 
     def add_traced_file(self, file_path: str):
         self.traced_files.add(file_path)

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -26,6 +26,8 @@ class ForwardContext:
     def __init__(self):
         self.forward_batch = None
         self.attention_layer = None
+        self.quant_config = None
+        self.moe_layers = None
 
     def set_forward_batch(self, forward_batch: ForwardBatch):
         self.forward_batch = forward_batch
@@ -36,6 +38,9 @@ class ForwardContext:
     def set_quant_config(self, quant_config: Any):
         self.quant_config = quant_config
 
+    def set_moe_layers(self, layers: List[Any]):
+        self.moe_layers = layers
+
 
 _forward_context: Optional[ForwardContext] = None
 
@@ -48,13 +53,17 @@ def get_forward_context() -> Optional[ForwardContext]:
 
 @contextmanager
 def set_forward_context(
-    forward_batch: ForwardBatch, attention_layers: List[Any], quant_config: Any
+    forward_batch: ForwardBatch,
+    attention_layers: List[Any],
+    quant_config: Any,
+    moe_layers: List[Any],
 ):
     global _forward_context
     _forward_context = ForwardContext()
     _forward_context.set_forward_batch(forward_batch)
     _forward_context.set_attention_layers(attention_layers)
     _forward_context.set_quant_config(quant_config)
+    _forward_context.set_moe_layers(moe_layers)
     try:
         yield
     finally: