diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index d04802179..aca442102 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/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( diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index c5d0b9ac3..f7d767da1 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 881ad2075..25c99cac9 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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], diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 6ca7632bf..88488bf88 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index c8f112360..974d2a669 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/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()