minke.yu/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
python/sglang/srt/disaggregation/common/conn.py
+4 -3
View File
@@ -168,12 +168,13 @@ class CommonKVManager(BaseKVManager):
num_kv_layers < dst_num_total_layers num_kv_layers < dst_num_total_layers
and dst_num_total_layers % num_kv_layers != 0 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) # Case: 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_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] dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer]
v_ptr_offset = num_kv_layers * multiplier_ratio
dst_v_ptrs = dst_kv_ptrs[ 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: else:
# Decode pp size should be equal to prefill pp size or 1 # Decode pp size should be equal to prefill pp size or 1