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, 11:00 UTC
[Kernel] Unified Triton kernel that doesn't distinguish between prefill + decode (#16828)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Thomas Parnell committed
May 6, 2025 at 18:21 UTC
2f925e5777cce9d574292bb6c91ff9f92de3fe62
3 files changed
+566
-27
tests/kernels/test_triton_unified_attention.py
new
+189
@@ -0,0 +1,189 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from typing import Optional
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.attention.ops.triton_unified_attention import unified_attention
9
+from vllm.platforms import current_platform
10
+
11
+NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
12
+HEAD_SIZES = [128, 256]
13
+BLOCK_SIZES = [16, 32]
14
+
15
+DTYPES = [torch.float16, torch.bfloat16]
16
+QDTYPES = [None, torch.float8_e4m3fn]
17
+# one value large enough to test overflow in index calculation.
18
+# one value small enough to test the schema op check
19
+NUM_BLOCKS = [32768, 2048]
20
+
21
+
22
+def ref_paged_attn(
23
+ query: torch.Tensor,
24
+ key_cache: torch.Tensor,
25
+ value_cache: torch.Tensor,
26
+ query_lens: list[int],
27
+ kv_lens: list[int],
28
+ block_tables: torch.Tensor,
29
+ scale: float,
30
+ sliding_window: Optional[int] = None,
31
+ soft_cap: Optional[float] = None,
32
+) -> torch.Tensor:
33
+ num_seqs = len(query_lens)
34
+ block_tables = block_tables.cpu().numpy()
35
+ _, block_size, num_kv_heads, head_size = key_cache.shape
36
+
37
+ outputs: list[torch.Tensor] = []
38
+ start_idx = 0
39
+ for i in range(num_seqs):
40
+ query_len = query_lens[i]
41
+ kv_len = kv_lens[i]
42
+ q = query[start_idx:start_idx + query_len]
43
+ q *= scale
44
+
45
+ num_kv_blocks = (kv_len + block_size - 1) // block_size
46
+ block_indices = block_tables[i, :num_kv_blocks]
47
+
48
+ k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
49
+ k = k[:kv_len]
50
+ v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
51
+ v = v[:kv_len]
52
+
53
+ if q.shape[1] != k.shape[1]:
54
+ k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
55
+ v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
56
+ attn = torch.einsum("qhd,khd->hqk", q, k).float()
57
+ empty_mask = torch.ones(query_len, kv_len)
58
+ mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
59
+ if sliding_window is not None:
60
+ sliding_window_mask = torch.triu(empty_mask,
61
+ diagonal=kv_len -
62
+ (query_len + sliding_window) +
63
+ 1).bool().logical_not()
64
+ mask |= sliding_window_mask
65
+ if soft_cap is not None and soft_cap > 0:
66
+ attn = soft_cap * torch.tanh(attn / soft_cap)
67
+ attn.masked_fill_(mask, float("-inf"))
68
+ attn = torch.softmax(attn, dim=-1).to(v.dtype)
69
+ out = torch.einsum("hqk,khd->qhd", attn, v)
70
+
71
+ outputs.append(out)
72
+ start_idx += query_len
73
+
74
+ return torch.cat(outputs, dim=0)
75
+
76
+
77
+@pytest.mark.parametrize("seq_lens",
78
+ [[(1, 1328), (5, 18),
79
+ (129, 463)], [(1, 523), (1, 37), (1, 2011)]])
80
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
81
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
82
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
83
+@pytest.mark.parametrize("sliding_window", [None, 256])
84
+@pytest.mark.parametrize("dtype", DTYPES)
85
+@pytest.mark.parametrize("soft_cap", [None, 10.0, 50.0])
86
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
87
+@pytest.mark.parametrize("q_dtype", QDTYPES)
88
+@torch.inference_mode()
89
+def test_triton_unified_attn(
90
+ seq_lens: list[tuple[int, int]],
91
+ num_heads: tuple[int, int],
92
+ head_size: int,
93
+ sliding_window: Optional[int],
94
+ dtype: torch.dtype,
95
+ block_size: int,
96
+ soft_cap: Optional[float],
97
+ num_blocks: int,
98
+ q_dtype: Optional[torch.dtype],
99
+) -> None:
100
+ torch.set_default_device("cuda")
101
+
102
+ current_platform.seed_everything(0)
103
+ num_seqs = len(seq_lens)
104
+ query_lens = [x[0] for x in seq_lens]
105
+ kv_lens = [x[1] for x in seq_lens]
106
+ num_query_heads = num_heads[0]
107
+ num_kv_heads = num_heads[1]
108
+ assert num_query_heads % num_kv_heads == 0
109
+ max_query_len = max(query_lens)
110
+ max_kv_len = max(kv_lens)
111
+ window_size = ((sliding_window - 1, 0) if sliding_window is not None else
112
+ (-1, -1))
113
+ scale = head_size**-0.5
114
+
115
+ query = torch.randn(sum(query_lens),
116
+ num_query_heads,
117
+ head_size,
118
+ dtype=dtype)
119
+ key_cache = torch.randn(num_blocks,
120
+ block_size,
121
+ num_kv_heads,
122
+ head_size,
123
+ dtype=dtype)
124
+ value_cache = torch.randn_like(key_cache)
125
+ cu_query_lens = torch.tensor([0] + query_lens,
126
+ dtype=torch.int32).cumsum(dim=0,
127
+ dtype=torch.int32)
128
+ kv_lens = torch.tensor(kv_lens, dtype=torch.int32)
129
+
130
+ max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
131
+ block_tables = torch.randint(0,
132
+ num_blocks,
133
+ (num_seqs, max_num_blocks_per_seq),
134
+ dtype=torch.int32)
135
+
136
+ output = torch.empty_like(query)
137
+
138
+ maybe_quantized_query = query
139
+ maybe_quantized_key_cache = key_cache
140
+ maybe_quantized_value_cache = value_cache
141
+ q_descale = None
142
+ k_descale = None
143
+ v_descale = None
144
+ if q_dtype is not None:
145
+ # QKV are drawn from N(0, 1): no need for a fp8 scaling factor
146
+ maybe_quantized_query = query.to(q_dtype)
147
+ maybe_quantized_key_cache = key_cache.to(q_dtype)
148
+ maybe_quantized_value_cache = value_cache.to(q_dtype)
149
+
150
+ scale_shape = (num_seqs, num_kv_heads)
151
+ q_descale = None # Not yet supported
152
+ k_descale = torch.rand(scale_shape, dtype=torch.float32)
153
+ v_descale = torch.rand(scale_shape, dtype=torch.float32)
154
+
155
+ unified_attention(
156
+ q=maybe_quantized_query,
157
+ k=maybe_quantized_key_cache,
158
+ v=maybe_quantized_value_cache,
159
+ out=output,
160
+ cu_seqlens_q=cu_query_lens,
161
+ seqused_k=kv_lens,
162
+ max_seqlen_q=max_query_len,
163
+ max_seqlen_k=max_kv_len,
164
+ softmax_scale=scale,
165
+ causal=True,
166
+ window_size=window_size,
167
+ block_table=block_tables,
168
+ softcap=soft_cap if soft_cap is not None else 0,
169
+ q_descale=q_descale,
170
+ k_descale=k_descale,
171
+ v_descale=v_descale,
172
+ )
173
+
174
+ ref_output = ref_paged_attn(
175
+ query=query,
176
+ key_cache=key_cache,
177
+ value_cache=value_cache,
178
+ query_lens=query_lens,
179
+ kv_lens=kv_lens,
180
+ block_tables=block_tables,
181
+ scale=scale,
182
+ sliding_window=sliding_window,
183
+ soft_cap=soft_cap,
184
+ )
185
+ atol, rtol = 1.5e-2, 1e-2
186
+ if q_dtype is not None:
187
+ atol, rtol = 1.5e-1, 1.5e-1
188
+ torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), \
189
+ f"{torch.max(torch.abs(output - ref_output))}"
vllm/attention/ops/triton_unified_attention.py
new
+333
@@ -0,0 +1,333 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+# Authors:
4
+# - Burkhard Ringlein <ngl@zurich.ibm.com>
5
+# - Jan van Lunteren <jvl@zurich.ibm.com>
6
+# - Chih-Chieh Yang <chih.chieh.yang@ibm.com>
7
+# - Thomas Parnell <tpa@zurich.ibm.com>
8
+
9
+import triton
10
+import triton.language as tl
11
+
12
+from vllm.logger import init_logger
13
+
14
+logger = init_logger(__name__)
15
+
16
+
17
+@triton.jit
18
+def cdiv_fn(x, y):
19
+ return (x + y - 1) // y
20
+
21
+
22
+@triton.jit
23
+def apply_softcap(S, x):
24
+ Sdiv = S / x
25
+ p1 = tl.exp(Sdiv)
26
+ p2 = tl.exp(-Sdiv)
27
+ return x * (p1 - p2) / (p1 + p2)
28
+
29
+
30
+@triton.jit
31
+def kernel_unified_attention_2d(
32
+ output_ptr, # [num_tokens, num_query_heads, head_size]
33
+ query_ptr, # [num_tokens, num_query_heads, head_size]
34
+ key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]
35
+ value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]
36
+ block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
37
+ seq_lens_ptr, # [num_seqs]
38
+ alibi_slopes_ptr, # [num_query_heads]
39
+ scale, # float32
40
+ k_scale, # float32
41
+ v_scale, # float32
42
+ softcap, # float32
43
+ num_query_heads: tl.constexpr, # int
44
+ num_queries_per_kv: tl.constexpr, # int
45
+ block_table_stride: tl.int64, # int
46
+ query_stride_0: tl.int64, # int
47
+ query_stride_1: tl.int64, # int, should be equal to head_size
48
+ output_stride_0: tl.int64, # int
49
+ output_stride_1: tl.int64, # int, should be equal to head_size
50
+ BLOCK_SIZE: tl.constexpr, # int
51
+ HEAD_SIZE: tl.constexpr, # int
52
+ HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
53
+ USE_ALIBI_SLOPES: tl.constexpr, # bool
54
+ USE_SOFTCAP: tl.constexpr, # bool
55
+ SLIDING_WINDOW: tl.constexpr, # int
56
+ stride_k_cache_0: tl.int64, # int
57
+ stride_k_cache_1: tl.int64, # int
58
+ stride_k_cache_2: tl.int64, # int
59
+ stride_k_cache_3: tl.int64, # int
60
+ stride_v_cache_0: tl.int64, # int
61
+ stride_v_cache_1: tl.int64, # int
62
+ stride_v_cache_2: tl.int64, # int
63
+ stride_v_cache_3: tl.int64, # int
64
+ query_start_len_ptr, # [num_seqs+1]
65
+ BLOCK_Q: tl.constexpr, # int
66
+ num_seqs: tl.int32,
67
+):
68
+
69
+ q_block_global_idx = tl.program_id(0)
70
+ kv_head_idx = tl.program_id(1)
71
+
72
+ left: tl.int32 = 0
73
+ right = num_seqs
74
+ while left < right:
75
+ mid = (left + right) // 2
76
+ mid_val = tl.load(query_start_len_ptr + mid) // BLOCK_Q + mid
77
+ if mid_val <= q_block_global_idx:
78
+ left = mid + 1
79
+ else:
80
+ right = mid
81
+
82
+ seq_idx = left - 1
83
+ q_block_start_idx = tl.load(query_start_len_ptr +
84
+ seq_idx) // BLOCK_Q + seq_idx
85
+
86
+ q_block_local_idx = q_block_global_idx - q_block_start_idx
87
+
88
+ cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
89
+ cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
90
+
91
+ cur_batch_query_len = cur_batch_in_all_stop_index \
92
+ - cur_batch_in_all_start_index
93
+
94
+ if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
95
+ return
96
+
97
+ offs_m = tl.arange(0, BLOCK_Q * num_queries_per_kv)
98
+ offs_d = tl.arange(0, HEAD_SIZE_PADDED)
99
+
100
+ query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
101
+
102
+ query_offset_0 = cur_batch_in_all_start_index + query_pos
103
+ query_offset_1 = kv_head_idx * num_queries_per_kv + \
104
+ offs_m % num_queries_per_kv
105
+
106
+ query_offset = (query_offset_0[:, None] * query_stride_0 +
107
+ query_offset_1[:, None] * query_stride_1 + offs_d[None, :])
108
+
109
+ dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)
110
+ query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
111
+ query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
112
+
113
+ # Q : (BLOCK_Q * num_queries_per_kv, HEAD_SIZE,)
114
+ Q = tl.load(
115
+ query_ptr + query_offset,
116
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
117
+ other=0.0,
118
+ )
119
+
120
+ block_table_offset = seq_idx * block_table_stride
121
+
122
+ M = tl.full([BLOCK_Q * num_queries_per_kv],
123
+ float("-inf"),
124
+ dtype=tl.float32)
125
+ L = tl.full([BLOCK_Q * num_queries_per_kv], 1.0, dtype=tl.float32)
126
+ acc = tl.zeros([BLOCK_Q * num_queries_per_kv, HEAD_SIZE_PADDED],
127
+ dtype=tl.float32)
128
+
129
+ # sequence len for this particular sequence
130
+ seq_len = tl.load(seq_lens_ptr + seq_idx)
131
+
132
+ # context length for this particular sequences
133
+ context_len = seq_len - cur_batch_query_len
134
+
135
+ # alibi slope for this head
136
+ if USE_ALIBI_SLOPES:
137
+ alibi_slope = tl.load(alibi_slopes_ptr + query_offset_1,
138
+ mask=query_mask_1,
139
+ other=0.0)
140
+
141
+ num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)
142
+
143
+ # iterate through tiles
144
+ for j in range(0, num_blocks):
145
+
146
+ physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j)
147
+
148
+ offs_n = tl.arange(0, BLOCK_SIZE)
149
+
150
+ v_offset = (physical_block_idx * stride_v_cache_0 +
151
+ kv_head_idx * stride_v_cache_2 +
152
+ offs_d[None, :] * stride_v_cache_3 +
153
+ offs_n[:, None] * stride_v_cache_1)
154
+
155
+ k_offset = (physical_block_idx * stride_k_cache_0 +
156
+ kv_head_idx * stride_k_cache_2 +
157
+ offs_d[:, None] * stride_k_cache_3 +
158
+ offs_n[None, :] * stride_k_cache_1)
159
+
160
+ # K : (HEAD_SIZE, BLOCK_SIZE)
161
+ K_load = tl.load(key_cache_ptr + k_offset,
162
+ mask=dim_mask[:, None],
163
+ other=0.0)
164
+
165
+ if K_load.dtype.is_fp8():
166
+ if Q.dtype.is_fp8():
167
+ K = K_load
168
+ else:
169
+ K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
170
+ else:
171
+ K = K_load
172
+
173
+ # V : (BLOCK_SIZE, HEAD_SIZE)
174
+ V_load = tl.load(value_cache_ptr + v_offset,
175
+ mask=dim_mask[None, :],
176
+ other=0.0)
177
+
178
+ if V_load.dtype.is_fp8():
179
+ if Q.dtype.is_fp8():
180
+ V = V_load
181
+ else:
182
+ V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
183
+ else:
184
+ V = V_load
185
+
186
+ seq_offset = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
187
+
188
+ seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1
189
+
190
+ # S : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
191
+ S = tl.zeros(shape=(BLOCK_Q * num_queries_per_kv, BLOCK_SIZE),
192
+ dtype=tl.float32)
193
+
194
+ S += scale * tl.dot(Q, K)
195
+
196
+ if USE_SOFTCAP:
197
+ S = apply_softcap(S, softcap)
198
+
199
+ S = tl.where(query_mask_1[:, None] & query_mask_0[:, None] & seq_mask,
200
+ S, float("-inf"))
201
+
202
+ if SLIDING_WINDOW > 0:
203
+ S = tl.where((context_len + query_pos[:, None] - seq_offset)
204
+ < SLIDING_WINDOW, S, float("-inf"))
205
+
206
+ if USE_ALIBI_SLOPES:
207
+ S += alibi_slope[:, None] * (seq_offset - context_len)
208
+
209
+ # compute running maximum
210
+ # m_j : (BLOCK_Q * num_queries_per_kv,)
211
+ m_j = tl.maximum(M, tl.max(S, axis=1))
212
+ # For sliding window there's a chance the max is -inf due to masking of
213
+ # the entire row. In this case we need to set m_j 0 to avoid NaN
214
+ m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
215
+
216
+ # P : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
217
+ P = tl.exp(S - m_j[:, None])
218
+
219
+ # l_j : (BLOCK_Q * num_queries_per_kv,)
220
+ l_j = tl.sum(P, axis=1)
221
+
222
+ # alpha : (BLOCK_Q * num_queries_per_kv, )
223
+ alpha = tl.exp(M - m_j)
224
+
225
+ # acc : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
226
+ acc = acc * alpha[:, None]
227
+
228
+ # update constants
229
+ L = L * alpha + l_j
230
+ M = m_j
231
+
232
+ # acc : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
233
+ acc += tl.dot(P.to(V.dtype), V)
234
+
235
+ # epilogue
236
+ acc = acc / L[:, None]
237
+
238
+ output_offset = (query_offset_0[:, None] * output_stride_0 +
239
+ query_offset_1[:, None] * output_stride_1 +
240
+ offs_d[None, :])
241
+
242
+ tl.store(
243
+ output_ptr + output_offset,
244
+ acc,
245
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
246
+ )
247
+
248
+
249
+def unified_attention(
250
+ q,
251
+ k,
252
+ v,
253
+ out,
254
+ cu_seqlens_q,
255
+ max_seqlen_q,
256
+ seqused_k,
257
+ max_seqlen_k,
258
+ softmax_scale,
259
+ causal,
260
+ window_size,
261
+ block_table,
262
+ softcap,
263
+ q_descale,
264
+ k_descale,
265
+ v_descale,
266
+ alibi_slopes=None,
267
+):
268
+ assert causal, "Only causal attention is supported"
269
+ assert q_descale is None, "Q scales not supported"
270
+
271
+ use_alibi_slopes = alibi_slopes is not None
272
+
273
+ block_size = v.shape[1]
274
+ num_seqs = len(seqused_k)
275
+ num_query_heads = q.shape[1]
276
+ num_kv_heads = k.shape[2]
277
+ num_queries_per_kv = num_query_heads // num_kv_heads
278
+ head_size = q.shape[2]
279
+
280
+ BLOCK_M = 16
281
+ BLOCK_Q = BLOCK_M // num_queries_per_kv
282
+
283
+ # Ideally we would launch with kernel with:
284
+ # \sum_i[ceil(query_len[i] / BLOCK_Q)] blocks.
285
+ # However, it is slow to realize the query_lens on cpu.
286
+ # Instead we use upper-bound:
287
+ # \sum_i[ceil(query_len[i] / BLOCK_Q)]
288
+ # <= \sum_i[floor(query_len[i] / BLOCK_Q) + 1]
289
+ # = \sum_i[floor(query_len[i] / BLOCK_Q)] + num_seqs
290
+ # <= floor(\sum_i(query_len[i]) / BLOCK_Q) + num_seqs
291
+ # = floor(q.shape[0] / BLOCK_Q) + num_seqs
292
+ total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
293
+
294
+ kernel_unified_attention_2d[(
295
+ total_num_q_blocks,
296
+ num_kv_heads,
297
+ )](
298
+ output_ptr=out,
299
+ query_ptr=q,
300
+ key_cache_ptr=k,
301
+ value_cache_ptr=v,
302
+ block_tables_ptr=block_table,
303
+ seq_lens_ptr=seqused_k,
304
+ alibi_slopes_ptr=alibi_slopes,
305
+ scale=softmax_scale,
306
+ k_scale=k_descale,
307
+ v_scale=v_descale,
308
+ softcap=softcap,
309
+ num_query_heads=num_query_heads,
310
+ num_queries_per_kv=num_queries_per_kv,
311
+ block_table_stride=block_table.stride(0),
312
+ query_stride_0=q.stride(0),
313
+ query_stride_1=q.stride(1),
314
+ output_stride_0=out.stride(0),
315
+ output_stride_1=out.stride(1),
316
+ BLOCK_SIZE=block_size,
317
+ HEAD_SIZE=head_size,
318
+ HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
319
+ USE_ALIBI_SLOPES=use_alibi_slopes,
320
+ USE_SOFTCAP=(softcap > 0),
321
+ SLIDING_WINDOW=(1 + window_size[0]),
322
+ stride_k_cache_0=k.stride(0),
323
+ stride_k_cache_1=k.stride(1),
324
+ stride_k_cache_2=k.stride(2),
325
+ stride_k_cache_3=k.stride(3),
326
+ stride_v_cache_0=v.stride(0),
327
+ stride_v_cache_1=v.stride(1),
328
+ stride_v_cache_2=v.stride(2),
329
+ stride_v_cache_3=v.stride(3),
330
+ query_start_len_ptr=cu_seqlens_q,
331
+ BLOCK_Q=BLOCK_Q,
332
+ num_seqs=num_seqs,
333
+ )
vllm/v1/attention/backends/triton_attn.py
+44
-27
@@ -4,11 +4,10 @@ from typing import Any, Optional
4
5
import torch
6
7
+from vllm import _custom_ops as ops
8
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
9
AttentionMetadata, AttentionType)
9
-from vllm.attention.ops.chunked_prefill_paged_decode import (
10
- chunked_prefill_paged_decode)
11
-from vllm.attention.ops.paged_attn import PagedAttention
10
+from vllm.attention.ops.triton_unified_attention import unified_attention
11
from vllm.logger import init_logger
12
from vllm.v1.attention.backends.flash_attn import (
13
FlashAttentionMetadata, FlashAttentionMetadataBuilder)
@@ -87,6 +86,11 @@ class TritonAttentionImpl(AttentionImpl):
86
else:
87
self.sliding_window = (sliding_window - 1, 0)
88
self.kv_cache_dtype = kv_cache_dtype
89
+ if logits_soft_cap is None:
90
+ # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
91
+ logits_soft_cap = 0
92
+ self.logits_soft_cap = logits_soft_cap
93
+
94
self.use_irope = use_irope
95
96
assert self.num_heads % self.num_kv_heads == 0
@@ -143,11 +147,9 @@ class TritonAttentionImpl(AttentionImpl):
147
# performance to make sure it does not introduce any overhead.
148
149
num_actual_tokens = attn_metadata.num_actual_tokens
146
- key_cache, value_cache = PagedAttention.split_kv_cache(
147
- kv_cache, self.num_kv_heads, self.head_size)
150
149
- # Reshape the input keys and values and store them in the cache.
150
- PagedAttention.write_to_paged_cache(
151
+ key_cache, value_cache = kv_cache.unbind(0)
152
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
153
key,
154
value,
155
key_cache,
@@ -158,6 +160,18 @@ class TritonAttentionImpl(AttentionImpl):
160
layer._v_scale,
161
)
162
163
+ if self.kv_cache_dtype.startswith("fp8"):
164
+ key_cache = key_cache.view(torch.float8_e4m3fn)
165
+ value_cache = value_cache.view(torch.float8_e4m3fn)
166
+ num_tokens, num_heads, head_size = query.shape
167
+ assert layer._q_scale == 1.0, \
168
+ "A non 1.0 q_scale is not currently supported."
169
+ query, _ = ops.scaled_fp8_quant(
170
+ query.reshape(
171
+ (num_tokens, num_heads * head_size)).contiguous(),
172
+ layer._q_scale)
173
+ query = query.reshape((num_tokens, num_heads, head_size))
174
+
175
use_local_attn = \
176
(self.use_irope and attn_metadata.local_attn_metadata is not None)
177
@@ -165,34 +179,37 @@ class TritonAttentionImpl(AttentionImpl):
179
assert attn_metadata.local_attn_metadata is not None
180
local_metadata = attn_metadata.local_attn_metadata
181
cu_seqlens_q = local_metadata.local_query_start_loc
168
- sequesd_k = local_metadata.local_seqused_k
182
+ seqused_k = local_metadata.local_seqused_k
183
max_seqlen_q = local_metadata.local_max_query_len
184
max_seqlen_k = local_metadata.local_max_seq_len
185
block_table = local_metadata.local_block_table
186
else:
187
cu_seqlens_q = attn_metadata.query_start_loc
174
- sequesd_k = attn_metadata.seq_lens
188
+ seqused_k = attn_metadata.seq_lens
189
max_seqlen_q = attn_metadata.max_query_len
190
max_seqlen_k = attn_metadata.max_seq_len
191
block_table = attn_metadata.block_table
192
179
- # Compute attention and update output up to `num_actual_tokens`.
180
- chunked_prefill_paged_decode(query=query[:num_actual_tokens],
181
- key=key[:num_actual_tokens],
182
- value=value[:num_actual_tokens],
183
- output=output[:num_actual_tokens],
184
- kv_cache_dtype=self.kv_cache_dtype,
185
- key_cache=key_cache,
186
- value_cache=value_cache,
187
- block_table=block_table,
188
- query_start_loc=cu_seqlens_q,
189
- seq_lens=sequesd_k,
190
- max_seq_len=max_seqlen_k,
191
- max_query_len=max_seqlen_q,
192
- k_scale=layer._k_scale,
193
- v_scale=layer._v_scale,
194
- alibi_slopes=self.alibi_slopes,
195
- sliding_window=self.sliding_window[0],
196
- sm_scale=self.scale)
193
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
194
+
195
+ unified_attention(
196
+ q=query[:num_actual_tokens],
197
+ k=key_cache,
198
+ v=value_cache,
199
+ out=output[:num_actual_tokens],
200
+ cu_seqlens_q=cu_seqlens_q,
201
+ max_seqlen_q=max_seqlen_q,
202
+ seqused_k=seqused_k,
203
+ max_seqlen_k=max_seqlen_k,
204
+ softmax_scale=self.scale,
205
+ causal=True,
206
+ alibi_slopes=self.alibi_slopes,
207
+ window_size=self.sliding_window,
208
+ block_table=block_table,
209
+ softcap=self.logits_soft_cap,
210
+ q_descale=None, # Not supported
211
+ k_descale=layer._k_scale.expand(descale_shape),
212
+ v_descale=layer._v_scale.expand(descale_shape),
213
+ )
214
215
return output