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, 21:00 UTC
[Kernel] [V1] Improved performance for V1 Triton (ROCm) backend (#14152)
Thomas Parnell committed
Mar 6, 2025 at 16:39 UTC
6bd1dd9d2625ac0ac81e8b5c740aa855399990b9
4 files changed
+398
-77
tests/kernels/test_prefix_prefill.py
+76
-59
@@ -3,6 +3,7 @@
3
import math
4
import random
5
import time
6
+from collections.abc import Callable
7
8
import pytest
9
import torch
@@ -10,6 +11,8 @@ from xformers import ops as xops
11
from xformers.ops.fmha.attn_bias import BlockDiagonalCausalFromBottomRightMask
12
13
from vllm.attention.backends.xformers import _make_alibi_bias
14
+from vllm.attention.ops.chunked_prefill_paged_decode import (
15
+ chunked_prefill_paged_decode)
16
from vllm.attention.ops.prefix_prefill import context_attention_fwd
17
from vllm.platforms import current_platform
18
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
@@ -24,6 +27,8 @@ CUDA_DEVICES = [
27
SLIDING_WINDOW = [0, 16, 64, 128, 256, 512, 2048]
28
KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
29
30
+OPS = [chunked_prefill_paged_decode, context_attention_fwd]
31
+
32
33
@pytest.mark.parametrize("num_heads", NUM_HEADS)
34
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
@@ -32,6 +37,7 @@ KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
37
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
38
@pytest.mark.parametrize("device", CUDA_DEVICES)
39
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
40
+@pytest.mark.parametrize("op", OPS)
41
@torch.inference_mode()
42
def test_contexted_kv_attention(
43
num_heads: int,
@@ -41,6 +47,7 @@ def test_contexted_kv_attention(
47
dtype: torch.dtype,
48
kv_cache_dtype: str,
49
device: str,
50
+ op: Callable,
51
) -> None:
52
53
if 'fp8' in kv_cache_dtype and not current_platform.has_device_capability(
@@ -65,6 +72,9 @@ def test_contexted_kv_attention(
72
block_size = 32
73
max_block_per_request = 64
74
query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
75
+ # ensure one sequence in batch is a decode
76
+ query_lens[-1] = 1
77
+
78
ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
79
seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
80
num_kv_heads = num_heads // num_queries_per_kv
@@ -144,36 +154,36 @@ def test_contexted_kv_attention(
154
155
# Warm up the Triton kernel by calling it once before actually measuring
156
# generation time
147
- context_attention_fwd(query,
148
- k,
149
- v,
150
- output,
151
- kv_cache_dtype,
152
- k_cache,
153
- v_cache,
154
- block_table,
155
- b_start_loc,
156
- b_seq_len,
157
- max_input_len,
158
- k_scale,
159
- v_scale,
160
- sliding_window=sliding_window)
157
+ op(query,
158
+ k,
159
+ v,
160
+ output,
161
+ kv_cache_dtype,
162
+ k_cache,
163
+ v_cache,
164
+ block_table,
165
+ b_start_loc,
166
+ b_seq_len,
167
+ max_input_len,
168
+ k_scale,
169
+ v_scale,
170
+ sliding_window=sliding_window)
171
torch.cuda.synchronize()
172
start_time = time.time()
163
- context_attention_fwd(query,
164
- k,
165
- v,
166
- output,
167
- kv_cache_dtype,
168
- k_cache,
169
- v_cache,
170
- block_table,
171
- b_start_loc,
172
- b_seq_len,
173
- max_input_len,
174
- k_scale,
175
- v_scale,
176
- sliding_window=sliding_window)
173
+ op(query,
174
+ k,
175
+ v,
176
+ output,
177
+ kv_cache_dtype,
178
+ k_cache,
179
+ v_cache,
180
+ block_table,
181
+ b_start_loc,
182
+ b_seq_len,
183
+ max_input_len,
184
+ k_scale,
185
+ v_scale,
186
+ sliding_window=sliding_window)
187
torch.cuda.synchronize()
188
end_time = time.time()
189
print(f"triton Time: {(end_time - start_time)*1000:.2f} ms")
@@ -228,7 +238,7 @@ def test_contexted_kv_attention(
238
end_time = time.time()
239
print(f"xformers Time: {(end_time - start_time)*1000:.2f} ms")
240
output_ref = output_ref.reshape(output.shape)
231
- atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-6
241
+ atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4
242
torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
243
244
@@ -238,6 +248,7 @@ def test_contexted_kv_attention(
248
@pytest.mark.parametrize("dtype", DTYPES)
249
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
250
@pytest.mark.parametrize("device", CUDA_DEVICES)
251
+@pytest.mark.parametrize("op", OPS)
252
@torch.inference_mode()
253
def test_contexted_kv_attention_alibi(
254
num_heads: int,
@@ -246,6 +257,7 @@ def test_contexted_kv_attention_alibi(
257
dtype: torch.dtype,
258
kv_cache_dtype: str,
259
device: str,
260
+ op: Callable,
261
) -> None:
262
263
if 'fp8' in kv_cache_dtype and not current_platform.has_device_capability(
@@ -375,36 +387,36 @@ def test_contexted_kv_attention_alibi(
387
388
# Warm up the Triton kernel by calling it once before actually measuring
389
# generation time
378
- context_attention_fwd(query,
379
- k,
380
- v,
381
- output,
382
- kv_cache_dtype,
383
- k_cache,
384
- v_cache,
385
- block_table,
386
- b_start_loc,
387
- b_seq_len,
388
- max_input_len,
389
- k_scale,
390
- v_scale,
391
- alibi_slopes=alibi_slopes)
390
+ op(query,
391
+ k,
392
+ v,
393
+ output,
394
+ kv_cache_dtype,
395
+ k_cache,
396
+ v_cache,
397
+ block_table,
398
+ b_start_loc,
399
+ b_seq_len,
400
+ max_input_len,
401
+ k_scale,
402
+ v_scale,
403
+ alibi_slopes=alibi_slopes)
404
torch.cuda.synchronize()
405
start_time = time.time()
394
- context_attention_fwd(query,
395
- k,
396
- v,
397
- output,
398
- kv_cache_dtype,
399
- k_cache,
400
- v_cache,
401
- block_table,
402
- b_start_loc,
403
- b_seq_len,
404
- max_input_len,
405
- k_scale,
406
- v_scale,
407
- alibi_slopes=alibi_slopes)
406
+ op(query,
407
+ k,
408
+ v,
409
+ output,
410
+ kv_cache_dtype,
411
+ k_cache,
412
+ v_cache,
413
+ block_table,
414
+ b_start_loc,
415
+ b_seq_len,
416
+ max_input_len,
417
+ k_scale,
418
+ v_scale,
419
+ alibi_slopes=alibi_slopes)
420
torch.cuda.synchronize()
421
end_time = time.time()
422
print(f"triton Time: {(end_time - start_time)*1000:.2f} ms")
@@ -503,6 +515,7 @@ def test_contexted_kv_attention_alibi(
515
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
516
@pytest.mark.parametrize("device", CUDA_DEVICES)
517
@pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
518
+@pytest.mark.parametrize("op", OPS)
519
@torch.inference_mode()
520
def test_contexted_kv_attention_f32(
521
num_heads: int,
@@ -512,9 +525,11 @@ def test_contexted_kv_attention_f32(
525
dtype: torch.dtype,
526
kv_cache_dtype: str,
527
device: str,
528
+ op: Callable,
529
) -> None:
530
test_contexted_kv_attention(num_heads, num_queries_per_kv, head_size,
517
- sliding_window, dtype, kv_cache_dtype, device)
531
+ sliding_window, dtype, kv_cache_dtype, device,
532
+ op)
533
534
535
@pytest.mark.optional
@@ -524,6 +539,7 @@ def test_contexted_kv_attention_f32(
539
@pytest.mark.parametrize("dtype", [torch.float32])
540
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
541
@pytest.mark.parametrize("device", CUDA_DEVICES)
542
+@pytest.mark.parametrize("op", OPS)
543
@torch.inference_mode()
544
def test_contexted_kv_attention_alibi_f32(
545
num_heads: int,
@@ -532,6 +548,7 @@ def test_contexted_kv_attention_alibi_f32(
548
dtype: torch.dtype,
549
kv_cache_dtype: str,
550
device: str,
551
+ op: Callable,
552
) -> None:
553
test_contexted_kv_attention_alibi(num_heads, num_queries_per_kv, head_size,
537
- dtype, kv_cache_dtype, device)
554
+ dtype, kv_cache_dtype, device, op)
vllm/attention/ops/chunked_prefill_paged_decode.py
new
+289
@@ -0,0 +1,289 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import torch
4
+import triton
5
+import triton.language as tl
6
+
7
+from .prefix_prefill import context_attention_fwd
8
+
9
+
10
+@triton.jit
11
+def cdiv_fn(x, y):
12
+ return (x + y - 1) // y
13
+
14
+
15
+@triton.jit
16
+def kernel_paged_attention_2d(
17
+ output_ptr, # [num_tokens, num_query_heads, head_size]
18
+ query_ptr, # [num_tokens, num_query_heads, head_size]
19
+ key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]
20
+ value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]
21
+ block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
22
+ seq_lens_ptr, # [num_seqs]
23
+ alibi_slopes_ptr, # [num_query_heads]
24
+ scale, # float32
25
+ k_scale, # float32
26
+ v_scale, # float32
27
+ num_query_heads: tl.constexpr, # int
28
+ num_queries_per_kv: tl.constexpr, # int
29
+ block_table_stride: tl.constexpr, # int
30
+ query_stride_0: tl.constexpr, # int
31
+ query_stride_1: tl.constexpr, # int, should be equal to head_size
32
+ output_stride_0: tl.constexpr, # int
33
+ output_stride_1: tl.constexpr, # int, should be equal to head_size
34
+ BLOCK_SIZE: tl.constexpr, # int
35
+ HEAD_SIZE: tl.constexpr, # int
36
+ HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
37
+ USE_ALIBI_SLOPES: tl.constexpr, # bool
38
+ SLIDING_WINDOW: tl.constexpr, # int
39
+ x: tl.constexpr, # int
40
+ stride_k_cache_0: tl.constexpr, # int
41
+ stride_k_cache_1: tl.constexpr, # int
42
+ stride_k_cache_2: tl.constexpr, # int
43
+ stride_k_cache_3: tl.constexpr, # int
44
+ stride_k_cache_4: tl.constexpr, # int
45
+ stride_v_cache_0: tl.constexpr, # int
46
+ stride_v_cache_1: tl.constexpr, # int
47
+ stride_v_cache_2: tl.constexpr, # int
48
+ stride_v_cache_3: tl.constexpr, # int
49
+ filter_by_query_len: tl.constexpr, # bool
50
+ query_start_len_ptr, # [num_seqs+1]
51
+):
52
+ seq_idx = tl.program_id(0)
53
+ query_head_idx = tl.program_id(1)
54
+ kv_head_idx = query_head_idx // num_queries_per_kv
55
+
56
+ if filter_by_query_len:
57
+ cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
58
+ cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx +
59
+ 1)
60
+ cur_batch_query_len = cur_batch_in_all_stop_index \
61
+ - cur_batch_in_all_start_index
62
+ if cur_batch_query_len > 1:
63
+ return
64
+ else:
65
+ cur_batch_in_all_start_index = seq_idx
66
+
67
+ query_offset = (cur_batch_in_all_start_index * query_stride_0 +
68
+ query_head_idx * query_stride_1)
69
+
70
+ dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1,
71
+ 0).to(tl.int1)
72
+
73
+ # Q : (HEAD_SIZE,)
74
+ Q = tl.load(
75
+ query_ptr + query_offset + tl.arange(0, HEAD_SIZE_PADDED),
76
+ mask=dim_mask,
77
+ other=0.0,
78
+ )
79
+
80
+ block_table_offset = seq_idx * block_table_stride
81
+
82
+ M = tl.full([1], float("-inf"), dtype=tl.float32)
83
+ L = tl.full([1], 1.0, dtype=tl.float32)
84
+ acc = tl.zeros([HEAD_SIZE_PADDED], dtype=tl.float32)
85
+
86
+ # sequence len for this particular sequence
87
+ seq_len = tl.load(seq_lens_ptr + seq_idx)
88
+
89
+ # alibi slope for this head
90
+ if USE_ALIBI_SLOPES:
91
+ alibi_slope = tl.load(alibi_slopes_ptr + query_head_idx)
92
+
93
+ num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)
94
+
95
+ # iterate through tiles
96
+ for j in range(0, num_blocks):
97
+
98
+ physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j)
99
+
100
+ offs_n = tl.arange(0, BLOCK_SIZE)
101
+ offs_d = tl.arange(0, HEAD_SIZE_PADDED)
102
+
103
+ v_offset = (physical_block_idx * stride_v_cache_0 +
104
+ kv_head_idx * stride_v_cache_1 +
105
+ offs_d[:, None] * stride_v_cache_2 +
106
+ offs_n[None, :] * stride_v_cache_3)
107
+
108
+ k_offset = (physical_block_idx * stride_k_cache_0 +
109
+ kv_head_idx * stride_k_cache_1 +
110
+ (offs_d[:, None] // x) * stride_k_cache_2 +
111
+ offs_n[None, :] * stride_k_cache_3 +
112
+ (offs_d[:, None] % x) * stride_k_cache_4)
113
+
114
+ # K : (HEAD_SIZE, BLOCK_SIZE)
115
+ K_load = tl.load(key_cache_ptr + k_offset,
116
+ mask=dim_mask[:, None],
117
+ other=0.0)
118
+
119
+ if K_load.dtype.is_fp8():
120
+ K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
121
+ else:
122
+ K = K_load
123
+
124
+ # V : (HEAD_SIZE, BLOCK_SIZE)
125
+ V_load = tl.load(value_cache_ptr + v_offset,
126
+ mask=dim_mask[:, None],
127
+ other=0.0)
128
+
129
+ if V_load.dtype.is_fp8():
130
+ V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
131
+ else:
132
+ V = V_load
133
+
134
+ tmp = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
135
+ boundary = tl.full([BLOCK_SIZE], seq_len, dtype=tl.int32)
136
+ mask_new = tmp < boundary
137
+ # S : (BLOCK_SIZE,)
138
+ S = tl.where(mask_new, 0.0, float("-inf")).to(tl.float32)
139
+ S += scale * tl.sum(K * Q[:, None], axis=0)
140
+
141
+ if SLIDING_WINDOW > 0:
142
+ S = tl.where((seq_len - 1 - tmp) < SLIDING_WINDOW, S, -10000)
143
+
144
+ if USE_ALIBI_SLOPES:
145
+ S += alibi_slope * (tmp - seq_len + 1)
146
+
147
+ # compute running maximum
148
+ # m_j : (1,)
149
+ m_j = tl.maximum(M, tl.max(S, axis=0))
150
+
151
+ # P : (BLOCK_SIZE,)
152
+ P = tl.exp(S - m_j)
153
+
154
+ # l_j : (1,)
155
+ l_j = tl.sum(P, axis=0)
156
+
157
+ # alpha : (1, )
158
+ alpha = tl.exp(M - m_j)
159
+
160
+ # acc : (BLOCK_SIZE,)
161
+ acc = acc * alpha
162
+
163
+ # update constants
164
+ L = L * alpha + l_j
165
+ M = m_j
166
+
167
+ # acc : (BLOCK_SIZE,)
168
+ acc += tl.sum(V * P[None, :], axis=1)
169
+
170
+ # epilogue
171
+ acc = acc / L
172
+
173
+ output_offset = (cur_batch_in_all_start_index * output_stride_0 +
174
+ query_head_idx * output_stride_1)
175
+
176
+ tl.store(output_ptr + output_offset + tl.arange(0, HEAD_SIZE_PADDED),
177
+ acc,
178
+ mask=dim_mask)
179
+
180
+
181
+def chunked_prefill_paged_decode(
182
+ query,
183
+ key,
184
+ value,
185
+ output,
186
+ kv_cache_dtype,
187
+ key_cache,
188
+ value_cache,
189
+ block_table,
190
+ query_start_loc,
191
+ seq_lens,
192
+ max_query_len,
193
+ k_scale,
194
+ v_scale,
195
+ alibi_slopes=None,
196
+ sliding_window=None,
197
+ sm_scale=None,
198
+):
199
+
200
+ if sm_scale is None:
201
+ sm_scale = 1.0 / (query.shape[1]**0.5)
202
+
203
+ use_alibi_slopes = alibi_slopes is not None
204
+
205
+ if sliding_window is None or sliding_window <= 0:
206
+ sliding_window = 0
207
+
208
+ if max_query_len > 1:
209
+ context_attention_fwd(
210
+ q=query,
211
+ k=key,
212
+ v=value,
213
+ o=output,
214
+ kv_cache_dtype=kv_cache_dtype,
215
+ k_cache=key_cache,
216
+ v_cache=value_cache,
217
+ b_loc=block_table,
218
+ b_start_loc=query_start_loc,
219
+ b_seq_len=seq_lens,
220
+ max_input_len=max_query_len,
221
+ k_scale=k_scale,
222
+ v_scale=v_scale,
223
+ alibi_slopes=alibi_slopes,
224
+ sliding_window=sliding_window,
225
+ sm_scale=sm_scale,
226
+ skip_decode=True,
227
+ )
228
+
229
+ block_size = value_cache.shape[3]
230
+ num_seqs = len(seq_lens)
231
+ num_query_heads = query.shape[1]
232
+ num_queries_per_kv = query.shape[1] // key.shape[1]
233
+ head_size = query.shape[2]
234
+
235
+ # Conversion of FP8 Tensor from uint8 storage to
236
+ # appropriate torch.dtype for interpretation by Triton
237
+ if "fp8" in kv_cache_dtype:
238
+ assert key_cache.dtype == torch.uint8
239
+ assert value_cache.dtype == torch.uint8
240
+
241
+ if kv_cache_dtype in ("fp8", "fp8_e4m3"):
242
+ target_dtype = torch.float8_e4m3fn
243
+ elif kv_cache_dtype == "fp8_e5m2":
244
+ target_dtype = torch.float8_e5m2
245
+ else:
246
+ raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
247
+
248
+ key_cache = key_cache.view(target_dtype)
249
+ value_cache = value_cache.view(target_dtype)
250
+
251
+ kernel_paged_attention_2d[(
252
+ num_seqs,
253
+ num_query_heads,
254
+ )](
255
+ output_ptr=output,
256
+ query_ptr=query,
257
+ key_cache_ptr=key_cache,
258
+ value_cache_ptr=value_cache,
259
+ block_tables_ptr=block_table,
260
+ seq_lens_ptr=seq_lens,
261
+ alibi_slopes_ptr=alibi_slopes,
262
+ scale=sm_scale,
263
+ k_scale=k_scale,
264
+ v_scale=v_scale,
265
+ num_query_heads=num_query_heads,
266
+ num_queries_per_kv=num_queries_per_kv,
267
+ block_table_stride=block_table.stride(0),
268
+ query_stride_0=query.stride(0),
269
+ query_stride_1=query.stride(1),
270
+ output_stride_0=output.stride(0),
271
+ output_stride_1=output.stride(1),
272
+ BLOCK_SIZE=block_size,
273
+ HEAD_SIZE=head_size,
274
+ HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
275
+ USE_ALIBI_SLOPES=use_alibi_slopes,
276
+ SLIDING_WINDOW=sliding_window,
277
+ x=key_cache.shape[4],
278
+ stride_k_cache_0=key_cache.stride(0),
279
+ stride_k_cache_1=key_cache.stride(1),
280
+ stride_k_cache_2=key_cache.stride(2),
281
+ stride_k_cache_3=key_cache.stride(3),
282
+ stride_k_cache_4=key_cache.stride(4),
283
+ stride_v_cache_0=value_cache.stride(0),
284
+ stride_v_cache_1=value_cache.stride(1),
285
+ stride_v_cache_2=value_cache.stride(2),
286
+ stride_v_cache_3=value_cache.stride(3),
287
+ filter_by_query_len=True,
288
+ query_start_len_ptr=query_start_loc,
289
+ )
vllm/attention/ops/prefix_prefill.py
+13
-1
@@ -64,7 +64,9 @@ if triton.__version__ >= "2.1.0":
64
BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
65
BLOCK_N: tl.constexpr,
66
SLIDING_WINDOW: tl.constexpr,
67
+ SKIP_DECODE: tl.constexpr,
68
):
69
+
70
cur_batch = tl.program_id(0)
71
cur_head = tl.program_id(1)
72
start_m = tl.program_id(2)
@@ -78,6 +80,9 @@ if triton.__version__ >= "2.1.0":
80
cur_batch_in_all_start_index)
81
cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
82
83
+ if SKIP_DECODE and cur_batch_query_len == 1:
84
+ return
85
+
86
# start position inside of the query
87
# generally, N goes over kv, while M goes over query_len
88
block_start_loc = BLOCK_M * start_m
@@ -500,6 +505,7 @@ if triton.__version__ >= "2.1.0":
505
BLOCK_DMODEL: tl.constexpr, # head size
506
BLOCK_DMODEL_PADDED: tl.constexpr, # head size padded to a power of 2
507
BLOCK_N: tl.constexpr,
508
+ SKIP_DECODE: tl.constexpr,
509
):
510
# attn_bias[]
511
cur_batch = tl.program_id(0)
@@ -518,6 +524,9 @@ if triton.__version__ >= "2.1.0":
524
cur_batch_in_all_start_index)
525
cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
526
527
+ if SKIP_DECODE and cur_batch_query_len == 1:
528
+ return
529
+
530
block_start_loc = BLOCK_M * start_m
531
532
# initialize offsets
@@ -721,7 +730,8 @@ if triton.__version__ >= "2.1.0":
730
v_scale: torch.Tensor,
731
alibi_slopes=None,
732
sliding_window=None,
724
- sm_scale=None):
733
+ sm_scale=None,
734
+ skip_decode=False):
735
736
q_dtype_is_f32 = q.dtype is torch.float32
737
# need to reduce num. blocks when using fp32
@@ -823,6 +833,7 @@ if triton.__version__ >= "2.1.0":
833
BLOCK_DMODEL=Lk,
834
BLOCK_DMODEL_PADDED=Lk_padded,
835
BLOCK_N=BLOCK,
836
+ SKIP_DECODE=skip_decode,
837
num_warps=NUM_WARPS,
838
num_stages=1,
839
)
@@ -875,6 +886,7 @@ if triton.__version__ >= "2.1.0":
886
BLOCK_DMODEL_PADDED=Lk_padded,
887
BLOCK_N=BLOCK,
888
SLIDING_WINDOW=sliding_window,
889
+ SKIP_DECODE=skip_decode,
890
num_warps=NUM_WARPS,
891
num_stages=1,
892
)
vllm/v1/attention/backends/rocm_attn.py
+20
-17
@@ -6,8 +6,9 @@ import torch
6
7
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
8
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.prefix_prefill import context_attention_fwd
12
from vllm.logger import init_logger
13
from vllm.v1.attention.backends.flash_attn import (
14
FlashAttentionMetadata, FlashAttentionMetadataBuilder)
@@ -156,20 +157,22 @@ class ROCmAttentionImpl(AttentionImpl):
157
)
158
159
# Compute attention and update output up to `num_actual_tokens`.
159
- context_attention_fwd(q=query[:num_actual_tokens],
160
- k=key[:num_actual_tokens],
161
- v=value[:num_actual_tokens],
162
- o=output[:num_actual_tokens],
163
- kv_cache_dtype=self.kv_cache_dtype,
164
- k_cache=key_cache,
165
- v_cache=value_cache,
166
- b_loc=attn_metadata.block_table,
167
- b_start_loc=attn_metadata.query_start_loc,
168
- b_seq_len=attn_metadata.seq_lens,
169
- max_input_len=attn_metadata.max_query_len,
170
- k_scale=layer._k_scale,
171
- v_scale=layer._v_scale,
172
- alibi_slopes=self.alibi_slopes,
173
- sliding_window=self.sliding_window[0],
174
- sm_scale=self.scale)
160
+ chunked_prefill_paged_decode(
161
+ query=query[:num_actual_tokens],
162
+ key=key[:num_actual_tokens],
163
+ value=value[:num_actual_tokens],
164
+ output=output[:num_actual_tokens],
165
+ kv_cache_dtype=self.kv_cache_dtype,
166
+ key_cache=key_cache,
167
+ value_cache=value_cache,
168
+ block_table=attn_metadata.block_table,
169
+ query_start_loc=attn_metadata.query_start_loc,
170
+ seq_lens=attn_metadata.seq_lens,
171
+ max_query_len=attn_metadata.max_query_len,
172
+ k_scale=layer._k_scale,
173
+ v_scale=layer._v_scale,
174
+ alibi_slopes=self.alibi_slopes,
175
+ sliding_window=self.sliding_window[0],
176
+ sm_scale=self.scale)
177
+
178
return output