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 = (