From 01d6ceff8f129efb4239f07060ecfd2f28429cc7 Mon Sep 17 00:00:00 2001 From: wxiwnd Date: Fri, 26 Jun 2026 22:50:35 +0800 Subject: [PATCH] fix(nsa): guard cp split to extend batches --- .../sglang/srt/layers/attention/nsa/utils.py | 10 +++++++- .../unit/layers/test_nsa_cp_utils.py | 25 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 524adc27b..5c7d23a94 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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: diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 795e38d7a..8e3c6ceab 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -77,6 +77,31 @@ class TestPageAlignedCacheExtent(unittest.TestCase): class TestNSAInSeqCPUtils(unittest.TestCase): + def test_can_cp_split_ignores_decode_batch_without_extend_lens(self): + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.DECODE, + extend_seq_lens_cpu=None, + uses_cp_shared_kv=True, + ) + + self.assertFalse( + can_cp_split( + seq_len=200, cp_size=8, use_nsa=True, forward_batch=forward_batch + ) + ) + + def test_can_cp_split_fails_extend_batch_without_extend_lens(self): + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + extend_seq_lens_cpu=None, + uses_cp_shared_kv=True, + ) + + with self.assertRaisesRegex(RuntimeError, "missing_extend_seq_lens_cpu"): + can_cp_split( + seq_len=200, cp_size=8, use_nsa=True, forward_batch=forward_batch + ) + def test_contiguous_valid_cp_query_count(self): from sglang.srt.layers.attention.nsa.nsa_indexer import ( _compute_contiguous_valid_cp_query_count,