attention_kernels.cu 49.5 KB
Newer Older
1
#include <torch/all.h>
Woosuk Kwon's avatar
Woosuk Kwon committed
2
#include <ATen/cuda/CUDAContext.h>
3
#include <c10/cuda/CUDAGuard.h>
4
#include <algorithm>
Woosuk Kwon's avatar
Woosuk Kwon committed
5

Woosuk Kwon's avatar
Woosuk Kwon committed
6
#include "attention_dtypes.h"
Woosuk Kwon's avatar
Woosuk Kwon committed
7
#include "attention_utils.cuh"
8
9
10

#ifdef USE_ROCM
  #include <hip/hip_bf16.h>
11
  #include "../quantization/fp8/amd/quant_utils.cuh"
12
typedef __hip_bfloat16 __nv_bfloat16;
13
14
#else
  #include "../quantization/fp8/nvidia/quant_utils.cuh"
15
16
17
#endif

#ifndef USE_ROCM
18
  #define WARP_SIZE 32
19
#else
20
  #define WARP_SIZE warpSize
21
22
#endif

23

zhangshao's avatar
zhangshao committed
24
#include "static_switch.h"
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
27
#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
Woosuk Kwon's avatar
Woosuk Kwon committed
28

Woosuk Kwon's avatar
Woosuk Kwon committed
29
namespace vllm {
Woosuk Kwon's avatar
Woosuk Kwon committed
30
31

// Utility function for attention softmax.
32
template <int NUM_WARPS>
Woosuk Kwon's avatar
Woosuk Kwon committed
33
34
35
36
37
38
39
40
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) {
41
    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
  }

  // 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) {
60
    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
Woosuk Kwon's avatar
Woosuk Kwon committed
61
62
63
  }

  // Broadcast to other threads.
64
  return VLLM_SHFL_SYNC(sum, 0);
Woosuk Kwon's avatar
Woosuk Kwon committed
65
66
}

zhuwenwen's avatar
zhuwenwen committed
67
68
// TODO(woosuk): Merge the last two dimensions of the grid.
// Grid: (num_heads, num_seqs, max_num_partitions).
69
70
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
flyingdown's avatar
flyingdown committed
71
          bool IS_BLOCK_SPARSE, int REUSE_KV_TIMES = 1, bool odd_nheads = false,
zhangshao's avatar
zhangshao committed
72
          int PARTITION_SIZE = 0>  // Zero means no partitioning.
zhangshao's avatar
zhangshao committed
73
74
75
76
77
78
79
80
81
82
83
__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]
    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]
84
    const int num_heads,                   // [num_heads]
zhangshao's avatar
zhangshao committed
85
    const int num_kv_heads,               // [num_kv_heads]
86
87
88
89
90
91
92
93
94
    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,
    const float kv_scale, const int tp_rank, const int blocksparse_local_blocks,
    const int blocksparse_vert_stride, const int blocksparse_block_size,
    const int blocksparse_head_sliding_step) {
zhangshao's avatar
zhangshao committed
95
96
97
  const int seq_idx = blockIdx.z;
  const int partition_idx = blockIdx.y;
  const int max_num_partitions = gridDim.y;
98
99
100
101
102
103
  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;
  }
