Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/unittest/_torch/executor/test_executor_request_queue.py (11)

456-459: LGTM! Token tracking fixture added correctly.

The new all_ranks_num_active_tokens fixture properly complements the existing all_ranks_num_active_requests fixture and supports the token-aware scheduling feature being tested.

---

484-500: Test parameter updated correctly for token-aware scheduling.

The test now properly passes the all_ranks_num_active_tokens parameter to _schedule_attention_dp_requests, aligning with the production code's new token-aware scheduling logic.

---

662-673: Test correctly validates token parameter propagation.

The test properly verifies that _balance_requests_across_ranks is called with both the request counts and token counts parameters.

---

869-962: Test cases comprehensively cover token-based distribution scenarios.

The parameterized test cases thoroughly validate the token-aware distribution logic across various scenarios including balanced distribution, capacity limits, targeted vs. relaxed requests, and mixed cases.

---

995-1001: Good integration test enhancement for token tracking.

The test correctly creates mock token counts and passes them through the scheduling workflow, ensuring proper integration of the token-aware scheduling feature.

---

1011-1028: Comprehensive test for empty requests scenario.

The test properly validates that _balance_requests_across_ranks handles empty request lists correctly while maintaining the per-rank structure.

---

1030-1053: Test validates single request assignment to lowest token count rank.

The test correctly verifies that a single request is assigned to the rank with the lowest active token count (rank 2), demonstrating proper heap-based load balancing.

---

1055-1098: Well-structured test for multi-request load balancing.

The test correctly validates that multiple requests are distributed based on token counts and current rank loads. The expected distribution pattern follows the heap-based algorithm correctly.

---

1100-1130: Good test for capacity limit enforcement.

The test properly validates that the balancing algorithm respects the expected_num_active_requests capacity limit for each rank.

---

1132-1175: Test validates heap ordering correctly.

The test properly verifies that requests are distributed based on heap ordering (lowest active count first) when token counts are equal.

---

1177-1217: Excellent test for token count sorting behavior.

The test correctly validates that requests are sorted by token count in descending order before distribution, ensuring proper load balancing across ranks.

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (6)

331-342: LGTM! Token-aware scheduling implementation looks correct.

The implementation properly:

1. Collects token counts from active requests
2. Uses tp_allgather to distribute both request counts and token counts across ranks
3. Unpacks and stores the gathered data correctly

---

355-357: Good integration of token counts into scheduling.

The token counts are properly passed to _schedule_attention_dp_requests for token-aware load balancing.

---

373-376: Method signature correctly updated for token-aware scheduling.

The additional all_ranks_num_active_tokens parameter enables token-based load balancing in the scheduling logic.

---

421-423: Token counts properly integrated into balancing logic.

The token counts are correctly passed to _balance_requests_across_ranks for improved load distribution.

---

478-482: Method signature updated for token-aware balancing.

The additional all_ranks_num_active_tokens parameter enables the balancing algorithm to consider token loads when distributing requests.

---

490-491: Good initialization of heap with token counts.

Using actual token counts instead of 0 for initial heap values provides better load balancing from the start.

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (1)

305-312: Public API change: document and align naming.

The fetch_new_requests signature now accepts a List[LlmRequest]. Please:

• Document this breaking change in release notes/migration guide, and update the method docstring to reflect the new parameter.
• Prefer the parameter name active_requests (not activate_requests) for clarity/consistency across the codebase.

Suggested rename (non-functional, improves readability):

-    def fetch_new_requests(
-            self, activate_requests: List[LlmRequest]) -> List[LlmRequest]:
+    def fetch_new_requests(
+            self, active_requests: List[LlmRequest]) -> List[LlmRequest]:
@@
-        if self.enable_attention_dp:
-            return self._fetch_new_requests_attention_dp(activate_requests)
+        if self.enable_attention_dp:
+            return self._fetch_new_requests_attention_dp(active_requests)
-        else:
-            return self._fetch_new_requests_attention_tp(len(activate_requests))
+        else:
+            return self._fetch_new_requests_attention_tp(len(active_requests))

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

1187-1189: Call site updated to pass active requests list — good.

Internal usage now matches the new API and unblocks ADP token balancing. Ensure any other call sites are updated.

Run to verify no remaining int-based calls:

#!/bin/bash
# Find fetch_new_requests call sites that are not passing a list named active_requests
rg -nP '\.fetch_new_requests\((?!\s*self\.active_requests\s*\))' --type py -g '!**/test/**'

tests/unittest/_torch/executor/test_executor_request_queue.py (1)

642-644: Fix the outdated arithmetic comment to match the implementation.

Use the exact formula used in code: ceil-div via (x + tp_size - 1) // tp_size. Prior review flagged this.

-    # 2 + 1 + 3 + 0 = 6, 6 + 2 = 8, (8 + 3) // 4 = 2, max(2, 2, 1, 3, 0) = 3
-    # expected_num_active_requests = max((6 + 2 + 3) // 4, 3) = max(2, 3) = 3
+    # 2 + 1 + 3 + 0 = 6, 6 + 2 = 8; tp_size = 4
+    # expected_num_active_requests = max(((6 + 2) + 4 - 1) // 4, max(2, 1, 3, 0))
+    # = max(11 // 4 = 2, 3) = 3

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py (2)

1-2: Add NVIDIA copyright header.

Per repository guidelines, prepend the NVIDIA copyright header to all .py sources.

Apply at file top:

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

---

476-534: Heap-based balancing: edge-case safety and maintainability nits.

• If every rank is already at expected_num_active_requests, the filtered heap is empty; current logic only executes when new_requests is truthy, but an explicit early return after filtering would be clearer.
• Micro-nit: use clearer tuple fields or a dataclass to self-document the heap item.

Example early guard:

             all_ranks_new_requests_heap = [
                 val for val in all_ranks_new_requests_heap
                 if val.num_requests < self.expected_num_active_requests
             ]
+            if not all_ranks_new_requests_heap:
+                return all_ranks_new_requests

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

1-2: Add NVIDIA copyright header.

Please prepend the NVIDIA header to comply with coding guidelines.

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

tests/unittest/_torch/executor/test_executor_request_queue.py (3)

1-2: Add NVIDIA copyright header in tests.

Tests are also Python sources; add the required header.

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

---

661-665: Remove no-op statement in the test.

The standalone all_ranks_new_requests[0] has no effect; drop it to keep tests clean.

-        all_ranks_new_requests[0]

---

1003-1004: Avoid printing from tests.

Printing can add noise to CI logs. Prefer assertions only.

-    print("all_ranks_new_requests:", all_ranks_new_requests)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/unittest/_torch/executor/test_executor_request_queue.py (1)

642-644: Fix the explanatory comment to match the formula (duplicate of prior feedback)

The second line is correct now, but the first line still shows an extra “2” inside the max(...) tuple. Suggest tightening the two-line explanation to mirror the code’s ceil-div and max operations precisely.

Apply:

-    # 2 + 1 + 3 + 0 = 6, 6 + 2 = 8, (8 + 3) // 4 = 2, max(2, 2, 1, 3, 0) = 3
-    # expected_num_active_requests = max((6 + 2 + 3) // 4, 3) = max(2, 3) = 3
+    # total_active = 2 + 1 + 3 + 0 = 6
+    # num_new = 2, tp_size = 4 -> ceil((6 + 2) / 4) = (8 + 4 - 1) // 4 = 11 // 4 = 2
+    # max_active_per_rank = max(2, 1, 3, 0) = 3
+    # expected_num_active_requests = max(2, 3) = 3

tests/unittest/_torch/executor/test_executor_request_queue.py (2)

1004-1004: Remove debug print in tests

Avoid stdout noise in test runs; assertions already validate behavior.

Apply:

-    print("all_ranks_new_requests:", all_ranks_new_requests)

---

1100-1130: Capacity limit test is valuable, but it validates ‘expected’ capacity, not max per-rank capacity

This test ensures ranks don’t exceed expected_num_active_requests. Add a complementary test to ensure scheduling never exceeds max_num_active_requests when expected_num_active_requests > max_num_active_requests.

Proposed additional test:

def test_balance_requests_respects_max_capacity_even_if_expected_higher(attention_dp_queue):
    # Simulate expected > max per-rank capacity
    attention_dp_queue.max_num_active_requests = 2
    attention_dp_queue.expected_num_active_requests = 5  # deliberately higher
    all_ranks_num_active_requests = [2, 0, 0, 0]  # rank 0 already at max
    all_ranks_num_active_tokens = [0, 0, 0, 0]
    # 3 new relaxed requests (no target)
    reqs = []
    for i in range(3):
        r = RequestQueueItem(i, create_mock_request_with_py_schedule_params(attention_dp_rank=None))
        r.request.input_token_ids = [1, 2, 3]
        reqs.append(r)

    # Balance
    all_ranks_new_requests = {0: [], 1: [], 2: [], 3: []}
    result = attention_dp_queue._balance_requests_across_ranks(
        reqs, all_ranks_new_requests, all_ranks_num_active_requests, all_ranks_num_active_tokens
    )

    # Rank 0 must stay at capacity; others can receive at most 2 total per rank
    assert len(result[0]) == 0
    assert len(result[1]) <= 2
    assert len(result[2]) <= 2
    assert len(result[3]) <= 2

Note: This may currently fail if _balance_requests_across_ranks does not enforce max_num_active_requests. If so, production code should cap per-rank assignments using max_num_active_requests while balancing.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.