biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[v1] Add encoder-only/cross attention support to Triton Attention backend (#31406)

Browse Source 
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
This commit is contained in:
Isotr0py
2026-01-06 00:00:23 +08:00
committed by GitHub
parent 911d38ed99
commit 6aa5b18e1d
6 changed files with 627 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

225
tests/kernels/attention/test_triton_prefill_attention.py Normal file
View File

@@ -0,0 +1,225 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.attention.ops.triton_prefill_attention import context_attention_fwd
def ref_masked_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
is_causal: bool = True,
sliding_window_q: int | None = None,
sliding_window_k: int | None = None,
) -> torch.Tensor:
"""Reference implementation using PyTorch SDPA."""
# q, k, v: [total_tokens, num_heads, head_dim]
# SDPA expects [batch, num_heads, seq_len, head_dim]
total_tokens = q.shape[0]
# Add batch dimension and transpose
q = q.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
k = k.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
v = v.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
# Create attention mask if needed
attn_mask = None
use_causal = is_causal
# If we have sliding window or need custom masking, create explicit mask
sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
if (sliding_window_q > 0) or (sliding_window_k > 0):
# Position indices
pos_q = torch.arange(total_tokens, device=q.device).unsqueeze(1)
pos_k = torch.arange(total_tokens, device=q.device).unsqueeze(0)
# Start with valid mask (False = no masking)
mask = torch.ones(
(total_tokens, total_tokens), dtype=torch.bool, device=q.device
)
# Apply causal mask
if is_causal:
mask = mask & (pos_q >= pos_k)
# Apply sliding window masks
sliding_window_mask = torch.ones_like(mask)
if sliding_window_q > 0:
sliding_window_mask &= pos_q - pos_k <= sliding_window_q
if sliding_window_k > 0:
sliding_window_mask &= pos_k - pos_q <= sliding_window_k
mask = mask & sliding_window_mask
attn_mask = torch.where(mask, 0.0, float("-inf")).to(q.dtype)
use_causal = False # Don't use is_causal when providing explicit mask
# Use SDPA
output = F.scaled_dot_product_attention(
q, k, v, attn_mask=attn_mask, is_causal=use_causal, dropout_p=0.0
)
# Convert back to original shape: [total_tokens, num_heads, head_dim]
output = output.transpose(1, 2).squeeze(0)
return output
@pytest.mark.parametrize("B", [5])
@pytest.mark.parametrize("max_seq_len", [1024])
@pytest.mark.parametrize("H_Q", [32])
@pytest.mark.parametrize("H_KV", [32, 8])
@pytest.mark.parametrize("D", [128])
@pytest.mark.parametrize("is_causal", [True, False])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_context_attention(
B: int,
max_seq_len: int,
H_Q: int,
H_KV: int,
D: int,
is_causal: bool,
dtype: torch.dtype,
):
"""Test basic context attention without sliding window."""
torch.manual_seed(42)
# Generate random sequence lengths for each batch
seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
total_tokens = seq_lens.sum().item()
# Create batch start locations
b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
# Create input tensors
q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
o = torch.zeros_like(q)
# Call Triton kernel
context_attention_fwd(
q,
k,
v,
o,
b_start_loc,
seq_lens,
max_seq_len,
is_causal=is_causal,
sliding_window_q=None,
sliding_window_k=None,
)
# Compute reference output for each sequence in batch
o_ref = torch.zeros_like(q)
for i in range(B):
start = b_start_loc[i].item()
end = start + seq_lens[i].item()
q_seq = q[start:end]
k_seq = k[start:end]
v_seq = v[start:end]
# Expand KV heads if using GQA
if H_Q != H_KV:
kv_group_num = H_Q // H_KV
k_seq = k_seq.repeat_interleave(kv_group_num, dim=1)
v_seq = v_seq.repeat_interleave(kv_group_num, dim=1)
o_ref[start:end] = ref_masked_attention(
q_seq,
k_seq,
v_seq,
is_causal=is_causal,
sliding_window_q=None,
sliding_window_k=None,
)
# Compare outputs
torch.testing.assert_close(o, o_ref, rtol=1e-2, atol=1e-2)
@pytest.mark.parametrize("B", [4])
@pytest.mark.parametrize("max_seq_len", [1024])
@pytest.mark.parametrize("H_Q", [32])
@pytest.mark.parametrize("H_KV", [32, 8])
@pytest.mark.parametrize("D", [128])
@pytest.mark.parametrize("sliding_window", [(32, 32), (32, 0), (0, 32)])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_context_attention_sliding_window(
B: int,
max_seq_len: int,
H_Q: int,
H_KV: int,
D: int,
sliding_window: tuple[int, int],
dtype: torch.dtype,
):
sliding_window_q, sliding_window_k = sliding_window
"""Test context attention with sliding window."""
torch.manual_seed(42)
# Generate random sequence lengths for each batch
seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
total_tokens = seq_lens.sum().item()
# Create batch start locations
b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
# Create input tensors
q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
o = torch.zeros_like(q)
# Call Triton kernel
context_attention_fwd(
q,
k,
v,
o,
b_start_loc,
seq_lens,
max_seq_len,
is_causal=False,
sliding_window_q=sliding_window_q,
sliding_window_k=sliding_window_k,
)
# Compute reference output for each sequence in batch
o_ref = torch.zeros_like(q)
for i in range(B):
start = b_start_loc[i].item()
end = start + seq_lens[i].item()
q_seq = q[start:end]
k_seq = k[start:end]
v_seq = v[start:end]
# Expand KV heads if using GQA
if H_Q != H_KV:
kv_group_num = H_Q // H_KV
k_seq = k_seq.repeat_interleave(kv_group_num, dim=1)
v_seq = v_seq.repeat_interleave(kv_group_num, dim=1)
o_ref[start:end] = ref_masked_attention(
q_seq,
k_seq,
v_seq,
is_causal=False,
sliding_window_q=sliding_window_q if sliding_window_q > 0 else None,
sliding_window_k=sliding_window_k if sliding_window_k > 0 else None,
)
# Compare outputs
torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2)

