opencl: general flash attention decode performance optimizations (#25366)
* opencl: vec flash-attention decode kernels for f16/q8_0/q4_0 KV * opencl: improve non FA KQ mv kernels * opencl: tweaks for multiquery FA * opencl: some tweaks for FA q1 kernels * opencl: FA with DK=DV=512 for gemma-4 * opencl: various fixes * opencl: cleanup * opencl: fix FA decode crash for DK=512 (gemma-4) The DK=512 decode-only program does not create the f32_f16 prefill kernel, so the compiled check in ensure_fa_variant never hit and supports_op gave inconsistent answers for the same op. block_n is also unset for DK=512 decode; guard it to avoid an out-of-range read at dispatch. * opencl: run DK=512 FA decode on CPU DK=512 decode is bandwidth-bound and faster on the CPU than the GPU, increasingly so with depth. Decline it in supports_op; prefill stays on the GPU. * opencl: compile MQ_GQA=8 FA kernels in a minimal program The full program compiled with -D MQ_GQA=8 runs the Adreno compiler out of memory at DK>=256. Only the vec_mq kernels are used from this program, so compile it with FA_MQ_ONLY, which excludes everything else. Also include the program name in the compile error log. * opencl: remove stray token in flash_attn_f32_f16.cl A stray "." broke the f32_f16 program build. * opencl: split f16-KV FA decode finer (FD_KV_PER_SPLIT_F16) The 2048 default under-fills the GPU on single-query f16-KV decode; use 512 for f16 KV to get more splits. Quantized KV keeps 2048. --------- Co-authored-by: Li He <lih@qti.qualcomm.com>
This commit is contained in:
Hongqiang Wang
@@ -18,6 +18,14 @@
|
||||
#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
|
||||
#endif
|
||||
|
||||
#ifdef cl_khr_subgroup_shuffle
|
||||
#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
|
||||
#define HAS_SUBGROUP_SHUFFLE 1
|
||||
#elif defined(cl_qcom_subgroup_shuffle)
|
||||
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
|
||||
#define HAS_SUBGROUP_SHUFFLE 1
|
||||
#endif
|
||||
|
||||
// Assumes row size (ne00) is a multiple of 4
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
@@ -378,3 +386,848 @@ kernel void kernel_mul_mat_f16_f32_l4_dr_lq(
|
||||
}
|
||||
}
|
||||
#endif // ADRENO_GPU
|
||||
|
||||
#define N_ROWS_PER_WG 8
|
||||
#define N_OUTS_PER_WG 8
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_ROWS_PER_WG;
|
||||
const int im = get_group_id(2);
|
||||
|
||||
const int i12 = im % ne12;
|
||||
const int i13 = im / ne12;
|
||||
|
||||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||||
|
||||
__local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
|
||||
if (sgs_lid < ne00 / 4) {
|
||||
q_loc[sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_ROWS_PER_WG; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
if (r0 >= ne01) return;
|
||||
|
||||
const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
const half4 k4 = x4[i];
|
||||
const float4 q = q_loc[i];
|
||||
sumf += convert_float(k4.s0) * q.s0
|
||||
+ convert_float(k4.s1) * q.s1
|
||||
+ convert_float(k4.s2) * q.s2
|
||||
+ convert_float(k4.s3) * q.s3;
|
||||
}
|
||||
|
||||
const float all_sum = sub_group_reduce_add(sumf);
|
||||
if (sgs_lid == 0) {
|
||||
dst[im * ne1 * ne0 + r0] = all_sum; // ne11 == 1, so r1==0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_y8(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_OUTS_PER_WG;
|
||||
const int im = get_group_id(2);
|
||||
|
||||
const int i12 = im % ne12;
|
||||
const int i13 = im / ne12;
|
||||
|
||||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||||
|
||||
global half4 * x4_o[N_OUTS_PER_WG];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||||
const ulong off = r0c * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||||
x4_o[o] = (global half4 *)(src0 + off);
|
||||
}
|
||||
|
||||
float sum[N_OUTS_PER_WG] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
|
||||
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
const float4 q4 = y4[i];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||||
const half4 v4 = x4_o[o][i];
|
||||
sum[o] += convert_float(v4.s0) * q4.s0
|
||||
+ convert_float(v4.s1) * q4.s1
|
||||
+ convert_float(v4.s2) * q4.s2
|
||||
+ convert_float(v4.s3) * q4.s3;
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const float s = sub_group_reduce_add(sum[o]);
|
||||
if (sgs_lid == 0 && r0 < ne01) {
|
||||
dst[im * ne1 * ne0 + r0] = s;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_OUTS_PAIR 8
|
||||
#define N_PAIRS_PAIR (N_OUTS_PAIR / 2)
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_pair(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int half_id = sgs_lid >> 5; // 0 = lower half, 1 = upper half
|
||||
const int lane_h = sgs_lid & 31; // lane 0..31 within half
|
||||
|
||||
const int r0_base = get_group_id(0) * N_OUTS_PAIR;
|
||||
const int im = get_group_id(2);
|
||||
|
||||
const int i12 = im % ne12;
|
||||
const int i13 = im / ne12;
|
||||
|
||||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||||
|
||||
__local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
|
||||
if (sgs_lid < ne00 / 4) {
|
||||
q_loc[sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int dk_vec = ne00 / 4;
|
||||
|
||||
#pragma unroll
|
||||
for (int p = 0; p < N_PAIRS_PAIR; ++p) {
|
||||
const int r0 = r0_base + 2 * p + half_id;
|
||||
|
||||
const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||||
|
||||
float sumf = 0.0f;
|
||||
for (int i = lane_h; i < dk_vec; i += 32) {
|
||||
const half4 k4 = x4[i];
|
||||
const float4 q = q_loc[i];
|
||||
sumf += convert_float(k4.s0) * q.s0
|
||||
+ convert_float(k4.s1) * q.s1
|
||||
+ convert_float(k4.s2) * q.s2
|
||||
+ convert_float(k4.s3) * q.s3;
|
||||
}
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 16);
|
||||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_h == 0) {
|
||||
dst[im * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_K_ROWS_GQA 16
|
||||
#define GQA_RATIO_GQA 8
|
||||
#define LANES_PER_QH 8 // 64 / GQA_RATIO_GQA
|
||||
#define DK_VEC_GQA 32 // DK / 4 for DK=128
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
|
||||
const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02; // K-head index (also K2 batch)
|
||||
const int i03 = im_kv / ne02; // n13 batch index
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_GQA;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA]; // 4 × 32 = 128 float4
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
|
||||
if (sgs_lid < DK_VEC_GQA) {
|
||||
q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
// K base offset for this WG. All 8 K-rows × 4 Q-heads share this K-head.
|
||||
const ulong offset_src0_base = (i02) * nb02 + (i03 / r3) * nb03;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
|
||||
const ulong offset_src0 = r0 * nb01 + offset_src0_base;
|
||||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||||
|
||||
float sumf = 0.0f;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < 4; ++t) {
|
||||
const int i = lane_q + t * LANES_PER_QH; // 8, 16, 24-step
|
||||
const half4 k4 = x4[i];
|
||||
const float4 q = q_loc[q_id * DK_VEC_GQA + i];
|
||||
sumf += convert_float(k4.s0) * q.s0
|
||||
+ convert_float(k4.s1) * q.s1
|
||||
+ convert_float(k4.s2) * q.s2
|
||||
+ convert_float(k4.s3) * q.s3;
|
||||
}
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_DV_ROWS_Y8GQA 8
|
||||
#define GQA_RATIO_Y8GQA 8
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_y8_gqa(
|
||||
global char * src0,
|
||||
ulong offset0,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb00,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src0 = (global char *)((global char *)src0 + offset0);
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02; // K-head index
|
||||
const int i03 = im_kv / ne02; // n13 batch index
|
||||
|
||||
// GQA Q-heads sharing this K-head.
|
||||
const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
|
||||
|
||||
global float4 * y4_q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
}
|
||||
|
||||
global half4 * x4_o[N_DV_ROWS_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||||
const ulong off = r0c * nb01 + (i02) * nb02 + (i03 / r3) * nb03;
|
||||
x4_o[o] = (global half4 *)(src0 + off);
|
||||
}
|
||||
|
||||
float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
|
||||
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
// load 8 V values (one per DV row), same K-head, K-pos = i.
|
||||
half4 v[N_DV_ROWS_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
v[o] = x4_o[o][i];
|
||||
}
|
||||
|
||||
// load 8 softmax values (one per Q-head).
|
||||
float4 q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
q[qh] = y4_q[qh][i];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const float4 vf = (float4)(convert_float(v[o].s0),
|
||||
convert_float(v[o].s1),
|
||||
convert_float(v[o].s2),
|
||||
convert_float(v[o].s3));
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
sum[o][qh] += vf.s0 * q[qh].s0
|
||||
+ vf.s1 * q[qh].s1
|
||||
+ vf.s2 * q[qh].s2
|
||||
+ vf.s3 * q[qh].s3;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const float s = sub_group_reduce_add(sum[o][qh]);
|
||||
if (sgs_lid == 0 && r0 < ne01) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + qh);
|
||||
dst[im_out * ne1 * ne0 + r0] = s;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
|
||||
const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_GQA;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
if (sgs_lid < DK_VEC_GQA) {
|
||||
q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 4);
|
||||
const int pitch_px_head = (int)(nb02 >> 4);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||||
|
||||
float sumf = 0.0f;
|
||||
#pragma unroll
|
||||
for (int t = 0; t < 2; ++t) {
|
||||
const int p = lane_q + t * LANES_PER_QH; // pixel idx in row, 0..15
|
||||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||||
const int i0 = 2 * p; // first half4 idx
|
||||
const float4 qa = q_loc[q_id * DK_VEC_GQA + i0 ];
|
||||
const float4 qb = q_loc[q_id * DK_VEC_GQA + i0 + 1];
|
||||
sumf += convert_float(k8.s0) * qa.s0
|
||||
+ convert_float(k8.s1) * qa.s1
|
||||
+ convert_float(k8.s2) * qa.s2
|
||||
+ convert_float(k8.s3) * qa.s3
|
||||
+ convert_float(k8.s4) * qb.s0
|
||||
+ convert_float(k8.s5) * qb.s1
|
||||
+ convert_float(k8.s6) * qb.s2
|
||||
+ convert_float(k8.s7) * qb.s3;
|
||||
}
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_y8_gqa_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int sgs_sz = get_max_sub_group_size();
|
||||
|
||||
const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
|
||||
|
||||
// Q (= softmax(KQ)) base pointers per Q-head
|
||||
global float4 * y4_q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
}
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 3);
|
||||
const int pitch_px_head = (int)(nb02 >> 3);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 3);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
// per-DV-row pixel base
|
||||
int row_px_base[N_DV_ROWS_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||||
row_px_base[o] = r0c * pitch_px_row + head_px_base;
|
||||
}
|
||||
|
||||
float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
|
||||
|
||||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||||
half4 v[N_DV_ROWS_Y8GQA];
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
v[o] = read_imageh(src0_img, row_px_base[o] + i);
|
||||
}
|
||||
|
||||
float4 q[GQA_RATIO_Y8GQA];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
q[qh] = y4_q[qh][i];
|
||||
}
|
||||
// 64 mads.
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const float4 vf = (float4)(convert_float(v[o].s0),
|
||||
convert_float(v[o].s1),
|
||||
convert_float(v[o].s2),
|
||||
convert_float(v[o].s3));
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
sum[o][qh] += vf.s0 * q[qh].s0
|
||||
+ vf.s1 * q[qh].s1
|
||||
+ vf.s2 * q[qh].s2
|
||||
+ vf.s3 * q[qh].s3;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||||
const int r0 = r0_base + o;
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||||
const float s = sub_group_reduce_add(sum[o][qh]);
|
||||
if (sgs_lid == 0 && r0 < ne01) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + qh);
|
||||
dst[im_out * ne1 * ne0 + r0] = s;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_K_ROWS_GQA_R4 16
|
||||
#define GQA_RATIO_R4 4
|
||||
#define LANES_PER_QH_R4 16 // = 64 / GQA_RATIO_R4
|
||||
#define DK_VEC_R4 32 // DK / 4 for DK=128
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 4; // 0..3
|
||||
const int lane_q = sgs_lid & 15; // 0..15
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R4;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_R4;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_R4 * DK_VEC_R4];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_R4; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
if (sgs_lid < DK_VEC_R4) {
|
||||
q_loc[qh * DK_VEC_R4 + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 4);
|
||||
const int pitch_px_head = (int)(nb02 >> 4);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA_R4; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||||
|
||||
const int p = lane_q;
|
||||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||||
const int i0 = 2 * p;
|
||||
const float4 qa = q_loc[q_id * DK_VEC_R4 + i0 ];
|
||||
const float4 qb = q_loc[q_id * DK_VEC_R4 + i0 + 1];
|
||||
|
||||
float sumf =
|
||||
convert_float(k8.s0) * qa.s0
|
||||
+ convert_float(k8.s1) * qa.s1
|
||||
+ convert_float(k8.s2) * qa.s2
|
||||
+ convert_float(k8.s3) * qa.s3
|
||||
+ convert_float(k8.s4) * qb.s0
|
||||
+ convert_float(k8.s5) * qb.s1
|
||||
+ convert_float(k8.s6) * qb.s2
|
||||
+ convert_float(k8.s7) * qb.s3;
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define N_K_ROWS_GQA_R2_DK256 16
|
||||
#define GQA_RATIO_R2 2
|
||||
#define LANES_PER_QH_R2 32 // = 64 / GQA_RATIO_R2
|
||||
#define DK_VEC_DK256 64 // DK / 4 for DK=256
|
||||
|
||||
#ifdef ADRENO_GPU
|
||||
REQD_SUBGROUP_SIZE_64
|
||||
#endif
|
||||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img(
|
||||
__read_only image1d_buffer_t src0_img,
|
||||
global char * src1,
|
||||
ulong offset1,
|
||||
global float * dst,
|
||||
ulong offsetd,
|
||||
int ne00,
|
||||
int ne01,
|
||||
int ne02,
|
||||
ulong nb01,
|
||||
ulong nb02,
|
||||
ulong nb03,
|
||||
int ne10,
|
||||
int ne11,
|
||||
int ne12,
|
||||
ulong nb10,
|
||||
ulong nb11,
|
||||
ulong nb12,
|
||||
ulong nb13,
|
||||
int ne0,
|
||||
int ne1,
|
||||
int r2,
|
||||
int r3
|
||||
) {
|
||||
src1 = (global char *)((global char *)src1 + offset1);
|
||||
dst = (global float*)((global char *)dst + offsetd);
|
||||
|
||||
const int sgs_lid = get_sub_group_local_id();
|
||||
const int q_id = sgs_lid >> 5; // 0..1
|
||||
const int lane_q = sgs_lid & 31; // 0..31
|
||||
|
||||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R2_DK256;
|
||||
const int im_kv = get_group_id(2);
|
||||
|
||||
const int i02 = im_kv % ne02;
|
||||
const int i03 = im_kv / ne02;
|
||||
|
||||
const int q_head_lo = i02 * GQA_RATIO_R2;
|
||||
|
||||
__local float4 q_loc[GQA_RATIO_R2 * DK_VEC_DK256];
|
||||
#pragma unroll
|
||||
for (int qh = 0; qh < GQA_RATIO_R2; ++qh) {
|
||||
const int qh_idx = q_head_lo + qh;
|
||||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||||
q_loc[qh * DK_VEC_DK256 + sgs_lid] = y4[sgs_lid];
|
||||
}
|
||||
barrier(CLK_LOCAL_MEM_FENCE);
|
||||
|
||||
const int pitch_px_row = (int)(nb01 >> 4);
|
||||
const int pitch_px_head = (int)(nb02 >> 4);
|
||||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||||
|
||||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||||
|
||||
#pragma unroll
|
||||
for (int dr = 0; dr < N_K_ROWS_GQA_R2_DK256; ++dr) {
|
||||
const int r0 = r0_base + dr;
|
||||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||||
|
||||
const int p = lane_q;
|
||||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||||
const int i0 = 2 * p;
|
||||
const float4 qa = q_loc[q_id * DK_VEC_DK256 + i0 ];
|
||||
const float4 qb = q_loc[q_id * DK_VEC_DK256 + i0 + 1];
|
||||
|
||||
float sumf =
|
||||
convert_float(k8.s0) * qa.s0
|
||||
+ convert_float(k8.s1) * qa.s1
|
||||
+ convert_float(k8.s2) * qa.s2
|
||||
+ convert_float(k8.s3) * qa.s3
|
||||
+ convert_float(k8.s4) * qb.s0
|
||||
+ convert_float(k8.s5) * qb.s1
|
||||
+ convert_float(k8.s6) * qb.s2
|
||||
+ convert_float(k8.s7) * qb.s3;
|
||||
|
||||
sumf += sub_group_shuffle_xor(sumf, 16);
|
||||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||||
|
||||
if (lane_q == 0) {
|
||||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user