MLC-LLM 部署RWKV World系列模型实战（3B模型Mac M2解码可达26tokens/s）

0x0. 前言

我的 ChatRWKV 学习笔记和使用指南 这篇文章是学习RWKV的第一步，然后学习了一下之后决定自己应该做一些什么。所以就在RWKV社区看到了这个将RWKV World系列模型通过MLC-LLM部署在各种硬件平台的需求，然后我就开始了解MLC-LLM的编译部署流程和RWKV World模型相比于MLC-LLM已经支持的Raven系列模型的特殊之处。

MLC-LLM的编译部署流程在MLC-LLM的官方文档已经比较详细了，但这部分有一些隐藏的坑点需要你去发现，比如现在要支持的RWKV-World模型它的Tokenizer是自定义的，并不是Huggingface的格式，这就导致我们不能使用MLC-LLM去直接编译这个模型，也不能使用预编译好的MLC-LLM二进制库去运行这个模型了。另外，在编译MLC-LLM仓库之前我们需要先编译Relax仓库而不是原始的TVM仓库，Relax可以认为是TVM的一个fork，在此基础上支持了Relax这个新一代的IR，这部分背景建议读者看一下我这个仓库的相关链接：

https://github.com/BBuf/tvm_mlir_learn

这个仓库已经揽下1.4k star，谢谢读者们支持。

从RWKV社区了解到，RWKV-World系列模型相比于Raven系列，推理代码和模型都是完全一样，不一样的地方主要是tokenizer是自定义的，并且system prompt不同。

在编译Relax的时候需要按需选择自己的编译平台进行编译，编译完之后 MLC-LLM 会通过 TVM_HOME 这个环境变量来感知 Relax 的位置，并且Relax编译时开启的选项要和MLC-LLM编译的选项匹配上，这样才可以在指定平台上进行正确的编译和推理。

在适配 RWKV-World 1.5B时，由于模型比较小对逗号比较敏感，导致第一层就炸了精度，最终挂在sampler里面，这个地方我定位2个晚上，后来mlc-ai官方的冯思远告诉我在 MLC-LLM 里如何逐层打印精度之后，我最终定位到了问题。并且在 RWKV 社区里面了解到了这个现象之前就出现过，那就是1.5B的模型第一层需要用FP32来计算，不然会炸精度，我后续实验了RWKV-4-World 3B/7B，这个现象就没有了。

另外，模型的组织格式也是值得注意的一点，并不是在任意位置编译好模型都可以在运行时被 MLC-LLM 正确发现。我大概花了快一周工作外时间在 MLC-LLM 上来支持 RWKV-World 系列模型，工作内容主要为：

• 将大缺弦的 https://github.com/daquexian/faster-rwkv 仓库中的 RWKV World模型tokenizer实现挂到 mlc-ai 的 tokenizers.cpp 中，作为一个 3rd 库提供给MLC-LLM。合并的PR为：https://github.com/mlc-ai/tokenizers-cpp/pull/14。
• 在上面的基础上，在MLC-LLM中支持 RWKV World系列模型的部署，对齐 World 系列模型的 Prompt ，获得良好的对话效果。分别在 Apple M2和A800显卡上进行了部署和测试。PR为：https://github.com/mlc-ai/mlc-llm/pull/848 ，这个pr还wip，如果你现在要使用的话可以直接切到这个pr对应的分支就可以了。
• debug到1.5B RWKV World小模型会炸精度的bug，相当于踩了个大坑。

我要特别感谢 mlc-ai 官方的冯思远在我部署过程中提供的支持以及帮我Review让代码合并到 mlc-ai 社区，以及感谢大缺弦的 RWKV World Tokenizer c++实现以及在编译第三方库时帮我解决的一个bug。

以下是MLC-LLM 部署RWKV World系列模型教程，尽量提供大家部署最不踩坑的实践。

效果：

0x1. 将RWKV-4-World-7B部署在A800上

准备工作

• RWKV-4-World模型地址：https://huggingface.co/StarRing2022/RWKV-4-World-7B
• 下载这里：https://github.com/BBuf/rwkv-world-tokenizer/releases/tag/v1.0.0 的 tokenizer_model.zip并解压为tokenizer_model文件，这是RWKV World系列模型的Tokenizer文件。
• 克隆好 https://github.com/mlc-ai/mlc-llm 和 https://github.com/mlc-ai/relax ，注意克隆的时候一定要加上 –recursive 参数，这样才会把它们依赖的第三方库也添加上。

编译Relax

git clone --recursive git@github.com:mlc-ai/relax.git
cd relax
mkdir build
cd build
cp .