diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 1bcfe66bd..ff1d1ac13 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index c852d2e83..267410e6c 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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,