本项目演示 Qwen3-0.6B ONNX 模型通过 onnxruntime 加载并进行自回归推理的完整实现。
模型下载&部署源码获取
shell
pip install modelscope==1.37.1
modelscope download --model KeanuX/Qwen3-0.6B-ONNX --local_dir ./模型目录结构
Qwen3-0.6B-ONNX/
├── model.onnx # 原始 ONNX 模型（~1.22 GB，fp16）
├── model_no_isnan.onnx # 移除 IsNaN 算子后的 ONNX 模型（已优化，推理用）
├── config.json # 模型架构配置
├── generation_config.json # 生成参数默认配置
├── tokenizer.json # Tokenizer 词表 & 编码规则
├── tokenizer_config.json # Tokenizer 行为配置
├── special_tokens_map.json # 特殊 token 映射
├── added_tokens.json # 附加 token 定义
├── vocab.json # 词表
├── merges.txt # BPE merge 规则
└── chat_template.jinja # Chat 模板（用于构建 prompt）ONNX 模型：
| 文件 | 说明 |
|---|---|
model.onnx |
optimum 导出的原始图，attention 层 Softmax 后含 IsNaN + Where 节点 |
model_no_isnan.onnx |
已移除 IsNaN/Where 冗余算子，直接以 Softmax → MatMul 连接，推理时使用 |
Tokenizer 文件： 可直接传给 transformers.AutoTokenizer.from_pretrained(...) 加载。
部署算法实现
推理引擎由 llm_loader.py 中的 QwenLoader 类实现。
初始化
python
model = QwenLoader(model_path, logger, device="cuda", system_prompt="你是智能助手小智")初始化过程：
- 通过
AutoTokenizer.from_pretrained加载同目录下的 tokenizer - 通过
onnxruntime.InferenceSession加载model_no_isnan.onnx，支持 CPU/CUDA 两种执行后端 - 自动解析模型的 IO 规格：输入/输出名称列表、层数
num_layers、head 数num_heads、head 维度head_dim
自回归生成流程
generate(user_text, max_new_tokens, enable_thinking) 分为三个阶段：
第一阶段：Prompt 编码 & Prefill
将用户输入包装为 system + user 格式的 chat template，通过 tokenizer 编码得到 input_ids（shape 1 × seq_len）。连同 attention_mask、position_ids、以及各层全零的 past_key_values 送入模型，一次性计算完整 prompt 的 logits，并获得首轮的 present KV-Cache。
输入: input_ids(1, seq_len) + attention_mask + position_ids + past KV (zeros)
输出: logits(1, seq_len, vocab_size) + present KV × 28 层取 logits[0, -1, :] 的 argmax 作为第一个生成 token。
第二阶段：自回归 Decode
循环执行，每步：
-
以当前 token 构造
input_ids（shape1 × 1），attention_mask随已生成长度逐步扩展，position_ids为当前绝对位置 -
同时传入上一轮输出的
presentKV-Cache 作为本轮past_key_values -
模型输出当前步的 logits 和更新后的 KV-Cache
-
argmax 选取下一 token，继续循环，直到命中 EOS 或达到
max_new_tokens循环:
input: input_ids(1,1) + attention_mask + position_ids + past KV
output: logits(1,1,vocab_size) + present KV
token = argmax(logits) → 更新 past KV ← present KV
第三阶段：Decode & 后处理
- 将生成的 token id 序列通过 tokenizer 解码
- 解析
</think>特殊 token（id=151668），将输出分离为 thinking content 和最终回答
输入/输出规格
| 类型 | 名称 | Shape | 说明 |
|---|---|---|---|
| 输入 | input_ids |
(1, seq_len) |
token id 序列 |
| 输入 | attention_mask |
(1, seq_len) |
注意力掩码 |
| 输入 | position_ids |
(1, seq_len) |
位置编码索引 |
| 输入 | past_key_values.{i}.key |
(1, num_heads, past_len, head_dim) |
第 i 层 key 缓存 |
| 输入 | past_key_values.{i}.value |
(1, num_heads, past_len, head_dim) |
第 i 层 value 缓存 |
| 输出 | logits |
(1, seq_len, vocab_size) |
预测 logits |
| 输出 | present.{i}.key |
(1, num_heads, total_len, head_dim) |
更新后的 key 缓存 |
| 输出 | present.{i}.value |
(1, num_heads, total_len, head_dim) |
更新后的 value 缓存 |
- 层数
i = 0 ... 27（共 28 层） num_heads = 8（GQA，key/value heads），head_dim = 128- 输入共 59 个，输出共 57 个
运行测试
test_run.py 加载模型并执行一次推理：
bash
python test_run.py
python
from model_classes.llm_loader import QwenLoader
model_loader = QwenLoader("./hf_models/Qwen3-0.6B-ONNX/", logger, device="cuda")
thk_con, cus_con = model_loader.generate("你是谁", max_new_tokens=256, enable_thinking=False)日志输出到 ./logs/run.log。
性能测试
对 model.onnx 和 model_no_isnan.onnx 在 CPU / CUDA 两种后端下进行基准对比（prompt 长度 29 tokens，decode 128 步，warmup 2 轮后取 3 轮均值）。
CPU 结果
| 指标 | model.onnx |
model_no_isnan.onnx |
差异 |
|---|---|---|---|
| Prefill 耗时 | 64.32 ms | 64.20 ms | -0.18% |
| Prefill 吞吐 | 450.89 tok/s | 451.72 tok/s | +0.18% |
| Decode 每步耗时 | 52.99 ms | 50.95 ms | -3.85% |
| Decode 吞吐 | 18.87 tok/s | 19.63 tok/s | +4.00% |
CUDA 结果
| 指标 | model.onnx |
model_no_isnan.onnx |
差异 |
|---|---|---|---|
| Prefill 耗时 | 8.23 ms | 8.23 ms | +0.02% |
| Prefill 吞吐 | 3522.76 tok/s | 3522.15 tok/s | -0.02% |
| Decode 每步耗时 | 8.78 ms | 8.62 ms | -1.87% |
| Decode 吞吐 | 113.88 tok/s | 116.05 tok/s | +1.91% |
分析：
- CPU 下 Decode 提升 ~4%，每步节省 ~2ms；Prefill 几乎无变化。
- CUDA 下 GPU 算力更强，单步延迟远低于 CPU （~8.6ms vs ~51ms），但相对提升幅度缩小到 ~2%，因为
IsNaN/Where在 GPU 上的计算成本比例更低。 - 两种后端下
model_no_isnan.onnx均有正向收益，无精度损失。
模型配置参数
| 参数 | 值 |
|---|---|
| 架构 | Qwen3ForCausalLM |
| 层数 | 28 |
| 隐藏维度 | 1024 |
| 中间层维度 | 3072 |
| Attention Heads | 16 (Q) / 8 (KV, GQA) |
| Head Dim | 128 |
| 词表大小 | 151,936 |
| 激活函数 | SiLU |
| Norm | RMSNorm (eps=1e-6) |
| RoPE theta | 1,000,000 |
| 最大位置 | 40,960 |
| 精度 | float16 |