From 6bdd77f13cf11b264b4231d320afc404f48d576e Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Fri, 17 Jul 2026 12:02:27 -0700 Subject: [PATCH] opencl: read/write MoE dp4a activation tiles to local memory as 128-bit (vectorized LD/ST perf opt) for Adreno GPUs (#25810) * opencl: read MoE dp4a activation tile as 128-bit local loads * opencl: vectorize MoE dp4a activation staging as 128-bit loads --- .../kernels/gemm_moe_mxfp4_q8_1_dp4a.cl | 30 +++++++++------- .../kernels/gemm_moe_q4_0_q8_1_dp4a.cl | 30 +++++++++------- .../kernels/gemm_moe_q4_k_q8_1_dp4a.cl | 35 +++++++++++-------- .../kernels/gemm_moe_q6_k_q8_1_dp4a.cl | 18 ++++++---- .../ggml-opencl/kernels/gemm_moe_q8_1_dp4a.cl | 17 +++++---- 5 files changed, 77 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-opencl/kernels/gemm_moe_mxfp4_q8_1_dp4a.cl b/ggml/src/ggml-opencl/kernels/gemm_moe_mxfp4_q8_1_dp4a.cl index 95d0638134..97fdc8e18c 100644 --- a/ggml/src/ggml-opencl/kernels/gemm_moe_mxfp4_q8_1_dp4a.cl +++ b/ggml/src/ggml-opencl/kernels/gemm_moe_mxfp4_q8_1_dp4a.cl @@ -37,15 +37,17 @@ static inline float e8m0_to_fp32(uchar x) { // One token's dp4a dot (8 uints = 32 K elems) + mxfp4 block-scale epilogue. // blk_scale already carries the 0.5 factor (== 0.5 * 2^e). #define MOE_MXFP4_DP4A_T(t) do { \ + uint4 a0 = vload4(0, &sh_qa[t][0]); \ + uint4 a1 = vload4(0, &sh_qa[t][4]); \ int raw = 0; \ - raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, raw); \ acc[t] += blk_scale * (float)sh_d[t] * (float)raw; \ } while (0) @@ -133,11 +135,13 @@ kernel void kernel_gemm_moe_mxfp4_q8_1_dp4a( qw[6] = mxfp4_pack((ushort)(r3)); qw[7] = mxfp4_pack((ushort)(r3 >> 16)); // cooperatively stage the n_real-token x 32-K int8 activations - const uint stage_lim = (uint)n_real * 8; - for (uint idx = lid; idx < stage_lim; idx += 64) { - const uint t = idx >> 3; - const uint u = idx & 7; - sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u]; + // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores. + const uint vlim = (uint)n_real * 2; + for (uint idx = lid; idx < vlim; idx += 64) { + const uint t = idx >> 1; + const uint h = (idx & 1) << 2; // 0 or 4 + uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]); + vstore4(v, 0, &sh_qa[t][h]); } if (lid < (uint)n_real) { sh_d[lid] = src1_da[(col + lid) * num_blocks + sub]; diff --git a/ggml/src/ggml-opencl/kernels/gemm_moe_q4_0_q8_1_dp4a.cl b/ggml/src/ggml-opencl/kernels/gemm_moe_q4_0_q8_1_dp4a.cl index 86ff943c51..502472049a 100644 --- a/ggml/src/ggml-opencl/kernels/gemm_moe_q4_0_q8_1_dp4a.cl +++ b/ggml/src/ggml-opencl/kernels/gemm_moe_q4_0_q8_1_dp4a.cl @@ -17,15 +17,17 @@ // One token's dp4a dot (8 uints = 32 K elems) + q4_0 scale/zero-point epilogue. #define MOE_Q40_DP4A_T(t) do { \ + uint4 a0 = vload4(0, &sh_qa[t][0]); \ + uint4 a1 = vload4(0, &sh_qa[t][4]); \ int raw = 0; \ - raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, raw); \ acc[t] += d_val * ((float)sh_d[t] * (float)raw - 8.0f * (float)sh_s[t]); \ } while (0) @@ -112,11 +114,13 @@ kernel void kernel_gemm_moe_q4_0_q8_1_dp4a( qw[6] = EXP4(r3); qw[7] = EXP4(r3 >> 16); // cooperatively stage the n_real-token x 32-K int8 activations - const uint stage_lim = (uint)n_real * 8; - for (uint idx = lid; idx < stage_lim; idx += 64) { - const uint t = idx >> 3; - const uint u = idx & 7; - sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u]; + // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores. + const uint vlim = (uint)n_real * 2; + for (uint idx = lid; idx < vlim; idx += 64) { + const uint t = idx >> 1; + const uint h = (idx & 1) << 2; // 0 or 4 + uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]); + vstore4(v, 0, &sh_qa[t][h]); } if (lid < (uint)n_real) { sh_d[lid] = src1_da[(col + lid) * num_blocks + sub]; diff --git a/ggml/src/ggml-opencl/kernels/gemm_moe_q4_k_q8_1_dp4a.cl b/ggml/src/ggml-opencl/kernels/gemm_moe_q4_k_q8_1_dp4a.cl index 5408975440..9d968f32ed 100644 --- a/ggml/src/ggml-opencl/kernels/gemm_moe_q4_k_q8_1_dp4a.cl +++ b/ggml/src/ggml-opencl/kernels/gemm_moe_q4_k_q8_1_dp4a.cl @@ -37,16 +37,21 @@ inline void get_scale_min_k4( (((uint)((u) & 0xF000u)) << 12) ) // One token's dp4a dot (8 uints = 32 K elems) + q4_K scale/min epilogue into acc[t]. +// The 8 activation uints are read as two 128-bit uint4 loads staged to private (Adreno +// wants 128-bit local reads, and a __local operand fed straight to the dp4a builtin is +// slower and can miscompile). #define MOE_Q4K_DP4A_T(t) do { \ + uint4 a0 = vload4(0, &sh_qa[t][0]); \ + uint4 a1 = vload4(0, &sh_qa[t][4]); \ int raw = 0; \ - raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \ - raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw); \ + raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, raw); \ acc[t] += scale * (float)sh_d[t] * (float)raw - minv * (float)sh_s[t]; \ } while (0) @@ -145,12 +150,14 @@ kernel void kernel_gemm_moe_q4_k_q8_1_dp4a( qw[4] = EXP4(r2); qw[5] = EXP4(r2 >> 16); qw[6] = EXP4(r3); qw[7] = EXP4(r3 >> 16); - // --- cooperatively stage the n_real-token x 32-K int8 activations to LDS --- - const uint stage_lim = (uint)n_real * 8; - for (uint idx = lid; idx < stage_lim; idx += 64) { - const uint t = idx >> 3; - const uint u = idx & 7; - sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u]; + // cooperatively stage the n_real-token x 32-K int8 activations to lm + // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores. + const uint vlim = (uint)n_real * 2; + for (uint idx = lid; idx < vlim; idx += 64) { + const uint t = idx >> 1; + const uint h = (idx & 1) << 2; // 0 or 4 + uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]); + vstore4(v, 0, &sh_qa[t][h]); } if (lid < (uint)n_real) { sh_d[lid] = src1_da[(col + lid) * ne00_b + sub]; diff --git a/ggml/src/ggml-opencl/kernels/gemm_moe_q6_k_q8_1_dp4a.cl b/ggml/src/ggml-opencl/kernels/gemm_moe_q6_k_q8_1_dp4a.cl index 35e63dcab0..4ffe9f8e66 100644 --- a/ggml/src/ggml-opencl/kernels/gemm_moe_q6_k_q8_1_dp4a.cl +++ b/ggml/src/ggml-opencl/kernels/gemm_moe_q6_k_q8_1_dp4a.cl @@ -41,8 +41,10 @@ inline int dp4a_q6(uint qw0, uint qw1, uint qw2, uint qw3, // One token's q6_K dp4a dot (two halves, per-16 scales) + epilogue into acc[t]. #define MOE_Q6K_DP4A_T(t) do { \ - const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], sh_qa[t][0], sh_qa[t][1], sh_qa[t][2], sh_qa[t][3]); \ - const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], sh_qa[t][4], sh_qa[t][5], sh_qa[t][6], sh_qa[t][7]); \ + uint4 a0 = vload4(0, &sh_qa[t][0]); \ + uint4 a1 = vload4(0, &sh_qa[t][4]); \ + const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], a0.s0, a0.s1, a0.s2, a0.s3); \ + const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], a1.s0, a1.s1, a1.s2, a1.s3); \ const float a_d = (float)sh_d[t]; \ acc[t] += scale0 * a_d * (float)raw1 + scale1 * a_d * (float)raw2; \ } while (0) @@ -144,11 +146,13 @@ kernel void kernel_gemm_moe_q6_k_q8_1_dp4a( qw[6] = SIGN6(EXP4(r3) | EXP2((qh2 >> 16) & 0xFFu)); qw[7] = SIGN6(EXP4(r3 >> 16) | EXP2((qh2 >> 24) & 0xFFu)); - const uint stage_lim = (uint)n_real * 8; - for (uint idx = lid; idx < stage_lim; idx += 64) { - const uint t = idx >> 3; - const uint u = idx & 7; - sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u]; + // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores. + const uint vlim = (uint)n_real * 2; + for (uint idx = lid; idx < vlim; idx += 64) { + const uint t = idx >> 1; + const uint h = (idx & 1) << 2; // 0 or 4 + uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]); + vstore4(v, 0, &sh_qa[t][h]); } if (lid < (uint)n_real) { sh_d[lid] = src1_da[(col + lid) * ne00_b + sub]; diff --git a/ggml/src/ggml-opencl/kernels/gemm_moe_q8_1_dp4a.cl b/ggml/src/ggml-opencl/kernels/gemm_moe_q8_1_dp4a.cl index 39bf5d8321..d0b191e183 100644 --- a/ggml/src/ggml-opencl/kernels/gemm_moe_q8_1_dp4a.cl +++ b/ggml/src/ggml-opencl/kernels/gemm_moe_q8_1_dp4a.cl @@ -102,8 +102,10 @@ inline int dp4a4(uint w0,uint w1,uint w2,uint w3,uint a0,uint a1,uint a2,uint a3 // One token's two-half dp4a + uniform scale/min epilogue into acc[t]. #define MOE_DP4A_T(t) do { \ - const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], sh_qa[t][0],sh_qa[t][1],sh_qa[t][2],sh_qa[t][3]); \ - const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], sh_qa[t][4],sh_qa[t][5],sh_qa[t][6],sh_qa[t][7]); \ + uint4 a0 = vload4(0, &sh_qa[t][0]); \ + uint4 a1 = vload4(0, &sh_qa[t][4]); \ + const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], a0.s0,a0.s1,a0.s2,a0.s3); \ + const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], a1.s0,a1.s1,a1.s2,a1.s3); \ const float a_d = (float)sh_d[t]; \ acc[t] += sc0*a_d*(float)raw1 + sc1*a_d*(float)raw2 - mn*(float)sh_s[t]; \ } while (0) @@ -178,10 +180,13 @@ kernel void kernel_gemm_moe_q8_1_dp4a( LOAD_QW(step, sub) - const uint stage_lim = (uint)n_real * 8; - for (uint idx = lid; idx < stage_lim; idx += 64) { - const uint t = idx >> 3, u = idx & 7; - sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u]; + // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores. + const uint vlim = (uint)n_real * 2; + for (uint idx = lid; idx < vlim; idx += 64) { + const uint t = idx >> 1; + const uint h = (idx & 1) << 2; // 0 or 4 + uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]); + vstore4(v, 0, &sh_qa[t][h]); } if (lid < (uint)n_real) { sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];