/*************************************************************************
 * 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_linear_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_LINEAR_PERHEAD kernels::MLUDequantFromLinearCacheKernelPerHead
#define DEQUANT_LINEAR_PERCHANNEL kernels::MLUDequantFromLinearCacheKernelPerChannel
#define DEQUANT_FUNC_LEN (24)
#define DEQUANT_BATCH_NUM (1024)

__wram__ int8_t wbuf[DEQUANT_WRAM_SIZE];
__nram__ int8_t nbuf[DEQUANT_NRAM_SIZE];
__nram__ uint8_t pre_table_nram[TRANS_TABLE_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];

__mlu_func__ void calcu_offsets_per_channel(int32_t &cache_id,
                                            size_t &context_offset,
                                            size_t &cache_offset,
                                            size_t &scale_offset,
                                            const int32_t *cache_bs_id,
                                            const int32_t *cache_seq_offsets,
                                            const int32_t cache_mem_len,
                                            const int32_t seq_len,
                                            const int32_t seq_begin,
                                            const int32_t seq_offset,
                                            const int32_t batch_idx,
                                            const size_t context_seq_stride,
                                            const size_t cache_bs_stride,
                                            const size_t cache_seq_stride,
                                            const size_t scale_bs_stride) {
  cache_id = cache_bs_id == nullptr ? batch_idx : __load_gdram((int32_t *)cache_bs_id + batch_idx);
  int32_t cache_seq_offset =
      cache_seq_offsets == nullptr ? 0 : __load_gdram((int32_t *)cache_seq_offsets + batch_idx);
  if (cache_id >= 0 && cache_seq_offset >= 0 && (cache_seq_offset + seq_len) <= cache_mem_len) {
    context_offset = context_seq_stride * (seq_offset + seq_begin);
    cache_offset = cache_bs_stride * cache_id + cache_seq_stride * (cache_seq_offset + seq_begin);
    scale_offset = scale_bs_stride * cache_id;
  } else {
    cache_id = -1;
  }
}

__mlu_func__ void calcu_offsets_per_head(int32_t &cache_id,
                                         size_t &context_offset,
                                         size_t &key_cache_offset,
                                         size_t &value_cache_offset,
                                         size_t &scale_offset,
                                         const int32_t *cache_bs_id,
                                         const int32_t *cache_seq_offsets,
                                         const int32_t cache_mem_len,
                                         const int32_t seq_len,
                                         const int32_t seq_begin,
                                         const int32_t seq_offset,
                                         const int32_t batch_idx,
                                         const size_t context_seq_stride,
                                         const size_t cache_bs_stride,
                                         const size_t key_cache_seq_stride,
                                         const size_t value_cache_seq_stride,
                                         const size_t scale_bs_stride) {
  cache_id = cache_bs_id == nullptr ? batch_idx : __load_gdram((int32_t *)cache_bs_id + batch_idx);
  int32_t cache_seq_offset =
      cache_seq_offsets == nullptr ? 0 : __load_gdram((int32_t *)cache_seq_offsets + batch_idx);
  if (cache_id >= 0 && cache_seq_offset >= 0 && (cache_seq_offset + seq_len) <= cache_mem_len) {
    context_offset = context_seq_stride * (seq_offset + seq_begin);
    key_cache_offset =
        cache_bs_stride * cache_id + key_cache_seq_stride * (cache_seq_offset + seq_begin);
    value_cache_offset =
        cache_bs_stride * cache_id + value_cache_seq_stride * (cache_seq_offset + seq_begin);
    scale_offset = cache_seq_offset + seq_begin + scale_bs_stride * cache_id;
  } else {
    cache_id = -1;
  }
}

template <typename T, typename Tc, typename Ts>
__mlu_func__ void dequantize_per_channel(T *output_nram,
                                         Tc *input_nram,
                                         Ts *scale_nram,
                                         T *data,
                                         Tc *cache,
                                         Ts *scale,
                                         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 cache_offset,
                                         const size_t scale_offset,
                                         const size_t context_seq_stride,
                                         const size_t context_head_stride,
                                         const size_t cache_head_stride,
                                         const size_t cache_seq_stride,
                                         const size_t scale_bs_stride,
                                         const size_t scale_head_stride) {
  if (scale_bs_stride != 0) {
    __memcpy((Ts *)scale_nram, (Ts *)scale + scale_offset, head_size * sizeof_(Ts), GDRAM2NRAM,
             head_size * sizeof_(Ts), scale_head_stride * sizeof_(Ts), head_num - 1);
  }

  if (std::is_same<Tc, int4x2_t>::value) {
    __memcpy((Tc *)input_nram, (Tc *)cache + cache_offset, head_size >> 1, GDRAM2NRAM,
             head_size >> 1, head_num - 1, scale_num >> 1, seq_num - 1, cache_head_stride,
             head_num - 1, cache_seq_stride, seq_num - 1);
  } else {
    __memcpy((Tc *)input_nram, (Tc *)cache + cache_offset, head_size * sizeof_(Tc), GDRAM2NRAM,
             head_size * sizeof_(Tc), head_num - 1, scale_num * sizeof_(Tc), seq_num - 1,
             cache_head_stride * sizeof_(Tc), head_num - 1, cache_seq_stride * sizeof_(Tc),
             seq_num - 1);
  }

  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_value_per_channel(T *output_nram,
                                               Tc *input_nram,
                                               Ts *scale_nram,
                                               int8_t *temp_nram,
                                               T *data,
                                               Tc *cache,
                                               Ts *scale,
                                               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 cache_offset,
                                               const size_t scale_offset,
                                               const size_t context_seq_stride,
                                               const size_t context_head_stride,
                                               const size_t cache_head_stride,
                                               const size_t cache_seq_stride,
                                               const size_t scale_bs_stride,
                                               const size_t scale_head_stride,
                                               const bool pad_front) {
  /* Step 1. load scale [head_num, head_size]*/
  if (scale_bs_stride != 0) {
    __memcpy((Ts *)scale_nram, (Ts *)scale + scale_offset, head_size * sizeof_(Ts), GDRAM2NRAM,
             head_size * sizeof_(Ts), scale_head_stride * sizeof_(Ts), head_num - 1);
  }

  /* Step 2. load cache [load_seq_num, head_num, head_size] */
  int32_t load_seq_num = (seq_num >> 1) + int32_t(seq_num % 2);
  int32_t deal_seq_num = load_seq_num << 1;
  __memcpy((Tc *)input_nram, (Tc *)cache + cache_offset, head_size, GDRAM2NRAM, head_size,
           head_num - 1, scale_num, load_seq_num - 1, cache_head_stride, head_num - 1,
           cache_seq_stride, load_seq_num - 1);
  /* Step 3. convert into int8 [load_seq_num, head_num, head_size, 2] */
  convert((int8_t *)output_nram, (int4x2_t *)input_nram, deal_seq_num * scale_num);
  /* Step 4. transpose to [deal_seq_num (load_seq_num, 2), head_num, head_size] */
  trans_nhwc2nchw_smallc((int8_t *)temp_nram, (int8_t *)output_nram, (uint8_t *)pre_table_nram,
                         load_seq_num, head_num, head_size, 2);
  /* Step 5. dequantize [save_seq_num, head_num, head_size] */
  int save_seq_num = pad_front ? seq_num - 1 : seq_num;
  dequantize<T, int8_t, Ts>((T *)output_nram, (int8_t *)temp_nram + (pad_front ? scale_num : 0),
                            (Ts *)scale_nram, (Ts *)nbuf, save_seq_num * scale_num, scale_num);
  /* Step 6. store [save_seq_num, head_num, head_size]*/
  __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),
           save_seq_num - 1, head_size * sizeof_(T), head_num - 1, scale_num * sizeof_(T),
           save_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,
                                      Tc *cache,
                                      Ts *scale,
                                      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 cache_offset,
                                      const size_t scale_offset,
                                      const size_t context_seq_stride,
                                      const size_t context_head_stride,
                                      const size_t cache_head_stride,
                                      const size_t cache_seq_stride,
                                      const size_t scale_head_stride) {
  __memcpy((Ts *)scale_nram, (Ts *)scale + scale_offset, seq_num * sizeof_(Ts), GDRAM2NRAM,
           seq_num * sizeof_(Ts), scale_head_stride * sizeof_(Ts), head_num - 1);
  if (std::is_same<Tc, int4x2_t>::value) {
    __memcpy((Tc *)input_nram, (Tc *)cache + cache_offset, head_size >> 1, GDRAM2NRAM,
             head_size >> 1, seq_num - 1, seq_num * (head_size >> 1), head_num - 1,
             cache_seq_stride, seq_num - 1, cache_head_stride, head_num - 1);
  } else {
    __memcpy((Tc *)input_nram, (Tc *)cache + cache_offset, head_size * sizeof_(Tc), GDRAM2NRAM,
             head_size * sizeof_(Tc), seq_num - 1, seq_num * head_size * sizeof_(Tc), head_num - 1,
             cache_seq_stride * sizeof_(Tc), seq_num - 1, cache_head_stride * sizeof_(Tc),
             head_num - 1);
  }

  convert((float *)output_nram, (Tc *)input_nram, head_num * seq_num * head_size);
  if (std::is_same<T, float>::value) {
    conv_fuse_mul_cvt((T *)output_nram, (float *)scale_nram, (float *)wbuf, (float *)output_nram,
                      head_num * seq_num, head_size, 1);
  } else {
    conv_fuse_mul_cvt((T *)temp_nram, (float *)scale_nram, (float *)wbuf, (float *)output_nram,
                      head_num * seq_num, head_size, 1);
    output_nram = (T *)temp_nram;
  }

  __memcpy((T *)data + context_offset, (T *)output_nram, head_size * sizeof_(T), NRAM2GDRAM,
           context_seq_stride * sizeof_(T), seq_num - 1, context_head_stride * sizeof_(T),
           head_num - 1, head_size * sizeof_(T), seq_num - 1, head_size * seq_num * sizeof_(T),
           head_num - 1);
}

