RWKV v6: RWKV_WKV op CUDA implementation

[MollySophia]

Added the RWKV_WKV CUDA impl and a test_case in test-backend-ops.cpp.
Also added unary op exp for cuda so that the rwkv v6 graph can be less splited when running on a gpu.

The kernel is modified from https://github.com/BlinkDL/ChatRWKV/blob/main/rwkv_pip_package/src/rwkv/cuda/rwkv6.cu and added support for batched inference.
Gonna add some speed and other test results later tomorrow.

Prompt:

RWKV (pronounced as RWaKuV) is an RNN with GPT-level LLM performance, which can also be directly trained like a GPT transformer (parallelizable).\nRWKV is an Open Source, non profit group, under the linux foundation. Supported by our sponsors.\nSo it's combining the best of RNN and transformer - great performance, fast inference, fast training, saves VRAM, "infinite" ctxlen, and free sentence embedding. Moreover it's 100% attention-free.\n

Here's the speed comparasion between the original and the PR version.
The test is done on my weird 12900HK ES + RTX4090 PC, which is relatively CPU-bound. The tests are all using FP16, offloading all layers to GPU. Prompt length = 107, generation length = 1000.

Parameter count	Prefill (before)	Prefill (after)	Decode (before)	Decode (after)
World-1.6B	678.72 tps	1375.48 tps	57.92 tps	142.13 tps
World-3B	481.70 tps	1262.74 tps	39.46 tps	86.00 tps
World-7B	340.25 tps	1053.80 tps	26.38 tps	47.39 tps

Here's the perplexity comparasion between the original and the PR version. Tested on wikitext-2 using FP16, offloading all layers to GPU.

Parameter count	Perplexity (before)	Perplexity (after)
World-1.6B	10.8599 +/- 0.07657	10.8604 +/- 0.07657
World-3B	9.3254 +/- 0.06322	9.3256 +/- 0.06322
World-7B	7.9571 +/- 0.05213	7.9570 +/- 0.05213

test-backend-ops perf tests:

Backend 1/2 (CPU)
  Backend name: CPU
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=1,n_seqs=1):             7826 runs -    40.90 us/run -     1072 kB/run -   24.99 GB/s
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=32,n_seqs=1):            3629 runs -   511.39 us/run -     2312 kB/run -    4.31 GB/s
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=32,n_seqs=4):             910 runs -  2195.17 us/run -     9224 kB/run -    4.01 GB/s
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=128,n_seqs=4):            342 runs -  8184.47 us/run -    24584 kB/run -    2.86 GB/s
  Backend CPU: OK

Backend 2/2 (CUDA0)
  Backend name: CUDA0
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=1,n_seqs=1):             8192 runs -     6.86 us/run -     1072 kB/run -  149.05 GB/s
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=32,n_seqs=1):            8192 runs -    22.27 us/run -     2312 kB/run -   99.00 GB/s
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=32,n_seqs=4):            3638 runs -    19.77 us/run -     9224 kB/run -  444.87 GB/s
  RWKV_WKV(type=f32,head_count=32,head_size=64,n_seq_tokens=128,n_seqs=4):           1365 runs -    71.58 us/run -    24584 kB/run -  327.55 GB/s
  Backend CUDA0: OK

• I have read the contributing guidelines
• Self-reported review complexity:
  • Medium

TODO:

• Add speed comparision
• Add perplexity comparision

[uniartisan]

Hi, Molly.
I have some small suggestions. Can you rename the wkv kernel to wkv6? As far as I know, wkv7 will be released soon.
Maybe we can move forward and implement fp16 calculations（Maybe next pr）

[MollySophia]

Hi, Molly. I have some small suggestions. Can you rename the wkv kernel to wkv6? As far as I know, wkv7 will be released soon. Maybe we can move forward and implement fp16 calculations（Maybe next pr）

Yes. However I guess this is not that urgent. That can also be done after RWKV v7 is released, in the initial rwkv v7 support PR in the future.

[MollySophia]

Hi! @ggerganov
In case you forgot about this PR :D