layernorm.cpp

#include "layernorm.h"
#include "kernels/layernorm_kernels.h"

LayerNorm::LayerNorm(int hidden_size, float eps, bool elementwise_affine, Tensor::ScalarType dtype, Device device) : 
    hidden_size(hidden_size), eps(eps)
{
    if (elementwise_affine) {
        weight = Tensor::allocate({hidden_size}, dtype, device);
        bias = Tensor::allocate({hidden_size}, dtype, device);
    }

    registerParams
        (weight, "weight")
        (bias, "bias")
    ;
}

Tensor LayerNorm::forward(Tensor x) {
    Tensor out = Tensor::empty(x.shape, x.scalar_type(), x.device());
    layernorm_general(out, x, this->weight, this->bias, this->eps);
    return out;
}

Tensor RMSNorm::forward(Tensor x) {
    Tensor out = Tensor::empty(x.shape, use_quant ? Tensor::INT8 : x.scalar_type(), x.device());
    rms_norm(out, x, this->weight, this->variance_epsilon, this->use_quant);
    return out;
}

void RMSNormGeneral::forward_with_act_sum(Tensor x, Tensor quantized_hidden_states_buffer, Tensor quantized_scale_buffer, Tensor quantized_sum_buffer) {
    rms_norm_general_fuse_sum(quantized_hidden_states_buffer, x, this->weight, quantized_sum_buffer, quantized_scale_buffer, variance_epsilon, use_per_token_quant);
}

void RMSNormGeneral::forward_wo_act_sum(Tensor x, Tensor quantized_hidden_states_buffer, Tensor quantized_scale_buffer, Tensor quantized_sum_buffer) {
    rms_norm_general(quantized_hidden_states_buffer, x, this->weight, quantized_scale_buffer, variance_epsilon, use_per_token_quant);
}