template <typename T, typename Tc, typename Ts>
__mlu_func__ void dequantize_value_per_head(T *output_nram,
                                            Tc *input_nram,
                                            Ts *scale_nram,
                                            T *temp_nram,
                                            T *data,
                                            Tc *cache,
                                            Ts *scale,
                                            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 cache_offset,
                                            const size_t scale_offset,
                                            const size_t context_seq_stride,
                                            const size_t context_head_stride,
                                            const size_t cache_head_stride,
                                            const size_t cache_seq_stride,
                                            const size_t scale_bs_stride,
                                            const size_t scale_head_stride,
                                            const bool pad_front) {
  int32_t load_seq_num = (seq_num >> 1) + int32_t(seq_num % 2);
  int32_t deal_seq_num = load_seq_num << 1;
  /* Step1. load scale first, [head_num, deal_seq_num] */
  __memcpy((Ts *)scale_nram, (Ts *)scale + scale_offset, deal_seq_num * sizeof_(Ts), GDRAM2NRAM,
           deal_seq_num * sizeof_(Ts), scale_head_stride * sizeof_(Ts), head_num - 1);
  /* Step2. load cache input, [head_num, load_seq_num, head_size, 2] for int4 */
  __memcpy((Tc *)input_nram, (Tc *)cache + cache_offset, head_size, GDRAM2NRAM, head_size,
           load_seq_num - 1, load_seq_num * head_size, head_num - 1, cache_seq_stride,
           load_seq_num - 1, cache_head_stride, head_num - 1);
  convert((int8_t *)output_nram, (Tc *)input_nram, head_num * head_size * deal_seq_num);
  /* Step3. trans to [head_num, load_seq_num, 2, head_size]*/
  trans_nhwc2nchw_smallc((int8_t *)temp_nram, (int8_t *)output_nram, (uint8_t *)pre_table_nram,
                         head_num * load_seq_num, head_size, 1, 2);
  /* Step4. dequant to T [head_num, deal_seq_num, head_size] */
  convert((float *)output_nram, (int8_t *)temp_nram, head_num * deal_seq_num * head_size);
  if (std::is_same<T, float>::value) {
    conv_fuse_mul_cvt((T *)output_nram, (float *)scale_nram, (float *)wbuf, (float *)output_nram,
                      head_num * deal_seq_num, head_size, 1);
  } else {
    conv_fuse_mul_cvt((T *)temp_nram, (float *)scale_nram, (float *)wbuf, (float *)output_nram,
                      head_num * deal_seq_num, head_size, 1);
    output_nram = (T *)temp_nram;
  }

  /* Step5. save [head_num, save_seq_num, head_size]*/
  int32_t save_seq_num = pad_front ? seq_num - 1 : seq_num;
  __memcpy((T *)data + context_offset, (T *)output_nram + (pad_front ? head_size : 0),
           head_size * sizeof_(T), NRAM2GDRAM, context_seq_stride * sizeof_(T), save_seq_num - 1,
           context_head_stride * sizeof_(T), head_num - 1, head_size * sizeof_(T), save_seq_num - 1,
           head_size * deal_seq_num * sizeof_(T), head_num - 1);
}

