Merge old mul_mat patterns into the lane-based approach

This commit is contained in:
Oliver Simons
2026-06-02 15:57:11 +02:00
parent d6190b97a5
commit b07d0f03d3
2 changed files with 173 additions and 490 deletions
+77 -39
View File
@@ -5789,20 +5789,21 @@ struct test_mul_mat_vec_fusion : public test_case {
const bool with_bias;
const bool with_gate;
const bool with_lane_scale;
const bool gate_first;
std::array<int64_t, 2> batch_dims;
test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k,
bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true,
bool with_lane_scale = false, std::array<int64_t, 2> batch_dims = {4, 2})
bool with_lane_scale = false, bool gate_first = false, std::array<int64_t, 2> batch_dims = {4, 2})
: type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias),
with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims) {
with_gate(with_gate), with_lane_scale(with_lane_scale), gate_first(gate_first), batch_dims(batch_dims) {
if (use_id) {
GGML_ASSERT(n_used <= n_mats);
}
}
std::string vars() override {
return VARS_TO_STR13(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims);
return VARS_TO_STR14(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, gate_first, batch_dims);
}
std::string op_desc(ggml_tensor * t) override {
@@ -5850,24 +5851,40 @@ 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);
if (with_bias) {
std::array<int64_t, 4> 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);
}
if (with_lane_scale) {
ffn_up = build_dense_lane_scale(ctx, ffn_up);
}
auto build_up_lane = [&]() {
ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur);
if (with_bias) {
std::array<int64_t, 4> 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);
}
if (with_lane_scale) {
ffn_up = build_dense_lane_scale(ctx, ffn_up);
}
return ffn_up;
};
ggml_tensor * ffn_gate = with_gate ? ggml_mul_mat(ctx, gate, cur) : nullptr;
if (with_bias && with_gate) {
std::array<int64_t, 4> bias_ne = { ffn_gate->ne[0], 1, channels, samples };
ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data());
ffn_gate = ggml_add(ctx, ffn_gate, gate_bias);
}
if (with_lane_scale && with_gate) {
ffn_gate = build_dense_lane_scale(ctx, ffn_gate);
auto build_gate_lane = [&]() {
ggml_tensor * ffn_gate = ggml_mul_mat(ctx, gate, cur);
if (with_bias) {
std::array<int64_t, 4> bias_ne = { ffn_gate->ne[0], 1, channels, samples };
ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data());
ffn_gate = ggml_add(ctx, ffn_gate, gate_bias);
}
if (with_lane_scale) {
ffn_gate = build_dense_lane_scale(ctx, ffn_gate);
}
return ffn_gate;
};
ggml_tensor * ffn_up = nullptr;
ggml_tensor * ffn_gate = nullptr;
if (with_gate && gate_first) {
ffn_gate = build_gate_lane();
ffn_up = build_up_lane();
} else {
ffn_up = build_up_lane();
ffn_gate = with_gate ? build_gate_lane() : nullptr;
}
ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up;
@@ -5890,22 +5907,38 @@ 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);
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);
}
if (with_lane_scale) {
ffn_up = build_id_lane_scale(ctx, ffn_up, ids);
}
auto build_up_lane = [&]() {
ggml_tensor * ffn_up = 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);
}
if (with_lane_scale) {
ffn_up = build_id_lane_scale(ctx, ffn_up, ids);
}
return ffn_up;
};
ggml_tensor * ffn_gate = with_gate? 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);
}
if (with_lane_scale && with_gate) {
ffn_gate = build_id_lane_scale(ctx, ffn_gate, ids);
auto build_gate_lane = [&]() {
ggml_tensor * ffn_gate = ggml_mul_mat_id(ctx, gates, cur, ids);
if (with_bias) {
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);
}
if (with_lane_scale) {
ffn_gate = build_id_lane_scale(ctx, ffn_gate, ids);
}
return ffn_gate;
};
ggml_tensor * ffn_up = nullptr;
ggml_tensor * ffn_gate = nullptr;
if (with_gate && gate_first) {
ffn_gate = build_gate_lane();
ffn_up = build_up_lane();
} else {
ffn_up = build_up_lane();
ffn_gate = with_gate ? build_gate_lane() : nullptr;
}
ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up;
@@ -9111,10 +9144,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
continue;
}
for (bool with_lane_scale : {false, true}) {
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale));
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
for (bool gate_first : {false, true}) {
if (!with_gate && gate_first) {
continue;
}
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, gate_first));
test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, gate_first, {1, 1}));
}
}
}
}