[None][fix] Fix disagg KV cache transfer crash on reuse budget skip

cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp

@@ -24,6 +24,25 @@ namespace tensorrt_llm::batch_manager
 
 using SizeType32 = MicroBatchScheduler::SizeType32;
 
+/// Return the forward-pass token cost for a context chunk with KV cache reuse.
+///
+/// setPrepopulatedPromptLen shifts the chunk window right by the prepopulated
+/// amount rather than shrinking it.  For non-last chunks the model still
+/// processes approximately @p chunkSize tokens; only for the last chunk is the
+/// cost @p contextRemaining - @p reusable.
+static SizeType32 reuse_adjusted_compute(SizeType32 chunkSize, SizeType32 reusable, SizeType32 contextRemaining)
+{
+    if (reusable <= 0)
+    {
+        return chunkSize;
+    }
+    if (reusable + chunkSize < contextRemaining)
+    {
+        return chunkSize;
+    }
+    return std::max<SizeType32>(0, contextRemaining - reusable);
+}
+
 MicroBatchScheduler::MicroBatchScheduler(std::optional<batch_scheduler::ContextChunkingConfig> ctxChunkConfig,
     std::optional<SizeType32> maxContextLength, LlmRequestState noScheduleUntilState,
     LlmRequestState noScheduleAfterState)
@@ -38,16 +57,15 @@ void MicroBatchScheduler::fitDraftTokens(RequestVector& contextsToBeChunked,
     std::optional<SizeType32> ctxTokensCapacity, SizeType32 const chunkUnitSize,
     std::optional<SizeType32> const& maxContextLength)
 {
-    // How many compute tokens (chunk - reusable) are in this batch already?
+    // How many compute tokens are in this batch already?
     SizeType32 numCtxTokens{0};
     for (auto const& llmReq : contextsToBeChunked)
     {
         SizeType32 const chunkSize = llmReq->getContextChunkSize();
-        // contextRemaining = P for first chunk; used to compute actual model token count.
         SizeType32 const contextRemaining = llmReq->getContextRemainingLength();
         SizeType32 const reusable
             = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), contextRemaining) : 0;
-        numCtxTokens += std::min(chunkSize, std::max<SizeType32>(0, contextRemaining - reusable));
+        numCtxTokens += reuse_adjusted_compute(chunkSize, reusable, contextRemaining);
     }
 
     // Discard draft tokens that won't fit into the existing chunk unit, max
@@ -125,14 +143,13 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
             SizeType32 actualChunkSize = llmReq->getContextChunkSize();
             SizeType32 actualIncrement = actualChunkSize - pastChunkSize;
 
-            // Compute-aware budget: reusable tokens are served from cache and do not
-            // consume forward-pass capacity. Only the tokens beyond the reusable prefix count.
-            SizeType32 const reusable = llmReq->isFirstContextChunk()
-                ? std::min(llmReq->getEstimatedReusableTokens(), llmReq->getContextRemainingLength())
-                : 0;
-            SizeType32 const pastCompute = std::max<SizeType32>(0, pastChunkSize - std::min(reusable, pastChunkSize));
-            SizeType32 const actualCompute
-                = std::max<SizeType32>(0, actualChunkSize - std::min(reusable, actualChunkSize));
+            // Compute-aware budget accounting for setPrepopulatedPromptLen's
+            // chunk-shift behaviour (non-last chunks keep their full size).
+            SizeType32 const contextRemaining = llmReq->getContextRemainingLength();
+            SizeType32 const reusable
+                = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), contextRemaining) : 0;
+            SizeType32 const pastCompute = reuse_adjusted_compute(pastChunkSize, reusable, contextRemaining);
+            SizeType32 const actualCompute = reuse_adjusted_compute(actualChunkSize, reusable, contextRemaining);
             SizeType32 const computeIncrement = actualCompute - pastCompute;
 
             if ((ctxTokensCapacity && numCtxTokens + computeIncrement > ctxTokensCapacity.value())
@@ -181,24 +198,21 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
     for (auto& llmReq : contextsToBeChunked)
     {
         SizeType32 const suggestedChunkSize = llmReq->getContextRemainingLength();
-        // Reusable tokens are "free" — they don't consume forward-pass compute budget.
         SizeType32 const reusable
             = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), suggestedChunkSize) : 0;
