Commit ca89968a2 for llama.cpp
commit ca89968a23a13e571a6b802a9891d22763153b61
Author: Ehsan Bateni <ebateni@qti.qualcomm.com>
Date: Fri Nov 6 19:06:59 2026 -0500
hexagon: fuse post-norm residual add into RMS_NORM_MUL (#31822)
* hexagon: fuse post-norm residual add into RMS_NORM_MUL
* hex-fusion: refactor RMS_NORM fusion and add dedicated RMS_NORM_MUL_ADD opcode
Improves performance by keeping things in HVX registers.
Also simplifies overall fusion logic and logging.
---------
Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>
diff --git a/ggml/src/ggml-hexagon/htp-fusion.cpp b/ggml/src/ggml-hexagon/htp-fusion.cpp
index 4de6dd8b7..7a3bdd46b 100644
--- a/ggml/src/ggml-hexagon/htp-fusion.cpp
+++ b/ggml/src/ggml-hexagon/htp-fusion.cpp
@@ -157,6 +157,30 @@ static bool fuse_rms_norm_mul(const ggml_hexagon_session * sess, htp_opnode & pr
return true;
}
+// post-norm residual: dst = rms_norm(x) * w + r
+static bool fuse_rms_norm_mul_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+ if (prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
+
+ const ggml_tensor * r = nullptr;
+ if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &r)) return false;
+
+ const ggml_tensor * src0 = prev.src0();
+ const ggml_tensor * weight = prev.inputs[1];
+ const ggml_tensor * dst = cur.dst();
+ if (!ggml_are_same_layout(src0, r)) return false;
+ if (htp_inputs_overlap_unsafe(prev, dst, r)) return false;
+
+ struct htp_unary_kernel_params kparams;
+ if (!htp_kparams_solve_rms_norm_mul_add(sess, src0, weight, dst, &kparams)) return false;
+
+ prev.opcode = HTP_OP_RMS_NORM_MUL_ADD;
+ prev.inputs.push_back(r);
+ prev.outputs[0] = dst;
+ prev.fused.push_back(cur.node);
+ prev.set_kparams(kparams);
+ return true;
+}
+
static bool fuse_rms_norm_scale(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
GGML_UNUSED(sess);
if (prev.opcode != HTP_OP_RMS_NORM && prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
@@ -620,6 +644,7 @@ bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_op
case HTP_OP_ADD:
return (htp_fusion_enabled(HTP_FUSE_ALLREDUCE_ADD) && fuse_allreduce_add(sess, prev, cur)) ||
(htp_fusion_enabled(HTP_FUSE_MUL_MAT_ADD) && fuse_mul_mat_add(sess, prev, cur)) ||
+ (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL_ADD) && fuse_rms_norm_mul_add(sess, prev, cur)) ||
(htp_fusion_enabled(HTP_FUSE_BINARY_NX) && fuse_binary_nx(sess, prev, cur));
case HTP_OP_MUL:
return (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL) && fuse_rms_norm_mul(sess, prev, cur)) ||
diff --git a/ggml/src/ggml-hexagon/htp-fusion.h b/ggml/src/ggml-hexagon/htp-fusion.h
index e794c1daa..e4f711c87 100644
--- a/ggml/src/ggml-hexagon/htp-fusion.h
+++ b/ggml/src/ggml-hexagon/htp-fusion.h
@@ -16,6 +16,7 @@ enum htp_fusion_flags {
HTP_FUSE_BINARY_NX = (1 << 10), // 1024
HTP_FUSE_RMS_NORM_SCALE = (1 << 11), // 2048
HTP_FUSE_SSM_CONV_SILU = (1 << 12), // 4096
+ HTP_FUSE_RMS_NORM_MUL_ADD = (1 << 13), // 8192
};
static inline bool htp_fusion_enabled(int flag) {
diff --git a/ggml/src/ggml-hexagon/htp-kparams.cpp b/ggml/src/ggml-hexagon/htp-kparams.cpp
index 022532328..49b673e85 100644
--- a/ggml/src/ggml-hexagon/htp-kparams.cpp
+++ b/ggml/src/ggml-hexagon/htp-kparams.cpp
@@ -1583,11 +1583,15 @@ static bool htp_kparams_solve_unary_impl(
kparams->n_threads = n_threads;
- const bool broadcast_weight = (op == HTP_OP_RMS_NORM_MUL) && (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
+ const bool has_weight = (op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_RMS_NORM_MUL_ADD);
+ const bool broadcast_weight = has_weight && (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
+ if (op == HTP_OP_RMS_NORM_MUL_ADD && !broadcast_weight) {
+ return false;
+ }
kparams->broadcast_weight = broadcast_weight;
const bool src0_contig = ggml_is_contiguous_1(src0);
- const bool src1_contig = (op == HTP_OP_RMS_NORM_MUL && !broadcast_weight) ? ggml_is_contiguous_1(src1) : true;
+ const bool src1_contig = (has_weight && !broadcast_weight) ? ggml_is_contiguous_1(src1) : true;
const uint32_t src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
const uint32_t min_block_rows = (uint32_t) (std::max<size_t>(1, 2048 / dst_data_row_size));
uint32_t max_block = (src0_contig && src1_contig) ? src0_nrows : (uint32_t) src0->ne[1];
@@ -1603,7 +1607,7 @@ static bool htp_kparams_solve_unary_impl(
uint32_t vtcm_row_per_thread = 0;
htp_unary_solve_layout(&L, op, src0->ne[0], dst->ne[0],
- op == HTP_OP_RMS_NORM_MUL ? src1->ne[0] : 0,
+ has_weight ? src1->ne[0] : 0,
max_block, broadcast_weight, n_threads, vtcm_size, elem_size,
&col_tile, &vtcm_row_per_thread);
@@ -1642,6 +1646,16 @@ bool htp_kparams_solve_rms_norm_mul(
return htp_kparams_solve_unary_impl(sess, HTP_OP_RMS_NORM_MUL, src0, src1, dst, kparams);
}
+bool htp_kparams_solve_rms_norm_mul_add(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * src0,
+ const struct ggml_tensor * src1,
+ const struct ggml_tensor * dst,
+ struct htp_unary_kernel_params * kparams
+) {
+ return htp_kparams_solve_unary_impl(sess, HTP_OP_RMS_NORM_MUL_ADD, src0, src1, dst, kparams);
+}
+
bool htp_kparams_solve_repeat_get_rows(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
diff --git a/ggml/src/ggml-hexagon/htp-kparams.h b/ggml/src/ggml-hexagon/htp-kparams.h
index bac338314..653554be4 100644
--- a/ggml/src/ggml-hexagon/htp-kparams.h
+++ b/ggml/src/ggml-hexagon/htp-kparams.h
@@ -95,6 +95,14 @@ bool htp_kparams_solve_rms_norm_mul(
struct htp_unary_kernel_params * kparams
);
+bool htp_kparams_solve_rms_norm_mul_add(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * src0,
+ const struct ggml_tensor * src1,
+ const struct ggml_tensor * dst,
+ struct htp_unary_kernel_params * kparams
+);
+
bool htp_kparams_solve_matmul_add(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h
index bd7812486..a3ff51ebe 100644
--- a/ggml/src/ggml-hexagon/htp/htp-ops.h
+++ b/ggml/src/ggml-hexagon/htp/htp-ops.h
@@ -51,6 +51,7 @@ enum htp_op_code {
HTP_OP_MUL_MAT_ADD,
HTP_OP_RMS_NORM,
HTP_OP_RMS_NORM_MUL,
+ HTP_OP_RMS_NORM_MUL_ADD,
HTP_OP_UNARY_SILU,
HTP_OP_UNARY_GELU,
HTP_OP_UNARY_SIGMOID,
diff --git a/ggml/src/ggml-hexagon/htp/hvx-norm.h b/ggml/src/ggml-hexagon/htp/hvx-norm.h
index 145c3086e..5d13acc87 100644
--- a/ggml/src/ggml-hexagon/htp/hvx-norm.h
+++ b/ggml/src/ggml-hexagon/htp/hvx-norm.h
@@ -141,6 +141,71 @@ static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
}
}
+static inline void hvx_fast_rms_norm_mul_add_f32(const uint8_t * restrict src,
+ const uint8_t * restrict weight,
+ const uint8_t * restrict add,
+ uint8_t * restrict dst,
+ const int num_elems,
+ float epsilon,
+ float scale) {
+ const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
+ const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
+ const HVX_Vector * restrict v_add = (const HVX_Vector *) add;
+ HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
+
+ const int nvec = num_elems / VLEN_FP32;
+ const int nloe = num_elems % VLEN_FP32;
+
+ HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
+ HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+ }
+
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+ sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+ }
+
+ sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+
+ HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
+ HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
+ HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
+ HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
+
+ HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+ if (scale != 0.0f) {
+ scale_v = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(scale_v, hvx_vec_splat_f32(scale)));
+ }
+
+ #pragma unroll(4)
+ for (int i = 0; i < nvec; i++) {
+ HVX_Vector v1 = v_src[i];
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+ HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
+ HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[i]));
+ v_dst[i] = hvx_vec_add_f32_f32(v4, v_add[i]);
+ }
+
+ if (nloe > 0) {
+ HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+ HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+ HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+ HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
+ HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[nvec]));
+ HVX_Vector res_v = hvx_vec_add_f32_f32(v4, v_add[nvec]);
+
+ hvx_vec_store_a(&v_dst[nvec], nloe * 4, res_v);
+ }
+}
+
static inline void hvx_fast_norm_f32(const uint8_t * restrict src,
uint8_t * restrict dst,
const int num_elems,
diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c
index 565673212..1424f1213 100644
--- a/ggml/src/ggml-hexagon/htp/main.c
+++ b/ggml/src/ggml-hexagon/htp/main.c
@@ -931,6 +931,7 @@ static int execute_op(struct htp_ops_context * octx) {
case HTP_OP_NORM:
case HTP_OP_RMS_NORM:
case HTP_OP_RMS_NORM_MUL:
+ case HTP_OP_RMS_NORM_MUL_ADD:
case HTP_OP_SCALE:
case HTP_OP_CLAMP:
case HTP_OP_LEAKY_RELU:
diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c
index 3b63ac674..3a6840169 100644
--- a/ggml/src/ggml-hexagon/htp/unary-ops.c
+++ b/ggml/src/ggml-hexagon/htp/unary-ops.c
@@ -36,6 +36,13 @@ typedef void (*unary_rms_norm_mul_compute_fn_t)(const void * restrict src,
uint32_t num_rows,
const struct htp_unary_context * uctx);
+typedef void (*unary_rms_norm_mul_add_compute_fn_t)(const void * restrict src,
+ const void * restrict weight,
+ const void * restrict add,
+ void * restrict dst,
+ uint32_t num_rows,
+ const struct htp_unary_context * uctx);
+
typedef void (*unary_tri_compute_fn_t)(const void * restrict src,
void * restrict dst,
uint32_t num_rows,
@@ -63,6 +70,7 @@ struct htp_unary_context {
dma_addr_t data_src0;
dma_addr_t data_src1; // weight/scale tensor for RMS_NORM_MUL
+ dma_addr_t data_src2; // residual tensor for RMS_NORM_MUL_ADD
dma_addr_t data_dst;
size_t src0_data_row_size; // actual data bytes per row
@@ -87,10 +95,12 @@ struct htp_unary_context {
uint8_t * vtcm_src0;
uint8_t * vtcm_src1;
+ uint8_t * vtcm_src2;
uint8_t * vtcm_dst;
size_t vtcm_src0_size_per_thread;
size_t vtcm_src1_size_per_thread;
+ size_t vtcm_src2_size_per_thread;
size_t vtcm_dst_size_per_thread;
};
@@ -244,6 +254,29 @@ static void rms_norm_mul_f32(const void * restrict src,
}
}
+static void rms_norm_mul_add_f32(const void * restrict src,
+ const void * restrict weight,
+ const void * restrict add,
+ void * restrict dst,
+ const uint32_t num_rows,
+ const struct htp_unary_context * uctx) {
+ htp_unary_op_preamble;
+ float epsilon = 0.f;
+ float scale = 0.f;
+ memcpy(&epsilon, &op_params[0], sizeof(float));
+ memcpy(&scale, &op_params[1], sizeof(float));
+
+ const uint8_t * restrict w_local = (const uint8_t *) weight;
+
+ for (uint32_t ir = 0; ir < num_rows; ir++) {
+ const uint8_t * restrict src_local = (const uint8_t *) src + (ir * src0_row_size_aligned);
+ const uint8_t * restrict add_local = (const uint8_t *) add + (ir * src0_row_size_aligned);
+ uint8_t * restrict dst_local = (uint8_t *) dst + (ir * dst_row_size_aligned);
+
+ hvx_fast_rms_norm_mul_add_f32(src_local, w_local, add_local, dst_local, ne0, epsilon, scale);
+ }
+}
+
static void norm_f32(const void * restrict src,
void * restrict dst,
const uint32_t num_rows,
@@ -1430,6 +1463,118 @@ static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, vo
dma_queue_flush(dma_q);
}
+// 2b. RMS_NORM_MUL_ADD row-block task with broadcast weight and residual add.
+static void unary_thread_rms_norm_mul_add_f32(unsigned int nth, unsigned int ith, void * data) {
+ (void) nth;
+ const struct htp_unary_context * uctx = (const struct htp_unary_context *) data;
+ struct htp_ops_context * octx = uctx->octx;
+ const struct htp_tensor * src = octx->src[0];
+ const struct htp_tensor * dst = octx->dst;
+ struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+ htp_unary_preamble;
+
+ const uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread;
+ const size_t src0_data_row_size = uctx->src0_data_row_size;
+ const size_t dst_data_row_size = uctx->dst_data_row_size;
+ const size_t src0_row_size_aligned = uctx->src0_row_size_aligned;
+ const size_t dst_row_size_aligned = uctx->dst_row_size_aligned;
+
+ const uint32_t src0_nrows = uctx->src0_nrows;
+ const uint32_t src0_start_row = uctx->row_start + src0_nrows_per_thread * ith;
+ const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows);
+
+ if (src0_start_row >= src0_end_row) {
+ return;
+ }
+
+ const dma_addr_t data_src = uctx->data_src0;
+ const dma_addr_t data_src1 = uctx->data_src1;
+ const dma_addr_t data_src2 = uctx->data_src2;
+ const dma_addr_t data_dst = uctx->data_dst;
+
+ uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread);
+ uint8_t * src1_vtcm_data = uctx->vtcm_src1 + (ith * uctx->vtcm_src1_size_per_thread);
+ uint8_t * src2_vtcm_data = uctx->vtcm_src2 + (ith * uctx->vtcm_src2_size_per_thread);
+ uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread);
+
+ const size_t src0_vtcm_half_size = uctx->src0_vtcm_half_size;
+ const size_t dst_vtcm_half_size = uctx->dst_vtcm_half_size;
+
+ const bool src0_contig = uctx->kparams->src0_contig;
+ const uint32_t BLOCK = uctx->block;
+
+ const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01;
+ const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02;
+ const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012;
+
+ dma_queue * dma_q = octx->ctx->dma[ith];
+
+ dma_queue_push(dma_q, dma_make_data(src1_vtcm_data, data_src1),
+ uctx->src1_row_size_aligned, 0, uctx->src1_data_row_size, 1);
+ dma_queue_flush(dma_q);
+
+ for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) {
+ const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, ne01, div_ne01);
+
+ dma_queue_push(dma_q,
+ dma_make_data(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)),
+ nb1, dst_row_size_aligned, dst_data_row_size, 0);
+
+ const size_t src0_off = src0_contig ? ((size_t) ir * nb01) :
+ unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03);
+ dma_queue_push(dma_q,
+ dma_make_data(src0_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src + src0_off),
+ src0_row_size_aligned, nb01, src0_data_row_size, block_size);
+ dma_queue_push(dma_q,
+ dma_make_data(src2_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src2 + src0_off),
+ src0_row_size_aligned, nb01, src0_data_row_size, block_size);
+
+ ir += block_size;
+ }
+
+ unary_rms_norm_mul_add_compute_fn_t compute = (unary_rms_norm_mul_add_compute_fn_t) uctx->compute;
+
+ for (uint32_t ir = src0_start_row; ir < src0_end_row; ) {
+ const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, ne01, div_ne01);
+
+ void * dst_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).src;
+ void * src0_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst;
+ void * src2_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst;
+
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);
+ compute(src0_vtcm, src1_vtcm_data, src2_vtcm, dst_vtcm, block_size, uctx);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);
+
+ const size_t dst_off = (size_t) ir * nb1;
+ dma_queue_push(dma_q,
+ dma_make_data(data_dst + dst_off, dst_vtcm),
+ nb1, dst_row_size_aligned, dst_data_row_size, block_size);
+
+ const uint32_t next_ir = ir + block_size;
+ if (next_ir < src0_end_row) {
+ const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig,
+ ne01, div_ne01);
+ const uint32_t pref_ir = next_ir + next_block_size;
+ if (pref_ir < src0_end_row) {
+ const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig,
+ ne01, div_ne01);
+ const size_t src0_pref_off = src0_contig ? ((size_t) pref_ir * nb01) :
+ unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03);
+ dma_queue_push(dma_q,
+ dma_make_data(src0_vtcm, data_src + src0_pref_off),
+ src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size);
+ dma_queue_push(dma_q,
+ dma_make_data(src2_vtcm, data_src2 + src0_pref_off),
+ src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size);
+ }
+ }
+ ir += block_size;
+ }
+
+ dma_queue_flush(dma_q);
+}
+
// 3. TRI row-block task with row index ir.
static void unary_thread_tri_f32(unsigned int nth, unsigned int ith, void * data) {
(void) nth;
@@ -1749,9 +1894,10 @@ static int execute_op_unary(struct htp_ops_context * octx) {
const char * op_type = NULL;
switch (octx->op) {
- case HTP_OP_NORM: op_type = is_f16 ? "norm-f16" : "norm-f32"; break;
- case HTP_OP_RMS_NORM: op_type = is_f16 ? "rmsnorm-f16" : "rmsnorm-f32"; break;
- case HTP_OP_RMS_NORM_MUL: op_type = "rmsnorm-mul-f32"; break;
+ case HTP_OP_NORM: op_type = is_f16 ? "norm-f16" : "norm-f32"; break;
+ case HTP_OP_RMS_NORM: op_type = is_f16 ? "rmsnorm-f16" : "rmsnorm-f32"; break;
+ case HTP_OP_RMS_NORM_MUL: op_type = "rmsnorm-mul-f32"; break;
+ case HTP_OP_RMS_NORM_MUL_ADD: op_type = "rmsnorm-mul-add-f32"; break;
case HTP_OP_SCALE: op_type = is_f16 ? "scale-f16" : "scale-f32"; break;
case HTP_OP_CLAMP: op_type = is_f16 ? "clamp-f16" : "clamp-f32"; break;
case HTP_OP_LEAKY_RELU: op_type = "leaky-relu-f32"; break;
@@ -1817,9 +1963,12 @@ static int execute_op_unary(struct htp_ops_context * octx) {
bool broadcast_weight = kparams->broadcast_weight;
const struct htp_tensor * src1 = NULL;
- // RMS_NORM_MUL fusion is F32-only (its weight tensor is always F32; see
+ const bool has_weight = (octx->op == HTP_OP_RMS_NORM_MUL || octx->op == HTP_OP_RMS_NORM_MUL_ADD);
+ const bool has_src2 = (octx->op == HTP_OP_RMS_NORM_MUL_ADD);
+
+ // RMS_NORM_MUL* fusion is F32-only (its weight tensor is always F32; see
// try_fuse_node()'s type guard), so this never triggers when is_f16 is true.
- if (octx->op == HTP_OP_RMS_NORM_MUL) {
+ if (has_weight) {
src1 = octx->src[1];
src1_data_row_size = src1->ne[0] * sizeof(float);
}
@@ -1850,7 +1999,8 @@ static int execute_op_unary(struct htp_ops_context * octx) {
.row_start = row_start,
.data_src0 = src0->data,
- .data_src1 = (octx->op == HTP_OP_RMS_NORM_MUL) ? src1->data : 0,
+ .data_src1 = has_weight ? src1->data : 0,
+ .data_src2 = has_src2 ? octx->src[2]->data : 0,
.data_dst = dst->data,
.src0_data_row_size = src0_data_row_size,
@@ -1862,7 +2012,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
.dst_row_size_aligned = L.dst_row_size_aligned,
.src0_vtcm_half_size = L.src0_bytes >> 1,
- .src1_vtcm_half_size = (octx->op == HTP_OP_RMS_NORM_MUL) ? (L.src1_bytes >> (broadcast_weight ? 0 : 1)) : 0,
+ .src1_vtcm_half_size = has_weight ? (L.src1_bytes >> (broadcast_weight ? 0 : 1)) : 0,
.dst_vtcm_half_size = L.dst_bytes >> 1,
.block = kparams->block,
@@ -1872,10 +2022,12 @@ static int execute_op_unary(struct htp_ops_context * octx) {
.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0),
.vtcm_src1 = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_src1, L.src1_bytes > 0),
+ .vtcm_src2 = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_src2, L.src2_bytes > 0),
.vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst),
.vtcm_src0_size_per_thread = L.src0_bytes,
.vtcm_src1_size_per_thread = L.src1_bytes,
+ .vtcm_src2_size_per_thread = L.src2_bytes,
.vtcm_dst_size_per_thread = L.dst_bytes,
};
@@ -1940,6 +2092,10 @@ static int execute_op_unary(struct htp_ops_context * octx) {
task_func = unary_thread_rms_norm_mul_f32;
compute_func = (void *) rms_norm_mul_f32;
break;
+ case HTP_OP_RMS_NORM_MUL_ADD:
+ task_func = unary_thread_rms_norm_mul_add_f32;
+ compute_func = (void *) rms_norm_mul_add_f32;
+ break;
case HTP_OP_SCALE: compute_func = (void *) scale_f32; break;
case HTP_OP_CLAMP: compute_func = (void *) clamp_f32; break;
case HTP_OP_LEAKY_RELU: compute_func = (void *) leaky_relu_f32; break;
diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.h b/ggml/src/ggml-hexagon/htp/unary-ops.h
index 713e40d50..8699a1403 100644
--- a/ggml/src/ggml-hexagon/htp/unary-ops.h
+++ b/ggml/src/ggml-hexagon/htp/unary-ops.h
@@ -39,6 +39,7 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
case HTP_OP_NORM:
case HTP_OP_RMS_NORM:
case HTP_OP_RMS_NORM_MUL:
+ case HTP_OP_RMS_NORM_MUL_ADD:
case HTP_OP_SCALE:
case HTP_OP_SQR:
case HTP_OP_SQRT:
@@ -74,6 +75,7 @@ static inline bool htp_unary_is_elementwise(uint32_t opcode) {
case HTP_OP_NORM:
case HTP_OP_RMS_NORM:
case HTP_OP_RMS_NORM_MUL:
+ case HTP_OP_RMS_NORM_MUL_ADD:
case HTP_OP_L2_NORM:
case HTP_OP_HADAMARD:
case HTP_OP_TRI:
@@ -87,6 +89,7 @@ struct htp_unary_vtcm_layout {
size_t total_bytes;
size_t off_src0;
size_t off_src1;
+ size_t off_src2;
size_t off_dst;
size_t src0_row_size_aligned;
@@ -95,6 +98,7 @@ struct htp_unary_vtcm_layout {
size_t src0_bytes;
size_t src1_bytes;
+ size_t src2_bytes;
size_t dst_bytes;
};
@@ -110,19 +114,23 @@ static inline void htp_unary_vtcm_layout_build(
bool broadcast_weight,
uint32_t n_threads
) {
+ const bool has_weight = (op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_RMS_NORM_MUL_ADD);
+
L->src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
L->dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
- L->src1_row_size_aligned = (op == HTP_OP_RMS_NORM_MUL) ? hex_round_up(src1_data_row_size, 128) : 0;
+ L->src1_row_size_aligned = has_weight ? hex_round_up(src1_data_row_size, 128) : 0;
if (col_tile > 0) {
const size_t col_tile_bytes = col_tile * elem_size;
L->src0_bytes = col_tile_bytes * 2;
L->dst_bytes = col_tile_bytes * 2;
L->src1_bytes = 0;
+ L->src2_bytes = 0;
} else {
L->src0_bytes = L->src0_row_size_aligned * block * 2;
+ L->src2_bytes = (op == HTP_OP_RMS_NORM_MUL_ADD) ? L->src0_bytes : 0;
L->dst_bytes = L->dst_row_size_aligned * block * 2;
- if (op == HTP_OP_RMS_NORM_MUL) {
+ if (has_weight) {
L->src1_bytes = broadcast_weight ? L->src1_row_size_aligned : (L->src1_row_size_aligned * block * 2);
} else {
L->src1_bytes = 0;
@@ -132,6 +140,7 @@ static inline void htp_unary_vtcm_layout_build(
size_t off = 0;
VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes * n_threads);
VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_src1, L->src1_bytes * n_threads, L->src1_bytes > 0);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_src2, L->src2_bytes * n_threads, L->src2_bytes > 0);
VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes * n_threads);
L->total_bytes = off;
}
@@ -150,19 +159,26 @@ static inline void htp_unary_solve_layout(
uint32_t * out_col_tile,
uint32_t * out_vtcm_row_per_thread
) {
+ const bool has_weight = (op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_RMS_NORM_MUL_ADD);
+
const size_t dst_elem_size = (op == HTP_OP_HADAMARD) ? sizeof(float) : elem_size;
const size_t src0_data_row_size = ne00 * elem_size;
const size_t dst_data_row_size = ne10 * dst_elem_size;
- const size_t src1_data_row_size = (op == HTP_OP_RMS_NORM_MUL) ? (ne11 * sizeof(float)) : 0;
+ const size_t src1_data_row_size = has_weight ? (ne11 * sizeof(float)) : 0;
const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
- const size_t src1_row_size_aligned = (op == HTP_OP_RMS_NORM_MUL) ? hex_round_up(src1_data_row_size, 128) : 0;
+ const size_t src1_row_size_aligned = has_weight ? hex_round_up(src1_data_row_size, 128) : 0;
size_t vtcm_size_per_row = 0;
size_t vtcm_row_per_thread = 0;
- if (op == HTP_OP_RMS_NORM_MUL) {
+ if (op == HTP_OP_RMS_NORM_MUL_ADD) {
+ size_t src1_vtcm_total = n_threads * src1_row_size_aligned;
+ size_t available_vtcm = (vtcm_size > src1_vtcm_total) ? (vtcm_size - src1_vtcm_total) : 0;
+ vtcm_size_per_row = 2 * (2 * src0_row_size_aligned + dst_row_size_aligned);
+ vtcm_row_per_thread = available_vtcm / (n_threads * vtcm_size_per_row);
+ } else if (op == HTP_OP_RMS_NORM_MUL) {
if (broadcast_weight) {
size_t available_vtcm = vtcm_size;
size_t src1_vtcm_total = n_threads * src1_row_size_aligned;
@@ -187,8 +203,8 @@ static inline void htp_unary_solve_layout(
}
const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
- op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM ||
- op == HTP_OP_HADAMARD);
+ op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_RMS_NORM_MUL_ADD ||
+ op == HTP_OP_L2_NORM || op == HTP_OP_HADAMARD);
// The tiled fallback path below only has F32 task functions (unary_task_f32_tiled_*);
// F16 has no tiled kernels, so it must stay on the row-block path like reduction ops.
// NOTE: if F16 ends up with vtcm_row_per_thread == 0 here (row too large for the VTCM