[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>
This commit is contained in:
Aleksandr Malyshev 2025-04-02 22:28:44 -07:00 committed by GitHub
parent 8b664706aa
commit 57a810db9c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -109,7 +109,8 @@ def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
ON_MI250_MI300 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942"])
# rocm custom page attention not support on navi (gfx1*)
return (ON_MI250_MI300 and not ON_NAVI and (sliding_window == 0)
return (ON_MI250_MI300 and not ON_NAVI
and (sliding_window == 0 or sliding_window == (-1, -1))
and (qtype == torch.half or qtype == torch.bfloat16)
and (head_size == 64 or head_size == 128)
and (block_size == 16 or block_size == 32)