-        SizeType32 const computeCost = suggestedChunkSize - reusable;
+        SizeType32 const computeCost = reuse_adjusted_compute(suggestedChunkSize, reusable, suggestedChunkSize);
         SizeType32 actualChunkSize = suggestedChunkSize;
         if (ctxTokensCapacity && computeCost > ctxTokensCapacity.value())
         {
-            // Model processes min(chunk_size, P - reusable) tokens starting from position reusable.
-            // To keep model tokens within budget: chunk_size <= capacity (not reusable + capacity).
             actualChunkSize = ctxTokensCapacity.value();
         }
         if (maxContextLength)
         {
-            // maxContextLength limits compute tokens, not total tokens.
-            SizeType32 const actualCompute = std::max<SizeType32>(0, actualChunkSize - reusable);
+            SizeType32 const actualCompute = reuse_adjusted_compute(actualChunkSize, reusable, suggestedChunkSize);
             if (actualCompute > maxContextLength.value())
             {
-                actualChunkSize = std::min<SizeType32>(reusable + maxContextLength.value(), suggestedChunkSize);
+                actualChunkSize = maxContextLength.value();
+                actualChunkSize = std::min(actualChunkSize, suggestedChunkSize);
             }
         }
         if (actualChunkSize != suggestedChunkSize)
@@ -208,10 +222,7 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
         llmReq->setContextChunkSize(actualChunkSize);
         if (ctxTokensCapacity)
         {
-            // Decrement by actual model token count: min(chunk_size, P - reusable).
-            // This equals min(actualChunkSize, computeCost) since computeCost = suggestedChunkSize - reusable.
-            SizeType32 const modelCost
-                = std::min(actualChunkSize, std::max<SizeType32>(0, suggestedChunkSize - reusable));
+            SizeType32 const modelCost = reuse_adjusted_compute(actualChunkSize, reusable, suggestedChunkSize);
             ctxTokensCapacity = ctxTokensCapacity.value() - modelCost;
         }
     }
