EP Support for Piecewise Cuda Graph (#14164)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>
2025-12-19 09:59:27 -08:00
parent 05eb0bcc61
commit 9d0347b33a
11 changed files with 224 additions and 56 deletions
--- a/python/sglang/srt/compilation/compilation_config.py
+++ b/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)