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

[Disagg][NIXL] Fix heterogeneous TP KV transfer for non-MLA models (same logic with mooncake, Step 1/2 for Qwen3.5 support) (#22145)

Browse Source
This commit is contained in:
YAMY
2026-04-07 14:52:02 +08:00
committed by GitHub
parent ba78f6e0ef
commit 3148742ddb
1 changed files with 20 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/nixl/conn.py
+20 -8
View File
@@ -477,25 +477,35 @@ class NixlKVManager(CommonKVManager):
# Get configuration from kv_args
local_tp_rank_in_group = self.kv_args.engine_rank % prefill_tp_size
dst_tp_rank_in_group = decode_tp_rank % decode_tp_size
num_kv_heads = self.kv_args.kv_head_num
# Calculate head distribution
src_heads_per_rank = num_kv_heads
dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size
src_kv_item_len = self.kv_args.kv_item_lens[0]
page_size = self.kv_args.page_size
# Use total KV head count (not per-rank) for correct head distribution.
# Per-rank kv_head_num is max(1, total//tp) which loses info when total < tp.
total_kv_heads = getattr(self.kv_args, "total_kv_head_num", 0)
if total_kv_heads <= 0:
total_kv_heads = self.kv_args.kv_head_num * prefill_tp_size
src_heads_per_rank = max(1, total_kv_heads // prefill_tp_size)
dst_heads_per_rank = max(1, total_kv_heads // decode_tp_size)
bytes_per_head_slice_to_send = (
dst_kv_item_len // page_size // dst_heads_per_rank
)
# GQA replication: how many prefill ranks share the same KV head
src_replication = max(1, prefill_tp_size // total_kv_heads)
# Determine which heads to send
if prefill_tp_size > decode_tp_size:
# Multiple prefill ranks to one decode rank
src_head_start_offset = 0
num_heads_to_send = src_heads_per_rank
dst_head_start_offset = local_tp_rank_in_group * src_heads_per_rank
unique_head_idx = local_tp_rank_in_group // src_replication
dst_head_start_offset = (
unique_head_idx * src_heads_per_rank
) % dst_heads_per_rank
else:
# Send KVCache from 1 prefill instance to multiple decode instances
src_head_start_offset = (
@@ -748,7 +758,9 @@ class NixlKVManager(CommonKVManager):
assert len(chunked_dst_kv_indice) == len(kv_indices)
assert req.agent_name in self.decode_kv_args_table
notif = f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.pp_rank}"
notif = (
f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.engine_rank}"
)
decode_tp_size = self.decode_kv_args_table[req.agent_name].decode_tp_size
if self.is_mla_backend or (decode_tp_size == self.attn_tp_size):
@@ -789,7 +801,7 @@ class NixlKVManager(CommonKVManager):
dst_info.dst_state_data_ptrs,
req.dst_state_indices,
dst_info.gpu_id,
f"{req.room}_state_{self.kv_args.pp_rank}",
f"{req.room}_state_{self.kv_args.engine_rank}",
decode_tp_size,
)
if state_xfer_handle is not None: