[4/N][KV-Cache Layout Refactor] Promote local KV cache specs via a class-changing replace helper (#51612)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
Lucas Wilkinson committed
Aug 11, 2026 at 17:29 UTC
61874f9842bd084cc5e9e348de4b3b6a8f46cdc8
2 files changed
+35
-36
vllm/v1/core/kv_cache_utils.py
+20
-35
@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (
33
SlidingWindowMLASpec,
34
SlidingWindowSpec,
35
UniformTypeKVCacheSpecs,
36
+ replace_as,
37
)
38
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
39
from vllm.v1.request import Request
@@ -1449,43 +1450,27 @@ def _promote_local_kv_cache_specs(
1450
)
1451
return max(spec.page_size_padded, unpadded_page_size)
1452
1453
+ promotions: dict[type[AttentionSpec], type[AttentionSpec]] = {
1454
+ SlidingWindowMLASpec: MLAAttentionSpec,
1455
+ SlidingWindowSpec: FullAttentionSpec,
1456
+ ChunkedLocalAttentionSpec: FullAttentionSpec,
1457
+ }
1458
+
1459
if has_full_attention and (has_sliding_window or has_chunked_local_attention):
1460
for layer_name, spec in kv_cache_spec.items():
1454
- if isinstance(spec, SlidingWindowMLASpec):
1455
- block_size = full_attention_block_size or spec.block_size
1456
- promoted_specs[layer_name] = MLAAttentionSpec(
1457
- block_size=block_size,
1458
- num_kv_heads=spec.num_kv_heads,
1459
- head_size=spec.head_size,
1460
- dtype=spec.dtype,
1461
- page_size_padded=promoted_page_size_padded(spec, block_size),
1462
- cache_dtype_str=spec.cache_dtype_str,
1463
- alignment=spec.alignment,
1464
- compress_ratio=spec.compress_ratio,
1465
- model_version=spec.model_version,
1466
- )
1467
- elif isinstance(spec, SlidingWindowSpec):
1468
- block_size = full_attention_block_size or spec.block_size
1469
- promoted_specs[layer_name] = FullAttentionSpec(
1470
- block_size=block_size,
1471
- num_kv_heads=spec.num_kv_heads,
1472
- head_size=spec.head_size,
1473
- head_size_v=spec.head_size_v,
1474
- dtype=spec.dtype,
1475
- kv_quant_mode=spec.kv_quant_mode,
1476
- sliding_window=spec.sliding_window,
1477
- page_size_padded=promoted_page_size_padded(spec, block_size),
1478
- )
1479
- elif isinstance(spec, ChunkedLocalAttentionSpec):
1480
- block_size = full_attention_block_size or spec.block_size
1481
- promoted_specs[layer_name] = FullAttentionSpec(
1482
- block_size=block_size,
1483
- num_kv_heads=spec.num_kv_heads,
1484
- head_size=spec.head_size,
1485
- dtype=spec.dtype,
1486
- attention_chunk_size=spec.attention_chunk_size,
1487
- page_size_padded=promoted_page_size_padded(spec, block_size),
1488
- )
1461
+ target_cls = next(
1462
+ (promotions[c] for c in type(spec).__mro__ if c in promotions), None
1463
+ )
1464
+ if target_cls is None:
1465
+ continue
1466
+ assert isinstance(spec, AttentionSpec)
1467
+ block_size = full_attention_block_size or spec.block_size
1468
+ promoted_specs[layer_name] = replace_as(
1469
+ spec,
1470
+ target_cls,
1471
+ block_size=block_size,
1472
+ page_size_padded=promoted_page_size_padded(spec, block_size),
1473
+ )
1474
1475
if not (
1476
is_kv_cache_spec_uniform(promoted_specs)
vllm/v1/kv_cache_interface.py
+15
-1
@@ -8,7 +8,7 @@ from collections import Counter
8
from dataclasses import dataclass, fields, replace
9
from enum import Enum, IntEnum
10
from math import prod
11
-from typing import TYPE_CHECKING
11
+from typing import TYPE_CHECKING, TypeVar
12
13
import torch
14
from typing_extensions import Self
@@ -24,6 +24,8 @@ if TYPE_CHECKING:
24
25
logger = init_logger(__name__)
26
27
+_SpecT = TypeVar("_SpecT", bound="KVCacheSpec")
28
+
29
30
# ---------------------------------------------------------------------------
31
# KV cache quantization mode
@@ -86,6 +88,18 @@ def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
88
return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE
89
90
91
+def replace_as(spec: KVCacheSpec, target_cls: type[_SpecT], **changes) -> _SpecT:
92
+ """``dataclasses.replace``, but rebuilding *spec* as *target_cls*
93
+ e.g. ``SlidingWindowSpec`` -> ``FullAttentionSpec``
94
+
95
+ Every field of *spec* must exist on *target_cls*; fields only *target_cls* has keep
96
+ their default values.
97
+ """
98
+ kwargs = {f.name: getattr(spec, f.name) for f in fields(spec) if f.init}
99
+ kwargs.update(changes)
100
+ return target_cls(**kwargs)
101
+
102
+
103
def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
104
"""Return True if *kv_cache_dtype* needs per-token-head scales."""
105
return get_kv_quant_mode(kv_cache_dtype).is_per_token_head