[TRTLLM-12499][perf] Bounded block-ID retrieval for pipelined KV transfer chunks

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -2170,6 +2170,15 @@ class BaseKVCacheManager
         std::vector<LlmRequest::RequestIdType> const& requestIds, SizeType32 windowSize) const
         = 0;
 
+    //! \brief Get resident block ids in the absolute request block range [blockBegin, blockEnd), per beam.
+    //! \details This non-virtual convenience wrapper copies only the requested range out of the sequence's persistent
+    //! block table, dropping front-detached SWA blocks (which keep their recycled ids in the raw table). The result is
+    //! always the *contiguous* run ending at blockEnd, i.e. ordinals [blockEnd - result.size(), blockEnd), so a caller
+    //! can recover each id's block ordinal from blockEnd and the result size alone. Requesting a blockEnd past the
+    //! sequence's allocated blocks would break that guarantee and is rejected.
+    [[nodiscard]] std::vector<std::vector<SizeType32>> getCacheBlockIdsRange(
+        LlmRequest::RequestIdType requestId, SizeType32 windowSize, SizeType32 blockBegin, SizeType32 blockEnd) const;
+
     /// @brief Get the last block id (beam 0) for a given sequence and window size
     [[nodiscard]] virtual std::optional<KVCacheBlock::IdType> getLastBlockId(LlmRequest::RequestIdType requestId) const
         = 0;

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -38,6 +38,7 @@
 #include "tensorrt_llm/runtime/worldConfig.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <limits>
 #include <map>
 #include <optional>
@@ -4629,6 +4630,34 @@ std::vector<std::vector<SizeType32>> const& KVCacheManager::getCacheBlockIds(
     return getSequence(requestId).getCacheBlockIds(windowSize);
 }
 
+std::vector<std::vector<SizeType32>> BaseKVCacheManager::getCacheBlockIdsRange(
+    LlmRequest::RequestIdType requestId, SizeType32 windowSize, SizeType32 blockBegin, SizeType32 blockEnd) const
+{
+    TLLM_CHECK_WITH_INFO(blockBegin >= 0, "blockBegin must be non-negative");
+    TLLM_CHECK_WITH_INFO(blockEnd >= 0, "blockEnd must be non-negative");
+    TLLM_CHECK_WITH_INFO(blockBegin <= blockEnd, "blockBegin must not exceed blockEnd");
+    // Read the block table and the eviction count off the same sequence, so a manager that overrides one but not the
+    // other cannot hand back a block table and a front-eviction count that disagree.
+    auto const& sequence = getSequence(requestId);
+    auto const& blockIdsPerBeam = sequence.getCacheBlockIds(windowSize);
+    auto const firstResidentBlock = static_cast<size_t>(sequence.getNumFrontBlocksRemoved(windowSize));
+    auto const end = static_cast<size_t>(blockEnd);
+    std::vector<std::vector<SizeType32>> result;
+    result.reserve(blockIdsPerBeam.size());
+    for (auto const& blockIds : blockIdsPerBeam)
+    {
+        TLLM_CHECK_WITH_INFO(end <= blockIds.size(),
+            "blockEnd=%d exceeds the %zu blocks allocated for request %lu at windowSize=%d; the result would not end "
+            "at blockEnd, and callers recover block ordinals from its size",
+            blockEnd, blockIds.size(), static_cast<unsigned long>(requestId), windowSize);
+        auto const begin = std::min(end, std::max(firstResidentBlock, static_cast<size_t>(blockBegin)));
+        auto const beginOffset = static_cast<std::ptrdiff_t>(begin);
+        auto const endOffset = static_cast<std::ptrdiff_t>(end);
+        result.emplace_back(blockIds.begin() + beginOffset, blockIds.begin() + endOffset);
+    }
+    return result;
+}
+
 std::vector<executor::IdType> KVCacheManager::commitAndGetBlockHashesForRequest(
     LlmRequest const& llmRequest, SizeType32 windowSize)
 {

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp

@@ -649,6 +649,11 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
             nb::arg("request_id"), nb::arg("window_size"), nb::call_guard<nb::gil_scoped_release>())
         .def("get_batch_cache_block_ids", &BaseKVCacheManager::getBatchCacheBlockIds,
             nb::call_guard<nb::gil_scoped_release>())
+        // Deliberately keeps the GIL, unlike its neighbours: it reaches the virtual getSequence, whose trampoline can
+        // call back into a Python subclass.  The copy it makes is bounded by the requested range, so there is little
+        // to gain from releasing.
+        .def("get_cache_block_ids_range", &BaseKVCacheManager::getCacheBlockIdsRange, nb::arg("request_id"),
+            nb::arg("window_size"), nb::arg("block_begin"), nb::arg("block_end"))
         .def("flush_iteration_events", &BaseKVCacheManager::flushIterationEvents,
             nb::call_guard<nb::gil_scoped_release>())
         .def("sync_transfer_manager_with_buffer_manager", &BaseKVCacheManager::syncTransferManagerWithBufferManager,

cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

@@ -9703,6 +9703,66 @@ TEST_F(KVCacheManagerTest, VSWABlockStoredDuringGeneration)
     EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool);
 }
 
