minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Fix][BCG][Spec] Restore EAGLE prefill plumbing dropped by #23906 (#28870)

Browse Source
This commit is contained in:
Yuwei An
2026-06-22 01:54:54 -07:00
committed by GitHub
parent 93553a67a3
commit 2ce32366a0
3 changed files with 124 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+9 -13
View File
@@ -2594,25 +2594,21 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if self.is_draft_worker and not force_for_draft_worker:
return
# EAGLE-family target worker: the prefill graph captures
# CaptureHiddenMode.NULL, but target prefill needs FULL hidden states to
# feed the draft, so can_run_graph always rejects it and prefill runs
# eagerly. The graph is therefore never used; capturing it (now before
# the decode graph) can perturb backend state on FP4 / TRTLLM-MoE paths
# and corrupt decode replay, so skip its capture and route prefill
# through the eager runner — the runtime path either way. With
# enable_return_hidden_states the prefill graph is FULL and usable, so
# only skip when it would capture NULL.
# Skip prefill CG for EAGLE target on tc_piecewise: that backend
# captures CaptureHiddenMode.NULL while runtime requests FULL, so
# the captured graph is dead, and capturing it perturbs FP4 /
# TRTLLM-MoE state and corrupts decode replay (see #28386). BCG
# captures FULL for EAGLE target in PrefillCudaGraphRunner.__init__
# (restored from #25795), so it does NOT need this skip.
if (
self.spec_algorithm.is_eagle()
and not self.is_draft_worker
and not self.server_args.enable_return_hidden_states
and not check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
):
logger.info(
"Disable prefill CUDA graph for the EAGLE target worker: target "
"prefill needs FULL hidden states but the prefill graph captures "
"NULL, so the graph is unused; skipping its capture keeps decode "
"graph capture clean."
"Disable prefill CUDA graph for EAGLE target on tc_piecewise "
"to avoid FP4/MoE decode-replay corruption (#28386)."
)
self.prefill_cuda_graph_runner = self.eager_runner
return
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+69 -6
View File
@@ -47,6 +47,7 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import (
CudaGraphBufferRegistry,
build_prefill_registry,
)
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
@@ -148,6 +149,22 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
or model_runner.spec_algorithm.is_dflash()
):
self.capture_hidden_mode = CaptureHiddenMode.FULL
# EAGLE captures FULL hidden states for the target and LAST for the
# draft (can_run_graph rejects on mismatch). BCG only; tc_piecewise
# EAGLE is routed to eager in ModelRunner.init_prefill_cuda_graph.
_cg_cfg = model_runner.server_args.cuda_graph_config
_prefill_backend_name = (
_cg_cfg.prefill.backend if _cg_cfg is not None else Backend.TC_PIECEWISE
)
if (
_prefill_backend_name == Backend.BREAKABLE
and model_runner.spec_algorithm.is_eagle()
):
self.capture_hidden_mode = (
CaptureHiddenMode.LAST
if model_runner.is_draft_worker
else CaptureHiddenMode.FULL
)
self.mamba_track_enabled = self._is_mamba_track_enabled()
@@ -196,12 +213,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# Initialize the slot to None BEFORE constructing the backend:
# TcPiecewise runs its compile pass during __init__ which calls
# _run_dummy_forward -> capture_prepare, and capture_prepare reads
# self._prefill_static_buffers. self.layer_model has the same
# ordering requirement: _run_forward checks `self.layer_model is
# not None` to decide whether to call the inner stack or outer
# model.forward, and that check fires inside TcPiecewise's
# _run_compile_pass before backend resolution returns.
# self._prefill_static_buffers and self.static_draft_hidden_states.
# self.layer_model has the same ordering requirement: _run_forward
# checks `self.layer_model is not None` to decide whether to call
# the inner stack or outer model.forward, and that check fires
# inside TcPiecewise's _run_compile_pass before backend resolution
# returns.
self._prefill_static_buffers: Optional[Dict[str, torch.Tensor]] = None
self.static_draft_hidden_states: Optional[torch.Tensor] = None
self.layer_model = None
self.backend = resolve_prefill_backend(self)
if isinstance(self.backend, BreakableCudaGraphBackend):
@@ -211,6 +230,32 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
for name in _PREFILL_STATIC_FIELDS
}
# Static hidden_states buffer giving the captured graph a stable
# address; load_batch refreshes it from live spec_info at replay.
# Draft consumes the aux-concatenated hidden states from the target
# (e.g. EAGLE3 stacks 3 target layers), so read the dim from the
# draft model's input fc when available; fall back to the per-layer
# dim for arches without an fc projection.
if (
isinstance(self.backend, BreakableCudaGraphBackend)
and model_runner.is_draft_worker
and model_runner.spec_algorithm.is_eagle()
):
from sglang.srt.speculative.eagle_utils import get_draft_hidden_dim
inner = getattr(model_runner.model, "model", model_runner.model)
fc = getattr(inner, "fc", None)
hidden_dim = (
fc.in_features
if fc is not None and hasattr(fc, "in_features")
else get_draft_hidden_dim(model_runner)
)
with torch.device(self.device):
self.static_draft_hidden_states = torch.zeros(
(self.max_num_tokens, hidden_dim),
dtype=model_runner.dtype,
)
# Some attention backends (e.g. DSV4) opt into a captured-metadata
# contract under BCG: capture-time builds a per-bucket metadata
# object the backend then refreshes in place at replay. We honor
@@ -462,6 +507,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# logits_processor eagerly on top with live multi-req metadata.
return True
def _build_capture_spec_info(self, num_tokens: int):
if self.static_draft_hidden_states is None:
return None
from sglang.srt.speculative.eagle_info import EagleDraftInput
return EagleDraftInput(
hidden_states=self.static_draft_hidden_states[:num_tokens],
)
def capture_prepare(self, num_tokens: int) -> tuple[ForwardBatch, AttentionBackend]:
"""Build a dummy prefill ForwardBatch for capture/warmup at this shape.
@@ -572,7 +626,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
else None
),
spec_algorithm=None,
spec_info=None,
spec_info=self._build_capture_spec_info(num_tokens),
# Use self.capture_hidden_mode so dflash spec (which needs
# FULL aux hidden states) captures with the right mode.
# Ported from main #27468.
@@ -778,6 +832,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
if forward_batch.orig_seq_lens is not None:
s["orig_seq_lens"][:bs].copy_(forward_batch.orig_seq_lens)
# Refresh the static buffer the captured graph reads from.
if (
self.static_draft_hidden_states is not None
and forward_batch.spec_info is not None
):
self.static_draft_hidden_states[:num_tokens].copy_(
forward_batch.spec_info.hidden_states
)
self._prepare_forward_metadata_for_replay(
forward_batch, static_forward_batch, static_num_tokens
)