fix(nsa): guard cp split to extend batches

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:

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,