From fe3ca5e3dbacc74c7458f3d748ac141efb6cb85c Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:27:04 -0700 Subject: [PATCH 1/4] [None][perf] Allocate DSA indexer k-cache only for layers that own an indexer (KVCacheManager V1) GLM 5.2 (glm_moe_dsa) uses DeepSeek Sparse Attention with cross-layer indexer sharing: only 21 of 78 layers own an indexer; the rest reuse the previous full layer's top-k and never touch the indexer k-cache. The V1 KVCacheManager allocated the indexer k-cache stride for every layer, wasting 13.6% of the KV cache footprint (55,224 -> 47,700 B/token, +15.8% effective KV token capacity per rank). Changes: - C++: add an optional per-local-layer indexerKCacheLayerMask to the KVCacheManager/BlockManager/WindowBlockManager ctors; createIndexerKCachePools now builds an indexer pool with one row per masked-in layer (rows follow the KV pool's layer order) and records the layer -> pool-row map (getIndexerKCachePoolLayerIdx). - nanobind: new indexer_k_cache_layer_mask ctor kwarg; get_indexer_k_cache_pool_data translates local layer -> pool row and rejects masked-out layers; get_indexer_k_cache_pool rejects a null pool. - Python: DSACacheManager derives the mask from to_sparse_params(pretrained_config, layer_idx).is_full_indexer_layer (the same source of truth MLA uses to construct Indexer modules, including trailing spec/MTP layers, which are always full), stores None pool entries for shared layers, and asserts if a shared layer requests indexer buffers. Both estimators (get_cache_size_per_token / get_cache_bytes_per_token) charge indexer bytes for full-indexer layers only, PP-sliced consistently. - Disagg is gated with clear fail-fast errors (C++ cache transceiver and Python NIXL extractor) until the transfer layouts understand the per-layer masked pool; dense models (DeepSeek V3.2) are unaffected. - Tests: two C++ pool-shape tests covering the masked and dense-default indexer pool layouts. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../batch_manager/kvCacheManager.h | 37 +++++- .../batch_manager/cacheTransceiver.cpp | 18 +++ .../batch_manager/kvCacheManager.cpp | 68 +++++++++-- .../nanobind/batch_manager/kvCacheManager.cpp | 27 ++++- .../batch_manager/kvCacheManagerTest.cpp | 109 +++++++++++++++++- .../_torch/attention_backend/sparse/dsa.py | 71 ++++++++++-- .../disaggregation/resource/kv_extractor.py | 31 +++++ .../_torch/pyexecutor/resource_manager.py | 14 +++ 8 files changed, 343 insertions(+), 32 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 01adf276f878..183e2b15860e 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -882,6 +882,7 @@ class WindowBlockManager radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr loopbackAgent = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, + std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional 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 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> mIndexerKCacheLayerMask; + // layerIdx -> row within the masked indexer K cache pool (only masked-in layers). + std::unordered_map mLayerToIndexerPoolRow; std::optional mLinearAttentionMetadata; }; @@ -1489,7 +1510,8 @@ class BlockManager std::shared_ptr kvCacheConnectorManager = nullptr, std::optional agentConfig = std::nullopt, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, - bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, + bool indexerKCacheUseFp4 = false, std::optional> const& indexerKCacheLayerMask = std::nullopt, + std::optional linearAttentionMetadata = std::nullopt, std::vector 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 kvCacheConnectorManager = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, + std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}); @@ -2299,6 +2329,7 @@ class KVCacheManager : public BaseKVCacheManager std::shared_ptr kvCacheConnectorManager = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, + std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}); @@ -2313,6 +2344,7 @@ class KVCacheManager : public BaseKVCacheManager std::shared_ptr kvCacheConnectorManager = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, + std::optional> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, std::vector 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> const& indexerKCacheLayerMask = std::nullopt, std::optional linearAttentionMetadata = std::nullopt, std::vector const& poolConfigurations = {}); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f730fb2aaf1e..7ac980883b2a 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -432,6 +432,24 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::make_unique(cacheManager, maxNumTokens)); if (isMLA && cacheManager->isEnableIndexerKCache()) { + // The indexer K cache split/concat path strides indexer blocks by the per-rank + // attention layer count. A masked indexer pool (per-layer indexer mask, e.g. + // GLM 5.2 cross-layer indexer sharing) holds fewer rows than that -- possibly + // zero, in which case no pool exists at all -- which the transfer kernels do + // not understand yet. Fail fast instead of corrupting ctx->gen transfers (or + // dereferencing a null pool in CacheTransBufferManager). + auto const indexerPool = cacheManager->getIndexerKCachePool(); + TLLM_CHECK_WITH_INFO(indexerPool != nullptr, + "The KV cache transceiver does not support a per-layer masked indexer K cache pool yet: " + "this rank owns no indexer K cache rows (every local layer uses cross-layer indexer sharing). Disable " + "the cache transceiver for models with cross-layer indexer sharing (e.g. GLM 5.2)."); + auto const numIndexerLayers = indexerPool->getShape().d[1]; + auto const numLocalLayers = cacheManager->getBlockManager().getNumLayers(); + TLLM_CHECK_WITH_INFO(numIndexerLayers == static_cast(numLocalLayers), + "The KV cache transceiver does not support a per-layer masked indexer K cache pool yet: " + "the indexer pool holds %ld layer rows but this rank manages %d attention layers. Disable the cache " + "transceiver for models with cross-layer indexer sharing (e.g. GLM 5.2).", + static_cast(numIndexerLayers), numLocalLayers); mCacheTransBufferManagers.push_back( std::make_unique(cacheManager, maxNumTokens, true)); } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 266cdafc8d06..30f5b07d31eb 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -593,6 +593,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si std::shared_ptr kvCacheConnectorManager, std::optional agentConfig, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, + std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : mNumLayers{static_cast(numKvHeadsPerLayer.size())} @@ -625,6 +626,11 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si else mLoopbackAgent = nullptr; + TLLM_CHECK_WITH_INFO( + !indexerKCacheLayerMask.has_value() || static_cast(indexerKCacheLayerMask->size()) == mNumLayers, + "indexerKCacheLayerMask must have one entry per layer (expected %d, got %zu)", mNumLayers, + indexerKCacheLayerMask.has_value() ? indexerKCacheLayerMask->size() : 0); + auto const uniqueWindowSizeToLayers = BaseKVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, mNumLayers); @@ -695,6 +701,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si allottedSecondaryBlocks, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, + indexerKCacheLayerMask, LinearAttentionMetadata::hasLinearCache(windowSize) ? linearAttentionMetadata : std::nullopt, numPlaceholderBlocks); } @@ -750,8 +757,8 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 std::shared_ptr kvCacheConnectorManager, radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr loopbackAgent, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, - bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, - SizeType32 numPlaceholderBlocks) + bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, + std::optional linearAttentionMetadata, SizeType32 numPlaceholderBlocks) : mDataType{dtype} , mWindowSize{windowSize} , mNumPrimaryBlocks{blocksInPrimaryPool} @@ -791,6 +798,7 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 , mIndexerKCacheQuantBlockSize{indexerKCacheQuantBlockSize} , mIndexerKCacheIndexHeadDim{indexerKCacheIndexHeadDim} , mIndexerKCacheUseFp4{indexerKCacheUseFp4} + , mIndexerKCacheLayerMask{indexerKCacheLayerMask} , mLinearAttentionMetadata{std::move(linearAttentionMetadata)} { TLLM_LOG_DEBUG("Creating WindowBlockManager for windowSize=%d", windowSize); @@ -849,7 +857,7 @@ WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 if (mEnableIndexerKCache) { - createIndexerKCachePools(); + createIndexerKCachePools(managedLayers); } // Create free blocks @@ -1057,7 +1065,7 @@ void WindowBlockManager::createBlockScalePools(SizeType32 quantBlockSize) } } -void WindowBlockManager::createIndexerKCachePools() +void WindowBlockManager::createIndexerKCachePools(std::vector const& managedLayers) { SizeType32 numPools = mPools.size(); for (SizeType32 i = 0; i < numPools; ++i) @@ -1067,6 +1075,38 @@ void WindowBlockManager::createIndexerKCachePools() { continue; } + // Count the layers of this KV pool that own an indexer K cache and assign each its + // pool row. Rows follow managedLayers order restricted to this pool, matching the + // KV pool's own layer-row order (mLayerToIndexWithinPool). Without a mask every layer + // owns a row and the indexer pool mirrors the KV pool layer count (dense, legacy + // layout); with a mask (e.g. GLM 5.2 cross-layer indexer sharing) only full-indexer + // layers get a row, so shared layers cost no indexer pool memory. + SizeType32 numIndexerLayers = 0; + for (auto const layerIdx : managedLayers) + { + if (mLayerToPoolIndex.at(layerIdx) != i) + { + continue; + } + if (mIndexerKCacheLayerMask.has_value() && !mIndexerKCacheLayerMask->at(layerIdx)) + { + continue; + } + mLayerToIndexerPoolRow[layerIdx] = numIndexerLayers++; + } + if (numIndexerLayers == 0) + { + TLLM_LOG_WARNING( + "[%s] Indexer K cache is enabled but every layer of pool %d is masked out; skipping indexer " + "K cache pool creation for this pool.", + mLogPrefix.c_str(), i); + continue; + } + if (numIndexerLayers < kvPool.numLayers) + { + TLLM_LOG_INFO("[%s] Indexer K cache pool: %d of %d layers own an indexer K cache.", mLogPrefix.c_str(), + numIndexerLayers, kvPool.numLayers); + } // scaleSize evaluates to 4 at indexHeadDim=128 / quantBlockSize=128 for both // FP8 (float32 scale) and FP4 (packed UE8M0 x4 stored as one int32). The // data size is either indexHeadDim bytes (one FP8 byte per element) or @@ -1075,7 +1115,7 @@ void WindowBlockManager::createIndexerKCachePools() SizeType32 dataSize = mIndexerKCacheUseFp4 ? mIndexerKCacheIndexHeadDim / 2 : mIndexerKCacheIndexHeadDim; SizeType32 perTokenSize = scaleSize + dataSize; - mPools.emplace_back(kvPool.numLayers, kvPool.kvFactor, 1, perTokenSize, kvPool.tokensPerBlock, + mPools.emplace_back(numIndexerLayers, kvPool.kvFactor, 1, perTokenSize, kvPool.tokensPerBlock, /*primaryPool=*/nullptr, /*secondaryPool=*/nullptr, /*containsBlockScales=*/false, @@ -3234,14 +3274,15 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, - bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, + bool indexerKCacheUseFp4, std::optional> const& indexerKCacheLayerMask, + std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, std::nullopt, nullptr, enablePartialReuse, copyOnPartialReuse, nullptr, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, - linearAttentionMetadata, poolConfigurations) + indexerKCacheLayerMask, linearAttentionMetadata, poolConfigurations) { } @@ -3253,6 +3294,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, + std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, @@ -3260,7 +3302,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::make_shared(reinterpret_cast(stream)), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, secondaryOffloadMinPriority, eventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, - indexerKCacheUseFp4, linearAttentionMetadata, poolConfigurations) + indexerKCacheUseFp4, indexerKCacheLayerMask, linearAttentionMetadata, poolConfigurations) { } @@ -3272,6 +3314,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, + std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : mMaxBeamWidth(maxBeamWidth) @@ -3285,8 +3328,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::move(stream), maxSequenceLength, maxBeamWidth, maxAttentionWindowVec, dtype, mSinkBubbleLength, mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), std::nullopt, enableIndexerKCache, - indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata, - poolConfigurations) + indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, indexerKCacheLayerMask, + linearAttentionMetadata, poolConfigurations) // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { @@ -3314,14 +3357,15 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, + std::optional> const& indexerKCacheLayerMask, std::optional linearAttentionMetadata, std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::move(stream), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), enableIndexerKCache, - indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata, - poolConfigurations) + indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, indexerKCacheLayerMask, + linearAttentionMetadata, poolConfigurations) { } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index b1c4391c0e26..559934e1ef72 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -551,13 +551,31 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) "get_indexer_k_cache_pool_data", [](tbk::BaseKVCacheManager& self, SizeType32 layer_idx) -> at::Tensor { + // Translate the (local) layer index to its row within the masked indexer + // K cache pool. Without a per-layer indexer mask the mapping is the + // identity for a single-pool manager (dense, legacy layout). + auto const row = self.getBlockManager().getIndexerKCachePoolLayerIdx(layer_idx); + TLLM_CHECK_WITH_INFO(row >= 0, + "Layer %d owns no indexer K cache (masked out by the per-layer indexer mask).", layer_idx); auto pool = tr::Torch::tensor(self.getIndexerKCachePool()); - return pool.index({torch::indexing::Slice(), layer_idx}); + return pool.index({torch::indexing::Slice(), row}); }, nb::call_guard()) + .def( + "get_indexer_k_cache_pool_layer_idx", + [](tbk::BaseKVCacheManager& self, SizeType32 layer_idx) -> SizeType32 + { return self.getBlockManager().getIndexerKCachePoolLayerIdx(layer_idx); }, + nb::call_guard()) .def( "get_indexer_k_cache_pool", - [](tbk::BaseKVCacheManager& self) -> at::Tensor { return tr::Torch::tensor(self.getIndexerKCachePool()); }, + [](tbk::BaseKVCacheManager& self) -> at::Tensor + { + auto const pool = self.getIndexerKCachePool(); + TLLM_CHECK_WITH_INFO(pool != nullptr, + "No indexer K cache pool exists on this rank (indexer K cache disabled, or every local layer " + "is masked out by the per-layer indexer mask)."); + return tr::Torch::tensor(pool); + }, nb::call_guard()) .def( "get_unique_primary_pool", [](tbk::BaseKVCacheManager& self) { return self.getUniquePrimaryPool(); }, @@ -665,8 +683,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) std::vector const&, tensorrt_llm::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, - bool, SizeType32, SizeType32, bool, std::optional, - std::vector const&>(), + bool, SizeType32, SizeType32, bool, std::optional> const&, + std::optional, std::vector const&>(), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"), nb::arg("blocks_per_window"), nb::arg("max_num_sequences"), nb::arg("max_beam_width"), nb::arg("max_attention_window_vec"), nb::arg("dtype"), nb::arg("sink_token_length"), nb::arg("stream"), @@ -676,6 +694,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("copy_on_partial_reuse") = true, nb::arg("kv_connector_manager") = nullptr, nb::arg("enable_indexer_k_cache") = false, nb::arg("indexer_k_cache_quant_block_size") = 128, nb::arg("indexer_k_cache_index_head_dim") = 0, nb::arg("indexer_k_cache_use_fp4") = false, + nb::arg("indexer_k_cache_layer_mask").none() = std::nullopt, nb::arg("linear_attention_metadata").none() = std::nullopt, nb::arg("pool_configurations") = std::vector{}, nb::call_guard()) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 98c2232b4062..b5cbe90b4de3 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -836,7 +836,8 @@ TEST_F(KVCacheManagerTest, FP4AttentionWithHalfRecurrentStatesPoolTest) /*sinkTokenLength=*/0, stream, maxAttentionWindow, /*chunkSize=*/0, /*enableBlockReuse=*/false, CacheType::kSELF, std::nullopt, nullptr, /*enablePartialReuse=*/false, /*copyOnPartialReuse=*/true, nullptr, /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, - /*indexerKCacheUseFp4=*/false, linearAttentionMetadata, poolConfigurations); + /*indexerKCacheUseFp4=*/false, /*indexerKCacheLayerMask=*/std::nullopt, linearAttentionMetadata, + poolConfigurations); kvCacheManager.allocatePools(/*useUvm=*/false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -887,6 +888,100 @@ TEST_F(KVCacheManagerTest, FP4AttentionWithHalfRecurrentStatesPoolTest) } #endif +TEST_F(KVCacheManagerTest, IndexerKCachePoolLayerMaskTest) +{ + // Masked indexer K cache pool: only masked-in (full-indexer) layers own a + // row, mirroring GLM 5.2 cross-layer indexer sharing (freq=4, offset=2 over + // 6 layers -> full layers {0, 1, 5}). + auto constexpr numLayers = 6; + auto constexpr numHeads = 1; + auto constexpr sizePerHead = 576; + auto constexpr tokensPerBlock = 4; + auto constexpr maxBlocksPerSeq = 4; + auto constexpr maxNumSequences = 8; + auto constexpr blocksInPrimaryPool = 16; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr indexerKCacheIndexHeadDim = 128; + auto constexpr indexerKCacheQuantBlockSize = 128; + + auto const stream = std::make_shared(); + auto constexpr beamWidth = 1; + auto const maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + auto const indexerLayerMask = std::vector{true, true, false, false, false, true}; + + KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, + /*sinkTokenLength=*/0, stream, maxAttentionWindow, maxAttentionWindow, /*enableBlockReuse=*/false, + CacheType::kSELFKONLY, std::nullopt, nullptr, /*enablePartialReuse=*/true, /*copyOnpartialReuse=*/true, nullptr, + /*enableIndexerKCache=*/true, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, + /*indexerKCacheUseFp4=*/false, indexerLayerMask); + + kvCacheManager.allocatePools(/*useUvm=*/false); + auto const& blockManager = kvCacheManager.getBlockManager(); + + // One KV pool plus one (masked) indexer pool. + EXPECT_EQ(blockManager.getNumPools(), 2); + EXPECT_EQ(blockManager.getNumPools(/*includeBlockScalePools=*/true, /*includeIndexerKCachePools=*/false), 1); + + // The indexer pool holds one row per masked-in layer only. + auto const indexerPool = kvCacheManager.getIndexerKCachePool(); + ASSERT_NE(indexerPool, nullptr); + EXPECT_EQ(indexerPool->getShape().d[1], 3); + + // Layer -> pool-row mapping: full layers get consecutive rows in layer + // order, shared layers map to -1. + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(0), 0); + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(1), 1); + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(2), -1); + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(3), -1); + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(4), -1); + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(5), 2); + + // Indexer pool block size: 1 head x (dataBytes + scaleBytes) x tokensPerBlock. + auto constexpr perTokenBytes + = indexerKCacheIndexHeadDim + indexerKCacheIndexHeadDim / indexerKCacheQuantBlockSize * 4; + EXPECT_EQ(blockManager.getBlockSize(1), perTokenBytes * tokensPerBlock); +} + +TEST_F(KVCacheManagerTest, IndexerKCachePoolDenseDefaultTest) +{ + // Without a per-layer indexer mask the indexer pool keeps the dense legacy + // layout: one row per layer, identity layer -> row mapping. + auto constexpr numLayers = 4; + auto constexpr numHeads = 1; + auto constexpr sizePerHead = 576; + auto constexpr tokensPerBlock = 4; + auto constexpr maxBlocksPerSeq = 4; + auto constexpr maxNumSequences = 8; + auto constexpr blocksInPrimaryPool = 16; + auto constexpr blocksInSecondaryPool = 0; + + auto const stream = std::make_shared(); + auto constexpr beamWidth = 1; + auto const maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; + + KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, + /*sinkTokenLength=*/0, stream, maxAttentionWindow, maxAttentionWindow, /*enableBlockReuse=*/false, + CacheType::kSELFKONLY, std::nullopt, nullptr, /*enablePartialReuse=*/true, /*copyOnpartialReuse=*/true, nullptr, + /*enableIndexerKCache=*/true, /*indexerKCacheQuantBlockSize=*/128, + /*indexerKCacheIndexHeadDim=*/128, /*indexerKCacheUseFp4=*/false); + + kvCacheManager.allocatePools(/*useUvm=*/false); + auto const& blockManager = kvCacheManager.getBlockManager(); + + auto const indexerPool = kvCacheManager.getIndexerKCachePool(); + ASSERT_NE(indexerPool, nullptr); + EXPECT_EQ(indexerPool->getShape().d[1], numLayers); + for (SizeType32 layerIdx = 0; layerIdx < numLayers; ++layerIdx) + { + EXPECT_EQ(blockManager.getIndexerKCachePoolLayerIdx(layerIdx), layerIdx); + } +} + TEST_F(KVCacheManagerTest, BlockManagerReuseTest) { auto constexpr numLayers = 12; @@ -7546,7 +7641,7 @@ void testBlockManagerLinearAttention_ContextNoReuse(int beamWidth, int numTokens std::vector{linearWindowSizeCode, maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, /*chunkSize*/ 0, CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, std::nullopt, false, 128, 0, - false, linearAttentionMetadata); + false, /*indexerKCacheLayerMask=*/std::nullopt, linearAttentionMetadata); blockManager.allocatePools(false); ASSERT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -7692,7 +7787,7 @@ void testBlockManagerLinearAttention_ContextReuse(int beamWidth, int numTokens0, std::vector{linearWindowSizeCode, maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, /*chunkSize*/ 0, CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, std::nullopt, false, 128, 0, - false, linearAttentionMetadata); + false, /*indexerKCacheLayerMask=*/std::nullopt, linearAttentionMetadata); blockManager.allocatePools(false); auto inputTokens0 = std::make_shared(); @@ -7930,6 +8025,7 @@ void testKVCacheManagerLinearAttention_DecodingBlockGrowth( /*indexerKCacheQuantBlockSize*/ 128, /*indexerKCacheIndexHeadDim*/ 0, /*indexerKCacheUseFp4=*/false, + /*indexerKCacheLayerMask=*/std::nullopt, /*linearAttentionMetadata*/ linearAttentionMetadata); auto inputTokens0 = std::make_shared(); @@ -8027,7 +8123,8 @@ void testKVCacheManagerLinearAttention_BlockCopying( KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{linearWindowSizeCode, maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLen, stream, maxAttentionWindow, /*chunkSize*/ 0, enableContextReuse, - CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, false, 128, 0, false, linearAttentionMetadata); + CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, false, 128, 0, false, + /*indexerKCacheLayerMask=*/std::nullopt, linearAttentionMetadata); kvCacheManager.allocatePools(false); char* poolBaseAddr @@ -10558,7 +10655,7 @@ TEST_F(KVCacheManagerTest, VswaMixedHeadDimReuseSmoke) /*enablePartialReuse=*/true, /*copyOnpartialReuse=*/true, /*kvCacheConnectorManager=*/nullptr, /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, - /*indexerKCacheUseFp4=*/false, + /*indexerKCacheUseFp4=*/false, /*indexerKCacheLayerMask=*/std::nullopt, /*linearAttentionMetadata=*/std::nullopt, poolConfigurations); kvCacheManager.allocatePools(/*useUvm=*/false); @@ -10659,7 +10756,7 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) /*enablePartialReuse=*/true, /*copyOnpartialReuse=*/true, /*kvCacheConnectorManager=*/nullptr, /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, - /*indexerKCacheUseFp4=*/false, + /*indexerKCacheUseFp4=*/false, /*indexerKCacheLayerMask=*/std::nullopt, /*linearAttentionMetadata=*/std::nullopt, poolConfigurations); kvCacheManager->allocatePools(/*useUvm=*/false); diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index af85390cf72e..5c12e6e3c12d 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -3100,6 +3100,21 @@ def mla_rope_append_paged_kv_assign_q( ) +def derive_indexer_k_cache_layer_mask( + sparse_attention_config: "SparseAttentionConfig", + pretrained_config, + num_layers: int, +) -> List[bool]: + return [ + bool( + getattr( + sparse_attention_config.to_sparse_params( + pretrained_config=pretrained_config, layer_idx=layer_idx), + "is_full_indexer_layer", True)) + for layer_idx in range(num_layers) + ] + + class DSACacheManager(KVCacheManager): """KV cache manager for DSA with additional indexer K-cache pools.""" @@ -3146,6 +3161,14 @@ def __init__( # allocates the pool with this smaller stride when the flag is set. self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" + from tensorrt_llm._torch.speculative import get_num_spec_layers + total_num_layers = (len(layer_mask) + if layer_mask is not None else num_layers) + if spec_config is not None and layer_mask is None: + total_num_layers += get_num_spec_layers(spec_config) + indexer_k_cache_layer_mask = derive_indexer_k_cache_layer_mask( + sparse_attention_config, pretrained_config, total_num_layers) + super().__init__( kv_cache_config, kv_cache_type, @@ -3166,6 +3189,7 @@ def __init__( indexer_k_cache_quant_block_size=128, indexer_k_cache_index_head_dim=self.index_head_dim, indexer_k_cache_use_fp4=self.use_fp4, + indexer_k_cache_layer_mask=indexer_k_cache_layer_mask, **kwargs, ) self.num_blocks = self.blocks_in_primary_pool @@ -3173,11 +3197,20 @@ def __init__( # Indexer K cache pool for DSA attention # Shape: [num_blocks, self.tokens_per_block * (index_head_dim + scale_size)] # Non-interleaved layout: [fp8_tok0 | fp8_tok1 | ... | scale_tok0 | scale_tok1 | ...] - # Store FP8-quantized k values from the indexer + # Store FP8-quantized k values from the indexer. + # One entry per local layer; None for shared-indexer layers, which own + # no row in the masked indexer pool (the C++ binding raises if asked). + local_mask = self.indexer_k_cache_local_layer_mask self.indexer_k_cache_pool_per_layer = [ - self.get_indexer_k_cache_pool_data(layer_idx) - for layer_idx in range(self.num_local_layers) + self.get_indexer_k_cache_pool_data(local_offset) + if local_mask[local_offset] else None + for local_offset in range(self.num_local_layers) ] + num_full = sum(local_mask) + if num_full < self.num_local_layers: + logger.info( + f"[DSACacheManager] Indexer k-cache: {num_full} of " + f"{self.num_local_layers} local layers own an indexer k-cache.") def get_indexer_k_cache_buffers(self, layer_idx: int): """Get indexer k cache buffer from a specific layer pool.""" @@ -3185,8 +3218,11 @@ def get_indexer_k_cache_buffers(self, layer_idx: int): data_bytes = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim per_token_size = data_bytes + self.index_head_dim // self.quant_block_size * 4 layer_offset = self.layer_offsets[layer_idx] - return self.indexer_k_cache_pool_per_layer[layer_offset].view( - self.num_blocks, block_size, 1, per_token_size) + pool = self.indexer_k_cache_pool_per_layer[layer_offset] + assert pool is not None, ( + f"Layer {layer_idx} is a shared-indexer layer and owns no indexer " + f"k-cache; only full-indexer layers may access it.") + return pool.view(self.num_blocks, block_size, 1, per_token_size) def get_batch_indexer_k_cache_indices( self, request_ids: List[int]) -> List[List[int]]: @@ -3241,12 +3277,23 @@ def get_cache_size_per_token(model_config: ModelConfig, # MLA latent K cache: stored at the KV cache dtype (BF16/FP8). mem_per_token *= num_attention_layers * head_dim + if num_layers is not None: + num_indexer_layers = max(num_layers, 1) + else: + local_layer_ids = mapping.pp_layers( + model_config.get_num_attention_layers()) + num_indexer_layers = sum( + 1 for layer_id in local_layer_ids if getattr( + sparse_attention_config.to_sparse_params( + pretrained_config=config, layer_idx=layer_id), + "is_full_indexer_layer", True)) + # Indexer K cache: physically allocated as raw UINT8 in # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike # the latent above). The data-portion byte count already reflects fp8 vs # fp4 via indexer_data_dim. - indexer_bytes_per_token = num_attention_layers * ( + indexer_bytes_per_token = num_indexer_layers * ( indexer_data_dim + index_head_dim // quant_block_size * 4) mem_per_token += indexer_bytes_per_token return mem_per_token @@ -3274,9 +3321,17 @@ def get_cache_bytes_per_token(self): # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike # the latent above). Under FP4 the indexer data portion is halved (two - # E2M1 codes per byte); the scale bytes are unchanged. + # E2M1 codes per byte); the scale bytes are unchanged. Only + # full-indexer local layers own a row in the masked indexer pool, so + # shared layers (cross-layer indexer sharing) contribute no bytes. indexer_data_dim = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim - indexer_bytes_per_token = sum(self.num_kv_heads_per_layer) * ( + local_mask = self.indexer_k_cache_local_layer_mask + if local_mask is not None: + num_indexer_layers = sum(kv_heads for kv_heads, has_indexer in zip( + self.num_kv_heads_per_layer, local_mask) if has_indexer) + else: + num_indexer_layers = sum(self.num_kv_heads_per_layer) + indexer_bytes_per_token = num_indexer_layers * ( indexer_data_dim + self.index_head_dim // self.quant_block_size * 4) cache_size_bytes_per_token += indexer_bytes_per_token diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 92f0a9c2f4ad..b7c18f419ebc 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -199,9 +199,40 @@ def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: # Indexer K cache support if getattr(kv_cache_manager, "enable_indexer_k_cache", False): + # The FLAT PoolView contract assumes the pool covers ALL layers of + # the layer group with a uniform per-layer stride. A masked + # indexer pool (per-layer indexer mask, e.g. GLM 5.2 cross-layer + # indexer sharing) holds fewer rows than layers — possibly zero, + # in which case no pool exists at all — so the stride math would + # corrupt transfers. Fail fast until the extractor publishes + # per-layer buffer_entries for the indexer pool. Check the mask + # BEFORE fetching the pool: a fully masked-out rank has no pool + # to fetch. + local_indexer_mask = getattr(kv_cache_manager, "indexer_k_cache_local_layer_mask", None) + if local_indexer_mask is not None and not all( + local_indexer_mask[lid] for lid in local_layer_ids + ): + raise NotImplementedError( + "Disaggregated KV transfer does not support a per-layer " + "masked indexer k-cache pool yet: " + f"{sum(local_indexer_mask[lid] for lid in local_layer_ids)}" + f" of {len(local_layer_ids)} layers in this layer group " + "own an indexer k-cache. Disable disaggregated serving " + "for models with cross-layer indexer sharing (e.g. " + "GLM 5.2)." + ) indexer_pool = kv_cache_manager.impl.get_indexer_k_cache_pool() # indexer_pool shape: (numBlocks, numLayers, kvFactor, blockSize), dtype=UINT8 # slot_bytes = numLayers * kvFactor * blockSize * element_size + if indexer_pool.shape[1] != len(local_layer_ids): + raise NotImplementedError( + "Disaggregated KV transfer does not support a per-layer " + "masked indexer k-cache pool yet: the indexer " + f"pool holds {indexer_pool.shape[1]} layer rows but the " + f"layer group has {len(local_layer_ids)} layers. Disable " + "disaggregated serving for models with cross-layer " + "indexer sharing (e.g. GLM 5.2)." + ) per_block_elems = 1 for d in indexer_pool.shape[1:]: # skip numBlocks dim per_block_elems *= d diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 157e2d7d720a..11fa2eb8e018 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -291,6 +291,8 @@ def __init__( indexer_k_cache_quant_block_size: int = 128, indexer_k_cache_index_head_dim: int = 0, indexer_k_cache_use_fp4: bool = False, + # Boolean array indicating whether each layer has an indexer K cache. + indexer_k_cache_layer_mask: Optional[List[bool]] = None, is_estimating_kv_cache: bool = False, execution_stream: Optional[torch.cuda.Stream] = None, linear_attention_metadata: Optional[LinearAttentionMetadata] = None, @@ -321,6 +323,17 @@ def __init__( for offset, idx in enumerate(self.pp_layers) } + if indexer_k_cache_layer_mask is not None: + assert len(indexer_k_cache_layer_mask) >= max(self.pp_layers) + 1, ( + f"indexer_k_cache_layer_mask covers " + f"{len(indexer_k_cache_layer_mask)} layers but this rank " + f"manages global layer {max(self.pp_layers)}") + self.indexer_k_cache_local_layer_mask: Optional[List[bool]] = [ + bool(indexer_k_cache_layer_mask[i]) for i in self.pp_layers + ] + else: + self.indexer_k_cache_local_layer_mask = None + self.kv_connector_manager = kv_connector_manager tp_size = mapping.tp_size @@ -614,6 +627,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], indexer_k_cache_quant_block_size, 'indexer_k_cache_index_head_dim': indexer_k_cache_index_head_dim, 'indexer_k_cache_use_fp4': indexer_k_cache_use_fp4, + 'indexer_k_cache_layer_mask': self.indexer_k_cache_local_layer_mask, 'linear_attention_metadata': linear_attention_metadata, # Forward the (possibly remapped) per-pool configurations. # window_size values are aligned with the post-clamp sizes. From 779bc7a60c52970fdaa5fbb4c5de4c1c25518fa9 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:27:42 -0700 Subject: [PATCH 2/4] [None][feat] Support the per-layer masked DSA indexer k-cache pool in the C++ cache transceiver Follow-up to the per-layer masked indexer k-cache pool: ctx->gen disaggregated transfer of the indexer k-cache now understands the masked layout (only full-indexer layers own a pool row, e.g. GLM 5.2 cross-layer indexer sharing) instead of failing fast. - CacheState carries per-PP indexer layer counts (indexerLayerNumPerPP; empty = dense fallback to the attention counts), wired through the ctors, operator==, toString, and serialization. - New targetIRanksForIndexerKCache keeps the attention pass's rank topology but recomputes per-peer layer counts in indexer layer space (interval intersections of the per-PP indexer prefix sums). splitKVCache/concatKVCache use it, plus the self indexer count as the kernel layer count, when transferring the indexer pool. - MLACacheFormatter sizes indexer-pass buffers with the indexer counts, skips zero-sized receive targets, and inquireSupport validates indexer layout equality, equal indexer layer totals, and rejects PP partitions that would produce zero-sized indexer exchanges (rank or overlap without a full-indexer layer). - The NIXL/Mooncake agent path computes a dedicated indexer-pass write-offset ratio in indexer layer space (previously shared the attention-space KV ratio, which would corrupt masked-pool transfers under PP resharding). - CacheTransceiver takes indexerLayerNumPerPP, validates it against the local pool row count, and the Python transceiver gathers it over PP from the manager's indexer layer mask. - CacheTransBufferManager sizes the indexer transfer buffer from the pool's actual layer count. - Tests: CacheState serialization round-trip of the new field (serializeUtilsTest) and a multi-GPU indexer-layer-mask transceiver test with ctx/gen PP resharding (cacheTransceiverTest). Also fixes a positional off-by-one in the transceiver test's CacheState construction (isIndexerKCache previously landed on enablePartialReuse). Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 7 +- .../executor/dataTransceiverState.h | 31 ++- .../batch_manager/cacheTransBuffer.cpp | 35 ++-- .../batch_manager/cacheTransceiver.cpp | 39 ++-- .../batch_manager/mlaCacheFormatter.cpp | 136 +++++++++++-- .../agent_utils/connection.cpp | 20 +- .../cache_transmission/cacheSplitConcat.cu | 65 +++++- .../cache_transmission/cacheSplitConcat.h | 3 + cpp/tensorrt_llm/executor/serialization.cpp | 5 +- .../batch_manager/cacheTransceiver.cpp | 5 +- .../executor/serializeUtilsTest.cpp | 9 +- .../multi_gpu/cacheTransceiverTest.cpp | 187 +++++++++++++++++- .../disaggregation/resource/kv_extractor.py | 17 +- .../_torch/pyexecutor/kv_cache_transceiver.py | 28 ++- 14 files changed, 498 insertions(+), 89 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 9d28fa26c4ee..44025f7e9277 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -243,7 +243,7 @@ class CacheTransceiver : public BaseCacheTransceiver executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, - std::vector const& rnnLayerNumPerPP = {}); + std::vector const& rnnLayerNumPerPP = {}, std::vector const& indexerLayerNumPerPP = {}); CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, @@ -251,10 +251,11 @@ class CacheTransceiver : public BaseCacheTransceiver executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, - std::vector const& rnnLayerNumPerPP = {}) + std::vector const& rnnLayerNumPerPP = {}, std::vector const& indexerLayerNumPerPP = {}) : CacheTransceiver(cacheManager, executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig, - attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP) + attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP, + indexerLayerNumPerPP) { } diff --git a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h index b00d44d129e7..7c56eca31032 100644 --- a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h +++ b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h @@ -54,7 +54,8 @@ class CacheState final std::vector 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 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 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 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 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 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 mIndexerLayerNumPerPP; // RNN/Mamba cache state (optional, for hybrid models) std::optional mRnnCacheState; }; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index e06198f9ea60..20ce219c6302 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/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(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(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(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; } } diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 7ac980883b2a..fc4a51ea2159 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -348,7 +348,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, - std::vector const& rnnLayerNumPerPP) + std::vector const& rnnLayerNumPerPP, std::vector const& indexerLayerNumPerPP) : mCacheTransceiverConfig{cacheTransceiverConfig} { using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter; @@ -402,11 +402,11 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa { kvFactor = 1; } - mCacheState - = std::make_unique(cacheStateModelCfg, worldConfig, attentionLayerNumPerPP, - dataType, attentionType, kvFactor, cacheManager->isEnableBlockReuse(), cacheManager->isEnablePartialReuse(), - cacheManager->isEnableIndexerKCache(), cacheManager->getIndexerKCacheIndexHeadDim(), - cacheManager->getIndexerKCacheQuantBlockSize(), cacheManager->getIndexerKCacheUseFp4()); + mCacheState = std::make_unique(cacheStateModelCfg, worldConfig, + attentionLayerNumPerPP, dataType, attentionType, kvFactor, cacheManager->isEnableBlockReuse(), + cacheManager->isEnablePartialReuse(), cacheManager->isEnableIndexerKCache(), + cacheManager->getIndexerKCacheIndexHeadDim(), cacheManager->getIndexerKCacheQuantBlockSize(), + cacheManager->getIndexerKCacheUseFp4(), indexerLayerNumPerPP); if (mCacheState->getParallelConfig().mEnableAttentionDP) { @@ -432,24 +432,21 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::make_unique(cacheManager, maxNumTokens)); if (isMLA && cacheManager->isEnableIndexerKCache()) { - // The indexer K cache split/concat path strides indexer blocks by the per-rank - // attention layer count. A masked indexer pool (per-layer indexer mask, e.g. - // GLM 5.2 cross-layer indexer sharing) holds fewer rows than that -- possibly - // zero, in which case no pool exists at all -- which the transfer kernels do - // not understand yet. Fail fast instead of corrupting ctx->gen transfers (or - // dereferencing a null pool in CacheTransBufferManager). + // The advertised per-PP indexer layer counts must match the local pool. A rank + // without any indexer rows is not supported by the transfer machinery, so reject + // it up front. auto const indexerPool = cacheManager->getIndexerKCachePool(); TLLM_CHECK_WITH_INFO(indexerPool != nullptr, - "The KV cache transceiver does not support a per-layer masked indexer K cache pool yet: " - "this rank owns no indexer K cache rows (every local layer uses cross-layer indexer sharing). Disable " - "the cache transceiver for models with cross-layer indexer sharing (e.g. GLM 5.2)."); + "The KV cache transceiver does not support a rank without indexer K cache rows (every local layer " + "uses cross-layer indexer sharing)."); + auto const selfPPRank = worldConfig.getPipelineParallelRank(); + auto const advertisedIndexerLayerNum = indexerLayerNumPerPP.empty() ? attentionLayerNumPerPP.at(selfPPRank) + : indexerLayerNumPerPP.at(selfPPRank); auto const numIndexerLayers = indexerPool->getShape().d[1]; - auto const numLocalLayers = cacheManager->getBlockManager().getNumLayers(); - TLLM_CHECK_WITH_INFO(numIndexerLayers == static_cast(numLocalLayers), - "The KV cache transceiver does not support a per-layer masked indexer K cache pool yet: " - "the indexer pool holds %ld layer rows but this rank manages %d attention layers. Disable the cache " - "transceiver for models with cross-layer indexer sharing (e.g. GLM 5.2).", - static_cast(numIndexerLayers), numLocalLayers); + TLLM_CHECK_WITH_INFO(numIndexerLayers == static_cast(advertisedIndexerLayerNum), + "The indexer K cache pool holds %ld layer rows but indexerLayerNumPerPP advertises %d for PP rank " + "%d. Pass the per-PP indexer layer counts matching the (masked) indexer pool layout.", + static_cast(numIndexerLayers), advertisedIndexerLayerNum, selfPPRank); mCacheTransBufferManagers.push_back( std::make_unique(cacheManager, maxNumTokens, true)); } diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 24b7182fe628..c042e59329cb 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -30,9 +30,11 @@ #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include #include #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager { @@ -225,7 +227,13 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses return; } - auto targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx); + // The indexer K cache pass shares the attention pass's connections/topology but runs in + // "indexer layer space": with a masked indexer pool only full-indexer layers own a row, + // so both the self layer count and the per-peer layer counts come from the indexer + // layer counts (dense models fall back to the attention counts). + auto targetInfo = transferIndexerKCache + ? executor::kv_cache::targetIRanksForIndexerKCache(destConfig, selfConfig, selfIdx) + : executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx); size_t pPDomainSize = targetInfo.mDomainPPSize; size_t cPDomainSize = targetInfo.mDomainCPSize; @@ -234,20 +242,22 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses auto const ppRank = selfIdx / (selfConfig.getParallelConfig().mTensorParallelism * selfConfig.getParallelConfig().mContextParallelism); - auto const selfAttentionLayerNum = selfConfig.getParallelConfig().mAttentionLayerNumPerPP.at(ppRank); + auto const selfLayerNum = transferIndexerKCache + ? selfConfig.getIndexerLayerNumPerPP().at(ppRank) + : selfConfig.getParallelConfig().mAttentionLayerNumPerPP.at(ppRank); auto const cacheBlockSize = inputKvCacheBlocks.at(0)->getSize(); - auto const blockSizePerLayer = cacheBlockSize / selfAttentionLayerNum; + auto const blockSizePerLayer = cacheBlockSize / selfLayerNum; std::vector bufferSizeForTarget(pPDomainSize * cPDomainSize, 0); for (size_t ppDomainIdx = 0; ppDomainIdx < pPDomainSize; ppDomainIdx++) { - auto const peerAttentionLayerNum = targetInfo.getPeerPPDomainLayerNum(ppDomainIdx); + auto const peerLayerNum = targetInfo.getPeerPPDomainLayerNum(ppDomainIdx); for (size_t cpDomainIdx = 0; cpDomainIdx < cPDomainSize; cpDomainIdx++) { auto const idx = cpDomainIdx * pPDomainSize + ppDomainIdx; // Note: contextCP is always 1. So, cpDomainSize == genCPSize and cpDomainIdx == genCPRank. auto const peerBlockNum = executor::kv_cache::getBlockNumAccountingForCP(cpDomainIdx, cPDomainSize, blockNum); - bufferSizeForTarget[idx] = blockSizePerLayer * peerAttentionLayerNum * peerBlockNum; + bufferSizeForTarget[idx] = blockSizePerLayer * peerLayerNum * peerBlockNum; } } return bufferSizeForTarget; @@ -542,20 +552,28 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s auto getBufferSizeForTarget = [&]() { - auto const targetInfo = executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx); + // Indexer K cache pass: sizes are computed in "indexer layer space" (masked + // pool: only full-indexer layers own a row); dense models fall back to the + // attention counts. Peers whose overlap holds no full-indexer layer get a + // zero-sized buffer and the receive is skipped, mirroring the sender. + auto const targetInfo = transferIndexerKCache + ? executor::kv_cache::targetIRanksForIndexerKCache(destConfig, selfConfig, selfIdx) + : executor::kv_cache::targetIRanks(destConfig, selfConfig, selfIdx); auto const cacheBlockSize = outputBuffers.at(0)->getSize(); auto const ppRank = selfIdx / (selfConfig.getParallelConfig().mTensorParallelism * selfConfig.getParallelConfig().mContextParallelism); - auto const selfAttentionLayerNum = selfConfig.getParallelConfig().mAttentionLayerNumPerPP.at(ppRank); - TLLM_CHECK_WITH_INFO(selfAttentionLayerNum != 0, "selfAttentionLayerNum should not be 0"); + auto const selfLayerNum = transferIndexerKCache + ? selfConfig.getIndexerLayerNumPerPP().at(ppRank) + : selfConfig.getParallelConfig().mAttentionLayerNumPerPP.at(ppRank); + TLLM_CHECK_WITH_INFO(selfLayerNum != 0, "selfLayerNum should not be 0"); std::vector bufferEleSizes(targetNum, 0); - auto const cacheSizePerLayer = cacheBlockSize * blockNum / selfAttentionLayerNum; + auto const cacheSizePerLayer = cacheBlockSize * blockNum / selfLayerNum; for (size_t i = 0; i < targetNum; i++) { - auto const peerAttentionLayerNum + auto const peerLayerNum = targetInfo.getPeerPPDomainLayerNum(static_cast(localRankIndices[i])); - bufferEleSizes[i] = cacheSizePerLayer * peerAttentionLayerNum; + bufferEleSizes[i] = cacheSizePerLayer * peerLayerNum; } return bufferEleSizes; }; @@ -587,6 +605,12 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); auto startTime = LlmRequest::getSteadyClockNow(); size_t size = 0; + // Zero-sized target (e.g. an indexer K cache peer whose layer overlap holds no + // full-indexer layer): the sender skips the matching send, so skip the receive. + if (recvSplitCaches.at(processIdx) != nullptr && recvSplitCaches.at(processIdx)->getSize() == 0) + { + return; + } if (processIdx >= remainNoCoverTargetNum) { auto& buffer = recvSplitCaches.at(processIdx); @@ -764,6 +788,96 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s TLLM_LOG_WARNING("MLACacheFormatter::inquireSupport: TP size must be equal to DP size"); return false; } + if (selfConfig.getHasIndexerKCache() != destConfig.getHasIndexerKCache()) + { + TLLM_LOG_WARNING("MLACacheFormatter::inquireSupport: both sides must agree on the indexer K cache"); + return false; + } + if (selfConfig.getHasIndexerKCache()) + { + if (selfConfig.getIndexerDimPerHead() != destConfig.getIndexerDimPerHead() + || selfConfig.getIndexerKCacheQuantBlockSize() != destConfig.getIndexerKCacheQuantBlockSize() + || selfConfig.getIndexerKCacheUseFp4() != destConfig.getIndexerKCacheUseFp4()) + { + TLLM_LOG_WARNING( + "MLACacheFormatter::inquireSupport: indexer K cache layout (dim per head, quant block size, fp4) " + "must match on both sides"); + return false; + } + // With a masked indexer pool (per-layer indexer mask) both sides must derive the same + // global full-indexer-layer schedule: the totals must agree even when the per-PP splits + // differ (PP resharding). + auto const& selfIndexerPerPP = selfConfig.getIndexerLayerNumPerPP(); + auto const& destIndexerPerPP = destConfig.getIndexerLayerNumPerPP(); + auto const selfIndexerTotal = std::accumulate(selfIndexerPerPP.begin(), selfIndexerPerPP.end(), SizeType32{0}); + auto const destIndexerTotal = std::accumulate(destIndexerPerPP.begin(), destIndexerPerPP.end(), SizeType32{0}); + if (selfIndexerTotal != destIndexerTotal) + { + TLLM_LOG_WARNING( + "MLACacheFormatter::inquireSupport: both sides must own the same total number of indexer K cache " + "layers (self: %d, dest: %d)", + selfIndexerTotal, destIndexerTotal); + return false; + } + // Zero-layer indexer exchanges (a rank or a PP-domain peer whose attention overlap holds + // no full-indexer layer) are not supported yet: the agent-backend handshake and the + // partial-coverage bounce-buffer accounting cannot represent zero-sized targets. Require + // every PP rank on both sides to own at least one full-indexer layer, and every + // attention-domain overlap to contain at least one full-indexer layer. + auto hasZeroEntry = [](std::vector const& perPP) + { return std::any_of(perPP.begin(), perPP.end(), [](SizeType32 layerNum) { return layerNum == 0; }); }; + if (hasZeroEntry(selfIndexerPerPP) || hasZeroEntry(destIndexerPerPP)) + { + TLLM_LOG_WARNING( + "MLACacheFormatter::inquireSupport: a PP rank without indexer K cache layers (all local layers " + "use cross-layer indexer sharing) is not supported"); + return false; + } + // Every attention-domain overlap must contain at least one full-indexer layer; a zero + // overlap would produce a zero-sized transfer target. + auto overlapsHaveIndexerLayers = [](CacheState const& a, CacheState const& b) + { + auto const& aAttention = a.getParallelConfig().mAttentionLayerNumPerPP; + auto const& bAttention = b.getParallelConfig().mAttentionLayerNumPerPP; + auto const& aIndexer = a.getIndexerLayerNumPerPP(); + auto const& bIndexer = b.getIndexerLayerNumPerPP(); + int64_t aAttentionStart = 0; + int64_t aIndexerStart = 0; + for (size_t aRank = 0; aRank < aAttention.size(); aRank++) + { + int64_t const aAttentionEnd = aAttentionStart + aAttention[aRank]; + int64_t const aIndexerEnd = aIndexerStart + aIndexer[aRank]; + int64_t bAttentionStart = 0; + int64_t bIndexerStart = 0; + for (size_t bRank = 0; bRank < bAttention.size(); bRank++) + { + int64_t const bAttentionEnd = bAttentionStart + bAttention[bRank]; + int64_t const bIndexerEnd = bIndexerStart + bIndexer[bRank]; + auto const attentionOverlap + = std::min(aAttentionEnd, bAttentionEnd) - std::max(aAttentionStart, bAttentionStart); + auto const indexerOverlap + = std::min(aIndexerEnd, bIndexerEnd) - std::max(aIndexerStart, bIndexerStart); + if (attentionOverlap > 0 && indexerOverlap <= 0) + { + return false; + } + bAttentionStart = bAttentionEnd; + bIndexerStart = bIndexerEnd; + } + aAttentionStart = aAttentionEnd; + aIndexerStart = aIndexerEnd; + } + return true; + }; + if (!overlapsHaveIndexerLayers(selfConfig, destConfig)) + { + TLLM_LOG_WARNING( + "MLACacheFormatter::inquireSupport: a pipeline-parallel layer overlap between the two sides " + "holds no full-indexer layer (zero-sized indexer K cache transfer target); this PP partition " + "combination is not supported"); + return false; + } + } return true; } diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index 2be557ac1f54..1e73ded18903 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -562,6 +562,7 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( auto bufferKinds = std::move(requestAndBufferInfo.mBufferKinds); std::optional> kvOffsetRatio; + std::optional> indexerOffsetRatio; std::optional> rnnOffsetRatio; std::vector> offsetRatios; offsetRatios.reserve(bufferDescs.size()); @@ -572,7 +573,6 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( switch (kind) { case batch_manager::BufferKind::kKV: - case batch_manager::BufferKind::kKV_INDEXER: { if (!kvOffsetRatio) { @@ -585,6 +585,24 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( offsetRatios.push_back(*kvOffsetRatio); break; } + case batch_manager::BufferKind::kKV_INDEXER: + { + if (!indexerOffsetRatio) + { + // The indexer K cache pass runs in "indexer layer space": with a + // masked indexer pool (per-layer indexer mask) the receiver + // slices its recv buffer at indexer-layer prefix offsets, so the + // write offset must use the indexer layer counts (identical to + // the attention counts for dense models). + auto indexerTargetInfo = targetIRanksForIndexerKCache(mCacheState, + requestInfo.getTransState().getCacheState().value(), + requestInfo.getTransState().getCommState()->getSelfIdx()); + validateConnectionIdx(connectionIdx, indexerTargetInfo, "IndexerK", remoteAgentName); + indexerOffsetRatio = computeSendOffsetRatio(indexerTargetInfo, connectionIdx); + } + offsetRatios.push_back(*indexerOffsetRatio); + break; + } case batch_manager::BufferKind::kRNN: { if (!rnnOffsetRatio) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu index d57ed7530e28..dff6627f3624 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu +++ b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu @@ -274,6 +274,55 @@ TargetRanksInfo targetIRanksForRnn( return targetInfo; } +TargetRanksInfo targetIRanksForIndexerKCache( + kv_cache::CacheState const& peerCacheState, kv_cache::CacheState const& selfCacheState, int selfRank) +{ + auto targetInfo = targetIRanks(peerCacheState, selfCacheState, selfRank); + if (targetInfo.mIRanks.empty()) + { + return targetInfo; + } + auto const& peerIndexerPerPP = peerCacheState.getIndexerLayerNumPerPP(); + auto const& selfIndexerPerPP = selfCacheState.getIndexerLayerNumPerPP(); + auto const& peerParConfig = peerCacheState.getParallelConfig(); + auto const& selfParConfig = selfCacheState.getParallelConfig(); + TLLM_CHECK(static_cast(peerIndexerPerPP.size()) == peerParConfig.mPipelineParallelism); + TLLM_CHECK(static_cast(selfIndexerPerPP.size()) == selfParConfig.mPipelineParallelism); + + auto const selfPPRank = selfRank / (selfParConfig.mTensorParallelism * selfParConfig.mContextParallelism); + int64_t selfStart = 0; + for (int ppRank = 0; ppRank < selfPPRank; ppRank++) + { + selfStart += selfIndexerPerPP[ppRank]; + } + int64_t const selfEnd = selfStart + selfIndexerPerPP[selfPPRank]; + + // First peer PP rank of the attention domain (mIRanks is ordered CP-major, then TP, then + // PP ascending from the domain start). + auto const peerPPRankStart + = targetInfo.mIRanks.front() / (peerParConfig.mTensorParallelism * peerParConfig.mContextParallelism); + int64_t peerStart = 0; + for (int ppRank = 0; ppRank < peerPPRankStart; ppRank++) + { + peerStart += peerIndexerPerPP[ppRank]; + } + + int64_t peerLayerNumSum = 0; + for (int i = 0; i < targetInfo.mDomainPPSize; i++) + { + int64_t const peerEnd = peerStart + peerIndexerPerPP.at(peerPPRankStart + i); + auto const overlap = std::max(0, std::min(peerEnd, selfEnd) - std::max(peerStart, selfStart)); + targetInfo.mPeerLayerNumInDomainPP.at(i) = static_cast(overlap); + peerLayerNumSum += overlap; + peerStart = peerEnd; + } + TLLM_CHECK_WITH_INFO(peerLayerNumSum == selfIndexerPerPP[selfPPRank], + "Indexer K cache layer counts are inconsistent between the two sides: the attention-domain peers cover " + "%ld indexer layers but this rank owns %d. Both sides must derive the same per-layer indexer schedule.", + static_cast(peerLayerNumSum), selfIndexerPerPP[selfPPRank]); + return targetInfo; +} + template struct BlockInfo { @@ -1107,7 +1156,8 @@ void splitKVCache(std::map> { inputBlockNumSum += blocks.size(); } - auto targetRankInfo = targetIRanks(destCacheState, selfCacheState, selfIdx); + auto targetRankInfo = isIndexerKCache ? targetIRanksForIndexerKCache(destCacheState, selfCacheState, selfIdx) + : targetIRanks(destCacheState, selfCacheState, selfIdx); TLLM_CHECK(targetRankInfo.mIRanks.size() == (static_cast( targetRankInfo.mDomainPPSize * targetRankInfo.mDomainTPSize * targetRankInfo.mDomainCPSize))); @@ -1231,7 +1281,10 @@ void splitKVCache(std::map> int const tokensPerBlock = selfModelConfig.mTokensPerBlock; int const selfPPRank = selfIdx / (selfParallelConfig.mTensorParallelism * selfParallelConfig.mContextParallelism); - int const numLayers = selfParallelConfig.mAttentionLayerNumPerPP.at(selfPPRank); + // For the indexer K cache pass only full-indexer layers own a pool row (masked layout), + // so the per-rank layer count comes from the indexer layer counts. + int const numLayers = isIndexerKCache ? selfCacheState.getIndexerLayerNumPerPP().at(selfPPRank) + : selfParallelConfig.mAttentionLayerNumPerPP.at(selfPPRank); int const headNum = selfModelConfig.mNbKvHeadsPerLayer[0]; int const dimsPerHead = computeDimsPerHead(selfCacheState, isIndexerKCache); @@ -1454,7 +1507,8 @@ void concatKVCache(std::vector const& inputSplitBlo outputBlockNumSum += blocks.size(); } - auto targetRankInfo = targetIRanks(destCacheState, selfCacheState, selfIdx); + auto targetRankInfo = isIndexerKCache ? targetIRanksForIndexerKCache(destCacheState, selfCacheState, selfIdx) + : targetIRanks(destCacheState, selfCacheState, selfIdx); TLLM_CHECK(targetRankInfo.mIRanks.size() == (static_cast(targetRankInfo.mDomainPPSize * targetRankInfo.mDomainTPSize))); @@ -1560,7 +1614,10 @@ void concatKVCache(std::vector const& inputSplitBlo = static_cast(PtrsDeviceBuffer->data()) + outputBlockNumSum + inputSplitBlocks.size(); int const tokensPerBlock = selfModelConfig.mTokensPerBlock; int const selfPPRank = selfIdx / (selfParallelConfig.mTensorParallelism * selfParallelConfig.mContextParallelism); - int const numLayers = selfParallelConfig.mAttentionLayerNumPerPP.at(selfPPRank); + // For the indexer K cache pass only full-indexer layers own a pool row (masked layout), + // so the per-rank layer count comes from the indexer layer counts. + int const numLayers = isIndexerKCache ? selfCacheState.getIndexerLayerNumPerPP().at(selfPPRank) + : selfParallelConfig.mAttentionLayerNumPerPP.at(selfPPRank); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "concatKVCache numLayers:%d", numLayers); int const headNum = selfModelConfig.mNbKvHeadsPerLayer[0]; int const dimsPerHead = computeDimsPerHead(selfCacheState, isIndexerKCache); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h index 80816bc5c3a0..9963d5f7c1f5 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h @@ -59,6 +59,9 @@ TargetRanksInfo targetIRanks( TargetRanksInfo targetIRanksForRnn( kv_cache::CacheState const& peerCacheState, kv_cache::CacheState const& selfCacheState, int selfRank); +TargetRanksInfo targetIRanksForIndexerKCache( + kv_cache::CacheState const& peerCacheState, kv_cache::CacheState const& selfCacheState, int selfRank); + /** * @brief Calculate the number of blocks allocated to a specific Context Parallelism (CP) rank. * diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 4c2998c25d26..4bfb36406eb3 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -623,6 +623,7 @@ kv_cache::CacheState Serialization::deserializeCacheState(std::istream& is) auto indexerDimPerHead = su::deserialize(is); auto indexerKCacheQuantBlockSize = su::deserialize(is); auto indexerKCacheUseFp4 = su::deserialize(is); + auto indexerLayerNumPerPP = su::deserialize>(is); // RNN config (optional) auto hasRnnConfig = su::deserialize(is); std::optional rnnModelConfig; @@ -650,7 +651,7 @@ kv_cache::CacheState Serialization::deserializeCacheState(std::istream& is) CacheState cacheState{nbKvHeadsPerLayer, sizePerHead, tokensPerBlock, tensorParallelism, pipelineParallelism, contextParallelism, attentionLayerNumPerPP, dataType, attentionType, kvFactor, enableAttentionDP, DPrank, DPsize, enableBlockReuse, enablePartialReuse, hasIndexerKCache, indexerDimPerHead, indexerKCacheQuantBlockSize, - indexerKCacheUseFp4}; + indexerKCacheUseFp4, indexerLayerNumPerPP}; if (rnnModelConfig.has_value()) { cacheState.setRnnConfig( @@ -680,6 +681,7 @@ void Serialization::serialize(kv_cache::CacheState const& state, std::ostream& o su::serialize(state.getIndexerDimPerHead(), os); su::serialize(state.getIndexerKCacheQuantBlockSize(), os); su::serialize(state.getIndexerKCacheUseFp4(), os); + su::serialize(state.mIndexerLayerNumPerPP, os); // RNN config (optional) su::serialize(state.mRnnCacheState.has_value(), os); if (state.mRnnCacheState.has_value()) @@ -722,6 +724,7 @@ size_t Serialization::serializedSize(kv_cache::CacheState const& state) totalSize += su::serializedSize(state.getIndexerDimPerHead()); totalSize += su::serializedSize(state.getIndexerKCacheQuantBlockSize()); totalSize += su::serializedSize(state.getIndexerKCacheUseFp4()); + totalSize += su::serializedSize(state.mIndexerLayerNumPerPP); // RNN config (optional) totalSize += su::serializedSize(state.mRnnCacheState.has_value()); if (state.mRnnCacheState.has_value()) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index d11838f68af5..cc2ec8645e3b 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -123,11 +123,12 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def(nb::init, SizeType32, SizeType32, runtime::WorldConfig, std::vector, tensorrt_llm::DataType, executor::kv_cache::CacheState::AttentionType, std::optional, - std::vector>(), + std::vector, std::vector>(), nb::arg("cache_manager"), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"), nb::arg("world_config"), nb::arg("attention_layer_num_per_pp"), nb::arg("dtype"), nb::arg("attention_type"), nb::arg("cache_transceiver_config") = std::nullopt, - nb::arg("rnn_layer_num_per_pp") = std::vector{}); + nb::arg("rnn_layer_num_per_pp") = std::vector{}, + nb::arg("indexer_layer_num_per_pp") = std::vector{}); nb::class_(m, "CacheTransceiverComm") .def( diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index 3c78659c87c4..e3c0071756d7 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -1566,10 +1566,13 @@ TEST(SerializeUtilsTest, CacheStateIndexerKCache) bool hasIndexerKCache = true; texec::SizeType32 indexerDimPerHead = 96; texec::SizeType32 indexerKCacheQuantBlockSize = 128; + bool indexerKCacheUseFp4 = false; + // Masked indexer pool: fewer indexer layers than attention layers. + std::vector indexerLayerNumPerPP{1}; CacheState state{nbKvHeadsPerLayer, sizePerHead, tokensPerBlock, tp, pp, cp, attentionLayerNumPerPP, dataType, attentionType, kvFactor, enableAttentionDP, dpRank, dpSize, enableBlockReuse, enablePartialReuse, - hasIndexerKCache, indexerDimPerHead, indexerKCacheQuantBlockSize}; + hasIndexerKCache, indexerDimPerHead, indexerKCacheQuantBlockSize, indexerKCacheUseFp4, indexerLayerNumPerPP}; std::ostringstream oss; texec::Serialization::serialize(state, oss); @@ -1591,4 +1594,8 @@ TEST(SerializeUtilsTest, CacheStateIndexerKCache) EXPECT_EQ(state.getHasIndexerKCache(), state2.getHasIndexerKCache()); EXPECT_EQ(state.getIndexerDimPerHead(), state2.getIndexerDimPerHead()); EXPECT_EQ(state.getIndexerKCacheQuantBlockSize(), state2.getIndexerKCacheQuantBlockSize()); + EXPECT_EQ(state.getIndexerKCacheUseFp4(), state2.getIndexerKCacheUseFp4()); + EXPECT_EQ(state.getIndexerLayerNumPerPP(), state2.getIndexerLayerNumPerPP()); + EXPECT_EQ(state.getIndexerLayerNumPerPP(), indexerLayerNumPerPP); + EXPECT_EQ(state, state2); } diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index e294cdff9d49..4e8db47133e9 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -582,7 +582,8 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam> const& indexerKCacheLayerMask = std::nullopt) { mIsWindowAttention = isWindow; @@ -612,6 +613,51 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam(indexerKCacheLayerMask->size()) == numLayers); + } + auto countIndexerLayers = [&](int startLayer, int layerCount) + { + if (!indexerKCacheLayerMask.has_value()) + { + return layerCount; + } + int count = 0; + for (int layerId = startLayer; layerId < startLayer + layerCount; layerId++) + { + count += indexerKCacheLayerMask->at(layerId) ? 1 : 0; + } + return count; + }; + mIndexerLayerNumPerPP = std::vector(mPpSize, 0); + std::optional> localIndexerKCacheLayerMask = std::nullopt; + { + int startLayer = 0; + for (int ppRank = 0; ppRank < mPpSize; ppRank++) + { + mIndexerLayerNumPerPP[ppRank] = countIndexerLayers(startLayer, mAttentionLayerNumPerPP[ppRank]); + if (ppRank == mPpRank && indexerKCacheLayerMask.has_value()) + { + localIndexerKCacheLayerMask = std::vector(indexerKCacheLayerMask->begin() + startLayer, + indexerKCacheLayerMask->begin() + startLayer + mAttentionLayerNumPerPP[ppRank]); + } + startLayer += mAttentionLayerNumPerPP[ppRank]; + } + } + auto contextIndexerLayerNumPerPP = std::vector(mContextPpSize, 0); + { + int startLayer = 0; + for (int ppRank = 0; ppRank < mContextPpSize; ppRank++) + { + contextIndexerLayerNumPerPP[ppRank] + = countIndexerLayers(startLayer, contextAttentionLayerNumPerPP[ppRank]); + startLayer += contextAttentionLayerNumPerPP[ppRank]; + } + } + if (!isMLA) { // ASSERT_EQ(numHeads % mTpSize , 0); @@ -695,17 +741,21 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam(numLayers, numHeadsPerRank, sizePerHead, tokensPerBlock, mTpSize, mPpSize, mCpSize, mAttentionLayerNumPerPP, dataType, attentionType, kvFactor, - enableDPAttention, DPrank, DPsize, false, isIndexerKCache, indexerDimPerHead, indexerKCacheQuantBlockSize); + enableDPAttention, DPrank, DPsize, /*enableBlockReuse=*/false, /*enablePartialReuse=*/false, + /*hasIndexerKCache=*/isIndexerKCache, indexerDimPerHead, indexerKCacheQuantBlockSize, + /*indexerKCacheUseFp4=*/false, mIndexerLayerNumPerPP); mContextCacheState = std::make_unique(numLayers, numHeadsPerRankForContext, sizePerHead, tokensPerBlock, mContextTpSize, mContextPpSize, mContextCpSize, contextAttentionLayerNumPerPP, - dataType, attentionType, kvFactor, mContextDP, DPrank, mContextTpSize, false, isIndexerKCache, - indexerDimPerHead, indexerKCacheQuantBlockSize); + dataType, attentionType, kvFactor, mContextDP, DPrank, mContextTpSize, /*enableBlockReuse=*/false, + /*enablePartialReuse=*/false, /*hasIndexerKCache=*/isIndexerKCache, indexerDimPerHead, + indexerKCacheQuantBlockSize, /*indexerKCacheUseFp4=*/false, contextIndexerLayerNumPerPP); // UVM seems to be incompatible with MPI, and it is continuing to investigate. bool constexpr useUvm = false; @@ -1071,9 +1121,12 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam mAttentionLayerNumPerPP; + // Per-PP indexer K cache layer counts (== attention counts without a mask). + std::vector mIndexerLayerNumPerPP; SizeType32 mMaxNumSequences{}; std::unique_ptr mManager; @@ -1806,6 +1864,121 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0ForMLAWithIndexerKCache, Asymmetrical testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(true), testing::Values(256), testing::Values(128))); +// Masked indexer K cache pool (per-layer indexer mask, GLM 5.2-style cross-layer indexer +// sharing): only full-indexer layers own a pool row, so the indexer pass runs in "indexer +// layer space". Exercises PP resharding of masked indexer blocks end to end. +class AsymmetricalCacheTestWithIndexerLayerMask : public AsymmetricalCacheTest +{ +}; + +TEST_P(AsymmetricalCacheTestWithIndexerLayerMask, TestCase) +{ + AsymmetricTestParam param = GetParam(); + int contextTp = std::get<0>(param); + int contextPp = std::get<1>(param); + int contextCp = std::get<2>(param); + int genTp = std::get<3>(param); + int genPp = std::get<4>(param); + int genCp = std::get<5>(param); + int numLayers = std::get<6>(param); + int numHeads = std::get<7>(param); + int sizePerHead = std::get<8>(param); + int tokensPerBlock = std::get<9>(param); + tensorrt_llm::DataType dataType = std::get<10>(param); + int kvFactor = std::get<11>(param); + bool isMLA = std::get<12>(param); + bool contextDP = std::get<13>(param); + bool generationDP = std::get<14>(param); + bool isWindow = std::get<15>(param); + bool isIndexerKCache = std::get<16>(param); + int indexerDimPerHead = std::get<17>(param); + int indexerKCacheQuantBlockSize = std::get<18>(param); + + if (isIndexerKCache && tensorrt_llm::common::getEnvUseMooncakeKvCache()) + { + // https://nvbugs/5760737 + GTEST_SKIP() << "Temporarily skipping cache transceiver tests with Mooncake backend for Indexer KCache."; + } + + // GLM 5.2-style schedule (index_topk_freq=4, index_skip_topk_offset=2): full-indexer + // layers are {0, 1, 5, 9, ...}; every PP slice below keeps at least one full layer. + std::vector indexerKCacheLayerMask(numLayers); + for (int layerId = 0; layerId < numLayers; layerId++) + { + indexerKCacheLayerMask[layerId] = (std::max(layerId - 2 + 1, 0) % 4) == 0; + } + + std::vector lenList = {30, 10, 60, 80}; + + setUpCommunicator(contextTp, contextPp, contextCp, genTp, genPp, genCp, isMLA, contextDP, generationDP); + + if (mIsContext || mIsGeneration) + { + setUpCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, dataType, kvFactor, isMLA, false, isWindow, + isIndexerKCache, indexerDimPerHead, indexerKCacheQuantBlockSize, indexerKCacheLayerMask); + setUpCacheTransceiver(); + std::vector> requests; + + // the second loop is for cache reuse + for (int i = 0; i < 2; i++) + { + for (auto len : lenList) + { + requests.emplace_back(makeLlmRequest(len)); + } + + if (mIsContext) + { + std::vector> contextFutures; + for (auto&& request : requests) + { + contextFutures.push_back(addRequestAndTransportCacheForContext(request)); + } + mComm->barrier(); + for (auto&& cfuture : contextFutures) + { + cfuture.get(); + } + } + else + { + std::vector> generationFutures; + mComm->barrier(); + for (auto&& request : requests) + { + generationFutures.push_back(addRequestAndTransportCacheForGeneration(request)); + } + + for (auto&& gfuture : generationFutures) + { + gfuture.get(); + } + for (auto&& request : requests) + { + generationVerifyKVCache(request); + } + } + for (auto&& request : requests) + { + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*request->mLlmRequest); + mManager->removeSequence(request->mLlmRequest->mRequestId, request->mLlmRequest); + } + requests.clear(); + mComm->barrier(); + } + } + tensorrt_llm::mpi::MpiComm::world().barrier(); +} + +// 8 layers with full-indexer layers {0, 1, 5}: PP=2 slices hold {2, 1} indexer rows, so +// ctx/gen PP resharding exchanges masked indexer blocks with asymmetric per-PP counts. +INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0ForMLAWithIndexerKCacheLayerMask, AsymmetricalCacheTestWithIndexerLayerMask, + testing::Combine(testing::Values(1), testing::Values(1, 2), testing::Values(1), testing::Values(1), + testing::Values(1, 2), testing::Values(1), testing::Values(8), testing::Values(1), testing::Values(4), + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT), testing::Values(1), testing::Values(true), + testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(true), + testing::Values(256), testing::Values(128))); + // Tests cases where there's non-trivial TP and PP on context side but only CP on gen side. INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0WithCPForMLA, AsymmetricalCacheTest, testing::Combine(/*contextTp*/ testing::Values(1, 2), diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index b7c18f419ebc..937e83f37e08 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -199,25 +199,16 @@ def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: # Indexer K cache support if getattr(kv_cache_manager, "enable_indexer_k_cache", False): - # The FLAT PoolView contract assumes the pool covers ALL layers of - # the layer group with a uniform per-layer stride. A masked - # indexer pool (per-layer indexer mask, e.g. GLM 5.2 cross-layer - # indexer sharing) holds fewer rows than layers — possibly zero, - # in which case no pool exists at all — so the stride math would - # corrupt transfers. Fail fast until the extractor publishes - # per-layer buffer_entries for the indexer pool. Check the mask - # BEFORE fetching the pool: a fully masked-out rank has no pool - # to fetch. local_indexer_mask = getattr(kv_cache_manager, "indexer_k_cache_local_layer_mask", None) if local_indexer_mask is not None and not all( local_indexer_mask[lid] for lid in local_layer_ids ): raise NotImplementedError( - "Disaggregated KV transfer does not support a per-layer " - "masked indexer k-cache pool yet: " + "The Python KV transceiver runtime does not support a " + "per-layer masked indexer k-cache pool yet: " f"{sum(local_indexer_mask[lid] for lid in local_layer_ids)}" f" of {len(local_layer_ids)} layers in this layer group " - "own an indexer k-cache. Disable disaggregated serving " + "own an indexer k-cache. Use the C++ cache transceiver " "for models with cross-layer indexer sharing (e.g. " "GLM 5.2)." ) @@ -231,7 +222,7 @@ def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: f"pool holds {indexer_pool.shape[1]} layer rows but the " f"layer group has {len(local_layer_ids)} layers. Disable " "disaggregated serving for models with cross-layer " - "indexer sharing (e.g. GLM 5.2)." + "indexer sharing." ) per_block_elems = 1 for d in indexer_pool.shape[1:]: # skip numBlocks dim diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 348d47a9659c..2251347dc91f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -278,6 +278,21 @@ def __init__(self, _is_disagg_inflight_cancel_config_supported( cache_transceiver_config)) + # Per-PP 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 the counts can be + # smaller than the attention layer counts. + indexer_layer_num_per_pp_rank = [] + if getattr(kv_cache_manager, "enable_indexer_k_cache", False): + local_indexer_mask = getattr(kv_cache_manager, + "indexer_k_cache_local_layer_mask", + None) + local_indexer_layer_num = (sum(local_indexer_mask) + if local_indexer_mask is not None else + pp_layer_num) + indexer_layer_num_per_pp_rank = dist.pp_allgather( + local_indexer_layer_num) + # Get RNN layer distribution if mamba_cache_manager is provided. rnn_layer_num_per_pp_rank = [] if mamba_cache_manager is not None: @@ -294,13 +309,12 @@ def __init__(self, f"RNN state transfer enabled: rnn_layer_num_per_pp={rnn_layer_num_per_pp_rank}" ) - self.impl = CacheTransceiverCpp(kv_cache_manager.impl, - total_num_kv_heads_per_layer, head_dim, - tokens_per_block, world_config, - pp_layer_num_per_pp_rank, dtype, - attention_type, - cache_transceiver_config._to_pybind(), - rnn_layer_num_per_pp_rank) + self.impl = CacheTransceiverCpp( + kv_cache_manager.impl, total_num_kv_heads_per_layer, head_dim, + tokens_per_block, world_config, + pp_layer_num_per_pp_rank, dtype, attention_type, + cache_transceiver_config._to_pybind(), rnn_layer_num_per_pp_rank, + indexer_layer_num_per_pp_rank) def respond_and_send_async(self, req: LlmRequest): return self.impl.respond_and_send_async(req) From 64de65fef9fb88f1239c6fa1735dd2b92f42ad6e Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:28:50 -0700 Subject: [PATCH 3/4] [None][test] Run GLM 5.2 disaggregated serving through the C++ cache transceiver The Python KV transceiver runtime does not support the per-layer masked DSA indexer k-cache pool yet, so switch TestGLM52NVFP4::test_nvfp4_nixl_python to the C++ NIXL transceiver (renamed to test_nvfp4_nixl), exercising ctx->gen transfer of the masked indexer pool end to end (MTP draft layer included). The l0_dgx_b200 and QA list entries are updated to the new test id. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tests/integration/defs/accuracy/test_disaggregated_serving.py | 3 +-- tests/integration/test_lists/qa/llm_function_core.txt | 2 +- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 59f4bcaa87f9..492c6c1fe14d 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2352,7 +2352,7 @@ class TestGLM52NVFP4(LlmapiAccuracyTestHarness): @pytest.mark.skip_less_device(8) @pytest.mark.parametrize("use_kv_cache_manager_v2", [False], ids=["cache_mgr_v1"]) - def test_nvfp4_nixl_python(self, use_kv_cache_manager_v2): + def test_nvfp4_nixl(self, use_kv_cache_manager_v2): kv_cache_config = { "free_gpu_memory_fraction": 0.7, "enable_block_reuse": False, @@ -2360,7 +2360,6 @@ def test_nvfp4_nixl_python(self, use_kv_cache_manager_v2): } cache_transceiver_config = { "backend": "NIXL", - "transceiver_runtime": "PYTHON", } moe_config = {"backend": "CUTEDSL"} speculative_config = { diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 02c43b89f8dd..9a3bb57ca4ad 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -20,7 +20,7 @@ accuracy/test_disaggregated_serving.py::TestGPTOSS::test_auto_dtype[False] accuracy/test_disaggregated_serving.py::TestGPTOSS::test_auto_dtype[True] accuracy/test_disaggregated_serving.py::TestGPTOSS::test_kv_cache_v2_nixl_python[cache_mgr_v1] accuracy/test_disaggregated_serving.py::TestGPTOSS::test_kv_cache_v2_nixl_python[cache_mgr_v2] -accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl_python[cache_mgr_v1] +accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl[cache_mgr_v1] accuracy/test_disaggregated_serving.py::TestKimiK25::test_nvfp4 accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[False-False-False-False] accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[False-False-False-True] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index cc96800d204a..2dbc3fe8d245 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -175,7 +175,7 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=False] TIMEOUT (60) - - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl_python[cache_mgr_v1] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl[cache_mgr_v1] TIMEOUT (60) # DeepSeek-V4 EPLB pre-merge sanity (uncomment once DeepSeek-V4-Flash/Flash-Base # checkpoints are staged under llm_models_root()). # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_8gpus_static_eplb[moe_backend=WIDEEP] TIMEOUT (120) From 643f1330e2c1b2c523803110859170e83f884841 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Fri, 24 Jul 2026 16:24:28 -0700 Subject: [PATCH 4/4] [None][test] Waive test_cache_transceiver[8proc-ucx_kvcache-90] (nvbugs/5838199) The full cacheTransceiverTest binary hangs under 8-process UCX due to a pre-existing cumulative resource issue in the transceiver test suite, independent of this PR: on an H100 node the same binary hangs at a non-indexer base case on both this branch and main (09e5d0c318). This mirrors the already-waived sibling variant test_cache_transceiver[8proc-mooncake_kvcache-90] and is tracked by the same NVBug. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 4aebb18910db..4f616140dcb8 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -124,6 +124,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] SKIP (https:/ accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6248827) accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray SKIP (https://nvbugs/6427411) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199) +cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] SKIP (https://nvbugs/5838199) disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6105768) disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6427411) disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6427411)