From 2ded8a6aeaebe235e7f72a4580b5d6effed09bd1 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Thu, 27 Aug 2026 17:31:41 +0800 Subject: [PATCH] [AMD] Enable moe_a2a_backend=mori for DeepSeek-V4 prefill context parallelism (#35611) --- python/sglang/srt/arg_groups/deepseek_v4_hook.py | 6 +++--- python/sglang/srt/models/deepseek_v4.py | 11 +++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 0359d2a4c..8bf9800f8 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -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( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 8052e150d..831ad0eb3 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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 = (