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

Fix spec v2+dp attention in nsa backend (#19134)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-02-21 21:46:15 -08:00
committed by GitHub
parent 8612358ee9
commit 8cf003c44b
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -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]