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:
@@ -20,6 +20,7 @@ static const ggml_opencl_fa_dim g_fa_dims_adreno_default[] = {
|
|||||||
{192, 128, 16, 16, 1, 0},
|
{192, 128, 16, 16, 1, 0},
|
||||||
{192, 192, 16, 16, 1, 0},
|
{192, 192, 16, 16, 1, 0},
|
||||||
{256, 256, 16, 16, 16, 0},
|
{256, 256, 16, 16, 16, 0},
|
||||||
|
{512, 512, 8, 16, 64, 0},
|
||||||
};
|
};
|
||||||
|
|
||||||
struct ggml_opencl_fa_dim_table {
|
struct ggml_opencl_fa_dim_table {
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -10,7 +10,12 @@
|
|||||||
#define DK_VEC (DK/4)
|
#define DK_VEC (DK/4)
|
||||||
#define DV_VEC (DV/4)
|
#define DV_VEC (DV/4)
|
||||||
#define WG_SIZE (BLOCK_M)
|
#define WG_SIZE (BLOCK_M)
|
||||||
#define Q1_WG_SIZE 64
|
// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
|
||||||
|
// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
|
||||||
|
#ifndef FA_SG
|
||||||
|
#define FA_SG 64
|
||||||
|
#endif
|
||||||
|
#define Q1_WG_SIZE FA_SG
|
||||||
|
|
||||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||||
|
|||||||
@@ -11,7 +11,12 @@
|
|||||||
#define DK_VEC (DK/4)
|
#define DK_VEC (DK/4)
|
||||||
#define DV_VEC (DV/4)
|
#define DV_VEC (DV/4)
|
||||||
#define WG_SIZE (BLOCK_M)
|
#define WG_SIZE (BLOCK_M)
|
||||||
#define Q1_WG_SIZE 64
|
// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
|
||||||
|
// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
|
||||||
|
#ifndef FA_SG
|
||||||
|
#define FA_SG 64
|
||||||
|
#endif
|
||||||
|
#define Q1_WG_SIZE FA_SG
|
||||||
|
|
||||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||||
@@ -114,6 +119,15 @@ __kernel void flash_attn_f32(
|
|||||||
__local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
|
__local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
|
||||||
|
|
||||||
for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
|
for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
|
||||||
|
#if FA_SG < 64
|
||||||
|
// WAR on l_k/l_v: threads with my_query_row >= n_q skip the compute below
|
||||||
|
// (continue) and would race ahead to reload the tiles while active threads
|
||||||
|
// still read them. A single 64-wide Adreno subgroup (WG == sg) runs lockstep
|
||||||
|
// and hides this; a WG that spans multiple narrower subgroups (Intel sg=32)
|
||||||
|
// corrupts the result. All threads reach this each iteration (no-op on the
|
||||||
|
// first), so it does not diverge with the continue. Compiled out at sg=64.
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
#endif
|
||||||
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
|
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
|
||||||
const int row = i / DK_VEC;
|
const int row = i / DK_VEC;
|
||||||
const int col = i % DK_VEC;
|
const int col = i % DK_VEC;
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -27,7 +27,11 @@
|
|||||||
|
|
||||||
#define DK_VEC (DK/4)
|
#define DK_VEC (DK/4)
|
||||||
#define DV_VEC (DV/4)
|
#define DV_VEC (DV/4)
|
||||||
#define Q1_WG_SIZE 64
|
|
||||||
|
#ifndef FA_SG
|
||||||
|
#define FA_SG 64
|
||||||
|
#endif
|
||||||
|
#define Q1_WG_SIZE FA_SG
|
||||||
|
|
||||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||||
@@ -365,6 +369,263 @@ __kernel void flash_attn_f32_q4_0_q1(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#ifdef cl_intel_subgroups
|
||||||
|
#pragma OPENCL EXTENSION cl_intel_subgroups : enable
|
||||||
|
#else
|
||||||
|
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#ifdef cl_qcom_reqd_sub_group_size
|
||||||
|
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
|
||||||
|
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
|
||||||
|
#else
|
||||||
|
#define REQD_SUBGROUP_SIZE_64
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#define VEC_NSG 4
|
||||||
|
#define VEC_WG_SIZE (Q1_WG_SIZE * VEC_NSG)
|
||||||
|
#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
|
||||||
|
|
||||||
|
// Dequant one float4 lane (0..7) from a q4_0 block.
|
||||||
|
// Lanes 0..3 → low nibbles of qs[0..15], lanes 4..7 → high nibbles.
|
||||||
|
inline float4 dequant_q4_0_lane(const global char * block_ptr, int lane) {
|
||||||
|
const float d = vload_half(0, (const global half *)block_ptr);
|
||||||
|
const global uchar * qs = (const global uchar *)(block_ptr + 2);
|
||||||
|
const int g = lane & 3;
|
||||||
|
const int shift = (lane < 4) ? 0 : 4;
|
||||||
|
return d * (float4)((float)((qs[g*4+0] >> shift) & 0x0F) - 8.0f,
|
||||||
|
(float)((qs[g*4+1] >> shift) & 0x0F) - 8.0f,
|
||||||
|
(float)((qs[g*4+2] >> shift) & 0x0F) - 8.0f,
|
||||||
|
(float)((qs[g*4+3] >> shift) & 0x0F) - 8.0f);
|
||||||
|
}
|
||||||
|
|
||||||
|
REQD_SUBGROUP_SIZE_64
|
||||||
|
__kernel void flash_attn_f32_q4_0_q1_vec(
|
||||||
|
const global void * q_void, ulong q_offset,
|
||||||
|
const global void * k_void, ulong k_offset,
|
||||||
|
const global void * v_void, ulong v_offset,
|
||||||
|
global void * o_void, ulong o_offset,
|
||||||
|
const float scale,
|
||||||
|
const int n_q,
|
||||||
|
const int n_kv,
|
||||||
|
const int is_causal,
|
||||||
|
const int n_head,
|
||||||
|
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||||
|
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||||
|
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||||
|
const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
|
||||||
|
const float max_bias,
|
||||||
|
const float m0,
|
||||||
|
const float m1,
|
||||||
|
const int n_head_log2,
|
||||||
|
const float logit_softcap,
|
||||||
|
const int n_head_kv,
|
||||||
|
const global void* mask_void,
|
||||||
|
const ulong mask_offset,
|
||||||
|
const ulong mask_nb1,
|
||||||
|
const ulong mask_nb2,
|
||||||
|
const ulong mask_nb3,
|
||||||
|
const int mask_ne2,
|
||||||
|
const int mask_ne3,
|
||||||
|
const global void* sinks_void,
|
||||||
|
const ulong sinks_offset
|
||||||
|
) {
|
||||||
|
const int tid = get_local_id(0);
|
||||||
|
const int sgid = tid / Q1_WG_SIZE;
|
||||||
|
const int tid_sg = tid % Q1_WG_SIZE;
|
||||||
|
const int head_batch_idx = get_global_id(1);
|
||||||
|
|
||||||
|
const int batch_idx = head_batch_idx / n_head;
|
||||||
|
const int head_idx = head_batch_idx % n_head;
|
||||||
|
|
||||||
|
const int gqa_ratio = n_head / n_head_kv;
|
||||||
|
const int head_kv_idx = head_idx / gqa_ratio;
|
||||||
|
|
||||||
|
const global char * q_base = (const global char *) q_void + q_offset;
|
||||||
|
const global char * k_base = (const global char *) k_void + k_offset;
|
||||||
|
const global char * v_base = (const global char *) v_void + v_offset;
|
||||||
|
global char * o_base = (global char *) o_void + o_offset;
|
||||||
|
|
||||||
|
const global char * mask_base = NULL;
|
||||||
|
if (mask_void != NULL) {
|
||||||
|
const int mask_head_idx = head_idx % mask_ne2;
|
||||||
|
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||||
|
mask_base = (const global char *) mask_void + mask_offset +
|
||||||
|
mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
|
||||||
|
}
|
||||||
|
|
||||||
|
__local ACC_TYPE4 q_shared[DK_VEC];
|
||||||
|
{
|
||||||
|
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
|
||||||
|
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||||
|
for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
|
||||||
|
q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
#ifdef FA_HAVE_INT_DOT
|
||||||
|
// quantize Q to int8-packed uints + per-block (qd, q_sum) once per WG for dp4a
|
||||||
|
// one thread per Q block, remaining threads idle this step
|
||||||
|
__local uint q_packed_shared[DK_Q4_BLOCKS * 8];
|
||||||
|
__local float q_d_shared[DK_Q4_BLOCKS];
|
||||||
|
__local int q_sum_shared[DK_Q4_BLOCKS];
|
||||||
|
if (tid < DK_Q4_BLOCKS) {
|
||||||
|
ACC_TYPE4 q_block[8];
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < 8; ++i) q_block[i] = q_shared[tid * 8 + i];
|
||||||
|
uint packed[8];
|
||||||
|
q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < 8; ++i) q_packed_shared[tid * 8 + i] = packed[i];
|
||||||
|
q_d_shared[tid] = info.qd;
|
||||||
|
q_sum_shared[tid] = info.q_sum;
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
#endif
|
||||||
|
|
||||||
|
const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
|
||||||
|
|
||||||
|
const global ACC_TYPE * sinks_ptr = NULL;
|
||||||
|
if (sinks_void != NULL) {
|
||||||
|
sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
|
||||||
|
|
||||||
|
ACC_TYPE m_i = FA_M_INIT;
|
||||||
|
ACC_TYPE l_i = 0.0f;
|
||||||
|
|
||||||
|
const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
|
||||||
|
const int kv_start = sgid * kv_per_sg;
|
||||||
|
const int kv_end = min(n_kv, kv_start + kv_per_sg);
|
||||||
|
|
||||||
|
for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
|
||||||
|
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||||
|
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||||
|
|
||||||
|
#ifdef FA_HAVE_INT_DOT
|
||||||
|
// per-lane dp4a: each lane packs 4 raw q4_0 nibbles into a uint,
|
||||||
|
// then dot_acc_sat_4x8packed_ss_int against the matching uint.
|
||||||
|
ACC_TYPE lane_contrib = 0.0f;
|
||||||
|
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||||
|
const int block_idx = qk / 8;
|
||||||
|
const int lane_in_block = qk % 8;
|
||||||
|
const int g = lane_in_block & 3;
|
||||||
|
const int shift = (lane_in_block < 4) ? 0 : 4;
|
||||||
|
const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
|
||||||
|
const float kd = vload_half(0, (const global half *)k_block);
|
||||||
|
const global uchar * k_qs = (const global uchar *)(k_block + 2);
|
||||||
|
const uchar b0 = k_qs[g*4 + 0];
|
||||||
|
const uchar b1 = k_qs[g*4 + 1];
|
||||||
|
const uchar b2 = k_qs[g*4 + 2];
|
||||||
|
const uchar b3 = k_qs[g*4 + 3];
|
||||||
|
const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
|
||||||
|
((uint)((b1 >> shift) & 0x0F)) << 8 |
|
||||||
|
((uint)((b2 >> shift) & 0x0F)) << 16 |
|
||||||
|
((uint)((b3 >> shift) & 0x0F)) << 24;
|
||||||
|
const uint q_packed_lane = q_packed_shared[block_idx * 8 + lane_in_block];
|
||||||
|
const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
|
||||||
|
const float qd = q_d_shared[block_idx];
|
||||||
|
const float block_scale = qd * kd;
|
||||||
|
float contrib = (float)raw_dot * block_scale;
|
||||||
|
if (lane_in_block == 0) {
|
||||||
|
// block bias correction is per-block
|
||||||
|
const int q_sum_b = q_sum_shared[block_idx];
|
||||||
|
contrib -= 8.0f * block_scale * (float)q_sum_b;
|
||||||
|
}
|
||||||
|
lane_contrib += contrib;
|
||||||
|
}
|
||||||
|
ACC_TYPE score = sub_group_reduce_add(lane_contrib) * scale;
|
||||||
|
#else
|
||||||
|
ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
|
||||||
|
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||||
|
const int block_idx = qk / 8;
|
||||||
|
const int lane = qk % 8;
|
||||||
|
const float4 k_v = dequant_q4_0_lane(k_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||||
|
dot4 = mad(q_shared[qk], k_v, dot4);
|
||||||
|
}
|
||||||
|
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
|
||||||
|
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
|
||||||
|
#endif
|
||||||
|
|
||||||
|
if (mask_base != NULL) {
|
||||||
|
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
|
||||||
|
score += slope * (ACC_TYPE) mask_ptr[k_idx];
|
||||||
|
}
|
||||||
|
if (logit_softcap > 0.0f) {
|
||||||
|
score = logit_softcap * tanh(score / logit_softcap);
|
||||||
|
}
|
||||||
|
|
||||||
|
const ACC_TYPE m_new = max(m_i, score);
|
||||||
|
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
|
||||||
|
const ACC_TYPE p = native_exp(score - m_new);
|
||||||
|
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
const int block_idx = dv / 8;
|
||||||
|
const int lane = dv % 8;
|
||||||
|
const float4 v_v = dequant_q4_0_lane(v_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||||
|
o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
|
||||||
|
}
|
||||||
|
l_i = l_i * scale_prev + p;
|
||||||
|
m_i = m_new;
|
||||||
|
}
|
||||||
|
|
||||||
|
__local ACC_TYPE sg_m[VEC_NSG];
|
||||||
|
__local ACC_TYPE sg_l[VEC_NSG];
|
||||||
|
__local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
|
||||||
|
|
||||||
|
if (tid_sg == 0) {
|
||||||
|
sg_m[sgid] = m_i;
|
||||||
|
sg_l[sgid] = l_i;
|
||||||
|
}
|
||||||
|
{
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
sg_o[sgid][dv] = o_acc[idx];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
if (sgid == 0) {
|
||||||
|
ACC_TYPE m_final = sg_m[0];
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 1; s < VEC_NSG; ++s) {
|
||||||
|
m_final = max(m_final, sg_m[s]);
|
||||||
|
}
|
||||||
|
if (sinks_ptr != NULL) {
|
||||||
|
m_final = max(m_final, sinks_ptr[head_idx]);
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE l_final = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < VEC_NSG; ++s) {
|
||||||
|
l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
|
||||||
|
}
|
||||||
|
if (sinks_ptr != NULL) {
|
||||||
|
l_final += native_exp(sinks_ptr[head_idx] - m_final);
|
||||||
|
}
|
||||||
|
const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
|
||||||
|
|
||||||
|
const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
|
||||||
|
global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
|
||||||
|
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < VEC_NSG; ++s) {
|
||||||
|
const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
|
||||||
|
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
|
||||||
|
}
|
||||||
|
o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Flash-decoding split pass for q4_0 KV. Merge kernel is type-agnostic and
|
// Flash-decoding split pass for q4_0 KV. Merge kernel is type-agnostic and
|
||||||
// shared with the f16/q8_0 FA kernels.
|
// shared with the f16/q8_0 FA kernels.
|
||||||
#define FA_PARTIAL_FLOATS (2 + DV)
|
#define FA_PARTIAL_FLOATS (2 + DV)
|
||||||
@@ -583,6 +844,319 @@ __kernel void flash_attn_f32_q4_0_q1_split(
|
|||||||
#define WG_SIZE BLOCK_M
|
#define WG_SIZE BLOCK_M
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
#ifndef MQ_GQA
|
||||||
|
#define MQ_GQA 4
|
||||||
|
#endif
|
||||||
|
#ifndef MQ_NSG_SPLIT
|
||||||
|
#define MQ_NSG_SPLIT 4
|
||||||
|
#endif
|
||||||
|
#define MQ_SPLIT_WG_SIZE_Q4 (Q1_WG_SIZE * MQ_NSG_SPLIT)
|
||||||
|
|
||||||
|
REQD_SUBGROUP_SIZE_64
|
||||||
|
__kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
|
||||||
|
const global void * q_void, ulong q_offset,
|
||||||
|
const global void * k_void, ulong k_offset,
|
||||||
|
const global void * v_void, ulong v_offset,
|
||||||
|
const float scale,
|
||||||
|
const int n_q,
|
||||||
|
const int n_kv,
|
||||||
|
const int n_head,
|
||||||
|
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||||
|
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||||
|
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||||
|
const float max_bias,
|
||||||
|
const float m0,
|
||||||
|
const float m1,
|
||||||
|
const int n_head_log2,
|
||||||
|
const float logit_softcap,
|
||||||
|
const int n_head_kv,
|
||||||
|
const global void * mask_void,
|
||||||
|
const ulong mask_offset,
|
||||||
|
const ulong mask_nb1,
|
||||||
|
const ulong mask_nb2,
|
||||||
|
const ulong mask_nb3,
|
||||||
|
const int mask_ne2,
|
||||||
|
const int mask_ne3,
|
||||||
|
global float * partial_void,
|
||||||
|
const int n_splits,
|
||||||
|
const int kv_per_split
|
||||||
|
) {
|
||||||
|
const int tid = get_local_id(0);
|
||||||
|
const int sgid = tid / Q1_WG_SIZE;
|
||||||
|
const int tid_sg = tid % Q1_WG_SIZE;
|
||||||
|
const int kvhead_batch_idx = get_global_id(1);
|
||||||
|
const int split_q_idx = get_global_id(2);
|
||||||
|
const int split_idx = split_q_idx % n_splits;
|
||||||
|
const int q_idx = split_q_idx / n_splits;
|
||||||
|
|
||||||
|
const int batch_idx = kvhead_batch_idx / n_head_kv;
|
||||||
|
const int head_kv_idx = kvhead_batch_idx % n_head_kv;
|
||||||
|
|
||||||
|
const int kv_start = split_idx * kv_per_split;
|
||||||
|
const int kv_end = min(kv_start + kv_per_split, n_kv);
|
||||||
|
|
||||||
|
const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
|
||||||
|
|
||||||
|
if (kv_start >= kv_end) {
|
||||||
|
if (tid == 0) {
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||||
|
* n_splits + split_idx);
|
||||||
|
global float * rec = partial_void + rec_idx * record_stride;
|
||||||
|
rec[0] = FA_M_INIT;
|
||||||
|
rec[1] = 0.0f;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const global char * q_base = (const global char *) q_void + q_offset;
|
||||||
|
const global char * k_base = (const global char *) k_void + k_offset;
|
||||||
|
const global char * v_base = (const global char *) v_void + v_offset;
|
||||||
|
|
||||||
|
__local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
|
||||||
|
for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q4) {
|
||||||
|
const int h = i / DK_VEC;
|
||||||
|
const int k = i % DK_VEC;
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
|
||||||
|
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||||
|
q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
#ifdef FA_HAVE_INT_DOT
|
||||||
|
__local uint q_packed_shared[MQ_GQA * DK_Q4_BLOCKS * 8];
|
||||||
|
__local float q_d_shared[MQ_GQA * DK_Q4_BLOCKS];
|
||||||
|
__local int q_sum_shared[MQ_GQA * DK_Q4_BLOCKS];
|
||||||
|
{
|
||||||
|
const int active = MQ_GQA * DK_Q4_BLOCKS;
|
||||||
|
if (tid < active) {
|
||||||
|
const int h = tid / DK_Q4_BLOCKS;
|
||||||
|
const int block_id = tid % DK_Q4_BLOCKS;
|
||||||
|
ACC_TYPE4 q_block[8];
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < 8; ++i) q_block[i] = q_shared[h * DK_VEC + block_id * 8 + i];
|
||||||
|
uint packed[8];
|
||||||
|
q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < 8; ++i) q_packed_shared[(h * DK_Q4_BLOCKS + block_id) * 8 + i] = packed[i];
|
||||||
|
q_d_shared[h * DK_Q4_BLOCKS + block_id] = info.qd;
|
||||||
|
q_sum_shared[h * DK_Q4_BLOCKS + block_id] = info.q_sum;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
#endif
|
||||||
|
|
||||||
|
float slope[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
|
||||||
|
}
|
||||||
|
|
||||||
|
const global char * mask_base[MQ_GQA];
|
||||||
|
if (mask_void != NULL) {
|
||||||
|
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||||
|
const global char * mask_base_b = (const global char *) mask_void + mask_offset +
|
||||||
|
mask_batch_idx * mask_nb3 +
|
||||||
|
(ulong) q_idx * mask_nb1;
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
const int mask_head_idx = head_idx % mask_ne2;
|
||||||
|
mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
|
||||||
|
ACC_TYPE m_i[MQ_GQA];
|
||||||
|
ACC_TYPE l_i[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
m_i[h] = FA_M_INIT;
|
||||||
|
l_i[h] = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
|
||||||
|
}
|
||||||
|
|
||||||
|
const int kv_len = kv_end - kv_start;
|
||||||
|
const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
|
||||||
|
const int kv_lo = kv_start + sgid * kv_per_sg;
|
||||||
|
const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
|
||||||
|
|
||||||
|
for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
|
||||||
|
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||||
|
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||||
|
|
||||||
|
#ifdef FA_HAVE_INT_DOT
|
||||||
|
ACC_TYPE lane_contrib[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
|
||||||
|
|
||||||
|
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||||
|
const int block_idx = qk / 8;
|
||||||
|
const int lane_in_block = qk % 8;
|
||||||
|
const int g = lane_in_block & 3;
|
||||||
|
const int shift = (lane_in_block < 4) ? 0 : 4;
|
||||||
|
const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
|
||||||
|
const float kd = vload_half(0, (const global half *)k_block);
|
||||||
|
const global uchar * k_qs = (const global uchar *)(k_block + 2);
|
||||||
|
const uchar b0 = k_qs[g*4 + 0];
|
||||||
|
const uchar b1 = k_qs[g*4 + 1];
|
||||||
|
const uchar b2 = k_qs[g*4 + 2];
|
||||||
|
const uchar b3 = k_qs[g*4 + 3];
|
||||||
|
const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
|
||||||
|
((uint)((b1 >> shift) & 0x0F)) << 8 |
|
||||||
|
((uint)((b2 >> shift) & 0x0F)) << 16 |
|
||||||
|
((uint)((b3 >> shift) & 0x0F)) << 24;
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const uint q_packed_lane = q_packed_shared[(h * DK_Q4_BLOCKS + block_idx) * 8 + lane_in_block];
|
||||||
|
const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
|
||||||
|
const float qd = q_d_shared[h * DK_Q4_BLOCKS + block_idx];
|
||||||
|
const float block_scale = qd * kd;
|
||||||
|
float contrib = (float) raw_dot * block_scale;
|
||||||
|
if (lane_in_block == 0) {
|
||||||
|
const int q_sum_b = q_sum_shared[h * DK_Q4_BLOCKS + block_idx];
|
||||||
|
contrib -= 8.0f * block_scale * (float) q_sum_b;
|
||||||
|
}
|
||||||
|
lane_contrib[h] += contrib;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE score[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * scale;
|
||||||
|
if (mask_base[h] != NULL) {
|
||||||
|
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
|
||||||
|
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
|
||||||
|
}
|
||||||
|
if (logit_softcap > 0.0f) {
|
||||||
|
s = logit_softcap * tanh(s / logit_softcap);
|
||||||
|
}
|
||||||
|
score[h] = s;
|
||||||
|
}
|
||||||
|
#else
|
||||||
|
// fallback float-dequant K dot
|
||||||
|
ACC_TYPE4 dot4[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
|
||||||
|
|
||||||
|
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||||
|
const int block_idx = qk / 8;
|
||||||
|
const int lane = qk % 8;
|
||||||
|
const float4 k_v = dequant_q4_0_lane(k_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE score[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
|
||||||
|
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
|
||||||
|
if (mask_base[h] != NULL) {
|
||||||
|
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
|
||||||
|
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
|
||||||
|
}
|
||||||
|
if (logit_softcap > 0.0f) {
|
||||||
|
s = logit_softcap * tanh(s / logit_softcap);
|
||||||
|
}
|
||||||
|
score[h] = s;
|
||||||
|
}
|
||||||
|
#endif
|
||||||
|
|
||||||
|
ACC_TYPE p_h[MQ_GQA];
|
||||||
|
ACC_TYPE sp_h[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const ACC_TYPE m_new = max(m_i[h], score[h]);
|
||||||
|
sp_h[h] = native_exp(m_i[h] - m_new);
|
||||||
|
p_h[h] = native_exp(score[h] - m_new);
|
||||||
|
l_i[h] = l_i[h] * sp_h[h] + p_h[h];
|
||||||
|
m_i[h] = m_new;
|
||||||
|
}
|
||||||
|
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
const int block_idx = dv / 8;
|
||||||
|
const int lane = dv % 8;
|
||||||
|
const float4 v_v = dequant_q4_0_lane(v_row + block_idx * Q4_0_BLOCK_SIZE, lane);
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// per-h cross-subgroup merge
|
||||||
|
__local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
|
||||||
|
__local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
|
||||||
|
__local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
|
||||||
|
|
||||||
|
if (tid_sg == 0) {
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
sg_m[h][sgid] = m_i[h];
|
||||||
|
sg_l[h][sgid] = l_i[h];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
{
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
|
||||||
|
sg_o[sgid][dv_idx] = o_acc[h][idx];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
if (sgid == 0) {
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
|
||||||
|
ACC_TYPE m_c = sg_m[h][0];
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
|
||||||
|
m_c = max(m_c, sg_m[h][s]);
|
||||||
|
}
|
||||||
|
ACC_TYPE l_c = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||||
|
l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
|
||||||
|
}
|
||||||
|
|
||||||
|
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||||
|
* n_splits + split_idx);
|
||||||
|
global float * rec = partial_void + rec_idx * record_stride;
|
||||||
|
global float4 * rec_o = (global float4 *) (rec + 2);
|
||||||
|
|
||||||
|
if (tid_sg == 0) {
|
||||||
|
rec[0] = (float) m_c;
|
||||||
|
rec[1] = (float) l_c;
|
||||||
|
}
|
||||||
|
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
|
||||||
|
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||||
|
const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
|
||||||
|
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
|
||||||
|
}
|
||||||
|
rec_o[dv_idx] = o_merged;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
__kernel void flash_attn_f32_q4_0(
|
__kernel void flash_attn_f32_q4_0(
|
||||||
const global void * q_void, ulong q_offset,
|
const global void * q_void, ulong q_offset,
|
||||||
const global void * k_void, ulong k_offset,
|
const global void * k_void, ulong k_offset,
|
||||||
|
|||||||
@@ -24,7 +24,11 @@
|
|||||||
|
|
||||||
#define DK_VEC (DK/4)
|
#define DK_VEC (DK/4)
|
||||||
#define DV_VEC (DV/4)
|
#define DV_VEC (DV/4)
|
||||||
#define Q1_WG_SIZE 64
|
|
||||||
|
#ifndef FA_SG
|
||||||
|
#define FA_SG 64
|
||||||
|
#endif
|
||||||
|
#define Q1_WG_SIZE FA_SG
|
||||||
|
|
||||||
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
|
||||||
// infinite operand can cause undefined behavior and miscompilation for exp.
|
// infinite operand can cause undefined behavior and miscompilation for exp.
|
||||||
@@ -310,6 +314,201 @@ __kernel void flash_attn_f32_q8_0_q1(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#ifdef cl_intel_subgroups
|
||||||
|
#pragma OPENCL EXTENSION cl_intel_subgroups : enable
|
||||||
|
#else
|
||||||
|
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#ifdef cl_qcom_reqd_sub_group_size
|
||||||
|
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
|
||||||
|
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
|
||||||
|
#else
|
||||||
|
#define REQD_SUBGROUP_SIZE_64
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#define VEC_NSG 4
|
||||||
|
#define VEC_WG_SIZE (Q1_WG_SIZE * VEC_NSG)
|
||||||
|
#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
|
||||||
|
|
||||||
|
inline float4 dequant_q8_0_lane(const global char * block_ptr, int lane) {
|
||||||
|
const float d = vload_half(0, (const global half *)block_ptr);
|
||||||
|
const global char * qs = block_ptr + 2 + lane * 4;
|
||||||
|
return d * (float4)((float)qs[0], (float)qs[1], (float)qs[2], (float)qs[3]);
|
||||||
|
}
|
||||||
|
|
||||||
|
REQD_SUBGROUP_SIZE_64
|
||||||
|
__kernel void flash_attn_f32_q8_0_q1_vec(
|
||||||
|
const global void * q_void, ulong q_offset,
|
||||||
|
const global void * k_void, ulong k_offset,
|
||||||
|
const global void * v_void, ulong v_offset,
|
||||||
|
global void * o_void, ulong o_offset,
|
||||||
|
const float scale,
|
||||||
|
const int n_q,
|
||||||
|
const int n_kv,
|
||||||
|
const int is_causal,
|
||||||
|
const int n_head,
|
||||||
|
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||||
|
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||||
|
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||||
|
const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
|
||||||
|
const float max_bias,
|
||||||
|
const float m0,
|
||||||
|
const float m1,
|
||||||
|
const int n_head_log2,
|
||||||
|
const float logit_softcap,
|
||||||
|
const int n_head_kv,
|
||||||
|
const global void* mask_void,
|
||||||
|
const ulong mask_offset,
|
||||||
|
const ulong mask_nb1,
|
||||||
|
const ulong mask_nb2,
|
||||||
|
const ulong mask_nb3,
|
||||||
|
const int mask_ne2,
|
||||||
|
const int mask_ne3,
|
||||||
|
const global void* sinks_void,
|
||||||
|
const ulong sinks_offset
|
||||||
|
) {
|
||||||
|
const int tid = get_local_id(0);
|
||||||
|
const int sgid = tid / Q1_WG_SIZE;
|
||||||
|
const int tid_sg = tid % Q1_WG_SIZE;
|
||||||
|
const int head_batch_idx = get_global_id(1);
|
||||||
|
|
||||||
|
const int batch_idx = head_batch_idx / n_head;
|
||||||
|
const int head_idx = head_batch_idx % n_head;
|
||||||
|
|
||||||
|
const int gqa_ratio = n_head / n_head_kv;
|
||||||
|
const int head_kv_idx = head_idx / gqa_ratio;
|
||||||
|
|
||||||
|
const global char * q_base = (const global char *) q_void + q_offset;
|
||||||
|
const global char * k_base = (const global char *) k_void + k_offset;
|
||||||
|
const global char * v_base = (const global char *) v_void + v_offset;
|
||||||
|
global char * o_base = (global char *) o_void + o_offset;
|
||||||
|
|
||||||
|
const global char * mask_base = NULL;
|
||||||
|
if (mask_void != NULL) {
|
||||||
|
const int mask_head_idx = head_idx % mask_ne2;
|
||||||
|
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||||
|
mask_base = (const global char *) mask_void + mask_offset +
|
||||||
|
mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
|
||||||
|
}
|
||||||
|
|
||||||
|
__local ACC_TYPE4 q_shared[DK_VEC];
|
||||||
|
{
|
||||||
|
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
|
||||||
|
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||||
|
for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
|
||||||
|
q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
|
||||||
|
|
||||||
|
const global ACC_TYPE * sinks_ptr = NULL;
|
||||||
|
if (sinks_void != NULL) {
|
||||||
|
sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
|
||||||
|
|
||||||
|
ACC_TYPE m_i = FA_M_INIT;
|
||||||
|
ACC_TYPE l_i = 0.0f;
|
||||||
|
|
||||||
|
const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
|
||||||
|
const int kv_start = sgid * kv_per_sg;
|
||||||
|
const int kv_end = min(n_kv, kv_start + kv_per_sg);
|
||||||
|
|
||||||
|
for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
|
||||||
|
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||||
|
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||||
|
|
||||||
|
ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
|
||||||
|
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||||
|
const int block_idx = qk / 8;
|
||||||
|
const int lane = qk % 8;
|
||||||
|
const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||||
|
dot4 = mad(q_shared[qk], k_v, dot4);
|
||||||
|
}
|
||||||
|
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
|
||||||
|
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
|
||||||
|
|
||||||
|
if (mask_base != NULL) {
|
||||||
|
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
|
||||||
|
score += slope * (ACC_TYPE) mask_ptr[k_idx];
|
||||||
|
}
|
||||||
|
if (logit_softcap > 0.0f) {
|
||||||
|
score = logit_softcap * tanh(score / logit_softcap);
|
||||||
|
}
|
||||||
|
|
||||||
|
const ACC_TYPE m_new = max(m_i, score);
|
||||||
|
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
|
||||||
|
const ACC_TYPE p = native_exp(score - m_new);
|
||||||
|
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
const int block_idx = dv / 8;
|
||||||
|
const int lane = dv % 8;
|
||||||
|
const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||||
|
o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
|
||||||
|
}
|
||||||
|
l_i = l_i * scale_prev + p;
|
||||||
|
m_i = m_new;
|
||||||
|
}
|
||||||
|
|
||||||
|
__local ACC_TYPE sg_m[VEC_NSG];
|
||||||
|
__local ACC_TYPE sg_l[VEC_NSG];
|
||||||
|
__local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
|
||||||
|
|
||||||
|
if (tid_sg == 0) {
|
||||||
|
sg_m[sgid] = m_i;
|
||||||
|
sg_l[sgid] = l_i;
|
||||||
|
}
|
||||||
|
{
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
sg_o[sgid][dv] = o_acc[idx];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
if (sgid == 0) {
|
||||||
|
ACC_TYPE m_final = sg_m[0];
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 1; s < VEC_NSG; ++s) {
|
||||||
|
m_final = max(m_final, sg_m[s]);
|
||||||
|
}
|
||||||
|
if (sinks_ptr != NULL) {
|
||||||
|
m_final = max(m_final, sinks_ptr[head_idx]);
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE l_final = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < VEC_NSG; ++s) {
|
||||||
|
l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
|
||||||
|
}
|
||||||
|
if (sinks_ptr != NULL) {
|
||||||
|
l_final += native_exp(sinks_ptr[head_idx] - m_final);
|
||||||
|
}
|
||||||
|
const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
|
||||||
|
|
||||||
|
const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
|
||||||
|
global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
|
||||||
|
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < VEC_NSG; ++s) {
|
||||||
|
const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
|
||||||
|
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
|
||||||
|
}
|
||||||
|
o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Flash-decoding split pass for q8_0 KV. Partial record: [m, l, O[DV]].
|
// Flash-decoding split pass for q8_0 KV. Partial record: [m, l, O[DV]].
|
||||||
// Merge kernel from flash_attn_f32_f16.cl is type-agnostic and reused.
|
// Merge kernel from flash_attn_f32_f16.cl is type-agnostic and reused.
|
||||||
#define FA_PARTIAL_FLOATS (2 + DV)
|
#define FA_PARTIAL_FLOATS (2 + DV)
|
||||||
@@ -533,6 +732,244 @@ __kernel void flash_attn_f32_q8_0_q1_split(
|
|||||||
#define FA_V_STRATEGY 0
|
#define FA_V_STRATEGY 0
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
#ifndef MQ_GQA
|
||||||
|
#define MQ_GQA 4
|
||||||
|
#endif
|
||||||
|
#ifndef MQ_NSG_SPLIT
|
||||||
|
#define MQ_NSG_SPLIT 4
|
||||||
|
#endif
|
||||||
|
#define MQ_SPLIT_WG_SIZE_Q8 (Q1_WG_SIZE * MQ_NSG_SPLIT)
|
||||||
|
|
||||||
|
REQD_SUBGROUP_SIZE_64
|
||||||
|
__kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
|
||||||
|
const global void * q_void, ulong q_offset,
|
||||||
|
const global void * k_void, ulong k_offset,
|
||||||
|
const global void * v_void, ulong v_offset,
|
||||||
|
const float scale,
|
||||||
|
const int n_q,
|
||||||
|
const int n_kv,
|
||||||
|
const int n_head,
|
||||||
|
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
|
||||||
|
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
|
||||||
|
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
|
||||||
|
const float max_bias,
|
||||||
|
const float m0,
|
||||||
|
const float m1,
|
||||||
|
const int n_head_log2,
|
||||||
|
const float logit_softcap,
|
||||||
|
const int n_head_kv,
|
||||||
|
const global void * mask_void,
|
||||||
|
const ulong mask_offset,
|
||||||
|
const ulong mask_nb1,
|
||||||
|
const ulong mask_nb2,
|
||||||
|
const ulong mask_nb3,
|
||||||
|
const int mask_ne2,
|
||||||
|
const int mask_ne3,
|
||||||
|
global float * partial_void,
|
||||||
|
const int n_splits,
|
||||||
|
const int kv_per_split
|
||||||
|
) {
|
||||||
|
const int tid = get_local_id(0);
|
||||||
|
const int sgid = tid / Q1_WG_SIZE;
|
||||||
|
const int tid_sg = tid % Q1_WG_SIZE;
|
||||||
|
const int kvhead_batch_idx = get_global_id(1);
|
||||||
|
const int split_q_idx = get_global_id(2);
|
||||||
|
const int split_idx = split_q_idx % n_splits;
|
||||||
|
const int q_idx = split_q_idx / n_splits;
|
||||||
|
|
||||||
|
const int batch_idx = kvhead_batch_idx / n_head_kv;
|
||||||
|
const int head_kv_idx = kvhead_batch_idx % n_head_kv;
|
||||||
|
|
||||||
|
const int kv_start = split_idx * kv_per_split;
|
||||||
|
const int kv_end = min(kv_start + kv_per_split, n_kv);
|
||||||
|
|
||||||
|
const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
|
||||||
|
|
||||||
|
if (kv_start >= kv_end) {
|
||||||
|
// Empty split — write sentinel for each of the MQ_GQA Q-heads.
|
||||||
|
if (tid == 0) {
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||||
|
* n_splits + split_idx);
|
||||||
|
global float * rec = partial_void + rec_idx * record_stride;
|
||||||
|
rec[0] = FA_M_INIT;
|
||||||
|
rec[1] = 0.0f;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const global char * q_base = (const global char *) q_void + q_offset;
|
||||||
|
const global char * k_base = (const global char *) k_void + k_offset;
|
||||||
|
const global char * v_base = (const global char *) v_void + v_offset;
|
||||||
|
|
||||||
|
__local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
|
||||||
|
for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q8) {
|
||||||
|
const int h = i / DK_VEC;
|
||||||
|
const int k = i % DK_VEC;
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
|
||||||
|
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
|
||||||
|
q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
float slope[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
|
||||||
|
}
|
||||||
|
|
||||||
|
const global char * mask_base[MQ_GQA];
|
||||||
|
if (mask_void != NULL) {
|
||||||
|
const int mask_batch_idx = batch_idx % mask_ne3;
|
||||||
|
const global char * mask_base_b = (const global char *) mask_void + mask_offset +
|
||||||
|
mask_batch_idx * mask_nb3 +
|
||||||
|
(ulong) q_idx * mask_nb1;
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
const int mask_head_idx = head_idx % mask_ne2;
|
||||||
|
mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
|
||||||
|
ACC_TYPE m_i[MQ_GQA];
|
||||||
|
ACC_TYPE l_i[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
m_i[h] = FA_M_INIT;
|
||||||
|
l_i[h] = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
|
||||||
|
}
|
||||||
|
|
||||||
|
const int kv_len = kv_end - kv_start;
|
||||||
|
const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
|
||||||
|
const int kv_lo = kv_start + sgid * kv_per_sg;
|
||||||
|
const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
|
||||||
|
|
||||||
|
for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
|
||||||
|
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
|
||||||
|
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
|
||||||
|
|
||||||
|
ACC_TYPE4 dot4[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
|
||||||
|
|
||||||
|
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
|
||||||
|
const int block_idx = qk / 8;
|
||||||
|
const int lane = qk % 8;
|
||||||
|
const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE score[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
|
||||||
|
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
|
||||||
|
if (mask_base[h] != NULL) {
|
||||||
|
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
|
||||||
|
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
|
||||||
|
}
|
||||||
|
if (logit_softcap > 0.0f) {
|
||||||
|
s = logit_softcap * tanh(s / logit_softcap);
|
||||||
|
}
|
||||||
|
score[h] = s;
|
||||||
|
}
|
||||||
|
|
||||||
|
ACC_TYPE p_h[MQ_GQA];
|
||||||
|
ACC_TYPE sp_h[MQ_GQA];
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
const ACC_TYPE m_new = max(m_i[h], score[h]);
|
||||||
|
sp_h[h] = native_exp(m_i[h] - m_new);
|
||||||
|
p_h[h] = native_exp(score[h] - m_new);
|
||||||
|
l_i[h] = l_i[h] * sp_h[h] + p_h[h];
|
||||||
|
m_i[h] = m_new;
|
||||||
|
}
|
||||||
|
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
|
||||||
|
const int block_idx = dv / 8;
|
||||||
|
const int lane = dv % 8;
|
||||||
|
const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
__local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
|
||||||
|
__local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
|
||||||
|
__local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
|
||||||
|
|
||||||
|
if (tid_sg == 0) {
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
sg_m[h][sgid] = m_i[h];
|
||||||
|
sg_l[h][sgid] = l_i[h];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#pragma unroll
|
||||||
|
for (int h = 0; h < MQ_GQA; ++h) {
|
||||||
|
{
|
||||||
|
int idx = 0;
|
||||||
|
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
|
||||||
|
sg_o[sgid][dv_idx] = o_acc[h][idx];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
|
||||||
|
if (sgid == 0) {
|
||||||
|
const int head_idx = head_kv_idx * MQ_GQA + h;
|
||||||
|
|
||||||
|
ACC_TYPE m_c = sg_m[h][0];
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
|
||||||
|
m_c = max(m_c, sg_m[h][s]);
|
||||||
|
}
|
||||||
|
ACC_TYPE l_c = 0.0f;
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||||
|
l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
|
||||||
|
}
|
||||||
|
|
||||||
|
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
|
||||||
|
* n_splits + split_idx);
|
||||||
|
global float * rec = partial_void + rec_idx * record_stride;
|
||||||
|
global float4 * rec_o = (global float4 *) (rec + 2);
|
||||||
|
|
||||||
|
if (tid_sg == 0) {
|
||||||
|
rec[0] = (float) m_c;
|
||||||
|
rec[1] = (float) l_c;
|
||||||
|
}
|
||||||
|
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
|
||||||
|
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
|
||||||
|
#pragma unroll
|
||||||
|
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
|
||||||
|
const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
|
||||||
|
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
|
||||||
|
}
|
||||||
|
rec_o[dv_idx] = o_merged;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
barrier(CLK_LOCAL_MEM_FENCE);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
__kernel void flash_attn_f32_q8_0(
|
__kernel void flash_attn_f32_q8_0(
|
||||||
const global void * q_void, ulong q_offset,
|
const global void * q_void, ulong q_offset,
|
||||||
const global void * k_void, ulong k_offset,
|
const global void * k_void, ulong k_offset,
|
||||||
|
|||||||
@@ -18,6 +18,14 @@
|
|||||||
#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
|
#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
|
||||||
#endif
|
#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
|
// Assumes row size (ne00) is a multiple of 4
|
||||||
#ifdef ADRENO_GPU
|
#ifdef ADRENO_GPU
|
||||||
REQD_SUBGROUP_SIZE_64
|
REQD_SUBGROUP_SIZE_64
|
||||||
@@ -378,3 +386,848 @@ kernel void kernel_mul_mat_f16_f32_l4_dr_lq(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
#endif // ADRENO_GPU
|
#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