diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 99331e7b354e..5b0626eb3abf 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -28,14 +28,16 @@ from vllm.lora.layers import LoRAMappingType from vllm.lora.request import LoRARequest from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange from vllm.platforms import current_platform +from vllm.platforms.interface import Platform from vllm.sampling_params import SamplingParams 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, @@ -322,6 +324,53 @@ def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch) assert selected_size == 16 +def test_update_block_size_for_backend_honors_strictest_backend(): + # A backend with a fixed block size requirement (e.g. the DeepSeek V3.2 + # indexer, which only supports 64) must win over more permissive sibling + # backends when picking the cache block size. + class _FlexibleBackend(AttentionBackend): + @staticmethod + def get_name(): + return "FLEXIBLE" + + @staticmethod + def get_supported_kernel_block_sizes(): + return [MultipleOf(16)] + + class _StrictBackend(AttentionBackend): + @staticmethod + def get_name(): + return "STRICT" + + @staticmethod + def get_supported_kernel_block_sizes(): + return [64] + + class _FakeLayer(AttentionLayerBase): + def __init__(self, backend): + self._backend = backend + + def get_attn_backend(self): + return self._backend + + def get_kv_cache_spec(self, vllm_config): + return None + + vllm_config = SimpleNamespace( + cache_config=CacheConfig(), + model_config=SimpleNamespace(is_hybrid=False), + compilation_config=SimpleNamespace( + static_forward_context={ + "layers.0.attn": _FakeLayer(_FlexibleBackend), + "layers.0.indexer": _FakeLayer(_StrictBackend), + } + ), + ) + + Platform.update_block_size_for_backend(vllm_config) + assert vllm_config.cache_config.block_size == 64 + + def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids( dist_init, ): diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index b0b522882994..b56ca5c1a727 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -458,9 +458,10 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: return # reconcile attention and mamba page sizes - backend_cls = cls._find_non_ssm_backend(vllm_config) - if backend_cls is None: + backend_classes = cls._find_non_ssm_backends(vllm_config) + if not backend_classes: return + backend_cls = backend_classes[0] cls._align_hybrid_block_size(vllm_config, backend_cls) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7ad0b42da774..9b42f7f5946f 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -586,10 +586,10 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: pass @classmethod - def _find_non_ssm_backend( + def _find_non_ssm_backends( cls, vllm_config: "VllmConfig" - ) -> "type[AttentionBackend] | None": - """Find the first non-SSM attention backend from model layers.""" + ) -> "list[type[AttentionBackend]]": + """Find all distinct non-SSM attention backends from model layers.""" from vllm.config.vllm import get_layers_from_vllm_config from vllm.model_executor.layers.attention_layer_base import ( AttentionLayerBase, @@ -599,11 +599,12 @@ def _find_non_ssm_backend( vllm_config, AttentionLayerBase, # type: ignore[type-abstract] ) + backends: list[type[AttentionBackend]] = [] for layer in attn_layers.values(): b = layer.get_attn_backend() - if not b.is_ssm(): - return b - return None + if not b.is_ssm() and b not in backends: + backends.append(b) + return backends @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: @@ -621,21 +622,31 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: if not model_config: return - backend_cls = cls._find_non_ssm_backend(vllm_config) - if backend_cls is None: + backend_classes = cls._find_non_ssm_backends(vllm_config) + if not backend_classes: return + backend_cls = backend_classes[0] - # Phase 1: Pick block size from backend (skip if user set --block-size) + # Phase 1: Pick block size from backends (skip if user set --block-size). + # Honor the strictest backend: e.g. the DeepSeek V3.2 indexer only + # supports block size 64 while the main sparse attention backend may + # accept the default. if not cache_config.user_specified_block_size: + preferred = CacheConfig.DEFAULT_BLOCK_SIZE + preferred_backend_cls = backend_cls with set_current_vllm_config(vllm_config): - preferred = backend_cls.get_preferred_block_size( - CacheConfig.DEFAULT_BLOCK_SIZE - ) + for backend in backend_classes: + backend_preferred = backend.get_preferred_block_size( + CacheConfig.DEFAULT_BLOCK_SIZE + ) + if backend_preferred > preferred: + preferred = backend_preferred + preferred_backend_cls = backend if preferred != CacheConfig.DEFAULT_BLOCK_SIZE: logger.info( "Setting kv cache block size to %d for %s backend.", preferred, - backend_cls.get_name(), + preferred_backend_cls.get_name(), ) cache_config.block_size = preferred