MNN 新 LLM 架构支持实战从自定义算子导出到 C 后端的全栈实现指南【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN导读本文是 MNN 支持全新 LLM 架构Tier 6的完整落地指南。当模型结构超出标准 Attention 的覆盖范围如混合 conv/attention 架构、全新 LinearAttention 变体时本文将带你走通Python 侧自定义算子导出 → MNN Converter 解析 → C 后端算子实现的整条链路。读完你将掌握FusedLinearAttention统一算子接口的用法、CPULinearAttention的扩展点以及从model_mapper.py注册到llm_demo构建验证的完整 Checklist。一、背景什么是 Tier 6特殊架构支持在 MNN 的 LLM 支持流程中参见 skills/support-new-llm/SKILL.md经过 step1-analyze.md 的架构分析后模型会被判定为不同的支持等级。当分析结果为Tier 6时意味着需要对全新架构提供支持典型场景包括新的 Attention 类型如 gated delta rule、short conv attention 等 LinearAttention 变体混合 conv / attention 架构标准 Attention 层与 conv 层交替出现。Tier 6 的核心目标非常聚焦把模型中的新组件用 MNN 现有框架表达出来而不是从零搭建推理框架。由于 MNN 已经内置了FusedLinearAttention自定义算子与CPULinearAttention后端执行类绝大多数新架构都能通过扩展 attn_type的方式落地而无需触碰框架底层。Tier 6 的两大子类子类特征典型模型主要工作混合架构layer_types中有非 Attention 层conv/mamba/rwkv与 full_attention 交替lfm2新 LinearAttention 变体 C attn_type全新 Attention所有层都使用非标准 Attention如 gated delta ruleqwen3_5新 LinearAttention 变体 C attn_type共同点两者都通过FusedLinearAttention自定义算子导出通过CPULinearAttentionC 类执行。区别仅在于是否与标准 Attention 层混合。文档明确鼓励全栈实现Python 导出 MNN Converter 解析 C 后端算子。只有遇到极度复杂的特殊硬件或底层系统问题才考虑向人工求助。二、6.1 识别新组件动手前的第一件事在写任何代码之前先基于 step1 的分析结果明确需要新增的组件。用下面的模板把问题钉死需要新增的组件类型____如 LinearAttention 变体、新的 conv 层、新的 MLP 等 HF 源码中的类名 ____ 该组件替代的是 ____替代 self_attn / 替代 mlp / 全新位置 该组件的输入/输出 ____通常是 [B, L, hidden_size] → [B, L, hidden_size] 该组件是否有状态 ____如 conv state / recurrent state这里是否有状态尤其关键它直接决定了 C 侧onResize需要分配什么样的持久 buffer详见下文 6.2.3 的StateCache设计。三、6.2 LinearAttention 架构参考Tier 6 的核心知识本节是 Tier 6 的核心知识。大多数新架构最终都通过 LinearAttention 框架实现。3.1 整体数据流一个全新的 LinearAttention 变体从 Hugging Face 模型到 MNN 后端执行走的是同一条流水线Python 侧transformers.py ONNX 导出 C 侧CPULinearAttention.cpp ┌─────────────────────────┐ ┌─────────────────────────┐ │ ShortConvAttention │ │ CPULinearAttention │ │ / LinearAttention │ │ │ │ │ FusedLinearAttention │ onResize(): │ │ __init__: │ custom op (ONNX) │ 分配 conv/rnn state │ │ ModelMapper.do_map() │ ────────────────────── │ 分配临时 buffer │ │ FusedLinearAttention()│ │ │ │ │ │ onExecute(): │ │ forward (test path): │ │ dispatch by │ │ 完整的 Python 计算 │ │ mAttentionType │ │ │ │ → short_conv() │ │ forward (ONNX path): │ │ → gated_delta_rule() │ │ 调用 fused_attn op │ │ → new_type() │ └─────────────────────────┘ └─────────────────────────┘值得强调的是 Python 侧的双路径设计test pathtorch.onnx.is_in_onnx_export()为 False 时执行完整的 Python 计算逻辑用于与 HF 原始模型的 hook 对齐验证对应 step3-test-python.mdONNX path导出时只调用FusedLinearAttention自定义算子把计算责任完全交给 C 后端。3.2 FusedLinearAttention 自定义算子接口文件transformers/llm/export/utils/custom_op.pyFusedLinearAttentionOp定义在 L176-244FusedLinearAttention包装模块在 L246-265算子名LlmExporter::FusedLinearAttention 输入 Tensor4 个顺序固定 [0] qkv [B, D, L] 投影输出conv 之前D 总维度 [1] gate [B, L, H] decay / gate不需要时传 zeros [2] beta [B, L, H] learning rate不需要时传 zeros [3] conv_weight [C, 1, K] depthwise conv 权重C conv 通道数 属性Attributes attn_type string 算子子类型如 gated_delta_rule / short_conv num_k_heads int K/Q 的 head 数short_conv 可设 1 num_v_heads int V 的 head 数short_conv 可设 1 head_k_dim int 每个 K head 的维度 head_v_dim int 每个 V head 的维度 use_qk_l2norm int 是否对 Q/K 做 L2 归一化0/1 输出 Tensor1 个 [0] attn_out [B, L, num_v_heads, head_v_dim]关键设计不同attn_type共享同一个算子接口通过属性区分行为。这样 Converter 只需注册一次新类型只需在 C Execution 中增加 dispatch 分支。从源码看该算子还有两个文档未展开的进阶属性值得注意gate_fold导出期 gate/beta 折叠symbolic中支持gate_fold_i、gate_coef_f、gate_bias_f属性custom_op.py。当开启后inputs[1]/[2]携带的是原始a/b投影gate/beta 在 C 侧用mGateCoef/mGateBias内联计算从而在导出图中省掉整条 gate 计算链。LinearAttention通过export_args.fuse_linear_attn_gate控制是否启用transformers.py。conv_bias 输入注释中定义了可选的conv_bias [D]输入custom_op.py当前主流程以conv_weight为主。在symbolic中算子通过g.op(LlmExporter::FusedLinearAttention, *inputs, **kwargs)导出输出 shape 由[batch_size, seq_len, num_v_heads, head_v_dim]推导custom_op.pyforward是一个 dummy 实现仅返回正确 shape 的零张量custom_op.py。3.3 CPULinearAttention C 架构文件source/backend/cpu/CPULinearAttention.hppsource/backend/cpu/CPULinearAttention.cpp两个文件均在MNN_SUPPORT_TRANSFORMER_FUSE宏保护下编译类是Execution的子类通过REGISTER_CPU_OP_CREATOR_TRANSFORMER(CPULinearAttentionCreator, OpType_LinearAttention)注册CPULinearAttention.cpp。类结构struct StateCache { std::shared_ptrTensor mConvState; // Conv1D padding state std::shared_ptrTensor mRecurrentState; // 递归状态仅部分类型需要 }; class CPULinearAttention : public Execution { // 算子参数从 FlatBuffers 读取对应 ONNX 属性 std::string mAttentionType; int mHeadKDim, mHeadVDim, mNumKHeads, mNumVHeads; bool mUseQKL2Norm; // 持久状态通过 onClone 在 prefill/decode Execution 间共享 std::shared_ptrStateCache mStateCache; // 临时 buffer每次 onResize 重新分配 std::shared_ptrTensor mConvPadded, mConvOut; std::shared_ptrTensor mTempVPred, mTempDelta; // 仅 gated_delta_rule 需要 };实际源码中的StateCache比文档示例更丰富CPULinearAttention.hpp除了mConvState[B, D, kernel_size - 1]和mRecurrentState[B, H, d_k, d_v]之外还包含mConvStateSnapshot/mRecurrentStateSnapshot/mSnapshotValidpost-prefix 快照。LinearAttention 状态不是按 token 索引的eraseHistory无法逐 token 截断因此下一次 prefill 从快照恢复mPrefixLayerIndexprefix-cache 文件索引每个 session 捕获一次在混合模型中若每次都重新推进layer_index会越过 Full Attention 层导致 SIGBUS所以 chunks 2..N 复用该索引。这些细节提示我们为混合架构添加新类型时务必考虑 prefix-cache 交互。参数来源与 Creator 白名单构造函数从op-main_as_LinearAttentionParam()读取attn_type/num_k_heads/num_v_heads/head_k_dim/head_v_dim/use_qk_l2norm/gate_fold/gate_coef/gate_bias等全部参数并按mBytesfp32 为 4、Arm82 fp16 为 2记录字节宽度CPULinearAttention.cpp。一个容易被忽略的扩展点CPULinearAttentionCreator::onCreate只接受gated_delta_rule和short_conv两种 attn_type其他类型会打印unsupported attn_type并返回 nullptrCPULinearAttention.cpp。因此新增类型时必须同步把新 attn_type 加入这个白名单否则模型在转换阶段就会被拒。onResize 的统一 buffer 分配模式// ─── Per-type 参数添加新类型只需在这里加分支─── int convChannels convDim; // 默认conv 覆盖所有通道 bool needRecurrentState false; // 默认不需要递归状态 if (mAttentionType short_conv) { convChannels mHeadVDim; // conv 只覆盖 hidden_size 通道 } else if (mAttentionType gated_delta_rule) { needRecurrentState true; // 需要 [B, H, dk, dv] 递归状态 } // 新类型在此添加 else if ... // ─── 以下是共用逻辑不需要修改 ─── // 1. 分配 conv state [B, convChannels, kernelSize-1]STATIC跨 decode 保持 // 2. 如果 needRecurrentState分配 recurrent stateSTATIC // 3. 分配 mConvPadded / mConvOut 临时 bufferDYNAMIC // 4. 如果 needRecurrentState分配 mTempVPred / mTempDeltaDYNAMIC这个模式的价值在于把类型差异收敛到一小段 per-type 参数区其余分配逻辑对所有类型通用。新增类型时只需确定 conv 通道覆盖范围与是否需要递归状态。onExecute 的 dispatchErrorCode onExecute(...) { if (mAttentionType short_conv) { short_conv(inputs, outputs); } else if (mAttentionType gated_delta_rule) { gated_delta_rule_mnn(inputs, outputs); } // 新类型在此添加 else if ... return NO_ERROR; }实际源码中 dispatch 位于 CPULinearAttention.cppgated_delta_rule 一侧还有gated_delta_rule_ref参考实现与gated_delta_rule_decodedecode 专用const 方法等变体。onClone 的状态共享bool onClone(Backend* bn, const Op* op, Execution** dst) { auto tmp new CPULinearAttention(bn, op); tmp-mStateCache mStateCache; // 共享持久状态 *dst tmp; return true; }prefill 与 decode 阶段会生成多个 Execution 实例onClone通过共享mStateCache保证 conv/recurrent 状态在两者之间连续CPULinearAttention.cpp。这是有状态算子的关键机制新类型如果带状态必须遵守这个共享约定。多线程实际实现使用MNN_CONCURRENCY_BEGIN(tId, threadNum) / MNN_CONCURRENCY_END()宏按B * channels或类似维度并行CPULinearAttention.cpp。新方法的实现应当沿用同样的并行骨架而不是写单线程串行代码。3.4 Python 侧 create_linear_attention 工厂文件transformers/llm/export/utils/transformers.pydef create_linear_attention(attn, layer_id, config, rotary, mapper): Factory function for creating LinearAttention variants based on config. if hasattr(config, conv_L_cache) and config.conv_L_cache 0: return ShortConvAttention(attn, layer_id, config, mapper) # 新类型在此添加 elif ... return LinearAttention(attn, layer_id, config, rotary, mapper)Decoder.__init__通过linear_attn槽位触发# Decoder.__init__ 中的关键逻辑 if hasattr(self, linear_attn) and self.linear_attn is not None: self.self_attn create_linear_attention(self.linear_attn, layer_id, config, rotary, mapper) self.layer_type linear_attention工厂函数定义于 transformers.pyDecoder 触发逻辑位于 L1477-1478。映射侧在 transformers/llm/export/utils/model_mapper.py 的 decoder 映射中将 HF 模型的非标准层映射到linear_attn槽位decoder { self_attn: self_attn, # 标准 Attention 层有的层有有的层没有 linear_attn: conv, # 非标准层 → linear_attn 槽位 mlp: feed_forward, # ... }对于混合架构如 lfm2同一个 decoder 映射同时包含self_attn和linear_attn。ModelMapper.do_map会对不存在的属性设置None所以conv 层self_attnNone,linear_attnconvModule→ShortConvAttentionattention 层self_attnattnModule,linear_attnNone→Attention四、6.3 添加新 LinearAttention 变体的 Checklist以short_convLFM2为实际案例说明。整个流程共涉及 4 个 Python 文件 2 个 C 文件。Python 侧4 个文件1.model_mapper.py— 注册映射文件transformers/llm/export/utils/model_mapper.py。仓库中regist_lfm2L924-969、regist_lfm2_moeL971-1026、regist_lfm2_vlL1028-1075、regist_lfm2_audioL1077 起四个注册函数展示了 lfm2 系列含 MoE / 多模态 / 音频变体的统一模式def regist_lfm2(self): # config 映射添加模型特有的配置字段 lfm2_config { hidden_size: hidden_size, # ... conv_L_cache: conv_L_cache, # ← 新字段 } # linear_attention 映射新组件的子模块名 lfm2_linear_attention { in_proj: in_proj, conv: conv, out_proj: out_proj, } # decoder 映射linear_attn 指向 HF 模型中的非标准层 lfm2_decoder { self_attn: self_attn, linear_attn: conv, # ← HF 的 conv 层 → linear_attn 槽位 mlp: feed_forward, # ... } lfm2_map { config: lfm2_config, model: lfm2_model, decoder: lfm2_decoder, attention: lfm2_attention, linear_attention: lfm2_linear_attention, # ← 新组件映射 } self.regist(lfm2, lfm2_map)2.config.py— 注册新配置字段如需要文件transformers/llm/export/utils/config.pyself.conv_L_cache kwargs.pop(conv_L_cache, 0)新字段必须提供默认值避免旧配置加载时 KeyError。3.transformers.py— 新组件类 工厂注册文件transformers/llm/export/utils/transformers.py。每个 LinearAttention 变体都是一个torch.nn.Module子类需要实现两条路径class ShortConvAttention(torch.nn.Module): def __init__(self, attn, layer_id, config, mapper): super().__init__() # 1. 用 ModelMapper.do_map 提取子模块 ModelMapper.do_map(self, attn, mapper[linear_attention]) # 2. 创建 FusedLinearAttention 实例用于 ONNX 导出 self.fused_attn FusedLinearAttention( namef/layers.{layer_id}/self_attn/FusedLinearAttention, attn_typeshort_conv, # ← 新的 attn_type 字符串 num_k_heads1, num_v_heads1, head_k_dimself.hidden_size, head_v_dimself.hidden_size, use_qk_l2normFalse ) # 3. 初始化内部状态用于 test path 的 decode 推理 self.conv_state None def forward(self, hidden_states, attention_maskNone): if torch.onnx.is_in_onnx_export(): # ONNX 路径调用 FusedLinearAttention 自定义算子 # 输入投影后的 tensor [B, D, L] # 不需要的输入传 zeros attn_out self.fused_attn(bcx_t, gate_zeros, beta_zeros, self.conv.weight) return self.out_proj(attn_out.view(B, L, -1)) # Test 路径完整的 Python 计算用于 hook 对齐验证 # 对照 HF 源码实现完整的前向逻辑 # 维护 self.conv_state 等内部状态 ...仓库中ShortConvAttention的真实实现transformers.py可以当作模板细读ONNX 路径bcx self.in_proj(hidden_states)得到[B, L, 3H]转置为[B, 3H, L]后传入算子gate/beta 用torch.zeros占位short_conv 不使用self.conv.weight.data.detach()作为 conv 权重输出[B, L, 1, H]reshape 回[B, L, H]再过out_projTest 路径B_, C_, x_ bcx.chunk(3, dim-1)拆三份 →Bx B_ * x_→ 转置成[B, H, L]后与self.conv_state拼接做groupsself.hidden_size的 depthwiseconv1d→无 SiLU与 gated_delta_rule 的关键区别→y C_ * conv_out→out_projmask-freeShortConvAttention不接受attention_mask参与计算参数仅为保持与Attention.forward的签名一致。在create_linear_attention工厂中注册def create_linear_attention(attn, layer_id, config, rotary, mapper): if hasattr(config, conv_L_cache) and config.conv_L_cache 0: return ShortConvAttention(attn, layer_id, config, mapper) # elif some_other_condition: # return NewTypeAttention(attn, layer_id, config, mapper) return LinearAttention(attn, layer_id, config, rotary, mapper)4.custom_op.py— 通常不需要修改transformers/llm/export/utils/custom_op.py 中FusedLinearAttention已支持任意attn_type字符串新类型不需要修改此文件。这是整个方案扩展成本最低的关键设计算子接口一次注册行为完全由 attn_type 属性驱动。C 侧2 个文件5.CPULinearAttention.hpp— 添加方法声明文件source/backend/cpu/CPULinearAttention.hppvoid short_conv(const std::vectorTensor* inputs, const std::vectorTensor* outputs);6.CPULinearAttention.cpp— 3 处修改文件source/backend/cpu/CPULinearAttention.cpp修改 1onResize顶部的 per-type 参数分支if (mAttentionType short_conv) { convChannels mHeadVDim; } else if (mAttentionType gated_delta_rule) { needRecurrentState true; } else if (mAttentionType new_type) { // 设置 convChannels 和 needRecurrentState }修改 2onExecute的 dispatch 分支if (mAttentionType short_conv) { short_conv(inputs, outputs); } else if (mAttentionType new_type) { new_type(inputs, outputs); } else { gated_delta_rule_mnn(inputs, outputs); }修改 3易遗漏CPULinearAttentionCreator::onCreate白名单加新类型。当前实现只放行gated_delta_rule与short_convCPULinearAttention.cpp新类型不加入白名单会直接转换失败。修改 4新方法实现void CPULinearAttention::short_conv(const std::vectorTensor* inputs, const std::vectorTensor* outputs) { // 输入qkv [B, D, L], conv_weight [C, 1, K] // 输出attn_out [B, L, num_v_heads, head_v_dim] // // 核心步骤 // 1. 从 qkv 中提取需要的分量 // 2. 使用 mStateCache-mConvState 做 depthwise conv1d带状态管理 // 3. 后处理如元素乘法 // 4. 写入 output tensor // // 多线程使用 MNN_CONCURRENCY_BEGIN/END按 B*channels 并行 }short_conv的真实实现位于 CPULinearAttention.cpp包含两段MNN_CONCURRENCY_BEGIN/END并行区间可作参考。构建与测试# 构建只需要编译 LLM 相关目标 cmake --build build --target llm_demo -j$(nproc) # 导出 cd transformers/llm/export python3 llmexport.py --path /path/to/model --export mnn --dst_path /tmp/MODEL # C 推理测试 echo 你好 /tmp/prompt.txt ./build/llm_demo /tmp/MODEL/llm_config.json /tmp/prompt.txt五、6.4 已有 attn_type 实现参考当前仓库CPULinearAttentionCreator 白名单支持的两种 attn_type 对比如下attn_type模型conv 通道递归状态SiLU核心逻辑gated_delta_ruleqwen3_5D (全部)是 [B,H,dk,dv]是conv→SiLU→split QKV→L2norm→scale→delta rule recurrenceshort_convlfm2H (部分)否否split BCx→Bx→conv→Cconv_outPython 侧对应关系gated_delta_rule由LinearAttention类导出transformers.pynum_k_heads/num_v_heads从 config 的linear_num_key_heads/linear_num_value_heads读取use_qk_l2normTrueshort_conv由ShortConvAttention类导出单 head、L2norm 关闭。新类型实现时的关键决策设计一个全新 attn_type 时只需回答下面五个问题就能确定全部扩展点conv 覆盖多少通道→ 决定convChannels值onResize 分支是否需要递归状态→ 决定needRecurrentState决定是否分配[B, H, dk, dv]状态conv 后是否有激活函数→ gated_delta_rule 有 SiLUshort_conv 没有Python test path 与 C 实现必须一致qkv tensor 的语义是什么→ 不同类型的 split 方式不同如 short_conv 按B_/C_/x_三等分gated_delta_rule 按 QKV head 拆分gate / beta 输入是否使用→ short_conv 传 zerosgated_delta_rule 实际使用若不需要可在 ONNX 路径传 zeros 占位六、步骤 6 测试标准Python 侧通过标准新组件类实现完成test path 和 ONNX path 都有model_mapper.py 中新组件映射已添加包含linear_attention子映射config.py 中新配置字段已注册如需要create_linear_attention 工厂已注册新类型Python test 推理输出正确与 HF 原始模型一致ONNX 导出不报错C 侧通过标准CPULinearAttention.hpp 中新方法已声明onResize 的 per-type 参数已添加onExecute 的 dispatch 已添加Creator 白名单已加入新 attn_type容易遗漏新方法实现完成含 conv state 管理 多线程C 推理输出正确部分完成的评估机制若全栈实现顺利则直接报告完全成功。在极端情况下若底层算子实现卡住允许分段交付✅ 已完成 - Python 侧新组件实现test path ONNX path - model_mapper.py 映射 - Python test 验证通过 - ONNX 导出成功 ⏳ 待解决 - CPULinearAttention 中的新 attn_type 实现七、下一步衔接新架构支持并不是流程的终点而是与整个支持流水线闭环Python 侧完成后→ 回到 step3-test-python.md 重新验证重点确认 test path 输出与 HF 原始模型逐层对齐全部完成后→ 回到 step4-export.md 完成最终导出验证导出产物需通过 llm_demo 端到端确认C 侧无法完成→ 总结工作请求人工协助按部分完成的评估机制分段交付保证已完成的 Python 侧成果可复用八、小结Tier 6 特殊架构支持的核心方法论可以浓缩为一句话把新架构抽象成新的 attn_type 字符串。得益于FusedLinearAttention统一算子接口Converter 只注册一次与CPULinearAttention的 per-type 参数分支设计onResize 只改一小段、onExecute 只加一个分支新增一个 LinearAttention 变体的工作量被压缩到Python 侧 1 个新类 1 个工厂分支 1 组映射注册C 侧 1 个方法声明 1 个 onResize 分支 1 个 dispatch 分支 1 个白名单条目。理解 conv 通道覆盖、递归状态、激活函数、qkv 语义与 gate/beta 用法这五个关键决策点就能把绝大多数全新 LLM 架构稳定接入 MNN。【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考