[https://nvbugs/6479324][fix] Keep KV dtype for hybrid recurrent-states pool to unstall disagg transfers

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -729,12 +729,16 @@ 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();
+        // TODO(disagg-multi-dtype): the canonical KV transport pool's dtype is treated
+        // as canonical for the wire transport here.  KV 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.  Pool 0 may be a recurrent-states pool on hybrid models
+        // (its sentinel window sorts first) whose dtype is unrelated to KV transport,
+        // hence the explicit canonical-pool lookup.
+        auto dataType
+            = mCacheManager->getPrimaryPool(CacheTransBufferManager::kvTransportPoolIdx(mCacheManager))->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");
@@ -191,6 +191,29 @@ bool FabricMemory::supportFabricMemory()
 // CacheTransBufferManager Implementation
 // ============================================================================
 
+SizeType32 CacheTransBufferManager::kvTransportPoolIdx(KVCacheManager::BaseKVCacheManager* cacheManager)
+{
+    auto const& blockManager = cacheManager->getBlockManager();
+    auto const numPools = blockManager.getNumPools();
+    for (SizeType32 poolIdx = 0; poolIdx < numPools; ++poolIdx)
+    {
+        if (LinearAttentionMetadata::hasLinearCache(blockManager.getPoolWindowSize(poolIdx)))
+        {
+            // Recurrent/linear-attention state pools transfer through
+            // RnnCacheTransBufferManager, not the KV path, and may carry a
+            // dtype different from the KV pools (e.g. BF16 SSM state next to
+            // FP8 KV). They must not define the KV wire dtype or sizing.
+            continue;
+        }
+        if (blockManager.containsBlockScales(poolIdx))
+        {
+            continue;
+        }
+        return poolIdx;
+    }
+    return 0;
+}
+
 size_t CacheTransBufferManager::computeTransferBufferSize(
     KVCacheManager::BaseKVCacheManager* cacheManager, std::optional<size_t> maxNumTokens, bool transferIndexerKCache)
 {
@@ -201,7 +224,7 @@ size_t CacheTransBufferManager::computeTransferBufferSize(
     }
     else
     {
-        dataType = cacheManager->getPrimaryPool(0)->getDataType();
+        dataType = cacheManager->getPrimaryPool(kvTransportPoolIdx(cacheManager))->getDataType();
     }
 
     auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock();
@@ -219,7 +242,7 @@ size_t CacheTransBufferManager::computeTransferBufferSize(
         }
         else
         {
-            auto primaryPool = cacheManager->getPrimaryPool(0);
+            auto primaryPool = cacheManager->getPrimaryPool(kvTransportPoolIdx(cacheManager));
             kvCacheByteSizePerTokenPerLayer
                 = primaryPool->getDimension<-1>() * primaryPool->getDimension<2>() * dataSize / tokensPerBlock;
         }
@@ -246,33 +269,43 @@ 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(),
+                              : cacheManager->getPrimaryPool(kvTransportPoolIdx(cacheManager))->getDataType(),
         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.
+    // not yet implemented.  Disagg picks the canonical KV transport pool's dtype (above),
+    // and the formatter's split/concat path moves every non-scale pool's blocks —
+    // including recurrent-states blocks on hybrid models — over the wire with that
+    // single dtype.  A pool with a differing dtype trips dtype/shape assertions
+    // mid-transfer; the sender swallows the exception and the receiver then waits
+    // forever for a ready signal (https://nvbugs/6479324).  Fail loudly at
+    // construction instead.  Block-scale / indexer-K pools legitimately have their
+    // own dtypes and travel through their own code paths.  The loop below iterates
+    // the absolute pool index domain — the same domain getPrimaryPool() uses —
+    // rather than the compacted count from getNumPools(false, false), whose
+    // index-domain mismatch previously made this check vacuous for hybrid models.
     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& blockManager = mCacheManager->getBlockManager();
+        auto const canonicalPoolIdx = kvTransportPoolIdx(mCacheManager);
+        auto const canonicalDtype = mCacheManager->getPrimaryPool(canonicalPoolIdx)->getDataType();
+        auto const numPools = blockManager.getNumPools();
+        for (SizeType32 i = 0; i < numPools; ++i)
         {
+            if (i == canonicalPoolIdx || blockManager.containsBlockScales(i))
+            {
+                continue;
+            }
             auto const dtypeI = mCacheManager->getPrimaryPool(i)->getDataType();
-            TLLM_CHECK_WITH_INFO(dtypeI == dtype0,
+            TLLM_CHECK_WITH_INFO(dtypeI == canonicalDtype,
                 "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));
+                "(canonical pool %d dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool "
+                "dtype dispatch in formatter.",
+                canonicalPoolIdx, static_cast<int>(canonicalDtype), 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");
@@ -68,6 +68,15 @@ class CacheTransBufferManager : public BaseTransBufferManager
         SizeType32 tokensPerBlock,
         std::optional<executor::CacheTransceiverConfig> const& cacheTransceiverConfig = std::nullopt);
 
+    /// @brief Index of the first pool that carries regular KV data for disagg transport.
+    /// @details Hybrid (Mamba/linear-attention) models host a recurrent-states pool whose sentinel
+    /// window sorts first, making it pool 0. Its contents travel through RnnCacheTransBufferManager
+    /// and may use a different dtype than the KV pools (e.g. BF16 SSM state alongside FP8 KV), so
+    /// it must not define the canonical wire dtype or per-token sizing of KV transfers.
+    /// Block-scale companion pools are skipped for the same reason. Falls back to pool 0 when
+    /// every pool is filtered out.
+    [[nodiscard]] static SizeType32 kvTransportPoolIdx(KVCacheManager::BaseKVCacheManager* cacheManager);
+
     /// @brief Get the KV cache manager.
     [[nodiscard]] KVCacheManager::BaseKVCacheManager* getCacheManager() const noexcept
     {

cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.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");
@@ -17,6 +17,7 @@
 
 #include "tensorrt_llm/batch_manager/cacheTransBuffer.h"
 #include "tensorrt_llm/batch_manager/kvCacheManager.h"
+#include "tensorrt_llm/common/assert.h"
 #include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/common/tllmDataType.h"
 #include "tensorrt_llm/executor/executor.h"
@@ -437,3 +438,72 @@ TEST_F(CacheTransBufferTest, TestForNullOptAndDefaultTransSize)
 }
 
 // TODO: pybinding
+
+TEST_F(CacheTransBufferTest, TestHybridRecurrentStatesPool)
+{
+    // Regression test for https://nvbugs/6479324: on hybrid (Mamba/linear-attention)
+    // models the recurrent-states pool's sentinel window sorts first, making it pool 0.
+    // The disagg transfer path must (a) not use it as the canonical KV transport pool
+    // for wire dtype / per-token sizing, and (b) fail loudly at construction when its
+    // dtype differs from the KV pools, since the formatter transfers its blocks over
+    // the KV wire with the canonical dtype (a silent mismatch previously surfaced as
+    // a permanent generation-side transfer stall).
+    auto constexpr numKvHeads = 2;
+    auto constexpr sizePerHead = 16;
+    auto constexpr tokensPerBlock = 4;
+    auto constexpr blocksInPrimaryPool = 8;
+    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,
+    };
+    using BlocksPerWindow = std::map<SizeType32, std::tuple<SizeType32, SizeType32>>;
+    auto const blocksPerWindow = BlocksPerWindow{
+        {recurrentStatesWindow, {blocksInPrimaryPool, blocksInSecondaryPool}},
+        {maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}},
+    };
+    auto const stream = std::make_shared<CudaStream>();
+
+    auto makeCacheManager = [&](tensorrt_llm::DataType recurrentStatesDtype)
+    {
+        auto const poolConfigurations = std::vector<PoolConfiguration>{
+            {recurrentStatesWindow, sizePerHead, recurrentStatesDtype},
+            {maxAttentionWindow, sizePerHead, tensorrt_llm::DataType::kFP8},
+        };
+        auto cacheManager = 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);
+        cacheManager->allocatePools(/*useUvm=*/false);
+        return cacheManager;
+    };
+
+    {
+        // Uniform dtype: the canonical KV transport pool must be the attention pool,
+        // and construction must succeed.
+        auto cacheManager = makeCacheManager(tensorrt_llm::DataType::kFP8);
+        auto const kvPoolIdx = CacheTransBufferManager::kvTransportPoolIdx(cacheManager.get());
+        EXPECT_EQ(cacheManager->getBlockManager().getPoolWindowSize(kvPoolIdx), maxAttentionWindow);
+        EXPECT_EQ(cacheManager->getPrimaryPool(kvPoolIdx)->getDataType(), tensorrt_llm::DataType::kFP8);
+        auto transBufferManager
+            = std::make_unique<CacheTransBufferManager>(cacheManager.get(), std::optional<size_t>{tokensPerBlock * 4});
+    }
+    {
+        // Differing recurrent-states dtype: unsupported on the disagg KV wire; must
+        // fail loudly at construction instead of stalling transfers at runtime.
+        auto cacheManager = makeCacheManager(tensorrt_llm::DataType::kHALF);
+        EXPECT_THROW(
+            std::make_unique<CacheTransBufferManager>(cacheManager.get(), std::optional<size_t>{tokensPerBlock * 4}),
+            tensorrt_llm::common::TllmException);
+    }
+}

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -1454,12 +1454,20 @@ def __init__(
             if mamba_layer_mask[layer_idx] else max_seq_len
             for layer_idx in self.pp_layers
         }
+        # The recurrent-states pool is a byte blob (allRecurrentStatesBytes); its
+        # declared dtype only matters when the KV dtype has container/scale
+        # semantics that would corrupt the blob (NVFP4: 2 elements per container
+        # plus block-scale companion pools).  Keep the manager dtype otherwise:
+        # the disagg KV wire transfers recurrent-window blocks with the KV pools'
+        # dtype, and a differing dtype breaks split/concat and stalls transfers
+        # (https://nvbugs/6479324).
         kwargs["pool_configurations"] = [
             PoolConfiguration(
                 window_size=window_size,
                 head_dim=head_dim,
                 dtype=torch_dtype_to_binding(self.ssm_state_dtype)
-                if window_size == recurrent_states_window else dtype,
+                if window_size == recurrent_states_window
+                and dtype == DataType.NVFP4 else dtype,
             ) for window_size in sorted(local_windows)
         ]
 

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -467,6 +467,30 @@ def _build_hybrid_with_mamba_layer(
     )
 
 
+@skip_no_cuda
+def test_cpp_hybrid_keeps_kv_dtype_for_recurrent_pool_on_non_nvfp4_kv_cache():
+    """https://nvbugs/6479324: the disagg KV wire transfers recurrent-window
+    blocks with the KV pools' dtype; a differing recurrent-pool dtype breaks
+    split/concat mid-transfer and stalls generation-side onboarding. The SSM
+    dtype override is only needed for NVFP4 KV caches, whose container/scale
+    semantics would corrupt the byte-blob recurrent state."""
+    mgr = _build_hybrid_with_mamba_layer(
+        dtype=DataType.FP8,
+        mamba_ssm_cache_dtype=torch.bfloat16,
+    )
+
+    expected_dtypes = [
+        (LinearCacheType.RECURRENT_STATES.value, DataType.FP8),
+        (128, DataType.FP8),
+    ]
+    assert [
+        (config.window_size, config.dtype) for config in mgr.pool_configurations
+    ] == expected_dtypes
+    assert [
+        (config.window_size, config.dtype) for config in mgr.impl.pool_configurations
+    ] == expected_dtypes
+
+
 @skip_no_cuda
 @pytest.mark.parametrize(
     "mamba_ssm_cache_dtype",