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, 21:00 UTC
[ROCm][V1][Bugfix] Add get_builder_cls method to the ROCmAttentionBackend class (#14065)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Feb 28, 2025 at 23:18 UTC
b28246f6ff1684ef166d04cc9185e113a8474696
1 file changed
+6
-1
vllm/v1/attention/backends/rocm_attn.py
+6
-1
@@ -9,7 +9,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
9
from vllm.attention.ops.paged_attn import PagedAttention
10
from vllm.attention.ops.prefix_prefill import context_attention_fwd
11
from vllm.logger import init_logger
12
-from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
12
+from vllm.v1.attention.backends.flash_attn import (
13
+ FlashAttentionMetadata, FlashAttentionMetadataBuilder)
14
15
logger = init_logger(__name__)
16
@@ -49,6 +50,10 @@ class ROCmAttentionBackend(AttentionBackend):
50
def use_cascade_attention(*args, **kwargs) -> bool:
51
return False
52
53
+ @staticmethod
54
+ def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]:
55
+ return FlashAttentionMetadataBuilder
56
+
57
58
class ROCmAttentionImpl(AttentionImpl):
59