Skip to content

Commit 28526fe

Browse files
authored
[https://nvbugs/5449218][fix] Fix KvCacheConfig error in test_perf (#6937)
Signed-off-by: peaceh <[email protected]>
1 parent 055fdd9 commit 28526fe

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

tests/integration/defs/perf/pytorch_model_config.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,6 @@
1717
Model pytorch yaml config for trtllm-bench perf tests
1818
"""
1919

20-
from tensorrt_llm.llmapi import KvCacheConfig
21-
2220

2321
def recursive_update(d, u):
2422
for k, v in u.items():
@@ -202,9 +200,10 @@ def get_model_yaml_config(model_label: str,
202200
lora_config['lora_config']['max_lora_rank'] = 320
203201
base_config.update(lora_config)
204202

205-
kv_cache_config = base_config.get('kv_cache_config', KvCacheConfig())
203+
kv_cache_config = base_config.get('kv_cache_config', {})
206204
if 'kv_cache_dtype' in base_config:
207-
kv_cache_config.dtype = base_config.pop('kv_cache_dtype', 'auto')
205+
kv_cache_dtype = base_config.pop('kv_cache_dtype', 'auto')
206+
kv_cache_config['dtype'] = kv_cache_dtype
208207
base_config.update({'kv_cache_config': kv_cache_config})
209208

210209
return base_config

0 commit comments

Comments
 (0)