From 85a3f1f75c5dea5571a2c7e8d2c5a32a95d8d422 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 24 Jun 2026 15:36:54 -0700 Subject: [PATCH] [Spec] Unify the overlap stash relay behind a RelayPayload dataclass (#29124) --- .../decode_schedule_batch_mixin.py | 5 +- python/sglang/srt/managers/overlap_utils.py | 90 +++++++++++-------- python/sglang/srt/managers/scheduler.py | 34 ++++--- .../sglang/srt/managers/scheduler_pp_mixin.py | 5 +- .../srt/speculative/eagle_disaggregation.py | 5 +- 5 files changed, 91 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 68a06ce8d..309f060bf 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, List import torch +from sglang.srt.managers.overlap_utils import RelayPayload from sglang.srt.mem_cache.common import maybe_cache_unfinished_req from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo @@ -153,5 +154,7 @@ class ScheduleBatchDisaggregationDecodeMixin: else: # Non-spec: stash last token into the relay so the first DECODE's # resolve_forward_inputs gathers it like any other decode iter. - future_map.stash(self.req_pool_indices, last_tokens_tensor) + future_map.stash( + self.req_pool_indices, RelayPayload(bonus_tokens=last_tokens_tensor) + ) self.input_ids = None diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 043341cfe..bee8fb5e2 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -1,12 +1,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Sequence, Union +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Sequence import torch from sglang.srt.environ import envs -from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -from sglang.srt.speculative.spec_utils import spec_need_hidden_states from sglang.srt.speculative.triton_ops.gather_spec_extras import gather_spec_extras from sglang.srt.utils import is_cuda, is_hip, is_npu @@ -16,6 +15,7 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.eagle_info import EagleDraftInput + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm def decide_needs_cpu_seq_lens( @@ -28,6 +28,10 @@ def decide_needs_cpu_seq_lens( CPU mirror outside the backend layer to split the batch) or ngram (its USE_FULL_MASK verify path reads the host mirror regardless of backend). """ + # Local import: keep overlap_utils' module-level deps leaf-only so it stays + # importable everywhere; spec_info pulls in the spec/schedule_batch graph. + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + if server_args.enable_two_batch_overlap: # FIXME: support TBO without seq lens cpu value return True @@ -96,6 +100,29 @@ def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None: future_map._resolve_spec_extras(batch) +@dataclass +class RelayPayload: + """Per-iteration stash payload for the FutureMap bufs. Non-spec fills only + `bonus_tokens`; which spec extras get relayed is decided by + `FutureMap.spec_algo`, not by this payload's shape.""" + + bonus_tokens: torch.Tensor + topk_p: Optional[torch.Tensor] = None + topk_index: Optional[torch.Tensor] = None + hidden_states: Optional[torch.Tensor] = None + draft_probs: Optional[torch.Tensor] = None + + @classmethod + def from_draft_input(cls, draft_input: EagleDraftInput) -> RelayPayload: + return cls( + bonus_tokens=draft_input.bonus_tokens, + topk_p=draft_input.topk_p, + topk_index=draft_input.topk_index, + hidden_states=draft_input.hidden_states, + draft_probs=getattr(draft_input, "draft_probs", None), + ) + + class FutureMap: """Always-on pool-indexed relay for cross-iter values. Forward writes via publish/stash; next iter reads via resolve_forward_inputs / resolve_seq_lens_cpu. @@ -139,23 +166,29 @@ class FutureMap: else: self.new_seq_lens_cpu_pinned = None self.fwd_prepare_d2h_stream = None - if self.spec_algo.is_some(): - self._forward_buf_initialized = False + # Lazy-inited on the first non-empty stash (peeks tensor shapes); non-spec's is a no-op. + self._forward_buf_initialized = False self.publish_ready = None # lazy device.Event(); only spec_v2 needs it - def _lazy_init_forward_buf(self, draft_input: EagleDraftInput): + def _lazy_init_forward_buf(self, payload: RelayPayload): + # Local import (see decide_needs_cpu_seq_lens): keep module-level deps leaf. + from sglang.srt.speculative.spec_utils import spec_need_hidden_states + self._forward_buf_initialized = True - self.need_topk = self.spec_algo.need_topk() + # Spec extras are gated by spec_algo, not by the payload's shape, so a + # non-spec stash allocates no extra bufs (only output_tokens_buf). + self.need_topk = self.spec_algo.is_some() and self.spec_algo.need_topk() self.need_hidden_states = ( - spec_need_hidden_states() - and getattr(draft_input, "hidden_states", None) is not None + self.spec_algo.is_some() + and spec_need_hidden_states() + and payload.hidden_states is not None ) if self.need_topk: - topk_p0 = draft_input.topk_p[0] - topk_index0 = draft_input.topk_index[0] + topk_p0 = payload.topk_p[0] + topk_index0 = payload.topk_index[0] self.topk_p_buf = torch.empty( (self.req_pool_size, *topk_p0.shape), dtype=topk_p0.dtype, @@ -167,7 +200,7 @@ class FutureMap: device=self.device, ) if self.need_hidden_states: - hidden_states0 = draft_input.hidden_states[0] + hidden_states0 = payload.hidden_states[0] self.hidden_states_buf = torch.empty( (self.req_pool_size, *hidden_states0.shape), dtype=hidden_states0.dtype, @@ -175,8 +208,8 @@ class FutureMap: ) self.draft_probs_buf = None - if getattr(draft_input, "draft_probs", None) is not None: - draft_probs0 = draft_input.draft_probs[0] + if payload.draft_probs is not None: + draft_probs0 = payload.draft_probs[0] self.draft_probs_buf = torch.empty( (self.req_pool_size, *draft_probs0.shape), dtype=draft_probs0.dtype, @@ -296,11 +329,7 @@ class FutureMap: self.publish_ready = torch.get_device_module(self.device).Event() self.publish_ready.record() - def stash( - self, - future_indices: torch.Tensor, - payload: Union[torch.Tensor, EagleDraftInput], - ) -> None: + def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None: if self.spec_algo.is_ngram(): # FIXME: remove once precomputed draft is supported. return @@ -308,29 +337,20 @@ class FutureMap: if indices.shape[0] == 0: # DP idle: payload is empty stub; lazy-init shape peek would IndexError. return - # Dispatch by payload type, not spec_algo: non-spec decode passes a - # token Tensor here. - # FIXME(lsyin): unify this relay path with a dataclass instead of the - # Tensor / EagleDraftInput type switch. - if isinstance(payload, torch.Tensor): - self.output_tokens_buf[indices] = payload.to(torch.int64) - return - - draft_input: EagleDraftInput = payload if not self._forward_buf_initialized: - self._lazy_init_forward_buf(draft_input) - self.output_tokens_buf[indices] = draft_input.bonus_tokens.to( + self._lazy_init_forward_buf(payload) + self.output_tokens_buf[indices] = payload.bonus_tokens.to( self.output_tokens_buf.dtype ) if self.need_topk: - 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_p_buf[indices] = payload.topk_p.to(self.topk_p_buf.dtype) + self.topk_index_buf[indices] = payload.topk_index.to( self.topk_index_buf.dtype ) if self.need_hidden_states: - self.hidden_states_buf[indices] = draft_input.hidden_states.to( + self.hidden_states_buf[indices] = payload.hidden_states.to( self.hidden_states_buf.dtype ) - if self.draft_probs_buf is not None and draft_input.draft_probs is not None: - self.draft_probs_buf[indices] = draft_input.draft_probs + if self.draft_probs_buf is not None and payload.draft_probs is not None: + self.draft_probs_buf[indices] = payload.draft_probs diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 92b5a35e8..fdc99616e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -156,6 +156,7 @@ from sglang.srt.managers.min_free_slots_delayer import ( ) from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors from sglang.srt.managers.overlap_utils import ( + RelayPayload, decide_needs_cpu_seq_lens, resolve_forward_inputs, ) @@ -2579,7 +2580,9 @@ class Scheduler( last_tokens = torch.tensor( [r.output_ids[-1] for r in reqs], dtype=torch.int64, device=device ) - self.future_map.stash(batch.req_pool_indices, last_tokens) + self.future_map.stash( + batch.req_pool_indices, RelayPayload(bonus_tokens=last_tokens) + ) batch.input_ids = None if batch.return_logprob: @@ -3247,12 +3250,20 @@ class Scheduler( # FIXME(lsyin): maybe move this to forward_batch_generation batch_result.copy_done = self.device_module.Event() if batch_result.delay_sample_func is None: - stash_payload = ( - batch_result.next_draft_input - if not batch.spec_algorithm.is_none() - else batch_result.next_token_ids - ) - self.future_map.stash(future_indices, stash_payload) + # ngram precomputes its draft and does not relay + # through the FutureMap (stash() no-ops for it); its + # verify input also has no bonus_tokens to project. + if not batch.spec_algorithm.is_ngram(): + stash_payload = ( + RelayPayload.from_draft_input( + batch_result.next_draft_input + ) + if not batch.spec_algorithm.is_none() + else RelayPayload( + bonus_tokens=batch_result.next_token_ids + ) + ) + self.future_map.stash(future_indices, stash_payload) # Result D2H on copy_stream overlaps the next forward # instead of serializing on forward_stream; it's a leaf # gated by copy_done, so nothing on forward_stream waits. @@ -3276,7 +3287,8 @@ class Scheduler( batch_result = self.tp_worker.forward_batch_split_prefill(batch) if isinstance(batch_result.next_token_ids, torch.Tensor): self.future_map.stash( - batch.req_pool_indices, batch_result.next_token_ids + batch.req_pool_indices, + RelayPayload(bonus_tokens=batch_result.next_token_ids), ) batch.input_ids = None elif not batch.spec_algorithm.is_none(): @@ -3314,7 +3326,8 @@ class Scheduler( if isinstance(batch_result.next_token_ids, torch.Tensor): # Non-spec: relay via future_map, gathered next iter. self.future_map.stash( - batch.req_pool_indices, batch_result.next_token_ids + batch.req_pool_indices, + RelayPayload(bonus_tokens=batch_result.next_token_ids), ) batch.input_ids = None self.update_cache_from_scheduler(batch, batch_result) @@ -3387,7 +3400,8 @@ class Scheduler( assert _batch_result is batch_result # Delay-sample is non-spec only; stash takes next_token_ids tensor. self.future_map.stash( - batch_result.future_indices, batch_result.next_token_ids + batch_result.future_indices, + RelayPayload(bonus_tokens=batch_result.next_token_ids), ) batch_result.copy_to_cpu( return_logprob=self.cur_batch.return_logprob, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 242afde9b..cac3bf91a 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -23,6 +23,7 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, set_is_extend_in_batch, ) +from sglang.srt.managers.overlap_utils import RelayPayload from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.utils import ( GenerationBatchResult, @@ -1109,7 +1110,9 @@ class SchedulerPPMixin: # PP rank 0 also relays into output_tokens_buf so the next iter's # resolve_forward_inputs finds these tokens for the decode portion # of mixed-chunk batches (which gather via mix_running_indices). - self.future_map.stash(batch.req_pool_indices, batch.input_ids) + self.future_map.stash( + batch.req_pool_indices, RelayPayload(bonus_tokens=batch.input_ids) + ) output_result = GenerationBatchResult( logits_output=logits_output, pp_hidden_states_proxy_tensors=None, diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index dc9493f28..410ab4817 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING import torch +from sglang.srt.managers.overlap_utils import RelayPayload from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode from sglang.srt.speculative.eagle_info import EagleDraftInput @@ -61,6 +62,8 @@ def build_eagle_disagg_draft_input( if batch.enable_overlap: spec_info.future_indices = batch.req_pool_indices future_map.publish(spec_info.future_indices, batch.seq_lens) - future_map.stash(spec_info.future_indices, spec_info) + future_map.stash( + spec_info.future_indices, RelayPayload.from_draft_input(spec_info) + ) return spec_info