[None][perf] Allocate DSA indexer k-cache only for layers that own an indexer

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -243,18 +243,19 @@ class CacheTransceiver : public BaseCacheTransceiver
         executor::kv_cache::CacheState::AttentionType attentionType
         = executor::kv_cache::CacheState::AttentionType::kDEFAULT,
         std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig = std::nullopt,
-        std::vector<SizeType32> const& rnnLayerNumPerPP = {});
+        std::vector<SizeType32> const& rnnLayerNumPerPP = {}, std::vector<SizeType32> const& indexerLayerNumPerPP = {});
 
     CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector<SizeType32> numKvHeadsPerLayer,
         SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig,
         std::vector<SizeType32> const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType,
         executor::kv_cache::CacheState::AttentionType attentionType
         = executor::kv_cache::CacheState::AttentionType::kDEFAULT,
         std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig = std::nullopt,
-        std::vector<SizeType32> const& rnnLayerNumPerPP = {})
+        std::vector<SizeType32> const& rnnLayerNumPerPP = {}, std::vector<SizeType32> const& indexerLayerNumPerPP = {})
         : CacheTransceiver(cacheManager,
             executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig,
-            attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP)
+            attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP,
+            indexerLayerNumPerPP)
     {
     }
 

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -882,6 +882,7 @@ class WindowBlockManager
         radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr<kvc::BaseLoopbackAgent> loopbackAgent = nullptr,
         bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
         SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
+        std::optional<std::vector<bool>> const& indexerKCacheLayerMask = std::nullopt,
         std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
         SizeType32 numPlaceholderBlocks = 0);
 
@@ -911,7 +912,19 @@ class WindowBlockManager
 
     void releasePools();
 
-    void createIndexerKCachePools();
+    //! \brief Create the indexer K cache pools mirroring each KV pool.
+    //! \details When an indexer layer mask is set, only masked-in layers get a row in the
+    //! indexer pool (masked layout); rows follow the order of `managedLayers` restricted to
+    //! the KV pool's layers, matching the KV pool's own layer-row order.
+    void createIndexerKCachePools(std::vector<SizeType32> const& managedLayers);
+
+    //! \brief Row of `layerIdx` within the masked indexer K cache pool, or -1 when the layer
+    //! is masked out (owns no indexer K cache).
+    [[nodiscard]] SizeType32 getIndexerKCachePoolLayerIdx(SizeType32 layerIdx) const
+    {
+        auto const it = mLayerToIndexerPoolRow.find(layerIdx);
+        return it != mLayerToIndexerPoolRow.end() ? it->second : -1;
+    }
 
     void startScheduling();
 
@@ -1457,6 +1470,14 @@ class WindowBlockManager
     // Whether the indexer K cache stores FP4-packed data (half the byte count
     // per token vs. FP8). Drives the createIndexerKCachePools() formula.
     bool mIndexerKCacheUseFp4{false};
+    // Optional per-layer indexer K cache mask, indexed like numKvHeadsPerLayer.
+    // nullopt = every layer owns an indexer K cache row (dense, legacy behavior).
+    // When set, only masked-in layers get a row in the (masked) indexer pool —
+    // e.g. GLM 5.2 cross-layer indexer sharing where "shared" layers reuse the
+    // previous full layer's top-k and never touch the indexer K cache.
+    std::optional<std::vector<bool>> mIndexerKCacheLayerMask;
+    // layerIdx -> row within the masked indexer K cache pool (only masked-in layers).
+    std::unordered_map<SizeType32, SizeType32> mLayerToIndexerPoolRow;
 
     std::optional<LinearAttentionMetadata> mLinearAttentionMetadata;
 };
@@ -1489,7 +1510,8 @@ class BlockManager
         std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
         std::optional<kvc::BaseAgentConfig> agentConfig = std::nullopt, bool enableIndexerKCache = false,
         SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0,
-        bool indexerKCacheUseFp4 = false, std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
+        bool indexerKCacheUseFp4 = false, std::optional<std::vector<bool>> const& indexerKCacheLayerMask = std::nullopt,
+        std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
         std::vector<PoolConfiguration> const& poolConfigurations = {});
 
     [[nodiscard]] bool isEnableIndexerKCache() const
@@ -1725,6 +1747,13 @@ class BlockManager
         return windowManagerByLayer(layerIdx).getPoolLayerIdx(layerIdx);
     }
 
+    //! \brief Row of `layerIdx` within the masked indexer K cache pool, or -1 when the layer
+    //! owns no indexer K cache (masked out by the per-layer indexer mask).
+    [[nodiscard]] SizeType32 getIndexerKCachePoolLayerIdx(SizeType32 layerIdx) const
+    {
+        return windowManagerByLayer(layerIdx).getIndexerKCachePoolLayerIdx(layerIdx);
+    }
+
     [[nodiscard]] bool isPoolLayerFirst(SizeType32 layerIdx) const
     {
         auto const& manager = windowManagerByLayer(layerIdx);
@@ -2285,6 +2314,7 @@ class KVCacheManager : public BaseKVCacheManager
         std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
         bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
         SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
+        std::optional<std::vector<bool>> const& indexerKCacheLayerMask = std::nullopt,
         std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
         std::vector<PoolConfiguration> const& poolConfigurations = {});
 
@@ -2299,6 +2329,7 @@ class KVCacheManager : public BaseKVCacheManager
         std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
         bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
         SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
+        std::optional<std::vector<bool>> const& indexerKCacheLayerMask = std::nullopt,
         std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
         std::vector<PoolConfiguration> const& poolConfigurations = {});
 
@@ -2313,6 +2344,7 @@ class KVCacheManager : public BaseKVCacheManager
         std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
         bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
         SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
+        std::optional<std::vector<bool>> const& indexerKCacheLayerMask = std::nullopt,
         std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
         std::vector<PoolConfiguration> const& poolConfigurations = {});
 
@@ -2323,6 +2355,7 @@ class KVCacheManager : public BaseKVCacheManager
         CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, bool copyOnpartialReuse = true,
         bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
         SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
+        std::optional<std::vector<bool>> const& indexerKCacheLayerMask = std::nullopt,
         std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
         std::vector<PoolConfiguration> const& poolConfigurations = {});
 

cpp/include/tensorrt_llm/executor/dataTransceiverState.h

@@ -54,7 +54,8 @@ class CacheState final
         std::vector<SizeType32> const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType,
         AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableBlockReuse = false,
         bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0,
-        SizeType32 indexerKCacheQuantBlockSize = 128, bool indexerKCacheUseFp4 = false)
+        SizeType32 indexerKCacheQuantBlockSize = 128, bool indexerKCacheUseFp4 = false,
+        std::vector<SizeType32> const& indexerLayerNumPerPP = {})
         : mModelConfig(std::move(modelConfig))
         , mParallelConfig{worldConfig.getTensorParallelism(), worldConfig.getPipelineParallelism(),
               worldConfig.getContextParallelism(), worldConfig.enableAttentionDP(), worldConfig.getTensorParallelRank(),
@@ -68,6 +69,7 @@ class CacheState final
         mIndexerDimPerHead = indexerDimPerHead;
         mIndexerKCacheQuantBlockSize = indexerKCacheQuantBlockSize;
         mIndexerKCacheUseFp4 = indexerKCacheUseFp4;
+        mIndexerLayerNumPerPP = indexerLayerNumPerPP;
     }
 
     CacheState(std::vector<SizeType32> nbKvHeadPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
@@ -76,7 +78,7 @@ class CacheState final
         AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false,
         int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false,
         bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128,
-        bool indexerKCacheUseFp4 = false)
+        bool indexerKCacheUseFp4 = false, std::vector<SizeType32> const& indexerLayerNumPerPP = {})
         : mModelConfig{std::move(nbKvHeadPerLayer), sizePerHead, tokensPerBlock}
         , mParallelConfig{tensorParallelism, pipelineParallelism, contextParallelism, enableAttentionDP, DPrank, DPsize,
               attentionLayerNumPerPP}
@@ -89,6 +91,7 @@ class CacheState final
         mIndexerDimPerHead = indexerDimPerHead;
         mIndexerKCacheQuantBlockSize = indexerKCacheQuantBlockSize;
         mIndexerKCacheUseFp4 = indexerKCacheUseFp4;
+        mIndexerLayerNumPerPP = indexerLayerNumPerPP;
     }
 
     CacheState(SizeType32 nbAttentionLayers, SizeType32 nbKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
@@ -97,7 +100,7 @@ class CacheState final
         AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false,
         int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false,
         bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128,
-        bool indexerKCacheUseFp4 = false)
+        bool indexerKCacheUseFp4 = false, std::vector<SizeType32> const& indexerLayerNumPerPP = {})
         : mModelConfig{std::vector(nbAttentionLayers, nbKvHeads), sizePerHead, tokensPerBlock}
         , mParallelConfig{tensorParallelism, pipelineParallelism, contextParallelism, enableAttentionDP, DPrank, DPsize,
               attentionLayerNumPerPP}
@@ -110,6 +113,7 @@ class CacheState final
         mIndexerDimPerHead = indexerDimPerHead;
         mIndexerKCacheQuantBlockSize = indexerKCacheQuantBlockSize;
         mIndexerKCacheUseFp4 = indexerKCacheUseFp4;
+        mIndexerLayerNumPerPP = indexerLayerNumPerPP;
     }
 
     [[nodiscard]] bool operator==(kv_cache::CacheState const& other) const noexcept
@@ -120,7 +124,8 @@ class CacheState final
             && mEnablePartialReuse == other.mEnablePartialReuse && mHasIndexerKCache == other.mHasIndexerKCache
             && mIndexerDimPerHead == other.mIndexerDimPerHead
             && mIndexerKCacheQuantBlockSize == other.mIndexerKCacheQuantBlockSize
