qwen600.cu —— 从零构建的轻量级 CUDA 推理引擎

qwen600 banner

在学习和实践 CUDA 与 GPGPU 的过程中,有了这样一个想法:能不能自己从头实现一个推理引擎?于是我选择了 QWEN3-0.6B 模型 —— 一个小巧但功能完整的模型,可以在 RTX 3050 8GB 显存上顺畅运行。

这个项目的目标很明确:
一方面,作为学习 大语言模型(LLMs)与 Transformer 架构的实验项目;另一方面,也能在不断练习 CUDA 编程的同时,探索极致性能的可能性。

最终产物是一个 静态编译的 QWEN3-0.6B 指令推理引擎(bf16 精度),在基准测试中表现优于一些现有实现:

  • 相比 llama.cpp,快约 8.5%
  • 相比 Hugging Face + flash-attn,快约 292%

功能特性

qwen600 提供了一组简洁但实用的功能:

  • 单批量推理(single batch inference)
  • 静态常量化,编译期优化
  • 全部使用 CUDA C/C++ 实现(Python 仅用于 tokenizer)
  • 依赖极简:仅需 cuBLAS、CUB 和标准 IO
  • 高效的内存管理:

    • 使用 mmap
    • 单 GPU block
    • 异步内存拷贝
  • GPU 权重管理采用零开销指针操作

灵感来源

该项目受到以下开源实现的启发:

架构图

设计理念

qwen600.cu 遵循 suckless philosophy 的精神:
极简、直接、高效,避免无谓的功能堆砌与复杂抽象。

  • 配置集中在源码 config.h
  • 尽可能减少依赖
  • 让用户在阅读与修改时一目了然

快速上手指南

1. 初始准备

首先,下载 QWEN3-0.6B 模型。
官方提供了 huggingface 文档,可快速入门。

下载完成后,请校验权重文件的完整性:

sha256sum <model_dir>/<safetensors-file-name>

正确的输出应为:

f47f71177f32bcd101b7573ec9171e6a57f4f4d31148d38e382306f42996874b

接着,获取项目代码:

git clone https://github.com/yassa9/qwen600
cd qwen600

假设 Hugging Face 模型路径为 <model_dir>,需要将其 tokenizer 转换为本项目可用格式:

python export.py <model_dir>

输出结果中最重要的文件是:tokenizer.bin


2. 构建项目

构建只需要以下环境:

  • CUDA + nvcc
  • cuBLAS + CUB

执行编译:

mkdir build && cd build
cmake .. && make -j$(nproc)

无需额外依赖,过程简洁高效。


3. 运行模型

进入 qwen600/build,直接运行:

./qwen600

会输出完整的参数说明:

usage:   ./qwen600 <model_dir> [options]
example: ./qwen600 <model_dir> -r 1

参数说明

参数 类型 说明
-r int 推理模式,0=普通,1=思维链(reasoning)
-s int 随机种子
-k int top-k 采样,默认 20
-t float temperature,默认 0.6
-p float top-p 采样,默认 0.95
-i string 输入 prompt
-y string 系统提示(chat 模式),默认无

示例运行:

./qwen600 <model_dir> -r 1 -t 0.65 -p 0.9 -k 20

或者直接使用默认参数:

./qwen600 <model_dir> -r 1

官方推荐参数

根据 Hugging Face 官方建议:

  • 思维链模式(enable_thinking=True)
    推荐:Temperature=0.6, TopP=0.95, TopK=20
    避免贪婪解码,否则可能性能下降、出现无限重复
  • 非思维链模式(enable_thinking=False)
    推荐:Temperature=0.7, TopP=0.8, TopK=20, MinP=0

实验展示

示例 1:无思维链模式

./qwen600 <model_dir> -r 0

输入:

>> what is capital of Greece ?

输出:

The capital of Greece is Athens
[231.71 tk/s, 19 tokens in 0.08s]

示例 2:开启思维链模式

示例截图
./qwen600 <model_dir> -r 1

输入:

>> what are llms used for ?

输出会展示推理过程,并最终总结回答,覆盖 LLMs 在客服、写作、教育、医疗等领域的应用。
同时,还会显示生成速度与 tokens 数量。


性能测试

测试环境:

  • 显卡:RTX 3050 8GB + CUDA 13.0
  • CPU:AMD Ryzen 5 3500
  • 内存:16GB RAM
  • 系统:Void Linux

测试问题统一为 what are llms used for ?,运行模式为 THINKING,参数 temp=0

引擎 tokens/sec
hf + flash-attn 29.57
llama.cpp 107.19
qwen600 116.15

结果显示,qwen600 在同等条件下表现最优。


待办事项

项目仍在不断完善中,当前进展如下:

  • [x] RMSnorm 内核融合
  • [x] 跳跃连接融合 cuBLAS
  • [ ] 修复 Softmax 内核与调度器
  • [ ] 探索 RoPE 预计算优化

常见问题(FAQ)

Q: 这个项目适合什么人使用?
A: 更适合学习 CUDA、LLM 内部机制和推理引擎原理的开发者,而非追求开箱即用的应用用户。

Q: 为什么要自己实现,而不是用现成的框架?
A: 一方面为了深入理解原理,另一方面避免不必要的依赖,实现极简和高性能。

Q: 是否支持多 GPU?
A: 当前版本只支持单 GPU,未来可能探索扩展。

Q: 性能为什么会优于 llama.cpp?
A: 静态编译优化与 CUDA 内核融合带来了额外性能提升。