// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once

#include "helper.h"
#include "mem_util.cuh"
#include "mma_tensor_op.cuh"
#include "utils.cuh"

template <int VecSize = 4, int HeadDim = 128>
__global__ void append_clear_cache_int8_block(
    uint8_t* __restrict__ key_cache,    // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    uint8_t* __restrict__ value_cache,  // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    const int* __restrict__ seq_lens,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens_encoder,  // [bsz]
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int block_size,
    const int gqa_group_size) {
  static_assert(HeadDim == 128, "just support HeadDim be 128 now!");
  static_assert(VecSize == 4, "just support VecSize be 4 now, 32 * 4!");
  constexpr int NUM_WARPS = 4;
  const int tid = threadIdx.x;
  const int wid = tid / 32;
  const int lane_id = tid % 32;
  const int token_id = blockIdx.x;
  
  const int bid = batch_id_per_token[token_id];

  const int start_token_idx = cu_seqlens_q[bid];
  const int head_idx = blockIdx.y * NUM_WARPS + wid;

  if (seq_lens_encoder[bid] > 0) return;
  const int write_seq_id = seq_lens[bid] + token_id - start_token_idx;
  if (write_seq_id == 0) return;
  const int* block_table_now = block_tables + bid * max_blocks_per_seq;
  const int block_idx = __ldg(&block_table_now[write_seq_id / block_size]);
  const int block_offset = write_seq_id % block_size;

  if (head_idx >= num_heads && head_idx < num_heads + 2 * gqa_group_size) {
    // k
    constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t);  // 16
    using LoadPadKVT = AlignedVector<uint8_t, KV_VEC_SIZE>;
    const uint32_t kv_head_idx = (head_idx - num_heads) % gqa_group_size;
    if (block_offset == 0) {
      // pad zero for this kv_head_idx for this block
      LoadPadKVT pad_cache_vec;
      *(reinterpret_cast<uint4*>(pad_cache_vec.val)) = make_uint4(0, 0, 0, 0);
      if (head_idx < num_heads + gqa_group_size) {
        constexpr int num_vecs_per_head_dim = HeadDim / KV_VEC_SIZE;
        constexpr int num_token_each_time = 32 / num_vecs_per_head_dim;
        const uint32_t tgt_idx =
            (block_idx * gqa_group_size + kv_head_idx) * block_size * HeadDim +
            lane_id % num_vecs_per_head_dim * KV_VEC_SIZE;
        for (int block_i = lane_id / num_vecs_per_head_dim;
             block_i < block_size;
             block_i += num_token_each_time) {
          Store<uint8_t, KV_VEC_SIZE>(pad_cache_vec,
                                      &key_cache[tgt_idx + block_i * HeadDim]);
        }
      } else {
        const int num_vecs_per_head_dim = block_size / KV_VEC_SIZE;
        const int num_token_each_time = 32 / num_vecs_per_head_dim;
        const uint32_t tgt_idx =
            (block_idx * gqa_group_size + kv_head_idx) * HeadDim * block_size +
            lane_id % num_vecs_per_head_dim * KV_VEC_SIZE;
        for (int block_i = lane_id / num_vecs_per_head_dim; block_i < HeadDim;
             block_i += num_token_each_time) {
          Store<uint8_t, KV_VEC_SIZE>(
              pad_cache_vec, &value_cache[tgt_idx + block_i * block_size]);
        }
      }
    }
  }
}


template <int VecSize = 4, int HeadDim = 128>
__global__ void append_clear_cache_int4_block(
    uint8_t* __restrict__ key_cache,    // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    uint8_t* __restrict__ value_cache,  // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    const int* __restrict__ seq_lens,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens_encoder,  // [bsz]
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int block_size,
    const int gqa_group_size) {
  static_assert(HeadDim == 128, "just support HeadDim be 128 now!");
  static_assert(VecSize == 4, "just support VecSize be 4 now, 32 * 4!");
  constexpr int NUM_WARPS = 4;
  const int tid = threadIdx.x;
  const int wid = tid / 32;
  const int lane_id = tid % 32;
  const int token_id = blockIdx.x;
  
  const int bid = batch_id_per_token[token_id];

  const int start_token_idx = cu_seqlens_q[bid];
  const int head_idx = blockIdx.y * NUM_WARPS + wid;

  if (seq_lens_encoder[bid] > 0) return;
  const int write_seq_id = seq_lens[bid] + token_id - start_token_idx;
  if (write_seq_id == 0) return;
  const int* block_table_now = block_tables + bid * max_blocks_per_seq;
  const int block_idx = __ldg(&block_table_now[write_seq_id / block_size]);
  const int block_offset = write_seq_id % block_size;

  constexpr int half_head_size = HeadDim / 2;
  const int half_block_size = block_size / 2;

  if (head_idx >= num_heads && head_idx < num_heads + 2 * gqa_group_size) {
    // k
    constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t);  // 16
    using LoadPadKVT = AlignedVector<uint8_t, KV_VEC_SIZE>;
    const uint32_t kv_head_idx = (head_idx - num_heads) % gqa_group_size;
    if (block_offset == 0) {
      // pad zero for this kv_head_idx for this block
      LoadPadKVT pad_cache_vec;
      *(reinterpret_cast<uint4*>(pad_cache_vec.val)) = make_uint4(0, 0, 0, 0);
      if (head_idx < num_heads + gqa_group_size) {
        constexpr int num_vecs_per_head_dim = half_head_size / KV_VEC_SIZE;
        constexpr int num_token_each_time = 32 / num_vecs_per_head_dim;
        const uint32_t tgt_idx = (block_idx * gqa_group_size + kv_head_idx) *
                                     block_size * half_head_size +
                                 lane_id % num_vecs_per_head_dim * KV_VEC_SIZE;
        for (int block_i = lane_id / num_vecs_per_head_dim;
             block_i < block_size;
             block_i += num_token_each_time) {
          Store<uint8_t, KV_VEC_SIZE>(
              pad_cache_vec, &key_cache[tgt_idx + block_i * half_head_size]);
        }
      } else {
        const int num_vecs_per_head_dim = half_block_size / KV_VEC_SIZE;
        const int num_token_each_time = 32 / num_vecs_per_head_dim;
        const uint32_t tgt_idx = (block_idx * gqa_group_size + kv_head_idx) *
                                     HeadDim * half_block_size +
                                 lane_id % num_vecs_per_head_dim * KV_VEC_SIZE;
        for (int block_i = lane_id / num_vecs_per_head_dim; block_i < HeadDim;
             block_i += num_token_each_time) {
          Store<uint8_t, KV_VEC_SIZE>(
              pad_cache_vec, &value_cache[tgt_idx + block_i * half_block_size]);
        }
      }
    }
  }
}

template <typename T, int VecSize = 1, typename InT = T>
__global__ void append_speculate_cache_rope_kernel(
    const InT* __restrict__ qkv,  // [token_num, num_heads + 2 * gqa_group_size,
                                  // head_size]
    T* __restrict__ key_cache,    // [num_blocks, gqa_group_size, block_size,
                                  // head_size // 2]
    T* __restrict__ value_cache,  // [num_blocks, gqa_group_size, block_size,
                                  // head_size // 2]
    T* __restrict__ q_out,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens_decoder,  // [bsz]
    const float* __restrict__ cos_emb,
    const float* __restrict__ sin_emb,
    const float*
        qkv_out_scales,   // [(num_heads + 2 * gqa_group_size) * head_size]
    const T* qkv_biases,  // [num_head + 2 * gqa_group_size, dim_head]
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int output_inner_dim,
    const int head_size,
    const int block_size,
    const int elem_cnt,
    const int gqa_group_size) {
  using LoadT = AlignedVector<T, VecSize>;
  using LoadFloat = AlignedVector<float, VecSize>;
  using LoadInT = AlignedVector<InT, VecSize>;
  constexpr int HalfVecSize = VecSize / 2;
  using LoadEmbT = AlignedVector<float, HalfVecSize>;
  LoadInT src_vec;
  LoadFloat scale_vec;
  LoadT bias_vec;
  LoadEmbT cos_emb_vec;
  LoadEmbT sin_emb_vec;

  int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
  const int64_t hidden_size = (num_heads + 2 * gqa_group_size) * head_size;
  // const int64_t offset = 2 * hidden_size;
  const int half_head_size = head_size / 2;
  for (int32_t linear_index = global_thread_idx * VecSize,
               step = gridDim.x * blockDim.x * VecSize;
       linear_index < elem_cnt;
       linear_index += step) {
    const int token_id = linear_index / hidden_size;
    const int ori_bi = batch_id_per_token[token_id];
    if (seq_lens_decoder[ori_bi] == 0) continue;
    const int bias = linear_index % hidden_size;
    const int hi = bias / head_size;  // q + k + v
    const int h_bias = bias % head_size;
    const int start_token_idx = cu_seqlens_q[ori_bi];
    const int write_seq_id =
        seq_lens_decoder[ori_bi] + token_id - start_token_idx;
    if (write_seq_id == 0) continue;

    const int* block_table_now = block_tables + ori_bi * max_blocks_per_seq;
    const int block_idx = block_table_now[write_seq_id / block_size];
    if (block_idx < 0) {
      printf(
          "Fatal Error!!!, block idx %d when write_seq_id is %d\n some key var "
          "%d %d %d %d\n",
          block_idx,
          write_seq_id,
          ori_bi,
          seq_lens_decoder[ori_bi],
          token_id,
          cu_seqlens_q[ori_bi]);
    }
    const int block_offset = write_seq_id % block_size;

    const int write_q_idx =
        token_id * output_inner_dim * head_size + hi * head_size + h_bias;

    const int bias_idx = hi * head_size + h_bias;
    Load<InT, VecSize>(&qkv[linear_index], &src_vec);
    if (qkv_biases) {
      Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
    }
    if (qkv_out_scales) {
      Load<float, VecSize>(&qkv_out_scales[bias_idx], &scale_vec);
    }
    if (hi < num_heads + gqa_group_size) {
      // q k rope
      const int64_t emb_idx = write_seq_id * half_head_size + h_bias / 2;
      Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
      Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
    }
#pragma unroll
    for (int i = 0; i < HalfVecSize; i++) {
      // add_bias + rope
      float input_left = static_cast<float>(src_vec[2 * i]);
      float input_right = static_cast<float>(src_vec[2 * i + 1]);
      if (qkv_out_scales) {
        input_left *= scale_vec[2 * i];
        input_right *= scale_vec[2 * i + 1];
      }
      if (qkv_biases) {
        input_left = input_left + static_cast<float>(bias_vec[2 * i]);
        input_right = input_right + static_cast<float>(bias_vec[2 * i + 1]);
      }
      if (hi < num_heads + gqa_group_size) {
        const float cos_tmp = cos_emb_vec[i];
        const float sin_tmp = sin_emb_vec[i];
        bias_vec[2 * i] =
            static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
        bias_vec[2 * i + 1] =
            static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
      } else {
        bias_vec[2 * i] = static_cast<T>(input_left);
        bias_vec[2 * i + 1] = static_cast<T>(input_right);
      }
    }
    if (hi < num_heads) {
      // write q
      Store<T, VecSize>(bias_vec, &q_out[write_q_idx]);
    } else {
      //  write k/v
      const int kv_head_idx = (hi - num_heads) % gqa_group_size;
      const int tgt_idx = (block_idx * gqa_group_size * block_size * head_size +
                           kv_head_idx * block_size * head_size +
                           block_offset * head_size + h_bias);
      // write

      if (hi < num_heads + gqa_group_size) {
        Store<T, VecSize>(bias_vec, &key_cache[tgt_idx]);
      } else {
        Store<T, VecSize>(bias_vec, &value_cache[tgt_idx]);
      }
    }
  }
}

