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)
This commit is contained in:
laoyao0822
@@ -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:
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user