Merge branch 'v0.9.2-dev-pdpp2' into 'v0.9.2-dev'

add PD P do pp

See merge request dcutoolkit/deeplearing/vllm!215

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

@@ -55,6 +55,11 @@ class ReqMeta:
             slot_mapping=slot_mapping,
         )
 
+        self.parallel_config = vllm_config.parallel_config
+        self.model_config = vllm_config.model_config
+        self.total_num_hidden_layers = getattr(self.model_config.hf_text_config,
+                                              "num_hidden_layers", 0)
+        self.pp_size = self.parallel_config.pipeline_parallel_size
         
 @dataclass
 class P2pNcclConnectorMetadata(KVConnectorMetadata):
@@ -285,8 +290,29 @@ class P2pNcclConnector(KVConnectorBase_V1):
             ip, port = self.parse_request_id(request_id, True)
             remote_address = ip + ":" + str(port + self._rank)
             kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)
+
+            pp_rank = (self.parallel_config.rank // self.parallel_config.tensor_parallel_size
+                   ) % self.parallel_config.pipeline_parallel_size
+            if (self.pp_size == 1):
+                self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
+                                             kv_cache, remote_address)
+            elif (self.pp_size == 2):
+                if (pp_rank == 0):
+                    self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
+                                                kv_cache, remote_address)
+                    self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
+                                                kv_cache, ip + ":" + str(port + self._rank + 4))
+                else:
                     self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
                                                 kv_cache, remote_address)
+                    self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
+                                                kv_cache, ip + ":" + str(port + self._rank - 4))
+            elif (self.pp_size == 8):
+                for i in range(8):
+                    self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
+                                                kv_cache, ip + ":" + str(port + i))
+            else:
+                print("Error: only suppprt pp1 pp2 pp8!!!!!!")
 
     def wait_for_save(self):
         pass