[NPU] support PD disaggregation on ascend when using PP (#14908)

Co-authored-by: iridiumine <42236072+iridiumine@users.noreply.github.com>

python/sglang/srt/disaggregation/ascend/conn.py

@@ -48,15 +48,36 @@ class AscendKVManager(MooncakeKVManager):
             prefill_kv_indices, dst_kv_indices
         )
 
-        num_layers = len(self.kv_args.kv_data_ptrs)
+        if self.pp_size > 1:
-        layers_params = [
+            src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
-            (
+                self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
-                self.kv_args.kv_data_ptrs[layer_id],
-                dst_kv_ptrs[layer_id],
-                self.kv_args.kv_item_lens[layer_id],
             )
-            for layer_id in range(num_layers)
+
-        ]
+            layers_params = [
+                (
+                    src_k_ptrs[layer_id],
+                    dst_k_ptrs[layer_id],
+                    self.kv_args.kv_item_lens[layer_id],
+                )
+                for layer_id in range(layers_current_pp_stage)
+            ] + [
+                (
+                    src_v_ptrs[layer_id],
+                    dst_v_ptrs[layer_id],
+                    self.kv_args.kv_item_lens[layers_current_pp_stage + layer_id],
+                )
+                for layer_id in range(layers_current_pp_stage)
+            ]
+        else:
+            num_layers = len(self.kv_args.kv_data_ptrs)
+            layers_params = [
+                (
+                    self.kv_args.kv_data_ptrs[layer_id],
+                    dst_kv_ptrs[layer_id],
+                    self.kv_args.kv_item_lens[layer_id],
+                )
+                for layer_id in range(num_layers)
+            ]
 
         def set_transfer_blocks(
             src_ptr: int, dst_ptr: int, item_len: int

python/sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -51,9 +51,9 @@ class AscendTransferEngine(MooncakeTransferEngine):
         self.initialize()
 
     def initialize(self) -> None:
-        from sglang.srt.layers.dp_attention import (
+        from sglang.srt.distributed.parallel_state import (
-            get_tensor_model_parallel_world_size,
+            get_world_group,
-            get_tp_group,
+            get_world_size,
         )
 
         transfer_protocol = self._get_transfer_protocol()
@@ -64,12 +64,11 @@ class AscendTransferEngine(MooncakeTransferEngine):
             """with device RDMA for PD transfer"""
             tmp_tensor = torch.zeros(1, device="npu")
             output_tensor_list = [
-                torch.empty_like(tmp_tensor)
+                torch.empty_like(tmp_tensor) for _ in range(get_world_size())
-                for _ in range(get_tensor_model_parallel_world_size())
             ]
             # Initialize hccl in advance through all_gather to avoid conflicts with rdma initialization.
             torch.distributed.all_gather(
-                output_tensor_list, tmp_tensor, group=get_tp_group().device_group
+                output_tensor_list, tmp_tensor, group=get_world_group().device_group
             )
         """Initialize the ascend transfer instance."""
         ret_value = self.engine.initialize(