"components/backends/vllm/deepseek-r1.md" did not exist on "3b3d0f2e09feba4335439951c7911c6c69c95041"
attention_kernels.cu 39.7 KB
Newer Older
1
2
/*
 * 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
Woosuk Kwon's avatar
Woosuk Kwon committed
3
 * Copyright (c) 2023, The vLLM team.
4
5
6
7
8
9
10
11
12
13
14
15
16
17
 * 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.
 */
18

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

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

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

#ifndef USE_ROCM
#define WARP_SIZE 32
#else
#define WARP_SIZE warpSize
#endif

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

Woosuk Kwon's avatar
Woosuk Kwon committed
45
namespace vllm {
Woosuk Kwon's avatar
Woosuk Kwon committed
46
47
48
49
50
51
52
53
54
55
56

// Utility function for attention softmax.
template<int NUM_WARPS>
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) {
57
    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
  }

  // 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) {
76
    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
77
78
79
  }

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

83
84
// TODO(woosuk): Merge the last two dimensions of the grid.
// Grid: (num_heads, num_seqs, max_num_partitions).
Woosuk Kwon's avatar
Woosuk Kwon committed
85
86
template<
  typename scalar_t,
87
  typename cache_t,
Woosuk Kwon's avatar
Woosuk Kwon committed
88
89
  int HEAD_SIZE,
  int BLOCK_SIZE,
90
  int NUM_THREADS,
91
  vllm::Fp8KVCacheDataType KV_DTYPE,
92
93
94
95
96
  int PARTITION_SIZE = 0> // Zero means no partitioning.
__device__ void paged_attention_kernel(
  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]
Woosuk Kwon's avatar
Woosuk Kwon committed
97
  const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
98
99
  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]
100
  const int num_kv_heads,                 // [num_heads]
Woosuk Kwon's avatar
Woosuk Kwon committed
101
102
  const float scale,
  const int* __restrict__ block_tables,   // [num_seqs, max_num_blocks_per_seq]
103
  const int* __restrict__ seq_lens,   // [num_seqs]
Woosuk Kwon's avatar
Woosuk Kwon committed
104
  const int max_num_blocks_per_seq,
Woosuk Kwon's avatar
Woosuk Kwon committed
105
  const float* __restrict__ alibi_slopes, // [num_heads]
Zhuohan Li's avatar
Zhuohan Li committed
106
107
  const int q_stride,
  const int kv_block_stride,
108
109
  const int kv_head_stride,
  const float kv_scale) {
110
111
112
113
  const int seq_idx = blockIdx.y;
  const int partition_idx = blockIdx.z;
  const int max_num_partitions = gridDim.z;
  constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0;
114
115
  const int seq_len = seq_lens[seq_idx];
  if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) {
116
117
118
119
    // No work to do. Terminate the thread block.
    return;
  }

120
121
  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;
122
123
124

  // [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;
125
  const int end_block_idx = MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks);
126
127
128
129
  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;
130
  const int end_token_idx = MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len);
131
132
  const int num_tokens = end_token_idx - start_token_idx;

Woosuk Kwon's avatar
Woosuk Kwon committed
133
  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
134
135
  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);
136
  constexpr int NUM_TOKENS_PER_THREAD_GROUP = DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE);
Woosuk Kwon's avatar
Woosuk Kwon committed
137
138
139
140
141
142
143
  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  const int thread_idx = threadIdx.x;
  const int warp_idx = thread_idx / WARP_SIZE;
  const int lane = thread_idx % WARP_SIZE;

  const int head_idx = blockIdx.x;
  const int num_heads = gridDim.x;
144
145
  const int num_queries_per_kv = num_heads / num_kv_heads;
  const int kv_head_idx = head_idx / num_queries_per_kv;
Woosuk Kwon's avatar
Woosuk Kwon committed
146
  const float alibi_slope = alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
Woosuk Kwon's avatar
Woosuk Kwon committed
147
148
149
150
151
152
153
154
155

  // 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.
  constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
  using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
  using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
156
  using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;
Woosuk Kwon's avatar
Woosuk Kwon committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170

  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.
  const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
171
  __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
Woosuk Kwon's avatar
Woosuk Kwon committed
172
#pragma unroll
173
  for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD; i += NUM_THREAD_GROUPS) {
Woosuk Kwon's avatar
Woosuk Kwon committed
174
    const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
175
    q_vecs[thread_group_offset][i] = *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
Woosuk Kwon's avatar
Woosuk Kwon committed
176
  }
