diff --git a/backend-golang/rwkv.go b/backend-golang/rwkv.go index 7dd7e88..563f5eb 100644 --- a/backend-golang/rwkv.go +++ b/backend-golang/rwkv.go @@ -28,7 +28,7 @@ func (a *App) StartServer(python string, port int, host string, webui bool, rwkv args = append(args, "--webui") } if rwkvBeta { - args = append(args, "--rwkv-beta") + // args = append(args, "--rwkv-beta") } if rwkvcpp { args = append(args, "--rwkv.cpp") diff --git a/backend-python/main.py b/backend-python/main.py index 92d6b20..a38285b 100644 --- a/backend-python/main.py +++ b/backend-python/main.py @@ -27,11 +27,6 @@ def get_args(args: Union[Sequence[str], None] = None): action="store_true", help="whether to enable WebUI (default: False)", ) - group.add_argument( - "--rwkv-beta", - action="store_true", - help="whether to use rwkv-beta (default: False)", - ) group.add_argument( "--rwkv.cpp", action="store_true", diff --git a/backend-python/rwkv_pip/beta/cuda/att_one.cu b/backend-python/rwkv_pip/beta/cuda/att_one.cu deleted file mode 100644 index 743fc12..0000000 --- a/backend-python/rwkv_pip/beta/cuda/att_one.cu +++ /dev/null @@ -1,124 +0,0 @@ -#include "ATen/ATen.h" -#include -#include -#include - -#include "element_wise.h" -#include "util.h" - -// Equivalent Python code: -// ww = t_first + k -// p = torch.maximum(pp, ww) -// e1 = torch.exp(pp - p) -// e2 = torch.exp(ww - p) -// wkv = ((e1 * aa + e2 * v) / (e1 * bb + e2)).to(dtype=x.dtype) -// ww = t_decay + pp -// p = torch.maximum(ww, k) -// e1 = torch.exp(ww - p) -// e2 = torch.exp(k - p) -// t1 = e1 * aa + e2 * v -// t2 = e1 * bb + e2 -// r = r * wkv -// return t1, t2, p, r -struct WkvForwardOne { - const float *t_first; - const float *k; - const float *pp; - const float *aa; - const float *bb; - const float *t_decay; - const float *v; - /* out */ float *t1; - /* out */ float *t2; - /* out */ float *p; - /* in & out */ half *r; - - __device__ void operator()(int i) const { - float ww = t_first[i] + k[i]; - float pp_ = pp[i]; - float p_ = (pp_ > ww) ? pp_ : ww; - float e1 = expf(pp_ - p_); - float e2 = expf(ww - p_); - float aa_ = aa[i]; - float bb_ = bb[i]; - float v_ = v[i]; - r[i] = __hmul(r[i], __float2half(((e1 * aa_ + e2 * v_) / (e1 * bb_ + e2)))); - ww = t_decay[i] + pp_; - float k_ = k[i]; - p_ = (ww > k_) ? ww : k_; - e1 = expf(ww - p_); - e2 = expf(k_ - p_); - t1[i] = e1 * aa_ + e2 * v_; - t2[i] = e1 * bb_ + e2; - p[i] = p_; - } -}; - -/* - Equivalent Python code: - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) -*/ - -struct Mix { - const half *xx; - const half *sx; - const half *k_mix; - const half *v_mix; - const half *r_mix; - /* out */ half *kx; - /* out */ half *vx; - /* out */ half *rx; - - __device__ void operator()(int i) const { - half xx_ = xx[i]; - half sx_ = sx[i]; - half k_mix_ = k_mix[i]; - half v_mix_ = v_mix[i]; - half r_mix_ = r_mix[i]; - kx[i] = __hadd(__hmul(xx_, k_mix_), - __hmul(sx_, __hsub(__float2half(1), k_mix_))); - vx[i] = __hadd(__hmul(xx_, v_mix_), - __hmul(sx_, __hsub(__float2half(1), v_mix_))); - rx[i] = __hadd(__hmul(xx_, r_mix_), - __hmul(sx_, __hsub(__float2half(1), r_mix_))); - } -}; - -using torch::Tensor; - -void gemm_fp16_cublas_tensor(Tensor a, Tensor b, Tensor c); - -Tensor att_one(Tensor x, Tensor ln_w, Tensor ln_b, Tensor sx, Tensor k_mix, - Tensor v_mix, Tensor r_mix, Tensor kw, - /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx, Tensor rw, - /* imm */ Tensor rx, Tensor ow, Tensor t_first, - /* imm */ Tensor k, Tensor pp, Tensor ww, Tensor aa, Tensor bb, - Tensor t_decay, /* imm */ Tensor v, /* in & out */ Tensor r, - /* out */ Tensor x_plus_out, /* out */ Tensor t1, - /* out */ Tensor t2, /* out */ Tensor p) { - Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b); - element_wise(Mix{data_ptr(xx), data_ptr(sx), - data_ptr(k_mix), data_ptr(v_mix), - data_ptr(r_mix), data_ptr(kx), - data_ptr(vx), data_ptr(rx)}, - x.numel()); - - gemm_fp16_cublas_tensor(kx, kw, k); - gemm_fp16_cublas_tensor(vx, vw, v); - gemm_fp16_cublas_tensor(rx, rw, r); - at::sigmoid_(r); - - element_wise(WkvForwardOne{data_ptr(t_first), data_ptr(k), - data_ptr(pp), data_ptr(aa), - data_ptr(bb), data_ptr(t_decay), - data_ptr(v), data_ptr(t1), - data_ptr(t2), data_ptr(p), - data_ptr(r)}, - x.numel()); - - gemm_fp16_cublas_tensor(r, ow, x_plus_out); - x_plus_out += x; - return xx; -} diff --git a/backend-python/rwkv_pip/beta/cuda/att_one_v5.cu b/backend-python/rwkv_pip/beta/cuda/att_one_v5.cu deleted file mode 100644 index 98f811e..0000000 --- a/backend-python/rwkv_pip/beta/cuda/att_one_v5.cu +++ /dev/null @@ -1,109 +0,0 @@ -#include "ATen/ATen.h" -#include -#include -#include - -#include "element_wise.h" -#include "util.h" - -// Equivalent Python code: -// s1 = t_first * a + s -// s2 = a + t_decay * s -struct Fused1 { - const float *t_first; - const float *t_decay; - const float *a; - const float *s; - const int32_t inner_size; - /* out */ float *s1; - /* out */ float *s2; - - __device__ void operator()(int i) const { - const int j = i / inner_size; - s1[i] = t_first[j] * a[i] + s[i]; - s2[i] = a[i] + t_decay[j] * s[i]; - } -}; - -/* - Equivalent Python code: - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) -*/ - -struct Mix { - const half *xx; - const half *sx; - const half *k_mix; - const half *v_mix; - const half *r_mix; - /* out */ half *kx; - /* out */ half *vx; - /* out */ half *rx; - - __device__ void operator()(int i) const { - half xx_ = xx[i]; - half sx_ = sx[i]; - half k_mix_ = k_mix[i]; - half v_mix_ = v_mix[i]; - half r_mix_ = r_mix[i]; - kx[i] = __hadd(__hmul(xx_, k_mix_), - __hmul(sx_, __hsub(__float2half(1), k_mix_))); - vx[i] = __hadd(__hmul(xx_, v_mix_), - __hmul(sx_, __hsub(__float2half(1), v_mix_))); - rx[i] = __hadd(__hmul(xx_, r_mix_), - __hmul(sx_, __hsub(__float2half(1), r_mix_))); - } -}; - -using torch::Tensor; - -void gemm_fp16_cublas_tensor(Tensor a, Tensor b, Tensor c); - -Tensor att_one_v5(Tensor x, Tensor sx, Tensor s, Tensor ln_w, Tensor ln_b, - Tensor lx_w, Tensor lx_b, Tensor k_mix, Tensor v_mix, - Tensor r_mix, Tensor kw, - /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx, - Tensor rw, - /* imm */ Tensor rx, Tensor ow, Tensor t_first, - /* imm */ Tensor k, Tensor t_decay, /* imm */ Tensor v, - /* imm */ Tensor r, /* imm */ Tensor s1, - /* out */ Tensor x_plus_out, /* out */ Tensor s2) { - Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b); - element_wise(Mix{data_ptr(xx), data_ptr(sx), - data_ptr(k_mix), data_ptr(v_mix), - data_ptr(r_mix), data_ptr(kx), - data_ptr(vx), data_ptr(rx)}, - x.numel()); - - int H = t_decay.size(0); - int S = x.size(-1) / H; - gemm_fp16_cublas_tensor(rx, rw, r); - r = at::reshape(r, {H, 1, S}); - gemm_fp16_cublas_tensor(kx, kw, k); - k = at::reshape(k, {H, S, 1}); - gemm_fp16_cublas_tensor(vx, vw, v); - v = at::reshape(v, {H, 1, S}); - - { - Tensor a = at::matmul(k, v); - - // s1 = t_first * a + s - // s2 = a + t_decay * s - element_wise(Fused1{data_ptr(t_first), data_ptr(t_decay), - data_ptr(a), data_ptr(s), - static_cast(a.size(1) * a.size(2)), - data_ptr(s1), data_ptr(s2)}, - a.numel()); - } - - Tensor out = at::matmul(r, s1); - out = at::flatten(out); - out = at::squeeze(at::group_norm(at::unsqueeze(out, 0), H, lx_w, lx_b), 0); - out = at::_cast_Half(out); - - gemm_fp16_cublas_tensor(out, ow, x_plus_out); - x_plus_out += x; - return xx; -} diff --git a/backend-python/rwkv_pip/beta/cuda/att_seq.cu b/backend-python/rwkv_pip/beta/cuda/att_seq.cu deleted file mode 100644 index c8db033..0000000 --- a/backend-python/rwkv_pip/beta/cuda/att_seq.cu +++ /dev/null @@ -1,178 +0,0 @@ -#include "ATen/ATen.h" -#include -#include -#include - -#include "util.h" -#include "element_wise.h" - -using torch::Tensor; - -void gemm_fp16_cublas(const void *a, const void *b, void *c, int m, - int n, int k, bool output_fp32); - -// based on `kernel_wkv_forward`, fusing more operations -__global__ void kernel_wkv_forward_new( - const int B, const int T, const int C, const float *__restrict__ const _w, - const float *__restrict__ const _u, const float *__restrict__ const _k, - const float *__restrict__ const _v, const half *__restrict__ const r, - half *__restrict__ const _y, float *__restrict__ const _aa, - float *__restrict__ const _bb, float *__restrict__ const _pp) { - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - const int _b = idx / C; - const int _c = idx % C; - const int _offset = _b * T * C + _c; - const int _state_offset = _b * C + _c; - - float u = _u[_c]; - float w = _w[_c]; - const float *__restrict__ const k = _k + _offset; - const float *__restrict__ const v = _v + _offset; - half *__restrict__ const y = _y + _offset; - - float aa = _aa[_state_offset]; - float bb = _bb[_state_offset]; - float pp = _pp[_state_offset]; - for (int i = 0; i < T; i++) { - const int ii = i * C; - const float kk = k[ii]; - const float vv = v[ii]; - float ww = u + kk; - float p = max(pp, ww); - float e1 = exp(pp - p); - float e2 = exp(ww - p); - y[ii] = __float2half((e1 * aa + e2 * vv) / (e1 * bb + e2)); - ww = w + pp; - p = max(ww, kk); - e1 = exp(ww - p); - e2 = exp(kk - p); - aa = e1 * aa + e2 * vv; - bb = e1 * bb + e2; - pp = p; - } - _aa[_state_offset] = aa; - _bb[_state_offset] = bb; - _pp[_state_offset] = pp; -} - -void cuda_wkv_forward_new(int B, int T, int C, float *w, float *u, float *k, - float *v, half *r, half *y, float *aa, float *bb, - float *pp) { - dim3 threadsPerBlock(min(C, 32)); - assert(B * C % threadsPerBlock.x == 0); - dim3 numBlocks(B * C / threadsPerBlock.x); - kernel_wkv_forward_new<<>>(B, T, C, w, u, k, v, r, - y, aa, bb, pp); -} - -__global__ void _att_mix(const half *xx, const half *sx, const half *k_mix, - const half *v_mix, const half *r_mix, - const int outer_size, const int inner_size, half *kx, - half *vx, half *rx) { - for (int idx2 = blockIdx.x * blockDim.x + threadIdx.x; idx2 < inner_size; - idx2 += blockDim.x * gridDim.x) { - half k_mix_ = k_mix[idx2]; - half v_mix_ = v_mix[idx2]; - half r_mix_ = r_mix[idx2]; - for (int row = 0; row < outer_size; ++row) { - int idx1 = row * inner_size + idx2; - half xx_ = xx[idx1]; - half sx_ = sx[idx1]; - kx[idx1] = __hadd(__hmul(xx_, k_mix_), - __hmul(sx_, __hsub(__float2half(1), k_mix_))); - vx[idx1] = __hadd(__hmul(xx_, v_mix_), - __hmul(sx_, __hsub(__float2half(1), v_mix_))); - rx[idx1] = __hadd(__hmul(xx_, r_mix_), - __hmul(sx_, __hsub(__float2half(1), r_mix_))); - } - } -} - -void att_mix(const half *xx, const half *sx, const half *k_mix, - const half *v_mix, const half *r_mix, const int outer_size, - const int inner_size, half *kx, half *vx, half *rx) { - // 256 is good enough on most GPUs - const int32_t BLOCK_SIZE = 256; - assert(inner_size % BLOCK_SIZE == 0); - _att_mix<<>>( - xx, sx, k_mix, v_mix, r_mix, outer_size, inner_size, kx, vx, rx); -} - -struct InplaceSigmoid { - __device__ __forceinline__ half operator()(int i) const { - ptr[i] = __float2half(1.0 / (1.0 + exp(-__half2float(ptr[i])))); - } - half *ptr; -}; - -struct InplaceMul { - __device__ __forceinline__ half operator()(int i) const { - y[i] = __hmul(x[i], y[i]); - } - half *y; - half *x; -}; - -/* - Equivalent Python code: - - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1,:])) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - k = gemm(kx, kw, output_dtype=torch.float32) - v = gemm(vx, vw, output_dtype=torch.float32) - - T = x.shape[0] - for t in range(T): - kk = k[t] - vv = v[t] - ww = t_first + kk - p = torch.maximum(pp, ww) - e1 = torch.exp(pp - p) - e2 = torch.exp(ww - p) - sx[t] = ((e1 * aa + e2 * vv) / (e1 * bb + e2)).to(dtype=x.dtype) - ww = t_decay + pp - p = torch.maximum(ww, kk) - e1 = torch.exp(ww - p) - e2 = torch.exp(kk - p) - aa = e1 * aa + e2 * vv - bb = e1 * bb + e2 - pp = p - out = gemm(r * sx, ow) - return x + out, xx[-1,:], aa, bb, pp -*/ -Tensor att_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, - Tensor v_mix, Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, - Tensor ow, Tensor t_first, Tensor pp, Tensor aa, Tensor bb, - Tensor t_decay, /* imm */ Tensor buf, /* out */ Tensor x_plus_out) { - Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b); - sx = at::cat({sx.unsqueeze(0), xx.slice(0, 0, -1)}, 0); - char* buf_ptr = (char*)buf.data_ptr(); - half* kx = (half*)buf_ptr; - half* vx = kx + x.numel(); - half* rx = vx + x.numel(); - half* wkv_y = rx + x.numel(); - att_mix(data_ptr(xx), data_ptr(sx), data_ptr(k_mix), - data_ptr(v_mix), data_ptr(r_mix), xx.size(0), xx.size(1), - kx, vx, rx); - float* k = reinterpret_cast(wkv_y + x.numel()); - float* v = k + x.size(0) * kw.size(1); - half* r = reinterpret_cast(v + x.size(0) * vw.size(1)); - - gemm_fp16_cublas(kx, kw.data_ptr(), k, x.size(0), kw.size(1), kw.size(0), true); - gemm_fp16_cublas(vx, vw.data_ptr(), v, x.size(0), vw.size(1), vw.size(0), true); - gemm_fp16_cublas(rx, rw.data_ptr(), r, x.size(0), rw.size(1), rw.size(0), false); - element_wise(InplaceSigmoid{r}, x.size(0) * rw.size(1)); - cuda_wkv_forward_new(1, x.size(0), x.size(1), data_ptr(t_decay), - data_ptr(t_first), k, v, r, - wkv_y, data_ptr(aa), - data_ptr(bb), data_ptr(pp)); - element_wise(InplaceMul{wkv_y, r}, x.numel()); - gemm_fp16_cublas(wkv_y, ow.data_ptr(), x_plus_out.data_ptr(), x.size(0), ow.size(1), ow.size(0), false); - x_plus_out += x; - return xx; -} diff --git a/backend-python/rwkv_pip/beta/cuda/element_wise.h b/backend-python/rwkv_pip/beta/cuda/element_wise.h deleted file mode 100644 index eedc2f9..0000000 --- a/backend-python/rwkv_pip/beta/cuda/element_wise.h +++ /dev/null @@ -1,21 +0,0 @@ -#include -#include -#include - -template __global__ void _element_wise(Func func, int n) { - for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; - i += blockDim.x * gridDim.x) { - func(i); - } -} - -// NOTE: packed data type (e.g. float4) is a overkill for current sizes -// (4096 in 7B model and 768 in 0.1B model), -// and is not faster than the plain float version. -template -void element_wise(Func func, int n) { - // 256 is good enough on most GPUs - const int32_t BLOCK_SIZE = 256; - assert(n % BLOCK_SIZE == 0); - _element_wise<<>>(func, n); -} diff --git a/backend-python/rwkv_pip/beta/cuda/ffn.cu b/backend-python/rwkv_pip/beta/cuda/ffn.cu deleted file mode 100644 index c1c2c80..0000000 --- a/backend-python/rwkv_pip/beta/cuda/ffn.cu +++ /dev/null @@ -1,165 +0,0 @@ -#include "ATen/ATen.h" -#include -#include -#include - -#include "element_wise.h" -#include "util.h" - -using torch::Tensor; - -void gemm_fp16_cublas(const void *a, const void *b, void *c, int ori_m, - int ori_n, int ori_k, bool output_fp32); - -__global__ void _ffn_seq_mix(const half *xx, const half *sx, const half *k_mix, - const half *r_mix, const int outer_size, - const int inner_size, half *kx, half *rx) { - for (int idx2 = blockIdx.x * blockDim.x + threadIdx.x; idx2 < inner_size; - idx2 += blockDim.x * gridDim.x) { - half k_mix_ = k_mix[idx2]; - half r_mix_ = r_mix[idx2]; - for (int row = 0; row < outer_size; ++row) { - int idx1 = row * inner_size + idx2; - half xx_ = xx[idx1]; - half sx_ = sx[idx1]; - kx[idx1] = __hadd(__hmul(xx_, k_mix_), - __hmul(sx_, __hsub(__float2half(1), k_mix_))); - rx[idx1] = __hadd(__hmul(xx_, r_mix_), - __hmul(sx_, __hsub(__float2half(1), r_mix_))); - } - } -} - -void ffn_seq_mix(const half *xx, const half *sx, const half *k_mix, - const half *r_mix, const int outer_size, const int inner_size, - half *kx, half *rx) { - // 256 is good enough on most GPUs - const int32_t BLOCK_SIZE = 256; - assert(inner_size % BLOCK_SIZE == 0); - _ffn_seq_mix<<>>( - xx, sx, k_mix, r_mix, outer_size, inner_size, kx, rx); -} - -struct InplaceSigmoid { - __device__ __forceinline__ void operator()(int i) const { - ptr[i] = __float2half(1.0 / (1.0 + exp(-__half2float(ptr[i])))); - } - half *ptr; -}; - -struct InplaceReLUAndSquare { - __device__ __forceinline__ void operator()(int i) const { - // __hmax is not defined in old cuda - if (__hgt(ptr[i], __float2half(0))) { - ptr[i] = __hmul(ptr[i], ptr[i]); - } else { - ptr[i] = __float2half(0); - } - } - half *ptr; -}; - -struct InplaceFma { - __device__ __forceinline__ void operator()(int i) const { - a[i] = __hfma(a[i], b[i], c[i]); - } - half *a; - const half *b; - const half *c; -}; - -/* - Equivalent Python code: - - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1,:])) - kx = xx * k_mix + sx * (1 - k_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - vx = torch.square(torch.relu(gemm(kx, kw))) - out = r * gemm(vx, vw) - return x + out, xx[-1,:] -*/ -Tensor ffn_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, - Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, - /* imm */ Tensor buf, - /* out */ Tensor x_plus_out) { - Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b); - sx = at::cat({sx.unsqueeze(0), xx.slice(0, 0, -1)}, 0); - char *buf_ptr = (char *)buf.data_ptr(); - half *kx = (half *)buf_ptr; - half *rx = kx + x.numel(); - half *vx = rx + x.numel(); - half *r = vx + x.size(0) * kw.size(1); - ffn_seq_mix(data_ptr(xx), data_ptr(sx), data_ptr(k_mix), - data_ptr(r_mix), xx.size(0), xx.size(1), kx, rx); - - gemm_fp16_cublas(rx, rw.data_ptr(), r, x.size(0), rw.size(1), x.size(1), - false); - element_wise(InplaceSigmoid{r}, x.size(0) * rw.size(1)); - gemm_fp16_cublas(kx, kw.data_ptr(), vx, x.size(0), kw.size(1), x.size(1), - false); - element_wise(InplaceReLUAndSquare{vx}, x.size(0) * kw.size(1)); - gemm_fp16_cublas(vx, vw.data_ptr(), x_plus_out.data_ptr(), x.size(0), - vw.size(1), vw.size(0), false); - element_wise(InplaceFma{data_ptr(x_plus_out), r, data_ptr(x)}, - x_plus_out.numel()); - return xx; -} - -struct FfnOneMix { - __device__ __forceinline__ void operator()(int idx) { - half k_mix_ = k_mix[idx]; - half r_mix_ = r_mix[idx]; - half xx_ = xx[idx]; - half sx_ = sx[idx]; - kx[idx] = __hadd(__hmul(xx_, k_mix_), - __hmul(sx_, __hsub(__float2half(1), k_mix_))); - rx[idx] = __hadd(__hmul(xx_, r_mix_), - __hmul(sx_, __hsub(__float2half(1), r_mix_))); - } - half *k_mix; - half *r_mix; - half *xx; - half *sx; - half *kx; - half *rx; -}; - -/* - Equivalent Python code: - - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - kx = xx * k_mix + sx * (1 - k_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - vx = torch.square(torch.relu(gemm(kx, kw))) - out = r * gemm(vx, vw) - return x + out, xx -*/ -Tensor ffn_one(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, - Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, - /* imm */ Tensor buf, - /* out */ Tensor x_plus_out) { - Tensor xx = at::layer_norm(x, {x.size(-1)}, ln_w, ln_b); - char *buf_ptr = (char *)buf.data_ptr(); - half *kx = (half *)buf_ptr; - half *rx = kx + x.numel(); - half *vx = rx + x.numel(); - half *r = vx + x.size(0) * kw.size(1); - element_wise(FfnOneMix{data_ptr(k_mix), data_ptr(r_mix), - data_ptr(xx), data_ptr(sx), kx, rx}, - x.numel()); - // vector * matrix, so m = 1 - gemm_fp16_cublas(rx, rw.data_ptr(), r, 1, rw.size(1), rw.size(0), false); - element_wise(InplaceSigmoid{r}, rw.size(1)); - gemm_fp16_cublas(kx, kw.data_ptr(), vx, 1, kw.size(1), kw.size(0), false); - element_wise(InplaceReLUAndSquare{vx}, kw.size(1)); - gemm_fp16_cublas(vx, vw.data_ptr(), x_plus_out.data_ptr(), 1, vw.size(1), - vw.size(0), false); - element_wise(InplaceFma{data_ptr(x_plus_out), r, data_ptr(x)}, - x_plus_out.numel()); - return xx; -} diff --git a/backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp b/backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp deleted file mode 100644 index e1162ad..0000000 --- a/backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp +++ /dev/null @@ -1,128 +0,0 @@ -#include -#include -#include -#include -#include - -#define CUBLAS_CHECK(condition) \ - for (cublasStatus_t _cublas_check_status = (condition); \ - _cublas_check_status != CUBLAS_STATUS_SUCCESS;) \ - throw std::runtime_error("cuBLAS error " + \ - std::to_string(_cublas_check_status) + " at " + \ - std::to_string(__LINE__)); - -#define CUDA_CHECK(condition) \ - for (cudaError_t _cuda_check_status = (condition); \ - _cuda_check_status != cudaSuccess;) \ - throw std::runtime_error( \ - "CUDA error " + std::string(cudaGetErrorString(_cuda_check_status)) + \ - " at " + std::to_string(__LINE__)); - -cublasHandle_t get_cublas_handle() { - static cublasHandle_t cublas_handle = []() { - cublasHandle_t handle = nullptr; - CUBLAS_CHECK(cublasCreate(&handle)); -#if CUDA_VERSION < 11000 - CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH)); -#else - CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH)); -#endif // CUDA_VERSION < 11000 - return handle; - }(); - return cublas_handle; -} - -/* - NOTE: blas gemm is column-major by default, but we need row-major output. - The data of row-major, transposed matrix is exactly the same as the - column-major, non-transposed matrix, and C = A * B ---> C^T = B^T * A^T - */ -void gemm_fp16_cublas(const void *a, const void *b, void *c, int ori_m, - int ori_n, int ori_k, bool output_fp32) { - const auto cuda_data_type = CUDA_R_16F; - const auto cuda_c_data_type = output_fp32 ? CUDA_R_32F : CUDA_R_16F; - const auto compute_type = CUDA_R_32F; - const float sp_alpha = 1.f; - // use CUBLAS_OP_N. see the notes above - const cublasOperation_t cublas_trans_a = CUBLAS_OP_N; - const cublasOperation_t cublas_trans_b = CUBLAS_OP_N; - // m = (B^T).size(0) = B.size(1) = n; - const int cublas_m = ori_n; - const int cublas_k = ori_k; - // comptiable with rwkv one mode, where 1-D tensor * 2-D tensor - // const int n = a.dense_dim() == 1 ? 1 : a.size(0); - const int cublas_n = ori_m; - const int cublas_lda = cublas_m; - const int cublas_ldb = cublas_k; - const int cublas_ldc = cublas_m; - cublasHandle_t cublas_handle = get_cublas_handle(); - -#if CUDA_VERSION >= 11000 - cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT; -#else - cublasGemmAlgo_t algo = CUBLAS_GEMM_DFALT_TENSOR_OP; -#endif - const float sp_beta = 0.f; - CUBLAS_CHECK(cublasGemmEx( - cublas_handle, cublas_trans_a, cublas_trans_b, cublas_m, cublas_n, - cublas_k, &sp_alpha, b, cuda_data_type, cublas_lda, - a, cuda_data_type, cublas_ldb, &sp_beta, c, - cuda_c_data_type, cublas_ldc, compute_type, algo)); -} - -/* - NOTE: blas gemm is column-major by default, but we need row-major output. - The data of row-major, transposed matrix is exactly the same as the - column-major, non-transposed matrix, and C = A * B ---> C^T = B^T * A^T - */ -void gemm_fp16_cublas_tensor(torch::Tensor a, torch::Tensor b, torch::Tensor c) { - if (a.sizes().size() == 1) { - assert(b.sizes().size() == 2); - a = at::unsqueeze(a, 0); - } - const auto cuda_data_type = CUDA_R_16F; - const auto cuda_c_data_type = - c.dtype() == torch::kFloat32 ? CUDA_R_32F : CUDA_R_16F; - const auto compute_type = CUDA_R_32F; - const float sp_alpha = 1.f; - // swap a and b, and use CUBLAS_OP_N. see the notes above - std::swap(a, b); - const cublasOperation_t cublas_trans_a = CUBLAS_OP_N; - const cublasOperation_t cublas_trans_b = CUBLAS_OP_N; - // m = (B^T).size(0) = B.size(1), and = A.size(1) after swap, - // negative axis is used because of the existence of batch matmul. - const int m = a.size(-1); - const int k = a.size(-2); - const int n = b.size(-2); - const int cublas_lda = m; - const int cublas_ldb = k; - const int cublas_ldc = m; - cublasHandle_t cublas_handle = get_cublas_handle(); - -#if CUDA_VERSION >= 11000 - cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT; -#else - cublasGemmAlgo_t algo = CUBLAS_GEMM_DFALT_TENSOR_OP; -#endif - const float sp_beta = 0.f; - if (a.sizes().size() == 2 && b.sizes().size() == 2) { - CUBLAS_CHECK(cublasGemmEx( - cublas_handle, cublas_trans_a, cublas_trans_b, m, n, k, &sp_alpha, - a.data_ptr(), cuda_data_type, cublas_lda, b.data_ptr(), cuda_data_type, - cublas_ldb, &sp_beta, c.data_ptr(), cuda_c_data_type, cublas_ldc, - compute_type, algo)); - } else { - // batch matmul - assert(a.sizes().size() == 3 && b.sizes().size() == 3); - - const long long int cublas_stride_a = m * k; - const long long int cublas_stride_b = k * n; - const long long int cublas_stride_c = m * n; - CUBLAS_CHECK(cublasGemmStridedBatchedEx( - cublas_handle, cublas_trans_a, cublas_trans_b, m, - n, k, &sp_alpha, a.data_ptr(), cuda_data_type, cublas_lda, - cublas_stride_a, b.data_ptr(), cuda_data_type, cublas_ldb, cublas_stride_b, - &sp_beta, c.data_ptr(), cuda_c_data_type, cublas_ldc, cublas_stride_c, - a.size(0), compute_type, algo)); - } -} diff --git a/backend-python/rwkv_pip/beta/cuda/operators.cu b/backend-python/rwkv_pip/beta/cuda/operators.cu deleted file mode 100644 index fa5a44f..0000000 --- a/backend-python/rwkv_pip/beta/cuda/operators.cu +++ /dev/null @@ -1,246 +0,0 @@ -#include -#include -#include "ATen/ATen.h" -#include -#define MIN_VALUE (-1e38) -typedef at::Half fp16; -__half *cast(fp16 *ptr) { - return reinterpret_cast<__half *>(ptr); -} - -template -__global__ void kernel_wkv_forward(const int B, const int T, const int C, - const float *__restrict__ const _w, const float *__restrict__ const _u, const F *__restrict__ const _k, const F *__restrict__ const _v, - F *__restrict__ const _y, float *__restrict__ const _aa, float *__restrict__ const _bb, float *__restrict__ const _pp) { - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - const int _b = idx / C; - const int _c = idx % C; - const int _offset = _b * T * C + _c; - const int _state_offset = _b * C + _c; - - float u = _u[_c]; - float w = _w[_c]; - const F *__restrict__ const k = _k + _offset; - const F *__restrict__ const v = _v + _offset; - F *__restrict__ const y = _y + _offset; - - float aa = _aa[_state_offset]; - float bb = _bb[_state_offset]; - float pp = _pp[_state_offset]; - for (int i = 0; i < T; i++) { - const int ii = i * C; - const float kk = float(k[ii]); - const float vv = float(v[ii]); - float ww = u + kk; - float p = max(pp, ww); - float e1 = exp(pp - p); - float e2 = exp(ww - p); - y[ii] = F((e1 * aa + e2 * vv) / (e1 * bb + e2)); - ww = w + pp; - p = max(ww, kk); - e1 = exp(ww - p); - e2 = exp(kk - p); - aa = e1 * aa + e2 * vv; - bb = e1 * bb + e2; - pp = p; - } - _aa[_state_offset] = aa; - _bb[_state_offset] = bb; - _pp[_state_offset] = pp; -} - -template -void cuda_wkv_forward(int B, int T, int C, float *w, float *u, F *k, F *v, F *y, float *aa, float *bb, float *pp) { - dim3 threadsPerBlock( min(C, 32) ); - assert(B * C % threadsPerBlock.x == 0); - dim3 numBlocks(B * C / threadsPerBlock.x); - kernel_wkv_forward<<>>(B, T, C, w, u, k, v, y, aa, bb, pp); -} - -template void cuda_wkv_forward( - int B, int T, int C, - float *w, float *u, fp16 *k, fp16 *v, fp16 *y, - float *aa, float *bb, float *pp); -template void cuda_wkv_forward( - int B, int T, int C, - float *w, float *u, float *k, float *v, float *y, - float *aa, float *bb, float *pp); - -__global__ void kernel_mm_seq_fp32i8( - const int B, const int N, const int M, - const float *__restrict__ const x, const int x_stride, - const uint8_t *__restrict__ const w, const int w_stride, - const float *__restrict__ const mx, - const float *__restrict__ const rx, - const float *__restrict__ const my, - const float *__restrict__ const ry, - float *__restrict__ const y, const int y_stride) { - - const int i = blockIdx.x * blockDim.x + threadIdx.x; - const int k = blockIdx.y * blockDim.y + threadIdx.y; - - if (i < B && k < M) { - float y_local = 0; - for (int j = 0; j < N; ++j) { - y_local += x[i * x_stride + j] * ( - (float(w[j * w_stride + k]) + 0.5f) - * rx[k] * ry[j] + mx[k] + my[j] - ); - } - y[i * y_stride + k] = y_local; - } -} - -template -void cuda_mm8_seq(int B, int N, int M, - F *x, int x_stride, - uint8_t *w, int w_stride, - F *mx, F *rx, - F *my, F *ry, - F *y, int y_stride); - -template <> -void cuda_mm8_seq(int B, int N, int M, - float *x, int x_stride, - uint8_t *w, int w_stride, - float *mx, float *rx, - float *my, float *ry, - float *y, int y_stride) { - dim3 blockSize(1, 128); - dim3 gridSize((B + blockSize.x - 1) / blockSize.x, (M + blockSize.y - 1) / blockSize.y); - kernel_mm_seq_fp32i8<<>>( - B, N, M, x, x_stride, w, w_stride, - mx, rx, my, ry, y, y_stride); -} - -__global__ void kernel_mm_seq_fp16i8( - const int B, const int N, const int M, - const __half *__restrict__ const x, const int x_stride, - const uint8_t *__restrict__ const w, const int w_stride, - const __half *__restrict__ const mx, - const __half *__restrict__ const rx, - const __half *__restrict__ const my, - const __half *__restrict__ const ry, - __half *__restrict__ const y, const int y_stride) { - - const int i = blockIdx.x * blockDim.x + threadIdx.x; - const int k = blockIdx.y * blockDim.y + threadIdx.y; - - if (i < B && k < M) { - float y_local = 0; - for (int j = 0; j < N; ++j) { - y_local += __half2float(x[i * x_stride + j]) * ( - (float(w[j * w_stride + k]) + 0.5f) - * __half2float(rx[k]) * __half2float(ry[j]) - + __half2float(mx[k]) + __half2float(my[j]) - ); - } - y[i * y_stride + k] = __float2half(y_local); - } -} - -template <> -void cuda_mm8_seq(int B, int N, int M, - fp16 *x, int x_stride, - uint8_t *w, int w_stride, - fp16 *mx, fp16 *rx, - fp16 *my, fp16 *ry, - fp16 *y, int y_stride) { - dim3 blockSize(1, 128); - dim3 gridSize((B + blockSize.x - 1) / blockSize.x, (M + blockSize.y - 1) / blockSize.y); - kernel_mm_seq_fp16i8<<>>( - B, N, M, cast(x), x_stride, w, w_stride, - cast(mx), cast(rx), cast(my), cast(ry), cast(y), y_stride); -} - -#define MM8_ONE_JSPLIT 24 -#define MM8_ONE_TILE 1024 - -__global__ void kernel_mm_one_fp32i8( - const int N, const int M, - const float *__restrict__ const x, - const uint8_t *__restrict__ const w, const int w_stride, - const float *__restrict__ const mx, - const float *__restrict__ const rx, - const float *__restrict__ const my, - const float *__restrict__ const ry, - float *__restrict__ const y) { - - const int k = blockIdx.y * blockDim.y + threadIdx.y; - const int j0 = min(N, blockIdx.x * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT)); - const int j1 = min(N, (blockIdx.x + 1) * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT)); - - if (k < M) { - float y_local = 0; - for (int j = j0; j < j1; ++j) { - y_local += x[j] * ( - (float(w[j * w_stride + k]) + 0.5f) - * rx[k] * ry[j] + mx[k] + my[j] - ); - } - atomicAdd(&y[k], y_local); - } -} - -template -void cuda_mm8_one(int N, int M, - F *x, - uint8_t *w, int w_stride, - F *mx, F *rx, - F *my, F *ry, - float *y); - -template <> -void cuda_mm8_one(int N, int M, - float *x, - uint8_t *w, int w_stride, - float *mx, float *rx, - float *my, float *ry, - float *y) { - dim3 blockSize(1, MM8_ONE_TILE); - dim3 gridSize(MM8_ONE_JSPLIT, (M + blockSize.y - 1) / blockSize.y); - kernel_mm_one_fp32i8<<>>( - N, M, x, w, w_stride, - mx, rx, my, ry, y); -} - -__global__ void kernel_mm_one_fp16i8( - const int N, const int M, - const __half *__restrict__ const x, - const uint8_t *__restrict__ const w, const int w_stride, - const __half *__restrict__ const mx, - const __half *__restrict__ const rx, - const __half *__restrict__ const my, - const __half *__restrict__ const ry, - float *__restrict__ const y) { - - const int k = blockIdx.y * blockDim.y + threadIdx.y; - const int j0 = min(N, blockIdx.x * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT)); - const int j1 = min(N, (blockIdx.x + 1) * ((N + MM8_ONE_JSPLIT - 1) / MM8_ONE_JSPLIT)); - - if (k < M) { - float y_local = 0; - for (int j = j0; j < j1; ++j) { - y_local += __half2float(x[j]) * ( - (float(w[j * w_stride + k]) + 0.5f) - * __half2float(rx[k]) * __half2float(ry[j]) - + __half2float(mx[k]) + __half2float(my[j]) - ); - } - atomicAdd(&y[k], y_local); - } -} - -template <> -void cuda_mm8_one(int N, int M, - fp16 *x, - uint8_t *w, int w_stride, - fp16 *mx, fp16 *rx, - fp16 *my, fp16 *ry, - float *y) { - dim3 blockSize(1, MM8_ONE_TILE); - dim3 gridSize(MM8_ONE_JSPLIT, (M + blockSize.y - 1) / blockSize.y); - kernel_mm_one_fp16i8<<>>( - N, M, cast(x), w, w_stride, - cast(mx), cast(rx), cast(my), cast(ry), y); -} diff --git a/backend-python/rwkv_pip/beta/cuda/util.h b/backend-python/rwkv_pip/beta/cuda/util.h deleted file mode 100644 index f00af22..0000000 --- a/backend-python/rwkv_pip/beta/cuda/util.h +++ /dev/null @@ -1,7 +0,0 @@ -#include "ATen/ATen.h" -#include - -template T *data_ptr(torch::Tensor x) { return x.data_ptr(); } -template <> inline half *data_ptr(torch::Tensor x) { - return reinterpret_cast(x.data_ptr()); -} diff --git a/backend-python/rwkv_pip/beta/cuda/wrapper.cpp b/backend-python/rwkv_pip/beta/cuda/wrapper.cpp deleted file mode 100644 index 5121499..0000000 --- a/backend-python/rwkv_pip/beta/cuda/wrapper.cpp +++ /dev/null @@ -1,181 +0,0 @@ -#include -#include "ATen/ATen.h" -#include -#include - -typedef at::Half fp16; - -template -void cuda_wkv_forward(int B, int T, int C, - float *w, float *u, F *k, F *v, F *y, - float *aa, float *bb, float *pp); -template -void cuda_mm8_seq(int B, int N, int M, - F *x, int x_stride, - uint8_t *w, int w_stride, - F *mx, F *rx, - F *my, F *ry, - F *y, int y_stride); -template -void cuda_mm8_one(int N, int M, - F *x, - uint8_t *w, int w_stride, - F *mx, F *rx, - F *my, F *ry, - float *y); - -void wkv_forward(int64_t B, int64_t T, int64_t C, - torch::Tensor &w, torch::Tensor &u, - torch::Tensor &k, torch::Tensor &v, torch::Tensor &y, - torch::Tensor &aa, torch::Tensor &bb, torch::Tensor &pp) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(w)); - switch (k.scalar_type()) { - case c10::ScalarType::Half: - cuda_wkv_forward(B, T, C, - w.data_ptr(), u.data_ptr(), - k.data_ptr(), v.data_ptr(), y.data_ptr(), - aa.data_ptr(), bb.data_ptr(), pp.data_ptr()); - break; - case c10::ScalarType::Float: - cuda_wkv_forward(B, T, C, - w.data_ptr(), u.data_ptr(), - k.data_ptr(), v.data_ptr(), y.data_ptr(), - aa.data_ptr(), bb.data_ptr(), pp.data_ptr()); - break; - default: - assert(false && "Only FP16 and FP32 are currently supported"); - } -} - -void mm8_seq(int64_t B, int64_t N, int64_t M, - torch::Tensor &x, torch::Tensor &w, - torch::Tensor &mx, torch::Tensor &rx, - torch::Tensor &my, torch::Tensor &ry, - torch::Tensor &y) { - assert(x.stride(1) == 1); - assert(w.stride(1) == 1); - assert(mx.stride(0) == 1 && rx.stride(0) == 1); - assert(my.stride(0) == 1 && ry.stride(0) == 1); - assert(y.stride(1) == 1); - const at::cuda::OptionalCUDAGuard device_guard(device_of(w)); - switch (x.scalar_type()) { - case c10::ScalarType::Half: - cuda_mm8_seq( - B, N, M, - x.data_ptr(), x.stride(0), - w.data_ptr(), w.stride(0), - mx.data_ptr(), rx.data_ptr(), - my.data_ptr(), ry.data_ptr(), - y.data_ptr(), y.stride(0)); - break; - case c10::ScalarType::Float: - cuda_mm8_seq( - B, N, M, - x.data_ptr(), x.stride(0), - w.data_ptr(), w.stride(0), - mx.data_ptr(), rx.data_ptr(), - my.data_ptr(), ry.data_ptr(), - y.data_ptr(), y.stride(0)); - break; - default: - assert(false && "Only FP16 and FP32 are currently supported"); - } -} -void mm8_one(int64_t N, int64_t M, - torch::Tensor &x, torch::Tensor &w, - torch::Tensor &mx, torch::Tensor &rx, - torch::Tensor &my, torch::Tensor &ry, - torch::Tensor &y) { - assert(x.stride(0) == 1); - assert(w.stride(1) == 1); - assert(mx.stride(0) == 1 && rx.stride(0) == 1); - assert(my.stride(0) == 1 && ry.stride(0) == 1); - assert(y.stride(0) == 1); - const at::cuda::OptionalCUDAGuard device_guard(device_of(w)); - switch (x.scalar_type()) { - case c10::ScalarType::Half: - cuda_mm8_one( - N, M, - x.data_ptr(), - w.data_ptr(), w.stride(0), - mx.data_ptr(), rx.data_ptr(), - my.data_ptr(), ry.data_ptr(), - y.data_ptr()); - break; - case c10::ScalarType::Float: - cuda_mm8_one( - N, M, - x.data_ptr(), - w.data_ptr(), w.stride(0), - mx.data_ptr(), rx.data_ptr(), - my.data_ptr(), ry.data_ptr(), - y.data_ptr()); - break; - default: - assert(false && "Only FP16 and FP32 are currently supported"); - } -} - -using torch::Tensor; - -#ifndef DISABLE_CUBLAS_GEMM -void gemm_fp16_cublas_tensor(Tensor a, Tensor b, Tensor c); -#endif - -Tensor att_one(Tensor x, Tensor ln_w, Tensor ln_b, Tensor sx, Tensor k_mix, - Tensor v_mix, Tensor r_mix, Tensor kw, - /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx, Tensor rw, - /* imm */ Tensor rx, Tensor ow, Tensor t_first, - /* imm */ Tensor k, Tensor pp, Tensor ww, Tensor aa, Tensor bb, - Tensor t_decay, /* imm */ Tensor v, /* in & out */ Tensor r, - /* out */ Tensor x_plus_out, /* out */ Tensor t1, - /* out */ Tensor t2, /* out */ Tensor p); - -Tensor att_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, - Tensor v_mix, Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, - Tensor ow, Tensor t_first, Tensor pp, Tensor aa, Tensor bb, - Tensor t_decay, /* imm */ Tensor buf, /* out */ Tensor x_plus_out); - -Tensor att_one_v5(Tensor x, Tensor sx, Tensor s, Tensor ln_w, Tensor ln_b, - Tensor lx_w, Tensor lx_b, Tensor k_mix, Tensor v_mix, - Tensor r_mix, Tensor kw, - /* imm */ Tensor kx, Tensor vw, /* imm */ Tensor vx, - Tensor rw, - /* imm */ Tensor rx, Tensor ow, Tensor t_first, - /* imm */ Tensor k, Tensor t_decay, /* imm */ Tensor v, - /* imm */ Tensor r, /* imm */ Tensor s1, - /* out */ Tensor x_plus_out, /* out */ Tensor s2); - -Tensor ffn_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, - Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, - /* imm */ Tensor buf, - /* out */ Tensor x_plus_out); - -Tensor ffn_one(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, - Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, - /* imm */ Tensor buf, - /* out */ Tensor x_plus_out); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("wkv_forward", &wkv_forward, "wkv forward"); - m.def("mm8_seq", &mm8_seq, "mm8 seq"); - m.def("mm8_one", &mm8_one, "mm8 one"); - m.def("gemm_fp16_cublas", &gemm_fp16_cublas_tensor, "gemv fp16 cublas"); - m.def("att_one", &att_one, "att one"); - m.def("att_one_v5", &att_one_v5, "att one v5"); - m.def("att_seq", &att_seq, "att seq"); - m.def("ffn_seq", &ffn_seq, "ffn seq"); - m.def("ffn_one", &ffn_one, "ffn one"); -} - -TORCH_LIBRARY(rwkv, m) { - m.def("wkv_forward", wkv_forward); - m.def("mm8_seq", mm8_seq); - m.def("mm8_one", mm8_one); - m.def("gemm_fp16_cublas", gemm_fp16_cublas_tensor); - m.def("att_one", att_one); - m.def("att_one_v5", &att_one_v5); - m.def("att_seq", att_seq); - m.def("ffn_seq", ffn_seq); - m.def("ffn_one", ffn_one); -} diff --git a/backend-python/rwkv_pip/beta/model.py b/backend-python/rwkv_pip/beta/model.py deleted file mode 100644 index f38038f..0000000 --- a/backend-python/rwkv_pip/beta/model.py +++ /dev/null @@ -1,1821 +0,0 @@ -######################################################################################################## -# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM -######################################################################################################## - -from typing import Optional -import types, gc, os, time, re, platform -import torch -from torch.nn import functional as F - -torch.backends.cudnn.benchmark = True -torch.backends.cudnn.allow_tf32 = True -torch.backends.cuda.matmul.allow_tf32 = True -current_path = os.path.dirname(os.path.abspath(__file__)) - - -# https://zhuanlan.zhihu.com/p/612879065 -def LoadPreCompileLibrary(file): - import importlib - import os - - import torch - - # load the custom_op_library and register the custom ops - lib_dir = os.path.dirname(__file__) - if os.name == "nt": - # Register the main torchvision library location on the default DLL path - import ctypes - import sys - - kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) - with_load_library_flags = hasattr(kernel32, "AddDllDirectory") - prev_error_mode = kernel32.SetErrorMode(0x0001) - - if with_load_library_flags: - kernel32.AddDllDirectory.restype = ctypes.c_void_p - - if sys.version_info >= (3, 8): - os.add_dll_directory(lib_dir) - elif with_load_library_flags: - res = kernel32.AddDllDirectory(lib_dir) - if res is None: - err = ctypes.WinError(ctypes.get_last_error()) - err.strerror += f' Error adding "{lib_dir}" to the DLL directories.' - raise ValueError(err) - - kernel32.SetErrorMode(prev_error_mode) - - loader_details = ( - importlib.machinery.ExtensionFileLoader, - importlib.machinery.EXTENSION_SUFFIXES, - ) - - extfinder = importlib.machinery.FileFinder(lib_dir, loader_details) - ext_specs = extfinder.find_spec(file) - if ext_specs is None: - return False - - try: - torch.ops.load_library(ext_specs.origin) - except OSError as exc: - return False - return True - - -######################################################################################################## - -if os.environ.get("RWKV_JIT_ON") != "0": - os.environ["RWKV_JIT_ON"] = "1" - MyModule = torch.jit.ScriptModule - MyFunction = torch.jit.script_method - MyStatic = torch.jit.script -else: - MyModule = torch.nn.Module - - def __nop(ob): - return ob - - MyFunction = __nop - MyStatic = __nop - -if os.environ.get("RWKV_CUDA_ON") == "1": - if LoadPreCompileLibrary("wkv_cuda") is False: - from torch.utils.cpp_extension import load - - load( - name=f"wkv_cuda", - sources=[ - f"{current_path}/cuda/wrapper.cpp", - f"{current_path}/cuda/operators.cu", - f"{current_path}/cuda/gemm_fp16_cublas.cpp", - f"{current_path}/cuda/att_one.cu", - f"{current_path}/cuda/att_seq.cu", - f"{current_path}/cuda/ffn.cu", - f"{current_path}/cuda/att_one_v5.cu", - ], - verbose=True, - extra_ldflags=["cublas.lib" if os.name == "nt" else ""], - extra_cuda_cflags=[ - "-t 4", - "-std=c++17", - "--use_fast_math", - "-O3", - "--extra-device-vectorization", - ], - is_python_module=False, - ) - - @MyStatic - def cuda_wkv(T: int, C: int, w, u, k, v, aa, bb, pp): - assert 1 * C % min(C, 32) == 0 - assert ( - k.dtype == v.dtype == torch.float16 or k.dtype == v.dtype == torch.float32 - ) - assert w.dtype == u.dtype == aa.dtype == bb.dtype == pp.dtype == torch.float32 - w = w.contiguous() - u = u.contiguous() - k = k.contiguous() - v = v.contiguous() - y = torch.empty( - (T, C), - device=w.device, - memory_format=torch.contiguous_format, - dtype=k.dtype, - ) - torch.ops.rwkv.wkv_forward(1, T, C, w, u, k, v, y, aa, bb, pp) - return y, aa, bb, pp - - @MyStatic - def cuda_mm8_seq(B: int, N: int, M: int, x, w, mx, rx, my, ry): - assert x.dtype == mx.dtype == rx.dtype == my.dtype == ry.dtype - assert x.dtype == torch.float32 or x.dtype == torch.float16 - assert w.dtype == torch.uint8 - assert x.shape == (B, N) - assert w.shape == (N, M) - assert rx.shape == mx.shape == (M,) - assert ry.shape == my.shape == (N, 1) - y = torch.empty((B, M), device=w.device, dtype=x.dtype) - torch.ops.rwkv.mm8_seq(B, N, M, x, w, mx, rx, my, ry, y) - return y - - @MyStatic - def cuda_mm8_one(N: int, M: int, x, w, mx, rx, my, ry): - assert x.dtype == mx.dtype == rx.dtype == my.dtype == ry.dtype - assert x.dtype == torch.float32 or x.dtype == torch.float16 - assert w.dtype == torch.uint8 - assert x.shape == (N,) - assert w.shape == (N, M) - assert rx.shape == mx.shape == (M,) - assert ry.shape == my.shape == (N, 1) - y = torch.zeros((M,), device=w.device, dtype=torch.float32) - torch.ops.rwkv.mm8_one(N, M, x, w, mx, rx, my, ry, y) - return y.to(dtype=x.dtype) - -else: - os.environ["RWKV_CUDA_ON"] = "0" - -if os.environ.get("RWKV_CUDA_ON") == "1": - - @MyStatic - def gemm(a, b, output_dtype: Optional[torch.dtype] = None): - if output_dtype is None: - output_dtype = a.dtype - if a.dtype == b.dtype == torch.float16 and a.device.type == "cuda": - if len(a.shape) == 1: - assert len(b.shape) == 2 - c = torch.empty((b.shape[-1],), dtype=output_dtype, device=a.device) - a = a.unsqueeze(0) - else: - assert len(a.shape) == len(b.shape) - assert len(a.shape) == 2 or len(a.shape) == 3 - # torch.empty((*a.shape[:-1], b.shape[-1])) doesn't work with jit - if len(a.shape) == 2: - c = torch.empty( - (a.shape[0], b.shape[-1]), dtype=output_dtype, device=a.device - ) - else: - c = torch.empty( - (a.shape[0], a.shape[1], b.shape[-1]), - dtype=output_dtype, - device=a.device, - ) - torch.ops.rwkv.gemm_fp16_cublas(a, b, c) - return c - else: - return (a @ b).to(output_dtype) - -else: - - def gemm(a, b, output_dtype: Optional[torch.dtype] = None): - if output_dtype is None: - output_dtype = a.dtype - return (a @ b).to(output_dtype) - - -######################################################################################################## - - -class RWKV(MyModule): - def __init__(self, model, strategy, verbose=True, convert_and_save_and_exit=None): - super().__init__() - if verbose: - prxxx = lambda *args, **kwargs: print(*args, **kwargs) - else: - prxxx = lambda *args, **kwargs: None - - STRATEGY_REGEX = r"^(?:(?:^|->) *(?:cuda(?::[\d]+)?|cpu|mps) (?:fp(?:16|32)|bf16)(?:i8|i4|i3)?(?: \*[\d]+\+?)? *)+$" - if not re.match(STRATEGY_REGEX, strategy): - raise ValueError( - "Invalid strategy. Please read https://pypi.org/project/rwkv/" - ) - - strategy = ("->".join([x.strip() for x in strategy.split("->")])).replace( - "->", " -> " - ) - self.args = types.SimpleNamespace() - args = self.args - args.MODEL_NAME = model - args.strategy_string = strategy - - # Rescale for fp16 mode: set x = x/2 every X layer (to avoid fp16 overflow) - self.RESCALE_LAYER = 6 if "fp16" in strategy else 0 - prxxx( - f'RWKV_JIT_ON {os.environ["RWKV_JIT_ON"]} RWKV_CUDA_ON {os.environ["RWKV_CUDA_ON"]} RESCALE_LAYER {self.RESCALE_LAYER}\n' - ) - - args.MODEL_NAME = args.MODEL_NAME.strip() - if not args.MODEL_NAME.endswith(".pth"): - args.MODEL_NAME += ".pth" - prxxx(f"Loading {args.MODEL_NAME} ...") - with torch.no_grad(): - self.w = torch.load( - args.MODEL_NAME, map_location="cpu" - ) # load model to CPU first - # it is supported to load a pure meta-tensor state dict (e.g. for quick testing) - for k, v in self.w.items(): - if isinstance(v, torch.Tensor) and v.is_meta: - # torch.zeros_like(v, device='cpu') doesn't produce an all-zero tensor - # if v is a meta tensor - self.w[k] = torch.zeros(v.shape, dtype=v.dtype, device="cpu") - gc.collect() - w = self.w - - ALREADY_CONVERTED = False - if "_strategy" in w: - ALREADY_CONVERTED = True - assert ( - convert_and_save_and_exit == None - ) # you should only convert a raw model - prxxx( - f"Converted model: strategy {w['_strategy']}, version {w['_version']}\n" - ) - assert ( - w["_strategy"] == args.strategy_string - ), "model has been converted and does not match current strategy; if you are using a new strategy, re-convert the model" - assert ( - float(w["_version"]) >= 0.7 - ) # sometimes you should re-convert using latest convert_model.py - assert w["_rescale_layer"] == self.RESCALE_LAYER - del w["_strategy"] - del w["_version"] - del w["_rescale_layer"] - - args.n_embd = w["emb.weight"].shape[1] - args.n_layer = 0 - keys = list(w.keys()) - self.version = 4 - for x in keys: - layer_id = int(x.split(".")[1]) if ("blocks." in x) else 0 - args.n_layer = max(args.n_layer, layer_id + 1) - if "ln_x" in x: - self.version = 5 - if self.version == 5 and "att.time_decay" in x: - args.n_head = w[x].shape[0] - - ####################### Compute strategy - - s = [x.strip().split(" ") for x in strategy.split("->")] - plan = [0] * len(s) - stream_i = -1 - stream_count = 0 - to_allocate = args.n_layer + 1 - allocated = 0 - free_slots = 0 - for i in range(len(s)): - si = s[i] - si1 = si[1] - if si1.startswith("fp32"): - si[1] = [torch.float] - elif si1.startswith("fp16"): - si[1] = [torch.float16] - elif si1.startswith("bf16"): - si[1] = [torch.bfloat16] - if si1.endswith("i8"): - si[1] += [torch.uint8] - else: - si[1] += [si[1][0]] - if len(si) > 2: - ss = si[2] - assert ss.startswith("*") - if ss.endswith("+"): - plan[i] = int(ss[1:-1]) - stream_i = i - else: - plan[i] = int(ss[1:]) - allocated += plan[i] - if allocated >= to_allocate: - plan[i] += to_allocate - allocated - break - else: - free_slots += 1 - if stream_i < 0: - if free_slots > 0 and to_allocate > allocated: - for i in range(len(s)): - if plan[i] == 0: - plan[i] = (to_allocate - allocated) // free_slots - allocated += plan[i] - free_slots -= 1 - if to_allocate > allocated: - plan[len(s) - 1] += to_allocate - allocated - else: - if to_allocate > allocated: - stream_count = to_allocate - allocated - plan[stream_i] += stream_count - prxxx(f"Strategy: (total {args.n_layer}+1={args.n_layer+1} layers)") - for i in range(len(s)): - ss = s[i] - if i != stream_i: - prxxx( - f'* {ss[0]} {str(ss[1]).replace("torch.","")}, store {plan[i]} layers' - ) - else: - prxxx( - f'* {ss[0]} {str(ss[1]).replace("torch.","")}, store {plan[i]-stream_count} layers, stream {stream_count} layers' - ) - plan[i] += 0 if i == 0 else plan[i - 1] - self.strategy = [None] * (args.n_layer + 1) - strategy = self.strategy - for n in range(args.n_layer + 1): - for i in range(len(s)): - if n < plan[i]: - strategy[n] = types.SimpleNamespace() - strategy[n].device = s[i][0] - strategy[n].atype = s[i][1][0] - strategy[n].wtype = s[i][1][1] - strategy[n].stream = False - if i == stream_i and n >= (plan[i] - stream_count): - strategy[n].stream = True - break - prxxx( - f"{n}-{strategy[n].device}-{str(strategy[n].atype).replace('torch.','')}-{str(strategy[n].wtype).replace('torch.','')}{'-stream' if strategy[n].stream else ''}", - end=" ", - ) - prxxx() - - ####################### Load weights to self.w - - if not ALREADY_CONVERTED: - try: # precompute embedding - w["emb.weight"] = F.layer_norm( - w["emb.weight"], - (args.n_embd,), - weight=w["blocks.0.ln0.weight"], - bias=w["blocks.0.ln0.bias"], - ) - except: - w["emb.weight"] = F.layer_norm( - w["emb.weight"].float(), - (args.n_embd,), - weight=w["blocks.0.ln0.weight"].float(), - bias=w["blocks.0.ln0.bias"].float(), - ) - del w["blocks.0.ln0.weight"] - del w["blocks.0.ln0.bias"] - - print_need_newline = False - - REAL_TIME_FIRST = False - for x in list(w.keys()): - if ".time_faaaa" in x: - REAL_TIME_FIRST = True - if REAL_TIME_FIRST: - w = { - k.replace(".time_faaaa", ".time_first") - if ".time_faaaa" in k - else k: v - for k, v in w.items() - } - self.w = w - - keys = list(w.keys()) - for x in keys: - w[x].requires_grad = False - layer_id = int(x.split(".")[1]) if ("blocks." in x) else 0 - if ("ln_out." in x) or ("head." in x): - layer_id = args.n_layer - dd = strategy[layer_id] - DEVICE = dd.device - ATYPE = dd.atype - WTYPE = dd.wtype - - if not ALREADY_CONVERTED: - if self.RESCALE_LAYER > 0: - if "att.output.weight" in x: - w[x] = w[x] / (2 ** int(layer_id // self.RESCALE_LAYER)) - if "ffn.value.weight" in x: - w[x] = w[x] / (2 ** int(layer_id // self.RESCALE_LAYER)) - - if ".time_" in x: - w[x] = w[x].squeeze() - if ( - "key.weight" in x - or "value.weight" in x - or "receptance.weight" in x - or "output.weight" in x - or "head.weight" in x - ): - w[x] = w[x].t() - - if ".time_decay" in x: # need fp32 for this - if self.version == 4: - w[x] = -torch.exp(w[x].float()) - elif self.version == 5: - w[x] = torch.exp(-torch.exp(w[x].float())).reshape(-1, 1, 1) - elif ".time_first" in x: # need fp32 for this - if self.version == 4: - w[x] = w[x].float() - elif self.version == 5: - if REAL_TIME_FIRST: - w[x] = w[x].float().reshape(-1, 1, 1) - else: - w[x] = torch.exp(w[x].float()).reshape(-1, 1, 1) - elif ".ln_x" in x: # need fp32 for group_norm - w[x] = w[x].float() - else: - if (len(w[x].shape) == 2) and ("emb" not in x): - if WTYPE != torch.uint8: - w[x] = w[x].to(dtype=WTYPE) - else: - w[x] = w[x].float() - - if w[x].shape[0] > w[x].shape[1]: - w[x + "_my"] = torch.amin(w[x], dim=1).unsqueeze(1) - w[x] = w[x] - w[x + "_my"] - w[x + "_mx"] = torch.amin(w[x], dim=0) - w[x] = w[x] - w[x + "_mx"] - w[x + "_rx"] = torch.amax(w[x], dim=0) - w[x] = w[x] / w[x + "_rx"] - w[x + "_ry"] = torch.amax(w[x], dim=1).unsqueeze(1) - w[x] = w[x] / w[x + "_ry"] - else: - w[x + "_mx"] = torch.amin(w[x], dim=0) - w[x] = w[x] - w[x + "_mx"] - w[x + "_my"] = torch.amin(w[x], dim=1).unsqueeze(1) - w[x] = w[x] - w[x + "_my"] - w[x + "_rx"] = torch.amax(w[x], dim=0) - w[x] = w[x] / w[x + "_rx"] - w[x + "_ry"] = torch.amax(w[x], dim=1).unsqueeze(1) - w[x] = w[x] / w[x + "_ry"] - - w[x] = torch.clip( - torch.floor(w[x] * 256), min=0, max=255 - ).to(dtype=torch.uint8) - w[x + "_mx"] = w[x + "_mx"].to(dtype=ATYPE).contiguous() - w[x + "_rx"] = ( - (w[x + "_rx"] / 16).to(dtype=ATYPE).contiguous() - ) - w[x + "_my"] = w[x + "_my"].to(dtype=ATYPE).contiguous() - w[x + "_ry"] = ( - (w[x + "_ry"] / 16).to(dtype=ATYPE).contiguous() - ) - else: - w[x] = w[x].to(dtype=ATYPE) - - if convert_and_save_and_exit == None: - if "emb." in x: - w[x] = w[x].contiguous() - elif (dd.stream) and ( - x.endswith("key.weight") - or x.endswith("value.weight") - or x.endswith("receptance.weight") - or x.endswith("output.weight") - ): - try: - w[x] = ( - w[x].contiguous().pin_memory() - ) # if you see "CUDA error: out of memory" here, that's out of CPU RAM, not VRAM. Get more RAM :) - except: - print( - "Note: You are running out of RAM. Get more CPU RAM. Now this will run much slower." - ) - elif DEVICE != "cpu": - w[x] = w[x].to(device=DEVICE).contiguous() - - if (dd.stream) or (DEVICE != "cpu"): - try: - w[x + "_mx"] = w[x + "_mx"].to(device=DEVICE).contiguous() - w[x + "_rx"] = w[x + "_rx"].to(device=DEVICE).contiguous() - w[x + "_my"] = w[x + "_my"].to(device=DEVICE).contiguous() - w[x + "_ry"] = w[x + "_ry"].to(device=DEVICE).contiguous() - except: - pass - - if "ffn.value.weight" in x: - gc.collect() - if "cuda" in args.strategy_string: - torch.cuda.empty_cache() - - shape = [i for i in w[x].shape if i != 1] - if len(shape) > 1: - shape = f" {str(shape[0]).rjust(5)} {str(shape[1]).rjust(5)}" - else: - shape = f" {str(shape[0]).rjust(5)} " - if layer_id == 0 or layer_id >= args.n_layer - 1: - if print_need_newline: - prxxx("\n", end="") - print_need_newline = False - dt = str(w[x].dtype).replace("torch.", "") - dt = ( - dt.replace("float32", "f32") - .replace("bfloat16", "bf16") - .replace("float16", "f16") - .replace("uint8", "i8") - ) - prxxx( - x.ljust(32), - dt.rjust(4), - str(w[x].device).rjust(8), - shape, - " (pinned)" if w[x].is_pinned() else "", - ) - else: - print_need_newline = True - prxxx(".", end="", flush=True) - - if convert_and_save_and_exit: - w["_strategy"] = args.strategy_string - w["_rescale_layer"] = self.RESCALE_LAYER - w["_version"] = "0.7" - if not convert_and_save_and_exit.endswith(".pth"): - convert_and_save_and_exit += ".pth" - prxxx(f"Saving to {convert_and_save_and_exit}...") - torch.save(w, convert_and_save_and_exit) - prxxx(f"Converted and saved. Now this will exit.") - exit(0) - - gc.collect() - if "cuda" in args.strategy_string: - torch.cuda.empty_cache() - - @MyFunction - def torch_mm8_seq(self, x, w, mx, rx, my, ry): - return x @ ((w.to(dtype=x.dtype) + 0.5) * ry * rx + my + mx) - - @MyFunction - def torch_mm8_one(self, x, w, mx, rx, my, ry): - return x @ ((w.to(dtype=x.dtype) + 0.5) * ry * rx + my + mx) - - if os.environ.get("RWKV_CUDA_ON") == "1": - - @MyFunction - def mm8_seq(self, x, w, mx, rx, my, ry): - if w.device.type == "cuda" and x.dtype == torch.float16: - B, N, M = x.shape[0], w.shape[0], w.shape[1] - return cuda_mm8_seq(B, N, M, x, w, mx, rx, my, ry) - else: - return self.torch_mm8_seq(x, w, mx, rx, my, ry) - - @MyFunction - def mm8_one(self, x, w, mx, rx, my, ry): - if w.device.type == "cuda": - N, M = w.shape[0], w.shape[1] - return cuda_mm8_one(N, M, x, w, mx, rx, my, ry) - else: - return self.torch_mm8_one(x, w, mx, rx, my, ry) - - else: - - @MyFunction - def mm8_seq(self, x, w, mx, rx, my, ry): - return self.torch_mm8_seq(x, w, mx, rx, my, ry) - - @MyFunction - def mm8_one(self, x, w, mx, rx, my, ry): - return self.torch_mm8_one(x, w, mx, rx, my, ry) - - ######################################################################################################## - - @MyFunction - def ffn_one( - self, - x, - sx, - ln_w, - ln_b, - k_mix, - r_mix, - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - kx = xx * k_mix + sx * (1 - k_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - vx = torch.square(torch.relu(gemm(kx, kw))) - out = r * gemm(vx, vw) - return x + out, xx - - @MyFunction - def ffn_one_i8( - self, - x, - sx, - ln_w, - ln_b, - k_mix, - r_mix, - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - kx = xx * k_mix + sx * (1 - k_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(self.mm8_one(rx, rw, rmx, rrx, rmy, rry)) - vx = torch.square(torch.relu(self.mm8_one(kx, kw, kmx, krx, kmy, kry))) - out = r * (self.mm8_one(vx, vw, vmx, vrx, vmy, vry)) - return x + out, xx - - ######################################################################################################## - - @MyFunction - def ffn_seq( - self, - x, - sx, - ln_w, - ln_b, - k_mix, - r_mix, - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - vx = torch.square(torch.relu(gemm(kx, kw))) - out = r * gemm(vx, vw) - return x + out, xx[-1, :] - - @MyFunction - def ffn_seq_i8( - self, - x, - sx, - ln_w, - ln_b, - k_mix, - r_mix, - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(self.mm8_seq(rx, rw, rmx, rrx, rmy, rry)) - vx = torch.square(torch.relu(self.mm8_seq(kx, kw, kmx, krx, kmy, kry))) - out = r * (self.mm8_seq(vx, vw, vmx, vrx, vmy, vry)) - return x + out, xx[-1, :] - - ######################################################################################################## - - @MyFunction - def att_one( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - k = gemm(kx, kw, output_dtype=torch.float32) - v = gemm(vx, vw, output_dtype=torch.float32) - - ww = t_first + k - p = torch.maximum(pp, ww) - e1 = torch.exp(pp - p) - e2 = torch.exp(ww - p) - wkv = ((e1 * aa + e2 * v) / (e1 * bb + e2)).to(dtype=x.dtype) - ww = t_decay + pp - p = torch.maximum(ww, k) - e1 = torch.exp(ww - p) - e2 = torch.exp(k - p) - - out = gemm(r * wkv, ow) - return x + out, xx, e1 * aa + e2 * v, e1 * bb + e2, p - - @MyFunction - def att_one_i8( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(self.mm8_one(rx, rw, rmx, rrx, rmy, rry)) - k = (self.mm8_one(kx, kw, kmx, krx, kmy, kry)).float() - v = (self.mm8_one(vx, vw, vmx, vrx, vmy, vry)).float() - - ww = t_first + k - p = torch.maximum(pp, ww) - e1 = torch.exp(pp - p) - e2 = torch.exp(ww - p) - wkv = ((e1 * aa + e2 * v) / (e1 * bb + e2)).to(dtype=x.dtype) - ww = t_decay + pp - p = torch.maximum(ww, k) - e1 = torch.exp(ww - p) - e2 = torch.exp(k - p) - - out = self.mm8_one(r * wkv, ow, omx, orx, omy, ory) - return x + out, xx, e1 * aa + e2 * v, e1 * bb + e2, p - - ######################################################################################################## - - @MyFunction - def att_seq( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - k = gemm(kx, kw, output_dtype=torch.float32) - v = gemm(vx, vw, output_dtype=torch.float32) - - T = x.shape[0] - for t in range(T): - kk = k[t] - vv = v[t] - ww = t_first + kk - p = torch.maximum(pp, ww) - e1 = torch.exp(pp - p) - e2 = torch.exp(ww - p) - sx[t] = ((e1 * aa + e2 * vv) / (e1 * bb + e2)).to(dtype=x.dtype) - ww = t_decay + pp - p = torch.maximum(ww, kk) - e1 = torch.exp(ww - p) - e2 = torch.exp(kk - p) - aa = e1 * aa + e2 * vv - bb = e1 * bb + e2 - pp = p - out = gemm(r * sx, ow) - return x + out, xx[-1, :], aa, bb, pp - - @MyFunction - def att_seq_i8( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(self.mm8_seq(rx, rw, rmx, rrx, rmy, rry)) - k = self.mm8_seq(kx, kw, kmx, krx, kmy, kry).float() - v = self.mm8_seq(vx, vw, vmx, vrx, vmy, vry).float() - - T = x.shape[0] - for t in range(T): - kk = k[t] - vv = v[t] - ww = t_first + kk - p = torch.maximum(pp, ww) - e1 = torch.exp(pp - p) - e2 = torch.exp(ww - p) - sx[t] = ((e1 * aa + e2 * vv) / (e1 * bb + e2)).to(dtype=x.dtype) - ww = t_decay + pp - p = torch.maximum(ww, kk) - e1 = torch.exp(ww - p) - e2 = torch.exp(kk - p) - aa = e1 * aa + e2 * vv - bb = e1 * bb + e2 - pp = p - out = self.mm8_seq(r * sx, ow, omx, orx, omy, ory) - return x + out, xx[-1, :], aa, bb, pp - - ######################################################################################################## - - @MyFunction - def att_one_v5( - self, - x, - sx, - s, - ln_w, - ln_b, - lx_w, - lx_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - H = t_decay.shape[0] - S = x.shape[-1] // H - - r = gemm(rx, rw, output_dtype=torch.float32).view(H, 1, S) - k = gemm(kx, kw, output_dtype=torch.float32).view(H, S, 1) - v = gemm(vx, vw, output_dtype=torch.float32).view(H, 1, S) - - a = gemm(k, v) - out = r @ (t_first * a + s) - s = a + t_decay * s - - out = out.flatten() - out = F.group_norm( - out.unsqueeze(0), num_groups=H, weight=lx_w, bias=lx_b - ).squeeze(0) - out = out.to(dtype=x.dtype) - out = gemm(out, ow) - - return x + out, xx, s - - @MyFunction - def att_seq_v5( - self, - x, - sx, - s, - ln_w, - ln_b, - lx_w, - lx_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - xx = F.layer_norm(x, (x.shape[-1],), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - H = t_decay.shape[0] - S = x.shape[-1] // H - T = x.shape[0] - - w = t_decay.reshape(-1, 1) - u = t_first.reshape(-1, 1) - ws = w.pow(T).reshape(H, 1, 1) - ind = torch.arange(T - 1, -1, -1, device=w.device).unsqueeze(0).repeat(H, 1) - w = w.repeat(1, T).pow(ind) - wk = w.reshape(H, 1, T) - wb = wk.transpose(-2, -1).flip(1) - w = torch.cat([w[:, 1:], u], dim=1) - w = F.pad(w, (0, T)) - w = torch.tile(w, [T]) - w = w[:, :-T].reshape(-1, T, 2 * T - 1) - w = w[:, :, T - 1 :].reshape(H, T, T) - - r = gemm(rx, rw, output_dtype=torch.float32).view(T, H, S).transpose(0, 1) - k = ( - gemm(kx, kw, output_dtype=torch.float32) - .view(T, H, S) - .transpose(0, 1) - .transpose(-2, -1) - ) - v = gemm(vx, vw, output_dtype=torch.float32).view(T, H, S).transpose(0, 1) - - out = ((r @ k) * w) @ v + (r @ s) * wb - s = ws * s + (k * wk) @ v - - out = out.transpose(0, 1).contiguous().reshape(T, H * S) - out = F.group_norm(out, num_groups=H, weight=lx_w, bias=lx_b) - out = out.to(dtype=x.dtype) - out = gemm(out, ow) - - return x + out, xx[-1, :], s - - ######################################################################################################## - - if os.environ["RWKV_CUDA_ON"] == "1": - - @MyFunction - def cuda_att_seq_fp16( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - seq_len = x.shape[0] - kvrx_and_y_bytes = x.numel() * 2 - k_bytes = seq_len * kw.shape[1] * 4 - v_bytes = seq_len * vw.shape[1] * 4 - r_bytes = seq_len * rw.shape[1] * 2 - buf = torch.empty( - (kvrx_and_y_bytes * 4 + k_bytes + v_bytes + r_bytes,), - device=x.device, - dtype=torch.int8, - ) - x_plus_out_t = torch.empty_like(x) - xx = torch.ops.rwkv.att_seq( - x, - sx, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - kw, - vw, - rw, - ow, - t_first, - pp, - aa, - bb, - t_decay, - buf, - x_plus_out_t, - ) - - return x_plus_out_t, xx[-1, :], aa, bb, pp - - @MyFunction - def cuda_att_seq_naive( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - T, C = x.size() - xx = F.layer_norm(x, (C,), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(gemm(rx, rw)) - k = gemm(kx, kw, output_dtype=torch.float32) - v = gemm(vx, vw, output_dtype=torch.float32) - y, aa, bb, pp = cuda_wkv(T, C, t_decay, t_first, k, v, aa, bb, pp) - - out = gemm(r * y.to(x.dtype), ow) - return x + out, xx[-1, :], aa, bb, pp - - @MyFunction - def cuda_att_seq_i8( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - T, C = x.size() - xx = F.layer_norm(x, (C,), weight=ln_w, bias=ln_b) - sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) - kx = xx * k_mix + sx * (1 - k_mix) - vx = xx * v_mix + sx * (1 - v_mix) - rx = xx * r_mix + sx * (1 - r_mix) - - r = torch.sigmoid(self.mm8_seq(rx, rw, rmx, rrx, rmy, rry)) - k = self.mm8_seq(kx, kw, kmx, krx, kmy, kry) - v = self.mm8_seq(vx, vw, vmx, vrx, vmy, vry) - y, aa, bb, pp = cuda_wkv(T, C, t_decay, t_first, k, v, aa, bb, pp) - - out = self.mm8_seq(r * y, ow, omx, orx, omy, ory) - return x + out, xx[-1, :], aa, bb, pp - - @MyFunction - def cuda_ffn_seq_fp16( - self, - x, - sx, - ln_w, - ln_b, - k_mix, - r_mix, - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ): - krx_bytes = x.numel() * x.element_size() - vx_bytes = x.shape[0] * kw.shape[1] * x.element_size() - r_bytes = x.shape[0] * rw.shape[1] * x.element_size() - buf = torch.empty( - (krx_bytes * 2 + vx_bytes + r_bytes,), device=x.device, dtype=torch.int8 - ) - x_plus_out = torch.empty_like(x) - xx = torch.ops.rwkv.ffn_seq( - x, sx, ln_w, ln_b, k_mix, r_mix, kw, vw, rw, buf, x_plus_out - ) - return x_plus_out, xx[-1, :] - - @MyFunction - def cuda_att_one_fp16( - self, - x, - sx, - aa, - bb, - pp, - ln_w, - ln_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - kx = torch.empty_like(x) - vx = torch.empty_like(x) - rx = torch.empty_like(x) - - k_t = torch.empty((kw.shape[0],), dtype=torch.float32, device=x.device) - v_t = torch.empty((vw.shape[0],), dtype=torch.float32, device=x.device) - r_t = torch.empty((rw.shape[0],), dtype=torch.float16, device=x.device) - x_plus_out_t = torch.empty_like(x) - t1_t = torch.empty_like(x, dtype=torch.float32) - t2_t = torch.empty_like(x, dtype=torch.float32) - p_t = torch.empty_like(x, dtype=torch.float32) - xx = torch.ops.rwkv.att_one( - x, - ln_w, - ln_b, - sx, - k_mix, - v_mix, - r_mix, - kw, - kx, - vw, - vx, - rw, - rx, - ow, - t_first, - k_t, - pp, - ow, - aa, - bb, - t_decay, - v_t, - r_t, - x_plus_out_t, - t1_t, - t2_t, - p_t, - ) - return x_plus_out_t, xx, t1_t, t2_t, p_t - - @MyFunction - def cuda_att_one_v5_fp16( - self, - x, - sx, - s, - ln_w, - ln_b, - lx_w, - lx_b, - k_mix, - v_mix, - r_mix, - t_decay, - t_first, - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ): - kx = torch.empty_like(x) - vx = torch.empty_like(x) - rx = torch.empty_like(x) - - H = t_decay.shape[0] - S = x.shape[-1] // H - - r = torch.empty((H * S,), dtype=torch.float32, device=x.device) - k = torch.empty((H * S,), dtype=torch.float32, device=x.device) - v = torch.empty((H * S,), dtype=torch.float32, device=x.device) - s1 = torch.empty((H, S, S), dtype=torch.float32, device=x.device) - s2 = torch.empty((H, S, S), dtype=torch.float32, device=x.device) - x_plus_out = torch.empty_like(x) - - xx = torch.ops.rwkv.att_one_v5( - x, - sx, - s, - ln_w, - ln_b, - lx_w, - lx_b, - k_mix, - v_mix, - r_mix, - kw, - kx, - vw, - vx, - rw, - rx, - ow, - t_first, - k, - t_decay, - v, - r, - s1, - x_plus_out, - s2, - ) - - return x_plus_out, xx, s2 - - @MyFunction - def cuda_ffn_one_fp16( - self, - x, - sx, - ln_w, - ln_b, - k_mix, - r_mix, - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ): - krx_bytes = x.numel() * x.element_size() - vx_bytes = x.shape[0] * kw.shape[1] * x.element_size() - r_bytes = x.shape[0] * rw.shape[1] * x.element_size() - buf = torch.empty( - (krx_bytes * 2 + vx_bytes + r_bytes,), device=x.device, dtype=torch.int8 - ) - x_plus_out = torch.empty_like(x) - xx = torch.ops.rwkv.ffn_one( - x, sx, ln_w, ln_b, k_mix, r_mix, kw, vw, rw, buf, x_plus_out - ) - return x_plus_out, xx - - ######################################################################################################## - - def forward(self, tokens, state, full_output=False): - with torch.no_grad(): - w = self.w - args = self.args - - if state == None: - if self.version == 4: - state = [None] * args.n_layer * 5 - for i in range( - args.n_layer - ): # state: 0=att_xx 1=att_aa 2=att_bb 3=att_pp 4=ffn_xx - dd = self.strategy[i] - dev = dd.device - atype = dd.atype - state[i * 5 + 0] = torch.zeros( - args.n_embd, dtype=atype, requires_grad=False, device=dev - ).contiguous() - state[i * 5 + 1] = torch.zeros( - args.n_embd, - dtype=torch.float, - requires_grad=False, - device=dev, - ).contiguous() - state[i * 5 + 2] = torch.zeros( - args.n_embd, - dtype=torch.float, - requires_grad=False, - device=dev, - ).contiguous() - state[i * 5 + 3] = ( - torch.zeros( - args.n_embd, - dtype=torch.float, - requires_grad=False, - device=dev, - ).contiguous() - - 1e30 - ) - state[i * 5 + 4] = torch.zeros( - args.n_embd, dtype=atype, requires_grad=False, device=dev - ).contiguous() - elif self.version == 5: - state = [None] * args.n_layer * 3 - for i in range(args.n_layer): # state: 0=att_xx 1=att_kv 2=ffn_xx - dd = self.strategy[i] - dev = dd.device - atype = dd.atype - state[i * 3 + 0] = torch.zeros( - args.n_embd, dtype=atype, requires_grad=False, device=dev - ).contiguous() - state[i * 3 + 1] = torch.zeros( - ( - args.n_head, - args.n_embd // args.n_head, - args.n_embd // args.n_head, - ), - dtype=torch.float, - requires_grad=False, - device=dev, - ).contiguous() - state[i * 3 + 2] = torch.zeros( - args.n_embd, dtype=atype, requires_grad=False, device=dev - ).contiguous() - - seq_mode = len(tokens) > 1 - - x = w["emb.weight"][tokens if seq_mode else tokens[0]] - - for i in range(args.n_layer): - bbb = f"blocks.{i}." - att = f"blocks.{i}.att." - ffn = f"blocks.{i}.ffn." - dd = self.strategy[i] - dev = dd.device - atype = dd.atype - wtype = dd.wtype - if seq_mode: - ATT = self.att_seq if wtype != torch.uint8 else self.att_seq_i8 - FFN = self.ffn_seq if wtype != torch.uint8 else self.ffn_seq_i8 - if "cuda" in str(dev) and os.environ["RWKV_CUDA_ON"] == "1": - if wtype == torch.float16: - ATT = self.cuda_att_seq_fp16 - FFN = self.cuda_ffn_seq_fp16 - elif wtype == torch.uint8: - ATT = self.cuda_att_seq_i8 - else: - ATT = self.cuda_att_seq_naive - if self.version == 5: - ATT = self.att_seq_v5 - else: - ATT = self.att_one if wtype != torch.uint8 else self.att_one_i8 - FFN = self.ffn_one if wtype != torch.uint8 else self.ffn_one_i8 - if self.version == 5: - ATT = self.att_one_v5 - if ( - "cuda" in str(dev) - and os.environ["RWKV_CUDA_ON"] == "1" - and wtype == torch.float16 - ): - ATT = self.cuda_att_one_fp16 - FFN = self.cuda_ffn_one_fp16 - if self.version == 5: - ATT = self.cuda_att_one_v5_fp16 - - x = x.to(dtype=atype, device=dev) - - kw = w[f"{att}key.weight"] - vw = w[f"{att}value.weight"] - rw = w[f"{att}receptance.weight"] - ow = w[f"{att}output.weight"] - if dd.stream: - kw = kw.to(device=dev, non_blocking=True) - vw = vw.to(device=dev, non_blocking=True) - rw = rw.to(device=dev, non_blocking=True) - ow = ow.to(device=dev, non_blocking=True) - kmx = w[f"{att}key.weight_mx"] if wtype == torch.uint8 else x - krx = w[f"{att}key.weight_rx"] if wtype == torch.uint8 else x - kmy = w[f"{att}key.weight_my"] if wtype == torch.uint8 else x - kry = w[f"{att}key.weight_ry"] if wtype == torch.uint8 else x - vmx = w[f"{att}value.weight_mx"] if wtype == torch.uint8 else x - vrx = w[f"{att}value.weight_rx"] if wtype == torch.uint8 else x - vmy = w[f"{att}value.weight_my"] if wtype == torch.uint8 else x - vry = w[f"{att}value.weight_ry"] if wtype == torch.uint8 else x - rmx = w[f"{att}receptance.weight_mx"] if wtype == torch.uint8 else x - rrx = w[f"{att}receptance.weight_rx"] if wtype == torch.uint8 else x - rmy = w[f"{att}receptance.weight_my"] if wtype == torch.uint8 else x - rry = w[f"{att}receptance.weight_ry"] if wtype == torch.uint8 else x - omx = w[f"{att}output.weight_mx"] if wtype == torch.uint8 else x - orx = w[f"{att}output.weight_rx"] if wtype == torch.uint8 else x - omy = w[f"{att}output.weight_my"] if wtype == torch.uint8 else x - ory = w[f"{att}output.weight_ry"] if wtype == torch.uint8 else x - if self.version == 4: - ( - x, - state[i * 5 + 0], - state[i * 5 + 1], - state[i * 5 + 2], - state[i * 5 + 3], - ) = ATT( - x, - state[i * 5 + 0], - state[i * 5 + 1], - state[i * 5 + 2], - state[i * 5 + 3], - w[f"{bbb}ln1.weight"], - w[f"{bbb}ln1.bias"], - w[f"{att}time_mix_k"], - w[f"{att}time_mix_v"], - w[f"{att}time_mix_r"], - w[f"{att}time_decay"], - w[f"{att}time_first"], - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ) - elif self.version == 5: - x, state[i * 3 + 0], state[i * 3 + 1] = ATT( - x, - state[i * 3 + 0], - state[i * 3 + 1], - w[f"{bbb}ln1.weight"], - w[f"{bbb}ln1.bias"], - w[f"{att}ln_x.weight"], - w[f"{att}ln_x.bias"], - w[f"{att}time_mix_k"], - w[f"{att}time_mix_v"], - w[f"{att}time_mix_r"], - w[f"{att}time_decay"], - w[f"{att}time_first"], - kw, - vw, - rw, - ow, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - omx, - orx, - omy, - ory, - ) - if dd.stream: - del kw, vw, rw, ow - - kw = w[f"{ffn}key.weight"] - vw = w[f"{ffn}value.weight"] - rw = w[f"{ffn}receptance.weight"] - if dd.stream: - kw = kw.to(device=dev, non_blocking=True) - vw = vw.to(device=dev, non_blocking=True) - rw = rw.to(device=dev, non_blocking=True) - kmx = w[f"{ffn}key.weight_mx"] if wtype == torch.uint8 else x - krx = w[f"{ffn}key.weight_rx"] if wtype == torch.uint8 else x - kmy = w[f"{ffn}key.weight_my"] if wtype == torch.uint8 else x - kry = w[f"{ffn}key.weight_ry"] if wtype == torch.uint8 else x - vmx = w[f"{ffn}value.weight_mx"] if wtype == torch.uint8 else x - vrx = w[f"{ffn}value.weight_rx"] if wtype == torch.uint8 else x - vmy = w[f"{ffn}value.weight_my"] if wtype == torch.uint8 else x - vry = w[f"{ffn}value.weight_ry"] if wtype == torch.uint8 else x - rmx = w[f"{ffn}receptance.weight_mx"] if wtype == torch.uint8 else x - rrx = w[f"{ffn}receptance.weight_rx"] if wtype == torch.uint8 else x - rmy = w[f"{ffn}receptance.weight_my"] if wtype == torch.uint8 else x - rry = w[f"{ffn}receptance.weight_ry"] if wtype == torch.uint8 else x - if self.version == 4: - offset = i * 5 + 4 - elif self.version == 5: - offset = i * 3 + 2 - x, state[offset] = FFN( - x, - state[offset], - w[f"{bbb}ln2.weight"], - w[f"{bbb}ln2.bias"], - w[f"{ffn}time_mix_k"], - w[f"{ffn}time_mix_r"], - kw, - vw, - rw, - kmx, - krx, - kmy, - kry, - vmx, - vrx, - vmy, - vry, - rmx, - rrx, - rmy, - rry, - ) - if dd.stream: - del kw, vw, rw - - if self.RESCALE_LAYER > 0: - if (i + 1) % self.RESCALE_LAYER == 0: - x = x / 2 - - dd = self.strategy[args.n_layer] - x = x[-1, :] if (seq_mode and (not full_output)) else x - x = x.to(dtype=dd.atype, device=dd.device) - - x = F.layer_norm( - x, (args.n_embd,), weight=w["ln_out.weight"], bias=w["ln_out.bias"] - ) - if w["head.weight"].dtype != torch.uint8: - x = x @ w["head.weight"] - else: - if seq_mode and full_output: - x = self.mm8_seq( - x, - w["head.weight"], - w["head.weight_mx"], - w["head.weight_rx"], - w["head.weight_my"], - w["head.weight_ry"], - ) - else: - x = self.mm8_one( - x, - w["head.weight"], - w["head.weight_mx"], - w["head.weight_rx"], - w["head.weight_my"], - w["head.weight_ry"], - ) - - return x.float(), state diff --git a/backend-python/rwkv_pip/beta/wkv_cuda.pyd b/backend-python/rwkv_pip/beta/wkv_cuda.pyd deleted file mode 100644 index 7734f12..0000000 Binary files a/backend-python/rwkv_pip/beta/wkv_cuda.pyd and /dev/null differ diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py index 430fee4..7b1190f 100644 --- a/backend-python/utils/rwkv.py +++ b/backend-python/utils/rwkv.py @@ -617,7 +617,6 @@ def get_model_path(model_path: str) -> str: def RWKV(model: str, strategy: str, tokenizer: Union[str, None]) -> AbstractRWKV: model_path = get_model_path(model) - rwkv_beta = global_var.get(global_var.Args).rwkv_beta rwkv_cpp = getattr(global_var.get(global_var.Args), "rwkv.cpp") webgpu = global_var.get(global_var.Args).webgpu @@ -625,12 +624,7 @@ def RWKV(model: str, strategy: str, tokenizer: Union[str, None]) -> AbstractRWKV os.environ["RWKV_RESCALE_LAYER"] = "999" # dynamic import to make RWKV_CUDA_ON work - if rwkv_beta: - print("Using rwkv-beta") - from rwkv_pip.beta.model import ( - RWKV as Model, - ) - elif rwkv_cpp: + if rwkv_cpp: print("Using rwkv.cpp, strategy is ignored") from rwkv_pip.cpp.model import ( RWKV as Model,