Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
107cf8e9
Unverified
Commit
107cf8e9
authored
Jan 08, 2026
by
Rabi Mishra
Committed by
GitHub
Jan 08, 2026
Browse files
fix(rocm): Add get_supported_kernel_block_sizes() to ROCM_ATTN (#31712)
Signed-off-by:
rabi
<
ramishra@redhat.com
>
parent
63baa28c
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
8 additions
and
0 deletions
+8
-0
vllm/v1/attention/backends/rocm_attn.py
vllm/v1/attention/backends/rocm_attn.py
+8
-0
No files found.
vllm/v1/attention/backends/rocm_attn.py
View file @
107cf8e9
...
@@ -11,6 +11,7 @@ from vllm.attention.backends.abstract import (
...
@@ -11,6 +11,7 @@ from vllm.attention.backends.abstract import (
AttentionBackend
,
AttentionBackend
,
AttentionImpl
,
AttentionImpl
,
AttentionType
,
AttentionType
,
MultipleOf
,
)
)
from
vllm.attention.ops.chunked_prefill_paged_decode
import
chunked_prefill_paged_decode
from
vllm.attention.ops.chunked_prefill_paged_decode
import
chunked_prefill_paged_decode
from
vllm.attention.ops.paged_attn
import
PagedAttention
from
vllm.attention.ops.paged_attn
import
PagedAttention
...
@@ -158,6 +159,13 @@ class RocmAttentionBackend(AttentionBackend):
...
@@ -158,6 +159,13 @@ class RocmAttentionBackend(AttentionBackend):
torch
.
float32
,
torch
.
float32
,
]
]
@
staticmethod
def
get_supported_kernel_block_sizes
()
->
list
[
int
|
MultipleOf
]:
# ROCM paged attention kernel only supports block sizes 16 and 32
# due to shared memory (LDS) constraints on AMD GPUs.
# See csrc/rocm/attention.cu CALL_CUSTOM_LAUNCHER_BLK macro.
return
[
16
,
32
]
@
classmethod
@
classmethod
def
get_supported_head_sizes
(
cls
)
->
list
[
int
]:
def
get_supported_head_sizes
(
cls
)
->
list
[
int
]:
return
[
32
,
64
,
96
,
128
,
160
,
192
,
224
,
256
]
return
[
32
,
64
,
96
,
128
,
160
,
192
,
224
,
256
]
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment