[PD] Deduplicate PD logprob normalization (#27085)

python/sglang/srt/disaggregation/prefill.py

@@ -529,26 +529,10 @@ class SchedulerDisaggregationPrefillMixin:
         logprob_pt = 0
         # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
         next_token_ids = result.next_token_ids.tolist()
-        if batch.return_logprob:
+        self.batch_result_processor.move_logprobs_to_cpu(
-            if logits_output.next_token_logprobs is not None:
+            batch=batch,
-                logits_output.next_token_logprobs = (
+            logits_output=logits_output,
-                    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 advance_logprob_pt(i: int, req: Req) -> None:
             nonlocal logprob_pt

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

@@ -206,7 +206,7 @@ class SchedulerBatchResultProcessor:
 
             # Move next_token_ids and logprobs to cpu
             next_token_ids = next_token_ids.tolist()
-            self._move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
+            self.move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
 
             self._validate_pp_skip_output_comm(batch, result)
 
@@ -356,7 +356,7 @@ class SchedulerBatchResultProcessor:
                 embeddings = [tensor.tolist() for tensor in embeddings]
         return embeddings
 
-    def _move_logprobs_to_cpu(
+    def move_logprobs_to_cpu(
         self,
         *,
         batch: ScheduleBatch,

test/registered/disaggregation/test_disaggregation_basic.py

@@ -95,8 +95,9 @@ class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServer
             None,
         )
         self.assertIsNotNone(first_top_logprobs)
-        self.assertGreater(len(first_top_logprobs), 0)
+        self.assertEqual(len(first_top_logprobs), 5)
         self.assertIsInstance(first_top_logprobs[0].token, str)
+        self.assertIsInstance(first_top_logprobs[0].logprob, float)
 
     def test_structured_output(self):
         json_schema = json.dumps(