diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index 661a0cc4e..1eff99404 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/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) - layers_params = [ - ( - self.kv_args.kv_data_ptrs[layer_id], - dst_kv_ptrs[layer_id], - self.kv_args.kv_item_lens[layer_id], + if self.pp_size > 1: + 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) ) - 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 diff --git a/python/sglang/srt/disaggregation/ascend/transfer_engine.py b/python/sglang/srt/disaggregation/ascend/transfer_engine.py index 93e08726d..53c402df1 100644 --- a/python/sglang/srt/disaggregation/ascend/transfer_engine.py +++ b/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 ( - get_tensor_model_parallel_world_size, - get_tp_group, + from sglang.srt.distributed.parallel_state import ( + get_world_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) - for _ in range(get_tensor_model_parallel_world_size()) + torch.empty_like(tmp_tensor) for _ in range(get_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(