Delete PagedAttention (#47361)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 2, 2026 at 15:31 UTC
d715b3aa1ea6af3f663eb6d3cd8f5b6bb15770e9
9 files changed
+90
-1472
CMakeLists.txt
-2
@@ -399,8 +399,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
399
"csrc/libtorch_stable/sampler.cu"
400
"csrc/libtorch_stable/topk.cu"
401
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
402
- "csrc/libtorch_stable/attention/paged_attention_v1.cu"
403
- "csrc/libtorch_stable/attention/paged_attention_v2.cu"
402
"csrc/libtorch_stable/cache_kernels.cu"
403
"csrc/libtorch_stable/cache_kernels.cu"
404
"csrc/libtorch_stable/cache_kernels_fused.cu"
benchmarks/kernels/benchmark_paged_attention.py
+39
-94
@@ -19,13 +19,11 @@ from vllm.utils.torch_utils import (
19
logger = init_logger(__name__)
20
21
NUM_BLOCKS = 128 * 1024
22
-PARTITION_SIZE = 512
22
PARTITION_SIZE_ROCM = 256
23
24
25
@torch.inference_mode()
26
def main(
28
- version: str,
27
num_seqs: int,
28
seq_len: int,
29
num_query_heads: int,
@@ -82,27 +80,20 @@ def main(
80
81
# Prepare for the paged attention kernel.
82
output = torch.empty_like(query)
85
- if version == "v2":
86
- if current_platform.is_rocm():
87
- global PARTITION_SIZE
88
- if not args.custom_paged_attn and not current_platform.is_navi():
89
- PARTITION_SIZE = 1024
90
- else:
91
- PARTITION_SIZE = PARTITION_SIZE_ROCM
92
- num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
93
- tmp_output = torch.empty(
94
- size=(num_seqs, num_query_heads, num_partitions, head_size),
95
- dtype=output.dtype,
96
- device=output.device,
97
- )
98
- exp_sums = torch.empty(
99
- size=(num_seqs, num_query_heads, num_partitions),
100
- dtype=torch.float32,
101
- device=output.device,
102
- )
103
- max_logits = torch.empty_like(exp_sums)
104
-
105
- def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
83
+ num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
84
+ tmp_output = torch.empty(
85
+ size=(num_seqs, num_query_heads, num_partitions, head_size),
86
+ dtype=output.dtype,
87
+ device=output.device,
88
+ )
89
+ exp_sums = torch.empty(
90
+ size=(num_seqs, num_query_heads, num_partitions),
91
+ dtype=torch.float32,
92
+ device=output.device,
93
+ )
94
+ max_logits = torch.empty_like(exp_sums)
95
+
96
+ def run_benchmark(num_iters: int, profile: bool = False) -> float:
97
torch.accelerator.synchronize()
98
if profile:
99
torch.cuda.cudart().cudaProfilerStart()
@@ -112,67 +103,26 @@ def main(
103
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
104
105
for _ in range(num_iters):
115
- if version == "v1":
116
- ops.paged_attention_v1(
117
- output,
118
- query,
119
- key_cache,
120
- value_cache,
121
- num_kv_heads,
122
- scale,
123
- block_tables,
124
- seq_lens,
125
- block_size,
126
- max_seq_len,
127
- alibi_slopes,
128
- kv_cache_dtype,
129
- k_scale,
130
- v_scale,
131
- )
132
- elif version == "v2":
133
- if not args.custom_paged_attn:
134
- ops.paged_attention_v2(
135
- output,
136
- exp_sums,
137
- max_logits,
138
- tmp_output,
139
- query,
140
- key_cache,
141
- value_cache,
142
- num_kv_heads,
143
- scale,
144
- block_tables,
145
- seq_lens,
146
- block_size,
147
- max_seq_len,
148
- alibi_slopes,
149
- kv_cache_dtype,
150
- k_scale,
151
- v_scale,
152
- )
153
- else:
154
- ops.paged_attention_rocm(
155
- output,
156
- exp_sums,
157
- max_logits,
158
- tmp_output,
159
- query,
160
- key_cache,
161
- value_cache,
162
- num_kv_heads,
163
- scale,
164
- block_tables,
165
- seq_lens,
166
- None,
167
- block_size,
168
- max_seq_len,
169
- alibi_slopes,
170
- kv_cache_dtype,
171
- k_scale,
172
- v_scale,
173
- )
174
- else:
175
- raise ValueError(f"Invalid version: {version}")
106
+ ops.paged_attention_rocm(
107
+ output,
108
+ exp_sums,
109
+ max_logits,
110
+ tmp_output,
111
+ query,
112
+ key_cache,
113
+ value_cache,
114
+ num_kv_heads,
115
+ scale,
116
+ block_tables,
117
+ seq_lens,
118
+ None,
119
+ block_size,
120
+ max_seq_len,
121
+ alibi_slopes,
122
+ kv_cache_dtype,
123
+ k_scale,
124
+ v_scale,
125
+ )
126
torch.accelerator.synchronize()
127
128
end_time = time.perf_counter()
@@ -182,7 +132,6 @@ def main(
132
133
# Warmup.
134
print("Warming up...")
185
- run_benchmark = run_cuda_benchmark
135
run_benchmark(num_iters=3, profile=False)
136
137
# Benchmark.
@@ -195,12 +144,13 @@ def main(
144
145
if __name__ == "__main__":
146
logger.warning(
198
- "This script benchmarks the paged attention kernel. "
147
+ "This script benchmarks the ROCm paged attention kernel. "
148
"By default this is no longer used in vLLM inference."
149
)
150
+ if not current_platform.is_rocm():
151
+ raise RuntimeError("This benchmark requires the ROCm platform.")
152
153
parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
203
- parser.add_argument("--version", type=str, choices=["v1", "v2"], default="v2")
154
parser.add_argument("--batch-size", type=int, default=8)
155
parser.add_argument("--seq-len", type=int, default=4096)
156
parser.add_argument("--num-query-heads", type=int, default=64)
@@ -208,7 +158,7 @@ if __name__ == "__main__":
158
parser.add_argument(
159
"--head-size",
160
type=int,
211
- choices=[64, 80, 96, 112, 120, 128, 192, 256],
161
+ choices=[64, 128],
162
default=128,
163
)
164
parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
@@ -224,11 +174,7 @@ if __name__ == "__main__":
174
choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"],
175
default="auto",
176
help="Data type for kv cache storage. If 'auto', will use model "
227
- "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. "
228
- "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
229
- )
230
- parser.add_argument(
231
- "--custom-paged-attn", action="store_true", help="Use custom paged attention"
177
+ "data type. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
178
)
179
args = parser.parse_args()
180
print(args)
@@ -236,7 +182,6 @@ if __name__ == "__main__":
182
if args.num_query_heads % args.num_kv_heads != 0:
183
raise ValueError("num_query_heads must be divisible by num_kv_heads")
184
main(
239
- version=args.version,
185
num_seqs=args.batch_size,
186
seq_len=args.seq_len,
187
num_query_heads=args.num_query_heads,
csrc/libtorch_stable/attention/attention_kernels.cuh
deleted
-667
@@ -1,667 +0,0 @@
1
-/*
2
- * Adapted from
3
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
4
- * Copyright (c) 2023, The vLLM team.
5
- * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
6
- *
7
- * Licensed under the Apache License, Version 2.0 (the "License");
8
- * you may not use this file except in compliance with the License.
9
- * You may obtain a copy of the License at
10
- *
11
- * http://www.apache.org/licenses/LICENSE-2.0
12
- *
13
- * Unless required by applicable law or agreed to in writing, software
14
- * distributed under the License is distributed on an "AS IS" BASIS,
15
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
- * See the License for the specific language governing permissions and
17
- * limitations under the License.
18
- */
19
-
20
-#include <algorithm>
21
-
22
-#include "../../attention/attention_dtypes.h"
23
-#include "attention_utils.cuh"
24
-#include "../../cuda_compat.h"
25
-
26
-#ifdef USE_ROCM
27
- #include <hip/hip_bf16.h>
28
- #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
29
-typedef __hip_bfloat16 __nv_bfloat16;
30
-#else
31
- #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
32
-#endif
33
-
34
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
35
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
36
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
37
-
38
-namespace vllm {
39
-
40
-// Utility function for attention softmax.
41
-template <int NUM_WARPS>
42
-inline __device__ float block_sum(float* red_smem, float sum) {
43
- // Decompose the thread index into warp / lane.
44
- int warp = threadIdx.x / WARP_SIZE;
45
- int lane = threadIdx.x % WARP_SIZE;
46
-
47
- // Compute the sum per warp.
48
-#pragma unroll
49
- for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
50
- sum += VLLM_SHFL_XOR_SYNC(sum, mask);
51
- }
52
-
53
- // Warp leaders store the data to shared memory.
54
- if (lane == 0) {
55
- red_smem[warp] = sum;
56
- }
57
-
58
- // Make sure the data is in shared memory.
59
- __syncthreads();
60
-
61
- // The warps compute the final sums.
62
- if (lane < NUM_WARPS) {
63
- sum = red_smem[lane];
64
- }
65
-
66
- // Parallel reduction inside the warp.
67
-#pragma unroll
68
- for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
69
- sum += VLLM_SHFL_XOR_SYNC(sum, mask);
70
- }
71
-
72
- // Broadcast to other threads.
73
- return VLLM_SHFL_SYNC(sum, 0);
74
-}
75
-
76
-// TODO(woosuk): Merge the last two dimensions of the grid.
77
-// Grid: (num_heads, num_seqs, max_num_partitions).
78
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
79
- int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
80
- bool IS_BLOCK_SPARSE,
81
- int PARTITION_SIZE = 0> // Zero means no partitioning.
82
-__device__ void paged_attention_kernel(
83
- float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
84
- float* __restrict__ max_logits, // [num_seqs, num_heads,
85
- // max_num_partitions]
86
- scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
87
- // head_size]
88
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
89
- const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
90
- // head_size/x, block_size, x]
91
- const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
92
- // head_size, block_size]
93
- const int num_kv_heads, // [num_heads]
94
- const float scale,
95
- const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
96
- const int* __restrict__ seq_lens, // [num_seqs]
97
- const int max_num_blocks_per_seq,
98
- const float* __restrict__ alibi_slopes, // [num_heads]
99
- const int q_stride, const int kv_block_stride, const int kv_head_stride,
100
- const float* k_scale, const float* v_scale, const int tp_rank,
101
- const int blocksparse_local_blocks, const int blocksparse_vert_stride,
102
- const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
103
- const int seq_idx = blockIdx.y;
104
- const int partition_idx = blockIdx.z;
105
- const int max_num_partitions = gridDim.z;
106
- constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0;
107
- const int seq_len = seq_lens[seq_idx];
108
- if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) {
109
- // No work to do. Terminate the thread block.
110
- return;
111
- }
112
-
113
- const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
114
- const int num_blocks_per_partition =
115
- USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks;
116
-
117
- // [start_block_idx, end_block_idx) is the range of blocks to process.
118
- const int start_block_idx =
119
- USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0;
120
- const int end_block_idx =
121
- MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks);
122
- const int num_blocks = end_block_idx - start_block_idx;
123
-
124
- // [start_token_idx, end_token_idx) is the range of tokens to process.
125
- const int start_token_idx = start_block_idx * BLOCK_SIZE;
126
- const int end_token_idx =
127
- MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len);
128
- const int num_tokens = end_token_idx - start_token_idx;
129
-
130
- constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
131
- constexpr int NUM_THREAD_GROUPS =
132
- NUM_THREADS / THREAD_GROUP_SIZE; // Note: This assumes THREAD_GROUP_SIZE
133
- // divides NUM_THREADS
134
- assert(NUM_THREADS % THREAD_GROUP_SIZE == 0);
135
- constexpr int NUM_TOKENS_PER_THREAD_GROUP =
136
- DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE);
137
- constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
138
- const int thread_idx = threadIdx.x;
139
- const int warp_idx = thread_idx / WARP_SIZE;
140
- const int lane = thread_idx % WARP_SIZE;
141
-
142
- const int head_idx = blockIdx.x;
143
- const int num_heads = gridDim.x;
144
- const int num_queries_per_kv = num_heads / num_kv_heads;
145
- const int kv_head_idx = head_idx / num_queries_per_kv;
146
- const float alibi_slope =
147
- alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
148
-
149
- // A vector type to store a part of a key or a query.
150
- // The vector size is configured in such a way that the threads in a thread
151
- // group fetch or compute 16 bytes at a time. For example, if the size of a
152
- // thread group is 4 and the data type is half, then the vector size is 16 /
153
- // (4 * sizeof(half)) == 2.
154
- constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
155
- using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
156
- using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
157
- using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;
158
-
159
- constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE;
160
- constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE;
161
-
162
- const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE;
163
- const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE;
164
-
165
- // Load the query to registers.
166
- // Each thread in a thread group has a different part of the query.
167
- // For example, if the thread group size is 4, then the first thread in
168
- // the group has 0, 4, 8, ... th vectors of the query, and the second thread
169
- // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because
170
- // q is split from a qkv tensor, it may not be contiguous.
171
- const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
172
- __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
173
-#pragma unroll
174
- for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD;
175
- i += NUM_THREAD_GROUPS) {
176
- const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
177
- q_vecs[thread_group_offset][i] =
178
- *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
179
- }
180
- __syncthreads(); // TODO(naed90): possible speedup if this is replaced with a
181
- // memory wall right before we use q_vecs
182
-
183
- // Memory planning.
184
- extern __shared__ char shared_mem[];
185
- // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy.
186
- float* logits = reinterpret_cast<float*>(shared_mem);
187
- // Workspace for reduction.
188
- __shared__ float red_smem[2 * NUM_WARPS];
189
-
190
- // x == THREAD_GROUP_SIZE * VEC_SIZE
191
- // Each thread group fetches x elements from the key at a time.
192
- constexpr int x = 16 / sizeof(cache_t);
193
- float qk_max = -FLT_MAX;
194
-
195
- // Iterate over the key blocks.
196
- // Each warp fetches a block of keys for each iteration.
197
- // Each thread group in a warp fetches a key from the block, and computes
198
- // dot product with the query.
199
- const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
200
-
201
- // blocksparse specific vars
202
- int bs_block_offset;
203
- int q_bs_block_id;
204
- if constexpr (IS_BLOCK_SPARSE) {
205
- // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len,
206
- // blocksparse_block_size);
207
- q_bs_block_id = (seq_len - 1) / blocksparse_block_size;
208
- if (blocksparse_head_sliding_step >= 0)
209
- // sliding on q heads
210
- bs_block_offset =
211
- (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1;
212
- else
213
- // sliding on kv heads
214
- bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) *
215
- (-blocksparse_head_sliding_step) +
216
- 1;
217
- }
218
-
219
- for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
220
- block_idx += NUM_WARPS) {
221
- // NOTE(woosuk): The block number is stored in int32. However, we cast it to
222
- // int64 because int32 can lead to overflow when this variable is multiplied
223
- // by large numbers (e.g., kv_block_stride).
224
- // For blocksparse attention: skip computation on blocks that are not
225
- // attended
226
- if constexpr (IS_BLOCK_SPARSE) {
227
- const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
228
- const bool is_remote =
229
- ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0);
230
- const bool is_local =
231
- (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks);
232
- if (!is_remote && !is_local) {
233
- for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
234
- const int physical_block_offset =
235
- (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
236
- const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
237
-
238
- if (thread_group_offset == 0) {
239
- // NOTE(linxihui): assign very large number to skipped tokens to
240
- // avoid contribution to the sumexp softmax normalizer. This will
241
- // not be used at computing sum(softmax*v) as the blocks will be
242
- // skipped.
243
- logits[token_idx - start_token_idx] = -FLT_MAX;
244
- }
245
- }
246
- continue;
247
- }
248
- }
249
- const int64_t physical_block_number =
250
- static_cast<int64_t>(block_table[block_idx]);
251
-
252
- // Load a key to registers.
253
- // Each thread in a thread group has a different part of the key.
254
- // For example, if the thread group size is 4, then the first thread in
255
- // the group has 0, 4, 8, ... th vectors of the key, and the second thread
256
- // has 1, 5, 9, ... th vectors of the key, and so on.
257
- for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
258
- const int physical_block_offset =
259
- (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
260
- const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
261
- K_vec k_vecs[NUM_VECS_PER_THREAD];
262
-
263
-#pragma unroll
264
- for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
265
- const cache_t* k_ptr =
266
- k_cache + physical_block_number * kv_block_stride +
267
- kv_head_idx * kv_head_stride + physical_block_offset * x;
268
- const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE;
269
- const int offset1 = (vec_idx * VEC_SIZE) / x;
270
- const int offset2 = (vec_idx * VEC_SIZE) % x;
271
-
272
- if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
273
- k_vecs[j] = *reinterpret_cast<const K_vec*>(
274
- k_ptr + offset1 * BLOCK_SIZE * x + offset2);
275
- } else {
276
- // Vector conversion from Quant_vec to K_vec.
277
- Quant_vec k_vec_quant = *reinterpret_cast<const Quant_vec*>(
278
- k_ptr + offset1 * BLOCK_SIZE * x + offset2);
279
- k_vecs[j] = fp8::scaled_convert<K_vec, Quant_vec, KV_DTYPE>(
280
- k_vec_quant, *k_scale);
281
- }
282
- }
283
-
284
- // Compute dot product.
285
- // This includes a reduction across the threads in the same thread group.
286
- float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(
287
- q_vecs[thread_group_offset], k_vecs);
288
- // Add the ALiBi bias if slopes are given.
289
- qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
290
-
291
- if (thread_group_offset == 0) {
292
- // Store the partial reductions to shared memory.
293
- // NOTE(woosuk): It is required to zero out the masked logits.
294
- const bool mask = token_idx >= seq_len;
295
- logits[token_idx - start_token_idx] = mask ? 0.f : qk;
296
- // Update the max value.
297
- qk_max = mask ? qk_max : fmaxf(qk_max, qk);
298
- }
299
- }
300
- }
301
-
302
- // Perform reduction across the threads in the same warp to get the
303
- // max qk value for each "warp" (not across the thread block yet).
304
- // The 0-th thread of each thread group already has its max qk value.
305
-#pragma unroll
306
- for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
307
- qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
308
- }
309
- if (lane == 0) {
310
- red_smem[warp_idx] = qk_max;
311
- }
312
- __syncthreads();
313
-
314
- // TODO(woosuk): Refactor this part.
315
- // Get the max qk value for the sequence.
316
- qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
317
-#pragma unroll
318
- for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
319
- qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
320
- }
321
- // Broadcast the max qk value to all threads.
322
- qk_max = VLLM_SHFL_SYNC(qk_max, 0);
323
-
324
- // Get the sum of the exp values.
325
- float exp_sum = 0.f;
326
- for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
327
- float val = __expf(logits[i] - qk_max);
328
- logits[i] = val;
329
- exp_sum += val;
330
- }
331
- exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
332
-
333
- // Compute softmax.
334
- const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
335
- for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
336
- logits[i] *= inv_sum;
337
- }
338
- __syncthreads();
339
-
340
- // If partitioning is enabled, store the max logit and exp_sum.
341
- if (USE_PARTITIONING && thread_idx == 0) {
342
- float* max_logits_ptr = max_logits +
343
- seq_idx * num_heads * max_num_partitions +
344
- head_idx * max_num_partitions + partition_idx;
345
- *max_logits_ptr = qk_max;
346
- float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions +
347
- head_idx * max_num_partitions + partition_idx;
348
- *exp_sums_ptr = exp_sum;
349
- }
350
-
351
- // Each thread will fetch 16 bytes from the value cache at a time.
352
- constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE);
353
- using V_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
354
- using L_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
355
- using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
356
- using Float_L_vec = typename FloatVec<L_vec>::Type;
357
-
358
- constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE;
359
- constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW;
360
- constexpr int NUM_ROWS_PER_THREAD =
361
- DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
362
-
363
- // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
364
- float accs[NUM_ROWS_PER_THREAD];
365
-#pragma unroll
366
- for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
367
- accs[i] = 0.f;
368
- }
369
-
370
- scalar_t zero_value;
371
- zero(zero_value);
372
- for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
373
- block_idx += NUM_WARPS) {
374
- // NOTE(woosuk): The block number is stored in int32. However, we cast it to
375
- // int64 because int32 can lead to overflow when this variable is multiplied
376
- // by large numbers (e.g., kv_block_stride).
377
- // For blocksparse attention: skip computation on blocks that are not
378
- // attended
379
- if constexpr (IS_BLOCK_SPARSE) {
380
- int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
381
- if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) &&
382
- !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) {
383
- continue;
384
- }
385
- }
386
- const int64_t physical_block_number =
387
- static_cast<int64_t>(block_table[block_idx]);
388
- const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE;
389
- const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
390
- L_vec logits_vec;
391
- from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + token_idx -
392
- start_token_idx));
393
-
394
- const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride +
395
- kv_head_idx * kv_head_stride;
396
-#pragma unroll
397
- for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
398
- const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
399
- if (row_idx < HEAD_SIZE) {
400
- const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
401
- V_vec v_vec;
402
-
403
- if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
404
- v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
405
- } else {
406
- V_quant_vec v_quant_vec =
407
- *reinterpret_cast<const V_quant_vec*>(v_ptr + offset);
408
- // Vector conversion from V_quant_vec to V_vec.
409
- v_vec = fp8::scaled_convert<V_vec, V_quant_vec, KV_DTYPE>(v_quant_vec,
410
- *v_scale);
411
- }
412
- if (block_idx == num_seq_blocks - 1) {
413
- // NOTE(woosuk): When v_vec contains the tokens that are out of the
414
- // context, we should explicitly zero out the values since they may
415
- // contain NaNs. See
416
- // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
417
- scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
418
-#pragma unroll
419
- for (int j = 0; j < V_VEC_SIZE; j++) {
420
- v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
421
- }
422
- }
423
- accs[i] += dot(logits_vec, v_vec);
424
- }
425
- }
426
- }
427
-
428
- // Perform reduction within each warp.
429
-#pragma unroll
430
- for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
431
- float acc = accs[i];
432
-#pragma unroll
433
- for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
434
- acc += VLLM_SHFL_XOR_SYNC(acc, mask);
435
- }
436
- accs[i] = acc;
437
- }
438
-
439
- // NOTE(woosuk): A barrier is required because the shared memory space for
440
- // logits is reused for the output.
441
- __syncthreads();
442
-
443
- // Perform reduction across warps.
444
- float* out_smem = reinterpret_cast<float*>(shared_mem);
445
-#pragma unroll
446
- for (int i = NUM_WARPS; i > 1; i /= 2) {
447
- int mid = i / 2;
448
- // Upper warps write to shared memory.
449
- if (warp_idx >= mid && warp_idx < i) {
450
- float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
451
-#pragma unroll
452
- for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
453
- const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
454
- if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
455
- dst[row_idx] = accs[i];
456
- }
457
- }
458
- }
459
- __syncthreads();
460
-
461
- // Lower warps update the output.
462
- if (warp_idx < mid) {
463
- const float* src = &out_smem[warp_idx * HEAD_SIZE];
464
-#pragma unroll
465
- for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
466
- const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
467
- if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
468
- accs[i] += src[row_idx];
469
- }
470
- }
471
- }
472
- __syncthreads();
473
- }
474
-
475
- // Write the final output.
476
- if (warp_idx == 0) {
477
- scalar_t* out_ptr =
478
- out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
479
- head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE;
480
-#pragma unroll
481
- for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
482
- const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
483
- if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
484
- from_float(*(out_ptr + row_idx), accs[i]);
485
- }
486
- }
487
- }
488
-}
489
-
490
-// Grid: (num_heads, num_seqs, 1).
491
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
492
- int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
493
- bool IS_BLOCK_SPARSE>
494
-__global__ void paged_attention_v1_kernel(
495
- scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size]
496
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
497
- const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
498
- // head_size/x, block_size, x]
499
- const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
500
- // head_size, block_size]
501
- const int num_kv_heads, // [num_heads]
502
- const float scale,
503
- const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
504
- const int* __restrict__ seq_lens, // [num_seqs]
505
- const int max_num_blocks_per_seq,
506
- const float* __restrict__ alibi_slopes, // [num_heads]
507
- const int q_stride, const int kv_block_stride, const int kv_head_stride,
508
- const float* k_scale, const float* v_scale, const int tp_rank,
509
- const int blocksparse_local_blocks, const int blocksparse_vert_stride,
510
- const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
511
- paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
512
- KV_DTYPE, IS_BLOCK_SPARSE>(
513
- /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache,
514
- v_cache, num_kv_heads, scale, block_tables, seq_lens,
515
- max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride,
516
- kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks,
517
- blocksparse_vert_stride, blocksparse_block_size,
518
- blocksparse_head_sliding_step);
519
-}
520
-
521
-// Grid: (num_heads, num_seqs, max_num_partitions).
522
-template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
523
- int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
524
- bool IS_BLOCK_SPARSE,
525
- int PARTITION_SIZE>
526
-__global__ void paged_attention_v2_kernel(
527
- float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
528
- float* __restrict__ max_logits, // [num_seqs, num_heads,
529
- // max_num_partitions]
530
- scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
531
- // max_num_partitions, head_size]
532
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
533
- const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
534
- // head_size/x, block_size, x]
535
- const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
536
- // head_size, block_size]
537
- const int num_kv_heads, // [num_heads]
538
- const float scale,
539
- const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
540
- const int* __restrict__ seq_lens, // [num_seqs]
541
- const int max_num_blocks_per_seq,
542
- const float* __restrict__ alibi_slopes, // [num_heads]
543
- const int q_stride, const int kv_block_stride, const int kv_head_stride,
544
- const float* k_scale, const float* v_scale, const int tp_rank,
545
- const int blocksparse_local_blocks, const int blocksparse_vert_stride,
546
- const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
547
- paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
548
- KV_DTYPE, IS_BLOCK_SPARSE, PARTITION_SIZE>(
549
- exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
550
- block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride,
551
- kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank,
552
- blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size,
553
- blocksparse_head_sliding_step);
554
-}
555
-
556
-// Grid: (num_heads, num_seqs).
557
-template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
558
- int PARTITION_SIZE>
559
-__global__ void paged_attention_v2_reduce_kernel(
560
- scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size]
561
- const float* __restrict__ exp_sums, // [num_seqs, num_heads,
562
- // max_num_partitions]
563
- const float* __restrict__ max_logits, // [num_seqs, num_heads,
564
- // max_num_partitions]
565
- const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
566
- // max_num_partitions, head_size]
567
- const int* __restrict__ seq_lens, // [num_seqs]
568
- const int max_num_partitions) {
569
- const int num_heads = gridDim.x;
570
- const int head_idx = blockIdx.x;
571
- const int seq_idx = blockIdx.y;
572
- const int seq_len = seq_lens[seq_idx];
573
- const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
574
- if (num_partitions == 1) {
575
- // No need to reduce. Only copy tmp_out to out.
576
- scalar_t* out_ptr =
577
- out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
578
- const scalar_t* tmp_out_ptr =
579
- tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
580
- head_idx * max_num_partitions * HEAD_SIZE;
581
- for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) {
582
- out_ptr[i] = tmp_out_ptr[i];
583
- }
584
- // Terminate the thread block.
585
- return;
586
- }
587
-
588
- constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
589
- const int warp_idx = threadIdx.x / WARP_SIZE;
590
- const int lane = threadIdx.x % WARP_SIZE;
591
-
592
- // Size: 2 * num_partitions.
593
- extern __shared__ char shared_mem[];
594
- // Workspace for reduction.
595
- __shared__ float red_smem[2 * NUM_WARPS];
596
-
597
- // Load max logits to shared memory.
598
- float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
599
- const float* max_logits_ptr = max_logits +
600
- seq_idx * num_heads * max_num_partitions +
601
- head_idx * max_num_partitions;
602
- float max_logit = -FLT_MAX;
603
- for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
604
- const float l = max_logits_ptr[i];
605
- shared_max_logits[i] = l;
606
- max_logit = fmaxf(max_logit, l);
607
- }
608
- __syncthreads();
609
-
610
- // Get the global max logit.
611
- // Reduce within the warp.
612
-#pragma unroll
613
- for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
614
- max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
615
- }
616
- if (lane == 0) {
617
- red_smem[warp_idx] = max_logit;
618
- }
619
- __syncthreads();
620
- // Reduce across warps.
621
- max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
622
-#pragma unroll
623
- for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
624
- max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
625
- }
626
- // Broadcast the max value to all threads.
627
- max_logit = VLLM_SHFL_SYNC(max_logit, 0);
628
-
629
- // Load rescaled exp sums to shared memory.
630
- float* shared_exp_sums =
631
- reinterpret_cast<float*>(shared_mem + sizeof(float) * num_partitions);
632
- const float* exp_sums_ptr = exp_sums +
633
- seq_idx * num_heads * max_num_partitions +
634
- head_idx * max_num_partitions;
635
- float global_exp_sum = 0.0f;
636
- for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
637
- float l = shared_max_logits[i];
638
- float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit);
639
- global_exp_sum += rescaled_exp_sum;
640
- shared_exp_sums[i] = rescaled_exp_sum;
641
- }
642
- __syncthreads();
643
- global_exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], global_exp_sum);
644
- const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f);
645
-
646
- // Aggregate tmp_out to out.
647
- const scalar_t* tmp_out_ptr =
648
- tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
649
- head_idx * max_num_partitions * HEAD_SIZE;
650
- scalar_t* out_ptr =
651
- out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
652
-#pragma unroll
653
- for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) {
654
- float acc = 0.0f;
655
- for (int j = 0; j < num_partitions; ++j) {
656
- acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] *
657
- inv_global_exp_sum;
658
- }
659
- from_float(out_ptr[i], acc);
660
- }
661
-}
662
-
663
-} // namespace vllm
664
-
665
-#undef MAX
666
-#undef MIN
667
-#undef DIVIDE_ROUND_UP
csrc/libtorch_stable/attention/paged_attention_v1.cu
deleted
-190
@@ -1,190 +0,0 @@
1
-/*
2
- * Adapted from
3
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
4
- * Copyright (c) 2023, The vLLM team.
5
- * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
6
- *
7
- * Licensed under the Apache License, Version 2.0 (the "License");
8
- * you may not use this file except in compliance with the License.
9
- * You may obtain a copy of the License at
10
- *
11
- * http://www.apache.org/licenses/LICENSE-2.0
12
- *
13
- * Unless required by applicable law or agreed to in writing, software
14
- * distributed under the License is distributed on an "AS IS" BASIS,
15
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
- * See the License for the specific language governing permissions and
17
- * limitations under the License.
18
- */
19
-#include "../torch_utils.h"
20
-#include "attention_kernels.cuh"
21
-#include "../../cuda_compat.h"
22
-
23
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
24
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
25
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
26
-
27
-#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE) \
28
- VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( \
29
- ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, \
30
- BLOCK_SIZE, NUM_THREADS, \
31
- KV_DTYPE, IS_BLOCK_SPARSE>), \
32
- shared_mem_size); \
33
- vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE, \
34
- NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE> \
35
- <<<grid, block, shared_mem_size, stream>>>( \
36
- out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
37
- scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq, \
38
- alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \
39
- k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks, \
40
- blocksparse_vert_stride, blocksparse_block_size, \
41
- blocksparse_head_sliding_step);
42
-
43
-// TODO(woosuk): Tune NUM_THREADS.
44
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
45
- vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
46
- int NUM_THREADS = 128>
47
-void paged_attention_v1_launcher(
48
- torch::stable::Tensor& out, torch::stable::Tensor& query,
49
- torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
50
- int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
51
- torch::stable::Tensor& seq_lens, int max_seq_len,
52
- const std::optional<torch::stable::Tensor>& alibi_slopes,
53
- torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
54
- const int tp_rank, const int blocksparse_local_blocks,
55
- const int blocksparse_vert_stride, const int blocksparse_block_size,
56
- const int blocksparse_head_sliding_step) {
57
- int num_seqs = query.size(0);
58
- int num_heads = query.size(1);
59
- int head_size = query.size(2);
60
- int max_num_blocks_per_seq = block_tables.size(1);
61
- int q_stride = query.stride(0);
62
- int kv_block_stride = key_cache.stride(0);
63
- int kv_head_stride = key_cache.stride(1);
64
-
65
- // NOTE: alibi_slopes is optional.
66
- const float* alibi_slopes_ptr =
67
- alibi_slopes
68
- ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
69
- : nullptr;
70
-
71
- T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
72
- T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
73
- CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
74
- CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
75
- int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
76
- int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
77
- const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
78
- const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
79
-
80
- const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
81
- int padded_max_seq_len =
82
- DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
83
- int logits_size = padded_max_seq_len * sizeof(float);
84
- int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
85
- // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
86
- // Keep that in sync with the logic here!
87
- int shared_mem_size = std::max(logits_size, outputs_size);
88
-
89
- dim3 grid(num_heads, num_seqs, 1);
90
- dim3 block(NUM_THREADS);
91
- const torch::stable::accelerator::DeviceGuard device_guard(
92
- query.get_device_index());
93
- const cudaStream_t stream = get_current_cuda_stream();
94
- switch (head_size) {
95
- // NOTE(woosuk): To reduce the compilation time, we only compile for the
96
- // head sizes that we use in the model. However, we can easily extend this
97
- // to support any head size which is a multiple of 16.
98
- case 32:
99
- LAUNCH_PAGED_ATTENTION_V1(32);
100
- break;
101
- case 64:
102
- LAUNCH_PAGED_ATTENTION_V1(64);
103
- break;
104
- case 80:
105
- LAUNCH_PAGED_ATTENTION_V1(80);
106
- break;
107
- case 96:
108
- LAUNCH_PAGED_ATTENTION_V1(96);
109
- break;
110
- case 112:
111
- LAUNCH_PAGED_ATTENTION_V1(112);
112
- break;
113
- case 120:
114
- LAUNCH_PAGED_ATTENTION_V1(120);
115
- break;
116
- case 128:
117
- LAUNCH_PAGED_ATTENTION_V1(128);
118
- break;
119
- case 192:
120
- LAUNCH_PAGED_ATTENTION_V1(192);
121
- break;
122
- case 256:
123
- LAUNCH_PAGED_ATTENTION_V1(256);
124
- break;
125
- default:
126
- STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
127
- break;
128
- }
129
-}
130
-
131
-#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \
132
- paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE, \
133
- IS_BLOCK_SPARSE>( \
134
- out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
135
- seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank, \
136
- blocksparse_local_blocks, blocksparse_vert_stride, \
137
- blocksparse_block_size, blocksparse_head_sliding_step);
138
-
139
-#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
140
- if (is_block_sparse) { \
141
- CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \
142
- } else { \
143
- CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \
144
- }
145
-
146
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
147
-// 1, 2, 4, 64, 128, 256.
148
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
149
- switch (block_size) { \
150
- case 8: \
151
- CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
152
- break; \
153
- case 16: \
154
- CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
155
- break; \
156
- case 32: \
157
- CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
158
- break; \
159
- default: \
160
- STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
161
- break; \
162
- }
163
-
164
-void paged_attention_v1(
165
- torch::stable::Tensor& out, // [num_seqs, num_heads, head_size]
166
- torch::stable::Tensor& query, // [num_seqs, num_heads, head_size]
167
- torch::stable::Tensor&
168
- key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
169
- torch::stable::Tensor&
170
- value_cache, // [num_blocks, num_heads, head_size, block_size]
171
- int64_t num_kv_heads, // [num_heads]
172
- double scale,
173
- torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
174
- torch::stable::Tensor& seq_lens, // [num_seqs]
175
- int64_t block_size, int64_t max_seq_len,
176
- const std::optional<torch::stable::Tensor>& alibi_slopes,
177
- const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
178
- torch::stable::Tensor& v_scale, const int64_t tp_rank,
179
- const int64_t blocksparse_local_blocks,
180
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
181
- const int64_t blocksparse_head_sliding_step) {
182
- const bool is_block_sparse = (blocksparse_vert_stride > 1);
183
-
184
- DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
185
- CALL_V1_LAUNCHER_BLOCK_SIZE)
186
-}
187
-
188
-#undef MAX
189
-#undef MIN
190
-#undef DIVIDE_ROUND_UP
csrc/libtorch_stable/attention/paged_attention_v2.cu
deleted
-202
@@ -1,202 +0,0 @@
1
-/*
2
- * Adapted from
3
- * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
4
- * Copyright (c) 2023, The vLLM team.
5
- * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
6
- *
7
- * Licensed under the Apache License, Version 2.0 (the "License");
8
- * you may not use this file except in compliance with the License.
9
- * You may obtain a copy of the License at
10
- *
11
- * http://www.apache.org/licenses/LICENSE-2.0
12
- *
13
- * Unless required by applicable law or agreed to in writing, software
14
- * distributed under the License is distributed on an "AS IS" BASIS,
15
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
- * See the License for the specific language governing permissions and
17
- * limitations under the License.
18
- */
19
-#include "../torch_utils.h"
20
-#include "attention_kernels.cuh"
21
-#include "../../cuda_compat.h"
22
-
23
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
24
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
25
-#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
26
-
27
-#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE) \
28
- vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE, \
29
- NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE, \
30
- PARTITION_SIZE> \
31
- <<<grid, block, shared_mem_size, stream>>>( \
32
- exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
33
- value_cache_ptr, num_kv_heads, scale, block_tables_ptr, \
34
- seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, \
35
- kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank, \
36
- blocksparse_local_blocks, blocksparse_vert_stride, \
37
- blocksparse_block_size, blocksparse_head_sliding_step); \
38
- vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS, \
39
- PARTITION_SIZE> \
40
- <<<reduce_grid, block, reduce_shared_mem_size, stream>>>( \
41
- out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \
42
- max_num_partitions);
43
-
44
-template <typename T, typename CACHE_T, int BLOCK_SIZE,
45
- vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
46
- int NUM_THREADS = 128, int PARTITION_SIZE = 512>
47
-void paged_attention_v2_launcher(
48
- torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
49
- torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
50
- torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
51
- torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
52
- torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
53
- int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
54
- torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
55
- const int tp_rank, const int blocksparse_local_blocks,
56
- const int blocksparse_vert_stride, const int blocksparse_block_size,
57
- const int blocksparse_head_sliding_step) {
58
- int num_seqs = query.size(0);
59
- int num_heads = query.size(1);
60
- int head_size = query.size(2);
61
- int max_num_blocks_per_seq = block_tables.size(1);
62
- int q_stride = query.stride(0);
63
- int kv_block_stride = key_cache.stride(0);
64
- int kv_head_stride = key_cache.stride(1);
65
-
66
- // NOTE: alibi_slopes is optional.
67
- const float* alibi_slopes_ptr =
68
- alibi_slopes
69
- ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
70
- : nullptr;
71
-
72
- T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
73
- float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
74
- float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
75
- T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
76
- T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
77
- CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
78
- CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
79
- int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
80
- int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
81
- const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
82
- const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
83
-
84
- const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
85
- int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
86
- int logits_size = PARTITION_SIZE * sizeof(float);
87
- int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
88
-
89
- // For paged attention v2 kernel.
90
- dim3 grid(num_heads, num_seqs, max_num_partitions);
91
- int shared_mem_size = std::max(logits_size, outputs_size);
92
- // For paged attention v2 reduce kernel.
93
- dim3 reduce_grid(num_heads, num_seqs);
94
- int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
95
-
96
- dim3 block(NUM_THREADS);
97
- const torch::stable::accelerator::DeviceGuard device_guard(
98
- query.get_device_index());
99
- const cudaStream_t stream = get_current_cuda_stream();
100
- switch (head_size) {
101
- // NOTE(woosuk): To reduce the compilation time, we only compile for the
102
- // head sizes that we use in the model. However, we can easily extend this
103
- // to support any head size which is a multiple of 16.
104
- case 32:
105
- LAUNCH_PAGED_ATTENTION_V2(32);
106
- break;
107
- case 64:
108
- LAUNCH_PAGED_ATTENTION_V2(64);
109
- break;
110
- case 80:
111
- LAUNCH_PAGED_ATTENTION_V2(80);
112
- break;
113
- case 96:
114
- LAUNCH_PAGED_ATTENTION_V2(96);
115
- break;
116
- case 112:
117
- LAUNCH_PAGED_ATTENTION_V2(112);
118
- break;
119
- case 120:
120
- LAUNCH_PAGED_ATTENTION_V2(120);
121
- break;
122
- case 128:
123
- LAUNCH_PAGED_ATTENTION_V2(128);
124
- break;
125
- case 192:
126
- LAUNCH_PAGED_ATTENTION_V2(192);
127
- break;
128
- case 256:
129
- LAUNCH_PAGED_ATTENTION_V2(256);
130
- break;
131
- default:
132
- STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
133
- break;
134
- }
135
-}
136
-
137
-#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \
138
- paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE, \
139
- IS_BLOCK_SPARSE>( \
140
- out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
141
- num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \
142
- k_scale, v_scale, tp_rank, blocksparse_local_blocks, \
143
- blocksparse_vert_stride, blocksparse_block_size, \
144
- blocksparse_head_sliding_step);
145
-
146
-#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
147
- if (is_block_sparse) { \
148
- CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \
149
- } else { \
150
- CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \
151
- }
152
-
153
-// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
154
-// 1, 2, 4, 64, 128, 256.
155
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \
156
- switch (block_size) { \
157
- case 8: \
158
- CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \
159
- break; \
160
- case 16: \
161
- CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \
162
- break; \
163
- case 32: \
164
- CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \
165
- break; \
166
- default: \
167
- STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
168
- break; \
169
- }
170
-
171
-void paged_attention_v2(
172
- torch::stable::Tensor& out, // [num_seqs, num_heads, head_size]
173
- torch::stable::Tensor&
174
- exp_sums, // [num_seqs, num_heads, max_num_partitions]
175
- torch::stable::Tensor&
176
- max_logits, // [num_seqs, num_heads, max_num_partitions]
177
- torch::stable::Tensor&
178
- tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
179
- torch::stable::Tensor& query, // [num_seqs, num_heads, head_size]
180
- torch::stable::Tensor&
181
- key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
182
- torch::stable::Tensor&
183
- value_cache, // [num_blocks, num_heads, head_size, block_size]
184
- int64_t num_kv_heads, // [num_heads]
185
- double scale,
186
- torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
187
- torch::stable::Tensor& seq_lens, // [num_seqs]
188
- int64_t block_size, int64_t max_seq_len,
189
- const std::optional<torch::stable::Tensor>& alibi_slopes,
190
- const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
191
- torch::stable::Tensor& v_scale, const int64_t tp_rank,
192
- const int64_t blocksparse_local_blocks,
193
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
194
- const int64_t blocksparse_head_sliding_step) {
195
- const bool is_block_sparse = (blocksparse_vert_stride > 1);
196
- DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
197
- CALL_V2_LAUNCHER_BLOCK_SIZE)
198
-}
199
-
200
-#undef MAX
201
-#undef MIN
202
-#undef DIVIDE_ROUND_UP
csrc/libtorch_stable/ops.h
-26
@@ -452,32 +452,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
452
void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm,
453
int64_t bit);
454
455
-void paged_attention_v1(
456
- torch::stable::Tensor& out, torch::stable::Tensor& query,
457
- torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
458
- int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
459
- torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
460
- const std::optional<torch::stable::Tensor>& alibi_slopes,
461
- const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
462
- torch::stable::Tensor& v_scale, const int64_t tp_rank,
463
- const int64_t blocksparse_local_blocks,
464
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
465
- const int64_t blocksparse_head_sliding_step);
466
-
467
-void paged_attention_v2(
468
- torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
469
- torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
470
- torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
471
- torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
472
- torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
473
- int64_t block_size, int64_t max_seq_len,
474
- const std::optional<torch::stable::Tensor>& alibi_slopes,
475
- const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
476
- torch::stable::Tensor& v_scale, const int64_t tp_rank,
477
- const int64_t blocksparse_local_blocks,
478
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
479
- const int64_t blocksparse_head_sliding_step);
480
-
455
// Cache ops (shared CUDA/ROCm)
456
void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
457
int64_t block_size_in_bytes,
csrc/libtorch_stable/torch_bindings.cpp
-30
@@ -598,33 +598,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
598
"Tensor? initial_state_idx,"
599
"Tensor? cu_chunk_seqlen,"
600
"Tensor? last_chunk_indices) -> ()");
601
-
602
- // Attention ops
603
- // Compute the attention between an input query and the cached
604
- // keys/values using PagedAttention.
605
- ops.def(
606
- "paged_attention_v1("
607
- " Tensor! out, Tensor query, Tensor key_cache,"
608
- " Tensor value_cache, int num_kv_heads, float scale,"
609
- " Tensor block_tables, Tensor seq_lens, int block_size,"
610
- " int max_seq_len, Tensor? alibi_slopes,"
611
- " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
612
- " int tp_rank, int blocksparse_local_blocks,"
613
- " int blocksparse_vert_stride, int blocksparse_block_size,"
614
- " int blocksparse_head_sliding_step) -> ()");
615
-
616
- // PagedAttention V2.
617
- ops.def(
618
- "paged_attention_v2("
619
- " Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
620
- " Tensor! tmp_out, Tensor query, Tensor key_cache,"
621
- " Tensor value_cache, int num_kv_heads, float scale,"
622
- " Tensor block_tables, Tensor seq_lens, int block_size,"
623
- " int max_seq_len, Tensor? alibi_slopes,"
624
- " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
625
- " int tp_rank, int blocksparse_local_blocks,"
626
- " int blocksparse_vert_stride, int blocksparse_block_size,"
627
- " int blocksparse_head_sliding_step) -> ()");
601
}
602
603
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -747,9 +720,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
720
721
// Mamba kernels
722
ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
750
-
751
- ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
752
- ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
723
}
724
725
STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
tests/kernels/attention/test_attention.py
+51
-167
@@ -21,16 +21,14 @@ MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
21
# There may not be enough gpu memory due to large NUM_BLOCKS.
22
# Reduce NUM_BLOCKS when it happens.
23
NUM_BLOCKS = 4321 # Arbitrary values for testing
24
-PARTITION_SIZE = 512
24
PARTITION_SIZE_ROCM = 256
25
DTYPES = [torch.bfloat16]
26
NUM_GEN_SEQS = [7] # Arbitrary values for testing
27
NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
28
NUM_HEADS = [(32, 8), (40, 40), (64, 8)] # Arbitrary values for testing
29
31
-# This should be sync with get_supported_head_sizes() in
32
-# vllm.v1.attention.ops.paged_attn.PagedAttention
33
-HEAD_SIZES = [32, 64, 80, 128, 256]
30
+# Head sizes supported by the ROCm paged attention kernel.
31
+HEAD_SIZES = [64, 128]
32
33
BLOCK_SIZES = [16, 32]
34
USE_ALIBI = [False, True]
@@ -111,8 +109,8 @@ def ref_single_query_cached_kv_attention(
109
output[i].copy_(out, non_blocking=True)
110
111
114
-@pytest.mark.parametrize(
115
- "version", ["v1", "v2"] if not current_platform.is_rocm() else ["v1", "v2", "rocm"]
112
+@pytest.mark.skipif(
113
+ not current_platform.is_rocm(), reason="ROCm-only paged attention kernel"
114
)
115
@pytest.mark.parametrize("num_seqs", NUM_GEN_SEQS)
116
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@@ -125,7 +123,6 @@ def ref_single_query_cached_kv_attention(
123
@pytest.mark.parametrize("device", CUDA_DEVICES)
124
def test_paged_attention(
125
kv_cache_factory,
128
- version: str,
126
num_seqs: int,
127
num_heads: tuple[int, int],
128
head_size: int,
@@ -136,22 +133,11 @@ def test_paged_attention(
133
seed: int,
134
device: str,
135
) -> None:
139
- if (kv_cache_dtype == "fp8" and head_size % 16) or (
140
- version == "rocm" and head_size not in (64, 128)
136
+ if current_platform.is_navi() and (
137
+ kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi
138
):
139
pytest.skip()
140
144
- if (
145
- version == "rocm"
146
- and current_platform.is_navi()
147
- and (
148
- kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi
149
- )
150
- ):
151
- pytest.skip()
152
-
153
- global PARTITION_SIZE
154
-
141
set_random_seed(seed)
142
torch.set_default_device(device)
143
scale = float(1.0 / (head_size**0.5))
@@ -200,9 +186,46 @@ def test_paged_attention(
186
187
# Call the paged attention kernel.
188
output = torch.empty_like(query)
203
- if version == "v1":
204
- ops.paged_attention_v1(
189
+ num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
190
+ assert PARTITION_SIZE_ROCM % block_size == 0
191
+ num_seqs, num_heads, head_size = output.shape
192
+ tmp_output = torch.empty(
193
+ size=(num_seqs, num_heads, num_partitions, head_size),
194
+ dtype=output.dtype,
195
+ )
196
+ exp_sums = torch.empty(
197
+ size=(num_seqs, num_heads, num_partitions),
198
+ dtype=torch.float32,
199
+ )
200
+ max_logits = torch.empty_like(exp_sums)
201
+ ops.paged_attention_rocm(
202
+ output,
203
+ exp_sums,
204
+ max_logits,
205
+ tmp_output,
206
+ query,
207
+ key_cache,
208
+ value_cache,
209
+ num_kv_heads,
210
+ scale,
211
+ block_tables,
212
+ seq_lens,
213
+ None,
214
+ block_size,
215
+ max_seq_len,
216
+ alibi_slopes,
217
+ kv_cache_dtype,
218
+ k_scale,
219
+ v_scale,
220
+ )
221
+
222
+ opcheck(
223
+ torch.ops._rocm_C.paged_attention,
224
+ (
225
output,
226
+ exp_sums,
227
+ max_logits,
228
+ tmp_output,
229
query,
230
key_cache,
231
value_cache,
@@ -210,157 +233,18 @@ def test_paged_attention(
233
scale,
234
block_tables,
235
seq_lens,
236
+ None,
237
block_size,
238
max_seq_len,
239
alibi_slopes,
240
kv_cache_dtype,
241
k_scale,
242
v_scale,
219
- )
220
-
221
- opcheck(
222
- torch.ops._C.paged_attention_v1,
223
- (
224
- output,
225
- query,
226
- key_cache,
227
- value_cache,
228
- num_kv_heads,
229
- scale,
230
- block_tables,
231
- seq_lens,
232
- block_size,
233
- max_seq_len,
234
- alibi_slopes,
235
- kv_cache_dtype,
236
- k_scale,
237
- v_scale,
238
- 0,
239
- 0,
240
- 0,
241
- 64,
242
- 0,
243
- ),
244
- cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]),
245
- )
246
-
247
- elif version in ("v2", "rocm"):
248
- if current_platform.is_rocm() and version == "rocm":
249
- PARTITION_SIZE = PARTITION_SIZE_ROCM
250
-
251
- num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
252
- assert PARTITION_SIZE % block_size == 0
253
- num_seqs, num_heads, head_size = output.shape
254
- tmp_output = torch.empty(
255
- size=(num_seqs, num_heads, num_partitions, head_size),
256
- dtype=output.dtype,
257
- )
258
- exp_sums = torch.empty(
259
- size=(num_seqs, num_heads, num_partitions),
260
- dtype=torch.float32,
261
- )
262
- max_logits = torch.empty_like(exp_sums)
263
- if version == "v2":
264
- ops.paged_attention_v2(
265
- output,
266
- exp_sums,
267
- max_logits,
268
- tmp_output,
269
- query,
270
- key_cache,
271
- value_cache,
272
- num_kv_heads,
273
- scale,
274
- block_tables,
275
- seq_lens,
276
- block_size,
277
- max_seq_len,
278
- alibi_slopes,
279
- kv_cache_dtype,
280
- k_scale,
281
- v_scale,
282
- )
283
-
284
- opcheck(
285
- torch.ops._C.paged_attention_v2,
286
- (
287
- output,
288
- exp_sums,
289
- max_logits,
290
- tmp_output,
291
- query,
292
- key_cache,
293
- value_cache,
294
- num_kv_heads,
295
- scale,
296
- block_tables,
297
- seq_lens,
298
- block_size,
299
- max_seq_len,
300
- alibi_slopes,
301
- kv_cache_dtype,
302
- k_scale,
303
- v_scale,
304
- 0,
305
- 0,
306
- 0,
307
- 64,
308
- 0,
309
- ),
310
- cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]),
311
- )
312
-
313
- else:
314
- ops.paged_attention_rocm(
315
- output,
316
- exp_sums,
317
- max_logits,
318
- tmp_output,
319
- query,
320
- key_cache,
321
- value_cache,
322
- num_kv_heads,
323
- scale,
324
- block_tables,
325
- seq_lens,
326
- None,
327
- block_size,
328
- max_seq_len,
329
- alibi_slopes,
330
- kv_cache_dtype,
331
- k_scale,
332
- v_scale,
333
- )
334
-
335
- opcheck(
336
- torch.ops._rocm_C.paged_attention,
337
- (
338
- output,
339
- exp_sums,
340
- max_logits,
341
- tmp_output,
342
- query,
343
- key_cache,
344
- value_cache,
345
- num_kv_heads,
346
- scale,
347
- block_tables,
348
- seq_lens,
349
- None,
350
- block_size,
351
- max_seq_len,
352
- alibi_slopes,
353
- kv_cache_dtype,
354
- k_scale,
355
- v_scale,
356
- None,
357
- "f16",
358
- ),
359
- cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
360
- )
361
-
362
- else:
363
- raise AssertionError(f"Unknown version: {version}")
243
+ None,
244
+ "f16",
245
+ ),
246
+ cond=(head_size == 64 and block_size == BLOCK_SIZES[0]),
247
+ )
248
249
# Run the reference implementation.
250
if kv_cache_dtype == "fp8":
vllm/_custom_ops.py
-94
@@ -117,100 +117,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
117
118
119
# page attention ops
120
-def paged_attention_v1(
121
- out: torch.Tensor,
122
- query: torch.Tensor,
123
- key_cache: torch.Tensor,
124
- value_cache: torch.Tensor,
125
- num_kv_heads: int,
126
- scale: float,
127
- block_tables: torch.Tensor,
128
- seq_lens: torch.Tensor,
129
- block_size: int,
130
- max_seq_len: int,
131
- alibi_slopes: torch.Tensor | None,
132
- kv_cache_dtype: str,
133
- k_scale: torch.Tensor,
134
- v_scale: torch.Tensor,
135
- tp_rank: int = 0,
136
- blocksparse_local_blocks: int = 0,
137
- blocksparse_vert_stride: int = 0,
138
- blocksparse_block_size: int = 64,
139
- blocksparse_head_sliding_step: int = 0,
140
-) -> None:
141
- torch.ops._C.paged_attention_v1(
142
- out,
143
- query,
144
- key_cache,
145
- value_cache,
146
- num_kv_heads,
147
- scale,
148
- block_tables,
149
- seq_lens,
150
- block_size,
151
- max_seq_len,
152
- alibi_slopes,
153
- kv_cache_dtype,
154
- k_scale,
155
- v_scale,
156
- tp_rank,
157
- blocksparse_local_blocks,
158
- blocksparse_vert_stride,
159
- blocksparse_block_size,
160
- blocksparse_head_sliding_step,
161
- )
162
-
163
-
164
-def paged_attention_v2(
165
- out: torch.Tensor,
166
- exp_sum: torch.Tensor,
167
- max_logits: torch.Tensor,
168
- tmp_out: torch.Tensor,
169
- query: torch.Tensor,
170
- key_cache: torch.Tensor,
171
- value_cache: torch.Tensor,
172
- num_kv_heads: int,
173
- scale: float,
174
- block_tables: torch.Tensor,
175
- seq_lens: torch.Tensor,
176
- block_size: int,
177
- max_seq_len: int,
178
- alibi_slopes: torch.Tensor | None,
179
- kv_cache_dtype: str,
180
- k_scale: torch.Tensor,
181
- v_scale: torch.Tensor,
182
- tp_rank: int = 0,
183
- blocksparse_local_blocks: int = 0,
184
- blocksparse_vert_stride: int = 0,
185
- blocksparse_block_size: int = 64,
186
- blocksparse_head_sliding_step: int = 0,
187
-) -> None:
188
- torch.ops._C.paged_attention_v2(
189
- out,
190
- exp_sum,
191
- max_logits,
192
- tmp_out,
193
- query,
194
- key_cache,
195
- value_cache,
196
- num_kv_heads,
197
- scale,
198
- block_tables,
199
- seq_lens,
200
- block_size,
201
- max_seq_len,
202
- alibi_slopes,
203
- kv_cache_dtype,
204
- k_scale,
205
- v_scale,
206
- tp_rank,
207
- blocksparse_local_blocks,
208
- blocksparse_vert_stride,
209
- blocksparse_block_size,
210
- blocksparse_head_sliding_step,
211
- )
212
-
213
-
120
def paged_attention_rocm(
121
out: torch.Tensor,
122
exp_sum: torch.Tensor,