From c7b62ede8b9307de52b9947576f0064f6ed63991 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 23 Jul 2026 13:53:42 -0700 Subject: [PATCH 1/2] [https://nvbugs/6479324][fix] Keep KV dtype for hybrid recurrent-states pool to unstall disagg transfers Since #16304, CppMambaHybridCacheManager labels the recurrent-states pool with the SSM state dtype (e.g. BF16) while the KV pools stay FP8. The recurrent pool's sentinel window sorts first, so it is pool 0, and the disagg KV wire moves every non-scale pool's blocks - including recurrent-states blocks - with a single canonical dtype. The dtype mismatch trips the split/concat dtype assertion mid-transfer on the context side; the sender swallows the exception, the generation side waits forever for a ready signal, and the gen worker wedges with num_scheduled_requests=0 while requests pile up (mass HTTP 500s in the disagg stress tests). Fix: - Only override the recurrent-states pool dtype when the KV cache dtype is NVFP4, which is the case #16304 addressed (FP4 container/scale semantics corrupting the byte-blob state). For all other KV dtypes the pool keeps the manager dtype, restoring the pre-#16304 wire behavior. - Derive the disagg transfer buffer sizing and wire dtype from the first real KV pool (CacheTransBufferManager::kvTransportPoolIdx) instead of pool 0, which is the recurrent-states pool on hybrid models. - Fix the differing-dtype guard in CacheTransBufferManager: it iterated a compacted pool count while indexing the absolute pool domain, which made it vacuous for hybrid models. It now fails loudly at construction (e.g. NVFP4-KV hybrid + disagg) instead of stalling at runtime. Validated on 2xB200 with test_disaggregated_stress_test[...qwen3_5_4b_fp8_stress]: stalls before (gen worker pinned at 0/138 requests, ~3000 HTTP 500s), passes after (~9 min end-to-end including the accuracy phase). Signed-off-by: Brian Nguyen --- .../batch_manager/cacheFormatter.cpp | 16 ++-- .../batch_manager/cacheTransBuffer.cpp | 69 +++++++++++++----- .../batch_manager/cacheTransBuffer.h | 11 ++- .../batch_manager/cacheTransBufferTest.cpp | 73 ++++++++++++++++++- .../_torch/pyexecutor/mamba_cache_manager.py | 10 ++- .../executor/test_mamba_cache_manager.py | 24 ++++++ 6 files changed, 176 insertions(+), 27 deletions(-) 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..e812d78e97ac 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,73 @@ 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", From 6aff28b732f2a66e10bcdf62d14a5c552bcf84eb Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Sun, 2 Aug 2026 08:47:27 -0700 Subject: [PATCH 2/2] [https://nvbugs/6479324][fix] Apply clang-format Signed-off-by: Brian Nguyen --- .../unit_tests/batch_manager/cacheTransBufferTest.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp index e812d78e97ac..55d8d8a259f6 100644 --- a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp @@ -481,8 +481,7 @@ TEST_F(CacheTransBufferTest, TestHybridRecurrentStatesPool) 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, + CacheType::kSELF, std::nullopt, nullptr, /*enablePartialReuse=*/false, /*copyOnPartialReuse=*/true, nullptr, /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, /*indexerKCacheUseFp4=*/false, linearAttentionMetadata, poolConfigurations); cacheManager->allocatePools(/*useUvm=*/false); @@ -503,8 +502,8 @@ TEST_F(CacheTransBufferTest, TestHybridRecurrentStatesPool) // 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}), + EXPECT_THROW( + std::make_unique(cacheManager.get(), std::optional{tokensPerBlock * 4}), tensorrt_llm::common::TllmException); } }