+// getCacheBlockIdsRange must return the contiguous run of resident blocks ending at blockEnd.
+// detachFrontBlock leaves the recycled physical id of an out-of-window block in the raw block
+// table, so a range starting before the eviction count must skip it: the pipelined KV transceiver
+// would otherwise read KV out of a block that already belongs to another request.
+TEST_F(KVCacheManagerTest, VSWAGetCacheBlockIdsRangeExcludesDetachedFrontBlocks)
+{
+    auto constexpr blocksInPrimaryPool = 10;
+    auto const stream = std::make_shared<tr::CudaStream>();
+    tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH};
+    auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream);
+
+    // Seq 0: 11 input tokens covering B0=[1000..1003], B1=[1004..1007], B2=[1008..1010] (partial).
+    auto inputTokens0 = std::make_shared<VecTokens>(11);
+    std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN);
+    auto llmRequest0
+        = std::make_shared<LlmRequest>(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING);
+    addSequenceForTest(*kvCacheManager, 0, 11, kVSWA_BEAM_WIDTH, llmRequest0);
+    tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest0);
+    kvCacheManager->storeContextBlocks(*llmRequest0);
+
+    auto const& rawBlockIds = kvCacheManager->getCacheBlockIds(0, kVSWA_ATTENTION_WINDOW).at(kVSWA_BEAM_IDX);
+    ASSERT_EQ(rawBlockIds.size(), 3U);
+
+    // Before any eviction the range is exactly what was asked for.
+    EXPECT_EQ(kvCacheManager->getSequence(0).getNumFrontBlocksRemoved(kVSWA_ATTENTION_WINDOW), 0);
+    EXPECT_THAT(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 0, 3).at(kVSWA_BEAM_IDX),
+        ::testing::ElementsAre(rawBlockIds.at(0), rawBlockIds.at(1), rawBlockIds.at(2)));
+    EXPECT_THAT(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 1, 3).at(kVSWA_BEAM_IDX),
+        ::testing::ElementsAre(rawBlockIds.at(1), rawBlockIds.at(2)));
+    // One entry per beam, even when the range is empty.
+    EXPECT_EQ(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 2, 2).size(),
+        static_cast<size_t>(kVSWA_BEAM_WIDTH));
+    EXPECT_THAT(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 2, 2).at(kVSWA_BEAM_IDX),
+        ::testing::IsEmpty());
+
+    // Generation step: numTokens becomes 12; adjustBlocksIfNeeded detaches B0 (12-0*4 >= 8+4).
+    llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX);
+    kvCacheManager->addToken(0);
+    ASSERT_EQ(kvCacheManager->getSequence(0).getNumFrontBlocksRemoved(kVSWA_ATTENTION_WINDOW), 1);
+    // B0's recycled id is still in the raw table, so the range query is the only thing standing
+    // between the transceiver and another request's data.
+    ASSERT_EQ(kvCacheManager->getCacheBlockIds(0, kVSWA_ATTENTION_WINDOW).at(kVSWA_BEAM_IDX).size(), 3U);
+
+    EXPECT_THAT(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 0, 1).at(kVSWA_BEAM_IDX),
+        ::testing::IsEmpty());
+    EXPECT_THAT(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 0, 3).at(kVSWA_BEAM_IDX),
+        ::testing::ElementsAre(rawBlockIds.at(1), rawBlockIds.at(2)));
+    EXPECT_THAT(kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 2, 3).at(kVSWA_BEAM_IDX),
+        ::testing::ElementsAre(rawBlockIds.at(2)));
+
+    // Bad bounds are programming errors, and so is reading past the allocated blocks: the result
+    // would no longer end at blockEnd, which is how callers recover each id's block ordinal.
+    EXPECT_THROW((void) kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, -1, 2), std::runtime_error);
+    EXPECT_THROW((void) kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 0, -1), std::runtime_error);
+    EXPECT_THROW((void) kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 2, 1), std::runtime_error);
+    EXPECT_THROW((void) kvCacheManager->getCacheBlockIdsRange(0, kVSWA_ATTENTION_WINDOW, 0, 4), std::runtime_error);
+
+    EXPECT_NO_THROW(static_cast<void>(kvCacheManager->removeSequence(0, llmRequest0)));
+}
+
 // Verify that when an OOW block is stolen by another sequence, storeBlocks does
 // not restore that missing anchor under the original sequence's key or corrupt
 // the acquiring sequence's trie, and all blocks are properly released.

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -1281,6 +1281,7 @@ def create_autodeploy_executor(
             attention_type_cpp,
             cache_transceiver_config,
             mamba_cache_manager=None,
+            enable_chunked_prefill=getattr(ad_config, "enable_chunked_prefill", False),
         )
 
     # Guided (structured) decoding.

