@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e18f0037a5

[Bugfix][KV cache] Support sparse-MLA targets with SWA drafts (#48776)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Michael Goin committed Jul 23, 2026 at 14:30 UTC e18f0037a5d54dc2ead5896af896305f2bf57496

2 files changed +181 -61

tests/v1/core/test_kv_cache_utils.py

+71 -6

@@ -4,6 +4,7 @@ import copy

4 import hashlib

5 import importlib

6 from collections.abc import Callable

7 +from types import SimpleNamespace

8 from typing import Any

10 import pytest

@@ -31,6 +32,7 @@ from vllm.v1.core.kv_cache_utils import (

32 generate_scheduler_kv_cache_config,

33 get_kv_cache_capacity,

34 get_kv_cache_configs,

35 + get_kv_cache_groups,

36 get_max_concurrency_for_kv_cache_config,

37 get_request_block_hasher,

38 group_and_unify_kv_cache_specs,

@@ -2067,10 +2069,10 @@ def test_generate_scheduler_kv_cache_config():

2069 )

2070

2071

2070 -def new_mla_spec(cache_dtype_str=None):

2072 +def new_mla_spec(cache_dtype_str=None, block_size=16):

2073 # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576

2074 return MLAAttentionSpec(

2073 - block_size=16,

2075 + block_size=block_size,

2076 num_kv_heads=1,

2077 head_size=576,

2078 dtype=torch.float32,

@@ -2120,6 +2122,66 @@ def test_group_and_unify_kv_cache_specs_mixed_page_size_groups():

2122 assert layer_names == {"mla.0", "mla.1", "swa.0"}

2123

2124

2125 +def new_indexer_mla_spec(block_size=16):

2126 + # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller

2127 + # page size than the main MLA attention (uint8, small head), so their pages

2128 + # cannot be unified.

2129 + return MLAAttentionSpec(

2130 + block_size=block_size,

2131 + num_kv_heads=1,

2132 + head_size=132,

2133 + dtype=torch.uint8,

2134 + )

2135 +

2136 +

2137 +def _grouping_config():

2138 + return SimpleNamespace(

2139 + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),

2140 + speculative_config=None,

2141 + )

2142 +

2143 +

2144 +def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified():

2145 + specs = {

2146 + "target.0.attn": new_mla_spec(),

2147 + "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288),

2148 + }

2149 + assert len({spec.page_size_bytes for spec in specs.values()}) == 1

2150 +

2151 + groups = get_kv_cache_groups(_grouping_config(), specs)

2152 +

2153 + assert len(groups) == 2

2154 + assert all(

2155 + not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups

2156 + )

2157 +

2158 +

2159 +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog):

2160 + # Sparse MLA pages cannot be padded safely. Keeping the draft's attention

2161 + # compute sliding-window while promoting only its allocation semantics lets

2162 + # every layer share the target's block table and remain contiguous.

2163 + draft = new_sliding_window_spec(block_size=16)

2164 + specs = {

2165 + "target.0.attn": new_mla_spec(block_size=64),

2166 + "target.0.indexer": new_indexer_mla_spec(block_size=64),

2167 + "draft.0": draft,

2168 + }

2169 +

2170 + groups = get_kv_cache_groups(_grouping_config(), specs)

2171 + assert len(groups) == 1

2172 + assert set(groups[0].layer_names) == set(specs)

2173 + group_spec = groups[0].kv_cache_spec

2174 + assert isinstance(group_spec, UniformTypeKVCacheSpecs)

2175 + assert group_spec.block_size == 64

2176 + promoted_draft = group_spec.kv_cache_specs["draft.0"]

2177 + assert isinstance(promoted_draft, FullAttentionSpec)

2178 + assert not isinstance(promoted_draft, SlidingWindowSpec)

2179 + assert promoted_draft.block_size == 64

2180 + assert promoted_draft.sliding_window == draft.sliding_window

2181 + assert specs["draft.0"] is draft

2182 + assert "attention compute is unchanged" in caplog.text

2183 +

2184 +

2185 def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():

2186 assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION

2187

@@ -2619,19 +2681,22 @@ def test_page_size_padded_wins():

2681

2682 def test_unify_hybrid_kv_cache_specs():

2683 # 1. has_full_attention and has_sliding_window

2622 - before_spec_1 = new_kv_cache_spec()

2684 + before_spec_1 = new_kv_cache_spec(block_size=64)

2685 before_spec_2 = new_sliding_window_spec(

2624 - page_size_padded=32 * 1024, sliding_window=1024

2686 + block_size=16, page_size_padded=32 * 1024, sliding_window=1024

2687 )

2688 kv_cache_spec = {

2689 "layer_1": before_spec_1,

2690 "layer_2": before_spec_2,

2691 }

2692 kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)

2631 - expected_spec_1 = new_kv_cache_spec()

2632 - expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024)

