Commit 6a88d590 authored by Isotr0py's avatar Isotr0py Committed by Kevin H. Luu
Browse files

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


Signed-off-by: default avatarIsotr0py <mozf@mail2.sysu.edu.cn>
(cherry picked from commit e94384bb)
parent ad8c0731
......@@ -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,
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment