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

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

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
         )

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()