diff --git a/backend-python/requirements.txt b/backend-python/requirements.txt index f570558..9075ce7 100644 Binary files a/backend-python/requirements.txt and b/backend-python/requirements.txt differ diff --git a/backend-python/requirements_without_cyac.txt b/backend-python/requirements_without_cyac.txt index 5950fac..b3f8834 100644 Binary files a/backend-python/requirements_without_cyac.txt and b/backend-python/requirements_without_cyac.txt differ diff --git a/backend-python/rwkv_pip/cuda/att_one.cu b/backend-python/rwkv_pip/cuda/att_one.cu new file mode 100644 index 0000000..f22858d --- /dev/null +++ b/backend-python/rwkv_pip/cuda/att_one.cu @@ -0,0 +1,124 @@ +#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 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(kx, kw, k); + gemm_fp16_cublas(vx, vw, v); + gemm_fp16_cublas(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(r, ow, x_plus_out); + x_plus_out += x; + return xx; +} diff --git a/backend-python/rwkv_pip/cuda/att_seq.cu b/backend-python/rwkv_pip/cuda/att_seq.cu new file mode 100644 index 0000000..4d506a3 --- /dev/null +++ b/backend-python/rwkv_pip/cuda/att_seq.cu @@ -0,0 +1,179 @@ +#include "ATen/ATen.h" +#include +#include +#include + +#include "util.h" +#include "element_wise.h" + +using torch::Tensor; + +void gemm_fp16_cublas(Tensor a, Tensor b, Tensor c); +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/cuda/element_wise.h b/backend-python/rwkv_pip/cuda/element_wise.h new file mode 100644 index 0000000..eedc2f9 --- /dev/null +++ b/backend-python/rwkv_pip/cuda/element_wise.h @@ -0,0 +1,21 @@ +#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/cuda/ffn.cu b/backend-python/rwkv_pip/cuda/ffn.cu new file mode 100644 index 0000000..c1c2c80 --- /dev/null +++ b/backend-python/rwkv_pip/cuda/ffn.cu @@ -0,0 +1,165 @@ +#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/cuda/gemm_fp16_cublas.cpp b/backend-python/rwkv_pip/cuda/gemm_fp16_cublas.cpp new file mode 100644 index 0000000..db48fcf --- /dev/null +++ b/backend-python/rwkv_pip/cuda/gemm_fp16_cublas.cpp @@ -0,0 +1,86 @@ +#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(torch::Tensor a, torch::Tensor b, torch::Tensor c) { + 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/cuda/operators.cu b/backend-python/rwkv_pip/cuda/operators.cu new file mode 100644 index 0000000..fa5a44f --- /dev/null +++ b/backend-python/rwkv_pip/cuda/operators.cu @@ -0,0 +1,246 @@ +#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/cuda/rwkv5.cu b/backend-python/rwkv_pip/cuda/rwkv5.cu new file mode 100644 index 0000000..a7f13c3 --- /dev/null +++ b/backend-python/rwkv_pip/cuda/rwkv5.cu @@ -0,0 +1,88 @@ +#include +#include +#include "ATen/ATen.h" +typedef at::BFloat16 bf16; +typedef at::Half fp16; +typedef float fp32; + +template +__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 _k, const F *__restrict__ const _v, const float *__restrict__ _w, const F *__restrict__ _u, + F *__restrict__ const _y) +{ + const int b = blockIdx.x / H; + const int h = blockIdx.x % H; + const int i = threadIdx.x; + _w += h*_N_; + _u += h*_N_; + _state += h*_N_*_N_ + i*_N_; // wrong if B > 1 !!! + + __shared__ float r[_N_], k[_N_], u[_N_], w[_N_]; + + float state[_N_]; + #pragma unroll + for (int j = 0; j < _N_; j++) + state[j] = _state[j]; + + __syncthreads(); + u[i] = float(_u[i]); + w[i] = _w[i]; + __syncthreads(); + + for (int t = b*T*C + h*_N_ + i; t < (b+1)*T*C + h*_N_ + i; t += C) + { + __syncthreads(); + r[i] = float(_r[t]); + k[i] = float(_k[t]); + __syncthreads(); + + const float v = float(_v[t]); + float y = 0; + + #pragma unroll + for (int j = 0; j < _N_; j+=4) + { + const float4& r_ = (float4&)(r[j]); + const float4& k_ = (float4&)(k[j]); + const float4& w_ = (float4&)(w[j]); + const float4& u_ = (float4&)(u[j]); + float4& s = (float4&)(state[j]); + float4 x; + + x.x = k_.x * v; + x.y = k_.y * v; + x.z = k_.z * v; + x.w = k_.w * v; + + y += r_.x * (u_.x * x.x + s.x); + y += r_.y * (u_.y * x.y + s.y); + y += r_.z * (u_.z * x.z + s.z); + y += r_.w * (u_.w * x.w + s.w); + + s.x = s.x * w_.x + x.x; + s.y = s.y * w_.y + x.y; + s.z = s.z * w_.z + x.z; + s.w = s.w * w_.w + x.w; + } + _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 *k, bf16 *v, float *w, bf16 *u, bf16 *y) +{ + assert(H*_N_ == C); + kernel_forward<<>>(B, T, C, H, state, r, k, v, w, u, y); +} +void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *k, fp16 *v, float *w, fp16 *u, fp16 *y) +{ + assert(H*_N_ == C); + kernel_forward<<>>(B, T, C, H, state, r, k, v, w, u, y); +} +void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *k, fp32 *v, float *w, fp32 *u, fp32 *y) +{ + assert(H*_N_ == C); + kernel_forward<<>>(B, T, C, H, state, r, k, v, w, u, y); +} diff --git a/backend-python/rwkv_pip/cuda/rwkv5_op.cpp b/backend-python/rwkv_pip/cuda/rwkv5_op.cpp new file mode 100644 index 0000000..5471bcf --- /dev/null +++ b/backend-python/rwkv_pip/cuda/rwkv5_op.cpp @@ -0,0 +1,30 @@ +#include +#include "ATen/ATen.h" +typedef at::BFloat16 bf16; +typedef at::Half fp16; +typedef float fp32; + +void cuda_forward_bf16(int B, int T, int C, int H, float *state, bf16 *r, bf16 *k, bf16 *v, float *w, bf16 *u, bf16 *y); +void cuda_forward_fp16(int B, int T, int C, int H, float *state, fp16 *r, fp16 *k, fp16 *v, float *w, fp16 *u, fp16 *y); +void cuda_forward_fp32(int B, int T, int C, int H, float *state, fp32 *r, fp32 *k, fp32 *v, float *w, fp32 *u, 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 &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) { + cuda_forward_bf16(B, T, C, H, state.data_ptr(), r.data_ptr(), k.data_ptr(), v.data_ptr(), w.data_ptr(), u.data_ptr(), y.data_ptr()); +} +void forward_fp16(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) { + cuda_forward_fp16(B, T, C, H, state.data_ptr(), r.data_ptr(), k.data_ptr(), v.data_ptr(), w.data_ptr(), u.data_ptr(), y.data_ptr()); +} +void forward_fp32(int64_t B, int64_t T, int64_t C, int64_t H, torch::Tensor &state, torch::Tensor &r, torch::Tensor &k, torch::Tensor &v, torch::Tensor &w, torch::Tensor &u, torch::Tensor &y) { + cuda_forward_fp32(B, T, C, H, state.data_ptr(), r.data_ptr(), k.data_ptr(), v.data_ptr(), w.data_ptr(), u.data_ptr(), y.data_ptr()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward_bf16", &forward_bf16, "rwkv5 forward_bf16"); + m.def("forward_fp16", &forward_fp16, "rwkv5 forward_fp16"); + m.def("forward_fp32", &forward_fp32, "rwkv5 forward_fp32"); +} +TORCH_LIBRARY(rwkv5, m) { + m.def("forward_bf16", forward_bf16); + m.def("forward_fp16", forward_fp16); + m.def("forward_fp32", forward_fp32); +} diff --git a/backend-python/rwkv_pip/cuda/util.h b/backend-python/rwkv_pip/cuda/util.h new file mode 100644 index 0000000..f00af22 --- /dev/null +++ b/backend-python/rwkv_pip/cuda/util.h @@ -0,0 +1,7 @@ +#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/cuda/wrapper.cpp b/backend-python/rwkv_pip/cuda/wrapper.cpp new file mode 100644 index 0000000..5e91eb5 --- /dev/null +++ b/backend-python/rwkv_pip/cuda/wrapper.cpp @@ -0,0 +1,141 @@ +#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 a, Tensor b, Tensor c); +#endif + +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"); +#ifndef DISABLE_CUBLAS_GEMM + m.def("gemm_fp16_cublas", &gemm_fp16_cublas, "gemv fp16 cublas"); +#endif +} + +TORCH_LIBRARY(rwkv, m) { + m.def("wkv_forward", wkv_forward); + m.def("mm8_seq", mm8_seq); + m.def("mm8_one", mm8_one); +#ifndef DISABLE_CUBLAS_GEMM + m.def("gemm_fp16_cublas", gemm_fp16_cublas); +#endif +} diff --git a/backend-python/rwkv_pip/model.py b/backend-python/rwkv_pip/model.py new file mode 100644 index 0000000..75c543e --- /dev/null +++ b/backend-python/rwkv_pip/model.py @@ -0,0 +1,1827 @@ +######################################################################################################## +# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM +######################################################################################################## + +from typing import Optional +import types, gc, os, time, re +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__)) + +######################################################################################################## + +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": + from torch.utils.cpp_extension import load + + try: + 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", + ], + verbose=True, + extra_cuda_cflags=[ + "--use_fast_math", + "-O3", + "--extra-device-vectorization", + ], + is_python_module=False, + ) + DISABLE_CUBLAS_GEMM = False + except: + print( + "Failed to build cuBLAS matmul, falling back to torch.matmul. Small model with fp16 will overflow." + ) + load( + name=f"wkv_cuda", + sources=[ + f"{current_path}/cuda/wrapper.cpp", + f"{current_path}/cuda/operators.cu", + ], + verbose=True, + extra_cuda_cflags=[ + "--use_fast_math", + "-O3", + "--extra-device-vectorization", + ], + extra_cflags=["-DDISABLE_CUBLAS_GEMM"], + is_python_module=False, + ) + DISABLE_CUBLAS_GEMM = True + + @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" and not DISABLE_CUBLAS_GEMM: + + @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) + try: + self.RESCALE_LAYER = int( + os.environ["RWKV_RESCALE_LAYER"] + ) # !!! NOTE: SEEMS YOU SHOULD SET IT TO 999 (disable) FOR RWKV-MUSIC MODELS !!! + except: + 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 + 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 + ) # 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 + ) # must use same RESCALE_LAYER to avoid mistakes + del w["_strategy"] + del w["_version"] + del w["_rescale_layer"] + + args.n_embd = w["emb.weight"].shape[1] + args.n_att = w["blocks.0.att.key.weight"].shape[ + 0 + ] # note: transposed matrix + args.n_ffn = w["blocks.0.ffn.key.weight"].shape[ + 0 + ] # note: transposed matrix + 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 = max(5, self.version) + if "gate.weight" in x: + self.version = max(5.1, self.version) + if int(self.version) == 5 and "att.time_decay" in x: + args.n_head = w[x].shape[0] + if len(w[x].shape) > 1: + if w[x].shape[1] > 1: + self.version = max(5.2, self.version) + + ####################### 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 "gate.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 int(self.version) == 5: + w[x] = torch.exp(-torch.exp(w[x].float())).reshape(-1, 1, 1) + if self.version == 5.2: + w[x] = w[x].reshape(args.n_head, -1, 1) + elif ".time_first" in x: # need fp32 for this + if self.version == 4: + w[x] = w[x].float() + elif int(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) + if self.version == 5.2: + w[x] = w[x].reshape(args.n_head, -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) + + if self.version == 5.2: + assert ( + os.environ["RWKV_CUDA_ON"] == "1" + ) # latest RWKV-5 requires os.environ["RWKV_CUDA_ON"] == '1' (will fix soon) + HEAD_SIZE = args.n_att // args.n_head + rwkv5 = load( + name="rwkv5", + sources=[ + f"{current_path}/cuda/rwkv5_op.cpp", + f"{current_path}/cuda/rwkv5.cu", + ], + verbose=True, + extra_cuda_cflags=[ + "-res-usage", + "--use_fast_math", + "-O3", + "-Xptxas -O3", + "--extra-device-vectorization", + f"-D_N_={HEAD_SIZE}", + ], + ) + + class RWKV_5(torch.autograd.Function): + @staticmethod + def forward(ctx, B, T, C, H, state, r, k, v, w, u): + with torch.no_grad(): + assert HEAD_SIZE == C // H + ctx.B = B + ctx.T = T + ctx.C = C + ctx.H = H + assert state.dtype == torch.float32 + assert w.dtype == torch.float32 + assert r.is_contiguous() + assert k.is_contiguous() + assert v.is_contiguous() + assert w.is_contiguous() + assert u.is_contiguous() + assert state.is_contiguous() + + y = torch.empty( + (B, T, C), + device=w.device, + dtype=r.dtype, + memory_format=torch.contiguous_format, + ) + if r.dtype == torch.bfloat16: + rwkv5.forward_bf16(B, T, C, H, state, r, k, v, w, u, y) + elif r.dtype == torch.float16: + rwkv5.forward_fp16(B, T, C, H, state, r, k, v, w, u, y) + elif r.dtype == torch.float32: + rwkv5.forward_fp32(B, T, C, H, state, r, k, v, w, u, y) + return y, state + + self.RWKV_5 = RWKV_5 + + gc.collect() + if "cuda" in args.strategy_string: + torch.cuda.empty_cache() + + def RUN_RWKV_5(self, B, T, C, H, state, r, k, v, w, u): + return self.RWKV_5.apply(B, T, C, H, state, r, k, v, w, u) + + @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 + + ######################################################################################################## + + @MyFunction + def att_one_v5_1( + self, + x, + sx, + s, + ln_w, + ln_b, + lx_w, + lx_b, + k_mix, + v_mix, + r_mix, + g_mix, + t_decay, + t_first, + kw, + vw, + rw, + gw, + 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) + gx = xx * g_mix + sx * (1 - g_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) + g = F.silu(gemm(gx, gw)) + + 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) * g + out = gemm(out, ow) + + return x + out, xx, s + + @MyFunction + def att_seq_v5_1( + self, + x, + sx, + s, + ln_w, + ln_b, + lx_w, + lx_b, + k_mix, + v_mix, + r_mix, + g_mix, + t_decay, + t_first, + kw, + vw, + rw, + gw, + 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) + gx = xx * g_mix + sx * (1 - g_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) + g = F.silu(gemm(gx, gw)) + + 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) * g + out = gemm(out, ow) + + return x + out, xx[-1, :], s + + ######################################################################################################## + + def att_seq_v5_2( + self, + x, + sx, + s, + ln_w, + ln_b, + lx_w, + lx_b, + k_mix, + v_mix, + r_mix, + g_mix, + t_decay, + t_first, + kw, + vw, + rw, + gw, + 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) + gx = xx * g_mix + sx * (1 - g_mix) + + H = t_decay.shape[0] + N = x.shape[-1] // H + T = x.shape[0] + + r = gemm(rx, rw, output_dtype=torch.float32) + k = gemm(kx, kw, output_dtype=torch.float32) + v = gemm(vx, vw, output_dtype=torch.float32) + g = F.silu(gemm(gx, gw)) + + out, s = self.RUN_RWKV_5( + 1, + T, + self.args.n_att, + H, + s.transpose(-1, -2).contiguous(), + r, + k, + v, + w=t_decay, + u=t_first, + ) + s = s.transpose(-1, -2) + + out = out.reshape(T, H * N) + out = F.group_norm(out, num_groups=H, weight=lx_w, bias=lx_b) + out = out.to(dtype=x.dtype) * g + out = gemm(out, ow) + + return x + out, xx[-1, :], s + + ######################################################################################################## + + if os.environ["RWKV_CUDA_ON"] == "1": + + @MyFunction + def cuda_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, + ): + T, C = x.shape + 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, aa.shape[0], 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.shape + 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 + + ######################################################################################################## + + 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_att, + dtype=torch.float, + requires_grad=False, + device=dev, + ).contiguous() + state[i * 5 + 2] = torch.zeros( + args.n_att, + dtype=torch.float, + requires_grad=False, + device=dev, + ).contiguous() + state[i * 5 + 3] = ( + torch.zeros( + args.n_att, + 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 int(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_att // args.n_head, + args.n_att // 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: + if "cuda" in str(dev) and os.environ["RWKV_CUDA_ON"] == "1": + ATT = ( + self.cuda_att_seq + if wtype != torch.uint8 + else self.cuda_att_seq_i8 + ) + else: + ATT = self.att_seq if wtype != torch.uint8 else self.att_seq_i8 + if self.version == 5: + ATT = self.att_seq_v5 + elif self.version == 5.1: + ATT = self.att_seq_v5_1 + elif self.version == 5.2: + ATT = self.att_seq_v5_2 + FFN = self.ffn_seq if wtype != torch.uint8 else self.ffn_seq_i8 + else: + ATT = self.att_one if wtype != torch.uint8 else self.att_one_i8 + if self.version == 5: + ATT = self.att_one_v5 + elif self.version == 5.1: + ATT = self.att_one_v5_1 + elif self.version == 5.2: + ATT = self.att_one_v5_1 # same as v5.1 + FFN = self.ffn_one if wtype != torch.uint8 else self.ffn_one_i8 + + 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 == 5.1 or self.version == 5.2: + gw = w[f"{att}gate.weight"] + if dd.stream: + gw = gw.to(device=dev, non_blocking=True) + gmx = w[f"{att}gate.weight_mx"] if wtype == torch.uint8 else x + grx = w[f"{att}gate.weight_rx"] if wtype == torch.uint8 else x + gmy = w[f"{att}gate.weight_my"] if wtype == torch.uint8 else x + gry = w[f"{att}gate.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, + ) + elif self.version == 5.1 or self.version == 5.2: + 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_mix_g"], + w[f"{att}time_decay"], + w[f"{att}time_first"], + kw, + vw, + rw, + gw, + 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 int(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/utils.py b/backend-python/rwkv_pip/utils.py index 0429fa7..8da5741 100644 --- a/backend-python/rwkv_pip/utils.py +++ b/backend-python/rwkv_pip/utils.py @@ -16,6 +16,7 @@ class PIPELINE_ARGS: top_k=0, alpha_frequency=0.2, alpha_presence=0.2, + alpha_decay=0.996, token_ban=[], token_stop=[], chunk_len=256, @@ -25,6 +26,7 @@ class PIPELINE_ARGS: self.top_k = top_k self.alpha_frequency = alpha_frequency # Frequency Penalty (as in GPT-3) self.alpha_presence = alpha_presence # Presence Penalty (as in GPT-3) + self.alpha_decay = alpha_decay # gradually decay the penalty self.token_ban = token_ban # ban the generation of some tokens self.token_stop = token_stop # stop generation whenever you see any token here self.chunk_len = ( @@ -84,7 +86,7 @@ class PIPELINE: sorted_ids = np.argsort(probs) sorted_probs = probs[sorted_ids][::-1] cumulative_probs = np.cumsum(sorted_probs) - cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)]) + cutoff = float(sorted_probs[np.argmax(cumulative_probs >= top_p)]) probs[probs < cutoff] = 0 if top_k < len(probs) and top_k > 0: probs[sorted_ids[:-top_k]] = 0 @@ -98,7 +100,7 @@ class PIPELINE: sorted_probs = probs[sorted_ids] sorted_probs = torch.flip(sorted_probs, dims=(0,)) cumulative_probs = torch.cumsum(sorted_probs, dim=-1).cpu().numpy() - cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)]) + cutoff = float(sorted_probs[np.argmax(cumulative_probs >= top_p)]) probs[probs < cutoff] = 0 if top_k < len(probs) and top_k > 0: probs[sorted_ids[:-top_k]] = 0 @@ -133,10 +135,13 @@ class PIPELINE: if token in args.token_stop: break all_tokens += [token] + for xxx in occurrence: + occurrence[xxx] *= args.alpha_decay if token not in occurrence: occurrence[token] = 1 else: occurrence[token] += 1 + # print(occurrence) # debug # output tmp = self.decode(all_tokens[out_last:]) diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py index ce025ec..7aa3042 100644 --- a/backend-python/utils/rwkv.py +++ b/backend-python/utils/rwkv.py @@ -36,7 +36,7 @@ class AbstractRWKV(ABC): RWKV as Model, ) else: - from rwkv.model import ( + from rwkv_pip.model import ( RWKV as Model, ) from rwkv_pip.utils import PIPELINE diff --git a/backend-python/wkv_cuda_utils/wkv_cuda10_30.pyd b/backend-python/wkv_cuda_utils/wkv_cuda10_30.pyd deleted file mode 100644 index 09927ff..0000000 Binary files a/backend-python/wkv_cuda_utils/wkv_cuda10_30.pyd and /dev/null differ diff --git a/backend-python/wkv_cuda_utils/wkv_cuda40.pyd b/backend-python/wkv_cuda_utils/wkv_cuda40.pyd deleted file mode 100644 index 442bae8..0000000 Binary files a/backend-python/wkv_cuda_utils/wkv_cuda40.pyd and /dev/null differ diff --git a/backend-python/wkv_cuda_utils/wkv_cuda_model.py b/backend-python/wkv_cuda_utils/wkv_cuda_model.py deleted file mode 100644 index 7d17727..0000000 --- a/backend-python/wkv_cuda_utils/wkv_cuda_model.py +++ /dev/null @@ -1,734 +0,0 @@ -######################################################################################################## -# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM -######################################################################################################## - -import types, gc, os, time, re -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"], - verbose=True, - 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' - -######################################################################################################## - -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 - 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 # 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()) - 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) - - ####################### 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 - 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 - w[x] = -torch.exp(w[x].float()) - elif '.time_first' in x: # need fp32 for this - 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(rx @ rw) - vx = torch.square(torch.relu(kx @ kw)) - out = r * (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(rx @ rw) - vx = torch.square(torch.relu(kx @ kw)) - out = r * (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(rx @ rw) - k = (kx @ kw).float() - v = (vx @ vw).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 = (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(rx @ rw) - k = (kx @ kw).float() - v = (vx @ vw).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 = (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 - - ######################################################################################################## - - if os.environ["RWKV_CUDA_ON"] == '1': - @MyFunction - def cuda_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): - 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(rx @ rw) - k = kx @ kw - v = vx @ vw - y, aa, bb, pp = cuda_wkv(T, C, t_decay, t_first, k, v, aa, bb, pp) - - out = (r * y) @ 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 - - ######################################################################################################## - - def forward(self, tokens, state, full_output=False): - with torch.no_grad(): - w = self.w - args = self.args - - if state == None: - 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() - - 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: - if 'cuda' in str(dev) and os.environ["RWKV_CUDA_ON"] == '1': - ATT = self.cuda_att_seq if wtype != torch.uint8 else self.cuda_att_seq_i8 - else: - 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 - 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 - - 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 - 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, - ) - 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 - x, state[i*5+4] = FFN( - x, state[i*5+4], - 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