177
  __syncthreads(); // TODO(naed90): possible speedup if this is replaced with a memory wall right before we use q_vecs
Woosuk Kwon's avatar
Woosuk Kwon committed
178
179
180
181
182
183
184
185
186
187

  // 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.
  __shared__ float red_smem[2 * NUM_WARPS];

  // x == THREAD_GROUP_SIZE * VEC_SIZE
  // Each thread group fetches x elements from the key at a time.
188
  constexpr int x = 16 / sizeof(cache_t);
Woosuk Kwon's avatar
Woosuk Kwon committed
189
190
191
192
193
194
  float qk_max = -FLT_MAX;

  // 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.
195
196
  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; block_idx += NUM_WARPS) {
197
198
199
200
    // 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).
    const int64_t physical_block_number = static_cast<int64_t>(block_table[block_idx]);
Woosuk Kwon's avatar
Woosuk Kwon committed
201
202
203
204
205
206
207
208
209
210
211
212
213

    // 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++) {
      const int physical_block_offset = (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
      const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
      K_vec k_vecs[NUM_VECS_PER_THREAD];

#pragma unroll
      for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
214
215
216
        const cache_t* k_ptr = k_cache + physical_block_number * kv_block_stride
                                       + kv_head_idx * kv_head_stride
                                       + physical_block_offset * x;
Woosuk Kwon's avatar
Woosuk Kwon committed
217
218
219
        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;
220
221

        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
222
          k_vecs[j] = *reinterpret_cast<const K_vec*>(k_ptr + offset1 * BLOCK_SIZE * x + offset2);
223
224
225
226
227
        } 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, kv_scale);
228
        }
Woosuk Kwon's avatar
Woosuk Kwon committed
229
230
231
232
      }

      // Compute dot product.
      // This includes a reduction across the threads in the same thread group.
233
      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
234
      // Add the ALiBi bias if slopes are given.
235
      qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
Woosuk Kwon's avatar
Woosuk Kwon committed
236

Woosuk Kwon's avatar
Woosuk Kwon committed
237
238
239
      if (thread_group_offset == 0) {
        // Store the partial reductions to shared memory.
        // NOTE(woosuk): It is required to zero out the masked logits.
240
        const bool mask = token_idx >= seq_len;
241
        logits[token_idx - start_token_idx] = mask ? 0.f : qk;
Woosuk Kwon's avatar
Woosuk Kwon committed
242
243
244
245
246
247
248
249
250
251
252
        // Update the max value.
        qk_max = mask ? qk_max : fmaxf(qk_max, qk);
      }
    }
  }

  // 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.
#pragma unroll
  for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
253
    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
Woosuk Kwon's avatar
Woosuk Kwon committed
254
255
256
257
258
259
260
261
262
263
264
  }
  if (lane == 0) {
    red_smem[warp_idx] = qk_max;
  }
  __syncthreads();

  // 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) {
265
    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
Woosuk Kwon's avatar
Woosuk Kwon committed
266
267
  }
  // Broadcast the max qk value to all threads.
268
  qk_max = VLLM_SHFL_SYNC(qk_max, 0);
Woosuk Kwon's avatar
Woosuk Kwon committed
269
270
271

  // Get the sum of the exp values.
  float exp_sum = 0.f;
272
  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
Woosuk Kwon's avatar
Woosuk Kwon committed
273
274
275
276
277
278
279
280
    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);

  // Compute softmax.
  const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
281
  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
Woosuk Kwon's avatar
Woosuk Kwon committed
282
283
284
285
    logits[i] *= inv_sum;
  }
  __syncthreads();

286
287
288
289
290
291
292
293
294
295
296
297
  // 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;
  }

Woosuk Kwon's avatar
Woosuk Kwon committed
298
299
300
301
  // 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;
302
  using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
Woosuk Kwon's avatar
Woosuk Kwon committed
303
304
305
306
  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;
307
  constexpr int NUM_ROWS_PER_THREAD = DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
Woosuk Kwon's avatar
Woosuk Kwon committed
308
309
310
311
312
313
314
315

  // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
  float accs[NUM_ROWS_PER_THREAD];
#pragma unroll
  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
    accs[i] = 0.f;
  }

316
317
  scalar_t zero_value;
  zero(zero_value);
