Address overlap future token map by request-pool index (#25862)

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -177,9 +177,9 @@ class ScheduleBatchDisaggregationDecodeMixin:
             )
             spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
             if self.enable_overlap:
-                spec_info.future_indices = future_map.alloc_future_indices(
-                    len(self.seq_lens)
-                )
+                from sglang.srt.managers.overlap_utils import FutureIndices
+
+                spec_info.future_indices = FutureIndices(indices=self.req_pool_indices)
                 future_map.store_to_map_for_new_batch(
                     spec_info.future_indices, spec_info
                 )

python/sglang/srt/managers/overlap_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -11,6 +11,7 @@ from sglang.srt.utils import is_cuda, is_hip
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
     from sglang.srt.managers.scheduler import GenerationBatchResult
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
     from sglang.srt.speculative.eagle_info import EagleDraftInput
     from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
@@ -39,73 +40,55 @@ else:
 @dataclass
 class FutureIndices:
     indices: torch.Tensor
-    interval: Optional[slice] = None
 
 
 class FutureMap:
     def __init__(
         self,
-        max_running_requests: int,
-        chunked_prefill_size: int,
-        context_len: int,
         device: torch.device,
         spec_algo: SpeculativeAlgorithm,
+        req_to_token_pool: ReqToTokenPool,
     ):
-        # FIXME: the calculation of future_limit and future_buffer_len maybe too conservative
-        self.future_ct = 0
-
-        # Circular buffer layout (wraps in this order):
-        # Running decode batch -> Prefill chunk 1 -> ... -> Prefill chunk N
-        # A running decode batch's result will be resolved after all prefill chunks are done.
-        # reserve `max_num_chunks` extra future slots on top of `max_running_requests * 3`.
-        max_num_chunks = (
-            (context_len + chunked_prefill_size - 1) // chunked_prefill_size
-            if chunked_prefill_size
-            else 0
-        )
-        self.future_limit = max_running_requests * (3 + max_num_chunks)
-        # Adding 2 * max_running_requests to future_limit ensures the buffer is sufficiently large.
-        self.future_buffer_len = self.future_limit + 2 * max_running_requests
+        # All buffers are indexed by req_pool_idx. Slot 0 mirrors the KV cache
+        # pool's padding row, so CUDA-graph padded batches (req_pool_idx == 0)
+        # read/write here harmlessly.
         self.device = device
         self.spec_algo = spec_algo
+        self.req_pool_size = req_to_token_pool.req_to_token.shape[0]
 
         if self.spec_algo.is_none():
-            # For non-speculative decoding, we only need to store the token ids.
             self.buf_initialized = True
             self.token_ids_buf = torch.empty(
-                (self.future_buffer_len,), dtype=torch.int64, device=self.device
+                (self.req_pool_size,), dtype=torch.int64, device=self.device
             )
         else:
-            # For speculative decoding, we lazily initialize the buffers
-            # This is to make the shape derivation easier.
             self.buf_initialized = False
 
     def _lazy_init_buf(self, draft_input: EagleDraftInput):
         self.buf_initialized = True
 
-        # Get a reference for each tensor
         topk_p0 = draft_input.topk_p[0]
         topk_index0 = draft_input.topk_index[0]
         bonus_token0 = draft_input.bonus_tokens[0]
         new_seq_lens0 = draft_input.new_seq_lens[0]
 
         self.topk_p_buf = torch.empty(
-            (self.future_buffer_len, *topk_p0.shape),
+            (self.req_pool_size, *topk_p0.shape),
             dtype=topk_p0.dtype,
             device=self.device,
         )
         self.topk_index_buf = torch.empty(
-            (self.future_buffer_len, *topk_index0.shape),
+            (self.req_pool_size, *topk_index0.shape),
             dtype=topk_index0.dtype,
             device=self.device,
         )
         self.bonus_tokens_buf = torch.empty(
-            (self.future_buffer_len, *bonus_token0.shape),
+            (self.req_pool_size, *bonus_token0.shape),
             dtype=bonus_token0.dtype,
             device=self.device,
         )
         self.new_seq_lens_buf = torch.empty(
-            (self.future_buffer_len, *new_seq_lens0.shape),
+            (self.req_pool_size, *new_seq_lens0.shape),
             dtype=new_seq_lens0.dtype,
             device=self.device,
         )
@@ -113,35 +96,23 @@ class FutureMap:
         if spec_need_hidden_states():
             hidden_states0 = draft_input.hidden_states[0]
             self.hidden_states_buf = torch.empty(
-                (self.future_buffer_len, *hidden_states0.shape),
+                (self.req_pool_size, *hidden_states0.shape),
                 dtype=hidden_states0.dtype,
                 device=self.device,
             )
 
-    def alloc_future_indices(self, bs: int) -> FutureIndices:
-        """Update the circular buffer pointer and allocate future indices."""
-        cur_future_ct = self.future_ct
-        self.future_ct = (cur_future_ct + bs) % self.future_limit
-        start = cur_future_ct + 1
-        end = cur_future_ct + 1 + bs
-        indices = torch.arange(start, end, dtype=torch.int64, device=self.device)
-        return FutureIndices(indices=indices, interval=slice(start, end))
-
     def resolve_future(self, batch: ScheduleBatch):
         if self.spec_algo.is_none():
             _resolve_future_token_ids(batch.input_ids, self.token_ids_buf)
         else:
-            # TODO(lsyin): write future indices into spec_info.future_indices
             draft_input: EagleDraftInput = batch.spec_info
             if draft_input is None:
                 # FIXME(lsyin): No future exists, only for prefill batch, not compatible with mixed mode
                 return
             indices = draft_input.future_indices.indices
