opencl: read/write MoE dp4a activation tiles to local memory as 128-bit (vectorized...

* opencl: read MoE dp4a activation tile as 128-bit local loads

* opencl: vectorize MoE dp4a activation staging as 128-bit loads

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 95d0638134e10ee9bdd492868353c814a4123a6c..97fdc8e18c821b37150b2a6756ccfc64c0eaf75e 100644 (file)
--- 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 86ff943c51b2cf46c7ca946bb4c38bcb40387fbc..502472049a9c0cfd564fa8c1ed23a459dc6eb4f0 100644 (file)
--- 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 540897544082c0216e0b707202f7a0a4c3443b6b..9d968f32ed5a34f5c0f41f41314d921be896d2d5 100644 (file)
--- 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 35e63dcab05ff817e73204ed67d242efb4851040..4ffe9f8e66c6b9e74d3f6c7445100565016950e2 100644 (file)
--- 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 39bf5d83211c1b2d0f8d6e4d1ed1721a8d5e00d3..d0b191e18363c461c43f5b1c68b6fbfefdd1b2c2 100644 (file)
--- 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];