EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Small refactor DeepEPMode to clean up code a bit (#4992)

Browse Source
This commit is contained in:
fzyzcjy
2025-04-03 17:56:44 +08:00
committed by GitHub
parent e8999b13b7
commit 8e10fec9a8
5 changed files with 44 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/models/deepseek_v2.py
View File

@@ -70,7 +70,7 @@ from sglang.srt.managers.expert_distribution import ExpertDistributionRecorder
from sglang.srt.managers.schedule_batch import global_server_args_dict
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix, is_cuda, is_hip
from sglang.srt.utils import DeepEPMode, add_prefix, is_cuda, is_hip
_is_hip = is_hip()
_is_cuda = is_cuda()
@@ -215,7 +215,7 @@ class DeepseekV2MoE(nn.Module):
topk_group=config.topk_group,
correction_bias=self.gate.e_score_correction_bias,
prefix=add_prefix("experts", prefix),
deepep_mode=global_server_args_dict["deepep_mode"],
deepep_mode=DeepEPMode[global_server_args_dict["deepep_mode"]],
)
if config.n_shared_experts is not None:
@@ -264,7 +264,7 @@ 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=DeepEPMode[global_server_args_dict["deepep_mode"]],
async_finish=True, # TODO
return_recv_hook=True,
)