drop FutureIndices wrapper class (#26085)

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

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?

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