rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/rwkv6_kernel/__init__.py

@@ -0,0 +1,120 @@
+# copy right from https://github.com/infiy-quine/RWKV6_Keras_Operator
+import os
+import keras
+from keras import ops
+from distutils.util import strtobool
+from packaging import version
+
+
+def get_rwkv6_kernel(KERNEL_TYPE="native"):
+    ops_kernel = True
+    if KERNEL_TYPE == "cuda":
+        if keras.config.backend() == "jax":
+            import jax
+
+            if version.parse(jax.__version__) < version.parse("0.6.0"):
+                from .jax_rwkv_kernel import RWKVKernelOperator as CudaOperator
+
+                ops_kernel = False
+            else:
+                CudaOperator = None
+        elif keras.config.backend() == "torch":
+            from .torch_rwkv_kernel import RWKVKernelOperator as CudaOperator
+
+            ops_kernel = False
+        else:
+            CudaOperator = None
+    else:
+        CudaOperator = None
+    from .ops_rwkv_kernel import RWKVKernelOperator as OpsOperator
+
+    class RWKVKernelOperator:
+        def __init__(self, head_size, max_sequence_length, ops_loop=False):
+            self.enbale_cuda = CudaOperator is not None
+
+            if self.enbale_cuda:
+                self.cuda_operator = CudaOperator(head_size, max_sequence_length)
+
+            self.ops_operator = OpsOperator(head_size, max_sequence_length)
+
+            self.ops_loop = ops_loop
+
+        def __call__(
+            self, r, k, v, w, u, with_state=False, init_state=None, state_map=None
+        ):
+            seq_len = r.shape[1]
+
+            def call_parallel():
+                if self.enbale_cuda:
+                    return self.cuda_operator(
+                        r=r,
+                        k=k,
+                        v=v,
+                        w=w,
+                        u=u,
+                        with_state=with_state,
+                        init_state=init_state,
+                        state_map=state_map,
+                    )
+                else:
+                    return self.ops_operator(
+                        r=r,
+                        k=k,
+                        v=v,
+                        w=w,
+                        u=u,
+                        with_state=with_state,
+                        init_state=init_state,
+                        state_map=state_map,
+                    )
+
+            def call_one_step():
+                return self.ops_operator(
+                    r=r,
+                    k=k,
+                    v=v,
+                    w=w,
+                    u=u,
+                    with_state=with_state,
+                    init_state=init_state,
+                    state_map=state_map,
+                )
+
+            if not self.ops_loop:
+                return ops.cond(
+                    seq_len != 1 and not ops_kernel, call_parallel, call_one_step
+                )
+            else:
+                return call_parallel()
+
+    return RWKVKernelOperator
+
+
+# from .ops_rwkv_kernal import RWKVKernelOperator as OPSKernelOperator
+
+
+"""
+新增三个参数
+return_state 布尔类型 是否返回最终的state,如果想自定义init_state也需要启用这个开关
+
+init_state
+    当init_state省缺时，则使用全零初始化BatchSize维度上的状态。
+    形状: (state_kinds,num_heads,head_size, head_size)， 其中state_kinds为小于等于Batch_Size的正整数
+    精度: 在r为fp16时 init_state为fp32 其余时候类型与r相同
+
+
+state_map
+    形状: (Batch_Size,)
+    精度: int64, list[int]
+    这个数组定义了state到r上每个Batch维度切片间的映射关系
+    取值范围: [0, state_kinds)
+
+返回:
+    output, output_state 
+
+def __call__(self,r, k, v, w, u, return_state=False, init_state=None, state_map=None):
+
+
+
+
+"""

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