/*************************************************************************
 * 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 <mlu.h>
#include <cassert>
#include <iostream>
#include <map>
#include "quant_to_linear_cache.mluh"

namespace tmo {

namespace kernels {

#define NRAM_BUFFER_SIZE (480 * 1024)
__nram__ int8_t nram_buffer[NRAM_BUFFER_SIZE];
__nram__ uint8_t post_table_nram[64];

#define sizeof_(T) (uint32_t)sizeof(T)
#define DIV_UP(x, y) ((x) / (y) + (int)((x) % (y) > 0))

#define __reshape_nhwc2nchw_smallhw(TYPE)                                                       \
  asm volatile(                                                                                 \
      "trans.tiling.nram.nram." TYPE                                                            \
      "[%[dst]], [%[src]], "                                                                    \
      "%[in0], %[in1], %[is1], %[in2], %[is2], %[in3], %[is3], %[in4], %[is4], %[in5], %[is5]," \
      "%[dn0], %[dn1], %[ds1], %[dn2], %[ds2], %[dn3], %[ds3], %[dn4], %[ds4], %[dn5], %[ds5]," \
      ".posttable.nram([%[post]]); \n\t" ::[dst] "r"(dst),                                      \
      [src] "r"(src), [post] "r"(post_table), [in0] "r"(in0), [in1] "r"(in1), [is1] "r"(is1),   \
      [in2] "i"(1), [is2] "i"(0), [in3] "r"(in3), [is3] "r"(is3), [in4] "r"(n),                 \
      [is4] "r"(batch_offset), [in5] "i"(1), [is5] "i"(0), [dn0] "r"(dn0), [dn1] "r"(dn1),      \
      [ds1] "r"(ds1), [dn2] "i"(1), [ds2] "i"(0), [dn3] "r"(dn3), [ds3] "r"(ds3), [dn4] "r"(n), \
      [ds4] "r"(batch_offset), [dn5] "i"(1), [ds5] "i"(0));

template <typename T>
__mlu_func__ void __reshape_nhwc2nchw_smallhw_init(uint8_t *post_table_nram, uint32_t hw) {
  uint32_t align_num = 64 / sizeof(T);
  int tmp = hw - 1;
  asm volatile("findlast1.gpr.b32 %[out], %[in];" : [out] "=r"(tmp) : [in] "r"(tmp));
  int align_hw = 1 << (tmp + 1);
  int repeat = align_num / align_hw;
  for (int i = 0; i < 64; i++) {
    int idx = i / sizeof_(T);
    int tmp_idx = (idx % hw) * repeat + idx / hw;
    int real_idx = tmp_idx * sizeof_(T) + i % sizeof_(T);
    int mask = idx < repeat * hw ? 0x80 : 0x0;
    post_table_nram[i] = (uint8_t)real_idx + mask;
  }
}

template <typename T>
__mlu_func__ void trans_nhwc2nchw(T *dst,
                                  const T *src,
                                  uint8_t *post_table,
                                  const uint32_t n,
                                  const uint32_t hw,
                                  const uint32_t c) {
  uint32_t align_num = 64 / sizeof_(T);
  int tmp = hw - 1;
  asm volatile("findlast1.gpr.b32 %[out], %[in];" : [out] "=r"(tmp) : [in] "r"(tmp));
  int align_hw = 1 << (tmp + 1);
  int repeat = align_num / align_hw;
  int in0 = 64;
  int in1 = hw;
  int is1 = c * sizeof(T);
  int in3 = c / align_num;
  int is3 = in0;
  int batch_offset = hw * c * sizeof(T);
  int dn0 = hw * repeat * sizeof_(T);
  int dn1 = align_hw;
  int ds1 = dn0;
  int dn3 = in3;
  int ds3 = dn0 * dn1;
  align_hw = in3 > 0 ? align_hw : 0;
  if (align_hw == 2) {
    __reshape_nhwc2nchw_smallhw("b256");
  } else if (align_hw == 4) {
    __reshape_nhwc2nchw_smallhw("b128");
  } else if (align_hw == 8) {
    __reshape_nhwc2nchw_smallhw("b64");
  } else if (align_hw == 16) {
    __reshape_nhwc2nchw_smallhw("b32");
  } else if (align_hw == 32) {
    __reshape_nhwc2nchw_smallhw("b16");
  }

  constexpr uint32_t bw = 8 * sizeof_(T);
  int in3_rem = c % align_num;
  int tail_in0 = in3_rem * sizeof_(T);
  int tail_dn0 = hw * sizeof_(T);
  if (in3_rem > 0) {
    asm volatile(
        "trans.tiling.nram.nram.b%[bw] [%[dst]], [%[src]], \
      %[in0], %[in1], %[is1], %[in2], %[is2], %[in3], %[is3], %[in4], %[is4], %[in5], %[is5], \
      %[dn0], %[dn1], %[ds1], %[dn2], %[ds2], %[dn3], %[ds3], %[dn4], %[ds4], %[dn5], %[ds5]; \n\t" ::
            [bw] "i"(bw),
        [dst] "r"(dst + dn3 * ds3 / sizeof_(T)), [src] "r"(src + is3 * in3 / sizeof_(T)),
        [in0] "r"(tail_in0), [in1] "r"(hw), [is1] "r"(c * sizeof_(T)), [in2] "i"(1), [is2] "i"(0),
        [in3] "i"(1), [is3] "i"(0), [in4] "r"(n), [is4] "r"(batch_offset), [in5] "i"(1),
        [is5] "i"(0), [dn0] "r"(tail_dn0), [dn1] "r"(in3_rem), [ds1] "r"(tail_dn0), [dn2] "i"(1),
        [ds2] "i"(0), [dn3] "i"(1), [ds3] "i"(0), [dn4] "r"(n), [ds4] "r"(batch_offset),
        [dn5] "i"(1), [ds5] "i"(0));
  }
}

template <typename T>
__mlu_func__ void quantify(T *src_input,
                           float *nram_input,
                           float *transpose,
                           float *scale,
                           float *scale_recip,
                           int high_dim,
                           int low_dim,
                           int quant_bit) {
  float recip_data = 1.0f / ((1 << (quant_bit - 1)) - 1);
  if (std::is_same<half, T>::value) {
    __bang_transpose((half *)transpose, (half *)src_input, high_dim, low_dim);
    __bang_abs((half *)src_input, (half *)transpose, high_dim * low_dim);
    __bang_maxpool((half *)scale, (half *)src_input, high_dim, low_dim, 1, low_dim, 1, 1, 1);
    __bang_half2float(scale_recip, (half *)scale, high_dim);
  } else if (std::is_same<bfloat16_t, T>::value) {
#if __BANG_ARCH__ > 500
    __bang_transpose((bfloat16_t *)transpose, (bfloat16_t *)src_input, high_dim, low_dim);
    __bang_abs((bfloat16_t *)src_input, (bfloat16_t *)transpose, high_dim * low_dim);
    __bang_maxpool((bfloat16_t *)scale, (bfloat16_t *)src_input, high_dim, low_dim, 1, low_dim, 1,
                   1, 1);
    __bang_bfloat162float(scale_recip, (bfloat16_t *)scale, high_dim);
#endif
  } else {
    __bang_transpose(transpose, (float *)src_input, high_dim, low_dim);
    __bang_abs((float *)src_input, transpose, high_dim * low_dim);
    __bang_maxpool(scale_recip, (float *)src_input, high_dim, low_dim, 1, low_dim, 1, 1, 1);
  }
  __bang_mul_scalar(scale, scale_recip, recip_data, high_dim);
  __bang_recip(scale_recip, scale, high_dim);
  if (std::is_same<half, T>::value) {
    __bang_half2float(nram_input, (half *)transpose, high_dim * low_dim);
    __bang_cycle_mul(transpose, nram_input, scale_recip, high_dim * low_dim, high_dim);
  } else if (std::is_same<bfloat16_t, T>::value) {
#if __BANG_ARCH__ > 500
    __bang_bfloat162float(nram_input, (bfloat16_t *)transpose, high_dim * low_dim);
    __bang_cycle_mul(transpose, nram_input, scale_recip, high_dim * low_dim, high_dim);
#endif
  } else {
    __bang_cycle_mul(transpose, transpose, scale_recip, high_dim * low_dim, high_dim);
  }
}

__mlu_func__ void castKeyToIntx(int8_t *dst, float *src, int high_dim, int low_dim, int quant_bit) {
  if (quant_bit == 8) {
    __bang_float2int8_rn((int8_t *)src, (float *)dst, high_dim * low_dim, 0);
    __bang_transpose((int8_t *)dst, (int8_t *)src, low_dim, high_dim);
  } else if (quant_bit == 4) {
    __bang_transpose(src, (float *)dst, low_dim, high_dim);
    __bang_float2int4_rn((int4x2_t *)dst, src, high_dim * low_dim, 0);
  }
}

__mlu_func__ void castValueToIntx(int8_t *value_cache_begin,
                                  int8_t *value_cache_end,
                                  uint8_t *post_table_nram,
                                  float *dst,
                                  float *src,
                                  size_t cache_head_stride,
                                  int seq,
                                  int head_num,
                                  int head_size,
                                  int group_num,
                                  int quant_bit,
                                  bool need_pad_front,
                                  bool need_pad_back) {
  if (quant_bit == 8) {
    __bang_float2int8_rn((int8_t *)src, dst, seq * head_num * head_size, 0);
    __bang_transpose((int8_t *)dst, (int8_t *)src, head_size / group_num,
                     seq * head_num * group_num);
  } else {
    __bang_transpose(src, dst, head_size / group_num, seq * head_num * group_num);
    __reshape_nhwc2nchw_smallhw_init<int8_t>(post_table_nram, 2);
    if (!(need_pad_front || need_pad_back)) {
      __bang_float2int8_rn((int8_t *)src, src, seq * head_num * head_size,
                           0);  // [head_num, seq, head_size]
      trans_nhwc2nchw((int8_t *)dst, (int8_t *)src, post_table_nram, seq / 2, 2,
                      head_num * head_size);
      __bang_int82int4_rn((int4x2_t *)dst, (int8_t *)dst, seq * head_num * head_size, 0, 0);
    } else {
      int origin_seq = seq;
      if (need_pad_front) {
        __memcpy((int8_t *)dst, value_cache_begin, head_size, GDRAM2NRAM, head_size,
                 cache_head_stride, head_num - 1);
        __bang_band_scalar((int8_t *)dst, (int8_t *)dst, 0x0F, head_num * head_size);
        seq += 1;
      }
      if (need_pad_back) {
        __memcpy((int8_t *)dst + seq * head_num * head_size, value_cache_end, head_size, GDRAM2NRAM,
                 head_size, cache_head_stride, head_num - 1);
        __bang_srl((int8_t *)dst + seq * head_num * head_size,
                   (int8_t *)dst + seq * head_num * head_size, 4, head_num * head_size);
        seq += 1;
      }
      __bang_float2int8_rn((int8_t *)dst + need_pad_front * head_num * head_size, src,
                           origin_seq * head_num * head_size, 0);  // [new_seq, head_num, head_size]
      trans_nhwc2nchw((int8_t *)src, (int8_t *)dst, post_table_nram, seq / 2, 2,
                      head_num * head_size);  // [seq / 2, 2, head_num, head_size]
      __bang_int82int4_rn((int4x2_t *)dst, (int8_t *)src, seq * head_num * head_size, 0, 0);
    }
  }
}

// [head_num, batch, seq_seg]
template <typename T>
__mlu_global__ void MLUQuantToLinearCacheKernel(int8_t *key_cache,
                                                int8_t *value_cache,
                                                float *key_cache_scale,
                                                float *value_cache_scale,
                                                int *cache_bs_offsets,
                                                int *cache_seq_offsets,
                                                T *key,
                                                T *value,
                                                int *context_seq_offsets,
                                                int *context_lens,
                                                int batch,
                                                int head_num,
                                                int head_size,
                                                int max_context_len,
                                                int cache_mem_len,
                                                size_t context_bs_stride,
                                                size_t context_head_stride,
                                                size_t context_seq_stride,
                                                size_t cache_bs_stride,
                                                size_t cache_head_stride,
                                                size_t key_cache_seq_stride,
                                                size_t value_cache_seq_stride,
                                                size_t cache_scale_bs_stride,
                                                size_t cache_scale_head_stride,
                                                bool packed,
                                                int seq_block,
                                                int quant_bit,
                                                int group_num) {
  float *nram_input = (float *)nram_buffer;
  T *src_input = (T *)nram_input;
  if (sizeof_(T) == sizeof_(half)) {
    src_input = (T *)(nram_buffer + seq_block * head_num * head_size * sizeof_(T));
  }

  float *nram_trans = nram_input + seq_block * head_num * head_size;
  float *nram_scale = nram_trans + seq_block * head_num * head_size;
  float *nram_scale_recip = nram_scale + seq_block * head_num * group_num;
  constexpr int dtype_size = sizeof_(T);
  int head_size_store = head_size * quant_bit / 8;

  for (int bs_idx = taskIdY; bs_idx < batch; bs_idx += taskDimY) {
    int seq_offset =
        (packed || context_seq_offsets == nullptr) ? 0 : __load_gdram(context_seq_offsets + bs_idx);
    int cache_seq_offset =
        cache_seq_offsets == nullptr ? 0 : __load_gdram(cache_seq_offsets + bs_idx);
    int context_len = __load_gdram(context_lens + bs_idx);
    int seq_len = packed ? (__load_gdram(context_lens + bs_idx + 1) - context_len) : context_len;
    int key_seq_begin = taskIdZ * seq_block;
    int first_value_seq = (quant_bit == 4 && cache_seq_offset % 2 == 1)
                              ? std::min(seq_len, seq_block - 1)
                              : std::min(seq_len, seq_block);
    int value_seq_begin = taskIdZ == 0 ? 0 : first_value_seq + (taskIdZ - 1) * seq_block;

    int key_seq = std::min(seq_len - key_seq_begin, seq_block);
    int value_seq = taskIdZ == 0 ? first_value_seq : std::min(seq_len - value_seq_begin, seq_block);
    size_t key_context_offset = 0;
    size_t value_context_offset = 0;
    if (packed) {
      key_context_offset += (context_len + key_seq_begin) * context_seq_stride * dtype_size;
      value_context_offset += (context_len + value_seq_begin) * context_seq_stride * dtype_size;
    } else {
      key_context_offset +=
          (bs_idx * context_bs_stride + (key_seq_begin + seq_offset) * context_seq_stride) *
          dtype_size;
      value_context_offset +=
          (bs_idx * context_bs_stride + (value_seq_begin + seq_offset) * context_seq_stride) *
          dtype_size;
    }
    int key_cache_seq_offset = cache_seq_offset;
    int value_cache_seq_offset = cache_seq_offset;
    int cache_bs_offset = cache_bs_offsets == nullptr ? bs_idx : cache_bs_offsets[bs_idx];
    if (key_cache_seq_offset < 0 || cache_bs_offset < 0) {
      continue;
    }
    key_cache_seq_offset += key_seq_begin;
    value_cache_seq_offset += value_seq_begin;

    size_t key_cache_offset =
        (cache_bs_offset * cache_bs_stride + key_cache_seq_offset * key_cache_seq_stride);
    size_t key_cache_scale_offset =
        cache_bs_offset * cache_scale_bs_stride + key_cache_seq_offset * group_num;

    if (key != nullptr && key_cache != nullptr && key_cache_scale != nullptr &&
        key_seq_begin < seq_len) {
      int8_t *key_cache_begin = key_cache + key_cache_offset;
      float *key_cache_scale_begin = key_cache_scale + key_cache_scale_offset;
      char *key_begin = (char *)key + key_context_offset;

      // nram_input: (head_num, seq, head_size)
      __memcpy(src_input, key_begin, head_size * dtype_size, GDRAM2NRAM,
               key_seq * head_size * dtype_size, head_num - 1, head_size * dtype_size, key_seq - 1,
               context_head_stride * dtype_size, head_num - 1, context_seq_stride * dtype_size,
               key_seq - 1);
      // [head_num, seq, head_size]
      quantify(src_input, nram_input, nram_trans, nram_scale, nram_scale_recip,
               key_seq * head_num * group_num, head_size / group_num, quant_bit);
      castKeyToIntx((int8_t *)nram_trans, nram_input, key_seq * head_num * group_num,
                    head_size / group_num, quant_bit);

      // after quantify: (head_num, seq, head_size * quant_bit / 8)
      __memcpy(key_cache_begin, nram_trans, head_size_store, NRAM2GDRAM, key_cache_seq_stride,
               key_seq - 1, cache_head_stride, head_num - 1, head_size_store, key_seq - 1,
               key_seq * head_size_store, head_num - 1);

      // nram_scale: (head_num, seq, group_num)
      __memcpy(key_cache_scale_begin, nram_scale, key_seq * group_num * sizeof_(float), NRAM2GDRAM,
               cache_scale_head_stride * sizeof_(float), key_seq * group_num * sizeof_(float),
               head_num - 1);
    }

    if (value != nullptr && value_cache != nullptr && value_cache_scale != nullptr &&
        value_seq_begin < seq_len) {
      size_t value_cache_offset =
          cache_bs_offset * cache_bs_stride + value_cache_seq_offset * value_cache_seq_stride;
      if (quant_bit == 4) {
        value_cache_offset = cache_bs_offset * cache_bs_stride +
                             (value_cache_seq_offset / 2) * value_cache_seq_stride;
      }
      size_t value_cache_scale_offset =
          cache_bs_offset * cache_scale_bs_stride + value_cache_seq_offset * group_num;

      int8_t *value_cache_begin = value_cache + value_cache_offset;
      int8_t *value_cache_end =
          value_cache + cache_bs_offset * cache_bs_stride +
          (DIV_UP(cache_seq_offset + seq_len, 2) - 1) * value_cache_seq_stride;

      float *value_cache_scale_begin = value_cache_scale + value_cache_scale_offset;
      char *value_begin = (char *)value + value_context_offset;
      bool need_pad_front = (quant_bit == 4) && (taskIdZ == 0) && (cache_seq_offset % 2 == 1);
      bool need_pad_back = (quant_bit == 4) && (value_seq_begin + value_seq >= seq_len) &&
                           ((cache_seq_offset + seq_len) % 2 == 1);
      // quant_bit == 8 : nram_input: (head_num, seq, head_size)
      if (quant_bit == 8) {
        __memcpy(src_input, value_begin, head_size * dtype_size, GDRAM2NRAM,
                 value_seq * head_size * dtype_size, head_num - 1, head_size * dtype_size,
                 value_seq - 1, context_head_stride * dtype_size, head_num - 1,
                 context_seq_stride * dtype_size, value_seq - 1);
      } else if (quant_bit == 4) {
        // quant_bit == 4 : nram_input: (seq, head_num. head_size)
        __memcpy(src_input, value_begin, head_size * dtype_size, GDRAM2NRAM, head_size * dtype_size,
                 head_num - 1, head_num * head_size * dtype_size, value_seq - 1,
                 context_head_stride * dtype_size, head_num - 1, context_seq_stride * dtype_size,
                 value_seq - 1);
      }
      quantify(src_input, nram_input, nram_trans, nram_scale, nram_scale_recip,
               value_seq * head_num * group_num, head_size / group_num, quant_bit);
      castValueToIntx(value_cache_begin, value_cache_end, post_table_nram, nram_trans, nram_input,
                      cache_head_stride, value_seq, head_num, head_size, group_num, quant_bit,
                      need_pad_front, need_pad_back);

      // [head_num, seq, head_size]
      if (quant_bit == 8) {
        __memcpy(value_cache_scale_begin, nram_scale, value_seq * group_num * sizeof_(float),
                 NRAM2GDRAM, cache_scale_head_stride * sizeof_(float),
                 value_seq * group_num * sizeof_(float), head_num - 1);
        __memcpy(value_cache_begin, nram_trans, head_size, NRAM2GDRAM, value_cache_seq_stride,
                 value_seq - 1, cache_head_stride, head_num - 1, head_size, value_seq - 1,
                 value_seq * head_size, head_num - 1);
      } else if (quant_bit == 4) {
        int new_seq = value_seq + ((int)need_pad_front + (int)need_pad_back);
        __sync();
        __memcpy_async(nram_scale_recip, nram_scale, group_num * sizeof_(float), NRAM2NRAM,
                       value_seq * group_num * sizeof_(float), head_num - 1,
                       group_num * sizeof_(float), value_seq - 1, group_num * sizeof_(float),
                       head_num - 1, head_num * group_num * sizeof_(float), value_seq - 1);
        __memcpy_async(value_cache_begin, nram_trans, head_size, NRAM2GDRAM, cache_head_stride,
                       head_num - 1, value_cache_seq_stride, new_seq / 2 - 1, head_size,
                       head_num - 1, head_num * head_size, new_seq / 2 - 1);
        __sync();
        __memcpy(value_cache_scale_begin, nram_scale_recip, value_seq * group_num * sizeof_(float),
                 NRAM2GDRAM, cache_scale_head_stride * sizeof_(float),
                 value_seq * group_num * sizeof_(float), head_num - 1);
      }
    }  // end if
  }  // end for
}
}  // namespace kernels

void invokeQuantToLinearCache(cnrtQueue_t queue,
                              void *key_cache,
                              void *value_cache,
                              void *key_cache_scale,
                              void *value_cache_scale,
                              const void *cache_bs_offsets,
                              const void *cache_seq_offsets,
                              void *key,
                              void *value,
                              const void *context_seq_offsets,
                              const void *context_lens,
                              const cnnlDataType_t dtype,
                              const int batch,
                              const int head_num,
                              const int head_size,
                              const int max_context_len,
                              const int cache_mem_len,
                              const size_t context_bs_stride,
                              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 cache_scale_bs_stride,
                              const size_t cache_scale_head_stride,
                              const bool packed,
                              const int quant_bit,
                              const int group_size) {
  constexpr int nram_size = 480 * 1024;
  int group_num = head_size / group_size;
  int hidden_bytes = head_num * (head_size + group_num) * sizeof(float) * 2;
  int seq_block = nram_size / hidden_bytes;
  if (seq_block <= 1) {
    std::cerr << __func__ << "," << __LINE__
              << " :head_num * (head_size + group_num) * sizeof(float) should be less than 120KB."
              << std::endl;
  }
  if (seq_block > 16) {
    seq_block = seq_block / 16 * 16;
  } else {
    seq_block = seq_block / 2 * 2;
  }
  int seq_seg = max_context_len / seq_block + 1;
  int cluster_num, core_dim;
  CNdev dev;
  cnCtxGetDevice(&dev);
  CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
  CNRT_CHECK(cnrtDeviceGetAttribute(&core_dim, cnrtAttrMcorePerCluster, dev));
  int core_num = core_dim * cluster_num;
  uint32_t task_y_dim = std::min(batch, core_num);
  cnrtDim3_t dim{1, task_y_dim, (uint32_t)seq_seg};

  if (dtype == CNNL_DTYPE_HALF) {
    kernels::MLUQuantToLinearCacheKernel<<<dim, cnrtFuncTypeBlock, queue>>>(
        (int8_t *)key_cache, (int8_t *)value_cache, (float *)key_cache_scale,
        (float *)value_cache_scale, (int *)cache_bs_offsets, (int *)cache_seq_offsets, (half *)key,
        (half *)value, (int *)context_seq_offsets, (int *)context_lens, batch, head_num, head_size,
        max_context_len, cache_mem_len, context_bs_stride, context_head_stride, context_seq_stride,
        cache_bs_stride, cache_head_stride, key_cache_seq_stride, value_cache_seq_stride,
        cache_scale_bs_stride, cache_scale_head_stride, packed, seq_block, quant_bit, group_num);
  } else if (dtype == CNNL_DTYPE_BFLOAT16) {
    kernels::MLUQuantToLinearCacheKernel<<<dim, cnrtFuncTypeBlock, queue>>>(
        (int8_t *)key_cache, (int8_t *)value_cache, (float *)key_cache_scale,
        (float *)value_cache_scale, (int *)cache_bs_offsets, (int *)cache_seq_offsets,
        (bfloat16_t *)key, (bfloat16_t *)value, (int *)context_seq_offsets, (int *)context_lens,
        batch, head_num, head_size, max_context_len, cache_mem_len, context_bs_stride,
        context_head_stride, context_seq_stride, cache_bs_stride, cache_head_stride,
        key_cache_seq_stride, value_cache_seq_stride, cache_scale_bs_stride,
        cache_scale_head_stride, packed, seq_block, quant_bit, group_num);
  } else {
    kernels::MLUQuantToLinearCacheKernel<<<dim, cnrtFuncTypeBlock, queue>>>(
        (int8_t *)key_cache, (int8_t *)value_cache, (float *)key_cache_scale,
        (float *)value_cache_scale, (int *)cache_bs_offsets, (int *)cache_seq_offsets, (float *)key,
        (float *)value, (int *)context_seq_offsets, (int *)context_lens, batch, head_num, head_size,
        max_context_len, cache_mem_len, context_bs_stride, context_head_stride, context_seq_stride,
        cache_bs_stride, cache_head_stride, key_cache_seq_stride, value_cache_seq_stride,
        cache_scale_bs_stride, cache_scale_head_stride, packed, seq_block, quant_bit, group_num);
  }
}
}  // namespace tmo