zhangshao's avatar
zhangshao committed
104
  if constexpr (sizeof(scalar_t)==2){
105
106
107
108

  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;
109
  const int partition_size = USE_PARTITIONING ? PARTITION_SIZE : num_seq_blocks * BLOCK_SIZE;
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
  // [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;
zhangshao's avatar
zhangshao committed
132
133
134
135
136
137
138
139
140
141
  // const int warp_idx_vec = thread_idx / WARP_SIZE;
  // int warp_idx =0;
  // asm volatile("v_readfirstlane_b32 %0,%1"
  //               : "=s"(warp_idx)
  //               : "v"(warp_idx_vec)
  //               :);
  // // const int warp_idx = thread_idx / WARP_SIZE;

  // const int lane = thread_idx % WARP_SIZE;

142
    //const int warp_idx = thread_idx / WARP_SIZE;
143
144
  const int lane = thread_idx % WARP_SIZE;

145
146
  int warp_id_vec = threadIdx.x / WARP_SIZE; //warp id in a block
  int warp_idx =0;
zhangshao's avatar
zhangshao committed
147
  asm volatile("v_readfirstlane_b32 %0,%1"
148
149
150
                : "=s"(warp_idx)
                : "v"(warp_id_vec)
                :);
zhangshao's avatar
zhangshao committed
151
152
153

  // const int head_idx = blockIdx.x;
  // const int num_heads = gridDim.x;
154
  const int num_queries_per_kv = num_heads / num_kv_heads;
zhangshao's avatar
zhangshao committed
155
156
  // 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.
zhangshao's avatar
zhangshao committed
163
  constexpr int VEC_SIZE = MAX(32 / (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.
zhangshao's avatar
zhangshao committed
180
181
182
  // const scalar_t* q_ptr = q + seq_idx * q_stride;
  const scalar_t* q_ptr_offset = q + seq_idx * q_stride;

183
184
185
186
187
188
189
190
191
192
  __shared__ Q_vec q_vecs[REUSE_KV_TIMES * 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
193
194
195
196
197
198

  // 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.
zhangshao's avatar
zhangshao committed
199
  __shared__ float red_smem[REUSE_KV_TIMES][2 * NUM_WARPS];
200
  // float (*red_smem)[2 * NUM_WARPS] = reinterpret_cast<float(*)[2 * NUM_WARPS]>(&shared_mem[10*1024]);
zhangshao's avatar
zhangshao committed
201
202
203
204

  // __shared__ char shared_mem[12 * 1024];
  // float* logits = reinterpret_cast<float*>(shared_mem);
  // __shared__ float red_smem[REUSE_KV_TIMES][2 * NUM_WARPS];
205
206
207
208

  // 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);
zhangshao's avatar
zhangshao committed
209
210
  float qk_max[REUSE_KV_TIMES];

211
212
  for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
      qk_max[reuse_kv_idx] = -FLT_MAX;
zhangshao's avatar
zhangshao committed
213
  }
214
215
216
   
  const int num_blocks_per_kv = ((num_queries_per_kv + REUSE_KV_TIMES -1) / REUSE_KV_TIMES);
  const int head_idx_soffset = (blockIdx.x / num_blocks_per_kv) * num_queries_per_kv + (blockIdx.x % num_blocks_per_kv) * REUSE_KV_TIMES;
zhangshao's avatar
zhangshao committed
217
  const int kv_head_idx = head_idx_soffset / num_queries_per_kv;
218
  const int q_boundary = (kv_head_idx + 1)* num_queries_per_kv;
zhangshao's avatar
zhangshao committed
219

220
221
222
  #pragma unroll
  for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
    const int head_idx = head_idx_soffset + reuse_kv_idx;//blockIdx.x * REUSE_KV_TIMES + reuse_kv_idx;
zhangshao's avatar
zhangshao committed
223
    const scalar_t* q_ptr = q_ptr_offset + head_idx * HEAD_SIZE;
224
225
    #pragma unroll
    for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD; i += NUM_THREAD_GROUPS) {
zhangshao's avatar
zhangshao committed
226
      const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
227
      q_vecs[reuse_kv_idx*THREAD_GROUP_SIZE + thread_group_offset][i] = *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
zhangshao's avatar
zhangshao committed
228
229
    }
  }
230
  __syncthreads(); // TODO(naed90): possible speedup if this is replaced with a memory wall right before we use q_vecs
231
232
233
234
235
236

  // 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;
237
  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; block_idx += NUM_WARPS) {
238
239
240
241
242
    // 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
243
244
245
    for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
    const int head_idx = head_idx_soffset + reuse_kv_idx;//blockIdx.x * REUSE_KV_TIMES + reuse_kv_idx;
    if(!odd_nheads || head_idx < q_boundary) {
zhangshao's avatar
zhangshao committed
246
        // blocksparse specific vars
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
    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;
    }
    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) {
270
271
272
273
        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;
274

275
          if (thread_group_offset == 0) {
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
            // 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;
      }
    }
    const float alibi_slope = alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
    const int64_t physical_block_number = static_cast<int64_t>(block_table[block_idx]);

    // 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];
      if(reuse_kv_idx == 0) {
        #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, kv_scale);
317
318
319
          }
        }
      }
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
      __builtin_amdgcn_sched_barrier(0);
      // Compute dot product.
      // This includes a reduction across the threads in the same thread group.
      float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(q_vecs[reuse_kv_idx*THREAD_GROUP_SIZE + thread_group_offset], k_vecs);
      // Add the ALiBi bias if slopes are given.
      qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
      __builtin_amdgcn_sched_barrier(0);
      if (thread_group_offset == 0) {
        // Store the partial reductions to shared memory.
        // NOTE(woosuk): It is required to zero out the masked logits.
        const bool mask = token_idx >= seq_len;
        logits[(reuse_kv_idx * partition_size) + (token_idx - start_token_idx)] = mask ? 0.f : qk;
        // Update the max value.
        qk_max[reuse_kv_idx] = mask ? qk_max[reuse_kv_idx] : fmaxf(qk_max[reuse_kv_idx], qk);
      }
