cann-recipes-infer：Hy3 295B MoE 大模型在 Atlas A3 平台的端到端推理优化实践

cann-recipes-inferHy3 295B MoE 大模型在 Atlas A3 平台的端到端推理优化实践【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本文以 Hy3 模型优化报告 为核心完整复现并解析腾讯混元 Hy3295B 总参数、每 token 激活约 21BMoE 模型在 Atlas A3Ascend 910C平台上 8 卡 16 rank、BF16 精度下的推理优化闭环从混合并行切分、KVCache 静态化与 FA 算子替换、torch_npu 融合算子到 ge_graph 图模式适配。读完本文你可以掌握一套并行化部署 → 静态 shape 化 → 算子融合 → 图模式编译的 MoE 模型推理优化方法论以及 Decode 单步时延从 293 ms/t 压到 28.44 ms/t10.3x各阶段的贡献来源与工程决策依据。1. 模型与平台概览Hy3 是腾讯混元团队发布的 MoE 语言模型采用 GQA 注意力与带共享专家的稀疏 MoE 结构decoder-only。在 Hy3 样例 中该模型已接入 cann-recipes-infer 统一推理框架模型代码位于 modeling_hy_v3.py配置类位于 configuration_hy_v3.py。本报告记录的优化环境为架构GQA 稀疏 MoE含共享专家decoder-only总参数量295B每 token 激活约 21B硬件平台Atlas A3Ascend 910C部署规模8 卡 16 rank每卡 2 die单节点量化模式BF16本报告为 Hy3 在 Atlas A3 平台的端到端基础优化闭环记录。Ascend 950 PR 平台上的进一步性能调优参见 Hy3 推理优化实践。1.1 模型结构概要Embedding (120832 × 4096, 495M params) └─ Transformer Block × 80 (layers 0-79) ├─ Pre-Attention RMSNorm (npu_add_rms_norm fused) ├─ Attention (GQA: 64Q heads, 8KV heads, head_dim128) │ ├─ QKV Projection: QKVParallelLinear (TP4) │ ├─ QK Norm: npu_rms_norm(128) per-head │ ├─ RoPE: npu_apply_rotary_pos_emb(layoutBSH) │ ├─ KV Cache: BSH [B,S,N_kv*D], scatter_update_(axis-2) │ ├─ Attention Core: FA v1 (eager) / torchair FA (ge_graph) │ └─ O Projection: RowParallelLinear (TP4) ├─ Post-Attention RMSNorm (npu_add_rms_norm fused) ├─ FFN / MoE │ ├─ Layer 0: Dense FFN (gate/up: 4096→13312, down: 13312→4096, SiLU) │ └─ Layers 1-79: MoE │ ├─ Router: npu_moe_gating_top_k(norm_type1)sigmoid 非 softmax │ ├─ EP dispatch/combine: MC2 (ge_graph) / all_to_all manual (eager) │ ├─ Expert FFN × 12 per-rank (GMM: npu_grouped_matmul split_item2) │ └─ Shared Expert × 1 (gate/up: 4096→1536, down: 1536→4096) └─ Residual Connection (fused via npu_add_rms_norm) └─ Final RMSNorm (npu_rms_norm) └─ LM Head: ColumnParallelLinear(4096, 120832, no bias)从源码结构看这一结构与 HYV3Attention、HYV3TopKRouter 和 HYV3MoE 三个核心类一一对应每个 Transformer 层由 HYV3DecoderLayer 组装。1.2 关键参数参数值参数值总参数量295B (295,033,543,488)激活参数量~21B总层数80 (excl. MTP layer 80)Hidden Size4096Attention Heads64 (Q), 8 (KV)Head Dim128Dense FFN Intermediate13312 (layer 0)MoE Expert Intermediate1536专家总数192每 token 激活专家8 (top-8)Shared Expert1 per MoE layer词表大小120,832最大上下文262,144 (256K)RoPE Theta11,158,8401.3 特殊设计特点特性说明适配影响QK NormQ 和 K 在 RoPE 前各经 RMSNorm(128)已通过 npu_rms_norm 优化Sigmoid Router使用 sigmoid 而非 softmax已通过 npu_moe_gating_top_k(norm_type1) 融合Expert Biaslearnablee_score_correction_bias参与 routing scoringRouter Scalingtop-k weights × 2.826已通过 routed_scaling_factor 参数处理MTP Layerlayer 80 (~3.8B params)推理跳过_keys_to_ignore_on_load_unexpectedLarge Vocab120K tokensembed/lmhead 各 ~1 GB BF16词表按 TP4 切分上述结构特殊点在后续适配中均通过对应融合算子或参数处理消化未引入定制算子。这一点在源码中可以得到印证sigmoid 路由直接由融合算子一行完成HYV3TopKRouter.forward 调用top_k_weights, top_k_index, _ torch_npu.npu_moe_gating_top_k( router_logits, kself.top_k, biase_score_correction_bias.float(), norm_type1, # sigmoid routed_scaling_factorself.router_scaling_factor, eps1e-20, )norm_type1对应 Hy3 的 sigmoid 激活而非默认 softmaxbias传入 learnable 的e_score_correction_biasrouted_scaling_factor消化了 top-k 权重 ×2.826 的缩放逻辑——7 个 Python 小算子sigmoidtopkgathernormscale被压缩为一次 NPU 算子调用。2. 性能基线本章记录 Hy3 完成并行化部署、尚未引入任何算子融合或图模式优化时的初始性能作为后续各项优化的对比基准。基线在 Atlas A3 16 die、BF16、eager 模式下采集。指标值测试条件Prefill 耗时2,070 ms1024 tokens, 16 die A3, eager mode, batch4Decode 单步耗时293 ms/teager mode, 16 die A3显存占用~54.5 GB/dieBF16, 295B 参数 EPTP 分布后执行模式eager初始部署该基线采集于并行化部署完成之后Decode 单步 293 ms/t、Prefill 2070 ms是后续所有性能对比的统一基准。3. 并行化改造在基线部署之上295B 参数的 Hy3 按 16 die 拓扑做混合并行切分这是所有后续优化的部署前提。切分需同时兼顾 GQA 的 KV 头数量约束与 MoE 专家规模避免碎矩阵与跨模块边界通信。3.1 并行策略参数值理由world_size16 die8 卡 × 2 die/卡attn_tp_size4N_kv8 约束, 均衡 Prefill/Decodedense_tp_size4与 attn_tp 对齐避免模块边界通信moe_tp_size1Expert intermediate1536 太小TP 导致碎矩阵moe_ep_size16全卡 EP192/1612 experts/rankembed_tp_size4大词表 120K 切分lmhead_tp_size4大词表切分这套策略在仓库 A3 配置文件 hy3_rank16_bf16_mtp.yaml 中有直接对应parallel_config段与报告逐项一致parallel_config: world_size: 16 attn_tp_size: 4 dense_tp_size: 4 moe_tp_size: 1 embed_tp_size: 4 lmhead_tp_size: 4moe_ep_size由框架在 EP1 场景下按 world_size 推导见 inference_config.py 与 model_worker.py 中的moe_ep_size处理逻辑experts_per_rank num_experts // moe_ep_size192/1612在 HYV3MoE.init中显式计算。3.2 关键实现通信组注意力、Dense、Embedding、LM Head 均按 TP4 切分各自建立 4 组 TP 通信组MoE 按 16 卡专家并行建立一个 EP 通信组。通信组的建立集中在 HYV3ForCausalLM.init_parallel_comm_group。EP 路由eager 模式使用手动 all_to_all 完成专家 dispatch/combine见 _moe_ep_manual图模式改用 MC2 dispatch/combine见 _moe_ep_mc2_decode第 6 章详述。padding_idx全局 padding 索引 120002 在 Embedding 切分后映射为每 rank 的局部索引。权重映射checkpoint 中分离的q/k/v_proj权重统一映射到merged_qkv_proj.weight该逻辑在 load_weights 中实现保证 QKV 投影权重正确加载。3.3 验收结果验收项结果编译16 die 编译通过多卡推理无 crash解码输出输出正常文本QKV 权重按 checkpoint 的 q/k/v_proj 映射到 merged_qkv_proj 后正确加载上述切分下模型编译与 16 die 推理均正常解码输出正确文本说明并行拆分与权重加载映射正确、精度未见异常可作为后续优化的部署基线。4. KVCache 静态化与 FA 算子替换完成并行化部署后本章将 Attention 的 KV 管理从动态实现改为静态连续缓存并把注意力核替换为 FA 融合算子。静态 shape 与融合注意力既提升当前 eager 性能也为后续图模式捕获提供前提。4.1 选型结果决策项选择理由KVCache 模式模式一连续缓存 (BSH)静态 batch, 无动态分配需求FA 算子FA v1 (npu_fused_infer_attention_score)GQA 无 MLA/量化需求KV Cache 写入scatter_update_(axis-2)BSH layout 标准写法LayoutBSH [B, S, N_kv*D]QK Norm per-head 兼容Prefill masksparse_mode3 ~tril boolDecoder-only causalDecode masksparse_mode0 None无需因果遮蔽4.2 实现印证从 HYV3Attention.forward 的实现可以看到该选型在代码中的落点QKV 合并投影merged_qkv_proj后按 rank 局部头数切分Q/K 经 per-head RMSNormq_norm/k_norm即 head_dim128 维的npu_rms_norm后调用torch_npu.npu_apply_rotary_pos_emb(..., layoutTND)完成旋转位置编码再通过npu_scatter_nd_update_按slot_mapping将 K/V 写入预分配的连续缓存。Prefill 与 Decode 两分支分别使用带因果 mask 与不带 mask 的 FA 调用路径与上表的 sparse_mode 选择一致。4.3 验收结果指标结果KVCache 静态化BSH 连续缓存scatter_update_FA 算子替换FA v1Prefill/Decode 分支正确精度 (vs HF)输出正常文本KVCache 改为 BSH 连续缓存并接入 FA v1 后Prefill 与 Decode 分支分别使用正确的 mask 与 sparse_mode输出正常文本、精度与 HF 对齐同时静态 shape 为后续图模式捕获提供了前提。5. 融合算子优化在 KVCache 与 FA 改造之上本章把模型中的 Norm、激活、MoE 路由等 Python 小算子模式替换为 torch_npu 融合算子。这些替换在 eager 模式下即可获得收益同时减少算子数量为图模式捕获铺路。5.1 实施项目P0零风险模块原实现替换为调用频率RMSNorm (all)Pythonpow(2).mean().rsqrt()npu_rms_norm321次/forwardResidual RMSNormx yRMSNorm(y)分离npu_add_rms_norm160次/forwardFFN SiLUF.silu(gate) * upnpu_swiglu80次/forwardP1精度敏感模块原实现替换为调用频率MoE RouterPython sigmoidtopkgathernormscale (7 ops)npu_moe_gating_top_k(norm_type1)79次/forwardP1 精度说明Hy3 的e_score_correction_bias在 sigmoid 后加仅影响 topk 选择而npu_moe_gating_top_k的 bias 在 sigmoid 内。该语义差异未造成可见输出质量退化。P2进阶候选暂缓模块候选算子未实施原因Dense FFN 全融合npu_ffn(activationswiglu)需绕过并行 Linear, 性能门槛待 profilingExpert FFN 全融合npu_ffn(expert_tokens...)forward_ordered 适配, 性能门槛待 profiling不适配算子算子不适配原因npu_kv_rmsnorm_rope_cacheMLA 专用, Hy3 为 GQA per-head QK Normnpu_mla_prolog_v3MLA absorb 模式专用npu_moe_distribute_combine_add_rms_norm需配合 dispatch_v2 使用, 且手动 EP 链路不匹配5.2 源码落点P0/P1 替换在模型代码中均可直接检索到Residual RMSNorm 融合HYV3DecoderLayer.forward 中Pre-Attention 与 Post-Attention 两处残差均走torch_npu.npu_add_rms_norm量化路径下切换为npu_add_rms_norm_dynamic_mx_quant/npu_add_rms_norm_cast一次调用同时输出归一化结果与残差SwiGLU 融合HYV3MLP.forward 中gate/up 合并投影后直接调用torch_npu.npu_swiglu(merged)量化路径下走npu_swiglu_group_quant见 modules/common.pyRouter 融合如 1.3 节所示npu_moe_gating_top_k一次完成 sigmoidtopk加权归一化2.826 缩放。5.3 验收结果eager modeP0P1指标改造前改造后改善Prefill (1024t)2,070 ms1,041 ms50%Decode avg293 ms/t246 ms/t16%P0P1 融合算子在 eager 模式下把 Prefill 从 2070 ms 降至 1041 ms约 50%、Decode 从 293 ms/t 降至 246 ms/t约 16%主要收益来自 Norm/激活/路由小算子合并MoE 路由的 bias 语义差异未造成可见输出质量退化精度保持不变。6. 图模式适配前几章改造均在 eager 模式下完成本章将 Decode 路径适配到 ge_graph 图模式通过消除 graph break、算子二进制融合与编译缓存进一步压低单步时延Prefill 因序列长度可变仍保持 eager图模式仅覆盖 Decode。6.1 实施方案项目方案说明图模式后端ge_graphtorchair CompilerConfigDecode 路由MC2 dispatch/combine GMM替代手动 all_to_all, 消除数据依赖 graph break专家计算npu_grouped_matmul(split_item2)替代 per-expert F.linear loopFA 接口torchair FA (ge_graph) / torch.ops.npu FA (eager Prefill)Prefill 保持 eagerSuperKernelenable_superkernelTrueDecode 层循环外包 superkernel_scopeCompile Cacheenable_cache_compileTrue持久化到 compile_cache/, 二次 warmup ~9s (10x)多流默认关闭 (enable_multi_streamsFalse)shared expert 计算量小无法有效 overlap实现保留待后续评估说明当前仓库 A3 配置文件 hy3_rank16_bf16_mtp.yaml 中exe_mode已支持[ge_graph, eager, npugraph_ex]三种模式当前默认npugraph_exenable_cache_compile作为开关保留在model_config段enable_multi_streams与enable_sp则放在custom_params字段下由 HYV3MoE 读取并据此创建 shared expert 独立流。6.2 关键实现要点MC2 路由消除 graph break。eager 路径的手动 all_to_all 依赖运行时张量值做切分数据依赖会打断图捕获图模式改用 CANN 的 MoE 分布式算子对见 _moe_ep_mc2_decodetorch_npu.npu_moe_distribute_dispatch_v2完成 dispatch中间由 GMM 专家矩阵乘再用torch_npu.npu_moe_distribute_combine_v2加权聚合。dispatch/combine 的参数中moe_expert_num192、ep_world_size/ep_rank_id来自moe_ep_group_mc2通信组均为编译期可静态确定的值无数据依赖分支。GMM 与权重布局。专家计算使用npu_grouped_matmul(split_item2)替代 per-expert 的F.linear循环split_item2的选择在 HYV3MoE 中显式固定Hy3 keeps the checkpoint expert layout and validated split_item2 GMM pathGMM 底层实现在 FusedMoEGMM 中其中npu_grouped_matmul的group_list即 dispatch 输出的每专家 token 数。Expert GMM 权重以 inline.transpose(1,2)view 复用避免.transpose().contiguous()产生的约 35.8 GB 权重副本不额外分配显存——对应 process_weights_after_loading 中is_transposeTrue的处理GE 会将常量转置折叠进编译产物使 Decode 不再每步重转置。权重加载映射。Checkpoint 中分离的q/k/v_proj权重映射到merged_qkv_proj.weight保证 QKV 投影权重正确加载见 load_weights。6.3 验收结果指标eager (P0P1)ge_graph (图模式阶段)改善Prefill (1024t)1,041 ms1,452 ms—Decode avg246 ms/t30.8 ms/t8.0x阶段值Decode warmup (首次)—91,151 ms—Decode warmup (cached)—8,016 ms10x图模式阶段将 Decode 从 eager 的 246 ms/t 降至 30.8 ms/t阶段加速约 8.0x主要收益来自 MC2 路由与 GMM 消除数据依赖 graph break、以及 SuperKernel 减少层循环开销compile cache 把二次 warmup 从约 91s 压到约 8s。此为图模式阶段结果最终经重编译进一步降至 28.44 ms/t见第 7 章。7. 累计优化效果本章汇总并行化、融合算子与图模式各阶段叠加后的累计效果标注最大收益来源与最终性能水平。指标原始基线 (eager)最终 (ge_graph)累计改善Prefill 耗时 (1024t)2,070 ms1,452 ms1.4xDecode 单步耗时293 ms/t28.44 ms/t10.3xvs 目标 (100ms/t)—28.44 ms/t3.5x margin累计来看Decode 从 293 ms/t 降至 28.44 ms/t10.3x最大收益来自图模式叠加 MC2 路由与 GMMPrefill 从 2070 ms 降至 1452 ms约 1.4x。最终 Decode 相对 100 ms/t 目标仍有约 3.5x 余量。7.1 各阶段贡献优化项Decode 性能增量累计改善并行化 (eager 基线)293 ms/t—1.0xP0P1 融合算子246 ms/t1.19x1.19xge_graph GMM33.3 ms/t7.4x8.8x compile cache31.5 ms/t1.06x9.3x SuperKernel30.8 ms/t1.02x9.5x最终重编译28.44 ms/t1.08x10.3x各阶段中ge_graph GMM 贡献了绝大部分收益7.4xcompile cache 与 SuperKernel 为增量优化最终重编译后 Decode 稳定在 28.44 ms/t。8. 遗留问题与后续建议本章记录本期优化未落地的候选方案及其阻塞原因并给出最终性能确认与已知限制作为后续迭代参考。8.1 P2 融合算子最终评估候选结论阻塞原因npu_ffn(expert_tokens...)不可行API 硬约束swiglu 模式不接受 expert_tokensnpu_ffn(activationswiglu)Dense不可行API 硬约束swiglu 仅支持 float16 (Hy3 用 bf16)npu_moe_distribute_combine_add_rms_norm不可行数学不等价API 计算norm(ABres)vs Hy3 deferred residualnorm(A)rescombine_v2 shared_expert_x暂不启用SuperKernel 编译失败融合链过长超过 2-stream 限制 (CANN 8.5.0)。待 CANN 更新后可重新启用上述四个进阶候选在当前 CANN 版本下均受 API 数据类型、语义等价性或融合链长度的硬约束限制本期不落地其中 combine_v2 shared_expert_x 待 CANN 更新后可重新评估。8.2 已验证最终性能指标数值备注Prefill (1024t)1,452 msge_graph compile cacheDecode avg28.44 ms/trange: 27.4-29.4, 32 tokensvs 原始 eager (293ms/t)10.3x 加速vs 目标 (100ms/t)3.5x margin输出质量正常The output is a weighted sum of the values...最终确认 Decode 28.44 ms/t32 token 区间 27.4–29.4、Prefill 1452 ms相对原始 eager 基线 293 ms/t 加速 10.3x输出质量正常达到本期优化目标并留有约 3.5x 余量。8.3 性能分析当前 Decode 28.44 ms/t 已远超 100ms/t 目标 (3.5x margin)。Prefill 1452ms (1024t) 仍有优化空间但 Prefill 非主要瓶颈单次推理仅执行 1 次 Prefill N 次 Decode。8.4 已知限制Multi-streamshared expert 计算量太小 (1536 intermediate)无法有效 overlap。若未来模型升级增大 intermediate 可重新评估长序列 (128K)KV Cache 增长到 ~10.7 GB/卡 (attn_tp4)余量减少。若扩展到 256K 需评估 KVP 叠加W8A8 量化若未来减至 8 卡部署需 W8A8 量化 (~295 GB int8)。9. 在仓库中复现与延伸阅读本报告的优化闭环在仓库中的落地形态如下可作为实操入口模型实现modeling_hy_v3.pyAttention/MoE/Router/MTP 全部 NPU 适配、configuration_hy_v3.py、modules/common.py并行配置hy3_rank16_bf16_mtp.yamlA3 16 rank BF16 MTP950 平台量化配置见 hy3_rank4_mxfp4_mtp.yaml 与 hy3_rank4_fp8_mtp.yamlYAML 通用参数说明参考 inference_config_guide.md推理入口按 Hy3 样例 README 准备环境后执行bash executor/scripts/infer.sh --model hy3 --yaml ci_a3/hy3_rank16_bf16_mtp.yaml即可拉起 8 卡 16 rank 推理Hy3 特有开关enable_multi_streams、enable_sp置于 YAML 的custom_params段。需要说明的适用前提本报告所有性能数字均为 Atlas A3Ascend 910C16 die、BF16 精度下的实测记录对应 CANN 8.5.0P2 评估部分及报告所述软件版本环境仓库当前 A3 配置已演进到npugraph_ex执行模式并引入 MTPnext_n1与报告阶段的ge_graph结果不完全等价。Ascend 950 PR 平台上的量化与进一步调优参见 Hy3 推理优化实践。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考