metal : opt-in PQ2_0 mat-vec for two columns (GGML_METAL_PQ2_MULTICOL=1)

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -882,6 +882,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta
     return res;
 }
 
+// PQ2_0 two-column mat-vec (GGML_METAL_PQ2_MULTICOL=1); from three columns mul_mv_ext is faster
+bool ggml_metal_pq2_multicol_enabled(const ggml_tensor * op) {
+    static const bool enabled = getenv("GGML_METAL_PQ2_MULTICOL") && atoi(getenv("GGML_METAL_PQ2_MULTICOL")) == 1;
+    return enabled && op->src[0]->type == GGML_TYPE_PQ2_0 && op->src[1]->type == GGML_TYPE_F32 &&
+           op->src[0]->ne[0] % ggml_blck_size(GGML_TYPE_PQ2_0) == 0 && op->src[1]->nb[0] == sizeof(float) &&
+           op->src[1]->ne[1] == 2;
+}
+
 bool ggml_metal_ptq1_multicol_enabled(const ggml_tensor * op) {
     static const bool enabled = getenv("GGML_METAL_PTQ1_MULTICOL") && atoi(getenv("GGML_METAL_PTQ1_MULTICOL")) == 1;
     return enabled && op->src[0]->type == GGML_TYPE_PTQ1_0 && op->src[1]->type == GGML_TYPE_F32 &&
@@ -953,6 +961,12 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
             {
                 nsg = N_SG_PQ2_0;
                 nr0 = N_R0_PQ2_0;
+                if (ggml_metal_pq2_multicol_enabled(op)) {
+                    nr0 = 2;
+                    nsg = 1;
+                    nr1 = 2;
+                    suffix = "_mc_c2";
+                }
             } break;
         case GGML_TYPE_PTQ1_0:
             {

ggml/src/ggml-metal/ggml-metal-device.h

@@ -7,6 +7,7 @@ extern "C" {
 #endif
 
 bool ggml_metal_ptq1_multicol_enabled(const struct ggml_tensor * op);
+bool ggml_metal_pq2_multicol_enabled(const struct ggml_tensor * op);
 
 struct ggml_metal_buffer_id {
     void * metal; // id<MTLBuffer>

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2799,7 +2799,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
            op->src[0]->type == GGML_TYPE_BF16 ||
            (op->src[0]->type == GGML_TYPE_Q1_0 && q1_0_ext_enable) ||
            op->src[0]->type == GGML_TYPE_Q2_0 ||
-           op->src[0]->type == GGML_TYPE_PQ2_0 ||
+           (op->src[0]->type == GGML_TYPE_PQ2_0 && !ggml_metal_pq2_multicol_enabled(op)) ||
            (op->src[0]->type == GGML_TYPE_PTQ1_0 && !ggml_metal_ptq1_multicol_enabled(op)) ||
            op->src[0]->type == GGML_TYPE_Q4_0 ||
            op->src[0]->type == GGML_TYPE_Q4_1 ||

ggml/src/ggml-metal/kernels/mul_mv.metal

@@ -1282,6 +1282,110 @@ kernel void kernel_mul_mv_pq2_0_f32(
     kernel_mul_mv_pq2_0_f32_impl<N_R0_PQ2_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
 }
 
+// PQ2_0 mat-vec for two columns: floors once per byte, collapse coefficients once per column (same arithmetic as kernel_mul_mv_pq2_0_f32)
+template<int nr0, int nr1>
+kernel void kernel_mul_mv_pq2_0_multicol(
+        constant ggml_metal_kargs_mul_mv & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    const short NSG = FC_mul_mv_nsg;
+
+    const int nb = args.ne00/QK_PQ2_0;
+
+    const int r0 = tgpig.x;
+    const int r1 = tgpig.y * nr1;
+    const int im = tgpig.z;
+
+    const int first_row = (r0 * NSG + sgitg) * nr0;
+
+    const uint i12 = im%FC_mul_mv_ne12;
+    const uint i13 = im/FC_mul_mv_ne12;
+
+    const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
+
+    device const float * y = (device const float *) (src1 + offset1);
+
+    device const block_pq2_0 * ax[nr0];
+    for (int row = 0; row < nr0; ++row) {
+        const uint64_t offset0 = min(first_row + row, args.ne01 - 1)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+        ax[row] = (device const block_pq2_0 *) ((device char *) src0 + offset0);
+    }
+
+    float yl[nr1][16];
+    float sumy[nr1];
+    float sumf[nr0][nr1] = {};
+
+    const short ix = (tiisg/8);
+    const short il = (tiisg%8)*16;
+
+    device const float * yb = y + ix*QK_PQ2_0 + il;
+
+    for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/8) {
+        FOR_UNROLL (short col = 0; col < nr1; ++col) {
+            device const float * yc = (device const float *) ((device const char *) yb + col*args.nb11);
+
+            sumy[col] = 0.f;
+
+            FOR_UNROLL (short j = 0; j < 4; j++) {
+                const float y0 = yc[4*j + 0];
+                const float y1 = yc[4*j + 1];
+                const float y2 = yc[4*j + 2];
+                const float y3 = yc[4*j + 3];
+                sumy[col] += (y0 + y1) + (y2 + y3);
+                yl[col][4*j + 0] = y3 - 4.0f*y2;
+                yl[col][4*j + 1] = y2 - 4.0f*y1;
+                yl[col][4*j + 2] = y1 - 4.0f*y0;
+                yl[col][4*j + 3] = y0;
+            }
+        }
+
+        FOR_UNROLL (short row = 0; row < nr0; row++) {
+            device const block_pq2_0 * qb = ax[row] + ib;
+            device const uint8_t * qs = qb->qs + (il / 4);
+
+            float acc[nr1] = {};
+
+            FOR_UNROLL (short j = 0; j < 4; j++) {
+                const float b  = (float) qs[j];
+                const float u  = b * (1.0f/256.0f);
+                const float g1 = floor( 4.0f*u);
+                const float g2 = floor(16.0f*u);
+                const float g3 = floor(64.0f*u);
+                FOR_UNROLL (short col = 0; col < nr1; ++col) {
+                    acc[col] += g1*yl[col][4*j + 0];
+                    acc[col] += g2*yl[col][4*j + 1];
+                    acc[col] += g3*yl[col][4*j + 2];
+                    acc[col] +=  b*yl[col][4*j + 3];
+                }
+            }
+
+            FOR_UNROLL (short col = 0; col < nr1; ++col) {
+                sumf[row][col] += qb->d * (acc[col] - sumy[col]);
+            }
+        }
+
+        yb += QK_PQ2_0 * (N_SIMDWIDTH/8);
+    }
+
+    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
+
+    for (int row = 0; row < nr0; ++row) {
+        FOR_UNROLL (short col = 0; col < nr1; ++col) {
+            const float tot = simd_sum(sumf[row][col]);
+            if (tiisg == 0 && first_row + row < args.ne01) {
+                dst_f32[(uint64_t) col*args.ne0 + first_row + row] = tot;
+            }
+        }
+    }
+}
+
+typedef decltype(kernel_mul_mv_pq2_0_multicol<2, 2>) mul_mv_pq2_multicol_t;
+template [[host_name("kernel_mul_mv_pq2_0_f32_mc_c2")]] kernel mul_mv_pq2_multicol_t kernel_mul_mv_pq2_0_multicol<2, 2>;
+
 kernel void kernel_mul_mv_q4_0_f32(
         constant ggml_metal_kargs_mul_mv & args,
         device const char * src0,

tests/test-backend-ops.cpp

@@ -9400,6 +9400,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // PQ2_0 two columns with a row tail and broadcast, and with a strided B
+    test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 7, 2, 384, {2, 2}, {2, 1}));
+    test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 7, 2, 384, {2, 3}, {1, 1}, {0, 2, 1, 3}));
+
     // PTQ1_0 small batches, row tails and broadcast dimensions.
     for (int n : {1, 2, 3, 4, 8}) {
         for (int k : {128, 384, 5120}) {