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 = [
+        self.split_ops = []
-            "sglang.unified_attention_with_output",
+        self.split_ops.extend(SPLIT_OPS)
-            "sglang.gdn_with_output",
-            "sglang.inplace_all_reduce",
-        ]
 
     def add_split_op(self, op: str):
         self.split_ops.append(op)

python/sglang/srt/distributed/parallel_state.py

@@ -39,6 +39,7 @@ import torch
 import torch.distributed
 from torch.distributed import Backend, ProcessGroup
 
+from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.environ import envs
 from sglang.srt.utils import (
     get_bool_env_var,
@@ -125,6 +126,7 @@ def _register_group(group: "GroupCoordinator") -> None:
 
 
 @register_custom_op(mutates_args=["tensor"])
+@register_split_op()
 def inplace_all_reduce(tensor: torch.Tensor, group_name: str) -> None:
     assert group_name in _groups, f"Group {group_name} is not found."
     group = _groups[group_name]()

python/sglang/srt/layers/radix_attention.py

@@ -20,6 +20,7 @@ from typing import TYPE_CHECKING, Optional
 import torch
 from torch import nn
 
+from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.compilation.piecewise_context_manager import get_forward_context
 from sglang.srt.utils.custom_op import register_custom_op
 
@@ -132,6 +133,7 @@ class RadixAttention(nn.Module):
 
 
 @register_custom_op(mutates_args=["output"])
+@register_split_op()
 def unified_attention_with_output(
     query: torch.Tensor,
     key: torch.Tensor,

python/sglang/srt/models/qwen3_next.py

@@ -5,6 +5,7 @@ from typing import Any, Iterable, Optional, Set, Tuple
 import torch
 from torch import nn
 
+from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.compilation.piecewise_context_manager import get_forward_context
 from sglang.srt.configs.qwen3_next import Qwen3NextConfig
 from sglang.srt.distributed import get_pp_group
@@ -1060,6 +1061,7 @@ EntryClass = Qwen3NextForCausalLM
 
 
 @register_custom_op(mutates_args=["output"])
+@register_split_op()
 def gdn_with_output(
     hidden_states: torch.Tensor,
     output: torch.Tensor,