[None][feat] Do not REVIEW: Stack PRs (v5 + updated_schedule)

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -162,6 +162,46 @@ struct LinearAttentionMetadata
         return false;
     }
 
+    [[nodiscard]] SizeType32 calcNumBlocksNeededForReq(
+        SizeType32 promptLen, SizeType32 tokensPerBlock, bool enableReuse) const
+    {
+        if (!enableReuse)
+        {
+            return 1;
+        }
+        SizeType32 count = 0;
+        if (statesSnapshotInterval > 0)
+        {
+            count += promptLen / statesSnapshotInterval; // round down
+        }
+        if (saveLastSnapshot
+            && (promptLen / tokensPerBlock * tokensPerBlock
+                != promptLen / statesSnapshotInterval * statesSnapshotInterval))
+        {
+            count += 1;
+        }
+        if (promptLen % tokensPerBlock == 0)
+        {
+            // corner case
+            count += 1;
+        }
+        return count;
+    }
+
+    [[nodiscard]] SizeType32 calcNumAdditionalBlocksNeededForReq(
+        SizeType32 numTokens, SizeType32 promptLen, SizeType32 tokensPerBlock, bool enableReuse) const
+    {
+        if (!enableReuse)
+        {
+            return 0;
+        }
+        if (promptLen % tokensPerBlock == 0 && numTokens <= promptLen + 1)
+        {
+            return 1;
+        }
+        return 0;
+    }
+
     [[nodiscard]] bool hasRecurrentStatesCache() const
     {
         return hasRecurrentStatesCache(cacheType);
@@ -872,7 +912,9 @@ class WindowBlockManager
 
     //! \brief According to request's current position, copy data from the last full block to the next block (ignoring
     //! the placeholder block). It should be called after every context chunk is processed.
-    void copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest);
+    //! \return true iff at least one async block transfer was actually issued. Callers can use this to decide
+    //! whether a subsequent refreshBlocks()/syncTransfers() is necessary.
+    bool copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest);
 
     void replaceSharedBlock(GenerationRequest& sequence, SizeType32 blockIdx);
 
@@ -1387,7 +1429,8 @@ class BlockManager
 
     //! \brief According to request's current position, copy data from the last full block to the next block (ignoring
     //! the placeholder block). It should be called after every context chunk is processed.
-    void copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest);
+    //! \return true iff at least one async block transfer was actually issued.
+    bool copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest);
 
     void replaceSharedBlock(GenerationRequest& sequence, SizeType32 windowSize, SizeType32 blockIdx);
 
@@ -2217,7 +2260,12 @@ class KVCacheManager : public BaseKVCacheManager
 
     //! \brief According to request's current position, copy data from the last full block to the next block (ignoring
     //! the placeholder block). It should be called before every forward step, after adding new tokens.
-    void copyLinearAttentionBlock(LlmRequest const& llmRequest);
+    //! \return true iff at least one async block transfer was actually issued for this request. The caller can
+    //! aggregate this across requests and skip refreshBlocks() (which performs a stream sync) when no copies happened.
+    bool copyLinearAttentionBlock(LlmRequest const& llmRequest);
+
+    //! \brief Batch variant of copyLinearAttentionBlock. Returns true iff at least one copy was issued.
+    bool copyLinearAttentionBlockBatch(std::vector<std::shared_ptr<LlmRequest>> const& llmRequests);
 
     void addSequenceBatch(
         std::vector<std::tuple<LlmRequest::RequestIdType, SizeType32, SizeType32>> const& requestInfos,

cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp

@@ -456,7 +456,7 @@ std::tuple<RequestVector, RequestVector> MaxUtilizationScheduler::operator()(
                 // Here we simulate freeing the kvCache blocks associated with that sequence
                 kvCacheManager.schedulingRemoveSequence((*lastStartedReqIt)->mRequestId);
                 pausedRequests.emplace_back(*lastStartedReqIt);
-                TLLM_LOG_DEBUG("MaxUtilizationScheduler: request ID %lu -> pause", (*lastStartedReqIt)->mRequestId);
+                TLLM_LOG_INFO("MaxUtilizationScheduler: request ID %lu -> pause", (*lastStartedReqIt)->mRequestId);
                 reqItEnd = std::next(lastStartedReqIt).base();
             }
             else

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -2222,12 +2222,16 @@ void BlockManager::allocateBlock(GenerationRequest& sequence, SizeType32 windowS
     mWindowBlockManagers.at(windowSize).allocateBlock(sequence, false);
 }
 
-void BlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest)
+bool BlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest)
 {
+    bool didCopy = false;
     for (auto& [windowSize, manager] : mWindowBlockManagers)
     {
-        manager.copyLinearAttentionBlock(sequence, llmRequest);
+        // Use a temp to avoid short-circuiting; every window must run.
+        bool const windowDidCopy = manager.copyLinearAttentionBlock(sequence, llmRequest);
+        didCopy = didCopy || windowDidCopy;
     }
+    return didCopy;
 }
 
 bool WindowBlockManager::tryAllocatePlaceholderForLinearAttention(GenerationRequest& sequence, bool shareAmongBeams)
