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
This commit is contained in:
Hongqiang Wang
2026-07-17 12:02:27 -07:00
committed by GitHub
parent 86d86ed439
commit 6bdd77f13c
5 changed files with 77 additions and 53 deletions
@@ -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];
@@ -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];
@@ -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];
@@ -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];
@@ -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];