Motif-Technologies
/

activation

Model card Files Files and versions

activation / torch-ext /torch_binding.h

iamwyldecat's picture

feat(rms-norm): Impl fused RMSNorm

f3b99fb about 2 months ago

history blame contribute delete

670 Bytes

	#pragma once

	#include <torch/torch.h>

	void poly_norm(torch::Tensor &out, const torch::Tensor &input, const torch::Tensor &weights, const torch::Tensor &bias, double eps);
	void poly_norm_backward(torch::Tensor& input_grad, torch::Tensor& weight_grad, torch::Tensor& bias_grad, const torch::Tensor& output_grad, const torch::Tensor& input, const torch::Tensor& weight, double eps);

	void rms_norm(torch::Tensor &out, const torch::Tensor &input, const torch::Tensor &weights, double eps);
	void rms_norm_backward(torch::Tensor& input_grad, torch::Tensor& weight_grad, const torch::Tensor& output_grad, const torch::Tensor& input, const torch::Tensor& weight, double eps);