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 )