xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 00:45:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
This commit is contained in:
Matt 2025-10-09 01:20:18 -05:00 committed by GitHub
parent 8bd696fa53
commit de253d63b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/platforms/rocm.py
View File

@ -276,6 +276,9 @@ class RocmPlatform(Platform):
) )
if envs.VLLM_USE_V1: 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 ( if (
envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9() envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()
) or selected_backend == _Backend.ROCM_AITER_FA: ) or selected_backend == _Backend.ROCM_AITER_FA: