ANE性能优化Forward Taps中间量输出避免CPU重计算【免费下载链接】ANETraining neural networks on Apple Neural Engine via reverse-engineered private APIs项目地址: https://gitcode.com/GitHub_Trending/ane2/ANE在 ANE 项目中开发者逆向 Apple 私有 API_ANEClient/_ANECompiler让 Apple Neural Engine 这块被限制为只推理的 NPU 也能跑训练。而在前向/反向传播的衔接处一项名为Forward Taps前向中间量输出的性能优化让反向传播避免了 CPU 重计算前向核在 ANE 上顺手把 Q、K、V、注意力分数、隐藏层状态等中间量搭车输出CPU 只需切块读取不再重复计算。本文用通俗的方式拆解这一技巧。先搞懂问题训练为什么要多要几个输出推理时模型只需输出最终结果logits但训练的反向传播必须知道前向计算过程中的中间值中间量作用不缓存的代价Q、K、V计算 dQ/dK/dV 梯度必须重新做 3 次 QKV 投影矩阵乘注意力分数 probs计算 dQ 的 softmax 梯度必须重新做 QKᵀ softmaxFFN 的 h1、h3、SiLU 输出SwiGLU 激活的导数需要它们必须重新做 3 次大矩阵乘RMSNorm 归一化后的输入Norm 反向公式依赖原值必须重跑归一化ANE 的 IOSurface 共享内存 I/O 带宽非常可观fp16 直传比 fp32 快约 37%多搬运一点内存远比在 CPU 上多算一遍便宜。于是项目采用了 Forward Taps。实现原理一个 concat 操作打包所有 Tap整个技巧只有一行 MIL 操作在前向计算图的末尾用concat沿通道维把所有需要的张量拼接成一个输出。以注意力前向核为例源码 stories_mil.htensorfp16, [1,6*DIM,1,SEQ] out concat(axiscax, values(oo, qf, kf, vf, af, xn))六个中间量被打包进同一个输出张量输出通道区间含义oo0 ~ DIM注意力输出模型继续用qf / kf / vfDIM ~ 4*DIMQ、K、V 投影结果af4*DIM ~ 5*DIM注意力后的隐状态xn5*DIM ~ 6*DIMRMSNorm 归一化后的输入FFN 前向核 stories_mil.h 同理输出(ffn_out, h1, h3, silu_out, x2norm)尺寸(2*DIM 3*HIDDEN) × SEQ。 为什么能合并成单输出因为 ANE 有单输入约束多输入请求会触发 0x1d 错误项目把所有张量都塞进通道维——输入用 slice 拆开输出用 concat 拼上风格完全对称。CPU 侧按 offset 切块读取反向传播零重算前向核执行完后训练主循环直接从同一个 IOSurface 输出缓冲区按通道偏移切出各段见 train_large.mio_read_fp16(kern[L].fwdAttn-ioOut, ac-o_out, 0, DIM, SEQ); io_read_fp16(kern[L].fwdAttn-ioOut, ac-attn_out, 4*DIM, DIM, SEQ); io_read_fp16(kern[L].fwdAttn-ioOut, ac-xnorm, 5*DIM, DIM, SEQ);更妙的一步在反向传播Q、K、V 根本不用先读回 CPU而是从 ANE 输出缓冲区直接拷进反向核的输入缓冲区见 train_large.mio_copy(kern[L].sdpaBwd1-ioIn, 0, kern[L].fwdAttn-ioOut, DIM, 3*DIM, SEQ);也就是说 QKV 三个 Tap 在内存里走了一趟ANE 输出 → ANE 输入的直通路径CPU 零参与。静态 vs 动态管线缓存还是重算项目有两套训练管线对同一问题的取舍不同见 README.md 的优化清单管线策略代表代码取舍静态train_large.mForward Taps 全量缓存stories_mil.h多占激活内存省掉全部重算动态training_dynamic/部分 Tap 注意力在 ANE 上重算mil_dynamic.h省内存但 QKᵀsoftmax 要再跑一遍动态管线的sdpaBwd1内核里明确写着 recompute attention——它只缓存 Q、K、V注意力分数让 ANE 在反向时重算。这组成了一个经典的光谱时间换空间 vs 空间换时间而 Forward Taps 是空间换时间这一端的最简实现一行concat零额外 FLOPsconcat 本身开销可忽略。性能收益与适用建议这套优化支撑了项目的整体训练性能M4 实测模型参数量每步耗时Stories110M12 层109M91 msQwen3-0.6B28 层 GQA596M412 ms如果你的工作负载也有加速器/异构设备只暴露最终输出的约束CoreML 推理模型、闭源 NPU 等可以借鉴 Forward Taps 的思路先列清单反向/调试/监控到底需要哪些中间量只 Tap 需要的通道维不会无限膨胀打包成单张量用通道维 concat 偏移切块兼容只有单 I/O 通道的设备让数据直通Tap 输出直接喂给下一个核的输入尽量避免读回 CPU 再写出去权衡重算对计算便宜、体积大的量如 SEQ×SEQ 的注意力分数选择重算对体积小的投影量DIM×SEQ选择缓存。小结Forward Taps 是 ANE 项目中最以小博大的优化之一一行concat把 Q、K、V、注意力分数、隐状态等中间量从 ANE 前向核带出来反向传播直接取用避免了 CPU 重计算省下的正是训练里最贵的几次矩阵乘。完整实现可从以下文件入手静态管线 Tap 生成training/stories_mil.h动态管线内核training/training_dynamic/mil_dynamic.h训练主循环与 Tap 读取training/train_large.m优化清单与性能数据README.md【免费下载链接】ANETraining neural networks on Apple Neural Engine via reverse-engineered private APIs项目地址: https://gitcode.com/GitHub_Trending/ane2/ANE创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 ANE性能优化：Forward Taps中间量输出避免CPU重计算
ANE性能优化Forward Taps中间量输出避免CPU重计算【免费下载链接】ANETraining neural networks on Apple Neural Engine via reverse-engineered private APIs项目地址: https://gitcode.com/GitHub_Trending/ane2/ANE在 ANE 项目中开发者逆向 Apple 私有 API_ANEClient/_ANECompiler让 Apple Neural Engine 这块被限制为只推理的 NPU 也能跑训练。而在前向/反向传播的衔接处一项名为Forward Taps前向中间量输出的性能优化让反向传播避免了 CPU 重计算前向核在 ANE 上顺手把 Q、K、V、注意力分数、隐藏层状态等中间量搭车输出CPU 只需切块读取不再重复计算。本文用通俗的方式拆解这一技巧。先搞懂问题训练为什么要多要几个输出推理时模型只需输出最终结果logits但训练的反向传播必须知道前向计算过程中的中间值中间量作用不缓存的代价Q、K、V计算 dQ/dK/dV 梯度必须重新做 3 次 QKV 投影矩阵乘注意力分数 probs计算 dQ 的 softmax 梯度必须重新做 QKᵀ softmaxFFN 的 h1、h3、SiLU 输出SwiGLU 激活的导数需要它们必须重新做 3 次大矩阵乘RMSNorm 归一化后的输入Norm 反向公式依赖原值必须重跑归一化ANE 的 IOSurface 共享内存 I/O 带宽非常可观fp16 直传比 fp32 快约 37%多搬运一点内存远比在 CPU 上多算一遍便宜。于是项目采用了 Forward Taps。实现原理一个 concat 操作打包所有 Tap整个技巧只有一行 MIL 操作在前向计算图的末尾用concat沿通道维把所有需要的张量拼接成一个输出。以注意力前向核为例源码 stories_mil.htensorfp16, [1,6*DIM,1,SEQ] out concat(axiscax, values(oo, qf, kf, vf, af, xn))六个中间量被打包进同一个输出张量输出通道区间含义oo0 ~ DIM注意力输出模型继续用qf / kf / vfDIM ~ 4*DIMQ、K、V 投影结果af4*DIM ~ 5*DIM注意力后的隐状态xn5*DIM ~ 6*DIMRMSNorm 归一化后的输入FFN 前向核 stories_mil.h 同理输出(ffn_out, h1, h3, silu_out, x2norm)尺寸(2*DIM 3*HIDDEN) × SEQ。 为什么能合并成单输出因为 ANE 有单输入约束多输入请求会触发 0x1d 错误项目把所有张量都塞进通道维——输入用 slice 拆开输出用 concat 拼上风格完全对称。CPU 侧按 offset 切块读取反向传播零重算前向核执行完后训练主循环直接从同一个 IOSurface 输出缓冲区按通道偏移切出各段见 train_large.mio_read_fp16(kern[L].fwdAttn-ioOut, ac-o_out, 0, DIM, SEQ); io_read_fp16(kern[L].fwdAttn-ioOut, ac-attn_out, 4*DIM, DIM, SEQ); io_read_fp16(kern[L].fwdAttn-ioOut, ac-xnorm, 5*DIM, DIM, SEQ);更妙的一步在反向传播Q、K、V 根本不用先读回 CPU而是从 ANE 输出缓冲区直接拷进反向核的输入缓冲区见 train_large.mio_copy(kern[L].sdpaBwd1-ioIn, 0, kern[L].fwdAttn-ioOut, DIM, 3*DIM, SEQ);也就是说 QKV 三个 Tap 在内存里走了一趟ANE 输出 → ANE 输入的直通路径CPU 零参与。静态 vs 动态管线缓存还是重算项目有两套训练管线对同一问题的取舍不同见 README.md 的优化清单管线策略代表代码取舍静态train_large.mForward Taps 全量缓存stories_mil.h多占激活内存省掉全部重算动态training_dynamic/部分 Tap 注意力在 ANE 上重算mil_dynamic.h省内存但 QKᵀsoftmax 要再跑一遍动态管线的sdpaBwd1内核里明确写着 recompute attention——它只缓存 Q、K、V注意力分数让 ANE 在反向时重算。这组成了一个经典的光谱时间换空间 vs 空间换时间而 Forward Taps 是空间换时间这一端的最简实现一行concat零额外 FLOPsconcat 本身开销可忽略。性能收益与适用建议这套优化支撑了项目的整体训练性能M4 实测模型参数量每步耗时Stories110M12 层109M91 msQwen3-0.6B28 层 GQA596M412 ms如果你的工作负载也有加速器/异构设备只暴露最终输出的约束CoreML 推理模型、闭源 NPU 等可以借鉴 Forward Taps 的思路先列清单反向/调试/监控到底需要哪些中间量只 Tap 需要的通道维不会无限膨胀打包成单张量用通道维 concat 偏移切块兼容只有单 I/O 通道的设备让数据直通Tap 输出直接喂给下一个核的输入尽量避免读回 CPU 再写出去权衡重算对计算便宜、体积大的量如 SEQ×SEQ 的注意力分数选择重算对体积小的投影量DIM×SEQ选择缓存。小结Forward Taps 是 ANE 项目中最以小博大的优化之一一行concat把 Q、K、V、注意力分数、隐状态等中间量从 ANE 前向核带出来反向传播直接取用避免了 CPU 重计算省下的正是训练里最贵的几次矩阵乘。完整实现可从以下文件入手静态管线 Tap 生成training/stories_mil.h动态管线内核training/training_dynamic/mil_dynamic.h训练主循环与 Tap 读取training/train_large.m优化清单与性能数据README.md【免费下载链接】ANETraining neural networks on Apple Neural Engine via reverse-engineered private APIs项目地址: https://gitcode.com/GitHub_Trending/ane2/ANE创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考