318
  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; block_idx += NUM_WARPS) {
319
320
321
322
    // 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).
    const int64_t physical_block_number = static_cast<int64_t>(block_table[block_idx]);
Woosuk Kwon's avatar
Woosuk Kwon committed
323
324
325
    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;
326
    from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + token_idx - start_token_idx));
Woosuk Kwon's avatar
Woosuk Kwon committed
327

328
329
    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
330
331
332
333
334
#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;
335
        V_vec v_vec;
336
337
338
339

        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
          v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
        } else {
340
341
          V_quant_vec v_quant_vec = *reinterpret_cast<const V_quant_vec*>(v_ptr + offset);
          // Vector conversion from V_quant_vec to V_vec.
342
          v_vec = fp8::scaled_convert<V_vec, V_quant_vec, KV_DTYPE>(v_quant_vec, kv_scale);
343
        }
344
        if (block_idx == num_seq_blocks - 1) {
345
346
347
348
349
          // 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
          scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
#pragma unroll
350
          for (int j = 0; j < V_VEC_SIZE; j++) {
351
            v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
352
353
          }
        }
Woosuk Kwon's avatar
Woosuk Kwon committed
354
355
356
357
358
359
360
361
362
363
364
        accs[i] += dot(logits_vec, v_vec);
      }
    }
  }

  // Perform reduction within each warp.
#pragma unroll
  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
    float acc = accs[i];
#pragma unroll
    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
365
      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
    }
    accs[i] = acc;
  }

  // NOTE(woosuk): A barrier is required because the shared memory space for logits
  // is reused for the output.
  __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) {
      float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
#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) {
          dst[row_idx] = accs[i];
        }
      }
    }
    __syncthreads();

    // Lower warps update the output.
    if (warp_idx < mid) {
      const float* src = &out_smem[warp_idx * HEAD_SIZE];
#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) {
          accs[i] += src[row_idx];
        }
      }
    }
    __syncthreads();
  }

  // Write the final output.
  if (warp_idx == 0) {
408
409
410
    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
411
412
413
414
415
416
417
418
419
420
#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) {
        from_float(*(out_ptr + row_idx), accs[i]);
      }
    }
  }
}

421
422
423
// Grid: (num_heads, num_seqs, 1).
template<
  typename scalar_t,
424
  typename cache_t,
425
426
  int HEAD_SIZE,
  int BLOCK_SIZE,
427
  int NUM_THREADS,
428
  vllm::Fp8KVCacheDataType KV_DTYPE>
429
430
431
__global__ void paged_attention_v1_kernel(
  scalar_t* __restrict__ out,             // [num_seqs, num_heads, head_size]
  const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
432
433
  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]
434
  const int num_kv_heads,                 // [num_heads]
435
436
  const float scale,
  const int* __restrict__ block_tables,   // [num_seqs, max_num_blocks_per_seq]
437
  const int* __restrict__ seq_lens,   // [num_seqs]
438
439
440
441
  const int max_num_blocks_per_seq,
  const float* __restrict__ alibi_slopes, // [num_heads]
  const int q_stride,
  const int kv_block_stride,
442
443
  const int kv_head_stride,
  const float kv_scale) {
444
  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS, KV_DTYPE>(
445
    /* exp_sums */ nullptr, /* max_logits */ nullptr,
446
    out, q, k_cache, v_cache, num_kv_heads, scale, block_tables, seq_lens,
447
    max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride, kv_head_stride, kv_scale);
448
449
450
451
452
}

// Grid: (num_heads, num_seqs, max_num_partitions).
template<
  typename scalar_t,
453
  typename cache_t,
454
455
456
  int HEAD_SIZE,
  int BLOCK_SIZE,
  int NUM_THREADS,
457
  vllm::Fp8KVCacheDataType KV_DTYPE,
458
459
460
461
462
463
  int PARTITION_SIZE>
__global__ void paged_attention_v2_kernel(
  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]
464
465
  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]
466
  const int num_kv_heads,                 // [num_heads]
467
468
  const float scale,
  const int* __restrict__ block_tables,   // [num_seqs, max_num_blocks_per_seq]
469
  const int* __restrict__ seq_lens,   // [num_seqs]
470
471
472
473
  const int max_num_blocks_per_seq,
  const float* __restrict__ alibi_slopes, // [num_heads]
  const int q_stride,
  const int kv_block_stride,
