[https://nvbugs/6465993][fix] use attention cache dtype for disaggregated transfer

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)
         {

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<tensorrt_llm::DataType> cacheDataType;
+    std::optional<tensorrt_llm::DataType> 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<int>(attentionDataType.value()), poolIdx, static_cast<int>(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<size_t> 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<size_t>(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<size_t>(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<size_t> 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<int>(dtype0), i, static_cast<int>(dtypeI));
-        }
-    }
     TLLM_LOG_INFO("CacheTransBufferManager created for KV cache");
 }
 

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;

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<PoolConfiguration>{
+        {recurrentStatesWindow, sizePerHead, tensorrt_llm::DataType::kHALF},
+        {maxAttentionWindow, sizePerHead, tensorrt_llm::DataType::kFP8},
+    };
+    auto const stream = std::make_shared<tr::CudaStream>();
+
+    auto kvCacheManager = std::make_unique<KVCacheManager>(std::vector<SizeType32>{0, numKvHeads}, sizePerHead,
+        tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth,
+        std::vector<SizeType32>{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<tensorrt_llm::DataType> 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;

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)