2693 + expected_spec_1 = new_kv_cache_spec(block_size=64)

2694 + expected_spec_2 = new_kv_cache_spec(

2695 + block_size=64, page_size_padded=64 * 1024, sliding_window=1024

2696 + )

2697 assert kv_cache_spec["layer_1"] == expected_spec_1

2698 assert kv_cache_spec["layer_2"] == expected_spec_2

2699 + assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024

2700

2701 # 2. has_full_attention and has_chunked_local_attention

2702 before_spec_1 = new_kv_cache_spec()

vllm/v1/core/kv_cache_utils.py

+110 -55

@@ -1422,100 +1422,150 @@ def get_kv_cache_config_from_groups(

1422 )

1423

1424

1425 -def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):

1426 - """

1427 - This function tries to convert the KV cache specs to one type if the model

1428 - is a hybrid model with multiple type of KV cache. It will convert all

1429 - SlidingWindowSpec to FullAttentionSpec if both types are present.

1425 +def _promote_local_kv_cache_specs(

1426 + kv_cache_spec: dict[str, KVCacheSpec],

1427 +) -> dict[str, KVCacheSpec]:

1428 + """Use full-attention allocation for local-attention cache specs.

1429

1431 - Args:

1432 - kv_cache_spec: The kv cache spec of each attention layer in the model

1430 + The returned specs affect KV cache management only. Attention modules keep

1431 + their original sliding-window or chunked-local compute behavior.

1432 """

1433 + promoted_specs = kv_cache_spec.copy()

1434

1435 if is_kv_cache_spec_uniform(

1436 - kv_cache_spec

1437 - ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):

1438 - return

1439 -

1440 - logger.warning(

1441 - "Hybrid KV cache manager is disabled for this hybrid model, "

1442 - "This means we do not enable any optimizations for saving KV cache "

1443 - "memory (e.g., dropping the KV cache outside the sliding window). "

1444 - "The compute of layers like sliding window is still saved."

1445 - )

1436 + promoted_specs

1437 + ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs):

1438 + return promoted_specs

1439

1440 has_full_attention = any(

1448 - isinstance(spec, FullAttentionSpec) for spec in kv_cache_spec.values()

1441 + isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values()

1442 )

1443 has_sliding_window = any(

1451 - isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values()

1444 + isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values()

1445 )

1446 has_chunked_local_attention = any(

1454 - isinstance(spec, ChunkedLocalAttentionSpec) for spec in kv_cache_spec.values()

1447 + isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values()

1448 )

1456 - has_swa_mla = any(

1457 - isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values()

1449 + full_block_sizes = {

1450 + spec.block_size

1451 + for spec in promoted_specs.values()

1452 + if isinstance(spec, FullAttentionSpec)

1453 + }

1454 + full_attention_block_size = (

1455 + next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None

1456 )

1457

1460 - uniform_block_size: int | None = None

1461 - if has_swa_mla:

1462 - # For DeepseekV4, block sizes can be different for different KV cache groups.

1463 - # E.g., Full MLA: 256; SWA MLA: 64; C4 partial states: 4, C128 states: 8.

1464 - assert has_full_attention

1465 - any_full_spec = next(

1466 - iter(

1467 - spec

1468 - for spec in kv_cache_spec.values()

1469 - if isinstance(spec, FullAttentionSpec)

1470 - )

1458 + def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None:

1459 + if spec.page_size_padded is None:

1460 + return None

1461 + unpadded_page_size = (

1462 + spec.unpadded_page_size_bytes * block_size // spec.block_size

1463 )

1472 - uniform_block_size = any_full_spec.block_size

1464 + return max(spec.page_size_padded, unpadded_page_size)

1465

1466 if has_full_attention and (has_sliding_window or has_chunked_local_attention):

1467 for layer_name, spec in kv_cache_spec.items():

1468 if isinstance(spec, SlidingWindowMLASpec):

1477 - kv_cache_spec[layer_name] = MLAAttentionSpec(

1478 - block_size=uniform_block_size

1479 - if uniform_block_size is not None

1480 - else spec.block_size,

1469 + block_size = full_attention_block_size or spec.block_size

1470 + promoted_specs[layer_name] = MLAAttentionSpec(

1471 + block_size=block_size,

1472 num_kv_heads=spec.num_kv_heads,

1473 head_size=spec.head_size,

1474 dtype=spec.dtype,

1484 - page_size_padded=spec.page_size_padded,

1475 + page_size_padded=promoted_page_size_padded(spec, block_size),

1476 cache_dtype_str=spec.cache_dtype_str,

1477 alignment=spec.alignment,

1478 compress_ratio=spec.compress_ratio,

1479 model_version=spec.model_version,

1480 )

1481 elif isinstance(spec, SlidingWindowSpec):

1491 - kv_cache_spec[layer_name] = FullAttentionSpec(

1492 - block_size=spec.block_size,

1482 + block_size = full_attention_block_size or spec.block_size

1483 + promoted_specs[layer_name] = FullAttentionSpec(

1484 + block_size=block_size,

1485 num_kv_heads=spec.num_kv_heads,

1486 head_size=spec.head_size,

1487 head_size_v=spec.head_size_v,

1488 dtype=spec.dtype,

1489 kv_quant_mode=spec.kv_quant_mode,

1490 sliding_window=spec.sliding_window,

1499 - page_size_padded=spec.page_size_padded,

1491 + page_size_padded=promoted_page_size_padded(spec, block_size),

1492 )

1493 elif isinstance(spec, ChunkedLocalAttentionSpec):

1502 - kv_cache_spec[layer_name] = FullAttentionSpec(

1503 - block_size=spec.block_size,

1494 + block_size = full_attention_block_size or spec.block_size

1495 + promoted_specs[layer_name] = FullAttentionSpec(

1496 + block_size=block_size,

1497 num_kv_heads=spec.num_kv_heads,

1498 head_size=spec.head_size,

1499 dtype=spec.dtype,

1500 attention_chunk_size=spec.attention_chunk_size,

1508 - page_size_padded=spec.page_size_padded,

1501 + page_size_padded=promoted_page_size_padded(spec, block_size),

1502 )

1503

1504 if not (

1512 - is_kv_cache_spec_uniform(kv_cache_spec)

1513 - or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec)

1505 + is_kv_cache_spec_uniform(promoted_specs)

1506 + or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs)

1507 ):

1515 - raise ValueError(

1516 - "Hybrid KV cache manager is disabled but failed to "

1517 - "convert the KV cache specs to one unified type."

1518 - )

1508 + raise ValueError("Failed to promote local KV cache specs to one unified type.")

1509 +

1510 + return promoted_specs

1511 +

1512 +

1513 +def _try_get_full_allocation_fallback_groups(

1514 + kv_cache_spec: dict[str, KVCacheSpec],

1515 +) -> list[KVCacheGroupSpec] | None:

1516 + """Try a supported full-allocation fallback for local-attention layers."""

1517 + if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()):

1518 + return None

1519 + if any(

1520 + isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec))

1521 + for spec in kv_cache_spec.values()

1522 + ):

1523 + return None

1524 +

1525 + has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values())

1526 + has_regular_swa = any(

1527 + isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values()

1528 + )

1529 + if not (has_mla and has_regular_swa):

1530 + return None

1531 +

1532 + try:

1533 + promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec)

1534 + except ValueError:

1535 + return None

1536 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs)

1537 + if uniform_spec is None:

1538 + return None

1539 + logger.warning(

1540 + "KV cache page sizes cannot be unified; treating sliding-window "

1541 + "layers as full attention for cache allocation. Sliding-window "

1542 + "attention compute is unchanged."

1543 + )

1544 + return _get_kv_cache_groups_uniform_type(uniform_spec)

1545 +

1546 +

1547 +def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):

1548 + """

