From bb6055b43c146d626b26f300c1b358292dbc92ae Mon Sep 17 00:00:00 2001 From: Ch3ngY1 <91232537+Ch3ngY1@users.noreply.github.com> Date: Sun, 18 Jan 2026 23:37:58 +0800 Subject: [PATCH] [Bugfix] Fix PD accuracy when MTP is not configured on the prefill node (#17212) Co-authored-by: Shangming Cai --- python/sglang/srt/disaggregation/common/conn.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 595bb3ef1..ef0ba5245 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -164,6 +164,17 @@ class CommonKVManager(BaseKVManager): if num_kv_layers == dst_num_total_layers: dst_k_ptrs = dst_kv_ptrs[:dst_num_total_layers] dst_v_ptrs = dst_kv_ptrs[dst_num_total_layers:] + elif ( + num_kv_layers < dst_num_total_layers + and dst_num_total_layers % num_kv_layers != 0 + ): + # Case: Decode has more layers than Prefill (e.g., Decode has draft model KV while Prefill is deployed without speculative decoding) + # To prevent empty Value Cache, which leads to wrong response + # dst_kv_ptrs layout: [K_main..., V_main..., draft_K..., draft_V...] + dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer] + dst_v_ptrs = dst_kv_ptrs[ + num_kv_layers + start_layer : num_kv_layers + end_layer + ] else: # Decode pp size should be equal to prefill pp size or 1 dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer]