前阵子帮人调 Mixtral-8x7B 在昇腾 NPU 上的推理性能，发现一个怪事：同样的 FlashAttention 算子，在 Llama-2-7B 上跑得飞快，在 Mixtral 上却慢了将近一倍。查了一圈，发现瓶颈不在 FlashAttention 本身——FlashAttention 算完注意力之后，输出的 token 要送进 8 个专家网络（Expert），路由选择和专家计算之间有一大段显存读写，这些读写才是慢的元凶。

ops-transformer 仓库里有个专门解决这个问题的算子：MoE 融合算子。它把路由选择和注意力计算的输出搬运融合到一起，省掉了中间的显存来回搬。今天咱们就把 MoE 模型里 FlashAttention 的特殊之处聊清楚。

MoE 模型跟普通 Transformer 有什么区别？

先花两分钟搞懂 MoE 的结构，不然后面的优化看不懂。

普通的 Transformer 模型（比如 Llama-2-ShiftB），每一层有两组 FFN（前馈网络），所有 token 都走同一组 FFN 计算：

token → Attention → FFN → 输出
         ↑
   所有 token 共享同一个 FFN

MoE 模型不一样，它有多个 FFN（叫“专家”），每个 token 只送给其中 1-2 个专家算：

token → Attention → 路由选择 → Expert2 + Expert5 → 输出
         ↑             ↑            ↑
      所有token共享   选top-2     只算2个专家

Mixtral-8x7B 有 8 个专家，每个 token 只选 top-2，所以实际参与计算的是 2 个专家的参数，另外 6 个闲着。

• MoE 的好处：模型总参数量大（8x7B=47B），但每个 token 只激活 2 个专家（14B），推理成本跟 14B 模型差不多。
• MoE 的麻烦：路由选择和专家计算之间，要把 token 按专家分组、搬过去算、再搬回来合并。这个“搬来搬去”就是性能瓶颈。

FlashAttention 在 MoE 模型里的位置

在 MoE 模型的一层里，计算流程是这样的：

1. FlashAttention（所有 token 共享，跟普通模型一样）
2. 路由选择（每个 token 算一个路由分数，选 top-2 专家）
3. 按 expert 分组（把分配给同一个 expert 的 token 挑出来）
4. 专家计算（8 个 expert 分别算自己的 FFN）
5. 合并结果（把 8 个 expert 的输出按路由权重加权合并）

FlashAttention 在第 1 步，它本身跟普通模型没有任何区别——输入是所有 token 的 hidden_states，输出也是所有 token 的 hidden_states。FlashAttention 不知道也不关心后面有 MoE。

问题出在第 2-5 步。

瓶颈在哪：FlashAttention 输出之后的三次显存搬运

标准实现里，FlashAttention 算完之后，到 MoE 的专家计算完成，中间要经历这些显存操作：

• 步骤 2（路由选择）：
  • 读 FlashAttention 的输出（全部 token）→ HBM 读 1 次
  • 算路由分数 → SRAM 里算
  • 写路由分数 → HBM 写 1 次
• 步骤 3（按 expert 分组）：
  • 读路由分数 → HBM 读 1 次
  • 把 token 按 expert 分组，写成分散的 tensor → HBM 写 8 次（每个 expert 1 次）
• 步骤 4（专家计算）：
  • 每个 expert 读自己的 token → HBM 读 8 次
  • 每个 expert 算 FFN → SRAM 里算
  • 每个 expert 写结果 → HBM 写 8 次
• 步骤 5（合并结果）：
  • 读 8 个 expert 的输出 → HBM 读 8 次
  • 按路由权重加权合并 → SRAM 里算
  • 写最终结果 → HBM 写 1 次

总计：18 次 HBM 读 + 18 次 HBM 写（路由分数 1 读+1 写，分组 1 读+8 写，专家 8 读+8 写，合并 8 读+1 写）。

而 FlashAttention 本身只要 3 次读 + 1 次写。MoE 的显存操作是 FlashAttention 的 9 倍，这就是为什么 MoE 模型的 FlashAttention 跑起来慢——不是 FlashAttention 慢，是 FlashAttention 之后的那堆搬运拖了后腿。

ops-transformer 的 MoE 融合算子：把三次搬运合成一次

ops-transformer 仓库里的 MoE 融合算子，核心思路是：把路由选择、按 expert 分组、合并结果这三个步骤融合成一个 Kernel，避免中间结果写回 HBM。

具体做了三件事：

融合一：路由选择 + 按 expert 分组

标准实现里，路由选择和分组是两个步骤：先算路由分数，写回 HBM，再读路由分数来分组。
融合后，路由分数直接在 SRAM 里算，算完立刻用路由分数做分组，不写回 HBM。分组的结果直接写到每个 expert 对应的 SRAM 区域里。

# 标准：两步，中间写 HBM
路由分数 = Router(FlashAttention输出)     # 算完写 HBM
分组结果 = GroupByExpert(路由分数, 输出)  # 从 HBM 读路由分数，分组后写 HBM

# 融合：一步，不写 HBM
分组结果 = FusedRouterAndGroup(FlashAttention输出)  # 路由+分组在 SRAM 里完成

省掉了 2 次读 + 8 次写（路由分数的 1 读 + 分组的 1 读 + 8 写 → 0）。

融合二：expert 计算的输出 + 合并结果

标准实现里，每个 expert 算完 FFN 之后，结果写回 HBM，然后合并步骤再从 HBM 读回来。
融合后，expert 的输出直接写到 SRAM 里的一个“累加缓冲区”，合并步骤在 SRAM 里就地完成。

# 标准：两步，中间写 HBM
for expert in experts:
    expert_output = FFN(expert_input)  # 写 HBM
merged = Merge(expert_outputs)         # 从 HBM 读 8 次

# 融合：一步，不写 HBM
accum = zeros(sram)
for expert in experts:
    expert_output = FFN(expert_input)
    accum += route_weight * expert_output  # 直接在 SRAM 里累加

省掉了 8 次读 + 8 次写（expert 输出的 8 写 + 合并的 8 读 → 0）。

融合三：FlashAttention 输出 → 路由选择的衔接

这个是最精细的融合。FlashAttention 的输出本来要写回 HBM，然后路由选择再从 HBM 读。
ops-transformer 的实现里，FlashAttention 的输出直接留在 SRAM 里，路由选择的 Kernel 从 SRAM 里直接读，不用走 HBM。

# 标准：FlashAttention 写 HBM，路由选择从 HBM 读
attn_output = FlashAttention(Q, K, V)  # 写 HBM
route_scores = Router(attn_output)       # 从 HBM 读

# 融合：FlashAttention 输出留在 SRAM，路由选择从 SRAM 读
attn_output = FlashAttention(Q, K, V)  # 留在 SRAM
route_scores = Router(attn_output)       # 从 SRAM 读（20倍快）

省掉了 1 次读 + 1 次写（FlashAttention 输出的 1 写 + 路由选择的 1 读 → 0）。

融合后的总效果

操作	标准实现	融合后	省了多少
HBM 读	18 次	5 次	72%
HBM 写	18 次	2 次	89%

HBM 读写次数从 36 次降到 7 次，减少了 80%。在显存带宽瓶颈的场景下，这直接等于 80% 的性能提升。

在昇腾 NPU 上实际跑出来的性能数据

我测了一组 Mixtral-8x7B 在 Atlas 800T A2 上的数据（8 卡 Tensor Parallel，FP16）：

配置	延迟 (ms/token)	吞吐 (tokens/s)	显存占用 (GB/卡)
标准实现（FlashAttention + 3步MoE）	38	26	9.2
MoE 融合（FlashAttention + 融合MoE）	21	48	6.1
MoE 融合 + INT8 量化	15	67	3.8

结论：MoE 融合让吞吐提升了 85%，显存省了 34%。加上 INT8 量化，吞吐能到 67 tokens/s，显存只占 3.8GB/卡（8 卡总共 30.4GB，32GB 的卡刚好能跑）。

跟 Llama-2-7B 的对比

模型	激活参数	吞吐 (tokens/s)	吞吐/参数
Llama-2-7B	7B	89	12.7
Mixtral-8x7B (融合)	14B	48 (8卡)	3.4/卡

MoE 模型的每卡吞吐比密集模型低 73%，但考虑到 Mixtral 的效果接近 47B 密集模型，这个 trade-off 是划算的。

跟 NVIDIA A100 的对比

我也在 A100 上跑了一组对比数据（8 卡，FP16，Mixtral-8x7B）：

指标	Ascend 910 (MoE 融合)	A100 80GB (MoE 融合)	比例
吞吐 (tokens/s)	48	85	0.56x
显存占用 (GB/卡)	6.1	5.8	1.05x
最大 batch_size	16	24	0.67x

差距分析：吞吐差 44%，主要还是 HBM 带宽的差距（1200 vs 1935 GB/s）。MoE 融合算子虽然是带宽密集型的，但 80% 的 HBM 读写已经被融合省掉了，剩下的 20% 还是受带宽限制。

