[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 transformers import PretrainedConfig
from vllm.attention.backends.registry import AttentionBackendEnum 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 ( from vllm.distributed import (
get_tensor_model_parallel_rank, get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size, get_tensor_model_parallel_world_size,
@ -88,16 +88,10 @@ def get_vit_attn_backend(
""" """
Get the available attention backend for Vision Transformer. 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( return current_platform.get_vit_attn_backend(
head_size, head_size,
dtype, dtype,
backend=attn_backend, backend=attn_backend_override,
) )