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, 22:15 UTC
Upstreaming aiter triton attention backend as a new backend (#28701)
Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com>
Aleksandr Malyshev committed
Nov 19, 2025 at 11:59 UTC
ac10fd3c6900228e3c0a8fae20d039668c132446
3 files changed
+80
-1
vllm/attention/backends/registry.py
+3
@@ -46,6 +46,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
46
XFORMERS = "vllm.v1.attention.backends.xformers.XFormersAttentionBackend"
47
ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
48
ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"
49
+ ROCM_AITER_TRITON_MLA = (
50
+ "vllm.v1.attention.backends.mla.aiter_triton_mla.AiterTritonMLABackend"
51
+ )
52
ROCM_AITER_FA = (
53
"vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
54
)
vllm/platforms/rocm.py
+3
-1
@@ -234,7 +234,6 @@ class RocmPlatform(Platform):
234
if rocm_aiter_ops.is_mla_enabled() or block_size == 1
235
else AttentionBackendEnum.TRITON_MLA
236
)
237
-
237
if selected_backend == AttentionBackendEnum.TRITON_MLA:
238
if block_size != 1:
239
logger.info_once("Using Triton MLA backend.")
@@ -246,6 +245,9 @@ class RocmPlatform(Platform):
245
if selected_backend == AttentionBackendEnum.ROCM_AITER_MLA:
246
logger.info("Using AITER MLA backend.")
247
return AttentionBackendEnum.ROCM_AITER_MLA.get_path()
248
+ if selected_backend == AttentionBackendEnum.ROCM_AITER_TRITON_MLA:
249
+ logger.info("Using AITER TRITON MLA backend.")
250
+ return AttentionBackendEnum.ROCM_AITER_TRITON_MLA.get_path()
251
252
raise ValueError(
253
f" The selected backend, {selected_backend.name},"
vllm/v1/attention/backends/mla/aiter_triton_mla.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from vllm.v1.attention.backends.mla.common import MLACommonBackend
4
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
5
+ AiterMLAImpl,
6
+ AiterMLAMetadataBuilder,
7
+)
8
+
9
+
10
+class AiterTritonMLABackend(MLACommonBackend):
11
+ @staticmethod
12
+ def get_name() -> str:
13
+ return "AITER_TRITON_MLA"
14
+
15
+ @staticmethod
16
+ def get_impl_cls() -> type["AiterTritonMLAImpl"]:
17
+ return AiterTritonMLAImpl
18
+
19
+ @staticmethod
20
+ def get_builder_cls() -> type["AiterMLAMetadataBuilder"]:
21
+ return AiterMLAMetadataBuilder
22
+
23
+
24
+class AiterTritonMLAImpl(AiterMLAImpl):
25
+ def __init__(
26
+ self,
27
+ num_heads: int,
28
+ head_size: int,
29
+ scale: float,
30
+ num_kv_heads: int,
31
+ alibi_slopes: list[float] | None,
32
+ sliding_window: int | None,
33
+ kv_cache_dtype: str,
34
+ logits_soft_cap: float | None,
35
+ attn_type: str,
36
+ kv_sharing_target_layer_name: str | None,
37
+ # MLA Specific Arguments
38
+ **mla_args,
39
+ ) -> None:
40
+ super().__init__(
41
+ num_heads,
42
+ head_size,
43
+ scale,
44
+ num_kv_heads,
45
+ alibi_slopes,
46
+ sliding_window,
47
+ kv_cache_dtype,
48
+ logits_soft_cap,
49
+ attn_type,
50
+ kv_sharing_target_layer_name,
51
+ **mla_args,
52
+ )
53
+ from aiter.ops.triton.mha import flash_attn_varlen_func
54
+
55
+ self.flash_attn_varlen_func = flash_attn_varlen_func
56
+
57
+ def _flash_attn_varlen_diff_headdims(
58
+ self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
59
+ ):
60
+ result = self.flash_attn_varlen_func(
61
+ q,
62
+ k,
63
+ v,
64
+ softmax_scale=softmax_scale,
65
+ return_lse=return_softmax_lse,
66
+ **kwargs,
67
+ )
68
+ # Transpose the LSE if Triton MHA is used:
69
+ # (q.shape[0], num_q_heads) to (num_q_heads, q.shape[0])
70
+ if type(result) is tuple and return_softmax_lse:
71
+ output, lse = result
72
+ lse = lse.T.contiguous()
73
+ return (output, lse)
74
+ return result