474
475
  const int kv_head_stride,
  const float kv_scale) {
476
  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS, KV_DTYPE, PARTITION_SIZE>(
477
    exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
478
    block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes,
479
    q_stride, kv_block_stride, kv_head_stride, kv_scale);
480
481
482
483
484
485
486
487
488
489
490
491
492
}

// Grid: (num_heads, num_seqs).
template<
  typename scalar_t,
  int HEAD_SIZE,
  int NUM_THREADS,
  int PARTITION_SIZE>
__global__ void paged_attention_v2_reduce_kernel(
  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]
493
  const int* __restrict__ seq_lens,   // [num_seqs]
494
495
496
497
  const int max_num_partitions) {
  const int num_heads = gridDim.x;
  const int head_idx = blockIdx.x;
  const int seq_idx = blockIdx.y;
498
499
  const int seq_len = seq_lens[seq_idx];
  const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
  if (num_partitions == 1) {
    // No need to reduce. Only copy tmp_out to out.
    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;
    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);
  const float* max_logits_ptr = max_logits + seq_idx * num_heads * max_num_partitions
                                           + head_idx * max_num_partitions;
  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) {
537
    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
538
539
540
541
542
543
544
545
546
  }
  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) {
547
    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
548
549
  }
  // Broadcast the max value to all threads.
550
  max_logit = VLLM_SHFL_SYNC(max_logit, 0);
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580

  // Load rescaled exp sums to shared memory.
  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;
  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.
  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;
#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) {
      acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] * inv_global_exp_sum;
    }
    from_float(out_ptr[i], acc);
  }
}

Woosuk Kwon's avatar
Woosuk Kwon committed
581
} // namespace vllm
Woosuk Kwon's avatar
Woosuk Kwon committed
582

583
#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                                  \
584
  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                                       \
585
    ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,   \
586
      KV_DTYPE>), shared_mem_size);                                                           \
587
  vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,             \
588
  KV_DTYPE><<<grid, block, shared_mem_size, stream>>>(                                        \
Woosuk Kwon's avatar
Woosuk Kwon committed
589
590
591
592
    out_ptr,                                                                                  \
    query_ptr,                                                                                \
    key_cache_ptr,                                                                            \
    value_cache_ptr,                                                                          \
593
    num_kv_heads,                                                                             \
Woosuk Kwon's avatar
Woosuk Kwon committed
594
595
    scale,                                                                                    \
    block_tables_ptr,                                                                         \
596
    seq_lens_ptr,                                                                              \
Woosuk Kwon's avatar
Woosuk Kwon committed
597
    max_num_blocks_per_seq,                                                                   \
Woosuk Kwon's avatar
Woosuk Kwon committed
598
    alibi_slopes_ptr,                                                                         \
Zhuohan Li's avatar
Zhuohan Li committed
599
600
    q_stride,                                                                                 \
    kv_block_stride,                                                                          \
601
602
    kv_head_stride,                                                                           \
    kv_scale);
Woosuk Kwon's avatar
Woosuk Kwon committed
603
604
605
606

// TODO(woosuk): Tune NUM_THREADS.
template<
  typename T,
607
  typename CACHE_T,
Woosuk Kwon's avatar
Woosuk Kwon committed
608
  int BLOCK_SIZE,
609
  vllm::Fp8KVCacheDataType KV_DTYPE,
Woosuk Kwon's avatar
Woosuk Kwon committed
610
  int NUM_THREADS = 128>
611
void paged_attention_v1_launcher(
Woosuk Kwon's avatar
Woosuk Kwon committed
612
613
614
615
  torch::Tensor& out,
  torch::Tensor& query,
  torch::Tensor& key_cache,
  torch::Tensor& value_cache,
616
  int num_kv_heads,
Woosuk Kwon's avatar
Woosuk Kwon committed
617
618
  float scale,
  torch::Tensor& block_tables,
619
620
  torch::Tensor& seq_lens,
  int max_seq_len,
621
622
  const c10::optional<torch::Tensor>& alibi_slopes,
  float kv_scale) {
Woosuk Kwon's avatar
Woosuk Kwon committed
623
624
625
626
  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
627
628
629
  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
630
631
632
633

  int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
  assert(head_size % thread_group_size == 0);

Woosuk Kwon's avatar
Woosuk Kwon committed
634
635
636
637
638
  // NOTE: alibi_slopes is optional.
  const float* alibi_slopes_ptr = alibi_slopes ?
    reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
    : nullptr;

Woosuk Kwon's avatar
Woosuk Kwon committed
639
640
  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
641
642
  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
643
  int* block_tables_ptr = block_tables.data_ptr<int>();
644
  int* seq_lens_ptr = seq_lens.data_ptr<int>();
Woosuk Kwon's avatar
Woosuk Kwon committed
645
646

  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
647
648
  int padded_max_seq_len = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
  int logits_size = padded_max_seq_len * sizeof(float);
Woosuk Kwon's avatar
Woosuk Kwon committed
649
  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
650
651
  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
  // Keep that in sync with the logic here!
Woosuk Kwon's avatar
Woosuk Kwon committed
652
653
  int shared_mem_size = std::max(logits_size, outputs_size);

654
655
  dim3 grid(num_heads, num_seqs, 1);
  dim3 block(NUM_THREADS);
656
  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
  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 128:
      LAUNCH_PAGED_ATTENTION_V1(128);
      break;
    case 256:
      LAUNCH_PAGED_ATTENTION_V1(256);
      break;
    default:
      TORCH_CHECK(false, "Unsupported head size: ", head_size);
      break;
  }
}

