Prevent batched CP draft from silently leaving the target path

EAGLE draft shared-KV is supposed to mirror the target CP layout, so bs>1 must not fall back to legacy full-input or padded-hidden behavior when required batch metadata is missing or inconsistent. This change keeps the existing bs=1 compatibility path but makes batched CP draft fail fast on missing/mismatched spec hidden states, embedding pad metadata, or input embed shapes. The docs record the current W7 boundary: draft prefill follows target metadata, while scheduler admission and ETE remain gated.

Constraint: CP draft KV must mirror target layout and must not silently diverge under bs>1 shared-KV.

Rejected: Allow bs>1 to use the old full-input fallback | it can hide wrong owner/page metadata and corrupt accept length.

Confidence: medium

Scope-risk: moderate

Directive: Do not open the scheduler bs>1 CP gate until EAGLE accept length/output length are verified with this fail-fast path enabled.

Tested: Remote g0034 targeted EAGLE fail-fast unit test passed; remote full test/registered/unit/layers/test_nsa_cp_utils.py passed 70 tests.

Not-tested: EAGLE bs>1 ETE, because scheduler CP bs>1 admission gate remains closed.

docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md

@@ -557,6 +557,13 @@ cp_size=8
 3. draft 不生成自己的 owner plan，一切跟随 target。
 4. target bs>1 ETE 通过后，再打开 draft local path。
 
+当前实现状态：
+
+- batch-aware input/position/spec-hidden split 已接到 draft local path；valid CP-local hidden 直接使用，full hidden 走 batch-aware split。
+- per-request last-token collect 已在 `cp_collect_last_token_hidden()` 中覆盖 batch plan。
+- 为避免打开 gate 后退回不正确的 full/draft fallback，bs>1 + `SGLANG_CP_DRAFT_SHARED_KV=1` + CP shared-KV 时，draft local path 对 missing/mismatch spec hidden、embedding pad metadata、input embeds shape mismatch 均 fail-fast。
+- admission gate 仍保持关闭；需要 EAGLE enabled ETE 验证 accept length/output length 后再打开。
+
 测试重点：
 
 - `cp_collect_last_token_hidden()` 对 bs=2 返回两个 request 的 last hidden。

docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md

@@ -555,6 +555,18 @@ target path 正确后，再恢复 EAGLE/draft，并做远端 ETE/perf 验证。
 4. `cp_collect_last_token_hidden()` 返回每个 request 的 last hidden。
 5. EAGLE accept length 不能因为 cache hit/bs>1 掉到 1。
 
+### 当前状态
+
+- target 侧 batch metadata、batch split helper、per-request last-token collect 已有单测覆盖。
+- EAGLE/draft prefill local path 复用上述 batch-aware split helper：
+  - `cp_split_and_rebuild_1d()` 处理 draft input ids；
+  - `cp_split_and_rebuild_position()` 处理 positions；
+  - `_get_cp_local_spec_hidden_states()` 接受已 CP-local 的 draft hidden，或对 full hidden 走 batch-aware split；
+  - `cp_collect_last_token_hidden()` 按 request owner/offset collect last hidden。
+- bs>1 时不允许 silent fallback：如果 CP draft shared-KV 已开启但 spec hidden、embedding pad metadata、input embeds 形状不满足 batch fast path，直接
+  `[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_*]` 报错。bs=1 兼容 fallback 暂时保留。
+- scheduler 的 bs>1 admission gate 仍未打开；打开前必须完成下面 ETE 场景，尤其是 EAGLE accept length 与 output len。
+
 ### ETE 验证场景
 
 远端：

python/sglang/srt/models/deepseek_nextn.py

@@ -151,15 +151,65 @@ class DeepseekModelNextN(nn.Module):
         if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
             logger.info("[CP_DRAFT_SHARED_KV] %s", message)
 
