Fix chunked prefill with output logprob (#2083)

2024-11-18 13:16:28 -08:00
parent 66318ffe96
commit 80e2c4a8de
2 changed files with 38 additions and 9 deletions
--- a/python/sglang/srt/managers/schedule_policy.py
+++ b/python/sglang/srt/managers/schedule_policy.py
@@ -302,7 +302,11 @@ class PrefillAdder:
            if (
                self.rem_chunk_tokens is None
                or input_tokens <= self.rem_chunk_tokens
-                or (req.return_logprob and req.normalized_prompt_logprob is None)
+                or (
+                    req.return_logprob
+                    and req.normalized_prompt_logprob is None
+                    and req.logprob_start_len != len(req.origin_input_ids) - 1
+                )
            ):
                # Non-chunked prefill
                self.can_run_list.append(req)