From df2a3696cd76c72c40557d882974ff5c536c8eb5 Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Tue, 9 Jun 2026 19:54:59 +0800
Subject: [PATCH] Gate CP split validation to eligible extend forwards

CP shared-KV marks all forwards with uses_cp_shared_kv, but TARGET_VERIFY/decode style forwards may legitimately carry no extend prefix page-plan. The CP split helper previously ran the hard page-plan validator before checking context-parallel extend eligibility, so non-extend spec paths could fail before the later no-split decision.

Hoist the eligibility check and only validate page-plan metadata for real context-parallel EXTEND or shared-KV draft extend forwards. Add a regression that TARGET_VERIFY with no prefix metadata returns no-split instead of failing.

Constraint: Absorb syh 7fea88278 only; MQA logits chunk and per-forward budget changes are intentionally excluded.
Rejected: Relax the validator globally | real prefill page-plan violations must remain fail-fast.
Confidence: high
Scope-risk: narrow
Directive: Do not run CP shared-KV page-plan validation for non-context-parallel forward modes without proving those modes own extend_prefix_lens_cpu.
Tested: Remote g0034 container py_compile for utils/test file.
Tested: Remote g0034 container pytest test_nsa_cp_utils.py -k can_cp_split: 8 passed, 92 deselected.
Not-tested: Full ETE speculative non-deepep MoE path.
---
 .../sglang/srt/layers/attention/nsa/utils.py  | 17 ++++++++-------
 .../unit/layers/test_nsa_cp_utils.py          | 21 +++++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

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):