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 (
cfg.tp_size <= 8
), "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(
"DeepSeekV4 CP supports moe_a2a_backend in "
"('none', 'deepep', 'megamoe'), "
f"DeepSeekV4 CP supports moe_a2a_backend in {supported_a2a_backends}, "
f"got {cfg.moe_a2a_backend!r}."
)
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():
hidden_states = dsa_cp_gather_hidden_states(hidden_states)
else:
assert moe_a2a_backend.is_deepep() or moe_a2a_backend.is_megamoe(), (
"CP requires DeepEP or megaMoE "
"(moe_a2a_backend == deepep or megamoe). "
f"Got {moe_a2a_backend.value}."
assert (
moe_a2a_backend.is_deepep()
or moe_a2a_backend.is_megamoe()
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:
hidden_states, local_hidden_states = (