[https://nvbugs/6374873][fix] Allow fp4 KV Cache + non-FP4 Mamba State

cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

@@ -799,6 +799,90 @@ TEST_F(KVCacheManagerTest, FP4BlockScaleManagementTest)
     // The expected block size of pool 1 should be the number of FP4 elements / vectorSize.
     EXPECT_EQ(blockManager.getBlockSize(0) * numFp4EltsPerContainer / vectorSize, blockManager.getBlockSize(1));
 }
+
+TEST_F(KVCacheManagerTest, FP4AttentionWithHalfRecurrentStatesPoolTest)
+{
+    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, nvinfer1::DataType::kHALF},
+        {maxAttentionWindow, sizePerHead, nvinfer1::DataType::kFP4},
+    };
+    auto const stream = std::make_shared<tr::CudaStream>();
+
+    KVCacheManager kvCacheManager(std::vector<SizeType32>{0, numKvHeads}, sizePerHead, tokensPerBlock, blocksPerWindow,
+        maxNumSequences, maxBeamWidth, std::vector<SizeType32>{recurrentStatesWindow, maxAttentionWindow},
+        nvinfer1::DataType::kFP4,
+        /*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);
+    auto const& blockManager = kvCacheManager.getBlockManager();
+
+    EXPECT_EQ(blockManager.getDataTypeForWindow(recurrentStatesWindow), nvinfer1::DataType::kHALF);
+    EXPECT_EQ(blockManager.getDataTypeForWindow(maxAttentionWindow), nvinfer1::DataType::kFP4);
+
+    auto const& recurrentStatesPool = blockManager.getRecurrentStatesPool();
+    ASSERT_NE(recurrentStatesPool.primaryPtr, nullptr);
+    EXPECT_EQ(recurrentStatesPool.primaryPtr->getDataType(), nvinfer1::DataType::kHALF);
+    auto const recurrentStatesElementsPerBlock = recurrentStatesBytes / tc::getDTypeSize(nvinfer1::DataType::kHALF);
+    EXPECT_EQ(recurrentStatesPool.blockSize, recurrentStatesElementsPerBlock);
+
+    SizeType32 numRecurrentScalePools = 0;
+    SizeType32 numAttentionScalePools = 0;
+    for (SizeType32 poolIdx = 0; poolIdx < blockManager.getNumPools(); ++poolIdx)
+    {
+        if (!blockManager.containsBlockScales(poolIdx))
+        {
+            continue;
+        }
+        if (blockManager.getPoolWindowSize(poolIdx) == recurrentStatesWindow)
+        {
+            ++numRecurrentScalePools;
+        }
+        else if (blockManager.getPoolWindowSize(poolIdx) == maxAttentionWindow)
+        {
+            ++numAttentionScalePools;
+        }
+    }
+    EXPECT_EQ(numRecurrentScalePools, 0);
+    EXPECT_EQ(numAttentionScalePools, 1);
+
+    auto const blockPoolPointers = kvCacheManager.getBlockPoolPointers();
+    auto const blockScalePoolPointers = kvCacheManager.getBlockScalePoolPointers();
+    ASSERT_NE(blockPoolPointers, nullptr);
+    ASSERT_NE(blockScalePoolPointers, nullptr);
+    EXPECT_EQ(blockPoolPointers->getShape().d[0], 2);
+    EXPECT_EQ(blockScalePoolPointers->getShape().d[0], 1);
+    auto const blockScalePoolPointersRange = tr::BufferRange<void*>(*blockScalePoolPointers);
+    EXPECT_NE(blockScalePoolPointersRange[0], nullptr);
+    EXPECT_EQ(blockScalePoolPointersRange[1], nullptr);
+
+    auto const layerToPoolMapping = kvCacheManager.getLayerToPoolMapping();
+    ASSERT_NE(layerToPoolMapping, nullptr);
+    auto const layerToPoolMappingRange = tr::BufferRange<SizeType32>(*layerToPoolMapping);
+    EXPECT_EQ(layerToPoolMappingRange[0], 0);
+    EXPECT_EQ(layerToPoolMappingRange[2], 1);
+}
 #endif
 
 TEST_F(KVCacheManagerTest, BlockManagerReuseTest)

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -1626,6 +1626,21 @@ def __init__(
                     LinearCacheType.RECURRENT_STATES.
                     value if mamba_layer_mask[i] else max_seq_len)
 
+        recurrent_states_window = LinearCacheType.RECURRENT_STATES.value
+        local_windows = {
+            recurrent_states_window
+            if mamba_layer_mask[layer_idx] else max_seq_len
+            for layer_idx in self.pp_layers
+        }
+        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,
+            ) for window_size in sorted(local_windows)
+        ]
+
         # Normalize num_kv_heads to a per-layer list and zero out mamba
         # layer positions: those layers carry SSM/conv state instead of KV
         # heads, so the parent KV cache should not allocate KV head storage

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -114,6 +114,27 @@ def _warn_if_unsupported_v1_kv_cache_event_hash_algo(hash_algo: str) -> None:
         "event hashes.")
 
 
+def _merge_kv_cache_pool_pointers(
+    kv_cache_pool_pointers: torch.Tensor,
+    block_scale_pool_pointers: torch.Tensor,
+    layer_to_pool_mapping: torch.Tensor,
+    layer_pool_dtypes: Sequence["DataType"],
+) -> torch.Tensor:
+    """Align compact scale rows with data rows in C++ physical-pool order."""
+    dtype_by_physical_pool = dict(
+        zip(layer_to_pool_mapping[:, 0].tolist(), layer_pool_dtypes))
+    fp4_pool_indices = [
+        compact_pool_idx for compact_pool_idx, physical_pool_idx in enumerate(
+            sorted(dtype_by_physical_pool))
+        if dtype_by_physical_pool[physical_pool_idx] == DataType.NVFP4
+    ]
+
+    aligned_scale_pool_pointers = torch.zeros_like(kv_cache_pool_pointers)
+    aligned_scale_pool_pointers[fp4_pool_indices] = block_scale_pool_pointers
+    return torch.stack([kv_cache_pool_pointers, aligned_scale_pool_pointers],
+                       dim=-1)
+
+
 class BaseResourceManager(ABC):
 
     @abstractmethod
@@ -621,15 +642,31 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
 
         self.impl.allocate_pools(False)
         self.kv_cache_pool_pointers = self.impl.get_block_pool_pointers()
+        self.kv_cache_pool_mapping = self.impl.get_layer_to_pool_mapping()
         kv_cache_block_scale_pool_pointers = self.impl.get_block_scale_pool_pointers(
         )
         if kv_cache_block_scale_pool_pointers.numel() > 0:
-            self.kv_cache_pool_pointers = torch.stack([
-                self.kv_cache_pool_pointers, kv_cache_block_scale_pool_pointers
-            ],
-                                                      dim=-1)
+            # C++ reports one effective configuration per actual window,
+            # including manager-level defaults when none were supplied.
+            dtype_by_window = {
+                config.window_size: config.dtype
+                for config in self.impl.pool_configurations
+            }
+            # Match the local layer order used by C++ to build the pointer
+            # mapping. The Python window helpers additionally account for
+            # global PP layer IDs and therefore do not describe these rows.
+            layer_pool_dtypes = [
+                dtype_by_window[self.max_attention_window_vec[
+                    layer_offset % len(self.max_attention_window_vec)]]
+                for layer_offset in range(self.num_local_layers)
+            ]
+            self.kv_cache_pool_pointers = _merge_kv_cache_pool_pointers(
+                self.kv_cache_pool_pointers,
+                kv_cache_block_scale_pool_pointers,
+                self.kv_cache_pool_mapping,
+                layer_pool_dtypes,
+            )
 
-        self.kv_cache_pool_mapping = self.impl.get_layer_to_pool_mapping()
         self.num_pools = self.impl.num_pools
         self.max_blocks_per_seq = self.impl.max_blocks_per_seq
         self.enable_block_reuse = kv_cache_config.enable_block_reuse

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -19,8 +19,9 @@
     PythonMambaCacheManager,
     _get_mamba_hybrid_pool_size,
 )
-from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp
+from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
+from tensorrt_llm._utils import torch_dtype_to_binding
 from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig
 from tensorrt_llm.mapping import Mapping
@@ -401,13 +402,17 @@ def _build_hybrid_with_mamba_layer(
     enable_block_reuse=False,
     mamba_state_cache_interval=256,
     is_estimating_kv_cache=False,
+    dtype=DataType.HALF,
+    mamba_layer_mask=None,
+    attention_layer_mask=None,
+    mamba_ssm_cache_dtype=torch.float16,
 ):
     """Construct a real CppMambaHybridCacheManager with one mamba layer +
     one full-attention layer so the parent KVCacheManager goes through the
     linear-attention pool sizing path."""
     # Layer 0: mamba; Layer 1: full attention. Single rank, no MPI.
-    mamba_mask = [True, False]
-    attn_mask = [False, True]
+    mamba_mask = mamba_layer_mask or [True, False]
+    attn_mask = attention_layer_mask or [False, True]
     mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1)
     # Cap max_tokens to keep the real C++ pool allocation tiny.
     kv_cache_config = KvCacheConfig(
@@ -421,13 +426,13 @@ def _build_hybrid_with_mamba_layer(
         mamba_num_heads=4,
         mamba_n_groups=1,
         mamba_head_dim=8,
-        mamba_num_layers=1,
+        mamba_num_layers=sum(mamba_mask),
         mamba_layer_mask=mamba_mask,
         mamba_cache_dtype=torch.float16,
-        mamba_ssm_cache_dtype=torch.float16,
+        mamba_ssm_cache_dtype=mamba_ssm_cache_dtype,
         kv_cache_config=kv_cache_config,
         kv_cache_type=CacheTypeCpp.SELF,
-        num_layers=1,
+        num_layers=sum(attn_mask),
         num_kv_heads=4,
         head_dim=64,
         tokens_per_block=32,
@@ -437,8 +442,62 @@ def _build_hybrid_with_mamba_layer(
         spec_config=spec_config,
         layer_mask=attn_mask,
         is_estimating_kv_cache=is_estimating_kv_cache,
+        dtype=dtype,
+    )
+
+
+@skip_no_cuda
+@pytest.mark.parametrize(
+    "mamba_ssm_cache_dtype",
+    [torch.float16, torch.float32, torch.bfloat16],
+)
+def test_cpp_hybrid_passes_per_window_pool_dtypes_for_nvfp4_kv_cache(
+    mamba_ssm_cache_dtype,
+):
+    mgr = _build_hybrid_with_mamba_layer(
+        dtype=DataType.NVFP4,
+        mamba_ssm_cache_dtype=mamba_ssm_cache_dtype,
+    )
+    recurrent_pool_dtype = torch_dtype_to_binding(mamba_ssm_cache_dtype)
+
+    expected_dtypes = [
+        (LinearCacheType.RECURRENT_STATES.value, recurrent_pool_dtype),
+        (128, DataType.NVFP4),
+    ]
+    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
+    assert mgr._layer_to_pool_idx == {0: 0, 1: 1}
+    assert mgr.recurrent_states_pool_index == 0
+    assert mgr.impl.get_recurrent_states_pool().dtype == mamba_ssm_cache_dtype
+
+    compact_scale_pointers = mgr.impl.get_block_scale_pool_pointers()
+    assert mgr.impl.get_block_pool_pointers().shape == (2, 2)
+    assert compact_scale_pointers.shape == (1, 2)
+    assert mgr.kv_cache_pool_pointers.shape == (2, 2, 2)
+    assert torch.count_nonzero(mgr.kv_cache_pool_pointers[0, :, 1]) == 0
+    assert torch.equal(mgr.kv_cache_pool_pointers[1, :, 1], compact_scale_pointers[0])
+
+
+@skip_no_cuda
+def test_cpp_hybrid_merges_compact_scale_rows_with_unmanaged_layers():
+    mgr = _build_hybrid_with_mamba_layer(
+        dtype=DataType.NVFP4,
+        mamba_layer_mask=[True, False, True, False],
+        attention_layer_mask=[False, False, False, True],
     )
 
+    assert mgr.pp_layers == [0, 2, 3]
+    assert mgr.kv_cache_pool_mapping[:, 0].tolist() == [0, 0, 1]
+    compact_scale_pointers = mgr.impl.get_block_scale_pool_pointers()
+    assert compact_scale_pointers.shape == (1, 2)
+    assert mgr.kv_cache_pool_pointers.shape == (2, 2, 2)
+    assert torch.count_nonzero(mgr.kv_cache_pool_pointers[0, :, 1]) == 0
+    assert torch.equal(mgr.kv_cache_pool_pointers[1, :, 1], compact_scale_pointers[0])
+
 
 @skip_no_cuda
 def test_cpp_hybrid_recurrent_pool_reserves_cuda_graph_padding_slot():
@@ -687,6 +746,14 @@ def test_cpp_hybrid_zero_local_mamba_layers():
     # On the early-exit branch, num_layers is forwarded as-is.
     assert mgr.num_layers == 4
     assert mgr.num_local_layers == 2
+    assert all(
+        config.window_size != LinearCacheType.RECURRENT_STATES.value
+        for config in mgr.pool_configurations
+    )
+    assert all(
+        config.window_size != LinearCacheType.RECURRENT_STATES.value
+        for config in mgr.impl.pool_configurations
+    )
 
     # No mamba-only state was allocated.
     for attr in (