A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 00:15 UTC
[MM Encoder] Default to use TORCH_SDPA backend for ViT on Volta/Turing GPU (#36472)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Mar 9, 2026 at 18:43 UTC
b0906d8b02681d8d8f0709f0cc730f5fe845b5b1
2 files changed
+30
-7
tests/kernels/attention/test_mha_attn.py
+15
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.attention import MMEncoderAttention
19
from vllm.platforms import current_platform
20
from vllm.platforms.cpu import CpuPlatform
21
from vllm.platforms.cuda import CudaPlatform
22
+from vllm.platforms.interface import DeviceCapability
23
from vllm.platforms.rocm import RocmPlatform
24
from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed
25
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -83,6 +84,20 @@ def test_mha_attn_platform(default_vllm_config, device: str):
84
attn = MMEncoderAttention(16, 72, scale=1)
85
assert attn.attn_backend == AttentionBackendEnum.TRITON_ATTN
86
87
+ # Test Turing (pre-Ampere, sm_75): FlashAttention requires sm>=80,
88
+ # and Triton no longer supports MMA on Turing, so we expect that
89
+ # TORCH_SDPA is used for MMEncoderAttention.
90
+ with (
91
+ patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
92
+ patch.object(
93
+ CudaPlatform,
94
+ "get_device_capability",
95
+ return_value=DeviceCapability(major=7, minor=5),
96
+ ),
97
+ ):
98
+ attn = MMEncoderAttention(16, 64, scale=1)
99
+ assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
100
+
101
102
def ref_attention(
103
query: torch.Tensor,
vllm/platforms/cuda.py
+15
-7
@@ -413,12 +413,20 @@ class CudaPlatformBase(Platform):
413
414
@classmethod
415
def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
416
- return [
417
- AttentionBackendEnum.FLASH_ATTN,
418
- AttentionBackendEnum.TRITON_ATTN,
419
- AttentionBackendEnum.TORCH_SDPA,
420
- AttentionBackendEnum.FLASHINFER,
421
- ]
416
+ if cls.has_device_capability(80):
417
+ return [
418
+ AttentionBackendEnum.FLASH_ATTN,
419
+ AttentionBackendEnum.TRITON_ATTN,
420
+ AttentionBackendEnum.TORCH_SDPA,
421
+ AttentionBackendEnum.FLASHINFER,
422
+ ]
423
+ else:
424
+ return [
425
+ AttentionBackendEnum.FLASH_ATTN,
426
+ AttentionBackendEnum.TORCH_SDPA,
427
+ AttentionBackendEnum.TRITON_ATTN,
428
+ AttentionBackendEnum.FLASHINFER,
429
+ ]
430
431
@classmethod
432
def get_vit_attn_backend(
@@ -438,7 +446,7 @@ class CudaPlatformBase(Platform):
446
cc = cls.get_device_capability()
447
for vit_attn_backend in cls.get_supported_vit_attn_backends():
448
if vit_attn_backend == AttentionBackendEnum.TORCH_SDPA:
441
- continue
449
+ return vit_attn_backend
450
try:
451
backend_class = vit_attn_backend.get_class()
452
is_backend_supported = backend_class.supports_head_size(