[pull] main from NVIDIA:main

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -237,7 +237,6 @@ class CacheTransceiver : public BaseCacheTransceiver
         executor::kv_cache::CacheState::AttentionType attentionType
         = executor::kv_cache::CacheState::AttentionType::kDEFAULT,
         std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig = std::nullopt,
-        rnn_state_manager::RnnStateManager* rnnStateManager = nullptr,
         std::vector<SizeType32> const& rnnLayerNumPerPP = {});
 
     CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector<SizeType32> numKvHeadsPerLayer,
@@ -246,11 +245,10 @@ class CacheTransceiver : public BaseCacheTransceiver
         executor::kv_cache::CacheState::AttentionType attentionType
         = executor::kv_cache::CacheState::AttentionType::kDEFAULT,
         std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig = std::nullopt,
-        rnn_state_manager::RnnStateManager* rnnStateManager = nullptr,
         std::vector<SizeType32> const& rnnLayerNumPerPP = {})
         : CacheTransceiver(cacheManager,
             executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig,
-            attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnStateManager, rnnLayerNumPerPP)
+            attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP)
     {
     }
 
@@ -306,7 +304,6 @@ class CacheTransceiver : public BaseCacheTransceiver
     std::vector<std::unique_ptr<kv_cache_manager::CacheTransBufferManager>> mCacheTransBufferManagers;
     std::vector<BaseTransBufferManager*> mCacheTransBufferManagerPtrs;
 
-    rnn_state_manager::RnnStateManager* mRnnStateManager{nullptr};
     // TODO(shreyasm): update this to use same container as kv by using base trans buffers instead
     std::unique_ptr<rnn_state_manager::RnnCacheTransBufferManager> mRnnCacheTransBufferManager{nullptr};
 

cpp/include/tensorrt_llm/batch_manager/rnnCacheFormatter.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");
@@ -34,27 +34,18 @@ class TransferSession;
 
 namespace rnn_state_manager
 {
-class RnnStateManager;
 class RnnCacheTransBufferManager;
 } // namespace rnn_state_manager
 
 /// @brief RNN Cache Formatter for formatting/unformatting RNN states during transfer.
-/// Supports two operating modes:
-/// - Slot mode: uses RnnStateManager (for CppMambaCacheManager, separate tensor storage)
-/// - Unified pool mode: uses BaseKVCacheManager (for CppMambaHybridCacheManager, block-indexed pool)
+/// Uses unified pool mode via BaseKVCacheManager (for CppMambaHybridCacheManager, block-indexed pool).
 class RnnCacheFormatter : public kv_cache_manager::BaseCacheFormatter
 {
 public:
     using SizeType32 = tensorrt_llm::runtime::SizeType32;
     using CacheState = executor::kv_cache::CacheState;
     using RequestIdType = tensorrt_llm::batch_manager::RequestIdType;
 
-    /// @brief Constructor for slot-based mode (CppMambaCacheManager with RnnStateManager).
-    /// @param rnnStateManager Pointer to the RNN state manager.
-    /// @param rnnCacheTransBufferManager Pointer to the RNN cache transfer buffer manager.
-    RnnCacheFormatter(rnn_state_manager::RnnStateManager* rnnStateManager,
-        rnn_state_manager::RnnCacheTransBufferManager* rnnCacheTransBufferManager);
-
     /// @brief Constructor for unified pool mode (CppMambaHybridCacheManager).
     /// @param kvCacheManager Pointer to the KV cache manager with unified pool.
     /// @param rnnCacheTransBufferManager Pointer to the RNN cache transfer buffer manager.
@@ -81,39 +72,13 @@ class RnnCacheFormatter : public kv_cache_manager::BaseCacheFormatter
         CacheState const& selfConfig, SizeType32 selfIdx, CacheState const& destConfig,
         std::vector<SizeType32> const& counterPartRanks) const override;
 
-    /// @brief Returns the KV cache manager (non-null in unified pool mode).
+    /// @brief Returns the KV cache manager.
     [[nodiscard]] kv_cache_manager::BaseKVCacheManager* getCacheManager() const noexcept override
     {
         return mKvCacheManager;
     }
 
-    /// @brief Get the RNN state manager (non-null in slot mode).
-    /// @return Pointer to the RNN state manager.
-    [[nodiscard]] rnn_state_manager::RnnStateManager* getRnnStateManager() const noexcept
-    {
-        return mRnnStateManager;
-    }
-
-    /// @brief Check if operating in unified pool mode.
-    [[nodiscard]] bool isUnifiedPoolMode() const noexcept
-    {
-        return mKvCacheManager != nullptr;
-    }
-
 private:
