From 8a13bd3c1e62934a12d99a6ef933e4671cd36923 Mon Sep 17 00:00:00 2001 From: josc146 Date: Mon, 14 Aug 2023 22:07:15 +0800 Subject: [PATCH] add rwkv-cuda-beta support (faster) --- backend-golang/rwkv.go | 9 +- backend-python/global_var.py | 1 + backend-python/main.py | 46 +- backend-python/rwkv_pip/beta/cuda/att_one.cu | 124 ++ backend-python/rwkv_pip/beta/cuda/att_seq.cu | 179 ++ .../rwkv_pip/beta/cuda/element_wise.h | 21 + backend-python/rwkv_pip/beta/cuda/ffn.cu | 165 ++ .../rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp | 80 + .../rwkv_pip/beta/cuda/operators.cu | 246 +++ backend-python/rwkv_pip/beta/cuda/util.h | 7 + backend-python/rwkv_pip/beta/cuda/wrapper.cpp | 167 ++ backend-python/rwkv_pip/beta/model.py | 1479 +++++++++++++++++ backend-python/utils/rwkv.py | 21 +- frontend/src/_locales/ja/main.json | 1 + frontend/src/_locales/zh-hans/main.json | 1 + frontend/src/components/RunButton.tsx | 6 +- frontend/src/pages/Configs.tsx | 10 +- frontend/src/utils/index.tsx | 1 + frontend/wailsjs/go/backend_golang/App.d.ts | 2 +- frontend/wailsjs/go/backend_golang/App.js | 4 +- 20 files changed, 2550 insertions(+), 20 deletions(-) create mode 100644 backend-python/rwkv_pip/beta/cuda/att_one.cu create mode 100644 backend-python/rwkv_pip/beta/cuda/att_seq.cu create mode 100644 backend-python/rwkv_pip/beta/cuda/element_wise.h create mode 100644 backend-python/rwkv_pip/beta/cuda/ffn.cu create mode 100644 backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp create mode 100644 backend-python/rwkv_pip/beta/cuda/operators.cu create mode 100644 backend-python/rwkv_pip/beta/cuda/util.h create mode 100644 backend-python/rwkv_pip/beta/cuda/wrapper.cpp create mode 100644 backend-python/rwkv_pip/beta/model.py diff --git a/backend-golang/rwkv.go b/backend-golang/rwkv.go index 27a3b5a..9e8d7f3 100644 --- a/backend-golang/rwkv.go +++ b/backend-golang/rwkv.go @@ -10,7 +10,7 @@ import ( "strings" ) -func (a *App) StartServer(python string, port int, host string) (string, error) { +func (a *App) StartServer(python string, port int, host string, rwkvBeta bool) (string, error) { var err error if python == "" { python, err = GetPython() @@ -18,7 +18,12 @@ func (a *App) StartServer(python string, port int, host string) (string, error) if err != nil { return "", err } - return Cmd(python, "./backend-python/main.py", strconv.Itoa(port), host) + args := []string{python, "./backend-python/main.py"} + if rwkvBeta { + args = append(args, "--rwkv-beta") + } + args = append(args, "--port", strconv.Itoa(port), "--host", host) + return Cmd(args...) } func (a *App) ConvertModel(python string, modelPath string, strategy string, outPath string) (string, error) { diff --git a/backend-python/global_var.py b/backend-python/global_var.py index 546880d..5f1215a 100644 --- a/backend-python/global_var.py +++ b/backend-python/global_var.py @@ -1,5 +1,6 @@ from enum import Enum, auto +Args = "args" Model = "model" Model_Status = "model_status" Model_Config = "model_config" diff --git a/backend-python/main.py b/backend-python/main.py index 8f65d25..f842b4b 100644 --- a/backend-python/main.py +++ b/backend-python/main.py @@ -1,5 +1,11 @@ +import time + +start_time = time.time() + import os import sys +import argparse +from typing import Sequence sys.path.append(os.path.dirname(os.path.realpath(__file__))) @@ -34,6 +40,11 @@ app.include_router(state_cache.router) @app.on_event("startup") def init(): global_var.init() + cmd_params = os.environ["RWKV_RUNNER_PARAMS"] + global_var.set( + global_var.Args, get_args(cmd_params.split(" ") if cmd_params else None) + ) + state_cache.init() set_torch() @@ -56,9 +67,34 @@ def exit(): parent.kill() -if __name__ == "__main__": - uvicorn.run( - "main:app", - port=8000 if len(sys.argv) < 2 else int(sys.argv[1]), - host="127.0.0.1" if len(sys.argv) < 3 else sys.argv[2], +def get_args(args: Union[Sequence[str], None] = None): + parser = argparse.ArgumentParser() + group = parser.add_argument_group(title="server arguments") + group.add_argument( + "--port", + type=int, + default=8000, + help="port to run the server on (default: 8000)", ) + group.add_argument( + "--host", + type=str, + default="127.0.0.1", + help="host to run the server on (default: 127.0.0.1)", + ) + group = parser.add_argument_group(title="mode arguments") + group.add_argument( + "--rwkv-beta", + action="store_true", + help="whether to use rwkv-beta (default: False)", + ) + args = parser.parse_args(args) + + return args + + +if __name__ == "__main__": + args = get_args() + os.environ["RWKV_RUNNER_PARAMS"] = " ".join(sys.argv[1:]) + print("--- %s seconds ---" % (time.time() - start_time)) + uvicorn.run("main:app", port=args.port, host=args.host, workers=1) diff --git a/backend-python/rwkv_pip/beta/cuda/att_one.cu b/backend-python/rwkv_pip/beta/cuda/att_one.cu new file mode 100644 index 0000000..f22858d --- /dev/null +++ b/backend-python/rwkv_pip/beta/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/beta/cuda/att_seq.cu b/backend-python/rwkv_pip/beta/cuda/att_seq.cu new file mode 100644 index 0000000..4d506a3 --- /dev/null +++ b/backend-python/rwkv_pip/beta/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/beta/cuda/element_wise.h b/backend-python/rwkv_pip/beta/cuda/element_wise.h new file mode 100644 index 0000000..eedc2f9 --- /dev/null +++ b/backend-python/rwkv_pip/beta/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/beta/cuda/ffn.cu b/backend-python/rwkv_pip/beta/cuda/ffn.cu new file mode 100644 index 0000000..c1c2c80 --- /dev/null +++ b/backend-python/rwkv_pip/beta/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/beta/cuda/gemm_fp16_cublas.cpp b/backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp new file mode 100644 index 0000000..6ec136d --- /dev/null +++ b/backend-python/rwkv_pip/beta/cuda/gemm_fp16_cublas.cpp @@ -0,0 +1,80 @@ +#include +#include +#include +#include +#include + +#define CUBLAS_CHECK(condition) \ + for (cublasStatus_t _cublas_check_status = (condition); \ + _cublas_check_status != CUBLAS_STATUS_SUCCESS;) \ + throw std::runtime_error("cuBLAS error " + \ + std::to_string(_cublas_check_status) + " at " + \ + std::to_string(__LINE__)); + +#define CUDA_CHECK(condition) \ + for (cudaError_t _cuda_check_status = (condition); \ + _cuda_check_status != cudaSuccess;) \ + throw std::runtime_error( \ + "CUDA error " + std::string(cudaGetErrorString(_cuda_check_status)) + \ + " at " + std::to_string(__LINE__)); + +cublasHandle_t get_cublas_handle() { + static cublasHandle_t cublas_handle = []() { + cublasHandle_t handle = nullptr; + CUBLAS_CHECK(cublasCreate(&handle)); +#if CUDA_VERSION < 11000 + CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH)); +#else + CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH)); +#endif // CUDA_VERSION < 11000 + return handle; + }(); + return cublas_handle; +} + +/* + NOTE: blas gemm is column-major by default, but we need row-major output. + The data of row-major, transposed matrix is exactly the same as the + column-major, non-transposed matrix, and C = A * B ---> C^T = B^T * A^T + */ +void gemm_fp16_cublas(const void *a, const void *b, void *c, int ori_m, + int ori_n, int ori_k, bool output_fp32) { + const auto cuda_data_type = CUDA_R_16F; + const auto cuda_c_data_type = output_fp32 ? CUDA_R_32F : CUDA_R_16F; + const auto compute_type = CUDA_R_32F; + const float sp_alpha = 1.f; + // use CUBLAS_OP_N. see the notes above + const cublasOperation_t cublas_trans_a = CUBLAS_OP_N; + const cublasOperation_t cublas_trans_b = CUBLAS_OP_N; + // m = (B^T).size(0) = B.size(1) = n; + const int cublas_m = ori_n; + const int cublas_k = ori_k; + // comptiable with rwkv one mode, where 1-D tensor * 2-D tensor + // const int n = a.dense_dim() == 1 ? 1 : a.size(0); + const int cublas_n = ori_m; + const int cublas_lda = cublas_m; + const int cublas_ldb = cublas_k; + const int cublas_ldc = cublas_m; + cublasHandle_t cublas_handle = get_cublas_handle(); + +#if CUDA_VERSION >= 11000 + cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT; +#else + cublasGemmAlgo_t algo = CUBLAS_GEMM_DFALT_TENSOR_OP; +#endif + const float sp_beta = 0.f; + CUBLAS_CHECK(cublasGemmEx( + cublas_handle, cublas_trans_a, cublas_trans_b, cublas_m, cublas_n, + cublas_k, &sp_alpha, b, cuda_data_type, cublas_lda, + a, cuda_data_type, cublas_ldb, &sp_beta, c, + cuda_c_data_type, cublas_ldc, compute_type, algo)); +} + +void gemm_fp16_cublas(torch::Tensor a, torch::Tensor b, torch::Tensor c) { + // comptiable with rwkv one mode, 1-D tensor * 2-D tensor + const int m = a.dense_dim() == 1 ? 1 : a.size(0); + const int n = b.size(1); + const int k = b.size(0); + gemm_fp16_cublas(a.data_ptr(), b.data_ptr(), c.data_ptr(), m, n, k, + c.dtype() == torch::kFloat32); +} diff --git a/backend-python/rwkv_pip/beta/cuda/operators.cu b/backend-python/rwkv_pip/beta/cuda/operators.cu new file mode 100644 index 0000000..fa5a44f --- /dev/null +++ b/backend-python/rwkv_pip/beta/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/beta/cuda/util.h b/backend-python/rwkv_pip/beta/cuda/util.h new file mode 100644 index 0000000..f00af22 --- /dev/null +++ b/backend-python/rwkv_pip/beta/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/beta/cuda/wrapper.cpp b/backend-python/rwkv_pip/beta/cuda/wrapper.cpp new file mode 100644 index 0000000..ee99bfc --- /dev/null +++ b/backend-python/rwkv_pip/beta/cuda/wrapper.cpp @@ -0,0 +1,167 @@ +#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; + +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 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 ffn_seq(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, + Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, + /* imm */ Tensor buf, + /* out */ Tensor x_plus_out); + +Tensor ffn_one(Tensor x, Tensor sx, Tensor ln_w, Tensor ln_b, Tensor k_mix, + Tensor r_mix, Tensor kw, Tensor vw, Tensor rw, + /* imm */ Tensor buf, + /* out */ Tensor x_plus_out); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("wkv_forward", &wkv_forward, "wkv forward"); + m.def("mm8_seq", &mm8_seq, "mm8 seq"); + m.def("mm8_one", &mm8_one, "mm8 one"); + m.def("gemm_fp16_cublas", &gemm_fp16_cublas, "gemv fp16 cublas"); + m.def("att_one", &att_one, "att one"); + m.def("att_seq", &att_seq, "att seq"); + m.def("ffn_seq", &ffn_seq, "ffn seq"); + m.def("ffn_one", &ffn_one, "ffn one"); +} + +TORCH_LIBRARY(rwkv, m) { + m.def("wkv_forward", wkv_forward); + m.def("mm8_seq", mm8_seq); + m.def("mm8_one", mm8_one); + m.def("gemm_fp16_cublas", gemm_fp16_cublas); + m.def("att_one", att_one); + m.def("att_seq", att_seq); + m.def("ffn_seq", ffn_seq); + m.def("ffn_one", ffn_one); +} diff --git a/backend-python/rwkv_pip/beta/model.py b/backend-python/rwkv_pip/beta/model.py new file mode 100644 index 0000000..808e456 --- /dev/null +++ b/backend-python/rwkv_pip/beta/model.py @@ -0,0 +1,1479 @@ +######################################################################################################## +# 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__)) + + +# https://zhuanlan.zhihu.com/p/612879065 +def LoadPreCompileLibrary(file): + import importlib + import os + + import torch + + # load the custom_op_library and register the custom ops + lib_dir = os.path.dirname(__file__) + if os.name == "nt": + # Register the main torchvision library location on the default DLL path + import ctypes + import sys + + kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True) + with_load_library_flags = hasattr(kernel32, "AddDllDirectory") + prev_error_mode = kernel32.SetErrorMode(0x0001) + + if with_load_library_flags: + kernel32.AddDllDirectory.restype = ctypes.c_void_p + + if sys.version_info >= (3, 8): + os.add_dll_directory(lib_dir) + elif with_load_library_flags: + res = kernel32.AddDllDirectory(lib_dir) + if res is None: + err = ctypes.WinError(ctypes.get_last_error()) + err.strerror += f' Error adding "{lib_dir}" to the DLL directories.' + raise ValueError(err) + + kernel32.SetErrorMode(prev_error_mode) + + loader_details = ( + importlib.machinery.ExtensionFileLoader, + importlib.machinery.EXTENSION_SUFFIXES, + ) + + extfinder = importlib.machinery.FileFinder(lib_dir, loader_details) + ext_specs = extfinder.find_spec(file) + if ext_specs is None: + return False + + try: + torch.ops.load_library(ext_specs.origin) + except OSError as exc: + return False + return True + + +######################################################################################################## + +if os.environ.get("RWKV_JIT_ON") != "0": + os.environ["RWKV_JIT_ON"] = "1" + MyModule = torch.jit.ScriptModule + MyFunction = torch.jit.script_method + MyStatic = torch.jit.script +else: + MyModule = torch.nn.Module + + def __nop(ob): + return ob + + MyFunction = __nop + MyStatic = __nop + +if os.environ.get("RWKV_CUDA_ON") == "1": + if LoadPreCompileLibrary("wkv_cuda") is False: + from torch.utils.cpp_extension import load + + load( + name=f"wkv_cuda", + sources=[ + f"{current_path}/cuda/wrapper.cpp", + f"{current_path}/cuda/operators.cu", + f"{current_path}/cuda/gemm_fp16_cublas.cpp", + f"{current_path}/cuda/att_one.cu", + f"{current_path}/cuda/att_seq.cu", + f"{current_path}/cuda/ffn.cu", + ], + 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) + + @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": + assert len(b.shape) == 2 + if len(a.shape) == 1: + c = torch.empty((b.shape[-1],), dtype=output_dtype, device=a.device) + a = a.unsqueeze(0) + else: + c = torch.empty( + (a.shape[0], 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: + os.environ["RWKV_CUDA_ON"] = "0" + + def gemm(a, b, output_dtype: Optional[torch.dtype] = None): + if output_dtype is None: + output_dtype = a.dtype + return (a @ b).to(output_dtype) + + +######################################################################################################## + + +class RWKV(MyModule): + def __init__(self, model, strategy, verbose=True, convert_and_save_and_exit=None): + super().__init__() + if verbose: + prxxx = lambda *args, **kwargs: print(*args, **kwargs) + else: + prxxx = lambda *args, **kwargs: None + + STRATEGY_REGEX = r"^(?:(?:^|->) *(?:cuda(?::[\d]+)?|cpu|mps) (?:fp(?:16|32)|bf16)(?:i8|i4|i3)?(?: \*[\d]+\+?)? *)+$" + if not re.match(STRATEGY_REGEX, strategy): + raise ValueError( + "Invalid strategy. Please read https://pypi.org/project/rwkv/" + ) + + strategy = ("->".join([x.strip() for x in strategy.split("->")])).replace( + "->", " -> " + ) + self.args = types.SimpleNamespace() + args = self.args + args.MODEL_NAME = model + args.strategy_string = strategy + + # Rescale for fp16 mode: set x = x/2 every X layer (to avoid fp16 overflow) + self.RESCALE_LAYER = 6 if "fp16" in strategy else 0 + prxxx( + f'RWKV_JIT_ON {os.environ["RWKV_JIT_ON"]} RWKV_CUDA_ON {os.environ["RWKV_CUDA_ON"]} RESCALE_LAYER {self.RESCALE_LAYER}\n' + ) + + args.MODEL_NAME = args.MODEL_NAME.strip() + if not args.MODEL_NAME.endswith(".pth"): + args.MODEL_NAME += ".pth" + prxxx(f"Loading {args.MODEL_NAME} ...") + with torch.no_grad(): + self.w = torch.load( + args.MODEL_NAME, map_location="cpu" + ) # load model to CPU first + # it is supported to load a pure meta-tensor state dict (e.g. for quick testing) + for k, v in self.w.items(): + if v.is_meta: + # torch.zeros_like(v, device='cpu') doesn't produce an all-zero tensor + # if v is a meta tensor + self.w[k] = torch.zeros(v.shape, dtype=v.dtype, device="cpu") + gc.collect() + w = self.w + + ALREADY_CONVERTED = False + if "_strategy" in w: + ALREADY_CONVERTED = True + assert ( + convert_and_save_and_exit == None + ) # you should only convert a raw model + prxxx( + f"Converted model: strategy {w['_strategy']}, version {w['_version']}\n" + ) + assert ( + w["_strategy"] == args.strategy_string + ) # 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(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 + + ######################################################################################################## + + if os.environ["RWKV_CUDA_ON"] == "1": + + @MyFunction + def cuda_att_seq_fp16( + self, + x, + sx, + aa, + bb, + pp, + ln_w, + ln_b, + k_mix, + v_mix, + r_mix, + t_decay, + t_first, + kw, + vw, + rw, + ow, + kmx, + krx, + kmy, + kry, + vmx, + vrx, + vmy, + vry, + rmx, + rrx, + rmy, + rry, + omx, + orx, + omy, + ory, + ): + seq_len = x.shape[0] + kvrx_and_y_bytes = x.numel() * 2 + k_bytes = seq_len * kw.shape[1] * 4 + v_bytes = seq_len * vw.shape[1] * 4 + r_bytes = seq_len * rw.shape[1] * 2 + buf = torch.empty( + (kvrx_and_y_bytes * 4 + k_bytes + v_bytes + r_bytes,), + device=x.device, + dtype=torch.int8, + ) + x_plus_out_t = torch.empty_like(x) + xx = torch.ops.rwkv.att_seq( + x, + sx, + ln_w, + ln_b, + k_mix, + v_mix, + r_mix, + kw, + vw, + rw, + ow, + t_first, + pp, + aa, + bb, + t_decay, + buf, + x_plus_out_t, + ) + + return x_plus_out_t, xx[-1, :], aa, bb, pp + + @MyFunction + def cuda_att_seq_naive( + self, + x, + sx, + aa, + bb, + pp, + ln_w, + ln_b, + k_mix, + v_mix, + r_mix, + t_decay, + t_first, + kw, + vw, + rw, + ow, + kmx, + krx, + kmy, + kry, + vmx, + vrx, + vmy, + vry, + rmx, + rrx, + rmy, + rry, + omx, + orx, + omy, + ory, + ): + T, C = x.size() + xx = F.layer_norm(x, (C,), weight=ln_w, bias=ln_b) + sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) + kx = xx * k_mix + sx * (1 - k_mix) + vx = xx * v_mix + sx * (1 - v_mix) + rx = xx * r_mix + sx * (1 - r_mix) + + r = torch.sigmoid(gemm(rx, rw)) + k = gemm(kx, kw, output_dtype=torch.float32) + v = gemm(vx, vw, output_dtype=torch.float32) + y, aa, bb, pp = cuda_wkv(T, C, t_decay, t_first, k, v, aa, bb, pp) + + out = gemm(r * y.to(x.dtype), ow) + return x + out, xx[-1, :], aa, bb, pp + + @MyFunction + def cuda_att_seq_i8( + self, + x, + sx, + aa, + bb, + pp, + ln_w, + ln_b, + k_mix, + v_mix, + r_mix, + t_decay, + t_first, + kw, + vw, + rw, + ow, + kmx, + krx, + kmy, + kry, + vmx, + vrx, + vmy, + vry, + rmx, + rrx, + rmy, + rry, + omx, + orx, + omy, + ory, + ): + T, C = x.size() + xx = F.layer_norm(x, (C,), weight=ln_w, bias=ln_b) + sx = torch.cat((sx.unsqueeze(0), xx[:-1, :])) + kx = xx * k_mix + sx * (1 - k_mix) + vx = xx * v_mix + sx * (1 - v_mix) + rx = xx * r_mix + sx * (1 - r_mix) + + r = torch.sigmoid(self.mm8_seq(rx, rw, rmx, rrx, rmy, rry)) + k = self.mm8_seq(kx, kw, kmx, krx, kmy, kry) + v = self.mm8_seq(vx, vw, vmx, vrx, vmy, vry) + y, aa, bb, pp = cuda_wkv(T, C, t_decay, t_first, k, v, aa, bb, pp) + + out = self.mm8_seq(r * y, ow, omx, orx, omy, ory) + return x + out, xx[-1, :], aa, bb, pp + + @MyFunction + def cuda_ffn_seq_fp16( + self, + x, + sx, + ln_w, + ln_b, + k_mix, + r_mix, + kw, + vw, + rw, + kmx, + krx, + kmy, + kry, + vmx, + vrx, + vmy, + vry, + rmx, + rrx, + rmy, + rry, + ): + krx_bytes = x.numel() * x.element_size() + vx_bytes = x.shape[0] * kw.shape[1] * x.element_size() + r_bytes = x.shape[0] * rw.shape[1] * x.element_size() + buf = torch.empty( + (krx_bytes * 2 + vx_bytes + r_bytes,), device=x.device, dtype=torch.int8 + ) + x_plus_out = torch.empty_like(x) + xx = torch.ops.rwkv.ffn_seq( + x, sx, ln_w, ln_b, k_mix, r_mix, kw, vw, rw, buf, x_plus_out + ) + return x_plus_out, xx[-1:] + + @MyFunction + def cuda_att_one_fp16( + self, + x, + sx, + aa, + bb, + pp, + ln_w, + ln_b, + k_mix, + v_mix, + r_mix, + t_decay, + t_first, + kw, + vw, + rw, + ow, + kmx, + krx, + kmy, + kry, + vmx, + vrx, + vmy, + vry, + rmx, + rrx, + rmy, + rry, + omx, + orx, + omy, + ory, + ): + kx = torch.empty_like(x) + vx = torch.empty_like(x) + rx = torch.empty_like(x) + + k_t = torch.empty((kw.shape[0],), dtype=torch.float32, device=x.device) + v_t = torch.empty((vw.shape[0],), dtype=torch.float32, device=x.device) + r_t = torch.empty((rw.shape[0],), dtype=torch.float16, device=x.device) + x_plus_out_t = torch.empty_like(x) + t1_t = torch.empty_like(x, dtype=torch.float32) + t2_t = torch.empty_like(x, dtype=torch.float32) + p_t = torch.empty_like(x, dtype=torch.float32) + xx = torch.ops.rwkv.att_one( + x, + ln_w, + ln_b, + sx, + k_mix, + v_mix, + r_mix, + kw, + kx, + vw, + vx, + rw, + rx, + ow, + t_first, + k_t, + pp, + ow, + aa, + bb, + t_decay, + v_t, + r_t, + x_plus_out_t, + t1_t, + t2_t, + p_t, + ) + return x_plus_out_t, xx, t1_t, t2_t, p_t + + @MyFunction + def cuda_ffn_one_fp16( + self, + x, + sx, + ln_w, + ln_b, + k_mix, + r_mix, + kw, + vw, + rw, + kmx, + krx, + kmy, + kry, + vmx, + vrx, + vmy, + vry, + rmx, + rrx, + rmy, + rry, + ): + krx_bytes = x.numel() * x.element_size() + vx_bytes = x.shape[0] * kw.shape[1] * x.element_size() + r_bytes = x.shape[0] * rw.shape[1] * x.element_size() + buf = torch.empty( + (krx_bytes * 2 + vx_bytes + r_bytes,), device=x.device, dtype=torch.int8 + ) + x_plus_out = torch.empty_like(x) + xx = torch.ops.rwkv.ffn_one( + x, sx, ln_w, ln_b, k_mix, r_mix, kw, vw, rw, buf, x_plus_out + ) + return x_plus_out, xx + + ######################################################################################################## + + def forward(self, tokens, state, full_output=False): + with torch.no_grad(): + w = self.w + args = self.args + + if state == None: + 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: + ATT = self.att_seq if wtype != torch.uint8 else self.att_seq_i8 + FFN = self.ffn_seq if wtype != torch.uint8 else self.ffn_seq_i8 + if "cuda" in str(dev) and os.environ["RWKV_CUDA_ON"] == "1": + if wtype == torch.float16: + ATT = self.cuda_att_seq_fp16 + FFN = self.cuda_ffn_seq_fp16 + elif wtype == torch.uint8: + ATT = self.cuda_att_seq_i8 + else: + ATT = self.cuda_att_seq_naive + else: + ATT = self.att_one if wtype != torch.uint8 else self.att_one_i8 + FFN = self.ffn_one if wtype != torch.uint8 else self.ffn_one_i8 + if ( + "cuda" in str(dev) + and os.environ["RWKV_CUDA_ON"] == "1" + and wtype == torch.float16 + ): + ATT = self.cuda_att_one_fp16 + FFN = self.cuda_ffn_one_fp16 + + 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 diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py index b24e4e1..ce025ec 100644 --- a/backend-python/utils/rwkv.py +++ b/backend-python/utils/rwkv.py @@ -10,6 +10,7 @@ from fastapi import HTTPException from pydantic import BaseModel, Field import numpy as np from routes import state_cache +import global_var END_OF_TEXT = 0 @@ -27,7 +28,17 @@ class RWKVType(Enum): class AbstractRWKV(ABC): def __init__(self, model: str, strategy: str, tokens_path: str): - from rwkv.model import RWKV as Model # dynamic import to make RWKV_CUDA_ON work + rwkv_beta = global_var.get(global_var.Args).rwkv_beta + + # dynamic import to make RWKV_CUDA_ON work + if rwkv_beta: + from rwkv_pip.beta.model import ( + RWKV as Model, + ) + else: + from rwkv.model import ( + RWKV as Model, + ) from rwkv_pip.utils import PIPELINE filename, _ = os.path.splitext(os.path.basename(model)) @@ -221,7 +232,7 @@ class AbstractRWKV(ABC): return state[0].tolist(), token_len def generate( - self, prompt: str, stop: Union[str, List[str]] = None + self, prompt: str, stop: Union[str, List[str], None] = None ) -> Iterable[Tuple[str, str, int, int]]: quick_log(None, None, "Generation Prompt:\n" + prompt) cache = None @@ -438,8 +449,10 @@ The following is a coherent verbose detailed conversation between a girl named { {bot} usually gives {user} kind, helpful and informative advices.\n """ if self.rwkv_type == RWKVType.Raven - else f"{user}{interface} hi\n\n{bot}{interface} Hi. " - + "I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.\n\n" + else ( + f"{user}{interface} hi\n\n{bot}{interface} Hi. " + + "I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.\n\n" + ) ) logits, _ = self.run_rnn(self.fix_tokens(self.pipeline.encode(preset_system))) try: diff --git a/frontend/src/_locales/ja/main.json b/frontend/src/_locales/ja/main.json index c4cac01..5f7e2c9 100644 --- a/frontend/src/_locales/ja/main.json +++ b/frontend/src/_locales/ja/main.json @@ -128,6 +128,7 @@ "Chinese Kongfu": "中国武術", "Allow external access to the API (service must be restarted)": "APIへの外部アクセスを許可する (サービスを再起動する必要があります)", "Custom": "カスタム", + "CUDA (Beta, Faster)": "CUDA (ベータ、高速)", "Reset All Configs": "すべての設定をリセット", "Cancel": "キャンセル", "Confirm": "確認", diff --git a/frontend/src/_locales/zh-hans/main.json b/frontend/src/_locales/zh-hans/main.json index 5f91f31..8481569 100644 --- a/frontend/src/_locales/zh-hans/main.json +++ b/frontend/src/_locales/zh-hans/main.json @@ -128,6 +128,7 @@ "Chinese Kongfu": "情境冒险", "Allow external access to the API (service must be restarted)": "允许外部访问API (必须重启服务)", "Custom": "自定义", + "CUDA (Beta, Faster)": "CUDA (Beta, 更快)", "Reset All Configs": "重置所有配置", "Cancel": "取消", "Confirm": "确认", diff --git a/frontend/src/components/RunButton.tsx b/frontend/src/components/RunButton.tsx index b43f0e2..cc72757 100644 --- a/frontend/src/components/RunButton.tsx +++ b/frontend/src/components/RunButton.tsx @@ -85,7 +85,9 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean await exit(1000).catch(() => { }); - StartServer(commonStore.settings.customPythonPath, port, commonStore.settings.host !== '127.0.0.1' ? '0.0.0.0' : '127.0.0.1').catch((e) => { + StartServer(commonStore.settings.customPythonPath, port, commonStore.settings.host !== '127.0.0.1' ? '0.0.0.0' : '127.0.0.1', + modelConfig.modelParameters.device === 'CUDA-Beta' + ).catch((e) => { const errMsg = e.message || e; if (errMsg.includes('path contains space')) toast(`${t('Error')} - ${t('File Path Cannot Contain Space')}`, { type: 'error' }); @@ -118,7 +120,7 @@ export const RunButton: FC<{ onClickRun?: MouseEventHandler, iconMode?: boolean const strategy = getStrategy(modelConfig); let customCudaFile = ''; - if ((modelConfig.modelParameters.device === 'CUDA' || modelConfig.modelParameters.device === 'Custom') + if ((modelConfig.modelParameters.device.includes('CUDA') || modelConfig.modelParameters.device === 'Custom') && modelConfig.modelParameters.useCustomCuda && !strategy.includes('fp32')) { if (commonStore.platform === 'windows') { customCudaFile = getSupportedCustomCudaFile(); diff --git a/frontend/src/pages/Configs.tsx b/frontend/src/pages/Configs.tsx index b7565f8..8d9a290 100644 --- a/frontend/src/pages/Configs.tsx +++ b/frontend/src/pages/Configs.tsx @@ -30,7 +30,7 @@ export type ApiParameters = { frequencyPenalty: number; } -export type Device = 'CPU' | 'CUDA' | 'MPS' | 'Custom'; +export type Device = 'CPU' | 'CUDA' | 'CUDA-Beta' | 'WebGPU' | 'MPS' | 'Custom'; export type Precision = 'fp16' | 'int8' | 'fp32'; export type ModelParameters = { @@ -284,6 +284,8 @@ export const Configs: FC = observer(() => { {commonStore.platform === 'darwin' && } + + } /> @@ -308,12 +310,12 @@ export const Configs: FC = observer(() => { } /> } { - selectedConfig.modelParameters.device == 'CUDA' && + selectedConfig.modelParameters.device.includes('CUDA') && {getStrategy(selectedConfig)} } /> } { - selectedConfig.modelParameters.device == 'CUDA' && + selectedConfig.modelParameters.device.includes('CUDA') && { } /> } { - selectedConfig.modelParameters.device == 'CUDA' &&
+ selectedConfig.modelParameters.device.includes('CUDA') &&
} { displayStrategyImg && diff --git a/frontend/src/utils/index.tsx b/frontend/src/utils/index.tsx index 917a9ed..1e49baf 100644 --- a/frontend/src/utils/index.tsx +++ b/frontend/src/utils/index.tsx @@ -177,6 +177,7 @@ export const getStrategy = (modelConfig: ModelConfig | undefined = undefined) => strategy += params.precision === 'int8' ? 'fp32i8' : 'fp32'; break; case 'CUDA': + case 'CUDA-Beta': if (avoidOverflow) strategy = 'cuda fp32 *1 -> '; strategy += 'cuda '; diff --git a/frontend/wailsjs/go/backend_golang/App.d.ts b/frontend/wailsjs/go/backend_golang/App.d.ts index 0414e3f..baa1dde 100755 --- a/frontend/wailsjs/go/backend_golang/App.d.ts +++ b/frontend/wailsjs/go/backend_golang/App.d.ts @@ -46,7 +46,7 @@ export function RestartApp():Promise; export function SaveJson(arg1:string,arg2:any):Promise; -export function StartServer(arg1:string,arg2:number,arg3:string):Promise; +export function StartServer(arg1:string,arg2:number,arg3:string,arg4:boolean):Promise; export function UpdateApp(arg1:string):Promise; diff --git a/frontend/wailsjs/go/backend_golang/App.js b/frontend/wailsjs/go/backend_golang/App.js index fa1f187..a182fb6 100755 --- a/frontend/wailsjs/go/backend_golang/App.js +++ b/frontend/wailsjs/go/backend_golang/App.js @@ -90,8 +90,8 @@ export function SaveJson(arg1, arg2) { return window['go']['backend_golang']['App']['SaveJson'](arg1, arg2); } -export function StartServer(arg1, arg2, arg3) { - return window['go']['backend_golang']['App']['StartServer'](arg1, arg2, arg3); +export function StartServer(arg1, arg2, arg3, arg4) { + return window['go']['backend_golang']['App']['StartServer'](arg1, arg2, arg3, arg4); } export function UpdateApp(arg1) {