diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 1d21b94a45ad..34cf8bc82def 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/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{ + {recurrentStatesWindow, sizePerHead, nvinfer1::DataType::kHALF}, + {maxAttentionWindow, sizePerHead, nvinfer1::DataType::kFP4}, + }; + auto const stream = std::make_shared(); + + KVCacheManager kvCacheManager(std::vector{0, numKvHeads}, sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, maxBeamWidth, std::vector{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(*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(*layerToPoolMapping); + EXPECT_EQ(layerToPoolMappingRange[0], 0); + EXPECT_EQ(layerToPoolMappingRange[2], 1); +} #endif TEST_F(KVCacheManagerTest, BlockManagerReuseTest) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 54be9fd150d5..b1558bec7529 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 2735622f6204..9c9935f11dac 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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 diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 5f73a1f86881..337ab4a24c88 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/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 ( diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index 5dcb363d9053..37c1fe0269d5 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import os import pathlib import subprocess @@ -13,7 +16,7 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ( - KVCacheManager, PeftCacheManager, + KVCacheManager, PeftCacheManager, _merge_kv_cache_pool_pointers, _warn_if_unsupported_v1_kv_cache_event_hash_algo) from tensorrt_llm.bindings import LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp @@ -66,6 +69,101 @@ def test_v1_kv_cache_event_hash_algo_no_warning_for_auto(): warning.assert_not_called() +class TestMergeKVCachePoolPointers(unittest.TestCase): + + def test_mixed_half_nvfp4_pools_align_scale_rows(self): + data_pointers = torch.tensor([[11, 12], [21, 22]]) + scale_pointers = torch.tensor([[31, 32]]) + layer_to_pool_mapping = torch.tensor([[0, 0], [1, 0]]) + + merged = _merge_kv_cache_pool_pointers( + data_pointers, + scale_pointers, + layer_to_pool_mapping, + [DataType.HALF, DataType.NVFP4], + ) + + expected = torch.tensor([ + [[11, 0], [12, 0]], + [[21, 31], [22, 32]], + ]) + self.assertEqual(merged.shape, (2, 2, 2)) + self.assertTrue(torch.equal(merged, expected)) + + def test_shared_nvfp4_pool_consumes_one_scale_row(self): + data_pointers = torch.tensor([[11, 12], [21, 22]]) + scale_pointers = torch.tensor([[31, 32]]) + layer_to_pool_mapping = torch.tensor([[0, 0], [0, 0], [1, 0]]) + + merged = _merge_kv_cache_pool_pointers( + data_pointers, + scale_pointers, + layer_to_pool_mapping, + [DataType.NVFP4, DataType.NVFP4, DataType.HALF], + ) + + expected = torch.tensor([ + [[11, 31], [12, 32]], + [[21, 0], [22, 0]], + ]) + self.assertTrue(torch.equal(merged, expected)) + + def test_physical_pool_ids_map_to_compact_data_rows(self): + data_pointers = torch.tensor([[11, 12], [21, 22]]) + scale_pointers = torch.tensor([[31, 32]]) + layer_to_pool_mapping = torch.tensor([[0, 0], [2, 0]]) + + merged = _merge_kv_cache_pool_pointers( + data_pointers, + scale_pointers, + layer_to_pool_mapping, + [DataType.NVFP4, DataType.HALF], + ) + + expected = torch.tensor([ + [[11, 31], [12, 32]], + [[21, 0], [22, 0]], + ]) + self.assertTrue(torch.equal(merged, expected)) + + +class TestKVCacheManagerPoolPointers(unittest.TestCase): + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_default_nvfp4_pool_configuration_merges_scale_pointers(self): + manager = KVCacheManager( + kv_cache_config=KvCacheConfig(max_tokens=64), + kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. + CacheType.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=128, + tokens_per_block=32, + max_seq_len=64, + max_batch_size=1, + mapping=Mapping(), + dtype=DataType.NVFP4, + ) + try: + data_pointers = manager.impl.get_block_pool_pointers() + scale_pointers = manager.impl.get_block_scale_pool_pointers() + + self.assertEqual(manager.pool_configurations, []) + self.assertEqual( + [(config.window_size, config.dtype) + for config in manager.impl.pool_configurations], + [(64, DataType.NVFP4)], + ) + self.assertTrue( + torch.equal(manager.kv_cache_pool_pointers[..., 0], + data_pointers)) + self.assertTrue( + torch.equal(manager.kv_cache_pool_pointers[..., 1], + scale_pointers)) + finally: + manager.shutdown() + + class TestResourceManager(unittest.TestCase): CPP_RESOURCES_DIR = os.path.join(str(root_dir), "cpp", "tests", "resources") CPP_DATA_DIR = os.path.join(CPP_RESOURCES_DIR, "data")