[Bugfix] Pick a KV block size supported by every attention backend (#49845)
Signed-off-by: Divy <divy@coralbricks.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Divy committed
Sep 23, 2026 at 13:36 UTC
e26547adfb4970f3e3958695256a84566683a683
3 files changed
+132
-21
tests/v1/worker/test_gpu_model_runner.py
+57
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.attention import Attention
31
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
32
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
33
from vllm.platforms import current_platform
34
+from vllm.platforms.interface import Platform
35
from vllm.sampling_params import SamplingParams
36
from vllm.utils.mem_constants import GiB_bytes
37
from vllm.utils.system_utils import update_environment_variables
@@ -425,6 +426,62 @@ def test_select_common_block_size_no_valid_option():
426
select_common_block_size(48, [backend_a, backend_b])
427
428
429
+def _mock_backend(supported: list, *, exact: bool = False):
430
+ """Backend accepting multiples of ``supported``, or only those exact sizes
431
+ when ``exact`` (as CPU_MLA does)."""
432
+ from vllm.v1.attention.backend import AttentionBackend
433
+
434
+ class _MockBackendCls(AttentionBackend):
435
+ @staticmethod
436
+ def get_name() -> str:
437
+ return "MOCK_EXACT" if exact else "MOCK"
438
+
439
+ @staticmethod
440
+ def get_supported_kernel_block_sizes():
441
+ return list(supported)
442
+
443
+ if exact:
444
+
445
+ @classmethod
446
+ def supports_block_size(cls, block_size: int | None) -> bool:
447
+ return block_size is None or block_size in supported
448
+
449
+ return _MockBackendCls
450
+
451
+
452
+@pytest.mark.parametrize(
453
+ "backends,expected",
454
+ [
455
+ # A lone backend keeps the default; MultipleOf(8) makes a
456
+ # minimum-based regression return 8 instead of 16.
457
+ ([[MultipleOf(8)]], 16),
458
+ # Sparse-MLA main backend (32 or 64) beside an indexer supporting only
459
+ # 64: the main backend alone picks 32, which the indexer rejects at
460
+ # select_common_block_size time.
461
+ ([[32, 64], [64]], 64),
462
+ ([[32, 64]], 32),
463
+ # Neither divides the other, so only the LCM satisfies both.
464
+ ([[32], [48]], 96),
465
+ ],
466
+)
467
+def test_preferred_block_size_satisfies_every_backend(backends, expected):
468
+ classes = [_mock_backend(s) for s in backends]
469
+ assert Platform._preferred_block_size_for_backends(classes, 16, None) == expected
470
+
471
+
472
+def test_preferred_block_size_searches_past_an_exact_size_backend():
473
+ # Extending greedily picks lcm(16, 32) = 32, which the exact backend
474
+ # rejects; 96 is accepted by both.
475
+ classes = [_mock_backend([16, 96], exact=True), _mock_backend([MultipleOf(32)])]
476
+ assert Platform._preferred_block_size_for_backends(classes, 16, None) == 96
477
+
478
+
479
+def test_preferred_block_size_rejects_backends_with_no_common_size():
480
+ classes = [_mock_backend([16], exact=True), _mock_backend([MultipleOf(64)])]
481
+ with pytest.raises(ValueError, match="share no supported KV cache block size"):
482
+ Platform._preferred_block_size_for_backends(classes, 16, None)
483
+
484
+
485
def test_set_active_mm_loras_builds_tower_and_connector_mappings():
486
model = Mock()
487
model.get_mm_lora_token_counts.side_effect = (
vllm/platforms/cpu.py
+3
-3
@@ -498,11 +498,11 @@ class CpuPlatform(Platform):
498
return
499
500
# reconcile attention and mamba page sizes
501
- backend_cls = cls._find_non_ssm_backend(vllm_config)
502
- if backend_cls is None:
501
+ backend_classes = cls._find_non_ssm_backends(vllm_config)
502
+ if not backend_classes:
503
return
504
505
- cls._align_hybrid_block_size(vllm_config, backend_cls)
505
+ cls._align_hybrid_block_size(vllm_config, backend_classes[0])
506
507
@classmethod
508
def discover_numa_topology(cls) -> list[list[int]]:
vllm/platforms/interface.py
+72
-18
@@ -3,6 +3,8 @@
3
import contextlib
4
import enum
5
import functools
6
+import itertools
7
+import math
8
import os
9
import platform
10
import sys
@@ -573,10 +575,10 @@ class Platform:
575
pass
576
577
@classmethod
576
- def _find_non_ssm_backend(
578
+ def _find_non_ssm_backends(
579
cls, vllm_config: "VllmConfig"
578
- ) -> "type[AttentionBackend] | None":
579
- """Find the first non-SSM attention backend from model layers."""
580
+ ) -> "list[type[AttentionBackend]]":
581
+ """Distinct non-SSM attention backends, in layer order."""
582
from vllm.config.vllm import get_layers_from_vllm_config
583
from vllm.model_executor.layers.attention_layer_base import (
584
AttentionLayerBase,
@@ -586,11 +588,62 @@ class Platform:
588
vllm_config,
589
AttentionLayerBase, # type: ignore[type-abstract]
590
)
591
+ backends: list[type[AttentionBackend]] = []
592
for layer in attn_layers.values():
593
b = layer.get_attn_backend()
591
- if not b.is_ssm():
592
- return b
593
- return None
594
+ if not b.is_ssm() and b not in backends:
595
+ backends.append(b)
596
+ return backends
597
+
598
+ @classmethod
599
+ def _preferred_block_size_for_backends(
600
+ cls,
601
+ backend_classes: "list[type[AttentionBackend]]",
602
+ default_block_size: int,
603
+ vllm_config: "VllmConfig",
604
+ ) -> int:
605
+ """Smallest block size every backend accepts.
606
+
607
+ ``supports_block_size`` may be overridden to accept exact sizes only
608
+ (CPU_MLA takes 16 and no multiple of it), so candidates are the LCMs of
609
+ one supported size per backend, each checked against every backend.
610
+
611
+ Raises:
612
+ ValueError: If the backends share no supported block size.
613
+
614
+ """
615
+ from vllm.config.vllm import set_current_vllm_config
616
+ from vllm.v1.attention.backend import MultipleOf
617
+
618
+ # Backends may read the current config to decide a preference.
619
+ with set_current_vllm_config(vllm_config):
620
+ if len(backend_classes) == 1:
621
+ return backend_classes[0].get_preferred_block_size(default_block_size)
622
+ if all(b.supports_block_size(default_block_size) for b in backend_classes):
623
+ return default_block_size
624
+ # A backend declaring no sizes accepts any, so it contributes 1.
625
+ per_backend_sizes = [
626
+ [
627
+ s.base if isinstance(s, MultipleOf) else s
628
+ for s in b.get_supported_kernel_block_sizes()
629
+ ]
630
+ or [1]
631
+ for b in backend_classes
632
+ ]
633
+ candidates = sorted(
634
+ {math.lcm(*sizes) for sizes in itertools.product(*per_backend_sizes)}
635
+ )
636
+ for candidate in candidates:
637
+ if all(b.supports_block_size(candidate) for b in backend_classes):
638
+ return candidate
639
+ raise ValueError(
640
+ "The attention backends share no supported KV cache block size ("
641
+ + "; ".join(
642
+ f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}"
643
+ for b in backend_classes
644
+ )
645
+ + ")."
646
+ )
647
648
@classmethod
649
def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -598,7 +651,6 @@ class Platform:
651
For hybrid models, also aligns block_size with mamba page sizes.
652
"""
653
from vllm.config.cache import CacheConfig
601
- from vllm.config.vllm import set_current_vllm_config
654
655
cache_config = vllm_config.cache_config
656
model_config = vllm_config.model_config
@@ -607,34 +659,36 @@ class Platform:
659
if not model_config:
660
return
661
610
- backend_cls = cls._find_non_ssm_backend(vllm_config)
611
- if backend_cls is None:
662
+ backend_classes = cls._find_non_ssm_backends(vllm_config)
663
+ if not backend_classes:
664
return
665
614
- # Phase 1: Pick block size from backend (skip if user set --block-size)
666
+ # Phase 1: Pick a block size every attention backend supports (skip if
667
+ # user set --block-size). Models can mix backends with disjoint
668
+ # preferences, and a size from the first alone later fails
669
+ # select_common_block_size().
670
if not cache_config.user_specified_block_size:
616
- with set_current_vllm_config(vllm_config):
617
- preferred = backend_cls.get_preferred_block_size(
618
- CacheConfig.DEFAULT_BLOCK_SIZE
619
- )
671
+ preferred = cls._preferred_block_size_for_backends(
672
+ backend_classes, CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config
673
+ )
674
if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
675
logger.info(
622
- "Setting kv cache block size to %d for %s backend.",
676
+ "Setting kv cache block size to %d for %s backend(s).",
677
preferred,
624
- backend_cls.get_name(),
678
+ "/".join(b.get_name() for b in backend_classes),
679
)
680
cache_config.block_size = preferred
681
682
# Phase 2: Align block/mamba sizes for hybrid models
683
# (may override user settings).
684
if model_config.is_hybrid:
631
- cls._align_hybrid_block_size(vllm_config, backend_cls)
685
+ cls._align_hybrid_block_size(vllm_config, backend_classes[0])
686
687
# Phase 3: Align block/page sizes when multiple KV dtypes share the
688
# block pool (e.g. nvfp4 primary + unquantized skip layers).
689
# May override the user's --block-size.
690
if cache_config.kv_cache_dtype_skip_layers:
637
- cls._align_heterogeneous_kv_block_size(vllm_config, backend_cls)
691
+ cls._align_heterogeneous_kv_block_size(vllm_config, backend_classes[0])
692
693
@classmethod
694
def _align_heterogeneous_kv_block_size(