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 10, 19:00 UTC
[v1] Add encoder-only/cross attention support to Triton Attention backend (#31406)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jan 6, 2026 at 00:00 UTC
6aa5b18e1dfd9c678fcdd12ca8b954343eb373eb
6 files changed
+627
-14
tests/kernels/attention/test_triton_prefill_attention.py
new
+225
@@ -0,0 +1,225 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+from vllm.attention.ops.triton_prefill_attention import context_attention_fwd
9
+
10
+
11
+def ref_masked_attention(
12
+ q: torch.Tensor,
13
+ k: torch.Tensor,
14
+ v: torch.Tensor,
15
+ is_causal: bool = True,
16
+ sliding_window_q: int | None = None,
17
+ sliding_window_k: int | None = None,
18
+) -> torch.Tensor:
19
+ """Reference implementation using PyTorch SDPA."""
20
+ # q, k, v: [total_tokens, num_heads, head_dim]
21
+ # SDPA expects [batch, num_heads, seq_len, head_dim]
22
+
23
+ total_tokens = q.shape[0]
24
+
25
+ # Add batch dimension and transpose
26
+ q = q.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
27
+ k = k.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
28
+ v = v.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim]
29
+
30
+ # Create attention mask if needed
31
+ attn_mask = None
32
+ use_causal = is_causal
33
+
34
+ # If we have sliding window or need custom masking, create explicit mask
35
+ sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
36
+ sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
37
+ if (sliding_window_q > 0) or (sliding_window_k > 0):
38
+ # Position indices
39
+ pos_q = torch.arange(total_tokens, device=q.device).unsqueeze(1)
40
+ pos_k = torch.arange(total_tokens, device=q.device).unsqueeze(0)
41
+
42
+ # Start with valid mask (False = no masking)
43
+ mask = torch.ones(
44
+ (total_tokens, total_tokens), dtype=torch.bool, device=q.device
45
+ )
46
+
47
+ # Apply causal mask
48
+ if is_causal:
49
+ mask = mask & (pos_q >= pos_k)
50
+
51
+ # Apply sliding window masks
52
+ sliding_window_mask = torch.ones_like(mask)
53
+ if sliding_window_q > 0:
54
+ sliding_window_mask &= pos_q - pos_k <= sliding_window_q
55
+
56
+ if sliding_window_k > 0:
57
+ sliding_window_mask &= pos_k - pos_q <= sliding_window_k
58
+
59
+ mask = mask & sliding_window_mask
60
+
61
+ attn_mask = torch.where(mask, 0.0, float("-inf")).to(q.dtype)
62
+ use_causal = False # Don't use is_causal when providing explicit mask
63
+
64
+ # Use SDPA
65
+ output = F.scaled_dot_product_attention(
66
+ q, k, v, attn_mask=attn_mask, is_causal=use_causal, dropout_p=0.0
67
+ )
68
+
69
+ # Convert back to original shape: [total_tokens, num_heads, head_dim]
70
+ output = output.transpose(1, 2).squeeze(0)
71
+
72
+ return output
73
+
74
+
75
+@pytest.mark.parametrize("B", [5])
76
+@pytest.mark.parametrize("max_seq_len", [1024])
77
+@pytest.mark.parametrize("H_Q", [32])
78
+@pytest.mark.parametrize("H_KV", [32, 8])
79
+@pytest.mark.parametrize("D", [128])
80
+@pytest.mark.parametrize("is_causal", [True, False])
81
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
82
+def test_context_attention(
83
+ B: int,
84
+ max_seq_len: int,
85
+ H_Q: int,
86
+ H_KV: int,
87
+ D: int,
88
+ is_causal: bool,
89
+ dtype: torch.dtype,
90
+):
91
+ """Test basic context attention without sliding window."""
92
+ torch.manual_seed(42)
93
+
94
+ # Generate random sequence lengths for each batch
95
+ seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
96
+ total_tokens = seq_lens.sum().item()
97
+
98
+ # Create batch start locations
99
+ b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
100
+ b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
101
+
102
+ # Create input tensors
103
+ q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
104
+ k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
105
+ v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
106
+ o = torch.zeros_like(q)
107
+
108
+ # Call Triton kernel
109
+ context_attention_fwd(
110
+ q,
111
+ k,
112
+ v,
113
+ o,
114
+ b_start_loc,
115
+ seq_lens,
116
+ max_seq_len,
117
+ is_causal=is_causal,
118
+ sliding_window_q=None,
119
+ sliding_window_k=None,
120
+ )
121
+
122
+ # Compute reference output for each sequence in batch
123
+ o_ref = torch.zeros_like(q)
124
+ for i in range(B):
125
+ start = b_start_loc[i].item()
126
+ end = start + seq_lens[i].item()
127
+
128
+ q_seq = q[start:end]
129
+ k_seq = k[start:end]
130
+ v_seq = v[start:end]
131
+
132
+ # Expand KV heads if using GQA
133
+ if H_Q != H_KV:
134
+ kv_group_num = H_Q // H_KV
135
+ k_seq = k_seq.repeat_interleave(kv_group_num, dim=1)
136
+ v_seq = v_seq.repeat_interleave(kv_group_num, dim=1)
137
+
138
+ o_ref[start:end] = ref_masked_attention(
139
+ q_seq,
140
+ k_seq,
141
+ v_seq,
142
+ is_causal=is_causal,
143
+ sliding_window_q=None,
144
+ sliding_window_k=None,
145
+ )
146
+
147
+ # Compare outputs
148
+ torch.testing.assert_close(o, o_ref, rtol=1e-2, atol=1e-2)
149
+
150
+
151
+@pytest.mark.parametrize("B", [4])
152
+@pytest.mark.parametrize("max_seq_len", [1024])
153
+@pytest.mark.parametrize("H_Q", [32])
154
+@pytest.mark.parametrize("H_KV", [32, 8])
155
+@pytest.mark.parametrize("D", [128])
156
+@pytest.mark.parametrize("sliding_window", [(32, 32), (32, 0), (0, 32)])
157
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
158
+def test_context_attention_sliding_window(
159
+ B: int,
160
+ max_seq_len: int,
161
+ H_Q: int,
162
+ H_KV: int,
163
+ D: int,
164
+ sliding_window: tuple[int, int],
165
+ dtype: torch.dtype,
166
+):
167
+ sliding_window_q, sliding_window_k = sliding_window
168
+ """Test context attention with sliding window."""
169
+ torch.manual_seed(42)
170
+
171
+ # Generate random sequence lengths for each batch
172
+ seq_lens = torch.randint(max_seq_len // 2, max_seq_len + 1, (B,), device="cuda")
173
+ total_tokens = seq_lens.sum().item()
174
+
175
+ # Create batch start locations
176
+ b_start_loc = torch.zeros(B, dtype=torch.int32, device="cuda")
177
+ b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
178
+
179
+ # Create input tensors
180
+ q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device="cuda")
181
+ k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
182
+ v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
183
+ o = torch.zeros_like(q)
184
+
185
+ # Call Triton kernel
186
+ context_attention_fwd(
187
+ q,
188
+ k,
189
+ v,
190
+ o,
191
+ b_start_loc,
192
+ seq_lens,
193
+ max_seq_len,
194
+ is_causal=False,
195
+ sliding_window_q=sliding_window_q,
196
+ sliding_window_k=sliding_window_k,
197
+ )
198
+
199
+ # Compute reference output for each sequence in batch
200
+ o_ref = torch.zeros_like(q)
201
+ for i in range(B):
202
+ start = b_start_loc[i].item()
203
+ end = start + seq_lens[i].item()
204
+
205
+ q_seq = q[start:end]
206
+ k_seq = k[start:end]
207
+ v_seq = v[start:end]
208
+
209
+ # Expand KV heads if using GQA
210
+ if H_Q != H_KV:
211
+ kv_group_num = H_Q // H_KV
212
+ k_seq = k_seq.repeat_interleave(kv_group_num, dim=1)
213
+ v_seq = v_seq.repeat_interleave(kv_group_num, dim=1)
214
+
215
+ o_ref[start:end] = ref_masked_attention(
216
+ q_seq,
217
+ k_seq,
218
+ v_seq,
219
+ is_causal=False,
220
+ sliding_window_q=sliding_window_q if sliding_window_q > 0 else None,
221
+ sliding_window_k=sliding_window_k if sliding_window_k > 0 else None,
222
+ )
223
+
224
+ # Compare outputs
225
+ torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2)
tests/models/multimodal/generation/test_whisper.py
+1
-1
@@ -114,7 +114,7 @@ def check_model_available(model: str) -> None:
114
@pytest.mark.core_model
115
@pytest.mark.cpu_model
116
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
117
-@pytest.mark.parametrize("dtype", ["half"])
117
+@pytest.mark.parametrize("dtype", ["half", "float"])
118
@pytest.mark.parametrize("num_logprobs", [5])
119
@pytest.mark.parametrize("enforce_eager", [True, False])
120
@create_new_process_for_each_test("spawn")
tests/v1/attention/test_attention_backends.py
+57
@@ -15,6 +15,7 @@ from tests.v1.attention.utils import (
15
create_vllm_config,
16
try_get_attention_backend,
17
)
18
+from vllm.attention.backends.abstract import AttentionType
19
from vllm.attention.backends.registry import AttentionBackendEnum
20
from vllm.config import ModelConfig
21
from vllm.platforms import current_platform
@@ -83,6 +84,13 @@ BATCH_SPECS = {
84
),
85
"single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
86
"single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
87
+ # encoder-only
88
+ "small_encoder_prefill": BatchSpec(
89
+ seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]
90
+ ),
91
+ "medium_encoder_prefill": BatchSpec(
92
+ seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]
93
+ ),
94
}
95
96
@@ -209,6 +217,7 @@ def run_attention_backend(
217
key: torch.Tensor,
218
value: torch.Tensor,
219
kv_cache: torch.Tensor,
220
+ attn_type: AttentionType = AttentionType.DECODER,
221
sliding_window: int | None = None,
222
) -> torch.Tensor:
223
"""Run attention computation using the specified backend's AttentionImpl."""
@@ -276,6 +285,7 @@ def run_attention_backend(
285
num_kv_heads=num_kv_heads,
286
alibi_slopes=None,
287
sliding_window=sliding_window,
288
+ attn_type=attn_type,
289
kv_cache_dtype="auto",
290
)
291
@@ -299,6 +309,7 @@ def _test_backend_correctness(
309
backend_to_test: list[AttentionBackendEnum | str],
310
mask_mod,
311
*,
312
+ attn_type: AttentionType = AttentionType.DECODER,
313
block_size: int = 16,
314
atol: float = 1e-2,
315
rtol: float = 1e-2,
@@ -436,6 +447,9 @@ def _test_backend_correctness(
447
common_attn_metadata = create_common_attn_metadata(
448
batch_spec, vllm_config.cache_config.block_size, device
449
)
450
+ if attn_type == AttentionType.ENCODER_ONLY:
451
+ # For encoder-only, all tokens are prefill tokens
452
+ common_attn_metadata.causal = False
453
454
# 3. Simulate Paged KV Cache and a realistic slot_mapping
455
kv_cache = create_and_prepopulate_kv_cache(
@@ -491,6 +505,7 @@ def _test_backend_correctness(
505
value_vllm,
506
kv_cache_for_backend,
507
sliding_window=sliding_window,
508
+ attn_type=attn_type,
509
)
510
finally:
511
if reset_kv_cache_layout:
@@ -676,3 +691,45 @@ def test_sliding_window_backend_correctness(
691
block_size=128,
692
tensor_parallel_size=tensor_parallel_size,
693
)
694
+
695
+
696
+@pytest.mark.parametrize(
697
+ "batch_spec_name",
698
+ [
699
+ "small_encoder_prefill",
700
+ "medium_encoder_prefill",
701
+ ],
702
+)
703
+@pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
704
+@pytest.mark.parametrize("tensor_parallel_size", [1, 2])
705
+def test_sliding_window_encoder_backend_correctness(
706
+ batch_spec_name: str, model: str, tensor_parallel_size: int
707
+):
708
+ """Test backend's correctness with sliding window attention."""
709
+
710
+ def bidi_sliding_window_mask_mod(
711
+ b: torch.Tensor,
712
+ h: torch.Tensor,
713
+ q_idx: torch.Tensor,
714
+ kv_idx: torch.Tensor,
715
+ *,
716
+ context_len: int,
717
+ sliding_window: int,
718
+ ):
719
+ return torch.abs(q_idx + context_len - kv_idx) < sliding_window
720
+
721
+ batch_spec = BATCH_SPECS[batch_spec_name]
722
+ model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
723
+ sliding_window = model_config.get_sliding_window()
724
+ sliding_window_mask_mod_fn = partial(
725
+ bidi_sliding_window_mask_mod, sliding_window=sliding_window
726
+ )
727
+
728
+ _test_backend_correctness(
729
+ batch_spec,
730
+ model,
731
+ SLIDING_WINDOW_BACKENDS_TO_TEST,
732
+ sliding_window_mask_mod_fn,
733
+ attn_type=AttentionType.ENCODER_ONLY,
734
+ tensor_parallel_size=tensor_parallel_size,
735
+ )
vllm/attention/ops/triton_prefill_attention.py
new
+271
@@ -0,0 +1,271 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Adapted from
5
+# https://github.com/sgl-project/sglang/blob/97cb762bb65ebf05025eb342de03c184660427a3/python/sglang/srt/layers/attention/triton_ops/prefill_attention.py
6
+# Changes:
7
+# - Add support for sliding window attention
8
+
9
+# Copyright 2023-2024 SGLang Team
10
+# Licensed under the Apache License, Version 2.0 (the "License");
11
+# you may not use this file except in compliance with the License.
12
+# You may obtain a copy of the License at
13
+#
14
+# http://www.apache.org/licenses/LICENSE-2.0
15
+#
16
+# Unless required by applicable law or agreed to in writing, software
17
+# distributed under the License is distributed on an "AS IS" BASIS,
18
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+# See the License for the specific language governing permissions and
20
+# limitations under the License.
21
+# ==============================================================================
22
+"""
23
+Memory-efficient attention for prefill.
24
+It supports page size = 1.
25
+"""
26
+
27
+# Adapted from
28
+# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1
29
+import torch
30
+
31
+from vllm.platforms import current_platform
32
+from vllm.triton_utils import tl, triton
33
+
34
+
35
+@triton.jit
36
+def _fwd_kernel(
37
+ Q,
38
+ K,
39
+ V,
40
+ sm_scale,
41
+ B_Start_Loc,
42
+ B_Seqlen,
43
+ Out,
44
+ stride_qbs,
45
+ stride_qh,
46
+ stride_kbs,
47
+ stride_kh,
48
+ stride_vbs,
49
+ stride_vh,
50
+ stride_obs,
51
+ stride_oh,
52
+ kv_group_num: tl.constexpr,
53
+ BLOCK_M: tl.constexpr,
54
+ BLOCK_DMODEL: tl.constexpr,
55
+ BLOCK_N: tl.constexpr,
56
+ IS_CAUSAL: tl.constexpr,
57
+ SLIDING_WINDOW_Q: tl.constexpr,
58
+ SLIDING_WINDOW_K: tl.constexpr,
59
+ Lk: tl.constexpr,
60
+):
61
+ cur_batch = tl.program_id(0)
62
+ cur_head = tl.program_id(1)
63
+ start_m = tl.program_id(2)
64
+
65
+ cur_kv_head = cur_head // kv_group_num
66
+
67
+ cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
68
+ cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
69
+
70
+ block_start_loc = BLOCK_M * start_m
71
+
72
+ # initialize offsets
73
+ offs_n = tl.arange(0, BLOCK_N)
74
+ offs_d = tl.arange(0, BLOCK_DMODEL)
75
+ offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
76
+ off_q = (
77
+ (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
78
+ + cur_head * stride_qh
79
+ + offs_d[None, :]
80
+ )
81
+ off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
82
+ off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :]
83
+
84
+ mask_d = offs_d < Lk
85
+
86
+ q = tl.load(
87
+ Q + off_q,
88
+ mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :]),
89
+ other=0.0,
90
+ )
91
+
92
+ k_ptrs = K + off_k
93
+ v_ptrs = V + off_v
94
+
95
+ # initialize pointer to m and l
96
+ m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
97
+ l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
98
+ acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
99
+
100
+ block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
101
+
102
+ # Calculate the end position for attention computation
103
+ end_n = cur_batch_seq_len
104
+
105
+ # Apply causal attention pruning and sliding window attention pruning
106
+ end_n = tl.minimum(end_n, (start_m + 1) * BLOCK_M) if IS_CAUSAL else end_n
107
+
108
+ # Calculate the start position for backward sliding window
109
+ start_n_limit = 0
110
+ end_n_limit = block_mask * end_n
111
+
112
+ for start_n in range(start_n_limit, end_n_limit, BLOCK_N):
113
+ start_n = tl.multiple_of(start_n, BLOCK_N)
114
+ # -- compute qk ----
115
+ k = tl.load(
116
+ k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
117
+ mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]),
118
+ other=0.0,
119
+ )
120
+
121
+ # Apply attention mask (causal + bidirectional sliding window)
122
+ # Position indices in the sequence
123
+ pos_q = offs_m[:, None] # Query positions [BLOCK_M, 1]
124
+ pos_k = start_n + offs_n[None, :] # Key positions [1, BLOCK_N]
125
+
126
+ # Valid sequence mask
127
+ mask = pos_k < cur_batch_seq_len
128
+ # Causal mask
129
+ if IS_CAUSAL:
130
+ mask &= pos_q >= pos_k
131
+
132
+ # Bidirectional sliding window masks
133
+ sliding_mask_q = (
134
+ pos_q - pos_k <= SLIDING_WINDOW_Q if SLIDING_WINDOW_Q > 0 else None
135
+ )
136
+ sliding_mask_k = (
137
+ pos_k - pos_q <= SLIDING_WINDOW_K if SLIDING_WINDOW_K > 0 else None
138
+ )
139
+ if sliding_mask_q is not None:
140
+ mask &= sliding_mask_q
141
+ if sliding_mask_k is not None:
142
+ mask &= sliding_mask_k
143
+
144
+ qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
145
+ qk += tl.where(mask, 0, float("-inf"))
146
+ qk += tl.dot(q, k)
147
+ qk *= sm_scale
148
+
149
+ # -- compute m_ij, p, l_ij
150
+ m_ij = tl.max(qk, 1)
151
+ # For sliding window there's a chance the max is -inf due to masking of
152
+ # the entire row. In this case we need to set m_j 0 to avoid NaN
153
+ m_ij_valid_mask = m_ij > float("-inf")
154
+ m_ij_masked = tl.where(m_ij_valid_mask, m_ij, 0.0)
155
+ # -- compute p and l_ij --
156
+ p = tl.exp(qk - m_ij_masked[:, None])
157
+ l_ij = tl.sum(p, 1)
158
+ # -- update m_i and l_i
159
+ m_i_new = tl.maximum(m_i, m_ij)
160
+ m_i_new_mask = m_i_new > float("-inf")
161
+ alpha = tl.exp(m_i - m_i_new)
162
+ beta = tl.exp(m_ij - m_i_new)
163
+ # mask alpha and beta for sliding window
164
+ alpha = tl.where(m_i_new_mask, alpha, 1.0)
165
+ beta = tl.where(m_i_new_mask, beta, 0.0)
166
+ l_i_new = alpha * l_i + beta * l_ij
167
+ # -- update output accumulator --
168
+ # scale p
169
+ # For sliding window there's a chance the l_i_new is 0 due to masking
170
+ # the entire row. We need to set l_i_new 1 to avoid zero division
171
+ l_i_new_mask = (l_i_new != 0.0) & (m_i_new_mask > float("-inf"))
172
+ l_i_new_safe = tl.where(l_i_new_mask, l_i_new, 1.0)
173
+ p_scale = beta / l_i_new_safe
174
+ p = p * p_scale[:, None]
175
+ # scale acc
176
+ acc_scale = l_i / l_i_new_safe * alpha
177
+ acc = acc * acc_scale[:, None]
178
+ # update acc
179
+ v = tl.load(
180
+ v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
181
+ mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]),
182
+ other=0.0,
183
+ )
184
+
185
+ p = p.to(v.dtype)
186
+ acc += tl.dot(p, v)
187
+ # update m_i and l_i
188
+ l_i = l_i_new
189
+ m_i = m_i_new
190
+ # initialize pointers to output
191
+ off_o = (
192
+ (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
193
+ + cur_head * stride_oh
194
+ + offs_d[None, :]
195
+ )
196
+ out_ptrs = Out + off_o
197
+ tl.store(
198
+ out_ptrs, acc, mask=(offs_m[:, None] < cur_batch_seq_len) & (mask_d[None, :])
199
+ )
200
+
201
+
202
+def get_block_size(dtype: torch.dtype) -> int:
203
+ if dtype == torch.float32:
204
+ return 32
205
+ elif (
206
+ current_platform.is_cuda_alike()
207
+ ) and current_platform.get_device_capability().major > 8:
208
+ return 128
209
+ else:
210
+ return 64
211
+
212
+
213
+def context_attention_fwd(
214
+ q,
215
+ k,
216
+ v,
217
+ o,
218
+ b_start_loc,
219
+ b_seq_len,
220
+ max_input_len,
221
+ is_causal=True,
222
+ sliding_window_q=None,
223
+ sliding_window_k=None,
224
+):
225
+ """
226
+ q, k, v: [b * s, head, head_dim]
227
+ b_start_loc: [b]
228
+ b_seq_len: [b]
229
+ out: [b * s, head, head_dim]
230
+ """
231
+ BLOCK = get_block_size(q.dtype)
232
+
233
+ Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1]
234
+
235
+ sm_scale = 1.0 / (Lq**0.5)
236
+ batch, head = b_seq_len.shape[0], q.shape[1]
237
+ kv_group_num = q.shape[1] // k.shape[1]
238
+
239
+ grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
240
+ num_warps = 4 if Lk <= 64 else 8
241
+
242
+ sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
243
+ sliding_window_k = sliding_window_k if sliding_window_k is not None else 0
244
+
245
+ _fwd_kernel[grid](
246
+ q,
247
+ k,
248
+ v,
249
+ sm_scale,
250
+ b_start_loc,
251
+ b_seq_len,
252
+ o,
253
+ q.stride(0),
254
+ q.stride(1),
255
+ k.stride(0),
256
+ k.stride(1),
257
+ v.stride(0),
258
+ v.stride(1),
259
+ o.stride(0),
260
+ o.stride(1),
261
+ kv_group_num=kv_group_num,
262
+ BLOCK_M=BLOCK,
263
+ BLOCK_DMODEL=triton.next_power_of_2(Lk),
264
+ BLOCK_N=BLOCK,
265
+ IS_CAUSAL=is_causal,
266
+ SLIDING_WINDOW_Q=sliding_window_q,
267
+ SLIDING_WINDOW_K=sliding_window_k,
268
+ num_warps=num_warps,
269
+ num_stages=1,
270
+ Lk=Lk,
271
+ )
vllm/platforms/rocm.py
-9
@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Optional
8
import torch
9
10
import vllm.envs as envs
11
-from vllm.attention.backends.abstract import AttentionType
11
from vllm.attention.backends.registry import AttentionBackendEnum
12
from vllm.logger import init_logger
13
from vllm.utils.torch_utils import cuda_device_count_stateless
@@ -289,14 +288,6 @@ class RocmPlatform(Platform):
288
logger.info("Using Aiter Flash Attention backend.")
289
return AttentionBackendEnum.ROCM_AITER_FA.get_path()
290
292
- # Priority 5: If model is Encoder-only self-attention type
293
- if (
294
- attn_selector_config.attn_type is not None
295
- and attn_selector_config.attn_type == AttentionType.ENCODER_ONLY
296
- ):
297
- logger.info("Using FlexAttention backend.")
298
- return AttentionBackendEnum.FLEX_ATTENTION.get_path()
299
-
291
# Default: Triton Unified Attention
292
logger.info("Using Triton Attention backend.")
293
return AttentionBackendEnum.TRITON_ATTN.get_path()
vllm/v1/attention/backends/triton_attn.py
+73
-4
@@ -13,6 +13,7 @@ from vllm.attention.backends.abstract import (
13
AttentionType,
14
MultipleOf,
15
)
16
+from vllm.attention.ops.triton_prefill_attention import context_attention_fwd
17
from vllm.attention.ops.triton_reshape_and_cache_flash import (
18
triton_reshape_and_cache_flash,
19
)
@@ -309,6 +310,16 @@ class TritonAttentionBackend(AttentionBackend):
310
def supports_sink(cls) -> bool:
311
return True
312
313
+ @classmethod
314
+ def supports_attn_type(cls, attn_type: str) -> bool:
315
+ """TritonAttention supports all attention types."""
316
+ return attn_type in (
317
+ AttentionType.DECODER,
318
+ AttentionType.ENCODER,
319
+ AttentionType.ENCODER_ONLY,
320
+ AttentionType.ENCODER_DECODER,
321
+ )
322
+
323
@classmethod
324
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
325
return True
@@ -341,6 +352,8 @@ class TritonAttentionImpl(AttentionImpl):
352
self.alibi_slopes = alibi_slopes
353
if sliding_window is None:
354
self.sliding_window = (-1, -1)
355
+ elif attn_type in (AttentionType.ENCODER, AttentionType.ENCODER_ONLY):
356
+ self.sliding_window = (sliding_window - 1, sliding_window - 1)
357
else:
358
self.sliding_window = (sliding_window - 1, 0)
359
self.kv_cache_dtype = kv_cache_dtype
@@ -352,10 +365,6 @@ class TritonAttentionImpl(AttentionImpl):
365
366
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
367
355
- if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
356
- raise NotImplementedError(
357
- "Encoder self-attention is not implemented for TritonAttentionImpl"
358
- )
368
self.attn_type = attn_type
369
self.fp8_dtype = current_platform.fp8_dtype()
370
@@ -417,6 +426,21 @@ class TritonAttentionImpl(AttentionImpl):
426
# performance to make sure it does not introduce any overhead.
427
428
num_actual_tokens = attn_metadata.num_actual_tokens
429
+
430
+ # Handle encoder attention differently - no KV cache needed
431
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
432
+ # For encoder attention,
433
+ # we use direct Q, K, V tensors without caching
434
+ return self._forward_encoder_attention(
435
+ query[:num_actual_tokens],
436
+ key[:num_actual_tokens],
437
+ value[:num_actual_tokens],
438
+ output[:num_actual_tokens],
439
+ attn_metadata,
440
+ layer,
441
+ )
442
+
443
+ # For decoder and cross-attention, use KV cache as before
444
key_cache, value_cache = kv_cache.unbind(1)
445
446
if (
@@ -495,3 +519,48 @@ class TritonAttentionImpl(AttentionImpl):
519
)
520
521
return output
522
+
523
+ def _forward_encoder_attention(
524
+ self,
525
+ query: torch.Tensor,
526
+ key: torch.Tensor,
527
+ value: torch.Tensor,
528
+ output: torch.Tensor,
529
+ attn_metadata: TritonAttentionMetadata,
530
+ layer: torch.nn.Module,
531
+ ) -> torch.Tensor:
532
+ """Forward pass for encoder attention without KV cache.
533
+
534
+ Args:
535
+ query: shape = [num_encoder_tokens, num_heads, head_size]
536
+ key: shape = [num_encoder_tokens, num_kv_heads, head_size]
537
+ value: shape = [num_encoder_tokens, num_kv_heads, head_size]
538
+ output: shape = [num_encoder_tokens, num_heads, head_size]
539
+ attn_metadata: Encoder attention metadata
540
+ layer: The attention layer
541
+ """
542
+ # For encoder attention, process FP8 quantization if needed
543
+ if self.kv_cache_dtype.startswith("fp8"):
544
+ raise NotImplementedError(
545
+ "quantization is not supported for encoder attention"
546
+ )
547
+
548
+ # Use encoder-specific metadata for sequence information
549
+ query_start_loc = attn_metadata.query_start_loc
550
+ seq_lens = attn_metadata.seq_lens
551
+ max_query_len = attn_metadata.max_query_len
552
+
553
+ # Call flash attention directly on Q, K, V tensors
554
+ context_attention_fwd(
555
+ q=query,
556
+ k=key,
557
+ v=value,
558
+ o=output,
559
+ b_start_loc=query_start_loc,
560
+ b_seq_len=seq_lens,
561
+ max_input_len=max_query_len,
562
+ is_causal=False, # Encoder attention is bidirectional
563
+ sliding_window_q=self.sliding_window[0],
564
+ sliding_window_k=self.sliding_window[1],
565
+ )
566
+ return output