template <typename T, typename Tc, typename Ts, bool ProcessOffsets>
__mlu_global__ void MLUDequantFromLinearCacheKernelPerChannel(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 *cache_bs_id,
                                                              const int32_t *cache_seq_offsets,
                                                              const int32_t max_context_len,
                                                              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 cache_mem_len,
                                                              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]|output/input[seq_block, head_num, head_size]|
   */
  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)));
  // temp_nram for nram to store int8 input
  int8_t *temp_nram = (int8_t *)output_nram + seq_block * scale_num * 3;
  int32_t seq_offset;
  int32_t seq_len;
  int32_t seq_begin;
  int32_t deal_seq_num;
  int32_t cache_id;
  int32_t cache_seq_offset;
  size_t context_offset;
  size_t cache_offset;
  size_t scale_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 = taskIdZ * seq_block;
      deal_seq_num = std::min(seq_len - seq_begin, seq_block);
      if (deal_seq_num <= 0 || seq_offset < 0) continue;
      calcu_offsets_per_channel(cache_id, context_offset, cache_offset, scale_offset,
                                (int32_t *)cache_bs_id, (int32_t *)cache_seq_offsets, cache_mem_len,
                                seq_len, seq_begin, seq_offset, batch_idx, context_seq_stride,
                                cache_bs_stride, key_cache_seq_stride, scale_bs_stride);
      if (cache_id < 0) continue;
      dequantize_per_channel((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)key,
                             (Tc *)key_cache, (Ts *)key_scale, scale_num, deal_seq_num, head_num,
                             head_size, context_offset, cache_offset, scale_offset,
                             context_seq_stride, context_head_stride, cache_head_stride,
                             key_cache_seq_stride, scale_bs_stride, scale_head_stride);
    }
  }

  if (has_value) {
    if (std::is_same<Tc, int4x2_t>::value) {
      __reshape_nhwc2nchw_smallc_init<int8_t>(pre_table_nram, 2);
    }
    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);
      if (std::is_same<Tc, int4x2_t>::value) {
        seq_begin = taskIdZ * seq_block;
        cache_id =
            cache_bs_id == nullptr ? batch_idx : __load_gdram((int32_t *)cache_bs_id + batch_idx);
        cache_seq_offset = cache_seq_offsets == nullptr
                               ? 0
                               : __load_gdram((int32_t *)cache_seq_offsets + batch_idx);
        // move seq_begin left by 1 when cache_seq_offset is odd
        seq_begin = cache_seq_offset % 2 ? seq_begin - 1 : seq_begin;
        deal_seq_num = std::min(seq_len - seq_begin, seq_block);
        if (deal_seq_num <= 0 || seq_offset < 0) continue;
        if (cache_id >= 0 && cache_seq_offset >= 0 &&
            (cache_seq_offset + seq_len) <= cache_mem_len) {
          context_offset =
              context_seq_stride * (seq_offset + seq_begin + ((seq_begin == -1) ? 1 : 0));
          // value cache is [max_batch_size, head_num, cache_mem_len/2, head_size] for int4x2_t
          cache_offset = cache_bs_stride * cache_id +
                         value_cache_seq_stride * ((cache_seq_offset + seq_begin) / 2);
          scale_offset = scale_bs_stride * cache_id;
        } else {
          cache_id = -1;
        }
        if (cache_id < 0) continue;
        dequantize_value_per_channel(
            (T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (int8_t *)temp_nram, (T *)value,
            (Tc *)value_cache, (Ts *)value_scale, scale_num, deal_seq_num, head_num, head_size,
            context_offset, cache_offset, scale_offset, context_seq_stride, context_head_stride,
            cache_head_stride, value_cache_seq_stride, scale_bs_stride, scale_head_stride,
            seq_begin == -1);
      } else {
        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;
        calcu_offsets_per_channel(
            cache_id, context_offset, cache_offset, scale_offset, (int32_t *)cache_bs_id,
            (int32_t *)cache_seq_offsets, cache_mem_len, seq_len, seq_begin, seq_offset, batch_idx,
            context_seq_stride, cache_bs_stride, value_cache_seq_stride, scale_bs_stride);
        if (cache_id < 0) continue;
        dequantize_per_channel((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)value,
                               (Tc *)value_cache, (Ts *)value_scale, scale_num, deal_seq_num,
                               head_num, head_size, context_offset, cache_offset, scale_offset,
                               context_seq_stride, context_head_stride, cache_head_stride,
                               value_cache_seq_stride, scale_bs_stride, scale_head_stride);
      }
    }
  }
}

