[refactor] Move iter_counter handling to PyExecutor

- Moved iter_counter in PyExecutor to ensure consistency in tracking iterations.
- This allows tracking of iteration where scheduled requests are empty.

Signed-off-by: Robin Kobus <[email protected]>

Author: Funatiq

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -153,7 +153,6 @@ def __init__(
         self.llm_args.batch_wait_timeout_iters = 0
         self.llm_args.batch_wait_max_tokens_ratio = 0.0
         self.llm_args.max_num_tokens = seq_info.max_num_tokens
-        self.iter_counter = 0
 
         # NOTE (lucaslie): not a declared base member in the base class; required by PyExecutor...
         self.max_beam_width = max_beam_width

tensorrt_llm/_torch/expert_statistic.py

@@ -29,11 +29,15 @@ def create(rank_id: int):
             rank_id, start, stop)
 
     @staticmethod
-    def set_iter(iter_id: int) -> bool:
+    def should_record() -> bool:
         if ExpertStatistic.expert_statistic_obj is not None:
-            return ExpertStatistic.expert_statistic_obj._set_iter(iter_id)
-        else:
-            return False
+            return ExpertStatistic.expert_statistic_obj._should_record
+        return False
+
+    @staticmethod
+    def set_iter(iter_id: int) -> None:
+        if ExpertStatistic.expert_statistic_obj is not None:
+            ExpertStatistic.expert_statistic_obj._set_iter(iter_id)
 
     @staticmethod
     def set_layer(layer_id: int) -> None:
@@ -57,10 +61,10 @@ def __init__(self, rank_id: int, start: int, stop: int) -> None:
         self._records = {}
 
     @property
-    def should_record(self) -> bool:
+    def _should_record(self) -> bool:
         return self.current_iter_id is not None and self.start <= self.current_iter_id < self.stop
 
-    def _set_iter(self, iter_id: int) -> bool:
+    def _set_iter(self, iter_id: int) -> None:
         self.current_iter_id = iter_id
         if iter_id == self.stop:
             logger.info(
@@ -74,14 +78,13 @@ def _set_iter(self, iter_id: int) -> bool:
                     json.dump(self._meta_info, f)
             safetensors.torch.save_file(
                 self._records, f"{path}/rank{self.rank_id}.safetensors")
-        return self.should_record
 
     def _set_layer(self, layer: int) -> None:
         self.current_layer = layer
 
     def _maybe_add_info(self, expert_count: int,
                         token_selected_experts: torch.Tensor) -> None:
-        if not self.should_record:
+        if not self._should_record:
             return
 
         if self._meta_info is None:

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -164,7 +164,6 @@ def __del__(self):
     def maybe_get_cuda_graph(
         self,
         batch: ScheduledRequests,
-        iter_counter: int,
         enable_spec_decode: bool,
         attn_metadata: Any,
         spec_metadata: Optional[Any] = None,
@@ -180,7 +179,7 @@ def maybe_get_cuda_graph(
         - The key for the graph, if applicable.
         """
         # disable when doing statistic
-        if ExpertStatistic.set_iter(iter_counter):
+        if ExpertStatistic.should_record():
             return None, None, None
 
         can_run_cuda_graph = batch.can_run_cuda_graph

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -364,7 +364,6 @@ def __init__(
         if self.use_mrope:
             self.mrope_position_ids_cuda = torch.empty(
                 (3, 1, self.max_num_tokens), dtype=torch.int, device='cuda')
-        self.iter_counter = 0
 
         # Pre-allocated buffers for draft model to avoid implicit synchronization
         # These are used to build index tensors without creating tensors from Python lists
@@ -2572,7 +2571,6 @@ def forward(self,
 
             maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph(
                 padded_requests,
-                iter_counter=self.iter_counter,
                 enable_spec_decode=self.enable_spec_decode,
                 attn_metadata=attn_metadata,
                 spec_metadata=spec_metadata,
@@ -2596,7 +2594,6 @@ def forward(self,
                 new_tensors_device, cache_indirection_buffer,
                 num_accepted_tokens_device, req_id_to_old_request)
 
-            self.iter_counter += 1
             with with_shared_pool(self.cuda_graph_runner.get_graph_pool()):
                 if not can_run_graph:
                     # Fallback to eager execution if graph was not used

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -11,6 +11,7 @@
 
 import torch
 
+from tensorrt_llm._torch.expert_statistic import ExpertStatistic
 from tensorrt_llm.serve.responses_utils import get_steady_clock_now_in_seconds
 
 try:
@@ -137,6 +138,7 @@ def __init__(self,
 
         self.peft_cache_config = peft_cache_config
 
+        self.iter_counter = 0
         # profile config
         self.profile_start_iters, self.profile_stop_iters = _load_iteration_indexes(
             PROFILE_START_STOP_ENV_VAR_NAME)
@@ -575,7 +577,7 @@ def profile_step():
                 formatted_timestamp = datetime.datetime.now().strftime(
                     "%Y-%m-%d %H:%M:%S")
                 logger.info(
-                    f"iter = {self.model_engine.iter_counter}, "
+                    f"iter = {self.iter_counter}, "
                     f"global_rank = {self.global_rank}, "
                     f"rank = {self.dist.rank}, "
                     f"currank_total_requests = {self.executor_request_queue.num_fetch_requests_cur_rank}/"
@@ -705,7 +707,7 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
         stats.cpu_mem_usage = 0
         stats.pinned_mem_usage = 0
 
-        stats.iter = self.model_engine.iter_counter
+        stats.iter = self.iter_counter
 
         kv_cache_manager = self.resource_manager.resource_managers.get(
             ResourceManagerType.KV_CACHE_MANAGER)
@@ -1004,6 +1006,8 @@ def _executor_loop_pp(self):
                                              self.active_requests,
                                              previous_batch)
 
+                self.iter_counter += 1
+
     def wait_on_pp_send_handles(self, microbatch_id):
         if self.send_handles[microbatch_id] is not None:
             self.send_handles[microbatch_id].wait()
@@ -1240,6 +1244,8 @@ def _executor_loop(self):
                                    iter_stats=iter_stats,
                                    iter_start_time=iter_start_time))
 
+                self.iter_counter += 1
+
     def _prepare_draft_requests(self):
         try:
             # Set draft tokens here to make the KV cache manager
@@ -1473,6 +1479,8 @@ def _executor_loop_overlap(self):
 
                 self._kv_connector_terminate_requests()
 
+                self.iter_counter += 1
+
     def _accept_draft_tokens(
         self, scheduled_batch: ScheduledRequests,
         target_outputs: SampleStateTensors,
@@ -1964,9 +1972,10 @@ def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0):
     def _forward_step(self,
                       scheduled_requests,
                       new_tensors_device: Optional[SampleStateTensors] = None):
+        ExpertStatistic.set_iter(self.iter_counter)
 
         @nvtx_range(
-            f"[Executor] _forward_step {self.model_engine.iter_counter + 1}: {len(scheduled_requests.context_requests)} ctx reqs, {len(scheduled_requests.generation_requests)} gen reqs"
+            f"[Executor] _forward_step {self.iter_counter}: {len(scheduled_requests.context_requests)} ctx reqs, {len(scheduled_requests.generation_requests)} gen reqs"
         )
         def forward(scheduled_requests, resource_manager, new_tensors_device,
                     gather_context_logits, cache_indirection_buffer):
@@ -2304,7 +2313,7 @@ def _handle_responses(self):
 
             # Skip active requests that are not scheduled
             if request.return_perf_metrics and request.py_decoding_iter >= 1:
-                request.update_perf_metrics(self.model_engine.iter_counter)
+                request.update_perf_metrics(self.iter_counter)
 
             request_done = False
             if request.py_decoding_iter == 1 or request.is_finished or \