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

[PD] Optimize MHA models pp util calculation logic (#17306)

Browse Source
This commit is contained in:
Shangming Cai
2026-01-19 11:23:15 +08:00
committed by GitHub
parent ad1b4e4728
commit 0227db8926
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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