LiteType-Engine

完全离线的中文文本预测引擎，面向输入法场景。隐私优先、超低延迟、CUDA 加速。

专为中文输入法设计：Batch=1，最大上下文 64 词，单 token 预测延迟 < 5ms。

特性

• 📦 体积极小：GPU 显存 < 10 MB（FP16），可执行文件 < 1 MB
• ⚡ 超低延迟：单 token 约 1-3 ms（远低于 5ms 目标）
• 🔒 完全离线：无需网络，无需云端，隐私数据不出设备
• 🀄 中文优化：jieba 分词 + 词级别预测，同时支持英文
• 🧠 词级别预测：词表包含 6200+ 单字和 1700+ 常用词，直接预测完整词汇
• 🚀 CUDA 优化：融合 Linear+Softmax 算子、静态环形 KV Cache、手写核函数
• 🧪 数值验证：所有 CUDA 算子与 PyTorch 参考实现逐元素对齐

架构

用户输入 → jieba分词 → Token Embed → Transformer×4 → Fused Linear+Softmax → Top-5 候选词
                            ↑ 静态环形 KV Cache（零运行时分配）↑

模型规格

组件	详情
架构	词级别 GPT（仅解码器 Transformer）
层数	4
隐藏维度	256
注意力头数	4（每头 64 维）
前馈维度	1024
最大序列长度	64
词表大小	8000（6200 单字 + 1800 常用词）
总参数量	~727 万
精度	FP16（权重文件约 13 MB）
激活函数	GELU（tanh 近似）
位置编码	可学习的绝对位置嵌入

CUDA 算子优化

算子	优化策略
融合 Linear+Softmax+TopK	单 kernel 完成：矩阵乘 → Softmax → Top-K，中间结果全部在寄存器，不写全局显存
多头自注意力	静态环形缓冲区 KV Cache，预分配，零运行时 malloc
LayerNorm	Warp 级规约 + 共享内存
前馈网络	单 kernel 融合 Linear→GELU→Linear

快速开始

环境要求

• CMake ≥ 3.18（或直接使用 nvcc）
• CUDA Toolkit ≥ 11.6
• GPU 计算能力 ≥ 7.5（Turing 及以上架构）
• Python 3.10+（仅训练需要）
• PyTorch ≥ 2.0、jieba、datasets（仅训练需要）

直接使用（交互式 IME）

# 安装依赖（仅首次）
pip install torch jieba

# 运行词级别 IME 交互程序
python scripts/ime_word.py

操作方式：

• 直接打字 → 每输入一个字，实时显示 5 个候选词
• 1-5 → 选择对应候选词
• Tab → 自动补全（选择第一个候选）
• 退格 → 删除上一个字
• 输入标点符号（。，！？） → 自动重置上下文
• Ctrl+C → 退出

训练自己的模型

# 1. 安装依赖
pip install -r scripts/requirements.txt
pip install jieba

# 2. 下载中文维基百科语料（约 500 万字符）
python scripts/download_corpus.py

# 3. 清洗语料（按标点分割句子，去除噪声）
# 自动生成 data/corpus_final.txt

# 4. 生成对话语料（补充维基百科缺失的日常表达）
python scripts/gen_conv.py

# 5. 合并语料
# 自动将对话语料与维基百科合并为 corpus_combined.txt

# 6. 构建词级别词表
python scripts/train_word_tokenizer.py --corpus data/corpus_combined.txt --prepare-data

# 7. 训练模型（GPU 约 30 分钟 / 20000 步）
python scripts/train_word.py --steps 20000 --device cuda

# 8. 导出为 C++ 推理格式（可选）
python scripts/export_weights.py

编译 C++ 推理引擎

# 直接用 nvcc 编译（开发推荐）
nvcc -std=c++17 -O2 -o build/litetype.exe \
  src/main.cpp src/tokenizer/tokenizer.cpp src/model/weights.cpp \
  src/cuda/kernels.cu src/cuda/inference.cu src/engine/engine.cpp \
  -I src

# 或通过 CMake
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
cmake --build . --config Release

项目结构

LiteType-Engine/
├── README.md
├── .gitignore
├── CMakeLists.txt
├── scripts/
│   ├── train.py                  # 字级别模型训练
│   ├── train_word.py             # 词级别模型训练（jieba 分词）
│   ├── train_word_tokenizer.py   # 词级别词表构建
│   ├── gen_conv.py               # 对话语料生成
│   ├── download_corpus.py        # 维基百科语料下载
│   ├── export_weights.py         # PyTorch → FP16 二进制导出
│   ├── ime_word.py              # ⭐ 词级别 IME 交互程序
│   ├── ime_demo.py              # 字级别 IME 交互程序
│   ├── interact.py              # 基础交互演示
│   └── requirements.txt         # Python 依赖
├── src/
│   ├── tokenizer/                # C++ UTF-8 字符级分词器
│   ├── model/                    # 权重加载 & 模型配置
│   ├── cuda/                     # CUDA 核函数 & 推理引擎
│   ├── engine/                   # 高层推理引擎封装
│   └── main.cpp                  # C++ 交互式命令行演示
├── tests/
│   ├── test_tokenizer.cpp        # 分词器单元测试
│   ├── test_math.cu              # CUDA 数值验证测试
│   └── test_inference.cpp        # 端到端集成测试
└── data/
    ├── vocab_word.txt            # 词级别词表（8000 token）
    ├── model_final.pt            # PyTorch 模型权重
    ├── model_fp16.bin            # FP16 二进制权重（C++ 推理用）
    └── corpus_combined.txt       # 合并训练语料

性能指标

指标	实测值	目标
单 token 延迟（CUDA）	~62 µs	< 5 ms ✅
单 token 延迟（PyTorch）	~1-3 ms	< 5 ms ✅
GPU 显存占用	~10 MB	< 20 MB ✅
模型权重文件	13 MB	< 20 MB ✅
可执行文件大小	< 1 MB	< 1 MB ✅
参数量	7,267,840	~7.3M

运行测试

# 分词器测试
nvcc -std=c++17 -O2 -o build/test_tokenizer.exe \
  tests/test_tokenizer.cpp src/tokenizer/tokenizer.cpp -I src
./build/test_tokenizer

# CUDA 算子数值验证
nvcc -std=c++17 -O2 -o build/test_math.exe \
  tests/test_math.cu src/cuda/kernels.cu src/model/weights.cpp -I src
./build/test_math

# 端到端集成测试
nvcc -std=c++17 -O2 -o build/test_inference.exe \
  tests/test_inference.cpp src/tokenizer/tokenizer.cpp \
  src/model/weights.cpp src/cuda/kernels.cu src/cuda/inference.cu \
  src/engine/engine.cpp -I src
./build/test_inference

核心技术决策

决策	选择	理由
数学库	手写 CUDA 核函数	零 cuBLAS 调用开销，支持算子融合
KV Cache	静态环形缓冲区，预分配	零运行时内存分配，无碎片
精度	FP16 计算 + FP32 累积	兼顾速度与数值稳定性
激活函数	GELU（tanh 近似）	与 PyTorch 推理模式一致
分词方式	jieba 词级别分词	支持单字 + 常用词混合预测
采样策略	Top-5 概率排序	输入法场景 5 个候选即可
训练语料	维基百科 + 对话模板	兼顾通用知识与日常表达

预测效果示例

输入	预测候选	置信度
我们	一起 / 以为 / 是	87%
你说	的话 / 是 / 的	100%
今天天气	不错 / 很有意思 / 太	36%
我觉得	很好 / 不错 / 人际关系	40%
他很	特别 / 实在 / 很	39%
我的	手机 / 父母 / 脾气 / 身体	4%

参考资料

• Karpathy nanoGPT — 字符级 GPT 训练
• NVIDIA Kernel Fusion — 算子融合原理
• HuggingFace KV Cache — KV 缓存优化
• 模型量化 — FP16/INT8 精度对比

许可证

MIT