From 2c7353d061230732021571d0da03980d9aa31c1a Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 13 Jun 2026 09:06:00 +0000 Subject: [PATCH] Fix spec_v2 CP draft extend hidden-state shape mismatch 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) --- .../sglang/srt/speculative/eagle_worker_v2.py | 44 ++++++++++++++++++- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index a086f7e09..567dbb979 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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 )