diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 9daccd870..04f542752 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -129,13 +129,15 @@ def dsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]): 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.utils.cp_utils import get_cp_padding_align_size + global_num_tokens = forward_batch.global_num_tokens_cpu.copy() sync_group_size = len(global_num_tokens) attn_cp_size = get_attention_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): - # Must match ForwardBatch.prepare_mlp_sync_batch, which pads to - # attn_cp_size * 2 (tokens are split into 2 * CP chunks for load balance). - global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size * 2) + global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size) dp_padding_mode = DpPaddingMode.get_dp_padding_mode( forward_batch.is_extend_in_batch, global_num_tokens ) diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index 1ca4da858..328aa383d 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/python/sglang/srt/layers/utils/cp_utils.py @@ -70,6 +70,21 @@ def is_prefill_cp_in_seq_split(): ) +def get_cp_padding_align_size() -> int: + """Token-count alignment for CP padding of global_num_tokens: 2 * cp_size + for zigzag (in-seq-split) CP, otherwise cp_size (1 when CP is off, so the + padding is a no-op; extra padding breaks EAGLE/MTP draft prefill, see + #23269). Keep prepare_mlp_sync_batch and cal_padded_tokens consistent + through this helper. + """ + from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split + + attn_cp_size = get_attention_cp_size() + if is_prefill_cp_in_seq_split() or is_dsa_prefill_cp_in_seq_split(): + return attn_cp_size * 2 + return attn_cp_size + + def is_mla_prefill_cp_enabled() -> bool: sa = get_global_server_args() return sa.enable_prefill_context_parallel and sa.use_mla_backend diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 79759ac86..8a9109e42 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -45,7 +45,6 @@ from sglang.srt.kv_canary.req_to_expected_token_ids_manager import ( ) from sglang.srt.layers.dp_attention import ( DpPaddingMode, - get_attention_cp_size, get_attention_dp_rank, get_attention_tp_rank, get_attention_tp_size, @@ -978,6 +977,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): def prepare_mlp_sync_batch(self, model_runner: ModelRunner): from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer + # Local import: a module-level cp_utils import here is circular (#27014). + from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size + assert self.global_num_tokens_cpu is not None assert self.global_num_tokens_for_logprob_cpu is not None @@ -990,11 +992,14 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # there is no reduce-scatter in LM logprob, so we do not need to adjust the padded length for logprob global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size) - # make sure that each rank has the same number of tokens to do collective communication and - # we can divide the tokens into 2 * CP chunks for load balance. - attn_cp_size = get_attention_cp_size() + # make sure that each rank has the same number of tokens to do collective communication. + # Zigzag (in-seq-split) CP pads to 2 * attn_cp_size for load balance; other CP modes + # pad to attn_cp_size; CP off pads nothing (extra padding breaks EAGLE/MTP draft + # prefill with NaN draft logits, see #23269). + # FIXME(kpham-sgl): revisit so draft prefill-extend tolerates padded dummy tokens. + 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], attn_cp_size * 2) + global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size) dp_padding_mode = DpPaddingMode.get_dp_padding_mode( self.is_extend_in_batch, global_num_tokens