From ad9a15bc7e5403f23413955bc54c72f540680a60 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 16 Jul 2026 09:19:22 -0700 Subject: [PATCH 1/2] [https://nvbugs/6465993][test] unwaive Nemotron 3 Super disaggregated tests Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 4aebb18910db..efc93925984c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -12,10 +12,7 @@ accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symm accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[GSM8K-tp2pp2] SKIP (https://nvbugs/6428069) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp1pp2] SKIP (https://nvbugs/6427411) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp2pp2] SKIP (https://nvbugs/6428069) -accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=False] SKIP (https://nvbugs/6465993) accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] SKIP (https://nvbugs/6478726) -accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False] SKIP (https://nvbugs/6465993) -accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4 SKIP (https://nvbugs/6465993) accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] SKIP (https://nvbugs/6427411) accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] SKIP (https://nvbugs/6478692) From fc9cd30f3d2acaf73e769db653557be8f1a36b8a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:17:59 -0700 Subject: [PATCH 2/2] [NVBUG 6465993][fix] use attention dtype for KV transfer buffers Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheFormatter.cpp | 7 +- .../batch_manager/cacheTransBuffer.cpp | 147 ++++++++++++------ .../batch_manager/cacheTransBuffer.h | 8 +- .../batch_manager/kvCacheManagerTest.cpp | 75 +++++++-- 4 files changed, 172 insertions(+), 65 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 9dc44531d409..3930c462f34b 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -729,12 +729,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { NVTX3_SCOPED_RANGE(formatInputRecvBuffer); - // TODO(disagg-multi-dtype): pool 0's dtype is treated as canonical for the wire - // transport here. Pools with differing dtypes are rejected up-front in - // CacheTransBufferManager's constructor (see cacheTransBuffer.cpp). When - // per-pool dtype dispatch lands, this single dataType variable must be replaced - // with a per-pool lookup keyed by the source pool of each block. - auto dataType = mCacheManager->getPrimaryPool(0)->getDataType(); + auto const dataType = mCacheTransBufferManager->getDataType(); bool layerWise = common::getEnvDisaggLayerwise() && numKvPools == 1; if (layerWise) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index e06198f9ea60..1e8bbcb766bd 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,6 +27,72 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager { +namespace +{ + +bool isCachePool(BlockManager const& blockManager, SizeType32 poolIdx) +{ + auto const& pool = blockManager.getPool(poolIdx); + return !pool.containsBlockScales && !pool.containsIndexerKCache; +} + +bool isAttentionCachePool(BlockManager const& blockManager, SizeType32 poolIdx) +{ + return isCachePool(blockManager, poolIdx) + && !LinearAttentionMetadata::hasLinearCache(blockManager.getPoolWindowSize(poolIdx)); +} + +tensorrt_llm::DataType getTransferDataType(KVCacheManager::BaseKVCacheManager* cacheManager, bool transferIndexerKCache) +{ + TLLM_CHECK(cacheManager); + if (transferIndexerKCache) + { + auto const indexerKCachePool = cacheManager->getIndexerKCachePool(); + TLLM_CHECK(indexerKCachePool); + return indexerKCachePool->getDataType(); + } + + auto const& blockManager = cacheManager->getBlockManager(); + std::optional cacheDataType; + std::optional attentionDataType; + SizeType32 firstPoolIdx = -1; + // Recurrent-state pools have a separate transfer manager and formatter. Only + // attention pools determine the KV transfer-buffer dtype. + for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx) + { + if (!isCachePool(blockManager, poolIdx)) + { + continue; + } + + auto const poolDataType = blockManager.getPrimaryPool(poolIdx)->getDataType(); + if (!cacheDataType.has_value()) + { + cacheDataType = poolDataType; + } + if (!isAttentionCachePool(blockManager, poolIdx)) + { + continue; + } + if (!attentionDataType.has_value()) + { + attentionDataType = poolDataType; + firstPoolIdx = poolIdx; + continue; + } + + TLLM_CHECK_WITH_INFO(poolDataType == attentionDataType.value(), + "Disaggregated KV cache transfer does not yet support attention pools with differing dtypes " + "(pool %d dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool dtype dispatch in formatter.", + firstPoolIdx, static_cast(attentionDataType.value()), poolIdx, static_cast(poolDataType)); + } + + TLLM_CHECK_WITH_INFO(cacheDataType.has_value(), "Disaggregated KV cache transfer requires a cache pool"); + return attentionDataType.value_or(cacheDataType.value()); +} + +} // namespace + // ============================================================================ // FabricMemory Implementation // ============================================================================ @@ -194,39 +260,38 @@ bool FabricMemory::supportFabricMemory() size_t CacheTransBufferManager::computeTransferBufferSize( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens, bool transferIndexerKCache) { - tensorrt_llm::DataType dataType; - if (transferIndexerKCache) - { - dataType = cacheManager->getIndexerKCachePool()->getDataType(); - } - else + auto const dataType = getTransferDataType(cacheManager, transferIndexerKCache); + + auto const& blockManager = cacheManager->getBlockManager(); + auto const tokensPerBlock = blockManager.getTokensPerBlock(); + bool hasAttentionCachePool = false; + for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx) { - dataType = cacheManager->getPrimaryPool(0)->getDataType(); + hasAttentionCachePool |= isAttentionCachePool(blockManager, poolIdx); } - - auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock(); size_t bufferSizeFromMaxNumToken = 0; if (maxNumTokens.has_value()) { TLLM_CHECK(maxNumTokens.value() % tokensPerBlock == 0); - auto dataSize = common::getDTypeSize(dataType); - SizeType32 kvCacheByteSizePerTokenPerLayer = 0; + auto const dataSize = common::getDTypeSize(dataType); + SizeType32 indexerCacheByteSizePerTokenPerLayer = 0; if (transferIndexerKCache) { - kvCacheByteSizePerTokenPerLayer + indexerCacheByteSizePerTokenPerLayer = cacheManager->getIndexerKCachePool()->getDimension<-1>() * dataSize / tokensPerBlock; } - else - { - auto primaryPool = cacheManager->getPrimaryPool(0); - kvCacheByteSizePerTokenPerLayer - = primaryPool->getDimension<-1>() * primaryPool->getDimension<2>() * dataSize / tokensPerBlock; - } - for (auto layerId = 0; layerId < cacheManager->getBlockManager().getNumLayers(); layerId++) + for (auto layerId = 0; layerId < blockManager.getNumLayers(); layerId++) { - auto poolIdx = cacheManager->getBlockManager().getLayerPoolIdx(layerId); - auto windowSize = static_cast(cacheManager->getBlockManager().getPoolWindowSize(poolIdx)); + auto const poolIdx = blockManager.getLayerPoolIdx(layerId); + auto const encodedWindowSize = blockManager.getPoolWindowSize(poolIdx); + if (!transferIndexerKCache && hasAttentionCachePool + && LinearAttentionMetadata::hasLinearCache(encodedWindowSize)) + { + continue; + } + + auto const windowSize = static_cast(encodedWindowSize); auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock; auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value()); if (common::getEnvKVCacheTransferAllBlocksForWindow()) @@ -235,7 +300,17 @@ size_t CacheTransBufferManager::computeTransferBufferSize( } validTokenNum += tokensPerBlock; // add one more block - bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer; + if (transferIndexerKCache) + { + bufferSizeFromMaxNumToken += validTokenNum * indexerCacheByteSizePerTokenPerLayer; + } + else + { + auto const primaryPool = blockManager.getPrimaryPool(poolIdx); + auto const kvCacheByteSizePerTokenPerLayer + = primaryPool->getDimension<-1>() * primaryPool->getDimension<2>() * dataSize / tokensPerBlock; + bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer; + } } } @@ -245,36 +320,12 @@ size_t CacheTransBufferManager::computeTransferBufferSize( CacheTransBufferManager::CacheTransBufferManager( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens, bool transferIndexerKCache) : BaseTransBufferManager(computeTransferBufferSize(cacheManager, maxNumTokens, transferIndexerKCache), - transferIndexerKCache ? cacheManager->getIndexerKCachePool()->getDataType() - : cacheManager->getPrimaryPool(0)->getDataType(), - maxNumTokens) + getTransferDataType(cacheManager, transferIndexerKCache), maxNumTokens) , mCacheManager{cacheManager} , mTransferIndexerKCache{transferIndexerKCache} { // TODO: FP4 dataSize TLLM_CHECK(mCacheManager); - // TODO(disagg-multi-dtype): Per-pool dtype dispatch in formatter / transfer buffer - // not yet implemented. Disagg currently picks pool 0's dtype as the canonical - // transport type (above), so any KV pool with a different dtype would be silently - // miscoerced on the wire. Fail loudly until per-pool dispatch lands. We restrict - // the comparison to KV pools (getNumPools(false, false)) since block-scale and - // indexer-K pools legitimately have their own dtypes and travel through their own - // code paths. - if (!transferIndexerKCache) - { - auto const numKvPools = mCacheManager->getBlockManager().getNumPools( - /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false); - auto const dtype0 = mCacheManager->getPrimaryPool(0)->getDataType(); - for (SizeType32 i = 1; i < numKvPools; ++i) - { - auto const dtypeI = mCacheManager->getPrimaryPool(i)->getDataType(); - TLLM_CHECK_WITH_INFO(dtypeI == dtype0, - "Disaggregated KV cache transfer does not yet support pools with differing dtypes " - "(pool 0 dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool dtype " - "dispatch in formatter.", - static_cast(dtype0), i, static_cast(dtypeI)); - } - } TLLM_LOG_INFO("CacheTransBufferManager created for KV cache"); } diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h index 1635c11bc673..72f67e8ca3c9 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -74,6 +74,12 @@ class CacheTransBufferManager : public BaseTransBufferManager return mCacheManager; } + /// @brief Get the data type used by KV cache transfer buffers. + [[nodiscard]] tensorrt_llm::DataType getDataType() const noexcept + { + return mDataType; + } + [[nodiscard]] BufferKind getBufferKind() const override { return mTransferIndexerKCache ? BufferKind::kKV_INDEXER : BufferKind::kKV; diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 98c2232b4062..7bacaf9a2575 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -10619,12 +10619,58 @@ TEST_F(KVCacheManagerTest, VswaMixedHeadDimReuseSmoke) } } +TEST_F(KVCacheManagerTest, HybridDisaggUsesAttentionPoolDtype) +{ + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 16; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimaryPool = 4; + auto constexpr blocksInSecondaryPool = 0; + auto constexpr maxNumSequences = 2; + auto constexpr maxBeamWidth = 1; + auto constexpr maxAttentionWindow = 16; + auto constexpr recurrentStatesBytes = 64; + SizeType32 constexpr recurrentStatesWindow = LinearAttentionMetadata::LinearCacheType::kRecurrentStates; + + LinearAttentionMetadata const linearAttentionMetadata{ + .linearLayerIndices = {0}, + .cacheType = recurrentStatesWindow, + .allRecurrentStatesBytes = recurrentStatesBytes, + }; + auto const blocksPerWindow = BlocksPerWindow{ + {recurrentStatesWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, + {maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, + }; + auto const poolConfigurations = std::vector{ + {recurrentStatesWindow, sizePerHead, tensorrt_llm::DataType::kHALF}, + {maxAttentionWindow, sizePerHead, tensorrt_llm::DataType::kFP8}, + }; + auto const stream = std::make_shared(); + + auto kvCacheManager = std::make_unique(std::vector{0, numKvHeads}, sizePerHead, + tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, + std::vector{recurrentStatesWindow, maxAttentionWindow}, tensorrt_llm::DataType::kFP8, + /*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); + kvCacheManager->allocatePools(/*useUvm=*/false); + + CacheTransBufferManager cacheTransBufferManager(kvCacheManager.get(), /*maxNumTokens=*/tokensPerBlock); + EXPECT_EQ(cacheTransBufferManager.getDataType(), tensorrt_llm::DataType::kFP8); + + auto const bufferId = cacheTransBufferManager.assignBufferIndexForSend(); + ASSERT_TRUE(bufferId.has_value()); + EXPECT_EQ(cacheTransBufferManager.getSendBuffer(bufferId)->getDataType(), tensorrt_llm::DataType::kFP8); + cacheTransBufferManager.freeBufferIndexForSend(bufferId); +} + // A6: VSWA + disagg dtype mismatch must fire the A4 guard. // -// The constructor of CacheTransBufferManager picks pool 0's dtype as canonical for -// the wire transport. When a KVCacheManager hosts pools with differing dtypes -// (mixed-precision per-window), that silent coercion would corrupt the wire format. -// The guard added in cacheTransBuffer.cpp must throw at construction time. +// CacheTransBufferManager uses a single dtype for the wire transport. When a +// KVCacheManager hosts attention pools with differing dtypes (mixed-precision +// per-window), that silent coercion would corrupt the wire format. The guard in +// cacheTransBuffer.cpp must throw at construction time. // // This test only exercises the helper / construction path that runs the guard; it // does not stand up a full disaggregated transfer (out of scope at unit-test @@ -10664,14 +10710,23 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) kvCacheManager->allocatePools(/*useUvm=*/false); // Sanity: the manager really does host KV pools with two different dtypes. - auto const numKvPools = kvCacheManager->getBlockManager().getNumPools( - /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false); - ASSERT_GE(numKvPools, 2); - auto const dtype0 = kvCacheManager->getPrimaryPool(0)->getDataType(); + auto const& blockManager = kvCacheManager->getBlockManager(); + ASSERT_GE(blockManager.getNumPools(/*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false), 2); + std::optional dtype0; bool foundMismatch = false; - for (SizeType32 i = 1; i < numKvPools; ++i) + for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx) { - if (kvCacheManager->getPrimaryPool(i)->getDataType() != dtype0) + auto const& pool = blockManager.getPool(poolIdx); + if (pool.containsBlockScales || pool.containsIndexerKCache) + { + continue; + } + auto const dataType = blockManager.getPrimaryPool(poolIdx)->getDataType(); + if (!dtype0.has_value()) + { + dtype0 = dataType; + } + else if (dataType != dtype0.value()) { foundMismatch = true; break;