Carry CP-local hidden contract into EAGLE v2 prefill draft extend

SPEC_V2 builds an EAGLE draft input during target prefill. Under CP shared-KV, the target model may expose draft_hidden_states as a CP-local side channel before CP output collection. The v2 path was still passing global hidden_states and never marked the draft input as CP-local, unlike the non-v2 path.

Thread cp_local_hidden_states through the v2 _draft_extend_for_prefill helper and prefer draft_hidden_states when present. This preserves the semantic marker consumed by the static MLP-sync padding guard without changing the non-CP path.

Constraint: Absorb syh 5562937cf only; SPEC_V2 remains opt-in and broader SPEC_V2-on-CP validation is still separate.
Rejected: Infer CP-local hidden from tensor length | tensor length is ambiguous under static padding and bs>1 compute padding.
Confidence: high
Scope-risk: narrow
Directive: Keep EagleDraftInput.cp_local_hidden_states as an explicit semantic contract; do not replace it with shape-based inference.
Tested: Remote g0034 container red test failed before implementation with unexpected cp_local_hidden_states kwarg.
Tested: Remote g0034 container py_compile for eagle_worker_v2.py and test_nsa_cp_utils.py.
Tested: Remote g0034 container pytest target EAGLE marker tests: 2 passed.
Tested: Remote g0034 container pytest test_nsa_cp_utils.py -k eagle: 2 passed, 99 deselected.
Not-tested: Full ETE with SGLANG_ENABLE_SPEC_V2=1 on CP prefill.

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

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