template <typename T, int VecSize = 1, typename InT = T>
__global__ void append_speculate_cache_neox_rope_kernel(
    const InT* __restrict__ qkv,  // [token_num, num_heads + 2 * gqa_group_size,
                                  // head_size]
    T* __restrict__ key_cache,    // [num_blocks, gqa_group_size, block_size,
                                  // head_size // 2]
    T* __restrict__ value_cache,  // [num_blocks, gqa_group_size, block_size,
                                  // head_size // 2]
    T* __restrict__ qkv_out,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens_decoder,  // [bsz]
    const float* __restrict__ cos_emb,
    const float* __restrict__ sin_emb,
    const float*
        qkv_out_scales,   // [(num_heads + 2 * gqa_group_size) * head_size]
    const T* qkv_biases,  // [num_head + 2 * gqa_group_size, dim_head]
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int output_inner_dim,
    const int head_size,
    const int block_size,
    const int elem_cnt,
    const int gqa_group_size) {
  using LoadT = AlignedVector<T, VecSize>;
  using LoadFloat = AlignedVector<float, VecSize>;
  using LoadInT = AlignedVector<InT, VecSize>;
  constexpr int HalfVecSize = VecSize / 2;
  using LoadEmbT = AlignedVector<float, VecSize>;
  LoadInT left_vec, right_vec;
  LoadT left_bias_vec, right_bias_vec;
  LoadFloat left_out_scale_vec, right_out_scale_vec;
  LoadEmbT cos_emb_vec;
  LoadEmbT sin_emb_vec;

  int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
  const int64_t hidden_size = (num_heads + 2 * gqa_group_size) * head_size;
  const int half_head_size = head_size / 2;
  const int64_t half_hidden_size = hidden_size / 2;
  for (int32_t linear_index = global_thread_idx * VecSize,
               step = gridDim.x * blockDim.x * VecSize;
       linear_index < elem_cnt;
       linear_index += step) {
    const int token_id = linear_index / half_hidden_size;
    const int ori_bi = batch_id_per_token[token_id];
    if (seq_lens_decoder[ori_bi] == 0) continue;
    const int bias = linear_index % half_hidden_size;
    const int hi = bias / half_head_size;  // q + k + v
    const int h_bias = bias % half_head_size;
    const int start_token_idx = cu_seqlens_q[ori_bi];
    const int write_seq_id =
        seq_lens_decoder[ori_bi] + token_id - start_token_idx;
    if (write_seq_id == 0) continue;

    const int* block_table_now = block_tables + ori_bi * max_blocks_per_seq;
    const int block_idx = block_table_now[write_seq_id / block_size];
    if (block_idx < 0) {
      printf(
          "Fatal Error!!!, block idx %d when write_seq_id is %d\n some key var "
          "%d %d %d %d\n",
          block_idx,
          write_seq_id,
          ori_bi,
          seq_lens_decoder[ori_bi],
          token_id,
          cu_seqlens_q[ori_bi]);
    }
    const int block_offset = write_seq_id % block_size;

    const int bias_idx_left = hi * head_size + h_bias;
    const int bias_idx_right = bias_idx_left + half_head_size;
    const int ori_idx_left = token_id * hidden_size + hi * head_size + h_bias;
    const int ori_idx_right = ori_idx_left + half_head_size;
    Load<InT, VecSize>(&qkv[ori_idx_left], &left_vec);
    Load<InT, VecSize>(&qkv[ori_idx_right], &right_vec);
    if (qkv_biases) {
      Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
      Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
    }
    if (qkv_out_scales) {
      Load<float, VecSize>(&qkv_out_scales[bias_idx_left], &left_out_scale_vec);
      Load<float, VecSize>(&qkv_out_scales[bias_idx_right],
                           &right_out_scale_vec);
    }
    if (hi < num_heads + gqa_group_size) {
      // q k rope
      const int64_t emb_idx = write_seq_id * head_size + h_bias;
      Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
      Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
    }
#pragma unroll
    for (int i = 0; i < VecSize; i++) {
      // add_bias + rope
      float input_left = static_cast<float>(left_vec[i]);
      float input_right = static_cast<float>(right_vec[i]);
      if (qkv_out_scales) {
        input_left *= left_out_scale_vec[i];
        input_right *= right_out_scale_vec[i];
      }
      if (qkv_biases) {
        input_left = input_left + static_cast<float>(left_bias_vec[i]);
        input_right = input_right + static_cast<float>(right_bias_vec[i]);
      }
      if (hi < num_heads + gqa_group_size) {
        const float cos_tmp = cos_emb_vec[i];
        const float sin_tmp = sin_emb_vec[i];
        left_bias_vec[i] =
            static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
        right_bias_vec[i] =
            static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
      } else {
        left_bias_vec[i] = static_cast<T>(input_left);
        right_bias_vec[i] = static_cast<T>(input_right);
      }
    }
    if (hi < num_heads) {
      // write q
      Store<T, VecSize>(left_bias_vec, &qkv_out[ori_idx_left]);
      Store<T, VecSize>(right_bias_vec, &qkv_out[ori_idx_right]);
    } else {
      //  write k/v
      const int kv_head_idx = (hi - num_heads) % gqa_group_size;
      const int tgt_idx_left =
          (block_idx * gqa_group_size * block_size * head_size +
           kv_head_idx * block_size * head_size + block_offset * head_size +
           h_bias);
      const uint32_t tgt_idx_right = tgt_idx_left + half_head_size;
      // write

      if (hi < num_heads + gqa_group_size) {
        Store<T, VecSize>(left_bias_vec, &key_cache[tgt_idx_left]);
        Store<T, VecSize>(right_bias_vec, &key_cache[tgt_idx_right]);
      } else {
        Store<T, VecSize>(left_bias_vec, &value_cache[tgt_idx_left]);
        Store<T, VecSize>(right_bias_vec, &value_cache[tgt_idx_right]);
      }
    }
  }
}

template <typename T,
          int VecSize = 4,
          int RoundType = 0,
          int HeadDim = 128,
          typename InT = int,
          bool IsFP8 = false>
