diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6d77f64d99..36ef22a29c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3847,8 +3847,8 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, // Matched MM lane forms: -// MUL_MAT [ADD] [MUL scalar_scale] -// MUL_MAT_ID [ADD_ID] [RESHAPE -> REPEAT -> GET_ROWS -> MUL expert_scale] +// MUL_MAT [MUL scalar_scale] [ADD] +// MUL_MAT_ID [RESHAPE -> REPEAT -> GET_ROWS -> MUL expert_scale] [ADD_ID] struct ggml_cuda_mm_lane { ggml_tensor * mm = nullptr; ggml_tensor * bias_node = nullptr; @@ -3894,6 +3894,40 @@ static bool ggml_cuda_parse_mul_mat_id_lane(const ggml_cgraph * cgraph, int i, g lane.out = mm; lane.n_nodes = 1; + if (ggml_is_quantized(mm->src[0]->type) && ggml_cuda_can_fuse_mm_lane_scale(mm) && i + lane.n_nodes + 3 < cgraph->n_nodes) { + ggml_tensor * reshape = cgraph->nodes[i + lane.n_nodes + 0]; + ggml_tensor * repeat = cgraph->nodes[i + lane.n_nodes + 1]; + ggml_tensor * getrows = cgraph->nodes[i + lane.n_nodes + 2]; + ggml_tensor * mul = cgraph->nodes[i + lane.n_nodes + 3]; + + if (reshape->op == GGML_OP_RESHAPE && repeat->op == GGML_OP_REPEAT && + getrows->op == GGML_OP_GET_ROWS && mul->op == GGML_OP_MUL) { + if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm->src[2]) { + return false; + } + + const bool mul_has_out = mul->src[0] == lane.out || mul->src[1] == lane.out; + const bool mul_has_scale = mul->src[0] == getrows || mul->src[1] == getrows; + if (!mul_has_out || !mul_has_scale) { + return false; + } + + const ggml_tensor * scale = reshape->src[0]; + if (scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm->src[0]->ne[2]) { + return false; + } + + if (mul->type != GGML_TYPE_F32 || !ggml_are_same_shape(mul, lane.out)) { + return false; + } + + lane.scale = scale; + lane.scale_view_node = i + lane.n_nodes; + lane.out = mul; + lane.n_nodes += 4; + } + } + if (i + lane.n_nodes < cgraph->n_nodes && cgraph->nodes[i + lane.n_nodes]->op == GGML_OP_ADD_ID) { ggml_tensor * add = cgraph->nodes[i + lane.n_nodes]; if (add->src[0] != lane.out || add->src[2] != mm->src[2] || add->type != GGML_TYPE_F32) { @@ -3908,43 +3942,6 @@ static bool ggml_cuda_parse_mul_mat_id_lane(const ggml_cgraph * cgraph, int i, g lane.n_nodes++; } - if (!ggml_is_quantized(mm->src[0]->type) || !ggml_cuda_can_fuse_mm_lane_scale(mm) || i + lane.n_nodes + 3 >= cgraph->n_nodes) { - return true; - } - - ggml_tensor * reshape = cgraph->nodes[i + lane.n_nodes + 0]; - ggml_tensor * repeat = cgraph->nodes[i + lane.n_nodes + 1]; - ggml_tensor * getrows = cgraph->nodes[i + lane.n_nodes + 2]; - ggml_tensor * mul = cgraph->nodes[i + lane.n_nodes + 3]; - - if (reshape->op != GGML_OP_RESHAPE || repeat->op != GGML_OP_REPEAT || - getrows->op != GGML_OP_GET_ROWS || mul->op != GGML_OP_MUL) { - return true; - } - - if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm->src[2]) { - return true; - } - - const bool mul_has_out = mul->src[0] == lane.out || mul->src[1] == lane.out; - const bool mul_has_scale = mul->src[0] == getrows || mul->src[1] == getrows; - if (!mul_has_out || !mul_has_scale) { - return true; - } - - const ggml_tensor * scale = reshape->src[0]; - if (scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm->src[0]->ne[2]) { - return false; - } - - if (mul->type != GGML_TYPE_F32 || !ggml_are_same_shape(mul, lane.out)) { - return false; - } - - lane.scale = scale; - lane.scale_view_node = i + lane.n_nodes; - lane.out = mul; - lane.n_nodes += 4; return true; } @@ -3963,6 +3960,29 @@ static bool ggml_cuda_parse_mul_mat_lane(const ggml_cgraph * cgraph, int i, ggml lane.out = mm; lane.n_nodes = 1; + if (ggml_is_quantized(mm->src[0]->type) && ggml_cuda_can_fuse_mm_lane_scale(mm) && + i + lane.n_nodes < cgraph->n_nodes && cgraph->nodes[i + lane.n_nodes]->op == GGML_OP_MUL) { + ggml_tensor * mul = cgraph->nodes[i + lane.n_nodes]; + const bool mul_lhs_out = mul->src[0] == lane.out; + const bool mul_rhs_out = mul->src[1] == lane.out; + if (!mul_lhs_out && !mul_rhs_out) { + return false; + } + + const ggml_tensor * scale = mul_lhs_out ? mul->src[1] : mul->src[0]; + if (scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1) { + return false; + } + + if (mul->type != GGML_TYPE_F32 || !ggml_are_same_shape(mul, lane.out)) { + return false; + } + + lane.scale = scale; + lane.out = mul; + lane.n_nodes++; + } + if (i + lane.n_nodes < cgraph->n_nodes && cgraph->nodes[i + lane.n_nodes]->op == GGML_OP_ADD) { ggml_tensor * add = cgraph->nodes[i + lane.n_nodes]; if (add->src[0] == lane.out) { @@ -3983,30 +4003,6 @@ static bool ggml_cuda_parse_mul_mat_lane(const ggml_cgraph * cgraph, int i, ggml lane.n_nodes++; } - if (!ggml_is_quantized(mm->src[0]->type) || !ggml_cuda_can_fuse_mm_lane_scale(mm) || - i + lane.n_nodes >= cgraph->n_nodes || cgraph->nodes[i + lane.n_nodes]->op != GGML_OP_MUL) { - return true; - } - - ggml_tensor * mul = cgraph->nodes[i + lane.n_nodes]; - const bool mul_lhs_out = mul->src[0] == lane.out; - const bool mul_rhs_out = mul->src[1] == lane.out; - if (!mul_lhs_out && !mul_rhs_out) { - return true; - } - - const ggml_tensor * scale = mul_lhs_out ? mul->src[1] : mul->src[0]; - if (scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1) { - return false; - } - - if (mul->type != GGML_TYPE_F32 || !ggml_are_same_shape(mul, lane.out)) { - return false; - } - - lane.scale = scale; - lane.out = mul; - lane.n_nodes++; return true; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 9c07ada06d..c2d06cbf48 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -659,24 +659,24 @@ static __global__ void mul_mat_vec_q( if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { float result = tmp[j][threadIdx.x]; if constexpr (has_fusion) { - if (use_bias) { - result += x_biases[j]; - } if constexpr (type == GGML_TYPE_NVFP4) { if (use_scale) { result *= x_scales; } } + if (use_bias) { + result += x_biases[j]; + } if (use_gate) { float gate_value = tmp_gate[j][threadIdx.x]; - if (use_gate_bias) { - gate_value += gate_biases[j]; - } if constexpr (type == GGML_TYPE_NVFP4) { if (use_gate_scale) { gate_value *= gate_scales; } } + if (use_gate_bias) { + gate_value += gate_biases[j]; + } switch (active_glu) { case GGML_GLU_OP_SWIGLU: result *= ggml_cuda_op_silu_single(gate_value); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 445575b62a..c770822776 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5853,27 +5853,27 @@ struct test_mul_mat_vec_fusion : public test_case { auto build_up_lane = [&]() { ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); + if (with_lane_scale) { + ffn_up = build_dense_lane_scale(ctx, ffn_up); + } 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); } - if (with_lane_scale) { - ffn_up = build_dense_lane_scale(ctx, ffn_up); - } return ffn_up; }; auto build_gate_lane = [&]() { ggml_tensor * ffn_gate = ggml_mul_mat(ctx, gate, cur); + if (with_lane_scale) { + ffn_gate = build_dense_lane_scale(ctx, ffn_gate); + } if (with_bias) { 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()); 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; }; @@ -5909,25 +5909,25 @@ struct test_mul_mat_vec_fusion : public test_case { auto build_up_lane = [&]() { ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, ids); + if (with_lane_scale) { + ffn_up = build_id_lane_scale(ctx, ffn_up, 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; }; auto build_gate_lane = [&]() { ggml_tensor * ffn_gate = ggml_mul_mat_id(ctx, gates, cur, ids); + if (with_lane_scale) { + ffn_gate = build_id_lane_scale(ctx, ffn_gate, 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; };