Skip to content

metal : PTQ1_0 multi-column mat-vec for 5-8 columns #277

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jasontitus wants to merge 1 commit into
base: prism
Choose a base branch
from
Open

metal : PTQ1_0 multi-column mat-vec for 5-8 columns #277

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 25 additions & 3 deletions ggml/src/ggml-metal/ggml-metal-device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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) {
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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) {
Expand All @@ -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);

Expand Down
10 changes: 8 additions & 2 deletions ggml/src/ggml-metal/kernels/mul_mv.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<typename block_q_type, short NR0, typename args_t>
void mul_vec_q_n_f32_impl(
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
Expand Down
6 changes: 5 additions & 1 deletion tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9401,11 +9401,15 @@ static std::vector<std::unique_ptr<test_case>> 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}));
Expand Down