__global__ void append_speculate_cache_int8_rope_kernel(
    const InT* __restrict__ quant_qkv,  // [num_head, num_heads + 2 *
                                        // gqa_group_size, head_size]
    uint8_t* __restrict__ key_cache,    // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    uint8_t* __restrict__ value_cache,  // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    T* __restrict__ qkv_out,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens,          // [bsz]
    const int* __restrict__ seq_lens_encoder,  // [bsz]
    const float* __restrict__ cos_emb,
    const float* __restrict__ sin_emb,
    const float* __restrict__ qkv_out_scales,  // [num_head + 2 *
                                               // gqa_group_size, dim_head]
    const T* __restrict__ qkv_biases,  // [num_head + 2 * gqa_group_size,
                                       // dim_head]
    const T* __restrict__ cache_k_scales,
    const T* __restrict__ cache_v_scales,
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int block_size,
    const float max_bound,
    const float min_bound,
    const int gqa_group_size) {
  static_assert(HeadDim == 128, "just support HeadDim be 128 now!");
  static_assert(VecSize == 4, "just support VecSize be 4 now, 32 * 4!");
  constexpr int NUM_WARPS = 4;
  const int tid = threadIdx.x;
  const int wid = tid / 32;
  const int lane_id = tid % 32;
  const int token_id = blockIdx.x;
  
  const int bid = batch_id_per_token[token_id];

  const int start_token_idx = cu_seqlens_q[bid];
  const int head_idx = blockIdx.y * NUM_WARPS + wid;
  int q_head_idx, k_head_idx, v_idx;
  const int64_t hidden_size = (num_heads + 2 * gqa_group_size) * HeadDim;
  constexpr int half_head_size = HeadDim / 2;
  if (seq_lens_encoder[bid] > 0) return;
  const int write_seq_id = seq_lens[bid] + token_id - start_token_idx;
  if (write_seq_id == 0) return;
  const int* block_table_now = block_tables + bid * max_blocks_per_seq;
  const int block_idx = __ldg(&block_table_now[write_seq_id / block_size]);
  const int block_offset = write_seq_id % block_size;

  if (head_idx < num_heads) {
    // q
    using LoadT = AlignedVector<InT, VecSize>;
    using LoadBiasT = AlignedVector<T, VecSize>;
    using LoadOutScaleT = AlignedVector<float, VecSize>;
    constexpr int HalfVecSize = VecSize / 2;
    using LoadEmbT = AlignedVector<float, HalfVecSize>;

    LoadT src_vec;
    LoadBiasT bias_vec;
    LoadOutScaleT out_scale_vec;
    LoadEmbT cos_emb_vec;
    LoadEmbT sin_emb_vec;
    const InT* qkv_now = quant_qkv + token_id * hidden_size;
    T* qkv_out_now = qkv_out + token_id * hidden_size;
#pragma unroll
    for (uint32_t head_bias = lane_id * VecSize; head_bias < HeadDim;
         head_bias += 32 * VecSize) {
      const int bias_idx = head_idx * HeadDim + head_bias;
      Load<InT, VecSize>(&qkv_now[bias_idx], &src_vec);

      // q rope
      const uint32_t emb_idx = write_seq_id * half_head_size + head_bias / 2;
      Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
      Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
      if (qkv_out_scales) {
        Load<float, VecSize>(&qkv_out_scales[bias_idx], &out_scale_vec);
      }

      if (qkv_biases) {
        Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
      }
#pragma unroll
      for (int i = 0; i < HalfVecSize; i++) {
        // dequant + add_bias + rope
        float input_left = static_cast<float>(src_vec[2 * i]);
        float input_right = static_cast<float>(src_vec[2 * i + 1]);
        if (qkv_out_scales) {
          input_left = input_left * out_scale_vec[2 * i];
          input_right = input_right * out_scale_vec[2 * i + 1];
        }

        if (qkv_biases) {
          input_left = input_left + static_cast<float>(bias_vec[2 * i]);
          input_right = input_right + static_cast<float>(bias_vec[2 * i + 1]);
        }
        const float cos_tmp = cos_emb_vec[i];
        const float sin_tmp = sin_emb_vec[i];
        bias_vec[2 * i] =
            static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
        bias_vec[2 * i + 1] =
            static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
      }
      Store<T, VecSize>(bias_vec, &qkv_out_now[bias_idx]);
    }
  } else if (head_idx < num_heads + 2 * gqa_group_size) {
    // k
    constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t);  // 16
    using LoadPadKVT = AlignedVector<uint8_t, KV_VEC_SIZE>;
    const uint32_t kv_head_idx = (head_idx - num_heads) % gqa_group_size;

    constexpr int K_VEC_SIZE = 4;
    constexpr int HALF_K_VEC_SIZE = 2;
    using LoadKVResT = AlignedVector<uint8_t, K_VEC_SIZE>;
    using LoadKVT = AlignedVector<uint8_t, HALF_K_VEC_SIZE>;
    using LoadT = AlignedVector<InT, HALF_K_VEC_SIZE>;
    using LoadBiasT = AlignedVector<T, HALF_K_VEC_SIZE>;
    using LoadOutScaleT = AlignedVector<float, HALF_K_VEC_SIZE>;
    using LoadEmbT = AlignedVector<float, 1>;
    LoadKVResT cache_vec;
    LoadT src_vec1, src_vec2;
    LoadBiasT bias_vec1, bias_vec2;
    LoadOutScaleT out_scale_vec1, out_scale_vec2;
    LoadEmbT cos_emb_vec1, cos_emb_vec2;
    LoadEmbT sin_emb_vec1, sin_emb_vec2;

    const InT* qkv_now = quant_qkv + token_id * hidden_size;
    const int head_bias = lane_id / 4 * 16 + lane_id % 4 * 2;
    const int bias_idx = head_idx * HeadDim + head_bias;
    Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx], &src_vec1);
    Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx + 8], &src_vec2);
    T scale;
    if (head_idx < num_heads + gqa_group_size) {
      const uint32_t emb_idx = write_seq_id * half_head_size + head_bias / 2;
      Load<float, 1>(&cos_emb[emb_idx], &cos_emb_vec1);
      Load<float, 1>(&cos_emb[emb_idx + 4], &cos_emb_vec2);
      Load<float, 1>(&sin_emb[emb_idx], &sin_emb_vec1);
      Load<float, 1>(&sin_emb[emb_idx + 4], &sin_emb_vec2);
      scale = __ldg(&cache_k_scales[kv_head_idx]);
    } else {
      scale = __ldg(&cache_v_scales[kv_head_idx]);
    }

    float input_left = static_cast<float>(src_vec1[0]);
    float input_right = static_cast<float>(src_vec1[1]);
    if (qkv_out_scales) {
      Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx], &out_scale_vec1);
      Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx + 8],
                                   &out_scale_vec2);
      input_left = input_left * out_scale_vec1[0];
      input_right = input_right * out_scale_vec1[1];
    }

    if (qkv_biases) {
      Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx], &bias_vec1);
      Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx + 8], &bias_vec2);
      input_left = input_left + static_cast<float>(bias_vec1[0]);
      input_right = input_right + static_cast<float>(bias_vec1[1]);
    }
    if (head_idx < num_heads + gqa_group_size) {
      float cos_tmp = cos_emb_vec1[0];
      float sin_tmp = sin_emb_vec1[0];
      bias_vec1[0] =
          static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
      bias_vec1[1] =
          static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
    } else {
      bias_vec1[0] = static_cast<T>(input_left);
      bias_vec1[1] = static_cast<T>(input_right);
    }

    input_left = static_cast<float>(src_vec2[0]);
    input_right = static_cast<float>(src_vec2[1]);
    if (qkv_out_scales) {
      input_left = input_left * out_scale_vec2[0];
      input_right = input_right * out_scale_vec2[1];
    }

    if (qkv_biases) {
      input_left = input_left + static_cast<float>(bias_vec2[0]);
      input_right = input_right + static_cast<float>(bias_vec2[1]);
    }
    if (head_idx < num_heads + gqa_group_size) {
      float cos_tmp = cos_emb_vec2[0];
      float sin_tmp = sin_emb_vec2[0];
      bias_vec2[0] =
          static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
      bias_vec2[1] =
          static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
    } else {
      bias_vec2[0] = static_cast<T>(input_left);
      bias_vec2[1] = static_cast<T>(input_right);
    }
#pragma unroll
    for (uint32_t i = 0; i < HALF_K_VEC_SIZE; i++) {
      cache_vec[i] = QuantToC8<T,true, IsFP8, RoundType>(scale, bias_vec1[i], max_bound, min_bound);
      cache_vec[i + HALF_K_VEC_SIZE] = QuantToC8<T,true, IsFP8, RoundType>(scale, bias_vec2[i], max_bound, min_bound);
    }
    if (head_idx < num_heads + gqa_group_size) {
      const int start_block_16 =
          block_offset / 16 * 16 + block_offset % 8 + lane_id / 4 % 2 * 8;
      const uint32_t tgt_cache_idx =
          block_idx * gqa_group_size * block_size * HeadDim +
          kv_head_idx * block_size * HeadDim + start_block_16 * HeadDim +
          lane_id / 4 / 2 * 32 + (block_offset % 16) / 8 * 16 + lane_id % 4 * 4;
      Store<uint8_t, K_VEC_SIZE>(cache_vec, &key_cache[tgt_cache_idx]);
    } else {
      const uint32_t base_tgt_cache_idx =
          block_idx * gqa_group_size * HeadDim * block_size +
          kv_head_idx * HeadDim * block_size +
          (lane_id / 4 * 16 + lane_id % 4 * 2) * block_size +
          block_offset / 16 % 2 * 8 * block_size + block_offset / 16 / 2 * 32;
      const uint32_t tgt_cache_idx1 = base_tgt_cache_idx +
                                      block_offset % 8 / 2 * 4     // per 4
                                      + block_offset % 16 / 8 * 2  // per 2
                                      + block_offset % 2;          // per 1
      const uint32_t tgt_cache_idx2 = tgt_cache_idx1 + block_size;
      const uint32_t tgt_cache_idx3 = tgt_cache_idx1 + 16;
      const uint32_t tgt_cache_idx4 = tgt_cache_idx3 + block_size;
      value_cache[tgt_cache_idx1] = cache_vec[0];
      value_cache[tgt_cache_idx2] = cache_vec[1];
      value_cache[tgt_cache_idx3] = cache_vec[2];
      value_cache[tgt_cache_idx4] = cache_vec[3];
    }
  }
}

template <typename T,
          int VecSize = 4,
          int RoundType = 0,
          int HeadDim = 128,
          typename InT = int>
