attention_kernels.cu 45 KB
Newer Older
zhuwenwen's avatar
zhuwenwen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/*
 * Adapted from
 * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
 * Copyright (c) 2023, The vLLM team.
 * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

20
#include <torch/all.h>
Woosuk Kwon's avatar
Woosuk Kwon committed
21
#include <ATen/cuda/CUDAContext.h>
22
#include <c10/cuda/CUDAGuard.h>
23
#include <algorithm>
Woosuk Kwon's avatar
Woosuk Kwon committed
24

Woosuk Kwon's avatar
Woosuk Kwon committed
25
#include "attention_dtypes.h"
Woosuk Kwon's avatar
Woosuk Kwon committed
26
#include "attention_utils.cuh"
27
28
29

#ifdef USE_ROCM
  #include <hip/hip_bf16.h>
30
  #include "../quantization/fp8/amd/quant_utils.cuh"
31
typedef __hip_bfloat16 __nv_bfloat16;
32
33
#else
  #include "../quantization/fp8/nvidia/quant_utils.cuh"
34
35
36
#endif

#ifndef USE_ROCM
37
  #define WARP_SIZE 32
38
#else
39
  #define WARP_SIZE warpSize
40
41
#endif

Woosuk Kwon's avatar
Woosuk Kwon committed
42
43
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
44
#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
Woosuk Kwon's avatar
Woosuk Kwon committed
45

Woosuk Kwon's avatar
Woosuk Kwon committed
46
namespace vllm {
Woosuk Kwon's avatar
Woosuk Kwon committed
47
48

// Utility function for attention softmax.
49
template <int NUM_WARPS>
Woosuk Kwon's avatar
Woosuk Kwon committed
50
51
52
53
54
55
56
57
inline __device__ float block_sum(float* red_smem, float sum) {
  // Decompose the thread index into warp / lane.
  int warp = threadIdx.x / WARP_SIZE;
  int lane = threadIdx.x % WARP_SIZE;

  // Compute the sum per warp.
#pragma unroll
  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
58
    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
  }

  // Warp leaders store the data to shared memory.
  if (lane == 0) {
    red_smem[warp] = sum;
  }

  // Make sure the data is in shared memory.
  __syncthreads();

  // The warps compute the final sums.
  if (lane < NUM_WARPS) {
    sum = red_smem[lane];
  }

  // Parallel reduction inside the warp.
#pragma unroll
  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
77
    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
78
79
80
  }

  // Broadcast to other threads.
81
  return VLLM_SHFL_SYNC(sum, 0);
Woosuk Kwon's avatar
Woosuk Kwon committed
82
83
}

zhuwenwen's avatar
zhuwenwen committed
84
85
// TODO(woosuk): Merge the last two dimensions of the grid.
// Grid: (num_heads, num_seqs, max_num_partitions).
86
87
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
zhuwenwen's avatar
zhuwenwen committed
88
          bool IS_BLOCK_SPARSE,
zhangshao's avatar
zhangshao committed
89
          int PARTITION_SIZE = 0>  // Zero means no partitioning.
zhuwenwen's avatar
zhuwenwen committed
90
__device__ void paged_attention_kernel(
zhangshao's avatar
zhangshao committed
91
92
93
94
95
96
97
98
99
100
    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
    float* __restrict__ max_logits,  // [num_seqs, num_heads,
                                     // max_num_partitions]
    scalar_t* __restrict__ out,  // [num_seqs, num_heads, max_num_partitions,
                                 // head_size]
    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
                                          // head_size/x, block_size, x]
    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
                                          // head_size, block_size]
zhuwenwen's avatar
zhuwenwen committed
101
    const int num_kv_heads,               // [num_heads]
102
103
104
105
106
107
    const float scale,
    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ seq_lens,      // [num_seqs]
    const int max_num_blocks_per_seq,
    const float* __restrict__ alibi_slopes,  // [num_heads]
    const int q_stride, const int kv_block_stride, const int kv_head_stride,
zhuwenwen's avatar
zhuwenwen committed
108
109
    const float k_scale, const float v_scale, const int tp_rank,
    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
110
111
    const int blocksparse_block_size, const int blocksparse_head_sliding_step,
    const int* __restrict__ attn_masks=nullptr, const int attn_masks_stride=0) {
zhuwenwen's avatar
zhuwenwen committed
112
113
114
  const int seq_idx = blockIdx.y;
  const int partition_idx = blockIdx.z;
  const int max_num_partitions = gridDim.z;
115
116
117
118
119
120
121
122
123
124
  constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0;
  const int seq_len = seq_lens[seq_idx];
  if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) {
    // No work to do. Terminate the thread block.
    return;
  }

  const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
  const int num_blocks_per_partition =
      USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks;
zhuwenwen's avatar
zhuwenwen committed
125

126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
  // [start_block_idx, end_block_idx) is the range of blocks to process.
  const int start_block_idx =
      USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0;
  const int end_block_idx =
      MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks);
  const int num_blocks = end_block_idx - start_block_idx;

  // [start_token_idx, end_token_idx) is the range of tokens to process.
  const int start_token_idx = start_block_idx * BLOCK_SIZE;
  const int end_token_idx =
      MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len);
  const int num_tokens = end_token_idx - start_token_idx;

  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
  constexpr int NUM_THREAD_GROUPS =
      NUM_THREADS / THREAD_GROUP_SIZE;  // Note: This assumes THREAD_GROUP_SIZE
                                        // divides NUM_THREADS
  assert(NUM_THREADS % THREAD_GROUP_SIZE == 0);
  constexpr int NUM_TOKENS_PER_THREAD_GROUP =
      DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE);
  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  const int thread_idx = threadIdx.x;
zhuwenwen's avatar
zhuwenwen committed
148
  const int warp_idx = thread_idx / WARP_SIZE;
149
150
  const int lane = thread_idx % WARP_SIZE;

zhuwenwen's avatar
zhuwenwen committed
151
152
  const int head_idx = blockIdx.x;
  const int num_heads = gridDim.x;
153
  const int num_queries_per_kv = num_heads / num_kv_heads;
zhuwenwen's avatar
zhuwenwen committed
154
155
156
  const int kv_head_idx = head_idx / num_queries_per_kv;
  const float alibi_slope =
      alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
157
158
159
160
161
162

  // A vector type to store a part of a key or a query.
  // The vector size is configured in such a way that the threads in a thread
  // group fetch or compute 16 bytes at a time. For example, if the size of a
  // thread group is 4 and the data type is half, then the vector size is 16 /
  // (4 * sizeof(half)) == 2.
zhuwenwen's avatar
zhuwenwen committed
163
  constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
  using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
  using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
  using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;

  constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE;
  constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE;

  const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE;
  const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE;

  // Load the query to registers.
  // Each thread in a thread group has a different part of the query.
  // For example, if the the thread group size is 4, then the first thread in
  // the group has 0, 4, 8, ... th vectors of the query, and the second thread
  // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because
  // q is split from a qkv tensor, it may not be contiguous.
zhuwenwen's avatar
zhuwenwen committed
180
181
182
183
184
185
186
187
188
189
190
  const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
  __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
#pragma unroll
  for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD;
       i += NUM_THREAD_GROUPS) {
    const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
    q_vecs[thread_group_offset][i] =
        *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
  }
  __syncthreads();  // TODO(naed90): possible speedup if this is replaced with a
                    // memory wall right before we use q_vecs
191
192
193
194
195
196

  // Memory planning.
  extern __shared__ char shared_mem[];
  // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy.
  float* logits = reinterpret_cast<float*>(shared_mem);
  // Workspace for reduction.
zhuwenwen's avatar
zhuwenwen committed
197
  __shared__ float red_smem[2 * NUM_WARPS];
198
199
200
201

  // x == THREAD_GROUP_SIZE * VEC_SIZE
  // Each thread group fetches x elements from the key at a time.
  constexpr int x = 16 / sizeof(cache_t);
zhuwenwen's avatar
zhuwenwen committed
202
  float qk_max = -FLT_MAX;
203
204
205
206
207
208

  // Iterate over the key blocks.
  // Each warp fetches a block of keys for each iteration.
  // Each thread group in a warp fetches a key from the block, and computes
  // dot product with the query.
  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
zhuwenwen's avatar
zhuwenwen committed
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229

  // blocksparse specific vars
  int bs_block_offset;
  int q_bs_block_id;
  if constexpr (IS_BLOCK_SPARSE) {
    // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len,
    // blocksparse_block_size);
    q_bs_block_id = (seq_len - 1) / blocksparse_block_size;
    if (blocksparse_head_sliding_step >= 0)
      // sliding on q heads
      bs_block_offset =
          (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1;
    else
      // sliding on kv heads
      bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) *
                            (-blocksparse_head_sliding_step) +
                        1;
  }

  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
       block_idx += NUM_WARPS) {
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
    // int64 because int32 can lead to overflow when this variable is multiplied
    // by large numbers (e.g., kv_block_stride).
    // For blocksparse attention: skip computation on blocks that are not
    // attended
    if constexpr (IS_BLOCK_SPARSE) {
      const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
      const bool is_remote =
          ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0);
      const bool is_local =
          (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks);
      if (!is_remote && !is_local) {
        for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
          const int physical_block_offset =
              (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
          const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;

          if (thread_group_offset == 0) {
            // NOTE(linxihui): assign very large number to skipped tokens to
            // avoid contribution to the sumexp softmax normalizer. This will
            // not be used at computing sum(softmax*v) as the blocks will be
            // skipped.
            logits[token_idx - start_token_idx] = -FLT_MAX;
          }
        }
        continue;
      }
    }
zhuwenwen's avatar
zhuwenwen committed
258
259
    const int64_t physical_block_number =
        static_cast<int64_t>(block_table[block_idx]);
260
261
262
263
264
265
266

    // Load a key to registers.
    // Each thread in a thread group has a different part of the key.
    // For example, if the the thread group size is 4, then the first thread in
    // the group has 0, 4, 8, ... th vectors of the key, and the second thread
    // has 1, 5, 9, ... th vectors of the key, and so on.
    for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
zhuwenwen's avatar
zhuwenwen committed
267
268
      const int physical_block_offset =
          (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
269
270
      const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
      K_vec k_vecs[NUM_VECS_PER_THREAD];
zhuwenwen's avatar
zhuwenwen committed
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289

#pragma unroll
      for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
        const cache_t* k_ptr =
            k_cache + physical_block_number * kv_block_stride +
            kv_head_idx * kv_head_stride + physical_block_offset * x;
        const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE;
        const int offset1 = (vec_idx * VEC_SIZE) / x;
        const int offset2 = (vec_idx * VEC_SIZE) % x;

        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
          k_vecs[j] = *reinterpret_cast<const K_vec*>(
              k_ptr + offset1 * BLOCK_SIZE * x + offset2);
        } else {
          // Vector conversion from Quant_vec to K_vec.
          Quant_vec k_vec_quant = *reinterpret_cast<const Quant_vec*>(
              k_ptr + offset1 * BLOCK_SIZE * x + offset2);
          k_vecs[j] = fp8::scaled_convert<K_vec, Quant_vec, KV_DTYPE>(
              k_vec_quant, k_scale);
290
        }
Woosuk Kwon's avatar
Woosuk Kwon committed
291
      }
zhuwenwen's avatar
zhuwenwen committed
292

Woosuk Kwon's avatar
Woosuk Kwon committed
293
294
      // Compute dot product.
      // This includes a reduction across the threads in the same thread group.
zhuwenwen's avatar
zhuwenwen committed
295
296
      float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(
                             q_vecs[thread_group_offset], k_vecs);
Woosuk Kwon's avatar
Woosuk Kwon committed
297
      // Add the ALiBi bias if slopes are given.
298
      qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
zhuwenwen's avatar
zhuwenwen committed
299

300
301
302
303
304
305
306
307
      // used for tree-style attention
      if (attn_masks != nullptr) {
        const int* attn_masks_ptr = attn_masks + seq_idx * attn_masks_stride;
        if (attn_masks_ptr[token_idx] == 0) {
          qk = -FLT_MAX;
        }
      }

Woosuk Kwon's avatar
Woosuk Kwon committed
308
309
310
      if (thread_group_offset == 0) {
        // Store the partial reductions to shared memory.
        // NOTE(woosuk): It is required to zero out the masked logits.
311
312
        const bool mask = token_idx >= seq_len;
        logits[token_idx - start_token_idx] = mask ? 0.f: qk;
Woosuk Kwon's avatar
Woosuk Kwon committed
313
        // Update the max value.
zhuwenwen's avatar
zhuwenwen committed
314
        qk_max = mask ? qk_max : fmaxf(qk_max, qk);
Woosuk Kwon's avatar
Woosuk Kwon committed
315
316
317
318
319
320
321
      }
    }
  }

  // Perform reduction across the threads in the same warp to get the
  // max qk value for each "warp" (not across the thread block yet).
  // The 0-th thread of each thread group already has its max qk value.
zhuwenwen's avatar
zhuwenwen committed
322
323
324
325
326
327
328
329
#pragma unroll
  for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
  }
  if (lane == 0) {
    red_smem[warp_idx] = qk_max;
  }
  __syncthreads();
Woosuk Kwon's avatar
Woosuk Kwon committed
330

zhuwenwen's avatar
zhuwenwen committed
331
332
333
334
335
336
337
338
339
  // TODO(woosuk): Refactor this part.
  // Get the max qk value for the sequence.
  qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
#pragma unroll
  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
  }
  // Broadcast the max qk value to all threads.
  qk_max = VLLM_SHFL_SYNC(qk_max, 0);
Woosuk Kwon's avatar
Woosuk Kwon committed
340

zhuwenwen's avatar
zhuwenwen committed
341
342
343
344
345
346
347
348
  // Get the sum of the exp values.
  float exp_sum = 0.f;
  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
    float val = __expf(logits[i] - qk_max);
    logits[i] = val;
    exp_sum += val;
  }
  exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
Woosuk Kwon's avatar
Woosuk Kwon committed
349

zhuwenwen's avatar
zhuwenwen committed
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
  // Compute softmax.
  const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
    logits[i] *= inv_sum;
  }
  __syncthreads();

  // If partitioning is enabled, store the max logit and exp_sum.
  if (USE_PARTITIONING && thread_idx == 0) {
    float* max_logits_ptr = max_logits +
                            seq_idx * num_heads * max_num_partitions +
                            head_idx * max_num_partitions + partition_idx;
    *max_logits_ptr = qk_max;
    float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions +
                          head_idx * max_num_partitions + partition_idx;
    *exp_sums_ptr = exp_sum;
366
  }
zhuwenwen's avatar
zhuwenwen committed
367

Woosuk Kwon's avatar
Woosuk Kwon committed
368
369
370
371
  // Each thread will fetch 16 bytes from the value cache at a time.
  constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE);
  using V_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
  using L_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
372
  using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
Woosuk Kwon's avatar
Woosuk Kwon committed
373
374
375
376
  using Float_L_vec = typename FloatVec<L_vec>::Type;

  constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE;
  constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW;
377
378
  constexpr int NUM_ROWS_PER_THREAD =
      DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
Woosuk Kwon's avatar
Woosuk Kwon committed
379
380

  // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
zhuwenwen's avatar
zhuwenwen committed
381
382
383
384
  float accs[NUM_ROWS_PER_THREAD];
#pragma unroll
  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
    accs[i] = 0.f;
zhangshao's avatar
zhangshao committed
385
  }
zhuwenwen's avatar
zhuwenwen committed
386

387
388
  scalar_t zero_value;
  zero(zero_value);
389
390
  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
       block_idx += NUM_WARPS) {
zhuwenwen's avatar
zhuwenwen committed
391
392
393
394
395
396
397
398
399
400
401
402
    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
    // int64 because int32 can lead to overflow when this variable is multiplied
    // by large numbers (e.g., kv_block_stride).
    // For blocksparse attention: skip computation on blocks that are not
    // attended
    if constexpr (IS_BLOCK_SPARSE) {
      int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
      if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) &&
          !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) {
        continue;
      }
    }
403
404
    const int64_t physical_block_number =
        static_cast<int64_t>(block_table[block_idx]);
Woosuk Kwon's avatar
Woosuk Kwon committed
405
406
407
    const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE;
    const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
    L_vec logits_vec;
zhuwenwen's avatar
zhuwenwen committed
408
409
410
411
412
    from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + token_idx -
                                                           start_token_idx));

    const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride +
                           kv_head_idx * kv_head_stride;
Woosuk Kwon's avatar
Woosuk Kwon committed
413
414
415
416
417
#pragma unroll
    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
      if (row_idx < HEAD_SIZE) {
        const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
zhuwenwen's avatar
zhuwenwen committed
418
        V_vec v_vec;
419
420
421
422

        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
          v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
        } else {
423
424
          V_quant_vec v_quant_vec =
              *reinterpret_cast<const V_quant_vec*>(v_ptr + offset);
425
          // Vector conversion from V_quant_vec to V_vec.
426
          v_vec = fp8::scaled_convert<V_vec, V_quant_vec, KV_DTYPE>(v_quant_vec,
427
                                                                    v_scale);
428
        }
429
        if (block_idx == num_seq_blocks - 1) {
430
431
432
433
          // NOTE(woosuk): When v_vec contains the tokens that are out of the
          // context, we should explicitly zero out the values since they may
          // contain NaNs. See
          // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
434
435
          scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
#pragma unroll
436
          for (int j = 0; j < V_VEC_SIZE; j++) {
437
            v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
438
439
          }
        }
zhuwenwen's avatar
zhuwenwen committed
440
        accs[i] += dot(logits_vec, v_vec);
Woosuk Kwon's avatar
Woosuk Kwon committed
441
442
443
444
445
446
447
      }
    }
  }

  // Perform reduction within each warp.
#pragma unroll
  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
zhuwenwen's avatar
zhuwenwen committed
448
    float acc = accs[i];
Woosuk Kwon's avatar
Woosuk Kwon committed
449
450
#pragma unroll
    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
451
      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
452
    }
zhuwenwen's avatar
zhuwenwen committed
453
    accs[i] = acc;
Woosuk Kwon's avatar
Woosuk Kwon committed
454
455
  }

456
457
  // NOTE(woosuk): A barrier is required because the shared memory space for
  // logits is reused for the output.
Woosuk Kwon's avatar
Woosuk Kwon committed
458
459
460
461
462
463
464
465
466
  __syncthreads();

  // Perform reduction across warps.
  float* out_smem = reinterpret_cast<float*>(shared_mem);
#pragma unroll
  for (int i = NUM_WARPS; i > 1; i /= 2) {
    int mid = i / 2;
    // Upper warps write to shared memory.
    if (warp_idx >= mid && warp_idx < i) {
zhuwenwen's avatar
zhuwenwen committed
467
      float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
Woosuk Kwon's avatar
Woosuk Kwon committed
468
469
470
471
#pragma unroll
      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
zhuwenwen's avatar
zhuwenwen committed
472
          dst[row_idx] = accs[i];
Woosuk Kwon's avatar
Woosuk Kwon committed
473
474
475
476
477
478
479
        }
      }
    }
    __syncthreads();

    // Lower warps update the output.
    if (warp_idx < mid) {
zhuwenwen's avatar
zhuwenwen committed
480
      const float* src = &out_smem[warp_idx * HEAD_SIZE];
Woosuk Kwon's avatar
Woosuk Kwon committed
481
482
483
484
#pragma unroll
      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
zhuwenwen's avatar
zhuwenwen committed
485
          accs[i] += src[row_idx];
Woosuk Kwon's avatar
Woosuk Kwon committed
486
487
488
489
490
491
492
493
        }
      }
    }
    __syncthreads();
  }

  // Write the final output.
  if (warp_idx == 0) {
494
495
496
    scalar_t* out_ptr =
        out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
        head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE;
Woosuk Kwon's avatar
Woosuk Kwon committed
497
#pragma unroll
zhuwenwen's avatar
zhuwenwen committed
498
499
500
501
    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
      if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
        from_float(*(out_ptr + row_idx), accs[i]);
Woosuk Kwon's avatar
Woosuk Kwon committed
502
503
504
505
506
      }
    }
  }
}

507
// Grid: (num_heads, num_seqs, 1).
508
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
509
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
zhuwenwen's avatar
zhuwenwen committed
510
          bool IS_BLOCK_SPARSE>
zhuwenwen's avatar
zhuwenwen committed
511
__global__ void paged_attention_v1_kernel(
512
513
514
515
516
517
518
519
520
521
522
523
524
    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
                                          // head_size/x, block_size, x]
    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
                                          // head_size, block_size]
    const int num_kv_heads,               // [num_heads]
    const float scale,
    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ seq_lens,      // [num_seqs]
    const int max_num_blocks_per_seq,
    const float* __restrict__ alibi_slopes,  // [num_heads]
    const int q_stride, const int kv_block_stride, const int kv_head_stride,
525
526
    const float k_scale, const float v_scale, const int tp_rank,
    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
527
528
    const int blocksparse_block_size, const int blocksparse_head_sliding_step,
    const int* __restrict__ attn_masks=nullptr, const int attn_masks_stride=0) {
zhuwenwen's avatar
zhuwenwen committed
529
  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
zhuwenwen's avatar
zhuwenwen committed
530
                         KV_DTYPE, IS_BLOCK_SPARSE>(
531
      /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache,
zhuwenwen's avatar
zhuwenwen committed
532
      v_cache, num_kv_heads, scale, block_tables, seq_lens,
533
      max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride,
534
      kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks,
535
      blocksparse_vert_stride, blocksparse_block_size,
536
      blocksparse_head_sliding_step, attn_masks, attn_masks_stride);
537
538
539
}

// Grid: (num_heads, num_seqs, max_num_partitions).
540
541
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
542
          bool IS_BLOCK_SPARSE,
zhuwenwen's avatar
zhuwenwen committed
543
          int PARTITION_SIZE>
zhuwenwen's avatar
zhuwenwen committed
544
__global__ void paged_attention_v2_kernel(
545
546
547
548
549
550
551
552
553
554
    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
    float* __restrict__ max_logits,       // [num_seqs, num_heads,
                                          // max_num_partitions]
    scalar_t* __restrict__ tmp_out,       // [num_seqs, num_heads,
                                          // max_num_partitions, head_size]
    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
                                          // head_size/x, block_size, x]
    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
                                          // head_size, block_size]
zhuwenwen's avatar
zhuwenwen committed
555
    const int num_kv_heads,               // [num_heads]
556
557
558
559
560
561
    const float scale,
    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ seq_lens,      // [num_seqs]
    const int max_num_blocks_per_seq,
    const float* __restrict__ alibi_slopes,  // [num_heads]
    const int q_stride, const int kv_block_stride, const int kv_head_stride,
562
563
    const float k_scale, const float v_scale, const int tp_rank,
    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
564
565
    const int blocksparse_block_size, const int blocksparse_head_sliding_step,
    const int* __restrict__ attn_masks=nullptr, const int attn_masks_stride=0) {
zhuwenwen's avatar
zhuwenwen committed
566
  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
zhuwenwen's avatar
zhuwenwen committed
567
568
                         KV_DTYPE, IS_BLOCK_SPARSE, PARTITION_SIZE>(
      exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
569
      block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride,
570
      kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank,
571
      blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size,
572
      blocksparse_head_sliding_step, attn_masks, attn_masks_stride);
573
574
575
}

// Grid: (num_heads, num_seqs).
576
577
template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
          int PARTITION_SIZE>
zhuwenwen's avatar
zhuwenwen committed
578
__global__ void paged_attention_v2_reduce_kernel(
579
580
581
582
583
584
585
586
587
    scalar_t* __restrict__ out,            // [num_seqs, num_heads, head_size]
    const float* __restrict__ exp_sums,    // [num_seqs, num_heads,
                                           // max_num_partitions]
    const float* __restrict__ max_logits,  // [num_seqs, num_heads,
                                           // max_num_partitions]
    const scalar_t* __restrict__ tmp_out,  // [num_seqs, num_heads,
                                           // max_num_partitions, head_size]
    const int* __restrict__ seq_lens,      // [num_seqs]
    const int max_num_partitions) {
588
589
590
  const int num_heads = gridDim.x;
  const int head_idx = blockIdx.x;
  const int seq_idx = blockIdx.y;
591
592
  const int seq_len = seq_lens[seq_idx];
  const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
593
594
  if (num_partitions == 1) {
    // No need to reduce. Only copy tmp_out to out.
595
596
597
598
599
    scalar_t* out_ptr =
        out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
    const scalar_t* tmp_out_ptr =
        tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
        head_idx * max_num_partitions * HEAD_SIZE;
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
    for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) {
      out_ptr[i] = tmp_out_ptr[i];
    }
    // Terminate the thread block.
    return;
  }

  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  const int warp_idx = threadIdx.x / WARP_SIZE;
  const int lane = threadIdx.x % WARP_SIZE;

  // Size: 2 * num_partitions.
  extern __shared__ char shared_mem[];
  // Workspace for reduction.
  __shared__ float red_smem[2 * NUM_WARPS];

  // Load max logits to shared memory.
  float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
618
619
620
  const float* max_logits_ptr = max_logits +
                                seq_idx * num_heads * max_num_partitions +
                                head_idx * max_num_partitions;
621
622
623
624
625
626
627
628
629
630
631
632
  float max_logit = -FLT_MAX;
  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
    const float l = max_logits_ptr[i];
    shared_max_logits[i] = l;
    max_logit = fmaxf(max_logit, l);
  }
  __syncthreads();

  // Get the global max logit.
  // Reduce within the warp.
#pragma unroll
  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
633
    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
634
635
636
637
638
639
640
641
642
  }
  if (lane == 0) {
    red_smem[warp_idx] = max_logit;
  }
  __syncthreads();
  // Reduce across warps.
  max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
#pragma unroll
  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
643
    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
644
645
  }
  // Broadcast the max value to all threads.
646
  max_logit = VLLM_SHFL_SYNC(max_logit, 0);
647
648

  // Load rescaled exp sums to shared memory.
649
650
651
652
653
  float* shared_exp_sums =
      reinterpret_cast<float*>(shared_mem + sizeof(float) * num_partitions);
  const float* exp_sums_ptr = exp_sums +
                              seq_idx * num_heads * max_num_partitions +
                              head_idx * max_num_partitions;
654
655
656
657
658
659
660
661
662
663
664
665
  float global_exp_sum = 0.0f;
  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
    float l = shared_max_logits[i];
    float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit);
    global_exp_sum += rescaled_exp_sum;
    shared_exp_sums[i] = rescaled_exp_sum;
  }
  __syncthreads();
  global_exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], global_exp_sum);
  const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f);

  // Aggregate tmp_out to out.
666
667
668
669
670
  const scalar_t* tmp_out_ptr =
      tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
      head_idx * max_num_partitions * HEAD_SIZE;
  scalar_t* out_ptr =
      out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
671
672
673
674
#pragma unroll
  for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) {
    float acc = 0.0f;
    for (int j = 0; j < num_partitions; ++j) {
675
676
      acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] *
             inv_global_exp_sum;
677
678
679
680
681
    }
    from_float(out_ptr[i], acc);
  }
}

682
683
684
685
}  // namespace vllm

#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
zhuwenwen's avatar
zhuwenwen committed
686
      ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
687
                                              BLOCK_SIZE, NUM_THREADS,      \
zhuwenwen's avatar
zhuwenwen committed
688
                                              KV_DTYPE, IS_BLOCK_SPARSE>),  \
689
      shared_mem_size);                                                     \
zhuwenwen's avatar
zhuwenwen committed
690
  vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
zhuwenwen's avatar
zhuwenwen committed
691
692
693
                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE>   \
      <<<grid, block, shared_mem_size, stream>>>(                           \
          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
694
695
          scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq,    \
          alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride,      \
696
          k_scale, v_scale, tp_rank, blocksparse_local_blocks,              \
697
          blocksparse_vert_stride, blocksparse_block_size,                  \
698
699
          blocksparse_head_sliding_step, attn_masks_ptr,                  \
          attn_masks_stride);
Woosuk Kwon's avatar
Woosuk Kwon committed
700
701

// TODO(woosuk): Tune NUM_THREADS.
702
template <typename T, typename CACHE_T, int BLOCK_SIZE,
zhuwenwen's avatar
zhuwenwen committed
703
704
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
          int NUM_THREADS = 128>
705
void paged_attention_v1_launcher(
706
707
708
    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
    torch::Tensor& value_cache, int num_kv_heads, float scale,
    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
709
710
    const c10::optional<torch::Tensor>& alibi_slopes, float k_scale,
    float v_scale, const int tp_rank, const int blocksparse_local_blocks,
711
    const int blocksparse_vert_stride, const int blocksparse_block_size,
712
713
714
    const int blocksparse_head_sliding_step,
    const c10::optional<torch::Tensor>& attn_masks,
    const int64_t attn_masks_stride=0) {
Woosuk Kwon's avatar
Woosuk Kwon committed
715
716
717
718
  int num_seqs = query.size(0);
  int num_heads = query.size(1);
  int head_size = query.size(2);
  int max_num_blocks_per_seq = block_tables.size(1);
Zhuohan Li's avatar
Zhuohan Li committed
719
720
721
  int q_stride = query.stride(0);
  int kv_block_stride = key_cache.stride(0);
  int kv_head_stride = key_cache.stride(1);
Woosuk Kwon's avatar
Woosuk Kwon committed
722

723
  [[maybe_unused]] int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
Woosuk Kwon's avatar
Woosuk Kwon committed
724
725
  assert(head_size % thread_group_size == 0);

Woosuk Kwon's avatar
Woosuk Kwon committed
726
  // NOTE: alibi_slopes is optional.
727
728
729
730
  const float* alibi_slopes_ptr =
      alibi_slopes
          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
          : nullptr;
Woosuk Kwon's avatar
Woosuk Kwon committed
731

Woosuk Kwon's avatar
Woosuk Kwon committed
732
733
  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
734
735
  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
Woosuk Kwon's avatar
Woosuk Kwon committed
736
  int* block_tables_ptr = block_tables.data_ptr<int>();
737
  int* seq_lens_ptr = seq_lens.data_ptr<int>();
Woosuk Kwon's avatar
Woosuk Kwon committed
738

739
740
741
742
743
744
  // NOTE: attn_masks is optional.
  const int* attn_masks_ptr =
      attn_masks
          ? attn_masks.value().data_ptr<int>()
          : nullptr;

zhuwenwen's avatar
zhuwenwen committed
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  int padded_max_seq_len =
      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
  int logits_size = padded_max_seq_len * sizeof(float);
  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
  // Keep that in sync with the logic here!
  int shared_mem_size = std::max(logits_size, outputs_size);

  dim3 grid(num_heads, num_seqs, 1);
  dim3 block(NUM_THREADS);
  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  switch (head_size) {
    // NOTE(woosuk): To reduce the compilation time, we only compile for the
    // head sizes that we use in the model. However, we can easily extend this
    // to support any head size which is a multiple of 16.
    case 64:
      LAUNCH_PAGED_ATTENTION_V1(64);
      break;
    case 80:
      LAUNCH_PAGED_ATTENTION_V1(80);
      break;
    case 96:
      LAUNCH_PAGED_ATTENTION_V1(96);
      break;
    case 112:
      LAUNCH_PAGED_ATTENTION_V1(112);
      break;
    case 120:
      LAUNCH_PAGED_ATTENTION_V1(120);
      break;
    case 128:
      LAUNCH_PAGED_ATTENTION_V1(128);
      break;
780
781
782
    case 160:
      LAUNCH_PAGED_ATTENTION_V1(160);
      break;
zhuwenwen's avatar
zhuwenwen committed
783
784
785
786
787
788
789
790
791
792
    case 192:
      LAUNCH_PAGED_ATTENTION_V1(192);
      break;
    case 256:
      LAUNCH_PAGED_ATTENTION_V1(256);
      break;
    default:
      TORCH_CHECK(false, "Unsupported head size: ", head_size);
      break;
  }
793
794
}

795
796
797
#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)  \
  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,              \
                              IS_BLOCK_SPARSE>(                              \
798
      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
799
      seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank,        \
800
      blocksparse_local_blocks, blocksparse_vert_stride,                     \
801
802
      blocksparse_block_size, blocksparse_head_sliding_step,                 \
      attn_masks, attn_masks_stride);
803
804
805
806
807
808
809
810
811
812

#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
  switch (is_block_sparse) {                                               \
    case true:                                                             \
      CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);     \
      break;                                                               \
    case false:                                                            \
      CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);    \
      break;                                                               \
  }
813
814
815

// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
816
817
#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
  switch (block_size) {                                           \
818
819
820
    case 8:                                                       \
      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
      break;                                                      \
821
    case 16:                                                      \
822
      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
823
      break;                                                      \
824
825
826
    case 32:                                                      \
      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
      break;                                                      \
827
828
829
    default:                                                      \
      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
      break;                                                      \
830
831
  }

zhuwenwen's avatar
zhuwenwen committed
832
void paged_attention_v1(
833
834
835
836
837
    torch::Tensor& out,    // [num_seqs, num_heads, head_size]
    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
    torch::Tensor&
        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
    torch::Tensor&
838
839
840
        value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
841
842
    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
    torch::Tensor& seq_lens,      // [num_seqs]
843
    int64_t block_size, int64_t max_seq_len,
844
    const c10::optional<torch::Tensor>& alibi_slopes,
845
846
    const std::string& kv_cache_dtype, double k_scale, double v_scale,
    const int64_t tp_rank, const int64_t blocksparse_local_blocks,
847
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
848
849
850
    const int64_t blocksparse_head_sliding_step,
    const c10::optional<torch::Tensor>& attn_masks, // [num_seqs, max_seq_len]
    const int64_t attn_masks_stride) {
851
852
853
854
855
  const bool is_block_sparse = (blocksparse_vert_stride > 1);

  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                             CALL_V1_LAUNCHER_BLOCK_SIZE)
}
856
857

#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \
zhuwenwen's avatar
zhuwenwen committed
858
  vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \
859
                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \
zhuwenwen's avatar
zhuwenwen committed
860
861
                                  PARTITION_SIZE>                              \
      <<<grid, block, shared_mem_size, stream>>>(                              \
862
          exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
zhuwenwen's avatar
zhuwenwen committed
863
          value_cache_ptr, num_kv_heads, scale, block_tables_ptr,              \
864
          seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride,    \
865
          kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank,          \
866
          blocksparse_local_blocks, blocksparse_vert_stride,                   \
867
868
          blocksparse_block_size, blocksparse_head_sliding_step,               \
          attn_masks_ptr, attn_masks_stride);              \
zhuwenwen's avatar
zhuwenwen committed
869
  vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \
zhuwenwen's avatar
zhuwenwen committed
870
871
                                         PARTITION_SIZE>                       \
      <<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                \
872
873
874
875
          out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \
          max_num_partitions);

template <typename T, typename CACHE_T, int BLOCK_SIZE,
876
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
zhuwenwen's avatar
zhuwenwen committed
877
          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
878
void paged_attention_v2_launcher(
879
880
881
882
    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
    torch::Tensor& value_cache, int num_kv_heads, float scale,
    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
883
884
    const c10::optional<torch::Tensor>& alibi_slopes, float k_scale,
    float v_scale, const int tp_rank, const int blocksparse_local_blocks,
885
    const int blocksparse_vert_stride, const int blocksparse_block_size,
886
887
888
    const int blocksparse_head_sliding_step,
    const c10::optional<torch::Tensor>& attn_masks,
    const int attn_masks_stride) {
889
890
891
892
893
894
895
896
  int num_seqs = query.size(0);
  int num_heads = query.size(1);
  int head_size = query.size(2);
  int max_num_blocks_per_seq = block_tables.size(1);
  int q_stride = query.stride(0);
  int kv_block_stride = key_cache.stride(0);
  int kv_head_stride = key_cache.stride(1);

897
  [[maybe_unused]] int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
898
899
900
  assert(head_size % thread_group_size == 0);

  // NOTE: alibi_slopes is optional.
901
902
903
904
  const float* alibi_slopes_ptr =
      alibi_slopes
          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
          : nullptr;
905
906
907
908
909
910

  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
  float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
  float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
  T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
911
912
  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
913
  int* block_tables_ptr = block_tables.data_ptr<int>();
914
  int* seq_lens_ptr = seq_lens.data_ptr<int>();
915

916
917
918
919
  // NOTE: attn_masks is optional.
  const int* attn_masks_ptr =
      attn_masks ? attn_masks.value().data_ptr<int>() : nullptr;

920
  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
921
  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
zhuwenwen's avatar
zhuwenwen committed
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
  int logits_size = PARTITION_SIZE * sizeof(float);
  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);

  // For paged attention v2 kernel.
  dim3 grid(num_heads, num_seqs, max_num_partitions);
  int shared_mem_size = std::max(logits_size, outputs_size);
  // For paged attention v2 reduce kernel.
  dim3 reduce_grid(num_heads, num_seqs);
  int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);

  dim3 block(NUM_THREADS);
  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  switch (head_size) {
    // NOTE(woosuk): To reduce the compilation time, we only compile for the
    // head sizes that we use in the model. However, we can easily extend this
    // to support any head size which is a multiple of 16.
    case 64:
      LAUNCH_PAGED_ATTENTION_V2(64);
      break;
    case 80:
      LAUNCH_PAGED_ATTENTION_V2(80);
      break;
    case 96:
      LAUNCH_PAGED_ATTENTION_V2(96);
      break;
    case 112:
      LAUNCH_PAGED_ATTENTION_V2(112);
      break;
    case 120:
      LAUNCH_PAGED_ATTENTION_V2(120);
      break;
    case 128:
      LAUNCH_PAGED_ATTENTION_V2(128);
      break;
957
958
959
    case 160:
      LAUNCH_PAGED_ATTENTION_V2(160);
      break;
zhuwenwen's avatar
zhuwenwen committed
960
961
962
963
964
965
966
967
968
969
    case 192:
      LAUNCH_PAGED_ATTENTION_V2(192);
      break;
    case 256:
      LAUNCH_PAGED_ATTENTION_V2(256);
      break;
    default:
      TORCH_CHECK(false, "Unsupported head size: ", head_size);
      break;
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
970
971
}

972
973
974
#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)   \
  paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,               \
                              IS_BLOCK_SPARSE>(                               \
975
976
      out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,      \
      num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \
977
978
      k_scale, v_scale, tp_rank, blocksparse_local_blocks,                    \
      blocksparse_vert_stride, blocksparse_block_size,                        \
979
      blocksparse_head_sliding_step, attn_masks, attn_masks_stride);
980
981
982
983
984
985
986
987
988
989

#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
  switch (is_block_sparse) {                                               \
    case true:                                                             \
      CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);     \
      break;                                                               \
    case false:                                                            \
      CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);    \
      break;                                                               \
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
990

Woosuk Kwon's avatar
Woosuk Kwon committed
991
992
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
993
994
#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
  switch (block_size) {                                           \
995
996
997
    case 8:                                                       \
      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
      break;                                                      \
998
    case 16:                                                      \
999
      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
1000
      break;                                                      \
1001
1002
1003
    case 32:                                                      \
      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
      break;                                                      \
1004
1005
1006
    default:                                                      \
      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
      break;                                                      \
Woosuk Kwon's avatar
Woosuk Kwon committed
1007
1008
  }

zhuwenwen's avatar
zhuwenwen committed
1009
void paged_attention_v2(
1010
1011
1012
1013
1014
1015
1016
1017
1018
    torch::Tensor& out,         // [num_seqs, num_heads, head_size]
    torch::Tensor& exp_sums,    // [num_seqs, num_heads, max_num_partitions]
    torch::Tensor& max_logits,  // [num_seqs, num_heads, max_num_partitions]
    torch::Tensor&
        tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
    torch::Tensor&
        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
    torch::Tensor&
1019
1020
1021
        value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
1022
1023
    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
    torch::Tensor& seq_lens,      // [num_seqs]
1024
    int64_t block_size, int64_t max_seq_len,
1025
    const c10::optional<torch::Tensor>& alibi_slopes,
1026
1027
    const std::string& kv_cache_dtype, double k_scale, double v_scale,
    const int64_t tp_rank, const int64_t blocksparse_local_blocks,
1028
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
1029
1030
1031
    const int64_t blocksparse_head_sliding_step,
    const c10::optional<torch::Tensor>& attn_masks, // [num_seqs, max_seq_len]
    const int64_t attn_masks_stride) {
1032
  const bool is_block_sparse = (blocksparse_vert_stride > 1);
1033
1034
  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                             CALL_V2_LAUNCHER_BLOCK_SIZE)
Woosuk Kwon's avatar
Woosuk Kwon committed
1035
1036
1037
1038
1039
}

#undef WARP_SIZE
#undef MAX
#undef MIN
zhuwenwen's avatar
zhuwenwen committed
1040
#undef DIVIDE_ROUND_UP