[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>

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",
     [

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,

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.