Unverified Commit 5490d633 authored by Ning Xie's avatar Ning Xie Committed by GitHub
Browse files

[UT] fix unify_kv_cache_configs when kv cache config needs sort (#23843)

parent 628d00cd
...@@ -601,8 +601,14 @@ def test_unify_kv_cache_configs(): ...@@ -601,8 +601,14 @@ def test_unify_kv_cache_configs():
] ]
unify_kv_cache_configs(need_sort_kv_cache_config) unify_kv_cache_configs(need_sort_kv_cache_config)
assert need_sort_kv_cache_config[0].num_blocks == 10 sorted_kv_cache_groups = [
assert need_sort_kv_cache_config[1].num_blocks == 10 KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
KVCacheGroupSpec(["layer2"], new_kv_cache_spec(num_kv_heads=4)),
]
assert (
need_sort_kv_cache_config[0].kv_cache_groups == sorted_kv_cache_groups)
assert (
need_sort_kv_cache_config[1].kv_cache_groups == sorted_kv_cache_groups)
diff_kv_cache_config = [ diff_kv_cache_config = [
KVCacheConfig( KVCacheConfig(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment