diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 41566d41aef3..6f85059b0722 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -756,6 +756,7 @@ extern "C" { GGML_API size_t ggml_element_size(const struct ggml_tensor * tensor); GGML_API bool ggml_is_quantized(enum ggml_type type); + GGML_API bool ggml_needs_scale_quantized(enum ggml_type type); // TODO: temporary until model loading of ggml examples is refactored GGML_API enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype); @@ -1419,6 +1420,13 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + GGML_API struct ggml_tensor * ggml_mul_mat_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations); + // change the precision of a matrix multiplication // set to GGML_PREC_F32 for higher precision (useful for phi-2) GGML_API void ggml_mul_mat_set_prec( @@ -1437,6 +1445,14 @@ extern "C" { struct ggml_tensor * b, struct ggml_tensor * ids); + GGML_API struct ggml_tensor * ggml_mul_mat_id_ext( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations); + // A: m columns, n rows, // B: p columns, n rows, // result is m columns, p rows @@ -2802,6 +2818,7 @@ extern "C" { int64_t blck_size_interleave; // interleave elements in blocks size_t type_size; bool is_quantized; + bool needs_scale; // whether the quantization type needs a scale factor for valid dequantization ggml_to_float_t to_float; ggml_from_float_t from_float_ref; }; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 476c30797956..625474a5c3a1 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -746,6 +746,7 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .blck_size = QK_NVFP4, .type_size = sizeof(block_nvfp4), .is_quantized = true, + .needs_scale = true, .to_float = (ggml_to_float_t) dequantize_row_nvfp4, .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, }, @@ -1335,6 +1336,13 @@ bool ggml_is_quantized(enum ggml_type type) { return type_traits[type].is_quantized; } +bool ggml_needs_scale_quantized(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + assert(!type_traits[type].needs_scale || type_traits[type].is_quantized); + return type_traits[type].needs_scale; +} + const char * ggml_op_name(enum ggml_op op) { return GGML_OP_NAME[op]; } @@ -3241,8 +3249,35 @@ struct ggml_tensor * ggml_mul_mat( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b) { + if (ggml_needs_scale_quantized(a->type) || ggml_needs_scale_quantized(b->type)) { + GGML_LOG_ERROR("%s: tensor types %s or %s requires explicit dequantization scales; use ggml_mul_mat_ext instead\n", + __func__, ggml_type_name(a->type), ggml_type_name(b->type)); + GGML_ABORT("fatal error"); + } + + return ggml_mul_mat_ext(ctx, a, b, NULL, NULL); +} + +struct ggml_tensor * ggml_mul_mat_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations) { GGML_ASSERT(ggml_can_mul_mat(a, b)); GGML_ASSERT(!ggml_is_transposed(a)); + if (ggml_needs_scale_quantized(a->type) && scale_weight == NULL) { + GGML_LOG_ERROR("%s: tensor type %s requires explicit dequantization scales; pass scale_weight to ggml_mul_mat_ext\n", + __func__, ggml_type_name(a->type)); + GGML_ABORT("fatal error"); + } + if (ggml_needs_scale_quantized(b->type)) { + GGML_LOG_ERROR("%s: scaled tensor type %s currently cannot be used as the activation tensor\n", + __func__, ggml_type_name(b->type)); + GGML_ABORT("fatal error"); + } + GGML_ASSERT(scale_weight == NULL || scale_weight->type == GGML_TYPE_F32); + GGML_ASSERT(scale_activations == NULL || scale_activations->type == GGML_TYPE_F32); const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); @@ -3250,6 +3285,15 @@ struct ggml_tensor * ggml_mul_mat( result->op = GGML_OP_MUL_MAT; result->src[0] = a; result->src[1] = b; + // TODO: decide during review whether scale_weight should be attached as matmul metadata + // or inferred only from the post-matmul multiply. + result->src[2] = scale_weight; + result->src[3] = scale_activations; + + if (scale_weight) { + GGML_ASSERT(ggml_can_repeat(scale_weight, result)); + result = ggml_mul(ctx, result, scale_weight); + } return result; } @@ -3257,6 +3301,10 @@ struct ggml_tensor * ggml_mul_mat( void ggml_mul_mat_set_prec( struct ggml_tensor * a, enum ggml_prec prec) { + if (a->op == GGML_OP_MUL && a->src[0] && a->src[0]->op == GGML_OP_MUL_MAT) { + a = a->src[0]; + } + GGML_ASSERT(a->op == GGML_OP_MUL_MAT); const int32_t prec_i32 = (int32_t) prec; @@ -3267,6 +3315,10 @@ void ggml_mul_mat_set_prec( void ggml_mul_mat_set_hint( struct ggml_tensor * a, enum ggml_op_hint hint) { + if (a->op == GGML_OP_MUL && a->src[0] && a->src[0]->op == GGML_OP_MUL_MAT) { + a = a->src[0]; + } + GGML_ASSERT(a->op == GGML_OP_MUL_MAT); const int32_t hint_i32 = (int32_t) hint; @@ -3293,8 +3345,24 @@ struct ggml_tensor * ggml_mul_mat_id( struct ggml_tensor * as, struct ggml_tensor * b, struct ggml_tensor * ids) { + GGML_ASSERT(!ggml_needs_scale_quantized(as->type) && !ggml_needs_scale_quantized(b->type)); + + return ggml_mul_mat_id_ext(ctx, as, b, ids, NULL, NULL); +} + +struct ggml_tensor * ggml_mul_mat_id_ext( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids, + struct ggml_tensor * scale_weight, + struct ggml_tensor * scale_activations) { GGML_ASSERT(!ggml_is_transposed(as)); GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(!ggml_needs_scale_quantized(as->type) || scale_weight != NULL); + GGML_ASSERT(!ggml_needs_scale_quantized(b->type)); + GGML_ASSERT(scale_weight == NULL || scale_weight->type == GGML_TYPE_F32); + GGML_ASSERT(scale_activations == NULL || scale_activations->type == GGML_TYPE_F32); GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) GGML_ASSERT(b->ne[3] == 1); // b is 3d @@ -3310,6 +3378,22 @@ struct ggml_tensor * ggml_mul_mat_id( result->src[0] = as; result->src[1] = b; result->src[2] = ids; + // TODO: decide during review whether scale_weight should be attached as matmul metadata + // or inferred only from the post-matmul multiply. + result->src[3] = scale_weight; + result->src[4] = scale_activations; + + if (scale_weight) { + struct ggml_tensor * s = scale_weight; + if (s->ne[0] == as->ne[2] && s->ne[1] == 1 && s->ne[2] == 1 && s->ne[3] == 1) { + s = ggml_reshape_3d(ctx, s, 1, as->ne[2], 1); + s = ggml_repeat_4d(ctx, s, 1, as->ne[2], b->ne[2], 1); + s = ggml_get_rows(ctx, s, ids); + } + + GGML_ASSERT(ggml_can_repeat(s, result)); + result = ggml_mul(ctx, result, s); + } return result; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 31cf41a1c2d2..918abe4aec72 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -977,8 +977,9 @@ ggml_tensor * llm_graph_context::build_cvec( ggml_tensor * llm_graph_context::build_lora_mm( ggml_tensor * w, ggml_tensor * cur, - ggml_tensor * w_s) const { - ggml_tensor * res = ggml_mul_mat(ctx0, w, cur); + ggml_tensor * w_s, + ggml_tensor * in_s) const { + ggml_tensor * res = ggml_mul_mat_ext(ctx0, w, cur, w_s, in_s); for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); @@ -998,18 +999,16 @@ ggml_tensor * llm_graph_context::build_lora_mm( res = ggml_add(ctx0, res, ab_cur); } - if (w_s) { - res = ggml_mul(ctx0, res, w_s); - } - return res; } ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b - ggml_tensor * ids) const { - ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids); + ggml_tensor * ids, + ggml_tensor * w_s, + ggml_tensor * in_s) const { + ggml_tensor * res = ggml_mul_mat_id_ext(ctx0, w, cur, ids, w_s, in_s); for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); if (lw == nullptr) { @@ -1083,7 +1082,7 @@ llm_graph_qkv llm_graph_context::build_qkv( if (layer.wqkv) { // fused QKV path - ggml_tensor * qkv = build_lora_mm(layer.wqkv, cur, layer.wqkv_s); + ggml_tensor * qkv = build_lora_mm(layer.wqkv, cur, layer.wqkv_s, layer.wqkv_in_s); cb(qkv, "wqkv", il); if (layer.wqkv_b) { qkv = ggml_add(ctx0, qkv, layer.wqkv_b); @@ -1103,7 +1102,7 @@ llm_graph_qkv llm_graph_context::build_qkv( ggml_row_size(qkv->type, n_embd_q + n_embd_kv)); } else { // separate Q/K/V path - Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); + Qcur = build_lora_mm(layer.wq, cur, layer.wq_s, layer.wq_in_s); cb(Qcur, "Qcur", il); if (layer.wq_b) { Qcur = ggml_add(ctx0, Qcur, layer.wq_b); @@ -1113,7 +1112,7 @@ llm_graph_qkv llm_graph_context::build_qkv( Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Qcur, "Qcur_clamped", il); } - Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); + Kcur = build_lora_mm(layer.wk, cur, layer.wk_s, layer.wk_in_s); cb(Kcur, "Kcur", il); if (layer.wk_b) { Kcur = ggml_add(ctx0, Kcur, layer.wk_b); @@ -1123,7 +1122,7 @@ llm_graph_qkv llm_graph_context::build_qkv( Kcur = ggml_clamp(ctx0, Kcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Kcur, "Kcur_clamped", il); } - Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + Vcur = build_lora_mm(layer.wv, cur, layer.wv_s, layer.wv_in_s); cb(Vcur, "Vcur", il); if (layer.wv_b) { Vcur = ggml_add(ctx0, Vcur, layer.wv_b); @@ -1160,8 +1159,18 @@ ggml_tensor * llm_graph_context::build_ffn( ggml_tensor * act_scales, llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, - int il) const { - ggml_tensor * tmp = up ? build_lora_mm(up, cur) : cur; + int il, + ggml_tensor * up_in_s, + ggml_tensor * gate_in_s, + ggml_tensor * down_in_s) const { + // TODO: Split the shared use of up/gate/down_s (Architecture-specific scaling factors vs. quantization scaling factors for derived quantization types) + const bool up_derived = up && ggml_needs_scale_quantized(up->type); + const bool gate_derived = gate && ggml_needs_scale_quantized(gate->type); + const bool down_derived = down && ggml_needs_scale_quantized(down->type); + + ggml_tensor * tmp = up ? build_lora_mm(up, cur, + up_derived ? up_s : nullptr, + up_derived ? up_in_s : nullptr) : cur; cb(tmp, "ffn_up", il); if (up_b) { @@ -1169,7 +1178,7 @@ ggml_tensor * llm_graph_context::build_ffn( cb(tmp, "ffn_up_b", il); } - if (up_s) { + if (up_s && !up_derived) { tmp = ggml_mul(ctx0, tmp, up_s); cb(tmp, "ffn_up_s", il); } @@ -1178,12 +1187,16 @@ ggml_tensor * llm_graph_context::build_ffn( switch (type_gate) { case LLM_FFN_SEQ: { - cur = build_lora_mm(gate, tmp); + cur = build_lora_mm(gate, tmp, + gate_derived ? gate_s : nullptr, + gate_derived ? gate_in_s : nullptr); cb(cur, "ffn_gate", il); } break; case LLM_FFN_PAR: { - cur = build_lora_mm(gate, cur); + cur = build_lora_mm(gate, cur, + gate_derived ? gate_s : nullptr, + gate_derived ? gate_in_s : nullptr); cb(cur, "ffn_gate", il); } break; } @@ -1193,7 +1206,7 @@ ggml_tensor * llm_graph_context::build_ffn( cb(cur, "ffn_gate_b", il); } - if (gate_s) { + if (gate_s && !gate_derived) { cur = ggml_mul(ctx0, cur, gate_s); cb(cur, "ffn_gate_s", il); } @@ -1287,7 +1300,9 @@ ggml_tensor * llm_graph_context::build_ffn( } if (down) { - cur = build_lora_mm(down, cur); + cur = build_lora_mm(down, cur, + down_derived ? down_s : nullptr, + down_derived ? down_in_s : nullptr); if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators ggml_mul_mat_set_prec(cur, GGML_PREC_F32); @@ -1302,7 +1317,7 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_add(ctx0, cur, down_b); } - if (down_s) { + if (down_s && !down_derived) { cur = ggml_mul(ctx0, cur, down_s); cb(cur, "ffn_down_s", il); } @@ -1328,7 +1343,10 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * gate_up_exps, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, - ggml_tensor * down_exps_s) const { + ggml_tensor * down_exps_s, + ggml_tensor * up_exps_in_s, + ggml_tensor * gate_exps_in_s, + ggml_tensor * down_exps_in_s) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -1348,7 +1366,10 @@ ggml_tensor * llm_graph_context::build_moe_ffn( /* gate_up_exps_b */ nullptr, up_exps_s, gate_exps_s, - down_exps_s + down_exps_s, + up_exps_in_s, + gate_exps_in_s, + down_exps_in_s ); } @@ -1375,7 +1396,10 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * gate_up_exps_b, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, - ggml_tensor * down_exps_s) const { + ggml_tensor * down_exps_s, + ggml_tensor * up_exps_in_s, + ggml_tensor * gate_exps_in_s, + ggml_tensor * down_exps_in_s) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN @@ -1514,12 +1538,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(cur, "ffn_moe_weighted", il); } - ggml_tensor * up = nullptr; + ggml_tensor * up = nullptr; ggml_tensor * experts = nullptr; if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views - ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts); // [n_ff*2, n_expert_used, n_tokens] + ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s, + up_exps_in_s); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il); if (gate_up_exps_b) { @@ -1527,23 +1552,16 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(gate_up, "ffn_moe_gate_up_biased", il); } - // apply per-expert scale2 to merged gate_up (use up_exps_s since gate and up are fused) - if (up_exps_s) { - ggml_tensor * s = ggml_reshape_3d(ctx0, up_exps_s, 1, n_expert, 1); - s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts); // [1, n_expert_used, n_tokens] - gate_up = ggml_mul(ctx0, gate_up, s); - cb(gate_up, "ffn_moe_gate_up_scaled", il); - } - const int64_t n_ff = gate_up->ne[0] / 2; cur = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], 0); cb(cur, "ffn_moe_gate", il); - up = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], n_ff * gate_up->nb[0]); + up = ggml_view_3d(ctx0, gate_up, n_ff, gate_up->ne[1], gate_up->ne[2], gate_up->nb[1], gate_up->nb[2], + n_ff * gate_up->nb[0]); cb(up, "ffn_moe_up", il); } else { // separate gate and up path - up = build_lora_mm_id(up_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens] + up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s, + up_exps_in_s); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il); if (up_exps_b) { @@ -1551,17 +1569,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(up, "ffn_moe_up_biased", il); } - // apply per-expert scale2 to up - if (up_exps_s) { - ggml_tensor * s = ggml_reshape_3d(ctx0, up_exps_s, 1, n_expert, 1); - s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts); // [1, n_expert_used, n_tokens] - up = ggml_mul(ctx0, up, s); - cb(up, "ffn_moe_up_scaled", il); - } - if (gate_exps) { - cur = build_lora_mm_id(gate_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens] + cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s, + gate_exps_in_s); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); } else { cur = up; @@ -1571,15 +1581,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cur = ggml_add_id(ctx0, cur, gate_exps_b, selected_experts); cb(cur, "ffn_moe_gate_biased", il); } - - // apply per-expert scale2 to gate - if (gate_exps_s) { - ggml_tensor * s = ggml_reshape_3d(ctx0, gate_exps_s, 1, n_expert, 1); - s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts); // [1, n_expert_used, n_tokens] - cur = ggml_mul(ctx0, cur, s); - cb(cur, "ffn_moe_gate_scaled", il); - } } const bool has_gate = gate_exps || gate_up_exps; @@ -1646,12 +1647,14 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cur = ggml_relu(ctx0, cur); cur = ggml_sqr(ctx0, cur); cb(cur, "ffn_moe_relu_sqr", il); - } break; + } + break; default: GGML_ABORT("fatal error"); } - experts = build_lora_mm_id(down_exps, cur, selected_experts); // [n_embd, n_expert_used, n_tokens] + experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, + down_exps_in_s); // [n_embd, n_expert_used, n_tokens] cb(experts, "ffn_moe_down", il); if (down_exps_b) { @@ -1659,15 +1662,6 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(experts, "ffn_moe_down_biased", il); } - // apply per-expert scale2 to down - if (down_exps_s) { - ggml_tensor * s = ggml_reshape_3d(ctx0, down_exps_s, 1, n_expert, 1); - s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, selected_experts); // [1, n_expert_used, n_tokens] - experts = ggml_mul(ctx0, experts, s); - cb(experts, "ffn_moe_down_scaled", il); - } - if (!weight_before_ffn) { experts = ggml_mul(ctx0, experts, weights); cb(experts, "ffn_moe_weighted", il); @@ -2235,15 +2229,10 @@ ggml_tensor * llm_graph_context::build_attn( } if (wo) { + cur = build_lora_mm(wo, cur, wo_s); if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE || arch == LLM_ARCH_JAIS2) { // GLM4, GLM4_MOE, and JAIS2 seem to have numerical issues with half-precision accumulators - cur = build_lora_mm(wo, cur); ggml_mul_mat_set_prec(cur, GGML_PREC_F32); - if (wo_s) { - cur = ggml_mul(ctx0, cur, wo_s); - } - } else { - cur = build_lora_mm(wo, cur, wo_s); } } @@ -2322,15 +2311,9 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (wo) { + cur = build_lora_mm(wo, cur, wo_s); if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) { - // GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators - cur = build_lora_mm(wo, cur); ggml_mul_mat_set_prec(cur, GGML_PREC_F32); - if (wo_s) { - cur = ggml_mul(ctx0, cur, wo_s); - } - } else { - cur = build_lora_mm(wo, cur, wo_s); } } diff --git a/src/llama-graph.h b/src/llama-graph.h index bf6778237e6f..c10502e967ca 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -789,13 +789,16 @@ struct llm_graph_context { ggml_tensor * build_lora_mm( ggml_tensor * w, ggml_tensor * cur, - ggml_tensor * w_s = nullptr) const; + ggml_tensor * w_s = nullptr, + ggml_tensor * in_s = nullptr) const; // do mat_mul_id, while optionally apply lora ggml_tensor * build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b - ggml_tensor * ids) const; + ggml_tensor * ids, + ggml_tensor * w_s = nullptr, + ggml_tensor * in_s = nullptr) const; ggml_tensor * build_norm( ggml_tensor * cur, @@ -829,7 +832,10 @@ struct llm_graph_context { ggml_tensor * act_scales, llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, - int il) const; + int il, + ggml_tensor * up_in_s = nullptr, + ggml_tensor * gate_in_s = nullptr, + ggml_tensor * down_in_s = nullptr) const; // build MoE FFN without bias tensors ggml_tensor * build_moe_ffn( @@ -850,7 +856,10 @@ struct llm_graph_context { ggml_tensor * gate_up_exps = nullptr, ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, - ggml_tensor * down_exps_s = nullptr) const; + ggml_tensor * down_exps_s = nullptr, + ggml_tensor * up_exps_in_s = nullptr, + ggml_tensor * gate_exps_in_s = nullptr, + ggml_tensor * down_exps_in_s = nullptr) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -875,7 +884,10 @@ struct llm_graph_context { ggml_tensor * gate_up_exps_b = nullptr, ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, - ggml_tensor * down_exps_s = nullptr) const; + ggml_tensor * down_exps_s = nullptr, + ggml_tensor * up_exps_in_s = nullptr, + ggml_tensor * gate_exps_in_s = nullptr, + ggml_tensor * down_exps_in_s = nullptr) const; // // inputs diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index c645d0785ab7..546f3fdfa4cc 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -908,7 +908,8 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w } ggml_context * ctx = ctx_ptr.get(); - ggml_tensor * op_tensor = nullptr; + ggml_tensor * op_tensor = nullptr; + ggml_tensor * op_tensor_companion = nullptr; switch (op) { case GGML_OP_GET_ROWS: @@ -919,7 +920,12 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w case GGML_OP_MUL_MAT: { ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]); - op_tensor = ggml_mul_mat(ctx, w, b); + ggml_tensor * s = ggml_needs_scale_quantized(w->type) ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1) : nullptr; + op_tensor = ggml_mul_mat_ext(ctx, w, b, s, nullptr); + if (s) { + op_tensor_companion = op_tensor; + op_tensor = op_tensor_companion->src[0]; + } } break; case GGML_OP_MUL_MAT_ID: { @@ -927,7 +933,12 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w GGML_ASSERT(n_expert_used > 0); ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512); ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512); - op_tensor = ggml_mul_mat_id(ctx, w, b, ids); + ggml_tensor * s = ggml_needs_scale_quantized(w->type) ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, w->ne[2]) : nullptr; + op_tensor = ggml_mul_mat_id_ext(ctx, w, b, ids, s, nullptr); + if (s) { + op_tensor_companion = op_tensor; + op_tensor = op_tensor_companion->src[0]; + } } break; case GGML_OP_ADD: { @@ -1022,6 +1033,9 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w GGML_ASSERT(w->buffer == nullptr); w->buffer = ggml_backend_buft_alloc_buffer(buft, 0); bool op_supported = ggml_backend_dev_supports_op(dev, op_tensor); + if (op_supported && op_tensor_companion != nullptr) { + op_supported = ggml_backend_dev_supports_op(dev, op_tensor_companion); + } ggml_backend_buffer_free(w->buffer); w->buffer = nullptr; diff --git a/src/models/llama.cpp b/src/models/llama.cpp index cef66d054b0c..d65fb0b95480 100644 --- a/src/models/llama.cpp +++ b/src/models/llama.cpp @@ -189,7 +189,10 @@ llama_model_llama::graph::graph(const llama_model & model, const llm_grap model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s, NULL, - LLM_FFN_SILU, LLM_FFN_PAR, il); + LLM_FFN_SILU, LLM_FFN_PAR, il, + model.layers[il].ffn_up_in_s, + model.layers[il].ffn_gate_in_s, + model.layers[il].ffn_down_in_s); cb(cur, "ffn_out", il); } else { // MoE branch @@ -212,7 +215,10 @@ llama_model_llama::graph::graph(const llama_model & model, const llm_grap nullptr, nullptr, model.layers[il].ffn_up_exps_s, model.layers[il].ffn_gate_exps_s, - model.layers[il].ffn_down_exps_s); + model.layers[il].ffn_down_exps_s, + model.layers[il].ffn_up_exps_in_s, + model.layers[il].ffn_gate_exps_in_s, + model.layers[il].ffn_down_exps_in_s); cb(cur, "ffn_moe_out", il); } cur = ggml_add(ctx0, cur, ffn_inp); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f54ab41c1953..4174113b1ae0 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4052,10 +4052,15 @@ struct test_mul_mat : public test_case { ggml_set_name(b, "b"); } - ggml_tensor * out = ggml_mul_mat(ctx, a, b); + ggml_tensor * scale_weight = nullptr; + if (ggml_needs_scale_quantized(type_a)) { + scale_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_set_name(scale_weight, "scale_weight"); + } + ggml_tensor * out = ggml_mul_mat_ext(ctx, a, b, scale_weight, nullptr); ggml_set_name(out, "out"); for (uint32_t i = 1; i < o; ++i) { - ggml_tensor * out2 = ggml_mul_mat(ctx, a, b); + ggml_tensor * out2 = ggml_mul_mat_ext(ctx, a, b, scale_weight, nullptr); ggml_set_name(out2, "out2"); out = ggml_add(ctx, out, out2); } @@ -4200,7 +4205,12 @@ struct test_mul_mat_id : public test_case { ggml_tensor * b = ggml_new_tensor_3d(ctx, type_b, k, this->b ? 1 : n_used, n); ggml_set_name(b, "b"); - ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); + ggml_tensor * scale_weight = nullptr; + if (ggml_needs_scale_quantized(type_a)) { + scale_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_mats); + ggml_set_name(scale_weight, "scale_weight"); + } + ggml_tensor * out = ggml_mul_mat_id_ext(ctx, as, b, ids, scale_weight, nullptr); ggml_set_name(out, "out"); return out; @@ -4209,6 +4219,11 @@ struct test_mul_mat_id : public test_case { void initialize_tensors(ggml_context * ctx) override { init_mul_mat_id_tensors(ctx, n_mats); } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return ggml_op_name(GGML_OP_MUL_MAT_ID); + } }; // GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL @@ -5803,14 +5818,26 @@ struct test_mul_mat_vec_fusion : public test_case { ggml_tensor * gate = with_gate ? ggml_new_tensor(ctx, type, 4, ne0.data()) : nullptr; ggml_tensor * up = ggml_new_tensor(ctx, type, 4, ne0.data()); - ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); + ggml_tensor * up_scale = nullptr; + if (ggml_needs_scale_quantized(type)) { + up_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_set_name(up_scale, "up_scale"); + } + + ggml_tensor * ffn_up = up_scale ? ggml_mul_mat_ext(ctx, up, cur, up_scale, nullptr) : ggml_mul_mat(ctx, up, cur); if (with_bias) { std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); ffn_up = ggml_add(ctx, ffn_up, up_bias); } - ggml_tensor * ffn_gate = with_gate ? ggml_mul_mat(ctx, gate, cur) : nullptr; + ggml_tensor * gate_scale = nullptr; + if (with_gate && ggml_needs_scale_quantized(type)) { + gate_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_set_name(gate_scale, "gate_scale"); + } + + ggml_tensor * ffn_gate = with_gate ? (gate_scale ? ggml_mul_mat_ext(ctx, gate, cur, gate_scale, nullptr) : ggml_mul_mat(ctx, gate, cur)) : nullptr; if (with_bias && with_gate) { std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); @@ -5837,13 +5864,25 @@ struct test_mul_mat_vec_fusion : public test_case { ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, this->b ? 1 : n_used, m); ggml_set_name(cur, "cur"); - ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, ids); + ggml_tensor * up_scale = nullptr; + if (ggml_needs_scale_quantized(type)) { + up_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_mats); + ggml_set_name(up_scale, "up_scale"); + } + + ggml_tensor * ffn_up = up_scale ? ggml_mul_mat_id_ext(ctx, ups, cur, ids, up_scale, nullptr) : ggml_mul_mat_id(ctx, ups, cur, ids); if (with_bias) { ggml_tensor * up_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_up->ne[0], n_mats); ffn_up = ggml_add_id(ctx, ffn_up, up_bias_param, ids); } - ggml_tensor * ffn_gate = with_gate? ggml_mul_mat_id(ctx, gates, cur, ids) : nullptr; + ggml_tensor * gate_scale = nullptr; + if (with_gate && ggml_needs_scale_quantized(type)) { + gate_scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_mats); + ggml_set_name(gate_scale, "gate_scale"); + } + + ggml_tensor * ffn_gate = with_gate ? (gate_scale ? ggml_mul_mat_id_ext(ctx, gates, cur, ids, gate_scale, nullptr) : ggml_mul_mat_id(ctx, gates, cur, ids)) : nullptr; if (with_bias && with_gate) { ggml_tensor * gate_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_gate->ne[0], n_mats); ffn_gate = ggml_add_id(ctx, ffn_gate, gate_bias_param, ids);