2
tests/models/multimodal/generation/test_whisper.py
View File

@@ -114,7 +114,7 @@ def check_model_available(model: str) -> None:
@pytest.mark.core_model
@pytest.mark.cpu_model
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
@pytest.mark.parametrize("dtype", ["half"])
@pytest.mark.parametrize("dtype", ["half", "float"])
@pytest.mark.parametrize("num_logprobs", [5])
@pytest.mark.parametrize("enforce_eager", [True, False])
@create_new_process_for_each_test("spawn")

57
tests/v1/attention/test_attention_backends.py
View File

@@ -15,6 +15,7 @@ from tests.v1.attention.utils import (
create_vllm_config,
try_get_attention_backend,
)
from vllm.attention.backends.abstract import AttentionType
from vllm.attention.backends.registry import AttentionBackendEnum
from vllm.config import ModelConfig
from vllm.platforms import current_platform
@@ -83,6 +84,13 @@ BATCH_SPECS = {
),
"single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
"single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
# encoder-only
"small_encoder_prefill": BatchSpec(
seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]
),
"medium_encoder_prefill": BatchSpec(
seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]
),
}
@@ -209,6 +217,7 @@ def run_attention_backend(
key: torch.Tensor,
value: torch.Tensor,
kv_cache: torch.Tensor,
attn_type: AttentionType = AttentionType.DECODER,
sliding_window: int | None = None,
) -> torch.Tensor:
"""Run attention computation using the specified backend's AttentionImpl."""
@@ -276,6 +285,7 @@ def run_attention_backend(
num_kv_heads=num_kv_heads,
alibi_slopes=None,
sliding_window=sliding_window,
attn_type=attn_type,
kv_cache_dtype="auto",
)
@@ -299,6 +309,7 @@ def _test_backend_correctness(
backend_to_test: list[AttentionBackendEnum | str],
mask_mod,
*,
attn_type: AttentionType = AttentionType.DECODER,
block_size: int = 16,
atol: float = 1e-2,
rtol: float = 1e-2,
@@ -436,6 +447,9 @@ def _test_backend_correctness(
common_attn_metadata = create_common_attn_metadata(
batch_spec, vllm_config.cache_config.block_size, device
)
if attn_type == AttentionType.ENCODER_ONLY:
# For encoder-only, all tokens are prefill tokens
common_attn_metadata.causal = False
# 3. Simulate Paged KV Cache and a realistic slot_mapping
kv_cache = create_and_prepopulate_kv_cache(
@@ -491,6 +505,7 @@ def _test_backend_correctness(
value_vllm,
kv_cache_for_backend,
sliding_window=sliding_window,
attn_type=attn_type,
)
finally:
if reset_kv_cache_layout:
@@ -676,3 +691,45 @@ def test_sliding_window_backend_correctness(
block_size=128,
tensor_parallel_size=tensor_parallel_size,
)
@pytest.mark.parametrize(
"batch_spec_name",
[
"small_encoder_prefill",
"medium_encoder_prefill",
],
)
@pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 2])
def test_sliding_window_encoder_backend_correctness(
batch_spec_name: str, model: str, tensor_parallel_size: int
):
"""Test backend's correctness with sliding window attention."""
def bidi_sliding_window_mask_mod(
b: torch.Tensor,
h: torch.Tensor,
q_idx: torch.Tensor,
kv_idx: torch.Tensor,
*,
context_len: int,
sliding_window: int,
):
return torch.abs(q_idx + context_len - kv_idx) < sliding_window
batch_spec = BATCH_SPECS[batch_spec_name]
model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
sliding_window = model_config.get_sliding_window()
sliding_window_mask_mod_fn = partial(
bidi_sliding_window_mask_mod, sliding_window=sliding_window
)
_test_backend_correctness(
batch_spec,
model,
SLIDING_WINDOW_BACKENDS_TO_TEST,
sliding_window_mask_mod_fn,
attn_type=AttentionType.ENCODER_ONLY,
tensor_parallel_size=tensor_parallel_size,
)