[None][perf] DSv4 follow-up: disagg routing improvements

cpp/include/tensorrt_llm/executor/executor.h

@@ -758,6 +758,7 @@ class Request
     ~Request();
 
     [[nodiscard]] VecTokens getInputTokenIds() const;
+    [[nodiscard]] SizeType32 getNumInputTokens() const;
     [[nodiscard]] SizeType32 getMaxTokens() const;
     [[nodiscard]] bool getStreaming() const;
     [[nodiscard]] SamplingConfig getSamplingConfig() const;

cpp/tensorrt_llm/executor/request.cpp

@@ -79,6 +79,11 @@ VecTokens Request::getInputTokenIds() const
     return mImpl->getInputTokenIds();
 }
 
+SizeType32 Request::getNumInputTokens() const
+{
+    return mImpl->getNumInputTokens();
+}
+
 SizeType32 Request::getMaxTokens() const
 {
     return mImpl->getMaxNewTokens();

cpp/tensorrt_llm/executor/requestImpl.h

@@ -135,6 +135,11 @@ class Request::Impl
         return mInputTokenIds;
     }
 
+    [[nodiscard]] SizeType32 getNumInputTokens() const
+    {
+        return static_cast<SizeType32>(mInputTokenIds.size());
+    }
+
     [[nodiscard]] SizeType32 getMaxNewTokens() const
     {
         return mMaxNewTokens;

cpp/tensorrt_llm/nanobind/executor/request.cpp

@@ -728,6 +728,7 @@ void initRequestBindings(nb::module_& m)
         nb::arg("cache_salt") = nb::none()
     )                // clang-format on
         .def_prop_ro("input_token_ids", &tle::Request::getInputTokenIds)
+        .def_prop_ro("num_input_tokens", &tle::Request::getNumInputTokens)
         .def_prop_ro("max_tokens", &tle::Request::getMaxTokens)
         .def_prop_rw("streaming", &tle::Request::getStreaming, &tle::Request::setStreaming)
         .def_prop_rw("sampling_config", &tle::Request::getSamplingConfig, &tle::Request::setSamplingConfig)

tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py

@@ -335,6 +335,25 @@ def _mismatch(self, field: str, local, peer) -> bool:
             local != peer, f"{field} mismatch", field=field, local=local, peer=peer
         )
 
+    def _tpb_check(self, local: int, peer: int) -> bool:
+        if local == peer:
+            return False
+        larger, smaller = max(local, peer), min(local, peer)
+        if larger % smaller != 0:
+            logger.warning(
+                "AttentionPolicy: incompatible: tokens_per_block not divisible; local=%d peer=%d",
+                local,
+                peer,
+            )
+            return True
+        logger.warning(
+            "AttentionPolicy: tokens_per_block mismatch (local=%d, peer=%d); "
+            "KV transfer proceeds — ensure block boundaries align during transfer.",
+            local,
+            peer,
+        )
+        return False
+
     def check_peer_compatible(self, peer_ri: RankInfo) -> bool:
         a = self._ri.attention
         b = peer_ri.attention
@@ -348,7 +367,7 @@ def check_peer_compatible(self, peer_ri: RankInfo) -> bool:
                 peer=peer_ri.cp_size,
             )
             or self._mismatch("element_bytes", a.element_bytes, b.element_bytes)
-            or self._mismatch("tokens_per_block", a.tokens_per_block, b.tokens_per_block)
+            or self._tpb_check(a.tokens_per_block, b.tokens_per_block)
             or self._mismatch("dims_per_head", a.dims_per_head, b.dims_per_head)
             or self._fail_if(
                 a.is_mla and (a.kv_heads_per_rank != 1 or b.kv_heads_per_rank != 1),

tensorrt_llm/_torch/distributed/communicator.py

@@ -312,18 +312,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/_util.py

@@ -455,6 +455,7 @@ def _create_dummy_context_requests(
                         torch.zeros(1, 1, dtype=torch.int32)
                     }
                 }
+            request.py_conversation_params = None
             requests.append(request)
             remaining_tokens -= input_seq_len
         if self._mapping.enable_attention_dp:

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -505,6 +505,9 @@ def __init__(
         from ..speculative import get_num_extra_kv_tokens
 
         self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config)
+        # Mirror V1: expose max_draft_len so the native disagg AuxBuffer
+        # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding.
+        self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0
         self.max_total_draft_tokens = (
             spec_config.max_total_draft_tokens if spec_config is not None else 0
         )
