[None][feat] Do not REVIEW: Stack PRs for sweep perfing and accuracy checking (v5)

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -872,7 +872,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 +1389,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 +2220,9 @@ 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);
 
     void addSequenceBatch(
         std::vector<std::tuple<LlmRequest::RequestIdType, SizeType32, SizeType32>> const& requestInfos,

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(
@@ -3548,10 +3556,10 @@ void KVCacheManager::addToken(RequestIdType requestId)
     mBlockManager.adjustBlocksIfNeeded(sequence);
 }
 
-void KVCacheManager::copyLinearAttentionBlock(LlmRequest const& llmRequest)
+bool KVCacheManager::copyLinearAttentionBlock(LlmRequest const& llmRequest)
 {
     auto& sequence = getSequence(llmRequest.mRequestId);
-    mBlockManager.copyLinearAttentionBlock(sequence, llmRequest);
+    return mBlockManager.copyLinearAttentionBlock(sequence, llmRequest);
 }
 
 void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence)

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;
             }
 

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)