From f0c6b892b6b00fcf3421e419e0ddfc5af34960f7 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Mon, 22 Jun 2026 01:05:49 +0800 Subject: [PATCH] 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) --- .../sglang/srt/speculative/eagle_worker_v2.py | 58 +++++-- .../test_eagle_worker_v2_cp_hidden.py | 164 ++++++++++++++++++ 2 files changed, 208 insertions(+), 14 deletions(-) create mode 100644 test/registered/unit/speculative/test_eagle_worker_v2_cp_hidden.py diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 567dbb979..13d4e6c81 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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: diff --git a/test/registered/unit/speculative/test_eagle_worker_v2_cp_hidden.py b/test/registered/unit/speculative/test_eagle_worker_v2_cp_hidden.py new file mode 100644 index 000000000..c1166a04f --- /dev/null +++ b/test/registered/unit/speculative/test_eagle_worker_v2_cp_hidden.py @@ -0,0 +1,164 @@ +import contextlib +import ast +from enum import IntEnum +from pathlib import Path +from types import SimpleNamespace + +import torch + + +class CaptureHiddenMode(IntEnum): + NULL = 0 + LAST = 1 + FULL = 2 + + +class _EnabledEnv: + def get(self): + return True + + +def _load_eagle_worker_v2_under_test(): + source_path = ( + Path(__file__).parents[4] + / "python" + / "sglang" + / "srt" + / "speculative" + / "eagle_worker_v2.py" + ) + module = ast.parse(source_path.read_text()) + class_node = next( + node + for node in module.body + if isinstance(node, ast.ClassDef) and node.name == "EAGLEWorkerV2" + ) + wanted = { + "_can_use_cp_draft_shared_kv", + "forward_batch_generation", + } + methods = [ + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + test_class = ast.ClassDef( + name="_EAGLEWorkerV2UnderTest", + bases=[], + keywords=[], + body=methods, + decorator_list=[], + ) + test_module = ast.Module(body=[test_class], type_ignores=[]) + ast.fix_missing_locations(test_module) + namespace = {} + exec( + compile(test_module, str(source_path), "exec"), + { + "CaptureHiddenMode": CaptureHiddenMode, + "ModelWorkerBatch": object, + "envs": SimpleNamespace(SGLANG_CP_DRAFT_SHARED_KV=_EnabledEnv()), + "speculative_moe_backend_context": contextlib.nullcontext, + "speculative_moe_a2a_backend_context": contextlib.nullcontext, + }, + namespace, + ) + return namespace["_EAGLEWorkerV2UnderTest"] + + +class _ExtendMode: + def is_extend(self): + return True + + +class _TargetWorker: + def __init__(self, logits_output): + self.model_runner = SimpleNamespace(uses_cp_shared_kv=True) + self._logits_output = logits_output + self.capture_state_seen = None + + def forward_batch_generation(self, model_worker_batch): + self.capture_state_seen = ( + model_worker_batch.capture_hidden_mode, + model_worker_batch.capture_draft_hidden_states, + ) + return SimpleNamespace( + logits_output=self._logits_output, + next_token_ids=torch.tensor([1], dtype=torch.int64), + ) + + +class _DraftWorker: + def __init__(self): + self.draft_runner = SimpleNamespace( + tp_group=object(), + uses_cp_shared_kv=True, + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + architectures=["DeepseekV3ForCausalLMNextN"], + ) + ), + ) + self.capture_state_seen = None + self.prefill_args = None + + def draft_tp_context(self, _tp_group): + return contextlib.nullcontext() + + def _draft_extend_for_prefill( + self, + model_worker_batch, + draft_hidden, + next_token_ids, + mm_input_embeds, + *, + cp_local_hidden_states, + ): + self.capture_state_seen = ( + model_worker_batch.capture_hidden_mode, + model_worker_batch.capture_draft_hidden_states, + ) + self.prefill_args = SimpleNamespace( + draft_hidden=draft_hidden, + cp_local_hidden_states=cp_local_hidden_states, + ) + return "next-draft-input" + + +def test_cp_draft_hidden_capture_is_target_only_and_released(): + draft_hidden = torch.randn(2, 4) + logits_output = SimpleNamespace( + draft_hidden_states=draft_hidden, + hidden_states=torch.randn(4, 4), + mm_input_embeds=None, + ) + target_worker = _TargetWorker(logits_output) + draft_worker = _DraftWorker() + worker_cls = _load_eagle_worker_v2_under_test() + worker = worker_cls.__new__(worker_cls) + worker.target_worker = target_worker + worker.draft_worker = draft_worker + + batch = SimpleNamespace( + forward_mode=_ExtendMode(), + is_extend_in_batch=False, + spec_info=None, + capture_hidden_mode=None, + capture_draft_hidden_states=False, + ) + + output = worker.forward_batch_generation(batch) + + assert target_worker.capture_state_seen == ( + CaptureHiddenMode.NULL, + True, + ) + assert draft_worker.capture_state_seen == ( + CaptureHiddenMode.LAST, + False, + ) + assert draft_worker.prefill_args.cp_local_hidden_states is True + assert draft_worker.prefill_args.draft_hidden is draft_hidden + assert output.next_draft_input == "next-draft-input" + assert logits_output.draft_hidden_states is None + assert batch.capture_draft_hidden_states is False