[Bugfix] Gate DP-attention even-token padding to CP-enabled configs (#26911)

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
     )

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

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