chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Move custom_ops under layers; move _custom_ops.py → custom_all_reduce_ops.py (#14326)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-12-03 10:33:37 -08:00
committed by GitHub
parent 20aad5b5ab
commit 46d7b35ec7
7 changed files with 4 additions and 313 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
View File

@@ -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,

0
python/sglang/srt/_custom_ops.py → python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py
View File

2
python/sglang/srt/distributed/device_communicators/pymscclpp.py
View File

@@ -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__)

2
python/sglang/srt/distributed/device_communicators/quick_all_reduce.py
View File

@@ -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,

98
python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py
View File

@@ -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

211
python/sglang/srt/operations_strategy.py
View File

@@ -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(),
],
)

2
test/manual/test_quick_allreduce.py
View File

@@ -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,