diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 87523e5d4..997f6f389 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1795,13 +1795,28 @@ class ServerArgs: "Context parallel feature is still under experiment. It has only been verified on Hopper platform." ) if self.nsa_prefill_cp_mode == "in-seq-split": - # TODO Supports moe_dense_tp_size != 1, kv cache dtype = "fp8",moe_a2a_backend non-deepep and cross-machine operation . + # TODO Supports moe_dense_tp_size != 1, kv cache dtype = "fp8", and cross-machine operation . self.enable_dp_attention = True self.moe_dense_tp_size = 1 - self.moe_a2a_backend = "deepep" + # B300 decouple (Strategy B): only force the historical + # deepep CP path when no MoE backend was selected + # explicitly. With an explicit runner (e.g. + # flashinfer_trtllm) keep a2a as "none"/"flashinfer" + # so CP runs on the standard/allgather path. + if ( + self.moe_a2a_backend == "none" + and self.moe_runner_backend == "auto" + ): + self.moe_a2a_backend = "deepep" + if self.moe_a2a_backend not in {"deepep", "flashinfer", "none"}: + raise ValueError( + "For NSA prefill CP in-seq split mode, moe_a2a_backend " + "must be one of {'deepep', 'flashinfer', 'none'}." + ) self.ep_size = self.tp_size logger.warning( - "For in-seq split mode, we have the following restrictions: moe_dense_tp_size == 1, moe_a2a_backend == deepep, ep_size == tp_size, batch_size == 1" + "For in-seq split mode: moe_dense_tp_size == 1, moe_a2a_backend in " + "{'deepep', 'flashinfer', 'none'}, ep_size == tp_size, batch_size == 1" ) else: self.enable_dp_attention = True @@ -1809,9 +1824,14 @@ class ServerArgs: assert ( self.dp_size == 1 ), "For round-robin split mode, dp attention is not supported." - assert ( - self.tp_size == 8 - ), "Current multi-machine CP support suffers from precision issues. So context parallel only support Single machine(tp_size == 8)" + # B300 decouple: allow single-node tp4/CP4 (target PD) in + # addition to tp8/CP8. CP code is parametric in cp_size; the + # tp_size%attn_cp_size and tp_size%(dp_size*attn_cp_size) + # divisibility asserts below still guard validity. + assert self.tp_size in (4, 8), ( + "Multi-machine CP has precision issues; context parallel only " + f"supports a single machine (tp_size in {{4, 8}}), got {self.tp_size}." + ) self.attn_cp_size = self.tp_size logger.warning(