[6/N] MoE Refactor: Cleanup MoE-related configs (#8849)

python/sglang/srt/models/gpt_oss.py

@@ -40,7 +40,6 @@ from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (
     get_attention_tp_rank,
     get_attention_tp_size,
-    get_local_attention_dp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -50,9 +49,10 @@ from sglang.srt.layers.linear import (
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.moe import get_moe_a2a_backend
 from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.moe.topk import TopK
-from sglang.srt.layers.moe.utils import DeepEPMode
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.fp8_utils import dequant_mxfp4
 from sglang.srt.layers.radix_attention import RadixAttention
@@ -110,16 +110,13 @@ class GptOssSparseMoeBlock(nn.Module):
         self.tp_size = get_tensor_model_parallel_world_size()
         self.layer_id = layer_id
         self.activation = config.hidden_act
-        self.activation_alpha = getattr(config, "hidden_act_alpha", 1.702)
-        self.swiglu_limit = config.swiglu_limit
+        self.gemm1_alpha = getattr(config, "hidden_act_alpha", 1.702)
+        self.gemm1_clamp_limit = config.swiglu_limit
 
-        if global_server_args_dict["enable_flashinfer_mxfp4_moe"]:
-            self.topk = None
-        else:
-            self.topk = TopK(
-                top_k=config.num_experts_per_tok,
-                renormalize=True,
-            )
+        self.topk = TopK(
+            top_k=config.num_experts_per_tok,
+            renormalize=True,
+        )
 
         self.top_k = config.num_experts_per_tok
         experts_type = get_moe_impl_class()
@@ -129,11 +126,9 @@ class GptOssSparseMoeBlock(nn.Module):
                 quant_config.get_name() if quant_config is not None else None
             )
             extra_kwargs = {
-                "enable_flashinfer_cutlass_moe": global_server_args_dict[
-                    "enable_flashinfer_cutlass_moe"
-                ],
                 # for moe gate_up_proj and down_proj and their bias loading
-                "use_weight_loader_fused": quant_config_name != "mxfp4",
+                "use_weight_loader_fused": quant_config_name
+                != "mxfp4"
             }
         self.experts = experts_type(
             num_experts=config.num_local_experts
@@ -144,15 +139,10 @@ class GptOssSparseMoeBlock(nn.Module):
             intermediate_size=config.intermediate_size,
             quant_config=quant_config,
             activation=self.activation,
-            activation_alpha=self.activation_alpha,
-            swiglu_limit=self.swiglu_limit,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_clamp_limit=self.gemm1_clamp_limit,
             with_bias=True,
             prefix=add_prefix("experts", prefix),
-            **(
-                dict(deepep_mode=DeepEPMode[global_server_args_dict["deepep_mode"]])
-                if global_server_args_dict["moe_a2a_backend"].is_deepep()
-                else {}
-            ),
             **extra_kwargs,
         )
 
@@ -171,7 +161,7 @@ class GptOssSparseMoeBlock(nn.Module):
         forward_batch: Optional[ForwardBatch] = None,
         should_allreduce_fusion: bool = False,
     ) -> torch.Tensor:
-        if not global_server_args_dict["moe_a2a_backend"].is_deepep():
+        if not get_moe_a2a_backend().is_deepep():
             return self.forward_normal(hidden_states, should_allreduce_fusion)
         else:
             raise Exception("forward_deepep branch not implemented yet")
@@ -189,17 +179,10 @@ class GptOssSparseMoeBlock(nn.Module):
         should_allreduce_fusion: bool = False,
     ) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
-        hidden_states = hidden_states.view(-1, hidden_dim)
 
-        # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.router(hidden_states)
-
-        kwargs = {"hidden_states": hidden_states}
-        if self.topk is not None:
-            kwargs["topk_output"] = self.topk(hidden_states, router_logits)
-        else:
-            kwargs["topk_output"] = (self.top_k, router_logits)
-        final_hidden_states = self.experts(**kwargs)
+        topk_output = self.topk(hidden_states, router_logits)
+        final_hidden_states = self.experts(hidden_states, topk_output)
 
         if self.tp_size > 1 and not should_allreduce_fusion:
             final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
@@ -436,7 +419,6 @@ class GptOssDecoderLayer(nn.Module):
 
         self.attn_tp_size = get_attention_tp_size()
         self.attn_tp_rank = get_attention_tp_rank()
-        self.local_dp_size = get_local_attention_dp_size()
 
         # GptOss all layers are sparse and have no nextn now
         self.is_layer_sparse = True
@@ -1060,7 +1042,7 @@ class GptOssForCausalLM(nn.Module):
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
         ]
-        expert_params_mapping = get_moe_impl_class().make_expert_params_mapping_fused(
+        expert_params_mapping = FusedMoE.make_expert_params_mapping_fused(
             ckpt_gate_up_proj_name="gate_up_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_gate_up_proj_bias_name="gate_up_proj_bias",