rwkv: add

rmsnorm/flake.nix

@@ -3,7 +3,8 @@
   inputs = {
     kernel-builder.url = "github:huggingface/kernel-builder";
   };
-  outputs = { self, kernel-builder, }:
+  outputs =
+    { self, kernel-builder }:
     kernel-builder.lib.genFlakeOutputs {
       path = ./.;
       rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;

rwkv/README.md

@@ -0,0 +1 @@
+RWKV kernel for transformers

rwkv/build.toml

@@ -0,0 +1,31 @@
+[general]
+name = "rwkv"
+universal = false
+
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+]
+
+[kernel.rwkv]
+depends = ["torch"]
+backend = "cuda"
+cuda-capabilities = [
+    "8.0",
+    "8.9",
+    "9.0",
+    "10.0",
+    "12.0",
+]
+include = ["."]
+src = [
+    "rwkv/wkv_cuda.cu",
+    "rwkv/wkv_cuda_bf16.cu",
+]
+cuda-flags = [
+        "-res-usage",
+        "--use_fast_math",
+        "-O3",
+        "--extra-device-vectorization",
+        "-DTmax=1024",
+]

rwkv/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for rwkv kernels";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}