disagg prebuilt: drop dead prepare_for_extend shift (#25819)

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -175,9 +175,6 @@ class ScheduleBatchDisaggregationDecodeMixin:
                 bonus_tokens=last_tokens_tensor,
                 new_seq_lens=self.seq_lens,
             )
-            # prepare_for_extend shifts batch.input_ids in place — keep it
-            # as the prefill prompt, not the [bs] last-token tensor.
-            spec_info.prepare_for_extend(self)
             spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
             if self.enable_overlap:
                 spec_info.future_indices = future_map.alloc_future_indices(

python/sglang/srt/speculative/eagle_info.py

@@ -708,22 +708,6 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
     def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
         return self.num_tokens_per_req, self.num_tokens_for_logprob_per_req
 
-    def prepare_for_extend(self, batch: ScheduleBatch):
-
-        if batch.forward_mode.is_idle():
-            return
-
-        # Prefill only generate 1 token.
-        assert len(self.bonus_tokens) == len(batch.seq_lens)
-
-        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:], self.bonus_tokens[i].reshape(1))
-            )
-            pt += extend_len
-
     @classmethod
     def hidden_size_for(cls, worker) -> Optional[int]:
         """Decode-phase `hidden_states` width: draft self-chain output

python/sglang/srt/speculative/eagle_utils.py

@@ -1,11 +1,16 @@
+from __future__ import annotations
+
 import math
 from enum import IntEnum
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional
 
 import torch
 
 from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
@@ -17,6 +22,29 @@ if _is_cuda or _is_hip or _is_musa:
     )
 
 
+def apply_eagle_prefill_input_rotation(
+    batch: ScheduleBatch, next_token_ids: torch.Tensor
+) -> None:
+    """EAGLE input rotation for draft prefill.
+
+    Each req's slice [t_0..t_{n-1}] -> [t_1..t_{n-1}, t_n] with
+    t_n = next_token_ids[i]. Aligns draft's position-i hidden with
+    target's label at i+1 — the basis of EAGLE chain prediction.
+    Vectorized: one whole-tensor left shift + scatter at segment tails.
+    """
+    if batch.forward_mode.is_idle():
+        return
+    assert len(next_token_ids) == len(batch.seq_lens)
+    extend_lens = torch.tensor(
+        batch.extend_lens, dtype=torch.int64, device=batch.input_ids.device
+    )
+    seg_ends = extend_lens.cumsum(0) - 1
+    rotated = torch.empty_like(batch.input_ids)
+    rotated[:-1] = batch.input_ids[1:]
+    rotated[seg_ends] = next_token_ids.to(batch.input_ids.dtype)
+    batch.input_ids = rotated
+
+
 def organize_draft_results(
     score_list: List[torch.Tensor],
     token_list: List[torch.Tensor],

python/sglang/srt/speculative/eagle_worker.py

@@ -51,6 +51,7 @@ from sglang.srt.speculative.eagle_info import (
     EagleVerifyOutput,
 )
 from sglang.srt.speculative.eagle_utils import (
+    apply_eagle_prefill_input_rotation,
     build_tree_kernel_efficient,
     organize_draft_results,
 )
@@ -1105,7 +1106,7 @@ class EAGLEWorker(TpModelWorker):
             num_tokens_for_logprob_per_req=1,
         )
         batch.return_hidden_states = False
-        batch.spec_info.prepare_for_extend(batch)
+        apply_eagle_prefill_input_rotation(batch, next_token_ids)
         capture_mode = (
             CaptureHiddenMode.NULL
             if self.speculative_algorithm.is_standalone()

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -41,6 +41,7 @@ from sglang.srt.speculative.eagle_info import (
     EagleVerifyOutput,
 )
 from sglang.srt.speculative.eagle_utils import (
+    apply_eagle_prefill_input_rotation,
     build_tree_kernel_efficient,
     organize_draft_results,
 )
@@ -651,7 +652,7 @@ class MultiLayerEagleWorker(TpModelWorker):
             num_tokens_for_logprob_per_req=1,
         )
         batch.return_hidden_states = False
-        batch.spec_info.prepare_for_extend(batch)
+        apply_eagle_prefill_input_rotation(batch, next_token_ids)
         capture_mode = (
             CaptureHiddenMode.NULL
             if self.speculative_algorithm.is_standalone()