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.
2026-06-09 22:42:54 +08:00
parent df2a3696cd
commit 5e22279670
2 changed files with 73 additions and 2 deletions
@@ -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