fix: support PP2+CP8+TP8 (PP with context parallelism)

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -859,7 +859,7 @@ def _pp_commit_send_output_work_and_preprocess_output_tensors(
 
     def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False):
         p2p_work = []
-        if self.attn_tp_rank == 0:
+        if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
             dp_offset = self.attn_dp_rank * self.attn_tp_size
             p2p_work = point_to_point_pyobj(
                 data,
@@ -872,7 +872,7 @@ def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False
         return p2p_work
 
     def _pp_recv_pyobj_from_prev_stage(self: Scheduler):
-        if self.attn_tp_rank == 0:
+        if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
             dp_offset = self.attn_dp_rank * self.attn_tp_size
             data = point_to_point_pyobj(
                 [],
@@ -892,6 +892,14 @@ def _pp_recv_pyobj_from_prev_stage(self: Scheduler):
                 src=self.attn_tp_group.ranks[0],
             )
 
+        if self.attn_cp_size > 1:
+            data = broadcast_pyobj(
+                data,
+                self.attn_cp_group.rank,
+                self.attn_cp_cpu_group,
+                src=self.attn_cp_group.ranks[0],
+            )
+
         return data
 
     def _pp_prepare_tensor_dict(

python/sglang/srt/server_args.py

@@ -1235,6 +1235,7 @@ def _handle_model_specific_adjustments(self):
                         assert (
                             self.tp_size == 8
                         ), "Current multi-machine CP support suffers from precision issues. So context parallel only support Single machine(tp_size == 8)"
+                        self.attn_cp_size = self.tp_size
 
                         logger.warning(
                             f"Enable Context Parallel opt for deeeseekv3.2-DSA, Setting dp_size == {self.dp_size} and moe_dense_tp_size == {self.moe_dense_tp_size}, ep_size == {self.ep_size}, tp_size == {self.tp_size}, kv_cache_dtype == {self.kv_cache_dtype}, moe_a2a_backend {self.moe_a2a_backend} "
@@ -2116,7 +2117,8 @@ def _handle_context_parallelism(self):
             assert (
                 self.tp_size % (self.dp_size * self.attn_cp_size) == 0
             ), "tp_size must be divisible by dp_size * attn_cp_size"
-            assert self.pp_size == 1, "PP is not supported with context parallelism"
+            if not self.enable_nsa_prefill_context_parallel:
+                assert self.pp_size == 1, "PP is not supported with context parallelism"
 
         if self.moe_dp_size > 1:
             # The tp_size is the world size, not the real tensor parallel size