diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 9dc44531d409..ece48e2c0cb5 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/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) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index e06198f9ea60..c7d0c820a6ae 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"); @@ -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 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,7 +269,7 @@ CacheTransBufferManager::CacheTransBufferManager( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional 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} @@ -254,25 +277,35 @@ CacheTransBufferManager::CacheTransBufferManager( // 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(dtype0), i, static_cast(dtypeI)); + "(canonical pool %d dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool " + "dtype dispatch in formatter.", + canonicalPoolIdx, static_cast(canonicalDtype), 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..5f65b930b8a9 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"); @@ -68,6 +68,15 @@ class CacheTransBufferManager : public BaseTransBufferManager SizeType32 tokensPerBlock, std::optional 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 { diff --git a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp index 2fa0477d2352..55d8d8a259f6 100644 --- a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp +++ b/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>; + auto const blocksPerWindow = BlocksPerWindow{ + {recurrentStatesWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, + {maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, + }; + auto const stream = std::make_shared(); + + auto makeCacheManager = [&](tensorrt_llm::DataType recurrentStatesDtype) + { + auto const poolConfigurations = std::vector{ + {recurrentStatesWindow, sizePerHead, recurrentStatesDtype}, + {maxAttentionWindow, sizePerHead, tensorrt_llm::DataType::kFP8}, + }; + auto cacheManager = 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); + 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(cacheManager.get(), std::optional{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(cacheManager.get(), std::optional{tokensPerBlock * 4}), + tensorrt_llm::common::TllmException); + } +} diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 16fe22d1c777..d236ffa3733e 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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) ] diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 0291ef3e74dd..e6829c833b80 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/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",