MoE 模型的 FlashAttention 跟普通模型有什么不一样？

MoE 模型的 FlashAttention 跟普通模型有什么不一样？

前阵子帮人调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 里就地完成。

# 标准：两步，中间写 HBMforexpertinexperts:expert_output=FFN(expert_input)# 写 HBMmerged=Merge(expert_outputs)# 从 HBM 读 8 次# 融合：一步，不写 HBMaccum=zeros(sram)forexpertinexperts: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)# 写 HBMroute_scores=Router(attn_output)# 从 HBM 读# 融合：FlashAttention 输出留在 SRAM，路由选择从 SRAM 读attn_output=FlashAttention(Q,K,V)# 留在 SRAMroute_scores=Router(attn_output)# 从 SRAM 读（20倍快）

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

融合后的总效果

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

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

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

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

跟 Llama-2-7B 的对比

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

跟 NVIDIA A100 的对比

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

差距分析：吞吐差 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 融合exportVLLM_USE_FUSED_MOE=1python-mvllm.entrypoints.openai.api_server\--model./models/Mixtral-8x7B-v0.1\--tensor-parallel-size8\--enable-flash-attn\--max-model-len4096

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

手动编译 MoE 融合算子

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

# 拉取仓库gitclone https://atomgit.com/cann/ops-transformer.gitcdops-transformer# 编译 MoE 融合算子cdsrc/moe_fusionbashbuild.sh--socAscend910--typrelease# 安装chmod+x ./output/moe_fusion_Ascend910.runsudo./output/moe_fusion_Ascend910.run

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

importtorchimporttorch_npufromtorch_npu.contrib.functionalimportnpu_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 加载fromtransformersimportMixtralForCausalLM model=MixtralForCausalLM.from_pretrained("./models/Mixtral-8x7B-v0.1")# 拼接 expert 权重expert_weights=torch.cat([model.model.layers[i].block_sparse_moe.experts[j].w1.weightforjinrange(8)],dim=0)# [8*hidden_dim, ffn_dim]

什么模型该用 MoE 融合？

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

判断一句话：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。