Web-RWKV

这是纯WebGPU实现的RWKV语言模型的推理引擎。

特点

• 不依赖于CUDA/Python。
• 支持Nvidia/AMD/Intel GPU，包括集成GPU。
• Vulkan/Dx12/OpenGL后端。
• WASM支持（可在浏览器中运行）。
• 批量推理。
• Int8和NF4量化。
• 非常快。
• 在加载时进行LoRA合并。
• 支持RWKV V4、V5和V6。
• 在任何点上干预推理过程的钩子。
• 模型（反）序列化。

请注意，web-rwkv仅是推理引擎。它只提供以下功能

• 分词器。
• 模型加载。
• 状态创建和更新。
• 模型实现了接受提示标记并返回logits的run函数，以及将logits转换为预测下一个标记概率的softmax函数。这两个函数都在GPU上执行。
• 模型量化和（反）序列化。
• WASM绑定。

它 不 提供以下内容

• OpenAI API或其他任何类型的API。
  • 如果您想部署API服务器，请查看AI00 RWKV Server，这是一个基于web-rwkv的完全功能的OpenAI兼容API服务器。
  • 如果您想要一些复杂的推理管道，包括无分类器指导（CFG）、巴科斯-诺尔范式（BNF）指导等，也可以查看web-rwkv-axum项目。
• 采样器，尽管在示例中实现了一个基本的核采样器，但这 不包括 在库本身中。
• 状态缓存或管理系统。
• Python 绑定。

编译

1. 安装 Rust.
2. 从 HuggingFace 下载模型，并使用 convert_safetensors.py 转换它。将 .st 模型放在 assets/models 目录下。
3. 编译

   $ cargo build --release --examples
   

示例

性能测试

该测试生成 500 个标记并测量时间成本。

$ cargo run --release --example rt-gen

聊天演示

要与模型聊天，运行

$ cargo run --release --example rt-chat

在此演示中，输入 + 重试上一轮的生成；输入 - 退出。

• 要指定您 safetensors 模型的位置，使用

  $ cargo run --release --example rt-chat -- --model /path/to/model
  

• 要为聊天加载自定义提示，使用

  $ cargo run --release --example rt-chat -- --prompt /path/to/prompt
  

  有关详细信息，请参阅 assets/prompt.json。

• 要指定层量化，使用 --quant <LAYERS> 或 --quant-nf4 <LAYERS> 对前 <LAYERS> 层进行量化。例如，使用

  $ cargo run --release --example rt-chat -- --quant 32
  

  量化所有 32 层。

• 使用 --turbo 标志在推断长提示时切换到替代的 GEMM 内核。 rt 示例强制执行 turbo。

批量推理

此演示展示了同时生成具有不同长度的 4 批文本。

$ cargo run --release --example rt-batch

检查器

检查器演示是高级使用方法（钩子）的指南。钩子允许用户将任何张量操作注入到模型的推理过程中，获取和修改运行时缓冲区、状态，甚至模型参数的内容。钩子使某些第三方实现（如动态 LoRA、控制网等）成为可能。

(反)序列化

所有版本的模型实现了 serde::ser::Serialize 和 serde::de::DeserializeSeed<'de>，这意味着可以将量化或 LoRA 合并的模型保存到文件中，然后加载它。

在您的项目中使用

要在您的 rust 项目中使用，只需将 web-rwkv = "0.8" 作为依赖项添加到您的 Cargo.toml 中。查看示例了解如何创建环境、分词器以及如何运行模型。

说明

推理运行时

从 v0.7 开始，该软件包具有一个 runtime 功能。当启用时，应用程序可以使用异步 runtime API 的基础设施。

一般来说，一个 runtime 是一个由 tokio 驱动的异步任务。它允许 CPU 和 GPU 并行工作，最大限度地提高 GPU 计算资源的利用率。

查看以 rt 开头的示例以获取更多信息，并比较它们的生成速度与它们的非 rt 对应物。

批量推理

从版本 v0.2.4 开始，该引擎支持批量推理，即并行推理一批提示（具有不同长度）。这是通过修改的 WKV 内核实现的。

在构建模型时，用户指定了 token_chunk_size（默认为32，但对于强大的GPU来说，这个值可能更高），它是引擎在一次 run 调用中可以处理的token的最大数量。

创建模型后，用户使用指定的 num_batch 创建一个 ModelState。这意味着有 num_batch 个槽位可以并行消耗输入。

在调用 run() 之前，用户在每个槽位中填充一些token作为提示。如果一个槽位为空，则不会对该槽位运行推理。

在调用 run() 之后，一些（但不一定全部）输入token被消耗，并且在本次运行期间如果该槽位的推理已完成，则相应的返回槽位中出现 logits。由于每次 run() 调用只处理 token_chunk_size 个token，结果中可能没有任何 logits。

钩子

钩子是定制模型推理过程的非常强大的工具。库提供了带有 Model::run_with_hooks 函数的功能，该函数接受一个 HookMap 作为参数。

HookMap 实际上是从 Model::Hook 到函数的哈希表。一个 Model::Hook 定义了钩子函数可以注入的特定位置。一个模型通常有几十个钩点。钩子函数是一个函数，其类型为 Fn(&Model<'_>, &ModelState, &Runtime) -> Result<TensorOp, TensorError>，您可以在其中创建读取/写入您在这里获得的全部张量的张量操作。

let info = model.info();
// create a buffer to store each layer's output
let buffer = Buffer::new(&context, &info);
let mut hooks = HookMap::default();
for layer in 0..info.num_layer {
   let buffer = buffer.clone();
   hooks.insert(
      v5::Hook::PostFfn(layer),
      Box::new(
            move |_model, _state, runtime: &v5::Runtime| -> Result<TensorOp, TensorError> {
               // figure out how many tokens this run has
               let shape = runtime.ffn_x.shape();
               let num_token = shape[1];
               // "steal" the layer's output (activation), and put it into our buffer
               TensorOp::blit(
                  runtime.ffn_x.view(.., num_token - 1, .., ..)?,
                  buffer.ffn_x.view(.., layer, .., ..)?,
               )
            },
      ),
   );
}
let output = model.run_with_hooks(&mut tokens, &state, &hooks).await?;

$ python convert_safetensors.py --input /path/to/model.pth --output /path/to/model.st

$ cd /path/to/web-rwkv-converter
$ cargo run --release --example converter -- --input /path/to/model.pth --output /path/to/model.st