686
687
688
689
690
691
692
693
694
695
696
697
#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE)            \
  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE>(      \
    out,                                                              \
    query,                                                            \
    key_cache,                                                        \
    value_cache,                                                      \
    num_kv_heads,                                                     \
    scale,                                                            \
    block_tables,                                                     \
    seq_lens,                                                         \
    max_seq_len,                                                      \
    alibi_slopes,                                                     \
698
    kv_scale);
699
700
701

// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
702
703
704
705
706
707
708
709
710
711
712
713
714
715
#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
  switch (block_size) {                                           \
    case 8:                                                       \
      CALL_V1_LAUNCHER(T, CACHE_T, 8, KV_DTYPE);                  \
      break;                                                      \
    case 16:                                                      \
      CALL_V1_LAUNCHER(T, CACHE_T, 16, KV_DTYPE);                 \
      break;                                                      \
    case 32:                                                      \
      CALL_V1_LAUNCHER(T, CACHE_T, 32, KV_DTYPE);                 \
      break;                                                      \
    default:                                                      \
      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
      break;                                                      \
716
717
718
719
720
721
722
  }

void paged_attention_v1(
  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& value_cache,     // [num_blocks, num_heads, head_size, block_size]
723
  int num_kv_heads,               // [num_heads]
724
725
  float scale,
  torch::Tensor& block_tables,    // [num_seqs, max_num_blocks_per_seq]
726
  torch::Tensor& seq_lens,    // [num_seqs]
727
  int block_size,
728
  int max_seq_len,
729
  const c10::optional<torch::Tensor>& alibi_slopes,
730
731
  const std::string& kv_cache_dtype,
  float kv_scale) {
732
733
  
  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype, CALL_V1_LAUNCHER_BLOCK_SIZE)
734
735
}

736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                           \
  vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,      \
  KV_DTYPE, PARTITION_SIZE>                                                            \
  <<<grid, block, shared_mem_size, stream>>>(                                          \
    exp_sums_ptr,                                                                      \
    max_logits_ptr,                                                                    \
    tmp_out_ptr,                                                                       \
    query_ptr,                                                                         \
    key_cache_ptr,                                                                     \
    value_cache_ptr,                                                                   \
    num_kv_heads,                                                                      \
    scale,                                                                             \
    block_tables_ptr,                                                                  \
    seq_lens_ptr,                                                                      \
    max_num_blocks_per_seq,                                                            \
    alibi_slopes_ptr,                                                                  \
    q_stride,                                                                          \
    kv_block_stride,                                                                   \
    kv_head_stride,                                                                    \
    kv_scale);                                                                         \
  vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS, PARTITION_SIZE>    \
  <<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                            \
    out_ptr,                                                                           \
    exp_sums_ptr,                                                                      \
    max_logits_ptr,                                                                    \
    tmp_out_ptr,                                                                       \
    seq_lens_ptr,                                                                      \
763
764
765
766
    max_num_partitions);

template<
  typename T,
767
  typename CACHE_T,
768
  int BLOCK_SIZE,
769
  vllm::Fp8KVCacheDataType KV_DTYPE,
770
771
772
773
774
775
776
777
778
779
  int NUM_THREADS = 128,
  int PARTITION_SIZE = 512>