+    def _cp_draft_shared_kv_batch_size(self, forward_batch: ForwardBatch) -> int:
+        metadata = getattr(forward_batch, "nsa_cp_metadata", None)
+        batch_plan = getattr(metadata, "batch_plan", None)
+        for owner in (batch_plan, metadata, forward_batch):
+            value = getattr(owner, "batch_size", None)
+            if value is None:
+                continue
+            try:
+                return int(value)
+            except (TypeError, ValueError):
+                continue
+
+        extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
+        if extend_seq_lens_cpu is not None:
+            return len(extend_seq_lens_cpu)
+
+        return 1
+
+    def _requires_cp_draft_batch_gt1_fast_path(
+        self, forward_batch: ForwardBatch
+    ) -> bool:
+        return (
+            envs.SGLANG_CP_DRAFT_SHARED_KV.get()
+            and bool(getattr(forward_batch, "uses_cp_shared_kv", False))
+            and self._cp_draft_shared_kv_batch_size(forward_batch) > 1
+        )
+
+    def _fail_cp_draft_batch_gt1(
+        self,
+        forward_batch: ForwardBatch,
+        reason: str,
+        message: str,
+    ):
+        error_msg = (
+            f"[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_{reason}] "
+            f"{message} "
+            f"batch_size={self._cp_draft_shared_kv_batch_size(forward_batch)}"
+        )
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)
+
     def _get_cp_local_spec_hidden_states(
         self,
         forward_batch: ForwardBatch,
-        spec_hidden_states: torch.Tensor,
+        spec_hidden_states: Optional[torch.Tensor],
         *,
         full_num_tokens: int,
         local_num_tokens: int,
     ) -> Optional[torch.Tensor]:
+        must_use_batch_fast_path = self._requires_cp_draft_batch_gt1_fast_path(
+            forward_batch
+        )
         if spec_hidden_states is None:
+            if must_use_batch_fast_path:
+                self._fail_cp_draft_batch_gt1(
+                    forward_batch,
+                    "missing_spec_hidden",
+                    f"full_tokens={full_num_tokens} local_tokens={local_num_tokens}",
+                )
             self._debug_cp_draft_shared_kv("fallback reason=missing_spec_hidden")
             return None
 
@@ -201,6 +251,14 @@ class DeepseekModelNextN(nn.Module):
 
         if spec_hidden_states.shape[0] < local_num_tokens:
             pad_rows = local_num_tokens - spec_hidden_states.shape[0]
+            if must_use_batch_fast_path:
+                self._fail_cp_draft_batch_gt1(
+                    forward_batch,
+                    "spec_hidden_shape_mismatch",
+                    f"spec_tokens={spec_hidden_states.shape[0]} "
+                    f"full_tokens={full_num_tokens} "
+                    f"local_tokens={local_num_tokens} pad_rows={pad_rows}",
+                )
             if pad_rows <= max(get_attention_cp_size(), 1):
                 _log_eagle_accept_cp_draft_hidden_debug(
                     "local_pad",
@@ -239,6 +297,13 @@ class DeepseekModelNextN(nn.Module):
             f"spec_tokens={spec_hidden_states.shape[0]} "
             f"full_tokens={full_num_tokens} local_tokens={local_num_tokens}"
         )
+        if must_use_batch_fast_path:
+            self._fail_cp_draft_batch_gt1(
+                forward_batch,
+                "spec_hidden_shape_mismatch",
+                f"spec_tokens={spec_hidden_states.shape[0]} "
+                f"full_tokens={full_num_tokens} local_tokens={local_num_tokens}",
+            )
         return None
 
     def _embed_cp_local_input_ids(
@@ -253,6 +318,13 @@ class DeepseekModelNextN(nn.Module):
             forward_batch, local_num_tokens
         )
         if padded_token_count is None:
+            if self._requires_cp_draft_batch_gt1_fast_path(forward_batch):
+                self._fail_cp_draft_batch_gt1(
+                    forward_batch,
+                    "missing_or_stale_embedding_pad_len",
+                    f"full_tokens={full_num_tokens} "
+                    f"local_tokens={local_num_tokens}",
+                )
             self._debug_cp_draft_shared_kv(
                 "fallback reason=missing_or_stale_embedding_pad_len "
                 f"full_tokens={full_num_tokens} local_tokens={local_num_tokens}"
@@ -292,6 +364,9 @@ class DeepseekModelNextN(nn.Module):
         use_cp = nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp)
         use_cp_local_draft = use_cp and envs.SGLANG_CP_DRAFT_SHARED_KV.get()
         if use_cp_local_draft:
