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

[Bugfix][DSA] Fix num_splits "(b+1)" crash on prefill-CP speculative decode (#34443)

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
NVShreyas
2026-08-11 19:13:49 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent a53d3636ce
commit 30cb848d4b
1 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/utils.py
+11 -5
View File
@@ -195,7 +195,7 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch, # Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode. # calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
from sglang.srt.layers.cp.padding import get_cp_padding_align_size from sglang.srt.layers.cp.padding import get_cp_padding_align_size
from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active
# CP-v2 already pads each rank-local shard to its physical size # CP-v2 already pads each rank-local shard to its physical size
if is_cp_v2_active(forward_batch): if is_cp_v2_active(forward_batch):
@@ -206,10 +206,16 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
global_num_tokens = forward_batch.global_num_tokens_cpu.copy() global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens) sync_group_size = len(global_num_tokens)
attn_cp_size = get_parallel().attn_cp_size attn_cp_size = get_parallel().attn_cp_size
# Must match the CP padding in ForwardBatch.prepare_mlp_sync_batch. # Must mirror ForwardBatch.prepare_mlp_sync_batch, which applies cp_align_size only when
cp_align_size = get_cp_padding_align_size() # CP-v2 is disabled. Under enable_cp_v2() the speculative forwards (TARGET_VERIFY /
for i in range(sync_group_size): # DRAFT_EXTEND_V2) reach here with is_cp_v2_active False, and q is padded to attn_tp_size only
global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size) # (not cp-aligned). Applying cp_align here over-pads the flashmla metadata past q, so
# num_splits ends up longer than q -> fwd_kvcache_mla fails "num_splits must have shape (b+1)".
# (attn_cp analog of the attn_tp fix in PR #30642 / issue #30296.)
if not enable_cp_v2():
cp_align_size = get_cp_padding_align_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size)
# Reuse the mode selected when the DP buffer was prepared. # Reuse the mode selected when the DP buffer was prepared.
dp_padding_mode = forward_batch.dp_padding_mode dp_padding_mode = forward_batch.dp_padding_mode
if dp_padding_mode is None: if dp_padding_mode is None: