diff options
Diffstat (limited to 'ggml/src/ggml.c')
| -rw-r--r-- | ggml/src/ggml.c | 15 |
1 files changed, 7 insertions, 8 deletions
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index a2bdc156..036bd8a8 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -21771,15 +21771,14 @@ struct ggml_cplan ggml_graph_plan(const struct ggml_cgraph * cgraph, int n_threa if (gcd_k > 1) { int nth_k = n_tasks/gcd_k; int rk2 = q->ne[2]/k->ne[2]; - if (rk2%nth_k == 0) { - size_t size = (Dv + 16)*rk2/nth_k*sizeof(float)*n_tasks; - if (ggml_is_quantized(k->type)) { - enum ggml_type vec_dot_type = type_traits[k->type].vec_dot_type; - size_t row_size = ggml_row_size(vec_dot_type, q->ne[0]); - size += q->ne[2]*row_size; - } - cur = MAX(cur, size); + int nq_per_thread = (rk2 + nth_k - 1)/nth_k; + size_t size = (Dv + 16)*nq_per_thread*sizeof(float)*n_tasks; + if (ggml_is_quantized(k->type)) { + enum ggml_type vec_dot_type = type_traits[k->type].vec_dot_type; + size_t row_size = ggml_row_size(vec_dot_type, q->ne[0]); + size += q->ne[2]*row_size; } + cur = MAX(cur, size); } } #endif |