From 545cd06903569de005f44873f844c3ed080bb34f Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Thu, 10 Apr 2025 14:05:15 -0500 Subject: [PATCH 1/2] enable custom PA for V0 Signed-off-by: Divakar Verma --- vllm/platforms/rocm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index d18b7c26f7ec..29b037d155a1 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -110,7 +110,8 @@ def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int, # rocm custom page attention not support on navi (gfx1*) return (ON_MI250_MI300 and not ON_NAVI - and (sliding_window == 0 or sliding_window == (-1, -1)) + and (not envs.VLLM_USE_V1 or 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) From 8e4add073d46458dd36feace469530b9ed796786 Mon Sep 17 00:00:00 2001 From: Divakar Verma Date: Mon, 14 Apr 2025 10:23:33 -0500 Subject: [PATCH 2/2] add comment for paged attn logic Signed-off-by: Divakar Verma --- vllm/platforms/rocm.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 29b037d155a1..ca6528313a19 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -109,6 +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*) + # custom paged attn always supported on V0. On V1, requires sliding window + # disabled due to observed numerical discrepancy. return (ON_MI250_MI300 and not ON_NAVI and (not envs.VLLM_USE_V1 or sliding_window == 0 or sliding_window == (-1, -1))