$ show BlinkDL/ChatRWKV
ChatRWKV is a collection of inference tools, demos, and a pip package for the RWKV language model family (v4-v7), a 100% RNN architecture that rivals transformers in performance. The repository provides interactive chat scripts (chat.py, v2/chat.py), API demos (API_DEMO_CHAT.py), model conversion utilities, state merging, music generation (MIDI/ABC), and benchmarking scripts. The core RWKV model class (rwkv_pip_package/src/rwkv/model.py) supports flexible deployment strategies (CPU/GPU, fp16/fp32/bf16, int8 quantization, streaming) and includes custom CUDA kernels for efficient WKV operations. The tokenizer module (tokenizer/rwkv_tokenizer.py) offers a fast trie-based byte-level tokenizer alongside a reference implementation. While the repository is primarily a research/demo project with many ad-hoc scripts and hardcoded paths, the pip package and core components are well-crafted and demonstrate significant engineering skill in GPU optimization and model inference. The codebase is actively maintained and supports the latest RWKV-7 architecture.
The LLM identified 3 files worth highlighting for reuse:
Implements a fast, trie-based byte-level tokenizer (TRIE_TOKENIZER) and a reference RWKV_TOKENIZER. The trie structure enables O(length) tokenization and detokenization, and the class is self-contained with no external dependencies beyond standard Python. It can be extracted and reused in any project requiring a lightweight, efficient tokenizer for byte-level or subword tokenization.
Provides the PIPELINE class that wraps a model and tokenizer into a high-level text generation interface. It handles tokenization, autoregressive generation with configurable sampling parameters (temperature, top-p, top-k, repetition penalty), and state management for multi-turn conversations. The class is well-structured and can be reused as a template for building generation pipelines for other autoregressive models.
Contains the CUDA kernel for the RWKV-7 WKV operation, a core recurrent computation. The kernel is optimized for GPU execution with block-level parallelism and memory coalescing. It demonstrates advanced CUDA programming techniques for implementing custom RNN-like operations and can serve as a reference for similar GPU kernel development in other sequence models.