qwen3-moe-aclnn
纯 C++ 实现的 Qwen3-235B-A22B-Instruct BF16 推理运行时，运行于 Ascend 910 × 16 NPU，直接基于 aclnn EAGER 单算子 API（无图编译、无 PyTorch、无 ggml 依赖）。
English version: README.md
性能表现
在 Ascend 910 初代 × 16 NPU（TP=16）上、Qwen3-235B-A22B-Instruct-2507 BF16 权重实测。
所有数字均为质量保持前提下的 TG（输出已人工核验），greedy temperature=0。
| 配置 | TG | 适用 prompt |
|---|---|---|
| 未调优基线 | 12 t/s | 全部 |
| 推荐默认（不开 PLD） | ~27 t/s | 全部 prompt，输出稳定 |
| PLD + degeneration guard | 29-45 t/s | 结构化文本（论述、长回答） |
| PLD 创意类 prompt | 25-40 t/s | 故事 / 多样生成 |
| PLD 事实 / 代码类 prompt | 不稳定（21-95 t/s，方差大） | 不推荐 |
参考：cann-recipes-infer GE graph 方案在同硬件约 54 t/s。本项目未超越该基线——以峰值速度换取 (a) 无需图编译、(b) 无 PyTorch 依赖、(c) 完整的算子调度控制权。
核心优化（按贡献排序）
| 排名 | 优化项 | 收益 | 位置 |
|---|---|---|---|
| 🥇 | HCCL env 调参（AIV + FFTS + TASK_QUEUE=2） |
+89%（12→23 t/s） | scripts/tp_launch.sh |
| 🥈 | Fused RoPE（aclnnApplyRotaryPosEmbV2） |
+17%（23→27 t/s） | include/rope.h |
| 🥉 | Prompt Lookup Decoding（PLD）+ degeneration guard | 适用 prompt 上 +10-60% | src/main_cli.cpp |
| ○ | Device-side topk-w normalize、MoE argsort、cos/sin cache | 累计 ~+15% | include/engine.h |
| ○ | WorkspacePool（thread-local + retain-old） | 降低 malloc 开销 | include/workspace_pool.h |
架构
模型：Qwen3-235B-A22B，94 层，128 experts（top-k=8），GQA（64 Q heads / 4 KV heads），BF16。
并行：TP=16，HCCL ring AllReduce。KV head 每 rank 1 份（4 个 KV heads < 16 ranks，因此每 rank 上的 Q head 0-3 共享 KV head 0）。
执行：aclnn EAGER 模式——所有算子走 aclnn* 单算子 API，配合 workspace 池；无 graph capture、无 GE IR。异步 stream 执行配合 TASK_QUEUE_ENABLE=2 实现 kernel 提交重叠。
Tokenizer：encode 通过 Python 子进程调用 HuggingFace transformers；decode 纯 C++，从导出的 vocab.bin 查表。
单层 forward 数据流
x_in [S, D=4096]
↓
┌── Attention 分支（TP：Q_DIM=512=4h×128，KV_DIM=128=1h×128） ──┐
│ RmsNorm(input_layernorm)
│ linear_hf q_proj / k_proj / v_proj → q, k, v
│ Per-head RmsNorm q_norm, k_norm
│ Fused RoPE：aclnnApplyRotaryPosEmbV2（layout=1, "half"）
│ K、V 追加到每层 KV cache
│ Mask 选择：
│ prefill： 2048×2048 causal + sparse_mode=3
│ decode S=1： mask=nullptr + sparse_mode=0
│ batch decode：[1,1,S,past+S] 自定义 bool mask + sparse_mode=0
│ FIAS（aclnnFusedInferAttentionScore）
│ o_proj linear_hf → per-rank partial
│ HCCL AllReduce（ring + AIV + FFTS） → full
└─────────┘
↓ residual add
┌── MoE 分支 ──┐
│ RmsNorm(post_attention_layernorm)
│ router linear_hf → logits [S, 128]
│ moe_gating_topk_softmax → topk_w[S,8], topk_idx[S,8]
│ Device-side normalize（reduce_sum + adds + cast + div）
│ moe_init_routing_v3 → expanded_x, expanded_ri, tokens_per_expert
│ grouped_matmul_v4 gate/up/down（SwiGLU 激活）
│ Device-side argsort × 2 → 前向置换（避免 host sync）
│ IndexSelect → packed
│ 广播 mul topk_w + ReduceSum axis=1
│ HCCL AllReduce → full
└─────────┘
↓ residual add
x_out
模型权重
目标模型：Qwen3-235B-A22B-Instruct-2507（BF16），safetensors 分片约 470 GB。
下载地址：
- HuggingFace：https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507
- ModelScope：https://www.modelscope.cn/models/Qwen/Qwen3-235B-A22B-Instruct-2507
通过 huggingface-cli 或 modelscope CLI 下载：
# HuggingFace
huggingface-cli download Qwen/Qwen3-235B-A22B-Instruct-2507 --local-dir /path/to/Qwen3-235B-A22B-Instruct-2507-BF16
# ModelScope
modelscope download --model Qwen/Qwen3-235B-A22B-Instruct-2507 --local_dir /path/to/Qwen3-235B-A22B-Instruct-2507-BF16
权重格式：二进制直接读 HuggingFace .safetensors 分片（多 shard mmap）、config.json、tokenizer.json。无需转换——--model-dir 指向下载目录即可。
目录结构：
Qwen3-235B-A22B-Instruct-2507-BF16/
├── config.json
├── tokenizer.json
├── tokenizer_config.json
├── model-00001-of-000XX.safetensors
├── ...
└── model.safetensors.index.json
编译
source /usr/local/Ascend/ascend-toolkit/set_env.sh
cmake -B build
cmake --build build -j8 --target qwen3-moe-aclnn
依赖：
- CANN 8.5.1 或兼容版本
- Python 3 +
transformers+torch_npu（仅 tokenizer 子进程和参考数据生成用） - C++17 编译器
- Ascend 910 × 16 NPU
- nlohmann/json（已打包在
external/json.hpp）
Python 环境设置 — Tokenizer 会调用 Python 子进程。如果 conda / venv 路径与默认不一致，用 QWEN3_PYENV_INIT 覆盖：
export QWEN3_PYENV_INIT="source /opt/my_conda/etc/profile.d/conda.sh && conda activate my_env && "
未设置时，默认尝试 ${HOME}/miniconda3 + env qwen3，并自动 source Ascend toolkit。
快速开始
# 1. 导出 tokenizer 词表到二进制（一次性）
python3 scripts/export_vocab.py /path/to/Qwen3-235B-A22B-Instruct-2507-BF16
# 2. 运行推理（TP=16）
./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn \
--model-dir /path/to/Qwen3-235B-A22B-Instruct-2507-BF16 \
--prompt "The capital of France is" \
--n-predict 100 \
--temperature 0 \
--vocab tokenizer_data/vocab.bin
预期：~27 t/s，输出连贯正确。
按场景推荐的参数
通用默认（稳定、任意 prompt） — 不开 PLD：
./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --temperature 0 --no-stream
结构化 / 长文本（论述、说明） — PLD + guard 提升 +60-90%：
./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --pld --temperature 0 --no-stream
交互式 REPL（多轮对话）：
./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... \
--interactive --chat --temperature 0.7 --top-p 0.8
PLD degeneration guard
Prompt Lookup Decoding 通过让模型一次性 batch verify 若干"draft"token 来加速生成，draft 来自生成历史里的 n-gram 匹配。
已知失效模式：在模型本身就有重复倾向的 prompt（事实问答、代码生成）上，n-gram 会把模型的重复 token 当作 draft 喂回给模型，形成正反馈循环——以 batch 速度加速退化输出。本项目早期曾误报由此类死循环得出的"高 TG"。
本项目的 guard 用两条启发式拦截可疑 draft：
- low-distinct：draft 中 distinct token 数 < 阈值 → 拒绝
- tail-echo：最后 N 个 hist token 全部等于 draft[0] → 拒绝
被拒 draft 走单 token decode fallback。生成末尾如出现 8 个连续相同 token，stderr 打印一次 [warn] 提示。
参数：
--pld 启用 PLD（opt-in）
--pld-k N draft 窗口（默认 10）
--pld-ngram N n-gram 匹配长度（默认 1，带多级回退）
--pld-min-hist N hist 达到 N 个 token 前跳过 PLD（默认 20）
--pld-no-guard 关闭 degeneration guard（危险，可能产生死循环）
--pld-guard-distinct N draft distinct token 最小值（默认 3）
--pld-guard-tail N tail-echo 检测窗口（默认 6）
--pld-loop-warn N 生成 N 个连续相同 token 时报警（默认 8）
诚实 benchmark：用 scripts/bench_pld_safe.sh，它会自动把每次 run 的输出分类为 OK / LOOP_N / LOW_DIVERSITY，并分别统计 OK-only 与 degraded 的 TG。
正确性验证
15+ 单元 / 集成测试，与 Python（HuggingFace Transformers）参考对比：
./build/test_attention_layer # rel=4.9e-4 vs Python prefill
./build/test_attention_decode # rel=0（bit-exact）
./build/test_moe_layer # rel=3.6e-3
./build/test_layer_forward # 完整单层
./build/test_runner # 多层 runner
./build/test_rope_fused # aclnnApplyRotaryPosEmbV2 vs 手写 HF rotate_half
./build/test_batch_decode # S=1..8 耗时
./build/test_batch_correctness # argmax 一致性
./build/test_op_support # 910 特定 op 可用性探针
# 集成冒烟：
./tests/test_chat_flow.sh # 7/7 PASS
测试期望 tests/<name>_data/ 下存放参考数据，由 scripts/gen_*_reference.py 生成。各脚本顶部有 docstring 说明。
环境变量调优（tp_launch.sh 自动应用）
HCCL_WHITELIST_DISABLE=1
HCCL_ALGO=level0:ring # ring，非 fullmesh（fullmesh 会让输出乱码）
HCCL_BUFFSIZE=200 # sweet spot；100 和 400 都更慢
HCCL_OP_EXPANSION_MODE=AIV # 关键：AI Vector cores 参与 reduce 调度
HCCL_OP_BASE_FFTS_MODE_ENABLE=1 # 关键：Fast Frequently-used Transfer Scheduling
TASK_QUEUE_ENABLE=2 # 关键：激进异步任务入队
三个"关键"env 任何一个去掉都会让 TG 降 20-40%。
目录结构
include/
├── acl_common.h RAII 包装、DeviceBuffer、make_contig_tensor
├── aclnn_ops.h 单算子 wrapper + WorkspacePool 集成
├── acl_runtime.h AclRuntime（device + stream 管理）
├── device_weights.h safetensors → device 加载 + TP 切分
├── engine.h attention_forward + moe_forward + RopeCache
├── hccl_comm.h HCCL init + allreduce + broadcast
├── model_config.h Qwen3 超参 + compute_derived
├── rope.h apply_rope_fused（aclnnApplyRotaryPosEmbV2 wrapper）
├── runner.h Runner 类（prefill/decode/decode_batch/rewind/profile）
├── safetensors_loader.h 多 shard safetensors mmap parser
├── tokenizer.h vocab decode + Python 子进程 encode
└── workspace_pool.h thread-local aclnn workspace 池（retain-old）
src/
├── device_weights.cpp load_attention（GQA 修复）、load_moe（permute sync 修复）
├── main_cli.cpp CLI 入口 + PLD 主循环 + degeneration guard + 多轮对话
├── model_config.cpp compute_derived（GQA KV 切分）
├── runner.cpp Runner 实现（build_batch_decode_mask_ 等）
├── safetensors_loader.cpp
└── tokenizer.cpp
scripts/
├── tp_launch.sh 产线启动器（自动设置 HCCL env）
├── bench_tg.sh 稳定 N-run TG 测量
├── bench_pld_safe.sh 带输出正确性分类器的 PLD benchmark
├── bench_hccl[_adv].sh HCCL 参数 sweep
├── bench_pld[_k].sh PLD K × ngram sweep（旧版，优先用 bench_pld_safe.sh）
├── export_vocab.py 从 HF tokenizer 导出 vocab.bin
└── gen_*_reference.py 逐 op 的 Python 参考数据生成器
tests/
├── test_attention_* attention 正确性（prefill / decode）
├── test_moe_layer MoE 正确性
├── test_layer_forward 完整单层
├── test_runner 多层 Runner
├── test_rope_fused fused RoPE vs 手写 HF
├── test_batch_* batch decode 耗时 + 正确性
├── test_op_support 910 特定 op 可用性探针
└── test_chat_flow.sh 端到端集成冒烟
CLI 参数参考
--model-dir <path> (必填) HF safetensors 目录
--prompt "<text>" prompt 文本
--prompt-file FILE 从文件读 prompt（避免 shell 转义）
--n-predict N 最大生成 token 数
--tp-size N tensor parallelism（也可通过 TP_SIZE env 设置）
--max-seq N KV cache + 上下文上限（默认 512）
--temperature F 0 = greedy；典型 0.7
--top-k N 0 = 禁用
--top-p F 1.0 = 禁用
--seed N 0 = 基于时间
--chat 应用 Qwen3 chat 模板
--system "<text>" system role（配合 --chat）
--interactive, -i REPL 模式（配合 --chat 实现多轮记忆）
--reset 强制无状态 REPL（每轮重置 KV）
--no-stream 批量打印最终文本，不逐 token 流式输出
--vocab <path> vocab.bin 路径（默认 tokenizer_data/vocab.bin）
--pld* 见上方 "PLD degeneration guard" 章节
已知限制
- 未达到 cann-recipes GE graph 54 t/s 基线（当前稳定 ~27 t/s，PLD 场景最高 ~45 t/s）。
要追上基线需要以下之一：(a) 真正的图编译；(b) 融合集合算子（
MatmulAllReduce、GroupedMatmulAllReduce）——910 初代没有；(c) 迁移到 910B / A2 / A3 硬件。 - **仅支持
tp_size∈ {1, 2, 4, 8, 16}**。不能整除 64 Q heads 的值会报错。 - PLD 在事实 / 代码类 prompt 上不可靠——要么产出基线 TG（guard 拒绝了绝大多数 draft），要么进入低强度退化（classifier 未必能抓到）。用
bench_pld_safe.sh诚实评估。 - Tokenizer 依赖 Python 子进程——首次 encode 有 ~1s 启动开销。默认 conda 路径不匹配时用
QWEN3_PYENV_INIT覆盖。 - NPU 性能 run-to-run 方差巨大（某些配置下高达 4×），源于 BF16 + MoE 固有非决定性与硬件资源共享。报告数字时用 ≥5 runs 的中位数。
下一步方向（按优先级）
- Draft Model Speculative Decoding（Qwen3-0.6B）——比 n-gram PLD 接受率稳定得多，预期跨 prompt 类型 +60-100% TG（1-2 周工程量）。
- HCCL AllReduce / compute 重叠——理论上 +10-15%，受 EAGER 路径串行依赖限制。
- KV cache INT8 量化——降低 memory-bandwidth 压力，长上下文场景 +15-25%（需先验证 910 初代算子支持）。
- W8 权重量化——若 910 初代有 aclnn 量化 kernel 可用，+10-20%。
不推荐：
aclmdlRIstream-capture 图记录（POC 证明上限仅 1.13×，工程成本不值）。- 自研 AscendC 融合算子（高维护成本，除非有专职 kernel 工程师）。
- torchair / torch.compile 迁移（破坏 pure-C++ 设计）。
文档
docs/optimization-summary-zh.md— 阶段性优化总结（中文）：关键优化原因、PLD 正确性边界、项目级教训docs/next-steps-draft-model-speculative.md— Draft Model Speculative Decoding（Qwen3-0.6B）执行规格：M1-M4 里程碑、正确性测试协议、风险兜底
License
Apache License 2.0，见 LICENSE。