From 512d164916226ed52f0a8f97c3aecc570e52c57d Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 20 May 2026 16:34:00 -0700 Subject: [PATCH] Address overlap future token map by request-pool index (#25862) --- .../decode_schedule_batch_mixin.py | 6 +- python/sglang/srt/managers/overlap_utils.py | 109 ++++++++---------- python/sglang/srt/managers/scheduler.py | 8 +- python/sglang/srt/speculative/spec_info.py | 12 +- 4 files changed, 53 insertions(+), 82 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index aca442102..d6db74ec0 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/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 ) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 118ef04b3..795f01b75 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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 + ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 204d3a248..8a08f03dd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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) diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 31814e416..ca2be5666 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/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()