ANE 可行性：等价变换与可测量的 10× 目标

研究日期：2026-09-20。硬件：M3 Max，40 核 GPU，128 GiB。本文描述的是假设和数学边界，不是声称有新的实测加速。起点是原始 Laya 权重和更快的 MLX FP16 运行时。下面的工程实验和测量可能会取代这里的起始观察。

最好的第一个实验是对整个 transformer 做固定形状、channel-first 的重写，然后检查执行计划。最有可能实现数量级提升的目标是 每个已完成决策的能耗，前提是模型保持可用的延迟和准确率。只改一个 Core ML 计算单元设置，不足以作为 Neural Engine 执行了该模型的证据。

随后的 工程实验 实现了该重写，实测的速度/能耗报告 现在记录了结果。未压缩的候选提升了效率，但没有达到 10× 目标。下面的假设和原始分母作为研究记录保留；实测性能说法请使用后来的编译 MLX 对比。

优化之前先定义目标

对于相同的输入、checkpoint、精度策略和已完成决策数，定义：

t = elapsed time / completed decisions
P = average measured power over that same interval
E = integrated energy / completed decisions = P × t

S = t_MLX / t_candidate                    speed gain
R = P_MLX / P_candidate                    power reduction factor
S × R = E_MLX / E_candidate                 energy efficiency gain

在这个恒等式中使用区块平均延迟，而不是延迟百分位数。P50 和 P95 分开报告。一个速度快 2×、功耗为五分之一的方案，就是 10× 的能耗改善。一个速度只有一半的方案，需要 20× 的功耗降低才能给出同样的 10× 能耗改善。把速度乘以已经算出的能耗改善，就是重复计算时间。

有两种不同的功耗测试：

1. 饱和的顺序推理：测量实际吞吐量、延迟和每决策焦耳数。功耗更低但更慢的候选并不自动更高效。
2. 等提供负载，例如相同的 Snake tick 速率：两个候选都必须在截止时间内完成相同的工作。报告平均功耗、区间总能耗、截止时间错失和已完成决策数。睡得更久或丢掉工作不是优化。

记录所测量的功耗域。CPU + GPU + ANE 遥测不一定是整机或电池功耗，不得被那样标注。报告原始能耗，以及（如果可用）配对的扣除空闲能耗。当负载减空闲接近噪声时，保留不确定性，而不是静默钳制它并报告一个巨大的比值。把分词、输入拷贝、CPU 回退和后处理都保留在端点的计账边界之内。

起始证据与分母

当前的 多语言 MLX 基准 测量一个 91-token 问题，为 P50 7.870 ms / P95 9.870 ms。英文 MLX 基准 测量一个 93-token 问题，为 13.334 / 13.734 ms。这些是端到端 predict 测量，不含模型加载和预热。新的对比必须重跑最强、适用的 MLX 路径，在工作负载允许时包含它需显式启用的 compile 和 prompt 缓存设置；历史的 eager 结果不是永久的分母。

现有的 Core ML 多语言导出测得 CPU + GPU 下 11.277 ms P50、ALL 下 78.037 ms、CPU + NE 下 81.336 ms。后一份计划记录 1,318 个 CPU 首选算子和零个 NE 首选算子；24 个 SDPA 算子的设备元数据未知。这不构成任何 NE 执行的说法。该计划把许多算子列为 NE-支持，这与 NE-首选不同。Apple 把 计算计划的设备使用 描述为预期的设备使用，所以即使是理想的计划，也应当由运行时剖析或可观察的 NE 活动来佐证。

现有的 FP16 回归门禁是 100% 夹具 argmax 一致、有限且确定性的输出，以及校准概率和动作概率的绝对漂移不超过 0.02。这是在一个小语料上的转换保真度检查。它不能证明通用任务准确率、Snake 能力或压缩模型的质量。

等价的图变换

Apple 的 Transformer 部署研究 启发了四维 BC1L 激活、1×1 卷积、把注意力拆成多头，以及减少布局拷贝。它发布的 10× 示例是另一个模型、设备和基线；不能移植到这里的 MLX 对比上。把这些布局建议当作在本操作系统和芯片上待测试的候选，而不是当下完整的硬件支持契约。

线性投影与门控 MLP

令 X[b,l,i] 是现有的激活，并定义 U[b,i,0,l] = X[b,l,i]。对于一个线性层，

Y[b,l,o] = sum_i W[o,i] X[b,l,i] + bias[o]
K[o,i,0,0] = W[o,i]
Conv2D(U,K)[b,o,0,l] = Y[b,l,o]

这改变了布局和算子表示，但不改变实值函数。把这个布局保持贯穿整个编码器和两个 decision-head transformer 层。在每个线性层周围来回转换会抹掉收益。QKV 可以保持为单个 D → 3D 卷积；把它的通道输出拆成 Q、K 和 V。类似地，保留现有的融合 D → 2I 编码器投影，把通道拆成 value 和 gate，对 value 施加原始的精确 GELU，乘以 gate，再投影 I → D。

