chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: iridiumine <42236072+iridiumine@users.noreply.github.com>
This commit is contained in:
Hexq0210
2026-03-02 21:33:16 +08:00
committed by GitHub
parent eaf18ebe8d
commit 714c53d609
2 changed files with 34 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
python/sglang/srt/disaggregation/ascend/conn.py
View File

@@ -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

11
python/sglang/srt/disaggregation/ascend/transfer_engine.py
View File

@@ -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(