A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Hardware][AMD] Enable FlexAttention backend on ROCm (#26439)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Oct 9, 2025 at 01:20 UTC
de253d63b7a4cb59bb6d9da107d26f8978e17107
1 file changed
+3
vllm/platforms/rocm.py
+3
@@ -276,6 +276,9 @@ class RocmPlatform(Platform):
276
)
277
278
if envs.VLLM_USE_V1:
279
+ if selected_backend == _Backend.FLEX_ATTENTION:
280
+ logger.info("Using FlexAttention backend on V1 engine.")
281
+ return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
282
if (
283
envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()
284
) or selected_backend == _Backend.ROCM_AITER_FA: