chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Bugfix] Fix PD accuracy when MTP is not configured on the prefill node (#17212)

Browse Source 
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Ch3ngY1
2026-01-18 23:37:58 +08:00
committed by GitHub
parent 4df74eb576
commit bb6055b43c
1 changed files with 11 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/disaggregation/common/conn.py
View File

@@ -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]