rwkv7 support

josStorer · Dec 11, 2024 · 8c12b39 · 8c12b39
1 parent 01528c4
commit 8c12b39
Show file tree

Hide file tree

Showing 10 changed files with 771 additions and 17 deletions.
diff --git a/.github/workflows/pre-release.yml b/.github/workflows/pre-release.yml
@@ -78,6 +78,7 @@ jobs:
           rm ./backend-python/rwkv_pip/wkv_cuda.pyd
           rm ./backend-python/rwkv_pip/rwkv5.pyd
           rm ./backend-python/rwkv_pip/rwkv6.pyd
+          rm ./backend-python/rwkv_pip/wkv7s.pyd
           rm ./backend-python/get-pip.py
           rm ./backend-python/rwkv_pip/cpp/librwkv.dylib
           rm ./backend-python/rwkv_pip/cpp/rwkv.dll
@@ -107,6 +108,7 @@ jobs:
           rm ./backend-python/rwkv_pip/wkv_cuda.pyd
           rm ./backend-python/rwkv_pip/rwkv5.pyd
           rm ./backend-python/rwkv_pip/rwkv6.pyd
+          rm ./backend-python/rwkv_pip/wkv7s.pyd
           rm ./backend-python/get-pip.py
           rm ./backend-python/rwkv_pip/cpp/rwkv.dll
           rm ./backend-python/rwkv_pip/cpp/librwkv.so

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -99,6 +99,7 @@ jobs:
           rm ./backend-python/rwkv_pip/wkv_cuda.pyd
           rm ./backend-python/rwkv_pip/rwkv5.pyd
           rm ./backend-python/rwkv_pip/rwkv6.pyd
+          rm ./backend-python/rwkv_pip/wkv7s.pyd
           rm ./backend-python/get-pip.py
           rm ./backend-python/rwkv_pip/cpp/librwkv.dylib
           rm ./backend-python/rwkv_pip/cpp/rwkv.dll
@@ -126,6 +127,7 @@ jobs:
           rm ./backend-python/rwkv_pip/wkv_cuda.pyd
           rm ./backend-python/rwkv_pip/rwkv5.pyd
           rm ./backend-python/rwkv_pip/rwkv6.pyd
+          rm ./backend-python/rwkv_pip/wkv7s.pyd
           rm ./backend-python/get-pip.py
           rm ./backend-python/rwkv_pip/cpp/rwkv.dll
           rm ./backend-python/rwkv_pip/cpp/librwkv.so

diff --git a/backend-python/requirements.txt b/backend-python/requirements.txt
@@ -2,7 +2,7 @@ torch
 torchvision
 torchaudio
 setuptools==69.5.1
-rwkv==0.8.26
+rwkv==0.8.28
 langchain==0.0.322
 fastapi==0.109.1
 uvicorn==0.23.2

diff --git a/backend-python/requirements_without_cyac.txt b/backend-python/requirements_without_cyac.txt
@@ -2,7 +2,7 @@ torch
 torchvision
 torchaudio
 setuptools==69.5.1
-rwkv==0.8.26
+rwkv==0.8.28
 langchain==0.0.322
 fastapi==0.109.1
 uvicorn==0.23.2

