CANN: add operator fusion support for ADD + RMS_NORM (#17512)
This commit implements operator fusion for ADD + RMS_NORM operations in the CANN backend to reduce memory access overhead and improve performance. The fusion is controlled by the GGML_CANN_OPERATOR_FUSION environment variable (default: false). Changes: - Implement ggml_cann_op_add_rms_norm_fused() using ACLNN AddRmsNorm - Add ggml_cann_can_fuse() to check fusion eligibility - Integrate fusion logic into computation graph evaluation - Add test cases for ADD + RMS_NORM fusion - Update documentation with new environment variable The fusion combines ADD and RMS_NORM into a single kernel call, which is more efficient than executing them separately.
This commit is contained in:
@@ -26,6 +26,7 @@
|
||||
#include "ggml.h"
|
||||
|
||||
#include <aclnnop/aclnn_add.h>
|
||||
#include <aclnnop/aclnn_add_rms_norm.h>
|
||||
#include <aclnnop/aclnn_addcdiv.h>
|
||||
#include <aclnnop/aclnn_argmax.h>
|
||||
#include <aclnnop/aclnn_avgpool2d.h>
|
||||
@@ -3805,3 +3806,57 @@ void ggml_cann_ssm_conv(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
|
||||
cubeMathType);
|
||||
}
|
||||
|
||||
|
||||
void ggml_cann_op_add_rms_norm_fused(ggml_backend_cann_context & ctx,
|
||||
ggml_tensor * add_node,
|
||||
ggml_tensor * rms_norm_node) {
|
||||
// Get the two input tensors for ADD operation
|
||||
ggml_tensor * x1 = add_node->src[0];
|
||||
ggml_tensor * x2 = add_node->src[1];
|
||||
|
||||
// Create ACL tensors for the two ADD inputs
|
||||
acl_tensor_ptr acl_x1 = ggml_cann_create_tensor(x1);
|
||||
acl_tensor_ptr acl_x2 = ggml_cann_create_tensor(x2);
|
||||
|
||||
// Get epsilon parameter from rms_norm_tensor
|
||||
float eps;
|
||||
memcpy(&eps, rms_norm_node->op_params, sizeof(float));
|
||||
|
||||
// Build gamma tensor (RMS normalization scaling factor)
|
||||
// Gamma should match the normalized dimensions (last dimension of x1)
|
||||
size_t acl_gamma_nb[GGML_MAX_DIMS];
|
||||
acl_gamma_nb[0] = ggml_type_size(rms_norm_node->type);
|
||||
for (int i = 1; i < GGML_MAX_DIMS; i++) {
|
||||
acl_gamma_nb[i] = acl_gamma_nb[i - 1] * x1->ne[i - 1];
|
||||
}
|
||||
acl_tensor_ptr acl_gamma =
|
||||
get_cache_acl_tensor(ctx, &ctx.rms_norm_one_tensor_cache.cache, ctx.rms_norm_one_tensor_cache.size, x1->ne,
|
||||
acl_gamma_nb, rms_norm_node->type,
|
||||
1, // dims - only the last dimension
|
||||
1.0f // value
|
||||
);
|
||||
|
||||
// Build rstdOut tensor (output for normalized standard deviation)
|
||||
// Shape should be the dimensions that are NOT normalized
|
||||
int64_t acl_rstd_ne[] = { 1, x1->ne[1], x1->ne[2], x1->ne[3] };
|
||||
size_t acl_rstd_nb[GGML_MAX_DIMS - 1];
|
||||
acl_rstd_nb[0] = sizeof(float);
|
||||
for (int i = 1; i < GGML_MAX_DIMS - 1; i++) {
|
||||
acl_rstd_nb[i] = acl_rstd_nb[i - 1] * acl_rstd_ne[i - 1];
|
||||
}
|
||||
acl_tensor_ptr acl_rstd =
|
||||
get_cache_acl_tensor(ctx, &ctx.rms_norm_zero_tensor_cache.cache, ctx.rms_norm_zero_tensor_cache.size,
|
||||
acl_rstd_ne, acl_rstd_nb, GGML_TYPE_F32, GGML_MAX_DIMS,
|
||||
0.0f // value
|
||||
);
|
||||
|
||||
acl_tensor_ptr acl_xout = ggml_cann_create_tensor(add_node);
|
||||
|
||||
// Create yOut tensor (final output after RMS normalization)
|
||||
acl_tensor_ptr acl_yout = ggml_cann_create_tensor(rms_norm_node);
|
||||
|
||||
// Call fused ADD + RMS_NORM operator
|
||||
GGML_CANN_CALL_ACLNN_OP(ctx, AddRmsNorm, acl_x1.get(), acl_x2.get(), acl_gamma.get(),
|
||||
eps, // double type
|
||||
acl_yout.get(), acl_rstd.get(), acl_xout.get());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user