diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h index 60422dea5910..c0a461e3c72d 100644 --- a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h +++ b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h b/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h index 9e316c0b4ebe..a62da181cbb5 100644 --- a/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h +++ b/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu index cbcb31f07111..cd33862afb4c 100644 --- a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu +++ b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu @@ -63,7 +63,7 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset( * @param out_scale Output scale data [num_tokens, scale_size] contiguous * @param k_token_start Start offset into slot_mapping arrays * @param num_tokens Number of tokens to gather - * @param head_dim Head dimension (must be 128) + * @param head_dim Payload byte width (128 for FP8, 64 for packed FP4) * @param scale_size Scale size in bytes (must be 4) * @param cache_stride_0 Stride for k_cache dimension 0 (in bytes) * @param cache_stride_1 Stride for k_cache dimension 1 (in bytes) diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu index f0a41a5009fc..f0c7c27a235c 100644 --- a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu +++ b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu @@ -61,7 +61,7 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset( * @param slot_mapping_fp8 Flat element index for FP8 data start position [num_tokens] * @param slot_mapping_scale Flat element index for scale data start position [num_tokens] * @param num_tokens Number of tokens - * @param head_dim Head dimension (must be 128) + * @param head_dim Payload byte width (128 for FP8, 64 for packed FP4) * @param scale_size Scale size in bytes (must be 4) * @param cache_stride_0 Stride for k_cache dimension 0 (in bytes) * @param cache_stride_1 Stride for k_cache dimension 1 (in bytes) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py index 6cf3aec0a6d1..80908e08b0ac 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py @@ -4,7 +4,7 @@ """DeepSeek Sparse Attention backend package.""" from .backend import DSATrtllmAttention -from .cache_manager import DSACacheManager +from .cache_manager import DSACacheManager, DSACacheManagerV2, is_dsa_cache_manager from .indexer import ( _DG_SCHEDULE_BLOCK_KV, HAS_FAST_HADAMARD, @@ -29,6 +29,7 @@ "HAS_FAST_HADAMARD", "DSABackendForwardArgs", "DSACacheManager", + "DSACacheManagerV2", "DSAMetadataParams", "DSAParams", "DSATrtllmAttention", @@ -44,6 +45,7 @@ "_select_indexer_compress_ratio", "build_req_idx_per_token", "compute_cu_seqlen_kv_bounds_with_cache", + "is_dsa_cache_manager", "rotate_activation", "split_prefill_chunks", "transform_local_topk_and_prepare_pool_view", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py index e6ff3d8b8b80..ed19c6ea869e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py @@ -5,17 +5,21 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING, List, Optional, Tuple, Union + +import torch import tensorrt_llm import tensorrt_llm.bindings -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._utils import get_size_in_bytes +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager, get_pp_layers +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole, PageIndexMode from .params import DSAParams @@ -25,6 +29,41 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, SparseAttentionConfig +def _get_indexer_k_cache_bytes_per_token( + index_head_dim: int, quant_block_size: int, use_fp4: bool +) -> int: + """Return the raw indexer K-cache footprint for one token.""" + data_bytes = index_head_dim // 2 if use_fp4 else index_head_dim + scale_bytes = index_head_dim // quant_block_size * 4 + return data_bytes + scale_bytes + + +def _get_indexer_k_cache_size_per_token( + model_config: ModelConfig, + mapping: Mapping, + num_layers: Optional[int] = None, +) -> int: + """Estimate the indexer-only cache cost across local attention layers.""" + sparse_attention_config = model_config.sparse_attention_config + if sparse_attention_config is None: + raise ValueError("sparse_attention_config is required for DSA cache") + sparse_params = sparse_attention_config.to_sparse_params( + pretrained_config=model_config.pretrained_config + ) + if not isinstance(sparse_params, DSAParams): + raise ValueError("DSA cache requires DSA sparse parameters") + + num_attention_layers = KVCacheManager._resolve_num_attention_layers( + model_config, mapping, num_layers + ) + bytes_per_layer = _get_indexer_k_cache_bytes_per_token( + sparse_params.index_head_dim, + 128, + sparse_params.indexer_k_dtype == "fp4", + ) + return num_attention_layers * bytes_per_layer + + def derive_indexer_k_cache_layer_mask( sparse_attention_config: "SparseAttentionConfig", pretrained_config, @@ -123,6 +162,8 @@ def __init__( **kwargs, ) self.num_blocks = self.blocks_in_primary_pool + # V1 stores one INDEX_KEY page per physical pool slot. + self.indexer_k_cache_page_scale = 1 # Indexer K cache pool for DSA attention # Shape: [num_blocks, self.tokens_per_block * (index_head_dim + scale_size)] @@ -141,11 +182,16 @@ def __init__( f"{self.num_local_layers} local layers own an indexer k-cache." ) + def get_primary_pool_page_index_params(self, local_layer_idx: int) -> Tuple[int, int]: + """Return V1 page scale and layer offset for sparse MLA indices.""" + return self.num_local_layers, local_layer_idx + def get_indexer_k_cache_buffers(self, layer_idx: int): """Get indexer k cache buffer from a specific layer pool.""" block_size = self.tokens_per_block - data_bytes = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim - per_token_size = data_bytes + self.index_head_dim // self.quant_block_size * 4 + per_token_size = _get_indexer_k_cache_bytes_per_token( + self.index_head_dim, self.quant_block_size, self.use_fp4 + ) layer_offset = self.layer_offsets[layer_idx] pool = self.indexer_k_cache_pool_per_layer[layer_offset] assert pool is not None, ( @@ -154,6 +200,20 @@ def get_indexer_k_cache_buffers(self, layer_idx: int): ) return pool.view(self.num_blocks, block_size, 1, per_token_size) + def get_pool_block_indices( + self, + num_seqs: int, + *, + request_ids: Optional[List[int]] = None, + num_contexts: int = 0, + beam_width: int = 1, + ) -> torch.Tensor: + """Decode V1 block offsets into physical memory-pool block indices.""" + del request_ids, num_contexts, beam_width + encoded = self.host_kv_cache_block_offsets[0, :num_seqs, 0, :] + max_pool_idx = self.blocks_in_primary_pool - 1 + return (encoded // self.num_local_layers).clamp(min=0, max=max_pool_idx).to(torch.int32) + def get_batch_indexer_k_cache_indices(self, request_ids: List[int]) -> List[List[int]]: """ Get the indices for the indexer k cache for a specific batch of requests. @@ -277,3 +337,286 @@ def get_cache_bytes_per_token(self): cache_size_bytes_per_token += indexer_bytes_per_token return cache_size_bytes_per_token + + +class DSACacheManagerV2(KVCacheManagerV2): + """KVCacheManagerV2-backed cache manager with a DSA indexer K-cache.""" + + def __init__( + self, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[List[bool]] = None, + max_num_tokens: int = 8192, + model_config: Optional[ModelConfig] = None, + max_beam_width: int = 1, + sparse_attention_config: Optional["SparseAttentionConfig"] = None, + pretrained_config=None, + **kwargs, + ) -> None: + if sparse_attention_config is None: + sparse_attention_config = kwargs.pop("sparse_attn_config", None) + if sparse_attention_config is None and model_config is not None: + sparse_attention_config = model_config.sparse_attention_config + if sparse_attention_config is None: + raise ValueError("sparse_attention_config is required for DSA cache") + sparse_params = sparse_attention_config.to_sparse_params( + pretrained_config=pretrained_config + ) + if not isinstance(sparse_params, DSAParams): + raise ValueError("DSA cache requires DSA sparse parameters") + + self.quant_block_size = 128 + self.index_head_dim = sparse_params.index_head_dim + self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" + self._unique_primary_pool: Optional[torch.Tensor] = None + + from tensorrt_llm._torch.speculative import get_num_spec_layers + + total_num_layers = len(layer_mask) if layer_mask is not None else num_layers + if spec_config is not None and layer_mask is None: + total_num_layers += get_num_spec_layers(spec_config) + indexer_k_cache_layer_mask = derive_indexer_k_cache_layer_mask( + sparse_attention_config, pretrained_config, total_num_layers + ) + pp_layers, _ = get_pp_layers( + num_layers, mapping, spec_config=spec_config, layer_mask=layer_mask + ) + self.indexer_k_cache_local_layer_mask = [ + indexer_k_cache_layer_mask[layer_idx] for layer_idx in pp_layers + ] + + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=layer_mask, + max_num_tokens=max_num_tokens, + model_config=model_config, + max_beam_width=max_beam_width, + **kwargs, + ) + if self.num_local_layers == 0: + raise ValueError("DSA requires at least one local attention layer") + primary_converters = [ + self.impl.get_page_index_converter(local_layer_idx, Role.KEY) + for local_layer_idx in range(self.num_local_layers) + ] + self._primary_pool_page_index_params = [ + (int(converter.scale), int(converter.layer_offset)) for converter in primary_converters + ] + full_indexer_local_layers = [ + local_layer_idx + for local_layer_idx, has_indexer in enumerate(self.indexer_k_cache_local_layer_mask) + if has_indexer + ] + indexer_converters = [ + self.impl.get_page_index_converter(local_layer_idx, Role.INDEX_KEY) + for local_layer_idx in full_indexer_local_layers + ] + if any(int(converter.expansion) != 1 for converter in indexer_converters): + raise ValueError("DSA indexer K-cache does not support page expansion") + # Each cache view starts at its SHARED base address, so the converter's + # layer offset is already represented by the base pointer. + self.indexer_k_cache_page_scale = ( + int(indexer_converters[0].scale) if indexer_converters else 1 + ) + if any( + int(converter.scale) != self.indexer_k_cache_page_scale + for converter in indexer_converters[1:] + ): + raise ValueError("DSA requires a uniform shared INDEX_KEY page mapping across layers") + self.num_blocks = self.blocks_in_primary_pool + self.indexer_k_cache_pool_per_layer = [ + self._get_indexer_k_cache_pool_data(local_layer_idx) + if self.indexer_k_cache_local_layer_mask[local_layer_idx] + else None + for local_layer_idx in range(self.num_local_layers) + ] + num_full = sum(self.indexer_k_cache_local_layer_mask) + if num_full < self.num_local_layers: + logger.info( + f"[DSACacheManagerV2] Indexer k-cache: {num_full} of " + f"{self.num_local_layers} local layers own an indexer k-cache." + ) + + def get_primary_pool_page_index_params(self, local_layer_idx: int) -> Tuple[int, int]: + """Return the formal V2 page scale and layer offset for sparse MLA.""" + return self._primary_pool_page_index_params[local_layer_idx] + + def _extra_buffers_per_layer(self, *, tokens_per_block: int) -> dict[int, List[BufferConfig]]: + return { + local_layer_idx: [ + BufferConfig( + role=Role.INDEX_KEY, + size=self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY) + * tokens_per_block, + ) + ] + for local_layer_idx in range(self.num_local_layers) + if self.indexer_k_cache_local_layer_mask[local_layer_idx] + } + + @property + def blocks_in_primary_pool(self) -> int: + """Return the physical slot count rather than the V2 page bound.""" + converter = self.impl.get_page_index_converter(0, Role.KEY) + page_upper = self.impl.get_page_index_upper_bound(0, Role.KEY) + expansion = int(converter.expansion) + assert page_upper % expansion == 0 + num_pages_with_offset = page_upper // expansion + int(converter.layer_offset) + scale = int(converter.scale) + assert num_pages_with_offset % scale == 0 + return num_pages_with_offset // scale + + def _get_indexer_k_cache_pool_data(self, local_layer_idx: int) -> torch.Tensor: + """Return a contiguous shared-page view for one indexer K-cache.""" + address = self.impl.get_mem_pool_base_address( + local_layer_idx, Role.INDEX_KEY, PageIndexMode.SHARED + ) + page_upper = self.impl.get_page_index_upper_bound(local_layer_idx, Role.INDEX_KEY) + flat_page_size = self.tokens_per_block * self.get_layer_bytes_per_token( + local_layer_idx, Role.INDEX_KEY + ) + return convert_to_torch_tensor( + TensorWrapper(address, torch.uint8, [page_upper, flat_page_size]) + ) + + def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor: + """Return the page-indexed indexer K-cache view for a global layer.""" + layer_offset = self.layer_offsets[layer_idx] + pool = self.indexer_k_cache_pool_per_layer[layer_offset] + assert pool is not None, ( + f"Layer {layer_idx} is a shared-indexer layer and owns no indexer " + f"k-cache; only full-indexer layers may access it." + ) + per_token_size = self.get_layer_bytes_per_token(layer_offset, Role.INDEX_KEY) + return pool.view(pool.shape[0], self.tokens_per_block, 1, per_token_size) + + def get_pool_block_indices( + self, + num_seqs: int, + *, + request_ids: Optional[List[int]] = None, + num_contexts: int = 0, + beam_width: int = 1, + ) -> torch.Tensor: + """Read V2 stable slots in current-batch order as physical block IDs.""" + if request_ids is None: + raise ValueError("DSACacheManagerV2 requires request_ids to map stable slots") + if len(request_ids) != num_seqs: + raise ValueError(f"Expected {num_seqs} request IDs, got {len(request_ids)}") + copy_idx = self.index_mapper.get_copy_index(list(request_ids), num_contexts, beam_width) + copy_idx = copy_idx.to(device="cpu", dtype=torch.long) + block_indices = self.host_kv_cache_block_offsets[0, copy_idx, 0, :] + return block_indices.clamp(min=0, max=self.num_blocks - 1).to(torch.int32) + + def get_unique_primary_pool(self) -> torch.Tensor: + """Return the uniform MLA K pool in the V1-compatible layout.""" + if self._unique_primary_pool is not None: + return self._unique_primary_pool + if self.num_local_layers == 0: + raise ValueError("DSA requires at least one local attention layer") + if self.kv_factor != 1: + raise ValueError("DSA requires a SELFKONLY KV cache") + + first_head_dim = self.head_dim_per_layer[0] + first_num_heads = self.num_kv_heads_per_layer[0] + if any(head_dim != first_head_dim for head_dim in self.head_dim_per_layer): + raise ValueError("DSA requires a uniform KV head dimension") + if any(num_heads != first_num_heads for num_heads in self.num_kv_heads_per_layer): + raise ValueError("DSA requires a uniform KV head count") + + first_converter = self.impl.get_page_index_converter(0, Role.KEY) + if int(first_converter.expansion) != 1: + raise ValueError("DSA MLA K-cache does not support page expansion") + if int(first_converter.layer_offset) != 0: + raise ValueError("The first DSA layer must start at pool offset 0") + if int(first_converter.scale) != self.num_local_layers: + raise ValueError("DSA requires one uniformly coalesced K page per local layer") + + page_stride = self.impl.get_page_stride(0, Role.KEY) + base_address = self.impl.get_mem_pool_base_address(0, Role.KEY, PageIndexMode.SHARED) + for local_layer_idx in range(1, self.num_local_layers): + converter = self.impl.get_page_index_converter(local_layer_idx, Role.KEY) + if int(converter.scale) != int(first_converter.scale) or int(converter.expansion) != 1: + raise ValueError("DSA requires a uniform page-index mapping across layers") + if int(converter.layer_offset) != local_layer_idx: + raise ValueError("DSA requires K pages to follow local-layer order") + address = self.impl.get_mem_pool_base_address( + local_layer_idx, Role.KEY, PageIndexMode.SHARED + ) + expected_address = base_address + int(converter.layer_offset) * page_stride + if int(address) != int(expected_address): + raise ValueError("DSA requires contiguous per-layer K pages in each slot") + + element_per_container = 2 if self.dtype == DataType.NVFP4 else 1 + dtype = torch.int8 if self.dtype == DataType.NVFP4 else self.dtype + elements_per_layer = ( + self.tokens_per_block * first_num_heads * first_head_dim // element_per_container + ) + shape = [ + self.blocks_in_primary_pool, + self.num_local_layers, + 1, + elements_per_layer, + ] + self._unique_primary_pool = convert_to_torch_tensor( + TensorWrapper(base_address, dtype, shape) + ) + return self._unique_primary_pool + + def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) -> int: + if data_role == Role.INDEX_KEY: + return _get_indexer_k_cache_bytes_per_token( + self.index_head_dim, self.quant_block_size, self.use_fp4 + ) + cache_bytes = super().get_layer_bytes_per_token(local_layer_idx, data_role) + if data_role == Role.ALL and self.indexer_k_cache_local_layer_mask[local_layer_idx]: + cache_bytes += self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY) + return cache_bytes + + def get_cache_bytes_per_token(self) -> int: + return sum( + self.get_layer_bytes_per_token(local_layer_idx, Role.ALL) + for local_layer_idx in range(self.num_local_layers) + ) + + @staticmethod + def get_cache_size_per_token( + model_config: ModelConfig, + mapping: Mapping, + num_layers: Optional[int] = None, + **kwargs, + ): + return DSACacheManager.get_cache_size_per_token( + model_config, mapping, num_layers=num_layers, **kwargs + ) + + def shutdown(self) -> None: + self.indexer_k_cache_pool_per_layer = [] + self._unique_primary_pool = None + super().shutdown() + + +def is_dsa_cache_manager(cache_manager: object) -> bool: + """Return whether a manager uses native DSA indexer page mapping.""" + return isinstance(cache_manager, (DSACacheManager, DSACacheManagerV2)) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index ded1bda7694f..eae55ac98ae7 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -294,6 +294,9 @@ def transform_local_topk_and_prepare_pool_view( assert topk_indices.dtype == torch.int32 attn_metadata._ensure_pool_view_cached() + page_index_scale, layer_offset = ( + attn_metadata.kv_cache_manager.get_primary_pool_page_index_params(layer_idx) + ) if is_generation: block_table = attn_metadata._cached_block_table_gen @@ -308,8 +311,8 @@ def transform_local_topk_and_prepare_pool_view( topk_indices, attn_metadata._cached_tokens_per_block, topk_indices.shape[1], - attn_metadata._cached_stride_factor, - layer_idx, + page_index_scale * attn_metadata._cached_tokens_per_block, + layer_offset, ) return global_indices, attn_metadata._cached_pool_view @@ -932,6 +935,86 @@ def recompute_slot_mappings( metadata.host_slot_mapping_scale[:total_new_kv_tokens], non_blocking=True ) + @staticmethod + def recompute_context_kv_gather_mappings( + metadata: DSAtrtllmAttentionMetadata, + indexer_params: Optional[IndexerParams] = None, + ) -> None: + """Recompute context KV-cache gather mappings for the active cache manager.""" + if not metadata.enable_context_mla_with_cached_kv: + metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 + metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + return + + if indexer_params is None: + indexer_params = Indexer.build_indexer_params(metadata) + if indexer_params is None: + return + + num_contexts = indexer_params.num_contexts + if num_contexts == 0: + # Generation does not consume full-KV gather mappings. In + # particular, avoid allocating temporary zero-length tensors + # while capturing generation CUDA graphs. + return + + if metadata.skip_indexer_for_ctx_reqs: + # The dense short-context path does not gather from the indexer K + # cache, so full-KV mappings would be allocated but never consumed. + metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 + metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + return + + cached_kv_tokens = indexer_params.cached_kv_tokens[:num_contexts] + if indexer_params.compress_ratio == 1 and cached_kv_tokens.count_nonzero().item() == 0: + # Without compression or cached tokens, the regular mappings cover + # every context KV token in the same order and remain equivalent + # when on_update_kv_lens() refreshes them in place. + metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 + metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + return + + total_kv_per_request = indexer_params.kv_lens[:num_contexts] + total_kv_len = total_kv_per_request.sum().item() + if total_kv_len == 0: + # No chunk can gather indexer KV, so avoid zero-length allocation + # and H2D work for short compressed contexts. + metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 + metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + return + + host_slot_mapping_fp8_fullkv = torch.empty( + total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned() + ) + host_slot_mapping_scale_fullkv = torch.empty( + total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned() + ) + + req_indices = torch.repeat_interleave( + torch.arange(num_contexts, dtype=torch.int64, device="cpu"), total_kv_per_request + ) + + cu_kv = torch.zeros(num_contexts + 1, dtype=torch.int64, device="cpu") + cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0) + kv_positions = torch.arange(total_kv_len, dtype=torch.int64, device="cpu") - cu_kv[ + :-1 + ].repeat_interleave(total_kv_per_request) + + fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( + kv_positions, + metadata.host_indexer_k_cache_block_offsets, + req_indices, + indexer_params.head_dim, + indexer_params.tokens_per_block, + indexer_params.quant_block_size, + data_bytes_per_token=indexer_params.data_bytes_per_token, + ) + + host_slot_mapping_fp8_fullkv.copy_(fp8_flat_indices) + host_slot_mapping_scale_fullkv.copy_(scale_flat_indices) + metadata.slot_mapping_fp8_fullkv = host_slot_mapping_fp8_fullkv.cuda(non_blocking=True) + metadata.slot_mapping_scale_fullkv = host_slot_mapping_scale_fullkv.cuda(non_blocking=True) + @staticmethod def prepare_for_update_k_cache( metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams @@ -953,8 +1036,6 @@ def prepare_for_chunked_prefill( """ num_contexts = indexer_params.num_contexts seq_lens = indexer_params.seq_lens - tokens_per_block = indexer_params.tokens_per_block - head_dim = indexer_params.head_dim # When MLA chunked prefill is active, it already handles chunking # Indexer should just process the current MLA chunk as a single chunk @@ -1001,57 +1082,9 @@ def prepare_for_chunked_prefill( else: metadata.indexer_prefill_chunks = None - # Chunked prefill and KV-cache reuse require the full KV for indexer - # logits. The indexer's own chunking gathers only the current chunk. - if metadata.enable_context_mla_with_cached_kv: - # Use kv_lens which correctly computes (raw_past + seq_lens) // compress_ratio. - total_kv_per_request = indexer_params.kv_lens[:num_contexts] - total_kv_len = total_kv_per_request.sum().item() - host_slot_mapping_fp8_fullkv = torch.empty( - total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned() - ) - host_slot_mapping_scale_fullkv = torch.empty( - total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned() - ) - - req_indices = torch.repeat_interleave( - torch.arange(num_contexts, dtype=torch.int64, device="cpu"), total_kv_per_request - ) - - cu_kv = torch.zeros(num_contexts + 1, dtype=torch.int64, device="cpu") - cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0) - kv_positions = torch.arange(total_kv_len, dtype=torch.int64, device="cpu") - cu_kv[ - :-1 - ].repeat_interleave(total_kv_per_request) - - fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( - kv_positions, - metadata.host_indexer_k_cache_block_offsets, - req_indices, - head_dim, - tokens_per_block, - indexer_params.quant_block_size, - data_bytes_per_token=head_dim // 2 - if metadata.kv_cache_manager.use_fp4 - else head_dim, - ) - - host_slot_mapping_fp8_fullkv[:total_kv_len] = fp8_flat_indices - host_slot_mapping_scale_fullkv[:total_kv_len] = scale_flat_indices - - assert len(fp8_flat_indices) == total_kv_len, ( - "host_slot_mapping_fp8_fullkv/host_slot_mapping_scale_fullkv " - f"length mismatch: {len(fp8_flat_indices)} != total_kv_len={total_kv_len}" - ) - - # Store extended mappings for indexer full KV gathering - metadata.slot_mapping_fp8_fullkv = host_slot_mapping_fp8_fullkv.cuda(non_blocking=True) - metadata.slot_mapping_scale_fullkv = host_slot_mapping_scale_fullkv.cuda( - non_blocking=True - ) - else: - metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 - metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + # Chunked prefill and KV-cache reuse require gather mappings that cover + # all cached and newly appended indexer tokens. + Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params) @staticmethod def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): @@ -1216,7 +1249,7 @@ def _gather_k_cache_for_chunk( k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) - head_dim = self.head_dim + head_dim = self.head_dim // 2 if self.use_fp4 else self.head_dim scale_size = 4 # float32 = 4 bytes # Extract slot mappings using chunk's k_token_start/end diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index c7b0dcafae38..a41631cda0e1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -17,6 +17,7 @@ from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata +from .cache_manager import is_dsa_cache_manager from .indexer import ( _DG_SCHEDULE_BLOCK_KV, Indexer, @@ -124,7 +125,6 @@ def __init__(self, *args, **kwargs): self._pool_cache_valid = False self._cached_kv_mgr_id = 0 self._cached_pool_view = None - self._cached_stride_factor = 0 self._cached_tokens_per_block = 0 self._cached_block_table_ctx = None self._cached_block_table_gen = None @@ -138,10 +138,8 @@ def __init__(self, *args, **kwargs): def __post_init__(self): """Allocate indexer K-cache buffers and heuristic TopK metadata.""" - from .cache_manager import DSACacheManager - super().__post_init__() - if not isinstance(self.kv_cache_manager, DSACacheManager): + if not is_dsa_cache_manager(self.kv_cache_manager): has_deepseek_v4_cache_interface = all( hasattr(self.kv_cache_manager, attr) for attr in ("compressed_block_sizes", "get_cache_indices") @@ -570,6 +568,48 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) + # Allocate separate indexer block-offset and slot-mapping buffers for + # the draft KV cache manager, mirroring draft_kv_cache_block_offsets: + # the draft-replay context swaps these in by rebinding, so CUDA graph + # capture bakes distinct addresses for the target and draft segments + # and both sides can be refreshed eagerly outside the graph. + if is_dsa_cache_manager(self.draft_kv_cache_manager): + self.draft_indexer_k_cache_block_offsets = self.get_empty( + self.cuda_graph_buffers, + [self.max_num_sequences, self.draft_kv_cache_manager.max_blocks_per_seq], + cache_name="draft_indexer_k_cache_block_offsets", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_draft_indexer_k_cache_block_offsets = torch.zeros_like( + self.draft_indexer_k_cache_block_offsets, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.draft_slot_mapping_fp8 = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="draft_slot_mapping_fp8", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_draft_slot_mapping_fp8 = torch.zeros_like( + self.draft_slot_mapping_fp8, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.draft_slot_mapping_scale = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="draft_slot_mapping_scale", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_draft_slot_mapping_scale = torch.zeros_like( + self.draft_slot_mapping_scale, + device="cpu", + pin_memory=prefer_pinned(), + ) # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory. if self.enable_context_mla_with_cached_kv: @@ -594,6 +634,17 @@ def create_buffers_for_indexer(self, capture_graph=False): dtype=torch.int32, capture_graph=capture_graph, ) + if is_dsa_cache_manager(self.draft_kv_cache_manager): + self.draft_block_table = self.get_empty( + self.cuda_graph_buffers, + [ + self.max_num_sequences, + self.draft_kv_cache_manager.max_blocks_per_seq, + ], + cache_name="draft_block_table", + dtype=torch.int32, + capture_graph=capture_graph, + ) self.scheduler_metadata_buffer = self.get_empty( self.cuda_graph_buffers, (self.num_sms + 1, 2), @@ -744,6 +795,22 @@ def create_expanded_buffers(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) + if is_dsa_cache_manager(self.draft_kv_cache_manager): + self.draft_block_table_expanded = self.get_empty( + self.cuda_graph_buffers, + [ + self.max_num_sequences * (1 + self.max_draft_tokens), + self.draft_kv_cache_manager.max_blocks_per_seq, + ], + cache_name="draft_block_table_expanded", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_draft_block_table_expanded = torch.zeros_like( + self.draft_block_table_expanded, + device="cpu", + pin_memory=prefer_pinned(), + ) self.scheduler_metadata_buffer_expanded = self.get_empty( self.cuda_graph_buffers, (self.num_sms + 1, 2), @@ -805,31 +872,26 @@ def update_spec_dec_param( # elements"). self._create_radix_aux_buffers(capture_graph=capture_graph) - def _get_pool_block_indices(self) -> torch.Tensor: - """Extract memory pool block indices from host_kv_cache_block_offsets. - - The C++ setOffsets() encodes offsets as: - encoded = memPoolBlockIndex * numLayers * kvFactor - For SELFKONLY (MLA/DSA), kvFactor=1, so: - memPoolBlockIndex = encoded // num_local_layers - - Returns a (num_seqs, max_blocks_per_seq) int32 CPU tensor with valid - pool indices clamped to [0, blocks_in_primary_pool - 1]. - """ - num_local_layers = self.kv_cache_manager.num_local_layers - max_pool_idx = self.kv_cache_manager.blocks_in_primary_pool - 1 - # DSA uses SELFKONLY mode where only key cache is stored (kv_factor=1). - # host_kv_cache_block_offsets shape: (num_pools, max_batch*beam, 2, max_blocks_per_seq) - # Note: dim=2 is always 2 in the tensor layout (K and V slots), but for - # SELFKONLY only the K slot (index 0) contains valid data. - assert self.kv_cache_manager.kv_factor == 1, ( - f"DSA requires SELFKONLY mode (kv_factor=1), got kv_factor={self.kv_cache_manager.kv_factor}" + def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: + """Refresh INDEX_KEY offsets and return their physical pool slots.""" + cache_manager = self.kv_cache_manager + # Raw block IDs can exceed the primary pool after host-cache offload; + # the manager resolves their current physical slots. + pool_indices = cache_manager.get_pool_block_indices( + self.num_seqs, + request_ids=self.request_ids, + num_contexts=self.num_contexts, + beam_width=self.beam_width, + ) + page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale + num_blocks = page_indices.shape[1] + self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices) + self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( + self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True ) - # Pool 0, first num_seqs entries, field 0 (key offsets) - encoded = self.kv_cache_manager.host_kv_cache_block_offsets[0, : self.num_seqs, 0, :] - pool_indices = encoded // num_local_layers - # Clamp for safety: handles garbage padding from torch.empty in uninitialized slots - pool_indices = pool_indices.clamp(min=0, max=max_pool_idx).to(torch.int32) + # Sanitize graph-padding entries that may be stale after cache + # eviction or host-cache onboarding. + self.indexer_k_cache_block_offsets.clamp_(min=0) return pool_indices def set_skip_topk(self, skip: bool) -> None: @@ -853,10 +915,10 @@ def _ensure_pool_view_cached(self): """Compute and cache values used by transform_local_topk_and_prepare_pool_view(). - These values (pool view, stride factor, block table slices, request - index slices) are constant across all layers sharing the same KV pool - and batch dimensions within a forward pass. Caching them avoids - redundant Python/CUDA overhead per layer. + These values (pool view, block table slices, and request index slices) + are constant across all layers sharing the same KV pool and batch + dimensions within a forward pass. Caching them avoids redundant + Python/CUDA overhead per layer. Safety: _invalidate_pool_view_cache() is called unconditionally at the start of every step (prepare() and on_update_kv_lens()), so the boolean @@ -867,11 +929,9 @@ def _ensure_pool_view_cached(self): pool = self.kv_cache_manager.get_unique_primary_pool() kv_cache_manager = self.kv_cache_manager - num_blocks, num_layers, _, _ = pool.shape self._cached_tokens_per_block = kv_cache_manager.tokens_per_block head_dim = kv_cache_manager.head_dim self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim) - self._cached_stride_factor = num_layers * self._cached_tokens_per_block self._cached_block_table_ctx = self.block_table[: self.num_contexts] self._cached_block_table_gen = self.block_table[self.num_contexts : self.num_seqs] self._cached_req_idx_ctx = self.req_idx_per_token[: self.num_ctx_tokens] @@ -969,24 +1029,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.kv_lens_expanded_host[:num_tokens], non_blocking=True ) - # Expand indexer_k_cache_block_offsets (only generation) - # host_indexer_k_cache_block_offsets already contains correct pool - # indices from _get_pool_block_indices() in prepare_for_indexer_k_cache(). - if self.kv_cache_manager is not None and self.num_generations > 0: - max_len = self.host_indexer_k_cache_block_offsets.shape[1] - gen_block_tensor = self.host_indexer_k_cache_block_offsets[ - self.num_contexts : self.num_seqs, :max_len - ] - expanded_blocks = gen_block_tensor.repeat_interleave( - 1 + self.max_draft_tokens, dim=0 - ) - self.host_block_table_expanded[:num_tokens, :max_len].copy_( - expanded_blocks, non_blocking=True - ) - self.block_table_expanded[:num_tokens].copy_( - self.host_block_table_expanded[:num_tokens], non_blocking=True - ) - self.block_table_expanded.clamp_(min=0) + self._refresh_expanded_block_table(1 + self.max_draft_tokens) self.expand_for_dsl = ( use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1 @@ -1012,39 +1055,42 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.kv_lens_expanded_cuda[:num_tokens].copy_( self.kv_lens_expanded_host[:num_tokens], non_blocking=True ) - max_len = self.host_indexer_k_cache_block_offsets.shape[1] - gen_block_tensor = self.host_indexer_k_cache_block_offsets[ - self.num_contexts : self.num_seqs, :max_len - ] - expanded_blocks = gen_block_tensor.repeat_interleave(expand_factor, dim=0) - self.host_block_table_expanded[:num_tokens, :max_len].copy_( - expanded_blocks, non_blocking=True - ) - self.block_table_expanded[:num_tokens].copy_( - self.host_block_table_expanded[:num_tokens], non_blocking=True - ) - self.block_table_expanded.clamp_(min=0) + self._refresh_expanded_block_table(expand_factor) else: self.dsl_expand_factor = 1 self.dsl_atom = 1 + self.max_draft_tokens + def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None): + """Refresh the active cache's expanded INDEX_KEY page table.""" + if self.kv_cache_manager is None or self.num_generations == 0: + return + if repeat_factor is None: + if self.use_expanded_buffers_for_mtp: + repeat_factor = 1 + self.max_draft_tokens + elif self.expand_for_dsl and self.dsl_expand_factor > 1: + repeat_factor = self.dsl_expand_factor + else: + return + + num_tokens = self.num_generations * repeat_factor + max_len = self.host_indexer_k_cache_block_offsets.shape[1] + gen_block_tensor = self.host_indexer_k_cache_block_offsets[ + self.num_contexts : self.num_seqs, :max_len + ] + expanded_blocks = gen_block_tensor.repeat_interleave(repeat_factor, dim=0) + self.host_block_table_expanded[:num_tokens, :max_len].copy_( + expanded_blocks, non_blocking=True + ) + self.block_table_expanded[:num_tokens].copy_( + self.host_block_table_expanded[:num_tokens], non_blocking=True + ) + self.block_table_expanded.clamp_(min=0) + def prepare_for_indexer_k_cache(self): - # Build indexer_k_cache_block_offsets using pool block indices derived - # from host_kv_cache_block_offsets (populated by super().prepare()). - # This correctly resolves block IDs to memory pool indices, which is - # required when host cache offload is enabled (block IDs != pool indices - # for onboarded secondary blocks). if self.kv_cache_manager is None: return - pool_indices = self._get_pool_block_indices() - self.host_indexer_k_cache_block_offsets[: self.num_seqs].copy_(pool_indices) - self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True - ) - # Safety clamp: prevent OOB from CUDA graph padding entries which - # may contain stale negative or out-of-range values after block - # eviction/onboarding with host cache offload. - self.indexer_k_cache_block_offsets.clamp_(min=0) + # Keep physical slots for the primary-pool TopK conversion below. + pool_indices = self._update_indexer_k_cache_block_offsets() # Build block_table for topk_indices conversion (actual block allocation) cached_token_lens = torch.tensor( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/registry.py b/tensorrt_llm/_torch/attention_backend/sparse/registry.py index cd2cea5d2db2..298462b0f409 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/registry.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/registry.py @@ -18,18 +18,21 @@ # loaded. -def get_sparse_attn_kv_cache_manager(sparse_attention_config: "SparseAttentionConfig") -> type: +def get_sparse_attn_kv_cache_manager( + sparse_attention_config: "SparseAttentionConfig", + use_kv_cache_manager_v2: bool = False, +) -> type: from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from .deepseek_v4 import DeepseekV4CacheManager - from .dsa import DSACacheManager + from .dsa import DSACacheManager, DSACacheManagerV2 from .minimax_m3 import MiniMaxM3KVCacheManagerV2 from .rocket import RocketKVCacheManager if sparse_attention_config.algorithm == "rocket": return RocketKVCacheManager elif sparse_attention_config.algorithm == "dsa": - return DSACacheManager + return DSACacheManagerV2 if use_kv_cache_manager_v2 else DSACacheManager elif sparse_attention_config.algorithm == "deepseek_v4": return DeepseekV4CacheManager elif sparse_attention_config.algorithm == "skip_softmax": diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index dd756d561849..d421e25ff5d6 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -158,11 +158,12 @@ def effective_beam_width(self) -> int: kv_cache_block_offsets: Optional[torch.Tensor] = None host_kv_cache_block_offsets: Optional[torch.Tensor] = None draft_kv_cache_block_offsets: Optional[torch.Tensor] = None - # Block IDs per sequence; populated in __post_init__ when a KV cache - # manager is present. Declared here so encoder-only metadata (no KV cache) - # still exposes the attribute. + # Active block-ID buffers, defaulting to the target cache. Separate draft + # storage lets CUDA graphs bind stable target and draft addresses. block_ids_per_seq: Optional[torch.Tensor] = None kv_block_ids_per_seq: Optional[torch.Tensor] = None + draft_block_ids_per_seq: Optional[torch.Tensor] = None + draft_kv_block_ids_per_seq: Optional[torch.Tensor] = None # Pre-computed FlashMLA tile-scheduler metadata and num_splits. # Computed once per forward pass in TrtllmAttention.forward() and reused across layers. @@ -335,6 +336,8 @@ def _post_init_with_buffers(self, buffers) -> None: self.host_kv_cache_block_offsets = self.kv_cache_manager.host_kv_cache_block_offsets self.block_ids_per_seq = None self.kv_block_ids_per_seq = None + self.draft_block_ids_per_seq = None + self.draft_kv_block_ids_per_seq = None # Allocate separate block offset tensors for draft KV cache manager # Used in one-model speculative decoding with different KV cache layouts @@ -374,6 +377,27 @@ def _post_init_with_buffers(self, buffers) -> None: dtype=torch.int32, capture_graph=capture_graph, ) + if self.draft_kv_cache_manager is not None: + self.draft_block_ids_per_seq = self.get_empty( + buffers, + [ + self.draft_kv_cache_manager.max_batch_size, + self.draft_kv_cache_manager.max_blocks_per_seq + ], + cache_name="draft_block_ids_per_seq", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.draft_kv_block_ids_per_seq = self.get_empty( + buffers, + [ + self.draft_kv_cache_manager.max_batch_size, + self.draft_kv_cache_manager.max_blocks_per_seq + ], + cache_name="draft_kv_block_ids_per_seq", + dtype=torch.int32, + capture_graph=capture_graph, + ) # Allocate fixed-size buffers for pre-computed FlashMLA metadata. # These are pre-allocated so their GPU addresses are stable across CUDA graph captures. sm_count = torch.cuda.get_device_properties( @@ -753,9 +777,12 @@ def prepare_encoder_cuda_graph_replay(self, seq_lens: List[int], self.host_total_kv_lens[0] = padded_num_tokens def prepare_flash_mla(self) -> None: - # Invalidate the pre-computed metadata so that forward() recomputes it - # for this forward pass before the first attention layer runs. self._flash_mla_metadata_valid = False + # Request-specific fills and H2D copies must happen before replay, not + # become fixed operations in the captured graph. + if torch.cuda.is_current_stream_capturing(): + return + block_ids_per_seq = maybe_pin_memory( self.kv_cache_manager.get_block_ids_per_seq(self.request_ids)) num_blocks = block_ids_per_seq.shape[1] @@ -1697,9 +1724,9 @@ def forward( self, q, k, metadata, forward_args) # Compute FlashMLA tile-scheduler metadata once per forward pass. - # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger - # recomputation when cache_seq_lens change. The metadata must always match the - # compacted generation sub-batch, which is also the layout used by block_ids_per_seq. + # The flag is invalidated whenever FlashMLA inputs change. The metadata + # must always match the compacted generation sub-batch, which is also + # the layout used by block_ids_per_seq. if (metadata.enable_flash_mla and forward_args.attention_input_type != AttentionInputType.context_only and metadata.num_generations > 0 diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 3569426f5849..9755f0fc2806 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1907,6 +1907,13 @@ def forward( class DeepseekV3ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV3Model, PretrainedConfig]): + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: Any = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for this model implementation.""" + return "V2" + @classmethod def get_preferred_transceiver_runtime( cls, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index a4260e15dbb0..7bbcd0ce9186 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -31,7 +31,7 @@ import copy import math import os -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Literal, Optional if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import TorchLlmArgs @@ -2537,11 +2537,17 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: return { "kv_cache_config": { "tokens_per_block": 128, - "use_kv_cache_manager_v2": True, "enable_swa_scratch_reuse": True, } } + @classmethod + def get_preferred_kv_cache_manager_version( + cls, pretrained_config: object | None = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for DeepSeek-V4.""" + return "V2" + def __init__(self, model_config: ModelConfig[PretrainedConfig]): model_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config) self.mapping_with_cp = None diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 6ba492929cd5..d333379c78a9 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -1,4 +1,7 @@ -from typing import TYPE_CHECKING, Any, Dict, Literal, Optional +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Dict, Literal, Optional import torch from torch import nn @@ -32,9 +35,6 @@ from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, filter_weights, register_auto_model -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import TorchLlmArgs - # Use TinyGEMM when the number of tokens is not larger than this threshold MIN_LATENCY_TINYGEMM_NUM_TOKENS = 128 @@ -556,8 +556,10 @@ def forward( class GptOssForCausalLM(SpecDecOneEngineForCausalLM[Transformer, GptOssConfig]): @classmethod - def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: - """Select KV cache manager V2 by default. + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: Any = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for the model's VSWA layout. GPT-OSS applies a sliding window to every other layer (see ``AttentionBlock.__init__``), so the KV cache is VSWA: two @@ -566,15 +568,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: sliding-window and full-attention pools independently instead of statically dividing memory between them. - Users keep full control: an explicit - ``kv_cache_config.use_kv_cache_manager_v2`` otherwise wins over this - default. Two-model speculative decoding is the exception, since V2 - sizes both the target and draft KV cache managers from the full - budget: ``auto`` demotes this default to V1 there and an explicit - ``True`` is rejected, both in - ``llm_utils._resolve_kv_cache_manager_v2_auto``. + The preference is adopted only when the user leaves + ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``. Two-model + speculative decoding demotes it to V1 because V2 sizes both the target + and draft KV cache managers from the full budget; an explicit ``True`` + is rejected by ``llm_utils._resolve_kv_cache_manager_v2_auto``. """ - return {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + return "V2" @classmethod def get_preferred_transceiver_runtime( diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 8eed050eafe2..723f242cb613 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -1518,9 +1518,12 @@ class KimiK25ForConditionalGeneration(PreTrainedModel): _LANG_PREFIX = "language_model." @classmethod - def get_model_defaults(cls, llm_args: Any) -> dict: - """Use the C++ KV cache manager V2 by default.""" - return {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None, + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for Kimi K2.5.""" + return "V2" @classmethod def get_preferred_transceiver_runtime( diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm2.py b/tensorrt_llm/_torch/models/modeling_minimaxm2.py index e35f8bf1b73b..2ca84ee1f2a0 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm2.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm2.py @@ -13,15 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import Dict, List, Literal, Optional import torch from torch import nn from transformers import PretrainedConfig -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import TorchLlmArgs - from tensorrt_llm._ipc_utils import can_access_peer from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType from tensorrt_llm.mapping import Mapping @@ -403,9 +400,11 @@ def forward( @register_auto_model("MiniMaxM2ForCausalLM") class MiniMaxM2ForCausalLM(DecoderModelForCausalLM[MiniMaxM2Model, PretrainedConfig]): @classmethod - def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: - """Use KV cache manager V2 by default.""" - return {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + def get_preferred_kv_cache_manager_version( + cls, pretrained_config: object | None = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for MiniMax M2.""" + return "V2" def __init__(self, model_config: ModelConfig[PretrainedConfig]): super().__init__( diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 4484ee8315b4..7bc4e125c968 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -985,16 +985,22 @@ def load_weights(self, def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Model-specific defaults for NemotronH. - Uses KV cache manager V2 for the hybrid state layout. Block reuse - remains opt-in because it also requires a Mamba snapshot policy. + Block reuse remains opt-in because it also requires a Mamba snapshot + policy. """ return { "kv_cache_config": { "enable_block_reuse": False, - "use_kv_cache_manager_v2": True, } } + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: object + | None = None) -> Literal["V2"]: + """Prefer KV cache manager V2 for the hybrid state layout.""" + return "V2" + @classmethod def get_preferred_transceiver_runtime(cls, pretrained_config: object diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 8acf89ab927b..b71fe8495928 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -701,10 +701,17 @@ def get_model_defaults(cls, llm_args): # model class (this VLM wrapper), not on the inner decoder. Both # inner LMs (`Qwen3_5MoeForCausalLM` / `Qwen3_5ForCausalLM`) inherit # `Qwen3NextForCausalLM`'s defaults unchanged, so delegate to it to - # propagate the hybrid manager selection and keep block reuse disabled - # until a recurrent-state snapshot policy is configured. + # keep block reuse disabled until a recurrent-state snapshot policy is + # configured. return Qwen3NextForCausalLM.get_model_defaults(llm_args) + @classmethod + def get_preferred_kv_cache_manager_version( + cls, pretrained_config: object | None = None + ) -> Literal["V2"]: + """Match the hybrid text decoder's KV cache manager preference.""" + return "V2" + @classmethod def get_preferred_transceiver_runtime( cls, pretrained_config: object | None = None diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 268429d038f3..aa2626bf1da7 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -1053,13 +1053,15 @@ def __init__( @classmethod def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict: - """Use V2 by default; explicit V1 selections remain supported.""" - return { - "kv_cache_config": { - "enable_block_reuse": False, - "use_kv_cache_manager_v2": True, - } - } + """Disable block reuse until a snapshot policy is configured.""" + return {"kv_cache_config": {"enable_block_reuse": False}} + + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: object + | None = None) -> Literal["V2"]: + """Prefer KV cache manager V2 for the hybrid state layout.""" + return "V2" @classmethod def get_preferred_transceiver_runtime(cls, diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1b4b5f84e52e..bb81db4cf27a 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -665,6 +665,23 @@ def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict: """ return {} + @classmethod + def get_preferred_kv_cache_manager_version( + cls, + pretrained_config: Any = None) -> Optional[Literal["V1", "V2"]]: + """Return the model's preferred KV cache manager version. + + The preference is adopted only when the user leaves + ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``. Return + ``None`` to use the built-in V1 fallback. + + Args: + pretrained_config: The loaded Hugging Face config. Shared model + implementations can inspect it to select a preference for the + original checkpoint architecture. + """ + return None + @classmethod def get_preferred_transceiver_runtime( cls, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c4fb115111fe..ea73b8e37479 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -129,6 +129,7 @@ def get_kv_cache_manager_cls( config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) + use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True if is_hybrid_linear(config): # Degenerate case: model is flagged as hybrid but the config has zero # mamba layers. Fall through to the standard non-hybrid routes. @@ -136,7 +137,8 @@ def get_kv_cache_manager_cls( logger.info("Hybrid linear model has 0 mamba layers; using " "KV cache manager without mamba caching") if sparse_attn_config is not None: - return get_sparse_attn_kv_cache_manager(sparse_attn_config) + return get_sparse_attn_kv_cache_manager( + sparse_attn_config, use_kv_cache_manager_v2=use_v2) return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) if (sparse_attn_config is not None @@ -149,8 +151,6 @@ def get_kv_cache_manager_cls( has_additional_snapshots = bool( state_config.additional_snapshot_offsets_from_start or state_config.additional_snapshot_offsets_from_end) - use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True - if has_additional_snapshots and not use_v2: raise ValueError("Mamba additional snapshot offsets require " "use_kv_cache_manager_v2=True; V1 supports only " @@ -262,7 +262,8 @@ def get_kv_cache_manager_cls( "yet model retained recurrent-state snapshots.") return MambaHybridCacheManagerV2 elif sparse_attn_config is not None: - return get_sparse_attn_kv_cache_manager(sparse_attn_config) + return get_sparse_attn_kv_cache_manager(sparse_attn_config, + use_kv_cache_manager_v2=use_v2) else: return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) @@ -586,7 +587,7 @@ def _get_model_kv_cache_manager_cls( kv_cache_config, is_disagg=self._is_disagg, cache_transceiver_config=self._cache_transceiver_config) - cls = self._fallback_if_unsupported_kv_cache_manager_v2( + cls = self._validate_or_fallback_kv_cache_manager_v2( cls, model_config, kv_cache_config) if is_hybrid_linear(model_config.pretrained_config): logger.info_once( @@ -604,7 +605,7 @@ def _get_model_kv_cache_manager_cls( f"when using non-V2 Mamba cache manager {cls.__name__}") return cls - def _fallback_if_unsupported_kv_cache_manager_v2( + def _validate_or_fallback_kv_cache_manager_v2( self, kv_cache_manager_cls, model_config: ModelConfig, @@ -619,25 +620,26 @@ def _fallback_if_unsupported_kv_cache_manager_v2( incompat.append("kv_connector_manager") if self._max_beam_width is not None and self._max_beam_width > 1: incompat.append("max_beam_width > 1") + sparse_attn_config = model_config.sparse_attention_config + if (sparse_attn_config is not None + and sparse_attn_config.algorithm == "dsa" + and self._mapping.cp_config.get("cp_type") == CpType.STAR): + incompat.append("STAR context parallelism") if incompat: incompat_str = ", ".join(incompat) - # Some models are structurally bound to V2 and cannot fall - # back to V1 without producing wrong outputs: - # * Sparse-attention models (e.g. MiniMax-M3) need V2's - # per-layer split-pool to allocate the per-sparse-layer - # INDEX_KEY pool with a different stride than the main - # K/V pool. V1's unified pool cannot represent that. - # * Gemma4 hybrid uses per-layer head_dim that V1 would - # coerce to ``max(head_dim)``, changing per-layer KV - # byte sizes — correctness bug, not just efficiency. - sparse_attn_config = model_config.sparse_attention_config + # Never silently replace a sparse V2 manager with V1. Some + # sparse models require V2 structurally; for models such as DSA + # that support both managers, fallback would ignore the user's + # explicit manager selection. if sparse_attn_config is not None: raise NotImplementedError( - f"Sparse-attention models " - f"(algorithm={sparse_attn_config.algorithm!r}) require " - f"KVCacheManagerV2, which is not yet supported with " - f"{incompat_str}. Disable these KvCacheConfig features " - f"to run sparse-attention models.") + f"KVCacheManagerV2 for sparse-attention models " + f"(algorithm={sparse_attn_config.algorithm!r}) is not " + f"supported with " + f"{incompat_str}. Disable the incompatible features to " + f"run sparse-attention models.") + # Gemma4 hybrid uses per-layer head_dim that V1 would coerce to + # ``max(head_dim)``, changing per-layer KV byte sizes. if is_gemma4_hybrid(config): raise NotImplementedError( f"Gemma4 hybrid attention requires KVCacheManagerV2, " @@ -649,7 +651,7 @@ def _fallback_if_unsupported_kv_cache_manager_v2( f"{incompat_str}; CppMambaHybridCacheManager does not " "provide a compatible fallback. Use max_beam_width=1 " "and disable the KV connector.") - # Plain V2 (explicitly enabled or selected by a model default): + # Plain V2 (explicitly enabled or selected by a model preference): # V2 was a preference, not a structural requirement, so we can # safely fall back to V1. logger.warning( @@ -1440,7 +1442,7 @@ def _create_one_model_draft_kv_cache_manager( # Get the appropriate KV cache manager class for the draft model draft_kv_cache_manager_cls = get_kv_cache_manager_cls( effective_draft_config, draft_kv_config, is_disagg=self._is_disagg) - draft_kv_cache_manager_cls = self._fallback_if_unsupported_kv_cache_manager_v2( + draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2( draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config) estimating_kv_cache = estimating_kv_cache and not self._skip_est diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2c567824a8cb..2852ac44fd57 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1080,11 +1080,23 @@ def append_to_kv_heads_per_layer( self._get_event_layer_group_ids(), ) + # Both backends build layer_grouping on demand, and the layer order + # within a group is not part of its contract. Cache a stable physical- + # layout representative for each role from the public pool descriptors. self.num_pools = len(self.impl.layer_grouping) - # num_pools is the physical pool count owned by the KV cache manager. - # 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. + self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {} + for pool_group in self.impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + pool_id = int(variant.layer_group_id) + for coalesced in variant.coalesced_buffers: + for buffer_id in coalesced.buffer_ids: + 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: self.num_attention_op_pools = self.num_local_layers else: @@ -1228,8 +1240,8 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_mapping_list.append([int(layer_id), 0]) else: for pool_id in range(self.num_pools): - layer_id = self.impl.layer_grouping[pool_id][0] role_a, _ = self._get_pool_roles(pool_id) + layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] key_base_addr = self.impl.get_mem_pool_base_address( layer_id, role_a, PageIndexMode.SHARED ) @@ -1342,8 +1354,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" ) for pool_id in range(self.num_pools): - layer_id = self.impl.layer_grouping[pool_id][0] role_a, role_b = self._get_pool_roles(pool_id) + layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, role_a) if role_b is not None: self.kv_offset[pool_id] = exact_div( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index bd832625d042..e4cf3fc129ae 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -7189,7 +7189,8 @@ def capture_postprocess_fn(inputs: Dict[str, Any]): if self.cuda_graph_runner.is_warmup_only: outputs = capture_outputs elif needs_capture: - # Pre-replay: set DSA slot mappings for current batch's draft cache (fixes 2nd warmup) + # Refresh attention metadata for the current batch's + # draft cache before replaying the captured graph. saved_draft = prepare_attn_metadata_for_draft_replay( attn_metadata, draft_kv_cache_manager) try: diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 7a852bf9bfd1..30e775a15d35 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -413,6 +413,7 @@ def load_config_and_apply_defaults( # No config to resolve a model class from; still resolve the # "auto" sentinel so it never leaks past config loading. _resolve_transceiver_runtime_auto(llm_args) + _resolve_kv_cache_manager_v2_auto(llm_args) return llm_args config_kwargs = { @@ -428,9 +429,19 @@ def load_config_and_apply_defaults( config = checkpoint_loader.load_config(checkpoint_dir, **config_kwargs) model_cls = AutoModelForCausalLM._resolve_class(config) - use_kv_cache_manager_v2 = ( + original_kv_cache_manager_setting = ( llm_args.kv_cache_config.use_kv_cache_manager_v2) + # Preferences follow the checkpoint's original architecture: + # _resolve_class may rewrite it to an execution class (e.g. + # MTPDraftModelForCausalLM), which must not drop the target model's + # preferences. + preference_cls = model_cls + architectures = getattr(config.pretrained_config, 'architectures', None) + if architectures: + preference_cls = get_registered_model_class( + architectures[0]) or model_cls + # model_cls is None when the architecture is unknown/unsupported. model_defaults = {} if model_cls and hasattr(model_cls, 'get_model_defaults'): @@ -452,23 +463,13 @@ def load_config_and_apply_defaults( update_spec_config_from_model_config(llm_args.speculative_config, config.pretrained_config) - # The transceiver preference follows the checkpoint's original - # architecture: _resolve_class may rewrite it to an execution class - # (e.g. MTPDraftModelForCausalLM), which must not drop the target - # model's preference. - preference_cls = model_cls - architectures = getattr(config.pretrained_config, 'architectures', None) - if architectures: - preference_cls = get_registered_model_class( - architectures[0]) or model_cls - # Resolve "auto" sentinel values after model defaults are applied. _resolve_transceiver_runtime_auto(llm_args, preference_cls, config.pretrained_config) - _resolve_kv_cache_manager_v2_auto( - llm_args, model_defaults, original_setting=use_kv_cache_manager_v2) + _resolve_kv_cache_manager_v2_auto(llm_args, preference_cls, + config.pretrained_config) _validate_and_adjust_mamba_snapshot_config(config, llm_args) - if use_kv_cache_manager_v2 == "auto": + if original_kv_cache_manager_setting == "auto": logger.info( "Resolved use_kv_cache_manager_v2='auto' to %s for %s", llm_args.kv_cache_config.use_kv_cache_manager_v2, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index c201c76592ed..34d8f150a866 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -939,6 +939,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, num_accepted_tokens, original_all_rank_num_tokens): """Linear draft loop, unified for Eagle3 and MTP Eagle.""" + from ..attention_backend.sparse.dsa import (DSAtrtllmAttentionMetadata, + is_dsa_cache_manager) + runtime_draft_len = spec_metadata.runtime_draft_len num_gens = batch_size - num_contexts next_draft_tokens = [] @@ -946,9 +949,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 position_ids = inputs["position_ids"] - reuse_mtp_topk = (self.is_mtp_eagle - and hasattr(attn_metadata, "set_skip_topk")) - if reuse_mtp_topk: + uses_dsa_mtp_metadata = self.is_mtp_eagle and isinstance( + attn_metadata, DSAtrtllmAttentionMetadata) + if uses_dsa_mtp_metadata: attn_metadata.set_in_mtp_draft_loop(True) # Accepted counts let the indexer stash each gen's last-accepted row. attn_metadata.set_mtp_num_accepted(num_accepted_tokens) @@ -957,8 +960,16 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, attn_metadata, draft_kv_cache_manager) as draft_attn_metadata: attn_metadata = draft_attn_metadata inputs["attn_metadata"] = draft_attn_metadata + if uses_dsa_mtp_metadata and is_dsa_cache_manager( + draft_kv_cache_manager): + # Overlap scheduling corrects kv_lens_cuda from the runtime + # accepted-token counts inside the captured graph. The target + # forward refreshes target slot mappings during that correction; + # refresh once more after rebinding so the first draft forward + # writes the separate DSA indexer cache at the same positions. + attn_metadata.on_update_kv_lens() for i in range(runtime_draft_len): - if reuse_mtp_topk: + if uses_dsa_mtp_metadata: attn_metadata.set_skip_topk(i > 0) # Run draft model (mode-specific via helper). The helper # passes ``all_rank_num_tokens`` as a kwarg so the draft model @@ -1156,7 +1167,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, } next_draft_tokens = torch.stack(next_draft_tokens, dim=1) - if reuse_mtp_topk: + if uses_dsa_mtp_metadata: attn_metadata.set_skip_topk(False) attn_metadata.set_in_mtp_draft_loop(False) attn_metadata.set_mtp_num_accepted(None) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index bdb122ac4457..8e51c1379a7f 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -127,9 +127,11 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool: def prepare_attn_metadata_for_draft_replay(attn_metadata, draft_kv_cache_manager): """ - Prepare attention metadata for CUDA graph replay when using separate draft KV cache. - Swaps to draft manager and (for DSA) re-prepares indexer slot mappings for the current - batch. Call restore_attn_metadata_after_draft_replay after replay in a finally block. + Prepare attention metadata for a draft forward or CUDA graph replay when using a + separate draft KV cache. Swaps cache-layout-dependent buffers, refreshes FlashMLA + block IDs outside capture, and (for DSA) re-prepares indexer slot mappings + for the current batch. + Call restore_attn_metadata_after_draft_replay in a finally block. Returns saved state or None if no-op. """ if draft_kv_cache_manager is None: @@ -149,6 +151,18 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata, 'target_host_kv_cache_block_offsets': attn_metadata.host_kv_cache_block_offsets, } + if attn_metadata.enable_flash_mla: + if (attn_metadata.draft_block_ids_per_seq is None + or attn_metadata.draft_kv_block_ids_per_seq is None): + raise RuntimeError( + "FlashMLA separate draft KV cache requires dedicated draft block-ID buffers" + ) + saved['target_block_ids_per_seq'] = attn_metadata.block_ids_per_seq + saved[ + 'target_kv_block_ids_per_seq'] = attn_metadata.kv_block_ids_per_seq + attn_metadata.block_ids_per_seq = attn_metadata.draft_block_ids_per_seq + attn_metadata.kv_block_ids_per_seq = ( + attn_metadata.draft_kv_block_ids_per_seq) attn_metadata.kv_cache_manager = draft_kv_cache_manager attn_metadata.kv_cache_block_offsets = attn_metadata.draft_kv_cache_block_offsets attn_metadata.host_kv_cache_block_offsets = ( @@ -157,45 +171,63 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata, attn_metadata.prepare_flash_mla() from ..attention_backend.sparse.dsa import (DSAtrtllmAttentionMetadata, - Indexer) + Indexer, is_dsa_cache_manager) + + # DeepSeek-V4 metadata inherits DSA metadata, but its cache manager uses a + # different dual-pool layout. Only native DSA cache managers use the DSA + # draft-replay buffers below. if (isinstance(attn_metadata, DSAtrtllmAttentionMetadata) - and hasattr(draft_kv_cache_manager, 'index_head_dim')): + and is_dsa_cache_manager(draft_kv_cache_manager)): m = attn_metadata saved['saved_dsa_state'] = { 'host_indexer_k_cache_block_offsets': - m.host_indexer_k_cache_block_offsets.clone(), - 'indexer_k_cache_block_offsets': - m.indexer_k_cache_block_offsets.clone(), - 'host_slot_mapping_fp8': - m.host_slot_mapping_fp8.clone(), - 'host_slot_mapping_scale': - m.host_slot_mapping_scale.clone(), - 'slot_mapping_fp8': - m.slot_mapping_fp8.clone(), - 'slot_mapping_scale': - m.slot_mapping_scale.clone(), + m.host_indexer_k_cache_block_offsets, + 'indexer_k_cache_block_offsets': m.indexer_k_cache_block_offsets, + 'host_slot_mapping_fp8': m.host_slot_mapping_fp8, + 'host_slot_mapping_scale': m.host_slot_mapping_scale, + 'slot_mapping_fp8': m.slot_mapping_fp8, + 'slot_mapping_scale': m.slot_mapping_scale, + 'block_table': m.block_table, + 'block_table_expanded': m.block_table_expanded, + 'host_block_table_expanded': m.host_block_table_expanded, } - # Derive pool indices from the draft manager's encoded block - # offsets (via _get_pool_block_indices) instead of using raw block - # IDs. With host cache offload, block IDs can exceed - # blocks_in_primary_pool after offload swaps (the block keeps its - # original high ID even though its memory now lives in the primary - # GPU pool). Using raw block IDs as pool indices causes OOB access - # in the indexer k-cache buffers. _get_pool_block_indices correctly - # decodes memPoolBlockIndex from the C++ encoded offsets. - # Note: kv_cache_manager was already swapped to draft above (line 67). - pool_indices = m._get_pool_block_indices() - num_blocks = pool_indices.shape[1] - m.host_indexer_k_cache_block_offsets[:m.num_seqs, :num_blocks].copy_( - pool_indices) - m.indexer_k_cache_block_offsets[:m.num_seqs].copy_( - m.host_indexer_k_cache_block_offsets[:m.num_seqs], - non_blocking=True) - # Safety clamp: sanitize stale padding entries beyond num_seqs - # that may contain negative or out-of-range values, matching the - # regular DSA prepare() flow. - m.indexer_k_cache_block_offsets.clamp_(min=0) - Indexer.recompute_slot_mappings(m) + # The cached-KV feature owns these references even when an optimized + # path aliases them to slot_mapping_*. With the feature disabled, the + # aliases are lazy and may not exist on the first generation replay. + if m.enable_context_mla_with_cached_kv: + saved['saved_dsa_state'].update({ + 'slot_mapping_fp8_fullkv': + m.slot_mapping_fp8_fullkv, + 'slot_mapping_scale_fullkv': + m.slot_mapping_scale_fullkv, + }) + # Rebind to the draft manager's dedicated buffers instead of + # overwriting the target tensors in place. Rebinding is invisible to + # CUDA graph capture, so the target and draft segments of the graph + # bake distinct addresses (like draft_kv_cache_block_offsets) and no + # graph-recorded copy from a transient host buffer is needed. + m.host_indexer_k_cache_block_offsets = ( + m.host_draft_indexer_k_cache_block_offsets) + m.indexer_k_cache_block_offsets = m.draft_indexer_k_cache_block_offsets + m.host_slot_mapping_fp8 = m.host_draft_slot_mapping_fp8 + m.slot_mapping_fp8 = m.draft_slot_mapping_fp8 + m.host_slot_mapping_scale = m.host_draft_slot_mapping_scale + m.slot_mapping_scale = m.draft_slot_mapping_scale + m.block_table = m.draft_block_table + m.block_table_expanded = m.draft_block_table_expanded + m.host_block_table_expanded = m.host_draft_block_table_expanded + m._invalidate_pool_view_cache() + # Recording a capture executes no kernels, so the draft mappings only + # need refreshing when the transfers actually run: eager forwards + # (warmup) and the pre-replay call from model_engine. The per-step + # advance inside the captured graph re-derives slot mappings on + # device from the rebound block-offset buffer. + # kv_cache_manager was already swapped to the draft manager above. + if not torch.cuda.is_current_stream_capturing(): + m.prepare_for_indexer_k_cache() + m._refresh_expanded_block_table() + Indexer.recompute_slot_mappings(m) + Indexer.recompute_context_kv_gather_mappings(m) return saved @@ -209,18 +241,37 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state): attn_metadata.host_kv_cache_block_offsets = ( saved_state['target_host_kv_cache_block_offsets']) if attn_metadata.enable_flash_mla: - attn_metadata.prepare_flash_mla() + attn_metadata.block_ids_per_seq = saved_state[ + 'target_block_ids_per_seq'] + attn_metadata.kv_block_ids_per_seq = saved_state[ + 'target_kv_block_ids_per_seq'] + # Target and draft block-ID buffers are independent. Restoring only + # needs to invalidate the scheduler metadata; refreshing the unchanged + # target buffers would repeat request-specific H2D work. + attn_metadata._flash_mla_metadata_valid = False saved_dsa = saved_state.get('saved_dsa_state') if saved_dsa is not None: m = attn_metadata - m.host_indexer_k_cache_block_offsets.copy_( - saved_dsa['host_indexer_k_cache_block_offsets'], non_blocking=True) - m.indexer_k_cache_block_offsets.copy_( - saved_dsa['indexer_k_cache_block_offsets'], non_blocking=True) - m.host_slot_mapping_fp8.copy_(saved_dsa['host_slot_mapping_fp8']) - m.host_slot_mapping_scale.copy_(saved_dsa['host_slot_mapping_scale']) - m.slot_mapping_fp8.copy_(saved_dsa['slot_mapping_fp8']) - m.slot_mapping_scale.copy_(saved_dsa['slot_mapping_scale']) + m.host_indexer_k_cache_block_offsets = saved_dsa[ + 'host_indexer_k_cache_block_offsets'] + m.indexer_k_cache_block_offsets = saved_dsa[ + 'indexer_k_cache_block_offsets'] + m.host_slot_mapping_fp8 = saved_dsa['host_slot_mapping_fp8'] + m.host_slot_mapping_scale = saved_dsa['host_slot_mapping_scale'] + m.slot_mapping_fp8 = saved_dsa['slot_mapping_fp8'] + m.slot_mapping_scale = saved_dsa['slot_mapping_scale'] + m.block_table = saved_dsa['block_table'] + m.block_table_expanded = saved_dsa['block_table_expanded'] + m.host_block_table_expanded = saved_dsa['host_block_table_expanded'] + m._invalidate_pool_view_cache() + if 'slot_mapping_fp8_fullkv' in saved_dsa: + m.slot_mapping_fp8_fullkv = saved_dsa['slot_mapping_fp8_fullkv'] + m.slot_mapping_scale_fullkv = saved_dsa['slot_mapping_scale_fullkv'] + else: + # The draft recomputation rebound the aliases to the draft tensors; + # point them back at the restored target tensors. + m.slot_mapping_fp8_fullkv = m.slot_mapping_fp8 + m.slot_mapping_scale_fullkv = m.slot_mapping_scale def get_force_num_accepted_tokens() -> int: @@ -2139,7 +2190,8 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): """ Select draft attention metadata for one-engine speculative decoding. - TRTLLM metadata temporarily swaps its manager and block offsets. + TRTLLM metadata temporarily swaps its manager and cache-layout-dependent + buffers, including DSA indexer offsets and slot mappings. FlashInfer uses an independently planned metadata view because its page tables and kernel wrappers are manager-specific. """ @@ -2158,35 +2210,16 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): yield attn_metadata return - # Check if draft KV cache block offsets are allocated - draft_block_offsets = getattr(attn_metadata, - 'draft_kv_cache_block_offsets', None) - if draft_block_offsets is None: - # Draft KV cache block offsets not allocated, skip switching + saved_state = prepare_attn_metadata_for_draft_replay( + attn_metadata, draft_kv_cache_manager) + if saved_state is None: yield attn_metadata return - # Save main KV cache manager and block offsets - target_kv_cache_manager = attn_metadata.kv_cache_manager - target_kv_cache_block_offsets = attn_metadata.kv_cache_block_offsets - target_host_kv_cache_block_offsets = attn_metadata.host_kv_cache_block_offsets - - # Switch to draft KV cache manager and its block offsets - attn_metadata.kv_cache_manager = draft_kv_cache_manager - attn_metadata.kv_cache_block_offsets = attn_metadata.draft_kv_cache_block_offsets - attn_metadata.host_kv_cache_block_offsets = draft_kv_cache_manager.host_kv_cache_block_offsets - if attn_metadata.enable_flash_mla: - attn_metadata.prepare_flash_mla() - try: yield attn_metadata finally: - # Restore main KV cache manager and block offsets - attn_metadata.kv_cache_manager = target_kv_cache_manager - attn_metadata.kv_cache_block_offsets = target_kv_cache_block_offsets - attn_metadata.host_kv_cache_block_offsets = target_host_kv_cache_block_offsets - if attn_metadata.enable_flash_mla: - attn_metadata.prepare_flash_mla() + restore_attn_metadata_after_draft_replay(attn_metadata, saved_state) def _sample_tokens_for_batch( self, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1afb2cc25994..b7527d0ffd95 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3913,8 +3913,8 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): status="prototype", description= "Whether to use the KV cache manager v2 (experimental). 'auto' uses " - "the model-specific default and falls back to False when the model " - "does not specify one.") + "the model-specific preference and falls back to False when the model " + "does not declare one.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. enable_swa_scratch_reuse: bool = Field( diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index a8bd94bdb797..a21241b6bcf0 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -555,7 +555,7 @@ def apply_model_defaults_to_llm_args( new_args = llm_args.__class__(**merged_state) - for field_name in llm_args.model_fields: + for field_name in type(llm_args).model_fields: setattr(llm_args, field_name, getattr(new_args, field_name)) def _compute_applied(defaults: Dict[str, Any], @@ -597,13 +597,12 @@ def _two_model_spec_dec_decoding_type( return getattr(spec_config, "decoding_type", "speculative decoding") -def _resolve_kv_cache_manager_v2_auto( - llm_args: 'TorchLlmArgs', - model_defaults_dict: Dict[str, Any], - original_setting: Optional[Union[bool, str]] = None) -> bool: - """Resolve the KV cache manager auto setting after model defaults are applied. +def _resolve_kv_cache_manager_v2_auto(llm_args: 'TorchLlmArgs', + model_cls: Optional[type] = None, + pretrained_config: Any = None) -> bool: + """Resolve the KV cache manager auto setting from the model preference. - A model default of V2 is demoted to V1 for routes V2 cannot serve; an + A model preference for V2 is demoted to V1 for routes V2 cannot serve; an explicit user value otherwise wins. The compatibility arms are: - Disaggregated serving: hybrid Mamba V2 requires the Python transceiver @@ -613,16 +612,15 @@ def _resolve_kv_cache_manager_v2_auto( with its own KV cache manager. ``build_managers`` hands that manager the target's ``kv_cache_config`` unsplit, and V2 capacity is governed solely by ``max_gpu_total_bytes``, so both managers size their pools from the - full budget. The model default falls back to V1, and an explicit ``True`` - is rejected rather than deferred to that allocation. + full budget. The model preference falls back to V1, and an explicit + ``True`` is rejected rather than deferred to that allocation. The fallback only reaches models whose manager class is selected by ``use_kv_cache_manager_v2``. Models routed to a V2 manager unconditionally -- sparse attention picks its class from the algorithm alone -- keep a V2 manager after the demotion, so the arm does not protect them. """ - setting = (llm_args.kv_cache_config.use_kv_cache_manager_v2 - if original_setting is None else original_setting) + setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 if setting != "auto": if setting is True: decoding_type = _two_model_spec_dec_decoding_type(llm_args) @@ -637,40 +635,43 @@ def _resolve_kv_cache_manager_v2_auto( "one-model variant of this decoding mode.") return setting - kv_cache_defaults = model_defaults_dict.get("kv_cache_config", {}) - model_default = (kv_cache_defaults.get("use_kv_cache_manager_v2", False) - if isinstance(kv_cache_defaults, dict) else False) - if model_default == "auto": - model_default = False - if not isinstance(model_default, bool): + preferred_version = None + if model_cls is not None: + get_preferred = getattr(model_cls, + 'get_preferred_kv_cache_manager_version', None) + if get_preferred is not None: + preferred_version = get_preferred(pretrained_config) + if preferred_version not in (None, "V1", "V2"): raise ValueError( - "Model default kv_cache_config.use_kv_cache_manager_v2 must be " - f"True, False, or 'auto', got {model_default!r}.") + f"{model_cls.__name__}.get_preferred_kv_cache_manager_version() " + f"must return 'V1', 'V2', or None, got {preferred_version!r}.") + + use_v2 = preferred_version == "V2" transceiver_config = llm_args.cache_transceiver_config - if (model_default and transceiver_config is not None + if (use_v2 and transceiver_config is not None and transceiver_config.backend is not None): effective_backend, _ = transceiver_config._resolve_default_backend() runtime = transceiver_config.transceiver_runtime if effective_backend != "NIXL" or runtime != "PYTHON": logger.info( - "KV cache manager V2 is the model default, but disaggregated " + "KV cache manager V2 is the model preference, but disaggregated " "serving uses transceiver_runtime=%r with backend=%r; " "falling back to V1.", runtime, effective_backend) - model_default = False + use_v2 = False - if model_default: + if use_v2: decoding_type = _two_model_spec_dec_decoding_type(llm_args) if decoding_type is not None: logger.info( - "KV cache manager V2 is the model default, but %s runs the " + "KV cache manager V2 is the model preference, but %s runs the " "draft model in a separate engine and V2 sizes both KV cache " "managers from the full max_gpu_total_bytes budget; falling " "back to V1.", decoding_type) - model_default = False + use_v2 = False - llm_args.kv_cache_config.use_kv_cache_manager_v2 = model_default - return model_default + llm_args.kv_cache_config.use_kv_cache_manager_v2 = use_v2 + return use_v2 def _resolve_transceiver_runtime_auto(llm_args: 'TorchLlmArgs', diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8b00d737c793..c2835eca30bf 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7977,6 +7977,7 @@ def test_nvfp4(self, tp_size, ep_size): kv_cache_config=kv_cache_config, max_seq_len=8192, **pytorch_config) as llm: + assert llm.args.kv_cache_config.use_kv_cache_manager_v2 is True assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 task = GSM8K(model_name) task.evaluate(llm) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index ddf2514b7a19..0722e0bf1618 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -16,7 +16,7 @@ starts. | Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. | | Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. | | Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). | -| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves against the model class's `get_model_defaults()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. | +| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves via `get_preferred_kv_cache_manager_version()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. | Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are supported: data is seeded per (request, **global** layer) and constant along diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 8219a2581d9c..53b79919a9cf 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -307,8 +307,8 @@ def resolve_plan(cfg, benchmark_mode="e2e"): plan[f"{role}_kv_dtype"] = str(kv_cfg.get("dtype", "auto")) # Tri-state, matching KvCacheConfig's pydantic default: explicit # True/False from the yaml wins; absent means "auto", which the - # driver resolves against the model class's get_model_defaults() at - # runtime — exactly like serving (_resolve_kv_cache_manager_v2_auto). + # driver resolves against the model class's manager preference at + # runtime, exactly like serving (_resolve_kv_cache_manager_v2_auto). plan[f"{role}_use_kv_cache_manager_v2"] = kv_cfg.get("use_kv_cache_manager_v2", "auto") plan["fingerprint"] = plan_fingerprint(plan) return plan diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 146c40173cc3..e8dcea152fc9 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -321,7 +321,7 @@ def resolve_model_prefs(model_dir, side, cache_cfg): """Mirror serving's model-preference resolution (PR #15823 semantics). - use_kv_cache_manager_v2 == "auto" (yaml absent): adopt the model - class's get_model_defaults() value, default False + class's get_preferred_kv_cache_manager_version() value, falling back to V1 (llm_utils._resolve_kv_cache_manager_v2_auto). - cache_cfg.transceiver_runtime == "auto": adopt model_cls.get_preferred_transceiver_runtime(), NIXL-gated, via the @@ -349,25 +349,17 @@ def resolve_model_prefs(model_dir, side, cache_cfg): setting = side["use_kv_cache_manager_v2"] if setting == "auto": - defaults = {} - if model_cls is not None: - try: - defaults = model_cls.get_model_defaults(None) or {} - except Exception as e: # noqa: BLE001 - model hooks may need llm_args - print( - f"[precheck] WARNING: get_model_defaults failed ({e!r}); assuming V1", - flush=True, - ) try: # The REAL serving resolver, via the same shim pattern as the # runtime resolution below -- one owner for the 'auto' semantics. # cache_transceiver_config feeds the resolver's disagg gating - # (a V2 model default requires the NIXL Python transceiver). + # (a V2 model preference requires the NIXL Python transceiver). shim = types.SimpleNamespace( kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto"), cache_transceiver_config=cache_cfg, + speculative_config=None, ) - use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, defaults)) + use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, model_cls, hf_view)) except Exception as e: # noqa: BLE001 - fall back like a missing model print( f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); assuming V1", flush=True @@ -820,7 +812,7 @@ def __init__(self, args, plan, side, comm): self.kvm = None self.xcvr = None self.runtime = "CPP" - # Resolved in setup(): "auto" needs the model class (get_model_defaults). + # Resolved in setup(): "auto" needs the model preference hook. self.use_v2 = False self.mapping = None self.llm_request_state = None diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 3cd06c02f863..6ae2544df242 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -26,7 +26,7 @@ import random import sys from pathlib import Path -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -43,6 +43,7 @@ from tensorrt_llm._torch.attention_backend.sparse.dsa import ( DSABackendForwardArgs, DSACacheManager, + DSACacheManagerV2, DSATrtllmAttention, DSAtrtllmAttentionMetadata, Indexer, @@ -50,20 +51,25 @@ _select_indexer_compress_ratio, compute_cu_seqlen_kv_bounds_with_cache, split_prefill_chunks, + transform_local_topk_and_prepare_pool_view, ) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, ) from tensorrt_llm._utils import get_sm_version from tensorrt_llm.bindings import DataType -from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.bindings.executor import KvCacheConfig as BindingKvCacheConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.llmapi.llm_args import ( DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, + KvCacheConfig, ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.utils import fp8_utils @@ -250,8 +256,10 @@ def create_dsa_cache_manager( num_layers: int = 1, indexer_k_dtype: str = "fp8", index_topk: int = 2048, + use_kv_cache_manager_v2: bool = False, + pretrained_config=None, ): - """Helper to create a DSACacheManager for testing.""" + """Helper to create a DSA cache manager for testing.""" sparse_attn_config = DeepSeekSparseAttentionConfig( index_head_dim=head_dim, @@ -261,7 +269,8 @@ def create_dsa_cache_manager( ) # Create KV cache config - kv_cache_config = KvCacheConfig( + kv_cache_config_cls = KvCacheConfig if use_kv_cache_manager_v2 else BindingKvCacheConfig + kv_cache_config = kv_cache_config_cls( enable_block_reuse=False, max_tokens=max_seq_len * batch_size, ) @@ -271,7 +280,8 @@ def create_dsa_cache_manager( # Create cache manager # Use SELFKONLY for DSA (similar to MLA usage in _util.py) - cache_manager = DSACacheManager( + cache_manager_cls = DSACacheManagerV2 if use_kv_cache_manager_v2 else DSACacheManager + cache_manager = cache_manager_cls( kv_cache_config=kv_cache_config, kv_cache_type=CacheTypeCpp.SELFKONLY, num_layers=num_layers, @@ -283,11 +293,108 @@ def create_dsa_cache_manager( mapping=mapping, dtype=DataType.HALF, sparse_attention_config=sparse_attn_config, + pretrained_config=pretrained_config, ) return cache_manager, sparse_attn_config +def test_transform_local_topk_uses_v2_page_mapping(): + """Use the V2 pool scale and offset when flattening DSA indices.""" + cache_manager = SimpleNamespace( + _primary_pool_page_index_params=[(3, 1)], + get_primary_pool_page_index_params=None, + ) + cache_manager.get_primary_pool_page_index_params = MethodType( + DSACacheManagerV2.get_primary_pool_page_index_params, + cache_manager, + ) + req_idx = torch.tensor([0], dtype=torch.int32) + block_table = torch.tensor([[3, 7]], dtype=torch.int32) + topk_indices = torch.tensor([[6]], dtype=torch.int32) + metadata = SimpleNamespace( + kv_cache_manager=cache_manager, + _ensure_pool_view_cached=Mock(), + _cached_block_table_ctx=block_table, + _cached_req_idx_ctx=req_idx, + _cached_tokens_per_block=4, + _cached_pool_view=torch.empty(0), + ) + expected_indices = torch.tensor([[90]], dtype=torch.int32) + + with patch( + "torch.ops.trtllm.convert_req_index_to_global", + return_value=expected_indices, + ) as convert: + actual, _ = transform_local_topk_and_prepare_pool_view(topk_indices, metadata, layer_idx=0) + + torch.testing.assert_close(actual, expected_indices) + convert.assert_called_once_with(req_idx, block_table, topk_indices, 4, 1, 12, 1) + + +def test_dsa_cache_manager_v2_respects_shared_indexer_layer_mask(): + """V2 registers INDEX_KEY storage only for layers that own an indexer.""" + num_layers = 3 + head_dim = 128 + tokens_per_block = 16 + pretrained_config = SimpleNamespace( + num_hidden_layers=num_layers, + index_topk_pattern=["F", "S", "F"], + ) + sparse_config = DeepSeekSparseAttentionConfig(index_head_dim=head_dim) + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(), + sparse_attention_config=sparse_config, + ) + assert ( + get_kv_cache_manager_cls( + model_config, + KvCacheConfig(use_kv_cache_manager_v2=True), + ) + is DSACacheManagerV2 + ) + + cache_manager, _ = create_dsa_cache_manager( + batch_size=2, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=64, + num_layers=num_layers, + use_kv_cache_manager_v2=True, + pretrained_config=pretrained_config, + ) + + try: + assert cache_manager.indexer_k_cache_local_layer_mask == [True, False, True] + assert cache_manager.indexer_k_cache_page_scale == 2 + assert cache_manager.get_indexer_k_cache_buffers(0).shape == ( + cache_manager.blocks_in_primary_pool * cache_manager.indexer_k_cache_page_scale, + tokens_per_block, + 1, + head_dim + 4, + ) + + indexer_bytes_per_token = head_dim + 4 + has_indexer_buffers = [] + for local_layer_idx in range(num_layers): + layer_config = cache_manager.kv_cache_manager_py_config.layers[local_layer_idx] + has_indexer_buffers.append( + any(buffer.role == Role.INDEX_KEY for buffer in layer_config.buffers) + ) + assert has_indexer_buffers == [True, False, True] + + key_bytes = sum( + cache_manager.get_layer_bytes_per_token(local_layer_idx, Role.KEY) + for local_layer_idx in range(num_layers) + ) + expected_cache_bytes = key_bytes + 2 * indexer_bytes_per_token + assert cache_manager.get_cache_bytes_per_token() == expected_cache_bytes + with pytest.raises(AssertionError, match="shared-indexer layer"): + cache_manager.get_indexer_k_cache_buffers(1) + finally: + cache_manager.shutdown() + + def create_indexer(sparse_attn_config, layer_idx=0): """Helper to create an Indexer for testing.""" # Create RopeParams @@ -316,7 +423,9 @@ def __init__(self, head_dim): mla_params = MLAParams(sparse_params.index_head_dim) # Mock RotaryEmbedding since we're only testing cache management, not rope functionality - with patch("tensorrt_llm._torch.attention_backend.sparse.dsa.RotaryEmbedding") as mock_rope: + with patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.RotaryEmbedding" + ) as mock_rope: # Create a mock instance with a simple forward method mock_rope_instance = Mock() mock_rope_instance.forward = Mock(side_effect=lambda pos_ids, tensors: tensors) @@ -3282,6 +3391,9 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): ) Indexer.prepare(metadata) + if not enable_indexer_skip: + assert metadata.slot_mapping_fp8_fullkv is not metadata.slot_mapping_fp8 + assert metadata.slot_mapping_scale_fullkv is not metadata.slot_mapping_scale indexer._update_k_cache(k_fp8, k_scale, metadata) # Test custom kernel @@ -3394,6 +3506,7 @@ class TestPrepareRestoreAttnMetadataForDraftReplay: def _make_mock_metadata(): """Create a mock attention metadata object with KV cache block offsets.""" meta = Mock() + meta.enable_flash_mla = False meta.kv_cache_manager = Mock(name="target_kv_cache_manager") meta.kv_cache_block_offsets = torch.tensor([10, 20, 30]) meta.host_kv_cache_block_offsets = torch.tensor([10, 20, 30]) @@ -3403,29 +3516,29 @@ def _make_mock_metadata(): @staticmethod def _make_mock_draft_manager(): """Create a mock draft KV cache manager with host block offsets.""" - mgr = Mock(name="draft_kv_cache_manager") - mgr.host_kv_cache_block_offsets = torch.tensor([100, 200, 300]) - return mgr + return SimpleNamespace(host_kv_cache_block_offsets=torch.tensor([100, 200, 300])) def test_prepare_swaps_and_restore_recovers(self): """Test that prepare swaps KV manager and restore recovers original state.""" - from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata - meta = self._make_mock_metadata() mgr = self._make_mock_draft_manager() + # DeepSeek-V4 metadata inherits DSA metadata, while its manager exposes + # index_head_dim without implementing DSA's full page-mapping contract. + mgr.index_head_dim = 128 + mgr.get_pool_block_indices = Mock() + mgr.indexer_k_cache_page_scale = 1 original_kv_mgr = meta.kv_cache_manager original_offsets = meta.kv_cache_block_offsets.clone() original_host_offsets = meta.host_kv_cache_block_offsets.clone() + def is_attn_metadata(obj, cls): + if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata): + return obj is meta + return builtins.isinstance(obj, cls) + with patch( "tensorrt_llm._torch.speculative.interface.isinstance", - side_effect=lambda obj, cls: ( - obj is meta - if cls is TrtllmAttentionMetadata - else False - if cls.__name__ == "DSAtrtllmAttentionMetadata" - else builtins.isinstance(obj, cls) - ), + side_effect=is_attn_metadata, ): saved = prepare_attn_metadata_for_draft_replay(meta, mgr) @@ -3440,6 +3553,57 @@ def test_prepare_swaps_and_restore_recovers(self): torch.testing.assert_close(meta.kv_cache_block_offsets, original_offsets) torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets) + def test_native_dsa_replay_swaps_and_restores_buffers(self): + """Switch native DSA metadata to draft buffers and restore it.""" + meta = self._make_mock_metadata() + mgr = object.__new__(DSACacheManagerV2) + mgr.host_kv_cache_block_offsets = torch.tensor([100, 200, 300]) + meta.enable_context_mla_with_cached_kv = False + meta._invalidate_pool_view_cache = Mock() + buffer_pairs = { + "host_indexer_k_cache_block_offsets": "host_draft_indexer_k_cache_block_offsets", + "indexer_k_cache_block_offsets": "draft_indexer_k_cache_block_offsets", + "host_slot_mapping_fp8": "host_draft_slot_mapping_fp8", + "host_slot_mapping_scale": "host_draft_slot_mapping_scale", + "slot_mapping_fp8": "draft_slot_mapping_fp8", + "slot_mapping_scale": "draft_slot_mapping_scale", + "block_table": "draft_block_table", + "block_table_expanded": "draft_block_table_expanded", + "host_block_table_expanded": "host_draft_block_table_expanded", + } + target_buffers = {} + draft_buffers = {} + for target_name, draft_name in buffer_pairs.items(): + target_buffers[target_name] = object() + draft_buffers[target_name] = object() + setattr(meta, target_name, target_buffers[target_name]) + setattr(meta, draft_name, draft_buffers[target_name]) + del meta.slot_mapping_fp8_fullkv + del meta.slot_mapping_scale_fullkv + + def is_attn_metadata(obj, cls): + if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata): + return obj is meta + return builtins.isinstance(obj, cls) + + with ( + patch( + "tensorrt_llm._torch.speculative.interface.isinstance", + side_effect=is_attn_metadata, + ), + patch( + "tensorrt_llm._torch.speculative.interface.torch.cuda.is_current_stream_capturing", + return_value=True, + ), + ): + saved = prepare_attn_metadata_for_draft_replay(meta, mgr) + + assert "saved_dsa_state" in saved + assert meta.slot_mapping_fp8 is draft_buffers["slot_mapping_fp8"] + restore_attn_metadata_after_draft_replay(meta, saved) + assert meta.slot_mapping_fp8 is target_buffers["slot_mapping_fp8"] + assert meta.slot_mapping_fp8_fullkv is target_buffers["slot_mapping_fp8"] + @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_blackwell diff --git a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py index 09f16a5e7666..5bc02ede7379 100644 --- a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py +++ b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py @@ -15,6 +15,7 @@ """Unit tests for DSA C++ custom ops: - ``torch.ops.trtllm.indexer_k_cache_gather_op`` +- ``torch.ops.trtllm.indexer_k_cache_scatter_op`` - ``torch.ops.trtllm.convert_req_index_to_global`` - ``torch.ops.trtllm.fused_cat_fp4`` - ``torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell`` diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 5c75c9186bd7..7ce5d626d944 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -754,7 +754,7 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: patch.object(creator, "_enable_kv_cache_stats", return_value=False), patch.object( creator, - "_fallback_if_unsupported_kv_cache_manager_v2", + "_validate_or_fallback_kv_cache_manager_v2", return_value=KVCacheManagerV2, ), patch( diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py index 073190eacf71..069c47841287 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py @@ -12,6 +12,7 @@ import gc import unittest +from unittest.mock import Mock import torch @@ -20,7 +21,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role 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 BufferConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode DataType = tensorrt_llm.bindings.DataType CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -156,6 +157,33 @@ def test_default_hook_keeps_nvfp4_scale_buffers(self): mgr.shutdown() del mgr + def test_page_table_uses_physical_pool_representative(self): + mgr = KVCacheManagerV2(**_make_kwargs(head_dim=[64, 192, 64, 192])) + 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 + + 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 + def test_subclass_registers_index_key_only_on_sparse_layers(self): # Sparse layer convention: layers 0-2 dense (no INDEX_KEY), 3+ sparse # (INDEX_KEY registered). Matches MiniMax-M3. diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 81db2977398d..3b519541fbf1 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -65,7 +65,6 @@ from tensorrt_llm.llmapi.llm_utils import ( _resolve_kv_cache_manager_v2_auto, _resolve_transceiver_runtime_auto, - apply_model_defaults_to_llm_args, ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( @@ -515,7 +514,7 @@ def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( kv_cache_config=kv_cache_config, ) - assert _resolve_kv_cache_manager_v2_auto(llm_args, {}) is False + assert _resolve_kv_cache_manager_v2_auto(llm_args) is False assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False with pytest.raises(ValueError, match="use_kv_cache_manager_v2=True"): get_kv_cache_manager_cls( @@ -675,7 +674,7 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2) ) -def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): +def test_hybrid_models_prefer_v2_and_python_transceiver(monkeypatch): from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM from tensorrt_llm._torch.models.modeling_qwen3_5 import Qwen3_5VLModel from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM @@ -693,12 +692,9 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): model="/tmp/dummy_model", cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), ) - model_defaults = model_cls.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, model_defaults) _resolve_transceiver_runtime_auto(llm_args, model_cls) - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto") + _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True - assert llm_args.kv_cache_config.enable_block_reuse is False assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" @@ -712,7 +708,7 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): (None, False, False), ], ) -def test_qwen3_gdn_replay_defaults_to_v2_cache_manager( +def test_qwen3_gdn_replay_uses_v2_preference( monkeypatch, replay_env, manager_setting, @@ -732,12 +728,9 @@ def test_qwen3_gdn_replay_defaults_to_v2_cache_manager( ), speculative_config=MTPDecodingConfig(max_draft_len=3), ) - model_defaults = Qwen3NextForCausalLM.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, model_defaults) _resolve_kv_cache_manager_v2_auto( llm_args, - model_defaults, - original_setting=manager_setting, + Qwen3NextForCausalLM, ) assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2 @@ -751,19 +744,23 @@ def test_qwen3_gdn_replay_defaults_to_v2_cache_manager( ) -def test_kimi_defaults_to_mixed_manager(monkeypatch: pytest.MonkeyPatch) -> None: - """Kimi K3 declares no V2 default: block reuse defaults off and the - Mixed manager (separate KV / recurrent-state pools) is the default - route, which SA speculative decoding requires.""" +def test_kimi_without_v2_preference_uses_mixed_manager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Kimi K3 uses separate KV and recurrent-state pools for SA decoding.""" from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - llm_args = TorchLlmArgs(model="/tmp/dummy_model") - model_defaults = KimiLinearForCausalLM.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, model_defaults) - resolved = _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto") + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + tokens_per_block=64, + ), + ) + resolved = _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM) assert resolved is False assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False @@ -841,7 +838,7 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( creator._max_beam_width = max_beam_width with pytest.raises(NotImplementedError, match=expected): - creator._fallback_if_unsupported_kv_cache_manager_v2( + creator._validate_or_fallback_kv_cache_manager_v2( MambaHybridCacheManagerV2, model_config, KvCacheConfig() ) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 1aba431efa05..6e80d1513ef8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -150,19 +150,16 @@ def test_deepseek_v4_fused_hc_default_enabled(monkeypatch): assert _resolve_enable_fused_hc(config) is False -def test_deepseek_v4_model_defaults(): - class LlmArgs: - pass - - defaults = DeepseekV4ForCausalLM.get_model_defaults(LlmArgs()) +def test_deepseek_v4_kv_cache_defaults_and_v2_preference(): + defaults = DeepseekV4ForCausalLM.get_model_defaults(None) assert defaults == { "kv_cache_config": { "tokens_per_block": 128, - "use_kv_cache_manager_v2": True, "enable_swa_scratch_reuse": True, } } + assert DeepseekV4ForCausalLM.get_preferred_kv_cache_manager_version() == "V2" def test_deepseek_v4_weight_remap_for_mxfp4_routed_experts(): diff --git a/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py b/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py index 00e3812ed0fb..9d4adc08987b 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py +++ b/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import json import os import shutil @@ -21,8 +24,7 @@ from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig) from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, - apply_model_defaults_to_llm_args) +from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto from tensorrt_llm.mapping import Mapping configs = """ @@ -56,31 +58,26 @@ def test_gpt_oss_prefers_python_transceiver() -> None: def _resolve_gpt_oss_kv_cache_manager_v2(**llm_args_kwargs) -> bool: - """Run GPT-OSS model defaults through the same path model loading uses.""" + """Resolve the GPT-OSS preference through the model-loading path.""" llm_args = TorchLlmArgs(model="/tmp/dummy_model", **llm_args_kwargs) - original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 - model_defaults = GptOssForCausalLM.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, model_defaults) - return _resolve_kv_cache_manager_v2_auto(llm_args, - model_defaults, - original_setting=original_setting) + return _resolve_kv_cache_manager_v2_auto(llm_args, GptOssForCausalLM) -def test_gpt_oss_model_defaults_select_v2(): +def test_gpt_oss_model_preference_selects_v2(): """GPT-OSS is VSWA, so "auto" resolves to KVCacheManagerV2.""" assert _resolve_gpt_oss_kv_cache_manager_v2() is True @pytest.mark.parametrize("user_setting", [False, True]) def test_gpt_oss_explicit_setting_wins(user_setting): - """An explicit user value is never overridden by the model default.""" + """An explicit user value is never overridden by the model preference.""" assert _resolve_gpt_oss_kv_cache_manager_v2(kv_cache_config=KvCacheConfig( use_kv_cache_manager_v2=user_setting)) is user_setting def test_gpt_oss_two_model_eagle3_falls_back_to_v1(): """Two-model Eagle3 builds a separate draft engine with its own KV cache - manager, and V2 sizes both from the full budget, so the model default is + manager, and V2 sizes both from the full budget, so the model preference is demoted to V1.""" assert _resolve_gpt_oss_kv_cache_manager_v2( speculative_config=Eagle3DecodingConfig( @@ -90,7 +87,7 @@ def test_gpt_oss_two_model_eagle3_falls_back_to_v1(): def test_gpt_oss_explicit_v2_rejects_two_model_eagle3(): - """The demotion above only applies to the model default. An explicit + """The demotion above only applies to the model preference. An explicit request for the same unsupported combination is rejected rather than silently honored.""" with pytest.raises(ValueError, diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py index 9aa8e7e11a06..36c9d8b0c9f7 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py @@ -29,7 +29,10 @@ from tensorrt_llm.inputs import ContentFormat from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args +from tensorrt_llm.llmapi.llm_utils import ( + _resolve_kv_cache_manager_v2_auto, + apply_model_defaults_to_llm_args, +) from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo @@ -187,6 +190,7 @@ def test_qwen35_moe_model_defaults( llm_args.quant_config = QuantConfig(quant_algo=quant_algo) defaults = model_cls.get_model_defaults(llm_args) apply_model_defaults_to_llm_args(llm_args, defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) assert llm_args.kv_cache_config.enable_block_reuse is False assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index b678377b2e6d..57571e4db563 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -1,8 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os import sys import tempfile import unittest +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -15,6 +31,8 @@ skip_pre_blackwell) from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + DSACacheManagerV2, DSAtrtllmAttentionMetadata) from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor._util import \ @@ -32,6 +50,70 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +def test_mtp_eagle_refreshes_dsa_metadata_before_draft_forward() -> None: + """Refresh DSA mappings after switching to the draft cache.""" + events = [] + draft_manager = object.__new__(DSACacheManagerV2) + + class _Metadata(DSAtrtllmAttentionMetadata): + + def __init__(self): + self._seq_lens_cuda = torch.tensor([1], dtype=torch.int32) + self._num_ctx_tokens = 0 + + def set_in_mtp_draft_loop(self, active): + pass + + def set_mtp_num_accepted(self, value): + pass + + def set_skip_topk(self, value): + pass + + def on_update_kv_lens(self): + events.append("refresh") + + metadata = _Metadata() + + @contextmanager + def draft_context(attn_metadata, manager): + assert attn_metadata is metadata + assert manager is draft_manager + events.append("switch") + yield metadata + + class _StopAfterFirstDraftForward(Exception): + pass + + def run_draft_forward(*args, **kwargs): + events.append("forward") + raise _StopAfterFirstDraftForward + + worker = SimpleNamespace( + is_mtp_eagle=True, + draft_kv_cache_context=draft_context, + _run_draft_forward=run_draft_forward, + ) + + from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + + with pytest.raises(_StopAfterFirstDraftForward): + Eagle3OneModelWorker._forward_linear_draft_loop( + worker, + {"position_ids": torch.tensor([0], dtype=torch.int32)}, + metadata, + SimpleNamespace(runtime_draft_len=1), + draft_model=object(), + draft_kv_cache_manager=draft_manager, + num_contexts=0, + batch_size=1, + num_accepted_tokens=torch.tensor([1], dtype=torch.int32), + original_all_rank_num_tokens=None, + ) + + assert events == ["switch", "refresh", "forward"] + + def test_dynamic_tree_metadata_forces_target_mask_prepare_each_step() -> None: metadata = TrtllmAttentionMetadata( seq_lens=None, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2fa8772fd073..ceef1c8de786 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -522,84 +522,6 @@ def test_compute_applied_llm_defaults_simple_field(self): applied = apply_model_defaults_to_llm_args(llm_args, model_defaults) assert applied == model_defaults - @pytest.mark.parametrize("explicit_auto", [False, True]) - def test_kv_cache_manager_v2_auto_uses_model_default(self, explicit_auto): - kv_cache_config = (KvCacheConfig(use_kv_cache_manager_v2="auto") - if explicit_auto else KvCacheConfig()) - llm_args = TorchLlmArgs(model="/tmp/dummy_model", - kv_cache_config=kv_cache_config) - model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} - - apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, - model_defaults, - original_setting="auto") - - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True - - def test_kv_cache_manager_v2_auto_falls_back_to_false(self): - llm_args = TorchLlmArgs(model="/tmp/dummy_model") - - _resolve_kv_cache_manager_v2_auto(llm_args, {}) - - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False - - @pytest.mark.parametrize( - ("backend", "runtime"), - [ - ("NIXL", "CPP"), - ("UCX", None), - ("MPI", None), - ], - ) - def test_kv_cache_manager_v2_auto_falls_back_for_incompatible_disagg( - self, backend, runtime): - llm_args = TorchLlmArgs( - model="/tmp/dummy_model", - cache_transceiver_config=CacheTransceiverConfig( - backend=backend, transceiver_runtime=runtime), - ) - model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} - - apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, - model_defaults, - original_setting="auto") - - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False - - def test_kv_cache_manager_v2_auto_keeps_python_nixl_model_default(self): - llm_args = TorchLlmArgs( - model="/tmp/dummy_model", - cache_transceiver_config=CacheTransceiverConfig( - backend="NIXL", transceiver_runtime="PYTHON"), - ) - model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} - - apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, - model_defaults, - original_setting="auto") - - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True - - @pytest.mark.parametrize("user_setting", [False, True]) - def test_kv_cache_manager_v2_explicit_value_overrides_model_default( - self, user_setting): - llm_args = TorchLlmArgs( - model="/tmp/dummy_model", - kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting)) - model_defaults = { - "kv_cache_config": { - "use_kv_cache_manager_v2": not user_setting - } - } - - apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) - - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting - @pytest.mark.parametrize( "defaults_dict,should_raise,error_contains", [ @@ -697,6 +619,111 @@ def get_model_defaults(cls, llm_args): assert "enable_block_reuse" in error_str or "max_tokens" in error_str +@pytest.mark.cpu_only +class TestKvCacheManagerV2AutoResolution: + """Test model preferences for the KV cache manager version.""" + + class _PreferV2: + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): + return "V2" + + class _PreferV1: + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): + return "V1" + + @pytest.mark.parametrize("explicit_auto", [False, True]) + def test_auto_uses_model_preference(self, explicit_auto): + kv_cache_config = (KvCacheConfig(use_kv_cache_manager_v2="auto") + if explicit_auto else KvCacheConfig()) + llm_args = TorchLlmArgs(model="/tmp/dummy_model", + kv_cache_config=kv_cache_config) + + _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2) + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + + def test_auto_without_preference_falls_back_to_v1(self): + llm_args = TorchLlmArgs(model="/tmp/dummy_model") + + _resolve_kv_cache_manager_v2_auto(llm_args) + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + + @pytest.mark.parametrize( + ("backend", "runtime"), + [ + ("NIXL", "CPP"), + ("UCX", None), + ("MPI", None), + ], + ) + def test_auto_v2_falls_back_for_incompatible_disagg(self, backend, runtime): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime), + ) + + _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2) + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + + def test_auto_v2_keeps_python_nixl_preference(self): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON"), + ) + + _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2) + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + + @pytest.mark.parametrize(("user_setting", "model_cls"), [ + pytest.param(False, _PreferV2, id="explicit-v1"), + pytest.param(True, _PreferV1, id="explicit-v2"), + ]) + def test_explicit_value_overrides_model_preference(self, user_setting, + model_cls): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting)) + + _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting + + def test_registered_models_prefer_v2(self): + from tensorrt_llm._torch.models.modeling_utils import \ + get_registered_model_class + + architectures = ( + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + "GlmMoeDsaForCausalLM", + "GptOssForCausalLM", + "MistralLarge3ForCausalLM", + "DeepseekV4ForCausalLM", + "KimiK25ForConditionalGeneration", + "MiniMaxM2ForCausalLM", + "NemotronHForCausalLM", + "NemotronHPuzzleForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3_5MoeForCausalLM", + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5ForConditionalGeneration", + ) + for architecture in architectures: + model_cls = get_registered_model_class(architecture) + assert model_cls is not None + assert model_cls.get_preferred_kv_cache_manager_version() == "V2" + + @pytest.mark.cpu_only def test_KvCacheConfig_declaration(): assert KvCacheConfig().mamba_state_cache_interval is None @@ -3572,12 +3599,16 @@ def test_rejects_non_bool(self): class _PreferPythonTransceiverModel: - """Fake model class opting into the Python transceiver.""" + """Fake model class preferring Python transceiver and KV manager V2.""" @classmethod def get_model_defaults(cls, llm_args): return {} + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): + return "V2" + @classmethod def get_preferred_transceiver_runtime(cls, pretrained_config=None): return "PYTHON" @@ -3694,6 +3725,11 @@ def test_hybrid_snapshot_policy_controls_block_reuse( assert "no Mamba state snapshot policy" in warnings[0] def test_hybrid_fixed_snapshot_rejects_auto_resolved_v1(self, monkeypatch): + from tensorrt_llm._torch.models.modeling_utils import \ + MODEL_CLASS_MAPPING + + monkeypatch.setitem(MODEL_CLASS_MAPPING, "Qwen3NextForCausalLM", + _NoModelDefaults) args = TorchLlmArgs( model="/tmp/dummy_model", kv_cache_config=KvCacheConfig( @@ -3793,13 +3829,9 @@ def test_default_backend_env_override_falls_back_to_v1_cpp( monkeypatch.delenv(env_var, raising=False) monkeypatch.setenv(backend_env, "1") args = self._disagg_args(backend="DEFAULT") - model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} - apply_model_defaults_to_llm_args(args, model_defaults) _resolve_transceiver_runtime_auto(args, _PreferPythonTransceiverModel) - _resolve_kv_cache_manager_v2_auto(args, - model_defaults, - original_setting="auto") + _resolve_kv_cache_manager_v2_auto(args, _PreferPythonTransceiverModel) assert args.cache_transceiver_config.transceiver_runtime is None assert args.kv_cache_config.use_kv_cache_manager_v2 is False @@ -3880,6 +3912,7 @@ def test_model_loader_resolves_auto_without_checkpoint_loader(self): ModelLoader.load_config_and_apply_defaults("/tmp/dummy_model", args, None) assert args.cache_transceiver_config.transceiver_runtime is None + assert args.kv_cache_config.use_kv_cache_manager_v2 is False @staticmethod def _fake_checkpoint_loader(architectures): @@ -3904,6 +3937,7 @@ def test_model_loader_full_chain_adopts_model_preference(self, monkeypatch): model_loader_mod.ModelLoader.load_config_and_apply_defaults( "/tmp/dummy_model", args, fake_loader) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + assert args.kv_cache_config.use_kv_cache_manager_v2 is True @pytest.mark.parametrize("arch,expected", [ ("ModelAForCausalLM", "PYTHON"), @@ -3963,6 +3997,7 @@ def get_preferred_transceiver_runtime(cls, pretrained_config=None): model_loader_mod.ModelLoader.load_config_and_apply_defaults( "/tmp/dummy_model", args, fake_loader) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + assert args.kv_cache_config.use_kv_cache_manager_v2 is True def test_model_loader_full_chain_aggregated_stays_none(self, monkeypatch): """Full chain in aggregated mode: no transceiver config appears.""" @@ -4004,8 +4039,8 @@ def test_resolve_default_backend_env_priority(self, monkeypatch): backend="UCX")._resolve_default_backend() == ("UCX", None) -class TestDeepseekTransceiverPreference: - """Per-architecture preferred KV-cache transceiver runtime. +class TestDeepseekRuntimePreferences: + """DeepSeek KV-cache manager and transceiver preferences. DeepseekV3ForCausalLM and DeepseekV32ForCausalLM prefer the Python KV-cache transceiver, while GlmMoeDsaForCausalLM (GLM 5.2) requires the C++ transceiver diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 99dc35706f03..252493d1dba2 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -280,7 +280,7 @@ def test_addr_file_owner_only(self, tmp_path): def test_use_kv_cache_manager_v2_flags(): # Absent -> "auto" (the driver resolves it against the model's - # get_model_defaults at runtime, like serving). + # get_preferred_kv_cache_manager_version at runtime, like serving). plan = pcfg.resolve_plan(_disagg_yaml()) assert plan["ctx_use_kv_cache_manager_v2"] == "auto" assert plan["gen_use_kv_cache_manager_v2"] == "auto" diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py index ed060d1e9053..d6b736a4d4bd 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_run.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py @@ -389,14 +389,34 @@ def test_kv_cache_manager_ctor_kwargs(self, api, manager_attr): def test_serving_resolvers(self, api): import inspect - # The driver calls resolve_kv_cache_manager_v2_auto(shim, defaults): - # the first two params are fixed, anything added later must default. + # The driver calls the serving resolver with a shim, model class, and + # pretrained config. Only the shim is required. v2 = inspect.signature(api.resolve_kv_cache_manager_v2_auto).parameters - assert list(v2)[:2] == ["llm_args", "model_defaults_dict"] - assert all(p.default is not inspect.Parameter.empty for p in list(v2.values())[2:]) + assert list(v2)[:3] == ["llm_args", "model_cls", "pretrained_config"] + assert all(p.default is not inspect.Parameter.empty for p in list(v2.values())[1:]) rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3 + def test_model_preference_resolver_shim_supports_v2(self, api, monkeypatch): + class _PreferV2: + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): + return "V2" + + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + monkeypatch.setattr( + rp, + "_lookup_model_cls", + lambda model_dir: (_PreferV2, types.SimpleNamespace()), + ) + cache_cfg = api.CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + + assert rp.resolve_model_prefs( + "/tmp/dummy_model", + {"use_kv_cache_manager_v2": "auto"}, + cache_cfg, + ) + def test_enum_members(self, api): for enum, members in ( (api.DataType, ("FP8", "HALF", "BF16")),