From f75abb4521cfa5a58a9876b3d71c5788302f751d Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Wed, 25 Feb 2026 00:26:44 -0800 Subject: [PATCH] [Fix][Qwen3.5] Fix KV cache slice transfer for GQA models with replicated KV heads (#19086) Co-authored-by: Shangming Cai --- python/sglang/srt/disaggregation/base/conn.py | 1 + .../srt/disaggregation/mooncake/conn.py | 20 ++++++++++++++----- python/sglang/srt/disaggregation/prefill.py | 3 +++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 3233c5e31..080079afc 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -32,6 +32,7 @@ class KVArgs: # for different tp decode_tp_size: int kv_head_num: int + total_kv_head_num: int page_size: int # for pp prefill prefill_pp_size: int diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 706d32fcc..cea0a822e 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -391,22 +391,32 @@ class MooncakeKVManager(CommonKVManager): local_tp_rank_in_group = self.kv_args.engine_rank % self.attn_tp_size src_kv_item_len = self.kv_args.kv_item_lens[0] dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size - num_kv_heads = self.kv_args.kv_head_num page_size = self.kv_args.page_size - # Calculate head distribution - src_heads_per_rank = num_kv_heads - dst_heads_per_rank = num_kv_heads * self.attn_tp_size // dst_attn_tp_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 * self.attn_tp_size + + src_heads_per_rank = max(1, total_kv_heads // self.attn_tp_size) + dst_heads_per_rank = max(1, total_kv_heads // dst_attn_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, self.attn_tp_size // total_kv_heads) + # Determine slicing parameters based on TP configuration if self.attn_tp_size > dst_attn_tp_size: # Send KVCache from multiple prefill instances to 1 decode instance 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 = ( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 66a5c87df..cddf469f7 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -157,6 +157,9 @@ class PrefillBootstrapQueue: kv_args.kv_item_lens = kv_item_lens if not self.is_mla_backend: kv_args.kv_head_num = self.token_to_kv_pool.head_num + kv_args.total_kv_head_num = ( + self.scheduler.model_config.get_total_num_kv_heads() + ) kv_args.page_size = self.token_to_kv_pool.page_size kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (