forked from EngineX-Cambricon/enginex-mlu370-vllm
44 lines
2.1 KiB
Plaintext
44 lines
2.1 KiB
Plaintext
/*************************************************************************
|
|
* 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_GENERATE_ALIBI_SLOPE_MLUH_
|
|
#define CSRC_KERNELS_GENERATE_ALIBI_SLOPE_MLUH_
|
|
|
|
#include "cnnl.h"
|
|
#include "kernel_utils.h"
|
|
namespace tmo {
|
|
/**
|
|
* @brief Generate causal mask for context satge.
|
|
* @param queue: The queue for mlu.
|
|
* @param alibi_slopes: Output. Pointer to the MLU memory that stores the output, the shape must be
|
|
* [batch_num, head_num].
|
|
* @param true_seq_lens: Input. Pointer to the MLU memory that stores the actual sequence length of
|
|
* each batch, the shape must be [batch_num].
|
|
* @param batch_num: Batch number.
|
|
* @param head_start: The index of first head.
|
|
* @param head_num: Head number in this card.
|
|
* @param head_num_total: Total head number in all cards.
|
|
* @param max_sequence_length: The maximum sequence length used during training.
|
|
* @param use_dynamic: A boolean value indicates whether to use dynamic NTK.
|
|
*/
|
|
KernelStatus invokeGenerateAlibiSlope(cnrtQueue_t queue,
|
|
void *alibi_slopes,
|
|
void *true_seq_lens,
|
|
int batch_num,
|
|
int head_start,
|
|
int head_num,
|
|
int head_num_total,
|
|
int max_sequence_length,
|
|
bool use_dynamic);
|
|
} // namespace tmo
|
|
|
|
#endif // CSRC_KERNELS_GENERATE_ALIBI_SLOPE_MLUH_
|