chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Prevent batched CP draft from silently leaving the target path

Browse Source 
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.
This commit is contained in:
laoyao0822
2026-06-04 03:47:56 +08:00
parent 3e3f1b776b
commit d7723aca07
4 changed files with 155 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

99
python/sglang/srt/models/deepseek_nextn.py
View File

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