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, 14:45 UTC
[rocm][V0] fix selection logic for custom PA in V0 (#16426)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Apr 16, 2025 at 21:52 UTC
95aca283b456c3aae09a37a7769cac07681fc585
1 file changed
+4
-1
vllm/platforms/rocm.py
+4
-1
@@ -109,8 +109,11 @@ 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
+ # custom paged attn always supported on V0. On V1, requires sliding window
113
+ # disabled due to observed numerical discrepancy.
114
return (ON_MI250_MI300 and not ON_NAVI
113
- and (sliding_window == 0 or sliding_window == (-1, -1))
115
+ and (not envs.VLLM_USE_V1 or sliding_window == 0
116
+ or sliding_window == (-1, -1))
117
and (qtype == torch.half or qtype == torch.bfloat16)
118
and (head_size == 64 or head_size == 128)
119
and (block_size == 16 or block_size == 32)