-    /// @brief Format logic for slot-based path (RnnStateManager).
-    void formatSlotMode(TransferSession& session);
-
-    /// @brief Unformat logic for slot-based path (RnnStateManager).
-    void unformatSlotMode(TransferSession& session);
-
-    /// @brief Format logic for unified pool path (BaseKVCacheManager).
-    void formatUnifiedPoolMode(TransferSession& session);
-
-    /// @brief Unformat logic for unified pool path (BaseKVCacheManager).
-    void unformatUnifiedPoolMode(TransferSession& session);
-
-    rnn_state_manager::RnnStateManager* mRnnStateManager{nullptr};
     rnn_state_manager::RnnCacheTransBufferManager* mRnnCacheTransBufferManager;
     kv_cache_manager::BaseKVCacheManager* mKvCacheManager{nullptr};
 };

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -293,9 +293,8 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
     std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
     executor::kv_cache::CacheState::AttentionType attentionType,
     std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig,
-    rnn_state_manager::RnnStateManager* rnnStateManager, std::vector<SizeType32> const& rnnLayerNumPerPP)
+    std::vector<SizeType32> const& rnnLayerNumPerPP)
     : mCacheTransceiverConfig{cacheTransceiverConfig}
-    , mRnnStateManager{rnnStateManager}
 {
     using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter;
     if (useMPI())
@@ -364,28 +363,9 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
             std::make_unique<kv_cache_manager::CacheTransBufferManager>(cacheManager, maxNumTokens, true));
     }
 
-    // RNN specific setup
-    if (mRnnStateManager != nullptr)
-    {
-        TLLM_LOG_DEBUG("Setting up RNN cache transfer components.");
-        TLLM_CHECK(!rnnLayerNumPerPP.empty());
-
-        mRnnCacheTransBufferManager
-            = std::make_unique<rnn_state_manager::RnnCacheTransBufferManager>(mRnnStateManager, maxNumTokens);
-
-        auto rnnModelCfg = mRnnStateManager->getRnnCacheStateModelConfig();
-
-        auto const convStateDataType = mRnnStateManager->getConvStateDataType();
-        auto const ssmStateDataType = mRnnStateManager->getSsmStateDataType();
-
-        mCacheState->setRnnConfig(rnnModelCfg, rnnLayerNumPerPP, convStateDataType, ssmStateDataType);
-
-        TLLM_LOG_INFO("RNN cache transfer components initialized.");
-    }
-
     // Unified pool path (CppMambaHybridCacheManager): build RnnModelConfig from
-    // LinearAttentionMetadata. Detected by rnnLayerNumPerPP set but no RnnStateManager.
-    if (mRnnStateManager == nullptr && !rnnLayerNumPerPP.empty())
+    // LinearAttentionMetadata. Detected by rnnLayerNumPerPP being non-empty.
+    if (!rnnLayerNumPerPP.empty())
     {
         auto const& blockManager = cacheManager->getBlockManager();
         auto const& linearMeta = blockManager.getLinearAttentionMetadata();
@@ -502,11 +482,6 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
 
     auto makeRnnFormatter = [this, cacheManager]() -> std::unique_ptr<RnnCacheFormatter>
     {
-        if (mRnnStateManager != nullptr && mRnnCacheTransBufferManager != nullptr)
-        {
-            // Slot-based path (CppMambaCacheManager)
-            return std::make_unique<RnnCacheFormatter>(mRnnStateManager, mRnnCacheTransBufferManager.get());
-        }
         // Unified pool path (CppMambaHybridCacheManager)
         if (mCacheState->hasRnnConfig() && mRnnCacheTransBufferManager != nullptr)
         {

cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp

@@ -54,8 +54,7 @@ void CacheTransferLayer::validateSupport(executor::DataTransceiverState const& p
 
     if (mRnnFormatter && selfHasRnn)
     {
-        // Both slot-based (CppMambaCacheManager) and unified pool (CppMambaHybridCacheManager)
-        // paths now use RnnCacheFormatter.
+        // Unified pool path (CppMambaHybridCacheManager) uses RnnCacheFormatter.
         if (peerHasRnn)
         {
             TLLM_CHECK_WITH_INFO(mRnnFormatter->inquireSupport(mCacheState, peerState.getCacheState().value()),