diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 098115d9d..524adc27b 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -1820,6 +1820,14 @@ def _validate_cp_shared_kv_cp_split_plan_inputs( def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): _fail_if_cp_shared_kv_round_robin(forward_batch, op="can_cp_split") + # Only genuine context-parallel EXTEND forwards are eligible for in-seq CP + # splitting. CP shared-KV spec/decode batches can be flagged + # uses_cp_shared_kv=True while legitimately carrying no prefix page-plan, + # so the hard page-plan validator below must be gated on this eligibility. + is_context_parallel_extend = ( + forward_batch.forward_mode.is_context_parallel_extend() + or _is_cp_shared_kv_draft_extend(forward_batch) + ) if is_nsa_prefill_cp_round_robin_split(): cur_cp_seq_len = seq_len // cp_size assert seq_len % cp_size == 0, ( @@ -1834,9 +1842,8 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): if _is_cp_shared_kv_forward_batch(forward_batch): cur_cp_seq_len = ( 1 - if _validate_cp_shared_kv_cp_split_plan_inputs( - forward_batch, cp_size - ) + if is_context_parallel_extend + and _validate_cp_shared_kv_cp_split_plan_inputs(forward_batch, cp_size) else 0 ) else: @@ -1846,10 +1853,6 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): min_extend_token_count = 1 else: min_extend_token_count = cp_size - is_context_parallel_extend = ( - forward_batch.forward_mode.is_context_parallel_extend() - or _is_cp_shared_kv_draft_extend(forward_batch) - ) if ( cur_cp_seq_len != 0 and cp_size > 1 diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index c0979726b..2a46a249d 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -427,6 +427,27 @@ class TestNSAInSeqCPUtils(unittest.TestCase): ): can_cp_split(1089, 8, True, forward_batch) + def test_can_cp_split_skips_page_plan_validator_for_target_verify(self): + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + extend_seq_lens_cpu=[64], + extend_prefix_lens_cpu=None, + token_to_kv_pool=SimpleNamespace(page_size=64), + forward_mode=ForwardMode.TARGET_VERIFY, + ) + + with ( + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp", + return_value=True, + ), + ): + self.assertFalse(can_cp_split(64, 8, True, forward_batch)) + def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self): class Mode: def is_context_parallel_extend(self):