From 46d7b35ec76bc8ca22d79037ab7d776342dc7e5e Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Wed, 3 Dec 2025 10:33:37 -0800 Subject: [PATCH] =?UTF-8?q?Move=20custom=5Fops=20under=20layers;=20move=20?= =?UTF-8?q?=5Fcustom=5Fops.py=20=E2=86=92=20custom=5Fall=5Freduce=5Fops.py?= =?UTF-8?q?=20(#14326)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../device_communicators/custom_all_reduce.py | 2 +- .../custom_all_reduce_ops.py} | 0 .../device_communicators/pymscclpp.py | 2 +- .../device_communicators/quick_all_reduce.py | 2 +- .../layers/attention/mamba/ops/mamba_ssm.py | 98 -------- python/sglang/srt/operations_strategy.py | 211 ------------------ test/manual/test_quick_allreduce.py | 2 +- 7 files changed, 4 insertions(+), 313 deletions(-) rename python/sglang/srt/{_custom_ops.py => distributed/device_communicators/custom_all_reduce_ops.py} (100%) delete mode 100644 python/sglang/srt/operations_strategy.py diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 5fecab5e6..fa0230fb6 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -10,7 +10,7 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup -from sglang.srt import _custom_ops as ops +import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import ( gpu_p2p_access_check, diff --git a/python/sglang/srt/_custom_ops.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py similarity index 100% rename from python/sglang/srt/_custom_ops.py rename to python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py diff --git a/python/sglang/srt/distributed/device_communicators/pymscclpp.py b/python/sglang/srt/distributed/device_communicators/pymscclpp.py index 78e1318fa..e45093c78 100644 --- a/python/sglang/srt/distributed/device_communicators/pymscclpp.py +++ b/python/sglang/srt/distributed/device_communicators/pymscclpp.py @@ -10,7 +10,7 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup, ReduceOp -from sglang.srt import _custom_ops as ops +import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops from sglang.srt.utils import is_hip logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py b/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py index 0113f02c3..f9d51246e 100644 --- a/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/quick_all_reduce.py @@ -10,7 +10,7 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup -from sglang.srt import _custom_ops as ops +import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import ( is_full_nvlink, is_weak_contiguous, diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py index 69a1ff9fb..5aca69397 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +++ b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py @@ -11,8 +11,6 @@ import triton import triton.language as tl from packaging import version -from sglang.srt import _custom_ops as ops - PAD_SLOT_ID = -1 TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0") @@ -344,99 +342,3 @@ def selective_state_update( BLOCK_SIZE_M, num_warps=num_warps, ) - - -def selective_scan_fn( - u, - ssm_states, - delta, - A, - B, - C, - D=None, - z=None, - delta_bias=None, - delta_softplus=False, - query_start_loc=None, - cache_indices=None, - has_initial_state=None, - pad_slot_id=PAD_SLOT_ID, -) -> torch.Tensor: - """ - u: (dim, total_length) for varlen or (batch, dim, seqlen) - applies changes in place. - ssm_states: (batch, dim, dstate) or (batch, nheads, dim, dstate) - applies changes in place. - delta: (dim, total_length) for varlen or (batch, dim, seqlen) - A: (dim, dstate) - B: (ngroups, dstate, total_length) for varlen or - (batch,ngroups,dstate,seqlen) - C: (ngroups, dstate, total_length) for varlen or - (batch,ngroups,dstate,seqlen) - D: (dim,) - z: (dim, total_length) for varlen or (batch, dim, seqlen) - dt_bias: (dim,) or (dim) - query_start_loc: (batch + 1) int32 - The cumulative sequence lengths of the sequences in - the batch, used to index into sequence. prepended with 0. - for example: query_start_loc = torch.Tensor([0,10,16,17]), - x.shape=(dim,17) - cache_indices: (batch) int32 - A tensor with each cell is a correspondent - input and output ssm_state index - has_initial_state: (batch) bool - A tensor populated with ones and zeros, - indicate if the ssm_state at the corresponding index should be - used as initial state. Not providing argument assumes - there's no initial state - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padding entries - that will not be processed, - for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id] - in this case, the kernel will not process entries at indices 0 and 3 - returns - output: (dim, total_length) for varlen or (batch, dim, seqlen) - supports inplace replacement - """ - if u.stride(-1) != 1: - u = u.contiguous() - if delta.stride(-1) != 1: - delta = delta.contiguous() - if D is not None: - D = D.contiguous() - if B.stride(-1) != 1: - B = B.contiguous() - if C.stride(-1) != 1: - C = C.contiguous() - if z is not None and z.stride(-1) != 1: - z = z.contiguous() - if B.dim() == 3 and query_start_loc is None: - B = B.unsqueeze(1) - if B.dim() == 2 and query_start_loc is not None: - B = B.unsqueeze(0) - if C.dim() == 3 and query_start_loc is None: - C = C.unsqueeze(1) - if C.dim() == 2 and query_start_loc is not None: - C = C.unsqueeze(0) - - ops.selective_scan_fwd( - u, - delta, - A, - B, - C, - D, - z, - delta_bias, - delta_softplus, - query_start_loc, - cache_indices, - has_initial_state, - ssm_states, - pad_slot_id, - ) - - if z is None: - return delta # output written inplace to delta - else: - return z # output written inplace to z diff --git a/python/sglang/srt/operations_strategy.py b/python/sglang/srt/operations_strategy.py deleted file mode 100644 index 152e4874d..000000000 --- a/python/sglang/srt/operations_strategy.py +++ /dev/null @@ -1,211 +0,0 @@ -from dataclasses import dataclass -from typing import List, Optional - -import torch - -from sglang.srt.batch_overlap import operations -from sglang.srt.batch_overlap.operations import Operation -from sglang.srt.layers.moe.token_dispatcher import DeepEPConfig -from sglang.srt.model_executor.forward_batch_info import ForwardMode - - -@dataclass -class OperationsStrategy: - operations: List[Operation] - deep_gemm_num_sms: Optional[int] = None - tbo_delta_stages: Optional[int] = None - - @classmethod - def concat(cls, items: List["OperationsStrategy"]) -> "OperationsStrategy": - return OperationsStrategy( - operations=[x for item in items for x in item.operations], - deep_gemm_num_sms=_assert_all_same( - [item.deep_gemm_num_sms for item in items] - ), - tbo_delta_stages=_assert_all_same( - [item.tbo_delta_stages for item in items] - ), - ) - - @staticmethod - def init_new_tbo( - layers: torch.nn.ModuleList, - forward_mode: ForwardMode, - ) -> "OperationsStrategy": - layer_name = layers[0].__class__.__name__ - if layer_name == "DeepseekV2DecoderLayer": - return OperationsStrategy.concat( - [ - _compute_moe_deepseek_layer_operations_strategy_tbo( - layer, forward_mode - ) - for layer in layers - ] - ) - elif layer_name == "Qwen3MoeDecoderLayer": - return OperationsStrategy.concat( - [ - _compute_moe_qwen3_layer_operations_strategy_tbo( - layer, forward_mode - ) - for layer in layers - ] - ) - else: - raise NotImplementedError - - -def _assert_all_same(items: List): - assert all(item == items[0] for item in items) - return items[0] - - -# -------------------------------- Strategy for DeepSeek --------------------------------------- - - -# TODO can refactor to make it more fancy if we have more complex strategies -def _compute_moe_deepseek_layer_operations_strategy_tbo( - layer: torch.nn.Module, - forward_mode: ForwardMode, -) -> OperationsStrategy: - assert layer.is_layer_sparse, "dense layer TBO not yet implemented" - if forward_mode == ForwardMode.EXTEND: - return _compute_moe_deepseek_blog_prefill(layer) - elif ( - forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY - ): - return _compute_moe_deepseek_blog_decode(layer) - else: - raise NotImplementedError(f"Unsupported {forward_mode=}") - - -def _compute_moe_deepseek_blog_prefill(layer): - device_properties = torch.cuda.get_device_properties(device="cuda") - total_num_sms = device_properties.multi_processor_count - deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms - - return OperationsStrategy( - deep_gemm_num_sms=deep_gemm_num_sms, - tbo_delta_stages=0, - operations=[ - layer.op_comm_prepare_attn, - layer.self_attn.op_prepare, - layer.self_attn.op_core, - layer.op_comm_prepare_mlp, - layer.mlp.op_gate, - layer.mlp.op_select_experts, - layer.mlp.op_dispatch_a, - operations.YieldOperation(), - layer.mlp.op_dispatch_b, - layer.mlp.op_experts, - layer.mlp.op_combine_a, - operations.YieldOperation(), - layer.mlp.op_shared_experts, - layer.mlp.op_combine_b, - layer.mlp.op_output, - layer.op_comm_postprocess_layer, - ], - ) - - -def _compute_moe_deepseek_blog_decode(layer): - return OperationsStrategy( - deep_gemm_num_sms=None, - tbo_delta_stages=2, - operations=[ - layer.op_comm_prepare_attn, - layer.self_attn.op_prepare, - operations.YieldOperation(), - layer.self_attn.op_core, - layer.op_comm_prepare_mlp, - layer.mlp.op_gate, - layer.mlp.op_select_experts, - operations.YieldOperation(), - layer.mlp.op_dispatch_a, - layer.mlp.op_shared_experts, - operations.YieldOperation(), - layer.mlp.op_dispatch_b, - layer.mlp.op_experts, - layer.mlp.op_combine_a, - operations.YieldOperation(), - layer.mlp.op_combine_b, - operations.YieldOperation(), - layer.mlp.op_output, - layer.op_comm_postprocess_layer, - ], - ) - - -# -------------------------------- Strategy for Qwen3 --------------------------------------- - - -# TODO: unstable, current strategy is almost the same as DeepSeek, keep redundant code here for -# convenience to adjust strategy -def _compute_moe_qwen3_layer_operations_strategy_tbo( - layer: torch.nn.Module, - forward_mode: ForwardMode, -) -> OperationsStrategy: - assert layer.is_layer_sparse, "qwen3 moe only support sparse layers" - if forward_mode == ForwardMode.EXTEND: - return _compute_moe_qwen3_prefill(layer) - elif ( - forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY - ): - return _compute_moe_qwen3_decode(layer) - else: - raise NotImplementedError(f"Unsupported {forward_mode=}") - - -def _compute_moe_qwen3_prefill(layer): - device_properties = torch.cuda.get_device_properties(device="cuda") - total_num_sms = device_properties.multi_processor_count - deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms - - return OperationsStrategy( - deep_gemm_num_sms=deep_gemm_num_sms, - tbo_delta_stages=0, - operations=[ - layer.op_comm_prepare_attn, - layer.self_attn.op_prepare, - layer.self_attn.op_core, - layer.op_comm_prepare_mlp, - layer.mlp.op_gate, - layer.mlp.op_select_experts, - layer.mlp.op_dispatch_a, - operations.YieldOperation(), - layer.mlp.op_dispatch_b, - layer.mlp.op_experts, - layer.mlp.op_combine_a, - operations.YieldOperation(), - layer.mlp.op_combine_b, - layer.mlp.op_output, - layer.op_comm_postprocess_layer, - ], - ) - - -def _compute_moe_qwen3_decode(layer): - return OperationsStrategy( - deep_gemm_num_sms=None, - tbo_delta_stages=2, - operations=[ - layer.op_comm_prepare_attn, - layer.self_attn.op_prepare, - operations.YieldOperation(), - layer.self_attn.op_core, - layer.op_comm_prepare_mlp, - layer.mlp.op_gate, - layer.mlp.op_select_experts, - operations.YieldOperation(), - layer.mlp.op_dispatch_a, - operations.YieldOperation(), - layer.mlp.op_dispatch_b, - layer.mlp.op_experts, - layer.mlp.op_combine_a, - operations.YieldOperation(), - layer.mlp.op_combine_b, - layer.mlp.op_output, - layer.op_comm_postprocess_layer, - operations.YieldOperation(), - ], - ) diff --git a/test/manual/test_quick_allreduce.py b/test/manual/test_quick_allreduce.py index 8789426a3..42bd4c9c0 100644 --- a/test/manual/test_quick_allreduce.py +++ b/test/manual/test_quick_allreduce.py @@ -9,7 +9,7 @@ import ray import torch import torch.distributed as dist -from sglang.srt import _custom_ops as ops +import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops from sglang.srt.distributed import init_distributed_environment from sglang.srt.distributed.communication_op import ( # noqa tensor_model_parallel_all_reduce,