[None][perf] Serving runtime: disagg/KV-transfer/ADP fixes, KV cache manager V2 prefix wait, ConversationRouter prefix placement, KV-balanced ADP admission

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/copyUtils.cu

@@ -31,6 +31,8 @@ using Grain = uint4;
 constexpr uint32_t ctaSize = 128;
 constexpr uint32_t nbBufs = 4;
 constexpr uint32_t grainBytes = sizeof(Grain);
+constexpr uint32_t lowBandwidthBytesPerCta = 8u << 20;
+constexpr uint32_t lowBandwidthMaxSplits = 16;
 
 using MMTask = Task<MemAddress, MemAddress>;
 
@@ -114,7 +116,8 @@ CUresult launchBatchedCopyImpl(
     {
         pTasks = reinterpret_cast<std::array<MMTask, N> const*>(tasks);
     }
-    uint32_t const nbSplits = lowBandwidth ? 1 : divUp(nbBytes, grainBytes * ctaSize * 2);
+    uint32_t const nbSplits = lowBandwidth ? std::min(divUp(nbBytes, lowBandwidthBytesPerCta), lowBandwidthMaxSplits)
+                                           : divUp(nbBytes, grainBytes * ctaSize * 2);
     void* args[] = {(void*) pTasks, (void*) &nbBytes};
     static CUkernel const kernel = []() -> CUkernel
     {
@@ -129,8 +132,9 @@ CUresult launchBatchedCopyImpl(
         stream, args, nullptr);
 }
 
-// When bandwidth is low, e.g. when host memory is involved, we avoid splitting as fewer CTAs should be enough to
-// saturate the bandwidth.
+// When bandwidth is low, e.g. when host memory is involved, few CTAs are enough to saturate it, so a task is split
+// over at most lowBandwidthMaxSplits CTAs of lowBandwidthBytesPerCta each. A single CTA per task would leave large
+// pages (e.g. a recurrent-state page holding every linear-attention layer) bound by one SM's bandwidth.
 CUresult launchBatchedCopy(bool lowBandwidth, std::vector<MMTask> const& tasks, uint32_t nbBytes, cudaStream_t stream)
 {
     constexpr uint32_t maxN = 256;

tensorrt_llm/_torch/disaggregation/native/transfer.py

@@ -130,6 +130,9 @@ class RecvReqInfo:
     aux_slot: Optional[int] = None
     slice_id: Optional[int] = None
     bounce_dst_base: Optional[int] = None
+    # Context-first receivers send REQUEST_DATA only after the context phase
+    # produced the KV cache, so a TxSession must already exist (or have ended).
+    context_first: bool = False
 
     def to_bytes(self) -> bytes:
         return msgpack.packb(
@@ -145,6 +148,7 @@ def to_bytes(self) -> bytes:
                 "aux_slot": self.aux_slot,
                 "slice_id": self.slice_id,
                 "bounce_dst_base": self.bounce_dst_base,
+                "context_first": self.context_first,
             }
         )
 
@@ -643,6 +647,11 @@ def sweep_stale_req_infos(self):
         Called opportunistically from the listener thread when a new REQUEST_DATA
         arrives. With gen-first ADP broadcast, non-assigned DP ranks accumulate
         entries that are never consumed; this sweep prevents unbounded growth.
+
+        A context-first entry that no TxSession claimed within the TTL never will
+        be (the session already ended, or the receiver names a request this rank
+        never had), so its receiver is told that the transfer failed instead of
+        waiting for data that will not arrive.
         """
         now = time.monotonic()
         with self._peer_requests_lock:
@@ -653,12 +662,32 @@ def sweep_stale_req_infos(self):
             ]
         if not stale_rids:
             return
+        orphaned_infos: list[RecvReqInfo] = []
         for rid in stale_rids:
             with self._sessions_lock, self._peer_requests_lock:
                 if rid not in self._sessions and rid in self._peer_requests:
-                    self._peer_requests.pop(rid, None)
+                    req_info_map = self._peer_requests.pop(rid)
                     self._peer_requests_timestamps.pop(rid, None)
+                    orphaned_infos.extend(
+                        info for info in req_info_map.values() if info.context_first
+                    )
                     logger.debug(f"Swept stale RecvReqInfo for rid={rid}")
+        for info in orphaned_infos:
+            self._fail_orphaned_receiver(info)
+
+    def _fail_orphaned_receiver(self, info: RecvReqInfo) -> None:
+        """Report FAILED for a context-first REQUEST_DATA that no TxSession claimed.
+
+        This resolves the receiver's task, so its RxSession fails through the
+        normal transfer-failure path. The sweep runs on the executor and listener
+        threads, so the message goes through a worker queue.
+        """
+        logger.warning(
+            f"KV transfer request rid={info.unique_rid} from receiver "
+            f"{info.instance_name} rank {info.instance_rank} has no TxSession on this "
+            f"rank after {self._STALE_REQ_INFO_TTL_S:.0f}s; failing it"
+        )
+        self._send_failed_result_to_receiver(info, defer_to_worker=True)
 
     def setup_session(self, tx_session: "TxSession"):
         unique_rid = tx_session.disagg_request_id
@@ -2754,6 +2783,7 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo:
             dst_start_token=None,
             aux_slot=task._aux_slot,
             slice_id=task.slice_id,
+            context_first=task._params.schedule_style != DisaggScheduleStyle.GENERATION_FIRST,
         )
 
     @staticmethod

tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py

@@ -13,8 +13,6 @@
 
 import torch
 
-from tensorrt_llm._utils import prefer_pinned
-
 if sys.version_info[:2] >= (3, 12):
     from typing import override
 else:
@@ -26,6 +24,7 @@
     MambaRole,
     MambaStateLayout,
     _get_local_mamba_cache_layout,
+    _host_slots_to_device,
     _mamba_effective_tp_size,
     _promote_intermediate_states,
     _stack_state_views,
@@ -116,6 +115,15 @@ def commit_conv_window(self, slot_indices: torch.Tensor, conv_pool: torch.Tensor
 _KDA_REPLAY_COUNT_BYTES = 16
 _KDA_REPLAY_COUNT_ROLE = DataRole("prev_num_accepted_tokens")
 _KDA_REPLAY_CONV_ROLES = ("kda_conv_q", "kda_conv_k", "kda_conv_v")
+# Replay buffers a transferred slot is seeded or cleared in. The drafts'
+# records (_KDA_DRAFT_RECORDS_ROLE) need no reset.
+_KDA_SEEDED_ROLES = (
+    *_KDA_REPLAY_CONV_ROLES,
+    "kda_k_cache",
+    "kda_g_cache",
+    "kda_v_cache",
+    "kda_beta_cache",
+)
 _KDA_DRAFT_RECORDS_ROLE = "kda_state_tok"
 
 
@@ -173,16 +181,6 @@ def _kda_replay_bytes_per_slot(
     return num_local_layers * layer_bytes + _KDA_REPLAY_COUNT_BYTES
 
 
-def _host_slots_to_device(slots: list[int], device: torch.device) -> torch.Tensor:
-    """Copy host slot IDs without synchronizing queued GPU work.
-
-    Each copy uses fresh pinned storage so later host writes cannot race it.
-    """
-    return torch.tensor(
-        slots, dtype=torch.long, pin_memory=device.type == "cuda" and prefer_pinned()
-    ).to(device, non_blocking=True)
-
-
 def _kda_beta_cache_padded_heads(num_heads: int) -> int:
     heads_per_alignment = _KDA_BETA_CACHE_ALIGNMENT_BYTES // torch.bfloat16.itemsize
     return (num_heads + heads_per_alignment - 1) // heads_per_alignment * heads_per_alignment
@@ -425,6 +423,11 @@ def __init__(
         self._seed_request_counter = 0
         self._seed_rank_offset = 0
         self._conv_section_dims: tuple[int, ...] = ()
+        # ``[num_layers, num_slots, ...]`` views of the conv states and of the
+        # seeded replay buffers when the per-layer views are affine (see
+        # _stack_state_views); None otherwise.
+        self._stacked_seed_conv: torch.Tensor | None = None
+        self._stacked_seed_replay: dict[str, torch.Tensor] | None = None
 
     def validate(self, context: MambaStateLayout) -> None:
         if context.spec_config is None:
@@ -597,8 +600,25 @@ def bind(
                 self._seed_rank_offset,
                 device,
             )
+        self._stack_seed_views()
         logger.info(f"Mamba Cache (kda-replay) is allocated for {cache_size} state slots")
 
+    def _stack_seed_views(self) -> None:
+        """Stack what seed_transferred_slots writes, so it runs on all layers at once."""
+        self._stacked_seed_conv = _stack_state_views(list(self._conv_states))
+        self._stacked_seed_replay = None
+        if self._stacked_seed_conv is None:
+            return
+        layers = [self._layer_cache_fields(layer) for layer in range(len(self._conv_states))]
+        stacked = {}
+        for name in _KDA_SEEDED_ROLES:
+            if layers[0].get(name) is None:
+                continue
+            stacked[name] = _stack_state_views([layer[name] for layer in layers])
+            if stacked[name] is None:
+                return
+        self._stacked_seed_replay = stacked
+
     def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]:
         if self.pool_layers is not None:
             return {
@@ -702,15 +722,24 @@ def seed_transferred_slots(self, state_indices: torch.Tensor) -> None:
             section_offsets.append(section_offsets[-1] + section_dim)
 
         self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0)
+        stacked_replay = self._stacked_seed_replay
+        if stacked_replay is not None:
+            # All layers at once: one launch per replay buffer and section.
+            for buffer in stacked_replay.values():
+                buffer.index_fill_(1, state_indices, 0)
+            selected_conv = self._stacked_seed_conv.index_select(1, state_indices)
+            for section_idx, name in enumerate(_KDA_REPLAY_CONV_ROLES):
+                buffer = stacked_replay[name]
+                section = selected_conv[
+                    :, :, section_offsets[section_idx] : section_offsets[section_idx + 1], :
+                ]
+                buffer[..., :committed_window].index_copy_(
+                    1, state_indices, section.to(buffer.dtype)
+                )
+            return
         for layer_offset, conv_state in enumerate(self._conv_states):
             replay = self._layer_cache_fields(layer_offset)
-            for name in (
-                *_KDA_REPLAY_CONV_ROLES,
-                "kda_k_cache",
-                "kda_g_cache",
-                "kda_v_cache",
-                "kda_beta_cache",
-            ):
+            for name in _KDA_SEEDED_ROLES:
                 buffer = replay.get(name)
                 if buffer is not None:
                     buffer.index_fill_(0, state_indices, 0)
@@ -775,7 +804,7 @@ def commit_pending_drafts(self, slot: int) -> None:
                 num_accepted,
                 state_tok=replay.get(_KDA_DRAFT_RECORDS_ROLE),
             )
-        self.prev_num_accepted_tokens[slot] = 0
+        self.prev_num_accepted_tokens[slot].zero_()
 
     def _layer_replay_buffers(self) -> tuple[torch.Tensor, ...]:
         buffers = (
@@ -796,6 +825,8 @@ def shutdown(self) -> None:
         self.pool_layers = None
         self._ssm_states = ()
         self._conv_states = ()
+        self._stacked_seed_conv = None
+        self._stacked_seed_replay = None
         self.suspended_requests.clear()
         self.acceptance_pending.clear()
         self.unseeded_requests.clear()

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -1795,6 +1795,10 @@ class KVCacheManagerV2(BaseResourceManager):
     # the class for instances built without running __init__.
     _commit_branch_snapshots = False
     _record_branch_snapshots = False
+    # Request id -> points of its prompt that requests waiting for its prefix
+    # need snapshotted (request_prefix_snapshot); kept apart from the request's
+    # own expect_snapshot_points, which other code rewrites. Created on first use.
+    _waiter_snapshot_points: Optional[Dict[int, set[int]]] = None
 
     def __init__(
         self,
@@ -2513,6 +2517,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
             and not self.enable_joint_kv_cache_reuse
             and not self._has_cp_helix
         )
+        self.can_share_inflight_prefix = self._can_share_inflight_prefix()
 
         # With pipeline parallelism, multiple microbatches can be in-flight
         # simultaneously, so we need slots for all concurrent sequences.
@@ -4942,18 +4947,97 @@ def _record_branch_snapshot_point(
         point = _swa_branch_snapshot_point(
             divergence, num_lookup_tokens, reused, req.prompt_len, self.tokens_per_block
         )
-        if point is None:
-            # Drop a point left by an earlier lookup of this request.
-            req.expect_snapshot_points = []
-            return
-        # Only recurrent-state managers rewrite this list each scheduler pass, so
-        # the point persists here until the request finishes.
-        req.expect_snapshot_points = [point]
-        logger.debug(
-            f"[KVCacheManagerV2] request {req.py_request_id} branch snapshot at {point} "
-            f"(divergence={divergence}, reused={reused}, prompt_len={req.prompt_len})"
+        # A fresh lookup replaces the point of an earlier lookup of this request.
+        # Points requested by waiting requests (request_prefix_snapshot) may
+        # already be in the list, so they are merged back in.
+        req.expect_snapshot_points = [] if point is None else [point]
+        self._merge_waiter_snapshot_points(req)
+        if point is not None:
+            logger.debug(
+                f"[KVCacheManagerV2] request {req.py_request_id} branch snapshot at {point} "
+                f"(divergence={divergence}, reused={reused}, prompt_len={req.prompt_len})"
+            )
+
+    def _can_share_inflight_prefix(self) -> bool:
+        """Whether a request can wait for another to make a shared prefix reusable.
+
+        Without SWA layers whose window reaches the radix tree only at a context
+        chunk end, every block a request computes is committed, at the latest
+        when its context ends. Otherwise the request computing the prefix must
+        end a chunk at the shared depth, which needs SWA branch snapshots.
+        Decided at construction; see :meth:`request_prefix_snapshot`.
+        """
+        if not (self.enable_block_reuse and self._can_publish_block_reuse):
+            return False
+        if self._record_branch_snapshots:
+            return True
+        # Only the SWA window at a context chunk end reaches the radix tree: SWA
+        # scratch reuse gives no other window real pages, and a policy other
+        # than all_reusable commits the window at a prompt end only.
+        swa_window_only_at_chunk_end = (
+            self.enable_swa_scratch_reuse
+            or self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE
+        )
+        return not (
+            swa_window_only_at_chunk_end
+            and any(
+                getattr(layer, "sliding_window_size", None)
+                for layer in self.impl.init_config.layers
+            )
         )
 
+    def request_prefix_snapshot(self, owner: LlmRequest, num_tokens: int) -> bool:
+        """Make the first ``num_tokens`` of ``owner``'s prompt reusable once computed.
+
+        ``owner`` is a context request that has not finished its context; a
+        request sharing that prefix waits for it (``PrefixWaitPlanner``). Only
+        called when :attr:`can_share_inflight_prefix`. Where only the SWA window
+        at a context chunk end reaches the radix tree (SWA branch snapshots, see
+        :meth:`_record_branch_snapshot_point`), ``owner`` ends a chunk at the
+        block-aligned depth. Returns False when that is no longer possible.
+        """
+        if not self._record_branch_snapshots:
+            return True
+        return self._add_waiter_snapshot_point(owner, num_tokens)
+
+    def _add_waiter_snapshot_point(self, owner: LlmRequest, point: int) -> bool:
+        """Add the block-aligned ``point`` to ``owner``'s snapshot points.
+
+        The point is also kept apart from the request, so that whatever rewrites
+        ``expect_snapshot_points`` later (a fresh lookup's branch point, or a
+        recurrent-state manager's per-pass placement) merges it back in with
+        :meth:`_merge_waiter_snapshot_points`. Returns False when the point is
+        not inside ``owner``'s prompt or not past the context ``owner`` has
+        computed or holds pages for: a chunk already allocated (an in-flight
+        chunk, or one grown under SWA scratch reuse whose forward was skipped)
+        cannot end at the point any more.
+        """
+        point = point // self.tokens_per_block * self.tokens_per_block
+        allocated_end = owner.context_current_position
+        kv_cache = self.kv_cache_map.get(owner.py_request_id)
+        if kv_cache is not None:
+            allocated_end = max(allocated_end, kv_cache.capacity - self.num_extra_kv_tokens)
+        if not allocated_end < point < owner.prompt_len:
+            return False
+        if self._waiter_snapshot_points is None:
+            self._waiter_snapshot_points = {}
+        self._waiter_snapshot_points.setdefault(owner.py_request_id, set()).add(point)
+        self._merge_waiter_snapshot_points(owner)
+        return True
+
+    def _merge_waiter_snapshot_points(self, request: LlmRequest) -> None:
+        """Merge into ``request``'s snapshot points the waiters' points it has not passed yet."""
+        points_by_id = self._waiter_snapshot_points
+        points = points_by_id.get(request.py_request_id) if points_by_id else None
+        if not points:
+            return
+        start = request.context_current_position
+        points.difference_update([point for point in points if point <= start])
+        if not points:
+            del points_by_id[request.py_request_id]
+            return
+        request.expect_snapshot_points = sorted({*request.expect_snapshot_points, *points})
+
     def probe_context_reuse(self, req: LlmRequest) -> int | None:
         """Reusable prefix depth for this pool, without claiming pages.
 
@@ -7029,6 +7113,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
             self.kv_connector_manager.reset_request_state(request)
         if self.conversation_manager is not None:
             self.conversation_manager.finish_request(request)
+        if self._waiter_snapshot_points is not None:
+            self._waiter_snapshot_points.pop(request.py_request_id, None)
         self._allocated_draft_lens.pop(request.py_request_id, None)
         self._inflight_draft_lens.pop(request.py_request_id, None)
         self._request_stats_enabled_ids.discard(request.py_request_id)
@@ -7371,6 +7457,8 @@ def shutdown(self):
         # its plan, which mutates manager state.
         if self.conversation_manager is not None:
             self.conversation_manager.clear()
+        if self._waiter_snapshot_points is not None:
+            self._waiter_snapshot_points.clear()
         self.impl.shutdown()
         if self._sharing is not None:
             # After the runtime's shutdown, which raises while any cache is open, in the map or not;