diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index abdacfd56..94f4f6167 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -195,7 +195,7 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"): # 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. 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 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() sync_group_size = len(global_num_tokens) attn_cp_size = get_parallel().attn_cp_size - # Must match the CP padding in ForwardBatch.prepare_mlp_sync_batch. - 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) + # Must mirror ForwardBatch.prepare_mlp_sync_batch, which applies cp_align_size only when + # CP-v2 is disabled. Under enable_cp_v2() the speculative forwards (TARGET_VERIFY / + # DRAFT_EXTEND_V2) reach here with is_cp_v2_active False, and q is padded to attn_tp_size only + # (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. dp_padding_mode = forward_batch.dp_padding_mode if dp_padding_mode is None: