chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Reduce prefill EAGLE memory pressure under CP shared KV

Browse Source 
Prefill CP only needs the local hidden shard for DeepSeek NextN draft extend. The change adds a draft shared-KV path that captures target hidden locally, feeds only the CP-local slice into the draft model, and keeps draft KV writes/transfers on the same shared logical-to-physical page mapping as target KV.\n\nDebug logs are gated behind SGLANG_CP_DRAFT_SHARED_KV_DEBUG and cover scheduler pool selection, KV manager buffer registration, local physical writes, prefill sender filtering, transfer pages, and decode commit metadata so ETE runs can prove draft KV is sharded rather than full-concatenated on a prefill rank.\n\nConstraint: Prefill runs CP while decode remains DP, so prefill must avoid full hidden/KV materialization but decode still receives full logical KV pages.\nRejected: Keep draft extend on full hidden state | preserves correctness but wastes prefill memory and defeats CP shared-KV intent.\nRejected: Transfer draft KV with a separate mapping | target and draft pools share req_to_token logical indices, so duplicating mapping adds risk without benefit.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the debug logs until ETE evidence confirms draft MLA/index writes and transfer pages are CP-sharded on all ranks.\nTested: Remote compileall for changed CP draft, transfer, scheduler, NSA index, MLA write, and EAGLE files.\nNot-tested: Full GLM-5 EAGLE ETE with SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1 after this logging addition; local pytest intentionally not run.
This commit is contained in:
laoyao0822
2026-05-13 22:29:18 +08:00
parent 3fc7a5c18c
commit 99b669f8b9
16 changed files with 951 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -5,6 +5,7 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.distributed import get_tp_group
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
@@ -276,6 +277,36 @@ class EAGLEWorker(TpModelWorker):
def draft_model_runner(self):
return self.model_runner
def _debug_cp_draft_shared_kv(self, message: str):
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
logger.info("[CP_DRAFT_SHARED_KV] %s", message)
def _can_use_cp_draft_shared_kv(self, batch: ScheduleBatch) -> bool:
if not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
return False
if not (batch.forward_mode.is_extend() or batch.is_extend_in_batch):
self._debug_cp_draft_shared_kv("fallback reason=non_extend_mode")
return False
draft_architectures = getattr(
self.draft_model_runner.model_config.hf_config, "architectures", []
)
if "DeepseekV3ForCausalLMNextN" not in (draft_architectures or []):
self._debug_cp_draft_shared_kv(
f"fallback reason=unsupported_arch architectures={draft_architectures}"
)
return False
if not getattr(self.target_worker.model_runner, "uses_cp_shared_kv", False):
self._debug_cp_draft_shared_kv(
"fallback reason=target_cp_shared_kv_disabled"
)
return False
if not getattr(self.draft_model_runner, "uses_cp_shared_kv", False):
self._debug_cp_draft_shared_kv(
"fallback reason=draft_cp_shared_kv_disabled"
)
return False
return True
def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult:
"""Run speculative decoding forward.
@@ -298,9 +329,21 @@ class EAGLEWorker(TpModelWorker):
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
draft_hidden_states = (
logits_output.draft_hidden_states
if logits_output.draft_hidden_states is not None
else logits_output.hidden_states
)
if (
envs.SGLANG_CP_DRAFT_SHARED_KV.get()
and draft_hidden_states is None
):
self._debug_cp_draft_shared_kv(
"fallback_failed reason=missing_target_hidden"
)
self.forward_draft_extend(
batch,
logits_output.hidden_states,
draft_hidden_states,
next_token_ids,
seq_lens_cpu,
logits_output.mm_input_embeds,
@@ -367,15 +410,22 @@ class EAGLEWorker(TpModelWorker):
batch: The batch to run. States could be modified.
Returns:
logits_output: The output of logits. It will contain the full hidden states.
logits_output: The output of logits. It contains full hidden states on the
legacy path, or CP-local draft hidden states in `draft_hidden_states`
when CP draft shared-KV is enabled.
next_token_ids: Next token ids generated.
seq_lens_cpu: CPU copy of sequence lengths for the draft prefill path.
can_run_cuda_graph: Whether the target prefill ran with cuda graph.
"""
# Forward with the target model and get hidden states.
# We need the full hidden states to prefill the KV cache of the draft model.
# Forward with the target model and get hidden states for draft prefill.
# CP draft shared-KV keeps this hidden side channel CP-local so target
# logits can still use the narrow output path.
model_worker_batch = batch.get_model_worker_batch()
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
if self._can_use_cp_draft_shared_kv(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_result = self.target_worker.forward_batch_generation(model_worker_batch)
logits_output, next_token_ids = (
batch_result.logits_output,