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

fix(nsa): guard cp split to extend batches

Browse Source
This commit is contained in:
wxiwnd
2026-06-26 22:50:35 +08:00
committed by wxiwnd
parent c739ea667d
commit 01d6ceff8f
2 changed files with 34 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1828,6 +1828,14 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
forward_batch.forward_mode.is_context_parallel_extend()
or _is_cp_shared_kv_draft_extend(forward_batch)
)
if not is_context_parallel_extend:
return False
extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
if extend_seq_lens_cpu is None:
raise RuntimeError(
"[NSA_CP_FAIL_FAST][missing_extend_seq_lens_cpu] "
"Context-parallel EXTEND batch is missing extend_seq_lens_cpu."
)
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert seq_len % cp_size == 0, (
@@ -1848,7 +1856,7 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
)
else:
cur_cp_seq_len = seq_len // (cp_size * 2)
extend_token_count = sum(forward_batch.extend_seq_lens_cpu)
extend_token_count = sum(extend_seq_lens_cpu)
if _is_cp_shared_kv_forward_batch(forward_batch):
min_extend_token_count = 1
else: