From 68da08ec992a2be6000f112a551316495d52ebd5 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 27 May 2025 22:29:37 +0000 Subject: [PATCH 01/15] Refactor and pipeclean, can pass gemma3 test. Need to run CI to find out whether there are more issues. Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 21 ++ tensorrt_llm/_torch/pyexecutor/_util.py | 26 +- .../_torch/pyexecutor/model_engine.py | 9 + .../_torch/pyexecutor/resource_manager.py | 253 +++++++++++++++--- .../defs/accuracy/test_llm_api_pytorch.py | 9 + 5 files changed, 282 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 97a816c3d478..53b17974e84c 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -8,6 +8,7 @@ from tensorrt_llm import logger from tensorrt_llm._utils import torch_dtype_to_binding +from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.functional import AllReduceStrategy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -323,6 +324,11 @@ def get_bindings_model_config(self) -> "ModelConfigCpp": model_config_cpp.mlp_hidden_size = mlp_hidden_size model_config_cpp.size_per_head = head_size + # NOTE: this method is not robust, for Gemma3ForCausalLM only + layer_types = self.get_layer_types() + if layer_types is not None: + model_config_cpp.layer_types = layer_types + return model_config_cpp def _infer_nemotron_ffn_mult(self): @@ -339,3 +345,18 @@ def _infer_nemotron_ffn_mult(self): biggest_ffn_mult, self.pretrained_config.hidden_size) return mlp_hidden_size + + def get_layer_types(self) -> Optional[List[LayerTypeCpp]]: + """ + This method is a hack to support the effort to switch to KvCacheManagerCpp. + Currently, it is only tested for Gemma3ForCausalLM. For other models, it will return None. + """ + if self.pretrained_config.architectures[0] in ["Gemma3ForCausalLM"]: + logger.debug( + f"Setting layer types for {self.pretrained_config.architectures}" + ) + return [ + LayerTypeCpp.ATTENTION, + ] * self.pretrained_config.num_hidden_layers + else: + return None diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 80213a85b7de..294e88c8a43e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -180,6 +180,20 @@ def try_prepare_estimation(self) -> bool: estimating_kv_cache = True self._executor_config.kv_cache_config.max_tokens = self._get_token_num_for_estimation( ) + maw = self._executor_config.kv_cache_config.max_attention_window + if maw is not None: + unique_windows = list(set(maw)) + if len(unique_windows) > 1: + # NOTE: For sliding window, we need to allocate more tokens(max_num_tokens + window_size) for each window, need to update if this requirement changes + assert len( + unique_windows + ) == 2, "Sliding window with more than 2 window sizes has not been tested" + # sliding window is the one that is not the executor_config.max_seq_len + sliding_window = unique_windows[0] if unique_windows[ + 0] != self._executor_config.max_seq_len else unique_windows[ + 1] + self._executor_config.kv_cache_config.max_tokens += ( + self._executor_config.max_num_tokens + sliding_window) return estimating_kv_cache def estimate_max_tokens(self, py_executor: PyExecutor) -> None: @@ -340,8 +354,15 @@ def _create_kv_cache_manager( spec_config=spec_config, ) else: + # NOTE: this is a workaround for VSWA to switch to calculate_max_num_blocks_from_cpp in KVCahceManager + is_vswa = executor_config.kv_cache_config.max_attention_window is not None and len( + set(executor_config.kv_cache_config.max_attention_window)) > 1 + binding_model_config = model_engine.model.model_config.get_bindings_model_config( + ) if is_vswa else None + kv_cache_manager = KVCacheManager( - executor_config.kv_cache_config, + # NOTE: from tensorrt_llm.bindings.executor.KvCacheConfig to tensorrt_llm.bindings.KvCacheConfig + KvCacheConfigCpp(executor_config.kv_cache_config), tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, num_layers=num_hidden_layers, num_kv_heads=num_key_value_heads, @@ -352,7 +373,8 @@ def _create_kv_cache_manager( mapping=mapping, dtype=kv_cache_dtype, spec_config=spec_config, - ) + max_num_tokens=executor_config.max_num_tokens, + model_config=binding_model_config) # KVCacheManager (Non-draft) modifies the max_seq_len field, update it to executor_config if model_engine.kv_cache_manager_key == ResourceManagerType.KV_CACHE_MANAGER: executor_config.max_seq_len = kv_cache_manager.max_seq_len diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c58b4ca266e8..aa63ee4deb48 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -588,6 +588,12 @@ def get_cuda_graph_warmup_request(batch_size): def get_torch_compile_warmup_request(batch_size, num_tokens_per_request): available_blocks = kv_cache_manager.get_num_free_blocks() + print( + f"====================get_torch_compile_warmup_request============================" + ) + print(f"batch_size: {batch_size}") + print(f"num_tokens_per_request: {num_tokens_per_request}") + print(f"available_blocks: {available_blocks}") if available_blocks >= batch_size * math.ceil( num_tokens_per_request / kv_cache_manager.tokens_per_block): # Should only need (at most) one more page per request. @@ -612,6 +618,9 @@ def get_torch_compile_warmup_request(batch_size, result.context_requests = requests else: result = None + print( + f"====================get_torch_compile_warmup_request result============================" + ) return result def get_autotune_warmup_request(): diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e58043dffb31..4b16279420c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1,7 +1,7 @@ import enum import math from abc import ABC, abstractmethod -from collections import OrderedDict +from collections import OrderedDict, defaultdict from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union import torch @@ -25,6 +25,8 @@ if TYPE_CHECKING: from ..speculative.interface import SpecConfig +ExecutorKvCacheConfig = tensorrt_llm.bindings.executor.KvCacheConfig +BufferManagerCpp = tensorrt_llm.bindings.internal.runtime.BufferManager KVCacheManagerCpp = tensorrt_llm.bindings.internal.batch_manager.KVCacheManager KvCacheConfigCpp = tensorrt_llm.bindings.executor.KvCacheConfig CacheTypeCpp = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -35,6 +37,7 @@ PeftCacheManagerCpp = tensorrt_llm.bindings.internal.batch_manager.PeftCacheManager PeftCacheConfig = tensorrt_llm.bindings.executor.PeftCacheConfig WorldConfig = tensorrt_llm.bindings.WorldConfig +TempAttentionWindowInputs = tensorrt_llm.bindings.internal.batch_manager.TempAttentionWindowInputs class ResourceManagerType(enum.Enum): @@ -124,7 +127,12 @@ def __init__( dtype: DataType = DataType.HALF, spec_config: Optional["SpecConfig"] = None, layer_mask: Optional[List[bool]] = None, + max_num_tokens: int = 8192, + model_config: Optional[ModelConfig] = None, ) -> None: + assert isinstance( + kv_cache_config, KvCacheConfigCpp + ), "kv_cache_config should be tensorrt_llm.bindings.KvCacheConfig" self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type @@ -149,7 +157,6 @@ def __init__( (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_local_layers) ] - else: assert len(num_kv_heads) == self.num_layers @@ -170,46 +177,79 @@ def __init__( self.max_batch_size = max_batch_size self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 # Some speculative decoding methods need to use different kv lengths for the - # draft/target layers. Add extra tokens to haddle this issue. + # draft/target layers. Add extra tokens to handle this issue. self.num_extra_kv_tokens = 0 if spec_config is None else spec_config.num_extra_kv_tokens self.event_buffer_max_size = kv_cache_config.event_buffer_max_size + # Determine max_attention_window_vec if kv_cache_config.max_attention_window is None: - max_attention_window = max_seq_len + # Use max_seq_len as default max_attention_window + self.max_attention_window_vec = [max_seq_len] else: - assert len( - kv_cache_config.max_attention_window - ) == 1, "Python KvCacheManager doesn't currently support variable window attention" - max_attention_window = kv_cache_config.max_attention_window[0] + self.max_attention_window_vec = kv_cache_config.max_attention_window.copy( + ) # Make a copy to avoid modifying original sink_token_length = (kv_cache_config.sink_token_length if kv_cache_config.sink_token_length is not None else 0) - self.blocks_in_primary_pool, self.blocks_in_secondary_pool = self.calculate_max_num_blocks( - kv_cache_config, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - mapping=mapping, - dtype=dtype, - kv_factor=self.kv_factor, - ) + # Determine if this is VSWA (Variable Sliding Window Attention) + is_vswa = len(self.max_attention_window_vec) > 1 - max_atten_window_upper_bound = self.get_max_atten_window_upper_bound( - blocks_in_primary_pool=self.blocks_in_primary_pool, + # Calculate blocks per window using appropriate method + if is_vswa: + # VSWA case: use C++ implementation for variable window sizes + if model_config is None: + raise ValueError( + "model_config is required for VSWA (Variable Sliding Window Attention)" + ) + blocks_per_window = self.calculate_max_num_blocks_from_cpp( + kv_cache_config=kv_cache_config, + model_config=model_config, + extra_cost_memory=0, + ) + else: + # Standard case: use original Python implementation + blocks_in_primary_pool, blocks_in_secondary_pool = self.calculate_max_num_blocks( + kv_cache_config=kv_cache_config, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + mapping=mapping, + dtype=dtype, + kv_factor=self.kv_factor, + ) + blocks_per_window = { + self.max_attention_window_vec[0]: + (blocks_in_primary_pool, blocks_in_secondary_pool) + } + + # Validate and adjust attention windows against their upper bounds if needed + blocks_per_window, self.max_seq_len, self.max_attention_window_vec = self._validate_and_adjust_attention_windows( + max_attention_window_vec=self.max_attention_window_vec, + blocks_per_window=blocks_per_window, tokens_per_block=tokens_per_block, - max_beam_width=1, - sink_token_len=sink_token_length, - max_seq_len=max_seq_len) + sink_token_length=sink_token_length, + max_seq_len=self.max_seq_len, + ) - if max_attention_window > max_atten_window_upper_bound: - logger.warning( - f"maxAttentionWindow and maxSequenceLen are too large for at least one sequence to fit in kvCache. They are reduced to {max_atten_window_upper_bound}" + if kv_cache_type != CacheTypeCpp.SELF: + assert len( + blocks_per_window + ) == 1, "Only one window size is supported for non-self KV cache" + # rewrite the attention window size in blocks_per_window + memory_pools = blocks_per_window[self.max_attention_window_vec[0]] + blocks_per_window = {self.max_seq_len: memory_pools} + logger.info( + f"Adjusted attention window size to {self.max_seq_len} in blocks_per_window" ) - max_attention_window = max_atten_window_upper_bound - self.max_seq_len = max_atten_window_upper_bound - self.max_attention_window = max_attention_window if kv_cache_type == CacheTypeCpp.SELF else self.max_seq_len + # Set up temp_attention_window_inputs only for VSWA or when needed + temp_attention_window_inputs = None + if is_vswa: # Only create when necessary + temp_attention_window_inputs = TempAttentionWindowInputs() + temp_attention_window_inputs.paged_context_fmha = True + temp_attention_window_inputs.max_input_len = self.max_seq_len - 1 + temp_attention_window_inputs.max_num_tokens = max_num_tokens # Note that this stream is unused for now. Will be used for copying to host # when that feature is enabled. @@ -218,14 +258,11 @@ def __init__( 'num_kv_heads_per_layer': self.num_kv_heads_per_layer, 'size_per_head': head_dim, 'tokens_per_block': tokens_per_block, - 'blocks_per_window': { - self.max_attention_window: - (self.blocks_in_primary_pool, self.blocks_in_secondary_pool) - }, + 'blocks_per_window': blocks_per_window, 'max_num_sequences': max_batch_size, 'max_beam_width': 1, # TODO: more than 1 beam? - 'max_attention_window_vec': [self.max_attention_window], - 'temp_attention_window_inputs': None, + 'max_attention_window_vec': self.max_attention_window_vec, + 'temp_attention_window_inputs': temp_attention_window_inputs, 'dtype': dtype, 'sink_token_length': sink_token_length, 'stream': self._stream.cuda_stream, @@ -236,6 +273,7 @@ def __init__( 'enable_partial_reuse': kv_cache_config.enable_partial_reuse, 'copy_on_partial_reuse': kv_cache_config.copy_on_partial_reuse, } + if self.event_buffer_max_size > 0: kwargs['event_manager'] = KVCacheEventManagerCpp( max_kv_event_entries=self.event_buffer_max_size) @@ -422,6 +460,9 @@ def calculate_max_num_blocks(self, ) else: max_tokens = kv_cache_config.max_tokens + logger.info( + f"max_tokens is set by kv_cache_config.max_tokens: {max_tokens}" + ) if mapping.world_size > 1: # make sure all ranks use same value for maxTokens @@ -442,7 +483,7 @@ def get_max_atten_window_upper_bound(self, blocks_in_primary_pool, token_capacity = blocks_in_primary_pool * tokens_per_block max_blocks_per_seq = math.floor(token_capacity / (max_beam_width * tokens_per_block)) - assert max_blocks_per_seq > 0, "Impossibe to fit in any sequence in kvCache" + assert max_blocks_per_seq > 0, "Impossible to fit in any sequence in kvCache" max_token_num = max_blocks_per_seq * tokens_per_block sink_tokens_in_last_block = sink_token_len % tokens_per_block @@ -513,6 +554,150 @@ def get_kv_cache_stats(self): def rewind_kv_cache(self, request: LlmRequest, rewind_len: int): self.impl.rewind_kv_cache(request.py_request_id, rewind_len) + def _get_window_size_to_layers(self) -> dict[int, list[int]]: + """ + Get the window size to layers mapping. + The returned map has window sizes as keys and lists of layer indices as values. + + max_attention_window_vec is treated as a repeating pattern. + """ + window_size_to_layers_map = defaultdict(list) + + if not self.max_attention_window_vec: + # This case should ideally be prevented by earlier config validation. + # If num_local_layers is 0, an empty map is fine. + if self.num_local_layers > 0: + raise ValueError( + "max_attention_window_vec cannot be empty if there are local layers." + ) + return { + } # Return an empty dict if no local layers or if somehow vec is empty and no layers. + + # Treat max_attention_window_vec as a repeating pattern. + # This also correctly handles the case where len(self.max_attention_window_vec) == 1. + pattern_len = len(self.max_attention_window_vec) + for local_layer_idx in range(self.num_local_layers): + window_size = self.max_attention_window_vec[local_layer_idx % + pattern_len] + window_size_to_layers_map[window_size].append(local_layer_idx) + return window_size_to_layers_map + + def calculate_max_num_blocks_from_cpp( + self, + kv_cache_config: KvCacheConfigCpp, + model_config: ModelConfig, + extra_cost_memory: int = 0) -> dict[int, tuple[int, int]]: + """ + This function is a wrapper of KVCacheManagerCpp.calculate_max_num_blocks. + The final goal is to switch to the C++ implementation of calculate_max_num_blocks. + Currently, this function is added to support *ONLY* VSWA. + + Args: + kv_cache_config: The KV cache configuration object. + model_config: The model configuration object. + extra_cost_memory: Extra memory in bytes to exclude from available memory. + + Returns: + A dict of (max_attention_window, (blocks_in_primary_pool, blocks_in_secondary_pool)). + """ + + # Construct WorldConfig from self.mapping + world_config_cpp = WorldConfig( + tensor_parallelism=self.mapping.tp_size, + pipeline_parallelism=self.mapping.pp_size, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node) + + window_size_to_layers = self._get_window_size_to_layers() + logger.debug(f"window_size_to_layers: {window_size_to_layers}") + + free_mem, total_mem = torch.cuda.mem_get_info() + primary_pool_memory_bytes = free_mem + secondary_pool_memory_bytes = 0 + logger.debug( + f"primary_pool_memory_bytes is set to {primary_pool_memory_bytes/1024**3}GB, \nsecondary_pool_memory_bytes is set to {secondary_pool_memory_bytes/1024**3}GB" + ) + + blocks_per_window = KVCacheManagerCpp.calculate_max_num_blocks( + config=kv_cache_config, + is_cross_attention=False, #TODO: support cross attention + dtype=self.dtype, + model_config=model_config, + world_config=world_config_cpp, + window_size_to_layers=window_size_to_layers, + allotted_primary_mem_bytes=primary_pool_memory_bytes, + allotted_secondary_mem_bytes=secondary_pool_memory_bytes, + extra_cost_memory=extra_cost_memory, + kvFactor=self.kv_factor, + ) + return blocks_per_window + + def _validate_and_adjust_attention_windows( + self, + max_attention_window_vec: List[int], + blocks_per_window: Dict[int, Tuple[int, int]], + tokens_per_block: int, + sink_token_length: int, + max_seq_len: int, + ) -> Tuple[Dict[int, Tuple[int, int]], int, List[int]]: + """ + Validate and adjust attention windows against their upper bounds. + It will change the self.max_attention_window_vec inside the function. + + Args: + max_attention_window_vec: List of attention window sizes + blocks_per_window: Dict mapping window size to (primary_blocks, secondary_blocks) + tokens_per_block: Number of tokens per block + sink_token_length: Length of sink tokens + max_seq_len: Maximum sequence length + + Returns: + Tuple of (adjusted_blocks_per_window, adjusted_max_seq_len, adjusted_max_attention_window_vec) + """ + max_beam_width = 1 # TODO: support more than 1 beam? + window_adjustments = {} + # Validate each window size in blocks_per_window against its upper bound + for window_size, (blocks_in_primary_pool, + _) in blocks_per_window.items(): + upper_bound = self.get_max_atten_window_upper_bound( + blocks_in_primary_pool=blocks_in_primary_pool, + tokens_per_block=tokens_per_block, + max_beam_width=max_beam_width, + sink_token_len=sink_token_length, + max_seq_len=max_seq_len) + if window_size > upper_bound: + logger.warning( + f"Attention window size {window_size} exceeds upper bound {upper_bound} " + f"for available blocks. Reducing to {upper_bound}.") + window_adjustments[window_size] = upper_bound + # Apply adjustments to the window vector if any were needed + if window_adjustments: + adjusted_window_vec = [ + window_adjustments.get(window, window) + for window in max_attention_window_vec + ] + logger.warning( + f"Adjusted max_attention_window_vec to {adjusted_window_vec}") + # update the window size in blocks_per_window if it is adjusted + adjusted_blocks_per_window = {} + for window_size, memory_pools in blocks_per_window.items(): + if window_size in window_adjustments: + adjusted_window_size = window_adjustments[window_size] + adjusted_blocks_per_window[ + adjusted_window_size] = memory_pools + logger.warning( + f"Adjusted window size {window_size} to {adjusted_window_size} in blocks_per_window" + ) + else: + adjusted_blocks_per_window[window_size] = memory_pools + # Update max_seq_len to the maximum of adjusted windows + adjusted_max_seq_len = max(adjusted_window_vec) + logger.warning(f"Adjusted max_seq_len to {adjusted_max_seq_len}") + + return adjusted_blocks_per_window, adjusted_max_seq_len, adjusted_window_vec + else: + return blocks_per_window, max_seq_len, max_attention_window_vec + class MambaCacheManager(BaseResourceManager): diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 41c0014ec9cf..d731557b3d12 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -450,6 +450,15 @@ def test_auto_dtype(self): task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) + def test_auto_dtype_vswa(self): + kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + free_gpu_memory_fraction=None, + max_attention_window=[512, 512, 512, 512, 512, 32768]) + with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config) as llm: + task = CnnDailymail(self.MODEL_NAME) + task.evaluate(llm) + class TestMixtral8x7B(LlmapiAccuracyTestHarness): MODEL_NAME = "mistralai/Mixtral-8x7B-v0.1" From 31847a85193d403ab611cd49449e13b4dca18ec6 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Sat, 7 Jun 2025 21:56:15 +0000 Subject: [PATCH 02/15] fix: - Many tests directly pass executor.KvCacheConfig to KvCacheManager, apply a conversion if this happen - add `self.blocks_in_primary_pool` and `self.blocks_in_secondary_pool` as these are used for standard case. Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- .../_torch/pyexecutor/resource_manager.py | 29 ++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 4b16279420c1..ccdd66651a2c 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -130,9 +130,18 @@ def __init__( max_num_tokens: int = 8192, model_config: Optional[ModelConfig] = None, ) -> None: - assert isinstance( - kv_cache_config, KvCacheConfigCpp - ), "kv_cache_config should be tensorrt_llm.bindings.KvCacheConfig" + if not isinstance(kv_cache_config, KvCacheConfigCpp): + # NOTE: There is a difference between + # tensorrt_llm.bindings.KvCacheConfig(KvCacheConfigCpp) and + # tensorrt_llm.bindings.executor.KvCacheConfig + # TODO: Remove this conversion in the future. + assert isinstance(kv_cache_config, ExecutorKvCacheConfig), ( + "Only ExecutorKvCacheConfig can be converted to KvCacheConfigCpp" + ) + kv_cache_config = KvCacheConfigCpp(kv_cache_config) + logger.warning( + "kv_cache_config is not a tensorrt_llm.bindings.KvCacheConfig, " + "converting it to a tensorrt_llm.bindings.KvCacheConfig") self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type @@ -210,7 +219,7 @@ def __init__( ) else: # Standard case: use original Python implementation - blocks_in_primary_pool, blocks_in_secondary_pool = self.calculate_max_num_blocks( + self.blocks_in_primary_pool, self.blocks_in_secondary_pool = self.calculate_max_num_blocks( kv_cache_config=kv_cache_config, head_dim=head_dim, tokens_per_block=tokens_per_block, @@ -220,7 +229,7 @@ def __init__( ) blocks_per_window = { self.max_attention_window_vec[0]: - (blocks_in_primary_pool, blocks_in_secondary_pool) + (self.blocks_in_primary_pool, self.blocks_in_secondary_pool) } # Validate and adjust attention windows against their upper bounds if needed @@ -495,8 +504,11 @@ def get_max_atten_window_upper_bound(self, blocks_in_primary_pool, return max_atten_window_upper_bound def get_cache_indices(self, request: LlmRequest) -> List[int]: + assert len(self.max_attention_window_vec + ) == 1, "Only support one attention window for now" + max_attention_window = max(self.max_attention_window_vec) result = self.impl.get_cache_block_ids(request.py_request_id, - self.max_attention_window) + max_attention_window) assert len(result) == 1 return result[0] @@ -504,8 +516,11 @@ def get_batch_cache_indices( self, request_ids: List[int], ) -> Dict[int, List[int]]: + assert len(self.max_attention_window_vec + ) == 1, "Only support one attention window for now" + max_attention_window = max(self.max_attention_window_vec) result = self.impl.get_batch_cache_block_ids(request_ids, - self.max_attention_window) + max_attention_window) for i in range(len(result)): assert (len(result[i])) == 1 result[i] = result[i][0] From 3401c61188f62a4973e3b60483a5fc639de55f12 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Mon, 9 Jun 2025 16:48:30 +0000 Subject: [PATCH 03/15] fix: correct the test that passes the kvcache config with the wrong type (the mirror class of executor.kvcache_config) Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tests/unittest/llmapi/test_llm_kv_cache_events.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index c21d12bd3af1..eeffe2d4ba68 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -33,7 +33,7 @@ def create_kv_cache_manager(): max_batch_size = 1 mapping = Mapping() return KVCacheManager( - kv_cache_config=global_kvcache_config, + kv_cache_config=global_kvcache_config._to_pybind(), kv_cache_type=tensorrt_llm.bindings.internal.batch_manager.CacheType. SELF, num_layers=num_layers, From 1ced7c1193164e4cc28a8e7b7587605692fc4e8d Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Wed, 11 Jun 2025 18:34:58 +0000 Subject: [PATCH 04/15] fix: - relax the type restirction for kv_cache_config and only require as needed. - update get_cache_indices and get_batch_cache_indices method accordingly - revert the change of test Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- .../_torch/pyexecutor/resource_manager.py | 54 +++++++++++-------- .../llmapi/test_llm_kv_cache_events.py | 2 +- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index ccdd66651a2c..af5c233cff32 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -130,18 +130,6 @@ def __init__( max_num_tokens: int = 8192, model_config: Optional[ModelConfig] = None, ) -> None: - if not isinstance(kv_cache_config, KvCacheConfigCpp): - # NOTE: There is a difference between - # tensorrt_llm.bindings.KvCacheConfig(KvCacheConfigCpp) and - # tensorrt_llm.bindings.executor.KvCacheConfig - # TODO: Remove this conversion in the future. - assert isinstance(kv_cache_config, ExecutorKvCacheConfig), ( - "Only ExecutorKvCacheConfig can be converted to KvCacheConfigCpp" - ) - kv_cache_config = KvCacheConfigCpp(kv_cache_config) - logger.warning( - "kv_cache_config is not a tensorrt_llm.bindings.KvCacheConfig, " - "converting it to a tensorrt_llm.bindings.KvCacheConfig") self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type @@ -208,10 +196,24 @@ def __init__( # Calculate blocks per window using appropriate method if is_vswa: # VSWA case: use C++ implementation for variable window sizes + # model config check if model_config is None: raise ValueError( "model_config is required for VSWA (Variable Sliding Window Attention)" ) + # kv cache config check + if not isinstance(kv_cache_config, KvCacheConfigCpp): + # NOTE: There is a difference between + # tensorrt_llm.bindings.KvCacheConfig(KvCacheConfigCpp) and + # tensorrt_llm.bindings.executor.KvCacheConfig + # calculate_max_num_blocks_from_cpp only accepts KvCacheConfigCpp + assert isinstance(kv_cache_config, ExecutorKvCacheConfig), ( + "Only ExecutorKvCacheConfig can be converted to KvCacheConfigCpp" + ) + kv_cache_config = KvCacheConfigCpp(kv_cache_config) + logger.warning( + "kv_cache_config is not a tensorrt_llm.bindings.KvCacheConfig, " + "converting it to a tensorrt_llm.bindings.KvCacheConfig") blocks_per_window = self.calculate_max_num_blocks_from_cpp( kv_cache_config=kv_cache_config, model_config=model_config, @@ -503,24 +505,30 @@ def get_max_atten_window_upper_bound(self, blocks_in_primary_pool, assert max_atten_window_upper_bound > 0, "Impossibe to fit in any sequence in kvCache" return max_atten_window_upper_bound - def get_cache_indices(self, request: LlmRequest) -> List[int]: - assert len(self.max_attention_window_vec - ) == 1, "Only support one attention window for now" - max_attention_window = max(self.max_attention_window_vec) + def get_cache_indices(self, + request: LlmRequest, + window_size: Optional[int] = None) -> List[int]: + if window_size is None: + if len(self.max_attention_window_vec) > 1: + raise ValueError("window_size must be provided for VSWA") + window_size = self.max_attention_window_vec[0] + result = self.impl.get_cache_block_ids(request.py_request_id, - max_attention_window) + window_size) assert len(result) == 1 return result[0] def get_batch_cache_indices( self, request_ids: List[int], - ) -> Dict[int, List[int]]: - assert len(self.max_attention_window_vec - ) == 1, "Only support one attention window for now" - max_attention_window = max(self.max_attention_window_vec) - result = self.impl.get_batch_cache_block_ids(request_ids, - max_attention_window) + window_size: Optional[int] = None, + ) -> List[List[int]]: + if window_size is None: + if len(self.max_attention_window_vec) > 1: + raise ValueError("window_size must be provided for VSWA") + window_size = self.max_attention_window_vec[0] + + result = self.impl.get_batch_cache_block_ids(request_ids, window_size) for i in range(len(result)): assert (len(result[i])) == 1 result[i] = result[i][0] diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index eeffe2d4ba68..c21d12bd3af1 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -33,7 +33,7 @@ def create_kv_cache_manager(): max_batch_size = 1 mapping = Mapping() return KVCacheManager( - kv_cache_config=global_kvcache_config._to_pybind(), + kv_cache_config=global_kvcache_config, kv_cache_type=tensorrt_llm.bindings.internal.batch_manager.CacheType. SELF, num_layers=num_layers, From 2010813958f6b4ce1dbb389fee0b35ad311db187 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Fri, 13 Jun 2025 00:40:08 +0000 Subject: [PATCH 05/15] fix: - update comments - clean up logging Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 10 +--------- tensorrt_llm/_torch/pyexecutor/resource_manager.py | 4 ++-- 2 files changed, 3 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index aa63ee4deb48..fa07065566a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -588,12 +588,6 @@ def get_cuda_graph_warmup_request(batch_size): def get_torch_compile_warmup_request(batch_size, num_tokens_per_request): available_blocks = kv_cache_manager.get_num_free_blocks() - print( - f"====================get_torch_compile_warmup_request============================" - ) - print(f"batch_size: {batch_size}") - print(f"num_tokens_per_request: {num_tokens_per_request}") - print(f"available_blocks: {available_blocks}") if available_blocks >= batch_size * math.ceil( num_tokens_per_request / kv_cache_manager.tokens_per_block): # Should only need (at most) one more page per request. @@ -618,9 +612,7 @@ def get_torch_compile_warmup_request(batch_size, result.context_requests = requests else: result = None - print( - f"====================get_torch_compile_warmup_request result============================" - ) + return result def get_autotune_warmup_request(): diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index af5c233cff32..0313229dc09b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -664,8 +664,8 @@ def _validate_and_adjust_attention_windows( max_seq_len: int, ) -> Tuple[Dict[int, Tuple[int, int]], int, List[int]]: """ - Validate and adjust attention windows against their upper bounds. - It will change the self.max_attention_window_vec inside the function. + Validate and adjust attention windows against their upper bounds if needed. + If there is no adjustment, the returned max_attention_window_vec will be the same as the input. Args: max_attention_window_vec: List of attention window sizes From 105acac0bf1e497db204b5e3ed6391f261278125 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Fri, 13 Jun 2025 17:58:22 +0000 Subject: [PATCH 06/15] test: Add GSM8K evaluation to Gemma3_1BInstruct tests - Included GSM8K task evaluation in both test cases for LLM API. - Updated gsm8k.yaml with accuracy reference for google/gemma-3-1b-it. Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tests/integration/defs/accuracy/references/gsm8k.yaml | 2 ++ tests/integration/defs/accuracy/test_llm_api_pytorch.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index b09e922668b0..3c1ce2f65b46 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -104,3 +104,5 @@ speakleash/Bielik-11B-v2.2-Instruct: - quant_algo: FP8 kv_cache_quant_algo: FP8 accuracy: 40.41 +google/gemma-3-1b-it: + - accuracy: 23.54 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index d731557b3d12..0995b39c8ec7 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -449,6 +449,8 @@ def test_auto_dtype(self): with LLM(self.MODEL_PATH) as llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) def test_auto_dtype_vswa(self): kv_cache_config = KvCacheConfig( @@ -458,6 +460,8 @@ def test_auto_dtype_vswa(self): with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config) as llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) class TestMixtral8x7B(LlmapiAccuracyTestHarness): From 72612d0906107cb835f6bb5039ca439bb9f5f384 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Mon, 16 Jun 2025 21:56:06 +0000 Subject: [PATCH 07/15] fix: - Explicitly set 'tokens_per_block' for Cpp binding model config. Previously 'tokens_per_block' was default value(64) which might leads to error. - Remove incorrect "fix" based on misunderstanding that each window size requires additional [max_num_tokens + window_size]. Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 16 +++++++++++++++- tensorrt_llm/_torch/pyexecutor/_util.py | 15 +-------------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 53b17974e84c..b39a927250c5 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -271,11 +271,19 @@ def from_pretrained(cls, model_config._frozen = True return model_config - def get_bindings_model_config(self) -> "ModelConfigCpp": + def get_bindings_model_config(self, + tokens_per_block: Optional[int] = None + ) -> "ModelConfigCpp": """ This method is used to construct the bindings config for the model. Currently it adheres to gptJsonConfig.cpp::createModelConfig, which assumes that an engine has been created. + + Args: + tokens_per_block: The number of tokens per block. Please note that in PyTorch flow tokens_per_block is not available in the model config, instead it is defined in the executor config. + + Returns: + The bindings model config. """ # TODO smor- this isn't robust, and currently tested for LlamaConfig only # TODO smor- currently assuming no rnn layers, no MOE @@ -294,6 +302,12 @@ def get_bindings_model_config(self) -> "ModelConfigCpp": hidden_size=hidden_size, data_type=torch_dtype_to_binding( self.pretrained_config.torch_dtype)) + if tokens_per_block is None: + logger.warning( + f"tokens_per_block is not set, using default value {model_config_cpp.tokens_per_block}" + ) + else: + model_config_cpp.tokens_per_block = tokens_per_block mlp_hidden_size = None if self.pretrained_config.intermediate_size is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 294e88c8a43e..829137df6446 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -180,20 +180,6 @@ def try_prepare_estimation(self) -> bool: estimating_kv_cache = True self._executor_config.kv_cache_config.max_tokens = self._get_token_num_for_estimation( ) - maw = self._executor_config.kv_cache_config.max_attention_window - if maw is not None: - unique_windows = list(set(maw)) - if len(unique_windows) > 1: - # NOTE: For sliding window, we need to allocate more tokens(max_num_tokens + window_size) for each window, need to update if this requirement changes - assert len( - unique_windows - ) == 2, "Sliding window with more than 2 window sizes has not been tested" - # sliding window is the one that is not the executor_config.max_seq_len - sliding_window = unique_windows[0] if unique_windows[ - 0] != self._executor_config.max_seq_len else unique_windows[ - 1] - self._executor_config.kv_cache_config.max_tokens += ( - self._executor_config.max_num_tokens + sliding_window) return estimating_kv_cache def estimate_max_tokens(self, py_executor: PyExecutor) -> None: @@ -358,6 +344,7 @@ def _create_kv_cache_manager( is_vswa = executor_config.kv_cache_config.max_attention_window is not None and len( set(executor_config.kv_cache_config.max_attention_window)) > 1 binding_model_config = model_engine.model.model_config.get_bindings_model_config( + tokens_per_block=executor_config.tokens_per_block ) if is_vswa else None kv_cache_manager = KVCacheManager( From a94798b7f542ca3f2b10ca6af68f17b2cc50b291 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 17 Jun 2025 06:59:19 +0000 Subject: [PATCH 08/15] fix: address comments Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- .../pybind/batch_manager/kvCacheManager.cpp | 2 +- .../_torch/pyexecutor/model_engine.py | 1 - .../_torch/pyexecutor/resource_manager.py | 49 +++++++++++++------ 3 files changed, 36 insertions(+), 16 deletions(-) diff --git a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp index 677180a45a76..608582ee8450 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp @@ -326,7 +326,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(py::module_& m) .def_static("calculate_max_num_blocks", &tbk::BaseKVCacheManager::calculateMaxNumBlocks, py::arg("config"), py::arg("is_cross_attention"), py::arg("dtype"), py::arg("model_config"), py::arg("world_config"), py::arg("window_size_to_layers"), py::arg("allotted_primary_mem_bytes"), - py::arg("allotted_secondary_mem_bytes"), py::arg("extra_cost_memory"), py::arg("kvFactor")) + py::arg("allotted_secondary_mem_bytes"), py::arg("extra_cost_memory"), py::arg("kv_factor")) .def("allocate_pools", &BaseKVCacheManager::allocatePools) .def("release_pools", &BaseKVCacheManager::releasePools) .def("start_scheduling", &BaseKVCacheManager::startScheduling) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index fa07065566a1..c58b4ca266e8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -612,7 +612,6 @@ def get_torch_compile_warmup_request(batch_size, result.context_requests = requests else: result = None - return result def get_autotune_warmup_request(): diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 0313229dc09b..cc512b8e3056 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -38,6 +38,9 @@ PeftCacheConfig = tensorrt_llm.bindings.executor.PeftCacheConfig WorldConfig = tensorrt_llm.bindings.WorldConfig TempAttentionWindowInputs = tensorrt_llm.bindings.internal.batch_manager.TempAttentionWindowInputs +BlocksPerWindow = Dict[int, Tuple[ + int, + int]] # window_size -> (blocks_in_primary_pool, blocks_in_secondary_pool) class ResourceManagerType(enum.Enum): @@ -177,6 +180,7 @@ def __init__( # draft/target layers. Add extra tokens to handle this issue. self.num_extra_kv_tokens = 0 if spec_config is None else spec_config.num_extra_kv_tokens self.event_buffer_max_size = kv_cache_config.event_buffer_max_size + self.max_num_tokens = max_num_tokens # Determine max_attention_window_vec if kv_cache_config.max_attention_window is None: @@ -254,13 +258,8 @@ def __init__( f"Adjusted attention window size to {self.max_seq_len} in blocks_per_window" ) - # Set up temp_attention_window_inputs only for VSWA or when needed - temp_attention_window_inputs = None - if is_vswa: # Only create when necessary - temp_attention_window_inputs = TempAttentionWindowInputs() - temp_attention_window_inputs.paged_context_fmha = True - temp_attention_window_inputs.max_input_len = self.max_seq_len - 1 - temp_attention_window_inputs.max_num_tokens = max_num_tokens + # Set up temp_attention_window_inputs + temp_attention_window_inputs = self._set_temp_attention_window_inputs() # Note that this stream is unused for now. Will be used for copying to host # when that feature is enabled. @@ -284,7 +283,6 @@ def __init__( 'enable_partial_reuse': kv_cache_config.enable_partial_reuse, 'copy_on_partial_reuse': kv_cache_config.copy_on_partial_reuse, } - if self.event_buffer_max_size > 0: kwargs['event_manager'] = KVCacheEventManagerCpp( max_kv_event_entries=self.event_buffer_max_size) @@ -590,15 +588,22 @@ def _get_window_size_to_layers(self) -> dict[int, list[int]]: # This case should ideally be prevented by earlier config validation. # If num_local_layers is 0, an empty map is fine. if self.num_local_layers > 0: - raise ValueError( + raise Exception( "max_attention_window_vec cannot be empty if there are local layers." ) return { } # Return an empty dict if no local layers or if somehow vec is empty and no layers. # Treat max_attention_window_vec as a repeating pattern. - # This also correctly handles the case where len(self.max_attention_window_vec) == 1. - pattern_len = len(self.max_attention_window_vec) + pattern_len = len( + self.max_attention_window_vec + ) # `sliding_window_pattern`, in HF config terms, e.g. https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json#L32 + # early return if max_attention_window_vec is a single value(SWA) + if pattern_len == 1: + return { + self.max_attention_window_vec[0]: + list(range(self.num_local_layers)) + } for local_layer_idx in range(self.num_local_layers): window_size = self.max_attention_window_vec[local_layer_idx % pattern_len] @@ -651,18 +656,18 @@ def calculate_max_num_blocks_from_cpp( allotted_primary_mem_bytes=primary_pool_memory_bytes, allotted_secondary_mem_bytes=secondary_pool_memory_bytes, extra_cost_memory=extra_cost_memory, - kvFactor=self.kv_factor, + kv_factor=self.kv_factor, ) return blocks_per_window def _validate_and_adjust_attention_windows( self, max_attention_window_vec: List[int], - blocks_per_window: Dict[int, Tuple[int, int]], + blocks_per_window: BlocksPerWindow, tokens_per_block: int, sink_token_length: int, max_seq_len: int, - ) -> Tuple[Dict[int, Tuple[int, int]], int, List[int]]: + ) -> Tuple[BlocksPerWindow, int, List[int]]: """ Validate and adjust attention windows against their upper bounds if needed. If there is no adjustment, the returned max_attention_window_vec will be the same as the input. @@ -721,6 +726,22 @@ def _validate_and_adjust_attention_windows( else: return blocks_per_window, max_seq_len, max_attention_window_vec + def _set_temp_attention_window_inputs( + self) -> Optional[TempAttentionWindowInputs]: + """ + Set up temp_attention_window_inputs for sliding window. + """ + is_sliding_window = min( + self.max_attention_window_vec) < self.max_seq_len + if is_sliding_window: + temp_attention_window_inputs = TempAttentionWindowInputs() + temp_attention_window_inputs.paged_context_fmha = True + temp_attention_window_inputs.max_input_len = self.max_seq_len - 1 + temp_attention_window_inputs.max_num_tokens = self.max_num_tokens + return temp_attention_window_inputs + else: + return None + class MambaCacheManager(BaseResourceManager): From 83266bc39c082caf60698800faae8590236f77d8 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 24 Jun 2025 22:31:03 +0000 Subject: [PATCH 09/15] chores: small fix for rebasing, will remove this after #5384 merged Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 829137df6446..4b6f3b0db4bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -23,9 +23,9 @@ from .llm_request import ExecutorResponse from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (KVCacheManager, MambaHybridCacheManager, - PeftCacheManager, ResourceManager, - ResourceManagerType) +from .resource_manager import (KvCacheConfigCpp, KVCacheManager, + MambaHybridCacheManager, PeftCacheManager, + ResourceManager, ResourceManagerType) from .sampler import EarlyStopSampler, TorchSampler, TRTLLMSampler from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, SimpleScheduler) From d150b297f11e642adb4211c3e4fca62cbdbb59e1 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Mon, 30 Jun 2025 02:49:29 +0000 Subject: [PATCH 10/15] when it is swa, we need to make sure chunk size is 256 Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index e627994d94d1..b7a0129367c2 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -319,6 +319,13 @@ def create_py_executor( if executor_config.enable_chunked_context: chunk_unit_size = executor_config.tokens_per_block + if (max_seq_len + > min(executor_config.kv_cache_config.max_attention_window)): + # maxKvStepSizeInFmha = 256 + chunk_unit_size = max(256, chunk_unit_size) + logger.info( + f"ChunkUnitSize is set to {chunk_unit_size} as sliding window attention is used." + ) chunking_policy = ( executor_config.scheduler_config.context_chunking_policy if executor_config.scheduler_config.context_chunking_policy From c153f6901715fce484a72b0e2fc6f79fea83aaa0 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Mon, 30 Jun 2025 22:00:48 +0000 Subject: [PATCH 11/15] fix: Update the test and GSM8K reference score Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- .../defs/accuracy/references/gsm8k.yaml | 2 +- .../defs/accuracy/test_llm_api_pytorch.py | 27 ++++++++++++++----- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 3c1ce2f65b46..0566cf8df4fd 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -105,4 +105,4 @@ speakleash/Bielik-11B-v2.2-Instruct: kv_cache_quant_algo: FP8 accuracy: 40.41 google/gemma-3-1b-it: - - accuracy: 23.54 + - accuracy: 25.52 # score getting from lm-eval with HF implementation diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0995b39c8ec7..00cced0c0a90 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -445,21 +445,34 @@ class TestGemma3_1BInstruct(LlmapiAccuracyTestHarness): MODEL_NAME = "google/gemma-3-1b-it" MODEL_PATH = f"{llm_models_root()}/gemma/gemma-3-1b-it/" + # NOTE: Disable block reuse for SWA window model. + kv_cache_config = KvCacheConfig(enable_block_reuse=False) + def test_auto_dtype(self): - with LLM(self.MODEL_PATH) as llm: + with LLM(self.MODEL_PATH, kv_cache_config=self.kv_cache_config) as llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) task.evaluate(llm) def test_auto_dtype_vswa(self): - kv_cache_config = KvCacheConfig( - enable_block_reuse=True, - free_gpu_memory_fraction=None, - max_attention_window=[512, 512, 512, 512, 512, 32768]) - with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config) as llm: - task = CnnDailymail(self.MODEL_NAME) + # NOTE: Test with VSWA kv cache config. + self.kv_cache_config.max_attention_window = [ + 512, 512, 512, 512, 512, 32768 + ] # Gemma3 1B attention window size pattern + + with LLM(self.MODEL_PATH, kv_cache_config=self.kv_cache_config) as llm: + task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + + # chunked prefill case or more features + extra_llm_config = dict( + enable_chunked_prefill=True, + max_num_tokens=1024, + ) + with LLM(self.MODEL_PATH, + kv_cache_config=self.kv_cache_config, + **extra_llm_config) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) From c5919f45a13b6181dace322b0a198aab0fc35254 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 1 Jul 2025 00:06:05 +0000 Subject: [PATCH 12/15] Remove `KvCacheConfig` conversion as #5384 merged. Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 9 ++++----- .../_torch/pyexecutor/resource_manager.py | 16 +++------------- 2 files changed, 7 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 4b6f3b0db4bf..74650602dd3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -23,9 +23,9 @@ from .llm_request import ExecutorResponse from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (KvCacheConfigCpp, KVCacheManager, - MambaHybridCacheManager, PeftCacheManager, - ResourceManager, ResourceManagerType) +from .resource_manager import (KVCacheManager, MambaHybridCacheManager, + PeftCacheManager, ResourceManager, + ResourceManagerType) from .sampler import EarlyStopSampler, TorchSampler, TRTLLMSampler from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, SimpleScheduler) @@ -348,8 +348,7 @@ def _create_kv_cache_manager( ) if is_vswa else None kv_cache_manager = KVCacheManager( - # NOTE: from tensorrt_llm.bindings.executor.KvCacheConfig to tensorrt_llm.bindings.KvCacheConfig - KvCacheConfigCpp(executor_config.kv_cache_config), + executor_config.kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, num_layers=num_hidden_layers, num_kv_heads=num_key_value_heads, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index cc512b8e3056..d938c944962c 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -25,7 +25,6 @@ if TYPE_CHECKING: from ..speculative.interface import SpecConfig -ExecutorKvCacheConfig = tensorrt_llm.bindings.executor.KvCacheConfig BufferManagerCpp = tensorrt_llm.bindings.internal.runtime.BufferManager KVCacheManagerCpp = tensorrt_llm.bindings.internal.batch_manager.KVCacheManager KvCacheConfigCpp = tensorrt_llm.bindings.executor.KvCacheConfig @@ -206,18 +205,9 @@ def __init__( "model_config is required for VSWA (Variable Sliding Window Attention)" ) # kv cache config check - if not isinstance(kv_cache_config, KvCacheConfigCpp): - # NOTE: There is a difference between - # tensorrt_llm.bindings.KvCacheConfig(KvCacheConfigCpp) and - # tensorrt_llm.bindings.executor.KvCacheConfig - # calculate_max_num_blocks_from_cpp only accepts KvCacheConfigCpp - assert isinstance(kv_cache_config, ExecutorKvCacheConfig), ( - "Only ExecutorKvCacheConfig can be converted to KvCacheConfigCpp" - ) - kv_cache_config = KvCacheConfigCpp(kv_cache_config) - logger.warning( - "kv_cache_config is not a tensorrt_llm.bindings.KvCacheConfig, " - "converting it to a tensorrt_llm.bindings.KvCacheConfig") + assert isinstance( + kv_cache_config, KvCacheConfigCpp + ), "calculate_max_num_blocks_from_cpp only accepts KvCacheConfigCpp" blocks_per_window = self.calculate_max_num_blocks_from_cpp( kv_cache_config=kv_cache_config, model_config=model_config, From a0d6a7c28f5214e171379a5405312facdc108f0b Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 1 Jul 2025 00:12:25 +0000 Subject: [PATCH 13/15] WAR for the kernel diff between Hopper and Blackwell Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- .../unfusedAttentionKernels_2_template.h | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index b4951e8c23ef..ff1240bb9ee9 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -530,6 +530,14 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams(logn_scale, q); } auto const channelIdx{tidx}; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + // Blackwell has already supported non-cyclic kv cache. + bool const useKVCache = params.kv_cache_buffer.data != nullptr; + auto token_idx_in_kv_cache = token_idx_in_seq; + bool valid_kv_cache_pos = useKVCache; +#else + // FIXME: This is a WAR for Hopper as the kernel has not been updated to support non-cyclic kv cache. Should + // be removed after Hopper kernel is updated. auto const tokenIdxLowerBound = max(cache_seq_len - params.cyclic_kv_cache_len + params.sink_token_len, params.sink_token_len); bool const useKVCache = params.kv_cache_buffer.data != nullptr; @@ -556,6 +564,7 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams= 1000)) + // Blackwell has already supported non-cyclic kv cache. + bool const useKVCache = GEN_PHASE || params.kv_cache_buffer.data != nullptr; + bool valid_kv_cache_pos = useKVCache; +#else + // FIXME: This is a WAR for Hopper as the kernel has not been updated to support non-cyclic kv cache. Should be + // removed after Hopper kernel is updated. auto const tokenIdxLowerBound = max(cache_seq_len - params.cyclic_kv_cache_len, 0); bool const useKVCache = GEN_PHASE || params.kv_cache_buffer.data != nullptr; bool valid_kv_cache_pos = useKVCache // In KV-cache-less mode. No need to store KV values @@ -971,6 +987,7 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams(params.kv_cache_buffer.getKBlockPtr(batch_idx, token_idx_in_kv_cache)) From 2665c44d69dbf79c944319521540963715a32deb Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 1 Jul 2025 04:02:30 +0000 Subject: [PATCH 14/15] fix: check whether max_attention_window is None Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index b7a0129367c2..4f79c070381c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -319,8 +319,8 @@ def create_py_executor( if executor_config.enable_chunked_context: chunk_unit_size = executor_config.tokens_per_block - if (max_seq_len - > min(executor_config.kv_cache_config.max_attention_window)): + max_attention_window = executor_config.kv_cache_config.max_attention_window + if max_attention_window and max_seq_len > min(max_attention_window): # maxKvStepSizeInFmha = 256 chunk_unit_size = max(256, chunk_unit_size) logger.info( From c2e7b9108b2e0935db0bfe76e85341e949921865 Mon Sep 17 00:00:00 2001 From: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> Date: Tue, 1 Jul 2025 04:59:15 +0000 Subject: [PATCH 15/15] Revert "WAR for the kernel diff between Hopper and Blackwell" Skip post blackwell for now. Signed-off-by: qixiang-99 <203170375+qixiang-99@users.noreply.github.com> --- .../unfusedAttentionKernels_2_template.h | 17 ----------------- .../defs/accuracy/test_llm_api_pytorch.py | 1 + 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index ff1240bb9ee9..b4951e8c23ef 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -530,14 +530,6 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams(logn_scale, q); } auto const channelIdx{tidx}; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) - // Blackwell has already supported non-cyclic kv cache. - bool const useKVCache = params.kv_cache_buffer.data != nullptr; - auto token_idx_in_kv_cache = token_idx_in_seq; - bool valid_kv_cache_pos = useKVCache; -#else - // FIXME: This is a WAR for Hopper as the kernel has not been updated to support non-cyclic kv cache. Should - // be removed after Hopper kernel is updated. auto const tokenIdxLowerBound = max(cache_seq_len - params.cyclic_kv_cache_len + params.sink_token_len, params.sink_token_len); bool const useKVCache = params.kv_cache_buffer.data != nullptr; @@ -564,7 +556,6 @@ __global__ void applyBiasRopeUpdateKVCache(QKVPreprocessingParams= 1000)) - // Blackwell has already supported non-cyclic kv cache. - bool const useKVCache = GEN_PHASE || params.kv_cache_buffer.data != nullptr; - bool valid_kv_cache_pos = useKVCache; -#else - // FIXME: This is a WAR for Hopper as the kernel has not been updated to support non-cyclic kv cache. Should be - // removed after Hopper kernel is updated. auto const tokenIdxLowerBound = max(cache_seq_len - params.cyclic_kv_cache_len, 0); bool const useKVCache = GEN_PHASE || params.kv_cache_buffer.data != nullptr; bool valid_kv_cache_pos = useKVCache // In KV-cache-less mode. No need to store KV values @@ -987,7 +971,6 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams(params.kv_cache_buffer.getKBlockPtr(batch_idx, token_idx_in_kv_cache)) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 00cced0c0a90..737399c93531 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -66,6 +66,7 @@ def test_nvfp4_streaming(self): task.evaluate(llm, streaming=True) +@skip_post_blackwell # TODO: remove this skip after this nvbug is fixed: https://nvbugspro.nvidia.com/bug/5295470 class TestLlama3_1_8BInstruct(LlmapiAccuracyTestHarness): MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct" MODEL_PATH = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct"