diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 0516e5a70..9197fc62a 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -176,9 +176,7 @@ class ScheduleBatchDisaggregationDecodeMixin: ) spec_info.capture_hidden_mode = CaptureHiddenMode.LAST if self.enable_overlap: - from sglang.srt.managers.overlap_utils import FutureIndices - - spec_info.future_indices = FutureIndices(indices=self.req_pool_indices) + spec_info.future_indices = self.req_pool_indices future_map.publish(spec_info.future_indices, self.seq_lens) future_map.stash(spec_info.future_indices, spec_info) self.spec_info = spec_info diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 52784d4c6..0330391a0 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -1,6 +1,5 @@ from __future__ import annotations -from dataclasses import dataclass from typing import TYPE_CHECKING, Union import torch @@ -36,11 +35,6 @@ else: _resolve_future_token_ids = _resolve_future_token_ids_native -@dataclass -class FutureIndices: - indices: torch.Tensor - - class FutureMap: """Cross-iter relay buffer for values the next iter's schedule cannot compute locally (e.g. spec_v2 seq_lens after accept_lens, sampled tokens). @@ -113,7 +107,7 @@ class FutureMap: if draft_input is None: # FIXME(lsyin): only prefill; not compatible with mixed mode return - indices = draft_input.future_indices.indices + indices = draft_input.future_indices # FIXME: indices = batch.req_pool_indices, pinned 2 iters via # record_batch_in_overlap; record_stream here is redundant. indices.record_stream(torch.get_device_module(self.device).current_stream()) @@ -124,12 +118,12 @@ class FutureMap: draft_input.hidden_states = self.hidden_states_buf[indices] def set_input_ids_sentinel( - self, batch: ScheduleBatch, future_indices: FutureIndices + self, batch: ScheduleBatch, future_indices: torch.Tensor ) -> None: # Sentinel for the decode portion so mixed batches can cat extend # (positive real tokens) + decode (negative sentinels) into one # input_ids; resolve_future translates negatives via output_tokens_buf. - batch.input_ids = -future_indices.indices + batch.input_ids = -future_indices def resolve_seq_lens_cpu(self, batch: ScheduleBatch) -> None: # Lazy pull from new_seq_lens_buf for spec_v2 (accept_lens not known to @@ -140,15 +134,13 @@ class FutureMap: return if self.publish_ready is not None: self.publish_ready.wait() - new_seq_lens = self.new_seq_lens_buf[fi.indices] + new_seq_lens = self.new_seq_lens_buf[fi] batch.seq_lens = new_seq_lens batch.seq_lens_cpu = new_seq_lens.cpu() batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) - def publish( - self, future_indices: FutureIndices, new_seq_lens: torch.Tensor - ) -> None: - indices = future_indices.indices + def publish(self, future_indices: torch.Tensor, new_seq_lens: torch.Tensor) -> None: + indices = future_indices if indices.shape[0] == 0: return # DP idle self.new_seq_lens_buf[indices] = new_seq_lens.to(self.new_seq_lens_buf.dtype) @@ -160,10 +152,10 @@ class FutureMap: def stash( self, - future_indices: FutureIndices, + future_indices: torch.Tensor, payload: Union[torch.Tensor, EagleDraftInput], ) -> None: - indices = future_indices.indices + indices = future_indices if indices.shape[0] == 0: # DP idle: payload is empty stub; lazy-init shape peek would IndexError. return diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 621941f3a..6c9a3e8e9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -144,7 +144,6 @@ from sglang.srt.managers.io_struct import ( UpdateWeightsFromTensorReqInput, ) from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors -from sglang.srt.managers.overlap_utils import FutureIndices from sglang.srt.managers.prefill_delayer import ( PrefillDelayer, PrefillDelayerSinglePassExecutor, @@ -2847,7 +2846,7 @@ class Scheduler( self.future_map.resolve_seq_lens_cpu(batch) with self._overlap_forward_isolation(batch): - future_indices = FutureIndices(indices=batch.req_pool_indices) + future_indices = batch.req_pool_indices # Spec_v2 fires on_publish mid-worker (between verify and # draft_extend) so schedule prep can overlap with draft_extend. diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index bffe856b9..ba545ccf7 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -10,7 +10,6 @@ import torch from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.managers.overlap_utils import FutureIndices from sglang.srt.managers.schedule_batch import Req from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.server_args import ServerArgs @@ -40,7 +39,7 @@ class GenerationBatchResult: # For overlap scheduling copy_done: Optional[torch.cuda.Event] = None delay_sample_func: Optional[callable] = None - future_indices: Optional[FutureIndices] = None + future_indices: Optional[torch.Tensor] = None speculative_num_draft_tokens: Optional[int] = None # FIXME(lsyin): maybe move to a better place? diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index b54f75a81..79f797b1e 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -16,7 +16,6 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.sampler import apply_custom_logit_processor -from sglang.srt.managers.overlap_utils import FutureIndices from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.common import ( @@ -693,8 +692,8 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): num_tokens_per_req: int = -1 num_tokens_for_logprob_per_req: int = -1 - # V2 overlap worker only - future_indices: Optional[FutureIndices] = None + # V2 overlap worker only: req_pool_indices used as buf slot keys. + future_indices: Optional[torch.Tensor] = None # V2 reuses `EagleDraftInput` across phases (V1 has a separate # `EagleDraftExtendInput` for these). Set during V2's draft-extend. num_correct_drafts: Optional[torch.Tensor] = None @@ -745,7 +744,7 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): def filter_batch(self, new_indices: torch.Tensor, has_been_filtered: bool = True): if self.future_indices is not None: - self.future_indices.indices = self.future_indices.indices[new_indices] + self.future_indices = self.future_indices[new_indices] return strict_check = envs.SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK.get() @@ -775,10 +774,8 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): def merge_batch(self, spec_info: "EagleDraftInput"): if self.future_indices is not None: assert spec_info.future_indices is not None - self.future_indices = FutureIndices( - indices=torch.cat( - [self.future_indices.indices, spec_info.future_indices.indices] - ) + self.future_indices = torch.cat( + [self.future_indices, spec_info.future_indices] ) return