Refactor batch_result_processor into per-step prefill/decode helpers (#25709)

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:
+            self._move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
-                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
-                    ]
 
             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
+                        logprob_pt = self._apply_prefill_logprobs(
-                        assert extend_input_len_per_req is not None
+                            req=req,
-                        extend_logprob_start_len = extend_logprob_start_len_per_req[i]
+                            i=i,
-                        extend_input_len = extend_input_len_per_req[i]
+                            logits_output=logits_output,
-
+                            extend_input_len_per_req=extend_input_len_per_req,
-                        num_input_logprobs = (
+                            extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
-                            self.logprob_result_processor.calculate_num_input_logprobs(
+                            next_token_ids=next_token_ids,
-                                req,
+                            logprob_pt=logprob_pt,
-                                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
 
                     if (
                         req.return_hidden_states
                         and logits_output.hidden_states is not None
                     ):
-                        req.hidden_states.append(
+                        hidden_state_offset = self._append_prefill_hidden_states(
-                            logits_output.hidden_states[
+                            req=req,
-                                hidden_state_offset : (
+                            logits_output=logits_output,
-                                    hidden_state_offset := hidden_state_offset
+                            hidden_state_offset=hidden_state_offset,
-                                    + len(req.origin_input_ids)
-                                )
-                            ]
-                            .cpu()
-                            .clone()
-                            .tolist()
                         )
 
                     if req.grammar is not None:
-                        # FIXME: this try-except block is for handling unexpected xgrammar issue.
+                        self._apply_prefill_grammar(
-                        try:
+                            req=req, next_token_id=next_token_id
-                            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()
 
                 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]
+                        logprob_pt = self._apply_chunked_prefill_logprobs(
-                        extend_input_len = extend_input_len_per_req[i]
+                            req=req,
-                        if extend_logprob_start_len < extend_input_len:
+                            i=i,
-                            # Update input logprobs.
+                            logits_output=logits_output,
-                            num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs(
+                            extend_input_len_per_req=extend_input_len_per_req,
-                                req,
+                            extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
-                                extend_input_len,
+                            logprob_pt=logprob_pt,
-                                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
 
                     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 = self._convert_embeddings(result=result)
-
-            embeddings = result.embeddings
             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:
+        next_token_ids, next_token_logprobs = self._normalize_decode_outputs(
-            if batch.is_spec_v2:
+            batch=batch,
-                next_token_ids = self._resolve_spec_overlap_tokens(result, batch)
+            result=result,
-            elif isinstance(next_token_ids, list):
+            logits_output=logits_output,
-                pass  # MLX path: already a list[int], skip torch round-trip
+            next_token_ids=next_token_ids,
-            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).
 
         self.metrics_reporter.num_generated_tokens += len(batch.reqs)
         if not batch.spec_algorithm.is_none():
@@ -571,6 +623,83 @@ class SchedulerBatchResultProcessor:
             self._handle_finished_req(req, i, logits_output)
 
             if req.return_logprob:
+                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(
+                    logits_output.hidden_states[i].cpu().clone().tolist()
+                )
+
+            if req.grammar is not None:
+                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()
+
+        self.metrics_reporter.forward_ct_decode = (
+            self.metrics_reporter.forward_ct_decode + 1
+        ) % (1 << 30)
+        self.metrics_reporter.report_decode_stats(
+            can_run_cuda_graph,
+            running_batch=batch,
+            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:
@@ -602,12 +731,13 @@ class SchedulerBatchResultProcessor:
                     logits_output.next_token_token_ids_logprobs_idx[flat_idx]
                 )
 
-            if req.return_hidden_states and logits_output.hidden_states is not None:
+    def _apply_decode_grammar(
-                req.hidden_states.append(
+        self,
-                    logits_output.hidden_states[i].cpu().clone().tolist()
+        *,
-                )
+        req: Req,
-
+        next_token_id: Union[int, List[int]],
-            if req.grammar is not None:
+        batch: ScheduleBatch,
+    ) -> None:
         # FIXME: this try-except block is for handling unexpected xgrammar issue.
         try:
             if batch.spec_algorithm.is_none():
@@ -626,18 +756,6 @@ class SchedulerBatchResultProcessor:
             self.abort_request(AbortReq(rid=req.rid))
         req.grammar.finished = req.finished()
 
-        self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
-        self.token_to_kv_pool_allocator.free_group_end()
-
-        self.metrics_reporter.forward_ct_decode = (
-            self.metrics_reporter.forward_ct_decode + 1
-        ) % (1 << 30)
-        self.metrics_reporter.report_decode_stats(
-            can_run_cuda_graph,
-            running_batch=batch,
-            num_correct_drafts=result.num_correct_drafts,
-        )
-
     def _handle_finished_req(
         self,
         req: Req,