minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Support PP x PD decode with nixl backend (#14392)

Browse Source
This commit is contained in:
Kevin Li
2025-12-04 12:22:17 +08:00
committed by GitHub
parent 4f73e53dcc
commit 04df80a9a1
1 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/nixl/conn.py
+4 -6
View File
@@ -322,12 +322,11 @@ class NixlKVManager(CommonKVManager):
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = ( src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
self.get_mla_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs) self.get_mla_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
) )
kv_item_len = self.kv_args.kv_item_lens[0]
layers_params = [ layers_params = [
( (
src_kv_ptrs[layer_id], src_kv_ptrs[layer_id],
dst_kv_ptrs[layer_id], dst_kv_ptrs[layer_id],
kv_item_len, self.kv_args.kv_item_lens[layer_id],
) )
for layer_id in range(layers_current_pp_stage) for layer_id in range(layers_current_pp_stage)
] ]
@@ -336,19 +335,18 @@ class NixlKVManager(CommonKVManager):
self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs) self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
) )
kv_item_len = self.kv_args.kv_item_lens[0]
layers_params = [ layers_params = [
( (
src_k_ptrs[layer_id], src_k_ptrs[layer_id],
dst_k_ptrs[layer_id], dst_k_ptrs[layer_id],
kv_item_len, self.kv_args.kv_item_lens[layer_id],
) )
for layer_id in range(layers_current_pp_stage) for layer_id in range(layers_current_pp_stage)
] + [ ] + [
( (
src_v_ptrs[layer_id], src_v_ptrs[layer_id],
dst_v_ptrs[layer_id], dst_v_ptrs[layer_id],
kv_item_len, self.kv_args.kv_item_lens[layer_id],
) )
for layer_id in range(layers_current_pp_stage) for layer_id in range(layers_current_pp_stage)
] ]
@@ -607,7 +605,7 @@ class NixlKVManager(CommonKVManager):
handles.append(kv_xfer_handle) handles.append(kv_xfer_handle)
# Only the last chunk we need to send the aux data. # Only the last chunk we need to send the aux data.
if is_last and self.pp_group.is_last_rank: if is_last:
assert aux_index is not None assert aux_index is not None
aux_xfer_handle = self.send_aux( aux_xfer_handle = self.send_aux(
req.agent_name, req.agent_name,