__global__ void append_speculate_cache_int8_neox_rope_kernel(
    const InT* __restrict__ quant_qkv,  // [num_head, num_heads + 2 *
                                        // gqa_group_size, head_size]
    uint8_t* __restrict__ key_cache,    // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    uint8_t* __restrict__ value_cache,  // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    T* __restrict__ qkv_out,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens,          // [bsz]
    const int* __restrict__ seq_lens_encoder,  // [bsz]
    const float* __restrict__ cos_emb,
    const float* __restrict__ sin_emb,
    const float* __restrict__ qkv_out_scales,  // [num_head + 2 *
                                               // gqa_group_size, dim_head]
    const T* __restrict__ qkv_biases,  // [num_head + 2 * gqa_group_size,
                                       // dim_head]
    const T* __restrict__ cache_k_scales,
    const T* __restrict__ cache_v_scales,
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int block_size,
    const float max_bound,
    const float min_bound,
    const int gqa_group_size) {
  static_assert(HeadDim == 128, "just support HeadDim be 128 now!");
  static_assert(VecSize == 4, "just support VecSize be 4 now, 32 * 4!");
  constexpr int NUM_WARPS = 4;
  const int tid = threadIdx.x;
  const int wid = tid / 32;
  const int lane_id = tid % 32;
  const int token_id = blockIdx.x;
  
  const int bid = batch_id_per_token[token_id];

  const int start_token_idx = cu_seqlens_q[bid];
  const int head_idx = blockIdx.y * NUM_WARPS + wid;
  int q_head_idx, k_head_idx, v_idx;

  const int64_t hidden_size = (num_heads + 2 * gqa_group_size) * HeadDim;
  constexpr int half_head_size = HeadDim / 2;
  if (seq_lens_encoder[bid] > 0) return;
  const int write_seq_id = seq_lens[bid] + token_id - start_token_idx;
  if (write_seq_id == 0) return;
  const int* block_table_now = block_tables + bid * max_blocks_per_seq;
  const int block_idx = __ldg(&block_table_now[write_seq_id / block_size]);
  const int block_offset = write_seq_id % block_size;

  if (head_idx < num_heads) {
    // q
    using LoadIntT = AlignedVector<InT, VecSize>;
    using LoadBiasT = AlignedVector<T, VecSize>;
    using LoadOutScaleT = AlignedVector<float, VecSize>;
    constexpr int HalfVecSize = VecSize / 2;
    using LoadEmbT = AlignedVector<float, VecSize>;

    LoadIntT left_vec, right_vec;
    LoadBiasT left_bias_vec, right_bias_vec;
    LoadOutScaleT left_out_scale_vec, right_out_scale_vec;
    LoadEmbT cos_emb_vec;
    LoadEmbT sin_emb_vec;
    const InT* qkv_now = quant_qkv + token_id * hidden_size;
    T* qkv_out_now = qkv_out + token_id * hidden_size;
#pragma unroll
    for (uint32_t head_bias = lane_id * VecSize; head_bias < half_head_size;
         head_bias += 32 * VecSize) {
      const int bias_idx_left = head_idx * HeadDim + head_bias;
      const int bias_idx_right = bias_idx_left + half_head_size;
      Load<InT, VecSize>(&qkv_now[bias_idx_left], &left_vec);
      Load<InT, VecSize>(&qkv_now[bias_idx_right], &right_vec);

      // q rope
      const uint32_t emb_idx = write_seq_id * HeadDim + head_bias;
      Load<float, VecSize>(&cos_emb[emb_idx], &cos_emb_vec);
      Load<float, VecSize>(&sin_emb[emb_idx], &sin_emb_vec);
      if (qkv_out_scales) {
        Load<float, VecSize>(&qkv_out_scales[bias_idx_left],
                             &left_out_scale_vec);
        Load<float, VecSize>(&qkv_out_scales[bias_idx_right],
                             &right_out_scale_vec);
      }

      if (qkv_biases) {
        Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
        Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
      }
#pragma unroll
      for (int i = 0; i < VecSize; i++) {
        // dequant + add_bias + rope
        float input_left = static_cast<float>(left_vec[i]);
        float input_right = static_cast<float>(right_vec[i]);
        if (qkv_out_scales) {
          input_left = input_left * left_out_scale_vec[i];
          input_right = input_right * right_out_scale_vec[i];
        }

        if (qkv_biases) {
          input_left = input_left + static_cast<float>(left_bias_vec[i]);
          input_right = input_right + static_cast<float>(right_bias_vec[i]);
        }
        const float cos_tmp = cos_emb_vec[i];
        const float sin_tmp = sin_emb_vec[i];
        left_bias_vec[i] =
            static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
        right_bias_vec[i] =
            static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
      }
      Store<T, VecSize>(left_bias_vec, &qkv_out_now[bias_idx_left]);
      Store<T, VecSize>(right_bias_vec, &qkv_out_now[bias_idx_right]);
    }
  } else if (head_idx < num_heads + 2 * gqa_group_size) {
    // k
    constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t);  // 16
    using LoadPadKVT = AlignedVector<uint8_t, KV_VEC_SIZE>;
    const uint32_t kv_head_idx = (head_idx - num_heads) % gqa_group_size;

    if (head_idx < num_heads + gqa_group_size) {
      // k
      const int head_bias = lane_id / 4 * 16 + lane_id % 4 * 2;
      if (head_bias < half_head_size) {
        constexpr int K_VEC_SIZE = 4;
        constexpr int HALF_K_VEC_SIZE = 2;
        using LoadKVResT = AlignedVector<uint8_t, K_VEC_SIZE>;
        using LoadKVT = AlignedVector<uint8_t, HALF_K_VEC_SIZE>;
        using LoadT = AlignedVector<InT, HALF_K_VEC_SIZE>;
        using LoadBiasT = AlignedVector<T, HALF_K_VEC_SIZE>;
        using LoadOutScaleT = AlignedVector<float, HALF_K_VEC_SIZE>;
        using LoadEmbT = AlignedVector<float, HALF_K_VEC_SIZE>;

        LoadKVResT left_cache_vec, right_cache_vec;
        LoadT left_src_vec1, left_src_vec2, right_src_vec1, right_src_vec2;
        LoadBiasT left_bias_vec1, left_bias_vec2, right_bias_vec1,
            right_bias_vec2;
        LoadOutScaleT left_out_scale_vec1, left_out_scale_vec2,
            right_out_scale_vec1, right_out_scale_vec2;
        LoadEmbT cos_emb_vec1, cos_emb_vec2;
        LoadEmbT sin_emb_vec1, sin_emb_vec2;

        const InT* qkv_now = quant_qkv + token_id * hidden_size;
        const int left_bias_idx = head_idx * HeadDim + head_bias;
        const int right_bias_idx = left_bias_idx + half_head_size;

        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[left_bias_idx], &left_src_vec1);
        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[left_bias_idx + 8], &left_src_vec2);
        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[right_bias_idx], &right_src_vec1);
        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[right_bias_idx + 8],
                                   &right_src_vec2);
        if (qkv_biases) {
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[left_bias_idx], &left_bias_vec1);
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[left_bias_idx + 8],
                                   &left_bias_vec2);
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[right_bias_idx],
                                   &right_bias_vec1);
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[right_bias_idx + 8],
                                   &right_bias_vec2);
        }
        if (qkv_out_scales) {
          Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[left_bias_idx],
                                       &left_out_scale_vec1);
          Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[left_bias_idx + 8],
                                       &left_out_scale_vec2);
          Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[right_bias_idx],
                                       &right_out_scale_vec1);
          Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[right_bias_idx + 8],
                                       &right_out_scale_vec2);
        }

        T scale;
        const uint32_t emb_idx = write_seq_id * HeadDim + head_bias;
        Load<float, HALF_K_VEC_SIZE>(&cos_emb[emb_idx], &cos_emb_vec1);
        Load<float, HALF_K_VEC_SIZE>(&cos_emb[emb_idx + 8], &cos_emb_vec2);
        Load<float, HALF_K_VEC_SIZE>(&sin_emb[emb_idx], &sin_emb_vec1);
        Load<float, HALF_K_VEC_SIZE>(&sin_emb[emb_idx + 8], &sin_emb_vec2);
        scale = __ldg(&cache_k_scales[kv_head_idx]);
#pragma unroll
        for (int i = 0; i < HALF_K_VEC_SIZE; i++) {
          float input_left = static_cast<float>(left_src_vec1[i]);
          float input_right = static_cast<float>(right_src_vec1[i]);
          if (qkv_out_scales) {
            input_left = input_left * left_out_scale_vec1[i];
            input_right = input_right * right_out_scale_vec1[i];
          }
          if (qkv_biases) {
            input_left = input_left + static_cast<float>(left_bias_vec1[i]);
            input_right = input_right + static_cast<float>(right_bias_vec1[i]);
          }

          float cos_tmp = cos_emb_vec1[i];
          float sin_tmp = sin_emb_vec1[i];
          left_bias_vec1[i] =
              static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
          right_bias_vec1[i] =
              static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);

          input_left = static_cast<float>(left_src_vec2[i]);
          input_right = static_cast<float>(right_src_vec2[i]);
          input_left = qkv_biases ? input_left * left_out_scale_vec2[i] +
                                        static_cast<float>(left_bias_vec2[i])
                                  : input_left * left_out_scale_vec2[i];
          input_right = qkv_biases ? input_right * right_out_scale_vec2[i] +
                                         static_cast<float>(right_bias_vec2[i])
                                   : input_right * right_out_scale_vec2[i];
          cos_tmp = cos_emb_vec2[i];
          sin_tmp = sin_emb_vec2[i];
          left_bias_vec2[i] =
              static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
          right_bias_vec2[i] =
              static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);

          float quant_value1 = static_cast<float>(scale * left_bias_vec1[i]);
          float quant_value2 = static_cast<float>(scale * left_bias_vec2[i]);
          if constexpr (RoundType == 0) {
            quant_value1 =
                static_cast<float>(roundWithTiesToEven(quant_value1));
            quant_value2 =
                static_cast<float>(roundWithTiesToEven(quant_value2));
          } else {
            quant_value1 = static_cast<float>(round(quant_value1));
            quant_value2 = static_cast<float>(round(quant_value2));
          }
          quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
          quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
          quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
          quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
          left_cache_vec[i] = static_cast<uint8_t>(quant_value1 + 128.0f);
          left_cache_vec[i + HALF_K_VEC_SIZE] =
              static_cast<uint8_t>(quant_value2 + 128.0f);

          quant_value1 = static_cast<float>(scale * right_bias_vec1[i]);
          quant_value2 = static_cast<float>(scale * right_bias_vec2[i]);
          if constexpr (RoundType == 0) {
            quant_value1 =
                static_cast<float>(roundWithTiesToEven(quant_value1));
            quant_value2 =
                static_cast<float>(roundWithTiesToEven(quant_value2));
          } else {
            quant_value1 = static_cast<float>(round(quant_value1));
            quant_value2 = static_cast<float>(round(quant_value2));
          }
          quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
          quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
          quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
          quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
          right_cache_vec[i] = static_cast<uint8_t>(quant_value1 + 128.0f);
          right_cache_vec[i + HALF_K_VEC_SIZE] =
              static_cast<uint8_t>(quant_value2 + 128.0f);
        }

        const int left_start_block_16 =
            block_offset / 16 * 16 + block_offset % 8 + lane_id / 4 % 2 * 8;
        const uint32_t left_tgt_cache_idx =
            block_idx * gqa_group_size * block_size * HeadDim +
            kv_head_idx * block_size * HeadDim + left_start_block_16 * HeadDim +
            lane_id / 4 / 2 * 32 + (block_offset % 16) / 8 * 16 +
            lane_id % 4 * 4;

        const int right_lane_id = lane_id + 16;
        const int right_start_block_16 = block_offset / 16 * 16 +
                                         block_offset % 8 +
                                         right_lane_id / 4 % 2 * 8;
        const uint32_t right_tgt_cache_idx =
            block_idx * gqa_group_size * block_size * HeadDim +
            kv_head_idx * block_size * HeadDim +
            right_start_block_16 * HeadDim + right_lane_id / 4 / 2 * 32 +
            (block_offset % 16) / 8 * 16 + right_lane_id % 4 * 4;
        Store<uint8_t, K_VEC_SIZE>(left_cache_vec,
                                   &key_cache[left_tgt_cache_idx]);
        Store<uint8_t, K_VEC_SIZE>(right_cache_vec,
                                   &key_cache[right_tgt_cache_idx]);
      }
    } else {
      // v
      constexpr int K_VEC_SIZE = 4;
      constexpr int HALF_K_VEC_SIZE = 2;
      using LoadKVResT = AlignedVector<uint8_t, K_VEC_SIZE>;
      using LoadKVT = AlignedVector<uint8_t, HALF_K_VEC_SIZE>;
      using LoadT = AlignedVector<InT, HALF_K_VEC_SIZE>;
      using LoadBiasT = AlignedVector<T, HALF_K_VEC_SIZE>;
      using LoadOutScaleT = AlignedVector<float, HALF_K_VEC_SIZE>;
      LoadKVResT cache_vec;
      LoadT src_vec1, src_vec2;
      LoadBiasT bias_vec1, bias_vec2;
      LoadOutScaleT out_scale_vec1, out_scale_vec2;

      const InT* qkv_now = quant_qkv + token_id * hidden_size;
      const int head_bias = lane_id / 4 * 16 + lane_id % 4 * 2;
      const int bias_idx = head_idx * HeadDim + head_bias;
      Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx], &src_vec1);
      Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx + 8], &src_vec2);
      if (qkv_biases) {
        Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx], &bias_vec1);
        Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx + 8], &bias_vec2);
      }
      if (qkv_out_scales) {
        Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx],
                                     &out_scale_vec1);
        Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx + 8],
                                     &out_scale_vec2);
      }

      T scale = __ldg(&cache_v_scales[kv_head_idx]);

      float input_left = static_cast<float>(src_vec1[0]);
      float input_right = static_cast<float>(src_vec1[1]);
      if (qkv_out_scales) {
        input_left = input_left * out_scale_vec1[0];
        input_right = input_right * out_scale_vec1[1];
      }
      if (qkv_biases) {
        input_left = input_left + static_cast<float>(bias_vec1[0]);
        input_right = input_right + static_cast<float>(bias_vec1[1]);
      }

      bias_vec1[0] = static_cast<T>(input_left);
      bias_vec1[1] = static_cast<T>(input_right);

      input_left = static_cast<float>(src_vec2[0]);
      input_right = static_cast<float>(src_vec2[1]);
      if (qkv_out_scales) {
        input_left = input_left * out_scale_vec2[0];
        input_right = input_right * out_scale_vec2[1];
      }
      if (qkv_biases) {
        input_left = input_left + static_cast<float>(bias_vec2[0]);
        input_right = input_right + static_cast<float>(bias_vec2[1]);
      }

      bias_vec2[0] = static_cast<T>(input_left);
      bias_vec2[1] = static_cast<T>(input_right);

