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

Fix mooncake dispatcher (#11908)

Browse Source
This commit is contained in:
Xun Sun
2025-10-22 16:11:49 +08:00
committed by GitHub
parent 80b2b3207a
commit e028af6998
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -86,7 +86,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
a2a_backend = get_moe_a2a_backend() a2a_backend = get_moe_a2a_backend()
if a2a_backend.is_none(): if a2a_backend.is_none():
return StandardDispatcher(moe_runner_config) return StandardDispatcher(moe_runner_config)
elif a2a_backend.is_deepep(): elif a2a_backend.is_deepep() or a2a_backend.is_mooncake():
return MaybeTboDeepEPDispatcher( return MaybeTboDeepEPDispatcher(
group=get_tp_group().device_group, group=get_tp_group().device_group,
router_topk=moe_runner_config.top_k, router_topk=moe_runner_config.top_k,

8
python/sglang/srt/layers/moe/token_dispatcher/mooncake.py
View File

@@ -36,7 +36,7 @@ class MooncakeDispatchOutput(NamedTuple):
"""Mooncake EP dispatch output.""" """Mooncake EP dispatch output."""
hidden_states: torch.Tensor hidden_states: torch.Tensor
hidden_states_scale: torch.Tensor hidden_states_scale: Optional[torch.Tensor]
topk_ids: torch.Tensor topk_ids: torch.Tensor
topk_weights: torch.Tensor topk_weights: torch.Tensor
masked_m: torch.Tensor masked_m: torch.Tensor
@@ -205,8 +205,14 @@ class _MooncakeEPDispatcherImpl:
masked_m masked_m
) )
if isinstance(hidden_states, tuple):
hidden_states, hidden_states_scale = hidden_states
else:
hidden_states_scale = None
return MooncakeDispatchOutput( return MooncakeDispatchOutput(
hidden_states, hidden_states,
hidden_states_scale,
topk_ids, topk_ids,
topk_weights, topk_weights,
masked_m, masked_m,