template <typename T, typename Tc, typename Ts, bool ProcessOffsets>
__mlu_global__ void MLUDequantFromLinearCacheKernelPerHead(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 *cache_bs_id,
                                                           const int32_t *cache_seq_offsets,
                                                           const int32_t max_context_len,
                                                           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 cache_mem_len,
                                                           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]|output/input[head_size, seq_block, head_num]|temp|
   */
  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 +
                   (std::is_same<Tc, int4x2_t>::value
                        ? (7 * head_size * (scale_num >> 1))
                        : head_size * scale_num * (sizeof_(float) - sizeof_(Tc)) / sizeof_(Tc));
  // temp_nram for nram to store converted output
  float *temp_nram = (float *)output_nram + head_size * scale_num;
  int32_t seq_offset;
  int32_t seq_len;
  int32_t seq_begin;
  int32_t deal_seq_num;
  int32_t cache_id;
  size_t context_offset;
  size_t key_cache_offset;
  size_t value_cache_offset;
  size_t scale_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);
  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;
    calcu_offsets_per_head(cache_id, context_offset, key_cache_offset, value_cache_offset,
                           scale_offset, (int32_t *)cache_bs_id, (int32_t *)cache_seq_offsets,
                           cache_mem_len, seq_len, seq_begin, seq_offset, batch_idx,
                           context_seq_stride, cache_bs_stride, key_cache_seq_stride,
                           value_cache_seq_stride, scale_bs_stride);
    if (cache_id < 0) continue;
    if (has_key) {
      dequantize_per_head((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)temp_nram,
                          (T *)key, (Tc *)key_cache, (Ts *)key_scale, scale_num, deal_seq_num,
                          head_num, head_size, context_offset, key_cache_offset, scale_offset,
                          context_seq_stride, context_head_stride, cache_head_stride,
                          key_cache_seq_stride, scale_head_stride);
    }
    if (has_value && std::is_same<Tc, int8_t>::value) {
      dequantize_per_head((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram, (T *)temp_nram,
                          (T *)value, (Tc *)value_cache, (Ts *)value_scale, scale_num, deal_seq_num,
                          head_num, head_size, context_offset, value_cache_offset, scale_offset,
                          context_seq_stride, context_head_stride, cache_head_stride,
                          value_cache_seq_stride, scale_head_stride);
    }
  }

  // process value int4 differently
  if (has_value && std::is_same<Tc, int4x2_t>::value) {
    int32_t cache_seq_offset;
    __reshape_nhwc2nchw_smallc_init<int8_t>(pre_table_nram, 2);
    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;
      cache_id =
          cache_bs_id == nullptr ? batch_idx : __load_gdram((int32_t *)cache_bs_id + batch_idx);
      cache_seq_offset =
          cache_seq_offsets == nullptr ? 0 : __load_gdram((int32_t *)cache_seq_offsets + batch_idx);
      // move seq_begin left by 1 when cache_seq_offset is odd
      seq_begin = cache_seq_offset % 2 ? seq_begin - 1 : seq_begin;
      deal_seq_num = std::min(seq_len - seq_begin, seq_block);
      if (deal_seq_num <= 0 || seq_offset < 0) continue;
      if (cache_id >= 0 && cache_seq_offset >= 0 && (cache_seq_offset + seq_len) <= cache_mem_len) {
        context_offset =
            context_seq_stride * (seq_offset + seq_begin + ((seq_begin == -1) ? 1 : 0));
        // value cache is [max_batch_size, head_num, cache_mem_len/2, head_size] for int4x2_t
        value_cache_offset = cache_bs_stride * cache_id +
                             value_cache_seq_stride * ((cache_seq_offset + seq_begin) / 2);
        scale_offset = cache_seq_offset + seq_begin + scale_bs_stride * cache_id;
      } else {
        cache_id = -1;
      }
      if (cache_id < 0) continue;
      dequantize_value_per_head((T *)output_nram, (Tc *)input_nram, (Ts *)scale_nram,
                                (T *)temp_nram, (T *)value, (Tc *)value_cache, (Ts *)value_scale,
                                scale_num, deal_seq_num, head_num, head_size, context_offset,
                                value_cache_offset, scale_offset, context_seq_stride,
                                context_head_stride, cache_head_stride, value_cache_seq_stride,
                                scale_bs_stride, scale_head_stride, seq_begin == -1);
    }
  }
}
}  // namespace kernels