对于 FP16，序列宽度能被 32 整除也与 Apple 研究中描述的 64 字节末轴对齐相符。初始形状是短 API 夹具的 B=1,L=96 和紧凑 Snake 的 B=3,L=64。这里建议 32 的倍数遵循该缓冲模型；这不是允许把每个工作负载都补齐到一个任意的大长度。Snake 不需要 96 token。

注意力与 RoPE

对每个头，把 Q 和 V 保持为 (B,d,1,L)，并把 K 转置为 (B,L,1,d)。计算：

score[b,k,0,q] = sum_c Q[b,c,0,q] K[b,k,0,c] / sqrt(d)
weight = softmax(score + additive_mask, axis=key)
out[b,c,0,q] = sum_k weight[b,k,0,q] V[b,c,0,k]

在这个表示中，key 轴是轴 1。把头输出拼接在通道轴上。这是同一个注意力函数；softmax 轴选错会静默地改变它。Apple 的 参考注意力实现 演示了对应的两次四维收缩。检查转换后的 MIL 算子：写一个 einsum 并不保证得到预期的设备或 lowering。

在 QK 之前对每个头的通道对施加 RoPE。保留 checkpoint 的 split-half 约定、原始位置和逐层 theta。在这个多语言 checkpoint 中，完整 RoPE 和局部 RoPE 都使用 theta 160000。把所有头拼在一起后的前后半分开是错误的；要在每个 64 通道的头内部拆分。依赖位置的旋转通常不能折叠进单个与位置无关的权重矩阵。

局部规则是双向的 abs(q-k) <= 64，含端点。保留 key padding 和现有的 padded-query 规则。在固定形状下，与位置相关的常量部分可以在追踪之前算好；改变样本 padding 仍必须影响 key mask。用有限的大负数掩码替换数学上的排除是一种数值近似，除非它与原实现的有限精度行为一致；要验证对抗性输入和补齐输入。

LayerNorm 不能与其他归一化互换

对每个 (b,l)，只在通道维度上归约：

mu = mean_c U
v = mean_c (U-mu)^2
normalized = (U-mu) / sqrt(v + epsilon)
output = normalized * gamma + beta

保留原始的 epsilon、仿射顺序、总体方差、第一层的恒等归一化、精确 GELU 和残差顺序。Apple 的 参考 LayerNorm 使用不同的仿射顺序；它的 DistilBERT adapter 通过变换 bias 来补偿。直接拷贝那个类并加载 Laya 的 state dict，在 bias 非零时会出错。显式写出原始顺序的仿射表达式，也避免了除以一个可能为零的 gamma。

钳制激活、把 GELU 换成 tanh，或用 RMSNorm 替换 LayerNorm，都会改变函数。如果平方值溢出，正向重缩放是一个数学等价的选项：

normalize(x/a, epsilon/a^2) = normalize(x, epsilon), for a > 0

有限精度累加仍需要一致性测试。FP32 归约可能带来拷贝或 CPU 回退的开销，所以要检查计划，而不是静默地放松数值行为。

把不支持的算子移到模型边界

如果 embedding 查找、动态标记 gather 或 action tail 阻碍了一个连续的 NE 区域，就用这个划分做一个单独的候选：

CPU: tokenizer → selected embedding rows → embedding LayerNorm
NE candidate: all encoder layers → type embedding addition → both heavy head layers
CPU: marker/CLS selection → small scorer → raw-probability features → action head

只有端点跨越引擎。不要把每一层的注意力或 LayerNorm 都卸载到 CPU。在多语言 B=1,L=96 下，一个 FP16 embedding 张量是 147,456 字节；在 B=3,L=64 下是 294,912 字节。把这些拷贝以及任何完整隐藏状态的输出拷贝都计入端到端测量。

多语言 token 表有 196,608,000 个参数，但每个请求只 gather 它的 token 行。不必把它送进 ANE transformer 子图，也不得把它算作每次预测都完整读取整张表。它的 LayerNorm 不依赖位置，所以对表行做离线的预归一化在实值算术上是等价的。它可能改变舍入和存储精度，需要自己的一致性检查。action head 消费的是 原始 标记 logits 的概率，在温度校准之前；从公开的校准概率重建它的特征会改变 checkpoint 的行为。

关于 10× 延迟说法的算术限制

令 D 为隐藏宽度，I 为门控编码器中间宽度，N 为编码器层数，H=2 为 decision-head 层数。主要的每 token 矩阵参数量和密集算术是：

A = N(4D^2 + 3DI) + 12HD^2
F_dense(B,L) = 2BLA + 4B(N+H)L^2D

乘法和加法分别计数。这些等式不含 norm、激活、embedding、scoring、mask、拷贝和运行时开销。它们不是 profiler。即使对带掩码的局部层，它们建模的仍是当前的密集注意力计算。

这些是所需达到的速率，不是断言的 ANE 峰值规格。布局重写会去掉开销，但不会去掉那些密集投影。硬件还必须执行一条由 24 或 30 个 attention/MLP 块构成的顺序链。

一个乐观的流式模型给出另一个有条件的下限：