@@ -2363,11 +2367,11 @@ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAm
     }
 }
 
-void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& request)
+bool WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& request)
 {
     if (!isRecurrentState())
     {
-        return;
+        return false;
     }
 
     auto const requestId = request.mRequestId;
@@ -2376,7 +2380,7 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L
     if (mAllocatedBlocksPerSeq.find(requestId) == mAllocatedBlocksPerSeq.end())
     {
         TLLM_LOG_WARNING("%s::copyLinearAttentionBlock - Request %lu not found", mLogPrefix.c_str(), requestId);
-        return;
+        return false;
     }
 
     // It points to the next token to be processed/generated
@@ -2391,9 +2395,10 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L
         if (TLLM_LIKELY(sequence.getBeamWidth() == 1))
         {
             // the block of beam0 is inherited from context phase, no need to copy
-            return;
+            return false;
         }
-        // copy beam 0 to other beams
+        // copy beam 0 to other beams: beamWidth >= 2 here, loop runs at least once
+        // with no skip path, so an onboard is always issued.
         auto beam0Block = getBlockById(sequence.getCacheBlockIds(mWindowSize).at(0).back());
         for (auto beamIdx = 1; beamIdx < sequence.getBeamWidth(); ++beamIdx)
         {
@@ -2406,17 +2411,18 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L
                                  // transfer manager to copy the entire block.
                 sequence.getTransferMode(), sequence.getDirectory());
         }
-        return;
+        return true;
     }
 
     // copy only happens in context phase or the corner case above
     if (currentPosition % mTokensPerBlock != 0 || currentPosition > request.getPromptLen() || currentPosition == 0)
     {
-        return;
+        return false;
     }
 
     auto prevBlockIndex = currentPosition / mTokensPerBlock - 1; // signed
     std::set<std::pair<KVCacheBlock::IdType, KVCacheBlock::IdType>> onboardedBlocks;
+    bool didCopy = false;
     for (auto beamIdx = 0; beamIdx < sequence.getBeamWidth(); ++beamIdx)
     {
         auto const& beamBlockIds = sequence.getCacheBlockIds(mWindowSize).at(beamIdx);
@@ -2455,7 +2461,9 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L
                              // manager to copy the entire block.
             sequence.getTransferMode(), sequence.getDirectory());
         onboardedBlocks.insert({prevBlockId, nextBlockId});
+        didCopy = true;
     }