#define DEQUANT_LINEAR_INIT(T, Tc, Ts, C, Name)                                                  \
  template __mlu_global__ void kernels::MLUDequantFromLinearCacheKernel##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 *cache_bs_id,                            \
      const int32_t *cache_seq_offsets, const int32_t max_context_len, 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 cache_mem_len, 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_LINEAR_INIT(half, int8_t, float, false, PerChannel)
DEQUANT_LINEAR_INIT(bfloat16_t, int8_t, float, false, PerChannel)
DEQUANT_LINEAR_INIT(float, int8_t, float, false, PerChannel)
DEQUANT_LINEAR_INIT(half, int4x2_t, float, false, PerChannel)
DEQUANT_LINEAR_INIT(bfloat16_t, int4x2_t, float, false, PerChannel)
DEQUANT_LINEAR_INIT(float, int4x2_t, float, false, PerChannel)
DEQUANT_LINEAR_INIT(half, int8_t, float, false, PerHead)
DEQUANT_LINEAR_INIT(bfloat16_t, int8_t, float, false, PerHead)
DEQUANT_LINEAR_INIT(float, int8_t, float, false, PerHead)
DEQUANT_LINEAR_INIT(half, int4x2_t, float, false, PerHead)
DEQUANT_LINEAR_INIT(bfloat16_t, int4x2_t, float, false, PerHead)
DEQUANT_LINEAR_INIT(float, int4x2_t, float, false, PerHead)
DEQUANT_LINEAR_INIT(half, int8_t, float, true, PerChannel)
DEQUANT_LINEAR_INIT(bfloat16_t, int8_t, float, true, PerChannel)
DEQUANT_LINEAR_INIT(float, int8_t, float, true, PerChannel)
DEQUANT_LINEAR_INIT(half, int4x2_t, float, true, PerChannel)
DEQUANT_LINEAR_INIT(bfloat16_t, int4x2_t, float, true, PerChannel)
DEQUANT_LINEAR_INIT(float, int4x2_t, float, true, PerChannel)
DEQUANT_LINEAR_INIT(half, int8_t, float, true, PerHead)
DEQUANT_LINEAR_INIT(bfloat16_t, int8_t, float, true, PerHead)
DEQUANT_LINEAR_INIT(float, int8_t, float, true, PerHead)
DEQUANT_LINEAR_INIT(half, int4x2_t, float, true, PerHead)
DEQUANT_LINEAR_INIT(bfloat16_t, int4x2_t, float, true, PerHead)
DEQUANT_LINEAR_INIT(float, int4x2_t, float, true, PerHead)

