[Bugfix] Fix broken ViT attention selection for Blackwell device (#30731)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
This commit is contained in:
Isotr0py 2025-12-16 13:24:32 +08:00 committed by GitHub
parent b9ff4f2a8d
commit e94384bbad
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -11,7 +11,7 @@ import torch
from transformers import PretrainedConfig
from vllm.attention.backends.registry import AttentionBackendEnum
from vllm.config import VllmConfig, get_current_vllm_config
from vllm.config import VllmConfig
from vllm.distributed import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
@ -88,16 +88,10 @@ def get_vit_attn_backend(
"""
Get the available attention backend for Vision Transformer.
"""
attn_backend = attn_backend_override
selected_backend = get_current_vllm_config().attention_config.backend
if attn_backend is None:
attn_backend = selected_backend
return current_platform.get_vit_attn_backend(
head_size,
dtype,
backend=attn_backend,
backend=attn_backend_override,
)