drop FutureIndices wrapper class (#26085)

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

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.

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?