Add attention dp scheduling logic

Signed-off-by: Shunkang <[email protected]>

Author: Shunkang

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -244,30 +244,105 @@ def _fetch_new_requests_attention_dp(
         new_requests = self._fetch_and_process_requests(
             total_num_active_requests, total_max_num_active_requests)
 
-        # Balance requests across ranks
-        num_new_requests_all_ranks = len(new_requests)
-        self.expected_num_active_requests = max(
-            (total_num_active_requests + num_new_requests_all_ranks +
-             self.dist.tp_size - 1) // self.dist.tp_size,
-            max(all_ranks_num_active_requests),
-        )
-
-        new_requests_cur_rank = self._balance_requests_across_ranks(
-            new_requests, all_ranks_num_active_requests)
+        # Schedule attention dp requests
+        new_requests_cur_rank = self._schedule_attention_dp_requests(
+            num_active_requests, new_requests, all_ranks_num_active_requests)
 
         # Update performance metrics
         if self.enable_iter_perf_stats and self.start_times:
             self._update_new_active_requests_queue_latency(
                 new_requests_cur_rank)
 
         # Update counters
-        self.num_fetch_requests += num_new_requests_all_ranks
+        self.num_fetch_requests += len(new_requests)
         self.num_fetch_requests_cur_rank += len(new_requests_cur_rank)
 
         # Merge requests and add to active list
         new_requests_cur_rank = self._merge_requests(new_requests_cur_rank)
         return new_requests_cur_rank
 
+    def _schedule_attention_dp_requests(
+            self, num_active_requests: int,
+            new_requests: List[RequestQueueItem],
+            all_ranks_num_active_requests: List[int]) -> List[RequestQueueItem]:
+        """Schedule attention dp requests."""
+        # Separate the requests into two groups
+        # 1. requests without schedule params or with schedule params that don't specify attention dp rank
+        # 2. requests with schedule params that specify attention dp rank
+        requests_specified_attention_dp_rank = []
+        for req_item in new_requests:
+            if req_item.request.schedule_params is not None and \
+                req_item.request.schedule_params.attention_dp_rank == self.dist.tp_rank:
+                requests_specified_attention_dp_rank.append(req_item)
+
+        # Routing requests to the corresponding attention dp without exceeding the max_num_active_requests
+        new_requests_cur_rank = []
+        new_requests_cur_rank_waiting = []
+        new_requests_cur_rank_relax = []
+
+        available_slots = self.max_num_active_requests - num_active_requests
+
+        for req_item in requests_specified_attention_dp_rank:
+            is_relax = req_item.request.schedule_params.attention_dp_relax
+
+            if len(new_requests_cur_rank) < available_slots:
+                # Prioritize the non-relax requests
+                target_list = new_requests_cur_rank_relax if is_relax else new_requests_cur_rank
+                target_list.append(req_item)
+            else:
+                # Add to waiting queue
+                target_list = new_requests_cur_rank_relax if is_relax else new_requests_cur_rank_waiting
+                target_list.append(req_item)
+
+        items_to_move = available_slots - len(new_requests_cur_rank)
+        if items_to_move > 0:
+            new_requests_cur_rank.extend(
+                new_requests_cur_rank_relax[:items_to_move])
+            new_requests_cur_rank_relax = new_requests_cur_rank_relax[
+                items_to_move:]
+
+        # Allgather the non-scheduled requests across ranks
+        # TODO: Remove the padding overhead
+        new_requests_cur_rank_relax_ids = [
+            req_item.id for req_item in new_requests_cur_rank_relax
+        ]
+        padding_num = self.max_num_active_requests - len(
+            new_requests_cur_rank_relax)
+        for _ in range(padding_num):
+            new_requests_cur_rank_relax_ids.append(None)
+        non_scheduled_requests_id = self.dist.tp_allgather(
+            new_requests_cur_rank_relax_ids)
+        non_scheduled_requests_id = [
+            req_id for req_id in non_scheduled_requests_id if req_id is not None
+        ]
+
+        # Non-scheduled requests should be same across ranks
+        non_scheduled_requests = []
+        for req_item in new_requests:
+            if req_item.id in non_scheduled_requests_id:
+                non_scheduled_requests.append(req_item)
+            elif req_item.request.schedule_params is None or \
+                req_item.request.schedule_params is not None and \
+                req_item.request.schedule_params.attention_dp_rank is None:
+                non_scheduled_requests.append(req_item)
+
+        # Put the request back to the waiting queue
+        self.waiting_queue.extendleft(new_requests_cur_rank_waiting)
+
+        # TODO: Balance the no attention dp rank requests and relax requests across ranks
+        num_new_requests_all_ranks = len(new_requests)
+        total_num_active_requests = sum(all_ranks_num_active_requests)
+        self.expected_num_active_requests = max(
+            (total_num_active_requests + num_new_requests_all_ranks +
+             self.dist.tp_size - 1) // self.dist.tp_size,
+            max(all_ranks_num_active_requests),
+        )
+
+        new_requests_cur_rank = self._balance_requests_across_ranks(
+            non_scheduled_requests, )
+
+        return new_requests_cur_rank
+
     def _handle_request_broadcasting(self,
                                      new_requests: List[RequestQueueItem]):
         """Handle broadcasting of requests and Python objects across ranks."""