335
    }
zhangshao's avatar
zhangshao committed
336
  }
337
338
  }
  }
zhangshao's avatar
zhangshao committed
339
340
  // Get the sum of the exp values.
  float exp_sum[REUSE_KV_TIMES] = {0.f};
Woosuk Kwon's avatar
Woosuk Kwon committed
341
342
343
344

  // 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.
345
  for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
zhangshao's avatar
zhangshao committed
346
    const int head_idx = head_idx_soffset + reuse_kv_idx;
347
348
    if(!odd_nheads || head_idx < q_boundary) {
      #pragma unroll
zhangshao's avatar
zhangshao committed
349
      for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
350
        qk_max[reuse_kv_idx] = fmaxf(qk_max[reuse_kv_idx], VLLM_SHFL_XOR_SYNC(qk_max[reuse_kv_idx], mask));
zhangshao's avatar
zhangshao committed
351
352
353
354
355
356
357
358
      }
      if (lane == 0) {
        red_smem[reuse_kv_idx][warp_idx] = qk_max[reuse_kv_idx];
      }
      __syncthreads();

      // TODO(woosuk): Refactor this part.
      // Get the max qk value for the sequence.
359
360
      qk_max[reuse_kv_idx] = lane < NUM_WARPS ? red_smem[reuse_kv_idx][lane] : -FLT_MAX;
    #pragma unroll
zhangshao's avatar
zhangshao committed
361
      for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
362
        qk_max[reuse_kv_idx] = fmaxf(qk_max[reuse_kv_idx], VLLM_SHFL_XOR_SYNC(qk_max[reuse_kv_idx], mask));
zhangshao's avatar
zhangshao committed
363
364
365
      }
      // Broadcast the max qk value to all threads.
      qk_max[reuse_kv_idx] = VLLM_SHFL_SYNC(qk_max[reuse_kv_idx], 0);
Woosuk Kwon's avatar
Woosuk Kwon committed
366

zhangshao's avatar
zhangshao committed
367
      for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
368
        float val = __expf(logits[(reuse_kv_idx * partition_size) + i] - qk_max[reuse_kv_idx]);
zhangshao's avatar
zhangshao committed
369
370
371
        logits[(reuse_kv_idx * partition_size) + i] = val;
        exp_sum[reuse_kv_idx] += val;
      }
372
      exp_sum[reuse_kv_idx] = block_sum<NUM_WARPS>(&red_smem[reuse_kv_idx][NUM_WARPS], exp_sum[reuse_kv_idx]);
Woosuk Kwon's avatar
Woosuk Kwon committed
373

zhangshao's avatar
zhangshao committed
374
375
376
377
378
379
      // Compute softmax.
      const float inv_sum = __fdividef(1.f, exp_sum[reuse_kv_idx] + 1e-6f);
      for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
        logits[(reuse_kv_idx * partition_size) + i] *= inv_sum;
      }
      __syncthreads();
Woosuk Kwon's avatar
Woosuk Kwon committed
380

zhangshao's avatar
zhangshao committed
381
382
383
384
385
386
      // 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[reuse_kv_idx];
387
        float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions +
zhangshao's avatar
zhangshao committed
388
389
390
391
                              head_idx * max_num_partitions + partition_idx;
        *exp_sums_ptr = exp_sum[reuse_kv_idx];
      }
    }
392
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
393
394
395
396
  // 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;
397
  using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
Woosuk Kwon's avatar
Woosuk Kwon committed
398
399
400
401
  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;
402
403
  constexpr int NUM_ROWS_PER_THREAD =
      DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
Woosuk Kwon's avatar
Woosuk Kwon committed
404
405

  // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
zhangshao's avatar
zhangshao committed
406
  float accs[REUSE_KV_TIMES][NUM_ROWS_PER_THREAD];
Woosuk Kwon's avatar
Woosuk Kwon committed
407

408
409
410
  #pragma unroll
  for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
    #pragma unroll
zhangshao's avatar
zhangshao committed
411
    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
412
        accs[reuse_kv_idx][i] = 0.f;
zhangshao's avatar
zhangshao committed
413
414
    }
  }
415
416
  scalar_t zero_value;
  zero(zero_value);
417
418
419
420
  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
       block_idx += NUM_WARPS) {
    const int64_t physical_block_number =
        static_cast<int64_t>(block_table[block_idx]);
Woosuk Kwon's avatar
Woosuk Kwon committed
421
422
423
    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;
flyingdown's avatar
flyingdown committed
424
#pragma unroll
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
    V_vec v_vec;
    for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
      // 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
      // blocksparse specific vars
      const int head_idx = head_idx_soffset + reuse_kv_idx;
      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;
zhangshao's avatar
zhangshao committed
450
      }
451
452
453
454
455
      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;
flyingdown's avatar
flyingdown committed
456
        }
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
      }
      if(!odd_nheads || head_idx < q_boundary) {


      const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride
                                   + kv_head_idx * kv_head_stride;

     from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + (reuse_kv_idx * partition_size) +  token_idx - start_token_idx));
      // scalar_t* logits_vec_ptr = reinterpret_cast<scalar_t*>(&logits_vec);
      // for(int i=0;i<8;++i){
      //   from_float(*(logits_vec_ptr+i), 1000);
      // }

      if(reuse_kv_idx==0) {
      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;

        if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
          v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
        } else {
          V_quant_vec v_quant_vec =
              *reinterpret_cast<const V_quant_vec*>(v_ptr + offset);
          // Vector conversion from V_quant_vec to V_vec.
          v_vec = fp8::scaled_convert<V_vec, V_quant_vec, KV_DTYPE>(v_quant_vec,
                                                                    kv_scale);
483
        }
484
485
486
487
488
489
        if (block_idx == num_seq_blocks - 1) {
          // 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);
490
#pragma unroll
491
492
          for (int j = 0; j < V_VEC_SIZE; j++) {
            v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
493
494
          }
        }
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
        // if(threadIdx.x==0){
        //   scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
        //   scalar_t* logits_vec_ptr = reinterpret_cast<scalar_t*>(&logits_vec);
        //   for(int i=0;i<8;++i){
        //     printf("v_vec[%d] = %f\n",i, half_to_float(v_vec_ptr[i]));
        //     // from_float(*(v_vec_ptr + i), 1000);
        //   }
        //   for(int i=0;i<8;++i){
        //     printf("logits_vec[%d] = %f\n",i,half_to_float(logits_vec_ptr[i]));
        //     // from_float(*(logits_vec_ptr + i), 1000);
        //   }
        // }
        // accs[reuse_kv_idx][i] += dot(logits_vec, v_vec);
      }
      } 
        accs[reuse_kv_idx][i] += dot(logits_vec, v_vec);
zhangshao's avatar
zhangshao committed
511
      }
Woosuk Kwon's avatar
Woosuk Kwon committed
512
513
514
515
      }
    }
  }

516
517
518
  // Perform reduction within each warp.
  #pragma unroll
  for(int reuse_kv_idx=0; reuse_kv_idx<REUSE_KV_TIMES; reuse_kv_idx++) {
zhangshao's avatar
zhangshao committed
519
520
    int head_idx = head_idx_soffset + reuse_kv_idx;

521
    if(!odd_nheads || head_idx < q_boundary) {
Woosuk Kwon's avatar
Woosuk Kwon committed
522
#pragma unroll
523
524
  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
    float acc = accs[reuse_kv_idx][i];
Woosuk Kwon's avatar
Woosuk Kwon committed
525
#pragma unroll
526
527
528
529
530
    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
    }
    accs[reuse_kv_idx][i] = acc;
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
531

532
533
534
  // NOTE(woosuk): A barrier is required because the shared memory space for
  // logits is reused for the output.
  __syncthreads();
Woosuk Kwon's avatar
Woosuk Kwon committed
535

536
537
  // Perform reduction across warps.
  float* out_smem = reinterpret_cast<float*>(shared_mem);
Woosuk Kwon's avatar
Woosuk Kwon committed
538
#pragma unroll
539
540
541
542
543
  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[(reuse_kv_idx * (NUM_WARPS / 2) * HEAD_SIZE) + (warp_idx - mid) * HEAD_SIZE];
Woosuk Kwon's avatar
Woosuk Kwon committed
544
#pragma unroll
545
546
547
548
      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[reuse_kv_idx][i];
Woosuk Kwon's avatar
Woosuk Kwon committed
549
        }
