From b6f23ee84863344bf4587ceadc74ae0337dbcadf Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 9 Sep 2026 08:13:31 +0000 Subject: [PATCH 1/4] [https://nvbugs/6739081][fix] Preserve per-layer KV page addressing for mixed head sizes Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 27 ++++-- .../test_kv_cache_v2_extra_buffers.py | 3 +- .../executor/test_per_layer_head_dim.py | 83 +++++++++++++++++-- 3 files changed, 96 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index cc598cbf67ca..4fe5655b2f01 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1369,11 +1369,22 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._pool_layer_ids_by_role.setdefault( (pool_id, buffer_id.role), buffer_id.layer_id ) - # num_pools is the logical layer-group count. With SWA scratch reuse, - # scratch slot IDs are only valid with per-layer page indices, so the - # attention op sees one virtual pool per local layer while the - # underlying manager can still group layers. - if self.enable_swa_scratch_reuse: + # A lifecycle group can contain different page sizes in separate + # physical pools (e.g. Gemma's heterogeneous head dimensions). Those + # layers cannot share a pool pointer or page-index scale. Use the + # per-layer conversion path, also required for SWA scratch slots, + # while keeping the underlying lifecycle grouping unchanged. + self._use_per_layer_page_tables = self.enable_swa_scratch_reuse or any( + len( + { + self.impl.get_page_stride(layer_id, self._get_pool_roles(pool_id)[0]) + for layer_id in layer_ids + } + ) + > 1 + for pool_id, layer_ids in enumerate(self.impl.layer_grouping) + ) + if self._use_per_layer_page_tables: self.num_attention_op_pools = self.num_local_layers else: self.num_attention_op_pools = self.num_pools @@ -1492,7 +1503,7 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] - if self.enable_swa_scratch_reuse: + if self._use_per_layer_page_tables: for layer_id in typed_range(LayerId(self.num_local_layers)): pool_id = self.impl.get_layer_group_id(layer_id) role_a, _ = self._get_pool_roles(pool_id) @@ -1658,7 +1669,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: pin_memory=prefer_pinned(), device="cpu", ) - if self.enable_swa_scratch_reuse: + if self._use_per_layer_page_tables: self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) def _kv_pool_mapping_offset( @@ -4363,7 +4374,7 @@ def copy_batch_block_offsets( copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) assert copy_idx.shape[0] == num_seqs - if self.enable_swa_scratch_reuse: + if self._use_per_layer_page_tables: self._copy_batch_block_offsets_per_layer( dst_tensor, request_ids, copy_idx, num_contexts, num_seqs ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py index d518d288581e..4dafe2804e85 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py @@ -158,7 +158,8 @@ def test_default_hook_keeps_nvfp4_scale_buffers(self): del mgr def test_page_table_uses_physical_pool_representative(self): - mgr = KVCacheManagerV2(**_make_kwargs(head_dim=[64, 192, 64, 192])) + # Equal page sizes exercise the shared page-table path. + mgr = KVCacheManagerV2(**_make_kwargs()) real_impl = mgr.impl try: pool_id = 0 diff --git a/tests/unittest/_torch/executor/test_per_layer_head_dim.py b/tests/unittest/_torch/executor/test_per_layer_head_dim.py index 851f56377081..51be7f0c711f 100644 --- a/tests/unittest/_torch/executor/test_per_layer_head_dim.py +++ b/tests/unittest/_torch/executor/test_per_layer_head_dim.py @@ -16,6 +16,7 @@ import gc import unittest +import pytest import torch import tensorrt_llm @@ -27,6 +28,7 @@ ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode DataType = tensorrt_llm.bindings.DataType CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -50,6 +52,7 @@ def _create_kv_cache_manager_v2( kv_cache_config = KvCacheConfigV2( max_tokens=max_tokens, enable_block_reuse=False, + dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) return KVCacheManagerV2( kv_cache_config, @@ -196,19 +199,13 @@ def test_per_layer_head_dim_with_equal_buffer_sizes(self): self.assertEqual(bytes_0, 4 * 64 * 2) self.assertEqual(bytes_1, 2 * 128 * 2) self.assertEqual(bytes_0, bytes_1) + self.assertEqual(mgr.num_attention_op_pools, mgr.num_pools) finally: mgr.shutdown() class TestPerLayerHeadDimHeterogeneous(unittest.TestCase): - """Tests with different buffer sizes per layer. - - All tests that create managers with heterogeneous per-layer buffer sizes - are consolidated into a single test method. This is necessary because - CUDA virtual memory addresses from destroyed managers may not be fully - reclaimed, causing subsequent managers to use addresses that lead to - large page offsets (exceeding int32 in pool mapping tensors). - """ + """Tests with different buffer sizes per layer.""" def setUp(self): torch.cuda.init() @@ -336,5 +333,75 @@ def test_per_layer_head_dim_heterogeneous(self): mgr_mixed.shutdown() +@pytest.mark.parametrize("dtype", [DataType.HALF, DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("is_gen", [False, True]) +def test_heterogeneous_attention_page_addresses(dtype, is_gen): + # Gemma's short-sequence configuration puts both head sizes in one + # lifecycle group. Unequal layer counts also give them different page + # index scales, so sharing a group's pointer or scale is incorrect. + mgr = _create_kv_cache_manager_v2( + num_layers=6, + head_dim=[256] * 5 + [512], + dtype=dtype, + ) + try: + request_ids = [11, 22, 33] + token_nums = [1, 9, 17] + requests = mgr.add_dummy_requests(request_ids, token_nums, is_gen=is_gen) + assert requests is not None + request_ids.reverse() + token_nums.reverse() + block_offsets = torch.empty( + mgr.num_attention_op_pools, + len(request_ids), + 2, + mgr.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + mgr.copy_batch_block_offsets( + block_offsets, + request_ids, + beam_width=1, + num_contexts=0 if is_gen else len(request_ids), + num_seqs=len(request_ids), + ) + torch.cuda.synchronize() + block_offsets = block_offsets.cpu() + for layer_id in range(mgr.num_local_layers): + pool_id, layer_offset = mgr.kv_cache_pool_mapping[layer_id].tolist() + lifecycle_id = mgr.impl.get_layer_group_id(layer_id) + roles = [(Role.KEY, Role.VALUE)] + if dtype == DataType.NVFP4: + roles.append((Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)) + for buffer_idx, role_pair in enumerate(roles): + pool_pointer = mgr.kv_cache_pool_pointers[pool_id, 0] + if dtype == DataType.NVFP4: + pool_pointer = pool_pointer[buffer_idx] + for kv_idx, role in enumerate(role_pair): + stride = mgr.impl.get_page_stride(layer_id, role) + layer_base = mgr.impl.get_mem_pool_base_address( + layer_id, role, PageIndexMode.SHARED + ) + scale = mgr.impl.get_page_index_scale(layer_id, role) + for seq_idx, (req_id, token_num) in enumerate(zip(request_ids, token_nums)): + base_indices = mgr.kv_cache_map[req_id].get_base_page_indices(lifecycle_id) + num_blocks = (token_num + mgr.tokens_per_block - 1) // mgr.tokens_per_block + for block_idx in range(num_blocks): + page_idx = int(block_offsets[pool_id, seq_idx, kv_idx, block_idx]) + actual = ( + int(pool_pointer) + + (layer_offset * mgr.kv_factor + page_idx) * stride + ) + expected = layer_base + int(base_indices[block_idx]) * scale * stride + assert actual == expected, ( + f"layer={layer_id}, role={role}, request={req_id}, " + f"block={block_idx}: attention address {actual:#x} " + f"does not match allocated address {expected:#x}" + ) + finally: + mgr.shutdown() + + if __name__ == "__main__": unittest.main() From dba57c5ad1b67d6b461ac33ce8c34c6a9ed6f247 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:27:24 +0000 Subject: [PATCH 2/4] [https://nvbugs/6739081][test] Cover heterogeneous page tables alongside pool mapping tests Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../test_kv_cache_v2_extra_buffers.py | 137 ++++++++++++++---- .../executor/test_per_layer_head_dim.py | 81 +---------- 2 files changed, 118 insertions(+), 100 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py index 4dafe2804e85..cf4d40f3f30f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -# Focused regression tests for the KVCacheManagerV2 per-layer extra-buffer -# registration hook used by MiniMax-M3 sparse index-K support. The hook lets +# Regression tests for KVCacheManagerV2 page-table mapping and the per-layer +# extra-buffer hook used by MiniMax-M3 sparse index-K support. The hook lets # subclasses register additional per-layer ``BufferConfig`` entries (e.g. # ``Role.INDEX_KEY``) alongside the standard K/V/NVFP4 scale buffers without # disturbing the existing K/V wiring or lifecycle. @@ -14,6 +14,7 @@ import unittest from unittest.mock import Mock +import pytest import torch import tensorrt_llm @@ -158,32 +159,41 @@ def test_default_hook_keeps_nvfp4_scale_buffers(self): del mgr def test_page_table_uses_physical_pool_representative(self): - # Equal page sizes exercise the shared page-table path. - mgr = KVCacheManagerV2(**_make_kwargs()) - real_impl = mgr.impl - try: - pool_id = 0 - physical_layer = mgr._pool_layer_ids_by_role[(pool_id, Role.KEY)] - other_layer = next( - layer_id - for layer_id in real_impl.layer_grouping[pool_id] - if int(layer_id) != int(physical_layer) - ) - impl_proxy = Mock(wraps=real_impl) - impl_proxy.layer_grouping = ((other_layer, physical_layer),) - mgr.impl = impl_proxy + for head_dim in (128, [64, 192, 64, 192]): + with self.subTest(head_dim=head_dim): + mgr = KVCacheManagerV2(**_make_kwargs(head_dim=head_dim)) + real_impl = mgr.impl + try: + pool_id = 0 + physical_layer = mgr._pool_layer_ids_by_role[(pool_id, Role.KEY)] + layers = real_impl.layer_grouping[pool_id] + other_layer = next(layer for layer in layers if layer != physical_layer) + impl_proxy = Mock(wraps=real_impl) + impl_proxy.layer_grouping = ( + (other_layer, *(layer for layer in layers if layer != other_layer)), + ) + mgr.impl = impl_proxy - mgr._prepare_page_table_tensor(index_mapper_capacity=1) + mgr._prepare_page_table_tensor(index_mapper_capacity=1) - self.assertEqual( - impl_proxy.get_mem_pool_base_address.call_args_list[0].args, - (physical_layer, Role.KEY, PageIndexMode.SHARED), - ) - impl_proxy.get_page_index_scale.assert_called_once_with(physical_layer, Role.KEY) - finally: - mgr.impl = real_impl - mgr.shutdown() - del mgr + # Both paths still use a physical representative for group + # metadata. Per-layer pointer lookups can precede it. + shared_key_calls = [ + call.args + for call in impl_proxy.get_mem_pool_base_address.call_args_list + if call.args[1:] == (Role.KEY, PageIndexMode.SHARED) + ] + self.assertEqual( + shared_key_calls[0], + (physical_layer, Role.KEY, PageIndexMode.SHARED), + ) + impl_proxy.get_page_index_scale.assert_called_once_with( + physical_layer, Role.KEY + ) + finally: + mgr.impl = real_impl + mgr.shutdown() + del mgr def test_subclass_registers_index_key_only_on_sparse_layers(self): # Sparse layer convention: layers 0-2 dense (no INDEX_KEY), 3+ sparse @@ -273,6 +283,81 @@ def test_duplicate_role_against_standard_buffer_asserts(self): _DuplicateRoleV2(**_make_kwargs(num_layers=4)) +@pytest.mark.parametrize("dtype", [DataType.HALF, DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("is_gen", [False, True]) +def test_heterogeneous_page_tables_match_allocated_addresses(dtype, is_gen): + # Gemma's short-sequence configuration puts both head sizes in one + # lifecycle group. Unequal layer counts also give them different page + # index scales, so sharing a group's pointer or scale is incorrect. + kwargs = _make_kwargs(num_layers=6, head_dim=[256] * 5 + [512], dtype=dtype) + if dtype == DataType.NVFP4: + kwargs["kv_cache_config"] = KvCacheConfigV2( + max_tokens=2048, enable_block_reuse=False, dtype="nvfp4" + ) + mgr = KVCacheManagerV2(**kwargs) + try: + request_ids = [11, 22, 33] + token_nums = [1, 9, 17] + requests = mgr.add_dummy_requests(request_ids, token_nums, is_gen=is_gen) + assert requests is not None + request_ids.reverse() + token_nums.reverse() + block_offsets = torch.empty( + mgr.num_attention_op_pools, + len(request_ids), + 2, + mgr.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + mgr.copy_batch_block_offsets( + block_offsets, + request_ids, + beam_width=1, + num_contexts=0 if is_gen else len(request_ids), + num_seqs=len(request_ids), + ) + torch.cuda.synchronize() + block_offsets = block_offsets.cpu() + for layer_id in range(mgr.num_local_layers): + pool_id, layer_offset = mgr.kv_cache_pool_mapping[layer_id].tolist() + lifecycle_id = mgr.impl.get_layer_group_id(layer_id) + roles = [(Role.KEY, Role.VALUE)] + if dtype == DataType.NVFP4: + roles.append((Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)) + for buffer_idx, role_pair in enumerate(roles): + pool_pointer = mgr.kv_cache_pool_pointers[pool_id, 0] + if dtype == DataType.NVFP4: + pool_pointer = pool_pointer[buffer_idx] + for kv_idx, role in enumerate(role_pair): + stride = mgr.impl.get_page_stride(layer_id, role) + layer_base = mgr.impl.get_mem_pool_base_address( + layer_id, role, PageIndexMode.SHARED + ) + scale = mgr.impl.get_page_index_scale(layer_id, role) + for seq_idx, (req_id, token_num) in enumerate(zip(request_ids, token_nums)): + base_indices = mgr.kv_cache_map[req_id].get_base_page_indices(lifecycle_id) + num_blocks = (token_num + mgr.tokens_per_block - 1) // mgr.tokens_per_block + page_indices = block_offsets[pool_id, seq_idx, kv_idx, :num_blocks].long() + actual = ( + int(pool_pointer) + + (layer_offset * mgr.kv_factor + page_indices) * stride + ) + allocated_indices = torch.tensor( + list(base_indices[:num_blocks]), dtype=torch.int64 + ) + expected = layer_base + allocated_indices * scale * stride + torch.testing.assert_close( + actual, + expected, + rtol=0, + atol=0, + msg=f"layer={layer_id}, role={role}, request={req_id}", + ) + finally: + mgr.shutdown() + + class TestIndexKeyBufferAccessor(unittest.TestCase): """CUDA/GPU regressions for :meth:`KVCacheManagerV2.get_index_k_buffer`. diff --git a/tests/unittest/_torch/executor/test_per_layer_head_dim.py b/tests/unittest/_torch/executor/test_per_layer_head_dim.py index 51be7f0c711f..5842ce3f903f 100644 --- a/tests/unittest/_torch/executor/test_per_layer_head_dim.py +++ b/tests/unittest/_torch/executor/test_per_layer_head_dim.py @@ -16,7 +16,6 @@ import gc import unittest -import pytest import torch import tensorrt_llm @@ -28,7 +27,6 @@ ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode DataType = tensorrt_llm.bindings.DataType CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -52,7 +50,6 @@ def _create_kv_cache_manager_v2( kv_cache_config = KvCacheConfigV2( max_tokens=max_tokens, enable_block_reuse=False, - dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) return KVCacheManagerV2( kv_cache_config, @@ -205,7 +202,13 @@ def test_per_layer_head_dim_with_equal_buffer_sizes(self): class TestPerLayerHeadDimHeterogeneous(unittest.TestCase): - """Tests with different buffer sizes per layer.""" + """Tests with different buffer sizes per layer. + + These cases were originally combined to avoid large offsets between + separately allocated pools. Heterogeneous pages now use per-layer base + pointers and page indices, so their mappings no longer depend on the + distance between pools. Explicit GPU cleanup is retained. + """ def setUp(self): torch.cuda.init() @@ -333,75 +336,5 @@ def test_per_layer_head_dim_heterogeneous(self): mgr_mixed.shutdown() -@pytest.mark.parametrize("dtype", [DataType.HALF, DataType.FP8, DataType.NVFP4]) -@pytest.mark.parametrize("is_gen", [False, True]) -def test_heterogeneous_attention_page_addresses(dtype, is_gen): - # Gemma's short-sequence configuration puts both head sizes in one - # lifecycle group. Unequal layer counts also give them different page - # index scales, so sharing a group's pointer or scale is incorrect. - mgr = _create_kv_cache_manager_v2( - num_layers=6, - head_dim=[256] * 5 + [512], - dtype=dtype, - ) - try: - request_ids = [11, 22, 33] - token_nums = [1, 9, 17] - requests = mgr.add_dummy_requests(request_ids, token_nums, is_gen=is_gen) - assert requests is not None - request_ids.reverse() - token_nums.reverse() - block_offsets = torch.empty( - mgr.num_attention_op_pools, - len(request_ids), - 2, - mgr.max_blocks_per_seq, - dtype=torch.int32, - device="cuda", - ) - mgr.copy_batch_block_offsets( - block_offsets, - request_ids, - beam_width=1, - num_contexts=0 if is_gen else len(request_ids), - num_seqs=len(request_ids), - ) - torch.cuda.synchronize() - block_offsets = block_offsets.cpu() - for layer_id in range(mgr.num_local_layers): - pool_id, layer_offset = mgr.kv_cache_pool_mapping[layer_id].tolist() - lifecycle_id = mgr.impl.get_layer_group_id(layer_id) - roles = [(Role.KEY, Role.VALUE)] - if dtype == DataType.NVFP4: - roles.append((Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)) - for buffer_idx, role_pair in enumerate(roles): - pool_pointer = mgr.kv_cache_pool_pointers[pool_id, 0] - if dtype == DataType.NVFP4: - pool_pointer = pool_pointer[buffer_idx] - for kv_idx, role in enumerate(role_pair): - stride = mgr.impl.get_page_stride(layer_id, role) - layer_base = mgr.impl.get_mem_pool_base_address( - layer_id, role, PageIndexMode.SHARED - ) - scale = mgr.impl.get_page_index_scale(layer_id, role) - for seq_idx, (req_id, token_num) in enumerate(zip(request_ids, token_nums)): - base_indices = mgr.kv_cache_map[req_id].get_base_page_indices(lifecycle_id) - num_blocks = (token_num + mgr.tokens_per_block - 1) // mgr.tokens_per_block - for block_idx in range(num_blocks): - page_idx = int(block_offsets[pool_id, seq_idx, kv_idx, block_idx]) - actual = ( - int(pool_pointer) - + (layer_offset * mgr.kv_factor + page_idx) * stride - ) - expected = layer_base + int(base_indices[block_idx]) * scale * stride - assert actual == expected, ( - f"layer={layer_id}, role={role}, request={req_id}, " - f"block={block_idx}: attention address {actual:#x} " - f"does not match allocated address {expected:#x}" - ) - finally: - mgr.shutdown() - - if __name__ == "__main__": unittest.main() From 596387be73bf14c49c6aa2097f0fecedfe990055 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:44:43 +0000 Subject: [PATCH 3/4] [https://nvbugs/6739081][test] Move heterogeneous page-table regression to head-dimension tests Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../test_kv_cache_v2_extra_buffers.py | 80 +------------------ .../executor/test_per_layer_head_dim.py | 73 +++++++++++++++++ 2 files changed, 75 insertions(+), 78 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py index cf4d40f3f30f..81622c5f0ee6 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # -# Regression tests for KVCacheManagerV2 page-table mapping and the per-layer -# extra-buffer hook used by MiniMax-M3 sparse index-K support. The hook lets +# Focused regression tests for the KVCacheManagerV2 per-layer extra-buffer +# registration hook used by MiniMax-M3 sparse index-K support. The hook lets # subclasses register additional per-layer ``BufferConfig`` entries (e.g. # ``Role.INDEX_KEY``) alongside the standard K/V/NVFP4 scale buffers without # disturbing the existing K/V wiring or lifecycle. @@ -14,7 +14,6 @@ import unittest from unittest.mock import Mock -import pytest import torch import tensorrt_llm @@ -283,81 +282,6 @@ def test_duplicate_role_against_standard_buffer_asserts(self): _DuplicateRoleV2(**_make_kwargs(num_layers=4)) -@pytest.mark.parametrize("dtype", [DataType.HALF, DataType.FP8, DataType.NVFP4]) -@pytest.mark.parametrize("is_gen", [False, True]) -def test_heterogeneous_page_tables_match_allocated_addresses(dtype, is_gen): - # Gemma's short-sequence configuration puts both head sizes in one - # lifecycle group. Unequal layer counts also give them different page - # index scales, so sharing a group's pointer or scale is incorrect. - kwargs = _make_kwargs(num_layers=6, head_dim=[256] * 5 + [512], dtype=dtype) - if dtype == DataType.NVFP4: - kwargs["kv_cache_config"] = KvCacheConfigV2( - max_tokens=2048, enable_block_reuse=False, dtype="nvfp4" - ) - mgr = KVCacheManagerV2(**kwargs) - try: - request_ids = [11, 22, 33] - token_nums = [1, 9, 17] - requests = mgr.add_dummy_requests(request_ids, token_nums, is_gen=is_gen) - assert requests is not None - request_ids.reverse() - token_nums.reverse() - block_offsets = torch.empty( - mgr.num_attention_op_pools, - len(request_ids), - 2, - mgr.max_blocks_per_seq, - dtype=torch.int32, - device="cuda", - ) - mgr.copy_batch_block_offsets( - block_offsets, - request_ids, - beam_width=1, - num_contexts=0 if is_gen else len(request_ids), - num_seqs=len(request_ids), - ) - torch.cuda.synchronize() - block_offsets = block_offsets.cpu() - for layer_id in range(mgr.num_local_layers): - pool_id, layer_offset = mgr.kv_cache_pool_mapping[layer_id].tolist() - lifecycle_id = mgr.impl.get_layer_group_id(layer_id) - roles = [(Role.KEY, Role.VALUE)] - if dtype == DataType.NVFP4: - roles.append((Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)) - for buffer_idx, role_pair in enumerate(roles): - pool_pointer = mgr.kv_cache_pool_pointers[pool_id, 0] - if dtype == DataType.NVFP4: - pool_pointer = pool_pointer[buffer_idx] - for kv_idx, role in enumerate(role_pair): - stride = mgr.impl.get_page_stride(layer_id, role) - layer_base = mgr.impl.get_mem_pool_base_address( - layer_id, role, PageIndexMode.SHARED - ) - scale = mgr.impl.get_page_index_scale(layer_id, role) - for seq_idx, (req_id, token_num) in enumerate(zip(request_ids, token_nums)): - base_indices = mgr.kv_cache_map[req_id].get_base_page_indices(lifecycle_id) - num_blocks = (token_num + mgr.tokens_per_block - 1) // mgr.tokens_per_block - page_indices = block_offsets[pool_id, seq_idx, kv_idx, :num_blocks].long() - actual = ( - int(pool_pointer) - + (layer_offset * mgr.kv_factor + page_indices) * stride - ) - allocated_indices = torch.tensor( - list(base_indices[:num_blocks]), dtype=torch.int64 - ) - expected = layer_base + allocated_indices * scale * stride - torch.testing.assert_close( - actual, - expected, - rtol=0, - atol=0, - msg=f"layer={layer_id}, role={role}, request={req_id}", - ) - finally: - mgr.shutdown() - - class TestIndexKeyBufferAccessor(unittest.TestCase): """CUDA/GPU regressions for :meth:`KVCacheManagerV2.get_index_k_buffer`. diff --git a/tests/unittest/_torch/executor/test_per_layer_head_dim.py b/tests/unittest/_torch/executor/test_per_layer_head_dim.py index 5842ce3f903f..38b1e0a12ef3 100644 --- a/tests/unittest/_torch/executor/test_per_layer_head_dim.py +++ b/tests/unittest/_torch/executor/test_per_layer_head_dim.py @@ -16,6 +16,7 @@ import gc import unittest +import pytest import torch import tensorrt_llm @@ -27,6 +28,7 @@ ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2 from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode DataType = tensorrt_llm.bindings.DataType CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -50,6 +52,7 @@ def _create_kv_cache_manager_v2( kv_cache_config = KvCacheConfigV2( max_tokens=max_tokens, enable_block_reuse=False, + dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) return KVCacheManagerV2( kv_cache_config, @@ -336,5 +339,75 @@ def test_per_layer_head_dim_heterogeneous(self): mgr_mixed.shutdown() +@pytest.mark.parametrize("dtype", [DataType.HALF, DataType.FP8, DataType.NVFP4]) +@pytest.mark.parametrize("is_gen", [False, True]) +def test_heterogeneous_page_tables_match_allocated_addresses(dtype, is_gen): + # Gemma's short-sequence configuration puts both head sizes in one + # lifecycle group. Unequal layer counts also give them different page + # index scales, so sharing a group's pointer or scale is incorrect. + mgr = _create_kv_cache_manager_v2(num_layers=6, head_dim=[256] * 5 + [512], dtype=dtype) + try: + request_ids = [11, 22, 33] + token_nums = [1, 9, 17] + requests = mgr.add_dummy_requests(request_ids, token_nums, is_gen=is_gen) + assert requests is not None + request_ids.reverse() + token_nums.reverse() + block_offsets = torch.empty( + mgr.num_attention_op_pools, + len(request_ids), + 2, + mgr.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + mgr.copy_batch_block_offsets( + block_offsets, + request_ids, + beam_width=1, + num_contexts=0 if is_gen else len(request_ids), + num_seqs=len(request_ids), + ) + torch.cuda.synchronize() + block_offsets = block_offsets.cpu() + for layer_id in range(mgr.num_local_layers): + pool_id, layer_offset = mgr.kv_cache_pool_mapping[layer_id].tolist() + lifecycle_id = mgr.impl.get_layer_group_id(layer_id) + roles = [(Role.KEY, Role.VALUE)] + if dtype == DataType.NVFP4: + roles.append((Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)) + for buffer_idx, role_pair in enumerate(roles): + pool_pointer = mgr.kv_cache_pool_pointers[pool_id, 0] + if dtype == DataType.NVFP4: + pool_pointer = pool_pointer[buffer_idx] + for kv_idx, role in enumerate(role_pair): + stride = mgr.impl.get_page_stride(layer_id, role) + layer_base = mgr.impl.get_mem_pool_base_address( + layer_id, role, PageIndexMode.SHARED + ) + scale = mgr.impl.get_page_index_scale(layer_id, role) + for seq_idx, (req_id, token_num) in enumerate(zip(request_ids, token_nums)): + base_indices = mgr.kv_cache_map[req_id].get_base_page_indices(lifecycle_id) + num_blocks = (token_num + mgr.tokens_per_block - 1) // mgr.tokens_per_block + page_indices = block_offsets[pool_id, seq_idx, kv_idx, :num_blocks].long() + actual = ( + int(pool_pointer) + + (layer_offset * mgr.kv_factor + page_indices) * stride + ) + allocated_indices = torch.tensor( + list(base_indices[:num_blocks]), dtype=torch.int64 + ) + expected = layer_base + allocated_indices * scale * stride + torch.testing.assert_close( + actual, + expected, + rtol=0, + atol=0, + msg=f"layer={layer_id}, role={role}, request={req_id}", + ) + finally: + mgr.shutdown() + + if __name__ == "__main__": unittest.main() From d025dfb96cc2c64166d97740e0dc962b29177b47 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 10 Sep 2026 05:17:20 +0000 Subject: [PATCH 4/4] [6739081][fix] Scope KV stride detection to default page tables Move heterogeneous-stride detection and attention-pool sizing into the default page-table initializer. DeepSeek V4 overrides that initializer with custom buffer roles, so its constructor no longer queries a nonexistent KEY buffer. The generic path retains per-layer addressing for heterogeneous KV pages and SWA scratch reuse. Existing DeepSeek V4 cache, compressor, and transfer tests cover the constructor regression. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 42 ++++++++++--------- 1 file changed, 22 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 4fe5655b2f01..98aba2f428d3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1369,26 +1369,6 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._pool_layer_ids_by_role.setdefault( (pool_id, buffer_id.role), buffer_id.layer_id ) - # A lifecycle group can contain different page sizes in separate - # physical pools (e.g. Gemma's heterogeneous head dimensions). Those - # layers cannot share a pool pointer or page-index scale. Use the - # per-layer conversion path, also required for SWA scratch slots, - # while keeping the underlying lifecycle grouping unchanged. - self._use_per_layer_page_tables = self.enable_swa_scratch_reuse or any( - len( - { - self.impl.get_page_stride(layer_id, self._get_pool_roles(pool_id)[0]) - for layer_id in layer_ids - } - ) - > 1 - for pool_id, layer_ids in enumerate(self.impl.layer_grouping) - ) - if self._use_per_layer_page_tables: - self.num_attention_op_pools = self.num_local_layers - else: - self.num_attention_op_pools = self.num_pools - num_layers = len(config.layers) self.layer_to_pool_mapping_dict: dict[int, int] = { layer_id: self.impl.get_layer_group_id(layer_id) @@ -1636,6 +1616,28 @@ def _build_pool_mapping_tensors(self): return kv_cache_pool_pointers, kv_cache_pool_mapping def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + # Keep role-dependent checks in this default page-table initializer: + # custom layouts such as DeepSeek V4 override it with their own roles. + # A lifecycle group can contain different page sizes in separate + # physical pools (e.g. Gemma's heterogeneous head dimensions). Those + # layers cannot share a pool pointer or page-index scale. Use the + # per-layer conversion path, also required for SWA scratch slots, + # while keeping the underlying lifecycle grouping unchanged. + self._use_per_layer_page_tables = self.enable_swa_scratch_reuse or any( + len( + { + self.impl.get_page_stride(layer_id, self._get_pool_roles(pool_id)[0]) + for layer_id in layer_ids + } + ) + > 1 + for pool_id, layer_ids in enumerate(self.impl.layer_grouping) + ) + if self._use_per_layer_page_tables: + self.num_attention_op_pools = self.num_local_layers + else: + self.num_attention_op_pools = self.num_pools + self.kv_cache_pool_pointers, self.kv_cache_pool_mapping = self._build_pool_mapping_tensors() self.index_scales = torch.empty( self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu"