Fix regex and logprob conflicts when chunked prefilling (sgl-project#…

…1703)
yukavio · Oct 18, 2024 · 9e0dac1 · 9e0dac1
1 parent a95d558
commit 9e0dac1
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
@@ -789,15 +789,15 @@ def process_batch_result_prefill(self, batch: ScheduleBatch, result):
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
                         self.tree_cache.cache_unfinished_req(req)
 
-                if req.regex_fsm is not None:
-                    req.regex_fsm_state = req.regex_fsm.get_next_state(
-                        req.regex_fsm_state, next_token_ids[i]
-                    )
+                    if req.regex_fsm is not None:
+                        req.regex_fsm_state = req.regex_fsm.get_next_state(
+                            req.regex_fsm_state, next_token_ids[i]
+                        )
 
-                if req.return_logprob:
-                    logprob_pt += self.add_logprob_return_values(
-                        i, req, logprob_pt, next_token_ids, logits_output
-                    )
+                    if req.return_logprob:
+                        logprob_pt += self.add_logprob_return_values(
+                            i, req, logprob_pt, next_token_ids, logits_output
+                        )
         else:  # embedding or reward model
             assert batch.extend_num_tokens != 0
             embeddings, bid = result