forked from EngineX-Cambricon/enginex-mlu370-vllm
659 lines
33 KiB
Plaintext
659 lines
33 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.
|
||
*************************************************************************/
|
||
#include "cnnl.h"
|
||
#include "cnrt.h"
|
||
#include "fused_rope.mluh"
|
||
// clang-format off
|
||
#include <mlu.h>
|
||
// clang-format on
|
||
|
||
using bf16 = bfloat16_t;
|
||
|
||
namespace tmo {
|
||
|
||
namespace kernels {
|
||
|
||
#ifndef PAD_UP
|
||
#define PAD_UP(x, y) ((x) / (y) + (int)((x) % (y) > 0) * (y))
|
||
#endif
|
||
|
||
#if __BANG_ARCH__ > 500
|
||
#include <bang_fusor.h>
|
||
template <typename T>
|
||
using bang_cycle_fusor = bang::experimental::cycle_fusor<T>;
|
||
#endif
|
||
|
||
#define NRAM_BUFFER_SIZE (480 * 1024)
|
||
__nram__ int8_t nram_buffer[NRAM_BUFFER_SIZE];
|
||
__nram__ float nram_mask[256];
|
||
__nram__ int nram_rope_offsets[256];
|
||
__nram__ float nram_zeros[1024] = {0.f};
|
||
|
||
template <typename T>
|
||
__mlu_func__ void toFloat(float *dst, T *src, int num) {
|
||
if (std::is_same<T, half>::value) {
|
||
__bang_half2float(dst, (half *)src, num);
|
||
} else if (std::is_same<T, bf16>::value) {
|
||
#if __BANG_ARCH__ > 500
|
||
__bang_bfloat162float(dst, (bf16 *)src, num);
|
||
#endif
|
||
}
|
||
}
|
||
|
||
template <typename T>
|
||
__mlu_func__ void floatTo(T *dst, float *src, int num) {
|
||
if (std::is_same<T, half>::value) {
|
||
__bang_float2half_rn((half *)dst, src, num);
|
||
} else if (std::is_same<T, bf16>::value) {
|
||
__bang_float2bfloat16_rn((bf16 *)dst, src, num);
|
||
}
|
||
}
|
||
|
||
__mlu_func__ void genScatterOffsetMask(int *cache_bs_id_begin,
|
||
int *cache_seq_offsets_begin,
|
||
int *slot_mapping_begin,
|
||
int *nram_k_cache_offsets,
|
||
int *nram_v_cache_offsets,
|
||
int *nram_v_onchip_offsets,
|
||
int *nram_kv_scale_offsets,
|
||
float *nram_cache_mask,
|
||
float *nram_zeros,
|
||
float *nram_temp,
|
||
int task_deal_batch,
|
||
int task_begin_batch,
|
||
int head_num_k,
|
||
int head_size,
|
||
int max_decode_len,
|
||
int block_size,
|
||
int kv_out_size,
|
||
int group_num,
|
||
bool discrete_batch,
|
||
bool paged_cache,
|
||
bool mixed_cache) {
|
||
// 目前先用标量化计算offset，便于理解(性能无影响)
|
||
int bh = task_deal_batch * head_num_k;
|
||
if (paged_cache) {
|
||
int cache_seq_stride = head_size;
|
||
int cache_head_stride = block_size * head_size;
|
||
int cache_scale_head_stride = block_size * group_num;
|
||
int cache_block_stride = head_num_k * cache_head_stride;
|
||
int cache_scale_block_stride = head_num_k * cache_scale_head_stride;
|
||
int *nram_slot_mapping = (int *)nram_mask;
|
||
__memcpy(nram_slot_mapping, slot_mapping_begin, task_deal_batch * sizeof(int), GDRAM2NRAM);
|
||
for (int i = 0; i < task_deal_batch; i++) {
|
||
int mapping_idx = __load_nram(nram_slot_mapping + i);
|
||
if (mapping_idx < 0) {
|
||
__bang_write_value(nram_k_cache_offsets + i * head_num_k, head_num_k, (int)-1);
|
||
continue;
|
||
}
|
||
int block_idx = mapping_idx / block_size;
|
||
int seq_idx = mapping_idx % block_size;
|
||
int k_seq_offset = block_idx * cache_block_stride + seq_idx * cache_seq_stride;
|
||
int v_seq_offset = block_idx * cache_block_stride / 2 + seq_idx / 2 * cache_seq_stride;
|
||
int scale_seq_offset = block_idx * cache_scale_block_stride + seq_idx * group_num;
|
||
int onchip_offset = i * head_num_k * head_size + seq_idx % 2 * bh * head_size;
|
||
|
||
for (int j = 0; j < head_num_k; j++) {
|
||
__store_nram(
|
||
nram_k_cache_offsets + i * head_num_k + j,
|
||
(int)((k_seq_offset + j * cache_head_stride) * kv_out_size / (mixed_cache + 1)));
|
||
if (mixed_cache) {
|
||
__store_nram(nram_v_cache_offsets + i * head_num_k + j,
|
||
(int)(v_seq_offset + j * cache_head_stride / 2));
|
||
__store_nram(nram_kv_scale_offsets + i * head_num_k + j,
|
||
(int)((scale_seq_offset + j * cache_scale_head_stride) * sizeof(float)));
|
||
__store_nram(nram_v_onchip_offsets + i * head_num_k + j, onchip_offset + j * head_size);
|
||
}
|
||
}
|
||
}
|
||
} else {
|
||
int *nram_seq_offsets = (int *)nram_mask;
|
||
int *nram_bs_id = nram_seq_offsets + 32;
|
||
int cache_seq_stride = head_size;
|
||
int cache_head_stride = max_decode_len * head_size;
|
||
int cache_scale_head_stride = max_decode_len * group_num;
|
||
int cache_bs_stride = head_num_k * cache_head_stride;
|
||
int cache_scale_bs_stride = head_num_k * cache_scale_head_stride;
|
||
__memcpy(nram_seq_offsets, cache_seq_offsets_begin, task_deal_batch * sizeof(int), GDRAM2NRAM);
|
||
if (discrete_batch) {
|
||
__memcpy(nram_bs_id, cache_bs_id_begin, task_deal_batch * sizeof(int), GDRAM2NRAM);
|
||
}
|
||
for (int i = 0; i < task_deal_batch; i++) {
|
||
int bs_idx = __load_nram(nram_bs_id + i);
|
||
int seq_idx = __load_nram(nram_seq_offsets + i);
|
||
int temp_bs_idx = discrete_batch ? bs_idx : task_begin_batch + i;
|
||
int temp_seq_idx = seq_idx;
|
||
bool masked = temp_bs_idx < 0 || temp_seq_idx < 0;
|
||
if (masked) {
|
||
__bang_write_value(nram_k_cache_offsets + i * head_num_k, head_num_k, (int)-1);
|
||
continue;
|
||
}
|
||
int k_seq_offset = temp_bs_idx * cache_bs_stride + temp_seq_idx * cache_seq_stride;
|
||
int scale_seq_offset = temp_bs_idx * cache_scale_bs_stride + temp_seq_idx * group_num;
|
||
int v_seq_offset = temp_bs_idx * cache_bs_stride / 2 + temp_seq_idx / 2 * cache_seq_stride;
|
||
int onchip_offset = i * head_num_k * head_size + temp_seq_idx % 2 * bh * head_size;
|
||
|
||
for (int j = 0; j < head_num_k; j++) {
|
||
__store_nram(
|
||
nram_k_cache_offsets + i * head_num_k + j,
|
||
(int)((k_seq_offset + j * cache_head_stride) * kv_out_size / (mixed_cache + 1)));
|
||
if (mixed_cache) {
|
||
__store_nram(nram_v_cache_offsets + i * head_num_k + j,
|
||
(int)(v_seq_offset + j * cache_head_stride / 2));
|
||
__store_nram(nram_kv_scale_offsets + i * head_num_k + j,
|
||
(int)((scale_seq_offset + j * cache_scale_head_stride) * sizeof(float)));
|
||
__store_nram(nram_v_onchip_offsets + i * head_num_k + j, onchip_offset + j * head_size);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 此处是为了做上scatter指令的 mask，如果bs offset或seq offset小于0则需要mask掉
|
||
__bang_int322float(nram_temp, nram_k_cache_offsets, bh, 0);
|
||
__bang_ge_bitindex(nram_cache_mask, nram_temp, nram_zeros, PAD_UP(bh, 8));
|
||
}
|
||
|
||
__mlu_func__ void layernormImpl(float *nram_k,
|
||
float *norm_params,
|
||
int task_deal_batch,
|
||
int k_hidden,
|
||
float eps) {
|
||
#if __BANG_ARCH__ > 500
|
||
float *buffer = nram_k + task_deal_batch * k_hidden;
|
||
for (int i = 0; i < task_deal_batch; i++) {
|
||
float *k_ = nram_k + i * k_hidden;
|
||
__bang_mul(buffer, k_, k_, k_hidden);
|
||
float mean = __bang_sum(k_, k_hidden);
|
||
mean = mean / k_hidden;
|
||
float rstd = __bang_sum(buffer, k_hidden);
|
||
rstd = rstd / k_hidden - mean * mean;
|
||
rstd = rstd < 0 ? eps : rstd + eps;
|
||
rstd = 1.f / std::sqrt(rstd);
|
||
__bang_fusion(FUSION_FSM, k_, k_, mean, rstd, k_hidden);
|
||
}
|
||
__bang_fusion(FUSION_FMA, nram_k, nram_k, norm_params, norm_params + k_hidden,
|
||
task_deal_batch * k_hidden, k_hidden);
|
||
#endif
|
||
}
|
||
|
||
__mlu_func__ void foldRotaryImpl(float *nram_qk,
|
||
float *nram_qk_rot,
|
||
float *nram_table,
|
||
int task_deal_batch,
|
||
int head_num_qk,
|
||
int head_size) {
|
||
int rotary_low_dim = task_deal_batch * head_size;
|
||
__bang_cycle_mul(nram_qk, nram_qk, nram_table, head_num_qk * rotary_low_dim, rotary_low_dim);
|
||
__bang_cycle_mul(nram_qk_rot, nram_qk_rot, nram_mask, head_num_qk * rotary_low_dim, head_size);
|
||
__bang_cycle_mul(nram_qk_rot, nram_qk_rot, nram_table + task_deal_batch * head_size,
|
||
head_num_qk * rotary_low_dim, rotary_low_dim);
|
||
__bang_add(nram_qk, nram_qk, nram_qk_rot, head_num_qk * rotary_low_dim);
|
||
}
|
||
|
||
template <typename T>
|
||
__mlu_func__ void quantify(T *input,
|
||
float *float_input,
|
||
void *output_hp,
|
||
void *output_lp,
|
||
float *nram_trans,
|
||
float *scale_hp,
|
||
float *scale_lp,
|
||
float *scale_lp_temp,
|
||
int batch,
|
||
int head_num,
|
||
int head_size,
|
||
int group_num,
|
||
int group_size,
|
||
bool quant_kv_hp,
|
||
bool mixed_cache) {
|
||
if (quant_kv_hp) {
|
||
int hidden = head_num * head_size;
|
||
int bh = batch * head_num;
|
||
toFloat<T>(float_input, input, batch * hidden);
|
||
__bang_recip(scale_hp, scale_hp, hidden);
|
||
#if __BANG_ARCH__ > 500
|
||
__asm__ __volatile__(
|
||
"fuse.nram.crn.s8.f32 "
|
||
"[%[dst]], %[num_long], %[num_short], [%[src0]], .mul.cycle([%[src1]]), .dstpos(%[pos])"
|
||
";\n\t" ::[dst] "r"(output_hp),
|
||
[num_long] "r"(batch * hidden), [num_short] "r"(hidden), [src0] "r"(float_input),
|
||
[src1] "r"(scale_hp), [pos] "i"(0));
|
||
#endif
|
||
if (mixed_cache) {
|
||
__bang_transpose(nram_trans, float_input, bh * group_num, group_size);
|
||
__bang_abs(float_input, nram_trans, bh * head_size);
|
||
__bang_maxpool(scale_lp_temp, float_input, bh * group_num, group_size, 1, group_size, 1, 1,
|
||
1);
|
||
__bang_mul_scalar(scale_lp, scale_lp_temp, 1 / 7.f, bh * group_num);
|
||
__bang_recip(scale_lp_temp, scale_lp, bh * group_num);
|
||
__bang_cycle_mul(nram_trans, nram_trans, scale_lp_temp, bh * group_num * group_size,
|
||
bh * group_num);
|
||
__bang_float2int8_rn((int8_t *)nram_trans, nram_trans, bh * head_size, 0);
|
||
__bang_transpose((int8_t *)output_lp, (int8_t *)nram_trans, group_size, bh * group_num);
|
||
}
|
||
}
|
||
}
|
||
|
||
template <typename T>
|
||
__mlu_func__ void fuseRopeImpl(T *input,
|
||
void *key_cache_hp,
|
||
void *value_cache_hp,
|
||
void *key_cache_lp,
|
||
void *value_cache_lp,
|
||
T *sin_table,
|
||
T *cos_table,
|
||
int *rope_offsets,
|
||
T *gamma,
|
||
T *beta,
|
||
float *key_scale_hp,
|
||
float *value_scale_hp,
|
||
float *key_scale_lp,
|
||
float *value_scale_lp,
|
||
int *cache_bs_id_hp,
|
||
int *cache_seq_offsets_hp,
|
||
int *cache_bs_id_lp,
|
||
int *cache_seq_offsets_lp,
|
||
int *slot_mapping_hp,
|
||
int *slot_mapping_lp,
|
||
int rotary_stride,
|
||
int task_deal_batch,
|
||
int task_begin_batch,
|
||
int head_num_q,
|
||
int head_num_k,
|
||
int head_size,
|
||
int max_decode_len_hp,
|
||
int max_decode_len_lp,
|
||
int block_size_hp,
|
||
int block_size_lp,
|
||
int group_size,
|
||
int batch_cap,
|
||
float eps) {
|
||
#if __BANG_ARCH__ > 500
|
||
/*
|
||
由于需要支持mixed cache，kernel支持cache的功能组合比较多，现规定只存在以下几种：
|
||
1.只存在hp_cache的情况(通过lp tensor不为0判断)，cache支持bf16，fp16，量化下支持离线perchannel int8
|
||
支持linear和paged，key和value cache形状一致，key/value_scale_hp形状为[head_num, head_size]
|
||
2.mixed cache的情况hp支持离线perchannel int8量化，支持linear和paged，key和value cache形状一致，
|
||
key/value_scale_hp 形状为[head_num, head_size]. lp支持int4在线pertoken group量化，
|
||
key_cache形状为 [batch, head_num_k, max_decode_len_lp, head_size / 2]
|
||
paged情况也是head_size / 2, value_cache的形状为[batch, head_num_l, max_decode_len_lp / 2,
|
||
head_size]，paged cache形状为 [num_blocks, head_num_k, block_size / 2,
|
||
head_size]，key/value_scale_lp形状为 [batch, head_num_k, max_decode_len_lp,
|
||
group_num]，paged_cache 为 [num_blocks, head_num_k, block_size, group_num]
|
||
*/
|
||
bool mixed_cache = key_cache_lp != nullptr && value_cache_lp != nullptr;
|
||
bool quant_kv_hp = key_scale_hp != nullptr && value_scale_hp != nullptr;
|
||
bool discrete_batch_hp = cache_bs_id_hp != nullptr;
|
||
bool discrete_batch_lp = cache_bs_id_lp != nullptr;
|
||
bool paged_cache_hp = slot_mapping_hp != nullptr;
|
||
bool paged_cache_lp = slot_mapping_lp != nullptr;
|
||
int head_num_qk = head_num_q + head_num_k;
|
||
int head_num_qkv = head_num_q + head_num_k * 2;
|
||
int qkv_hidden = head_num_qkv * head_size;
|
||
int qk_hidden = head_num_qk * head_size;
|
||
int q_hidden = head_num_q * head_size;
|
||
int k_hidden = head_num_k * head_size;
|
||
int float_size = sizeof(float);
|
||
int dtype_size = sizeof(T);
|
||
int kv_size_hp = quant_kv_hp ? sizeof(int8_t) : dtype_size;
|
||
int group_num = mixed_cache ? head_size / group_size : 1;
|
||
|
||
// task ddr offset
|
||
T *input_begin = input + task_begin_batch * qkv_hidden;
|
||
int *cache_bs_id_begin_hp = cache_bs_id_hp + task_begin_batch;
|
||
int *cache_seq_offsets_begin_hp = cache_seq_offsets_hp + task_begin_batch;
|
||
int *slot_mapping_begin_hp = slot_mapping_hp + task_begin_batch;
|
||
|
||
// nram_buffer
|
||
float *nram_qk = (float *)nram_buffer;
|
||
float *nram_qk_rot = nram_qk + batch_cap * qk_hidden;
|
||
float *nram_v = nram_qk_rot + batch_cap * qk_hidden;
|
||
float *nram_kv_trans = nram_v + batch_cap * k_hidden;
|
||
float *nram_table = nram_kv_trans + (int)mixed_cache * batch_cap * k_hidden;
|
||
float *norm_params = nram_table + 2 * batch_cap * head_size;
|
||
float *nram_k_scale_hp = norm_params + 2 * head_size;
|
||
float *nram_v_scale_hp = nram_k_scale_hp + (int)quant_kv_hp * k_hidden;
|
||
float *nram_k_scale_lp = nram_v_scale_hp + (int)quant_kv_hp * k_hidden;
|
||
float *nram_v_scale_lp = nram_k_scale_lp + (int)mixed_cache * batch_cap * head_num_k * group_num;
|
||
int8_t *nram_kv_hp =
|
||
(int8_t *)(nram_v_scale_lp + (int)mixed_cache * batch_cap * head_num_k * group_num);
|
||
int8_t *nram_kv_lp = nram_kv_hp + (int)quant_kv_hp * batch_cap * k_hidden;
|
||
int8_t *nram_cache_v = nram_kv_lp + (int)mixed_cache * batch_cap * k_hidden;
|
||
int *nram_kv_cache_offsets_hp =
|
||
(int *)(nram_cache_v + (int)mixed_cache * batch_cap * k_hidden * 2);
|
||
int *nram_k_cache_offsets_lp = nram_kv_cache_offsets_hp + batch_cap * head_num_k;
|
||
int *nram_v_cache_offsets_lp =
|
||
nram_k_cache_offsets_lp + (int)mixed_cache * batch_cap * head_num_k;
|
||
int *nram_kv_scale_offsets = nram_v_cache_offsets_lp + (int)mixed_cache * batch_cap * head_num_k;
|
||
int *nram_v_onchip_offsets = nram_kv_scale_offsets + (int)mixed_cache * batch_cap * head_num_k;
|
||
float *cache_mask_hp =
|
||
(float *)(nram_v_onchip_offsets + (int)mixed_cache * batch_cap * head_num_k);
|
||
float *cache_mask_lp = (float *)((int8_t *)cache_mask_hp + PAD_UP(batch_cap * head_num_k, 8) / 8);
|
||
|
||
// 这里将qk和qk_rot放在一起是为升位宽可以一起做，减少指令，同样还有sincostable和norm的gamma和beta
|
||
T *qk_in = (T *)nram_qk_rot;
|
||
T *qk_rot_in = (T *)((int8_t *)nram_qk_rot + (float_size - dtype_size) * batch_cap * qk_hidden);
|
||
T *v_in =
|
||
(T *)((int8_t *)nram_v + (int)quant_kv_hp * (float_size - dtype_size) * batch_cap * k_hidden);
|
||
T *norm_params_in = (T *)((int8_t *)norm_params + (float_size - dtype_size) * 2 * head_size);
|
||
T *table_in = (T *)((int8_t *)nram_table + (float_size - dtype_size) * 2 * batch_cap * head_size);
|
||
int8_t *nram_cache_v_in = nram_cache_v + batch_cap * k_hidden * sizeof(int8_t);
|
||
|
||
// 生成 kv cache的offset和mask，供scatter kv到kvcache使用
|
||
genScatterOffsetMask(cache_bs_id_begin_hp, cache_seq_offsets_begin_hp, slot_mapping_begin_hp,
|
||
nram_kv_cache_offsets_hp, nullptr, nullptr, nullptr, cache_mask_hp,
|
||
nram_zeros, nram_qk, task_deal_batch, task_begin_batch, head_num_k,
|
||
head_size, max_decode_len_hp, block_size_hp, kv_size_hp, 1,
|
||
discrete_batch_hp, paged_cache_hp, false);
|
||
if (mixed_cache) {
|
||
int *cache_bs_id_begin_lp = cache_bs_id_lp + task_begin_batch;
|
||
int *cache_seq_offsets_begin_lp = cache_seq_offsets_lp + task_begin_batch;
|
||
int *slot_mapping_begin_lp = slot_mapping_lp + task_begin_batch;
|
||
genScatterOffsetMask(cache_bs_id_begin_lp, cache_seq_offsets_begin_lp, slot_mapping_begin_lp,
|
||
nram_k_cache_offsets_lp, nram_v_cache_offsets_lp, nram_v_onchip_offsets,
|
||
nram_kv_scale_offsets, cache_mask_lp, nram_zeros, nram_qk, task_deal_batch,
|
||
task_begin_batch, head_num_k, head_size, max_decode_len_lp, block_size_lp,
|
||
1, group_num, discrete_batch_lp, paged_cache_lp, mixed_cache);
|
||
}
|
||
|
||
/*
|
||
-----------------------
|
||
load v |
|
||
-----------------------
|
||
load qk | quant v
|
||
-----------------------
|
||
store v | rope qk
|
||
-----------------------
|
||
store_q | layernorm k
|
||
| quant k
|
||
-----------------------
|
||
store k |
|
||
*/
|
||
// prepare v v_scale cache_v rope_offset
|
||
__memcpy_async(v_in, input_begin + qk_hidden, k_hidden * dtype_size, GDRAM2NRAM,
|
||
k_hidden * dtype_size, qkv_hidden * dtype_size, task_deal_batch - 1);
|
||
if (quant_kv_hp) {
|
||
__memcpy_async(nram_k_scale_hp, key_scale_hp, k_hidden * float_size, GDRAM2NRAM);
|
||
__memcpy_async(nram_v_scale_hp, value_scale_hp, k_hidden * float_size, GDRAM2NRAM);
|
||
}
|
||
__memcpy_async(nram_rope_offsets, rope_offsets + task_begin_batch, task_deal_batch * sizeof(int),
|
||
GDRAM2NRAM);
|
||
|
||
__sync_io();
|
||
if (mixed_cache) {
|
||
__gather(nram_cache_v_in, value_cache_lp, (uint32_t *)nram_v_cache_offsets_lp, cache_mask_lp,
|
||
head_size * sizeof(int8_t), GDRAM2NRAM, head_size * sizeof(int8_t),
|
||
task_deal_batch * head_num_k);
|
||
}
|
||
__bang_mul_scalar(nram_rope_offsets, nram_rope_offsets, rotary_stride * dtype_size,
|
||
task_deal_batch);
|
||
__sync_compute();
|
||
/*==============================================================================================*/
|
||
|
||
// load_qk,rope_table | quant v
|
||
__memcpy_async(qk_in, input_begin, head_size * dtype_size, GDRAM2NRAM, head_size * dtype_size,
|
||
task_deal_batch - 1, task_deal_batch * head_size * dtype_size, head_num_qk - 1,
|
||
qkv_hidden * dtype_size, task_deal_batch - 1, head_size * dtype_size,
|
||
head_num_qk - 1);
|
||
__gather_async(table_in, cos_table, (uint32_t *)nram_rope_offsets, head_size * dtype_size,
|
||
GDRAM2NRAM, head_size * dtype_size, task_deal_batch);
|
||
__gather_async(table_in + task_deal_batch * head_size, sin_table, (uint32_t *)nram_rope_offsets,
|
||
head_size * dtype_size, GDRAM2NRAM, head_size * dtype_size, task_deal_batch);
|
||
__memcpy_async(norm_params_in, gamma, head_size * dtype_size, GDRAM2NRAM);
|
||
__memcpy_async(norm_params_in + head_size, beta, head_size * dtype_size, GDRAM2NRAM);
|
||
|
||
int8_t *nram_temp = (int8_t *)nram_qk;
|
||
if (mixed_cache) {
|
||
__bang_int42int8(nram_cache_v, (int4x2_t *)nram_cache_v_in, task_deal_batch * k_hidden * 2, 0,
|
||
0);
|
||
__bang_transpose(nram_temp, nram_cache_v, task_deal_batch * k_hidden, 2);
|
||
}
|
||
quantify<T>(v_in, nram_v, nram_kv_hp, nram_kv_lp, nram_kv_trans, nram_v_scale_hp, nram_v_scale_lp,
|
||
nram_k_scale_lp /*lp_scale temp*/, task_deal_batch, head_num_k, head_size, group_num,
|
||
group_size, quant_kv_hp, mixed_cache);
|
||
if (mixed_cache) {
|
||
__scatter(nram_temp, nram_kv_lp, (uint32_t *)nram_v_onchip_offsets, cache_mask_lp,
|
||
head_size * sizeof(int8_t), NRAM2NRAM, head_size * sizeof(int8_t),
|
||
task_deal_batch * head_num_k);
|
||
__bang_transpose(nram_cache_v, nram_temp, 2, task_deal_batch * k_hidden);
|
||
__bang_int82int4_rn((int4x2_t *)nram_cache_v, nram_cache_v, task_deal_batch * k_hidden * 2, 0,
|
||
0);
|
||
}
|
||
__sync_io_move_compute();
|
||
/*==============================================================================================*/
|
||
|
||
// rope | store v
|
||
// 将qk的左右部分交换，用于生成qk_rot
|
||
__memcpy(qk_rot_in, qk_in + head_size / 2, head_size / 2 * dtype_size, NRAM2NRAM,
|
||
head_size * dtype_size, head_size * dtype_size, task_deal_batch * head_num_qk - 1);
|
||
__memcpy(qk_rot_in + head_size / 2, qk_in, head_size / 2 * dtype_size, NRAM2NRAM,
|
||
head_size * dtype_size, head_size * dtype_size, task_deal_batch * head_num_qk - 1);
|
||
toFloat<T>(nram_qk, qk_in, 2 * batch_cap * qk_hidden);
|
||
toFloat<T>(nram_table, table_in, 2 * task_deal_batch * head_size);
|
||
toFloat<T>(norm_params, norm_params_in, 2 * head_size);
|
||
__bang_write_value(nram_mask, head_size / 2, (float)-1);
|
||
__bang_write_value(nram_mask + head_size / 2, head_size / 2, (float)1);
|
||
foldRotaryImpl(nram_qk, nram_qk_rot, nram_table, task_deal_batch, head_num_qk, head_size);
|
||
floatTo<T>((T *)nram_qk, nram_qk, task_deal_batch * q_hidden);
|
||
|
||
int8_t *scatter_v_src = quant_kv_hp ? nram_kv_hp : (int8_t *)nram_v;
|
||
__scatter_async(value_cache_hp, scatter_v_src, (uint32_t *)nram_kv_cache_offsets_hp,
|
||
cache_mask_hp, head_size * kv_size_hp, NRAM2GDRAM, head_size * kv_size_hp,
|
||
head_num_k * task_deal_batch);
|
||
if (mixed_cache) {
|
||
__scatter_async(value_cache_lp, nram_cache_v, (uint32_t *)nram_v_cache_offsets_lp,
|
||
cache_mask_lp, head_size * sizeof(int8_t), NRAM2GDRAM,
|
||
head_size * sizeof(int8_t), head_num_k * task_deal_batch);
|
||
__scatter_async(value_scale_lp, nram_v_scale_lp, (uint32_t *)nram_kv_scale_offsets,
|
||
cache_mask_lp, group_num * sizeof(float), NRAM2GDRAM, group_num * sizeof(float),
|
||
head_num_k * task_deal_batch);
|
||
}
|
||
__sync_io_move_compute();
|
||
/*==============================================================================================*/
|
||
// layernrom k quant k | store q
|
||
// 从qk的nram buffer中提取出k做layernorm和量化
|
||
float *nram_k = nram_qk_rot;
|
||
__memcpy(nram_k, nram_qk + task_deal_batch * q_hidden, head_size * float_size, NRAM2NRAM,
|
||
head_size * float_size, head_num_k - 1, k_hidden * float_size, task_deal_batch - 1,
|
||
task_deal_batch * head_size * float_size, head_num_k - 1, head_size * float_size,
|
||
task_deal_batch - 1);
|
||
layernormImpl(nram_k, norm_params, task_deal_batch * head_num_k, head_size, eps);
|
||
quantify<float>(nram_k, nram_k, nram_kv_hp, nram_kv_lp, nram_kv_trans, nram_k_scale_hp,
|
||
nram_k_scale_lp, nram_v_scale_lp /*lp_scale temp*/, task_deal_batch, head_num_k,
|
||
head_size, group_num, group_size, quant_kv_hp, mixed_cache);
|
||
if (mixed_cache) {
|
||
__bang_int82int4_rn((int4x2_t *)nram_kv_lp, nram_kv_lp, task_deal_batch * k_hidden, 0, 0);
|
||
}
|
||
if (!quant_kv_hp) {
|
||
floatTo<T>((T *)nram_k, nram_k, task_deal_batch * k_hidden);
|
||
}
|
||
|
||
// store q
|
||
__memcpy_async(input_begin, nram_qk, head_size * dtype_size, NRAM2GDRAM, qkv_hidden * dtype_size,
|
||
task_deal_batch - 1, head_size * dtype_size, head_num_q - 1,
|
||
head_size * dtype_size, task_deal_batch - 1,
|
||
task_deal_batch * head_size * dtype_size, head_num_q - 1);
|
||
// ===============================================================================================
|
||
|
||
int8_t *scatter_k_src = quant_kv_hp ? nram_kv_hp : (int8_t *)nram_k;
|
||
__scatter(key_cache_hp, scatter_k_src, (uint32_t *)nram_kv_cache_offsets_hp, cache_mask_hp,
|
||
head_size * kv_size_hp, NRAM2GDRAM, head_size * kv_size_hp,
|
||
head_num_k * task_deal_batch);
|
||
if (mixed_cache) {
|
||
__scatter(key_cache_lp, nram_kv_lp, (uint32_t *)nram_k_cache_offsets_lp, cache_mask_lp,
|
||
head_size / 2 * sizeof(int8_t), NRAM2GDRAM, head_size / 2 * sizeof(int8_t),
|
||
head_num_k * task_deal_batch);
|
||
__scatter(key_scale_lp, nram_k_scale_lp, (uint32_t *)nram_kv_scale_offsets, cache_mask_lp,
|
||
group_num * sizeof(float), NRAM2GDRAM, group_num * sizeof(float),
|
||
head_num_k * task_deal_batch);
|
||
}
|
||
__sync_io_move_compute();
|
||
#endif
|
||
}
|
||
|
||
template <typename T>
|
||
__mlu_global__ void MLUFuseRope(T *input,
|
||
void *key_cache_hp,
|
||
void *value_cache_hp,
|
||
void *key_cache_lp,
|
||
void *value_cache_lp,
|
||
T *sin_table,
|
||
T *cos_table,
|
||
int *rope_offsets,
|
||
T *gamma,
|
||
T *beta,
|
||
float *key_scale_hp,
|
||
float *value_scale_hp,
|
||
float *key_scale_lp,
|
||
float *value_scale_lp,
|
||
int *cache_bs_id_hp,
|
||
int *cache_seq_offsets_hp,
|
||
int *cache_bs_id_lp,
|
||
int *cache_seq_offsets_lp,
|
||
int *slot_mapping_hp,
|
||
int *slot_mapping_lp,
|
||
int rotary_stride,
|
||
int batch,
|
||
int head_num_q,
|
||
int head_num_k,
|
||
int head_size,
|
||
int max_decode_len_hp,
|
||
int max_decode_len_lp,
|
||
int block_size_hp,
|
||
int block_size_lp,
|
||
int group_size,
|
||
int batch_cap,
|
||
int task_avg_batch,
|
||
float eps) {
|
||
int task_begin_batch = taskId * task_avg_batch;
|
||
int task_deal_batch = std::min(batch - task_begin_batch, task_avg_batch);
|
||
if (task_deal_batch <= 0 || __is_mpu()) {
|
||
return;
|
||
}
|
||
|
||
int task_loop = (task_deal_batch + batch_cap - 1) / batch_cap;
|
||
int once_batch = (task_deal_batch + task_loop - 1) / task_loop;
|
||
|
||
for (int i = 0; i < task_loop; i++) {
|
||
int cur_batch = std::min(task_deal_batch - i * once_batch, once_batch);
|
||
int batch_offset = task_begin_batch + once_batch * i;
|
||
|
||
fuseRopeImpl<T>(input, key_cache_hp, value_cache_hp, key_cache_lp, value_cache_lp, sin_table,
|
||
cos_table, rope_offsets, gamma, beta, key_scale_hp, value_scale_hp,
|
||
key_scale_lp, value_scale_lp, cache_bs_id_hp, cache_seq_offsets_hp,
|
||
cache_bs_id_lp, cache_seq_offsets_lp, slot_mapping_hp, slot_mapping_lp,
|
||
rotary_stride, cur_batch, batch_offset, head_num_q, head_num_k, head_size,
|
||
max_decode_len_hp, max_decode_len_lp, block_size_hp, block_size_lp, group_size,
|
||
once_batch, eps);
|
||
}
|
||
}
|
||
} // namespace kernels
|
||
|
||
KernelStatus invokeFusedRope(cnrtQueue_t queue,
|
||
void *input,
|
||
void *key_cache_hp,
|
||
void *value_cache_hp,
|
||
void *key_cache_lp,
|
||
void *value_cache_lp,
|
||
const void *sin_table,
|
||
const void *cos_table,
|
||
const void *rope_offsets,
|
||
const void *gamma,
|
||
const void *beta,
|
||
const void *key_scale_hp,
|
||
const void *value_scale_hp,
|
||
void *key_scale_lp,
|
||
void *value_scale_lp,
|
||
const void *cache_bs_id_hp,
|
||
const void *cache_seq_offsets_hp,
|
||
const void *cache_bs_id_lp,
|
||
const void *cache_seq_offsets_lp,
|
||
const void *slot_mapping_hp,
|
||
const void *slot_mapping_lp,
|
||
int rotary_stride,
|
||
int batch_size,
|
||
int head_num_q,
|
||
int head_num_kv,
|
||
int head_size,
|
||
int max_decode_len_hp,
|
||
int max_decode_len_lp,
|
||
int block_size_hp,
|
||
int block_size_lp,
|
||
int group_size,
|
||
cnnlDataType_t dtype,
|
||
float eps) {
|
||
if (is_arch300()) {
|
||
std::cerr << "[invokeFusedRope]: kernel does not support MLU300 devices." << std::endl;
|
||
return KernelStatus::KERNEL_STATUS_FAILED;
|
||
}
|
||
CNdev dev;
|
||
cnCtxGetDevice(&dev);
|
||
int cluster_num;
|
||
CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
|
||
int core_num;
|
||
CNRT_CHECK(cnrtDeviceGetAttribute(&core_num, cnrtAttrMcorePerCluster, dev));
|
||
uint32_t taskdimx = cluster_num * core_num;
|
||
|
||
int task_avg_batch = (batch_size + taskdimx - 1) / taskdimx;
|
||
int float_size = sizeof(float);
|
||
int group_num = head_size / group_size;
|
||
bool quant_kv_hp = key_scale_hp != nullptr && value_scale_hp != nullptr;
|
||
bool mixed_cache = key_cache_lp != nullptr && value_cache_lp != nullptr;
|
||
|
||
int nram_avalible_bytes = 480 * 1024;
|
||
int task_max_batch = 32;
|
||
int mask_bytes = PAD_UP(task_max_batch * head_num_kv, 8) / 8 * (mixed_cache + 1);
|
||
int nram_params_bytes = 2 * head_size * float_size;
|
||
int nram_kv_hp_scale_bytes = 2 * (int)quant_kv_hp * head_num_kv * head_size * float_size;
|
||
int nram_remain_bytes =
|
||
nram_avalible_bytes - nram_params_bytes - nram_kv_hp_scale_bytes - mask_bytes;
|
||
int nram_qk_bytes = (head_num_q + head_num_kv) * head_size * float_size * 2;
|
||
int nram_v_bytes = head_num_kv * head_size * float_size * (mixed_cache + 1);
|
||
int nram_table_bytes = 2 * head_size * float_size;
|
||
int nram_kv_lp_scale_bytes = 2 * (int)mixed_cache * head_num_kv * group_num * float_size;
|
||
int nram_kv_hp_bytes = (int)quant_kv_hp * head_num_kv * head_size;
|
||
int nram_kv_lp_bytes = (int)mixed_cache * head_num_kv * head_size;
|
||
int nram_cache_v_bytes = (int)mixed_cache * head_num_kv * head_size * 2;
|
||
int nram_cache_offsets_hp = head_num_kv * sizeof(int);
|
||
int nram_cache_offsets_lp = (int)mixed_cache * head_num_kv * 3 * sizeof(int);
|
||
int batch_cap =
|
||
nram_remain_bytes /
|
||
(nram_qk_bytes + nram_v_bytes + nram_table_bytes + nram_kv_lp_scale_bytes + nram_kv_hp_bytes +
|
||
nram_kv_lp_bytes + nram_cache_v_bytes + nram_cache_offsets_hp + nram_cache_offsets_lp);
|
||
batch_cap = batch_cap < task_avg_batch ? std::min(task_max_batch, batch_cap) : task_avg_batch;
|
||
|
||
cnrtDim3_t dim{taskdimx, 1, 1};
|
||
if (dtype == CNNL_DTYPE_HALF) {
|
||
kernels::MLUFuseRope<<<dim, cnrtFuncTypeBlock, queue>>>(
|
||
(half *)input, key_cache_hp, value_cache_hp, key_cache_lp, value_cache_lp,
|
||
(half *)sin_table, (half *)cos_table, (int *)rope_offsets, (half *)gamma, (half *)beta,
|
||
(float *)key_scale_hp, (float *)value_scale_hp, (float *)key_scale_lp,
|
||
(float *)value_scale_lp, (int *)cache_bs_id_hp, (int *)cache_seq_offsets_hp,
|
||
(int *)cache_bs_id_lp, (int *)cache_seq_offsets_lp, (int *)slot_mapping_hp,
|
||
(int *)slot_mapping_lp, rotary_stride, batch_size, head_num_q, head_num_kv, head_size,
|
||
max_decode_len_hp, max_decode_len_lp, block_size_hp, block_size_lp, group_size, batch_cap,
|
||
task_avg_batch, eps);
|
||
} else if (dtype == CNNL_DTYPE_BFLOAT16) {
|
||
kernels::MLUFuseRope<<<dim, cnrtFuncTypeBlock, queue>>>(
|
||
(bf16 *)input, key_cache_hp, value_cache_hp, key_cache_lp, value_cache_lp,
|
||
(bf16 *)sin_table, (bf16 *)cos_table, (int *)rope_offsets, (bf16 *)gamma, (bf16 *)beta,
|
||
(float *)key_scale_hp, (float *)value_scale_hp, (float *)key_scale_lp,
|
||
(float *)value_scale_lp, (int *)cache_bs_id_hp, (int *)cache_seq_offsets_hp,
|
||
(int *)cache_bs_id_lp, (int *)cache_seq_offsets_lp, (int *)slot_mapping_hp,
|
||
(int *)slot_mapping_lp, rotary_stride, batch_size, head_num_q, head_num_kv, head_size,
|
||
max_decode_len_hp, max_decode_len_lp, block_size_hp, block_size_lp, group_size, batch_cap,
|
||
task_avg_batch, eps);
|
||
}
|
||
return KernelStatus::KERNEL_STATUS_SUCCESS;
|
||
}
|
||
} // namespace tmo
|