diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 0e6089d529a6..16c1de048c27 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -79,7 +79,7 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size( backend = SimpleNamespace( get_name=lambda name=name: name, customize_spec=AttentionBackend.customize_spec, - get_supported_kernel_block_sizes=lambda sizes=sizes: sizes, + get_supported_kernel_block_sizes=lambda *_, sizes=sizes: sizes, ) layers[name] = SimpleNamespace( get_kv_cache_spec=lambda _, spec=specs[name]: spec, diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index fb4b7d1e58e4..e29c6de5bc6c 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -36,7 +36,7 @@ from vllm.utils.mem_constants import GiB_bytes from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed -from vllm.v1.attention.backend import MultipleOf +from vllm.v1.attention.backend import AttentionBackend, MultipleOf from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( ROCMAiterMLASparseBackend, @@ -303,7 +303,7 @@ def get_name(): return "MOCK" @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return supported_sizes return _MockBackend() @@ -441,13 +441,13 @@ def get_name() -> str: return "MOCK_EXACT" if exact else "MOCK" @staticmethod - def get_supported_kernel_block_sizes(): + def get_supported_kernel_block_sizes(kv_cache_spec=None): return list(supported) if exact: @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size, kv_cache_spec=None): return block_size is None or block_size in supported return _MockBackendCls @@ -469,21 +469,61 @@ def supports_block_size(cls, block_size: int | None) -> bool: ], ) def test_preferred_block_size_satisfies_every_backend(backends, expected): - classes = [_mock_backend(s) for s in backends] - assert Platform._preferred_block_size_for_backends(classes, 16, None) == expected + specs = [(_mock_backend(s), None) for s in backends] + assert Platform._preferred_block_size_for_backends(specs, 16, None) == expected def test_preferred_block_size_searches_past_an_exact_size_backend(): # Extending greedily picks lcm(16, 32) = 32, which the exact backend # rejects; 96 is accepted by both. - classes = [_mock_backend([16, 96], exact=True), _mock_backend([MultipleOf(32)])] - assert Platform._preferred_block_size_for_backends(classes, 16, None) == 96 + specs = [(_mock_backend([16, 96], exact=True), None), (_mock_backend([32]), None)] + assert Platform._preferred_block_size_for_backends(specs, 16, None) == 96 def test_preferred_block_size_rejects_backends_with_no_common_size(): - classes = [_mock_backend([16], exact=True), _mock_backend([MultipleOf(64)])] + specs = [(_mock_backend([16], exact=True), None), (_mock_backend([64]), None)] with pytest.raises(ValueError, match="share no supported KV cache block size"): - Platform._preferred_block_size_for_backends(classes, 16, None) + Platform._preferred_block_size_for_backends(specs, 16, None) + + +class _SpecAwareMockBackend(AttentionBackend): + """FlashAttention-like: a 128-token page for head size 256 (and model-wide), + multiples of 16 for other head sizes.""" + + @staticmethod + def get_name() -> str: + return "MOCK_SPEC_AWARE" + + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None): + if kv_cache_spec is None or kv_cache_spec.head_size == 256: + return [128] + return [MultipleOf(16)] + + +def _spec(head_size: int) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=head_size, dtype=torch.bfloat16 + ) + + +def test_preferred_block_size_resolves_constraints_per_spec(): + target = _mock_backend([16, 32, 64]) # FlashInfer-like head-size-256 target + drafter = _SpecAwareMockBackend # FlashAttention drafter, head size 128 + # Without specs, the drafter is held to the target's 128-token page. + backend_specs = [(target, None), (drafter, None)] + assert Platform._preferred_block_size_for_backends(backend_specs, 16, None) == 128 + backend_specs = [(target, _spec(256)), (drafter, _spec(128))] + assert Platform._preferred_block_size_for_backends(backend_specs, 16, None) == 16 + + +def test_select_common_block_size_resolves_sizes_per_spec(): + backend = _SpecAwareMockBackend + # 832 is a hybrid head-size-256 target's mamba-aligned manager block; its + # drafter's head-size-128 layers must not be held to the 128-token page. + with pytest.raises(ValueError): + select_common_block_size(832, [backend]) + assert select_common_block_size(832, [backend], [_spec(128)]) == 832 def test_set_active_mm_loras_builds_tower_and_connector_mappings(): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 769ffe5e8e25..33323f9181c1 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -27,6 +27,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.selector import AttentionSelectorConfig + from vllm.v1.kv_cache_interface import KVCacheSpec else: FlexibleArgumentParser = object @@ -592,10 +593,13 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: pass @classmethod - def _find_non_ssm_backends( + def _find_non_ssm_backend_specs( cls, vllm_config: "VllmConfig" - ) -> "list[type[AttentionBackend]]": - """Distinct non-SSM attention backends, in layer order.""" + ) -> "list[tuple[type[AttentionBackend], KVCacheSpec | None]]": + """Distinct (non-SSM attention backend, KV cache spec) pairs, in layer + order. A backend's kernel block sizes may depend on the spec it serves + (e.g. a drafter's head size), so one backend can appear with several. + """ from vllm.config.vllm import get_layers_from_vllm_config from vllm.model_executor.layers.attention_layer_base import ( AttentionLayerBase, @@ -605,21 +609,33 @@ def _find_non_ssm_backends( vllm_config, AttentionLayerBase, # type: ignore[type-abstract] ) - backends: list[type[AttentionBackend]] = [] + backend_specs: list[tuple[type[AttentionBackend], KVCacheSpec | None]] = [] for layer in attn_layers.values(): b = layer.get_attn_backend() - if not b.is_ssm() and b not in backends: - backends.append(b) - return backends + if b.is_ssm(): + continue + backend_spec = (b, layer.get_kv_cache_spec(vllm_config)) + if backend_spec not in backend_specs: + backend_specs.append(backend_spec) + return backend_specs + + @classmethod + def _find_non_ssm_backends( + cls, vllm_config: "VllmConfig" + ) -> "list[type[AttentionBackend]]": + """Distinct non-SSM attention backends, in layer order.""" + return list( + dict.fromkeys(b for b, _ in cls._find_non_ssm_backend_specs(vllm_config)) + ) @classmethod def _preferred_block_size_for_backends( cls, - backend_classes: "list[type[AttentionBackend]]", + backend_specs: "list[tuple[type[AttentionBackend], KVCacheSpec | None]]", default_block_size: int, vllm_config: "VllmConfig", ) -> int: - """Smallest block size every backend accepts. + """Smallest block size every backend accepts for the spec it serves. ``supports_block_size`` may be overridden to accept exact sizes only (CPU_MLA takes 16 and no multiple of it), so candidates are the LCMs of @@ -634,30 +650,36 @@ def _preferred_block_size_for_backends( # Backends may read the current config to decide a preference. with set_current_vllm_config(vllm_config): - if len(backend_classes) == 1: - return backend_classes[0].get_preferred_block_size(default_block_size) - if all(b.supports_block_size(default_block_size) for b in backend_classes): + # A lone backend keeps its preference, whichever specs it serves. + if len({b for b, _ in backend_specs}) == 1: + return backend_specs[0][0].get_preferred_block_size(default_block_size) + if all( + b.supports_block_size(default_block_size, spec) + for b, spec in backend_specs + ): return default_block_size # A backend declaring no sizes accepts any, so it contributes 1. per_backend_sizes = [ [ s.base if isinstance(s, MultipleOf) else s - for s in b.get_supported_kernel_block_sizes() + for s in b.get_supported_kernel_block_sizes(spec) ] or [1] - for b in backend_classes + for b, spec in backend_specs ] candidates = sorted( {math.lcm(*sizes) for sizes in itertools.product(*per_backend_sizes)} ) for candidate in candidates: - if all(b.supports_block_size(candidate) for b in backend_classes): + if all( + b.supports_block_size(candidate, spec) for b, spec in backend_specs + ): return candidate raise ValueError( "The attention backends share no supported KV cache block size (" + "; ".join( - f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" - for b in backend_classes + f"{b.get_name()}: {b.get_supported_kernel_block_sizes(spec)}" + for b, spec in backend_specs ) + ")." ) @@ -676,7 +698,8 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: if not model_config: return - backend_classes = cls._find_non_ssm_backends(vllm_config) + backend_specs = cls._find_non_ssm_backend_specs(vllm_config) + backend_classes = list(dict.fromkeys(b for b, _ in backend_specs)) if not backend_classes: return @@ -686,7 +709,7 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: # select_common_block_size(). if not cache_config.user_specified_block_size: preferred = cls._preferred_block_size_for_backends( - backend_classes, CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config + backend_specs, CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config ) if preferred != CacheConfig.DEFAULT_BLOCK_SIZE: logger.info( diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 319b9977f4a0..719f728a65c3 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -118,11 +118,15 @@ def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool: ) @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size( + cls, block_size: int | None, kv_cache_spec: "KVCacheSpec | None" = None + ) -> bool: if block_size is None: return True - supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes() + supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes( + kv_cache_spec + ) if not supported_kernel_block_sizes: return True diff --git a/vllm/v1/attention/backends/b12x.py b/vllm/v1/attention/backends/b12x.py index 833372ab446b..58d6a8f128c6 100644 --- a/vllm/v1/attention/backends/b12x.py +++ b/vllm/v1/attention/backends/b12x.py @@ -187,7 +187,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO return list(_B12X_SUPPORTED_PAGE_SIZES) @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool: return block_size is None or int(block_size) in _B12X_SUPPORTED_PAGE_SIZES @classmethod diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py index 0ba40e99f20a..f0847c5d62db 100644 --- a/vllm/v1/attention/backends/composite.py +++ b/vllm/v1/attention/backends/composite.py @@ -397,10 +397,10 @@ def get_supported_kernel_block_sizes( ) @classmethod - def supports_block_size(cls, block_size): + def supports_block_size(cls, block_size, kv_cache_spec=None): return bool( - cls.get_supported_kernel_block_sizes() - ) and super().supports_block_size(block_size) + cls.get_supported_kernel_block_sizes(kv_cache_spec) + ) and super().supports_block_size(block_size, kv_cache_spec) @classmethod def get_supported_head_sizes(cls): diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 16689887af0a..3755e8b774f7 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -134,7 +134,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO return [MultipleOf(32)] @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool: if block_size is None: return True return block_size % 32 == 0 diff --git a/vllm/v1/attention/backends/mla/cpu_mla.py b/vllm/v1/attention/backends/mla/cpu_mla.py index 431f6c318533..b500197a5c1f 100644 --- a/vllm/v1/attention/backends/mla/cpu_mla.py +++ b/vllm/v1/attention/backends/mla/cpu_mla.py @@ -68,7 +68,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO return [16] @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool: return block_size is None or block_size == 16 @staticmethod diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 950d798f2210..00220e459f18 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -111,7 +111,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO return [MultipleOf(16)] @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool: if block_size is None: return True return block_size % 16 == 0 diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index d57b420cd7da..ca574f5e2b9e 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -72,7 +72,7 @@ def get_preferred_block_size(cls, default_block_size: int) -> int: return 64 @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool: if block_size is None: return True return block_size % 16 == 0 diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index d8577b61f1c9..bc421ab3c128 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -344,7 +344,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO return [MultipleOf(16)] @classmethod - def supports_block_size(cls, block_size: int | None) -> bool: + def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool: if block_size is None: return True return block_size % 16 == 0 diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 422ba8885fd7..504504b6954d 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -330,6 +330,7 @@ def update_draft_decode_metadata( def select_common_block_size( kv_manager_block_size: int, backends: list[type[AttentionBackend]], + kv_cache_specs: list[KVCacheSpec | None] | None = None, ) -> int: """Select a block size that is supported by all backends and is a factor of kv_manager_block_size. @@ -340,6 +341,8 @@ def select_common_block_size( Args: kv_manager_block_size: Block size of KV cache. backends: List of attention backend classes. + kv_cache_specs: The KV cache spec each backend serves, for backends whose + kernel block sizes depend on it. Returns: The selected block size. @@ -348,14 +351,18 @@ def select_common_block_size( ValueError: If no valid block size found. """ - - def block_size_is_supported( - backends: list[type[AttentionBackend]], block_size: int - ) -> bool: + if kv_cache_specs is None: + kv_cache_specs = [None] * len(backends) + supported_sizes = [ + backend.get_supported_kernel_block_sizes(spec) + for backend, spec in zip(backends, kv_cache_specs, strict=True) + ] + + def block_size_is_supported(block_size: int) -> bool: """Check if the block size is supported by all backends.""" - for backend in backends: + for backend_sizes in supported_sizes: is_supported = False - for supported_size in backend.get_supported_kernel_block_sizes(): + for supported_size in backend_sizes: if isinstance(supported_size, int): if block_size == supported_size: is_supported = True @@ -368,25 +375,25 @@ def block_size_is_supported( return False return True - if block_size_is_supported(backends, kv_manager_block_size): + if block_size_is_supported(kv_manager_block_size): return kv_manager_block_size # MultipleOf constraints also accept the manager size if they accept a divisor. # Any remaining candidate must therefore be an explicit size from a backend. candidates = { size - for backend in backends - for size in backend.get_supported_kernel_block_sizes() + for backend_sizes in supported_sizes + for size in backend_sizes if isinstance(size, int) and kv_manager_block_size % size == 0 } for size in sorted(candidates, reverse=True): - if block_size_is_supported(backends, size): + if block_size_is_supported(size): return size raise ValueError( f"No common block size for {kv_manager_block_size} (" + "; ".join( - f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" for b in backends + f"{b.get_name()}: {sizes}" for b, sizes in zip(backends, supported_sizes) ) + ")." ) @@ -492,9 +499,11 @@ def prepare_kernel_block_sizes( elif isinstance(kv_cache_spec, AttentionSpec): # This is an attention backend that supports virtual block splitting. kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size - group_backends = [g.backend for g in attn_groups[kv_cache_gid]] + groups = attn_groups[kv_cache_gid] selected_kernel_size = select_common_block_size( - kv_manager_block_size, group_backends + kv_manager_block_size, + [g.backend for g in groups], + [g.kv_cache_spec for g in groups], ) kernel_block_sizes.append(selected_kernel_size) elif isinstance(kv_cache_spec, MambaSpec):