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

Fix spec_v2 CP draft extend hidden-state shape mismatch

Browse Source 
EAGLEWorkerV2's target prefill always captured the global FULL hidden,
so the bs>1 NSA CP shared-KV draft extend received a global hidden whose
row count (real total tokens) matched neither the per-rank CP-local count
nor the MLP-sync CP-aligned padded count, tripping
[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_spec_hidden_shape_mismatch]
under SGLANG_ENABLE_SPEC_V2=1.

Mirror the legacy EAGLEWorker contract: add _can_use_cp_draft_shared_kv
and, when CP draft shared-KV applies, capture the CP-local hidden side
channel (CaptureHiddenMode.NULL + capture_draft_hidden_states) instead of
FULL. The v2 consumer already prefers draft_hidden_states when present
(commit 5e22279670 added the consumer but not this producer side).

Fixes the shape at the source rather than loosening the NextN fail-fast.
CP-off keeps FULL; non-CP / bs=1 unaffected; legacy v1 path untouched.
Validated on g0033 PD (prefill + 2-node decode): bs>1 up to 17 incl. the
former-crashing bs=6, 0 non-400 errors, GSM8K 0.965 / 0 invalid.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-13 09:06:00 +00:00
parent 1d168d061f
commit 2c7353d061
1 changed files with 42 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -692,13 +692,53 @@ class EAGLEWorkerV2(BaseSpecWorker):
# allocator and kv cache pool are shared with target worker, which are cleared in scheduler
pass
def _can_use_cp_draft_shared_kv(
self, model_worker_batch: ModelWorkerBatch
) -> bool:
"""Whether the CP-local draft-hidden side channel applies to this prefill.
Mirrors ``EAGLEWorker._can_use_cp_draft_shared_kv`` (the legacy v1 path).
When True the v2 target prefill must capture a CP-local hidden (NULL
capture mode + ``capture_draft_hidden_states``) instead of a global FULL
hidden, because the NSA CP shared-KV bs>1 draft extend
(``DeepseekModelNextN``) requires the target hidden row count to equal the
per-rank CP-local token count. Without it the target emits a global FULL
hidden that MLP-sync static padding then CP-misaligns, tripping the bs>1
``[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_spec_hidden_shape_mismatch]``.
"""
if not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
return False
if not (
model_worker_batch.forward_mode.is_extend()
or model_worker_batch.is_extend_in_batch
):
return False
draft_architectures = getattr(
self.draft_worker.draft_runner.model_config.hf_config,
"architectures",
[],
)
if "DeepseekV3ForCausalLMNextN" not in (draft_architectures or []):
return False
if not getattr(self.target_worker.model_runner, "uses_cp_shared_kv", False):
return False
if not getattr(self.draft_worker.draft_runner, "uses_cp_shared_kv", False):
return False
return True
def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
if (
model_worker_batch.forward_mode.is_extend()
or model_worker_batch.is_extend_in_batch
):
# Target prefill
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
# Target prefill. Keep the target hidden side channel CP-local when CP
# draft shared-KV is enabled (NULL capture + draft-hidden flag), exactly
# like the legacy worker; otherwise capture the global FULL hidden.
if self._can_use_cp_draft_shared_kv(model_worker_batch):
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.NULL
model_worker_batch.capture_draft_hidden_states = True
else:
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
batch_output = self.target_worker.forward_batch_generation(
model_worker_batch
)