| # qwen3-moe-aclnn |
|
|
| 纯 C++ 实现的 **Qwen3-235B-A22B-Instruct** BF16 推理运行时，运行于 **Ascend 910 × 16 NPU**，直接基于 aclnn EAGER 单算子 API（无图编译、无 PyTorch、无 ggml 依赖）。 |
|
|
| English version: [README.md](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 下载： |
| ```bash |
| # 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 |
| ``` |
|
|
| --- |
|
|
| ## 编译 |
|
|
| ```bash |
| 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` 覆盖： |
| ```bash |
| export QWEN3_PYENV_INIT="source /opt/my_conda/etc/profile.d/conda.sh && conda activate my_env && " |
| ``` |
| 未设置时，默认尝试 `${HOME}/miniconda3` + env `qwen3`，并自动 source Ascend toolkit。 |
|
|
| --- |
|
|
| ## 快速开始 |
|
|
| ```bash |
| # 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： |
| ```bash |
| ./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --temperature 0 --no-stream |
| ``` |
|
|
| **结构化 / 长文本（论述、说明）** — PLD + guard 提升 +60-90%： |
| ```bash |
| ./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --pld --temperature 0 --no-stream |
| ``` |
|
|
| **交互式 REPL（多轮对话）**： |
| ```bash |
| ./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： |
|
|
| 1. **low-distinct**：draft 中 distinct token 数 < 阈值 → 拒绝 |
| 2. **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）参考对比： |
|
|
| ```bash |
| ./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` 自动应用） |
| |
| ```bash |
| 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 的中位数。 |
| |
| --- |
| |
| ## 下一步方向（按优先级） |
| |
| 1. **Draft Model Speculative Decoding**（Qwen3-0.6B）——比 n-gram PLD 接受率稳定得多，预期跨 prompt 类型 +60-100% TG（1-2 周工程量）。 |
| 2. **HCCL AllReduce / compute 重叠**——理论上 +10-15%，受 EAGER 路径串行依赖限制。 |
| 3. **KV cache INT8 量化**——降低 memory-bandwidth 压力，长上下文场景 +15-25%（需先验证 910 初代算子支持）。 |
| 4. **W8 权重量化**——若 910 初代有 aclnn 量化 kernel 可用，+10-20%。 |
| |
| 不推荐： |
| - `aclmdlRI` stream-capture 图记录（POC 证明上限仅 1.13×，工程成本不值）。 |
| - 自研 AscendC 融合算子（高维护成本，除非有专职 kernel 工程师）。 |
| - torchair / torch.compile 迁移（破坏 pure-C++ 设计）。 |
| |
| --- |
| |
| ## 文档 |
| |
| - [`docs/optimization-summary-zh.md`](docs/optimization-summary-zh.md) — 阶段性优化总结（中文）：关键优化原因、PLD 正确性边界、项目级教训 |
| - [`docs/next-steps-draft-model-speculative.md`](docs/next-steps-draft-model-speculative.md) — Draft Model Speculative Decoding（Qwen3-0.6B）执行规格：M1-M4 里程碑、正确性测试协议、风险兜底 |
| |
| --- |
| |
| ## License |
| |
| Apache License 2.0，见 `LICENSE`。 |
| |
