From 206db66f5c239bc13e6e4cfa7a4ea952612a2e59 Mon Sep 17 00:00:00 2001
From: Qiaolin Yu <liin1211@outlook.com>
Date: Sat, 10 Jan 2026 07:45:28 -0800
Subject: [PATCH] tiny refactor pcg split op registration (#16863)

---
 .../srt/compilation/compilation_config.py     | 20 +++++++++++++------
 .../sglang/srt/distributed/parallel_state.py  |  2 ++
 python/sglang/srt/layers/radix_attention.py   |  2 ++
 python/sglang/srt/models/qwen3_next.py        |  2 ++
 4 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/compilation/compilation_config.py b/python/sglang/srt/compilation/compilation_config.py
index d7687aaf1..0388bbeda 100644
--- a/python/sglang/srt/compilation/compilation_config.py
+++ b/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)
diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index 0478526ef..3e3e3bf0a 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/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]()
diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py
index 199b2aa36..f09e3474a 100644
--- a/python/sglang/srt/layers/radix_attention.py
+++ b/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,
diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py
index ae17e4319..6a662edfa 100644
--- a/python/sglang/srt/models/qwen3_next.py
+++ b/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,