diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index f0e78f45b..34bf39289 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -433,6 +433,9 @@ class Scheduler:
                         result = self.run_batch(batch)
                         self.process_batch_result(batch, result)
 
+                    if self.running_batch.is_empty():
+                        self.running_batch = None
+
                     if self.running_batch is None:
                         break
 
@@ -772,9 +775,6 @@ class Scheduler:
         if self.tp_rank == 0 and self.decode_forward_ct % 40 == 0:
             self.print_decode_stats()
 
-        if self.running_batch.is_empty():
-            self.running_batch = None
-
     def add_logprob_return_values(
         self,
         i: int,