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/model_config.py b/tensorrt_llm/_torch/model_config.py index 97a816c3d478..b39a927250c5 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 @@ -270,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 @@ -293,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: @@ -323,6 +338,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 +359,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..74650602dd3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -340,6 +340,13 @@ 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( + tokens_per_block=executor_config.tokens_per_block + ) if is_vswa else None + kv_cache_manager = KVCacheManager( executor_config.kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, @@ -352,7 +359,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/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index e627994d94d1..4f79c070381c 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 + 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( + 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 diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e58043dffb31..d938c944962c 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,7 @@ if TYPE_CHECKING: from ..speculative.interface import SpecConfig +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 +36,10 @@ 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 +BlocksPerWindow = Dict[int, Tuple[ + int, + int]] # window_size -> (blocks_in_primary_pool, blocks_in_secondary_pool) class ResourceManagerType(enum.Enum): @@ -124,6 +129,8 @@ 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: self.mapping = mapping self.dtype = dtype @@ -149,7 +156,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 +176,80 @@ 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 + self.max_num_tokens = max_num_tokens + # 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 + # model config check + if model_config is None: + raise ValueError( + "model_config is required for VSWA (Variable Sliding Window Attention)" + ) + # kv cache config check + 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, + extra_cost_memory=0, + ) + else: + # Standard case: use original Python implementation + 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, + mapping=mapping, + dtype=dtype, + kv_factor=self.kv_factor, + ) + blocks_per_window = { + self.max_attention_window_vec[0]: + (self.blocks_in_primary_pool, self.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 + 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. @@ -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, @@ -422,6 +459,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 +482,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 @@ -453,18 +493,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]: + 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, - self.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]]: - result = self.impl.get_batch_cache_block_ids(request_ids, - self.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] @@ -513,6 +565,173 @@ 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 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. + 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] + 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, + 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: BlocksPerWindow, + tokens_per_block: int, + sink_token_length: int, + max_seq_len: 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. + + 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 + + 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): diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index b09e922668b0..0566cf8df4fd 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: 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 41c0014ec9cf..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" @@ -445,10 +446,36 @@ 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): + # 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) class TestMixtral8x7B(LlmapiAccuracyTestHarness):