[https://nvbugs/5501557][fix] Fix nemotron build error

* Root cause is Nemotron-nas model includes some no-op attention layers.

Signed-off-by: Wanli Jiang <[email protected]>

Author: Wanli-Jiang

cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp

@@ -272,8 +272,8 @@ TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr<nvinfer
             auto [numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd] = modelConfig.getNumKvHeadsPerLayerLocalRange(
                 worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank(), isCrossAttention);
             auto numKvHeadsPerLayer = std::vector<SizeType32>(numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd);
-            auto windowSizeLayers
-                = BaseKVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, modelConfig.getNbLayers());
+            auto const numLayers = static_cast<SizeType32>(numKvHeadsPerLayer.size());
+            auto const windowSizeLayers = KVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, numLayers);
             std::map<SizeType32, SizeType32> cacheSizeBytesPerTokenPerWindow;
             for (auto const& [windowSize, managedLayers] : windowSizeLayers)
             {