#pragma unroll
      for (uint32_t i = 0; i < HALF_K_VEC_SIZE; i++) {
        float quant_value1 = static_cast<float>(scale * bias_vec1[i]);
        float quant_value2 = static_cast<float>(scale * bias_vec2[i]);
        if constexpr (RoundType == 0) {
          quant_value1 = static_cast<float>(roundWithTiesToEven(quant_value1));
          quant_value2 = static_cast<float>(roundWithTiesToEven(quant_value2));
        } else {
          quant_value1 = static_cast<float>(round(quant_value1));
          quant_value2 = static_cast<float>(round(quant_value2));
        }
        quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
        quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
        quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
        quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
        cache_vec[i] = static_cast<uint8_t>(quant_value1 + 128.0f);
        cache_vec[i + HALF_K_VEC_SIZE] =
            static_cast<uint8_t>(quant_value2 + 128.0f);
      }

      const uint32_t base_tgt_cache_idx =
          block_idx * gqa_group_size * HeadDim * block_size +
          kv_head_idx * HeadDim * block_size +
          (lane_id / 4 * 16 + lane_id % 4 * 2) * block_size +
          block_offset / 16 % 2 * 8 * block_size + block_offset / 16 / 2 * 32;
      const uint32_t tgt_cache_idx1 = base_tgt_cache_idx +
                                      block_offset % 8 / 2 * 4     // per 4
                                      + block_offset % 16 / 8 * 2  // per 2
                                      + block_offset % 2;          // per 1
      const uint32_t tgt_cache_idx2 = tgt_cache_idx1 + block_size;
      const uint32_t tgt_cache_idx3 = tgt_cache_idx1 + 16;
      const uint32_t tgt_cache_idx4 = tgt_cache_idx3 + block_size;
      value_cache[tgt_cache_idx1] = cache_vec[0];
      value_cache[tgt_cache_idx2] = cache_vec[1];
      value_cache[tgt_cache_idx3] = cache_vec[2];
      value_cache[tgt_cache_idx4] = cache_vec[3];
    }
  }
}

template <typename T,
          int VecSize = 4,
          int RoundType = 0,
          int HeadDim = 128,
          typename InT = int>
