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, 16:00 UTC
[ROCM][V0] PA kennel selection when no sliding window provided (#15982)
Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com>
Aleksandr Malyshev committed
Apr 2, 2025 at 22:28 UTC
57a810db9c3cdfff96edb004ed111df31859c775
1 file changed
+2
-1
vllm/platforms/rocm.py
+2
-1
@@ -109,7 +109,8 @@ def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
109
ON_MI250_MI300 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])
110
111
# rocm custom page attention not support on navi (gfx1*)
112
- return (ON_MI250_MI300 and not ON_NAVI and (sliding_window == 0)
112
+ return (ON_MI250_MI300 and not ON_NAVI
113
+ and (sliding_window == 0 or sliding_window == (-1, -1))
114
and (qtype == torch.half or qtype == torch.bfloat16)
115
and (head_size == 64 or head_size == 128)
116
and (block_size == 16 or block_size == 32)