Localize cur_batch field in Scheduler to avoid field-based state access (#29407)

python/sglang/srt/disaggregation/decode.py

@@ -1901,7 +1901,7 @@ class SchedulerDisaggregationDecodeMixin:
 
             # Get the next batch to run
             batch = self.get_next_disagg_decode_batch_to_run()
-            self.cur_batch = batch
+            self.cur_batch_for_debug = batch
 
             # Launch the current batch
             if batch:
@@ -1935,7 +1935,7 @@ class SchedulerDisaggregationDecodeMixin:
 
             # Get the next batch to run
             batch = self.get_next_disagg_decode_batch_to_run()
-            self.cur_batch = batch
+            self.cur_batch_for_debug = batch
             # overlap + spec + grammar is unsupported (would desync DP ranks).
             disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
 
@@ -1958,7 +1958,7 @@ class SchedulerDisaggregationDecodeMixin:
 
             # Run sample of the current batch
             # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed.
-            self.launch_batch_sample_if_needed(batch_result)
+            self.launch_batch_sample_if_needed(batch_result, batch)
 
             # Update last_batch
             self.last_batch = batch

python/sglang/srt/disaggregation/prefill.py

@@ -502,7 +502,7 @@ class SchedulerDisaggregationPrefillMixin:
 
             # Get the next batch to run
             batch = self.get_next_disagg_prefill_batch_to_run()
-            self.cur_batch = batch
+            self.cur_batch_for_debug = batch
 
             # Launch the current batch
             if batch:
@@ -536,7 +536,7 @@ class SchedulerDisaggregationPrefillMixin:
 
             # Get the next batch to run
             batch = self.get_next_disagg_prefill_batch_to_run()
-            self.cur_batch = batch
+            self.cur_batch_for_debug = batch
 
             # Launch the current batch
             if batch:
@@ -559,7 +559,7 @@ class SchedulerDisaggregationPrefillMixin:
 
             # Run sample of the current batch
             # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed.
-            self.launch_batch_sample_if_needed(batch_result)
+            self.launch_batch_sample_if_needed(batch_result, batch)
 
             # Update last_batch
             self.last_batch = batch

python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py

@@ -234,7 +234,7 @@ class SchedulerMlxOverlapMixin:
             ):
                 pending_curr = pending_next
                 pending_next = None
-                self.cur_batch = pending_curr.schedule_batch
+                self.cur_batch_for_debug = pending_curr.schedule_batch
                 self.last_batch = pending_curr.schedule_batch
                 if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
                     self.invariant_checker.self_check_during_busy()
@@ -247,7 +247,7 @@ class SchedulerMlxOverlapMixin:
                 self.result_queue.popleft()
                 pending_next = None
             next_batch = self.get_next_batch_to_run()
-            self.cur_batch = next_batch
+            self.cur_batch_for_debug = next_batch
             if next_batch:
                 pending_curr = _launch_fresh(next_batch)
                 self.result_queue.append(pending_curr)

python/sglang/srt/managers/scheduler.py

