/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#include <algorithm>
#include <cassert>
#include <climits>
#include <cstddef>
#include <iostream>
#include <type_traits>
#include "dequant_from_paged_cache.mluh"
#include "quant_utils.h"
// clang-format off
#include <mlu.h>
// clang-format on

namespace tmo {

namespace kernels {

#pragma bang walign(16)
#define REM_FOR_STACK (32 * 1024)
#define DEQUANT_WRAM_SIZE (__MLU_WRAM_SIZE__ * 1024)
#define DEQUANT_NRAM_SIZE (__MLU_NRAM_SIZE__ * 1024 - REM_FOR_STACK)
#define DEQUANT_PAGED_PERHEAD kernels::MLUDequantFromPagedCacheKernelPerHead
#define DEQUANT_PAGED_PERCHANNEL kernels::MLUDequantFromPagedCacheKernelPerChannel
#define DEQUANT_FUNC_LEN (24)
#define DEQUANT_BATCH_NUM (1024)

__wram__ int8_t wbuf[DEQUANT_WRAM_SIZE];
__nram__ int8_t nbuf[DEQUANT_NRAM_SIZE];
// Uses 8K = 1K * (4 + 4) to process offsets
__nram__ int32_t n_lens[DEQUANT_BATCH_NUM];
__nram__ int32_t n_offsets[DEQUANT_BATCH_NUM];

template <typename T, typename Tc, typename Ts>
__mlu_func__ void dequantize_per_channel(T *output_nram,
                                         Tc *input_nram,
                                         Ts *scale_nram,
                                         T *data,
                                         const int32_t scale_num,
                                         const int32_t seq_num,
                                         const int32_t head_num,
                                         const int32_t head_size,
                                         const size_t context_offset,
                                         const size_t context_seq_stride,
                                         const size_t context_head_stride) {
  dequantize<T, Tc, Ts>((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (Ts *)nbuf,
                        seq_num * scale_num, scale_num);
  __memcpy((T *)data + context_offset, (T *)output_nram, head_size * sizeof_(T), NRAM2GDRAM,
           context_head_stride * sizeof_(T), head_num - 1, context_seq_stride * sizeof_(T),
           seq_num - 1, head_size * sizeof_(T), head_num - 1, scale_num * sizeof_(T), seq_num - 1);
}

template <typename T, typename Tc, typename Ts>
__mlu_func__ void dequantize_per_head(T *output_nram,
                                      Tc *input_nram,
                                      Ts *scale_nram,
                                      T *temp_nram,
                                      T *data,
                                      const int32_t seq_num,
                                      const int32_t head_num,
                                      const int32_t block_size,
                                      const int32_t head_size,
                                      const size_t context_offset,
                                      const size_t context_seq_stride,
                                      const size_t context_head_stride) {
  int block_count = DIV_UP(seq_num, block_size);
  int rem_token = seq_num % block_size;
  convert((float *)output_nram, (Tc *)input_nram, block_count * head_num * block_size * head_size);
  T *res_nram = std::is_same<T, float>::value ? (T *)output_nram : (T *)temp_nram;
  // dequantize [block_count, head_num, block_size, head_size]
  conv_fuse_mul_cvt((T *)res_nram, (float *)scale_nram, (float *)wbuf, (float *)output_nram,
                    block_count * head_num * block_size, head_size, 1);
  // copy to [seq_num, head_num, head_size]
  int whole_block_count = block_count - int(rem_token > 0);
  if (whole_block_count) {
    for (int i = 0; i < head_num; ++i) {
      // copy from [whole_block_count, i, block_size, head_size]
      // to [whole_block_count * block_size, 1, head_size]
      __memcpy((T *)data + context_offset + i * context_head_stride,
               (T *)res_nram + i * block_size * head_size, head_size * sizeof_(T), NRAM2GDRAM,
               context_seq_stride * sizeof_(T), block_size - 1,
               block_size * context_seq_stride * sizeof_(T), whole_block_count - 1,
               head_size * sizeof_(T), block_size - 1,
               head_num * block_size * head_size * sizeof_(T), whole_block_count - 1);
    }
  }

  if (rem_token) {
    // copy from [last, head_num, block_size(rem_token), head_size]
    // to [rem_token, head_num, head_size]
    __memcpy((T *)data + context_offset + whole_block_count * block_size * context_seq_stride,
             (T *)res_nram + whole_block_count * head_num * block_size * head_size,
             head_size * sizeof_(T), NRAM2GDRAM, context_head_stride * sizeof_(T), head_num - 1,
             context_seq_stride * sizeof_(T), rem_token - 1, block_size * head_size * sizeof_(T),
             head_num - 1, head_size * sizeof_(T), rem_token - 1);
  }
}

template <typename Tc>
__mlu_func__ void load_input_per_channel(Tc *input_nram,
                                         Tc *cache,
                                         Tc *temp_nram,
                                         int32_t *block_offsets,
                                         uint32_t *cache_offsets,
                                         const int32_t *block_tables,
                                         const int32_t batch_idx,
                                         const int32_t scale_num,
                                         const int32_t max_block_num,
                                         const int32_t head_num,
                                         const int32_t block_size,
                                         const int32_t head_size,
                                         const int32_t seq_begin,
                                         const int32_t deal_seq_num,
                                         const size_t cache_bs_stride,
                                         const size_t cache_head_stride) {
  int32_t block_start = batch_idx * max_block_num + seq_begin / block_size;
  int32_t block_end = batch_idx * max_block_num + (seq_begin + deal_seq_num - 1) / block_size;
  int32_t block_count = block_end - block_start + 1;
  // make sure elements in block_tables >= 0
  __memcpy((int32_t *)block_offsets, (int32_t *)block_tables + block_start,
           block_count * sizeof_(int32_t), GDRAM2NRAM);
  __bang_mul_scalar((uint32_t *)cache_offsets, (uint32_t *)block_offsets,
                    (uint32_t)cache_bs_stride * sizeof(Tc), block_count);
#if __BANG_ARCH__ >= 500
  // gather [block_count, head_num, block_size, head_size]
  __gather((Tc *)input_nram, (Tc *)cache, (uint32_t *)cache_offsets,
           (uint32_t)cache_bs_stride * sizeof(Tc), GDRAM2NRAM,
           (uint32_t)cache_bs_stride * sizeof(Tc), block_count);
  if (head_num != 1 && block_size != 1) {
    // mv to [head_num, whole_block_count, block_size, head_size]
    __memcpy((Tc *)temp_nram, (Tc *)input_nram, block_size * head_size * sizeof(Tc), NRAM2NRAM,
             block_size * head_size * sizeof(Tc), block_count - 1,
             block_count * block_size * head_size * sizeof(Tc), head_num - 1,
             head_num * block_size * head_size * sizeof(Tc), block_count - 1,
             block_size * head_size * sizeof(Tc), head_num - 1);
    // mv to [whole_block_count, block_size, head_num, head_size]
    __memcpy((Tc *)input_nram, (Tc *)temp_nram, head_size * sizeof(Tc), NRAM2NRAM,
             head_size * sizeof(Tc), head_num - 1, head_num * head_size * sizeof(Tc),
             block_count * block_size - 1, block_count * block_size * head_size * sizeof(Tc),
             head_num - 1, head_size * sizeof(Tc), block_count * block_size - 1);
  }
#endif
}

template <typename Tc, typename Ts>
__mlu_func__ void load_input_per_head(Tc *input_nram,
                                      Ts *scale_nram,
                                      Tc *cache,
                                      Ts *scale,
                                      Tc *temp_nram,
                                      int32_t *block_offsets,
                                      uint32_t *cache_offsets,
                                      uint32_t *scale_offsets,
                                      const int32_t *block_tables,
                                      const int32_t batch_idx,
                                      const int32_t max_block_num,
                                      const int32_t head_num,
                                      const int32_t block_size,
                                      const int32_t head_size,
                                      const int32_t seq_begin,
                                      const int32_t deal_seq_num,
                                      const size_t cache_bs_stride,
                                      const size_t cache_head_stride,
                                      const size_t scale_bs_stride,
                                      const size_t scale_head_stride) {
  int32_t block_start = batch_idx * max_block_num + seq_begin / block_size;
  int32_t block_end = batch_idx * max_block_num + (seq_begin + deal_seq_num - 1) / block_size;
  int32_t block_count = block_end - block_start + 1;
  // make sure elements in block_tables >= 0
  __memcpy((int32_t *)block_offsets, (int32_t *)block_tables + block_start,
           block_count * sizeof_(int32_t), GDRAM2NRAM);
  __bang_mul_scalar((uint32_t *)cache_offsets, (uint32_t *)block_offsets,
                    (uint32_t)cache_bs_stride * sizeof(Tc), block_count);
  __bang_mul_scalar((uint32_t *)scale_offsets, (uint32_t *)block_offsets,
                    (uint32_t)scale_bs_stride * sizeof(Ts), block_count);
#if __BANG_ARCH__ >= 500
  // gather [block_count, head_num, block_size, head_size]
  __gather((Tc *)input_nram, (Tc *)cache, (uint32_t *)cache_offsets,
           (uint32_t)cache_bs_stride * sizeof(Tc), GDRAM2NRAM,
           (uint32_t)cache_bs_stride * sizeof(Tc), block_count);
  // gather [block_count, head_num, block_size]
  __gather((Ts *)scale_nram, (Ts *)scale, (uint32_t *)scale_offsets,
           (uint32_t)scale_bs_stride * sizeof(Ts), GDRAM2NRAM,
           (uint32_t)scale_bs_stride * sizeof(Ts), block_count);
#endif
}

template <typename T, typename Tc, typename Ts, bool ProcessOffsets>
__mlu_global__ void MLUDequantFromPagedCacheKernelPerChannel(void *key,
                                                             void *value,
                                                             const void *key_cache,
                                                             const void *value_cache,
                                                             const void *key_scale,
                                                             const void *value_scale,
                                                             const int32_t *context_lens,
                                                             const int32_t *context_seq_offsets,
                                                             const int32_t *block_tables,
                                                             const int32_t max_context_len,
                                                             const int32_t max_block_num,
                                                             const int32_t batch_size,
                                                             const int32_t head_num,
                                                             const int32_t key_group_num,
                                                             const int32_t value_group_num,
                                                             const int32_t block_size,
                                                             const int32_t head_size,
                                                             const int32_t seq_block,
                                                             const size_t context_head_stride,
                                                             const size_t context_seq_stride,
                                                             const size_t cache_bs_stride,
                                                             const size_t cache_head_stride,
                                                             const size_t key_cache_seq_stride,
                                                             const size_t value_cache_seq_stride,
                                                             const size_t scale_bs_stride,
                                                             const size_t scale_head_stride) {
  bool has_key = (key && key_cache && key_scale);
  bool has_value = (value && value_cache && value_scale);
  if (!(has_key || has_value)) {
    return;
  }

  /* *********************************nram space **************************************
   * NRAM |scale[head_num, head_size] fp32|output/input[seq_block, head_num, head_size] fp32|
   */
  int32_t scale_num = head_num * head_size;
  Ts *scale_nram = (Ts *)nbuf + (scale_num * (sizeof_(float) - sizeof_(Ts)) / sizeof_(Ts));
  float *output_nram = (float *)nbuf + scale_num;
  Tc *input_nram = (Tc *)output_nram +
                   (std::is_same<Tc, int4x2_t>::value
                        ? (7 * seq_block * (scale_num >> 1))
                        : (seq_block * scale_num * (sizeof_(float) - sizeof_(Tc)) / sizeof_(Tc)));
  int32_t *block_offsets = (int32_t *)((int8_t *)output_nram + seq_block * scale_num);
  uint32_t *cache_offsets = (uint32_t *)block_offsets + DIV_UP(seq_block, block_size);
  int32_t seq_offset;
  int32_t seq_len;
  int32_t seq_begin;
  int32_t deal_seq_num;
  size_t context_offset;
  process_offsets<ProcessOffsets>((int32_t *)n_lens, (int32_t *)n_offsets, (int32_t *)context_lens,
                                  (int32_t *)context_seq_offsets, batch_size);
  if (has_key) {
    load_scale_once((Ts *)scale_nram, (Ts *)key_scale, head_num, head_size, scale_bs_stride,
                    scale_head_stride);
    for (int32_t batch_idx = taskIdY; batch_idx < batch_size; batch_idx += taskDimY) {
      load_len_offset<ProcessOffsets>(seq_len, seq_offset, (int32_t *)n_lens, (int32_t *)n_offsets,
                                      (int32_t *)context_lens, (int32_t *)context_seq_offsets,
                                      batch_idx);
      // seq_begin % block_size != 0 only when seq_block < block_size
      seq_begin = taskIdZ * seq_block;
      deal_seq_num = std::min(seq_len - seq_begin, seq_block);
      if (deal_seq_num <= 0 || seq_offset < 0) continue;
      context_offset = context_seq_stride * (seq_offset + seq_begin);
      load_input_per_channel((Tc *)input_nram, (Tc *)key_cache, (Tc *)output_nram,
                             (int32_t *)block_offsets, (uint32_t *)cache_offsets,
                             (int32_t *)block_tables, batch_idx, scale_num, max_block_num, head_num,
                             block_size, head_size, seq_begin, deal_seq_num, cache_bs_stride,
                             cache_head_stride);
      dequantize_per_channel((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)key,
                             scale_num, deal_seq_num, head_num, head_size, context_offset,
                             context_seq_stride, context_head_stride);
    }
  }

  if (has_value) {
    load_scale_once((Ts *)scale_nram, (Ts *)value_scale, head_num, head_size, scale_bs_stride,
                    scale_head_stride);
    for (int32_t batch_idx = taskIdY; batch_idx < batch_size; batch_idx += taskDimY) {
      load_len_offset<ProcessOffsets>(seq_len, seq_offset, (int32_t *)n_lens, (int32_t *)n_offsets,
                                      (int32_t *)context_lens, (int32_t *)context_seq_offsets,
                                      batch_idx);
      // seq_begin % block_size != 0 only when seq_block < block_size
      seq_begin = taskIdZ * seq_block;
      deal_seq_num = std::min(seq_len - seq_begin, seq_block);
      if (deal_seq_num <= 0 || seq_offset < 0) continue;
      context_offset = context_seq_stride * (seq_offset + seq_begin);
      load_input_per_channel((Tc *)input_nram, (Tc *)value_cache, (Tc *)output_nram,
                             (int32_t *)block_offsets, (uint32_t *)cache_offsets,
                             (int32_t *)block_tables, batch_idx, scale_num, max_block_num, head_num,
                             block_size, head_size, seq_begin, deal_seq_num, cache_bs_stride,
                             cache_head_stride);
      dequantize_per_channel((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)value,
                             scale_num, deal_seq_num, head_num, head_size, context_offset,
                             context_seq_stride, context_head_stride);
    }
  }
}

template <typename T, typename Tc, typename Ts, bool ProcessOffsets>
__mlu_global__ void MLUDequantFromPagedCacheKernelPerHead(void *key,
                                                          void *value,
                                                          const void *key_cache,
                                                          const void *value_cache,
                                                          const void *key_scale,
                                                          const void *value_scale,
                                                          const int32_t *context_lens,
                                                          const int32_t *context_seq_offsets,
                                                          const int32_t *block_tables,
                                                          const int32_t max_context_len,
                                                          const int32_t max_block_num,
                                                          const int32_t batch_size,
                                                          const int32_t head_num,
                                                          const int32_t key_group_num,
                                                          const int32_t value_group_num,
                                                          const int32_t block_size,
                                                          const int32_t head_size,
                                                          const int32_t seq_block,
                                                          const size_t context_head_stride,
                                                          const size_t context_seq_stride,
                                                          const size_t cache_bs_stride,
                                                          const size_t cache_head_stride,
                                                          const size_t key_cache_seq_stride,
                                                          const size_t value_cache_seq_stride,
                                                          const size_t scale_bs_stride,
                                                          const size_t scale_head_stride) {
  bool has_key = (key && key_cache && key_scale);
  bool has_value = (value && value_cache && value_scale);
  if (!(has_key || has_value)) {
    return;
  }

  /* *********************************nram space **************************************
   * NRAM |scale[seq_block, head_num] * fp32|output/input[seq_block, head_num, head_size] * fp32|
   *      |temp[seq_block, head_num, head_size] * output_dtype|
   * WRAM |head_size * 64B|
   */
  int32_t scale_num = seq_block * head_num;
  Ts *scale_nram = (Ts *)nbuf + (scale_num * (sizeof_(float) - sizeof_(Ts)) / sizeof_(Ts));
  float *output_nram = (float *)nbuf + scale_num;
  Tc *input_nram =
      (Tc *)output_nram + (head_size * scale_num * (sizeof_(float) - sizeof_(Tc)) / sizeof_(Tc));
  // temp_nram for nram to store temp output
  float *temp_nram = (float *)output_nram + head_size * scale_num;
  int32_t *block_offsets = (int32_t *)((int8_t *)output_nram + head_size * scale_num);
  uint32_t *cache_offsets = (uint32_t *)block_offsets + DIV_UP(seq_block, block_size);
  uint32_t *scale_offsets = (uint32_t *)cache_offsets + DIV_UP(seq_block, block_size);
  int32_t seq_len;
  int32_t seq_begin;
  int32_t seq_offset;
  int32_t deal_seq_num;
  size_t context_offset;
  __bang_write_value((float *)nbuf, head_size * 16, 1.0f);
  mvNram2WramLT16<float>((int8_t *)wbuf, (int8_t *)nbuf, head_size, 16, 16);
  process_offsets<ProcessOffsets>((int32_t *)n_lens, (int32_t *)n_offsets, (int32_t *)context_lens,
                                  (int32_t *)context_seq_offsets, batch_size);
  if (has_key) {
    for (int32_t batch_idx = taskIdY; batch_idx < batch_size; batch_idx += taskDimY) {
      load_len_offset<ProcessOffsets>(seq_len, seq_offset, (int32_t *)n_lens, (int32_t *)n_offsets,
                                      (int32_t *)context_lens, (int32_t *)context_seq_offsets,
                                      batch_idx);
      seq_begin = taskIdZ * seq_block;
      deal_seq_num = std::min(seq_len - seq_begin, seq_block);
      if (deal_seq_num <= 0 || seq_offset < 0) continue;
      context_offset = context_seq_stride * (seq_offset + seq_begin);
      load_input_per_head((Tc *)input_nram, (Ts *)scale_nram, (Tc *)key_cache, (Ts *)key_scale,
                          (Tc *)temp_nram, (int32_t *)block_offsets, (uint32_t *)cache_offsets,
                          (uint32_t *)scale_offsets, (int32_t *)block_tables, batch_idx,
                          max_block_num, head_num, block_size, head_size, seq_begin, deal_seq_num,
                          cache_bs_stride, cache_head_stride, scale_bs_stride, scale_head_stride);
      dequantize_per_head((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)temp_nram,
                          (T *)key, deal_seq_num, head_num, block_size, head_size, context_offset,
                          context_seq_stride, context_head_stride);
    }
  }

  if (has_value) {
    for (int32_t batch_idx = taskIdY; batch_idx < batch_size; batch_idx += taskDimY) {
      load_len_offset<ProcessOffsets>(seq_len, seq_offset, (int32_t *)n_lens, (int32_t *)n_offsets,
                                      (int32_t *)context_lens, (int32_t *)context_seq_offsets,
                                      batch_idx);
      seq_begin = taskIdZ * seq_block;
      deal_seq_num = std::min(seq_len - seq_begin, seq_block);
      if (deal_seq_num <= 0 || seq_offset < 0) continue;
      context_offset = context_seq_stride * (seq_offset + seq_begin);
      load_input_per_head((Tc *)input_nram, (Ts *)scale_nram, (Tc *)value_cache, (Ts *)value_scale,
                          (Tc *)temp_nram, (int32_t *)block_offsets, (uint32_t *)cache_offsets,
                          (uint32_t *)scale_offsets, (int32_t *)block_tables, batch_idx,
                          max_block_num, head_num, block_size, head_size, seq_begin, deal_seq_num,
                          cache_bs_stride, cache_head_stride, scale_bs_stride, scale_head_stride);
      dequantize_per_head((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)temp_nram,
                          (T *)value, deal_seq_num, head_num, block_size, head_size, context_offset,
                          context_seq_stride, context_head_stride);
    }
  }
}
}  // namespace kernels

#define DEQUANT_PAGED_INIT(T, Tc, Ts, C, Name)                                              \
  template __mlu_global__ void kernels::MLUDequantFromPagedCacheKernel##Name<T, Tc, Ts, C>( \
      void *key, void *value, const void *key_cache, const void *value_cache,               \
      const void *key_scale, const void *value_scale, const int32_t *context_lens,          \
      const int32_t *context_seq_offsets, const int32_t *block_tables,                      \
      const int32_t max_context_len, const int32_t max_block_num, const int32_t batch_size, \
      const int32_t head_num, const int32_t key_group_num, const int32_t value_group_num,   \
      const int32_t block_size, const int32_t head_size, const int32_t seq_block,           \
      const size_t context_head_stride, const size_t context_seq_stride,                    \
      const size_t cache_bs_stride, const size_t cache_head_stride,                         \
      const size_t key_cache_seq_stride, const size_t value_cache_seq_stride,               \
      const size_t scale_bs_stride, const size_t scale_head_stride);

DEQUANT_PAGED_INIT(half, int8_t, float, false, PerChannel)
DEQUANT_PAGED_INIT(bfloat16_t, int8_t, float, false, PerChannel)
DEQUANT_PAGED_INIT(float, int8_t, float, false, PerChannel)
// DEQUANT_PAGED_INIT(half, int4x2_t, float, false, PerChannel)
// DEQUANT_PAGED_INIT(bfloat16_t, int4x2_t, float, false, PerChannel)
// DEQUANT_PAGED_INIT(float, int4x2_t, float, false, PerChannel)
DEQUANT_PAGED_INIT(half, int8_t, float, false, PerHead)
DEQUANT_PAGED_INIT(bfloat16_t, int8_t, float, false, PerHead)
DEQUANT_PAGED_INIT(float, int8_t, float, false, PerHead)
// DEQUANT_PAGED_INIT(half, int4x2_t, float, false, PerHead)
// DEQUANT_PAGED_INIT(bfloat16_t, int4x2_t, float, false, PerHead)
// DEQUANT_PAGED_INIT(float, int4x2_t, float, false, PerHead)
DEQUANT_PAGED_INIT(half, int8_t, float, true, PerChannel)
DEQUANT_PAGED_INIT(bfloat16_t, int8_t, float, true, PerChannel)
DEQUANT_PAGED_INIT(float, int8_t, float, true, PerChannel)
// DEQUANT_PAGED_INIT(half, int4x2_t, float, true, PerChannel)
// DEQUANT_PAGED_INIT(bfloat16_t, int4x2_t, float, true, PerChannel)
// DEQUANT_PAGED_INIT(float, int4x2_t, float, true, PerChannel)
DEQUANT_PAGED_INIT(half, int8_t, float, true, PerHead)
DEQUANT_PAGED_INIT(bfloat16_t, int8_t, float, true, PerHead)
DEQUANT_PAGED_INIT(float, int8_t, float, true, PerHead)
// DEQUANT_PAGED_INIT(half, int4x2_t, float, true, PerHead)
// DEQUANT_PAGED_INIT(bfloat16_t, int4x2_t, float, true, PerHead)
// DEQUANT_PAGED_INIT(float, int4x2_t, float, true, PerHead)

typedef void (*DequantFromPagedCachePointer)(void *,           // key
                                             void *,           // value
                                             const void *,     // key_cache
                                             const void *,     // value_cache
                                             const void *,     // key_scale
                                             const void *,     // value_scale
                                             const int32_t *,  // context_lens
                                             const int32_t *,  // context_seq_offsets
                                             const int32_t *,  // block_tables
                                             const int32_t,    // max_context_len
                                             const int32_t,    // max_block_num
                                             const int32_t,    // batch_size
                                             const int32_t,    // head_num
                                             const int32_t,    // key_group_num
                                             const int32_t,    // value_group_num
                                             const int32_t,    // block_size
                                             const int32_t,    // head_size
                                             const int32_t,    // seq_block
                                             const size_t,     // context_head_stride
                                             const size_t,     // context_seq_stride
                                             const size_t,     // cache_bs_stride
                                             const size_t,     // cache_head_stride
                                             const size_t,     // key_cache_seq_stride
                                             const size_t,     // value_cache_seq_stride
                                             const size_t,     // scale_bs_stride
                                             const size_t);    // scale_head_stride

static DequantFromPagedCachePointer DequantFromPagedCacheFuncArr[DEQUANT_FUNC_LEN] = {
    DEQUANT_PAGED_PERCHANNEL<half, int8_t, float, false>,
    DEQUANT_PAGED_PERCHANNEL<bfloat16_t, int8_t, float, false>,
    DEQUANT_PAGED_PERCHANNEL<float, int8_t, float, false>, nullptr, nullptr, nullptr,
    // DEQUANT_PAGED_PERCHANNEL<half, int4x2_t, float, false>,
    // DEQUANT_PAGED_PERCHANNEL<bfloat16_t, int4x2_t, float, false>,
    // DEQUANT_PAGED_PERCHANNEL<float, int4x2_t, float, false>,
    DEQUANT_PAGED_PERHEAD<half, int8_t, float, false>,
    DEQUANT_PAGED_PERHEAD<bfloat16_t, int8_t, float, false>,
    DEQUANT_PAGED_PERHEAD<float, int8_t, float, false>, nullptr, nullptr, nullptr,
    // DEQUANT_PAGED_PERHEAD<half, int4x2_t, float, false>,
    // DEQUANT_PAGED_PERHEAD<bfloat16_t, int4x2_t, float, false>,
    // DEQUANT_PAGED_PERHEAD<float, int4x2_t, float, false>,
    DEQUANT_PAGED_PERCHANNEL<half, int8_t, float, true>,
    DEQUANT_PAGED_PERCHANNEL<bfloat16_t, int8_t, float, true>,
    DEQUANT_PAGED_PERCHANNEL<float, int8_t, float, true>, nullptr, nullptr, nullptr,
    // DEQUANT_PAGED_PERCHANNEL<half, int4x2_t, float, true>,
    // DEQUANT_PAGED_PERCHANNEL<bfloat16_t, int4x2_t, float, true>,
    // DEQUANT_PAGED_PERCHANNEL<float, int4x2_t, float, true>,
    DEQUANT_PAGED_PERHEAD<half, int8_t, float, true>,
    DEQUANT_PAGED_PERHEAD<bfloat16_t, int8_t, float, true>,
    DEQUANT_PAGED_PERHEAD<float, int8_t, float, true>, nullptr, nullptr, nullptr};
// DEQUANT_PAGED_PERHEAD<half, int4x2_t, float, true>,
// DEQUANT_PAGED_PERHEAD<bfloat16_t, int4x2_t, float, true>,
// DEQUANT_PAGED_PERHEAD<float, int4x2_t, float, true>};

uint32_t getDequantPagedIdx(cnnlDataType_t dtype,
                            int32_t quant_mode,
                            int32_t quant_bit,
                            const void *context_seq_offset) {
  uint32_t idx = 0;
  idx += (quant_mode != 0) ? 6 : 0;
  idx += (quant_bit != 8) ? 3 : 0;
  idx += (dtype == CNNL_DTYPE_BFLOAT16) ? 1 : 0;
  idx += (dtype == CNNL_DTYPE_FLOAT) ? 2 : 0;
  idx += (context_seq_offset == nullptr) ? 12 : 0;

  return idx;
}

void getBlockAndDimForPaged(int32_t &seq_block,
                            cnrtDim3_t &task_dim,
                            cnrtFunctionType_t &task_type,
                            const int32_t max_context_len,
                            const int32_t head_num,
                            const int32_t batch_size,
                            const int32_t head_size,
                            const int32_t block_size,
                            const int32_t quant_mode,
                            const int32_t quant_bit,
                            const cnnlDataType_t dtype) {
  int32_t core_dim;
  int32_t cluster_dim;
  int32_t nram_size = 480 * 1024;
  int32_t wram_size = 512 * 1024;
  int32_t sram_size = 2016 * 1024;
  getDeviceCoreAndRam(cluster_dim, core_dim, nram_size, wram_size, sram_size, REM_FOR_STACK);
  if (quant_mode == 0) {
    seq_block = int32_t((int64_t)nram_size / ((int64_t)head_num * head_size * sizeof_(float)) - 1);
    if (seq_block < block_size) {
      std::cerr << __func__ << "," << __LINE__
                << " :head_num * head_size * sizeof_(float) should be less than "
                << nram_size / block_size << " when quant_mode is 0." << std::endl;
    }
  } else {
    int32_t dtype_size = dtype == CNNL_DTYPE_FLOAT ? sizeof_(float) : sizeof_(half);
    seq_block = int32_t((int64_t)nram_size / ((int64_t)head_num * (head_size + 1) * sizeof_(float) +
                                              (int64_t)head_num * head_size * dtype_size));
    if (seq_block < block_size) {
      std::cerr << __func__ << "," << __LINE__
                << " :head_num * sizeof_(float) + head_num * head_size * (sizeof_(float) + "
                   "context_dtype_size) "
                << "should be less than " << nram_size / block_size << " when quant_mode is 1."
                << std::endl;
    }

    /* head_size * 64B put in the wram. */
    if (head_size * ONE_LINE >= wram_size) {
      std::cerr << __func__ << "," << __LINE__ << " head_size * 64 " << "should be less than "
                << wram_size << " when quant_mode is 1." << std::endl;
    }
  }
  // seq_block should be a multiply of block_size
  seq_block = PAD_DOWN(seq_block, block_size);
  int seq_seg = DIV_UP(max_context_len, seq_block);
  int32_t core_num = cluster_dim * core_dim;
  if (batch_size * seq_seg <= (core_num / 2)) {
    int times = core_num / batch_size / seq_seg;
    seq_block = std::max(seq_block / times, 2);
    if (seq_block > block_size) {
      seq_block = PAD_DOWN(seq_block, block_size);
    } else {
      seq_block = block_size;
    }
    seq_seg = DIV_UP(max_context_len, seq_block);
  }

  task_dim.x = 1;
  task_dim.y = uint32_t(std::min(batch_size, core_num));
  task_dim.z = uint32_t(seq_seg);
  task_type = cnrtFuncTypeBlock;
}

KernelStatus invokeDequantFromPagedCache(cnrtQueue_t queue,
                                         void *key,
                                         void *value,
                                         const void *key_cache,
                                         const void *value_cache,
                                         const void *key_scale,
                                         const void *value_scale,
                                         const void *context_lens,
                                         const void *context_seq_offsets,
                                         const void *block_tables,
                                         const int32_t max_context_len,
                                         const int32_t max_block_num,
                                         const int32_t batch_size,
                                         const int32_t head_num,
                                         const int32_t key_group_num,
                                         const int32_t value_group_num,
                                         const int32_t block_size,
                                         const int32_t head_size,
                                         const int32_t quant_mode,
                                         const int32_t quant_bit,
                                         const size_t context_head_stride,
                                         const size_t context_seq_stride,
                                         const size_t cache_bs_stride,
                                         const size_t cache_head_stride,
                                         const size_t key_cache_seq_stride,
                                         const size_t value_cache_seq_stride,
                                         const size_t scale_bs_stride,
                                         const size_t scale_head_stride,
                                         const cnnlDataType_t dtype) {
  if (is_arch300()) {
    std::cerr << "[invokeDequantFromPagedCache]: kernel does not support MLU300 devices."
              << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  int32_t index;
  int32_t seq_block;
  cnrtDim3_t k_dim;
  cnrtFunctionType_t k_type;
  getBlockAndDimForPaged(seq_block, k_dim, k_type, max_context_len, head_num, batch_size, head_size,
                         block_size, quant_mode, quant_bit, dtype);
  index = getDequantPagedIdx(dtype, quant_mode, quant_bit, context_seq_offsets);
  auto dequant_paged_func = DequantFromPagedCacheFuncArr[index];
  dequant_paged_func<<<k_dim, k_type, queue>>>(
      (void *)key, (void *)value, (const void *)key_cache, (const void *)value_cache,
      (const void *)key_scale, (const void *)value_scale, (const int32_t *)context_lens,
      (const int32_t *)context_seq_offsets, (const int32_t *)block_tables, max_context_len,
      max_block_num, batch_size, head_num, key_group_num, value_group_num, block_size, head_size,
      seq_block, context_head_stride, context_seq_stride, cache_bs_stride, cache_head_stride,
      key_cache_seq_stride, value_cache_seq_stride, scale_bs_stride, scale_head_stride);

  return KernelStatus::KERNEL_STATUS_SUCCESS;
}
}  // namespace tmo
