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

[Spec] Unify spec/non-spec decode result handling and overlap relay-payload gating (#29225)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-25 00:44:38 -07:00
committed by GitHub
parent 32bae911f6
commit 2812a3c93a
3 changed files with 53 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+21 -29
View File
@@ -3251,20 +3251,7 @@ 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 RelayPayload(
bonus_tokens=batch_result.next_token_ids
)
)
self.future_map.stash(future_indices, stash_payload)
self._relay_forward_payload(future_indices, batch_result)
# 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.
@@ -3286,11 +3273,7 @@ class Scheduler(
elif self.enable_pdmux and batch.forward_mode.is_split_prefill():
resolve_forward_inputs(batch, self.future_map)
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,
RelayPayload(bonus_tokens=batch_result.next_token_ids),
)
self._relay_forward_payload(batch.req_pool_indices, batch_result)
batch.input_ids = None
elif not batch.spec_algorithm.is_none():
# Non-overlap: drive the V2 worker synchronously (no
@@ -3324,12 +3307,9 @@ class Scheduler(
batch_result = self.model_worker.forward_batch_generation(
batch, **kwargs
)
if isinstance(batch_result.next_token_ids, torch.Tensor):
if batch_result.has_sampled_token_ids:
# Non-spec: relay via future_map, gathered next iter.
self.future_map.stash(
batch.req_pool_indices,
RelayPayload(bonus_tokens=batch_result.next_token_ids),
)
self._relay_forward_payload(batch.req_pool_indices, batch_result)
batch.input_ids = None
self.update_cache_from_scheduler(batch, batch_result)
@@ -3388,6 +3368,21 @@ class Scheduler(
ActiveRanksOutput(status=dp_active_ranks.tolist())
)
def _relay_forward_payload(
self, future_indices: torch.Tensor, batch_result: GenerationBatchResult
) -> None:
"""Stash this iter's relay payload for next iter's resolve_forward_inputs.
ngram is skipped: it relays its draft via batch.spec_info, not the FutureMap."""
if self.spec_algorithm.is_ngram():
return
if batch_result.next_draft_input is not None:
payload = RelayPayload.from_draft_input(batch_result.next_draft_input)
elif batch_result.has_sampled_token_ids:
payload = RelayPayload(bonus_tokens=batch_result.next_token_ids)
else:
return
self.future_map.stash(future_indices, payload)
def launch_batch_sample_if_needed(
self, batch_result: GenerationBatchResult
) -> Union[GenerationBatchResult]:
@@ -3400,11 +3395,8 @@ class Scheduler(
self.forward_stream.wait_stream(self.schedule_stream)
_batch_result = batch_result.delay_sample_func()
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,
RelayPayload(bonus_tokens=batch_result.next_token_ids),
)
# Delay-sample is non-spec only; relays the sampled bonus tokens.
self._relay_forward_payload(batch_result.future_indices, batch_result)
batch_result.copy_to_cpu(
return_logprob=self.cur_batch.return_logprob,
return_hidden_states=self.cur_batch.return_hidden_states,
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+26 -33
View File
@@ -675,20 +675,13 @@ class SchedulerBatchResultProcessor:
# And all the over-allocated tokens will be freed in `release_kv_cache`.
continue
# Non-spec and Spec V2: full post-processing.
# next_token_id is a per-req list: 1 token for non-spec, the verified
# run for spec (already grammar-truncated in _resolve_spec_v2_tokens).
next_token_id = next_token_ids[i]
is_spec = not batch.spec_algorithm.is_none()
if not is_spec:
# Normal decode: a single sampled token.
req.output_ids.append(next_token_id)
new_accept_len = 1
else:
# Spec: accept the whole verified run. For grammar requests the
# run was already truncated at the grammar-terminating token in
# _resolve_spec_v2_tokens, so nothing is emitted past completion.
req.output_ids.extend(next_token_id)
new_accept_len = len(next_token_id)
req.output_ids.extend(next_token_id)
new_accept_len = len(next_token_id)
self._maybe_update_reasoning_tokens(req, next_token_id)
req.time_stats.set_last_decode_finish_time()
@@ -707,23 +700,16 @@ class SchedulerBatchResultProcessor:
)
if req.return_hidden_states and logits_output.hidden_states is not None:
if not is_spec:
req.hidden_states.append(
logits_output.hidden_states[i].cpu().clone().tolist()
)
else:
# Spec V2: hidden_states is [bs * speculative_num_draft_tokens, hidden_dim].
# One row per emitted token; next_token_id is already truncated
# at grammar termination, so this stays aligned with output_ids.
stride = result.speculative_num_draft_tokens
accept_len = len(next_token_id)
start = i * stride
req.hidden_states.extend(
logits_output.hidden_states[start : start + accept_len]
.cpu()
.clone()
.tolist()
)
# hidden_states is [bs * stride, hidden_dim], one row per emitted
# token; stride = speculative_num_draft_tokens for spec, 1 for non-spec.
stride = result.speculative_num_draft_tokens or 1
accept_len = len(next_token_id)
start = i * stride
req.hidden_states.extend(
logits_output.hidden_states[start : start + accept_len]
.cpu()
.tolist()
)
if req.grammar is not None:
if not is_spec:
@@ -753,12 +739,18 @@ class SchedulerBatchResultProcessor:
next_token_ids: Union[torch.Tensor, List[int]],
) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]:
next_token_logprobs = None
# Normalize to a uniform per-req list of accepted tokens (List[List[int]]):
# spec unpacks the padded verify output; non-spec wraps its single token.
if not batch.spec_algorithm.is_none():
next_token_ids = self._resolve_spec_v2_tokens(result, batch)
elif isinstance(next_token_ids, list):
pass # MLX path: already a list[int], skip torch round-trip
else:
next_token_ids = next_token_ids.tolist()
# CUDA workers return a device tensor, MLX a host list[int]; both -> list.
ids = (
next_token_ids.tolist()
if torch.is_tensor(next_token_ids)
else next_token_ids
)
next_token_ids = [[t] for t in ids]
if batch.return_logprob:
next_token_logprobs = logits_output.next_token_logprobs.tolist()
@@ -786,14 +778,15 @@ class SchedulerBatchResultProcessor:
next_token_logprobs: list,
logits_output: LogitsProcessorOutput,
) -> None:
# Normalize: non-spec has 1 token, spec decoding has multiple.
# accepted_ids is already a per-req list; non-spec logprobs are flat, so
# the scalar logprob still needs wrapping.
if not batch.spec_algorithm.is_none():
accepted_logprobs = next_token_logprobs[i]
accepted_ids = next_token_id
max_accept = len(accepted_logprobs)
else:
accepted_logprobs = [next_token_logprobs[i]]
accepted_ids = [next_token_id]
accepted_ids = next_token_id
max_accept = 1
for j, tok_id in enumerate(accepted_ids):
python/sglang/srt/managers/utils.py
+6
View File
@@ -85,6 +85,12 @@ class GenerationBatchResult:
fpm_start_event: Optional[torch.cuda.Event] = None
fpm_end_event: Optional[torch.cuda.Event] = None
@property
def has_sampled_token_ids(self) -> bool:
"""True when this iter sampled token ids; False when none were produced
this rank/split (a non-last PP rank or a non-final prefill split)."""
return isinstance(self.next_token_ids, torch.Tensor)
@torch.profiler.record_function("copy_result_to_cpu")
def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True):
"""Copy tensors to CPU in overlap scheduling.