diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index fe58ea3bb7a6..35b58ccb5a6a 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -550,6 +550,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( tensor->src[1]->ne[src_ss[0].axis] == 1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { return src_ss[0]; } + // feature-sharded activation times a full-width mirrored operand: each shard reads its + // slice of src1 (see init_tensor), which is only done for the four elementwise ops + if (tensor->op != GGML_OP_ADD_ID && src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && + src_ss[0].n_segments == 1 && src_ss[0].nr[0] == 1 && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + tensor->src[1]->ne[0] == tensor->src[0]->ne[0]) { + return src_ss[0]; + } if (src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[0].axis == src_ss[1].axis || (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) { return src_ss[0]; // GGML_OP_ADD_ID @@ -581,7 +589,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { ggml_backend_meta_split_state ret = src_ss[0]; ret.axis = GGML_BACKEND_SPLIT_AXIS_0; - ret.nr[0] = 1; + ret.nr[0] = src_ss[0].n_segments == 1 ? src_ss[0].nr[0] : 1; ret.n_segments = 1; return ret; } @@ -597,6 +605,10 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); return src_ss[0]; } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[1].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[1].axis < GGML_MAX_DIMS) { + return src_ss[1]; + } // batched matmul with the batches split across devices and a replicated activation if (src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[0].axis < GGML_MAX_DIMS && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { @@ -669,6 +681,15 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_view = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (tensor->src[0]->op == GGML_OP_GATED_DELTA_NET && tensor->view_offs == 0 && + src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { + const ggml_tensor * q = tensor->src[0]->src[0]; + const ggml_tensor * v = tensor->src[0]->src[2]; + if (tensor->ne[0] == v->ne[0] && tensor->ne[1] == v->ne[1] && + v->ne[1] > q->ne[1] && v->ne[1] % q->ne[1] == 0) { + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {uint32_t(v->ne[1] / q->ne[1])}, 1}; + } + } if (ggml_is_contiguous(tensor) && ggml_is_contiguous(tensor->src[0])) { return handle_reshape(src_ss); } @@ -1269,6 +1290,20 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m } } + if ((tensor->op == GGML_OP_MUL || tensor->op == GGML_OP_ADD || tensor->op == GGML_OP_SUB || tensor->op == GGML_OP_DIV) && + split_state.axis == GGML_BACKEND_SPLIT_AXIS_0 && split_state.n_segments == 1 && split_state.nr[0] == 1 && + tensor->src[1]->ne[0] == tensor->src[0]->ne[0] && + ggml_backend_buffer_is_meta(tensor->src[1]->buffer) && + ggml_backend_meta_get_split_state(tensor->src[1], true).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + int64_t offset = 0; + for (size_t k = 0; k < j; k++) { + offset += split_state.ne[k]; + } + ggml_tensor * source = t_ij->src[1]; + t_ij->src[1] = ggml_view_4d(simple_ctx, source, ne[0], source->ne[1], source->ne[2], source->ne[3], + source->nb[1], source->nb[2], source->nb[3], offset * source->nb[0]); + } + simple_tensors.push_back(t_ij); } diff --git a/src/llama-impl.h b/src/llama-impl.h index c6bc214ccba0..28c73193d9c3 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -59,13 +59,14 @@ static inline ggml_tensor * llama_mul_mat_hadamard( ggml_tensor * cur, ggml_tensor * rot) { const auto n = rot->ne[0]; + GGML_ASSERT(cur->ne[0] % n == 0); ggml_tensor * res; if (!ggml_is_contiguous(cur)) { - res = ggml_cont_2d(ctx, cur, n, ggml_nelements(cur)/n); + res = ggml_cont_4d(ctx, cur, n, cur->ne[0]/n, cur->ne[1], cur->ne[2]*cur->ne[3]); } else { - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); + res = ggml_reshape_4d(ctx, cur, n, cur->ne[0]/n, cur->ne[1], cur->ne[2]*cur->ne[3]); } res = ggml_mul_mat(ctx, rot, res); ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index a31e031fb1c3..3bf8245e406c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -605,6 +605,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str GGML_ASSERT(tensor->ne[axis] == 2*key_dim + value_dim); return {{key_dim, 2 + head_ratio}}; } + if (ud->model->hadamard_gdn_v_grouped && std::regex_match(tensor_name, pattern_ssm_out_weight)) { + return {{value_dim, 1}}; + } if (std::regex_match(tensor_name, pattern_attn_gate_weight) || std::regex_match(tensor_name, pattern_ssm_out_weight)) { return {{key_dim, head_ratio}}; } @@ -659,6 +662,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t head_dim = hparams.ssm_d_state; const int64_t blck_size_perf = std::lcm(blck_size, 128); const int64_t granularity_qkv = std::lcm(blck_size_perf, head_dim); + if (ud->model->hadamard_gdn_v_grouped && std::regex_match(tensor_name, pattern_ssm_out_weight)) { + return std::vector(segments.size(), granularity_qkv * (hparams.ssm_dt_rank / hparams.ssm_n_group)); + } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_attn_gate_weight) || std::regex_match(tensor_name, pattern_ssm_conv1d) || std::regex_match(tensor_name, pattern_ssm_out_weight)) { return std::vector(segments.size(), granularity_qkv); @@ -771,7 +777,11 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str tensor_config tc = get_tensor_config(); split_state.axis = tc.axis; if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { - const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); + int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); + if (const auto it = ud->model->hadamard_weight_blocks.find(tc.tensor_axis_0->name); + it != ud->model->hadamard_weight_blocks.end()) { + blck_size = std::lcm(blck_size, int64_t(it->second)); + } const float * tensor_split = ud->model->tensor_split(); std::vector tensor_split_scan; tensor_split_scan.reserve(ud->n_devices);