__global__ void append_speculate_cache_int4_rope_kernel(
    const InT* __restrict__ quant_qkv,  // [bsz, num_heads + 2 * gqa_group_size,
                                        // head_size]
    uint8_t* __restrict__ key_cache,    // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    uint8_t* __restrict__ value_cache,  // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    T* __restrict__ qkv_out,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens,          // [bsz]
    const int* __restrict__ seq_lens_encoder,  // [bsz]
    const float* __restrict__ cos_emb,
    const float* __restrict__ sin_emb,
    const float* __restrict__ qkv_out_scales,  // [num_head + 2 *
                                               // gqa_group_size, dim_head]
    const T* __restrict__ qkv_biases,  // [num_head + 2 * gqa_group_size,
                                       // dim_head]
    const T* __restrict__ cache_k_scales,
    const T* __restrict__ cache_v_scales,
    const T* __restrict__ cache_k_zero_points,
    const T* __restrict__ cache_v_zero_points,
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int block_size,
    const float max_bound,
    const float min_bound,
    const int gqa_group_size) {
  static_assert(HeadDim == 128, "just support HeadDim be 128 now!");
  static_assert(VecSize == 4, "just support VecSize be 4 now, 32 * 4!");
  constexpr int NUM_WARPS = 4;
  const int tid = threadIdx.x;
  const int wid = tid / 32;
  const int lane_id = tid % 32;

  const int token_id = blockIdx.x;
  
  const int bid = batch_id_per_token[token_id];

  const int start_token_idx = cu_seqlens_q[bid];
  const int head_idx = blockIdx.y * NUM_WARPS + wid;

  const int64_t hidden_size = (num_heads + 2 * gqa_group_size) * HeadDim;
  constexpr int half_head_size = HeadDim / 2;
  const int half_block_size = block_size / 2;
  if (seq_lens_encoder[bid] > 0) return;
  const int write_seq_id = seq_lens[bid] + token_id - start_token_idx;
  if (write_seq_id == 0) return;
  const int* block_table_now = nullptr;

  block_table_now = block_tables + bid * max_blocks_per_seq;

  const int block_idx = __ldg(&block_table_now[write_seq_id / block_size]);

  const int block_offset = write_seq_id % block_size;

  if (head_idx < num_heads) {
    // q
    using LoadInT = AlignedVector<InT, VecSize>;
    using LoadBiasT = AlignedVector<T, VecSize>;
    using LoadOutScaleT = AlignedVector<float, VecSize>;
    constexpr int HalfVecSize = VecSize / 2;
    using LoadEmbT = AlignedVector<float, HalfVecSize>;

    LoadInT src_vec;
    LoadBiasT bias_vec;
    LoadOutScaleT out_scale_vec;
    LoadEmbT cos_emb_vec;
    LoadEmbT sin_emb_vec;
    const InT* qkv_now = quant_qkv + token_id * hidden_size;
    T* qkv_out_now = qkv_out + token_id * hidden_size;
#pragma unroll
    for (uint32_t head_bias = lane_id * VecSize; head_bias < HeadDim;
         head_bias += 32 * VecSize) {
      const int bias_idx = head_idx * HeadDim + head_bias;
      Load<InT, VecSize>(&qkv_now[bias_idx], &src_vec);
      Load<T, VecSize>(&qkv_biases[bias_idx], &bias_vec);
      Load<float, VecSize>(&qkv_out_scales[bias_idx], &out_scale_vec);
      // q rope
      const uint32_t emb_idx = write_seq_id * half_head_size + head_bias / 2;
      Load<float, HalfVecSize>(&cos_emb[emb_idx], &cos_emb_vec);
      Load<float, HalfVecSize>(&sin_emb[emb_idx], &sin_emb_vec);
#pragma unroll
      for (int i = 0; i < HalfVecSize; i++) {
        // dequant + add_bias + rope
        float input_left = static_cast<float>(src_vec[2 * i]);
        float input_right = static_cast<float>(src_vec[2 * i + 1]);
        input_left = input_left * out_scale_vec[2 * i] +
                     static_cast<float>(bias_vec[2 * i]);
        input_right = input_right * out_scale_vec[2 * i + 1] +
                      static_cast<float>(bias_vec[2 * i + 1]);
        const float cos_tmp = cos_emb_vec[i];
        const float sin_tmp = sin_emb_vec[i];
        bias_vec[2 * i] =
            static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
        bias_vec[2 * i + 1] =
            static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
      }
      Store<T, VecSize>(bias_vec, &qkv_out_now[bias_idx]);
    }
  } else if (head_idx < num_heads + 2 * gqa_group_size) {
    // k
    constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t);  // 16
    using LoadPadKVT = AlignedVector<uint8_t, KV_VEC_SIZE>;
    const uint32_t kv_head_idx = (head_idx - num_heads) % gqa_group_size;

    constexpr int K_VEC_SIZE = 4;
    constexpr int HALF_K_VEC_SIZE = 2;
    using LoadKVResT = AlignedVector<uint8_t, K_VEC_SIZE>;
    using LoadInT = AlignedVector<InT, HALF_K_VEC_SIZE>;
    using LoadBiasT = AlignedVector<T, HALF_K_VEC_SIZE>;
    using LoadOutScaleT = AlignedVector<float, HALF_K_VEC_SIZE>;
    using LoadScaleT = AlignedVector<T, HALF_K_VEC_SIZE>;
    using LoadEmbT = AlignedVector<float, 1>;
    LoadInT src_vec1, src_vec2;
    LoadBiasT bias_vec1, bias_vec2;
    LoadOutScaleT out_scale_vec1, out_scale_vec2;
    LoadScaleT scale_vec1, scale_vec2;
    LoadScaleT zp_vec1, zp_vec2;
    LoadEmbT cos_emb_vec1, cos_emb_vec2;
    LoadEmbT sin_emb_vec1, sin_emb_vec2;

    const InT* qkv_now = quant_qkv + token_id * hidden_size;
    const int head_bias = lane_id / 4 * 16 + lane_id % 4 * 2;
    //////////
    const uint32_t cache_idx = kv_head_idx * HeadDim + head_bias;
    const int bias_idx = head_idx * HeadDim + head_bias;
    Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx], &src_vec1);
    Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx + 8], &src_vec2);
    /////
    Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx], &bias_vec1);
    Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx + 8], &bias_vec2);
    Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx], &out_scale_vec1);
    Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx + 8],
                                 &out_scale_vec2);
    if (head_idx < num_heads + gqa_group_size) {
      const uint32_t emb_idx = write_seq_id * half_head_size + head_bias / 2;
      Load<float, 1>(&cos_emb[emb_idx], &cos_emb_vec1);
      Load<float, 1>(&cos_emb[emb_idx + 4], &cos_emb_vec2);
      Load<float, 1>(&sin_emb[emb_idx], &sin_emb_vec1);
      Load<float, 1>(&sin_emb[emb_idx + 4], &sin_emb_vec2);
      Load<T, HALF_K_VEC_SIZE>(&cache_k_scales[cache_idx], &scale_vec1);
      Load<T, HALF_K_VEC_SIZE>(&cache_k_scales[cache_idx + 8], &scale_vec2);
      Load<T, HALF_K_VEC_SIZE>(&cache_k_zero_points[cache_idx], &zp_vec1);
      Load<T, HALF_K_VEC_SIZE>(&cache_k_zero_points[cache_idx + 8], &zp_vec2);
    } else {
      Load<T, HALF_K_VEC_SIZE>(&cache_v_scales[cache_idx], &scale_vec1);
      Load<T, HALF_K_VEC_SIZE>(&cache_v_scales[cache_idx + 8], &scale_vec2);
      Load<T, HALF_K_VEC_SIZE>(&cache_v_zero_points[cache_idx], &zp_vec1);
      Load<T, HALF_K_VEC_SIZE>(&cache_v_zero_points[cache_idx + 8], &zp_vec2);
    }

    float input_left = static_cast<float>(src_vec1[0]);
    float input_right = static_cast<float>(src_vec1[1]);
    input_left =
        input_left * out_scale_vec1[0] + static_cast<float>(bias_vec1[0]);
    input_right =
        input_right * out_scale_vec1[1] + static_cast<float>(bias_vec1[1]);
    if (head_idx < num_heads + gqa_group_size) {
      float cos_tmp = cos_emb_vec1[0];
      float sin_tmp = sin_emb_vec1[0];
      bias_vec1[0] =
          static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
      bias_vec1[1] =
          static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
    } else {
      bias_vec1[0] = static_cast<T>(input_left);
      bias_vec1[1] = static_cast<T>(input_right);
    }

    input_left = static_cast<float>(src_vec2[0]);
    input_right = static_cast<float>(src_vec2[1]);
    input_left =
        input_left * out_scale_vec2[0] + static_cast<float>(bias_vec2[0]);
    input_right =
        input_right * out_scale_vec2[1] + static_cast<float>(bias_vec2[1]);
    if (head_idx < num_heads + gqa_group_size) {
      float cos_tmp = cos_emb_vec2[0];
      float sin_tmp = sin_emb_vec2[0];
      bias_vec2[0] =
          static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
      bias_vec2[1] =
          static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
    } else {
      bias_vec2[0] = static_cast<T>(input_left);
      bias_vec2[1] = static_cast<T>(input_right);
    }
    if (head_idx < num_heads + gqa_group_size) {
      LoadKVResT cache_vec;
      const int start_block_16 =
          block_offset / 16 * 16 + block_offset % 8 + lane_id / 4 % 4 / 2 * 8;
      const uint32_t tgt_cache_idx =
          block_idx * gqa_group_size * block_size * half_head_size +
          kv_head_idx * block_size * half_head_size +
          start_block_16 * half_head_size + lane_id / 4 / 4 * 32 +
          lane_id / 4 % 2 * 16 + lane_id % 4 * 4;
      Load<uint8_t, K_VEC_SIZE>(&key_cache[tgt_cache_idx], &cache_vec);
#pragma unroll
      for (uint32_t i = 0; i < HALF_K_VEC_SIZE; i++) {
        float quant_value =
            static_cast<float>(scale_vec1[i] * bias_vec1[i] + zp_vec1[i]);
        if constexpr (RoundType == 0) {
          quant_value = roundWithTiesToEven(quant_value);
        } else {
          quant_value = round(quant_value);
        }
        quant_value = quant_value > max_bound ? max_bound : quant_value;
        quant_value = quant_value < min_bound ? min_bound : quant_value;
        uint8_t uint_quant_value = static_cast<uint8_t>(quant_value + 8.0f);
        uint8_t ano_uint_quant_value = 0;
        if (block_offset % 16 / 8 == 0) {
          cache_vec[i] &= 0xF0;
          cache_vec[i] |= ((ano_uint_quant_value) | (uint_quant_value & 0x0F));
        } else {
          cache_vec[i] &= 0x0F;
          cache_vec[i] |= ((uint_quant_value << 4) | (ano_uint_quant_value));
        }
      }
#pragma unroll
      for (uint32_t i = 0; i < HALF_K_VEC_SIZE; i++) {
        float quant_value =
            static_cast<float>(scale_vec2[i] * bias_vec2[i] + zp_vec2[i]);
        if constexpr (RoundType == 0) {
          quant_value = roundWithTiesToEven(quant_value);
        } else {
          quant_value = round(quant_value);
        }
        quant_value = quant_value > max_bound ? max_bound : quant_value;
        quant_value = quant_value < min_bound ? min_bound : quant_value;
        uint8_t uint_quant_value = static_cast<uint8_t>(quant_value + 8.0f);
        uint8_t ano_uint_quant_value = 0;
        if (block_offset % 16 / 8 == 0) {
          cache_vec[i + HALF_K_VEC_SIZE] &= 0xF0;
          cache_vec[i + HALF_K_VEC_SIZE] |=
              ((ano_uint_quant_value) | (uint_quant_value & 0x0F));
        } else {
          cache_vec[i + HALF_K_VEC_SIZE] &= 0x0F;
          cache_vec[i + HALF_K_VEC_SIZE] |=
              ((uint_quant_value << 4) | (ano_uint_quant_value));
        }
      }
      Store<uint8_t, K_VEC_SIZE>(cache_vec, &key_cache[tgt_cache_idx]);
    } else {

      const uint32_t base_tgt_cache_idx =
          block_idx * gqa_group_size * HeadDim * half_block_size +
          kv_head_idx * HeadDim * half_block_size +
          (lane_id / 4 * 16 + lane_id % 4 * 2) * half_block_size +
          block_offset / 16 % 4 / 2 * 8 * half_block_size +
          block_offset / 16 / 4 * 32 + block_offset / 16 % 2 * 16;
      const uint32_t tgt_cache_idx1 = base_tgt_cache_idx +
                                      block_offset % 8 / 2 * 4     // per 4
                                      + block_offset % 16 / 8 * 2  // per 2
                                      + block_offset % 2;          // per 1
      const uint32_t tgt_cache_idx2 = tgt_cache_idx1 + half_block_size;

      float quant_value1 =
          static_cast<float>(scale_vec1[0] * bias_vec1[0] + zp_vec1[0]);
      float quant_value2 =
          static_cast<float>(scale_vec2[0] * bias_vec2[0] + zp_vec2[0]);
      if constexpr (RoundType == 0) {
        quant_value1 = roundWithTiesToEven(quant_value1);
        quant_value2 = roundWithTiesToEven(quant_value2);
      } else {
        quant_value1 = round(quant_value1);
        quant_value2 = round(quant_value2);
      }
      quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
      quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
      quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
      quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
      uint8_t uint_quant_value1 = static_cast<uint8_t>(quant_value1 + 8.0f);
      uint8_t uint_quant_value2 = static_cast<uint8_t>(quant_value2 + 8.0f);
      value_cache[tgt_cache_idx1] =
          (uint_quant_value2 << 4) | (uint_quant_value1 & 0x0F);

      quant_value1 =
          static_cast<float>(scale_vec1[1] * bias_vec1[1] + zp_vec1[1]);
      quant_value2 =
          static_cast<float>(scale_vec2[1] * bias_vec2[1] + zp_vec2[1]);
      if constexpr (RoundType == 0) {
        quant_value1 = roundWithTiesToEven(quant_value1);
        quant_value2 = roundWithTiesToEven(quant_value2);
      } else {
        quant_value1 = round(quant_value1);
        quant_value2 = round(quant_value2);
      }
      quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
      quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
      quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
      quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
      uint_quant_value1 = static_cast<uint8_t>(quant_value1 + 8.0f);
      uint_quant_value2 = static_cast<uint8_t>(quant_value2 + 8.0f);
      value_cache[tgt_cache_idx2] =
          (uint_quant_value2 << 4) | (uint_quant_value1 & 0x0F);
    }
  }
}

template <typename T,
          int VecSize = 4,
          int RoundType = 0,
          int HeadDim = 128,
          typename InT = int>
