diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 3420ba85116c..cbcdb36c1f6a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -882,11 +882,21 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta return res; } +// widest batch the multi-column kernel takes, as column tiles of at most four; wider batches keep the default route. +// GGML_METAL_PTQ1_MULTICOL_MAX is clamped to 4..8 (default 8); 4 keeps the kernel to 2-4 columns +static int ggml_metal_ptq1_multicol_max(void) { + static const int max_cols = [] { + const char * env = getenv("GGML_METAL_PTQ1_MULTICOL_MAX"); + return env ? std::min(8, std::max(4, atoi(env))) : 8; + }(); + return max_cols; +} + 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 && op->src[0]->ne[0] % ggml_blck_size(GGML_TYPE_PTQ1_0) == 0 && op->src[1]->nb[0] == sizeof(float) && - op->src[1]->ne[1] >= 2 && op->src[1]->ne[1] <= 4; + op->src[1]->ne[1] >= 2 && op->src[1]->ne[1] <= ggml_metal_ptq1_multicol_max(); } ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) { @@ -907,6 +917,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta const char * suffix = ""; char ptq1_suffix[16]; + bool ptq1_mc = false; // use custom matrix x vector kernel switch (tsrc0) { @@ -959,9 +970,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta nsg = N_SG_PTQ1_0; nr0 = N_R0_PTQ1_0; if (ggml_metal_ptq1_multicol_enabled(op)) { + ptq1_mc = true; nr0 = 4; nsg = 1; - nr1 = ne11; + // at most four columns per tile, split evenly: 5 -> 3+2, 6 -> 3+3, 7 -> 4+3, 8 -> 4+4 + nr1 = (ne11 + (ne11 + 3)/4 - 1) / ((ne11 + 3)/4); snprintf(ptq1_suffix, sizeof(ptq1_suffix), "_mc_c%d", nr1); suffix = ptq1_suffix; } @@ -1091,7 +1104,13 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta const int16_t r3 = (int16_t) (ne13 / ne03); snprintf(base, 256, "kernel_mul_mv_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix); - snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, ne12, r2, r3); + // PTQ1_0 multi-column: every column tile complete (2-4, 6, 8 columns) skips the partial-tile checks + const bool ptq1_full_cols = ptq1_mc && ne11 % nr1 == 0; + if (ptq1_mc) { + snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d_full=%d", base, nsg, ne12, r2, r3, ptq1_full_cols); + } else { + snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, ne12, r2, r3); + } ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); if (!res.pipeline) { @@ -1101,6 +1120,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV + 2); ggml_metal_cv_set_int16(cv, r2, FC_MUL_MV + 3); ggml_metal_cv_set_int16(cv, r3, FC_MUL_MV + 4); + if (ptq1_mc) { + ggml_metal_cv_set_bool(cv, ptq1_full_cols, FC_MUL_MV + 5); + } res = ggml_metal_library_compile_pipeline(lib, base, name, cv); diff --git a/ggml/src/ggml-metal/kernels/mul_mv.metal b/ggml/src/ggml-metal/kernels/mul_mv.metal index 7d3bbc0bee89..b3d6ce7c17b8 100644 --- a/ggml/src/ggml-metal/kernels/mul_mv.metal +++ b/ggml/src/ggml-metal/kernels/mul_mv.metal @@ -262,6 +262,7 @@ constant short FC_mul_mv_nxpsg [[function_constant(FC_MUL_MV + 1)]]; constant short FC_mul_mv_ne12 [[function_constant(FC_MUL_MV + 2)]]; constant short FC_mul_mv_r2 [[function_constant(FC_MUL_MV + 3)]]; constant short FC_mul_mv_r3 [[function_constant(FC_MUL_MV + 4)]]; +constant bool FC_mul_mv_ptq1_full_cols [[function_constant(FC_MUL_MV + 5)]]; template void mul_vec_q_n_f32_impl( @@ -1061,6 +1062,9 @@ kernel void kernel_mul_mv_ptq1_0_multicol( const int r1 = tgpig.y * nr1; const int im = tgpig.z; + // columns of this tile that exist: nr1 except in a partial last tile (5 or 7 columns) + const short ncols = FC_mul_mv_ptq1_full_cols ? nr1 : (short) min(nr1, args.ne11 - r1); + const int first_row = (r0 * NSG + sgitg) * nr0; const uint i12 = im%FC_mul_mv_ne12; @@ -1097,7 +1101,9 @@ kernel void kernel_mul_mv_ptq1_0_multicol( // Reuse collapse coefficients across rows: c[k-1] = y_{k-1} - 3*y_k, c[4] = y_4. float sumy[nr1] = {}; FOR_UNROLL (short col = 0; col < nr1; ++col) { - device const float * yc = (device const float *) ((device const char *) yb + col*args.nb11); + // a partial last tile re-reads its final valid column; that result is not written + const short yc_col = FC_mul_mv_ptq1_full_cols ? col : min(col, (short) (ncols - 1)); + device const float * yc = (device const float *) ((device const char *) yb + yc_col*args.nb11); FOR_UNROLL (short k = 0; k < 2; ++k) { const short m = 2*it + k; @@ -1141,7 +1147,7 @@ kernel void kernel_mul_mv_ptq1_0_multicol( 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) { + if (tiisg == 0 && first_row + row < args.ne01 && (FC_mul_mv_ptq1_full_cols || col < ncols)) { dst_f32[(uint64_t) col*args.ne0 + first_row + row] = tot; } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 72e580d895c0..b6ac8ac988d3 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9401,11 +9401,15 @@ static std::vector> make_test_cases_eval() { } // PTQ1_0 small batches, row tails and broadcast dimensions. - for (int n : {1, 2, 3, 4, 8}) { + for (int n : {1, 2, 3, 4, 5, 6, 7, 8}) { for (int k : {128, 384, 5120}) { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PTQ1_0, GGML_TYPE_F32, 7, n, k, {2, 2}, {2, 1})); } } + // partial column tiles with a strided B + for (int n : {5, 7}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PTQ1_0, GGML_TYPE_F32, 7, n, 384, {2, 3}, {1, 1}, {0, 2, 1, 3})); + } // BF16 is absent from base_types: add the 3 standard non-contig permutations explicitly test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 1, 256, {2, 3}, {1, 1}, {0, 2, 1, 3}));