[Refactor] disaggregation: Commit the PD handoff token through the shared per-token commit

python/sglang/srt/disaggregation/decode.py

@@ -2929,13 +2929,38 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest):
                 )
             else:
                 committed_output_id = output_id[0].item()
-            decode_req.req.output_ids.append(committed_output_id)
-            # The handoff token is generated on the prefill worker, so it does
-            # not pass through the decode worker's normal batch-result path.
-            # Account for it here using the same request-selected reasoning
-            # terminator matcher as subsequent decode tokens.
-            self.scheduler.batch_result_processor._maybe_update_reasoning_tokens(
-                decode_req.req, committed_output_id
+            req = decode_req.req
+            logprob_token_id = logprob = top_logprobs_val = top_logprobs_idx = None
+            if req.return_logprob:
+                top_logprobs_num = req.logprob.top_logprobs_num
+                logprob_token_id = output_token_logprobs_idx[0].item()
+                logprob = output_token_logprobs_val[0].item()
+                top_logprobs_val = output_top_logprobs_val[:top_logprobs_num].tolist()
+                top_logprobs_idx = output_top_logprobs_idx[:top_logprobs_num].tolist()
+            sampling_mask_idx = sampling_mask_logprobs = None
+            if req.return_sampling_mask:
+                assert output_token_sampling_mask_idx is not None, (
+                    "sampling mask buffer disabled on decode side"
+                )
+                sampling_mask_len = int(output_token_sampling_mask_len[0].item())
+                num_logprobs = (
+                    sampling_mask_len if req.sampling_logprobs_mode == "support" else 1
+                )
+                sampling_mask_idx = (
+                    output_token_sampling_mask_idx[:sampling_mask_len].cpu().numpy()
+                )
+                sampling_mask_logprobs = (
+                    output_token_sampling_logprobs[:num_logprobs].cpu().numpy()
+                )
+            self.scheduler.batch_result_processor.commit_handoff_token(
+                req,
+                committed_output_id,
+                logprob_token_id=logprob_token_id,
+                logprob=logprob,
+                top_logprobs_val=top_logprobs_val,
+                top_logprobs_idx=top_logprobs_idx,
+                sampling_mask_idx=sampling_mask_idx,
+                sampling_mask_logprobs=sampling_mask_logprobs,
             )
         decode_req.req.cached_tokens = cached_tokens[0].item()
         # The prefill node already reported its prefix-cache hit in
@@ -2965,42 +2990,6 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest):
                 output_dsa_topk_indices = None
             decode_req.req.output_dsa_topk_indices = output_dsa_topk_indices
 
