minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-24 15:36:54 -07:00
committed by GitHub
parent c6822f81b1
commit 85a3f1f75c
5 changed files with 91 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
+4 -1
View File
@@ -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
+55 -35
View File
@@ -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
python/sglang/srt/managers/scheduler.py
+24 -10
View File
@@ -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,
python/sglang/srt/managers/scheduler_pp_mixin.py
+4 -1
View File
@@ -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 -1
View File
@@ -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