550
551
552
      }
    }
    __syncthreads();
Woosuk Kwon's avatar
Woosuk Kwon committed
553

554
555
556
    // Lower warps update the output.
    if (warp_idx < mid) {
      const float* src = &out_smem[(reuse_kv_idx * (NUM_WARPS / 2) * HEAD_SIZE) + warp_idx * HEAD_SIZE];
Woosuk Kwon's avatar
Woosuk Kwon committed
557
#pragma unroll
558
559
560
561
      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[reuse_kv_idx][i] += src[row_idx];
Woosuk Kwon's avatar
Woosuk Kwon committed
562
563
        }
      }
564
565
566
    }
    __syncthreads();
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
567

568
569
570
571
572
  // Write the final output.
  if (warp_idx == 0) {
    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
573
#pragma unroll
flyingdown's avatar
flyingdown committed
574
575
576
577
578
579
        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[reuse_kv_idx][i]);
          }
        }
Woosuk Kwon's avatar
Woosuk Kwon committed
580
581
582
      }
    }
  }
zhangshao's avatar
zhangshao committed
583
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
584
585
}

586

587
// Grid: (num_heads, num_seqs, 1).
588
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
589
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
590
591
592
593
          int REUSE_KV_TIMES = 1,
          bool IS_BLOCK_SPARSE,
          bool odd_nheads = false>
__global__ __launch_bounds__(256,1) void paged_attention_v1_kernel(
594
595
596
597
598
599
    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]
600
    const int num_heads,               // [num_heads]    
601
602
603
604
605
606
607
    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,
608
609
610
    const float kv_scale, const int tp_rank, const int blocksparse_local_blocks,
    const int blocksparse_vert_stride, const int blocksparse_block_size,
    const int blocksparse_head_sliding_step) {
611
612
613
614
615
616
617
618
      paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
                          KV_DTYPE, IS_BLOCK_SPARSE, REUSE_KV_TIMES, odd_nheads>(
        /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache,
        v_cache, num_heads, num_kv_heads, scale, block_tables, seq_lens,
        max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride,
        kv_head_stride, kv_scale, tp_rank, blocksparse_local_blocks,
        blocksparse_vert_stride, blocksparse_block_size,
        blocksparse_head_sliding_step);
619
620
621
}

// Grid: (num_heads, num_seqs, max_num_partitions).
622
623
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
624
625
626
          bool IS_BLOCK_SPARSE,
          int REUSE_KV_TIMES,
          int PARTITION_SIZE,
zhangshao's avatar
zhangshao committed
627
          bool odd_nheads = false>
628
__global__ __launch_bounds__(256,1) void paged_attention_v2_kernel(
629
630
631
632
633
634
635
636
637
638
    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]
639
    const int num_heads,               // [num_heads]                                      
zhangshao's avatar
zhangshao committed
640
    const int num_kv_heads,               // [num_kv_heads]
641
642
643
644
645
646
    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,
647
648
649
    const float kv_scale, const int tp_rank, const int blocksparse_local_blocks,
    const int blocksparse_vert_stride, const int blocksparse_block_size,
    const int blocksparse_head_sliding_step) {
650
651
652
653
654
655
656
        paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
                         KV_DTYPE, IS_BLOCK_SPARSE, REUSE_KV_TIMES, odd_nheads, PARTITION_SIZE>(
          exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_heads, num_kv_heads, scale,
          block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride,
          kv_block_stride, kv_head_stride, kv_scale, tp_rank,
          blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size,
          blocksparse_head_sliding_step);
657
658
659
}