__global__ void append_speculate_cache_int4_neox_rope_kernel(
    const InT* __restrict__ quant_qkv,  // [bsz, num_heads + 2 * gqa_group_size,
                                        // head_size]
    uint8_t* __restrict__ key_cache,    // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    uint8_t* __restrict__ value_cache,  // [num_blocks, gqa_group_size,
                                        // block_size, head_size // 2]
    T* __restrict__ qkv_out,
    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
    const int* __restrict__ batch_id_per_token,  // [num_tokens]
    const int* __restrict__ cu_seqlens_q,
    const int* __restrict__ seq_lens,          // [bsz]
    const int* __restrict__ seq_lens_encoder,  // [bsz]
    const float* __restrict__ cos_emb,
    const float* __restrict__ sin_emb,
    const float* __restrict__ qkv_out_scales,  // [num_head + 2 *
                                               // gqa_group_size, dim_head]
    const T* __restrict__ qkv_biases,  // [num_head + 2 * gqa_group_size,
                                       // dim_head]
    const T* __restrict__ cache_k_scales,
    const T* __restrict__ cache_v_scales,
    const T* __restrict__ cache_k_zero_points,
    const T* __restrict__ cache_v_zero_points,
    const int max_seq_len,
    const int max_blocks_per_seq,
    const int num_heads,
    const int block_size,
    const float max_bound,
    const float min_bound,
    const int gqa_group_size) {
  static_assert(HeadDim == 128, "just support HeadDim be 128 now!");
  static_assert(VecSize == 4, "just support VecSize be 4 now, 32 * 4!");
  constexpr int NUM_WARPS = 4;
  const int tid = threadIdx.x;
  const int wid = tid / 32;
  const int lane_id = tid % 32;

  const int token_id = blockIdx.x;
  
  const int bid = batch_id_per_token[token_id];

  const int start_token_idx = cu_seqlens_q[bid];
  const int head_idx = blockIdx.y * NUM_WARPS + wid;

  const int64_t hidden_size = (num_heads + 2 * gqa_group_size) * HeadDim;
  constexpr int half_head_size = HeadDim / 2;
  const int half_block_size = block_size / 2;
  if (seq_lens_encoder[bid] > 0) return;
  const int write_seq_id = seq_lens[bid] + token_id - start_token_idx;
  if (write_seq_id == 0) return;
  const int* block_table_now = nullptr;

  block_table_now = block_tables + bid * max_blocks_per_seq;

  const int block_idx = __ldg(&block_table_now[write_seq_id / block_size]);

  const int block_offset = write_seq_id % block_size;

  if (head_idx < num_heads) {
    // q
    using LoadIntT = AlignedVector<InT, VecSize>;
    using LoadBiasT = AlignedVector<T, VecSize>;
    using LoadOutScaleT = AlignedVector<float, VecSize>;
    constexpr int HalfVecSize = VecSize / 2;
    using LoadEmbT = AlignedVector<float, VecSize>;

    LoadIntT left_vec;
    LoadIntT right_vec;
    LoadBiasT left_bias_vec;
    LoadBiasT right_bias_vec;
    LoadOutScaleT left_out_scale_vec;
    LoadOutScaleT right_out_scale_vec;
    LoadEmbT cos_emb_vec;
    LoadEmbT sin_emb_vec;
    const InT* qkv_now = quant_qkv + token_id * hidden_size;
    T* qkv_out_now = qkv_out + token_id * hidden_size;
#pragma unroll
    for (uint32_t head_bias = lane_id * VecSize; head_bias < half_head_size;
         head_bias += 32 * VecSize) {
      const int bias_idx_left = head_idx * HeadDim + head_bias;
      const int bias_idx_right = bias_idx_left + half_head_size;
      Load<InT, VecSize>(&qkv_now[bias_idx_left], &left_vec);
      Load<InT, VecSize>(&qkv_now[bias_idx_right], &right_vec);

      if (qkv_biases) {
        Load<T, VecSize>(&qkv_biases[bias_idx_left], &left_bias_vec);
        Load<T, VecSize>(&qkv_biases[bias_idx_right], &right_bias_vec);
      }
      if (qkv_out_scales) {
        Load<float, VecSize>(&qkv_out_scales[bias_idx_left],
                             &left_out_scale_vec);
        Load<float, VecSize>(&qkv_out_scales[bias_idx_right],
                             &right_out_scale_vec);
      }
      // q rope
      const uint32_t emb_idx = write_seq_id * HeadDim + head_bias;
#pragma unroll
      for (int i = 0; i < VecSize; i++) {
        // dequant + add_bias + rope
        float input_left = static_cast<float>(left_vec[i]);
        float input_right = static_cast<float>(right_vec[i]);
        if (qkv_out_scales) {
          input_left = input_left * left_out_scale_vec[i];
          input_right = input_right * right_out_scale_vec[i];
        }

        if (qkv_biases) {
          input_left = input_left + static_cast<float>(left_bias_vec[i]);
          input_right = input_right + static_cast<float>(right_bias_vec[i]);
        }
        const float cos_tmp = cos_emb_vec[i];
        const float sin_tmp = sin_emb_vec[i];
        left_bias_vec[i] =
            static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
        right_bias_vec[i] =
            static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
      }
      Store<T, VecSize>(left_bias_vec, &qkv_out_now[bias_idx_left]);
      Store<T, VecSize>(right_bias_vec, &qkv_out_now[bias_idx_right]);
    }
  } else if (head_idx < num_heads + 2 * gqa_group_size) {
    // k
    constexpr int KV_VEC_SIZE = 16 / sizeof(uint8_t);  // 16
    using LoadPadKVT = AlignedVector<uint8_t, KV_VEC_SIZE>;
    const uint32_t kv_head_idx = (head_idx - num_heads) % gqa_group_size;

    if (head_idx < num_heads + gqa_group_size) {
      const int head_bias = lane_id / 4 * 16 + lane_id % 4 * 2;
      if (head_bias < half_head_size) {
        constexpr int K_VEC_SIZE = 4;
        constexpr int HALF_K_VEC_SIZE = 2;
        using LoadKVResT = AlignedVector<uint8_t, K_VEC_SIZE>;
        using LoadIntT = AlignedVector<InT, HALF_K_VEC_SIZE>;
        using LoadBiasT = AlignedVector<T, HALF_K_VEC_SIZE>;
        using LoadOutScaleT = AlignedVector<float, HALF_K_VEC_SIZE>;
        using LoadScaleT = AlignedVector<T, HALF_K_VEC_SIZE>;
        using LoadEmbT = AlignedVector<float, HALF_K_VEC_SIZE>;
        LoadIntT left_src_vec1, left_src_vec2, right_src_vec1, right_src_vec2;
        LoadBiasT left_bias_vec1, left_bias_vec2, right_bias_vec1,
            right_bias_vec2;
        LoadOutScaleT left_out_scale_vec1, left_out_scale_vec2,
            right_out_scale_vec1, right_out_scale_vec2;
        LoadScaleT left_scale_vec1, left_scale_vec2, right_scale_vec1,
            right_scale_vec2;
        LoadScaleT left_zp_vec1, left_zp_vec2, right_zp_vec1, right_zp_vec2;
        LoadEmbT cos_emb_vec1, cos_emb_vec2;
        LoadEmbT sin_emb_vec1, sin_emb_vec2;

        const InT* qkv_now = quant_qkv + token_id * hidden_size;
        const int left_bias_idx = head_idx * HeadDim + head_bias;
        const int right_bias_idx = left_bias_idx + half_head_size;

        const uint32_t left_cache_idx = kv_head_idx * HeadDim + head_bias;
        const uint32_t right_cache_idx = left_cache_idx + half_head_size;

        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[left_bias_idx], &left_src_vec1);
        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[left_bias_idx + 8], &left_src_vec2);
        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[right_bias_idx], &right_src_vec1);
        Load<InT, HALF_K_VEC_SIZE>(&qkv_now[right_bias_idx + 8],
                                   &right_src_vec2);
        if (qkv_biases) {
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[left_bias_idx], &left_bias_vec1);
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[left_bias_idx + 8],
                                   &left_bias_vec2);
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[right_bias_idx],
                                   &right_bias_vec1);
          Load<T, HALF_K_VEC_SIZE>(&qkv_biases[right_bias_idx + 8],
                                   &right_bias_vec2);
        }
        Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[left_bias_idx],
                                     &left_out_scale_vec1);
        Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[left_bias_idx + 8],
                                     &left_out_scale_vec2);
        Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[right_bias_idx],
                                     &right_out_scale_vec1);
        Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[right_bias_idx + 8],
                                     &right_out_scale_vec2);

        const uint32_t emb_idx = write_seq_id * HeadDim + head_bias;
        Load<float, HALF_K_VEC_SIZE>(&cos_emb[emb_idx], &cos_emb_vec1);
        Load<float, HALF_K_VEC_SIZE>(&cos_emb[emb_idx + 8], &cos_emb_vec2);
        Load<float, HALF_K_VEC_SIZE>(&sin_emb[emb_idx], &sin_emb_vec1);
        Load<float, HALF_K_VEC_SIZE>(&sin_emb[emb_idx + 8], &sin_emb_vec2);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_scales[left_cache_idx],
                                 &left_scale_vec1);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_scales[left_cache_idx + 8],
                                 &left_scale_vec2);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_zero_points[left_cache_idx],
                                 &left_zp_vec1);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_zero_points[left_cache_idx + 8],
                                 &left_zp_vec2);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_scales[right_cache_idx],
                                 &right_scale_vec1);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_scales[right_cache_idx + 8],
                                 &right_scale_vec2);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_zero_points[right_cache_idx],
                                 &right_zp_vec1);
        Load<T, HALF_K_VEC_SIZE>(&cache_k_zero_points[right_cache_idx + 8],
                                 &right_zp_vec2);

        for (int i = 0; i < HALF_K_VEC_SIZE; i++) {
          float input_left = static_cast<float>(left_src_vec1[i]);
          float input_right = static_cast<float>(right_src_vec1[i]);
          input_left = qkv_biases ? input_left * left_out_scale_vec1[i] +
                                        static_cast<float>(left_bias_vec1[i])
                                  : input_left * left_out_scale_vec1[i];
          input_right = qkv_biases ? input_right * right_out_scale_vec1[i] +
                                         static_cast<float>(right_bias_vec1[i])
                                   : input_right * right_out_scale_vec1[i];
          float cos_tmp = cos_emb_vec1[0];
          float sin_tmp = sin_emb_vec1[0];
          left_bias_vec1[i] =
              static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
          right_bias_vec1[i] =
              static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);


          input_left = static_cast<float>(left_src_vec2[i]);
          input_right = static_cast<float>(right_src_vec2[i]);
          cos_tmp = cos_emb_vec2[i];
          sin_tmp = sin_emb_vec2[i];
          left_bias_vec2[i] =
              static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
          right_bias_vec2[i] =
              static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
          // quant + write k
        }
        LoadKVResT left_cache_vec, right_cache_vec;
        const int left_start_block_16 =
            block_offset / 16 * 16 + block_offset % 8 + lane_id / 4 % 4 / 2 * 8;
        const uint32_t left_tgt_cache_idx =
            block_idx * gqa_group_size * block_size * half_head_size +
            kv_head_idx * block_size * half_head_size +
            left_start_block_16 * half_head_size + lane_id / 4 / 4 * 32 +
            lane_id / 4 % 2 * 16 + lane_id % 4 * 4;

        const int right_lane_id = lane_id + 16;
        const int right_start_block_16 = block_offset / 16 * 16 +
                                         block_offset % 8 +
                                         right_lane_id / 4 % 4 / 2 * 8;
        const uint32_t right_tgt_cache_idx =
            block_idx * gqa_group_size * block_size * half_head_size +
            kv_head_idx * block_size * half_head_size +
            right_start_block_16 * half_head_size + right_lane_id / 4 / 4 * 32 +
            right_lane_id / 4 % 2 * 16 + right_lane_id % 4 * 4;
        Load<uint8_t, K_VEC_SIZE>(&key_cache[left_tgt_cache_idx],
                                  &left_cache_vec);
        Load<uint8_t, K_VEC_SIZE>(&key_cache[right_tgt_cache_idx],
                                  &right_cache_vec);

