[Core] Let SWA layers take the primary block size to avoid inflating the KV block LCM (#53007)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
bnellnm committed
Sep 8, 2026 at 16:44 UTC
d29c88f162a3ec17ddb0afdf5f4f3b3ca6b84d59
3 files changed
+45
-11
tests/v1/attention/test_backend_per_kind.py
+22
-1
@@ -5,12 +5,33 @@
5
import pytest
6
7
from vllm.config.attention import AttentionConfig
8
-from vllm.v1.attention.backend import AttentionType
8
+from vllm.model_executor.layers.attention.attention import (
9
+ _largest_kernel_block_within,
10
+)
11
+from vllm.v1.attention.backend import AttentionType, MultipleOf
12
from vllm.v1.attention.backends.registry import AttentionBackendEnum
13
from vllm.v1.attention.selector import get_attn_spec_kind
14
from vllm.v1.kv_cache_interface import KVCacheSpecKind
15
16
17
+@pytest.mark.parametrize(
18
+ "supported_sizes,expected",
19
+ [
20
+ ([MultipleOf(16)], 1536),
21
+ ([16, 32], 32),
22
+ ([2048], 2048),
23
+ ([MultipleOf(2048)], 2048),
24
+ ],
25
+)
26
+def test_largest_kernel_block_within(supported_sizes, expected):
27
+ class Backend:
28
+ @staticmethod
29
+ def get_supported_kernel_block_sizes():
30
+ return supported_sizes
31
+
32
+ assert _largest_kernel_block_within(Backend, 1024, 1024 * 1536, 2048) == expected
33
+
34
+
35
@pytest.mark.parametrize(
36
"signals,expected",
37
[
vllm/model_executor/layers/attention/attention.py
+20
-10
@@ -98,7 +98,7 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
98
def _largest_kernel_block_within(
99
attn_backend: "type[AttentionBackend]",
100
per_token_bytes: int,
101
- page_budget: int | None,
101
+ page_budget: int,
102
fallback: int,
103
) -> int:
104
"""Largest supported kernel block size whose page fits in ``page_budget``.
@@ -106,20 +106,22 @@ def _largest_kernel_block_within(
106
A padded spec (e.g. skip-quant layer) that pads its page up to a large shared page
107
wastes ``page_budget - block*per_token`` bytes per block. Picking the largest kernel
108
block whose natural page still fits under ``page_budget`` minimizes that waste.
109
- Falls back to the smallest supported block when ``page_budget`` is None (no padding
110
- — the block is handled by ``unify``'s integer scaling instead) or nothing fits.
109
+ ``MultipleOf`` declarations are expanded to the largest aligned block that fits.
110
+ Falls back to the smallest supported block when nothing fits.
111
"""
112
from vllm.v1.attention.backend import MultipleOf
113
114
sizes = attn_backend.get_supported_kernel_block_sizes()
115
+ max_block_size = page_budget // per_token_bytes
116
candidates = [s for s in sizes if isinstance(s, int)]
116
- if not candidates:
117
- candidates = [s.base for s in sizes if isinstance(s, MultipleOf)]
117
+ candidates.extend(
118
+ max(s.base, max_block_size // s.base * s.base)
119
+ for s in sizes
120
+ if isinstance(s, MultipleOf)
121
+ )
122
if not candidates:
123
return fallback
124
smallest = min(candidates)
121
- if not page_budget or per_token_bytes <= 0:
122
- return smallest
125
fitting = [b for b in candidates if b * per_token_bytes <= page_budget]
126
return max(fitting) if fitting else smallest
127
@@ -621,8 +623,15 @@ class Attention(nn.Module, AttentionLayerBase):
623
# When this SW layer is a padded spec (skip-quant: its page is
624
# padded up to ``skip_page_size_padded``), pick the largest kernel
625
# block that still fits the shared page so we waste fewer padding
624
- # bytes per block. Otherwise (page_size_padded is None) the smallest
625
- # block is fine — ``unify`` scales it up by an integer ratio.
626
+ # bytes per block. Otherwise (page_size_padded is None) take the
627
+ # primary block size when the backend can run it unsplit: if this
628
+ # page does not divide the primary page, ``unify`` then pads it
629
+ # (padded pages cannot be split) instead of scaling a small block
630
+ # to a size coprime with the primary one, which inflates the
631
+ # scheduler LCM (e.g. a 1024 B/token SWA draft next to a 1152
632
+ # B/token MLA target: 1728 vs 1536 gives LCM 13824). Backends that
633
+ # cannot run the primary block start from their smallest block and
634
+ # ``unify`` scales it up by an integer ratio.
635
shared_page = vllm_config.cache_config.skip_page_size_padded
636
# The backend owns its packing
637
sw_per_token = self.attn_backend.customize_spec(
@@ -636,8 +645,9 @@ class Attention(nn.Module, AttentionLayerBase):
645
sliding_window=self.sliding_window,
646
)
647
).real_page_size_bytes
648
+ page_budget = shared_page or sw_per_token * block_size
649
sw_block_size = _largest_kernel_block_within(
640
- self.attn_backend, sw_per_token, shared_page, block_size
650
+ self.attn_backend, sw_per_token, page_budget, block_size
651
)
652
return SlidingWindowSpec(
653
block_size=sw_block_size,
vllm/v1/core/kv_cache_utils.py
+3
@@ -729,6 +729,9 @@ def resolve_kv_cache_block_sizes(
729
]
730
scheduler_block_size = math.lcm(*group_block_sizes)
731
732
+ logger.info("kv cache group sizes %s", group_block_sizes)
733
+ logger.info("kv lcm block sizes %s", scheduler_block_size)
734
+
735
# Block hashes are only consumed by prefix caching and KV connectors
736
# (P/D, offloading); when neither is active, keep hash_block_size equal
737
# to the scheduler block size.