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 10, 12:45 UTC
fix(rocm): Add get_supported_kernel_block_sizes() to ROCM_ATTN (#31712)
fix(rocm): Add get_supported_kernel_block_sizes() to ROCM_ATTN (#31712) Signed-off-by: rabi <ramishra@redhat.com>
Rabi Mishra committed
Jan 8, 2026 at 13:16 UTC
107cf8e92f88678712e3aa207f054243e56df330
1 file changed
+8
vllm/v1/attention/backends/rocm_attn.py
+8
@@ -11,6 +11,7 @@ from vllm.attention.backends.abstract import (
11
AttentionBackend,
12
AttentionImpl,
13
AttentionType,
14
+ MultipleOf,
15
)
16
from vllm.attention.ops.chunked_prefill_paged_decode import chunked_prefill_paged_decode
17
from vllm.attention.ops.paged_attn import PagedAttention
@@ -158,6 +159,13 @@ class RocmAttentionBackend(AttentionBackend):
159
torch.float32,
160
]
161
162
+ @staticmethod
163
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
164
+ # ROCM paged attention kernel only supports block sizes 16 and 32
165
+ # due to shared memory (LDS) constraints on AMD GPUs.
166
+ # See csrc/rocm/attention.cu CALL_CUSTOM_LAUNCHER_BLK macro.
167
+ return [16, 32]
168
+
169
@classmethod
170
def get_supported_head_sizes(cls) -> list[int]:
171
return [32, 64, 96, 128, 160, 192, 224, 256]