[Refactor] Rename CustomOp -> MultiPlatformOp (#16175)

python/sglang/srt/layers/activation.py

@@ -22,13 +22,13 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PretrainedConfig
 
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import (
     divide,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -59,7 +59,7 @@ if is_npu():
 logger = logging.getLogger(__name__)
 
 
-class SiluAndMul(CustomOp):
+class SiluAndMul(MultiPlatformOp):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if get_global_server_args().rl_on_policy_target is not None:
@@ -95,7 +95,7 @@ class SiluAndMul(CustomOp):
         return out
 
 
-class GeluAndMul(CustomOp):
+class GeluAndMul(MultiPlatformOp):
     def __init__(self, approximate="tanh"):
         super().__init__()
         self.approximate = approximate
@@ -140,7 +140,7 @@ class GeluAndMul(CustomOp):
         return y_npu
 
 
-class NewGELU(CustomOp):
+class NewGELU(MultiPlatformOp):
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         c = math.sqrt(2.0 / math.pi)
         return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0))))
@@ -161,7 +161,7 @@ class ReLU2(nn.Module):
         return x * x
 
 
-class QuickGELU(CustomOp):
+class QuickGELU(MultiPlatformOp):
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         return x * torch.sigmoid(1.702 * x)
 
@@ -177,7 +177,7 @@ class QuickGELU(CustomOp):
         return torch_npu.npu_fast_gelu(x)
 
 
-class XIELU(CustomOp):
+class XIELU(MultiPlatformOp):
     """
     Applies the xIELU activation function introduced in https://arxiv.org/abs/2411.13010
     If the user has installed the nickjbrowning/XIELU, we import xIELU CUDA

python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py

@@ -2,7 +2,6 @@ from typing import Union
 
 import torch
 
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed.communication_op import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
@@ -12,11 +11,12 @@ from sglang.srt.distributed.parallel_state import (
     get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.attention.fla.layernorm_gated import rms_norm_gated
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_loader.weight_utils import sharded_weight_loader
 from sglang.srt.utils.common import set_weight_attrs
 
 
-class Mixer2RMSNormGated(CustomOp):
+class Mixer2RMSNormGated(MultiPlatformOp):
     def __init__(
         self,
         full_hidden_size: int,

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -6,8 +6,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 import torch
 from einops import rearrange
 
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.layers.layernorm import LayerNorm
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
 
 global _use_multi_stream
@@ -93,7 +93,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
     return hadamard_transform(x, scale=hidden_size**-0.5)
 
 
-class Indexer(CustomOp):
+class Indexer(MultiPlatformOp):
     def __init__(
         self,
         hidden_size: int,

python/sglang/srt/layers/layernorm.py

@@ -24,7 +24,7 @@ from sglang.srt.batch_invariant_ops import (
     is_batch_invariant_mode_enabled,
     rms_norm_batch_invariant,
 )
-from sglang.srt.custom_op import CustomOp
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -77,7 +77,7 @@ if _is_npu:
     import torch_npu
 
 
-class RMSNorm(CustomOp):
+class RMSNorm(MultiPlatformOp):
     def __init__(
         self,
         hidden_size: int,
@@ -285,7 +285,7 @@ class RMSNorm(CustomOp):
         return self.forward(x, residual)
 
 
-class LayerNorm(CustomOp):
+class LayerNorm(MultiPlatformOp):
     def __init__(
         self,
         hidden_size: int,
@@ -357,7 +357,7 @@ class LayerNorm(CustomOp):
             return self.forward_native(x)
 
 
-class GemmaRMSNorm(CustomOp):
+class GemmaRMSNorm(MultiPlatformOp):
     def __init__(
         self,
         hidden_size: int,
@@ -444,7 +444,7 @@ class GemmaRMSNorm(CustomOp):
         return self._forward_impl(x, residual)
 
 
-class Gemma3RMSNorm(CustomOp):
+class Gemma3RMSNorm(MultiPlatformOp):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
         self.eps = eps

python/sglang/srt/layers/model_parallel.py

@@ -88,7 +88,7 @@ class RowwiseParallelMaybeWait(RowwiseParallel):
     A version of RowwiseParallel that waits for the output (establish dependency
     between comm stream and compute stream in CUDA sense) before going into the
     next op. This is needed to workaround the current interaction between
-    AsyncCollectiveTensor and custom ops, such as `class RMSNorm(CustomOp)`.
+    AsyncCollectiveTensor and multi-platform ops, such as `RMSNorm`.
     """
 
     def _partition_linear_fn(self, name, module, device_mesh):