tensorrt_llm/_torch/disaggregation/base/transfer.py

@@ -11,6 +11,36 @@
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 
 
+def project_blocks_to_global_chunk(
+    block_ids: np.ndarray,
+    chunk_block_offset: int,
+    chunk_block_count: int,
+    resident_block_end: int,
+) -> np.ndarray:
+    """Project a global block chunk into a suffix-resident block list.
+
+    ``block_ids`` represents the resident suffix of the logical range
+    ``[0, resident_block_end)``. ``chunk_block_offset`` and
+    ``chunk_block_count`` describe a chunk in that global coordinate space.
+    """
+    if chunk_block_count <= 0 or len(block_ids) == 0:
+        return block_ids[:0]
+
+    resident_start = max(0, resident_block_end - len(block_ids))
+    resident_end = resident_block_end
+    chunk_start = chunk_block_offset
+    chunk_end = chunk_start + chunk_block_count
+
+    overlap_start = max(chunk_start, resident_start)
+    overlap_end = min(chunk_end, resident_end)
+    if overlap_start >= overlap_end:
+        return block_ids[:0]
+
+    local_start = overlap_start - resident_start
+    local_end = overlap_end - resident_start
+    return block_ids[local_start:local_end]
+
+
 @dataclass
 class TokenRange:
     """Range of tokens in the sequence dimension."""
@@ -25,6 +55,30 @@ def __post_init__(self):
             raise ValueError(f"Invalid range: [{self.start}, {self.end})")
 
 
+def derive_chunk_block_coords(
+    token_range: Optional[TokenRange],
+    tokens_per_block: int,
+) -> tuple[int, int]:
+    """Derive global chunk block offset and count from a block-aligned token_range.
+
+    Producers of pipelined slices must align every non-final chunk boundary to
+    ``tokens_per_block``. The final partial block is represented by an aligned
+    range ending at its enclosing block boundary.
+    """
+    if token_range is None:
+        return 0, 0
+    if tokens_per_block <= 0:
+        raise ValueError("tokens_per_block must be positive")
+    if token_range.start % tokens_per_block != 0 or token_range.end % tokens_per_block != 0:
+        raise ValueError(
+            f"token_range [{token_range.start}, {token_range.end}) must be "
+            f"block-aligned with tokens_per_block={tokens_per_block}"
+        )
+    chunk_offset = token_range.start // tokens_per_block
+    chunk_block_count = (token_range.end - token_range.start) // tokens_per_block
+    return chunk_offset, chunk_block_count
+
+
 @dataclass
 class LayerRange:
     """Range of layers to transfer."""
@@ -44,7 +98,9 @@ class KVSlice:
     """A KV cache slice covering token_range = [start, end) of one request.
 
     Single-slice transfer uses [0, prompt_len) with is_last_slice=True;
-    multi-slice transfers split token_range and mark the last slice.
+    multi-slice (pipelined) transfers split token_range and mark the last slice.
+    For pipelined chunks, token_range is block-aligned and encodes the global
+    chunk position; derive block offset/count via derive_chunk_block_coords().
 
     Per-layer token starts are NOT encoded in token_range — they are derived
     from block count by the sender:
@@ -65,6 +121,7 @@ class KVSlice:
     )  # Physical block IDs per layer group, each np.ndarray(dtype=np.int64)
     is_last_slice: bool = False
     mamba_state_index: Optional[int] = None
+    total_blocks: Optional[int] = None
 
 
 class SessionStatus(Enum):
@@ -104,7 +161,7 @@ class SessionArgsBase:
 
     params: DisaggregatedParams
     # Captured from LlmRequest.prompt_len; needed for SWA stale_end derivation.
-    prompt_len: Optional[int] = None
+    prompt_len: int
     beam_width: int = 1
 
 
@@ -158,7 +215,16 @@ def __init__(self, sender: SenderBase, args: SessionArgsBase):
         self._sender = sender
 
     @abstractmethod
-    def send(self, slice: KVSlice) -> None: ...
+    def send(self, slice: KVSlice) -> None:
+        """Send a KV slice.
+
+        Args:
+            slice: The KV slice describing which source blocks to send.
+                For pipelined chunks, ``token_range`` is the shared sender-side
+                chunk cursor; each layer group projects it into its own
+                resident/windowed source and destination block ranges.
+        """
+        ...
 
     @abstractmethod
     def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: ...