[BugFix][KV Cache] Fix hybrid prefix caching with hidden-state extraction (#51108)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Canlin Guo committed
Aug 6, 2026 at 14:42 UTC
537046164944cbd34ac940d8655f91704e5a0bf1
2 files changed
+82
-1
tests/v1/core/test_kv_cache_utils.py
+60
@@ -45,6 +45,7 @@ from vllm.v1.core.kv_cache_utils import (
45
from vllm.v1.kv_cache_interface import (
46
ChunkedLocalAttentionSpec,
47
FullAttentionSpec,
48
+ HiddenStateCacheSpec,
49
KVCacheConfig,
50
KVCacheGroupSpec,
51
KVCacheSpec,
@@ -2196,6 +2197,65 @@ def _grouping_config():
2197
)
2198
2199
2200
+def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity():
2201
+ block_size = 544
2202
+ full_spec = FullAttentionSpec(
2203
+ block_size=block_size,
2204
+ num_kv_heads=1,
2205
+ # FullAttentionSpec stores both K and V, so this produces a
2206
+ # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page.
2207
+ head_size=512,
2208
+ dtype=torch.bfloat16,
2209
+ )
2210
+ mamba_spec = MambaSpec(
2211
+ block_size=block_size,
2212
+ shapes=((557056,),),
2213
+ dtypes=(torch.bfloat16,),
2214
+ mamba_cache_mode="align",
2215
+ )
2216
+ hidden_spec = HiddenStateCacheSpec(
2217
+ block_size=block_size,
2218
+ num_kv_heads=3,
2219
+ head_size=1024,
2220
+ dtype=torch.bfloat16,
2221
+ )
2222
+ assert full_spec.page_size_bytes == mamba_spec.page_size_bytes
2223
+
2224
+ groups = get_kv_cache_groups(
2225
+ _grouping_config(),
2226
+ {
2227
+ "model.full_attn": full_spec,
2228
+ "model.mamba": mamba_spec,
2229
+ "cache_only_layers.0": hidden_spec,
2230
+ },
2231
+ )
2232
+
2233
+ hidden_group = next(
2234
+ group
2235
+ for group in groups
2236
+ if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
2237
+ )
2238
+ assert hidden_group.kv_cache_spec.block_size == 136
2239
+
2240
+ kv_cache_config = KVCacheConfig(
2241
+ num_blocks=1,
2242
+ kv_cache_tensors=[],
2243
+ kv_cache_groups=groups,
2244
+ )
2245
+ vllm_config = SimpleNamespace(
2246
+ cache_config=SimpleNamespace(
2247
+ block_size=16,
2248
+ enable_prefix_caching=True,
2249
+ prefix_match_unit=None,
2250
+ ),
2251
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
2252
+ kv_transfer_config=object(),
2253
+ )
2254
+ assert kv_cache_utils.resolve_kv_cache_block_sizes(
2255
+ kv_cache_config, vllm_config
2256
+ ) == (544, 136)
2257
+
2258
+
2259
def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified():
2260
specs = {
2261
"target.0.attn": new_mla_spec(),
vllm/v1/core/kv_cache_utils.py
+22
-1
@@ -645,6 +645,9 @@ def resolve_kv_cache_block_sizes(
645
if not (cache_config.enable_prefix_caching or connector_enabled):
646
return scheduler_block_size, scheduler_block_size
647
648
+ # Mamba groups outside align mode break divisibility; back off to the
649
+ # scheduler block size. Read the mode from the resolved group spec because
650
+ # its block size may have been updated independently of cache_config.
651
if any(
652
isinstance(g.kv_cache_spec, MambaSpec)
653
and g.kv_cache_spec.mamba_cache_mode != "align"
@@ -1741,6 +1744,13 @@ def _annotate_eagle_groups_deepseek_v4(
1744
break
1745
1746
1747
+def _largest_divisor_at_most(value: int, limit: int) -> int:
1748
+ for candidate in range(min(value, limit), 0, -1):
1749
+ if value % candidate == 0:
1750
+ return candidate
1751
+ return 1
1752
+
1753
+
1754
def get_kv_cache_groups(
1755
vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec]
1756
) -> list[KVCacheGroupSpec]:
@@ -1806,9 +1816,20 @@ def get_kv_cache_groups(
1816
# Add hidden-state layers back with page aligned to the common page.
1817
if hidden_specs:
1818
common_page = get_uniform_page_size([g.kv_cache_spec for g in groups])
1819
+ group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups))
1820
for name, spec in hidden_specs.items():
1821
per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
1811
- new_bs = max(common_page // per_token, 1)
1822
+ max_block_size = max(common_page // per_token, 1)
1823
+ new_bs = _largest_divisor_at_most(group_block_size, max_block_size)
1824
+ wasted_bytes = common_page - new_bs * per_token
1825
+ logger.info(
1826
+ "Using block size %d for hidden-state cache layer %s; "
1827
+ "page alignment wastes %d bytes (%.2f%%) per block",
1828
+ new_bs,
1829
+ name,
1830
+ wasted_bytes,
1831
+ wasted_bytes / common_page * 100,
1832
+ )
1833
aligned = replace(spec, block_size=new_bs, page_size_padded=common_page)
1834
groups.append(KVCacheGroupSpec([name], aligned))
1835