A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 04:00 UTC
[BugFix] Assign page_size_padded when unifying kv cache spec. (#32283)
Signed-off-by: Lihao Ran <imlihao.ran@gmail.com>
Lumosis committed
Jan 14, 2026 at 12:10 UTC
66652e8082b69ba7d1e6aca7c234433de55f1b9b
2 files changed
+99
-1
tests/v1/core/test_kv_cache_utils.py
+97
-1
@@ -36,6 +36,7 @@ from vllm.v1.core.kv_cache_utils import (
36
tensor_data,
37
)
38
from vllm.v1.kv_cache_interface import (
39
+ ChunkedLocalAttentionSpec,
40
FullAttentionSpec,
41
KVCacheConfig,
42
KVCacheGroupSpec,
@@ -102,29 +103,57 @@ def new_kv_cache_spec(
103
num_kv_heads=2,
104
head_size=64,
105
dtype=torch.float32,
106
+ page_size_padded=None,
107
sliding_window=None,
108
+ attention_chunk_size=None,
109
):
110
return FullAttentionSpec(
111
block_size=block_size,
112
num_kv_heads=num_kv_heads,
113
head_size=head_size,
114
dtype=dtype,
115
+ page_size_padded=page_size_padded,
116
sliding_window=sliding_window,
117
+ attention_chunk_size=attention_chunk_size,
118
)
119
120
121
def new_sliding_window_spec(
117
- block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float32, sliding_window=1
122
+ block_size=16,
123
+ num_kv_heads=2,
124
+ head_size=64,
125
+ dtype=torch.float32,
126
+ page_size_padded=None,
127
+ sliding_window=1,
128
):
129
return SlidingWindowSpec(
130
block_size=block_size,
131
num_kv_heads=num_kv_heads,
132
head_size=head_size,
133
dtype=dtype,
134
+ page_size_padded=page_size_padded,
135
sliding_window=sliding_window,
136
)
137
138
139
+def new_chunked_local_attention_spec(
140
+ block_size=16,
141
+ num_kv_heads=2,
142
+ head_size=64,
143
+ dtype=torch.float32,
144
+ page_size_padded=None,
145
+ attention_chunk_size=4,
146
+):
147
+ return ChunkedLocalAttentionSpec(
148
+ block_size=block_size,
149
+ num_kv_heads=num_kv_heads,
150
+ head_size=head_size,
151
+ dtype=dtype,
152
+ page_size_padded=page_size_padded,
153
+ attention_chunk_size=attention_chunk_size,
154
+ )
155
+
156
+
157
@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
158
def test_none_hash(monkeypatch, hash_fn):
159
import vllm.v1.core.kv_cache_utils
@@ -1855,3 +1884,70 @@ def test_auto_fit_max_model_len_not_triggered():
1884
vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32]
1885
)
1886
assert vllm_config.model_config.max_model_len == 16
1887
+
1888
+
1889
+def test_unify_hybrid_kv_cache_specs():
1890
+ # 1. has_full_attention and has_sliding_window
1891
+ before_spec_1 = new_kv_cache_spec()
1892
+ before_spec_2 = new_sliding_window_spec(
1893
+ page_size_padded=32 * 1024, sliding_window=1024
1894
+ )
1895
+ kv_cache_spec = {
1896
+ "layer_1": before_spec_1,
1897
+ "layer_2": before_spec_2,
1898
+ }
1899
+ kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
1900
+ expected_spec_1 = new_kv_cache_spec()
1901
+ expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024)
1902
+ assert kv_cache_spec["layer_1"] == expected_spec_1
1903
+ assert kv_cache_spec["layer_2"] == expected_spec_2
1904
+
1905
+ # 2. has_full_attention and has_chunked_local_attention
1906
+ before_spec_1 = new_kv_cache_spec()
1907
+ before_spec_2 = new_chunked_local_attention_spec(
1908
+ page_size_padded=32 * 1024, attention_chunk_size=512
1909
+ )
1910
+ kv_cache_spec = {
1911
+ "layer_1": before_spec_1,
1912
+ "layer_2": before_spec_2,
1913
+ }
1914
+ kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
1915
+ expected_spec_1 = new_kv_cache_spec()
1916
+ expected_spec_2 = new_kv_cache_spec(
1917
+ page_size_padded=32 * 1024, attention_chunk_size=512
1918
+ )
1919
+
1920
+ assert kv_cache_spec["layer_1"] == expected_spec_1
1921
+ assert kv_cache_spec["layer_2"] == expected_spec_2
1922
+
1923
+ # 3. has_full_attention, has_sliding_window and has_chunked_local_attention
1924
+ before_spec_1 = new_kv_cache_spec()
1925
+ before_spec_2 = new_sliding_window_spec(
1926
+ page_size_padded=32 * 1024, sliding_window=1024
1927
+ )
1928
+ before_spec_3 = new_chunked_local_attention_spec(
1929
+ page_size_padded=32 * 1024, attention_chunk_size=512
1930
+ )
1931
+ kv_cache_spec = {
1932
+ "layer_1": before_spec_1,
1933
+ "layer_2": before_spec_2,
1934
+ "layer_3": before_spec_3,
1935
+ }
1936
+ kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
1937
+ expected_spec_1 = new_kv_cache_spec()
1938
+ expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024)
1939
+ expected_spec_3 = new_kv_cache_spec(
1940
+ page_size_padded=32 * 1024, attention_chunk_size=512
1941
+ )
1942
+ assert kv_cache_spec["layer_1"] == expected_spec_1
1943
+ assert kv_cache_spec["layer_2"] == expected_spec_2
1944
+ assert kv_cache_spec["layer_3"] == expected_spec_3
1945
+
1946
+ # 4. No FullAttentionSpec, should not convert
1947
+ kv_cache_spec = {
1948
+ "layer_1": new_sliding_window_spec(sliding_window=1024),
1949
+ "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512),
1950
+ }
1951
+
1952
+ with pytest.raises(ValueError):
1953
+ kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
vllm/v1/core/kv_cache_utils.py
+2
@@ -1185,6 +1185,7 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1185
head_size=spec.head_size,
1186
dtype=spec.dtype,
1187
sliding_window=spec.sliding_window,
1188
+ page_size_padded=spec.page_size_padded,
1189
)
1190
elif isinstance(spec, ChunkedLocalAttentionSpec):
1191
kv_cache_spec[layer_name] = FullAttentionSpec(
@@ -1193,6 +1194,7 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1194
head_size=spec.head_size,
1195
dtype=spec.dtype,
1196
attention_chunk_size=spec.attention_chunk_size,
1197
+ page_size_padded=spec.page_size_padded,
1198
)
1199
1200
if not (