From 26d6d23485df530dbfd40706b6faa3fb7c8e20d0 Mon Sep 17 00:00:00 2001 From: jianyi85 <40061533+jianyi85@users.noreply.github.com> Date: Sun, 23 Aug 2026 16:41:36 +0800 Subject: [PATCH 1/2] [Platform] Allow pinning the attention backend for components that auto-select Signed-off-by: jianyi85 <40061533+jianyi85@users.noreply.github.com> --- vllm/platforms/cuda.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 0aff4ff9bef7..008d2da720ac 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -432,6 +432,35 @@ def get_attn_backend_cls( logger.info("Using %s backend.", selected_backend) return _backend_cls_path(backend_class) + # Components that do not plumb through the user's --attention-backend + # (the spec-decode draft model in particular) arrive here with + # selected_backend=None and auto-select. On SM121 auto-selection picks + # FLASHINFER, whose kernels fault with MTP + fp8 KV (vllm#37754), so + # allow pinning the auto-selection path too. Invalid pins fall back to + # auto-selection instead of failing components with other constraints. + forced_name = os.environ.get("VLLM_FORCE_ATTN_BACKEND") + if forced_name: + forced = AttentionBackendEnum[forced_name] + try: + backend_class = _get_attn_backend_class(forced) + invalid_reasons = backend_class.validate_configuration( + device_capability=device_capability, + **attn_selector_config._asdict(), + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + logger.info( + "Using %s backend (VLLM_FORCE_ATTN_BACKEND).", forced + ) + return _backend_cls_path(backend_class) + logger.warning( + "VLLM_FORCE_ATTN_BACKEND=%s is not valid here (%s); " + "falling back to auto-selection.", + forced_name, + invalid_reasons, + ) + # No selected backend or the selected backend is invalid, # so we try finding a valid backend. valid_backends_priorities, all_invalid_reasons = cls.get_valid_backends( From a3ecdbe98b6bbe3bf3b8b57608e58f723bc0a617 Mon Sep 17 00:00:00 2001 From: jianyi85 <40061533+jianyi85@users.noreply.github.com> Date: Sun, 23 Aug 2026 17:03:12 +0800 Subject: [PATCH 2/2] [Platform] Guard unknown VLLM_FORCE_ATTN_BACKEND values; add selector tests Signed-off-by: jianyi85 <40061533+jianyi85@users.noreply.github.com> --- .../attention/test_forced_attn_backend_env.py | 120 ++++++++++++++++++ vllm/platforms/cuda.py | 43 ++++--- 2 files changed, 146 insertions(+), 17 deletions(-) create mode 100644 tests/v1/attention/test_forced_attn_backend_env.py diff --git a/tests/v1/attention/test_forced_attn_backend_env.py b/tests/v1/attention/test_forced_attn_backend_env.py new file mode 100644 index 000000000000..aa6181f00200 --- /dev/null +++ b/tests/v1/attention/test_forced_attn_backend_env.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the VLLM_FORCE_ATTN_BACKEND auto-selection pin. + +Components that do not plumb through the user's --attention-backend (the +spec-decode draft model in particular) reach get_attn_backend_cls with +selected_backend=None and auto-select. The env pins that path too; a forced +backend that is invalid for a component's configuration (or unknown) falls +back to auto-selection instead of failing the component. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.platforms.cuda import CudaPlatform +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.selector import AttentionSelectorConfig + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA-specific tests" +) + +SELECTOR_CONFIG = AttentionSelectorConfig( + head_size=64, + dtype=torch.float16, + kv_cache_dtype=None, + block_size=16, +) + +SM90 = DeviceCapability(major=9, minor=0) + + +class _AutoSelectionReached(Exception): + """Sentinel: get_valid_backends (the auto-selection path) was entered.""" + + +def _get_backend_cls(): + return CudaPlatform.get_attn_backend_cls( + selected_backend=None, + attn_selector_config=SELECTOR_CONFIG, + num_heads=32, + ) + + +def test_forced_backend_bypasses_auto_selection(monkeypatch): + monkeypatch.setenv("VLLM_FORCE_ATTN_BACKEND", "TRITON_ATTN") + healthy = MagicMock() + healthy.validate_configuration.return_value = [] + with ( + patch.object(CudaPlatform, "get_device_capability", return_value=SM90), + patch("vllm.platforms.cuda._get_attn_backend_class", return_value=healthy), + patch("vllm.platforms.cuda._backend_cls_path", return_value="forced.path"), + patch.object( + CudaPlatform, "get_valid_backends", side_effect=_AutoSelectionReached + ), + ): + assert _get_backend_cls() == "forced.path" + + +def test_invalid_forced_backend_falls_back_to_auto_selection(monkeypatch): + monkeypatch.setenv("VLLM_FORCE_ATTN_BACKEND", "TRITON_ATTN") + unfit = MagicMock() + unfit.validate_configuration.return_value = ["head_size not supported"] + with ( + patch.object(CudaPlatform, "get_device_capability", return_value=SM90), + patch("vllm.platforms.cuda._get_attn_backend_class", return_value=unfit), + patch.object( + CudaPlatform, "get_valid_backends", side_effect=_AutoSelectionReached + ), + pytest.raises(_AutoSelectionReached), + ): + _get_backend_cls() + + +def test_unknown_forced_backend_falls_back_to_auto_selection(monkeypatch): + monkeypatch.setenv("VLLM_FORCE_ATTN_BACKEND", "NO_SUCH_BACKEND") + with ( + patch.object(CudaPlatform, "get_device_capability", return_value=SM90), + patch.object( + CudaPlatform, "get_valid_backends", side_effect=_AutoSelectionReached + ), + pytest.raises(_AutoSelectionReached), + ): + _get_backend_cls() + + +def test_no_env_leaves_auto_selection_unchanged(monkeypatch): + monkeypatch.delenv("VLLM_FORCE_ATTN_BACKEND", raising=False) + with ( + patch.object(CudaPlatform, "get_device_capability", return_value=SM90), + patch.object( + CudaPlatform, "get_valid_backends", side_effect=_AutoSelectionReached + ), + pytest.raises(_AutoSelectionReached), + ): + _get_backend_cls() + + +def test_explicit_selected_backend_takes_precedence(monkeypatch): + """A component that DOES receive an explicit backend is unaffected by the + env: the selected_backend branch returns before the forced-env check.""" + monkeypatch.setenv("VLLM_FORCE_ATTN_BACKEND", "TRITON_ATTN") + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + healthy = MagicMock() + healthy.validate_configuration.return_value = [] + with ( + patch.object(CudaPlatform, "get_device_capability", return_value=SM90), + patch("vllm.platforms.cuda._get_attn_backend_class", return_value=healthy), + patch("vllm.platforms.cuda._backend_cls_path", return_value="explicit.path"), + ): + result = CudaPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.FLASH_ATTN, + attn_selector_config=SELECTOR_CONFIG, + num_heads=32, + ) + assert result == "explicit.path" diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 008d2da720ac..d9728f22990c 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -440,26 +440,35 @@ def get_attn_backend_cls( # auto-selection instead of failing components with other constraints. forced_name = os.environ.get("VLLM_FORCE_ATTN_BACKEND") if forced_name: - forced = AttentionBackendEnum[forced_name] try: - backend_class = _get_attn_backend_class(forced) - invalid_reasons = backend_class.validate_configuration( - device_capability=device_capability, - **attn_selector_config._asdict(), + forced = AttentionBackendEnum[forced_name] + except KeyError: + logger.warning( + "VLLM_FORCE_ATTN_BACKEND=%s is not a known backend; " + "falling back to auto-selection.", + forced_name, ) - except ImportError: - invalid_reasons = ["ImportError"] - if not invalid_reasons: - logger.info( - "Using %s backend (VLLM_FORCE_ATTN_BACKEND).", forced + forced = None + if forced is not None: + try: + backend_class = _get_attn_backend_class(forced) + invalid_reasons = backend_class.validate_configuration( + device_capability=device_capability, + **attn_selector_config._asdict(), + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + logger.info( + "Using %s backend (VLLM_FORCE_ATTN_BACKEND).", forced + ) + return _backend_cls_path(backend_class) + logger.warning( + "VLLM_FORCE_ATTN_BACKEND=%s is not valid here (%s); " + "falling back to auto-selection.", + forced_name, + invalid_reasons, ) - return _backend_cls_path(backend_class) - logger.warning( - "VLLM_FORCE_ATTN_BACKEND=%s is not valid here (%s); " - "falling back to auto-selection.", - forced_name, - invalid_reasons, - ) # No selected backend or the selected backend is invalid, # so we try finding a valid backend.