updates

build.toml

@@ -1,5 +1,6 @@
 [general]
 name = "rmsnorm_kernel"
+universal = false
 
 [torch]
 src = [
@@ -8,6 +9,7 @@ src = [
 ]
 
 [kernel.rmsnorm_kernel]
+backend = "cuda"
 src = [
   "rmsnorm_kernel/rmsnorm.cu",
 ]

torch-ext/torch_binding.cpp

@@ -4,8 +4,8 @@
 #include "torch_binding.h"
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  ops.def("rmsnorm_forward(Tensor input, Tensor gamma) -> ()");
+  ops.def("rmsnorm_forward(Tensor input, Tensor gamma) -> Tensor");
   ops.impl("rmsnorm_forward", torch::kCUDA, &rmsnorm_forward);
 }
 
-REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -2,4 +2,4 @@
 
 #include <torch/torch.h>
 
-void rmsnorm_forward(torch::Tensor const &input, torch::Tensor const &gamma);
+torch::Tensor rmsnorm_forward(torch::Tensor const &input, torch::Tensor const &gamma);