Fix overlap scheduler not take effect when outputing logprobs (#14096)

2025-11-28 18:15:56 +08:00
parent 0e8ce1e832
commit 45cf575852
2 changed files with 6 additions and 6 deletions
@@ -2003,7 +2003,7 @@ class Scheduler(
                    ).Event()
                    if batch_result.delay_sample_func is None:
                        self.future_map.store_to_map(future_indices, batch_result)
-                        batch_result.copy_to_cpu()
+                        batch_result.copy_to_cpu(return_logprob=batch.return_logprob)
                    else:
                        batch_result.future_indices = future_indices

@@ -2087,7 +2087,7 @@ class Scheduler(
            _batch_result = batch_result.delay_sample_func()
            assert _batch_result is batch_result
            self.future_map.store_to_map(batch_result.future_indices, batch_result)
-            batch_result.copy_to_cpu()
+            batch_result.copy_to_cpu(return_logprob=self.cur_batch.return_logprob)

    def process_batch_result(
        self,