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 5, 21:00 UTC
[ROCm][CI] Fix cross-attention dispatch for encoder-decoder models (#38450)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 29, 2026 at 00:08 UTC
43cc5138e5145752413235a2a8aa303886077327
6 files changed
+90
-19
docs/design/attention_backends.md
+2
-2
@@ -173,9 +173,9 @@ Priority is **1 = highest** (tried first).
173
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
174
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
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 |
176
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | 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 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | 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/entrypoints/openai/speech_to_text/test_transcription_validation_whisper.py
+52
-3
@@ -14,13 +14,62 @@ import pytest_asyncio
14
import soundfile as sf
15
16
from tests.utils import RemoteOpenAIServer
17
+from vllm.platforms import current_platform
18
19
MODEL_NAME = "openai/whisper-large-v3-turbo"
20
21
+# Disable prefix caching on ROCm to reduce non-determinism in
22
+# streaming-vs-non-streaming comparisons.
23
+_ROCM_ARGS = ["--no-enable-prefix-caching"] if current_platform.is_rocm() else []
24
21
-@pytest.fixture(scope="module")
22
-def server():
23
- with RemoteOpenAIServer(MODEL_NAME, []) as remote_server:
25
+
26
+def _get_attention_backend_params() -> list[str | None]:
27
+ """Return attention backends to parametrize the server fixture with.
28
+
29
+ On ROCm, we test multiple backends explicitly:
30
+ - None: default auto-selection (ROCM_ATTN for decoder self-attention,
31
+ falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for
32
+ cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER)
33
+ - TRITON_ATTN: always available on ROCm
34
+ - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950
35
+
36
+ On non-ROCm platforms, we just run with the default backend.
37
+ """
38
+ try:
39
+ from vllm.platforms import current_platform
40
+
41
+ if current_platform.is_rocm():
42
+ backends: list[str | None] = [None, "TRITON_ATTN"]
43
+ from vllm.platforms.rocm import _ON_MI3XX
44
+
45
+ if _ON_MI3XX:
46
+ backends.append("ROCM_AITER_UNIFIED_ATTN")
47
+ return backends
48
+ except Exception:
49
+ pass
50
+ return [None]
51
+
52
+
53
+# Aiter backends need VLLM_ROCM_USE_AITER=1 (and MHA=1 for ROCM_AITER_FA)
54
+# to be enabled in the server subprocess.
55
+_AITER_ENV = {
56
+ "VLLM_ROCM_USE_AITER": "1",
57
+ "VLLM_ROCM_USE_AITER_MHA": "1",
58
+}
59
+
60
+_ATTN_BACKENDS = _get_attention_backend_params()
61
+_ATTN_IDS = [b or "default" for b in _ATTN_BACKENDS]
62
+
63
+
64
+@pytest.fixture(scope="module", params=_ATTN_BACKENDS, ids=_ATTN_IDS)
65
+def server(request):
66
+ args = [*_ROCM_ARGS]
67
+ env_dict = None
68
+ if request.param is not None:
69
+ args += ["--attention-backend", request.param]
70
+ if "AITER" in request.param:
71
+ env_dict = _AITER_ENV
72
+ with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
73
yield remote_server
74
75
tools/pre_commit/generate_attention_backend_docs.py
+1
-1
@@ -446,7 +446,7 @@ def parse_attention_types(node: ast.ClassDef) -> str:
446
447
if not types:
448
return "Decoder"
449
- return "All" if len(types) >= 3 else ", ".join(sorted(types))
449
+ return "All" if types >= set(type_map.values()) else ", ".join(sorted(types))
450
451
452
def parse_impl_bool_attr(
vllm/platforms/rocm.py
+22
-6
@@ -439,7 +439,10 @@ class RocmPlatform(Platform):
439
f"this configuration. Reason: {invalid_reasons}"
440
)
441
else:
442
- logger.info("Using %s backend.", selected_backend)
442
+ logger.info_once(
443
+ "Using %s backend (selected via --attention-backend).",
444
+ selected_backend.name,
445
+ )
446
return selected_backend.get_path()
447
448
# No selected backend or the selected backend is invalid,
@@ -476,12 +479,25 @@ class RocmPlatform(Platform):
479
)
480
selected_index = sorted_indices[0]
481
selected_backend = valid_backends_priorities[selected_index][0]
479
- logger.info_once(
480
- "Using %s attention backend out of potential backends: %s.",
481
- selected_backend.name,
482
- "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
483
- scope="local",
482
+ valid_str = (
483
+ "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]"
484
)
485
+ if invalid_reasons:
486
+ rejected_str = ", ".join(b.name for b in invalid_reasons)
487
+ logger.info(
488
+ "Found incompatible backend(s) [%s] with %s. "
489
+ "Overriding with %s out of potential backends: %s.",
490
+ rejected_str,
491
+ attn_selector_config.attn_type,
492
+ selected_backend.name,
493
+ valid_str,
494
+ )
495
+ else:
496
+ logger.info_once(
497
+ "Using %s backend out of potential backends: %s.",
498
+ selected_backend.name,
499
+ valid_str,
500
+ )
501
502
return selected_backend.get_path()
503
vllm/v1/attention/backends/rocm_aiter_fa.py
+6
-5
@@ -758,11 +758,12 @@ class AiterFlashAttentionBackend(AttentionBackend):
758
759
@classmethod
760
def supports_attn_type(cls, attn_type: str) -> bool:
761
- """ROCM AITER FA supports decoder and encoder-decoder (cross) attention."""
762
- return attn_type in (
763
- AttentionType.DECODER,
764
- AttentionType.ENCODER_DECODER,
765
- )
761
+ """ENCODER_DECODER is not supported because the prefill path uses
762
+ flash_attn_varlen_func with cu_seqlens_k set to decoder
763
+ query_start_loc (not encoder seq lens) and causal=True, both of
764
+ which are incorrect for cross-attention layers.
765
+ """
766
+ return attn_type in (AttentionType.DECODER,)
767
768
@staticmethod
769
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
vllm/v1/attention/backends/rocm_attn.py
+7
-2
@@ -212,12 +212,17 @@ class RocmAttentionBackend(AttentionBackend):
212
213
@classmethod
214
def supports_attn_type(cls, attn_type: str) -> bool:
215
- """RocmAttention supports all attention types."""
215
+ """ENCODER_DECODER is not supported because
216
+ chunked_prefill_paged_decode's prefill kernel (context_attention_fwd)
217
+ assumes self-attention semantics: it treats passed K/V as new tokens
218
+ to mix with cached K/V. For cross-attention layers the encoder K/V
219
+ are already fully cached, so mixing them again produces incorrect
220
+ results when max_query_len > 1 (e.g. beam search).
221
+ """
222
return attn_type in (
223
AttentionType.DECODER,
224
AttentionType.ENCODER,
225
AttentionType.ENCODER_ONLY,
220
- AttentionType.ENCODER_DECODER,
226
)
227
228
@staticmethod