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 14, 05:30 UTC
[ROCm] Add attention sink to use_rocm_custom_paged_attention (#22329)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Minseok Lee <47620120+minseokl@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Woosuk Kwon committed
Aug 5, 2025 at 23:30 UTC
98a3a81024649985ed8814a4b7d083d2303fd73c
1 file changed
+6
-5
vllm/platforms/rocm.py
+6
-5
@@ -127,7 +127,8 @@ def use_rocm_custom_paged_attention(
127
max_seq_len: int,
128
sliding_window: int,
129
kv_cache_dtype: str,
130
- alibi_slopes: Optional[torch.Tensor] = None) -> bool:
130
+ alibi_slopes: Optional[torch.Tensor] = None,
131
+ sinks: Optional[torch.Tensor] = None) -> bool:
132
133
GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
134
ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
@@ -145,7 +146,7 @@ def use_rocm_custom_paged_attention(
146
and max_seq_len <= 128 * 1024
147
and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
148
and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
148
- and envs.VLLM_ROCM_USE_AITER))
149
+ and envs.VLLM_ROCM_USE_AITER) and sinks is None)
150
151
else:
152
return (ON_GFX11_GFX12 and (not envs.VLLM_USE_V1 or sliding_window == 0
@@ -155,7 +156,7 @@ def use_rocm_custom_paged_attention(
156
and (gqa_ratio >= 3 and gqa_ratio <= 16)
157
and max_seq_len <= 128 * 1024 and alibi_slopes is None
158
and kv_cache_dtype == "auto"
158
- and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
159
+ and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN and sinks is None)
160
161
162
class RocmPlatform(Platform):
@@ -170,7 +171,7 @@ class RocmPlatform(Platform):
171
172
supported_quantization: list[str] = [
173
"awq", "gptq", "fp8", "compressed-tensors", "fbgemm_fp8", "gguf",
173
- "quark", "ptpc_fp8"
174
+ "quark", "ptpc_fp8", "mxfp4"
175
]
176
177
@classmethod
@@ -469,4 +470,4 @@ class RocmPlatform(Platform):
470
471
@classmethod
472
def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str) -> bool:
472
- return True
\ No newline at end of file
473
+ return True