// Grid: (num_heads, num_seqs).
660
661
662
template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
          int PARTITION_SIZE>
__global__ __launch_bounds__(256,1) void paged_attention_v2_reduce_kernel(
663
664
665
666
667
668
669
670
671
    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) {
672
673
674
  const int num_heads = gridDim.x;
  const int head_idx = blockIdx.x;
  const int seq_idx = blockIdx.y;
675
676
  const int seq_len = seq_lens[seq_idx];
  const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
677
678
  if (num_partitions == 1) {
    // No need to reduce. Only copy tmp_out to out.
679
680
681
682
683
    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;
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
    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);
702
703
704
  const float* max_logits_ptr = max_logits +
                                seq_idx * num_heads * max_num_partitions +
                                head_idx * max_num_partitions;
705
706
707
708
709
710
711
712
713
714
715
716
  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) {
717
    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
718
719
720
721
722
723
724
725
726
  }
  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) {
727
    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
728
729
  }
  // Broadcast the max value to all threads.
730
  max_logit = VLLM_SHFL_SYNC(max_logit, 0);
731
732

  // Load rescaled exp sums to shared memory.
733
734
735
736
737
  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;
738
739
740
741
742
743
744
745
746
747
748
749
  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.
750
751
752
753
754
  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;
755
756
757
758
#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) {
759
760
      acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] *
             inv_global_exp_sum;
761
762
763
764
765
    }
    from_float(out_ptr[i], acc);
  }
}

766
767
}  // namespace vllm

768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
      ((void*)vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
                                              BLOCK_SIZE, NUM_THREADS,      \
                                              KV_DTYPE, REUSE_KV_TIMES, IS_BLOCK_SPARSE, odd_nheads>),  \
      shared_mem_size);                                                     \
 hipLaunchKernelGGL(( vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
                                  NUM_THREADS, KV_DTYPE, REUSE_KV_TIMES, IS_BLOCK_SPARSE, odd_nheads>)   \
      , dim3(grid), dim3(block), shared_mem_size, stream,                            \
          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_heads, 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, tp_rank, blocksparse_local_blocks,                      \
          blocksparse_vert_stride, blocksparse_block_size,                  \
          blocksparse_head_sliding_step);
Woosuk Kwon's avatar
Woosuk Kwon committed
783

zhangshao's avatar
zhangshao committed
784
785
786
787
788
789
790
791
792
793
794
// #define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
// vllm::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
//                                   NUM_THREADS, KV_DTYPE, REUSE_KV_TIMES, IS_BLOCK_SPARSE, odd_nheads>   \
//       <<<dim3(grid), dim3(block)>>>(                           \
//           out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_heads, 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, tp_rank, blocksparse_local_blocks,                      \
//           blocksparse_vert_stride, blocksparse_block_size,                  \
//           blocksparse_head_sliding_step);