t >= max(F / effective_compute, bytes_from_DRAM / effective_bandwidth)

40 GPU 核心的 M3 Max 标称 400 GB/s 统一内存带宽。如果每个主 FP16 矩阵每次请求都要从 DRAM 取一次，那么即使完全用满该带宽，多语言也至少需要 0.622 ms、英文至少需要 1.842 ms。实际的 ANE 带宽访问可能更小，而缓存或压缩权重会改变这个假设。这不是无条件的物理下限。它说明了为什么在未压缩流式下 10× 英文延迟尤其苛刻，也说明了为什么即使延迟只是小幅改善，测量能耗仍然有用。

对于端到端时间中被加速因子 s 改善的实测比例 f，Amdahl 定律给出 S = 1 / (1-f+f/s)。即使某个区域被无限加速，除非它至少占原始延迟的 90%，否则也达不到 10×。当目标是每决策焦耳数时，类似的边界用的是实测能耗的比例，而不是 FLOPs。final-head 的 selected-query 优化只去掉了模型算术的几个百分点；在 L<=64 时局部注意力稀疏性也可忽略，因为局部窗口覆盖所有位置。两者都不能提供一条可信的、独立的 10× 路径。

压缩与架构改动有不同的契约

Apple 当前的 优化概览 指出调色板化可为 NE 带来内存/延迟收益，并标识出 A17 Pro/M4 上更新的 W8A8 计算路径。不要把那个更新硬件的加速外推到这台 M3 Max。量化性能指南 也警告激活反量化可能拖慢 CPU/GPU 执行。先取得一个驻留 NE 的基线，然后从 8 位向下测试权重调色板化，同时按需保留敏感的 norm/scorer。

把 FP16 权重打包到八位或四位，在计入元数据之前给出 2× 或 4× 的理想权重存储比。这并不等于同样的延迟倍数：解压、激活搬运和计算仍在。剪枝只有在所选表示真的利用了这些零时才有助于运行时。移除任意的头/层或做低秩截断需要质量恢复，并且若无保留说明就不能保持原模型的身份。

对于逐输入的 logit 误差界 ||z'-z||_infinity <= delta，一个充分的 argmax 证明是 top1(z)-top2(z) > 2delta。在校准温度 T 相同且为正时，softmax 的无穷范数 Lipschitz 界给出 ||p'-p||_infinity <= delta/(2T)。这些是在被评估输入上有用的诊断，不是量化的全局证明。保留单独的 close-margin 和多语言评估切片；饱和样本会掩盖大的 logit 误差。

三个实验与验收门禁

1. 固定形状的等价 NE 图。 导出多语言 B=1,L=96,K=4 和 Snake B=3,L=64,K=4，带 BC1L 投影、正确的逐头 RoPE、显式注意力和原始 LayerNorm/GELU。把输入数组和各层输出与原图比较。检查哪些主要投影和注意力块是 NE 首选，然后检查实际的运行时 NE 活动。单凭受支持算子的数量不算成功。用交替区块重跑对应的优化 MLX 基线。
2. 一个连续的 transformer 孤岛。 如果第一个图碎片化，就把 embedding 和小的最终 tail 移到 CPU 边界。在相同的功耗测量下，把它与整图候选比较。只有当完整预测在延迟或能耗上优于观察到的运行间波动时，才保留一个候选。包含所有拷贝；一个快速但孤立的编码器是不够的。
3. 在放置生效之后，做以能耗为重点的压缩。 先筛选 8 位调色板化，然后 6/4 位作为单独的近似变体。运行不变的门禁夹具、留出的 choice/score/noul 任务、多语言输入、接近的并列情况和 Snake 轨迹。把准确率、概率漂移和校准与性能一并发布。如果激进压缩或蒸馏改变了已学到的行为，就应归入一个单独命名的模型。

在做出发布说法之前，使用相同的 checkpoint/输入哈希和交替的基线/候选顺序；排除冷编译，但单独报告它。每个入围方案至少用五个持续区块，并保留原始延迟、完成调用和功耗采样。要求原始夹具一致和现有的概率容差，不得为了让候选通过而放宽它们，此外还要求稳定有限的输出和有界内存。把长输入和形状边界输入与短的固定形状演示分开测量。

只有当相关的速度、等负载功耗或能耗比值至少为十，且其不确定性支持该说法，同时满足所述延迟和任务质量限制时，才宣称 10×。如果不确定性的下界达不到十，就报告实测比值。在验证了 NE 执行的情况下，较小的能耗提升仍是有用的证据；它不是数量级的结果。

文档来源

使用了必需的 Context7 CLI 工作流：把 Core ML Tools 解析为 /apple/coremltools，然后查询 transformer 算子/布局 lowering 和 NE 压缩行为（共三条命令）。查阅了 Apple 的研究文章、参考源码、当前的 Core ML 优化文档、计算计划文档，以及上面链接的设备规格。模型等式、参数计数和起始测量来自本仓库及其 MLX 姊妹项目。本研究分支没有运行任何可竞争的 GPU/ANE 基准。