Skip to content

[Attention][V1] Toggle for v1 attention backend #18275

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 29, 2025
Merged

[Attention][V1] Toggle for v1 attention backend #18275

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ac2b381
Toggle for v1 attention
gshtras May 16, 2025
9540f72
Caching the env variable in the __init__
gshtras May 19, 2025
dc0f0d1
Merge remote-tracking branch 'origin/main' into attention_toggle_upst…
gshtras May 19, 2025
f82da97
Better naming and logic extracted to a variable
gshtras May 22, 2025
3b264d5
Merge remote-tracking branch 'origin/main' into attention_toggle_upst…
gshtras May 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions vllm/attention/ops/chunked_prefill_paged_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,8 +264,8 @@ def chunked_prefill_paged_decode(
# Conversion of FP8 Tensor from uint8 storage to
# appropriate torch.dtype for interpretation by Triton
if "fp8" in kv_cache_dtype:
assert key_cache.dtype == torch.uint8
assert value_cache.dtype == torch.uint8
assert key_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
assert value_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]

if kv_cache_dtype in ("fp8", "fp8_e4m3"):
target_dtype = current_platform.fp8_dtype()
Expand Down
4 changes: 2 additions & 2 deletions vllm/attention/ops/prefix_prefill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -744,8 +744,8 @@ def context_attention_fwd(q,
# Conversion of FP8 Tensor from uint8 storage to
# appropriate torch.dtype for interpretation by Triton
if "fp8" in kv_cache_dtype:
assert (k_cache.dtype == torch.uint8)
assert (v_cache.dtype == torch.uint8)
assert k_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]
assert v_cache.dtype in [torch.uint8, current_platform.fp8_dtype()]

if kv_cache_dtype in ("fp8", "fp8_e4m3"):
target_dtype = current_platform.fp8_dtype()
Expand Down
12 changes: 10 additions & 2 deletions vllm/envs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
VLLM_NCCL_SO_PATH: Optional[str] = None
LD_LIBRARY_PATH: Optional[str] = None
VLLM_USE_TRITON_FLASH_ATTN: bool = False
VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False
VLLM_FLASH_ATTN_VERSION: Optional[int] = None
LOCAL_RANK: int = 0
CUDA_VISIBLE_DEVICES: Optional[str] = None
Expand Down Expand Up @@ -290,6 +291,13 @@ def get_vllm_port() -> Optional[int]:
lambda: (os.environ.get("VLLM_USE_TRITON_FLASH_ATTN", "True").lower() in
("true", "1")),

# Use separate prefill and decode kernels for V1 attention instead of
# the unified triton kernel.
"VLLM_V1_USE_PREFILL_DECODE_ATTENTION":
Copy link
Collaborator

@ProExpertProg ProExpertProg May 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: VLLM_V1_TRITON_ATTN_FORCE_PREFILL_DECODE sounds slightly more accurate to me, but feel free to use the name that works best in your eyes

lambda:
(os.getenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "False").lower() in
("true", "1")),

# Force vllm to use a specific flash-attention version (2 or 3), only valid
# when using the flash-attention backend.
"VLLM_FLASH_ATTN_VERSION":
Expand Down Expand Up @@ -323,8 +331,8 @@ def get_vllm_port() -> Optional[int]:

# Whether to log responses from API Server for debugging
"VLLM_DEBUG_LOG_API_SERVER_RESPONSE":
lambda: os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False").
lower() == "true",
lambda: os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False"
Copy link
Collaborator

@ProExpertProg ProExpertProg May 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: accidental change?

Copy link
Collaborator Author

@gshtras gshtras May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ruff or yapf change, reformatting the whole file now results in this

).lower() == "true",

# S3 access information, used for tensorizer to load model from S3
"S3_ACCESS_KEY_ID":
Expand Down
9 changes: 6 additions & 3 deletions vllm/v1/attention/backends/triton_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import torch

from vllm import _custom_ops as ops
from vllm import envs
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
AttentionMetadata, AttentionType)
from vllm.attention.ops.chunked_prefill_paged_decode import (
Expand Down Expand Up @@ -126,6 +127,8 @@ def __init__(
"TritonAttentionImpl")

self.fp8_dtype = current_platform.fp8_dtype()
self.force_prefill_decode_attn = \
envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION

def forward(
self,
Expand Down Expand Up @@ -166,9 +169,9 @@ def forward(
# performance to make sure it does not introduce any overhead.

num_queries_per_kv = query.shape[1] // key.shape[1]
use_prefill_decode_attn = (num_queries_per_kv &
(num_queries_per_kv - 1)) != 0

num_q_is_pow2 = (num_queries_per_kv & (num_queries_per_kv - 1)) == 0
use_prefill_decode_attn = (self.force_prefill_decode_attn
or not num_q_is_pow2)
num_actual_tokens = attn_metadata.num_actual_tokens

if use_prefill_decode_attn:
Expand Down