[NPU] fix pp 2 hang on npu (#38249)

python/sglang/srt/distributed/parallel_state.py

@@ -1811,6 +1811,161 @@ class GroupCoordinator:
                 tensor_dict[key] = value
         return tensor_dict
 
+    def send_recv_tensor_dict(
+        self,
+        send_tensor_dict: Dict[str, Union[torch.Tensor, Any]],
+        send_dst: Optional[int] = None,
+        recv_src: Optional[int] = None,
+        send_all_gather_group: Optional["GroupCoordinator"] = None,
+        recv_all_gather_group: Optional["GroupCoordinator"] = None,
+    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+        """Send tensor dict to *send_dst* and simultaneously recv from *recv_src*.
+
+        Uses ``batch_isend_irecv`` to submit all send/recv operations
+        atomically, avoiding deadlock on backends (e.g. NPU/HCCL) where
+        ``isend`` may block until a matching ``recv`` is posted.
+
+        NOTE: ``send_dst`` / ``recv_src`` are local ranks within this group.
+        """
+        if not torch.distributed.is_initialized() or self.world_size == 1:
+            return None
+
+        if send_dst is None:
+            send_dst = (self.rank_in_group + 1) % self.world_size
+        if recv_src is None:
+            recv_src = (self.rank_in_group - 1) % self.world_size
+
+        assert send_dst < self.world_size, f"Invalid send_dst rank ({send_dst})"
+        assert recv_src < self.world_size, f"Invalid recv_src rank ({recv_src})"
+
+        group = self.device_group
+        metadata_group = self.cpu_group
+
+        # ---- 1. Exchange metadata via batch_isend_irecv on CPU group ----
+        send_metadata_list, send_tensor_list = _split_tensor_dict(send_tensor_dict)
+
+        send_meta_bytes = pickle.dumps(send_metadata_list)
+        send_meta_size = torch.tensor([len(send_meta_bytes)], dtype=torch.long)
+        send_meta_data = torch.frombuffer(send_meta_bytes, dtype=torch.uint8)
+
+        recv_meta_size = torch.empty(1, dtype=torch.long)
+
+        meta_ops = [
+            torch.distributed.P2POp(
+                torch.distributed.isend,
+                send_meta_size,
+                self.ranks[send_dst],
+                group=metadata_group,
+            ),
+            torch.distributed.P2POp(
+                torch.distributed.irecv,
+                recv_meta_size,
+                self.ranks[recv_src],
+                group=metadata_group,
+            ),
+        ]
+        reqs = torch.distributed.batch_isend_irecv(meta_ops)
+        for req in reqs:
+            req.wait()
+
+        recv_meta_len = recv_meta_size.item()
+        recv_meta_data = torch.empty(recv_meta_len, dtype=torch.uint8)
+
+        meta_ops = [
+            torch.distributed.P2POp(
+                torch.distributed.isend,
+                send_meta_data,
+                self.ranks[send_dst],
+                group=metadata_group,
+            ),
+            torch.distributed.P2POp(
+                torch.distributed.irecv,
+                recv_meta_data,
+                self.ranks[recv_src],
+                group=metadata_group,
+            ),
+        ]
+        reqs = torch.distributed.batch_isend_irecv(meta_ops)
+        for req in reqs:
+            req.wait()
+
+        recv_metadata_list = pickle.loads(recv_meta_data.numpy())
+
+        # ---- 2. Prepare recv buffers and collect all tensor ops ----
+        recv_tensor_dict: Dict[str, Any] = {}
+        tensor_ops: List[torch.distributed.P2POp] = []
+        recv_tensor_info: List[
+            Tuple[str, torch.Tensor, bool, Optional[torch.Size]]
+        ] = []
+
+        for key, value in recv_metadata_list:
+            if isinstance(value, TensorMetadata):
+                tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)
+                if tensor.numel() == 0:
+                    recv_tensor_dict[key] = tensor
+                    continue
+
+                use_all_gather = (
+                    recv_all_gather_group is not None
+                    and tensor.numel() % recv_all_gather_group.world_size == 0
+                )
+                orig_shape = None
+                if use_all_gather:
+                    orig_shape = tensor.shape
+                    tensor = tensor.reshape(recv_all_gather_group.world_size, -1)[
+                        recv_all_gather_group.rank_in_group
+                    ]
+
+                comm_group = metadata_group if tensor.is_cpu else group
+                tensor_ops.append(
+                    torch.distributed.P2POp(
+                        torch.distributed.irecv,
+                        tensor,
+                        self.ranks[recv_src],
+                        group=comm_group,
+                    )
+                )
+                recv_tensor_info.append((key, tensor, use_all_gather, orig_shape))
+            else:
+                recv_tensor_dict[key] = value
+
+        # Add send ops
+        for tensor in send_tensor_list:
+            if tensor.numel() == 0:
+                continue
+            send_t = tensor
+            if (
+                send_all_gather_group is not None
+                and send_t.numel() % send_all_gather_group.world_size == 0
+            ):
+                send_t = send_t.reshape(send_all_gather_group.world_size, -1)[
+                    send_all_gather_group.rank_in_group
+                ]
+            comm_group = metadata_group if send_t.is_cpu else group
+            tensor_ops.append(
+                torch.distributed.P2POp(
+                    torch.distributed.isend,
+                    send_t,
+                    self.ranks[send_dst],
+                    group=comm_group,
+                )
+            )
+
+        # ---- 3. Batch exchange all tensors ----
+        if tensor_ops:
+            reqs = torch.distributed.batch_isend_irecv(tensor_ops)
+            for req in reqs:
+                req.wait()
+
+        # ---- 4. Post-process received tensors (all_gather if needed) ----
+        for key, tensor, use_all_gather, orig_shape in recv_tensor_info:
+            if use_all_gather:
+                tensor = recv_all_gather_group.all_gather(tensor, dim=0)
+                tensor = tensor.reshape(orig_shape)
+            recv_tensor_dict[key] = tensor
+
+        return recv_tensor_dict
+
     def barrier(self):
         """Barrier synchronization among the group.
         NOTE: don't use `device_group` here! `barrier` in NCCL is

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -32,8 +32,9 @@ from sglang.srt.sampling.sampling_observer_pp import (
     pop_auxiliary_output_from_pp_tensors,
 )
 from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj
-from sglang.srt.utils.common import is_xpu
+from sglang.srt.utils.common import is_npu, is_xpu
 
+_is_npu = is_npu()
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -1025,6 +1026,20 @@ class SchedulerPPMixin:
         batch_result = None
         send_output_work = []
 
+        # On NPU (HCCL), isend/irecv may block until a matching peer op is
+        # posted, so the parity-based send-first/recv-first ordering used
+        # for NPU is replaced by batch_isend_irecv which submits all
+        # send/recv operations atomically.
+        if _is_npu and self.ps.pp_size == 2:
+            return self._pp2_only_send_recv_output_tensors_npu(
+                next_first_rank_mb_id,
+                next_mb_id,
+                mbs,
+                mb_metadata,
+                last_rank_comm_queue,
+                pp_outputs,
+            )
+
         # On CUDA, isend is async: it enqueues to the stream and returns,
         # so every rank can send first safely. On some backends isend is
         # effectively blocking and does not return until the peer posts a
@@ -1075,6 +1090,118 @@ class SchedulerPPMixin:
 
         return next_pp_outputs, batch_result, d2h_event, send_output_work
 
+    def _pp2_only_send_recv_output_tensors_npu(
+        self: Scheduler,
+        next_first_rank_mb_id: int,
+        next_mb_id: int,
+        mbs: List[ScheduleBatch],
+        mb_metadata: List[PPBatchMetadata],
+        last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]],
+        pp_outputs: PPProxyTensors | None,
+    ) -> Tuple[
+        Optional[PPProxyTensors],
+        Optional[GenerationBatchResult],
+        Optional[torch.Event],
+        List[P2PWork],
+    ]:
+        """NPU-specific output tensor send/recv using batch_isend_irecv.
+
+        Pairs the send of output tensors to the next stage with the recv
+        of output tensors from the previous stage in a single
+        ``batch_isend_irecv`` call, avoiding the deadlock that can occur
+        on HCCL when separate isend/irecv calls block waiting for each
+        other.
+        """
+        next_pp_outputs = None
+        d2h_event = None
+        batch_result = None
+        send_output_work = []
+
+        all_gather_group = (
+            self.attn_tp_group if self.require_attn_tp_allgather else None
+        )
+
+        # ---- Prepare send dict ----
+        # On NPU, always send something (full output or a lightweight skip
+        # marker) so the peer's recv in batch_isend_irecv always has a
+        # matching send.  Without this, SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM
+        # would cause asymmetric skip decisions between adjacent ranks
+        # (send target and recv target are different micro-batches), leading
+        # to deadlock or forcing the user to disable the optimisation
+        # entirely (≈10 % throughput loss).
+        send_dict: Optional[Dict[str, torch.Tensor]] = None
+        if self.pp_group.is_last_rank:
+            target_send = mbs[next_first_rank_mb_id]
+            if target_send is not None:
+                q_event, pp_outputs_to_send = last_rank_comm_queue.popleft()
+                if not target_send.forward_mode.is_prebuilt():
+                    if _pp_can_skip_output_comm(target_send):
+                        send_dict = {"__msg_type__": "output", "__skip__": True}
+                    else:
+                        self.device_module.current_stream().wait_event(q_event)
+                        send_dict = dict(pp_outputs_to_send.tensors)
+                        send_dict["__msg_type__"] = "output"
+        elif pp_outputs:
+            if pp_outputs.tensors.get("__skip__"):
+                send_dict = {"__msg_type__": "output", "__skip__": True}
+            else:
+                send_dict = dict(pp_outputs.tensors)
+                send_dict["__msg_type__"] = "output"
+
+        # ---- Determine recv conditions ----
+        # NOTE: skip_recv is NOT computed here.  The skip decision is made
+        # by the sender (last rank) and communicated via the __skip__
+        # marker.  The receiver always participates in the
+        # batch_isend_irecv and inspects the marker after recv.
+        target_recv = mbs[next_mb_id]
+        should_recv = (
+            target_recv is not None and not target_recv.forward_mode.is_prebuilt()
+        )
+
+        def _handle_recv_dict(recv_dict):
+            nonlocal next_pp_outputs, batch_result, d2h_event
+            if recv_dict.get("__skip__"):
+                # _pp_make_skip_output_result returns next_pp_outputs=None
+                # (correct for the non-NPU path where the skip propagates
+                # via pp_outputs=None).  On NPU we must propagate the skip
+                # marker through the ring, so override it here.
+                _, batch_result, d2h_event = self._pp_make_skip_output_result(
+                    target_recv, mb_metadata[next_mb_id]
+                )
+                next_pp_outputs = PPProxyTensors(recv_dict)
+            else:
+                next_pp_outputs = PPProxyTensors(recv_dict)
+                with self.copy_stream_ctx:
+                    self.copy_stream.wait_stream(self.schedule_stream)
+                    batch_result = self._pp_prep_batch_result(
+                        target_recv, mb_metadata[next_mb_id], next_pp_outputs
+                    )
+                    d2h_event = self.device_module.Event()
+                    d2h_event.record(self.device_module.current_stream())
+
+        # ---- Execute communication ----
+        if send_dict is not None and should_recv:
+            # Paired send + recv via batch_isend_irecv
+            with torch.profiler.record_function("send_recv_res_dict"):
+                recv_dict = self.pp_group.send_recv_tensor_dict(
+                    send_tensor_dict=send_dict,
+                    send_all_gather_group=all_gather_group,
+                    recv_all_gather_group=all_gather_group,
+                )
+            _handle_recv_dict(recv_dict)
+        elif send_dict is not None:
+            # Send only (recv not needed — target is None or prebuilt)
+            send_output_work = self._pp_send_dict_to_next_stage(
+                send_dict, async_send=True, msg_type="output"
+            )
+        elif should_recv:
+            # Recv only (no send needed)
+            with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
+                recv_dict = self._pp_recv_dict_from_prev_stage()
+            _handle_recv_dict(recv_dict)
+
+        return next_pp_outputs, batch_result, d2h_event, send_output_work
+
     def _pp_launch_batch(
         self: Scheduler,
         mb_id: int,