1549 + This function tries to convert the KV cache specs to one type if the model

1550 + is a hybrid model with multiple type of KV cache. It will convert all

1551 + SlidingWindowSpec to FullAttentionSpec if both types are present.

1552 +

1553 + Args:

1554 + kv_cache_spec: The kv cache spec of each attention layer in the model

1555 + """

1556 +

1557 + if is_kv_cache_spec_uniform(

1558 + kv_cache_spec

1559 + ) or UniformTypeKVCacheSpecs.is_uniform_type(kv_cache_spec):

1560 + return

1561 +

1562 + logger.warning(

1563 + "Hybrid KV cache manager is disabled for this hybrid model, "

1564 + "This means we do not enable any optimizations for saving KV cache "

1565 + "memory (e.g., dropping the KV cache outside the sliding window). "

1566 + "The compute of layers like sliding window is still saved."

1567 + )

1568 + kv_cache_spec.update(_promote_local_kv_cache_specs(kv_cache_spec))

1569

1570

1571 def group_and_unify_kv_cache_specs(

@@ -1758,10 +1808,15 @@ def get_kv_cache_groups(

1808 if not isinstance(v, HiddenStateCacheSpec)

1809 }

1810

1761 - # As KVCacheManager can only allocate memory of one size, we need to unify

1762 - # the page size of the layers. For cases cannot be unified, this function

1763 - # will raise an error.

1764 - filtered_spec = unify_kv_cache_spec_page_size(filtered_spec)

1811 + # Prefer preserving each layer's cache semantics. If physical pages cannot

1812 + # be unified, try a supported allocation-only fallback before failing.

1813 + try:

1814 + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec)

1815 + except NotImplementedError:

1816 + fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec)

1817 + if fallback_groups is None:

1818 + raise

1819 + return fallback_groups

1820 groups = _get_kv_cache_groups_uniform_page_size(filtered_spec)

1821

1822 # Add hidden-state layers back with page aligned to the common page.