diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index 0d5273ebef92..f9ef18064e50 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -63,10 +63,11 @@ class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): V2 storage coalesces the buffers of one lifecycle by page size alone. The attention op addresses a layer's K/V pages and its NVFP4 block-scale pages - with one block-offset table, so a side page that lands in the K/V pool - (NVFP4 at 32 tokens per block: index-K 8192 B == K == V) changes the K/V - slot stride but not the block-scale one, and any side page in the K/V pool - breaks the K/V pool mapping. When a side page could be the size of a main + with one block-offset table. A side page lands in the K/V pool when it has + the size of a K/V page: at 32 tokens per block, index-K 8192 B == K == V for + FP8 at one KV head per rank and for NVFP4 at two. Any side page in the K/V + pool breaks the K/V pool mapping, and under NVFP4 it also changes the K/V + slot stride but not the block-scale one. When a side page could be the size of a main page, the side planes are therefore registered as one packed per-slot buffer whose size can be no main-cache page, and this manager serves their per-layer views itself. Both the choice and the packed size are the same on 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 54b9deae7fc3..10448941dd29 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 @@ -2815,12 +2815,22 @@ def _kv_pool_mapping_offset( self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) - key_base_addr ) - return exact_div( - addr_offset, + layer_stride = ( self.get_layer_bytes_per_token(layer_id, Role.KEY) * self._get_layer_kv_factor_by_local_id(layer_id) - * self.tokens_per_block, + * self.tokens_per_block ) + offset, remainder = divmod(addr_offset, layer_stride) + if remainder: + raise RuntimeError( + f"Layer {int(layer_id)}'s K page starts {addr_offset} B into its pool slot, " + f"which is not a multiple of one layer's K/V pages ({layer_stride} B). " + "V2 coalesces a lifecycle's buffers by page size, so another buffer of this " + "page size shares the K/V pool and the attention op cannot address the layer. " + "Register that buffer with a different page size or override " + "_kv_pool_mapping_offset." + ) + return offset def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: if self.layer_specs is not None: diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 5705d5ca2098..0dde5b0c1eac 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -320,6 +320,16 @@ Qwen/Qwen3.8-Flash-Next: kv_cache_quant_algo: NVFP4 spec_dec_algo: MTP accuracy: 95.26 + # FP8 KV cache, measured on the built-in implementation: TEP4 on GB300, no MTP, + # TRTLLM MoE, bf16 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + accuracy: 95.49 + # NVFP4 KV cache, measured on the built-in implementation: ADP4 on GB300, no MTP, + # TRTLLM MoE, bf16 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: NVFP4 + accuracy: 96.02 moonshotai/Kimi-K2.5: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index 9ec9c353f72d..b47f464251b9 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -136,6 +136,16 @@ Qwen/Qwen3.8-Flash-Next: - quant_algo: MIXED_PRECISION spec_dec_algo: MTP accuracy: 87.57 + # FP8 KV cache, measured on the built-in implementation: TEP4 on GB300, no MTP, + # TRTLLM MoE, bf16 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + accuracy: 87.33 + # NVFP4 KV cache, measured on the built-in implementation: ADP4 on GB300, no MTP, + # TRTLLM MoE, bf16 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: NVFP4 + accuracy: 87.21 Qwen/Qwen3.5-397B-A17B: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index c4c8464872a5..0ec6a27cb69b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -5877,7 +5877,8 @@ def _build_llm( enable_attention_dp: bool = False, cover_guided_decoding: bool = False, enable_block_reuse: bool = False, - sparse_attention_config: Optional[QSASparseAttentionConfig] = None + sparse_attention_config: Optional[QSASparseAttentionConfig] = None, + kv_cache_dtype: str = "auto", ) -> LLM: """Construct the engine shared by both evaluation tasks. @@ -5886,6 +5887,7 @@ def _build_llm( """ kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5, enable_block_reuse=enable_block_reuse, + dtype=kv_cache_dtype, mamba_ssm_cache_dtype="bfloat16") if enable_block_reuse: # Attention pages alone cannot restore GDN or PLE state, so the @@ -5933,7 +5935,8 @@ def _run_evals( cover_guided_decoding: bool = False, enable_block_reuse: bool = False, sparse_attention_config: Optional[QSASparseAttentionConfig] = None, - acceptance_length_key: Optional[str] = None) -> None: + acceptance_length_key: Optional[str] = None, + kv_cache_dtype: str = "auto") -> None: if not os.path.exists(model_path): pytest.skip(f"Model directory {model_path} does not exist") @@ -5942,17 +5945,17 @@ def _run_evals( # A case opts into the acceptance-length guard by naming its own entry # in references/acceptance_length.yaml. check_acceptance_length = acceptance_length_key is not None - with self._build_llm( - model_path, - tensor_parallel_size, - moe_backend, - max_draft_len, - check_acceptance_length, - moe_expert_parallel_size=moe_expert_parallel_size, - enable_attention_dp=enable_attention_dp, - cover_guided_decoding=cover_guided_decoding, - enable_block_reuse=enable_block_reuse, - sparse_attention_config=sparse_attention_config) as llm: + with self._build_llm(model_path, + tensor_parallel_size, + moe_backend, + max_draft_len, + check_acceptance_length, + moe_expert_parallel_size=moe_expert_parallel_size, + enable_attention_dp=enable_attention_dp, + cover_guided_decoding=cover_guided_decoding, + enable_block_reuse=enable_block_reuse, + sparse_attention_config=sparse_attention_config, + kv_cache_dtype=kv_cache_dtype) as llm: assert llm.args.quant_config.quant_algo == expected_quant_algo if cover_guided_decoding: assert_guided_decoding_regex(llm) @@ -6097,6 +6100,50 @@ def test_nvfp4_adp4_megamoe_cutedsl_ple_offload( moe_expert_parallel_size=4, enable_attention_dp=True) + @skip_pre_blackwell + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(70000) + @pytest.mark.skip_less_host_memory(131072) + def test_nvfp4_tep4_fp8kv_trtllm_ple_offload( + self, monkeypatch: pytest.MonkeyPatch, mocker) -> None: + """NVFP4 with TEP4 and an FP8 KV cache. + + Each rank holds one KV head, so an FP8 K/V page is the size of a QSA + index-K page. + """ + self._run_evals(f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4", + tensor_parallel_size=4, + moe_backend="TRTLLM", + max_draft_len=None, + expected_quant_algo=QuantAlgo.MIXED_PRECISION, + monkeypatch=monkeypatch, + mocker=mocker, + moe_expert_parallel_size=4, + kv_cache_dtype="fp8") + + @skip_pre_blackwell + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(70000) + @pytest.mark.skip_less_host_memory(131072) + def test_nvfp4_adp4_nvfp4kv_trtllm_ple_offload( + self, monkeypatch: pytest.MonkeyPatch, mocker) -> None: + """NVFP4 on four GPUs with attention DP and an NVFP4 KV cache. + + Each rank holds both KV heads, so an NVFP4 K/V page is the size of a + QSA index-K page. QSA attention runs on the regular backend for an + NVFP4 KV cache. + """ + self._run_evals(f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4", + tensor_parallel_size=4, + moe_backend="TRTLLM", + max_draft_len=None, + expected_quant_algo=QuantAlgo.MIXED_PRECISION, + monkeypatch=monkeypatch, + mocker=mocker, + moe_expert_parallel_size=4, + enable_attention_dp=True, + kv_cache_dtype="nvfp4") + @pytest.mark.skip_less_device_memory(80000) class TestQwen3_8_27B(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 1fbcfa32289f..205344ec1de0 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -611,6 +611,8 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_pre accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_1gpu_cutedsl_ple_offload accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_megamoe_cutedsl_ple_offload +accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_nvfp4kv_trtllm_ple_offload +accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_tep4_fp8kv_trtllm_ple_offload accuracy/test_llm_api_pytorch.py::TestSeedOss_36B::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestStep3p7SpeculativeDecoding::test_guided_decoding TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestStep3p7SpeculativeDecoding::test_step3p7_mtp TIMEOUT (60) diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py index 18abace17337..347022748819 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py @@ -6,8 +6,13 @@ offset) * page`` and its NVFP4 block-scale page with the same block index in the scale pool. These tests build the managers a Qwen3.8-Flash-Next engine builds -- the 48-layer trunk (12 QSA layers among 36 recurrent ones) and the -one-layer MTP draft manager -- and check that every view the manager serves is +one-layer MTP draft manager -- on rank 0 of each tensor, expert and +attention-data parallel layout, and check that every view the manager serves is where that formula lands, for bf16, fp8 and NVFP4 pools. + +With 2 KV heads of 256 dims, a 128-dim bf16 index-K page is the size of a K/V +page for FP8 at one head per rank (TP2, TEP4) and for NVFP4 at two (TP1, +attention DP), which V2 would coalesce into the K/V pool. """ from types import SimpleNamespace @@ -41,12 +46,33 @@ _INDEX_PAGE = 1 * 128 * 2 * _TPB _POSITION_PAGE = 3 * 4 * _TPB +#: (tp_size, moe_ep_size, attention_dp) of each layout. TP1 and attention DP hold both +#: KV heads per rank, TP2 and TEP4 one. +_LAYOUTS = { + "tp1": (1, 1, False), + "tp2": (2, 1, False), + "tep4": (4, 4, False), + "adp2": (2, 1, True), + "adp4": (4, 4, True), +} + + +def _kv_heads_per_rank(layout: str) -> int: + """KV heads a rank holds; with more TP ranks than KV heads, each head is replicated.""" + tp_size, _, attention_dp = _LAYOUTS[layout] + return _KV_HEADS if attention_dp else max(1, _KV_HEADS // tp_size) + def _manager( - dtype: DataType, *, draft: bool, tp_size: int = 1, attention_dp: bool = False + dtype: DataType, + *, + draft: bool, + layout: str = "tp1", + ssm_dtype: torch.dtype = torch.float32, ) -> QSAMambaHybridCacheManagerV2: - """The trunk or the separate MTP draft manager, as the engine builds them (rank 0 of - `tp_size`, with attention data parallelism when `attention_dp`).""" + """The trunk or the separate MTP draft manager, as the engine builds them on rank 0 + of `layout`.""" + tp_size, moe_ep_size, attention_dp = _LAYOUTS[layout] if draft: attention_mask = [False] * _TRUNK_LAYERS + [True] mamba_mask = [False] * (_TRUNK_LAYERS + 1) @@ -62,7 +88,7 @@ def _manager( mamba_num_layers=sum(mamba_mask), mamba_layer_mask=mamba_mask, mamba_cache_dtype=torch.bfloat16, - mamba_ssm_cache_dtype=torch.float32, + mamba_ssm_cache_dtype=ssm_dtype, kv_cache_config=KvCacheConfig(max_tokens=4096, enable_block_reuse=False), kv_cache_type=CacheTypeCpp.SELF, num_layers=sum(attention_mask), @@ -75,6 +101,7 @@ def _manager( world_size=tp_size, rank=0, tp_size=tp_size, + moe_ep_size=moe_ep_size, pp_size=1, enable_attention_dp=attention_dp, ), @@ -92,9 +119,9 @@ def _layers(draft: bool) -> list[int]: return [_DRAFT_LAYER] if draft else _QSA_LAYERS -def _kv_page_bytes(dtype: DataType) -> int: +def _kv_page_bytes(dtype: DataType, kv_heads: int) -> int: element_bits = {DataType.BF16: 16, DataType.FP8: 8, DataType.NVFP4: 4}[dtype] - return _KV_HEADS * _TPB * _HEAD_DIM * element_bits // 8 + return kv_heads * _TPB * _HEAD_DIM * element_bits // 8 def _thop_addresses(cm, layer: int, slot: int, page_bytes: int, plane: int) -> tuple[int, int]: @@ -120,19 +147,23 @@ def _require_cuda(): assert torch.cuda.is_available(), "the managers allocate their pools on the GPU" +@pytest.mark.parametrize("layout", list(_LAYOUTS)) @pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) @pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) -def test_views_are_where_the_attention_op_addresses(dtype: DataType, draft: bool) -> None: +def test_views_are_where_the_attention_op_addresses( + dtype: DataType, draft: bool, layout: str +) -> None: """Every K/V (and NVFP4 block-scale) view of every QSA layer starts, at every checked slot, at the address the attention op computes for it.""" - cm = _manager(dtype, draft=draft) + cm = _manager(dtype, draft=draft, layout=layout) try: - page = _kv_page_bytes(dtype) + kv_heads = _kv_heads_per_rank(layout) + page = _kv_page_bytes(dtype, kv_heads) for layer in _layers(draft): kv = cm.get_buffers(layer, kv_layout="HND") assert kv.shape[1:] == ( 2, - _KV_HEADS, + kv_heads, _TPB, _HEAD_DIM // (2 if dtype == DataType.NVFP4 else 1), ) @@ -144,7 +175,7 @@ def test_views_are_where_the_attention_op_addresses(dtype: DataType, draft: bool continue scales = cm.get_block_scale_buffers(layer, kv_layout="HND") assert scales.dtype == torch.uint8 - assert scales.shape == (kv.shape[0], 2, _KV_HEADS, _TPB, _HEAD_DIM // 16) + assert scales.shape == (kv.shape[0], 2, kv_heads, _TPB, _HEAD_DIM // 16) for slot in _slots_to_check(scales): key, value = _thop_addresses(cm, layer, slot, page // 8, plane=1) assert scales[slot, 0].data_ptr() == key, (layer, slot) @@ -153,14 +184,17 @@ def test_views_are_where_the_attention_op_addresses(dtype: DataType, draft: bool cm.shutdown() +@pytest.mark.parametrize("layout", list(_LAYOUTS)) @pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) @pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) -def test_side_planes_pack_when_a_page_could_collide(dtype: DataType, draft: bool) -> None: +def test_side_planes_pack_when_a_page_could_collide( + dtype: DataType, draft: bool, layout: str +) -> None: """An index page could be a K/V page here -- FP8 at one KV head a rank (1 x 256 x 32 tokens), NVFP4 at both -- so the index and position planes are packed into a buffer of their own instead of landing in the K/V pool, and under NVFP4 the K/V and block-scale pools hold the same number of pages per slot.""" - cm = _manager(dtype, draft=draft) + cm = _manager(dtype, draft=draft, layout=layout) try: layers = _layers(draft) assert cm.qsa_side_planes_size == len(layers) * _INDEX_PAGE + _POSITION_PAGE @@ -206,12 +240,20 @@ def test_bf16_keeps_per_layer_side_roles(draft: bool) -> None: cm.shutdown() +@pytest.mark.parametrize( + "ssm_dtype", + [torch.float32, torch.float16, torch.bfloat16], + ids=["ssm_fp32", "ssm_fp16", "ssm_bf16"], +) +@pytest.mark.parametrize("layout", ["tp1", "tep4"]) @pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) @pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) -def test_side_and_main_views_do_not_alias(dtype: DataType, draft: bool) -> None: +def test_side_and_main_views_do_not_alias( + dtype: DataType, draft: bool, layout: str, ssm_dtype: torch.dtype +) -> None: """One slot written through every view reads back unchanged through each: no K/V, scale, index or position page overlaps another.""" - cm = _manager(dtype, draft=draft) + cm = _manager(dtype, draft=draft, layout=layout, ssm_dtype=ssm_dtype) try: views = [] for layer in _layers(draft): @@ -243,22 +285,18 @@ def _logical_page(view: torch.Tensor) -> tuple: return tuple(page.shape), page.dtype, page.numel() * page.element_size(), page.is_contiguous() -#: (tp_size, attention_dp): the layouts a disaggregated pair can mix over one global -#: geometry (2 KV heads): TP1 and ADP2 hold both heads per rank, TP2 one. -_LAYOUTS = {"tp1": (1, False), "tp2": (2, False), "adp2": (2, True)} - - @pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) def test_transfer_views_do_not_depend_on_the_kv_head_split(dtype: DataType) -> None: """Disaggregated peers compare the replicated side roles each layer registers and copy their pages byte for byte, so roles, mapper kinds and every side view's page (shape, dtype, bytes) must not depend on how many KV heads a rank holds; where pages sit in - their pools may. At TP2 an FP8 K/V page (1 head x 256 x 32 tokens) is the size of an - index page; under NVFP4 the same holds at TP1 and ADP2 (2 heads x 128 bytes x 32 - tokens). Both pack on every layout; bf16, which can never collide, on none.""" + their pools may. At TP2 and TEP4 an FP8 K/V page (1 head x 256 x 32 tokens) is the + size of an index page; under NVFP4 the same holds at TP1 and attention DP (2 heads x + 128 bytes x 32 tokens). Both pack on every layout; bf16, which can never collide, on + none.""" seen = {} - for name, (tp_size, attention_dp) in _LAYOUTS.items(): - cm = _manager(dtype, draft=False, tp_size=tp_size, attention_dp=attention_dp) + for name in _LAYOUTS: + cm = _manager(dtype, draft=False, layout=name) try: buffers = cm._extra_buffers_per_layer(tokens_per_block=_TPB) roles = { @@ -273,8 +311,9 @@ def test_transfer_views_do_not_depend_on_the_kv_head_split(dtype: DataType) -> N seen[name] = (roles, logical, cm.get_disagg_role_mapper_kinds()) finally: cm.shutdown() - assert seen["tp1"] == seen["tp2"] == seen["adp2"] - roles = {role for layer in seen["tp1"][0].values() for role, _ in layer} + first, *rest = seen.values() + assert all(other == first for other in rest) + roles = {role for layer in first[0].values() for role, _ in layer} if dtype == DataType.BF16: assert roles == {str(Role.INDEX_KEY), str(QSA_INDEX_POSITION)} else: 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 81622c5f0ee6..8a8f0b47570f 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 @@ -274,6 +274,17 @@ def test_index_key_does_not_disturb_main_kv_buffers(self): mgr.shutdown() del mgr + def test_index_key_coalesced_into_kv_pool_needs_an_offset_override(self): + # One bf16 head of 128 dims makes K == V == INDEX_KEY bytes, so V2 + # puts all three in one pool slot and layer 1's K page lands three + # pages in: no uniform K/V layer offset exists. Without an override + # of _kv_pool_mapping_offset the construction must name the cause. + with self.assertRaisesRegex(RuntimeError, "_kv_pool_mapping_offset"): + _IndexKeyOnSparseLayersV2( + sparse_layer_ids=[0, 1, 2, 3], + **_make_kwargs(num_layers=4, num_kv_heads=1, dtype=DataType.BF16), + ) + def test_duplicate_role_against_standard_buffer_asserts(self): # Negative control: registering Role.KEY as an extra collides with # the standard buffer and must fail loud, both for the assertion