diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 483195c39..172c30adb 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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: - 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.batch_result_processor.move_logprobs_to_cpu( + batch=batch, + logits_output=logits_output, + ) def advance_logprob_pt(i: int, req: Req) -> None: nonlocal logprob_pt 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 4a6d0df47..98100df52 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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, diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index 0cdbf0340..50708530f 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/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(