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 6, 00:45 UTC
[ROCm] Attention selector reordering (#36702)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Micah Williamson <micah.williamson@amd.com>
Gregory Shtrasberg committed
Mar 25, 2026 at 04:42 UTC
189ddefbfdf74e676c8f74ebe89826e94bf44ccf
8 files changed
+29
-56
.buildkite/scripts/hardware_ci/run-amd-test.sh
+1
-1
@@ -282,7 +282,7 @@ apply_rocm_test_overrides() {
282
283
# --- LoRA: disable custom paged attention ---
284
if [[ $cmds == *"pytest -v -s lora"* ]]; then
285
- cmds=${cmds//"pytest -v -s lora"/"VLLM_ROCM_CUSTOM_PAGED_ATTN=0 pytest -v -s lora"}
285
+ cmds=${cmds//"pytest -v -s lora"/"pytest -v -s lora"}
286
fi
287
288
# --- Kernel ignores ---
docs/design/attention_backends.md
+1
-1
@@ -175,7 +175,7 @@ Priority is **1 = highest** (tried first).
175
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
176
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
177
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
178
-| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
178
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | All | N/A |
179
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
180
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
181
tests/v1/attention/test_rocm_attention_backends_selection.py
+6
-23
@@ -54,7 +54,7 @@ def mock_on_mi3xx():
54
(
55
{},
56
None,
57
- AttentionBackendEnum.TRITON_ATTN.get_path(),
57
+ AttentionBackendEnum.ROCM_ATTN.get_path(),
58
),
59
# Test Case 2: Explicit TRITON_ATTN backend
60
(
@@ -81,41 +81,24 @@ def mock_on_mi3xx():
81
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
82
),
83
# Test Case 6: VLLM_ROCM_USE_AITER=1
84
- # (defaults to AITER FA when MHA not explicitly disabled)
84
(
85
{"VLLM_ROCM_USE_AITER": "1"},
86
None,
88
- AttentionBackendEnum.ROCM_AITER_FA.get_path(),
89
- ),
90
- # Test Case 7: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=1
91
- (
92
- {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "1"},
93
- None,
94
- AttentionBackendEnum.ROCM_AITER_FA.get_path(),
95
- ),
96
- # Test Case 8: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1
97
- (
98
- {
99
- "VLLM_ROCM_USE_AITER": "1",
100
- "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": "1",
101
- },
102
- None,
103
- AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
87
+ AttentionBackendEnum.ROCM_ATTN.get_path(),
88
),
105
- # Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
89
+ # Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
90
(
91
{"VLLM_ROCM_USE_AITER": "1"},
92
"TRITON_ATTN",
93
AttentionBackendEnum.TRITON_ATTN.get_path(),
94
),
111
- # Test Case 10: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0
112
- # (explicitly disabled)
95
+ # Test Case 8: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0
96
(
97
{"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"},
98
None,
116
- AttentionBackendEnum.TRITON_ATTN.get_path(),
99
+ AttentionBackendEnum.ROCM_ATTN.get_path(),
100
),
118
- # Test Case 11: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
101
+ # Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
102
(
103
{"VLLM_ROCM_USE_AITER": "1"},
104
"ROCM_ATTN",
vllm/_aiter_ops.py
+2
-2
@@ -50,9 +50,9 @@ def is_aiter_found_and_supported() -> bool:
50
VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.
51
"""
52
if current_platform.is_rocm() and IS_AITER_FOUND:
53
- from vllm.platforms.rocm import on_gfx9
53
+ from vllm.platforms.rocm import on_mi3xx
54
55
- return on_gfx9()
55
+ return on_mi3xx()
56
return False
57
58
vllm/envs.py
-5
@@ -117,7 +117,6 @@ if TYPE_CHECKING:
117
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
118
VLLM_ROCM_FP8_PADDING: bool = True
119
VLLM_ROCM_MOE_PADDING: bool = True
120
- VLLM_ROCM_CUSTOM_PAGED_ATTN: bool = True
120
VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False
121
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
122
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
@@ -1001,10 +1000,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1000
"VLLM_ROCM_FP8_PADDING": lambda: bool(int(os.getenv("VLLM_ROCM_FP8_PADDING", "1"))),
1001
# Pad the weights for the moe kernel
1002
"VLLM_ROCM_MOE_PADDING": lambda: bool(int(os.getenv("VLLM_ROCM_MOE_PADDING", "1"))),
1004
- # custom paged attention kernel for MI3* cards
1005
- "VLLM_ROCM_CUSTOM_PAGED_ATTN": lambda: (
1006
- os.getenv("VLLM_ROCM_CUSTOM_PAGED_ATTN", "True").lower() in ("true", "1")
1007
- ),
1003
# Whether to use the shuffled kv cache layout
1004
"VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT": lambda: (
1005
os.getenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "False").lower() in ("true", "1")
vllm/platforms/rocm.py
+8
-23
@@ -265,7 +265,6 @@ def use_rocm_custom_paged_attention(
265
and (block_size == 16 or block_size == 32)
266
and (gqa_ratio >= 1 and gqa_ratio <= 16)
267
and max_seq_len <= 128 * 1024
268
- and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
268
and sinks is None
269
)
270
@@ -280,7 +279,6 @@ def use_rocm_custom_paged_attention(
279
and max_seq_len <= 128 * 1024
280
and alibi_slopes is None
281
and kv_cache_dtype == "auto"
283
- and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN
282
and sinks is None
283
)
284
@@ -311,7 +309,7 @@ def _get_backend_priorities(
309
use_mla: bool,
310
use_sparse: bool,
311
) -> list[AttentionBackendEnum]:
314
- from vllm._aiter_ops import rocm_aiter_ops
312
+ from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
313
314
if use_sparse:
315
return [AttentionBackendEnum.ROCM_AITER_MLA_SPARSE]
@@ -328,28 +326,15 @@ def _get_backend_priorities(
326
AttentionBackendEnum.TRITON_MLA,
327
]
328
331
- backends = []
332
-
333
- # Priority 1: Check for AITER Unified Attention (must check before MHA)
334
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION:
335
- backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
336
-
337
- # Priority 2: Check for AITER MHA (Flash Attention)
338
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA:
329
+ backends = [
330
+ AttentionBackendEnum.ROCM_ATTN,
331
+ ]
332
+ if rocm_aiter_ops.is_mha_enabled():
333
backends.append(AttentionBackendEnum.ROCM_AITER_FA)
340
-
341
- # Priority 3: Check for ROCM_ATTN (prefill-decode split)
342
- from vllm.config import get_current_vllm_config_or_none
343
-
344
- vllm_config = get_current_vllm_config_or_none()
345
- if (
346
- vllm_config is not None
347
- and vllm_config.attention_config.use_prefill_decode_attention
348
- ):
349
- backends.append(AttentionBackendEnum.ROCM_ATTN)
350
-
351
- # Default: Triton Unified Attention
334
+ if is_aiter_found_and_supported():
335
+ backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
336
backends.append(AttentionBackendEnum.TRITON_ATTN)
337
+
338
return backends
339
340
vllm/v1/attention/backends/rocm_attn.py
+4
-1
@@ -195,7 +195,10 @@ class RocmAttentionBackend(AttentionBackend):
195
196
@classmethod
197
def supports_sink(cls) -> bool:
198
- return True
198
+ # ROCM custom attention kernel does not support sinks.
199
+ # Callink this backend with sinks will cause it to fall back to the Triton
200
+ # kernel, which is less efficient than the proper triton backends.
201
+ return False
202
203
forward_includes_kv_cache_update: bool = False
204
vllm/v1/attention/ops/chunked_prefill_paged_decode.py
+7
@@ -10,11 +10,14 @@
10
import torch
11
12
from vllm import _custom_ops as ops
13
+from vllm.logger import init_logger
14
from vllm.platforms import current_platform
15
from vllm.triton_utils import tl, triton
16
17
from .prefix_prefill import context_attention_fwd
18
19
+logger = init_logger(__name__)
20
+
21
float8_info = torch.finfo(current_platform.fp8_dtype())
22
23
@@ -392,6 +395,10 @@ def chunked_prefill_paged_decode(
395
fp8_out_scale=output_scale,
396
)
397
else:
398
+ logger.warning_once(
399
+ "Cannot use ROCm custom paged attention kernel,"
400
+ " falling back to Triton implementation."
401
+ )
402
real_block_size = value_cache.shape[3]
403
# The standard model directly uses the original block_size.
404
# Non-standard 544 uses 32 to accommodate integer division logic.