diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 57c281bee..00ef96f9b 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -112,7 +112,11 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"): def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens): attn_cp_size = get_attention_cp_size() - if attn_cp_size == 1 or not can_nsa_prefill_cp_round_robin_split(forward_batch): + needs_cp_pad = attn_cp_size > 1 and can_nsa_prefill_cp_round_robin_split( + forward_batch + ) + needs_dp_pad = forward_batch.global_num_tokens_cpu is not None + if not needs_cp_pad and not needs_dp_pad: return nsa_cache_seqlens tokens = cal_padded_tokens(forward_batch) pad_len = tokens - nsa_cache_seqlens.shape[0]