@@ -319,7 +319,7 @@ class Scheduler(
         self.is_initializing = True
         # init_soft_watchdog starts a daemon thread that reads these on its first tick.
         self.forward_ct: int = 0
-        self.cur_batch: Optional[ScheduleBatch] = None
+        self.cur_batch_for_debug: Optional[ScheduleBatch] = None
         self.init_soft_watchdog(server_args)
 
         # Parse args
@@ -965,7 +965,7 @@ class Scheduler(
         # The running decoding batch for continuous batching
         self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False)
         # The current forward batch
-        self.cur_batch: Optional[ScheduleBatch] = None
+        self.cur_batch_for_debug: Optional[ScheduleBatch] = None
         # The last forward batch
         self.last_batch: Optional[ScheduleBatch] = None
         self.forward_ct = 0
@@ -1542,7 +1542,7 @@ class Scheduler(
 
             # Get the next batch to run
             batch = self.get_next_batch_to_run()
-            self.cur_batch = batch
+            self.cur_batch_for_debug = batch
 
             # Launch the current batch
             if batch:
@@ -1583,7 +1583,7 @@ class Scheduler(
 
             # Get the next batch to run
             batch = self.get_next_batch_to_run()
-            self.cur_batch = batch
+            self.cur_batch_for_debug = batch
             disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
 
             # If we do not need to overlap the current batch with the last batch,
@@ -1617,7 +1617,7 @@ class Scheduler(
             # Run sample of the current batch
             # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed.
             if self.is_generation:
-                self.launch_batch_sample_if_needed(batch_result)
+                self.launch_batch_sample_if_needed(batch_result, batch)
 
             # Update last_batch
             self.last_batch = batch
@@ -3415,7 +3415,7 @@ class Scheduler(
         self.future_map.stash(future_indices, payload)
 
     def launch_batch_sample_if_needed(
-        self, batch_result: GenerationBatchResult
+        self, batch_result: GenerationBatchResult, cur_batch: ScheduleBatch
     ) -> Union[GenerationBatchResult]:
         # TODO(lsyin): make the delayed sample a default behavior after
         # unifying the forward_batch_generation interface (related to spec V2).
@@ -3429,8 +3429,8 @@ class Scheduler(
             # Delay-sample is non-spec only; relays the sampled bonus tokens.
             self._relay_forward_payload(batch_result.future_indices, batch_result)
             batch_result.copy_to_cpu(
-                return_logprob=self.cur_batch.return_logprob,
-                return_hidden_states=self.cur_batch.return_hidden_states,
+                return_logprob=cur_batch.return_logprob,
+                return_hidden_states=cur_batch.return_hidden_states,
             )
 
         # Release the closure and large GPU tensors that are no longer needed.
@@ -3731,7 +3731,7 @@ class Scheduler(
     def flush_cache(self, empty_cache: bool = True):
         """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache."""
         if self.is_fully_idle():
-            self.cur_batch = None
+            self.cur_batch_for_debug = None
             self.last_batch = None
             self.tree_cache.reset()
             self.req_to_token_pool.clear()
@@ -4021,7 +4021,7 @@ class Scheduler(
                     self.running_batch.merge_batch(self.last_batch)
 
         self.last_batch = None
-        self.cur_batch = None
+        self.cur_batch_for_debug = None
 
         if recv_req.mode == "retract" and not self.running_batch.is_empty():
             self.running_batch.filter_batch()

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -470,14 +470,16 @@ def create_scheduler_watchdog(
             scheduler.pool_stats_observer.get_pool_stats(),
         )
         return (
-            f"{scheduler.cur_batch.batch_size()=}\n"
-            f"{scheduler.cur_batch.reqs=}\n" + "\n".join(messages)
+            f"{scheduler.cur_batch_for_debug.batch_size()=}\n"
+            f"{scheduler.cur_batch_for_debug.reqs=}\n" + "\n".join(messages)
         )
 
     return WatchdogRaw(
         debug_name="Scheduler",
         get_counter=lambda: scheduler.forward_ct,
-        is_active=lambda: scheduler.is_initializing or scheduler.cur_batch is not None,
+        is_active=lambda: (
+            scheduler.is_initializing or scheduler.cur_batch_for_debug is not None
+        ),
         watchdog_timeout=watchdog_timeout,
         soft=soft,
         dump_info=dump_info,

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -110,8 +110,9 @@ class SchedulerPPMixin:
                 with torch.profiler.record_function("get_next_batch_to_run"):
                     self.mbs[mb_id] = self.get_next_batch_to_run()
                 self.running_mbs[mb_id] = self.running_batch
-                self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id]
-                if self.cur_batch:
+                cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id]
+                self.cur_batch_for_debug = cur_batch
+                if cur_batch:
                     server_is_idle = False
                     pp_proxy_tensors = self._pp_recv_proxy_tensors()
                 next_pp_outputs = None
@@ -125,9 +126,10 @@ class SchedulerPPMixin:
                         )
                     )
                 self._pp_commit_comm_work(self.send_proxy_work)
-                if self.cur_batch:
+                if cur_batch:
                     result, self.launch_event = self._pp_launch_batch(
                         mb_id,
+                        cur_batch,
                         pp_proxy_tensors,
                         self.mb_metadata,
                         self.last_rank_comm_queue,
@@ -148,7 +150,7 @@ class SchedulerPPMixin:
                         )
                     self.last_mbs[next_mb_id] = self.mbs[next_mb_id]
                 if not self.pp_group.is_last_rank:
-                    if self.cur_batch:
+                    if cur_batch:
                         self.device_module.current_stream().wait_event(
                             self.launch_event
                         )
@@ -252,8 +254,9 @@ class SchedulerPPMixin:
                 self.mbs[mb_id] = batch
                 self.running_mbs[mb_id] = self.running_batch
 
-                self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id]
-                if self.cur_batch:
+                cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id]
+                self.cur_batch_for_debug = cur_batch
+                if cur_batch:
                     server_is_idle = False
                     pp_proxy_tensors = self._pp_recv_proxy_tensors()
 
@@ -265,9 +268,10 @@ class SchedulerPPMixin:
                         )
                     )
                 self._pp_commit_comm_work(self.send_proxy_work)
-                if self.cur_batch:
+                if cur_batch:
                     result, self.launch_event = self._pp_launch_batch(
                         mb_id,
+                        cur_batch,
                         pp_proxy_tensors,
                         self.mb_metadata,
                         self.last_rank_comm_queue,
@@ -325,7 +329,7 @@ class SchedulerPPMixin:
                     send_transfer_work = self._pp_send_pyobj_to_next_stage(
                         transferred_rids, async_send=True
                     )
-                    if self.cur_batch:
+                    if cur_batch:
                         self.device_module.current_stream().wait_event(
                             self.launch_event
                         )
@@ -402,11 +406,12 @@ class SchedulerPPMixin:
                 self.mbs[mb_id] = batch
                 self.running_mbs[mb_id] = self.running_batch
 
-                self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id]
-                if self.cur_batch:
+                cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id]
+                self.cur_batch_for_debug = cur_batch
+                if cur_batch:
                     server_is_idle = False
                     pp_proxy_tensors = None
