[Bugfix][KV Cache] Don't route uniform-page-size MLA+SWA models into DeepseekV4 packing (#48256)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Claude <noreply@anthropic.com>
Nicolò Lucchesi committed
Jul 13, 2026 at 10:16 UTC
56a357ed33dec39357344816ccba2732109a96a8
2 files changed
+48
tests/v1/core/test_kv_cache_utils.py
+43
@@ -32,6 +32,7 @@ from vllm.v1.core.kv_cache_utils import (
32
get_kv_cache_configs,
33
get_max_concurrency_for_kv_cache_config,
34
get_request_block_hasher,
35
+ group_and_unify_kv_cache_specs,
36
hash_block_tokens,
37
init_none_hash,
38
is_kv_cache_spec_uniform,
@@ -1951,6 +1952,48 @@ def new_mla_spec(cache_dtype_str=None):
1952
)
1953
1954
1955
+def new_swa_mla_spec(head_size=576, sliding_window=128):
1956
+ return SlidingWindowMLASpec(
1957
+ block_size=16,
1958
+ num_kv_heads=1,
1959
+ head_size=head_size,
1960
+ dtype=torch.float32,
1961
+ sliding_window=sliding_window,
1962
+ )
1963
+
1964
+
1965
+def test_group_and_unify_kv_cache_specs_no_swa_mla_returns_none():
1966
+ # Without any SlidingWindowMLASpec the function does not apply.
1967
+ specs = {"mla.0": new_mla_spec(), "mla.1": new_mla_spec()}
1968
+ assert group_and_unify_kv_cache_specs(specs) is None
1969
+
1970
+
1971
+def test_group_and_unify_kv_cache_specs_uniform_page_size_returns_none():
1972
+ # A non-DeepseekV4 model that mixes full MLA and sliding-window MLA layers
1973
+ # with a uniform page size must not fall into the DeepseekV4 tuple-packing
1974
+ # path; it should defer to the generic uniform-page-size grouping instead.
1975
+ mla_spec = new_mla_spec()
1976
+ swa_spec = new_swa_mla_spec()
1977
+ assert mla_spec.page_size_bytes == swa_spec.page_size_bytes
1978
+ specs = {"mla.0": mla_spec, "mla.1": new_mla_spec(), "swa.0": swa_spec}
1979
+ assert group_and_unify_kv_cache_specs(specs) is None
1980
+
1981
+
1982
+def test_group_and_unify_kv_cache_specs_mixed_page_size_groups():
1983
+ # DeepseekV4-style: differing page sizes across MLA and sliding-window MLA
1984
+ # layers do require tuple packing, so grouping must still be produced.
1985
+ mla_spec = new_mla_spec()
1986
+ swa_spec = new_swa_mla_spec(head_size=1024)
1987
+ assert mla_spec.page_size_bytes != swa_spec.page_size_bytes
1988
+ specs = {"mla.0": mla_spec, "mla.1": new_mla_spec(), "swa.0": swa_spec}
1989
+ grouped = group_and_unify_kv_cache_specs(specs)
1990
+ assert grouped is not None
1991
+ # One MLA group plus one sliding-window MLA group.
1992
+ assert len(grouped) == 2
1993
+ layer_names = {name for g in grouped for name in g.kv_cache_specs}
1994
+ assert layer_names == {"mla.0", "mla.1", "swa.0"}
1995
+
1996
+
1997
def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
1998
assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
1999
vllm/v1/core/kv_cache_utils.py
+5
@@ -1535,6 +1535,11 @@ def group_and_unify_kv_cache_specs(
1535
):
1536
return None
1537
1538
+ # SlidingWindowMLASpec models with uniform page sizes don't need tuple packing.
1539
+ page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()}
1540
+ if len(page_sizes) <= 1:
1541
+ return None
1542
+
1543
mla_specs: dict[str, KVCacheSpec] = {}
1544
grouped_swa_mla_specs: dict[tuple[int, int], dict[str, KVCacheSpec]] = defaultdict(
1545
dict