Kernels:

kernels-test
/

only-torch-2.4

Trusted publisher

drbh

Migrated from model repo kernels-test/only-torch-2.4

4087a2f unverified 28 days ago

1.35 kB

	#include <torch/library.h>

	#include "registration.h"
	#include "torch_binding.h"

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	// Activation ops
	// Activation function used in SwiGLU.
	ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
	ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);

	// Activation function used in GeGLU with `none` approximation.
	ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
	ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);

	// Activation function used in GeGLU with `tanh` approximation.
	ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
	ops.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);

	// FATReLU implementation.
	ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
	ops.impl("fatrelu_and_mul", torch::kCUDA, &fatrelu_and_mul);

	// GELU implementation used in GPT-2.
	ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
	ops.impl("gelu_new", torch::kCUDA, &gelu_new);

	// Approximate GELU implementation.
	ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
	ops.impl("gelu_fast", torch::kCUDA, &gelu_fast);

	// Quick GELU implementation.
	ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
	ops.impl("gelu_quick", torch::kCUDA, &gelu_quick);
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)