Woosuk Kwon's avatar
Woosuk Kwon committed
795
// TODO(woosuk): Tune NUM_THREADS.
796
template <typename T, typename CACHE_T, int BLOCK_SIZE,
zhangshao's avatar
zhangshao committed
797
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE>
798
void paged_attention_v1_launcher(
799
800
801
    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,
802
803
804
805
    const c10::optional<torch::Tensor>& alibi_slopes, float kv_scale,
    const int tp_rank, const int blocksparse_local_blocks,
    const int blocksparse_vert_stride, const int blocksparse_block_size,
    const int blocksparse_head_sliding_step) {
Woosuk Kwon's avatar
Woosuk Kwon committed
806
807
808
809
  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
810
811
812
  int q_stride = query.stride(0);
  int kv_block_stride = key_cache.stride(0);
  int kv_head_stride = key_cache.stride(1);
zhangshao's avatar
zhangshao committed
813
  int num_threads = 128;
814
815
  if(num_heads!=num_kv_heads){
    num_threads =256;
zhangshao's avatar
zhangshao committed
816
  }
Woosuk Kwon's avatar
Woosuk Kwon committed
817
818
819
  int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
  assert(head_size % thread_group_size == 0);

Woosuk Kwon's avatar
Woosuk Kwon committed
820
  // NOTE: alibi_slopes is optional.
821
822
823
824
  const float* alibi_slopes_ptr =
      alibi_slopes
          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
          : nullptr;
Woosuk Kwon's avatar
Woosuk Kwon committed
825

Woosuk Kwon's avatar
Woosuk Kwon committed
826
827
  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
828
829
  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
830
  int* block_tables_ptr = block_tables.data_ptr<int>();
831
  int* seq_lens_ptr = seq_lens.data_ptr<int>();
Woosuk Kwon's avatar
Woosuk Kwon committed
832

833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
  int padded_max_seq_len = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
  REUSEKV_SWITCH_V1(num_heads * num_seqs , [&] {
    BOOL_SWITCH((num_heads/num_kv_heads % REUSE_KV_TIMES != 0), odd_nheads, [&] {
      HEADSIZE_SWITCH(head_size, [&] {
        NUM_THREADS_SWITCH(num_threads, [&] {
          OPT_SWITCH(num_heads == num_kv_heads, [&] {
          constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
          int logits_size =  REUSE_KV_TIMES*padded_max_seq_len * sizeof(float);
          int outputs_size =  REUSE_KV_TIMES*(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 = ::max(logits_size, outputs_size);
          if(num_heads == num_kv_heads) shared_mem_size = ::max(12 * 1024, shared_mem_size);
          // int shared_mem_size = ::max(31*1024, ::max(logits_size, outputs_size));
          // std::cout<<"shared_mem_size = "<<shared_mem_size<<std::endl;
          dim3 grid((num_heads/num_kv_heads + REUSE_KV_TIMES - 1) / REUSE_KV_TIMES*num_kv_heads, 1, num_seqs);
          dim3 block(NUM_THREADS);
          const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(query));
          const hipStream_t stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
          LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE);
zhangshao's avatar
zhangshao committed
853
854
          });
        });
855
856
857
      });
    });
  }); 
858
859
}

860
861
862
#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>(                              \
863
      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
864
865
866
867
868
869
870
871
872
873
874
875
876
      seq_lens, max_seq_len, alibi_slopes, kv_scale, tp_rank,                \
      blocksparse_local_blocks, blocksparse_vert_stride,                     \
      blocksparse_block_size, blocksparse_head_sliding_step);

#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;                                                               \
  }
877
878
879

// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
880
881
#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
  switch (block_size) {                                           \
zhangshao's avatar
zhangshao committed
882
883
884
    case 8:                                                       \
      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
      break;                                                      \
885
    case 16:                                                      \
886
      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
887
      break;                                                      \
zhangshao's avatar
zhangshao committed
888
889
890
    case 32:                                                      \
      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
      break;                                                      \
891
892
893
    default:                                                      \
      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
      break;                                                      \
894
895
  }

zhangshao's avatar
zhangshao committed
896

897
void paged_attention_v1(
898
899
900
901
902
    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&
903
904
905
        value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
906
907
    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
    torch::Tensor& seq_lens,      // [num_seqs]
908
    int64_t block_size, int64_t max_seq_len,
909
    const c10::optional<torch::Tensor>& alibi_slopes,
910
911
912
913
    const std::string& kv_cache_dtype, double kv_scale, const int64_t tp_rank,
    const int64_t blocksparse_local_blocks,
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
    const int64_t blocksparse_head_sliding_step) {
914
915
916
917
918
  const bool is_block_sparse = (blocksparse_vert_stride > 1);

  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                             CALL_V1_LAUNCHER_BLOCK_SIZE)
}
919
920

#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
 hipLaunchKernelGGL(( vllm::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \
                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \
                                  REUSE_KV_TIMES, PARTITION_SIZE, odd_nheads>)                              \
      , dim3(grid), dim3(block), shared_mem_size, stream,                               \
          exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
          value_cache_ptr, num_heads, 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, tp_rank,                  \
          blocksparse_local_blocks, blocksparse_vert_stride,                   \
          blocksparse_block_size, blocksparse_head_sliding_step);              \
 hipLaunchKernelGGL(( vllm::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \
                                         PARTITION_SIZE>)                       \
      , dim3(reduce_grid), dim3(block), reduce_shared_mem_size, stream,                 \
          out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \
          max_num_partitions);
936
937

template <typename T, typename CACHE_T, int BLOCK_SIZE,
938
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
zhuwenwen's avatar
zhuwenwen committed
939
          int NUM_THREADS = 256, int PARTITION_SIZE = 512>
