Skip to content

GitLab

Unverified Commit de253d63 authored by Matt's avatar Matt Committed by GitHub
Browse files

[Hardware][AMD] Enable FlexAttention backend on ROCm (#26439) 


Signed-off-by: default avatarMatthew Wong <Matthew.Wong2@amd.com>
parent 8bd696fa
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
vllm/platforms/rocm.py
View file @ de253d63
......@@ -276,6 +276,9 @@ class RocmPlatform(Platform):
)
if envs.VLLM_USE_V1:
if selected_backend == _Backend.FLEX_ATTENTION:
logger.info("Using FlexAttention backend on V1 engine.")
return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
if (
envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()
) or selected_backend == _Backend.ROCM_AITER_FA:
......
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