[None][perf] reduce rank-0 GIL contention in disaggregated generation (#15133)

Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

Author: lancelly

tensorrt_llm/_torch/distributed/communicator.py

@@ -308,18 +308,12 @@ def safe_broadcast(comm, obj, root=0, chunk_size: int = 4 * 1024 * 1024):
         offset += cur
 
     # ---- Reconstruction and deserialization ----
-    # Validate the received byte count and unpickle.
     if rank == root:
-        # Root already has `serialized`
-        if len(serialized) != total_size:
-            raise RuntimeError(
-                f"Data size mismatch at root: expected {total_size}, got {len(serialized)}"
-            )
-        try:
-            return pickle.loads(serialized)  # nosec B301
-        except Exception as e:
-            raise RuntimeError(f"Deserialization failed: {str(e)}") from e
+        # Root already holds `obj`; rebuilding it from its own serialized bytes
+        # would be a needless deep copy.
+        return obj
     else:
+        # Validate the received byte count and unpickle.
         if len(dst_buf) != total_size:
             raise RuntimeError(
                 f"Data size mismatch at rank {rank}: expected {total_size}, got {len(dst_buf)}"

tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

@@ -369,17 +369,17 @@ def _balance_requests_across_ranks(
 
         heapq.heapify(all_ranks_new_requests_heap)
 
-        new_requests = sorted(
-            new_requests,
-            key=lambda x: len(getattr(x.request, "input_token_ids", [])) if x.request else 0,
-            reverse=True,
-        )
+        # input_token_ids is a C++ getter that copies the whole token list on every
+        # access, so read it once per request instead of in both the sort key and
+        # the loop below.
+        counted = [
+            (len(req.input_token_ids) if (req := item.request) is not None else 0, item)
+            for item in new_requests
+        ]
+        counted.sort(key=lambda item: item[0], reverse=True)
 
-        for req_item in new_requests:
+        for token_count, req_item in counted:
             val = heapq.heappop(all_ranks_new_requests_heap)
-            token_count = (
-                len(getattr(req_item.request, "input_token_ids", [])) if req_item.request else 0
-            )
             val = val._replace(
                 num_tokens=val.num_tokens + token_count,
                 num_requests=val.num_requests + 1,