diff --git a/backend-python/routes/completion.py b/backend-python/routes/completion.py
@@ -201,7 +201,10 @@ async def eval_rwkv(
             # torch_gc()
             requests_num = requests_num - 1
             completion_end_time = time.time()
-            completion_interval = completion_end_time - completion_start_time
+            if completion_start_time is not None:
+                completion_interval = completion_end_time - completion_start_time
+            else:
+                completion_interval = 0
             tps = 0
             if completion_interval > 0:
                 tps = completion_tokens / completion_interval

diff --git a/backend-python/rwkv_pip/cuda/rwkv7.cu b/backend-python/rwkv_pip/cuda/rwkv7.cu
@@ -0,0 +1,77 @@
+#include <stdio.h>
+#include <assert.h>
+#include "ATen/ATen.h"
+
+typedef at::Half fp16;
+typedef at::BFloat16 bf16;
+typedef float fp32;
+
+template <typename F>
+__global__ void kernel_forward(const int B, const int T, const int C, const int H,
+                               float *__restrict__ _state, const F *__restrict__ const _r, const F *__restrict__ const _w, const F *__restrict__ const _k, const F *__restrict__ const _v, const F *__restrict__ const _a, const F *__restrict__ const _b,
+                               F *__restrict__ const _y)
+{
+    const int e = blockIdx.x / H;
+    const int h = blockIdx.x % H;
+    const int i = threadIdx.x;
+    _state += h*_N_*_N_ + i*_N_; // wrong if B > 1 !!!
+
+    float state[_N_];
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        state[j] = _state[j];
+
+    __shared__ float r[_N_], k[_N_], w[_N_], a[_N_], b[_N_];
+
+    for (int _t = 0; _t < T; _t++)
+    {
+        const int t = e*T*C + h*_N_ + i + _t * C;
+        __syncthreads();
+        r[i] = float(_r[t]);
+        w[i] = __expf(-__expf(float(_w[t])));
+        k[i] = float(_k[t]);
+        a[i] = float(_a[t]);
+        b[i] = float(_b[t]);
+        __syncthreads();
+
+        float sa = 0;
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            sa += a[j] * state[j];
+        }
+
+        float vv = float(_v[t]);
+        float y = 0;
+        #pragma unroll
+        for (int j = 0; j < _N_; j++)
+        {
+            float& s = state[j];
+            s = s * w[j] + k[j] * vv + sa * b[j];
+            y += s * r[j];
+        }
+        _y[t] = F(y);
+    }
+    #pragma unroll
+    for (int j = 0; j < _N_; j++)
+        _state[j] = state[j];    
+}
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16* w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y)
+{
+    assert(H*_N_ == C);
+    assert(B == 1); // only for B=1
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
+}
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16* w, fp16 *k, fp16 *v, fp16 *a, fp16 *b, fp16 *y)
+{
+    assert(H*_N_ == C);
+    assert(B == 1); // only for B=1
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
+}
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32* w, fp32 *k, fp32 *v, fp32 *a, fp32 *b, fp32 *y)
+{
+    assert(H*_N_ == C);
+    assert(B == 1); // only for B=1
+    kernel_forward<<<dim3(B * H), dim3(_N_)>>>(B, T, C, H, state, r, w, k, v, a, b, y);
+}
diff --git a/backend-python/rwkv_pip/cuda/rwkv7_op.cpp b/backend-python/rwkv_pip/cuda/rwkv7_op.cpp
@@ -0,0 +1,26 @@
+#include <torch/extension.h>
+#include "ATen/ATen.h"
+
+typedef at::Half fp16;
+typedef at::BFloat16 bf16;
+typedef float fp32;
+
+void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *w, bf16 *k, bf16 *v, bf16 *a, bf16 *b, bf16 *y);
+void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *w, fp16 *k, fp16 *v, fp16 *a, fp16 *b, fp16 *y);
+void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *w, fp32 *k, fp32 *v, fp32 *a, fp32 *b, fp32 *y);
+
+void forward_bf16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y) {
+    cuda_forward_bf16(B, T, C, H, state.data_ptr<float>(), r.data_ptr<bf16>(), w.data_ptr<bf16>(), k.data_ptr<bf16>(), v.data_ptr<bf16>(), a.data_ptr<bf16>(), b.data_ptr<bf16>(), y.data_ptr<bf16>());
+}
+void forward_fp16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y) {
+    cuda_forward_fp16(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp16>(), w.data_ptr<fp16>(), k.data_ptr<fp16>(), v.data_ptr<fp16>(), a.data_ptr<fp16>(), b.data_ptr<fp16>(), y.data_ptr<fp16>());
+}
+void forward_fp32(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &w, torch::Tensor &k, torch::Tensor &v, torch::Tensor &a, torch::Tensor &b, torch::Tensor &y) {
+    cuda_forward_fp32(B, T, C, H, state.data_ptr<float>(), r.data_ptr<fp32>(), w.data_ptr<fp32>(), k.data_ptr<fp32>(), v.data_ptr<fp32>(), a.data_ptr<fp32>(), b.data_ptr<fp32>(), y.data_ptr<fp32>());
+}
+
+TORCH_LIBRARY(wkv7s, m) {
+    m.def("forward_bf16", forward_bf16);
+    m.def("forward_fp16", forward_fp16);
+    m.def("forward_fp32", forward_fp32);
+}