+    return didCopy;
 }
 
 std::pair<SizeType32, std::vector<KVCacheBlock::IdType>> WindowBlockManager::storeBlocks(
@@ -3311,6 +3319,11 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep(LlmRequest const& req, bool tw
             = std::min((isCrossKv() ? req.getEncoderOutputLen() : req.mPromptLen) + maxDraftTokensToAdd,
                   windowSize + mChunkSize)
             + mSinkBubbleLength;
+        if (LinearAttentionMetadata::hasLinearCache(windowSize))
+        {
+            return mBlockManager.getLinearAttentionMetadata()->calcNumBlocksNeededForReq(
+                promptCacheLen, getTokensPerBlock(), mEnableBlockReuse);
+        }
         auto const numSharedBlocks = promptCacheLen / getTokensPerBlock();
         auto const numUnSharedTokens = promptCacheLen % getTokensPerBlock();
         auto const numUnSharedBlocks
@@ -3355,6 +3368,12 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep(LlmRequest const& req, bool tw
         auto const maxTokensToAdd = std::min((twoStepsLookAhead ? 2 : 1) * tokensPerStep, maxTokensToAddToKVCache);
         auto const numNextTokens = numCurrTokens + maxTokensToAdd;
 
+        if (LinearAttentionMetadata::hasLinearCache(windowSize))
+        {
+            return mBlockManager.getLinearAttentionMetadata()->calcNumAdditionalBlocksNeededForReq(
+                numCurrTokens, req.getPromptLen(), getTokensPerBlock(), mEnableBlockReuse);
+        }
+
         if (numNextTokens > mBlockManager.getWindowSizeMetadata(windowSize).maxTokenNum)
         {
             return 0;
@@ -3389,7 +3408,8 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(
     {
         if (req.isGenerationInProgressState())
         {
-            return 0; // no need to allocate blocks for recurrent states during generation
+            return mBlockManager.getLinearAttentionMetadata()->calcNumAdditionalBlocksNeededForReq(
+                req.getNumTokens(0), req.getPromptLen(), getTokensPerBlock(), mEnableBlockReuse);
         }
         else if (!req.isContextFinished())
         {
@@ -3399,12 +3419,8 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(
             {
                 return 0;
             }
-            if (mEnableBlockReuse)
-            {
-                return req.getPromptLen() / mBlockManager.getLinearAttentionMetadata()->statesSnapshotInterval + 1
-                    + (mBlockManager.getLinearAttentionMetadata()->saveLastSnapshot ? 1 : 0);
-            }
-            return 1;
+            return mBlockManager.getLinearAttentionMetadata()->calcNumBlocksNeededForReq(
+                req.mPromptLen, getTokensPerBlock(), mEnableBlockReuse);
         }
     }
 
@@ -3548,10 +3564,29 @@ void KVCacheManager::addToken(RequestIdType requestId)
     mBlockManager.adjustBlocksIfNeeded(sequence);
 }
 
-void KVCacheManager::copyLinearAttentionBlock(LlmRequest const& llmRequest)
+bool KVCacheManager::copyLinearAttentionBlock(LlmRequest const& llmRequest)
 {
+    if (!mEnableBlockReuse)
+    {
+        // When block reuse is disabled, we always use a single slot for a sequence and move it to
+        // the end of blocks list, so copying is not needed.
+        return false;
+    }
     auto& sequence = getSequence(llmRequest.mRequestId);
-    mBlockManager.copyLinearAttentionBlock(sequence, llmRequest);
+    return mBlockManager.copyLinearAttentionBlock(sequence, llmRequest);
+}
+
+bool KVCacheManager::copyLinearAttentionBlockBatch(std::vector<std::shared_ptr<LlmRequest>> const& llmRequests)
+{
+    bool copiedAny = false;
+    for (auto const& req : llmRequests)
+    {
+        if (copyLinearAttentionBlock(*req))
+        {
+            copiedAny = true;
+        }
+    }
+    return copiedAny;
 }
 
 void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence)
@@ -3741,7 +3776,7 @@ void KVCacheManager::storeContextBlocks(LlmRequest const& llmRequest)
     }
     else
     {
-        TLLM_LOG_WARNING("[kv cache manager] storeContextBlocks: Can not find sequence for request %lu", requestId);
+        // TLLM_LOG_WARNING("[kv cache manager] storeContextBlocks: Can not find sequence for request %lu", requestId);
     }
 }
 

cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp

@@ -114,22 +114,38 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst,
             auto const& pool = pools[poolIdx];
 
             // For layer-first layout pools, block data is non-contiguous across layers.
-            // Copy each layer's block data separately.
+            // Pool shape: {numLayers, numBlocks, kvFactor, blockSize}. For a fixed block
+            // index, per-layer slices are contiguous rows of (kvFactor * blockSize) elements,
+            // separated by a stride of numBlocks rows between layers. Issue this as a single
+            // pitched cudaMemcpy2DAsync instead of one cudaMemcpyAsync per layer.
             if (pool.layerFirstLayout)
             {
                 auto srcPool = src->isPrimary() ? pool.primaryPtr : pool.secondaryPtr;
                 auto dstPool = dst->isPrimary() ? pool.primaryPtr : pool.secondaryPtr;
-                auto const srcBlockIdx = static_cast<tr::ITensor::DimType64>(src->getMemoryPoolBlockIndex());
-                auto const dstBlockIdx = static_cast<tr::ITensor::DimType64>(dst->getMemoryPoolBlockIndex());
-
-                for (SizeType32 layerIdx = 0; layerIdx < pool.numLayers; ++layerIdx)
-                {
-                    // pool shape: {numLayers, numBlocks, kvFactor, blockSize}
-                    // slice at {layerIdx, blockIdx} gives {1, kvFactor, blockSize}
-                    auto srcBlock = tr::ITensor::slice(srcPool, {layerIdx, srcBlockIdx}, 1);
-                    auto dstBlock = tr::ITensor::slice(dstPool, {layerIdx, dstBlockIdx}, 1);
-                    (isOffload ? mOffloadManager : mOnboardManager).copy(*srcBlock, *dstBlock);
-                }
+                auto const srcBlockIdx = static_cast<size_t>(src->getMemoryPoolBlockIndex());
+                auto const dstBlockIdx = static_cast<size_t>(dst->getMemoryPoolBlockIndex());
+
+                // Compute pitches from each pool independently: primary and secondary pools
+                // may have different block counts (mNumPrimaryBlocks vs mNumSecondaryBlocks),
+                // so their per-layer strides differ. Using the primary shape for both pitches
+                // would corrupt host-offloaded recurrent state on CPU<->GPU transfers.
+                auto const& srcShape = srcPool->getShape();
+                auto const& dstShape = dstPool->getShape();
+                TLLM_CHECK_WITH_INFO(srcShape.nbDims >= 2,
+                    "Expected layer-first KVCache pool to have at least 2 dims, got %d", srcShape.nbDims);
+                TLLM_CHECK_WITH_INFO(dstShape.nbDims >= 2,
+                    "Expected layer-first KVCache pool to have at least 2 dims, got %d", dstShape.nbDims);
+                auto const srcLayerStrideBytes = srcPool->getSizeInBytes() / static_cast<size_t>(pool.numLayers);
+                auto const dstLayerStrideBytes = dstPool->getSizeInBytes() / static_cast<size_t>(pool.numLayers);
+                // rowBytes is the per-block per-layer payload — identical for primary and secondary.
+                auto const rowBytes = srcLayerStrideBytes / static_cast<size_t>(srcShape.d[1]);
+
+                auto* srcBase = static_cast<char*>(srcPool->data()) + srcBlockIdx * rowBytes;
+                auto* dstBase = static_cast<char*>(dstPool->data()) + dstBlockIdx * rowBytes;
+
+                auto stream = (isOffload ? mOffloadManager : mOnboardManager).getStream().get();
+                TLLM_CUDA_CHECK(cudaMemcpy2DAsync(dstBase, dstLayerStrideBytes, srcBase, srcLayerStrideBytes, rowBytes,
+                    static_cast<size_t>(pool.numLayers), cudaMemcpyDefault, stream));
                 continue;
             }
 

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp

@@ -641,7 +641,9 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m)
         .def_prop_ro(
             "is_variable_window", [](tbk::KVCacheManager& self) { return self.getBlockManager().isVariableWindow(); })
         .def("copy_linear_attention_block", &tbk::KVCacheManager::copyLinearAttentionBlock, nb::arg("llm_request"),
-            nb::call_guard<nb::gil_scoped_release>());
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("copy_linear_attention_block_batch", &tbk::KVCacheManager::copyLinearAttentionBlockBatch,
+            nb::arg("llm_requests"), nb::call_guard<nb::gil_scoped_release>());
 }
 
 void tb::BasePeftCacheManagerBindings::initBindings(nb::module_& m)

tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py

@@ -25,6 +25,11 @@
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict, Optional, Tuple, Type
 
+from tensorrt_llm.quantization.modelopt_config import (
+    is_modelopt_quant_config,
+    parse_modelopt_quant_config,
+)
+
 from ..utils.logger import ad_logger
 
 
@@ -110,14 +115,23 @@ class ModelOPTQuantConfigReader(QuantConfigReader):
     DEFAULT_KV_CACHE_DTYPE = "fp8"
 
     def read_config(self, config: Dict) -> Dict:
-        producer = config.get("producer", {}).get("name")
-        # sanity check
-        if producer != "modelopt":
-            raise ValueError(f"Expected producer 'modelopt', got '{producer}'")
+        # Accept either modelopt shape: legacy (producer.name == "modelopt"
+        # with a "quantization" wrapper) or flat (quant_method == "modelopt").
+        # A bare producer check would reject flat configs that omit producer.
+        if not is_modelopt_quant_config(config):
+            raise ValueError(
+                "Expected a modelopt quant config "
+                f"(producer={config.get('producer')}, "
+                f"quant_method={config.get('quant_method')})"
+            )
 
-        quant_config = config.get("quantization", {})
+        # Downstream auto-deploy transforms read field-by-field from a dict
+        # using the legacy field names. Parse via the parallel readers and
+        # re-render as the legacy inner dict for those consumers.
+        parsed = parse_modelopt_quant_config(config)
+        quant_config = parsed.to_legacy_inner_dict()
 
-        quant_algo = quant_config.get("quant_algo", "").upper()
+        quant_algo = (quant_config.get("quant_algo") or "").upper()
 
         if quant_algo == "MIXED_PRECISION":
             self._read_mixed_precision_config(quant_config)