-                    if not self.cur_batch.forward_mode.is_prebuilt():
+                    if not cur_batch.forward_mode.is_prebuilt():
                         pp_proxy_tensors = self._pp_recv_proxy_tensors()
 
                 # early send output if possible
@@ -419,9 +424,10 @@ class SchedulerPPMixin:
                     )
                 self._pp_commit_comm_work(self.send_proxy_work)
 
-                if self.cur_batch:
+                if cur_batch:
                     result, self.launch_event = self._pp_launch_batch(
                         mb_id,
+                        cur_batch,
                         pp_proxy_tensors,
                         self.mb_metadata,
                         self.last_rank_comm_queue,
@@ -508,7 +514,7 @@ class SchedulerPPMixin:
                     send_transfer_work = self._pp_send_pyobj_to_next_stage(
                         transferred_rids, async_send=True
                     )
-                    if self.cur_batch and not self.cur_batch.forward_mode.is_prebuilt():
+                    if cur_batch and not cur_batch.forward_mode.is_prebuilt():
                         self.device_module.current_stream().wait_event(
                             self.launch_event
                         )
@@ -1240,6 +1246,7 @@ class SchedulerPPMixin:
     def _pp_launch_batch(
         self: Scheduler,
         mb_id: int,
+        cur_batch: ScheduleBatch,
         pp_proxy_tensors: PPProxyTensors,
         mb_metadata: List[Optional[PPBatchMetadata]],
         last_rank_comm_queue: deque,
@@ -1248,13 +1255,13 @@ class SchedulerPPMixin:
             with self.forward_stream_ctx:
                 self.forward_stream.wait_stream(self.schedule_stream)
                 set_time_batch(
-                    self.cur_batch.reqs,
+                    cur_batch.reqs,
                     "set_run_batch_cpu_start_time",
                     trace_only=True,
                 )
-                result = self.run_batch(self.cur_batch, pp_proxy_tensors)
+                result = self.run_batch(cur_batch, pp_proxy_tensors)
                 set_time_batch(
-                    self.cur_batch.reqs,
+                    cur_batch.reqs,
                     "set_run_batch_cpu_end_time",
                     trace_only=True,
                     attrs={"pp_mb_id": mb_id},
@@ -1270,7 +1277,7 @@ class SchedulerPPMixin:
                         (
                             event,
                             PPProxyTensors(
-                                self._pp_prepare_tensor_dict(result, self.cur_batch)
+                                self._pp_prepare_tensor_dict(result, cur_batch)
                             ),
                         )
                     )