From bad83ab427bb0b15c0db968861b45954f477007a Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Sun, 31 May 2026 09:48:30 +0800
Subject: [PATCH] Fix the EAGLE chunked-prefill next-token chain (#26329)
 (#26800)

---
 python/sglang/srt/managers/schedule_batch.py    | 15 +++++++++++++++
 python/sglang/srt/speculative/eagle_utils.py    | 17 +++++++++++++++++
 .../sglang/srt/speculative/eagle_worker_v2.py   |  5 ++++-
 .../speculative/multi_layer_eagle_worker_v2.py  |  1 +
 4 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py
index cf7e304e7..68a9c73a8 100755
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -1463,6 +1463,17 @@ def set_mamba_track_indices_from_reqs(batch):
     )
 
 
+def _compute_chunked_req_next_prompt_token(
+    chunked_req: Optional[Req],
+) -> Optional[int]:
+    if chunked_req is None:
+        return None
+    fill_len = len(chunked_req.fill_ids)
+    if fill_len >= len(chunked_req.origin_input_ids):
+        return None
+    return int(chunked_req.origin_input_ids[fill_len])
+
+
 @dataclasses.dataclass
 class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     """Store all information of a batch on the scheduler."""
@@ -1494,6 +1505,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
 
     # For chunked prefill in PP
     chunked_req: Optional[Req] = None
+    chunked_req_next_prompt_token: Optional[int] = None
     contains_last_prefill_chunk: bool = True
 
     # For DP attention
@@ -1660,6 +1672,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             return_hidden_states=any(req.return_hidden_states for req in reqs),
             is_prefill_only=all(req.is_prefill_only for req in reqs),
             chunked_req=chunked_req,
+            chunked_req_next_prompt_token=_compute_chunked_req_next_prompt_token(
+                chunked_req
+            ),
             dllm_config=dllm_config,
         )
         return batch
diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py
index 97b2f0fc3..db435075a 100644
--- a/python/sglang/srt/speculative/eagle_utils.py
+++ b/python/sglang/srt/speculative/eagle_utils.py
@@ -66,10 +66,27 @@ def apply_eagle_prefill_input_rotation(
     seg_ends = extend_lens.cumsum(0) - 1
     rotated = torch.empty_like(batch.input_ids)
     rotated[:-1] = batch.input_ids[1:]
+    # TODO: chunked-prefill chain divergence at non-final-chunk seg end; fix per PR #26329.
     rotated[seg_ends] = next_token_ids.to(batch.input_ids.dtype)
     batch.input_ids = rotated
 
 
+def _eagle_prefill_tail_tokens(
+    batch: ScheduleBatch, next_token_ids: torch.Tensor
+) -> torch.Tensor:
+    """Per-seq tail token for EAGLE prefill rotation; uses next prompt token for
+    non-final chunks (chunked-prefill chain consistency, see PR #26329)."""
+    tail_tokens = next_token_ids.to(batch.input_ids.dtype)
+    next_prompt_token = batch.chunked_req_next_prompt_token
+    if next_prompt_token is not None:
+        for i, r in enumerate(batch.reqs):
+            if r is batch.chunked_req:
+                tail_tokens = tail_tokens.clone()
+                tail_tokens[i] = next_prompt_token
+                break
+    return tail_tokens
+
+
 def organize_draft_results(
     score_list: List[torch.Tensor],
     token_list: List[torch.Tensor],
diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index 8672913af..311e7bd93 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/eagle_worker_v2.py
@@ -55,6 +55,7 @@ from sglang.srt.speculative.eagle_info_v2 import (
 )
 from sglang.srt.speculative.eagle_utils import (
     TreeMaskMode,
+    _eagle_prefill_tail_tokens,
     build_tree_kernel_efficient,
     per_step_draft_out_cache_loc,
 )
@@ -579,11 +580,13 @@ class EagleDraftWorker(BaseDraftWorker):
         """
         # Construct input_ids
         if not batch.forward_mode.is_idle():
+            # Chunked-prefill-aware tail tokens (see PR #26329).
+            tail_tokens = _eagle_prefill_tail_tokens(batch, next_token_ids)
             pt = 0
             for i, extend_len in enumerate(batch.extend_lens):
                 input_ids = batch.input_ids[pt : pt + extend_len]
                 batch.input_ids[pt : pt + extend_len] = torch.cat(
-                    (input_ids[1:], next_token_ids[i].reshape(1))
+                    (input_ids[1:], tail_tokens[i].reshape(1))
                 )
                 pt += extend_len
 
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
index 24a57375c..725142669 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
@@ -410,6 +410,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         forward_batch = ForwardBatch.init_new(batch, self.draft_runner_list[0])
 
         # Construct input_ids
+        # TODO: same chunked-prefill chain divergence as PR #26329.
         if not batch.forward_mode.is_idle():
             rotate_input_ids_triton(
                 forward_batch.input_ids,