diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 6db49012d..8fde6b80f 100644 --- a/python/sglang/srt/layers/activation.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py index 271394c8e..547b4aa99 100644 --- a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 86a740f29..51b9d4ff3 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/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, diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 68f94a090..8655b02a3 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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 diff --git a/python/sglang/srt/layers/model_parallel.py b/python/sglang/srt/layers/model_parallel.py index 53c8b622e..d08754f32 100644 --- a/python/sglang/srt/layers/model_parallel.py +++ b/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): diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 65a593991..7ecaa5cf4 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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). diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 630b60068..d625f8a7f 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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__( diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index e9f375138..5cbe8f179 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/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__( diff --git a/python/sglang/srt/layers/utils/__init__.py b/python/sglang/srt/layers/utils/__init__.py index 7881b2b49..e3101d534 100644 --- a/python/sglang/srt/layers/utils/__init__.py +++ b/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 diff --git a/python/sglang/srt/custom_op.py b/python/sglang/srt/layers/utils/multi_platform.py similarity index 93% rename from python/sglang/srt/custom_op.py rename to python/sglang/srt/layers/utils/multi_platform.py index 90cfa8e12..33033d4c8 100644 --- a/python/sglang/srt/custom_op.py +++ b/python/sglang/srt/layers/utils/multi_platform.py @@ -1,8 +1,4 @@ -""" -The definition of CustomOps for multi hardware dispatching. - -TODO: Move this to python/sglang/srt/layers/custom_op.py -""" +from typing import Callable 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 diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 6ef74ca80..e6391a748 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/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: diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 556ddcfdb..1087b3b7f 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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: