From 9ac8c408a33b04396880deba4a9f44470ff12156 Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Mon, 7 Sep 2026 01:08:28 -0500 Subject: [PATCH] vulkan: rms_norm fusion opportunities (#28024) Support RMS_NORM + MUL + ADD (+ MUL) and RMS_NORM + VIEW + SET_ROWS. Extend ROPE + VIEW + SET_ROWS to support IMROPE. Worth around 4% in gemma4 on my system. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 350 +++++++++++++++--- .../ggml-vulkan/vulkan-shaders/rms_norm.comp | 32 +- .../vulkan-shaders/rms_norm_partials.comp | 26 +- .../vulkan-shaders/vulkan-shaders-gen.cpp | 4 + tests/test-backend-ops.cpp | 148 +++++--- 5 files changed, 454 insertions(+), 106 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index efadd3663..9b47c6c95 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -671,6 +671,11 @@ static constexpr std::initializer_list> topk_qsa_edges { { 5, 1, 4 }, // add->src[1] == reshape { 6, 0, 5 }, // top_k->src[0] == add }; +static constexpr std::initializer_list rms_norm_mul_add_mul_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_MUL }; +static constexpr std::initializer_list rms_norm_mul_add_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }; +static constexpr std::initializer_list rms_norm_mul_rope_view_set_rows_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; +static constexpr std::initializer_list rms_norm_view_set_rows_pattern { GGML_OP_RMS_NORM, GGML_OP_VIEW, GGML_OP_SET_ROWS }; +static constexpr std::initializer_list rope_view_set_rows_pattern { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; //node #978 ( SOFT_MAX): ffn_moe_probs-15 ( 0K) [Vulka ] use=2: ffn_moe_logits-15 ( 0K) [Vulka ] //node #979 ( RESHAPE): ffn_moe_probs-15 (re ( 0K) [Vulka ] use=1: ffn_moe_probs-15 ( 0K) [Vulka ] @@ -770,6 +775,16 @@ enum topk_moe_mode { TOPK_MOE_COUNT, }; +enum rms_norm_mode { + RMS_NORM_MUL, + RMS_NORM_MUL_ADD, + RMS_NORM_MUL_ADD_MUL, + RMS_NORM_MUL_ROPE, + RMS_NORM_MUL_ROPE_VIEW_SET_ROWS, + RMS_NORM_VIEW_SET_ROWS, + RMS_NORM_COUNT, +}; + static constexpr std::initializer_list> rope_view_set_rows_edges { { 1, 0, 0 }, // view->src[0] == rope { 2, 0, 1 }, // set_rows->src[0] == view @@ -782,6 +797,11 @@ static constexpr std::initializer_list> rms_norm_mul_rope_vie { 4, 0, 3 }, // set_rows->src[0] == view }; +static constexpr std::initializer_list> rms_norm_view_set_rows_edges { + { 1, 0, 0 }, // view->src[0] == rms_norm + { 2, 0, 1 }, // set_rows->src[0] == view +}; + static constexpr std::array lightning_indexer_k_types = { GGML_TYPE_F32, GGML_TYPE_F16, @@ -1002,6 +1022,12 @@ struct vk_device_struct { vk_pipeline pipeline_group_norm_f32; vk_pipeline pipeline_rms_norm_f32; vk_pipeline pipeline_rms_norm_mul_f32; + vk_pipeline pipeline_rms_norm_mul_add_f32; + vk_pipeline pipeline_rms_norm_mul_add_mul_f32; + vk_pipeline pipeline_rms_norm_mul_add_partials_f32; + vk_pipeline pipeline_rms_norm_mul_add_mul_partials_f32; + vk_pipeline pipeline_rms_norm_set_rows_f32_f32; + vk_pipeline pipeline_rms_norm_set_rows_f32_f16; vk_pipeline pipeline_rms_norm_partials_f32; vk_pipeline pipeline_rms_norm_mul_partials_f32; vk_pipeline pipeline_rms_norm_mul_rope_f32_f32; @@ -2467,6 +2493,7 @@ struct ggml_backend_vk_context { bool fused_topk_moe_scale {}; // QSA indexer gather+add+top_k fused into one radix-select bool fused_topk_qsa {}; + rms_norm_mode fused_rms_norm_mode {RMS_NORM_COUNT}; // for GGML_VK_PERF_LOGGER std::unique_ptr perf_logger; @@ -5609,6 +5636,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_f32, "rms_norm_mul_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_f32, "rms_norm_mul_add_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_f32, "rms_norm_mul_add_mul_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_partials_f32, "rms_norm_mul_add_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_partials_f32, "rms_norm_mul_add_mul_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f32, "rms_norm_set_rows_f32_f32", rms_norm_set_rows_f32_f32_len, rms_norm_set_rows_f32_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f16, "rms_norm_set_rows_f32_f16", rms_norm_set_rows_f32_f16_len, rms_norm_set_rows_f32_f16_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_partials_f32, "rms_norm_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_partials_f32, "rms_norm_mul_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); @@ -11551,10 +11584,9 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const case GGML_OP_RMS_NORM: if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { if (ctx->do_add_rms_partials) { - return ctx->num_additional_fused_ops > 0 ? ctx->device->pipeline_rms_norm_mul_partials_f32 : ctx->device->pipeline_rms_norm_partials_f32; - } else { - return ctx->num_additional_fused_ops > 0 ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32; + return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_partials_f32 : ctx->device->pipeline_rms_norm_partials_f32; } + return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32; } return nullptr; case GGML_OP_RMS_NORM_BACK: @@ -13500,40 +13532,121 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor * return rope; } -static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, float * op_params) { - ggml_tensor * dst; - const ggml_tensor * src0; - const ggml_tensor * src1; - - if (ctx->num_additional_fused_ops > 0) { - // fused rms_norm + mul - ggml_tensor *mul = cgraph->nodes[node_idx + 1]; - ggml_tensor *other_src = mul->src[0] == cgraph->nodes[node_idx + 0] ? mul->src[1] : mul->src[0]; - dst = mul; - src0 = cgraph->nodes[node_idx]->src[0]; - src1 = other_src; - } else { - dst = cgraph->nodes[node_idx]; - src0 = src1 = dst->src[0]; - } - +static vk_op_binary_push_constants ggml_vk_rms_norm_push_constants( + const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, + float eps, uint32_t num_partials) { const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); const uint32_t dst_type_size = ggml_type_size(dst->type); - uint32_t param3 = ctx->do_add_rms_partials ? ggml_vk_rms_num_partials(ctx, dst) : 0; - - vk_op_binary_push_constants bin { + return { (uint32_t)ggml_nelements(src0), (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, 0, - op_params[0], 0.0f, (int32_t)param3, + eps, 0.0f, (int32_t)num_partials, }; +} - // more than one fused op means rms_norm+mul+rope - if (ctx->num_additional_fused_ops > 1) { +static void ggml_vk_rms_norm_finish(ggml_backend_vk_context * ctx, const ggml_tensor * src0) { + if (ctx->do_add_rms_partials_offset_calculation) { + ctx->prealloc_size_add_rms_partials_offset += ggml_vk_rms_partials_size(ctx, src0); + ctx->do_add_rms_partials = false; + ctx->do_add_rms_partials_offset_calculation = false; + } +} + +static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, float * op_params) { + ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * src0 = rms->src[0]; + + if (ctx->fused_rms_norm_mode == RMS_NORM_VIEW_SET_ROWS) { + GGML_ASSERT(ctx->num_additional_fused_ops == 2); + ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; + const ggml_tensor * indices = set_rows->src[1]; + vk_op_binary_push_constants pc = ggml_vk_rms_norm_push_constants(src0, src0, set_rows, op_params[0], 0); + init_pushconst_tensor_offsets(ctx, pc, src0, src0, nullptr, nullptr, set_rows); + + vk_pipeline pipeline = set_rows->type == GGML_TYPE_F16 ? + ctx->device->pipeline_rms_norm_set_rows_f32_f16 : ctx->device->pipeline_rms_norm_set_rows_f32_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, set_rows, true), + ggml_vk_tensor_subbuffer(ctx, indices), + }, pc, { (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3] }); + ggml_vk_rms_norm_finish(ctx, src0); + return; + } + + if (ctx->fused_rms_norm_mode == RMS_NORM_MUL_ADD || ctx->fused_rms_norm_mode == RMS_NORM_MUL_ADD_MUL) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + ggml_tensor * add = cgraph->nodes[node_idx + 2]; + const ggml_tensor * weight = mul->src[0] == rms ? mul->src[1] : mul->src[0]; + const ggml_tensor * residual = add->src[0] == mul ? add->src[1] : add->src[0]; + const bool do_post_multiply = ctx->fused_rms_norm_mode == RMS_NORM_MUL_ADD_MUL; + GGML_ASSERT(ctx->num_additional_fused_ops == (do_post_multiply ? 3 : 2)); + ggml_tensor * dst = do_post_multiply ? cgraph->nodes[node_idx + 3] : add; + const ggml_tensor * post_scale = do_post_multiply ? + (dst->src[0] == add ? dst->src[1] : dst->src[0]) : src0; + + const uint32_t num_partials = ctx->do_add_rms_partials ? ggml_vk_rms_num_partials(ctx, dst) : 0; + vk_op_binary_push_constants pc = ggml_vk_rms_norm_push_constants(src0, weight, dst, op_params[0], num_partials); + init_pushconst_tensor_offsets(ctx, pc, src0, weight, residual, post_scale, dst); + + vk_pipeline pipeline; + if (ctx->do_add_rms_partials) { + pipeline = do_post_multiply ? + ctx->device->pipeline_rms_norm_mul_add_mul_partials_f32 : ctx->device->pipeline_rms_norm_mul_add_partials_f32; + } else { + pipeline = do_post_multiply ? + ctx->device->pipeline_rms_norm_mul_add_mul_f32 : ctx->device->pipeline_rms_norm_mul_add_f32; + } + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + if (ctx->do_add_rms_partials) { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, weight, true), + ggml_vk_tensor_subbuffer(ctx, dst, true), + ggml_vk_subbuffer(ctx, ctx->prealloc_add_rms_partials, ctx->prealloc_size_add_rms_partials_offset), + ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, post_scale), + }, pc, { (uint32_t)CEIL_DIV(src0->ne[0], 128), 1, 1 }); + } else { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, weight, true), + ggml_vk_tensor_subbuffer(ctx, dst, true), + ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, post_scale), + }, pc, { (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3] }); + } + ggml_vk_rms_norm_finish(ctx, src0); + return; + } + + ggml_tensor * dst; + const ggml_tensor * src1; + + if (ctx->fused_rms_norm_mode != RMS_NORM_COUNT) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + dst = mul; + src1 = mul->src[0] == rms ? mul->src[1] : mul->src[0]; + } else { + dst = rms; + src1 = src0; + } + + const uint32_t num_partials = ctx->do_add_rms_partials ? ggml_vk_rms_num_partials(ctx, dst) : 0; + vk_op_binary_push_constants bin = ggml_vk_rms_norm_push_constants(src0, src1, dst, op_params[0], num_partials); + + if (ctx->fused_rms_norm_mode == RMS_NORM_MUL_ROPE || + ctx->fused_rms_norm_mode == RMS_NORM_MUL_ROPE_VIEW_SET_ROWS) { static constexpr uint32_t max_tensors = 7; const ggml_tensor *tensors[max_tensors] {}; @@ -13543,7 +13656,8 @@ static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor *other_src = mul->src[0] == rms ? mul->src[1] : mul->src[0]; - bool do_set_rows = ctx->num_additional_fused_ops == 4; + bool do_set_rows = ctx->fused_rms_norm_mode == RMS_NORM_MUL_ROPE_VIEW_SET_ROWS; + GGML_ASSERT(ctx->num_additional_fused_ops == (do_set_rows ? 4 : 2)); tensors[0] = rms->src[0]; tensors[1] = other_src; @@ -13610,14 +13724,11 @@ static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_vk_subbuffer(ctx, buf[6], offset[6]), }, pc, elements); } else { + GGML_ASSERT(ctx->fused_rms_norm_mode == RMS_NORM_MUL || ctx->fused_rms_norm_mode == RMS_NORM_COUNT); ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_RMS_NORM, std::move(bin)); } - if (ctx->do_add_rms_partials_offset_calculation) { - ctx->prealloc_size_add_rms_partials_offset += ggml_vk_rms_partials_size(ctx, src0); - ctx->do_add_rms_partials = false; - ctx->do_add_rms_partials_offset_calculation = false; - } + ggml_vk_rms_norm_finish(ctx, src0); } static void ggml_vk_rms_norm_back(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { @@ -16938,7 +17049,8 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g return false; } - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + if ((ops.size() == 2 || ops.size() == 3 || ops.size() == 4) && + ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { // additional constraints specific to this fusion const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; @@ -16960,6 +17072,43 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { return false; } + + if (ops.size() >= 3 && ops.begin()[2] == GGML_OP_ADD) { + const ggml_tensor *add = cgraph->nodes[node_idx + 2]; + const ggml_tensor *residual = add->src[0] == mul ? add->src[1] : add->src[0]; + if (add->src[0] != mul && add->src[1] != mul) { + return false; + } + if (residual->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || + !ggml_are_same_shape(add, residual) || !ggml_is_contiguous(residual) || + !ggml_is_contiguous(add) || get_misalign_bytes(ctx, residual) != 0) { + return false; + } + + const ggml_tensor *dst = add; + if (ops.size() == 4) { + if (ops.begin()[3] != GGML_OP_MUL) { + return false; + } + + const ggml_tensor *post_mul = cgraph->nodes[node_idx + 3]; + const ggml_tensor *scale = post_mul->src[0] == add ? post_mul->src[1] : post_mul->src[0]; + if (post_mul->src[0] != add && post_mul->src[1] != add) { + return false; + } + // The shader reads data_e[0], so the final multiply must use a scalar. + if (scale->type != GGML_TYPE_F32 || post_mul->type != GGML_TYPE_F32 || + ggml_nelements(scale) != 1 || !ggml_is_contiguous(post_mul) || + get_misalign_bytes(ctx, scale) != 0) { + return false; + } + dst = post_mul; + } + + if (get_misalign_bytes(ctx, dst) != 0) { + return false; + } + } } auto const &mm_add_ok = [&](const ggml_tensor *mul, const ggml_tensor *add) { const ggml_tensor *bias = add->src[0] == mul ? add->src[1] : add->src[0]; @@ -17341,12 +17490,11 @@ static bool ggml_vk_can_fuse_topk_qsa(ggml_backend_vk_context * ctx, const struc static bool ggml_vk_can_fuse_rope_set_rows(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { - GGML_UNUSED(ctx); const ggml_tensor *rope = cgraph->nodes[node_idx + 0]; const ggml_tensor *view = cgraph->nodes[node_idx + 1]; const ggml_tensor *set_rows = cgraph->nodes[node_idx + 2]; - // ne3 not tested + // The set_rows epilogue uses one index per ne2 slice and does not encode ne3. if (rope->src[0]->ne[3] != 1) { return false; } @@ -17355,19 +17503,50 @@ static bool ggml_vk_can_fuse_rope_set_rows(ggml_backend_vk_context * ctx, const return false; } - if (set_rows->src[1]->type != GGML_TYPE_I64) { + // The shader reads each aligned I64 index as a uvec2 and uses its low 32 bits. + if (set_rows->src[1]->type != GGML_TYPE_I64 || !ggml_is_contiguous(set_rows->src[1]) || + set_rows->nb[0] != ggml_type_size(set_rows->type) || get_misalign_bytes(ctx, set_rows->src[1]) != 0) { return false; } - // The view should flatten two dims of rope into one dim + // SET_ROWS consumes one flattened [ne0*ne1] row for each ne2 slice. if (!ggml_is_contiguous(view) || - view->ne[0] != rope->ne[0] * rope->ne[1]) { + view->ne[0] != rope->ne[0] * rope->ne[1] || view->ne[1] != rope->ne[2] || + view->ne[2] != 1 || view->ne[3] != 1 || + ggml_nelements(set_rows->src[1]) != rope->ne[2]) { return false; } - // Only norm/neox/mrope shaders have the fusion code + // Only norm/neox/mrope/imrope shaders have the fusion code const int mode = ((const int32_t *) rope->op_params)[2]; - if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_MROPE) { + if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && + mode != GGML_ROPE_TYPE_MROPE && mode != GGML_ROPE_TYPE_IMROPE) { + return false; + } + + return true; +} + +static bool ggml_vk_can_fuse_rms_norm_set_rows(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, + int node_idx) { + const ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * view = cgraph->nodes[node_idx + 1]; + const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; + + // The RMS kernel reads F32 and writes directly to the F32 or F16 SET_ROWS destination. + if (rms->src[0]->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || + (set_rows->type != GGML_TYPE_F32 && set_rows->type != GGML_TYPE_F16) || + set_rows->src[1]->type != GGML_TYPE_I64 || !ggml_is_contiguous(set_rows->src[1]) || + set_rows->nb[0] != ggml_type_size(set_rows->type) || get_misalign_bytes(ctx, set_rows->src[1]) != 0) { + return false; + } + // As with the ROPE epilogue, each ne2 slice supplies one flattened row and ne3 is not encoded. + if (rms->ne[3] != 1 || !ggml_is_contiguous(rms->src[0]) || !ggml_is_contiguous(view)) { + return false; + } + if (view->ne[0] != rms->ne[0] * rms->ne[1] || view->ne[1] != rms->ne[2] || + view->ne[2] != 1 || view->ne[3] != 1 || + ggml_nelements(set_rows->src[1]) != rms->ne[2]) { return false; } @@ -17462,7 +17641,6 @@ static bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b static bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { - GGML_UNUSED(ctx); const ggml_tensor *rms = cgraph->nodes[node_idx + 0]; const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; const ggml_tensor *rope = cgraph->nodes[node_idx + 2]; @@ -17719,6 +17897,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_topk_qsa = false; + ctx->fused_rms_norm_mode = RMS_NORM_COUNT; const char *fusion_string {}; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); @@ -17753,27 +17932,47 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; - } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 4 }) && + } else if (ggml_can_fuse_subgraph(cgraph, i, rms_norm_mul_rope_view_set_rows_pattern, { i + 4 }) && ggml_check_edges(cgraph, i, rms_norm_mul_rope_view_set_rows_edges) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i) && ggml_vk_can_fuse_rope_set_rows(ctx, cgraph, i + 2)) { ctx->num_additional_fused_ops = 4; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ROPE_VIEW_SET_ROWS; fusion_string = "RMS_NORM_MUL_ROPE_VIEW_SET_ROWS"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = false; op_srcs_fused_elementwise[2] = false; op_srcs_fused_elementwise[3] = false; op_srcs_fused_elementwise[4] = false; - } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE })&& + } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i)) { ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ROPE; fusion_string = "RMS_NORM_MUL_ROPE"; // rope is approximately elementwise - whole rows are done by a single workgroup and it's row-wise op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; op_srcs_fused_elementwise[2] = true; + } else if (ggml_vk_can_fuse(ctx, cgraph, i, rms_norm_mul_add_mul_pattern)) { + ctx->num_additional_fused_ops = 3; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ADD_MUL; + fusion_string = "RMS_NORM_MUL_ADD_MUL"; + std::fill_n(op_srcs_fused_elementwise, 4, true); + } else if (ggml_vk_can_fuse(ctx, cgraph, i, rms_norm_mul_add_pattern)) { + ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ADD; + fusion_string = "RMS_NORM_MUL_ADD"; + std::fill_n(op_srcs_fused_elementwise, 3, true); + } else if (ggml_can_fuse_subgraph(cgraph, i, rms_norm_view_set_rows_pattern, { i + 2 }) && + ggml_check_edges(cgraph, i, rms_norm_view_set_rows_edges) && + ggml_vk_can_fuse_rms_norm_set_rows(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_VIEW_SET_ROWS; + fusion_string = "RMS_NORM_VIEW_SET_ROWS"; + std::fill_n(op_srcs_fused_elementwise, 3, false); } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { ctx->num_additional_fused_ops = 1; + ctx->fused_rms_norm_mode = RMS_NORM_MUL; fusion_string = "RMS_NORM_MUL"; // rms_norm is not elementwise, but whole rows must be consumed and the scale factor computed before // they are overwritten, and one workgroup per row. So close enough. @@ -17792,7 +17991,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "SSM_CONV_SILU"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; - } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 2 }) && + } else if (ggml_can_fuse_subgraph(cgraph, i, rope_view_set_rows_pattern, { i + 2 }) && ggml_check_edges(cgraph, i, rope_view_set_rows_edges) && ggml_vk_can_fuse_rope_set_rows(ctx, cgraph, i)) { ctx->num_additional_fused_ops = 2; @@ -17930,6 +18129,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_topk_qsa = false; + ctx->fused_rms_norm_mode = RMS_NORM_COUNT; } } @@ -18130,6 +18330,22 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * continue; } + if (keep_pattern(rms_norm_mul_add_mul_pattern)) { + continue; + } + if (keep_pattern(rms_norm_mul_add_pattern)) { + continue; + } + if (keep_pattern(rms_norm_mul_rope_view_set_rows_pattern)) { + continue; + } + if (keep_pattern(rms_norm_view_set_rows_pattern)) { + continue; + } + if (keep_pattern(rope_view_set_rows_pattern)) { + continue; + } + // First, grab the next unused node. current_set.push_back(first_unused); @@ -18163,7 +18379,12 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * match_pattern(topk_moe_early_softmax, j) || match_pattern(topk_moe_late_softmax, j) || match_pattern(snake_pattern, j) || - in_qsa_pattern(j)) { + in_qsa_pattern(j) || + match_pattern(rms_norm_mul_add_mul_pattern, j) || + match_pattern(rms_norm_mul_add_pattern, j) || + match_pattern(rms_norm_mul_rope_view_set_rows_pattern, j) || + match_pattern(rms_norm_view_set_rows_pattern, j) || + match_pattern(rope_view_set_rows_pattern, j)) { continue; } bool ok = true; @@ -18203,30 +18424,41 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * } } } - // Look for ROPE + VIEW + SET_ROWS and make them consecutive - if (graph->nodes[rope_idx]->op == GGML_OP_ROPE) { + // Look for ROPE/RMS_NORM + VIEW + SET_ROWS and make them consecutive + if (graph->nodes[rope_idx]->op == GGML_OP_ROPE || graph->nodes[rope_idx]->op == GGML_OP_RMS_NORM) { int view_idx = -1; int set_rows_idx = -1; - for (int k = rope_idx+1; k < std::min(rope_idx + 10, graph->n_nodes); ++k) { - if (view_idx == -1 && - graph->nodes[k]->op == GGML_OP_VIEW && - graph->nodes[k]->src[0] == graph->nodes[rope_idx]) { + for (int k = rope_idx + 1; k < std::min(rope_idx + 15, graph->n_nodes); ++k) { + if (used[k]) { + continue; + } + if (view_idx == -1 && graph->nodes[k]->op == GGML_OP_VIEW && graph->nodes[k]->src[0] == graph->nodes[rope_idx]) { view_idx = k; continue; } - if (view_idx != -1 && - set_rows_idx == -1 && - graph->nodes[k]->op == GGML_OP_SET_ROWS && - graph->nodes[k]->src[0] == graph->nodes[view_idx]) { + if (view_idx != -1 && graph->nodes[k]->op == GGML_OP_SET_ROWS && graph->nodes[k]->src[0] == graph->nodes[view_idx]) { set_rows_idx = k; break; } } if (set_rows_idx != -1) { - current_set.push_back(view_idx); - current_set.push_back(set_rows_idx); - used[view_idx] = true; - used[set_rows_idx] = true; + const int node_idxs[] = { rope_idx, view_idx, set_rows_idx }; + const ggml_op ops[] = { graph->nodes[rope_idx]->op, GGML_OP_VIEW, GGML_OP_SET_ROWS }; + bool can_pull = ggml_can_fuse_subgraph_ext(graph, node_idxs, 3, ops, &set_rows_idx, 1); + + for (int c = rope_idx + 1; can_pull && c < set_rows_idx; ++c) { + if (!used[c] && c != view_idx && !is_empty(graph->nodes[c]) && + is_src_of(graph->nodes[set_rows_idx], graph->nodes[c])) { + can_pull = false; + } + } + + if (can_pull) { + current_set.push_back(view_idx); + current_set.push_back(set_rows_idx); + used[view_idx] = true; + used[set_rows_idx] = true; + } } } // Look for MUL_MAT_ID + ADD_ID + MUL diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp index 55b89f19a..ee813842c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp @@ -27,12 +27,24 @@ layout (binding = 6) readonly buffer R_I {uvec2 rope_data_i[];}; // indices for #define GGML_ROPE_TYPE_MROPE 8 #define GGML_ROPE_TYPE_VISION 24 +#elif RMS_NORM_ADD_FUSION + +layout (binding = 3) readonly buffer C {float data_c[];}; +layout (binding = 4) readonly buffer E {float data_e[];}; + +#elif RMS_NORM_SET_ROWS_FUSION + +layout (binding = 3) readonly buffer I {uvec2 data_i[];}; + #endif #extension GL_EXT_control_flow_attributes : enable #define BLOCK_SIZE 512 layout (constant_id = 1) const bool do_multiply = false; +#if RMS_NORM_ADD_FUSION +layout (constant_id = 2) const bool do_post_multiply = false; +#endif layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; @@ -57,6 +69,8 @@ void rms_norm(uint num_iters) { #if RMS_NORM_ROPE_FUSION // Per-row offset in shared memory uint32_t d_offset = 0; +#elif RMS_NORM_SET_ROWS_FUSION + uint32_t d_offset = data_i[channel].x*p.nb21 + row*ncols + get_doffset(); #else uint32_t d_offset = ((samp*nchannels + channel)*nrows + row)*ncols + get_doffset(); #endif @@ -91,14 +105,28 @@ void rms_norm(uint num_iters) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } } else { [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } } } else { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp index 4618b2c7e..cf7ab21f2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp @@ -10,11 +10,19 @@ #define BLOCK_SIZE 128 layout (constant_id = 1) const bool do_multiply = false; +#if RMS_NORM_ADD_FUSION +layout (constant_id = 2) const bool do_post_multiply = false; +#endif layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; layout (binding = 3, std430) readonly buffer PartialsBuf {float partial_sums[];}; +#if RMS_NORM_ADD_FUSION +layout (binding = 4) readonly buffer C {float data_c[];}; +layout (binding = 5) readonly buffer E {float data_e[];}; +#endif + shared FLOAT_TYPE sumsh[BLOCK_SIZE]; void main() { @@ -55,9 +63,23 @@ void main() { if (do_multiply) { if (ncols > p.ne10) { - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } else { - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } } else { data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 5d4b9c5fa..da0d54ab4 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -806,6 +806,10 @@ void process_shaders() { string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rms_norm_mul_add_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); + string_to_spv("rms_norm_mul_add_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); + string_to_spv("rms_norm_set_rows_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}})); + string_to_spv("rms_norm_set_rows_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}})); string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_mul_rope_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}})); string_to_spv("rms_norm_mul_rope_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 31593541f..6a8cfd46d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2668,13 +2668,16 @@ struct test_rope_set_rows : public test_case { } }; -// GGML_OP_RMS_NORM + GGML_OP_MUL + GGML_OP_ROPE (+ GGML_OP_VIEW + GGML_OP_SET_ROWS) +// GGML_OP_RMS_NORM with optional GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW and GGML_OP_SET_ROWS struct test_rms_norm_mul_rope : public test_case { const std::array ne; const float eps; const bool multi_add; // test a sequence of adds feeding into rms_norm + const bool mul; + const bool rope; const bool set_rows; const bool broadcast; // multiply by a 1D [ne0] weight, as model norm weights are + const ggml_type set_rows_type; int mode; std::string op_desc(ggml_tensor * t) override { @@ -2685,63 +2688,90 @@ struct test_rms_norm_mul_rope : public test_case { bool run_whole_graph() override { return true; } std::string vars() override { - return VARS_TO_STR6(ne, eps, multi_add, set_rows, broadcast, mode); + return VARS_TO_STR9(ne, eps, multi_add, mul, rope, set_rows, broadcast, mode, set_rows_type); } test_rms_norm_mul_rope(std::array ne, float eps = 1e-6f, bool multi_add = false, - bool set_rows = false, bool broadcast = false, int mode = GGML_ROPE_TYPE_NORMAL) - : ne(ne), eps(eps), multi_add(multi_add), set_rows(set_rows), broadcast(broadcast), mode(mode) {} + bool set_rows = false, bool broadcast = false, int mode = GGML_ROPE_TYPE_NORMAL, + bool mul = true, bool rope = true, ggml_type set_rows_type = GGML_TYPE_F16) + : ne(ne), eps(eps), multi_add(multi_add), mul(mul), rope(rope), set_rows(set_rows), broadcast(broadcast), + set_rows_type(set_rows_type), mode(mode) {} ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); - ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); - ggml_tensor * c = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); + ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], ne[3]); + + ggml_tensor * b = nullptr; + ggml_tensor * c = nullptr; + ggml_tensor * w = nullptr; + + if (multi_add || (mul && !broadcast)) { + b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); + } + if (multi_add) { + c = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); + } + if (mul) { + w = broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]) : b; + } if (multi_add) { a = ggml_add(ctx, ggml_add(ctx, a, b), c); } - ggml_tensor * w = broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]) : b; + a = ggml_rms_norm(ctx, a, eps); - a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), w); - - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[2]); - - ggml_tensor * rope = ggml_rope(ctx, a, pos, ne[0], mode); - - ggml_tensor * out; - - if (set_rows) { - ggml_tensor * view = ggml_view_2d(ctx, rope, ne[0] * ne[1], ne[2], rope->nb[2], 0); - - ggml_tensor * dst = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, ne[0] * ne[1], ne[2] * ne[3], 1, 1); - ggml_set_name(dst, "dst"); - - ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, GGML_TYPE_I64, ne[2], 1, 1); - ggml_set_name(row_idxs, "row_idxs"); - - out = ggml_set_rows(ctx, dst, view, row_idxs); - ggml_set_name(out, "out"); - } else { - out = rope; + if (mul) { + a = ggml_mul(ctx, a, w); } - return out; + if (rope) { + const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[2] * (is_mrope ? 4 : 1)); + + if (is_mrope) { + const int n_dims = ne[0]; + int sections[4] = { n_dims/3, n_dims/3, n_dims/3, 0 }; + a = ggml_rope_multi(ctx, a, pos, nullptr, n_dims, sections, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); + } else { + a = ggml_rope(ctx, a, pos, ne[0], mode); + } + } + + if (set_rows) { + ggml_tensor * view = ggml_view_2d(ctx, a, ne[0] * ne[1], ne[2], a->nb[2], 0); + + ggml_tensor * dst = ggml_new_tensor_2d(ctx, set_rows_type, ne[0] * ne[1], ne[2] * 2); + ggml_set_name(dst, "dst"); + + ggml_tensor * row_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, ne[2]); + ggml_set_name(row_idxs, "row_idxs"); + + a = ggml_set_rows(ctx, dst, view, row_idxs); + } + + ggml_set_name(a, "out"); + return a; } void initialize_tensors(ggml_context * ctx) override { for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I64 || t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { - continue; + if (t->type == GGML_TYPE_I64) { + init_set_rows_row_ids(t, ne[2] * 2); + } else if (t->type == GGML_TYPE_I32) { + std::vector data(ggml_nelements(t)); + for (int32_t & value : data) { + value = rand() % 512; } - - init_set_rows_row_ids(t, ne[2]); + ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); } else { init_tensor_uniform(t); } } } + + double max_nmse_err() override { + return ne[0] == 8192 ? 5e-6 : test_case::max_nmse_err(); + } }; // GGML_OP_ARGMAX @@ -3636,13 +3666,16 @@ struct test_rms_norm_back : public test_case { } }; -// GGML_OP_RMS_NORM + GGML_OP_MUL + GGML_OP_ADD +// GGML_OP_RMS_NORM + GGML_OP_MUL + GGML_OP_ADD (+ GGML_OP_MUL) struct test_rms_norm_mul_add : public test_case { const ggml_type type; const std::array ne; const float eps; const bool broadcast; const bool multi_add; // test a sequence of adds feeding into rms_norm + const bool post_mul; + const bool alias_rms_input; + const bool weight_broadcast; std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); @@ -3652,20 +3685,23 @@ struct test_rms_norm_mul_add : public test_case { bool run_whole_graph() override { return true; } std::string vars() override { - return VARS_TO_STR5(type, ne, eps, broadcast, multi_add); + return VARS_TO_STR8(type, ne, eps, broadcast, multi_add, post_mul, alias_rms_input, weight_broadcast); } test_rms_norm_mul_add(ggml_type type = GGML_TYPE_F32, std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f, bool broadcast = false, bool multi_add = false) - : type(type), ne(ne), eps(eps), broadcast(broadcast), multi_add(multi_add) {} + float eps = 1e-6f, bool broadcast = false, bool multi_add = false, bool post_mul = false, + bool alias_rms_input = false, bool weight_broadcast = false) + : type(type), ne(ne), eps(eps), broadcast(broadcast), multi_add(multi_add), post_mul(post_mul), + alias_rms_input(alias_rms_input), weight_broadcast(weight_broadcast) {} ggml_tensor * build_graph(ggml_context * ctx) override { std::array broadcast_dims = {ne[0]*2, ne[1]*3, ne[2]*3, ne[3]*4}; ggml_tensor * a = ggml_new_tensor(ctx, type, 4, broadcast ? broadcast_dims.data() : ne.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_tensor * b = weight_broadcast ? ggml_new_tensor_1d(ctx, type, ne[0]) : ggml_new_tensor(ctx, type, 4, ne.data()); ggml_tensor * c = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_tensor * d = nullptr; ggml_set_param(a); ggml_set_name(a, "a"); @@ -3676,10 +3712,20 @@ struct test_rms_norm_mul_add : public test_case { // Use a, b and c early, so we don't end up with an OP_NONE between rms_norm and mul a = ggml_add(ctx, ggml_add(ctx, a, b), c); + if (post_mul) { + d = ggml_new_tensor_1d(ctx, type, 1); + ggml_set_param(d); + ggml_set_name(d, "d"); + a = ggml_add(ctx, a, d); + } if (multi_add) { a = ggml_add(ctx, ggml_add(ctx, a, b), c); } - ggml_tensor * out = ggml_add(ctx, ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b), c); + ggml_tensor * mul = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b); + ggml_tensor * out = alias_rms_input ? ggml_add_inplace(ctx, a, mul) : ggml_add(ctx, mul, c); + if (post_mul) { + out = ggml_mul(ctx, out, d); + } ggml_set_name(out, "out"); return out; @@ -8848,7 +8894,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_set_rows(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, true)); test_cases.emplace_back(new test_set_rows(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, true)); - for (int mode : { GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX, GGML_ROPE_TYPE_MROPE, GGML_ROPE_TYPE_VISION }) { + for (int mode : { GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX, GGML_ROPE_TYPE_MROPE, GGML_ROPE_TYPE_VISION, GGML_ROPE_TYPE_IMROPE }) { for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { for (int ne2 : {1, 8, 512}) { test_cases.emplace_back(new test_rope_set_rows(type, GGML_TYPE_I64, { 128, 32, ne2, 1 }, mode)); @@ -8856,6 +8902,7 @@ static std::vector> make_test_cases_eval() { } } } + test_cases.emplace_back(new test_rope_set_rows(GGML_TYPE_F32, GGML_TYPE_I32, { 128, 32, 8, 1 }, GGML_ROPE_TYPE_IMROPE)); for (ggml_type type_input : {GGML_TYPE_F32}) { for (ggml_op_pool pool_type : {GGML_OP_POOL_AVG, GGML_OP_POOL_MAX}) { @@ -9437,6 +9484,11 @@ static std::vector> make_test_cases_eval() { // in-place tests test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true)); + for (ggml_type set_rows_type : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + test_cases.emplace_back(new test_rms_norm_mul_rope({ 256, 1, 1, 1 }, 1e-6f, false, true, false, GGML_ROPE_TYPE_NORMAL, false, false, set_rows_type)); + test_cases.emplace_back(new test_rms_norm_mul_rope({ 128, 4, 3, 1 }, 1e-6f, false, true, false, GGML_ROPE_TYPE_NORMAL, false, false, set_rows_type)); + } + for (float eps : { 0.0f, 1e-6f, 1e-4f, 1e-1f, 1.0f }) { for (uint32_t n : { 64, 1025 }) { test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false)); @@ -9462,10 +9514,20 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_add_add(GGML_TYPE_F16, GGML_TYPE_F32, { n, 5, 4, 3 }, true, false)); } + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 1536, 1, 1, 1 }, 1e-6f, false, false, true)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 1, 1 }, 1e-6f, false, false, true)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 3, 2 }, 1e-6f, false, false, true)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 3, 2 }, 1e-6f, false, false, true, false, true)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 1536, 1, 1, 1 }, 1e-6f, false, false, false, true)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 1, 1 }, 1e-6f, false, false, false, true)); + + test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 7, 2})); + test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 7, 2}, 1e-6f, false, true)); + for (auto multi_add : {false, true}) { for (auto set_rows : {false, true}) { for (auto broadcast : {false, true}) { - for (auto rope : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX}) { + for (auto rope : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX, GGML_ROPE_TYPE_IMROPE}) { test_cases.emplace_back(new test_rms_norm_mul_rope({768, 1, 1, 1}, 1e-6f, multi_add, set_rows, broadcast, rope)); test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 1, 1}, 1e-6f, multi_add, set_rows, broadcast, rope)); test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 5, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));