Skip to content

Commit 5490d63

Browse files
authored
[UT] fix unify_kv_cache_configs when kv cache config needs sort (#23843)
1 parent 628d00c commit 5490d63

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

tests/v1/core/test_kv_cache_utils.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -601,8 +601,14 @@ def test_unify_kv_cache_configs():
601601
]
602602

603603
unify_kv_cache_configs(need_sort_kv_cache_config)
604-
assert need_sort_kv_cache_config[0].num_blocks == 10
605-
assert need_sort_kv_cache_config[1].num_blocks == 10
604+
sorted_kv_cache_groups = [
605+
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
606+
KVCacheGroupSpec(["layer2"], new_kv_cache_spec(num_kv_heads=4)),
607+
]
608+
assert (
609+
need_sort_kv_cache_config[0].kv_cache_groups == sorted_kv_cache_groups)
610+
assert (
611+
need_sort_kv_cache_config[1].kv_cache_groups == sorted_kv_cache_groups)
606612

607613
diff_kv_cache_config = [
608614
KVCacheConfig(

0 commit comments

Comments
 (0)