factor out child req id generation logic

Signed-off-by: Jaedeok Kim <[email protected]>

Author: jaedeok-nvidia

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -325,6 +325,27 @@ def _get_request_id(self):
         self._next_req_id = (self._next_req_id + 1) & ((1 << 64) - 1)
         return self._next_req_id
 
+    def _generate_child_request_ids(
+            self, request: ExecutorRequest) -> List[int] | None:
+        """ Generate child request IDs if needed. """
+        child_req_ids = None
+        sampling_config = request.sampling_config
+        beam_width = (sampling_config.beam_width
+                      if sampling_config.beam_width else 1)
+        num_return_sequences = (sampling_config.num_return_sequences
+                                if sampling_config.num_return_sequences else 1)
+
+        # Create child requests if beam_width == 1 and num_return_sequences > 1.
+        if beam_width == 1 and num_return_sequences > 1:
+            child_req_ids = []
+            for _ in range(num_return_sequences - 1):
+                child_req_id = self._get_request_id()
+                if self.enable_iter_perf_stats:
+                    self.start_times[child_req_id] = time.time()
+                child_req_ids.append(child_req_id)
+
+        return child_req_ids
+
     def enqueue_requests(self, requests: List[ExecutorRequest]):
         """
         Enqueue new requests
@@ -339,21 +360,8 @@ def enqueue_requests(self, requests: List[ExecutorRequest]):
                 if self.enable_iter_perf_stats:
                     self.start_times[req_id] = time.time()
 
-                # Generate child request IDs if needed
-                child_req_ids = None
-                sampling_config = request.sampling_config
-                beam_width = sampling_config.beam_width
-                num_return_sequences = sampling_config.num_return_sequences or beam_width
-
-                if beam_width == 1 and num_return_sequences > 1:
-                    # Reserve request ids for child requests.
-                    child_req_ids = []
-                    for _ in range(num_return_sequences - 1):
-                        child_req_id = self._get_request_id()
-                        if self.enable_iter_perf_stats:
-                            self.start_times[child_req_id] = time.time()
-                        child_req_ids.append(child_req_id)
-
+                # Reserve child request ids if needed.
+                child_req_ids = self._generate_child_request_ids(request)
                 self.request_queue.put(
                     RequestQueueItem(req_id,
                                      request,
@@ -476,23 +484,8 @@ def enqueue_request(self,
             if self.enable_iter_perf_stats:
                 self.start_times[req_id] = time.time()
 
-            # Generate child request IDs if needed
-            child_req_ids = None
-            sampling_config = request.sampling_config
-            beam_width = (sampling_config.beam_width
-                          if sampling_config.beam_width else 1)
-            num_return_sequences = (sampling_config.num_return_sequences if
-                                    sampling_config.num_return_sequences else 1)
-
-            # Only create child requests if beam_width == 1 and num_return_sequences > 1
-            if beam_width == 1 and num_return_sequences > 1:
-                child_req_ids = []
-                for i in range(num_return_sequences - 1):
-                    child_req_id = self._get_request_id()
-                    if self.enable_iter_perf_stats:
-                        self.start_times[child_req_id] = time.time()
-                    child_req_ids.append(child_req_id)
-
+            # Reserve child request ids if needed.
+            child_req_ids = self._generate_child_request_ids(request)
             self.request_queue.put(
                 RequestQueueItem(req_id,
                                  request,