-            # The indices tensor was allocated on the default stream but is
-            # used here on the forward stream. Meanwhile, the old spec_info
-            # holding this tensor will lose all Python references (replaced at
-            # batch.spec_info), so the caching allocator (torch GC) could
-            # reclaim the memory before the GPU finishes reading it.
+            # FIXME: redundant. `indices` = batch.req_pool_indices, pinned via
+            # record_batch_in_overlap's attr_snapshot for 2 iters; refcount > 0
+            # across forward's read, allocator can't reclaim. Safe to remove.
             indices.record_stream(torch.get_device_module(self.device).current_stream())
             draft_input.topk_p = self.topk_p_buf[indices]
             draft_input.topk_index = self.topk_index_buf[indices]
@@ -150,22 +121,19 @@ class FutureMap:
             if spec_need_hidden_states():
                 draft_input.hidden_states = self.hidden_states_buf[indices]
 
-    def is_empty_slice(self, s: slice) -> bool:
-        start, stop, step = s.indices(self.future_buffer_len)
-        if step > 0:
-            return start >= stop
-        else:
-            return start <= stop
-
     def store_to_map(
         self, future_indices: FutureIndices, batch_result: GenerationBatchResult
     ):
         if self.spec_algo.is_none():
-            intv = future_indices.interval
-            if self.is_empty_slice(intv):
-                # idle indices in dp attention do not need store info
+            indices = future_indices.indices
+            if indices.shape[0] == 0:
+                # DP attention idle rank: indices is empty but next_token_ids
+                # may carry padded values from sibling ranks. Nothing to store
+                # for this rank.
                 return
-            self.token_ids_buf[intv] = batch_result.next_token_ids
+            # next_token_ids is int32; buf is int64. Slice assignment used to
+            # cast implicitly, but advanced indexing requires an explicit match.
+            self.token_ids_buf[indices] = batch_result.next_token_ids.to(torch.int64)
         else:
             draft_input: EagleDraftInput = batch_result.next_draft_input
             self.store_to_map_for_new_batch(future_indices, draft_input)
@@ -173,17 +141,30 @@ class FutureMap:
     def store_to_map_for_new_batch(
         self, future_indices: FutureIndices, draft_input: EagleDraftInput
     ):
-        intv = future_indices.interval
-        if self.is_empty_slice(intv):
-            # idle indices in dp attention do not need store info
+        indices = future_indices.indices
+        if indices.shape[0] == 0:
+            # DP idle rank: draft_input fields are empty stubs without a usable
+            # shape, so _lazy_init_buf's shape peek (draft_input.topk_p[0])
+            # would IndexError. Defer init until a real batch arrives.
             return
 
         if not self.buf_initialized:
             self._lazy_init_buf(draft_input)
 
-        self.topk_p_buf[intv] = draft_input.topk_p
-        self.topk_index_buf[intv] = draft_input.topk_index
-        self.bonus_tokens_buf[intv] = draft_input.bonus_tokens
-        self.new_seq_lens_buf[intv] = draft_input.new_seq_lens
+        # Slice assignment used to coerce src dtype to buf dtype implicitly;
+        # advanced index requires an explicit cast. bonus_tokens / new_seq_lens
+        # in particular differ across disagg (int64) and forward (int32) paths.
+        self.topk_p_buf[indices] = draft_input.topk_p.to(self.topk_p_buf.dtype)
+        self.topk_index_buf[indices] = draft_input.topk_index.to(
+            self.topk_index_buf.dtype
+        )
+        self.bonus_tokens_buf[indices] = draft_input.bonus_tokens.to(
+            self.bonus_tokens_buf.dtype
+        )
+        self.new_seq_lens_buf[indices] = draft_input.new_seq_lens.to(
+            self.new_seq_lens_buf.dtype
+        )
         if spec_need_hidden_states():
-            self.hidden_states_buf[intv] = draft_input.hidden_states
+            self.hidden_states_buf[indices] = draft_input.hidden_states.to(
+                self.hidden_states_buf.dtype
+            )

python/sglang/srt/managers/scheduler.py

@@ -143,6 +143,7 @@ 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,
@@ -1278,10 +1279,8 @@ class Scheduler(
             return
 
         self.future_map = self.spec_algorithm.create_future_map(
-            self.max_running_requests,
-            self.chunked_prefill_size,
-            self.model_config.context_len,
             self.device,
+            self.req_to_token_pool,
         )
         self.batch_record_buf = [None] * 2
         self.batch_record_ct = 0
@@ -2840,8 +2839,7 @@ class Scheduler(
                 batch.refresh_seq_lens_cpu()
 
                 with self._overlap_forward_isolation(batch):
-                    bs = len(batch.seq_lens)
-                    future_indices = self.future_map.alloc_future_indices(bs)
+                    future_indices = FutureIndices(indices=batch.req_pool_indices)
 
                     with self.forward_stream_ctx:
                         self.forward_stream.wait_stream(self.schedule_stream)

python/sglang/srt/speculative/spec_info.py

@@ -117,20 +117,12 @@ class SpeculativeAlgorithm(Enum):
 
     def create_future_map(
         self,
-        max_running_requests: int,
-        chunked_prefill_size: int,
-        context_len: int,
         device: torch.device,
+        req_to_token_pool,
     ) -> FutureMap:
         from sglang.srt.managers.overlap_utils import FutureMap
 
-        return FutureMap(
-            max_running_requests,
-            chunked_prefill_size,
-            context_len,
-            device,
-            self,
-        )
+        return FutureMap(device, self, req_to_token_pool)
 
     def supports_spec_v2(self) -> bool:
         return (self.is_eagle() and not self.is_frozen_kv_mtp()) or self.is_standalone()