[Spec] Unify the overlap stash relay behind a RelayPayload dataclass (#29124)

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

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():
+        # 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()
+            self.spec_algo.is_some()
-            and getattr(draft_input, "hidden_states", None) is not None
+            and spec_need_hidden_states()
+            and payload.hidden_states is not None
         )
 
         if self.need_topk:
-            topk_p0 = draft_input.topk_p[0]
+            topk_p0 = payload.topk_p[0]
-            topk_index0 = draft_input.topk_index[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:
+        if payload.draft_probs is not None:
-            draft_probs0 = draft_input.draft_probs[0]
+            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(
+    def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None:
-        self,
-        future_indices: torch.Tensor,
-        payload: Union[torch.Tensor, EagleDraftInput],
-    ) -> 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._lazy_init_forward_buf(payload)
-        self.output_tokens_buf[indices] = draft_input.bonus_tokens.to(
+        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_p_buf[indices] = payload.topk_p.to(self.topk_p_buf.dtype)
-            self.topk_index_buf[indices] = draft_input.topk_index.to(
+            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:
+        if self.draft_probs_buf is not None and payload.draft_probs is not None:
-            self.draft_probs_buf[indices] = draft_input.draft_probs
+            self.draft_probs_buf[indices] = payload.draft_probs

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,10 +3250,18 @@ 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:
+                            # 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 batch_result.next_token_ids
+                                    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
@@ -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,

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,

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