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, 17:15 UTC
[ROCm] Env variable to trigger custom PA (#15557)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Mar 27, 2025 at 01:46 UTC
ecff8309a3ca5159ac09ac9a7976516b9301f64d
2 files changed
+8
-1
vllm/attention/backends/rocm_flash_attn.py
+2
-1
@@ -908,4 +908,5 @@ def _use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
908
and (qtype == torch.half or qtype == torch.bfloat16)
909
and (head_size == 64 or head_size == 128)
910
and (block_size == 16 or block_size == 32)
911
- and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768)
911
+ and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768
912
+ and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
vllm/envs.py
+6
@@ -78,6 +78,7 @@ if TYPE_CHECKING:
78
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
79
VLLM_ROCM_FP8_PADDING: bool = True
80
VLLM_ROCM_MOE_PADDING: bool = True
81
+ VLLM_ROCM_CUSTOM_PAGED_ATTN: bool = True
82
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
83
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
84
VLLM_DISABLE_COMPILE_CACHE: bool = False
@@ -541,6 +542,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
542
"VLLM_ROCM_MOE_PADDING":
543
lambda: bool(int(os.getenv("VLLM_ROCM_MOE_PADDING", "1"))),
544
545
+ # custom paged attention kernel for MI3* cards
546
+ "VLLM_ROCM_CUSTOM_PAGED_ATTN":
547
+ lambda: (os.getenv("VLLM_ROCM_CUSTOM_PAGED_ATTN", "True").lower() in
548
+ ("true", "1")),
549
+
550
# Divisor for dynamic query scale factor calculation for FP8 KV Cache
551
"Q_SCALE_CONSTANT":
552
lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")),