void paged_attention_v2_launcher(
  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,
780
  int num_kv_heads,
781
782
  float scale,
  torch::Tensor& block_tables,
783
784
  torch::Tensor& seq_lens,
  int max_seq_len,
785
786
  const c10::optional<torch::Tensor>& alibi_slopes,
  float kv_scale) {
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
  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);

  int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
  assert(head_size % thread_group_size == 0);

  // NOTE: alibi_slopes is optional.
  const float* alibi_slopes_ptr = alibi_slopes ?
    reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
    : nullptr;

  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());
808
809
  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());
810
  int* block_tables_ptr = block_tables.data_ptr<int>();
811
  int* seq_lens_ptr = seq_lens.data_ptr<int>();
812
813

  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
814
  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
815
816
817
818
819
820
821
822
823
824
  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);

Woosuk Kwon's avatar
Woosuk Kwon committed
825
  dim3 block(NUM_THREADS);
826
  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
Woosuk Kwon's avatar
Woosuk Kwon committed
827
828
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  switch (head_size) {
829
830
831
    // 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.
Woosuk Kwon's avatar
Woosuk Kwon committed
832
    case 64:
833
      LAUNCH_PAGED_ATTENTION_V2(64);
Woosuk Kwon's avatar
Woosuk Kwon committed
834
835
      break;
    case 80:
836
      LAUNCH_PAGED_ATTENTION_V2(80);
Woosuk Kwon's avatar
Woosuk Kwon committed
837
838
      break;
    case 96:
839
      LAUNCH_PAGED_ATTENTION_V2(96);
Woosuk Kwon's avatar
Woosuk Kwon committed
840
      break;
841
    case 112:
842
      LAUNCH_PAGED_ATTENTION_V2(112);
843
      break;
Woosuk Kwon's avatar
Woosuk Kwon committed
844
    case 128:
845
      LAUNCH_PAGED_ATTENTION_V2(128);
Woosuk Kwon's avatar
Woosuk Kwon committed
846
      break;
847
    case 256:
848
      LAUNCH_PAGED_ATTENTION_V2(256);
849
      break;
Woosuk Kwon's avatar
Woosuk Kwon committed
850
851
852
853
854
855
    default:
      TORCH_CHECK(false, "Unsupported head size: ", head_size);
      break;
  }
}

856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE)                \
  paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE>(          \
    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,                                                         \
871
    kv_scale);
Woosuk Kwon's avatar
Woosuk Kwon committed
872

Woosuk Kwon's avatar
Woosuk Kwon committed
873
874
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
875
876
877
878
879
880
881
882
883
884
885
886
887
888
#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)            \
  switch (block_size) {                                              \
    case 8:                                                          \
      CALL_V2_LAUNCHER(T, CACHE_T, 8, KV_DTYPE);                     \
      break;                                                         \
    case 16:                                                         \
      CALL_V2_LAUNCHER(T, CACHE_T, 16, KV_DTYPE);                    \
      break;                                                         \
    case 32:                                                         \
      CALL_V2_LAUNCHER(T, CACHE_T, 32, KV_DTYPE);                    \
      break;                                                         \
    default:                                                         \
      TORCH_CHECK(false, "Unsupported block size: ", block_size);    \
      break;                                                         \
Woosuk Kwon's avatar
Woosuk Kwon committed
889
890
  }

891
void paged_attention_v2(
Woosuk Kwon's avatar
Woosuk Kwon committed
892
  torch::Tensor& out,             // [num_seqs, num_heads, head_size]
893
894
895
  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]
Woosuk Kwon's avatar
Woosuk Kwon committed
896
897
898
  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& value_cache,     // [num_blocks, num_heads, head_size, block_size]
899
  int num_kv_heads,               // [num_heads]
Woosuk Kwon's avatar
Woosuk Kwon committed
900
901
  float scale,
  torch::Tensor& block_tables,    // [num_seqs, max_num_blocks_per_seq]
902
  torch::Tensor& seq_lens,    // [num_seqs]
Woosuk Kwon's avatar
Woosuk Kwon committed
903
  int block_size,
904
  int max_seq_len,
905
  const c10::optional<torch::Tensor>& alibi_slopes,
906
907
  const std::string& kv_cache_dtype,
  float kv_scale) {
908
  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype, CALL_V2_LAUNCHER_BLOCK_SIZE)
Woosuk Kwon's avatar
Woosuk Kwon committed
909
910
911
912
913
}

#undef WARP_SIZE
#undef MAX
#undef MIN
914
#undef DIVIDE_ROUND_UP