From 0e198f0f4f4a53dd20fb4cde784ea7829e2d8006 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 19 May 2026 09:17:17 +0800 Subject: [PATCH] Refactor batch_result_processor into per-step prefill/decode helpers (#25709) --- .../batch_result_processor.py | 464 +++++++++++------- 1 file changed, 291 insertions(+), 173 deletions(-) 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 5b50e1bbd..baf09cfb7 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -7,6 +7,7 @@ from typing import ( Callable, List, Optional, + Tuple, Union, ) @@ -202,27 +203,7 @@ class SchedulerBatchResultProcessor: # Move next_token_ids and logprobs to cpu next_token_ids = next_token_ids.tolist() - if batch.return_logprob: - if logits_output.next_token_logprobs is not None: - logits_output.next_token_logprobs = ( - logits_output.next_token_logprobs.tolist() - ) - if logits_output.input_token_logprobs is not None: - logits_output.input_token_logprobs = tuple( - logits_output.input_token_logprobs.tolist() - ) - if logits_output.next_token_top_logprobs_val: - logits_output.next_token_top_logprobs_val = [ - v.tolist() for v in logits_output.next_token_top_logprobs_val - ] - logits_output.next_token_top_logprobs_idx = [ - x.tolist() for x in logits_output.next_token_top_logprobs_idx - ] - if logits_output.next_token_token_ids_logprobs_val: - logits_output.next_token_token_ids_logprobs_val = [ - v.tolist() - for v in logits_output.next_token_token_ids_logprobs_val - ] + self._move_logprobs_to_cpu(batch=batch, logits_output=logits_output) hidden_state_offset = 0 @@ -256,58 +237,30 @@ class SchedulerBatchResultProcessor: self._maybe_collect_customized_info(i, req, logits_output) if batch.return_logprob: - assert extend_logprob_start_len_per_req is not None - assert extend_input_len_per_req is not None - extend_logprob_start_len = extend_logprob_start_len_per_req[i] - extend_input_len = extend_input_len_per_req[i] - - num_input_logprobs = ( - self.logprob_result_processor.calculate_num_input_logprobs( - req, - extend_input_len, - extend_logprob_start_len, - ) + logprob_pt = self._apply_prefill_logprobs( + req=req, + i=i, + logits_output=logits_output, + extend_input_len_per_req=extend_input_len_per_req, + extend_logprob_start_len_per_req=extend_logprob_start_len_per_req, + next_token_ids=next_token_ids, + logprob_pt=logprob_pt, ) - if req.return_logprob: - self.logprob_result_processor.add_logprob_return_values( - i, - req, - logprob_pt, - next_token_ids, - num_input_logprobs, - logits_output, - ) - logprob_pt += num_input_logprobs - if ( req.return_hidden_states and logits_output.hidden_states is not None ): - req.hidden_states.append( - logits_output.hidden_states[ - hidden_state_offset : ( - hidden_state_offset := hidden_state_offset - + len(req.origin_input_ids) - ) - ] - .cpu() - .clone() - .tolist() + hidden_state_offset = self._append_prefill_hidden_states( + req=req, + logits_output=logits_output, + hidden_state_offset=hidden_state_offset, ) if req.grammar is not None: - # FIXME: this try-except block is for handling unexpected xgrammar issue. - try: - req.grammar.accept_token(next_token_id) - except ValueError as e: - # Grammar accept_token can raise ValueError if the token is not in the grammar. - # This can happen if the grammar is not set correctly or the token is invalid. - logger.error( - f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" - ) - self.abort_request(AbortReq(rid=req.rid)) - req.grammar.finished = req.finished() + self._apply_prefill_grammar( + req=req, next_token_id=next_token_id + ) else: # being chunked reqs' prefill is not finished @@ -319,25 +272,14 @@ class SchedulerBatchResultProcessor: # Incrementally update input logprobs. if batch.return_logprob: - extend_logprob_start_len = extend_logprob_start_len_per_req[i] - extend_input_len = extend_input_len_per_req[i] - if extend_logprob_start_len < extend_input_len: - # Update input logprobs. - num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs( - req, - extend_input_len, - extend_logprob_start_len, - ) - if req.return_logprob: - self.logprob_result_processor.add_input_logprob_return_values( - i, - req, - logits_output, - logprob_pt, - num_input_logprobs, - last_prefill_chunk=False, - ) - logprob_pt += num_input_logprobs + logprob_pt = self._apply_chunked_prefill_logprobs( + req=req, + i=i, + logits_output=logits_output, + extend_input_len_per_req=extend_input_len_per_req, + extend_logprob_start_len_per_req=extend_logprob_start_len_per_req, + logprob_pt=logprob_pt, + ) req.time_stats.set_last_chunked_prefill_finish_time() @@ -345,26 +287,9 @@ class SchedulerBatchResultProcessor: if result.copy_done is not None: result.copy_done.synchronize() - is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set() - - embeddings = result.embeddings + embeddings = self._convert_embeddings(result=result) phs = result.pooled_hidden_states - if is_sparse: - batch_ids, token_ids = embeddings.indices() - values = embeddings.values() - - embeddings = [{} for _ in range(embeddings.size(0))] - for i in range(batch_ids.shape[0]): - embeddings[batch_ids[i].item()][token_ids[i].item()] = values[ - i - ].item() - else: - if isinstance(embeddings, torch.Tensor): - embeddings = embeddings.tolist() - else: - embeddings = [tensor.tolist() for tensor in embeddings] - if phs is not None: if isinstance(phs, list): phs = [t.cpu().detach() for t in phs] @@ -407,6 +332,152 @@ class SchedulerBatchResultProcessor: dp_cooperation_info=batch.dp_cooperation_info, ) + def _convert_embeddings(self, *, result: EmbeddingBatchResult) -> list: + is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set() + + embeddings = result.embeddings + + if is_sparse: + batch_ids, token_ids = embeddings.indices() + values = embeddings.values() + + embeddings = [{} for _ in range(embeddings.size(0))] + for i in range(batch_ids.shape[0]): + embeddings[batch_ids[i].item()][token_ids[i].item()] = values[i].item() + else: + if isinstance(embeddings, torch.Tensor): + embeddings = embeddings.tolist() + else: + embeddings = [tensor.tolist() for tensor in embeddings] + return embeddings + + def _move_logprobs_to_cpu( + self, + *, + batch: ScheduleBatch, + logits_output: LogitsProcessorOutput, + ) -> None: + if batch.return_logprob: + if logits_output.next_token_logprobs is not None: + logits_output.next_token_logprobs = ( + logits_output.next_token_logprobs.tolist() + ) + if logits_output.input_token_logprobs is not None: + logits_output.input_token_logprobs = tuple( + logits_output.input_token_logprobs.tolist() + ) + if logits_output.next_token_top_logprobs_val: + logits_output.next_token_top_logprobs_val = [ + v.tolist() for v in logits_output.next_token_top_logprobs_val + ] + logits_output.next_token_top_logprobs_idx = [ + x.tolist() for x in logits_output.next_token_top_logprobs_idx + ] + if logits_output.next_token_token_ids_logprobs_val: + logits_output.next_token_token_ids_logprobs_val = [ + v.tolist() for v in logits_output.next_token_token_ids_logprobs_val + ] + + def _apply_prefill_logprobs( + self, + *, + req: Req, + i: int, + logits_output: LogitsProcessorOutput, + extend_input_len_per_req: Optional[List[int]], + extend_logprob_start_len_per_req: Optional[List[int]], + next_token_ids: List[int], + logprob_pt: int, + ) -> int: + assert extend_logprob_start_len_per_req is not None + assert extend_input_len_per_req is not None + extend_logprob_start_len = extend_logprob_start_len_per_req[i] + extend_input_len = extend_input_len_per_req[i] + + num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs( + req, + extend_input_len, + extend_logprob_start_len, + ) + + if req.return_logprob: + self.logprob_result_processor.add_logprob_return_values( + i, + req, + logprob_pt, + next_token_ids, + num_input_logprobs, + logits_output, + ) + logprob_pt += num_input_logprobs + return logprob_pt + + def _append_prefill_hidden_states( + self, + *, + req: Req, + logits_output: LogitsProcessorOutput, + hidden_state_offset: int, + ) -> int: + req.hidden_states.append( + logits_output.hidden_states[ + hidden_state_offset : ( + hidden_state_offset := hidden_state_offset + + len(req.origin_input_ids) + ) + ] + .cpu() + .clone() + .tolist() + ) + return hidden_state_offset + + def _apply_prefill_grammar(self, *, req: Req, next_token_id: int) -> None: + # FIXME: this try-except block is for handling unexpected xgrammar issue. + try: + req.grammar.accept_token(next_token_id) + except ValueError as e: + # Grammar accept_token can raise ValueError if the token is not in the grammar. + # This can happen if the grammar is not set correctly or the token is invalid. + logger.error( + f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" + ) + self.abort_request(AbortReq(rid=req.rid)) + req.grammar.finished = req.finished() + + def _apply_chunked_prefill_logprobs( + self, + *, + req: Req, + i: int, + logits_output: LogitsProcessorOutput, + extend_input_len_per_req: Optional[List[int]], + extend_logprob_start_len_per_req: Optional[List[int]], + logprob_pt: int, + ) -> int: + extend_logprob_start_len = extend_logprob_start_len_per_req[i] + extend_input_len = extend_input_len_per_req[i] + if extend_logprob_start_len < extend_input_len: + # Update input logprobs. + num_input_logprobs = ( + self.logprob_result_processor.calculate_num_input_logprobs( + req, + extend_input_len, + extend_logprob_start_len, + ) + ) + if req.return_logprob: + self.logprob_result_processor.add_input_logprob_return_values( + i, + req, + logits_output, + logprob_pt, + num_input_logprobs, + last_prefill_chunk=False, + ) + logprob_pt += num_input_logprobs + return logprob_pt + def _resolve_spec_overlap_tokens( self, result: GenerationBatchResult, @@ -488,31 +559,12 @@ class SchedulerBatchResultProcessor: result.can_run_cuda_graph, ) - if batch.spec_algorithm.is_none() or batch.is_spec_v2: - if batch.is_spec_v2: - next_token_ids = self._resolve_spec_overlap_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() - - if batch.return_logprob: - next_token_logprobs = logits_output.next_token_logprobs.tolist() - if logits_output.next_token_top_logprobs_val: - logits_output.next_token_top_logprobs_val = [ - v.tolist() for v in logits_output.next_token_top_logprobs_val - ] - logits_output.next_token_top_logprobs_idx = [ - x.tolist() for x in logits_output.next_token_top_logprobs_idx - ] - - if logits_output.next_token_token_ids_logprobs_val: - logits_output.next_token_token_ids_logprobs_val = [ - v.tolist() - for v in logits_output.next_token_token_ids_logprobs_val - ] - # else: Spec V1 — output_ids, check_finished, grammar, and reasoning tokens - # are already handled in the verify phase (eagle_info.py / ngram_info.py). + next_token_ids, next_token_logprobs = self._normalize_decode_outputs( + batch=batch, + result=result, + logits_output=logits_output, + next_token_ids=next_token_ids, + ) self.metrics_reporter.num_generated_tokens += len(batch.reqs) if not batch.spec_algorithm.is_none(): @@ -571,36 +623,14 @@ class SchedulerBatchResultProcessor: self._handle_finished_req(req, i, logits_output) if req.return_logprob: - # Spec v1 handles logprobs inside its own worker. - # Normalize: non-spec has 1 token, spec v2 has multiple. - if batch.is_spec_v2: - 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] - max_accept = 1 - - for j, tok_id in enumerate(accepted_ids): - req.output_token_logprobs_val.append(accepted_logprobs[j]) - req.output_token_logprobs_idx.append(tok_id) - if req.top_logprobs_num > 0: - flat_idx = i * max_accept + j - req.output_top_logprobs_val.append( - logits_output.next_token_top_logprobs_val[flat_idx] - ) - req.output_top_logprobs_idx.append( - logits_output.next_token_top_logprobs_idx[flat_idx] - ) - if req.token_ids_logprob is not None: - flat_idx = i * max_accept + j - req.output_token_ids_logprobs_val.append( - logits_output.next_token_token_ids_logprobs_val[flat_idx] - ) - req.output_token_ids_logprobs_idx.append( - logits_output.next_token_token_ids_logprobs_idx[flat_idx] - ) + self._apply_decode_logprobs( + req=req, + i=i, + batch=batch, + next_token_id=next_token_id, + next_token_logprobs=next_token_logprobs, + logits_output=logits_output, + ) if req.return_hidden_states and logits_output.hidden_states is not None: req.hidden_states.append( @@ -608,23 +638,9 @@ class SchedulerBatchResultProcessor: ) if req.grammar is not None: - # FIXME: this try-except block is for handling unexpected xgrammar issue. - try: - if batch.spec_algorithm.is_none(): - # Normal decode: single token - req.grammar.accept_token(next_token_id) - elif batch.is_spec_v2: - # Speculative decode: next_token_id is a list of accepted tokens - for token_id in next_token_id: - req.grammar.accept_token(token_id) - except ValueError as e: - # Grammar accept_token can raise ValueError if the token is not in the grammar. - # This can happen if the grammar is not set correctly or the token is invalid. - logger.error( - f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" - ) - self.abort_request(AbortReq(rid=req.rid)) - req.grammar.finished = req.finished() + self._apply_decode_grammar( + req=req, next_token_id=next_token_id, batch=batch + ) self.output_streamer.stream_output(batch.reqs, batch.return_logprob) self.token_to_kv_pool_allocator.free_group_end() @@ -638,6 +654,108 @@ class SchedulerBatchResultProcessor: num_correct_drafts=result.num_correct_drafts, ) + def _normalize_decode_outputs( + self, + *, + batch: ScheduleBatch, + result: GenerationBatchResult, + logits_output: LogitsProcessorOutput, + next_token_ids: Union[torch.Tensor, List[int]], + ) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]: + next_token_logprobs = None + if batch.spec_algorithm.is_none() or batch.is_spec_v2: + if batch.is_spec_v2: + next_token_ids = self._resolve_spec_overlap_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() + + if batch.return_logprob: + next_token_logprobs = logits_output.next_token_logprobs.tolist() + if logits_output.next_token_top_logprobs_val: + logits_output.next_token_top_logprobs_val = [ + v.tolist() for v in logits_output.next_token_top_logprobs_val + ] + logits_output.next_token_top_logprobs_idx = [ + x.tolist() for x in logits_output.next_token_top_logprobs_idx + ] + + if logits_output.next_token_token_ids_logprobs_val: + logits_output.next_token_token_ids_logprobs_val = [ + v.tolist() + for v in logits_output.next_token_token_ids_logprobs_val + ] + # else: Spec V1 — output_ids, check_finished, grammar, and reasoning tokens + # are already handled in the verify phase (eagle_info.py / ngram_info.py). + return next_token_ids, next_token_logprobs + + def _apply_decode_logprobs( + self, + *, + req: Req, + i: int, + batch: ScheduleBatch, + next_token_id: Union[int, List[int]], + next_token_logprobs: list, + logits_output: LogitsProcessorOutput, + ) -> None: + # Spec v1 handles logprobs inside its own worker. + # Normalize: non-spec has 1 token, spec v2 has multiple. + if batch.is_spec_v2: + 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] + max_accept = 1 + + for j, tok_id in enumerate(accepted_ids): + req.output_token_logprobs_val.append(accepted_logprobs[j]) + req.output_token_logprobs_idx.append(tok_id) + if req.top_logprobs_num > 0: + flat_idx = i * max_accept + j + req.output_top_logprobs_val.append( + logits_output.next_token_top_logprobs_val[flat_idx] + ) + req.output_top_logprobs_idx.append( + logits_output.next_token_top_logprobs_idx[flat_idx] + ) + if req.token_ids_logprob is not None: + flat_idx = i * max_accept + j + req.output_token_ids_logprobs_val.append( + logits_output.next_token_token_ids_logprobs_val[flat_idx] + ) + req.output_token_ids_logprobs_idx.append( + logits_output.next_token_token_ids_logprobs_idx[flat_idx] + ) + + def _apply_decode_grammar( + self, + *, + req: Req, + next_token_id: Union[int, List[int]], + batch: ScheduleBatch, + ) -> None: + # FIXME: this try-except block is for handling unexpected xgrammar issue. + try: + if batch.spec_algorithm.is_none(): + # Normal decode: single token + req.grammar.accept_token(next_token_id) + elif batch.is_spec_v2: + # Speculative decode: next_token_id is a list of accepted tokens + for token_id in next_token_id: + req.grammar.accept_token(token_id) + except ValueError as e: + # Grammar accept_token can raise ValueError if the token is not in the grammar. + # This can happen if the grammar is not set correctly or the token is invalid. + logger.error( + f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" + ) + self.abort_request(AbortReq(rid=req.rid)) + req.grammar.finished = req.finished() + def _handle_finished_req( self, req: Req,