-        if decode_req.req.return_logprob and not resumes_output:
-            decode_req.req.logprob.output_token_logprobs_val.append(
-                output_token_logprobs_val[0].item()
-            )
-            decode_req.req.logprob.output_token_logprobs_idx.append(
-                output_token_logprobs_idx[0].item()
-            )
-            decode_req.req.logprob.output_top_logprobs_val.append(
-                output_top_logprobs_val[
-                    : decode_req.req.logprob.top_logprobs_num
-                ].tolist()
-            )
-            decode_req.req.logprob.output_top_logprobs_idx.append(
-                output_top_logprobs_idx[
-                    : decode_req.req.logprob.top_logprobs_num
-                ].tolist()
-            )
-            if decode_req.req.logprob.token_ids_logprob is not None:
-                # Prefill omits these scores; reserve a row for the streaming offset.
-                decode_req.req.logprob.output_token_ids_logprobs_val.append([])
-                decode_req.req.logprob.output_token_ids_logprobs_idx.append([])
-        if decode_req.req.return_sampling_mask and not resumes_output:
-            assert output_token_sampling_mask_idx is not None, (
-                "sampling mask buffer disabled on decode side"
-            )
-            sampling_mask_len = int(output_token_sampling_mask_len[0].item())
-            num_logprobs = (
-                sampling_mask_len
-                if decode_req.req.sampling_logprobs_mode == "support"
-                else 1
-            )
-            decode_req.req.sampling_mask_rows.append(
-                output_token_sampling_mask_idx[:sampling_mask_len].cpu().numpy(),
-                output_token_sampling_logprobs[:num_logprobs].cpu().numpy(),
-            )
-
         decode_req.kv_receiver.clear()
         decode_req.kv_receiver = None
         decode_req.req.time_stats.set_wait_queue_entry_time()

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -1261,24 +1261,16 @@ def _apply_decode_logprobs(
             max_accept = 1
 
         for j, tok_id in enumerate(accepted_ids):
-            req.logprob.output_token_logprobs_val.append(accepted_logprobs[j])
-            req.logprob.output_token_logprobs_idx.append(tok_id)
-            if req.logprob.top_logprobs_num > 0:
-                flat_idx = i * max_accept + j
-                req.logprob.output_top_logprobs_val.append(
-                    logits_output.next_token_top_logprobs_val[flat_idx]
-                )
-                req.logprob.output_top_logprobs_idx.append(
-                    logits_output.next_token_top_logprobs_idx[flat_idx]
-                )
-            if req.logprob.token_ids_logprob is not None:
-                flat_idx = i * max_accept + j
-                req.logprob.output_token_ids_logprobs_val.append(
-                    logits_output.next_token_token_ids_logprobs_val[flat_idx]
-                )
-                req.logprob.output_token_ids_logprobs_idx.append(
-                    logits_output.next_token_token_ids_logprobs_idx[flat_idx]
-                )
+            self.logprob_result_processor.append_output_token_logprobs(
+                req,
+                tok_id,
+                accepted_logprobs[j],
+                row=i * max_accept + j,
+                top_logprobs_val=logits_output.next_token_top_logprobs_val,
+                top_logprobs_idx=logits_output.next_token_top_logprobs_idx,
+                token_ids_logprobs_val=logits_output.next_token_token_ids_logprobs_val,
+                token_ids_logprobs_idx=logits_output.next_token_token_ids_logprobs_idx,
+            )
 
     def add_sampling_mask_return_values(
         self,
@@ -1587,6 +1579,35 @@ def _mamba_can_cache_finished_req(
         )
         return len(retained[i]) == num_verified
 
+    def commit_handoff_token(
+        self,
+        req: Req,
+        token_id: int,
+        *,
+        logprob_token_id: Optional[int] = None,
+        logprob: Optional[float] = None,
+        top_logprobs_val: Optional[List[float]] = None,
+        top_logprobs_idx: Optional[List[int]] = None,
+        sampling_mask_idx: Optional[np.ndarray] = None,
+        sampling_mask_logprobs: Optional[np.ndarray] = None,
+    ) -> None:
+        """Commit the PD prefill worker's sampled token to a decode request."""
+        req.output_ids.append(token_id)
+        self._maybe_update_reasoning_tokens(req, token_id)
+        if req.return_logprob:
+            # Prefill omits token-ids scores; reserve a row for the streaming offset.
+            self.logprob_result_processor.append_output_token_logprobs(
+                req,
+                logprob_token_id,
+                logprob,
+                top_logprobs_val=[top_logprobs_val],
+                top_logprobs_idx=[top_logprobs_idx],
+                token_ids_logprobs_val=[[]],
+                token_ids_logprobs_idx=[[]],
+            )
+        if req.return_sampling_mask:
+            req.sampling_mask_rows.append(sampling_mask_idx, sampling_mask_logprobs)
+
     def _maybe_update_reasoning_tokens(
         self,
         req: Req,

python/sglang/srt/managers/scheduler_components/logprob_result_processor.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass
 from typing import (
     List,
+    Optional,
     Tuple,
 )
 
@@ -326,9 +327,20 @@ def add_logprob_return_values(
         output: LogitsProcessorOutput,
     ):
         """Attach logprobs to the return values."""
-        if output.next_token_logprobs is not None:
-            req.logprob.output_token_logprobs_val.append(output.next_token_logprobs[i])
-            req.logprob.output_token_logprobs_idx.append(next_token_ids[i])
+        self.append_output_token_logprobs(
+            req,
+            next_token_ids[i],
+            (
+                output.next_token_logprobs[i]
+                if output.next_token_logprobs is not None
+                else None
+            ),
+            row=i,
+            top_logprobs_val=output.next_token_top_logprobs_val,
+            top_logprobs_idx=output.next_token_top_logprobs_idx,
+            token_ids_logprobs_val=output.next_token_token_ids_logprobs_val,
+            token_ids_logprobs_idx=output.next_token_token_ids_logprobs_idx,
+        )
 
         # Only add input logprobs if there are input tokens to process
         # Note: For prefill-only requests with default logprob_start_len, this will be 0,
@@ -345,35 +357,40 @@ def add_logprob_return_values(
         else:
             self._initialize_empty_logprob_containers(req)
 
-        if (
-            req.logprob.top_logprobs_num > 0
-            and output.next_token_top_logprobs_val is not None
-        ):
-            # Guarded like next_token_logprobs above: a backend may leave the
-            # top-logprob fields unset even for a request that asked for them
-            # (indexing None raises TypeError).
-            req.logprob.output_top_logprobs_val.append(
-                output.next_token_top_logprobs_val[i]
-            )
-            req.logprob.output_top_logprobs_idx.append(
-                output.next_token_top_logprobs_idx[i]
-            )
+        return num_input_logprobs
 
+    @staticmethod
+    def append_output_token_logprobs(
+        req: Req,
+        token_id: int,
+        logprob: Optional[float],
+        *,
+        row: int = 0,
+        top_logprobs_val: Optional[list] = None,
+        top_logprobs_idx: Optional[list] = None,
+        token_ids_logprobs_val: Optional[list] = None,
+        token_ids_logprobs_idx: Optional[list] = None,
+    ) -> None:
+        """Append one output token's logprobs; list arguments are read at `row`,
+        and None (a backend may leave a field unset) appends nothing."""
+        if logprob is not None:
+            req.logprob.output_token_logprobs_val.append(logprob)
+            req.logprob.output_token_logprobs_idx.append(token_id)
+        if req.logprob.top_logprobs_num > 0 and top_logprobs_val is not None:
+            req.logprob.output_top_logprobs_val.append(top_logprobs_val[row])
+            req.logprob.output_top_logprobs_idx.append(top_logprobs_idx[row])
         if (
             req.logprob.token_ids_logprob is not None
-            and output.next_token_token_ids_logprobs_val is not None
+            and token_ids_logprobs_val is not None
         ):
-            # Convert GPU tensor to list if needed
-            logprobs_val = output.next_token_token_ids_logprobs_val[i]
+            logprobs_val = token_ids_logprobs_val[row]
             if isinstance(logprobs_val, torch.Tensor):
                 logprobs_val = logprobs_val.tolist()
             req.logprob.output_token_ids_logprobs_val.append(logprobs_val)
             req.logprob.output_token_ids_logprobs_idx.append(
-                output.next_token_token_ids_logprobs_idx[i]
+                token_ids_logprobs_idx[row]
             )
 
-        return num_input_logprobs
-
     def _initialize_empty_logprob_containers(self, req: Req) -> None:
         """
         Initialize logprob fields to empty lists if unset.

test/registered/core/test_srt_endpoint.py

@@ -887,7 +887,7 @@ def test_detokenize_invalid_tokens(self):
         r2 = self.session.post(
             self.detokenize_url, json={"model": self.model, "tokens": [1, -1, 2]}
         )
-        self.assertEqual(r2.status_code, 500)
+        self.assertEqual(r2.status_code, 400)
 
 
 # ---------------------------------------------------------------------------

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -1,4 +1,5 @@
 import concurrent.futures
+import dataclasses
 import threading
 import unittest
 from collections import defaultdict
@@ -49,6 +50,9 @@
 from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk
 from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload
 from sglang.srt.managers.schedule_batch import ReqKvInfo
+from sglang.srt.managers.scheduler_components.batch_result_processor import (
+    SchedulerBatchResultProcessor,
+)
 from sglang.srt.mem_cache.deepseek_v4_compress_state import (
     CompressStatePool,
     c4_state_transfer_indices,
@@ -1845,12 +1849,16 @@ def test_sampling_mask_row_reaches_decode_unchanged(self):
                         sampling_logprobs_mode=mode,
                     )
                 )
+                fields = dict.fromkeys(
+                    f.name for f in dataclasses.fields(SchedulerBatchResultProcessor)
+                )
+                fields["model_config"] = SimpleNamespace(
+                    self_opened_think_start_ids=None
+                )
                 queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
                 queue.scheduler = SimpleNamespace(
                     kv_checksum_computer=None,
-                    batch_result_processor=SimpleNamespace(
-                        _maybe_update_reasoning_tokens=lambda req, token_id: None
-                    ),
+                    batch_result_processor=SchedulerBatchResultProcessor(**fields),
                 )
                 queue.spec_algorithm = SimpleNamespace(is_none=lambda: True)
                 queue.metadata_buffers = buffers
@@ -1859,6 +1867,7 @@ def test_sampling_mask_row_reaches_decode_unchanged(self):
                     bootstrap_host="127.0.0.1",
                     bootstrap_room=9,
                     output_ids=[],
+                    require_reasoning=False,
                     return_logprob=False,
                     return_sampling_mask=True,
                     sampling_logprobs_mode=mode,