From 2ce32366a073135b462778359b01228c0a043aac Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Mon, 22 Jun 2026 01:54:54 -0700 Subject: [PATCH] [Fix][BCG][Spec] Restore EAGLE prefill plumbing dropped by #23906 (#28870) --- .../sglang/srt/model_executor/model_runner.py | 22 +++--- .../runner/prefill_cuda_graph_runner.py | 75 +++++++++++++++++-- .../test_bcg_with_speculative_decoding.py | 46 ++++++++++++ 3 files changed, 124 insertions(+), 19 deletions(-) create mode 100644 test/registered/cuda_graph/breakable/test_bcg_with_speculative_decoding.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f45f027d0..f48b60b0b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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 diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index dba9899e7..7f5fb616f 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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 ) diff --git a/test/registered/cuda_graph/breakable/test_bcg_with_speculative_decoding.py b/test/registered/cuda_graph/breakable/test_bcg_with_speculative_decoding.py new file mode 100644 index 000000000..faa760373 --- /dev/null +++ b/test/registered/cuda_graph/breakable/test_bcg_with_speculative_decoding.py @@ -0,0 +1,46 @@ +"""Test breakable CUDA graph (BCG) coexisting with EAGLE3 speculative +decoding. Sibling of test_pcg_with_speculative_decoding.py — same +target/draft pair, only flips the prefill backend from tc_piecewise to +breakable. Verifies the draft-side BCG plumbing in PrefillCudaGraphRunner +stays wired (capture_hidden_mode for EAGLE, static_draft_hidden_states +buffer sized from the draft's fc input, EagleDraftInput at capture, and +the load_batch refresh). +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.server_fixtures.pcg_spec_fixture import PCGSpecBase + +register_cuda_ci(est_time=531, stage="base-b", runner_config="2-gpu-large") + + +class TestBCGWithEAGLE3(PCGSpecBase, unittest.TestCase): + """BCG + EAGLE3 on Qwen3-30B-A3B-Instruct-2507.""" + + model = "Qwen/Qwen3-30B-A3B-Instruct-2507" + server_args = [ + "--tp", + "2", + "--trust-remote-code", + "--cuda-graph-backend-prefill=breakable", + "--mem-fraction-static", + "0.6", + "--speculative-algorithm", + "EAGLE3", + "--speculative-draft-model-path", + "lmsys/SGLang-EAGLE3-Qwen3-30B-A3B-Instruct-2507-SpecForge-Nex", + "--speculative-num-steps", + "5", + "--speculative-eagle-topk", + "4", + "--speculative-num-draft-tokens", + "8", + ] + timeout_mult = 3 + server_env = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"} + accuracy_threshold = 0.75 + + +if __name__ == "__main__": + unittest.main()