Commit 9ac8c408a for llama.cpp
commit 9ac8c408a33b04396880deba4a9f44470ff12156
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Mon Sep 7 01:08:28 2026 -0500
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.
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<std::array<int, 3>> topk_qsa_edges {
{ 5, 1, 4 }, // add->src[1] == reshape
{ 6, 0, 5 }, // top_k->src[0] == add
};
+static constexpr std::initializer_list<ggml_op> rms_norm_mul_add_mul_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_MUL };
+static constexpr std::initializer_list<ggml_op> rms_norm_mul_add_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD };
+static constexpr std::initializer_list<ggml_op> 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<ggml_op> rms_norm_view_set_rows_pattern { GGML_OP_RMS_NORM, GGML_OP_VIEW, GGML_OP_SET_ROWS };
+static constexpr std::initializer_list<ggml_op> 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<std::array<int, 3>> 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<std::array<int, 3>> rms_norm_mul_rope_vie
{ 4, 0, 3 }, // set_rows->src[0] == view
};
+static constexpr std::initializer_list<std::array<int, 3>> 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<ggml_type, 9> 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<vk_perf_logger> 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<vk_op_binary_push_constants>(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<int64_t, 4> 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<int64_t, 4> 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) {
- a = ggml_add(ctx, ggml_add(ctx, a, b), c);
+ 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;
}
- ggml_tensor * 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);
+ }
- a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), w);
+ a = ggml_rms_norm(ctx, a, eps);
- ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[2]);
+ if (mul) {
+ a = ggml_mul(ctx, a, w);
+ }
- ggml_tensor * rope = ggml_rope(ctx, a, pos, ne[0], mode);
+ 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));
- ggml_tensor * out;
+ 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, rope, ne[0] * ne[1], ne[2], rope->nb[2], 0);
+ ggml_tensor * view = ggml_view_2d(ctx, a, ne[0] * ne[1], ne[2], a->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_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_3d(ctx, GGML_TYPE_I64, ne[2], 1, 1);
+ ggml_tensor * row_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, ne[2]);
ggml_set_name(row_idxs, "row_idxs");
- out = ggml_set_rows(ctx, dst, view, row_idxs);
- ggml_set_name(out, "out");
- } else {
- out = rope;
+ a = ggml_set_rows(ctx, dst, view, row_idxs);
}
- return out;
+ 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<int32_t> 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<int64_t, 4> 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<int64_t, 4> 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<int64_t, 4> 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<std::unique_ptr<test_case>> 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<std::unique_ptr<test_case>> 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<std::unique_ptr<test_case>> 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<std::unique_ptr<test_case>> 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));