#pragma unroll
        for (uint32_t i = 0; i < HALF_K_VEC_SIZE; i++) {
          float quant_value1 = static_cast<float>(
              left_scale_vec1[i] * left_bias_vec1[i] + left_zp_vec1[i]);
          float quant_value2 = static_cast<float>(
              left_scale_vec2[i] * left_bias_vec2[i] + left_zp_vec2[i]);
          if constexpr (RoundType == 0) {
            quant_value1 = roundWithTiesToEven(quant_value1);
            quant_value2 = roundWithTiesToEven(quant_value2);
          } else {
            quant_value1 = round(quant_value1);
            quant_value2 = round(quant_value2);
          }
          quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
          quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
          quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
          quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
          uint8_t uint_quant_value1 = static_cast<uint8_t>(quant_value1 + 8.0f);
          uint8_t uint_quant_value2 = static_cast<uint8_t>(quant_value2 + 8.0f);
          uint8_t ano_uint_quant_value = 0;
          if (block_offset % 16 / 8 == 0) {
            left_cache_vec[i] |=
                ((ano_uint_quant_value) | (uint_quant_value1 & 0x0F));
            left_cache_vec[i + HALF_K_VEC_SIZE] |=
                ((ano_uint_quant_value) | (uint_quant_value2 & 0x0F));
          } else {
            left_cache_vec[i] |=
                ((uint_quant_value1 << 4) | (ano_uint_quant_value));
            left_cache_vec[i + HALF_K_VEC_SIZE] |=
                ((uint_quant_value2 << 4) | (ano_uint_quant_value));
          }

          quant_value1 = static_cast<float>(
              right_scale_vec1[i] * right_bias_vec1[i] + right_zp_vec1[i]);
          quant_value2 = static_cast<float>(
              right_scale_vec2[i] * right_bias_vec2[i] + right_zp_vec2[i]);
          if constexpr (RoundType == 0) {
            quant_value1 = roundWithTiesToEven(quant_value1);
            quant_value2 = roundWithTiesToEven(quant_value2);
          } else {
            quant_value1 = round(quant_value1);
            quant_value2 = round(quant_value2);
          }
          quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
          quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
          quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
          quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
          uint_quant_value1 = static_cast<uint8_t>(quant_value1 + 8.0f);
          uint_quant_value2 = static_cast<uint8_t>(quant_value2 + 8.0f);
          ano_uint_quant_value = 0;
          if (block_offset % 16 / 8 == 0) {
            right_cache_vec[i] |=
                ((ano_uint_quant_value) | (uint_quant_value1 & 0x0F));
            right_cache_vec[i + HALF_K_VEC_SIZE] |=
                ((ano_uint_quant_value) | (uint_quant_value2 & 0x0F));
          } else {
            right_cache_vec[i] |=
                ((uint_quant_value1 << 4) | (ano_uint_quant_value));
            right_cache_vec[i + HALF_K_VEC_SIZE] |=
                ((uint_quant_value2 << 4) | (ano_uint_quant_value));
          }
        }
        Store<uint8_t, K_VEC_SIZE>(left_cache_vec,
                                   &key_cache[left_tgt_cache_idx]);
        Store<uint8_t, K_VEC_SIZE>(right_cache_vec,
                                   &key_cache[right_tgt_cache_idx]);
      }
    } else {
      constexpr int K_VEC_SIZE = 4;
      constexpr int HALF_K_VEC_SIZE = 2;
      using LoadKVResT = AlignedVector<uint8_t, K_VEC_SIZE>;
      using LoadIntT = AlignedVector<InT, HALF_K_VEC_SIZE>;
      using LoadBiasT = AlignedVector<T, HALF_K_VEC_SIZE>;
      using LoadOutScaleT = AlignedVector<float, HALF_K_VEC_SIZE>;
      using LoadScaleT = AlignedVector<T, HALF_K_VEC_SIZE>;
      LoadIntT src_vec1, src_vec2;
      LoadBiasT bias_vec1, bias_vec2;
      LoadOutScaleT out_scale_vec1, out_scale_vec2;
      LoadScaleT scale_vec1, scale_vec2;
      LoadScaleT zp_vec1, zp_vec2;

      const InT* qkv_now = quant_qkv + token_id * hidden_size;
      const int head_bias = lane_id / 4 * 16 + lane_id % 4 * 2;
      const uint32_t cache_idx = kv_head_idx * HeadDim + head_bias;
      const int bias_idx = head_idx * HeadDim + head_bias;
      Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx], &src_vec1);
      Load<InT, HALF_K_VEC_SIZE>(&qkv_now[bias_idx + 8], &src_vec2);
      if (qkv_biases) {
        Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx], &bias_vec1);
        Load<T, HALF_K_VEC_SIZE>(&qkv_biases[bias_idx + 8], &bias_vec2);
      }
      Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx], &out_scale_vec1);
      Load<float, HALF_K_VEC_SIZE>(&qkv_out_scales[bias_idx + 8],
                                   &out_scale_vec2);

      Load<T, HALF_K_VEC_SIZE>(&cache_v_scales[cache_idx], &scale_vec1);
      Load<T, HALF_K_VEC_SIZE>(&cache_v_scales[cache_idx + 8], &scale_vec2);
      Load<T, HALF_K_VEC_SIZE>(&cache_v_zero_points[cache_idx], &zp_vec1);
      Load<T, HALF_K_VEC_SIZE>(&cache_v_zero_points[cache_idx + 8], &zp_vec2);

      float input_left = static_cast<float>(src_vec1[0]);
      float input_right = static_cast<float>(src_vec1[1]);
      input_left = qkv_biases ? input_left * out_scale_vec1[0] +
                                    static_cast<float>(bias_vec1[0])
                              : input_left * out_scale_vec1[0];
      input_right = qkv_biases ? input_right * out_scale_vec2[1] +
                                     static_cast<float>(bias_vec1[1])
                               : input_right * out_scale_vec2[1];

      bias_vec1[0] = static_cast<T>(input_left);
      bias_vec1[1] = static_cast<T>(input_right);

      input_left = static_cast<float>(src_vec2[0]);
      input_right = static_cast<float>(src_vec2[1]);
      input_left = qkv_biases ? input_left * out_scale_vec1[0] +
                                    static_cast<float>(bias_vec2[0])
                              : input_left * out_scale_vec1[0];
      input_right = qkv_biases ? input_right * out_scale_vec1[1] +
                                     static_cast<float>(bias_vec2[1])
                               : input_right * out_scale_vec1[1];

      bias_vec2[0] = static_cast<T>(input_left);
      bias_vec2[1] = static_cast<T>(input_right);

      const uint32_t base_tgt_cache_idx =
          block_idx * gqa_group_size * HeadDim * half_block_size +
          kv_head_idx * HeadDim * half_block_size +
          (lane_id / 4 * 16 + lane_id % 4 * 2) * half_block_size +
          block_offset / 16 % 4 / 2 * 8 * half_block_size +
          block_offset / 16 / 4 * 32 + block_offset / 16 % 2 * 16;
      const uint32_t tgt_cache_idx1 = base_tgt_cache_idx +
                                      block_offset % 8 / 2 * 4     // per 4
                                      + block_offset % 16 / 8 * 2  // per 2
                                      + block_offset % 2;          // per 1
      const uint32_t tgt_cache_idx2 = tgt_cache_idx1 + half_block_size;

      float quant_value1 =
          static_cast<float>(scale_vec1[0] * bias_vec1[0] + zp_vec1[0]);
      float quant_value2 =
          static_cast<float>(scale_vec2[0] * bias_vec2[0] + zp_vec2[0]);
      if constexpr (RoundType == 0) {
        quant_value1 = roundWithTiesToEven(quant_value1);
        quant_value2 = roundWithTiesToEven(quant_value2);
      } else {
        quant_value1 = round(quant_value1);
        quant_value2 = round(quant_value2);
      }
      quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
      quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
      quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
      quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
      uint8_t uint_quant_value1 = static_cast<uint8_t>(quant_value1 + 8.0f);
      uint8_t uint_quant_value2 = static_cast<uint8_t>(quant_value2 + 8.0f);
      value_cache[tgt_cache_idx1] =
          (uint_quant_value2 << 4) | (uint_quant_value1 & 0x0F);

      quant_value1 =
          static_cast<float>(scale_vec1[1] * bias_vec1[1] + zp_vec1[1]);
      quant_value2 =
          static_cast<float>(scale_vec2[1] * bias_vec2[1] + zp_vec2[1]);
      if constexpr (RoundType == 0) {
        quant_value1 = roundWithTiesToEven(quant_value1);
        quant_value2 = roundWithTiesToEven(quant_value2);
      } else {
        quant_value1 = round(quant_value1);
        quant_value2 = round(quant_value2);
      }
      quant_value1 = quant_value1 > max_bound ? max_bound : quant_value1;
      quant_value1 = quant_value1 < min_bound ? min_bound : quant_value1;
      quant_value2 = quant_value2 > max_bound ? max_bound : quant_value2;
      quant_value2 = quant_value2 < min_bound ? min_bound : quant_value2;
      uint_quant_value1 = static_cast<uint8_t>(quant_value1 + 8.0f);
      uint_quant_value2 = static_cast<uint8_t>(quant_value2 + 8.0f);
      value_cache[tgt_cache_idx2] =
          (uint_quant_value2 << 4) | (uint_quant_value1 & 0x0F);
    }
  }
}