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 15, 18:30 UTC
[V1] Enable Triton(ROCm) Attention backend for Nvidia GPUs (#14071)
Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Isotr0py committed
Mar 21, 2025 at 11:14 UTC
f8a08cb90dc0b5b45663cd2605d0c98c77efe009
5 files changed
+23
-16
vllm/engine/arg_utils.py
+1
-1
@@ -1588,7 +1588,7 @@ class EngineArgs:
1588
# No FlashInfer or XFormers so far.
1589
V1_BACKENDS = [
1590
"FLASH_ATTN_VLLM_V1", "FLASH_ATTN", "PALLAS", "PALLAS_VLLM_V1",
1591
- "TRITON_MLA", "FLASHMLA"
1591
+ "TRITON_ATTN_VLLM_V1", "TRITON_MLA", "FLASHMLA"
1592
]
1593
if (envs.is_set("VLLM_ATTENTION_BACKEND")
1594
and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
vllm/platforms/cuda.py
+8
-3
@@ -213,9 +213,14 @@ class CudaPlatformBase(Platform):
213
return ("vllm.attention.backends."
214
"flashmla.FlashMLABackend")
215
if use_v1:
216
- logger.info_once("Using Flash Attention backend on V1 engine.")
217
- return ("vllm.v1.attention.backends.flash_attn."
218
- "FlashAttentionBackend")
216
+ if selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
217
+ logger.info_once("Using Triton backend on V1 engine.")
218
+ return ("vllm.v1.attention.backends."
219
+ "triton_attn.TritonAttentionBackend")
220
+ if cls.has_device_capability(80):
221
+ logger.info_once("Using Flash Attention backend on V1 engine.")
222
+ return ("vllm.v1.attention.backends."
223
+ "flash_attn.FlashAttentionBackend")
224
if selected_backend == _Backend.FLASHINFER:
225
logger.info("Using FlashInfer backend.")
226
return "vllm.attention.backends.flashinfer.FlashInferBackend"
vllm/platforms/interface.py
+1
@@ -29,6 +29,7 @@ def in_wsl() -> bool:
29
class _Backend(enum.Enum):
30
FLASH_ATTN = enum.auto()
31
FLASH_ATTN_VLLM_V1 = enum.auto()
32
+ TRITON_ATTN_VLLM_V1 = enum.auto()
33
XFORMERS = enum.auto()
34
ROCM_FLASH = enum.auto()
35
TORCH_SDPA = enum.auto()
vllm/platforms/rocm.py
+3
-2
@@ -120,8 +120,9 @@ class RocmPlatform(Platform):
120
selected_backend = (_Backend.ROCM_FLASH if selected_backend
121
== _Backend.FLASH_ATTN else selected_backend)
122
if envs.VLLM_USE_V1:
123
- logger.info("Using ROCm Attention backend on V1 engine.")
124
- return "vllm.v1.attention.backends.rocm_attn.ROCmAttentionBackend"
123
+ logger.info("Using Triton Attention backend on V1 engine.")
124
+ return ("vllm.v1.attention.backends."
125
+ "triton_attn.TritonAttentionBackend")
126
if selected_backend == _Backend.ROCM_FLASH:
127
if not cls.has_device_capability(90):
128
# not Instinct series GPUs.
vllm/v1/attention/backends/triton_attn.py
renamed
+10
-10
@@ -1,5 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-"""Attention layer with PagedAttention on rocm"""
2
+"""Attention layer with PagedAttention and Triton prefix prefill."""
3
from typing import Any, Optional
4
5
import torch
@@ -16,7 +16,7 @@ from vllm.v1.attention.backends.flash_attn import (
16
logger = init_logger(__name__)
17
18
19
-class ROCmAttentionBackend(AttentionBackend):
19
+class TritonAttentionBackend(AttentionBackend):
20
21
accept_output_buffer: bool = True
22
@@ -26,11 +26,11 @@ class ROCmAttentionBackend(AttentionBackend):
26
27
@staticmethod
28
def get_name() -> str:
29
- return "ROCM_ATTN_VLLM_V1"
29
+ return "TRITON_ATTN_VLLM_V1"
30
31
@staticmethod
32
- def get_impl_cls() -> type["ROCmAttentionImpl"]:
33
- return ROCmAttentionImpl
32
+ def get_impl_cls() -> type["TritonAttentionImpl"]:
33
+ return TritonAttentionImpl
34
35
@staticmethod
36
def get_metadata_cls() -> type["AttentionMetadata"]:
@@ -56,7 +56,7 @@ class ROCmAttentionBackend(AttentionBackend):
56
return FlashAttentionMetadataBuilder
57
58
59
-class ROCmAttentionImpl(AttentionImpl):
59
+class TritonAttentionImpl(AttentionImpl):
60
61
def __init__(
62
self,
@@ -73,7 +73,7 @@ class ROCmAttentionImpl(AttentionImpl):
73
) -> None:
74
if blocksparse_params is not None:
75
raise ValueError(
76
- "ROCmAttention does not support block-sparse attention.")
76
+ "TritonAttention does not support block-sparse attention.")
77
self.num_heads = num_heads
78
self.head_size = head_size
79
self.scale = float(scale)
@@ -90,17 +90,17 @@ class ROCmAttentionImpl(AttentionImpl):
90
assert self.num_heads % self.num_kv_heads == 0
91
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
92
93
- support_head_sizes = ROCmAttentionBackend.get_supported_head_sizes()
93
+ support_head_sizes = TritonAttentionBackend.get_supported_head_sizes()
94
if head_size not in support_head_sizes:
95
raise ValueError(
96
- f"Head size {head_size} is not supported by ROCmAttention. "
96
+ f"Head size {head_size} is not supported by TritonAttention. "
97
f"Supported head sizes are: {support_head_sizes}.")
98
99
if attn_type != AttentionType.DECODER:
100
raise NotImplementedError("Encoder self-attention and "
101
"encoder/decoder cross-attention "
102
"are not implemented for "
103
- "ROCmAttentionImpl")
103
+ "TritonAttentionImpl")
104
105
def forward(
106
self,