simplify further

Signed-off-by: Balaram Buddharaju <[email protected]>

Author: brb-nv

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -362,33 +362,16 @@ def _fetch_new_requests_attention_dp(
         # Get active request counts across all ranks.
         all_ranks_num_active_requests = []
         all_ranks_num_active_tokens = []
-        num_active_tokens = sum(
-            [req.py_orig_prompt_len for req in activate_requests])
 
         if self.dist.has_cp_helix:
-            # When CP is enabled with Helix parallelism, we need to gather from all ranks
-            # in the TP x CP space. CP ranks within the same DP group (same tp_rank) handle
-            # the same requests with different token portions (sequence is split across CP ranks).
-            responses_list = self.dist.tp_cp_allgather(
-                [len(activate_requests), num_active_tokens])
-
-            aggregated_responses = []
-            for dp_group_idx in range(self.dist.tp_size):
-                # Get all entries for this DP group (cp_size entries per group).
-                group_start = dp_group_idx * self.dist.cp_size
-                group_end = (dp_group_idx + 1) * self.dist.cp_size
-                group_entries = responses_list[group_start:group_end]
-
-                # All CP ranks within a DP group should have the same number of requests.
-                assert all(entry[0] == group_entries[0][0] for entry in group_entries), \
-                    f"CP ranks within DP group {dp_group_idx} have mismatched request counts: " \
-                    f"{[entry[0] for entry in group_entries]}"
-                # Use token count from cp_rank0.
-                aggregated_responses.append(group_entries[0])
-            responses_list = aggregated_responses
+            num_active_tokens = sum(
+                [req.total_input_len_cp for req in activate_requests])
         else:
-            responses_list = self.dist.tp_allgather(
-                [len(activate_requests), num_active_tokens])
+            num_active_tokens = sum(
+                [req.py_orig_prompt_len for req in activate_requests])
+
+        responses_list = self.dist.tp_allgather(
+            [len(activate_requests), num_active_tokens])
 
         for num_active_requests, num_active_tokens in responses_list:
             all_ranks_num_active_requests.append(num_active_requests)