[0.9.1][PD] Added support for delay-free blocks in prefill nodes (#1691)

### What this PR does / why we need it?

PD Logic Analysis:
In the current implementation, the P-node immediately releases memory
blocks after completing inference. Under high concurrency scenarios, if
the P-node's inference speed significantly outpaces the D-node's block
pulling operations, this leads to memory block contention/corruption
issues.

Current Solution:
The D-node sends acknowledgment messages to the worker connector in the
P-node's driver worker after completing data reception. The P-node
maintains a counter to track these acknowledgments - memory blocks are
only released after receiving confirmations from all D-node worker
connectors involved in the KV cache transfer.

### Does this PR introduce _any_ user-facing change?


### How was this patch tested?


---------

Signed-off-by: underfituu <hzhucong@163.com>

Author: underfituu

vllm_ascend/distributed/llmdatadist_c_mgr_connector.py

@@ -3,6 +3,7 @@
 import math
 import threading
 import time
+from collections import defaultdict
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
@@ -268,9 +269,11 @@ def request_finished(
         # we just transfer any data that computed from prefill node
         # note: there might be some issue on this, check it if there is any unexpected result
         computed_block_ids = block_ids
-        # If prompt < block_size, no xfer so free blocks immediately.
-
-        return False, dict(
+        delay_free_blocks = len(computed_block_ids) > 0
+        if delay_free_blocks:
+            logger.info("Delaying free of %d blocks for request %s",
+                        len(computed_block_ids), request.request_id)
+        return delay_free_blocks, dict(
             do_remote_prefill=True,
             do_remote_decode=False,
             remote_block_ids=computed_block_ids,
@@ -334,6 +337,8 @@ def __init__(self, vllm_config: VllmConfig):
         self.init_llm_datadist()
         self.finished_reqs: set[str] = set()
         self.soc_info = NPUSocInfo()
+        self.done_receiving_counts: defaultdict[str,
+                                                set[int]] = defaultdict(set)
 
     def listen_for_agent_metadata_req(self, event: threading.Event):
         assert self.local_agent_metadata is not None
@@ -368,16 +373,41 @@ def listen_for_agent_metadata_req(self, event: threading.Event):
                         )
                 elif event_msg == LLMDataDistCMgrEvent.ReqForFinished:
                     finished_req_id = decode_msg[0]
+                    decode_tp_rank = decode_msg[1]
+                    decode_tp_size = decode_msg[2]
                     with self.thread_lock:
-                        logger.debug(
-                            f"LLMDataDistCMgrConnectorWorker: Receiving request {finished_req_id} finished"
-                        )
-                        self.finished_reqs.add(finished_req_id)
+                        if self._increment_task_count(finished_req_id,
+                                                      decode_tp_rank,
+                                                      decode_tp_size):
+                            logger.debug(
+                                f"LLMDataDistCMgrConnectorWorker: Receiving request {finished_req_id} finished"
+                            )
+                            self.finished_reqs.add(finished_req_id)
+                    sock.send_multipart(
+                        (identity, b"", b"receiving decode finished"))
                 else:
                     raise RuntimeError(
                         f"LLMDataDistCMgrConnectorWorker: Receiving unexpected request event {event_msg} from remote !"
                     )
 
+    def _increment_task_count(self, request_id: str, tp_rank: int,
+                              decode_tp_size: int):
+        if request_id not in self.done_receiving_counts:
+            self.done_receiving_counts[request_id] = set()
+        if tp_rank in self.done_receiving_counts[request_id]:
+            logger.warning(
+                f"Received duplicate done signal for request {request_id} "
+                f"from tp rank {tp_rank}. Ignoring.")
+            return False
+        self.done_receiving_counts[request_id].add(tp_rank)
+        if len(self.done_receiving_counts[request_id]) == decode_tp_size:
+            self.done_receiving_counts.pop(request_id)
+            logger.info("All transfers completed for request: "
+                        f"{request_id}. Total ranks: "
+                        f"{decode_tp_size}.")
+            return True
+        return False
+
     def init_llm_datadist(self):
         assert self.local_agent_metadata is not None
         llm_config = LLMConfig()
@@ -722,18 +752,21 @@ def connect_to_remote_agent(self, host: str, port: int) -> int:
             cluster_id = self.add_remote_agent(metadata)
         return cluster_id
 
-    def send_finsih_to_remote(self, host: str, port: int, request_id):
+    def send_finish_to_remote(self, host: str, port: int, request_id):
         url = f"tcp://{host}:{port}"
         logger.debug(f"Sending finished to remote: {url}")
         msg_encoder = msgspec.msgpack.Encoder()
-        msg_send = msg_encoder.encode(
-            [LLMDataDistCMgrEvent.ReqForFinished, [request_id]])
+        msg_send = msg_encoder.encode([
+            LLMDataDistCMgrEvent.ReqForFinished,
+            [request_id, self.tp_rank, self.tp_size]
+        ])
         with zmq_ctx(zmq.REQ, url) as sock:  # type: ignore[attr-defined]
             try:
                 sock.send(msg_send)
                 logger.debug(
                     f"Request id {request_id} finished message send to remote {url}"
                 )
+                _ = sock.recv()
             except Exception as e:
                 logger.error(
                     f"Failed to send reqest_id {request_id} to prefill: {e}")
@@ -803,8 +836,7 @@ def _read_blocks(
                 raise RuntimeError(
                     "LLMDataDistCMgrConnectorWorker: Timeout during pull_blocks, you can try to increase the sync_kv_timeout config or checking your connect status"
                 )
-        self.send_finsih_to_remote(remote_ip, remote_port + tp_offset,
-                                   request_id)
+        self.send_finish_to_remote(remote_ip, remote_port, request_id)
         with self.thread_lock:
             self.finished_reqs.add(request_id)