Release CP draft-hidden buffers after EAGLE target prefill

CP draft shared-KV target prefill only needs the CP-local draft hidden side channel until draft prefill consumes it.  The reused ModelWorkerBatch was left with capture_draft_hidden_states enabled, and LogitsProcessorOutput retained the extra draft_hidden_states tensor past that point.  Restore the capture flags after target prefill, disable draft-hidden capture for draft prefill, and drop the consumed reference.

Constraint: Large CP prefill batches are memory-sensitive; an extra hidden tensor held across the speculative step materially increases peak memory.

Rejected: Leave capture_draft_hidden_states sticky across draft prefill | it can capture or retain hidden state outside the target-side CP-local side channel.

Confidence: high

Scope-risk: narrow

Directive: Keep CP-local draft hidden as a target-prefill-only side channel unless the draft worker explicitly owns a new lifetime contract.

Tested: g0034 cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_token_slot_remap_cache_distinguishes_same_storage_views test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_paged_slot_remap_cache_distinguishes_same_storage_views test/registered/unit/speculative/test_eagle_worker_v2_cp_hidden.py

Not-tested: full speculative integration/E2E workload.
(cherry picked from commit f31ef2293ce0cdda2359a5f4713996d78b47f9df)

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -734,17 +734,38 @@ class EAGLEWorkerV2(BaseSpecWorker):
             # 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(
+            can_use_cp_draft_shared_kv = self._can_use_cp_draft_shared_kv(
                 model_worker_batch
             )
+            if can_use_cp_draft_shared_kv:
+                old_capture_hidden_mode = model_worker_batch.capture_hidden_mode
+                old_capture_draft_hidden_states = (
+                    model_worker_batch.capture_draft_hidden_states
+                )
+                model_worker_batch.capture_hidden_mode = CaptureHiddenMode.NULL
+                model_worker_batch.capture_draft_hidden_states = True
+                try:
+                    batch_output = self.target_worker.forward_batch_generation(
+                        model_worker_batch
+                    )
+                finally:
+                    # This side channel is target-prefill-only.  Do not let the
+                    # reused ModelWorkerBatch make the draft prefill capture an
+                    # additional hidden tensor and hold it live for the rest of the
+                    # speculative step.
+                    model_worker_batch.capture_hidden_mode = old_capture_hidden_mode
+                    model_worker_batch.capture_draft_hidden_states = (
+                        old_capture_draft_hidden_states
+                    )
+            else:
+                model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
+                batch_output = self.target_worker.forward_batch_generation(
+                    model_worker_batch
+                )
 
             # Draft prefill
             model_worker_batch.capture_hidden_mode = CaptureHiddenMode.LAST
+            model_worker_batch.capture_draft_hidden_states = False
             with self.draft_worker.draft_tp_context(
                 self.draft_worker.draft_runner.tp_group
             ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
@@ -755,15 +776,24 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 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,
-                        draft_hidden,
-                        batch_output.next_token_ids,
-                        lo.mm_input_embeds,
-                        cp_local_hidden_states=cp_local,
+                try:
+                    batch_output.next_draft_input = (
+                        self.draft_worker._draft_extend_for_prefill(
+                            model_worker_batch,
+                            draft_hidden,
+                            batch_output.next_token_ids,
+                            lo.mm_input_embeds,
+                            cp_local_hidden_states=cp_local,
+                        )
                     )
-                )
+                finally:
+                    if cp_local:
+                        # The CP-local hidden is consumed by draft prefill above.
+                        # Drop the extra LogitsProcessorOutput reference so the
+                        # caching allocator can reuse this target-side buffer
+                        # before the next large prefill.
+                        lo.draft_hidden_states = None
+                    draft_hidden = None
                 return batch_output
         else:
             if model_worker_batch.spec_info is None: