tiny refactor pcg split op registration (#16863)

python/sglang/srt/compilation/compilation_config.py

@@ -1,6 +1,17 @@
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/compilation_config.py
 
-from typing import List
+from typing import Callable, List, Optional
+
+SPLIT_OPS = []
+
+
+def register_split_op(op_name: Optional[str] = None):
+    def decorator(op_func: Callable):
+        name = op_name or op_func.__name__
+        SPLIT_OPS.append(f"sglang.{name}")
+        return op_func
+
+    return decorator
 
 
 # TODO(Yuwei): support better compile config support
@@ -15,11 +26,8 @@ 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",
-            "sglang.inplace_all_reduce",
-        ]
+        self.split_ops = []
+        self.split_ops.extend(SPLIT_OPS)
 
     def add_split_op(self, op: str):
         self.split_ops.append(op)