-            && mIndexerKCacheUseFp4 == other.mIndexerKCacheUseFp4;
+            && mIndexerKCacheUseFp4 == other.mIndexerKCacheUseFp4
+            && mIndexerLayerNumPerPP == other.mIndexerLayerNumPerPP;
     }
 
     struct ModelConfig
@@ -299,6 +304,15 @@ class CacheState final
         return mIndexerKCacheUseFp4;
     }
 
+    //! \brief Per-PP-rank indexer K cache layer counts. With a masked indexer pool (per-layer
+    //! indexer mask, e.g. GLM 5.2 cross-layer indexer sharing) only full-indexer layers own a
+    //! pool row, so these counts can be smaller than the attention layer counts. Falls back to
+    //! the attention layer counts when unset (dense legacy layout).
+    [[nodiscard]] std::vector<SizeType32> const& getIndexerLayerNumPerPP() const
+    {
+        return mIndexerLayerNumPerPP.empty() ? mParallelConfig.mAttentionLayerNumPerPP : mIndexerLayerNumPerPP;
+    }
+
     // =========================================================================
     // RNN/Mamba cache state (optional, present only for hybrid models)
     // =========================================================================
@@ -361,6 +375,12 @@ class CacheState final
         sstring << "indexerDimPerHead:" << mIndexerDimPerHead << "\n";
         sstring << "indexerKCacheQuantBlockSize:" << mIndexerKCacheQuantBlockSize << "\n";
         sstring << "indexerKCacheUseFp4:" << mIndexerKCacheUseFp4 << "\n";
+        sstring << "indexerLayerNumPerPP:";
+        for (auto layerNum : mIndexerLayerNumPerPP)
+        {
+            sstring << layerNum << ",";
+        }
+        sstring << "\n";
         if (mRnnCacheState.has_value())
         {
             auto const& rnn = mRnnCacheState.value();
@@ -404,6 +424,9 @@ class CacheState final
     SizeType32 mIndexerDimPerHead{0};
     SizeType32 mIndexerKCacheQuantBlockSize{128};
     bool mIndexerKCacheUseFp4{false};
+    // Per-PP-rank indexer K cache layer counts (masked indexer pool). Empty = every attention
+    // layer owns an indexer K cache row (dense legacy layout); see getIndexerLayerNumPerPP().
+    std::vector<SizeType32> mIndexerLayerNumPerPP;
     // RNN/Mamba cache state (optional, for hybrid models)
     std::optional<RnnCacheState> mRnnCacheState;
 };

cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp

@@ -214,28 +214,35 @@ size_t CacheTransBufferManager::computeTransferBufferSize(
         SizeType32 kvCacheByteSizePerTokenPerLayer = 0;
         if (transferIndexerKCache)
         {
-            kvCacheByteSizePerTokenPerLayer
-                = cacheManager->getIndexerKCachePool()->getDimension<-1>() * dataSize / tokensPerBlock;
+            auto const indexerPool = cacheManager->getIndexerKCachePool();
+            kvCacheByteSizePerTokenPerLayer = indexerPool->getDimension<-1>() * dataSize / tokensPerBlock;
+            // The (possibly masked) indexer pool holds one row per full-indexer layer; size
+            // the buffer from the pool's actual layer count instead of the attention layer
+            // count.
+            auto const numIndexerLayers = static_cast<SizeType32>(indexerPool->getDimension<1>());
+            auto const validTokenNum = maxNumTokens.value() + tokensPerBlock; // add one more block
+            bufferSizeFromMaxNumToken = validTokenNum * kvCacheByteSizePerTokenPerLayer * numIndexerLayers;
         }
         else
         {
             auto primaryPool = cacheManager->getPrimaryPool(0);
             kvCacheByteSizePerTokenPerLayer
                 = primaryPool->getDimension<-1>() * primaryPool->getDimension<2>() * dataSize / tokensPerBlock;
-        }
-        for (auto layerId = 0; layerId < cacheManager->getBlockManager().getNumLayers(); layerId++)
-        {
-            auto poolIdx = cacheManager->getBlockManager().getLayerPoolIdx(layerId);
-            auto windowSize = static_cast<size_t>(cacheManager->getBlockManager().getPoolWindowSize(poolIdx));
-            auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock;
-            auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value());
-            if (common::getEnvKVCacheTransferAllBlocksForWindow())
+            for (auto layerId = 0; layerId < cacheManager->getBlockManager().getNumLayers(); layerId++)
             {
-                validTokenNum = maxNumTokens.value();
+                auto poolIdx = cacheManager->getBlockManager().getLayerPoolIdx(layerId);
+                auto windowSize = static_cast<size_t>(cacheManager->getBlockManager().getPoolWindowSize(poolIdx));
+                auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock;
+                auto validTokenNum
+                    = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value());
+                if (common::getEnvKVCacheTransferAllBlocksForWindow())
+                {
+                    validTokenNum = maxNumTokens.value();
+                }
+                validTokenNum += tokensPerBlock; // add one more block
+
+                bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer;
             }
-            validTokenNum += tokensPerBlock; // add one more block
-
-            bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer;
         }
     }