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 12, 14:45 UTC
[Bugfix] [ROCm] [UX]: revert Flex attention backend (#29371)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Nov 25, 2025 at 14:56 UTC
64deead719cc181a1930982b0a5f4d280c284156
2 files changed
+10
tests/v1/attention/test_rocm_attention_backends_selection.py
+6
@@ -36,6 +36,12 @@ def mock_on_gfx9():
36
@pytest.mark.parametrize(
37
"env_vars, selected_backend, expected_backend_path",
38
[
39
+ # Test Case: Explicit FLEX_ATTENTION backend
40
+ (
41
+ {},
42
+ "FLEX_ATTENTION",
43
+ AttentionBackendEnum.FLEX_ATTENTION.get_path(),
44
+ ),
45
# Test Case 1: Default (no env vars, no explicit backend)
46
(
47
{},
vllm/platforms/rocm.py
+4
@@ -262,6 +262,10 @@ class RocmPlatform(Platform):
262
f"is not MLA type while requested for MLA backend."
263
)
264
265
+ if selected_backend == AttentionBackendEnum.FLEX_ATTENTION:
266
+ logger.info("Using FlexAttention backend.")
267
+ return AttentionBackendEnum.FLEX_ATTENTION.get_path()
268
+
269
if selected_backend == AttentionBackendEnum.TRITON_ATTN:
270
logger.info("Using Triton Attention backend on V1 engine.")
271
return AttentionBackendEnum.TRITON_ATTN.get_path()