/*************************************************************************
 * 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.
 *************************************************************************/
#ifndef CSRC_KERNELS_DEQUANT_FROM_LINEAR_CACHE_MLUH_
#define CSRC_KERNELS_DEQUANT_FROM_LINEAR_CACHE_MLUH_

#include "cnnl.h"
#include "kernel_utils.h"

namespace tmo {
/**
 * @brief De-quantizes the key and value tensors from the provided linear cache and scale.
 * @param queue: The queue for mlu.
 * @param key: Pointer to the MLU memory that stores the key tensor,
 *        with shape [total_seqlen, head_num, head_size]. Data type can be float32, half,
 *        or bfloat16. This parameter can be nullptr.
 * @param value: Pointer to the MLU memory that stores the value tensor,
 *        with shape [total_seqlen, head_num, head_size]. Data type can be float32,
 *        half, or bfloat16.This parameter can be nullptr.
 * @param key_cache: Pointer to the MLU memory that stores the key cache tensor,
 *        with shape [max_batch_size, head_num, cache_mem_len, head_size] for 8-bit quantization
 *        or [max_bs, head_num, cache_mem_len, head_size//2] for 4-bit quantization.
 *        Data type must be int8. This parameter can be nullptr.
 * @param value_cache: Pointer to the MLU memory that stores the value cache tensor,
 *        with shape [max_batch_size, head_num, cache_mem_len, head_size] for 8-bit quantization
 *        or [max_bs, head_num, cache_mem_len//2, head_size] for 4-bit quantization.
 *        Data type must be int8. This parameter can be nullptr.
 * @param key_scale: Pointer to the MLU memory that stores the key cache quantization scale.
 *        Shape depends on quantization mode:
 *        - For per-channel quantization (quant_mode = 0): [head_num, head_size].
 *        - For per-token quantization (quant_mode = 1): [max_batch_size, head_num, cache_mem_len].
 *        Data type must be float32. This parameter can be nullptr.
 * @param value_scale: Pointer to the MLU memory that stores the value cache quantization scale,
 *        with the same shape as key_scale. Data type must be float32. This parameter can be
 nullptr.
 * @param context_lens: Pointer to the MLU memory that stores the sequence lengths.
 *        The shape must be [batch].
 * @param context_seq_offset: Pointer to the MLU memory that stores the sequence offset in the
 context.
 *        The shape must be [batch]. If nullptr, the default value is the cumulative sum of
 context_lengths.
 * @param cache_bs_id: Pointer to the MLU memory that stores the batch index in the cache.
 *        The shape must be [batch]. If nullptr, the default value is {0, 1, 2, ..., batch - 1}.
 * @param cache_seq_offset: Pointer to the MLU memory that stores the sequence offset in the cache.
 *        The shape must be [batch]. If nullptr, the default value is 0 for every batch.
 * @param max_contxt_len: The maximum sequence length of context.
 * @param batch: Batch size.
 * @param head_num: Head number.
 * @param key_group_num: group number of key group-wise quantization.
 * @param value_group_num: group number of value group-wise quantization.
 * @param cache_mem_len: The maximum sequence length of cache.
 * @param head_size: Head size.
 * @param quant_mode: An integer value indicating the quantization mode:
 *        0 for per-channel quantization and 1 for per-token quantization.
 * @param quant_bit: An integer value indicating the quantization bit width:
 *        8 for 8-bit quantization and 4 for 4-bit quantization.
 * @param contxt_head_stride: The stride of head_num in context.
 * @param contxt_seq_stride: The stride of max_contxt_len in context.
 * @param cache_bs_stride: The stride of batch in cache.
 * @param cache_head_stride: The stride of head_num in cache.
 * @param key_cache_seq_stride: The stride of cache_mem_len in key cache.
 * @param value_cache_seq_stride: The stride of cache_mem_len in value cache.
 * @param cache_scale_bs_stride: The stride of batch in cache scale, only valid if quant_per_quant.
 * @param cache_scale_head_stride: The stride of head in cache scale.
 * @param dtype: The data type of the key and value tensors.
 * @note If any of key/key_cache/key_scale is nullptr, no operation is performed on the key.
 *       If any of value/value_cache/value_scale is nullptr, no operation is performed on the value.
 */
KernelStatus invokeDequantFromLinearCache(cnrtQueue_t queue,
                                          void *key,
                                          void *value,
                                          const void *key_cache,
                                          const void *value_cache,
                                          const void *key_scale,
                                          const void *value_scale,
                                          const void *context_lens,
                                          const void *context_seq_offsets,
                                          const void *cache_bs_ids,
                                          const void *cache_seq_offsets,
                                          const int32_t max_context_len,
                                          const int32_t batch,
                                          const int32_t head_num,
                                          const int32_t key_group_num,
                                          const int32_t value_group_num,
                                          const int32_t cache_mem_len,
                                          const int32_t head_size,
                                          const int32_t quant_mode,
                                          const int32_t quant_bit,
                                          const size_t context_head_stride,
                                          const size_t context_seq_stride,
                                          const size_t cache_bs_stride,
                                          const size_t cache_head_stride,
                                          const size_t key_cache_seq_stride,
                                          const size_t value_cache_seq_stride,
                                          const size_t cache_scale_bs_stride,
                                          const size_t cache_scale_head_stride,
                                          const cnnlDataType_t dtype);
}  // namespace tmo

#endif  // CSRC_KERNELS_DEQUANT_FROM_LINEAR_CACHE_MLUH_
