wehub-resource-sync
988 lines
54 KiB
Common Lisp
988 lines
54 KiB
Common Lisp
#ifdef MNN_SUPPORT_FP16
|
||
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||
#endif
|
||
|
||
#define GLOBAL_SIZE_3_DIMS \
|
||
__private const int global_size_dim0, __private const int global_size_dim1, __private const int global_size_dim2,
|
||
|
||
#define DEAL_NON_UNIFORM_DIM3(input1, input2, input3) \
|
||
if (input1 >= global_size_dim0 || input2 >= global_size_dim1 || input3 >= global_size_dim2) { \
|
||
return; \
|
||
}
|
||
|
||
#define GLOBAL_SIZE_2_DIMS \
|
||
__private const int global_size_dim0, __private const int global_size_dim1,
|
||
|
||
#define DEAL_NON_UNIFORM_DIM2(input1, input2) \
|
||
if (input1 >= global_size_dim0 || input2 >= global_size_dim1) { \
|
||
return; \
|
||
}
|
||
|
||
#define DEAL_OUTER_SEQLEN_NOT_ALIGN(length) \
|
||
if(4 * sl + 3 >= length) {\
|
||
temp_3 = (FLOAT4)0;\
|
||
}\
|
||
if(4 * sl + 2 >= length) {\
|
||
temp_2 = (FLOAT4)0;\
|
||
}\
|
||
if(4 * sl + 1 >= length) {\
|
||
temp_1 = (FLOAT4)0;\
|
||
}
|
||
|
||
#define DEAL_INNER_HEADDIM_NOT_ALIGN(length) \
|
||
if(hd * 4 + 3 >= length) {\
|
||
temp_0.w = (FLOAT)0;\
|
||
temp_1.w = (FLOAT)0;\
|
||
temp_2.w = (FLOAT)0;\
|
||
temp_3.w = (FLOAT)0;\
|
||
}\
|
||
if(hd * 4 + 2 >= length) {\
|
||
temp_0.z = (FLOAT)0;\
|
||
temp_1.z = (FLOAT)0;\
|
||
temp_2.z = (FLOAT)0;\
|
||
temp_3.z = (FLOAT)0;\
|
||
}\
|
||
if(hd * 4 + 1 >= length) {\
|
||
temp_0.y = (FLOAT)0;\
|
||
temp_1.y = (FLOAT)0;\
|
||
temp_2.y = (FLOAT)0;\
|
||
temp_3.y = (FLOAT)0;\
|
||
}
|
||
|
||
|
||
#ifdef VALUE_C4
|
||
static inline FLOAT load_c4_value(__global const FLOAT* value,
|
||
const int seq_storage,
|
||
const int token,
|
||
const int channel) {
|
||
return value[((channel >> 2) * seq_storage + token) * 4 + (channel & 3)];
|
||
}
|
||
|
||
static inline FLOAT4 load_c4_value4(__global const FLOAT* value,
|
||
const int seq_storage,
|
||
const int token,
|
||
const int channel,
|
||
const int head_dim_offset,
|
||
const int head_dim) {
|
||
return (FLOAT4)(
|
||
load_c4_value(value, seq_storage, token, channel),
|
||
(head_dim_offset + 1 >= head_dim) ? (FLOAT)0 : load_c4_value(value, seq_storage, token, channel + 1),
|
||
(head_dim_offset + 2 >= head_dim) ? (FLOAT)0 : load_c4_value(value, seq_storage, token, channel + 2),
|
||
(head_dim_offset + 3 >= head_dim) ? (FLOAT)0 : load_c4_value(value, seq_storage, token, channel + 3));
|
||
}
|
||
#endif
|
||
|
||
|
||
__kernel void rearrange_qkv(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *input_q, //[batch, seqLenQ/4, headNum, headDim, seqLenQ_4]
|
||
__global const FLOAT *input_k, // [batch, seqLenKV/4, headNum/group, headDim, seqLenKV_4]
|
||
__global const FLOAT *input_v, // [batch, seqLenKV/4, headNum/group, headDim, seqLenKV_4]
|
||
__global FLOAT *output_q, // [batch*headNum, ROUND_UP(headDim, mTileHDK), ROUND_UP(seqLenQ, mTileQ)]
|
||
__global FLOAT *output_k, // [batch*headNum/group, ROUND_UP(headDim, mTileHDK), ROUND_UP(seqLenKV, mTileKV)]
|
||
__global FLOAT *output_v, // [batch*headNum/group, ROUND_UP(seqLenKV, mTileKV), ROUND_UP(headDim, mTileHDN)]
|
||
#ifdef SAVE_KV
|
||
__global FLOAT *past_k, // [batch, headNum/group, headDim, seqLenKV_4]
|
||
__global FLOAT *past_v, // [batch, headNum/group, seqLenKV_4, headDim]
|
||
#endif
|
||
__private const int4 tile, // [mTileQ, mTileKV, mTileHDK, mTileHDN]
|
||
__private const int4 shape,// [seqLenQ, seqLenKV, headNum, headDim]
|
||
__private const int4 param, // [group, batch, max_len, past_len]
|
||
__private const int maxLenKV
|
||
) {
|
||
const int sl = get_global_id(0); // seqLen/4 : max(seqLenPackQ/4, seqLenPackKV/4)
|
||
const int hd = get_global_id(1); // headDim/4 : max(headDimPackQK/4, headDimPackV/4)
|
||
const int z = get_global_id(2); // batch * headNum
|
||
DEAL_NON_UNIFORM_DIM3(sl, hd, z);
|
||
|
||
const int seqLenQ = shape.x;
|
||
const int seqLenKV = shape.y;
|
||
const int headNum = shape.z;
|
||
const int headDim = shape.w;
|
||
const int group = param.x;
|
||
const int batch = param.y;
|
||
|
||
const int b = z % batch;
|
||
const int hn = z / batch;
|
||
|
||
const int seqLenQ_4 = (seqLenQ + 3) / 4;
|
||
//const int in_offset_q = (((b * seqLenQ_4 + sl) * headNum + hn) * headDim + 4 * hd) * 4;
|
||
const int in_offset_q = (((b * seqLenQ + sl * 4) * headNum + hn) * headDim + 4 * hd);
|
||
|
||
const int seqLenPackQ = ((seqLenQ + tile.x - 1) / tile.x) * tile.x;
|
||
const int headDimPackQK = ((headDim + tile.z - 1) / tile.z) * tile.z;
|
||
const int out_offset_q = (((b * headNum + hn) * headDimPackQK + hd * 4) * seqLenPackQ + sl * 4);
|
||
|
||
if(sl * 4 < seqLenPackQ && hd * 4 < headDimPackQK) {
|
||
if(sl * 4 >= seqLenQ || hd * 4 >= headDim) {
|
||
vstore4((FLOAT4)0, 0, output_q + out_offset_q);
|
||
vstore4((FLOAT4)0, 0, output_q + out_offset_q + seqLenPackQ);
|
||
vstore4((FLOAT4)0, 0, output_q + out_offset_q + 2 * seqLenPackQ);
|
||
vstore4((FLOAT4)0, 0, output_q + out_offset_q + 3 * seqLenPackQ);
|
||
} else {
|
||
FLOAT4 temp_0 = vload4(0, input_q + in_offset_q);
|
||
FLOAT4 temp_1 = (sl * 4 + 1 >= seqLenQ) ? (FLOAT4)0 : vload4(0, input_q + in_offset_q + headNum*headDim);
|
||
FLOAT4 temp_2 = (sl * 4 + 2 >= seqLenQ) ? (FLOAT4)0 : vload4(0, input_q + in_offset_q + 2*headNum*headDim);
|
||
FLOAT4 temp_3 = (sl * 4 + 3 >= seqLenQ) ? (FLOAT4)0 : vload4(0, input_q + in_offset_q + 3*headNum*headDim);
|
||
#ifdef HEADDIM_LEAVE
|
||
DEAL_INNER_HEADDIM_NOT_ALIGN(headDim)
|
||
#endif
|
||
#ifdef SEQLEN_LEAVE
|
||
DEAL_OUTER_SEQLEN_NOT_ALIGN(seqLenQ)
|
||
#endif
|
||
vstore4((FLOAT4)(temp_0.s0, temp_1.s0, temp_2.s0, temp_3.s0), 0, output_q + out_offset_q);
|
||
vstore4((FLOAT4)(temp_0.s1, temp_1.s1, temp_2.s1, temp_3.s1), 0, output_q + out_offset_q + seqLenPackQ);
|
||
vstore4((FLOAT4)(temp_0.s2, temp_1.s2, temp_2.s2, temp_3.s2), 0, output_q + out_offset_q + 2 * seqLenPackQ);
|
||
vstore4((FLOAT4)(temp_0.s3, temp_1.s3, temp_2.s3, temp_3.s3), 0, output_q + out_offset_q + 3 * seqLenPackQ);
|
||
}
|
||
}
|
||
|
||
if(hn >= headNum / group) {
|
||
return;
|
||
}
|
||
|
||
|
||
const int seqLenPackKV = ((seqLenKV + tile.y - 1) / tile.y) * tile.y;
|
||
const int headDimPackV = ((headDim + tile.w - 1) / tile.w) * tile.w;
|
||
const int seqLenKV_4 = (seqLenKV + 3) / 4;
|
||
const int in_offset_kv = (((b * seqLenKV + sl*4) * headNum/group + hn) * headDim + 4 * hd);
|
||
const int past_offset_k = (((b * headNum/group + hn) * headDim + hd * 4) * maxLenKV + sl*4);
|
||
const int past_offset_v = (((b * headNum/group + hn) * maxLenKV + sl*4) * headDim + 4 * hd);
|
||
if(sl * 4 < seqLenPackKV && hd * 4 < headDimPackQK) {
|
||
const int out_offset_k = (((b * headNum/group + hn) * headDimPackQK + hd * 4) * seqLenPackKV + sl * 4);
|
||
|
||
if(sl * 4 >= seqLenKV || hd * 4 >= headDim) {
|
||
vstore4((FLOAT4)0, 0, output_k + out_offset_k);
|
||
vstore4((FLOAT4)0, 0, output_k + out_offset_k + seqLenPackKV);
|
||
vstore4((FLOAT4)0, 0, output_k + out_offset_k + 2 * seqLenPackKV);
|
||
vstore4((FLOAT4)0, 0, output_k + out_offset_k + 3 * seqLenPackKV);
|
||
} else {
|
||
FLOAT4 temp_0 = vload4(0, input_k + in_offset_kv);
|
||
FLOAT4 temp_1 = (sl * 4 + 1 >= seqLenKV) ? (FLOAT4)0 : vload4(0, input_k + in_offset_kv + headNum*headDim/group);
|
||
FLOAT4 temp_2 = (sl * 4 + 2 >= seqLenKV) ? (FLOAT4)0 : vload4(0, input_k + in_offset_kv + 2*headNum*headDim/group);
|
||
FLOAT4 temp_3 = (sl * 4 + 3 >= seqLenKV) ? (FLOAT4)0 : vload4(0, input_k + in_offset_kv + 3*headNum*headDim/group);
|
||
#ifdef HEADDIM_LEAVE
|
||
DEAL_INNER_HEADDIM_NOT_ALIGN(headDim)
|
||
#endif
|
||
#ifdef SEQLEN_LEAVE
|
||
DEAL_OUTER_SEQLEN_NOT_ALIGN(seqLenKV)
|
||
#endif
|
||
FLOAT4 key0 = (FLOAT4)(temp_0.s0, temp_1.s0, temp_2.s0, temp_3.s0);
|
||
FLOAT4 key1 = (FLOAT4)(temp_0.s1, temp_1.s1, temp_2.s1, temp_3.s1);
|
||
FLOAT4 key2 = (FLOAT4)(temp_0.s2, temp_1.s2, temp_2.s2, temp_3.s2);
|
||
FLOAT4 key3 = (FLOAT4)(temp_0.s3, temp_1.s3, temp_2.s3, temp_3.s3);
|
||
vstore4(key0, 0, output_k + out_offset_k);
|
||
vstore4(key1, 0, output_k + out_offset_k + seqLenPackKV);
|
||
vstore4(key2, 0, output_k + out_offset_k + 2 * seqLenPackKV);
|
||
vstore4(key3, 0, output_k + out_offset_k + 3 * seqLenPackKV);
|
||
|
||
// pastK
|
||
#ifdef SAVE_KV
|
||
vstore4(key0, 0, past_k + past_offset_k);
|
||
vstore4(key1, 0, past_k + past_offset_k + maxLenKV);
|
||
vstore4(key2, 0, past_k + past_offset_k + 2*maxLenKV);
|
||
vstore4(key3, 0, past_k + past_offset_k + 3*maxLenKV);
|
||
#endif
|
||
}
|
||
|
||
}
|
||
|
||
if(sl * 4 < seqLenPackKV && hd * 4 < headDimPackV) {
|
||
const int out_offset_v = (((b * headNum/group + hn) * seqLenPackKV + sl * 4) * headDimPackV + hd * 4);
|
||
|
||
if(sl * 4 >= seqLenKV || hd * 4 >= headDim) {
|
||
vstore4((FLOAT4)0, 0, output_v + out_offset_v);
|
||
vstore4((FLOAT4)0, 0, output_v + out_offset_v + headDimPackV);
|
||
vstore4((FLOAT4)0, 0, output_v + out_offset_v + 2 * headDimPackV);
|
||
vstore4((FLOAT4)0, 0, output_v + out_offset_v + 3 * headDimPackV);
|
||
} else {
|
||
#ifdef VALUE_C4
|
||
const int value_seq_storage = batch * seqLenKV;
|
||
const int value_channel = hn * headDim + 4 * hd;
|
||
const int value_token = b * seqLenKV + sl * 4;
|
||
FLOAT4 temp_0 = load_c4_value4(input_v, value_seq_storage, value_token, value_channel, 4 * hd, headDim);
|
||
FLOAT4 temp_1 = (sl * 4 + 1 >= seqLenKV) ? (FLOAT4)0 :
|
||
load_c4_value4(input_v, value_seq_storage, value_token + 1, value_channel, 4 * hd, headDim);
|
||
FLOAT4 temp_2 = (sl * 4 + 2 >= seqLenKV) ? (FLOAT4)0 :
|
||
load_c4_value4(input_v, value_seq_storage, value_token + 2, value_channel, 4 * hd, headDim);
|
||
FLOAT4 temp_3 = (sl * 4 + 3 >= seqLenKV) ? (FLOAT4)0 :
|
||
load_c4_value4(input_v, value_seq_storage, value_token + 3, value_channel, 4 * hd, headDim);
|
||
#else
|
||
FLOAT4 temp_0 = vload4(0, input_v + in_offset_kv);
|
||
FLOAT4 temp_1 = (sl * 4 + 1 >= seqLenKV) ? (FLOAT4)0 : vload4(0, input_v + in_offset_kv + headNum*headDim/group);
|
||
FLOAT4 temp_2 = (sl * 4 + 2 >= seqLenKV) ? (FLOAT4)0 : vload4(0, input_v + in_offset_kv + 2*headNum*headDim/group);
|
||
FLOAT4 temp_3 = (sl * 4 + 3 >= seqLenKV) ? (FLOAT4)0 : vload4(0, input_v + in_offset_kv + 3*headNum*headDim/group);
|
||
#endif
|
||
#ifdef HEADDIM_LEAVE
|
||
DEAL_INNER_HEADDIM_NOT_ALIGN(headDim)
|
||
#endif
|
||
#ifdef SEQLEN_LEAVE
|
||
DEAL_OUTER_SEQLEN_NOT_ALIGN(seqLenKV)
|
||
#endif
|
||
vstore4(temp_0, 0, output_v + out_offset_v);
|
||
vstore4(temp_1, 0, output_v + out_offset_v + headDimPackV);
|
||
vstore4(temp_2, 0, output_v + out_offset_v + 2 * headDimPackV);
|
||
vstore4(temp_3, 0, output_v + out_offset_v + 3 * headDimPackV);
|
||
|
||
// pastV
|
||
#ifdef SAVE_KV
|
||
vstore4(temp_0, 0, past_v + past_offset_v);
|
||
vstore4(temp_1, 0, past_v + past_offset_v + headDim);
|
||
vstore4(temp_2, 0, past_v + past_offset_v + 2*headDim);
|
||
vstore4(temp_3, 0, past_v + past_offset_v + 3*headDim);
|
||
#endif
|
||
}
|
||
|
||
}
|
||
}
|
||
|
||
#ifndef MASK_DTYPE
|
||
#define MASK_DTYPE FLOAT
|
||
#define MASK_DTYPE4 FLOAT4
|
||
#endif
|
||
__kernel void rearrange_mask(GLOBAL_SIZE_3_DIMS
|
||
__global const MASK_DTYPE *input_mask, // [batch, 1, seqLenQ, seqLenKV, 4]
|
||
__global MASK_DTYPE *output_mask, // [batch, ROUND_UP(seqLenQ, mTileQ), ROUND_UP(seqLenKV, mTileKV)]
|
||
const int4 shape // [seqLenQ, seqLenKV, mTileQ, mTileKV]
|
||
) {
|
||
const int sl = get_global_id(0); // seqLen_4
|
||
const int sl_kv = get_global_id(1); // seqLenKV_4
|
||
const int b = get_global_id(2); // Batch
|
||
DEAL_NON_UNIFORM_DIM3(sl, sl_kv, b);
|
||
|
||
const int seq_len_pack = ((shape.x + shape.z - 1) / shape.z) * shape.z;
|
||
const int seq_len_kv_pack = ((shape.y + shape.w - 1) / shape.w) * shape.w;
|
||
|
||
int in_offset = ((b * shape.x + sl * 4) * shape.y + sl_kv * 4);
|
||
int out_offset = (b * seq_len_pack + sl * 4) * seq_len_kv_pack + sl_kv * 4;
|
||
|
||
if(sl * 4 >= shape.x || sl_kv * 4 >= shape.y) {
|
||
vstore4((MASK_DTYPE4)0, 0, output_mask + out_offset);
|
||
vstore4((MASK_DTYPE4)0, 0, output_mask + out_offset + seq_len_kv_pack);
|
||
vstore4((MASK_DTYPE4)0, 0, output_mask + out_offset + seq_len_kv_pack * 2);
|
||
vstore4((MASK_DTYPE4)0, 0, output_mask + out_offset + seq_len_kv_pack * 3);
|
||
} else {
|
||
int y_down_align4 = (shape.y / 4 * 4);
|
||
MASK_DTYPE4 temp_0, temp_1, temp_2, temp_3;
|
||
|
||
if(sl_kv * 4 < y_down_align4) {
|
||
temp_0 = vload4(0, input_mask + in_offset);
|
||
temp_1 = (sl * 4 + 1 >= shape.x) ? (MASK_DTYPE4)0 : vload4(0, input_mask + in_offset + shape.y);
|
||
temp_2 = (sl * 4 + 2 >= shape.x) ? (MASK_DTYPE4)0 : vload4(0, input_mask + in_offset + shape.y * 2);
|
||
temp_3 = (sl * 4 + 3 >= shape.x) ? (MASK_DTYPE4)0 : vload4(0, input_mask + in_offset + shape.y * 3);
|
||
} else if(sl_kv * 4 + 1 == shape.y){
|
||
temp_0 = (MASK_DTYPE4)(input_mask[in_offset], 0, 0, 0);
|
||
temp_1 = (sl * 4 + 1 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y], 0, 0, 0);//vload4(0, input_mask + in_offset + shape.y);
|
||
temp_2 = (sl * 4 + 2 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y*2], 0, 0, 0);//vload4(0, input_mask + in_offset + shape.y * 2);
|
||
temp_3 = (sl * 4 + 3 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y*3], 0, 0, 0);//vload4(0, input_mask + in_offset + shape.y * 3);
|
||
} else if(sl_kv * 4 + 2 == shape.y){
|
||
temp_0 = (MASK_DTYPE4)(input_mask[in_offset], input_mask[in_offset+1], 0, 0);
|
||
temp_1 = (sl * 4 + 1 >= shape.x) ? (MASK_DTYPE4)0 : (FLOAT4)(input_mask[in_offset + shape.y], input_mask[in_offset + shape.y + 1], 0, 0);//vload4(0, input_mask + in_offset + shape.y);
|
||
temp_2 = (sl * 4 + 2 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y*2], input_mask[in_offset + shape.y*2 + 1], 0, 0);//vload4(0, input_mask + in_offset + shape.y * 2);
|
||
temp_3 = (sl * 4 + 3 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y*3], input_mask[in_offset + shape.y*3 + 1], 0, 0);//vload4(0, input_mask + in_offset + shape.y * 3);
|
||
} else if(sl_kv * 4 + 3 == shape.y){
|
||
temp_0 = (MASK_DTYPE4)(input_mask[in_offset], input_mask[in_offset+1], input_mask[in_offset+2], 0);
|
||
temp_1 = (sl * 4 + 1 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y], input_mask[in_offset + shape.y + 1], input_mask[in_offset + shape.y + 2], 0);//vload4(0, input_mask + in_offset + shape.y);
|
||
temp_2 = (sl * 4 + 2 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y*2], input_mask[in_offset + shape.y*2 + 1], input_mask[in_offset + shape.y*2 + 2], 0);//vload4(0, input_mask + in_offset + shape.y * 2);
|
||
temp_3 = (sl * 4 + 3 >= shape.x) ? (MASK_DTYPE4)0 : (MASK_DTYPE4)(input_mask[in_offset + shape.y*3], input_mask[in_offset + shape.y*3 + 1], input_mask[in_offset + shape.y*3 + 2], 0);//vload4(0, input_mask + in_offset + shape.y * 3);
|
||
}
|
||
|
||
vstore4(temp_0, 0, output_mask + out_offset);
|
||
vstore4(temp_1, 0, output_mask + out_offset + seq_len_kv_pack);
|
||
vstore4(temp_2, 0, output_mask + out_offset + 2 * seq_len_kv_pack);
|
||
vstore4(temp_3, 0, output_mask + out_offset + 3 * seq_len_kv_pack);
|
||
}
|
||
|
||
}
|
||
|
||
__kernel void qkv_transpose_output(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *input, // [Batch * mNumHead, ROUND_UP(mHeadDim, mTileHDN), ROUND_UP(seqLen, mTileQ)]
|
||
__global FLOAT *output, // [Batch, seqLen/4, mNumHead， mHeadDim, 4] (or NC4HW4 when ATTENTION_C4)
|
||
__private const int tile_q,
|
||
__private const int tile_hdn,
|
||
__private const int seq_len,
|
||
__private const int head_num,
|
||
__private const int head_dim,
|
||
__private const int batch
|
||
) {
|
||
|
||
const int sl = get_global_id(0); // seqLen_4
|
||
const int hd = get_global_id(1); // mHeadDim_4
|
||
const int z = get_global_id(2); // Batch * mNumHead
|
||
DEAL_NON_UNIFORM_DIM3(sl, hd, z);
|
||
|
||
const int b = z / head_num;
|
||
const int hn = z % head_num;
|
||
|
||
const int seq_len_pack = ((seq_len + tile_q - 1) / tile_q) * tile_q;
|
||
const int head_dim_pack = ((head_dim + tile_hdn - 1) / tile_hdn) * tile_hdn;
|
||
|
||
const int offset_inp = ((b * head_num + hn) * head_dim_pack + 4 * hd) * seq_len_pack + 4 * sl;
|
||
|
||
// Q
|
||
FLOAT4 temp_0 = vload4(0, input + offset_inp);
|
||
FLOAT4 temp_1 = vload4(0, input + offset_inp + seq_len_pack);
|
||
FLOAT4 temp_2 = vload4(0, input + offset_inp + 2 * seq_len_pack);
|
||
FLOAT4 temp_3 = vload4(0, input + offset_inp + 3 * seq_len_pack);
|
||
|
||
#ifdef ATTENTION_C4
|
||
// output is NC4HW4: [(head_num*head_dim)/4, batch*seq_len, 4], channel = hn*head_dim + 4*hd.
|
||
// Must match matmul_qkv_prefill's ATTENTION_C4 output layout so o_proj reads it correctly.
|
||
const int channel4 = (hn * head_dim + 4 * hd) >> 2;
|
||
const int seq_storage = seq_len * batch;
|
||
const int output_offset = (channel4 * seq_storage + (b * seq_len + sl * 4)) * 4;
|
||
vstore4((FLOAT4)(temp_0.s0, temp_1.s0, temp_2.s0, temp_3.s0), 0, output + output_offset);
|
||
if(4 * sl + 1 >= seq_len) return;
|
||
vstore4((FLOAT4)(temp_0.s1, temp_1.s1, temp_2.s1, temp_3.s1), 0, output + output_offset + 4);
|
||
if(4 * sl + 2 >= seq_len) return;
|
||
vstore4((FLOAT4)(temp_0.s2, temp_1.s2, temp_2.s2, temp_3.s2), 0, output + output_offset + 8);
|
||
if(4 * sl + 3 >= seq_len) return;
|
||
vstore4((FLOAT4)(temp_0.s3, temp_1.s3, temp_2.s3, temp_3.s3), 0, output + output_offset + 12);
|
||
#else
|
||
const int offset_out = (((b * seq_len + sl*4) * head_num + hn) * head_dim + 4 * hd);
|
||
vstore4((FLOAT4)(temp_0.s0, temp_1.s0, temp_2.s0, temp_3.s0), 0, output + offset_out);
|
||
if(4 * sl + 1 >= seq_len) return;
|
||
vstore4((FLOAT4)(temp_0.s1, temp_1.s1, temp_2.s1, temp_3.s1), 0, output + offset_out + head_num*head_dim);
|
||
if(4 * sl + 2 >= seq_len) return;
|
||
vstore4((FLOAT4)(temp_0.s2, temp_1.s2, temp_2.s2, temp_3.s2), 0, output + offset_out + 2*head_num*head_dim);
|
||
if(4 * sl + 3 >= seq_len) return;
|
||
vstore4((FLOAT4)(temp_0.s3, temp_1.s3, temp_2.s3, temp_3.s3), 0, output + offset_out + 3*head_num*head_dim);
|
||
#endif
|
||
|
||
}
|
||
|
||
#ifndef NUMHEAD_GROUP_SIZE
|
||
#define NUMHEAD_GROUP_SIZE 1
|
||
#endif
|
||
|
||
__kernel void rearrange_q(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *query, // [batch query_seq_len head_num head_dim]
|
||
__global FLOAT *query_tmp, // [batch head_num head_dim_4 query_seq_len_4]
|
||
__private const int seq_len,
|
||
__private const int head_dim,
|
||
__private const int head_num) {
|
||
/*
|
||
the kernel assume head_dim is multiple of 4.
|
||
*/
|
||
const int x = get_global_id(0); // query_seq_len/4
|
||
const int y = get_global_id(1); // head_dim/4
|
||
int z = get_global_id(2);
|
||
DEAL_NON_UNIFORM_DIM3(x, y, z);
|
||
|
||
const int b = z / head_num;// batch
|
||
z = z % head_num;// head_num
|
||
|
||
const int x4 = x << 2;
|
||
const int y4 = y << 2;
|
||
const int seq_len4 = (seq_len + 3) / 4 * 4;;
|
||
const int stride = head_num * head_dim;
|
||
int query_offset = ((b * seq_len + x4) * head_num + z) * head_dim + y4;
|
||
FLOAT4 query_vec0 = vload4(0, query + query_offset); query_offset += stride;
|
||
FLOAT4 query_vec1 = (x4 + 1 >= seq_len) ? (FLOAT4)0 : vload4(0, query + query_offset); query_offset += stride;
|
||
FLOAT4 query_vec2 = (x4 + 2 >= seq_len) ? (FLOAT4)0 : vload4(0, query + query_offset); query_offset += stride;
|
||
FLOAT4 query_vec3 = (x4 + 3 >= seq_len) ? (FLOAT4)0 : vload4(0, query + query_offset);
|
||
|
||
const int queryout_offset = ((b * head_num + z) * head_dim + y4) * seq_len4 + x4;
|
||
vstore4((FLOAT4)(query_vec0.s0, query_vec1.s0, query_vec2.s0, query_vec3.s0), 0, query_tmp + queryout_offset);
|
||
vstore4((FLOAT4)(query_vec0.s1, query_vec1.s1, query_vec2.s1, query_vec3.s1), 0, query_tmp + queryout_offset + seq_len4);
|
||
vstore4((FLOAT4)(query_vec0.s2, query_vec1.s2, query_vec2.s2, query_vec3.s2), 0, query_tmp + queryout_offset + seq_len4 + seq_len4);
|
||
vstore4((FLOAT4)(query_vec0.s3, query_vec1.s3, query_vec2.s3, query_vec3.s3), 0, query_tmp + queryout_offset + seq_len4 + seq_len4 + seq_len4);
|
||
}
|
||
|
||
__kernel void rearrange_k(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *key, // [batch key_seq_len kv_head_num head_dim]
|
||
__global FLOAT *past_key, // [batch kv_head_num head_dim max_length]
|
||
__private const int past_len, // prefill = 0, decode = past_key len
|
||
__private const int max_len,
|
||
__private const int seq_len,
|
||
__private const int kv_head_num,
|
||
__private const int head_num,
|
||
__private const int head_dim) {
|
||
|
||
const int x = get_global_id(0); // seq_len decode = 1
|
||
const int y = get_global_id(1); // head_dim
|
||
int z = get_global_id(2); //
|
||
DEAL_NON_UNIFORM_DIM3(x, y, z);
|
||
|
||
const int b = z / kv_head_num;
|
||
z = z % kv_head_num;
|
||
const int y4 = y << 2;
|
||
|
||
#ifdef OPENCL_PREFILL_ATTENTION
|
||
const int x4 = x << 2;
|
||
const int stride = kv_head_num * head_dim;
|
||
int key_offset = ((b * seq_len + x4) * kv_head_num + z) * head_dim + y4;
|
||
FLOAT4 key_vec0 = vload4(0, key + key_offset); key_offset += stride;
|
||
FLOAT4 key_vec1 = (x4 + 1 >= seq_len) ? (FLOAT4)0 : vload4(0, key + key_offset); key_offset += stride;
|
||
FLOAT4 key_vec2 = (x4 + 2 >= seq_len) ? (FLOAT4)0 : vload4(0, key + key_offset); key_offset += stride;
|
||
FLOAT4 key_vec3 = (x4 + 3 >= seq_len) ? (FLOAT4)0 : vload4(0, key + key_offset);
|
||
const int output_offset = ((b * kv_head_num + z) * head_dim + y4) * max_len + past_len + x4;
|
||
vstore4((FLOAT4)(key_vec0.s0, key_vec1.s0, key_vec2.s0, key_vec3.s0), 0, past_key + output_offset);
|
||
vstore4((FLOAT4)(key_vec0.s1, key_vec1.s1, key_vec2.s1, key_vec3.s1), 0, past_key + output_offset + max_len);
|
||
vstore4((FLOAT4)(key_vec0.s2, key_vec1.s2, key_vec2.s2, key_vec3.s2), 0, past_key + output_offset + max_len + max_len);
|
||
vstore4((FLOAT4)(key_vec0.s3, key_vec1.s3, key_vec2.s3, key_vec3.s3), 0, past_key + output_offset + max_len + max_len + max_len);
|
||
#else
|
||
FLOAT4 key_vec = vload4(0, key + (b * kv_head_num + z) * head_dim + y4);
|
||
const int output_offset = ((b * kv_head_num + z) * head_dim + y4) * max_len + past_len;
|
||
past_key[output_offset] = key_vec.s0;
|
||
past_key[output_offset + max_len] = key_vec.s1;
|
||
past_key[output_offset + max_len + max_len] = key_vec.s2;
|
||
past_key[output_offset + max_len + max_len + max_len] = key_vec.s3;
|
||
#endif
|
||
}
|
||
|
||
__kernel void rearrange_v(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *value, // [batch value_seq_len kv_head_num head_dim]
|
||
__global FLOAT *past_value, // [batch kv_head_num max_length head_dim]
|
||
__private const int past_len,
|
||
__private const int max_len,
|
||
__private const int seq_len,
|
||
__private const int kv_head_num,
|
||
__private const int head_dim) {
|
||
|
||
const int x = get_global_id(0); // head_dim
|
||
const int y = get_global_id(1); // seq_len decode = 1
|
||
int z = get_global_id(2); // kv_head_num
|
||
DEAL_NON_UNIFORM_DIM3(x, y, z);
|
||
|
||
const int b = z / kv_head_num;
|
||
z = z % kv_head_num;
|
||
const int x4 = x << 2;
|
||
|
||
#ifdef OPENCL_PREFILL_ATTENTION
|
||
const int y4 = y << 2;
|
||
const int stride = kv_head_num * head_dim;
|
||
#ifdef VALUE_C4
|
||
const int value_seq_storage = (global_size_dim2 / kv_head_num) * seq_len;
|
||
const int value_channel = z * head_dim + x4;
|
||
const int value_token = b * seq_len + y4;
|
||
FLOAT4 value_vec0 = load_c4_value4(value, value_seq_storage, value_token, value_channel, x4, head_dim);
|
||
FLOAT4 value_vec1 = (y4 + 1 >= seq_len) ? (FLOAT4)0 :
|
||
load_c4_value4(value, value_seq_storage, value_token + 1, value_channel, x4, head_dim);
|
||
FLOAT4 value_vec2 = (y4 + 2 >= seq_len) ? (FLOAT4)0 :
|
||
load_c4_value4(value, value_seq_storage, value_token + 2, value_channel, x4, head_dim);
|
||
FLOAT4 value_vec3 = (y4 + 3 >= seq_len) ? (FLOAT4)0 :
|
||
load_c4_value4(value, value_seq_storage, value_token + 3, value_channel, x4, head_dim);
|
||
#else
|
||
int value_offset = ((b * seq_len + y4) * kv_head_num + z) * head_dim + x4;
|
||
FLOAT4 value_vec0 = vload4(0, value + value_offset); value_offset += stride;
|
||
FLOAT4 value_vec1 = (y4 + 1 >= seq_len) ? (FLOAT4)0 : vload4(0, value + value_offset); value_offset += stride;
|
||
FLOAT4 value_vec2 = (y4 + 2 >= seq_len) ? (FLOAT4)0 : vload4(0, value + value_offset); value_offset += stride;
|
||
FLOAT4 value_vec3 = (y4 + 3 >= seq_len) ? (FLOAT4)0 : vload4(0, value + value_offset);
|
||
#endif
|
||
const int output_offset = ((b * kv_head_num + z) * max_len + past_len + y4) * head_dim + x4;
|
||
vstore4(value_vec0, 0, past_value + output_offset);
|
||
vstore4(value_vec1, 0, past_value + output_offset + head_dim);
|
||
vstore4(value_vec2, 0, past_value + output_offset + head_dim + head_dim);
|
||
vstore4(value_vec3, 0, past_value + output_offset + head_dim + head_dim + head_dim);
|
||
#else
|
||
#ifdef VALUE_C4
|
||
const int value_seq_storage = (global_size_dim2 / kv_head_num) * seq_len;
|
||
const int value_channel = z * head_dim + x4;
|
||
const int value_token = b * seq_len;
|
||
FLOAT4 value_vec = load_c4_value4(value, value_seq_storage, value_token, value_channel, x4, head_dim);
|
||
#else
|
||
FLOAT4 value_vec = vload4(0, value + (b * kv_head_num + z) * head_dim + x4);
|
||
#endif
|
||
const int output_offset = ((b * kv_head_num + z) * max_len + past_len) * head_dim + x4;
|
||
vstore4(value_vec, 0, past_value + output_offset);
|
||
#endif
|
||
}
|
||
|
||
__kernel void rearrange_mask_shortprefill(GLOBAL_SIZE_3_DIMS
|
||
#ifdef ADD_MASK
|
||
__global const FLOAT* mask,
|
||
__global FLOAT* maskout,
|
||
#else
|
||
__global const int* mask, // [1 1 query_seq_len mask_key_seq_len4]
|
||
__global int* maskout, // [1 1 mask_key_seq_len4 query_seq_len4]
|
||
#endif
|
||
__private const int query_seq_len,
|
||
__private const int mask_key_seq_len){
|
||
const int x = get_global_id(0); // query_seq_len4
|
||
const int y = get_global_id(1); // mask_key_seq_len4
|
||
const int z = get_global_id(2); // batch
|
||
DEAL_NON_UNIFORM_DIM3(x, y, z);
|
||
|
||
const int x4 = x << 2;
|
||
const int y4 = y << 2;
|
||
float4 mask_tmp0, mask_tmp1, mask_tmp2, mask_tmp3;
|
||
float4 mask0, mask1, mask2, mask3;
|
||
int mask_offset = x4 * mask_key_seq_len + y4;
|
||
if(x4 + 3 < query_seq_len && y4 + 3 < mask_key_seq_len){
|
||
mask_tmp0 = convert_float4(vload4(0, mask + mask_offset)); mask_offset += mask_key_seq_len;
|
||
mask_tmp1 = convert_float4(vload4(0, mask + mask_offset)); mask_offset += mask_key_seq_len;
|
||
mask_tmp2 = convert_float4(vload4(0, mask + mask_offset)); mask_offset += mask_key_seq_len;
|
||
mask_tmp3 = convert_float4(vload4(0, mask + mask_offset));
|
||
} else{
|
||
if(y4 + 3 < mask_key_seq_len){
|
||
mask_tmp0 = convert_float4(vload4(0, mask + mask_offset)); mask_offset += mask_key_seq_len;
|
||
mask_tmp1 = (x4 + 1 >= query_seq_len) ? (float4)0 : convert_float4(vload4(0, mask + mask_offset)); mask_offset += mask_key_seq_len;
|
||
mask_tmp2 = (x4 + 2 >= query_seq_len) ? (float4)0 : convert_float4(vload4(0, mask + mask_offset)); mask_offset += mask_key_seq_len;
|
||
mask_tmp3 = (x4 + 3 >= query_seq_len) ? (float4)0 : convert_float4(vload4(0, mask + mask_offset));
|
||
} else if(y4 + 1 == mask_key_seq_len){
|
||
mask_tmp0 = (float4)(mask[mask_offset], 0, 0, 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp1 = (x4 + 1 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], 0, 0, 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp2 = (x4 + 2 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], 0, 0, 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp3 = (x4 + 3 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], 0, 0, 0);
|
||
}else if(y4 + 2 == mask_key_seq_len){
|
||
mask_tmp0 = (float4)(mask[mask_offset], mask[mask_offset + 1], 0, 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp1 = (x4 + 1 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], mask[mask_offset + 1], 0, 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp2 = (x4 + 2 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], mask[mask_offset + 1], 0, 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp3 = (x4 + 3 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], mask[mask_offset + 1], 0, 0);
|
||
}else if(y4 + 3 == mask_key_seq_len){
|
||
mask_tmp0 = (float4)(mask[mask_offset], mask[mask_offset + 1], mask[mask_offset + 2], 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp1 = (x4 + 1 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], mask[mask_offset + 1], mask[mask_offset + 2], 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp2 = (x4 + 2 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], mask[mask_offset + 1], mask[mask_offset + 2], 0); mask_offset += mask_key_seq_len;
|
||
mask_tmp3 = (x4 + 3 >= query_seq_len) ? (float4)0 : (float4)(mask[mask_offset], mask[mask_offset + 1], mask[mask_offset + 2], 0);
|
||
}
|
||
}
|
||
mask0 = (float4)(mask_tmp0.s0, mask_tmp1.s0, mask_tmp2.s0, mask_tmp3.s0);
|
||
mask1 = (float4)(mask_tmp0.s1, mask_tmp1.s1, mask_tmp2.s1, mask_tmp3.s1);
|
||
mask2 = (float4)(mask_tmp0.s2, mask_tmp1.s2, mask_tmp2.s2, mask_tmp3.s2);
|
||
mask3 = (float4)(mask_tmp0.s3, mask_tmp1.s3, mask_tmp2.s3, mask_tmp3.s3);
|
||
|
||
int query_seq_len4 = ((query_seq_len + 3) / 4) * 4;
|
||
int output_offset = y4 * query_seq_len4 + x4;
|
||
#ifdef ADD_MASK
|
||
vstore4(CONVERT_FLOAT4(mask0), 0, maskout + output_offset);
|
||
vstore4(CONVERT_FLOAT4(mask1), 0, maskout + output_offset + query_seq_len4);
|
||
vstore4(CONVERT_FLOAT4(mask2), 0, maskout + output_offset + query_seq_len4 + query_seq_len4);
|
||
vstore4(CONVERT_FLOAT4(mask3), 0, maskout + output_offset + query_seq_len4 + query_seq_len4 + query_seq_len4);
|
||
#else
|
||
vstore4(convert_int4(mask0), 0, maskout + output_offset);
|
||
vstore4(convert_int4(mask1), 0, maskout + output_offset + query_seq_len4);
|
||
vstore4(convert_int4(mask2), 0, maskout + output_offset + query_seq_len4 + query_seq_len4);
|
||
vstore4(convert_int4(mask3), 0, maskout + output_offset + query_seq_len4 + query_seq_len4 + query_seq_len4);
|
||
#endif
|
||
}
|
||
|
||
__kernel void matmul_qk_div_mask_prefill(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *query, // [batch head_num head_dim_4 query_seq_len_4]
|
||
__global const FLOAT *past_key, // [batch kv_head_num head_dim_4 kv_max_length]
|
||
#ifdef ADD_MASK
|
||
__global const FLOAT* mask,
|
||
#elif defined(SET_MASK)
|
||
__global const int* mask, // [1 1 query_seq_len mask_key_seq_len]
|
||
#else
|
||
__global const FLOAT* mask,
|
||
#endif
|
||
__global FLOAT *qk, // [batch head_num kv_seq_length query_seq_len_4]
|
||
__private const float scale,
|
||
__private const int query_seq_len,
|
||
__private const int mask_key_seq_len,
|
||
__private const int key_seq_len,
|
||
__private const int max_len,
|
||
__private const int head_num,
|
||
__private const int head_dim) {
|
||
|
||
const int x = get_global_id(0); // query_seq_len
|
||
const int y = get_global_id(1); // kv_seq_length
|
||
const int z = get_global_id(2); // head_num * batch
|
||
DEAL_NON_UNIFORM_DIM3(x, y, z);
|
||
|
||
const int x4 = x << 2;
|
||
const int y4 = y << 2;
|
||
|
||
const int query_seq_len4 = (query_seq_len + 3) / 4 * 4;;
|
||
const int query_offset = z * head_dim * query_seq_len4 + x4;
|
||
const int past_offset = (z / NUMHEAD_GROUP_SIZE) * head_dim * max_len + y4;
|
||
float4 out0 = 0, out1 = 0, out2 = 0, out3 = 0;
|
||
|
||
for(int i = 0; i < head_dim / 4; ++i){
|
||
int i4 = i << 2;
|
||
float4 query_vec0 = convert_float4(vload4(0, query + query_offset + i4 * query_seq_len4));
|
||
float4 query_vec1 = convert_float4(vload4(0, query + query_offset + (i4 + 1) * query_seq_len4));
|
||
float4 query_vec2 = convert_float4(vload4(0, query + query_offset + (i4 + 2) * query_seq_len4));
|
||
float4 query_vec3 = convert_float4(vload4(0, query + query_offset + (i4 + 3) * query_seq_len4));
|
||
|
||
float4 past_vec0 = convert_float4(vload4(0, past_key + past_offset + i4 * max_len));
|
||
float4 past_vec1 = convert_float4(vload4(0, past_key + past_offset + (i4 + 1) * max_len));
|
||
float4 past_vec2 = convert_float4(vload4(0, past_key + past_offset + (i4 + 2) * max_len));
|
||
float4 past_vec3 = convert_float4(vload4(0, past_key + past_offset + (i4 + 3) * max_len));
|
||
|
||
out0 = mad((float4)past_vec0.s0, query_vec0, out0);
|
||
out0 = mad((float4)past_vec1.s0, query_vec1, out0);
|
||
out0 = mad((float4)past_vec2.s0, query_vec2, out0);
|
||
out0 = mad((float4)past_vec3.s0, query_vec3, out0);
|
||
|
||
out1 = mad((float4)past_vec0.s1, query_vec0, out1);
|
||
out1 = mad((float4)past_vec1.s1, query_vec1, out1);
|
||
out1 = mad((float4)past_vec2.s1, query_vec2, out1);
|
||
out1 = mad((float4)past_vec3.s1, query_vec3, out1);
|
||
|
||
out2 = mad((float4)past_vec0.s2, query_vec0, out2);
|
||
out2 = mad((float4)past_vec1.s2, query_vec1, out2);
|
||
out2 = mad((float4)past_vec2.s2, query_vec2, out2);
|
||
out2 = mad((float4)past_vec3.s2, query_vec3, out2);
|
||
|
||
out3 = mad((float4)past_vec0.s3, query_vec0, out3);
|
||
out3 = mad((float4)past_vec1.s3, query_vec1, out3);
|
||
out3 = mad((float4)past_vec2.s3, query_vec2, out3);
|
||
out3 = mad((float4)past_vec3.s3, query_vec3, out3);
|
||
}
|
||
out0 *= (float4)scale;
|
||
out1 *= (float4)scale;
|
||
out2 *= (float4)scale;
|
||
out3 *= (float4)scale;
|
||
{
|
||
#if defined(ADD_MASK) || defined(SET_MASK)
|
||
int query_seq_len4 = ((query_seq_len + 3) / 4) * 4;
|
||
int mask_clp = y4 + mask_key_seq_len - key_seq_len;
|
||
int mask_offset = mask_clp * query_seq_len4 + x4;
|
||
float4 mask0 = mask_clp >= 0 && mask_clp < mask_key_seq_len ? convert_float4(vload4(0, mask + mask_offset)) : 0; mask_offset += query_seq_len4;
|
||
float4 mask1 = mask_clp + 1 >= 0 && mask_clp + 1 < mask_key_seq_len? convert_float4(vload4(0, mask + mask_offset)) : 0; mask_offset += query_seq_len4;
|
||
float4 mask2 = mask_clp + 2 >= 0 && mask_clp + 2 < mask_key_seq_len? convert_float4(vload4(0, mask + mask_offset)) : 0; mask_offset += query_seq_len4;
|
||
float4 mask3 = mask_clp + 3 >= 0 && mask_clp + 3 < mask_key_seq_len? convert_float4(vload4(0, mask + mask_offset)) : 0;
|
||
#endif
|
||
|
||
#ifdef ADD_MASK
|
||
out0 += mask0;
|
||
out1 += mask1;
|
||
out2 += mask2;
|
||
out3 += mask3;
|
||
#elif defined(SET_MASK)
|
||
out0 = (mask0 == (float4)0) ? (float4)(-FLT_MAX) : out0;
|
||
out1 = (mask1 == (float4)0) ? (float4)(-FLT_MAX) : out1;
|
||
out2 = (mask2 == (float4)0) ? (float4)(-FLT_MAX) : out2;
|
||
out3 = (mask3 == (float4)0) ? (float4)(-FLT_MAX) : out3;
|
||
#elif defined(DEFAULT_MASK)
|
||
{
|
||
int kv_valid_offset = key_seq_len - query_seq_len;
|
||
int k0 = y4 + 0;
|
||
int k1 = y4 + 1;
|
||
int k2 = y4 + 2;
|
||
int k3 = y4 + 3;
|
||
int q0 = x4 + 0;
|
||
int q1 = x4 + 1;
|
||
int q2 = x4 + 2;
|
||
int q3 = x4 + 3;
|
||
if (k0 > kv_valid_offset + q0) { out0.s0 = -FLT_MAX; }
|
||
if (k1 > kv_valid_offset + q0) { out1.s0 = -FLT_MAX; }
|
||
if (k2 > kv_valid_offset + q0) { out2.s0 = -FLT_MAX; }
|
||
if (k3 > kv_valid_offset + q0) { out3.s0 = -FLT_MAX; }
|
||
if (k0 > kv_valid_offset + q1) { out0.s1 = -FLT_MAX; }
|
||
if (k1 > kv_valid_offset + q1) { out1.s1 = -FLT_MAX; }
|
||
if (k2 > kv_valid_offset + q1) { out2.s1 = -FLT_MAX; }
|
||
if (k3 > kv_valid_offset + q1) { out3.s1 = -FLT_MAX; }
|
||
if (k0 > kv_valid_offset + q2) { out0.s2 = -FLT_MAX; }
|
||
if (k1 > kv_valid_offset + q2) { out1.s2 = -FLT_MAX; }
|
||
if (k2 > kv_valid_offset + q2) { out2.s2 = -FLT_MAX; }
|
||
if (k3 > kv_valid_offset + q2) { out3.s2 = -FLT_MAX; }
|
||
if (k0 > kv_valid_offset + q3) { out0.s3 = -FLT_MAX; }
|
||
if (k1 > kv_valid_offset + q3) { out1.s3 = -FLT_MAX; }
|
||
if (k2 > kv_valid_offset + q3) { out2.s3 = -FLT_MAX; }
|
||
if (k3 > kv_valid_offset + q3) { out3.s3 = -FLT_MAX; }
|
||
}
|
||
#endif
|
||
}
|
||
|
||
const int qk_offset = (z * key_seq_len + y4) * query_seq_len4 + x4;
|
||
vstore4(CONVERT_FLOAT4(out0), 0, qk + qk_offset);
|
||
if(y4 + 1 >= key_seq_len) return;
|
||
vstore4(CONVERT_FLOAT4(out1), 0, qk + qk_offset + query_seq_len4);
|
||
if(y4 + 2 >= key_seq_len) return;
|
||
vstore4(CONVERT_FLOAT4(out2), 0, qk + qk_offset + query_seq_len4 + query_seq_len4);
|
||
if(y4 + 3 >= key_seq_len) return;
|
||
vstore4(CONVERT_FLOAT4(out3), 0, qk + qk_offset + query_seq_len4 + query_seq_len4 + query_seq_len4);
|
||
}
|
||
|
||
__kernel void matmul_qk_decode(GLOBAL_SIZE_2_DIMS
|
||
__global const FLOAT *query, // key [1 head_num head_dim]
|
||
__global const FLOAT *past_key, // [1 head_num head_dim max_length]
|
||
__global FLOAT *qk, // [1 head_num key_seq_len 1]
|
||
__private const float scale,
|
||
__private const int seq_len,
|
||
__private const int max_len,
|
||
__private const int head_num,
|
||
__private const int head_dim) {
|
||
|
||
const int x = get_global_id(0); // key_seq_len
|
||
const int y = get_global_id(1); // head_num
|
||
DEAL_NON_UNIFORM_DIM2(x, y);
|
||
const int x4 = x << 2;
|
||
|
||
const int query_offset = y * head_dim;
|
||
const int past_offset = (y / NUMHEAD_GROUP_SIZE) * head_dim * max_len + x4;
|
||
float4 out0 = 0;
|
||
|
||
for(int i = 0; i < head_dim / 4; ++i){
|
||
int i4 = i << 2;
|
||
float4 query_vec = convert_float4(vload4(0, query + query_offset + i4));
|
||
|
||
float4 past_vec0 = convert_float4(vload4(0, past_key + past_offset + i4 * max_len));
|
||
float4 past_vec1 = convert_float4(vload4(0, past_key + past_offset + (i4 + 1) * max_len));
|
||
float4 past_vec2 = convert_float4(vload4(0, past_key + past_offset + (i4 + 2) * max_len));
|
||
float4 past_vec3 = convert_float4(vload4(0, past_key + past_offset + (i4 + 3) * max_len));
|
||
|
||
out0 = mad((float4)query_vec.s0, past_vec0, out0);
|
||
out0 = mad((float4)query_vec.s1, past_vec1, out0);
|
||
out0 = mad((float4)query_vec.s2, past_vec2, out0);
|
||
out0 = mad((float4)query_vec.s3, past_vec3, out0);
|
||
}
|
||
out0 *= (float4)scale;
|
||
const int qk_offset = y * seq_len + x4;
|
||
if(x4 + 3 < seq_len){
|
||
vstore4(CONVERT_FLOAT4(out0), 0, qk + qk_offset);
|
||
}else {
|
||
int remain = seq_len - x4;
|
||
if(remain == 3){
|
||
vstore3(CONVERT_FLOAT3((float3)(out0.s012)), 0, qk + qk_offset);
|
||
} else if(remain == 2){
|
||
vstore2(CONVERT_FLOAT2((float2)(out0.s01)), 0, qk + qk_offset);
|
||
}else if(remain == 1){
|
||
qk[qk_offset] = out0.s0;
|
||
}
|
||
}
|
||
}
|
||
|
||
__kernel void matmul_qkv_prefill(GLOBAL_SIZE_3_DIMS
|
||
__global const FLOAT *qk, // qk prefill [batch head_num kv_seq_length query_seq_len]
|
||
__global const FLOAT *past_value, // [batch kv_head_num max_len head_dim]
|
||
__global FLOAT *output, // [batch query_seq_len head_num head_dim]
|
||
__private const int query_seq_len,
|
||
__private const int kv_seq_len,
|
||
__private const int max_len,
|
||
__private const int head_num,
|
||
__private const int kv_head_num,
|
||
__private const int head_dim,
|
||
__private const int batch) {
|
||
|
||
const int x = get_global_id(0); // head_dim
|
||
const int y = get_global_id(1); // query_seq_len
|
||
int z = get_global_id(2); // head_num * batch
|
||
|
||
DEAL_NON_UNIFORM_DIM3(x, y, z);
|
||
const int b = z / head_num;
|
||
z = z % head_num;
|
||
const int x8 = x << 3;
|
||
const int y4 = y << 2;
|
||
|
||
const int query_seq_len4 = (query_seq_len + 3) / 4 * 4;
|
||
const int qk_offset = (b * head_num + z) * kv_seq_len * query_seq_len4 + y4;
|
||
const int past_offset = ((b * kv_head_num + z / NUMHEAD_GROUP_SIZE) * max_len) * head_dim + x8;
|
||
const int loop_end = max(kv_seq_len / 4 - 1, 0);
|
||
COMPUTE_FLOAT8 out0 = 0, out1 = 0, out2 = 0, out3 = 0;
|
||
|
||
for(int i = 0; i < loop_end; ++i){
|
||
int i4 = i << 2;
|
||
COMPUTE_FLOAT4 qk_vec0 = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + i4 * query_seq_len4));
|
||
COMPUTE_FLOAT4 qk_vec1 = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + (i4 + 1) * query_seq_len4));
|
||
COMPUTE_FLOAT4 qk_vec2 = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + (i4 + 2) * query_seq_len4));
|
||
COMPUTE_FLOAT4 qk_vec3 = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + (i4 + 3) * query_seq_len4));
|
||
|
||
COMPUTE_FLOAT8 past_vec0 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + i4 * head_dim));
|
||
COMPUTE_FLOAT8 past_vec1 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i4 + 1) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec2 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i4 + 2) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec3 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i4 + 3) * head_dim));
|
||
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec0.s0, past_vec0, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec1.s0, past_vec1, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec2.s0, past_vec2, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec3.s0, past_vec3, out0);
|
||
|
||
out1 = mad((COMPUTE_FLOAT8)qk_vec0.s1, past_vec0, out1);
|
||
out1 = mad((COMPUTE_FLOAT8)qk_vec1.s1, past_vec1, out1);
|
||
out1 = mad((COMPUTE_FLOAT8)qk_vec2.s1, past_vec2, out1);
|
||
out1 = mad((COMPUTE_FLOAT8)qk_vec3.s1, past_vec3, out1);
|
||
|
||
out2 = mad((COMPUTE_FLOAT8)qk_vec0.s2, past_vec0, out2);
|
||
out2 = mad((COMPUTE_FLOAT8)qk_vec1.s2, past_vec1, out2);
|
||
out2 = mad((COMPUTE_FLOAT8)qk_vec2.s2, past_vec2, out2);
|
||
out2 = mad((COMPUTE_FLOAT8)qk_vec3.s2, past_vec3, out2);
|
||
|
||
out3 = mad((COMPUTE_FLOAT8)qk_vec0.s3, past_vec0, out3);
|
||
out3 = mad((COMPUTE_FLOAT8)qk_vec1.s3, past_vec1, out3);
|
||
out3 = mad((COMPUTE_FLOAT8)qk_vec2.s3, past_vec2, out3);
|
||
out3 = mad((COMPUTE_FLOAT8)qk_vec3.s3, past_vec3, out3);
|
||
}
|
||
for(int i = (loop_end << 2); i < kv_seq_len; ++i){
|
||
COMPUTE_FLOAT4 qk_vec = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + i * query_seq_len4));
|
||
COMPUTE_FLOAT8 past_vec = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + i * head_dim));
|
||
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s0, past_vec, out0);
|
||
out1 = mad((COMPUTE_FLOAT8)qk_vec.s1, past_vec, out1);
|
||
out2 = mad((COMPUTE_FLOAT8)qk_vec.s2, past_vec, out2);
|
||
out3 = mad((COMPUTE_FLOAT8)qk_vec.s3, past_vec, out3);
|
||
}
|
||
|
||
#ifdef ATTENTION_C4
|
||
int output_offset = (z * head_dim + x8) * query_seq_len * batch + (b * query_seq_len + y4) * 4;
|
||
const int stride = query_seq_len * batch * 4;
|
||
vstore4(CONVERT_FLOAT4(out0.lo), 0, output + output_offset);
|
||
vstore4(CONVERT_FLOAT4(out0.hi), 0, output + output_offset + stride);
|
||
if(y4 + 1 >= query_seq_len) return;
|
||
output_offset += 4;
|
||
vstore4(CONVERT_FLOAT4(out1.lo), 0, output + output_offset);
|
||
vstore4(CONVERT_FLOAT4(out1.hi), 0, output + output_offset + stride);
|
||
if(y4 + 2 >= query_seq_len) return;
|
||
output_offset += 4;
|
||
vstore4(CONVERT_FLOAT4(out2.lo), 0, output + output_offset);
|
||
vstore4(CONVERT_FLOAT4(out2.hi), 0, output + output_offset + stride);
|
||
if(y4 + 3 >= query_seq_len) return;
|
||
output_offset += 4;
|
||
vstore4(CONVERT_FLOAT4(out3.lo), 0, output + output_offset);
|
||
vstore4(CONVERT_FLOAT4(out3.hi), 0, output + output_offset + stride);
|
||
#else
|
||
const int output_offset = ((b * query_seq_len + y4) * head_num + z) * head_dim + x8;
|
||
const int stride = head_num * head_dim;
|
||
vstore8(CONVERT_FLOAT8(out0), 0, output + output_offset);
|
||
if(y4 + 1 >= query_seq_len) return;
|
||
vstore8(CONVERT_FLOAT8(out1), 0, output + output_offset + stride);
|
||
if(y4 + 2 >= query_seq_len) return;
|
||
vstore8(CONVERT_FLOAT8(out2), 0, output + output_offset + stride + stride);
|
||
if(y4 + 3 >= query_seq_len) return;
|
||
vstore8(CONVERT_FLOAT8(out3), 0, output + output_offset + stride + stride + stride);
|
||
#endif
|
||
}
|
||
|
||
|
||
__kernel void matmul_qkv_decode_b8(GLOBAL_SIZE_2_DIMS
|
||
__global const FLOAT *qk, // qk [1 head_num qk_seq_len 1]
|
||
__global const FLOAT *past_value, // [1 head_num max_len head_dim]
|
||
__global FLOAT *output, // [1 1 head_num head_dim]
|
||
__private const int qk_seq_len,
|
||
__private const int max_len,
|
||
__private const int head_num,
|
||
__private const int kv_head_num,
|
||
__private const int head_dim) {
|
||
|
||
const int x = get_global_id(0); // head_dim
|
||
const int y = get_global_id(1); // head_num
|
||
|
||
DEAL_NON_UNIFORM_DIM2(x, y);
|
||
const int x8 = x << 3;
|
||
|
||
const int qk_offset = y * qk_seq_len;
|
||
const int past_offset = ((y / NUMHEAD_GROUP_SIZE) * max_len) * head_dim + x8;
|
||
COMPUTE_FLOAT8 out0 = 0;
|
||
#ifdef LOOP_UNROLL_4
|
||
const int loop_end = max((qk_seq_len + 3) / 4 - 1, 0);
|
||
for(int i = 0; i < loop_end; ++i){
|
||
int i4 = i << 2;
|
||
COMPUTE_FLOAT4 qk_vec = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + i4));
|
||
|
||
COMPUTE_FLOAT8 past_vec0 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + i4 * head_dim));
|
||
COMPUTE_FLOAT8 past_vec1 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i4 + 1) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec2 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i4 + 2) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec3 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i4 + 3) * head_dim));
|
||
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s0, past_vec0, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s1, past_vec1, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s2, past_vec2, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s3, past_vec3, out0);
|
||
}
|
||
for(int i = (loop_end << 2); i < qk_seq_len; ++i){
|
||
COMPUTE_FLOAT qk_vec = qk[qk_offset + i];
|
||
COMPUTE_FLOAT8 past_vec = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + i * head_dim));
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec, past_vec, out0);
|
||
}
|
||
#elif (defined LOOP_UNROLL_8)
|
||
const int loop_end = max((qk_seq_len + 7) / 8 - 1, 0);
|
||
for(int i = 0; i < loop_end; ++i){
|
||
int i8 = i << 3;
|
||
COMPUTE_FLOAT8 qk_vec = CONVERT_COMPUTE_FLOAT8(vload8(0, qk + qk_offset + i8));
|
||
|
||
COMPUTE_FLOAT8 past_vec0 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + i8 * head_dim));
|
||
COMPUTE_FLOAT8 past_vec1 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 1) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec2 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 2) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec3 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 3) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec4 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 4) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec5 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 5) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec6 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 6) * head_dim));
|
||
COMPUTE_FLOAT8 past_vec7 = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + (i8 + 7) * head_dim));
|
||
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s0, past_vec0, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s1, past_vec1, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s2, past_vec2, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s3, past_vec3, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s4, past_vec4, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s5, past_vec5, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s6, past_vec6, out0);
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec.s7, past_vec7, out0);
|
||
}
|
||
for(int i = (loop_end << 3); i < qk_seq_len; ++i){
|
||
COMPUTE_FLOAT qk_vec = qk[qk_offset + i];
|
||
COMPUTE_FLOAT8 past_vec = CONVERT_COMPUTE_FLOAT8(vload8(0, past_value + past_offset + i * head_dim));
|
||
out0 = mad((COMPUTE_FLOAT8)qk_vec, past_vec, out0);
|
||
}
|
||
#endif
|
||
|
||
const int output_offset = y * head_dim + x8;
|
||
vstore8(CONVERT_FLOAT8(out0), 0, output + output_offset);
|
||
}
|
||
|
||
__kernel void matmul_qkv_decode_b4(GLOBAL_SIZE_2_DIMS
|
||
__global const FLOAT *qk, // qk [1 head_num qk_seq_len 1]
|
||
__global const FLOAT *past_value, // [1 head_num max_len head_dim]
|
||
__global FLOAT *output, // [1 1 head_num head_dim]
|
||
__private const int qk_seq_len,
|
||
__private const int max_len,
|
||
__private const int head_num,
|
||
__private const int kv_head_num,
|
||
__private const int head_dim) {
|
||
|
||
const int x = get_global_id(0); // head_dim
|
||
const int y = get_global_id(1); // head_num
|
||
|
||
DEAL_NON_UNIFORM_DIM2(x, y);
|
||
const int x4 = x << 2;
|
||
|
||
const int qk_offset = y * qk_seq_len;
|
||
const int past_offset = ((y / NUMHEAD_GROUP_SIZE) * max_len) * head_dim + x4;
|
||
COMPUTE_FLOAT4 out0 = 0;
|
||
#ifdef LOOP_UNROLL_4
|
||
const int loop_end = max((qk_seq_len + 3) / 4 - 1, 0);
|
||
for(int i = 0; i < loop_end; ++i){
|
||
int i4 = i << 2;
|
||
COMPUTE_FLOAT4 qk_vec = CONVERT_COMPUTE_FLOAT4(vload4(0, qk + qk_offset + i4));
|
||
|
||
COMPUTE_FLOAT4 past_vec0 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + i4 * head_dim));
|
||
COMPUTE_FLOAT4 past_vec1 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i4 + 1) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec2 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i4 + 2) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec3 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i4 + 3) * head_dim));
|
||
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s0, past_vec0, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s1, past_vec1, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s2, past_vec2, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s3, past_vec3, out0);
|
||
}
|
||
for(int i = (loop_end << 2); i < qk_seq_len; ++i){
|
||
COMPUTE_FLOAT qk_vec = qk[qk_offset + i];
|
||
COMPUTE_FLOAT4 past_vec = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + i * head_dim));
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec, past_vec, out0);
|
||
}
|
||
#elif (defined LOOP_UNROLL_8)
|
||
const int loop_end = max((qk_seq_len + 7) / 8 - 1, 0);
|
||
for(int i = 0; i < loop_end; ++i){
|
||
int i8 = i << 3;
|
||
COMPUTE_FLOAT8 qk_vec = CONVERT_COMPUTE_FLOAT8(vload8(0, qk + qk_offset + i8));
|
||
|
||
COMPUTE_FLOAT4 past_vec0 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + i8 * head_dim));
|
||
COMPUTE_FLOAT4 past_vec1 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 1) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec2 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 2) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec3 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 3) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec4 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 4) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec5 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 5) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec6 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 6) * head_dim));
|
||
COMPUTE_FLOAT4 past_vec7 = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + (i8 + 7) * head_dim));
|
||
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s0, past_vec0, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s1, past_vec1, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s2, past_vec2, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s3, past_vec3, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s4, past_vec4, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s5, past_vec5, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s6, past_vec6, out0);
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec.s7, past_vec7, out0);
|
||
}
|
||
for(int i = (loop_end << 3); i < qk_seq_len; ++i){
|
||
COMPUTE_FLOAT qk_vec = qk[qk_offset + i];
|
||
COMPUTE_FLOAT4 past_vec = CONVERT_COMPUTE_FLOAT4(vload4(0, past_value + past_offset + i * head_dim));
|
||
out0 = mad((COMPUTE_FLOAT4)qk_vec, past_vec, out0);
|
||
}
|
||
#endif
|
||
|
||
const int output_offset = y * head_dim + x4;
|
||
vstore4(CONVERT_FLOAT4(out0), 0, output + output_offset);
|
||
}
|