/*************************************************************************
 * 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_OFFLINE_QUANT_TO_PAGED_CACHE_MLUH_
#define CSRC_KERNELS_OFFLINE_QUANT_TO_PAGED_CACHE_MLUH_

#include "kernel_utils.h"
namespace tmo {
/**
 * @brief Perform offline_quant_to_paged_cache operation.
 * @param queue[in]: The queue for mlu.
 * @param data_type[in]: The cnnl data type of key.
 * @param key[in]: Pointer to the MLU memory that stores the key tensor which has shape [num_tokens,
 * num_heads, head_size]. Data type of key must be half/bfloat16_t/float.
 * @param value[in]: Pointer to the MLU memory that stores the value tensor which has shape
 * [num_tokens, num_heads, head_size]. Data type of value must be half/bfloat16_t/float.
 * @param key_cache[out]: Pointer to the MLU memory that stores the key_cache tensor which has
 * shape [num_blocks, num_heads, block_size, head_size]. Data type of key cache must be int8_t.
 * @param value_cache[out]: Pointer to the MLU memory that stores the value_cache tensor which has
 * shape [num_blocks, num_heads, block_size, head_size]. Data type of value cache must be int8_t.
 * @param key_cache_scale[in]: Pointer to the MLU memory that stores the key_cache_scale tensor
 * which has shape [num_heads, head_size]. Data type of key cache scale must be float.
 * @param value_cache_scale[in]: Pointer to the MLU memory that stores the value_cache_scale tensor
 * which has shape [num_heads, head_size]. Data type of value cache scale must be float.
 * @param slot_mapping[in]: Pointer to the MLU memory that stores the slot_mapping tensor which has
 * shape [num_tokens]. Data type of slot mapping must be int32_t.
 * @param key_stride0[in]: The first dimension stride length of key_cache tensor.
 * @param value_stride0[in]: The first dimension stride length of value_cache tensor.
 * @param num_tokens[in]: Total number of tokens.
 * @param num_heads[in]: Head number.
 * @param block_num[in]: Total number of blocks.
 * @param block_size[in]: Number of tokens per block.
 * @param head_size[in]: Head size.
 * @note: offline_quant_to_paged_cache does not support MLU300 device.
 */
KernelStatus invokeOfflineQuantToPagedCache(cnrtQueue_t queue,
                                            cnnlDataType_t data_type,
                                            void *key,
                                            void *value,
                                            void *key_cache,
                                            void *value_cache,
                                            void *key_cache_scale,
                                            void *value_cache_scale,
                                            void *slot_mapping,
                                            size_t key_stride0,
                                            size_t value_stride0,
                                            int num_tokens,
                                            int num_heads,
                                            int block_num,
                                            int block_size,
                                            int head_size);
}  // namespace tmo

#endif  // CSRC_KERNELS_OFFLINE_QUANT_TO_PAGED_CACHE_MLUH_
