Revert "[DeepEP] Reduce routed scaling overhead"

This reverts commit 2a4fc7e.

Author: yuleil

python/sglang/srt/layers/moe/ep_moe/kernels.py

@@ -707,71 +707,3 @@ def grouped_gemm_triton(
         **config,
     )
     return c
-
-
-@triton.jit
-def _masked_scale_kernel(
-    in_ptr,
-    in_stride_0,
-    in_stride_1,
-    out_ptr,
-    out_stride_0,
-    out_stride_1,
-    masked_m_ptr,
-    scale: float,
-    D: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-    BLOCK_NUM_PER_EXPERT: tl.constexpr,
-):
-    pid_expert = tl.program_id(0)
-    pid_token = tl.program_id(1)
-    pid_dim = tl.program_id(2)
-
-    TOKENS_CUR_EXPERT = tl.load(masked_m_ptr + pid_expert).to(tl.int32)
-
-    offs_in_d = pid_dim * BLOCK_D + tl.arange(0, BLOCK_D)
-    mask = offs_in_d < D
-
-    in_ptr_offs = in_ptr + pid_expert * in_stride_0 + offs_in_d
-    out_ptr_offs = out_ptr + pid_expert * out_stride_0 + offs_in_d
-
-    for token_index in tl.range(pid_token, TOKENS_CUR_EXPERT, BLOCK_NUM_PER_EXPERT):
-        v = tl.load(in_ptr_offs + token_index * in_stride_1, mask)
-        tl.store(out_ptr_offs + token_index * out_stride_1, v * scale, mask)
-
-
-def masked_scale(
-    x: torch.Tensor, masked_m: torch.Tensor, scale: float, out: torch.Tensor = None
-):
-    assert x.stride(-1) == 1
-
-    if out is None:
-        out = torch.empty_like(x)
-
-    expert_num = len(masked_m)
-
-    if expert_num < 4:
-        BLOCK_NUM_PER_EXPERT = 64
-    else:
-        BLOCK_NUM_PER_EXPERT = 32
-
-    BLOCK_D = 512
-
-    grid = (len(masked_m), BLOCK_NUM_PER_EXPERT, triton.cdiv(x.size(-1), BLOCK_D))
-
-    _masked_scale_kernel[grid](
-        x,
-        x.stride(0),
-        x.stride(1),
-        out,
-        out.stride(0),
-        out.stride(1),
-        masked_m,
-        scale,
-        x.size(-1),
-        BLOCK_D=BLOCK_D,
-        BLOCK_NUM_PER_EXPERT=BLOCK_NUM_PER_EXPERT,
-        num_warps=1,
-        num_stages=6,
-    )
-    return out

python/sglang/srt/models/deepseek_v2.py

@@ -49,7 +49,6 @@
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
-from sglang.srt.layers.moe.ep_moe.kernels import masked_scale
 from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, EPMoE
 from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
 from sglang.srt.layers.moe.topk import select_experts
@@ -338,24 +337,17 @@ def forward_deepep(
                 topk_weights,
                 forward_mode=forward_mode,
             )
-        final_hidden_states = self.experts(
-            hidden_states=hidden_states,
-            reorder_topk_ids=reorder_topk_ids,
-            seg_indptr=seg_indptr,
-            masked_m=masked_m,
-            expected_m=expected_m,
-            forward_mode=forward_mode,
-        )
-        if self.ep_size > 1 and masked_m is not None:
-            final_hidden_states = masked_scale(
-                final_hidden_states,
-                masked_m,
-                self.routed_scaling_factor,
-                final_hidden_states,
+        final_hidden_states = (
+            self.experts(
+                hidden_states=hidden_states,
+                reorder_topk_ids=reorder_topk_ids,
+                seg_indptr=seg_indptr,
+                masked_m=masked_m,
+                expected_m=expected_m,
+                forward_mode=forward_mode,
             )
-        else:
-            final_hidden_states *= self.routed_scaling_factor
-
+            * self.routed_scaling_factor
+        )
         if self.ep_size > 1:
             final_hidden_states = self.deepep_dispatcher.combine(
                 final_hidden_states,