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

MoE: Skip SiLU/GELU activation for masked experts (#15539)

Browse Source 
Co-authored-by: Runkai Tao <rt572@physics.rutgers.edu>
This commit is contained in:
yuchengz816-bot
2025-12-22 22:08:31 -05:00
committed by GitHub
parent 82f1d6157f
commit 061f41affc
2 changed files with 132 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

@@ -25,6 +25,7 @@ from sglang.srt.utils import (
from .fused_moe_triton_config import get_config_dtype_str, try_get_optimal_moe_config
from .fused_moe_triton_kernels import (
act_and_mul_triton,
invoke_fused_moe_kernel,
moe_sum_reduce_triton,
support_tensor_descriptor,
@@ -544,6 +545,7 @@ def fused_experts_impl(
c_sorted=down_moe_use_tma,
filter_expert=filter_expert,
)
# Activation function with multiplication
if activation == "silu" and is_gated:
if gemm1_alpha is not None:
@@ -553,8 +555,19 @@ def fused_experts_impl(
gemm1_alpha,
gemm1_limit,
)
elif _is_cuda or _is_hip:
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
elif _is_hip or _is_cuda:
if not filter_expert:
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
else:
act_and_mul_triton(
intermediate_cache1.view(-1, N),
intermediate_cache2,
config,
topk_ids,
expert_ids,
down_moe_use_tma,
activation,
)
else:
vllm_ops.silu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)
@@ -562,8 +575,19 @@ def fused_experts_impl(
elif activation == "gelu" and is_gated:
assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu"
assert gemm1_limit is None, "gemm1_limit is not supported for gelu"
if _is_cuda or _is_hip:
gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
if _is_hip or _is_cuda:
if not filter_expert:
gelu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
else:
act_and_mul_triton(
intermediate_cache1.view(-1, N),
intermediate_cache2,
config,
topk_ids,
expert_ids,
down_moe_use_tma,
activation,
)
else:
vllm_ops.gelu_and_mul(
intermediate_cache2, intermediate_cache1.view(-1, N)

104
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py
View File

@@ -790,6 +790,110 @@ def invoke_fused_moe_kernel(
)
@triton.jit
def tanh(x):
return 2 * tl.sigmoid(2 * x) - 1
@triton.jit
def _apply_activation(x, ACTIVATION_TYPE: tl.constexpr):
"""
Apply activation function based on compile-time constant.
Args:
x: Input tensor (converted to float32 inside)
ACTIVATION_TYPE: Compile-time constant string ("silu" or "gelu")
Returns:
Activated output in the same dtype as input
"""
x = x.to(tl.float32)
if ACTIVATION_TYPE == "silu":
return x * tl.sigmoid(x)
else:
kAlpha = 0.7978845608028654
return 0.5 * x * (1 + tanh(kAlpha * (x + 0.044715 * x * x * x)))
@triton.jit
def act_and_mul_kernel(
gateup_output,
down_input,
hidden_size,
expert_ids_ptr,
expert_step: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
ACTIVATION_TYPE: tl.constexpr,
):
"""
Unified activation and multiply kernel that handles both sorted and unsorted routing,
and both SiLU and GELU activations using compile-time constants.
"""
InDtype = gateup_output.dtype.element_ty
OutDtype = down_input.dtype.element_ty
half_hidden_size = hidden_size // 2
pid = tl.program_id(0)
expert_id = tl.load(expert_ids_ptr + pid // expert_step)
if expert_id == -1:
return
gateup_output_ptr = gateup_output + pid * hidden_size
down_input_ptr = down_input + pid * half_hidden_size
gate_output_ptr = gateup_output_ptr
up_output_ptr = gateup_output_ptr + half_hidden_size
for start_offset in tl.range(0, half_hidden_size, BLOCK_SIZE):
offset = start_offset + tl.arange(0, BLOCK_SIZE)
mask = offset < half_hidden_size
gate_output = tl.load(gate_output_ptr + offset, mask=mask)
up_output = tl.load(up_output_ptr + offset, mask=mask)
gate_output_activated = _apply_activation(gate_output, ACTIVATION_TYPE)
gate_output_activated = gate_output_activated.to(InDtype)
act_mul_output = gate_output_activated * up_output
act_mul_output = act_mul_output.to(OutDtype)
tl.store(down_input_ptr + offset, act_mul_output, mask=mask)
def act_and_mul_triton(
gateup_output: torch.Tensor,
down_input: torch.Tensor,
config: Dict[str, Any],
topk_ids: Optional[torch.Tensor] = None,
expert_ids: Optional[torch.Tensor] = None,
down_moe_use_tma: bool = False,
activation: str = "silu",
) -> None:
"""
Args:
gateup_output: Input tensor containing gate and up outputs concatenated
down_input: Output tensor for the result
config: Configuration dictionary with BLOCK_SIZE_M and BLOCK_SIZE_N
topk_ids: Expert IDs for unsorted routing (used when down_moe_use_tma=False)
expert_ids: Expert IDs for sorted routing (used when down_moe_use_tma=True)
down_moe_use_tma: Whether to use sorted routing layout
activation: Activation type ("silu" or "gelu")
"""
grid = (down_input.shape[0],)
hidden_size = gateup_output.shape[1]
expert_ids_row = topk_ids.view(-1) if not down_moe_use_tma else expert_ids
expert_step = 1 if not down_moe_use_tma else config["BLOCK_SIZE_M"]
act_and_mul_kernel[grid](
gateup_output,
down_input,
hidden_size,
expert_ids_row,
expert_step,
BLOCK_SIZE=512,
ACTIVATION_TYPE=activation,
)
# _moe_sum_reduce_kernel kernel modified from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/moe_sum_reduce.py
@triton.jit
def _moe_sum_reduce_kernel(