python/sglang/srt/layers/moe/topk.py

@@ -35,7 +35,6 @@ try:
 except ImportError:
     pass
 
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
@@ -49,6 +48,7 @@ from sglang.srt.eplb.expert_location_dispatch import (
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
 from sglang.srt.layers.moe import get_moe_runner_backend
 from sglang.srt.layers.moe.routed_experts_capturer import get_global_experts_capturer
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -191,7 +191,7 @@ class BypassedTopKOutput(NamedTuple):
 # -------------------------------- TopK ---------------------------------------
 
 
-class TopK(CustomOp):
+class TopK(MultiPlatformOp):
     """
     Parameters:
     --top_k: The all number of top experts selected per token, including the fused shared expert(s).

python/sglang/srt/layers/quantization/unquant.py

@@ -6,7 +6,6 @@ import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
 from sglang.srt.layers.moe import (
     MoeRunner,
@@ -20,6 +19,7 @@ from sglang.srt.layers.quantization.base_config import (
     LinearMethodBase,
     QuantizeMethodBase,
 )
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -143,7 +143,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
         return F.linear(x, layer.weight, bias)
 
 
-class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
+class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
     """MoE method without quantization."""
 
     def __init__(

python/sglang/srt/layers/rotary_embedding.py

@@ -11,7 +11,7 @@ import torch.nn as nn
 import triton
 import triton.language as tl
 
-from sglang.srt.custom_op import CustomOp
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -89,7 +89,7 @@ def _apply_rotary_emb(
         return torch.stack((o1, o2), dim=-1).flatten(-2)
 
 
-class RotaryEmbedding(CustomOp):
+class RotaryEmbedding(MultiPlatformOp):
     """Original rotary positional embedding."""
 
     def __init__(
@@ -2298,7 +2298,7 @@ class MRotaryEmbedding(RotaryEmbedding):
         return llm_pos_ids
 
 
-class DualChunkRotaryEmbedding(CustomOp):
+class DualChunkRotaryEmbedding(MultiPlatformOp):
     """Rotary positional embedding for Dual Chunk Attention."""
 
     def __init__(

python/sglang/srt/layers/utils/__init__.py

@@ -1,2 +1,3 @@
 # Temp workaround, make layer utils more fine-grained later
 from sglang.srt.layers.utils.common import *
+from sglang.srt.layers.utils.multi_platform import MultiPlatformOp

python/sglang/srt/layers/utils/multi_platform.py

@@ -1,8 +1,4 @@
-"""
+from typing import Callable
-The definition of CustomOps for multi hardware dispatching.
-
-TODO: Move this to python/sglang/srt/layers/custom_op.py
-"""
 
 from torch import nn
 
@@ -23,10 +19,10 @@ _is_npu = is_npu()
 _is_xpu = is_xpu()
 
 
-class CustomOp(nn.Module):
+class MultiPlatformOp(nn.Module):
     def __init__(self):
         super().__init__()
-        self._forward_method = self.dispatch_forward()
+        self._forward_method: Callable = self.dispatch_forward()
 
         # States for torch.compile
         self._original_forward_method = None

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -30,7 +30,6 @@ from torch.profiler import ProfilerActivity, profile
 
 from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
 from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     set_graph_pool_id,
@@ -52,6 +51,7 @@ from sglang.srt.layers.dp_attention import (
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
 from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardBatch,
@@ -128,7 +128,7 @@ def freeze_gc(enable_cudagraph_gc: bool):
 
 def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
     for sub in model._modules.values():
-        if isinstance(sub, CustomOp):
+        if isinstance(sub, MultiPlatformOp):
             if reverse:
                 sub.leave_torch_compile()
             else:

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -33,7 +33,6 @@ from sglang.srt.compilation.piecewise_context_manager import (
     set_forward_context,
     set_pcg_capture_stream,
 )
-from sglang.srt.custom_op import CustomOp
 from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     set_graph_pool_id,
@@ -49,6 +48,7 @@ from sglang.srt.layers.dp_attention import (
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.layers.pooler import EmbeddingPoolerOutput
+from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardBatch,
@@ -83,7 +83,7 @@ def freeze_gc(enable_cudagraph_gc: bool):
 
 def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
     for sub in model._modules.values():
-        if isinstance(sub, CustomOp):
+        if isinstance(sub, MultiPlatformOp):
             if reverse:
                 sub.leave_torch_compile()
             else: