[6/N] MoE Refactor: Cleanup MoE-related configs (#8849)
This commit is contained in:
Cheng Wan
@@ -50,7 +50,6 @@ from sglang.srt.layers.communicator import (
|
||||
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
|
||||
@@ -61,9 +60,10 @@ from sglang.srt.layers.linear import (
|
||||
RowParallelLinear,
|
||||
)
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
from sglang.srt.layers.moe import get_deepep_mode, get_moe_a2a_backend
|
||||
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, 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 should_use_flashinfer_trtllm_moe
|
||||
from sglang.srt.layers.quantization import deep_gemm_wrapper
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.layers.quantization.fp8_kernel import (
|
||||
@@ -336,30 +336,6 @@ class DeepseekV2MoE(nn.Module):
|
||||
quant_config=quant_config,
|
||||
routed_scaling_factor=self.routed_scaling_factor,
|
||||
prefix=add_prefix("experts", prefix),
|
||||
**(
|
||||
dict(deepep_mode=global_server_args_dict["deepep_mode"])
|
||||
if global_server_args_dict["moe_a2a_backend"].is_deepep()
|
||||
else {}
|
||||
),
|
||||
# Additional args for FusedMoE
|
||||
**(
|
||||
dict(
|
||||
enable_flashinfer_cutlass_moe=True,
|
||||
)
|
||||
if global_server_args_dict["enable_flashinfer_cutlass_moe"]
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
dict(
|
||||
renormalize=config.norm_topk_prob,
|
||||
use_grouped_topk=True,
|
||||
num_expert_group=config.n_group,
|
||||
topk_group=config.topk_group,
|
||||
correction_bias=self.gate.e_score_correction_bias,
|
||||
)
|
||||
if should_use_flashinfer_trtllm_moe()
|
||||
else {}
|
||||
),
|
||||
)
|
||||
|
||||
self.shared_experts_is_int8 = False
|
||||
@@ -377,7 +353,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
prefix=add_prefix("shared_experts", prefix),
|
||||
**(
|
||||
dict(tp_rank=0, tp_size=1)
|
||||
if global_server_args_dict["moe_a2a_backend"].is_deepep()
|
||||
if get_moe_a2a_backend().is_deepep()
|
||||
else {}
|
||||
),
|
||||
)
|
||||
@@ -407,7 +383,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
|
||||
self.top_k = config.num_experts_per_tok
|
||||
|
||||
if global_server_args_dict["moe_a2a_backend"].is_deepep():
|
||||
if get_moe_a2a_backend().is_deepep():
|
||||
# TODO: we will support tp < ep in the future
|
||||
self.ep_size = get_moe_expert_parallel_world_size()
|
||||
self.num_experts = (
|
||||
@@ -431,12 +407,12 @@ class DeepseekV2MoE(nn.Module):
|
||||
num_local_experts=config.n_routed_experts // self.tp_size,
|
||||
hidden_size=config.hidden_size,
|
||||
params_dtype=config.torch_dtype,
|
||||
deepep_mode=global_server_args_dict["deepep_mode"],
|
||||
deepep_mode=get_deepep_mode(),
|
||||
async_finish=True,
|
||||
return_recv_hook=True,
|
||||
)
|
||||
|
||||
self._enable_deepep_moe = global_server_args_dict["moe_a2a_backend"].is_deepep()
|
||||
self._enable_deepep_moe = get_moe_a2a_backend().is_deepep()
|
||||
|
||||
def get_moe_weights(self):
|
||||
return [
|
||||
@@ -484,13 +460,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
# router_logits: (num_tokens, n_experts)
|
||||
router_logits = self.gate(hidden_states)
|
||||
kwargs = {"hidden_states": hidden_states}
|
||||
|
||||
# FlashInferFP4MoE (TRTLLM path) expects (TopK, router_logits) tuple
|
||||
# Regular FusedMoE (CUTLASS path) expects StandardTopKOutput
|
||||
if should_use_flashinfer_trtllm_moe():
|
||||
kwargs["topk_output"] = (self.topk, router_logits)
|
||||
else:
|
||||
kwargs["topk_output"] = self.topk(hidden_states, router_logits)
|
||||
kwargs["topk_output"] = self.topk(hidden_states, router_logits)
|
||||
|
||||
final_hidden_states = self.experts(**kwargs)
|
||||
if not _is_cuda:
|
||||
@@ -520,13 +490,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
# router_logits: (num_tokens, n_experts)
|
||||
router_logits = self.gate(hidden_states)
|
||||
kwargs = {"hidden_states": hidden_states}
|
||||
|
||||
# FlashInferFP4MoE (TRTLLM path) expects (TopK, router_logits) tuple
|
||||
# Regular FusedMoE (CUTLASS path) expects StandardTopKOutput
|
||||
if should_use_flashinfer_trtllm_moe():
|
||||
kwargs["topk_output"] = (self.topk, router_logits)
|
||||
else:
|
||||
kwargs["topk_output"] = self.topk(hidden_states, router_logits)
|
||||
kwargs["topk_output"] = self.topk(hidden_states, router_logits)
|
||||
|
||||
final_hidden_states = self.experts(**kwargs)
|
||||
if not _is_cuda and not _use_aiter:
|
||||
@@ -2478,17 +2442,15 @@ class DeepseekV2ForCausalLM(nn.Module):
|
||||
|
||||
# Params for weights, fp8 weight scales, fp8 activation scales
|
||||
# (param_name, weight_name, expert_id, shard_id)
|
||||
expert_params_mapping = get_moe_impl_class().make_expert_params_mapping(
|
||||
expert_params_mapping = FusedMoE.make_expert_params_mapping(
|
||||
ckpt_gate_proj_name="gate_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_up_proj_name="up_proj",
|
||||
num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
|
||||
)
|
||||
if self.quant_config and self.quant_config.get_name() == "w4afp8":
|
||||
expert_params_mapping += (
|
||||
get_moe_impl_class().make_expert_input_scale_params_mapping(
|
||||
num_experts=self.config.n_routed_experts
|
||||
)
|
||||
expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping(
|
||||
num_experts=self.config.n_routed_experts
|
||||
)
|
||||
|
||||
# Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
|
||||
|
||||
Reference in New Issue
Block a user