diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py index d60a18878086..0a55a5ed0ae4 100644 --- a/tests/v1/core/test_kv_cache_planning.py +++ b/tests/v1/core/test_kv_cache_planning.py @@ -1783,6 +1783,16 @@ def test_unify_hybrid_kv_cache_specs(): kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) +def test_unify_hybrid_kv_cache_specs_swa_mla(): + swa = replace(new_swa_mla_spec(), bounded_replay=True) + kv_cache_spec = {"full": new_mla_spec(), "swa": swa} + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + promoted = kv_cache_spec["swa"] + assert type(promoted) is MLAAttentionSpec + assert promoted.sliding_window == swa.sliding_window + assert promoted.page_size_bytes == swa.page_size_bytes + + def test_unify_kv_cache_spec_page_size_mamba(): """Regression test for https://github.com/vllm-project/vllm/issues/43626. diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py index 4d73f4303e2f..2b98d2e9ad75 100644 --- a/vllm/v1/core/kv_cache_planning.py +++ b/vllm/v1/core/kv_cache_planning.py @@ -801,8 +801,9 @@ def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | Non spec, target_cls, # Promoted specs allocate blocks for all tokens and never free - # below the window, so the trailing-edge extension is moot. - drop=("extra_retained_tokens",), + # below the window, so the trailing-edge extension and the + # window-bounded prefix replay are both moot. + drop=("extra_retained_tokens", "bounded_replay"), block_size=block_size, page_size_padded=promoted_page_size_padded(spec, block_size), )