rwkv-ops 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

rwkv_ops/__init__.py

@@ -1,17 +1,16 @@
-__version__ = "0.2.2"
+__version__ = "0.3.0"
 import os
 
-KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "triton")
+KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "cuda").lower()
 KERAS_BACKEND = os.environ.get("KERAS_BACKEND")
 BACKEND = os.environ.get("KERNEL_BACKEND")
 
 
 if KERAS_BACKEND is not None:
-    BACKEND = KERAS_BACKEND
+    BACKEND = KERAS_BACKEND.lower()
 elif BACKEND is not None:
-    os.environ["KERAS_BACKEND"] = BACKEND
+    os.environ["KERAS_BACKEND"] = BACKEND.lower()
 else:
-    import torch
     import keras
 
     BACKEND = "torch"
@@ -22,7 +21,7 @@ assert BACKEND in ["torch", "jax", "numpy", "tensorflow"]
 from .rwkv7_kernel import get_generalized_delta_rule
 from .rwkv6_kernel import get_rwkv6_kernel
 
-generalized_delta_rule, RWKV7_USE_KERNEL = get_generalized_delta_rule(
+generalized_delta_rule, RWKV7_USE_TRITON_KERNEL = get_generalized_delta_rule(
     KERNEL_TYPE=KERNEL_TYPE
 )
 rwkv7_op = generalized_delta_rule

rwkv_ops/rwkv6_kernel/__init__.py

@@ -17,12 +17,6 @@ def get_rwkv6_kernel(KERNEL_TYPE="native"):
 
                 ops_kernel = False
             else:
-                print(
-                    "The CUDA kernel of RWKV6 cannot be used in JAX > = 0.6 version. If you need to use the CUDA KERNEL of RWKV6 version, please downgrade the version. It is recommended to downgrade to 0.4.34"
-                )
-                print(
-                    "RWKV6的CUDA kernel在JAX> = 0.6版本无法使用，如果需要使用RWKV6版本的CUDA KERNEL，请降级版本，建议降级到0.4.34"
-                )
                 CudaOperator = None
         elif keras.config.backend() == "torch":
             from .torch_rwkv_kernel import RWKVKernelOperator as CudaOperator

rwkv_ops/rwkv6_kernel/jax_kernel_cuda/gpu_ops.cpp

@@ -0,0 +1,44 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "kernels.h"
+#include "pybind11_kernel_helpers.h"
+
+namespace {
+pybind11::dict WKVRegistrations() {
+  pybind11::dict dict;
+  dict["wkv_forward"] =
+      gpu_ops::EncapsulateFunction(gpu_ops::rwkv_forward_fn);
+  dict["wkv_backward"] =
+      gpu_ops::EncapsulateFunction(gpu_ops::rwkv_backward_fn);
+  dict["wkv_forward_with_state"] = 
+      gpu_ops::EncapsulateFunction(gpu_ops::rwkv_forward_with_state_fn);
+  return dict;
+}
+
+PYBIND11_MODULE(gpu_ops, m) {
+  m.def("get_rwkv_registrations", &WKVRegistrations);
+  m.def("create_rwkv_descriptor",
+        [](int B, int T,int C, int H,bool S, gpu_ops::ElementType input_type,gpu_ops::ElementType output_type) {
+          return gpu_ops::PackDescriptor(gpu_ops::WKVDescriptor{B, T, C, H, S, input_type, output_type});
+        });
+
+  pybind11::enum_<gpu_ops::ElementType>(m, "ElementType")
+      .value("BF16", gpu_ops::ElementType::BF16)
+      .value("F16", gpu_ops::ElementType::F16)
+      .value("F32", gpu_ops::ElementType::F32);
+
+}
+} // namespace

rwkv_ops/rwkv6_kernel/jax_kernel_cuda/kernel_helpers.h

@@ -0,0 +1,64 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This header is not specific to our application and you'll probably want
+// something like this for any extension you're building. This includes the
+// infrastructure needed to serialize descriptors that are used with the
+// "opaque" parameter of the GPU custom call. In our example we'll use this
+// parameter to pass the size of our problem.
+
+#ifndef _GPU_OPS_KERNEL_HELPERS_H_
+#define _GPU_OPS_KERNEL_HELPERS_H_
+
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <type_traits>
+
+#define JAX_APEX_WARP_SIZE 32
+
+namespace gpu_ops {
+
+// https://en.cppreference.com/w/cpp/numeric/bit_cast
+template <class To, class From>
+typename std::enable_if<sizeof(To) == sizeof(From) &&
+                            std::is_trivially_copyable<From>::value &&
+                            std::is_trivially_copyable<To>::value,
+                        To>::type
+bit_cast(const From &src) noexcept {
+  static_assert(std::is_trivially_constructible<To>::value,
+                "This implementation additionally requires destination type to "
+                "be trivially constructible");
+
+  To dst;
+  memcpy(&dst, &src, sizeof(To));
+  return dst;
+}
+
+template <typename T> std::string PackDescriptorAsString(const T &descriptor) {
+  return std::string(bit_cast<const char *>(&descriptor), sizeof(T));
+}
+
+template <typename T>
+const T *UnpackDescriptor(const char *opaque, std::size_t opaque_len) {
+  if (opaque_len != sizeof(T)) {
+    throw std::runtime_error("Invalid opaque object size");
+  }
+  return bit_cast<const T *>(opaque);
+}
+
+} // namespace gpu_ops
+
+#endif

rwkv_ops/rwkv6_kernel/jax_kernel_cuda/kernels.h

@@ -0,0 +1,56 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef _GPU_OPS_KERNELS_H_
+#define _GPU_OPS_KERNELS_H_
+
+#include <cuda_runtime_api.h>
+
+#include <cstddef>
+#include <cstdint>
+
+#ifndef _N_
+  #define _N_ 8
+#endif
+#ifndef _T_
+  #define _T_ 16
+#endif
+namespace gpu_ops {
+
+enum ElementType { BF16, F16, F32 };
+
+struct WKVDescriptor {
+  int B;
+  int T;
+  int C;
+  int H;
+  bool S;
+  ElementType x_type;
+  ElementType y_type;
+};
+
+void rwkv_forward_fn(cudaStream_t stream, void **buffers,
+                      const char *opaque,
+                      std::size_t opaque_len);
+void rwkv_backward_fn(cudaStream_t stream, void **buffers,
+                      const char *opaque,
+                      std::size_t opaque_len);
+
+void rwkv_forward_with_state_fn(cudaStream_t stream, void **buffers,
+                      const char *opaque,
+                      std::size_t opaque_len);
+} // namespace gpu_ops
+
+#endif

rwkv_ops/rwkv6_kernel/jax_kernel_cuda/pybind11_kernel_helpers.h

@@ -0,0 +1,41 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This header extends kernel_helpers.h with the pybind11 specific interface to
+// serializing descriptors. It also adds a pybind11 function for wrapping our
+// custom calls in a Python capsule. This is separate from kernel_helpers so
+// that the CUDA code itself doesn't include pybind11. I don't think that this
+// is strictly necessary, but they do it in jaxlib, so let's do it here too.
+
+#ifndef _GPU_OPS_PYBIND11_KERNEL_HELPERS_H_
+#define _GPU_OPS_PYBIND11_KERNEL_HELPERS_H_
+
+#include <pybind11/pybind11.h>
+
+#include "kernel_helpers.h"
+
+namespace gpu_ops {
+
+template <typename T> pybind11::bytes PackDescriptor(const T &descriptor) {
+  return pybind11::bytes(PackDescriptorAsString(descriptor));
+}
+
+template <typename T> pybind11::capsule EncapsulateFunction(T *fn) {
+  return pybind11::capsule(bit_cast<void *>(fn), "xla._CUSTOM_CALL_TARGET");
+}
+
+} // namespace gpu_ops
+
+#endif