[None][fix] Derive KV capacity once and select windowed blocks by ordinal

tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py

@@ -74,6 +74,20 @@ def get_block_ids(
         must translate before returning.
         """
 
+    #: V1 hands back a pre-eviction chain the legacy trim already understands.
+    exposes_block_ordinals: bool = False
+
+    @abstractmethod
+    def get_block_ordinals(
+        self,
+        req: LlmRequest,
+        group_idx: int,
+        lg: AttentionLayerGroup,
+    ) -> np.ndarray:
+        """Entry ``i`` is the slot for tokens ``[i * tokens_per_block, ...)``,
+        ``-1`` where no page is bound. Position must not come from a length.
+        """
+
     @abstractmethod
     def commit_blocks_for_reuse(self, req: LlmRequest) -> None:
         """Commit KV blocks to radix tree for future prefix reuse.
@@ -123,6 +137,10 @@ def get_block_ids(self, req, group_idx, lg):  # noqa: ARG002
         )
         return np.asarray(pool_indices, dtype=np.int64)
 
+    def get_block_ordinals(self, req, group_idx, lg):
+        # V1 exposes the whole chain in order, so ids are already positional.
+        return self.get_block_ids(req, group_idx, lg)
+
     def commit_blocks_for_reuse(self, req: LlmRequest) -> None:
         if not self.enable_block_reuse:
             return
@@ -165,6 +183,16 @@ def get_block_ids(self, req, group_idx, lg):  # noqa: ARG002
             dtype=np.int64,
         )
 
+    exposes_block_ordinals = True
+
+    def get_block_ordinals(self, req, group_idx, lg):  # noqa: ARG002
+        return np.fromiter(
+            self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices(
+                group_idx, valid_only=False
+            ),
+            dtype=np.int64,
+        )
+
     def commit_blocks_for_reuse(self, req: LlmRequest) -> None:
         self._mgr.try_commit_blocks(req)
 

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -50,7 +50,7 @@
     GenTransferStatus,
     KvCacheTransceiver,
 )
-from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
+from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_kv_capacity_tokens
 from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
     MambaHybridCacheManager,
     MambaHybridCacheManagerV2,
@@ -318,9 +318,20 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
             window_size = lg.sliding_window_size
 
             if window_size is not None:
-                allocated_blocks = (
-                    req.prompt_len + self._kv_cache_manager.num_extra_kv_tokens + tpb - 1
-                ) // tpb
+                stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb)
+                if req.py_beam_width == 1 and adapter.exposes_block_ordinals:
+                    # Position comes from the index, never from a length.
+                    ordinals = adapter.get_block_ordinals(req, idx, lg)
+                    in_window = ordinals[stale_end:total_blocks]
+                    groups.append(in_window[in_window >= 0].astype(np.int64))
+                    continue
+
+                capacity_tokens = get_kv_capacity_tokens(
+                    req,
+                    self._kv_cache_manager.num_extra_kv_tokens,
+                    include_drafts=is_gen_only,
+                )
+                allocated_blocks = (capacity_tokens + tpb - 1) // tpb
                 beam0_block_ids, tail_block_ids = self._split_packed_beam_block_ids(
                     block_ids,
                     req.py_beam_width,
@@ -333,24 +344,27 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
                         if tail_block_ids.size > 0
                         else beam0_block_ids
                     )
-                # Current PyExecutor cache managers disable KV-cache token sinks,
-                # so SWA block lists contain an evictable prompt prefix followed
-                # by the speculative scratch tail. If token sinks are enabled,
-                # this must use block-ordinal metadata to preserve the sink prefix.
-                # Remove scratch before trimming stale prompt blocks; otherwise a
-                # boundary-crossing allocation can displace initialized prompt KV.
-                scratch_blocks = max(0, allocated_blocks - total_blocks)
+                # Beam search still walks the packed-beam path below.
+                expected_valid = max(0, total_blocks - stale_end)
+                surplus_blocks = max(0, block_ids.size - expected_valid)
+                scratch_blocks = min(max(0, allocated_blocks - total_blocks), surplus_blocks)
                 if scratch_blocks > 0:
                     if req.py_beam_width != 1:
                         raise ValueError("speculative scratch blocks require beam_width == 1")
-                    block_ids = (
-                        block_ids[:-scratch_blocks]
-                        if scratch_blocks < block_ids.size
-                        else np.array([], dtype=np.int64)
+                    block_ids = block_ids[:-scratch_blocks]
+                # Only two sizes are explainable here: the manager kept the
+                # stale head (total_blocks) or pruned it (expected_valid).
+                size_is_unexplained = expected_valid < block_ids.size < total_blocks
+                if req.py_beam_width == 1 and size_is_unexplained:
+                    logger.warning_once(
+                        "disagg KV slice: windowed layer group has "
+                        f"{block_ids.size} pruned blocks but only {expected_valid} "
+                        f"are in window (prompt_len={req.prompt_len}, tpb={tpb}, "
+                        f"window={window_size}, allocated={allocated_blocks}). "
+                        "Some producer disagrees with get_kv_capacity_tokens; "
+                        "the trim is about to drop live KV.",
+                        key="disagg_kv_slice_capacity_divergence",
                     )
-                # Drop stale blocks the manager may still expose (V1 pre-eviction).
-                stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb)
-                expected_valid = max(0, total_blocks - stale_end)
                 # Stale prefix already pruned above; skip reuse-hit blocks that
                 # land inside the window. Clamp to 0: ctx side has cached_per_lg
                 # synthetically 0, and a reuse hit may fall entirely inside the

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -92,7 +92,13 @@
     KVCacheV2SsmLifeCycleIterationStats,
     KVCacheV2SsmSnapshotIterationStats,
 )
-from .llm_request import LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length
+from .llm_request import (
+    LlmRequest,
+    LlmRequestState,
+    SamplingConfig,
+    get_draft_token_length,
+    get_kv_capacity_tokens,
+)
 from .resource_manager import (
     BaseResourceManager,
     CacheTypeCpp,
@@ -2661,7 +2667,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool:
         # Helix requests carry the rank-local strided slice in prompt_len;
         # the global ledger sizes off the full prompt instead.
         prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len
-        target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens
+        target = get_kv_capacity_tokens(req, self.num_extra_kv_tokens, prompt_len=prompt_len)
         capacity = max(kv_cache.capacity, target)
         pre_cap = kv_cache.capacity
 

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -1702,3 +1702,17 @@ def get_draft_token_length(request: LlmRequest) -> int:
     if request.py_draft_tokens is not None:
         return len(request.py_draft_tokens)
     return 0
+
+
+def get_kv_capacity_tokens(request: LlmRequest,
+                           num_extra_kv_tokens: int,
+                           *,
+                           prompt_len: Optional[int] = None,
+                           include_drafts: bool = True) -> int:
+    """KV tokens a request occupies, defined once.
+
+    Consumers that each spell the sum out drift apart silently.
+    """
+    base = request.prompt_len if prompt_len is None else prompt_len
+    draft_len = get_draft_token_length(request) if include_drafts else 0
+    return base + draft_len + num_extra_kv_tokens