@@ -309,10 +320,12 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
             if (!mCtxChunkConfig) // skip chunking
             {
                 constexpr SizeType32 beam{0};
-                reqNumTokens
-                    = llmReq->getNumTokens(beam) + (llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0);
-                // Compute tokens = total - reusable (at least 1 to make progress)
-                SizeType32 const computeTokens = std::max(1, reqNumTokens - reusable);
+                SizeType32 const contextTokens = llmReq->getNumTokens(beam);
+                SizeType32 const draftTokens = llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0;
+                reqNumTokens = contextTokens + draftTokens;
+                SizeType32 const contextCompute
+                    = reuse_adjusted_compute(contextTokens, reusable, llmReq->getContextRemainingLength());
+                SizeType32 const computeTokens = std::max(1, contextCompute + draftTokens);
                 TLLM_CHECK_WITH_INFO(!mMaxContextLength || computeTokens <= mMaxContextLength.value(),
                     "Context compute tokens (%d) exceeds the limit value (%d)", computeTokens,
                     mMaxContextLength.value());
@@ -329,9 +342,8 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
                 llmReq->setContextChunkSize(llmReq->getContextRemainingLength());
                 auto const draftTokens
                     = (llmReq->isLastContextChunk() && llmReq->hasDraftTokens()) ? llmReq->getNumDraftTokens() : 0;
-                // Compute cost: context compute + draft tokens
-                // (reusable tokens only offset context tokens, not draft tokens)
-                SizeType32 const contextCompute = std::max(0, llmReq->getContextChunkSize() - reusable);
+                SizeType32 const contextCompute = reuse_adjusted_compute(
+                    llmReq->getContextChunkSize(), reusable, llmReq->getContextRemainingLength());
                 SizeType32 computeTokens = contextCompute + draftTokens;
 
                 if (mMaxContextLength)
@@ -398,10 +410,9 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
         if (llmReq->getContextChunkSize() > 0)
         {
             contextRequests.emplace_back(llmReq);
-            // Only count compute tokens (total - reusable).
-            // Reusable credit only applies to the first context chunk.
             SizeType32 const reusable = llmReq->isFirstContextChunk() ? llmReq->getEstimatedReusableTokens() : 0;
-            SizeType32 const computeTokens = std::max(0, llmReq->getContextChunkSize() - reusable);
+            SizeType32 const computeTokens
+                = reuse_adjusted_compute(llmReq->getContextChunkSize(), reusable, llmReq->getContextRemainingLength());
             batchNumTokens += computeTokens;
             TLLM_LOG_DEBUG("context request scheduled: ID %lu, chunk size %d%s", llmReq->mRequestId,
                 llmReq->getContextChunkSize(), reusable > 0 ? (", reusable " + std::to_string(reusable)).c_str() : "");

cpp/tensorrt_llm/kernels/IndexerKCacheGather.h

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mapping_fp8,
+    int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
+    int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
+    int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
+    cudaStream_t stream = 0);
+
+}
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.cu

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "convertReqIndexToGlobal.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+// Each thread handles one element at (token_id, col).
+// Grid: (num_tokens, ceil(numTopkTokens / blockDim.x))
+__global__ void convertReqIndexToGlobalKernel(int32_t const* __restrict__ reqId, int32_t const* __restrict__ blockTable,
+    int32_t const* __restrict__ tokenIndices, int32_t* __restrict__ output, int32_t numTopkTokens,
+    int32_t maxNumBlocksPerReq, int32_t blockSize, int32_t strideFactor, int32_t layerId, int64_t btStride0,
+    int64_t btStride1, int64_t tiStride0, int64_t tiStride1, int64_t outStride0, int64_t outStride1)
+{
+    int32_t const tokenId = blockIdx.x;
+    int32_t const col = blockIdx.y * blockDim.x + threadIdx.x;
+
+    if (col >= numTopkTokens)
+    {
+        return;
+    }
+
+    // Load request id for this token
+    int32_t const req = reqId[tokenId];
+
+    // Load token index
+    int32_t const tok = tokenIndices[tokenId * tiStride0 + col * tiStride1];
+
+    // Invalid token → output -1
+    if (tok < 0)
+    {
+        output[tokenId * outStride0 + col * outStride1] = -1;
+        return;
+    }
+
+    // Compute block id and in-block offset
+    int32_t const blockId = tok / blockSize;
+    int32_t const inblockOff = tok % blockSize + layerId * blockSize;
+
+    // Guard block_table access
+    if (blockId >= maxNumBlocksPerReq)
+    {
+        output[tokenId * outStride0 + col * outStride1] = -1;
+        return;
+    }
+
+    int32_t const base = blockTable[req * btStride0 + blockId * btStride1];
+
+    // Padding entry in block table
+    if (base < 0)
+    {
+        output[tokenId * outStride0 + col * outStride1] = -1;
+        return;
+    }
+
+    output[tokenId * outStride0 + col * outStride1] = base * strideFactor + inblockOff;
+}
+
+void invokeConvertReqIndexToGlobal(int32_t const* reqId, int32_t const* blockTable, int32_t const* tokenIndices,
+    int32_t* output, int32_t numTokens, int32_t numTopkTokens, int32_t maxNumBlocksPerReq, int32_t blockSize,
+    int32_t strideFactor, int32_t layerId, int64_t btStride0, int64_t btStride1, int64_t tiStride0, int64_t tiStride1,
+    int64_t outStride0, int64_t outStride1, cudaStream_t stream)
+{
+    if (numTokens == 0 || numTopkTokens == 0)
+    {
+        return;
+    }
+
+    constexpr int32_t kThreadsPerBlock = 256;
+    int32_t const tilesPerRow = (numTopkTokens + kThreadsPerBlock - 1) / kThreadsPerBlock;
+    dim3 const grid(numTokens, tilesPerRow);
+    dim3 const block(kThreadsPerBlock);
+
+    convertReqIndexToGlobalKernel<<<grid, block, 0, stream>>>(reqId, blockTable, tokenIndices, output, numTopkTokens,
+        maxNumBlocksPerReq, blockSize, strideFactor, layerId, btStride0, btStride1, tiStride0, tiStride1, outStride0,
+        outStride1);
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/convertReqIndexToGlobal.h

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+void invokeConvertReqIndexToGlobal(int32_t const* reqId, int32_t const* blockTable, int32_t const* tokenIndices,
+    int32_t* output, int32_t numTokens, int32_t numTopkTokens, int32_t maxNumBlocksPerReq, int32_t blockSize,
+    int32_t strideFactor, int32_t layerId, int64_t btStride0, int64_t btStride1, int64_t tiStride0, int64_t tiStride1,
+    int64_t outStride0, int64_t outStride1, cudaStream_t stream = 0);
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END