From 2812a3c93a7060a40a3e7e5fc4acd2d3ec6fdca2 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 25 Jun 2026 00:44:38 -0700 Subject: [PATCH] [Spec] Unify spec/non-spec decode result handling and overlap relay-payload gating (#29225) --- python/sglang/srt/managers/scheduler.py | 50 +++++++--------- .../batch_result_processor.py | 59 ++++++++----------- python/sglang/srt/managers/utils.py | 6 ++ 3 files changed, 53 insertions(+), 62 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 0d6891f4f..7055e25d9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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, diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index ec6827ddf..6e20e6532 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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): diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 286121e88..87fbb35de 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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.