940
void paged_attention_v2_launcher(
941
942
943
944
    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,
945
946
947
948
    const c10::optional<torch::Tensor>& alibi_slopes, float kv_scale,
    const int tp_rank, const int blocksparse_local_blocks,
    const int blocksparse_vert_stride, const int blocksparse_block_size,
    const int blocksparse_head_sliding_step) {
949
950
951
952
953
954
955
956
957
958
959
960
  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.
961
962
963
964
  const float* alibi_slopes_ptr =
      alibi_slopes
          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
          : nullptr;
965
966
967
968
969
970

  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());
971
972
  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());
973
  int* block_tables_ptr = block_tables.data_ptr<int>();
974
  int* seq_lens_ptr = seq_lens.data_ptr<int>();
975
976

  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
977
  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
  REUSEKV_SWITCH(num_heads * max_num_partitions * num_seqs , [&] {
    BOOL_SWITCH((num_heads/num_kv_heads % REUSE_KV_TIMES != 0), odd_nheads, [&] {
      HEADSIZE_SWITCH(head_size, [&] {
        OPT_SWITCH(num_heads == num_kv_heads, [&] {
        int logits_size = REUSE_KV_TIMES*PARTITION_SIZE * sizeof(float);
        int outputs_size = REUSE_KV_TIMES*(NUM_WARPS / 2) * head_size * sizeof(float);

        // For paged attention v2 kernel.
        // dim3 grid(num_heads, max_num_partitions, num_seqs);

        dim3 grid;
        grid.x = (num_heads/num_kv_heads + REUSE_KV_TIMES -1)/REUSE_KV_TIMES * num_kv_heads;
        grid.y = max_num_partitions;
        grid.z = num_seqs;
        // int shared_mem_size = ::max(1024*32, ::max(logits_size, outputs_size));
        int shared_mem_size = ::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::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(query));
        const hipStream_t stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
        LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE);
zhangshao's avatar
zhangshao committed
1001
        });
1002
1003
      });
    });
zhangshao's avatar
zhangshao committed
1004
  });
Woosuk Kwon's avatar
Woosuk Kwon committed
1005
1006
}

1007
1008
1009
#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>(                               \
1010
1011
      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, \
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
      kv_scale, tp_rank, blocksparse_local_blocks, blocksparse_vert_stride,   \
      blocksparse_block_size, blocksparse_head_sliding_step);

#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
1024

Woosuk Kwon's avatar
Woosuk Kwon committed
1025
1026
// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
// 1, 2, 4, 64, 128, 256.
1027
1028
#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
  switch (block_size) {                                           \
zhangshao's avatar
zhangshao committed
1029
1030
1031
    case 8:                                                       \
      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
      break;                                                      \
1032
    case 16:                                                      \
1033
      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
1034
      break;                                                      \
zhangshao's avatar
zhangshao committed
1035
1036
1037
    case 32:                                                      \
      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
      break;                                                      \
1038
1039
1040
    default:                                                      \
      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
      break;                                                      \
Woosuk Kwon's avatar
Woosuk Kwon committed
1041
1042
  }

zhangshao's avatar
zhangshao committed
1043

1044
void paged_attention_v2(
1045
1046
1047
1048
1049
1050
1051
1052
1053
    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&
1054
1055
1056
        value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
1057
1058
    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
    torch::Tensor& seq_lens,      // [num_seqs]
1059
    int64_t block_size, int64_t max_seq_len,
1060
    const c10::optional<torch::Tensor>& alibi_slopes,
1061
1062
1063
1064
    const std::string& kv_cache_dtype, double kv_scale, const int64_t tp_rank,
    const int64_t blocksparse_local_blocks,
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
    const int64_t blocksparse_head_sliding_step) {
1065
  const bool is_block_sparse = (blocksparse_vert_stride > 1);
1066
1067
  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                             CALL_V2_LAUNCHER_BLOCK_SIZE)
Woosuk Kwon's avatar
Woosuk Kwon committed
1068
1069
1070
1071
1072
}

#undef WARP_SIZE
#undef MAX
#undef MIN
zhangshao's avatar
zhangshao committed
1073
#undef DIVIDE_ROUND_UP