Repository: hpcaitech/ColossalAI
File: extensions/csrc/kernel/cuda/flash_decoding_attention_kernel.cu (CUDA)

Code before:

// We only support head size of { 64, 128, 256 }
// models like Phi-2, whose head size is 80, is not supported right now
template<typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS>
__global__ void flash_decoding_attention_kernel(
  scalar_t* __restrict__ out,                 // [num_tokens, num_heads, head_size]
  const scalar_t* __restrict__ q,             // [num_tokens, num_heads, head_size]
  const cache_t* __restrict__ k_cache,        // [num_blocks, num_kv_heads, block_size, head_size]
  const cache_t* __restrict__ v_cache,        // [num_blocks, num_kv_heads, block_size, head_size]
  const int* __restrict__ context_lens,       // [num_tokens]
  const int* __restrict__ block_tables,       // [num_tokens, max_num_blocks_per_seq]
  const int max_seq_len,
  const int num_kv_heads,
  const float scale,
  const int max_num_blocks_per_seq,
  const int q_stride,                         // num_heads * head_size
  const int kv_block_stride,
  const int kv_head_stride) {
  const int seq_idx = blockIdx.y;
  const int head_idx = blockIdx.x;
  const int thread_idx = threadIdx.x;
  const int lane = thread_idx % WARP_SIZE;
  const int warp_idx = thread_idx / WARP_SIZE;
  const int num_heads = gridDim.x;
  const int num_queries_per_kv = num_heads / num_kv_heads;
  const int kv_head_idx = head_idx / num_queries_per_kv;
  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  constexpr int Q_SHARED_SIZE = (HEAD_SIZE * sizeof(scalar_t)) / sizeof(float4);
  // here thread_group does not determine the number of threads responsible for a key
  // but only the VEC_SIZE of each thread
  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
  constexpr int VEC_SIZE = MIN(ROUND_DOWN_HIGHEST_POWER_OF_TWO((HEAD_SIZE / THREAD_GROUP_SIZE)), sizeof(float4) / sizeof(scalar_t));
  constexpr int NUM_VECS_PER_TOKEN = HEAD_SIZE / VEC_SIZE;
  constexpr int NUM_THREADS_PER_TOKEN = MIN(NUM_VECS_PER_TOKEN, WARP_SIZE);
  constexpr int NUM_ROUNDS_PER_TOKEN = NUM_VECS_PER_TOKEN / NUM_THREADS_PER_TOKEN;
  constexpr int WARP_STRIDE = WARP_SIZE * NUM_ROUNDS_PER_TOKEN;

  using K_vec = typename VecTypeTrait<scalar_t, VEC_SIZE>::Type;
  using V_vec = typename VecTypeTrait<scalar_t, VEC_SIZE>::Type;
  using L_vec = typename VecTypeTrait<scalar_t, VEC_SIZE>::Type;
  using Float_vec = typename FloatVecTypeTrait<L_vec>::Type;

  const int context_len = context_lens[seq_idx];
  const int thread_group_offset = thread_idx % NUM_THREADS_PER_TOKEN;
  const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;

  __shared__ float4 q_shared[Q_SHARED_SIZE];
  __shared__ float red_shared_mem[2 * NUM_WARPS];
  extern __shared__ char shared_mem[];
  float* logits = reinterpret_cast<float*>(shared_mem);

Review comment (suggestion):
const local variable should be named as kConstXxx, you can refer to [Google C++ Style](https://google-styleguide.readthedocs.io/zh-cn/latest/google-cpp-styleguide/naming.html), Things like this can be done in next PR 

Code after:

// We only support head size of { 64, 128, 256 }
// models like Phi-2, whose head size is 80, is not supported right now
template<typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS>
__global__ void flash_decoding_attention_kernel(
  scalar_t* __restrict__ out,                 // [num_tokens, num_heads, head_size]
  const scalar_t* __restrict__ q,             // [num_tokens, 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, block_size, head_size]
  const int* __restrict__ context_lens,       // [num_tokens]
  const int* __restrict__ block_tables,       // [num_tokens, max_num_blocks_per_seq]
  const int max_seq_len,
  const int num_kv_heads,
  const float scale,
  const int max_num_blocks_per_seq,
  const int q_stride,                         // num_heads * head_size
  const int kv_block_stride,
  const int kv_head_stride) {
  const int seq_idx = blockIdx.y;
  const int head_idx = blockIdx.x;
  const int thread_idx = threadIdx.x;
  const int lane = thread_idx % WARP_SIZE;
  const int warp_idx = thread_idx / WARP_SIZE;
  const int num_heads = gridDim.x;
  const int num_queries_per_kv = num_heads / num_kv_heads;
  const int kv_head_idx = head_idx / num_queries_per_kv;
  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  constexpr int x = sizeof(float4) / sizeof(scalar_t);
  constexpr int Q_SHARED_SIZE = HEAD_SIZE / x;
  // here thread_group does not determine the number of threads responsible for a key
  // but only the VEC_SIZE of each thread
  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
  constexpr int VEC_SIZE = MIN(ROUND_DOWN_HIGHEST_POWER_OF_TWO((HEAD_SIZE / THREAD_GROUP_SIZE)), x);
  constexpr int NUM_VECS_PER_TOKEN = HEAD_SIZE / VEC_SIZE;
  constexpr int NUM_THREADS_PER_TOKEN = MIN(NUM_VECS_PER_TOKEN, WARP_SIZE);
  constexpr int NUM_ROUNDS_PER_TOKEN = NUM_VECS_PER_TOKEN / NUM_THREADS_PER_TOKEN;
  constexpr int WARP_STRIDE = WARP_SIZE * NUM_ROUNDS_PER_TOKEN;
  constexpr int NUM_THREADS_PER_X = x / VEC_SIZE;
  constexpr int NUM_ROWS_PER_ROUNDS = MIN(WARP_SIZE / NUM_THREADS_PER_X, BLOCK_SIZE);
  constexpr int NUM_VECS_PER_THREAD = NUM_ROWS_PER_ROUNDS * NUM_VECS_PER_TOKEN / WARP_SIZE;

  using K_vec = typename VecTypeTrait<scalar_t, VEC_SIZE>::Type;
  using V_vec = typename VecTypeTrait<scalar_t, VEC_SIZE>::Type;
  using L_vec = typename VecTypeTrait<scalar_t, VEC_SIZE>::Type;
  using Float_vec = typename FloatVecTypeTrait<L_vec>::Type;

  const int context_len = context_lens[seq_idx];
  const int thread_group_offset = lane % NUM_THREADS_PER_X;
  const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
  const int shared_memory_offset = DIVIDE_ROUND_UP(max_num_blocks_per_seq * sizeof(int), sizeof(float4)) * sizeof(float4);