typedef void (*DequantFromLinearCachePointer)(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 *,  // cache_bs_id
                                              const int32_t *,  // cache_seq_offsets
                                              const int32_t,    // max_context_len
                                              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,    // cache_mem_len
                                              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 DequantFromLinearCachePointer DequantFromLinearCacheFuncArr[DEQUANT_FUNC_LEN] = {
    DEQUANT_LINEAR_PERCHANNEL<half, int8_t, float, false>,
    DEQUANT_LINEAR_PERCHANNEL<bfloat16_t, int8_t, float, false>,
    DEQUANT_LINEAR_PERCHANNEL<float, int8_t, float, false>,
    DEQUANT_LINEAR_PERCHANNEL<half, int4x2_t, float, false>,
    DEQUANT_LINEAR_PERCHANNEL<bfloat16_t, int4x2_t, float, false>,
    DEQUANT_LINEAR_PERCHANNEL<float, int4x2_t, float, false>,
    DEQUANT_LINEAR_PERHEAD<half, int8_t, float, false>,
    DEQUANT_LINEAR_PERHEAD<bfloat16_t, int8_t, float, false>,
    DEQUANT_LINEAR_PERHEAD<float, int8_t, float, false>,
    DEQUANT_LINEAR_PERHEAD<half, int4x2_t, float, false>,
    DEQUANT_LINEAR_PERHEAD<bfloat16_t, int4x2_t, float, false>,
    DEQUANT_LINEAR_PERHEAD<float, int4x2_t, float, false>,
    DEQUANT_LINEAR_PERCHANNEL<half, int8_t, float, true>,
    DEQUANT_LINEAR_PERCHANNEL<bfloat16_t, int8_t, float, true>,
    DEQUANT_LINEAR_PERCHANNEL<float, int8_t, float, true>,
    DEQUANT_LINEAR_PERCHANNEL<half, int4x2_t, float, true>,
    DEQUANT_LINEAR_PERCHANNEL<bfloat16_t, int4x2_t, float, true>,
    DEQUANT_LINEAR_PERCHANNEL<float, int4x2_t, float, true>,
    DEQUANT_LINEAR_PERHEAD<half, int8_t, float, true>,
    DEQUANT_LINEAR_PERHEAD<bfloat16_t, int8_t, float, true>,
    DEQUANT_LINEAR_PERHEAD<float, int8_t, float, true>,
    DEQUANT_LINEAR_PERHEAD<half, int4x2_t, float, true>,
    DEQUANT_LINEAR_PERHEAD<bfloat16_t, int4x2_t, float, true>,
    DEQUANT_LINEAR_PERHEAD<float, int4x2_t, float, true>};

uint32_t getDequantLinearIdx(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 getBlockAndDimForLinear(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 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 (quant_bit == 4) {
      if (seq_block <= 1) {
        std::cerr << __func__ << "," << __LINE__
                  << " :head_num * head_size * sizeof_(float) should be less than "
                  << (nram_size >> 1) << " when quant_mode is 0." << std::endl;
      }
    } else {
      if (seq_block <= 0) {
        std::cerr << __func__ << "," << __LINE__
                  << " :head_num * head_size * sizeof_(float) should be less than " << nram_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 (quant_bit == 4) {
      if (seq_block <= 1) {
        std::cerr << __func__ << "," << __LINE__
                  << " :head_num * sizeof_(float) + head_num * head_size * (sizeof_(float) + "
                     "context_dtype_size) "
                  << "should be less than " << (nram_size >> 1) << " when quant_mode is 1."
                  << std::endl;
      }
    } else {
      if (seq_block <= 0) {
        std::cerr << __func__ << "," << __LINE__
                  << " :head_num * sizeof_(float) + head_num * head_size * (sizeof_(float) + "
                     "context_dtype_size) "
                  << "should be less than " << nram_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 = std::min(seq_block, max_context_len);
  if (seq_block > 16 && seq_block < max_context_len) {
    seq_block = PAD_DOWN(seq_block, 16);
  }
  if (quant_bit == 4) {
    seq_block = PAD_DOWN(seq_block, 2);
  }
  int seq_seg = DIV_UP(max_context_len, seq_block);
  // need an extra seg block to dealwith int4 value_cache [...,seq_len/2, head_size]
  if (quant_bit == 4) {
    seq_seg += 1;
  }
  uint32_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 (quant_bit == 4) {
      seq_block = PAD_DOWN(seq_block, 2);
    }
    seq_seg = DIV_UP(max_context_len, seq_block);
    // same as above to dealwise int4 value_cache with an extra seg block
    if (quant_bit == 4) {
      seq_seg += 1;
    }
  }

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

KernelStatus invokeDequantFromLinearCache(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 *cache_bs_id,
                                          const void *cache_seq_offsets,
                                          const int32_t max_context_len,
                                          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 cache_mem_len,
                                          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 (dtype == CNNL_DTYPE_BFLOAT16 && !isBf16Supported()) {
    std::cerr << "[invokeDequantFromPagedCache]: "
                 "MLU300 devices do not support bfloat16."
              << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  int32_t index;
  int32_t seq_block;
  cnrtDim3_t k_dim;
  cnrtFunctionType_t k_type;
  getBlockAndDimForLinear(seq_block, k_dim, k_type, max_context_len, head_num, batch_size,
                          head_size, quant_mode, quant_bit, dtype);
  index = getDequantLinearIdx(dtype, quant_mode, quant_bit, context_seq_offsets);
  auto dequant_linear_func = DequantFromLinearCacheFuncArr[index];
  dequant_linear_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 *)cache_bs_id,
      (const int32_t *)cache_seq_offsets, max_context_len, batch_size, head_num, key_group_num,
      value_group_num, cache_mem_len, 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
