From 89bde0ce40af66e72157d78bd66489ffaf9f9c6c Mon Sep 17 00:00:00 2001 From: Jason Titus <870238+jasontitus@users.noreply.github.com> Date: Fri, 25 Sep 2026 21:51:41 -0700 Subject: [PATCH] metal : add PQ2_0 mat-vec for two columns (GGML_METAL_PQ2_MULTICOL=1) Assisted-by: Claude Opus 5.5 --- ggml/src/ggml-metal/ggml-metal-device.cpp | 14 +++ ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 2 +- ggml/src/ggml-metal/kernels/mul_mv.metal | 104 ++++++++++++++++++++++ tests/test-backend-ops.cpp | 4 + 5 files changed, 124 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 3420ba85116c..6dcd4c6b0feb 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/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: { diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 0d305888f925..8d62831d8a59 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/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 diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index c41bae469ac5..adab5a8d892c 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/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 || diff --git a/ggml/src/ggml-metal/kernels/mul_mv.metal b/ggml/src/ggml-metal/kernels/mul_mv.metal index 7d3bbc0bee89..2d9516630f67 100644 --- a/ggml/src/ggml-metal/kernels/mul_mv.metal +++ b/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(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 +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, diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 72e580d895c0..b0d4f6490023 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9400,6 +9400,10 @@ static std::vector> 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}) {