From d29c88f162a3ec17ddb0afdf5f4f3b3ca6b84d59 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:44:43 -0400 Subject: [PATCH] [Core] Let SWA layers take the primary block size to avoid inflating the KV block LCM (#53007) Signed-off-by: Bill Nell Signed-off-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- tests/v1/attention/test_backend_per_kind.py | 23 +++++++++++++- .../layers/attention/attention.py | 30 ++++++++++++------- vllm/v1/core/kv_cache_utils.py | 3 ++ 3 files changed, 45 insertions(+), 11 deletions(-) diff --git a/tests/v1/attention/test_backend_per_kind.py b/tests/v1/attention/test_backend_per_kind.py index 5b197f7d15..9d180d3aee 100644 --- a/tests/v1/attention/test_backend_per_kind.py +++ b/tests/v1/attention/test_backend_per_kind.py @@ -5,12 +5,33 @@ import pytest from vllm.config.attention import AttentionConfig -from vllm.v1.attention.backend import AttentionType +from vllm.model_executor.layers.attention.attention import ( + _largest_kernel_block_within, +) +from vllm.v1.attention.backend import AttentionType, MultipleOf from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import get_attn_spec_kind from vllm.v1.kv_cache_interface import KVCacheSpecKind +@pytest.mark.parametrize( + "supported_sizes,expected", + [ + ([MultipleOf(16)], 1536), + ([16, 32], 32), + ([2048], 2048), + ([MultipleOf(2048)], 2048), + ], +) +def test_largest_kernel_block_within(supported_sizes, expected): + class Backend: + @staticmethod + def get_supported_kernel_block_sizes(): + return supported_sizes + + assert _largest_kernel_block_within(Backend, 1024, 1024 * 1536, 2048) == expected + + @pytest.mark.parametrize( "signals,expected", [ diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index a72d36a718..a001c8c29c 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -98,7 +98,7 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool: def _largest_kernel_block_within( attn_backend: "type[AttentionBackend]", per_token_bytes: int, - page_budget: int | None, + page_budget: int, fallback: int, ) -> int: """Largest supported kernel block size whose page fits in ``page_budget``. @@ -106,20 +106,22 @@ def _largest_kernel_block_within( A padded spec (e.g. skip-quant layer) that pads its page up to a large shared page wastes ``page_budget - block*per_token`` bytes per block. Picking the largest kernel block whose natural page still fits under ``page_budget`` minimizes that waste. - Falls back to the smallest supported block when ``page_budget`` is None (no padding - — the block is handled by ``unify``'s integer scaling instead) or nothing fits. + ``MultipleOf`` declarations are expanded to the largest aligned block that fits. + Falls back to the smallest supported block when nothing fits. """ from vllm.v1.attention.backend import MultipleOf sizes = attn_backend.get_supported_kernel_block_sizes() + max_block_size = page_budget // per_token_bytes candidates = [s for s in sizes if isinstance(s, int)] - if not candidates: - candidates = [s.base for s in sizes if isinstance(s, MultipleOf)] + candidates.extend( + max(s.base, max_block_size // s.base * s.base) + for s in sizes + if isinstance(s, MultipleOf) + ) if not candidates: return fallback smallest = min(candidates) - if not page_budget or per_token_bytes <= 0: - return smallest fitting = [b for b in candidates if b * per_token_bytes <= page_budget] return max(fitting) if fitting else smallest @@ -621,8 +623,15 @@ class Attention(nn.Module, AttentionLayerBase): # When this SW layer is a padded spec (skip-quant: its page is # padded up to ``skip_page_size_padded``), pick the largest kernel # block that still fits the shared page so we waste fewer padding - # bytes per block. Otherwise (page_size_padded is None) the smallest - # block is fine — ``unify`` scales it up by an integer ratio. + # bytes per block. Otherwise (page_size_padded is None) take the + # primary block size when the backend can run it unsplit: if this + # page does not divide the primary page, ``unify`` then pads it + # (padded pages cannot be split) instead of scaling a small block + # to a size coprime with the primary one, which inflates the + # scheduler LCM (e.g. a 1024 B/token SWA draft next to a 1152 + # B/token MLA target: 1728 vs 1536 gives LCM 13824). Backends that + # cannot run the primary block start from their smallest block and + # ``unify`` scales it up by an integer ratio. shared_page = vllm_config.cache_config.skip_page_size_padded # The backend owns its packing sw_per_token = self.attn_backend.customize_spec( @@ -636,8 +645,9 @@ class Attention(nn.Module, AttentionLayerBase): sliding_window=self.sliding_window, ) ).real_page_size_bytes + page_budget = shared_page or sw_per_token * block_size sw_block_size = _largest_kernel_block_within( - self.attn_backend, sw_per_token, shared_page, block_size + self.attn_backend, sw_per_token, page_budget, block_size ) return SlidingWindowSpec( block_size=sw_block_size, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 2fce0c0f4e..38bae03909 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -729,6 +729,9 @@ def resolve_kv_cache_block_sizes( ] scheduler_block_size = math.lcm(*group_block_sizes) + logger.info("kv cache group sizes %s", group_block_sizes) + logger.info("kv lcm block sizes %s", scheduler_block_size) + # Block hashes are only consumed by prefix caching and KV connectors # (P/D, offloading); when neither is active, keep hash_block_size equal # to the scheduler block size.