@@ -2515,6 +2518,20 @@ def copy_batch_block_offsets(
             self._stream.cuda_stream,
         )
 
+    @staticmethod
+    def _derive_reuse_salt(cache_salt: str | None) -> int | None:
+        """Derive ``ReuseScope.salt`` (int|None) from the ``cache_salt`` string.
+
+        Deterministic so the same string yields the same reuse namespace across
+        processes (matches C++ blockKey hashing on cacheSalt). Shared by cache
+        creation, prefetch, and the cache-aware router probe so they all hit the
+        same radix-tree namespace.
+        """
+        if cache_salt is None:
+            return None
+        digest = hashlib.sha256(cache_salt.encode("utf-8")).digest()
+        return int.from_bytes(digest[:8], "little")
+
     def _create_kv_cache(
         self,
         request_id: int,
@@ -2538,14 +2555,7 @@ def _create_kv_cache(
                 self.index_mapper.size(),
             )
             return None
-        # ReuseScope.salt is int|None; derive a deterministic int from the
-        # cache_salt string so the same string yields the same reuse namespace
-        # across processes (matches C++ blockKey hashing on cacheSalt).
-        salt_int = (
-            int.from_bytes(hashlib.sha256(cache_salt.encode("utf-8")).digest()[:8], "little")
-            if cache_salt is not None
-            else None
-        )
+        salt_int = self._derive_reuse_salt(cache_salt)
         kv_cache = self.impl.create_kv_cache(
             ReuseScope(lora_id=lora_task_id, salt=salt_int),
             input_tokens,
@@ -2565,6 +2575,28 @@ def _create_kv_cache(
                 kv_cache.set_base_page_index_buf(i, pool_idx, memoryview(buffer.numpy()))
         return kv_cache
 
+    def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt=None):
+        """Probe the v2 KV cache radix tree for prefix match length.
+
+        Returns the number of prefix tokens already cached on this rank,
+        without acquiring page ownership. Mirrors the v1 ``KVCacheManager``
+        adapter so ``KVCacheAwareADPRouter`` works on both KV cache backends.
+
+        ``cache_salt`` (and ``lora_task_id``) must match the values used by
+        ``_create_kv_cache``; the salt is derived from the ``cache_salt``
+        string the same way, so the probe queries the same reuse namespace.
+        Otherwise the router would see an incorrect match length.
+        """
+        if not self.enable_block_reuse:
+            return 0
+        if not input_tokens:
+            return 0
+        salt_int = self._derive_reuse_salt(cache_salt)
+        return self.impl.probe_reuse(
+            ReuseScope(lora_id=lora_task_id, salt=salt_int),
+            input_tokens,
+        )
+
     def prefetch_for_context_tokens(self, requests: list) -> bool:
         """Prefetch radix-tree blocks from disk→host for upcoming context requests.
 
@@ -2580,14 +2612,9 @@ def prefetch_for_context_tokens(self, requests: list) -> bool:
         for req in requests:
             all_tokens = req.get_tokens(DEFAULT_BEAM_INDEX)
             tokens = self._augment_tokens_for_block_reuse(all_tokens, req, end=len(all_tokens) - 1)
-            # Match the ReuseScope salt derivation used in _create_kv_cache so
-            # the transient cache hits the same radix-tree blocks.
-            cache_salt = req.cache_salt
-            salt_int = (
-                int.from_bytes(hashlib.sha256(cache_salt.encode("utf-8")).digest()[:8], "little")
-                if cache_salt is not None
-                else None
-            )
+            # Use the same salt derivation as _create_kv_cache so the transient
+            # cache hits the same radix-tree blocks.
+            salt_int = self._derive_reuse_salt(req.cache_salt)
             kv_cache = self.impl.create_kv_cache(
                 ReuseScope(lora_id=req.lora_task_id, salt=salt_int), tokens
             )

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -801,6 +801,7 @@ def __init__(
         self.py_logprobs_mode = LogprobMode(
             logprobs_mode)  # handle passed a raw string
         self.py_disaggregated_params = None
+        self.py_conversation_params = None
 
         self.py_num_connector_matched_tokens = 0
 
@@ -1191,6 +1192,8 @@ def executor_request_to_llm_request(
     llm_request.py_disaggregated_params = getattr(executor_request,
                                                   "py_disaggregated_params",
                                                   None)
+    llm_request.py_conversation_params = getattr(executor_request,
+                                                 "py_conversation_params", None)
     if child_req_ids:
         for child_id in child_req_ids:
             llm_request.create_child_request(child_id)

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -558,9 +558,11 @@ def __init__(
         self._prefetched_request_ids: set[int] = set()
         self.enable_kv_cache_events = self.kv_cache_manager is not None and self.kv_cache_manager.event_buffer_max_size > 0
         self.enable_kv_cache_reuse = self.kv_cache_manager is not None and self.kv_cache_manager.enable_block_reuse
+        # AsyncTransferManager pin/unpin path is V1-only; V2 holds blocks via _KVCache refcount.
         self.enable_partial_reuse_for_disagg = (
             self.enable_kv_cache_reuse
-            and self.kv_cache_manager.enable_partial_reuse)
+            and self.kv_cache_manager.enable_partial_reuse
+            and not isinstance(self.kv_cache_manager, KVCacheManagerV2))
 
         self.max_input_len = max_input_len
         # _executor_loop private data
@@ -596,6 +598,12 @@ def __init__(
         # responses and flushing them at a synchronised point in the executor
         # loop avoids the mismatch.
         self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = []
+        # Same buffer-then-synced-flush pattern as _pending_transfer_responses
+        # above: _handle_responses and _append_iter_stats are reached from
+        # per-rank-divergent gates, so their tp_allgather collectives are
+        # lifted to _handle_kv_transfer_timeouts_synced / _flush_iter_stats_synced.
+        self._pending_timed_out_requests: List[LlmRequest] = []
+        self._pending_iter_stats_dict: Optional[Dict] = None
         # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen;
         # updated from observed request types.
         self._adp_dummy_is_gen: bool = True
@@ -983,6 +991,53 @@ def _flush_pending_transfer_responses(self):
             # can complete.
             self._enqueue_responses(responses)
 
+    def _handle_kv_transfer_timeouts_synced(self):
+        """ADP-safe drain of the KV-transfer-timeout consensus collective.
+
+        Lifted out of _handle_responses, which is reached from per-rank-
+        divergent gates (_process_previous_batch, _handle_executed_batch).
+        Non-ADP runs handle timeouts inline; the buffer is empty here.
+        """
+        if not (self.enable_attention_dp and self.dist.world_size != 1):
+            return
+        timed_out = self._pending_timed_out_requests
+        self._pending_timed_out_requests = []
+        any_timed_out = any(self.dist.tp_allgather(bool(timed_out)))
+        if any_timed_out:
+            self._handle_errors(error_msg="Request timed out (KV transfer)",
+                                requests=timed_out,
+                                charge_budget=False)
+
+    def _flush_iter_stats_synced(self):
+        """ADP-safe drain of the TLLM_METRICS_ALL_RANKS dict-gather collective.
+
+        Lifted out of _append_iter_stats (same divergent-gate problem as
+        _handle_kv_transfer_timeouts_synced).  Under ADP every rank sends
+        either the local dict or ``None`` every iter so the gather stays in
+        lockstep with peer collectives.  When the gate doesn't hold,
+        _append_iter_stats takes its legacy path and never populates the
+        buffer.
+        """
+        tp_size = self.dist.tp_size
+        gather_all_ranks = os.environ.get("TLLM_METRICS_ALL_RANKS", "0") == "1"
+        if not (gather_all_ranks and self.enable_iter_perf_stats and tp_size > 1
+                and self.enable_attention_dp):
+            return
+        local_dict = self._pending_iter_stats_dict
+        self._pending_iter_stats_dict = None
+        gathered = self.dist.tp_allgather(local_dict)
+        if self.dist.tp_rank != 0:
+            return
+        rank_dicts = [d for d in gathered if d is not None]
+        if not rank_dicts:
+            return
+        with self.stats_lock:
+            for d in rank_dicts:
+                self.stats.append(("per_rank_dict", d))
+            cap = self.max_stats_len * tp_size
+            if len(self.stats) > cap:
+                del self.stats[:len(self.stats) - cap]
+
     # Performance metrics methods are in PerfMetricsManager (self.perf_manager)
 
     def _event_loop_wrapper(self):
@@ -1905,23 +1960,9 @@ def _append_iter_stats(self,
                 local_dict["gpuForwardTimeMS"] = gpu_forward_time_ms
             local_dict["schedulerMode"] = scheduler_mode
             local_dict["rank"] = self.dist.tp_rank
-
-            gathered = self.dist.tp_allgather(local_dict)
-
-            if self.dist.tp_rank == 0:
-                with self.stats_lock:
-                    # Wrap as ("per_rank_dict", dict) so the serializer can
-                    # distinguish from the legacy (stats, req_stats, kv) tuple.
-                    # Trim before appending so every rank's entry for the
-                    # same iteration lands in the buffer atomically; evicting
-                    # during the append loop would let partial iterations
-                    # survive at the head of self.stats.
-                    cap = self.max_stats_len * tp_size
-                    overflow = max(0, len(self.stats) + len(gathered) - cap)
-                    if overflow:
-                        del self.stats[:overflow]
-                    for d in gathered:
-                        self.stats.append(("per_rank_dict", d))
+            # Buffer for _flush_iter_stats_synced; in-place tp_allgather would
+            # desync (reached from per-rank-divergent gates).
+            self._pending_iter_stats_dict = local_dict
             return
 
         # Legacy path: rank-0-only (single-rank or iter stats disabled).
@@ -2580,6 +2621,13 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]):
                 executed_batch.microbatch_id % self.dist.pp_size,
             )
 
+        # Drain the synced-collective buffers outside the per-rank-divergent
+        # ``executed_batch is not None`` branch.  handle_executed_batches
+        # invokes this helper the same number of times on every rank
+        # (broadcast count), so the flushes are rank-symmetric.
+        self._handle_kv_transfer_timeouts_synced()
+        self._flush_iter_stats_synced()
+
     @nvtx_range("wait_on_pp_send_handles")
     def wait_on_pp_send_handles(self, send_handles, microbatch_id):
         if send_handles[microbatch_id] is not None:
@@ -3311,6 +3359,10 @@ def _executor_loop(self):
                     if self.enable_kv_cache_events:
                         self._add_kv_cache_events()
 
+                # Drain timeout buffer outside ``if can_queue`` so the synced
+                # collective fires every iter regardless of future restructuring.
+                self._handle_kv_transfer_timeouts_synced()
+
                 if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests(
                 ):
                     self._check_kv_transfer_timeout()
@@ -3332,6 +3384,9 @@ def _executor_loop(self):
                 elif gpu_forward_events_from_perf_pool:
                     self.perf_manager.release_forward_timing_events(
                         gpu_forward_start, gpu_forward_end)
+                # Same lockstep guarantee for iter-stats; no-op when
+                # TLLM_METRICS_ALL_RANKS=0.
+                self._flush_iter_stats_synced()
 
                 self.iter_counter += 1
 
@@ -3758,6 +3813,12 @@ def _executor_loop_overlap(self):
                 else:
                     self._enqueue_responses([])
 
+                # Drain buffers from the (per-rank-divergent)
+                # _process_previous_batch above; rank-symmetric companion to
+                # the _enqueue_responses([]) call in the else branch.
+                self._handle_kv_transfer_timeouts_synced()
+                self._flush_iter_stats_synced()
+
                 # Call set_exclude_last_generation_logits after _process_previous_batch.
                 # If set before, the response of a request may be incorrect, as it will
                 # use the wrong indices for generation logits when streaming is enabled.
@@ -4526,8 +4587,9 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None:
             elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000
             if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out:
                 logger.warning(
-                    f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout"
-                )
+                    f"Terminating {type} request {req.py_request_id} due to KV "
+                    f"cache transfer timeout: elapsed {elapsed_time:.0f}ms > "
+                    f"kv_transfer_timeout_ms={timeout_ms}ms")
                 req.py_kv_transfer_timed_out = True
 
         for req in self.async_transfer_manager.requests_in_transfer().values():
@@ -5679,12 +5741,9 @@ def _handle_responses(self, emit_first_iter: bool = True):
             self._terminate_request(request)
         if (self.kv_cache_transceiver is not None and self.enable_attention_dp
                 and self.dist.world_size != 1):
-            any_timed_out = any(self.dist.tp_allgather(
-                bool(timed_out_requests)))
-            if any_timed_out:
-                self._handle_errors(error_msg="Request timed out (KV transfer)",
-                                    requests=timed_out_requests,
-                                    charge_budget=False)
+            # Buffer for _handle_kv_transfer_timeouts_synced; in-place
+            # tp_allgather would desync (reached from per-rank-divergent gates).
+            self._pending_timed_out_requests.extend(timed_out_requests)
         else:
             for req in timed_out_requests:
                 self._handle_errors(