有意思的发现：A100 上的 MoE 融合收益（相比标准实现）只有 60%，而 Ascend 910 上是 85%。因为 Ascend 910 的带宽更紧张，融合的收益更明显。带宽越低，减少 HBM 读写带来的性能提升越大。

在 vLLM 里开启 MoE 融合

vLLM 的昇腾适配已经支持 MoE 融合，启动的时候加一个环境变量：

# 开启 MoE 融合
export VLLM_USE_FUSED_MOE=1

python -m vllm.entrypoints.openai.api_server \
  --model ./models/Mixtral-8x7B-v0.1 \
  --tensor-parallel-size 8 \
  --enable-flash-attn \
  --max-model-len 4096

⚠️ 踩坑预警：VLLM_USE_FUSED_MOE=1 要求 ops-transformer 的 MoE 融合算子已经编译并安装。你要是没装，vLLM 会静默降级到标准实现，不会报错，但性能会差很多。启动的时候看日志里有没有 INFO: Using fused MoE kernel 这行，有就是开了，没有就是没开。

手动编译 MoE 融合算子

如果你不想用 vLLM，想直接调 MoE 融合算子，得手动编译 ops-transformer：

# 拉取仓库
git clone https://atomgit.com/cann/ops-transformer.git
cd ops-transformer

# 编译 MoE 融合算子
cd src/moe_fusion
bash build.sh --soc Ascend910 --typ release

# 安装
chmod +x ./output/moe_fusion_Ascend910.run
sudo ./output/moe_fusion_Ascend910.run

编译完之后，在 Python 里这样调：

import torch
import torch_npu
from torch_npu.contrib.functional import npu_moe_fusion

# FlashAttention 先算注意力
attn_output = npu_flash_attention(q, k, v, head_num=32, input_layout="BNSD")

# MoE 融合算子：路由+分组+FFN+合并，一步搞定
# router_weight: 路由权重矩阵 [hidden_dim, num_experts]
# expert_weights: 8 个 expert 的 FFN 权重
moe_output = npu_moe_fusion(
    attn_output,
    router_weight=router_weight,
    expert_weights=expert_weights,
    num_experts=8,
    top_k=2,
    activation="silu"
)

⚠️ 踩坑预警：npu_moe_fusion 的 expert_weights 参数需要是 8 个 expert 的权重拼在一起的大 tensor（形状 [8, hidden_dim, ffn_dim]），不能是 8 个独立的 tensor。你要是从 HuggingFace 的 Mixtral 模型里加载权重，得先把 8 个 expert 的权重 cat 起来：

# 从 HuggingFace 加载
from transformers import MixtralForCausalLM
model = MixtralForCausalLM.from_pretrained("./models/Mixtral-8x7B-v0.1")

# 拼接 expert 权重
expert_weights = torch.cat([
    model.model.layers[i].block_sparse_moe.experts[j].w1.weight
    for j in range(8)
], dim=0)  # [8*hidden_dim, ffn_dim]

什么模型该用 MoE 融合？

不是所有 MoE 模型都适合用 ops-transformer 的 MoE 融合算子。我的判断标准：

模型	expert 数	top-k	适合用 MoE 融合吗？	原因
Mixtral-8x7B	8	2	强烈推荐	8 个 expert 刚好适合昇腾的 8 卡 TP
DeepSeek-V2	160	6	不推荐	expert 太多，融合的 SRAM 开销太大
QWen-MoE	60	4	看情况	卡数是 expert 数的因数才行
Jamba-52B	16	2	推荐	expert 数适中

判断一句话：expert 数 ≤ 16，而且 top-k ≤ expert 数的 1/4，用 MoE 融合收益最大。expert 太多的话，分组本身的开销就超过了融合省下来的带宽。

完整排查清单

MoE 融合跑不起来，按这个清单查：

1. ops-transformer 的 MoE 融合算子装了吗？ls /usr/local/Ascend/ascend-toolkit/latest/op_api/moe_fusion/ 有东西吗？
2. 模型是 MoE 架构吗？config.json 里有 num_local_experts 字段才是 MoE。
3. expert 权重拼接对了吗？expert_weights 的形状应该是 [num_experts, ...]，不是独立的 tensor。
4. FlashAttention 开了吗？MoE 融合的前提是 FlashAttention 已经算完了。
5. vLLM 日志里有 Using fused MoE kernel 吗？没有就是静默降级了。
6. 卡数是 expert 数的因数吗？8 卡 + 8 expert 可以，8 卡 + 60 expert 不行。
7. 显存够吗？MoE 模型的专家权重很大，很容易 OOM。