[BugFix]Resolve the bubble in the PP pipeline of PP+MTP and fix the verification of PP in the pd separation scenario. (vllm-project#10855)

### What this PR does / why we need it?
In the pd separation + pp + mtp scenario, synchronization is avoided and
bubbles are eliminated by skipping the broadcast of the p node.
Fix the verification of PP in the pd separation scenario.
### Does this PR introduce _any_ user-facing change?
no
### How was this patch tested?
ran GLM5.1 gpqa test 87.37


- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@967c5c3

Signed-off-by: LostFox11 <wangziyue17@huawei.com>
Co-authored-by: LostFox11 <wangziyue17@huawei.com>

Author: LostFox11

vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py

@@ -1411,6 +1411,15 @@ def set_xfer_handshake_metadata(self, metadata: dict[int, KVConnectorHandshakeMe
         assert self.connector_scheduler is not None
         self.connector_scheduler.set_xfer_handshake_metadata(metadata)
 
+    def set_xfer_handshake_metadata_pp_aware(
+        self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
+    ) -> None:
+        tp_size = max(tp_rank for (_, tp_rank) in metadata) + 1
+        flat_metadata: dict[int, KVConnectorHandshakeMetadata] = {
+            pp_rank * tp_size + tp_rank: meta for (pp_rank, tp_rank), meta in metadata.items()
+        }
+        self.set_xfer_handshake_metadata(flat_metadata)
+
 
 class MooncakeConnectorScheduler:
     """Implementation of Scheduler side methods"""
@@ -1782,7 +1791,11 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: KVC
         device_index = (self.pp_rank + self.pcp_rank) * self.tp_size + self.tp_rank
         self.handshake_port = self.side_channel_port + device_index
         self.sockets: dict = {}
-        self.engine = global_te.get_transfer_engine(self.side_channel_host, device_name=None)
+        device_name = str(torch.npu.current_device()) if self.pp_size > 1 else None
+        self.engine = global_te.get_transfer_engine(
+            self.side_channel_host,
+            device_name=device_name,
+        )
         self.te_rpc_port = self.engine.get_rpc_port()
 
         # Background thread for sending or receiving KV caches.

vllm_ascend/worker/model_runner_v1.py

@@ -2394,13 +2394,15 @@ def sample_tokens(
     ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
         kv_connector_output = self.kv_connector_output
         self.kv_connector_output = None
+        pp = get_pp_group()
+        skip_pp_pd_broadcast = self.is_kv_producer and pp.world_size > 1
 
         if self.execute_model_state is None:
             # Nothing to do (PP non-final rank case), output isn't used.
             # receive sampled token ids from the last PP rank when using
             # async scheduling + pipeline parallelism so downstream code
             # (e.g., PCP input preparation) can access them.
-            if self.use_async_scheduling and get_pp_group().world_size > 1:
+            if self.use_async_scheduling and pp.world_size > 1 and not skip_pp_pd_broadcast:
                 self._pp_receive_prev_sampled_token_ids_to_input_batch()
             if not kv_connector_output:
                 return None  # noqa
@@ -2546,8 +2548,7 @@ def propose_draft_token_ids(sampled_token_ids):
         # last PP rank so other PP ranks can receive them without going
         # through the scheduler/engine IPC path.
         if self.use_async_scheduling:
-            pp = get_pp_group()
-            if pp.world_size > 1 and pp.is_last_rank:
+            if pp.world_size > 1 and pp.is_last_rank and not skip_pp_pd_broadcast:
                 self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids)
 
         if not self.use_async_scheduling: