minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Enable moe_a2a_backend=mori for DeepSeek-V4 prefill context parallelism (#35611)

Browse Source
This commit is contained in:
AMD-yanfeiwang
2026-08-27 02:31:41 -07:00
committed by GitHub
parent 636a6f7dba
commit 2ded8a6aea
2 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/deepseek_v4_hook.py
+3 -3
View File
@@ -204,10 +204,10 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
assert ( assert (
cfg.tp_size <= 8 cfg.tp_size <= 8
), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues."
if cfg.moe_a2a_backend not in ("none", "deepep", "megamoe"): supported_a2a_backends = ("none", "deepep", "megamoe", "mori")
if cfg.moe_a2a_backend not in supported_a2a_backends:
raise ValueError( raise ValueError(
"DeepSeekV4 CP supports moe_a2a_backend in " f"DeepSeekV4 CP supports moe_a2a_backend in {supported_a2a_backends}, "
"('none', 'deepep', 'megamoe'), "
f"got {cfg.moe_a2a_backend!r}." f"got {cfg.moe_a2a_backend!r}."
) )
logger.warning( logger.warning(
python/sglang/srt/models/deepseek_v4.py
+7 -4
View File
@@ -2304,10 +2304,13 @@ class DeepseekV4DecoderLayer(nn.Module):
if moe_a2a_backend.is_none(): if moe_a2a_backend.is_none():
hidden_states = dsa_cp_gather_hidden_states(hidden_states) hidden_states = dsa_cp_gather_hidden_states(hidden_states)
else: else:
assert moe_a2a_backend.is_deepep() or moe_a2a_backend.is_megamoe(), ( assert (
"CP requires DeepEP or megaMoE " moe_a2a_backend.is_deepep()
"(moe_a2a_backend == deepep or megamoe). " or moe_a2a_backend.is_megamoe()
f"Got {moe_a2a_backend.value}." or moe_a2a_backend.is_mori()
), (
"CP requires moe_a2a_backend in ('deepep', 'megamoe', 'mori'), "
f"got {moe_a2a_backend.value!r}."
) )
elif _use_tp_moe_gather: elif _use_tp_moe_gather:
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (