diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md index 29112caec..9ac605096 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md +++ b/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。 diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index d237e9894..df631d5f9 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/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 验证场景 远端： diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 5a0927ff8..66c7db029 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/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: diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index f284ac516..9eae9249f 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/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