Add FlexAttention backend on ROCm

mawong-amd · mawong-amd · commit f643ec2c37f8 · 2025-10-08T14:43:33.000-05:00
Signed-off-by: Matthew Wong &lt;Matthew.Wong2@amd.com&gt;
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
@@ -276,6 +276,9 @@ def get_attn_backend_cls(
             )
 
         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:

Original file line number	Diff line number	Diff line change
`@@ -276,6 +276,9 @@ def get_attn_backend_cls(`
`276`	`276`	`)`
`277`	`277`
`278`	`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"`
`279`	`282`	`if (`
`280`	`283`	`envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()`
`281`	`284`	`) or selected_backend == _Backend.ROCM_AITER_FA:`