diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index e7d5a67cc..aa38b8d74 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -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) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py index 8737c26b7..8ce947d43 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -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(