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]