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, 08:30 UTC
[Attention][V1] Toggle for v1 attention backend (#18275)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
May 29, 2025 at 10:48 UTC
da4b69d0b435e74c0a208cffbf2f430f5529ec64
4 files changed
+20
-9
vllm/attention/ops/chunked_prefill_paged_decode.py
+2
-2
@@ -264,8 +264,8 @@ def chunked_prefill_paged_decode(
264
# Conversion of FP8 Tensor from uint8 storage to
265
# appropriate torch.dtype for interpretation by Triton
266
if "fp8" in kv_cache_dtype:
267
- assert key_cache.dtype == torch.uint8
268
- assert value_cache.dtype == torch.uint8
267
+ assert key_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
268
+ assert value_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
269
270
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
271
target_dtype = current_platform.fp8_dtype()
vllm/attention/ops/prefix_prefill.py
+2
-2
@@ -744,8 +744,8 @@ def context_attention_fwd(q,
744
# Conversion of FP8 Tensor from uint8 storage to
745
# appropriate torch.dtype for interpretation by Triton
746
if "fp8" in kv_cache_dtype:
747
- assert (k_cache.dtype == torch.uint8)
748
- assert (v_cache.dtype == torch.uint8)
747
+ assert k_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
748
+ assert v_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
749
750
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
751
target_dtype = current_platform.fp8_dtype()
vllm/envs.py
+10
-2
@@ -15,6 +15,7 @@ if TYPE_CHECKING:
15
VLLM_NCCL_SO_PATH: Optional[str] = None
16
LD_LIBRARY_PATH: Optional[str] = None
17
VLLM_USE_TRITON_FLASH_ATTN: bool = False
18
+ VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False
19
VLLM_FLASH_ATTN_VERSION: Optional[int] = None
20
LOCAL_RANK: int = 0
21
CUDA_VISIBLE_DEVICES: Optional[str] = None
@@ -290,6 +291,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
291
lambda: (os.environ.get("VLLM_USE_TRITON_FLASH_ATTN", "True").lower() in
292
("true", "1")),
293
294
+ # Use separate prefill and decode kernels for V1 attention instead of
295
+ # the unified triton kernel.
296
+ "VLLM_V1_USE_PREFILL_DECODE_ATTENTION":
297
+ lambda:
298
+ (os.getenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "False").lower() in
299
+ ("true", "1")),
300
+
301
# Force vllm to use a specific flash-attention version (2 or 3), only valid
302
# when using the flash-attention backend.
303
"VLLM_FLASH_ATTN_VERSION":
@@ -323,8 +331,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
331
332
# Whether to log responses from API Server for debugging
333
"VLLM_DEBUG_LOG_API_SERVER_RESPONSE":
326
- lambda: os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False").
327
- lower() == "true",
334
+ lambda: os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False"
335
+ ).lower() == "true",
336
337
# S3 access information, used for tensorizer to load model from S3
338
"S3_ACCESS_KEY_ID":
vllm/v1/attention/backends/triton_attn.py
+6
-3
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Optional
5
import torch
6
7
from vllm import _custom_ops as ops
8
+from vllm import envs
9
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
10
AttentionMetadata, AttentionType)
11
from vllm.attention.ops.chunked_prefill_paged_decode import (
@@ -126,6 +127,8 @@ class TritonAttentionImpl(AttentionImpl):
127
"TritonAttentionImpl")
128
129
self.fp8_dtype = current_platform.fp8_dtype()
130
+ self.force_prefill_decode_attn = \
131
+ envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
132
133
def forward(
134
self,
@@ -166,9 +169,9 @@ class TritonAttentionImpl(AttentionImpl):
169
# performance to make sure it does not introduce any overhead.
170
171
num_queries_per_kv = query.shape[1] // key.shape[1]
169
- use_prefill_decode_attn = (num_queries_per_kv &
170
- (num_queries_per_kv - 1)) != 0
171
-
172
+ num_q_is_pow2 = (num_queries_per_kv & (num_queries_per_kv - 1)) == 0
173
+ use_prefill_decode_attn = (self.force_prefill_decode_attn
174
+ or not num_q_is_pow2)
175
num_actual_tokens = attn_metadata.num_actual_tokens
176
177
if use_prefill_decode_attn: