diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 248e7015c..a086f7e09 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -498,6 +498,7 @@ class EagleDraftWorker(BaseDraftWorker): target_hidden_states: torch.Tensor, next_token_ids: torch.Tensor, mm_input_embeds: Optional[torch.Tensor] = None, + cp_local_hidden_states: bool = False, ): """ Run draft model extend to correctly fill the KV cache. @@ -525,6 +526,10 @@ class EagleDraftWorker(BaseDraftWorker): # draft mode is same with decode mode, only 1 token per req num_tokens_per_req=1, num_tokens_for_logprob_per_req=1, + # Carry the CP-local contract so the static MLP-sync padding guard + # does not reject a draft hidden that is legitimately CP-local-sized. + # Mirrors the non-v2 path in eagle_worker.py. + cp_local_hidden_states=cp_local_hidden_states, ) batch.spec_info = next_draft_input @@ -703,12 +708,20 @@ class EAGLEWorkerV2(BaseSpecWorker): with self.draft_worker.draft_tp_context( self.draft_worker.draft_runner.tp_group ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): + # Prefer the target model's CP-local side channel when CP prefill + # emitted one. Without this marker, ForwardBatch static MLP-sync + # padding can treat a CP-local hidden tensor as globally padded and + # reject the v2 draft extend. + lo = batch_output.logits_output + cp_local = lo.draft_hidden_states is not None + draft_hidden = lo.draft_hidden_states if cp_local else lo.hidden_states batch_output.next_draft_input = ( self.draft_worker._draft_extend_for_prefill( model_worker_batch, - batch_output.logits_output.hidden_states, + draft_hidden, batch_output.next_token_ids, - batch_output.logits_output.mm_input_embeds, + lo.mm_input_embeds, + cp_local_hidden_states=cp_local, ) ) return batch_output diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 2a46a249d..7301d3666 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -4092,6 +4092,64 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertIs(draft_input.hidden_states, draft_output_hidden) self.assertFalse(draft_input.cp_local_hidden_states) + def test_eagle_v2_draft_extend_for_prefill_preserves_cp_local_hidden_marker( + self, + ): + import torch + + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.speculative import eagle_worker_v2 + + worker = object.__new__(eagle_worker_v2.EagleDraftWorker) + worker.topk = 1 + forward_batch = SimpleNamespace(mm_input_embeds=None) + + class DraftRunner: + def forward(self, _forward_batch): + return SimpleNamespace( + logits_output=LogitsProcessorOutput( + next_token_logits=torch.tensor([[0.1, 0.9]]), + hidden_states=torch.tensor([[3.0, 4.0]]), + ) + ) + + worker.draft_runner = DraftRunner() + batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + extend_seq_lens=[2], + input_ids=torch.tensor([10, 11]), + seq_lens=[2], + spec_info=None, + ) + target_hidden_states = torch.tensor([[1.0, 2.0]]) + next_token_ids = torch.tensor([99]) + init_new_checks = [] + + def init_new_side_effect(_batch, _draft_runner): + init_new_checks.append(True) + self.assertTrue(_batch.spec_info.cp_local_hidden_states) + self.assertIs(_batch.spec_info.hidden_states, target_hidden_states) + return forward_batch + + with patch.object( + eagle_worker_v2.ForwardBatch, + "init_new", + side_effect=init_new_side_effect, + ): + next_draft_input = ( + eagle_worker_v2.EagleDraftWorker._draft_extend_for_prefill( + worker, + batch, + target_hidden_states, + next_token_ids, + cp_local_hidden_states=True, + ) + ) + + self.assertEqual(init_new_checks, [True]) + self.assertTrue(next_draft_input.cp_local_hidden_states) + self.assertEqual(next_draft_input.hidden_states.tolist(), [[3.0, 4.0]]) + def test_indexer_in_seq_cp_pair_composes_current_only_index_reuse(self): import torch