+            must_use_batch_fast_path = self._requires_cp_draft_batch_gt1_fast_path(
+                forward_batch
+            )
             local_input_ids = cp_split_and_rebuild_1d(forward_batch, input_ids)
             local_num_tokens = local_input_ids.shape[0]
             local_positions = cp_split_and_rebuild_position(forward_batch, positions)
@@ -302,6 +377,13 @@ class DeepseekModelNextN(nn.Module):
                 local_num_tokens=local_num_tokens,
             )
             if spec_hidden_states is None:
+                if must_use_batch_fast_path:
+                    self._fail_cp_draft_batch_gt1(
+                        forward_batch,
+                        "missing_local_spec_hidden",
+                        f"full_tokens={input_ids.shape[0]} "
+                        f"local_tokens={local_num_tokens}",
+                    )
                 use_cp_local_draft = False
             else:
                 positions = local_positions
@@ -312,6 +394,13 @@ class DeepseekModelNextN(nn.Module):
                         full_num_tokens=input_ids.shape[0],
                     )
                     if hidden_states is None:
+                        if must_use_batch_fast_path:
+                            self._fail_cp_draft_batch_gt1(
+                                forward_batch,
+                                "missing_local_embedding",
+                                f"full_tokens={input_ids.shape[0]} "
+                                f"local_tokens={local_num_tokens}",
+                            )
                         # Conservative compatibility fallback: embed full input
                         # so all TP ranks all-reduce the same shape, then CP-split.
                         self._debug_cp_draft_shared_kv(
@@ -335,6 +424,14 @@ class DeepseekModelNextN(nn.Module):
                         f"full_tokens={input_ids.shape[0]} "
                         f"local_tokens={local_num_tokens}"
                     )
+                    if must_use_batch_fast_path:
+                        self._fail_cp_draft_batch_gt1(
+                            forward_batch,
+                            "input_embeds_shape_mismatch",
+                            f"input_embed_tokens={input_embeds.shape[0]} "
+                            f"full_tokens={input_ids.shape[0]} "
+                            f"local_tokens={local_num_tokens}",
+                        )
                     use_cp_local_draft = False
 
         if not use_cp_local_draft:

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -1,4 +1,5 @@
 import ast
+import os
 from pathlib import Path
 import unittest
 import sys
@@ -33,6 +34,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+from sglang.srt.models.deepseek_nextn import DeepseekModelNextN
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
@@ -714,6 +716,42 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         ):
             cp_collect_last_token_hidden(torch.zeros((8, 1)), forward_batch, 2)
 
+    def test_deepseek_nextn_cp_draft_bs_gt1_fails_fast_on_hidden_shape_fallback(
+        self,
+    ):
+        import torch
+
+        model = DeepseekModelNextN.__new__(DeepseekModelNextN)
+        model._debug_cp_draft_shared_kv = lambda _message: None
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            batch_size=2,
+            extend_seq_lens_cpu=[4, 9],
+            nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2),
+        )
+
+        with (
+            patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}),
+            patch(
+                "sglang.srt.models.deepseek_nextn.get_attention_cp_rank",
+                return_value=0,
+            ),
+            patch(
+                "sglang.srt.models.deepseek_nextn.get_attention_cp_size",
+                return_value=8,
+            ),
+            self.assertRaisesRegex(
+                RuntimeError,
+                r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_batch_gt1_spec_hidden_shape_mismatch\]",
+            ),
+        ):
+            model._get_cp_local_spec_hidden_states(
+                forward_batch,
+                torch.zeros((3, 2)),
+                full_num_tokens=13,
+                local_num_tokens=8,
+            )
+
     def test_full_rerange_fails_fast_for_batch_metadata(self):
         import torch