From 0227db892662ada1b4ae325c5522958a1c9b2690 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Mon, 19 Jan 2026 11:23:15 +0800 Subject: [PATCH] [PD] Optimize MHA models pp util calculation logic (#17306) --- python/sglang/srt/disaggregation/common/conn.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index ef0ba5245..67fe82ad6 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -168,12 +168,13 @@ class CommonKVManager(BaseKVManager): 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 + # Case: Decode has draft model KV while Prefill is deployed without speculative decoding # dst_kv_ptrs layout: [K_main..., V_main..., draft_K..., draft_V...] + multiplier_ratio = dst_num_total_layers // num_kv_layers dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer] + v_ptr_offset = num_kv_layers * multiplier_ratio dst_v_ptrs = dst_kv_ptrs[ - num_kv_layers + start_layer : num_kv_layers + end_layer + v_ptr_offset + start_layer : v_ptr_offset + end_layer ] else: # Decode pp size should be equal to prefill pp size or 1