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 7, 04:00 UTC
Reenable features for ROCm attention backends (#36185)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Mar 5, 2026 at 22:21 UTC
c5362c739fb31c171fd345ed4a83fb0127804aa3
8 files changed
+66
-33
docs/design/attention_backends.md
+5
-5
@@ -171,9 +171,9 @@ Priority is **1 = highest** (tried first).
171
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
172
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
173
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
174
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
175
-| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | All | N/A |
176
-| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
174
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
175
+| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
176
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 544 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
177
| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
178
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
179
@@ -210,7 +210,7 @@ configuration.
210
| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
211
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
212
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
213
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
-| `ROCM_AITER_MLA_SPARSE` | bf16 | `auto` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
213
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
214
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
215
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
216
| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
vllm/v1/attention/backend.py
+1
-1
@@ -252,7 +252,7 @@ class AttentionBackend(ABC):
252
else:
253
invalid_reasons.append("non-MLA not supported")
254
if has_sink and not cls.supports_sink():
255
- invalid_reasons.append("sink setting not supported")
255
+ invalid_reasons.append("attention sinks not supported")
256
if use_sparse != cls.is_sparse():
257
if use_sparse:
258
invalid_reasons.append("sparse not supported")
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+10
@@ -8,6 +8,7 @@ import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.config import VllmConfig
11
+from vllm.config.cache import CacheDType
12
from vllm.model_executor.layers.attention.mla_attention import (
13
MLACommonBackend,
14
MLACommonDecodeMetadata,
@@ -21,6 +22,15 @@ from vllm.v1.kv_cache_interface import AttentionSpec
22
23
24
class AiterMLABackend(MLACommonBackend):
25
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
26
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
27
+ "auto",
28
+ "bfloat16",
29
+ "fp8",
30
+ "fp8_e4m3",
31
+ "fp8_e5m2",
32
+ ]
33
+
34
@staticmethod
35
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
36
return [1]
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+11
-10
@@ -9,6 +9,7 @@ import torch
9
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig
12
+from vllm.config.cache import CacheDType
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.attention.mla_attention import (
15
get_mla_dims,
@@ -21,6 +22,7 @@ from vllm.v1.attention.backend import (
22
AttentionMetadata,
23
AttentionMetadataBuilder,
24
CommonAttentionMetadata,
25
+ MultipleOf,
26
SparseMLAAttentionImpl,
27
)
28
from vllm.v1.attention.backends.mla.flashmla_sparse import (
@@ -77,7 +79,15 @@ def fetch_id_to_ragged_triton(
79
80
class ROCMAiterMLASparseBackend(AttentionBackend):
81
accept_output_buffer: bool = True
80
- supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
82
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
83
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
84
+ "auto",
85
+ "bfloat16",
86
+ ]
87
+
88
+ @staticmethod
89
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
90
+ return [1]
91
92
@staticmethod
93
def get_name() -> str:
@@ -105,10 +115,6 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
115
) -> tuple[int, ...]:
116
return (num_blocks, block_size, head_size)
117
108
- @classmethod
109
- def get_supported_head_sizes(cls) -> list[int]:
110
- return [576]
111
-
118
@classmethod
119
def is_mla(cls) -> bool:
120
return True
@@ -117,11 +123,6 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
123
def is_sparse(cls) -> bool:
124
return True
125
120
- @classmethod
121
- def supports_block_size(cls, block_size: int | None) -> bool:
122
- # The only supported block_size is 1
123
- return block_size is None or block_size == 1
124
-
126
127
@dataclass
128
class ROCMAiterMLASparseMetadata(AttentionMetadata):
vllm/v1/attention/backends/mla/triton_mla.py
-5
@@ -45,11 +45,6 @@ class TritonMLABackend(MLACommonBackend):
45
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
46
return True
47
48
- @classmethod
49
- def supports_block_size(cls, block_size: int | None) -> bool:
50
- # The only unsupported block_size is 1
51
- return block_size is None or block_size != 1
52
-
48
49
class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
50
can_return_lse_for_decode: bool = True
vllm/v1/attention/backends/rocm_aiter_fa.py
+8
@@ -9,6 +9,7 @@ import torch
9
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
+from vllm.config.cache import CacheDType
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.attention import Attention
15
from vllm.platforms import current_platform
@@ -732,6 +733,13 @@ class AiterFlashAttentionMetadataBuilder(
733
class AiterFlashAttentionBackend(AttentionBackend):
734
accept_output_buffer: bool = True
735
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
736
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
737
+ "auto",
738
+ "bfloat16",
739
+ "fp8",
740
+ "fp8_e4m3",
741
+ "fp8_e5m2",
742
+ ]
743
744
@staticmethod
745
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+17
-1
@@ -11,7 +11,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
11
QuantKey,
12
kFp8StaticTensorSym,
13
)
14
-from vllm.v1.attention.backend import AttentionLayer, AttentionType
14
+from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf
15
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
16
from vllm.v1.attention.backends.rocm_attn import (
17
RocmAttentionBackend,
@@ -25,6 +25,22 @@ logger = init_logger(__name__)
25
class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
26
accept_output_buffer: bool = True
27
28
+ @staticmethod
29
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
30
+ return [MultipleOf(16)]
31
+
32
+ @classmethod
33
+ def supports_head_size(cls, head_size: int) -> bool:
34
+ return head_size >= 32
35
+
36
+ @classmethod
37
+ def supports_mm_prefix(cls) -> bool:
38
+ return True
39
+
40
+ @classmethod
41
+ def supports_sink(cls) -> bool:
42
+ return True
43
+
44
forward_includes_kv_cache_update: bool = False
45
46
@staticmethod
vllm/v1/attention/backends/rocm_attn.py
+14
-11
@@ -9,6 +9,7 @@ import torch
9
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig
12
+from vllm.config.cache import CacheDType
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
QuantKey,
@@ -163,6 +164,13 @@ class RocmAttentionBackend(AttentionBackend):
164
torch.bfloat16,
165
torch.float32,
166
]
167
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
168
+ "auto",
169
+ "bfloat16",
170
+ "fp8",
171
+ "fp8_e4m3",
172
+ "fp8_e5m2",
173
+ ]
174
175
@staticmethod
176
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
@@ -185,15 +193,12 @@ class RocmAttentionBackend(AttentionBackend):
193
return [32, 64, 80, 96, 128, 160, 192, 224, 256]
194
195
@classmethod
188
- def validate_head_size(cls, head_size: int) -> None:
189
- if not cls.supports_head_size(head_size):
190
- attn_type = cls.__name__.removesuffix("Backend")
191
- raise ValueError(
192
- f"Head size {head_size} is not supported by {attn_type}. "
193
- f"Supported head sizes are: {cls.get_supported_head_sizes()}. "
194
- "Set --attention-backend=FLEX_ATTENTION to use "
195
- "FlexAttention backend which supports all head sizes."
196
- )
196
+ def supports_mm_prefix(cls) -> bool:
197
+ return True
198
+
199
+ @classmethod
200
+ def supports_sink(cls) -> bool:
201
+ return True
202
203
forward_includes_kv_cache_update: bool = False
204
@@ -275,8 +280,6 @@ class RocmAttentionImpl(AttentionImpl):
280
281
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
282
278
- RocmAttentionBackend.validate_head_size(head_size)
279
-
283
self.fp8_dtype = current_platform.fp8_dtype()
284
285
self.sinks = sinks