工程探究：把 Laya 的 transformer 搬到 ANE 上

研究环境：Apple M3 Max（40 个 GPU 核心，128 GiB 统一内存），macOS 27.2，Core ML Tools 9.0，PyTorch 2.7.0 和 NumPy 2.1.3。这是 experiments/ane_engineering/ 下一个独立的原型；已发布的运行时未改动。

当前结果

一个固定的 B=1, L=96 多语言原型成功地在 Core ML 的预期计算计划中，把完整的编码器、decision head 和 scorer 分配给 Neural Engine：6,390 个非常量算子首选 ANE，估计开销权重合计约为 1。其余 3,809 个条目是常量。原始的 enumerated-shape/SDPA 导出在 CPU_AND_NE 下对全部 1,318 个已分配算子都首选 CPU，尽管有 988 个单独算子把 ANE 列为受支持设备。

原型的完整预测路径在 50 次筛选调用中测得 p50 5.167 ms，包含分词、embedding 查找、attention mask、ANE 推理、CPU action-head 计算、校准和格式化。孤立的 Core ML 主体用合成 embedding 在 30 次调用中测得 p50 4.403 ms。后者是组件测量，不是端到端速度说法。两次热计时都不含模型加载和编译。

在原始 FP32 golden reference 的固定长度子集上，59/59 个答案比较一致，涵盖八种语言和 choice/score/noul 问题。最大的校准概率变化是 0.002925，100 次重复的公开调用都是有限的，并返回相同的四舍五入结果。原始的 63 问题夹具包含三个 1,024-token 长输入和一个 147-token 的 20 选项输入；这四次评估被 L96 导出显式跳过。该夹具中有重复的评分表。这些是回归比较，不是 59 个独立的带标签样本，也不能证明通用任务准确率未变。

单独的 L192 和 L1024 导出也把全部 6,390 个已分配的主体算子放在 ANE 上。L192 通过 60/60 个比较；L1024 通过完整的 63/63 golden 夹具。两者都通过 100 次重复的公开调用，最大校准概率误差保持为 0.002925。长输入子集本身最大误差为 0.001128。所有被评估的 token 使用计数都与参考一致。

这些是串行筛选运行，不是跨后端的配对比较。主体计时使用 5 次预热和 30 次测量调用；完整短问题计时使用 10 次预热和 50 次测量调用。完整路径包含主机端工作和输入检查。L96 筛选早于最后的额外输入校验检查；最终的受控对比使用当前的 adapter 并记录其源指纹。编译/加载时间是在转换后于每个进程内测得的，不是对框架缓存为空时首次系统启动的承诺。大的固定图即使对短请求也会做补齐的工作。一个实用的 adapter 会选择独立的长度桶；让每个请求都走 L1024 会丢弃短输入的优势。

对实际的 workload(1, long=True) 夹具单独跑一次，确认这是一个 1,024-token 请求，而不是被补齐到该大小的短请求。它在十次预热调用后、50 次完整预测中测得 p50 91.703 ms / p95 94.776 ms；所有四舍五入输出都保持稳定。历史的 MLX 长输入基准 是 p50 51.98 ms。它们不是同一轮的配对测量，但这次筛选没有提供任何证据表明当前的 ANE 图会加速长输入。输入哈希、实际 token 长度、当前实验指纹和原始计时都保留在 long1024-performance.json。

原始证据：

• 完整主体的计算计划与组件计时
• 真实输入的 FP32 参考验证与完整预测计时
• L192 验证
• L1024 完整夹具验证
• 单层计划与数值检查
• 单 MLP 计划与数值检查

MLComputePlan 描述的是预期的放置，不是硬件执行 trace。CPU_AND_NE 允许 CPU 和 ANE；它不是仅 ANE 的开关。在本实验中，所有已分配的 heavy-body 算子都首选 ANE，但运行时硬件遥测是单独评估的。随后的 Instruments 诊断记录了 Neural Engine 硬件活动；它的表是全局的，不能把每个事件都归给这个模型。配对的 MLX 对比、功耗积分结果和 trace 限制报告在 ANE_BENCHMARKS.md。如果不在本操作系统/设备上验证那个计数器，监控工具给出的 ANE 计数为零就不能证明 ANE 活动不存在。

为什么原图是一个糟糕的 ANE 目标

基线保留传统的 B×L×C transformer 布局、动态形状算子、整头批处理注意力和 Core ML 的 SDPA 算子。在 CPU/ANE 选择下，它的计划包含 24 个未报告设备分配的 SDPA 算子，以及许多 cast、slice、shape query、gather 和 transpose。有些单独算子支持 ANE，但整个图并未被划分到它上面。因此，单个算子受设备支持，不足以作为存在有用 ANE 执行路径的证据。

成功的原型同时改动了若干东西。它证明这个组合能够实现 ANE 放置，而不是一次已完成、能指出唯一罪魁算子的消融。固定形状、原始布局的 SDPA 和显式注意力的对照，是接下来有用的区分性实验。

Apple 已发布的 Transformer 指南推荐 channel-first 的 4D 张量、用于投影的 1×1 卷积、逐头注意力和更少的布局拷贝。这些原则启发了该实现；Apple 历史上的 DistilBERT 加速，并不能证明相对本项目已经很快的 MLX FP16 基线有 10× 改进。Apple 的 ANE Transformer 文章、Apple 的参考实现。

原型架构与数值契约

model.py 是一个单独的导出模型，由原始 checkpoint 参数构建：

• 隐藏激活使用 B,C,1,L。每个密集权重 W[out,in] 变成一个 1×1 卷积核 K[out,in,0,0]，无需重新训练或权重近似。
• 注意力被拆成单独的 64 通道头。key 张量只转置一次，两次显式 einsum 计算 QK 和 AV，同时保持 4D 布局。Softmax 在 key 轴（维度 1）上进行。
• RoPE 把每个头拆成两个 32 通道的半部分。它的余弦、正弦和基底来自原始模型，包括多语言局部 theta 160000。
• 通道归一化保留原始的 normalized * weight + bias 顺序和 epsilon。它不拷贝 Apple 参考 LayerNorm 中顺序不同的仿射表达式或可选裁剪。
• 第一个编码器注意力 norm 保持为恒等；编码器使用精确的 erf GELU，而两个 decision-head FFN 保留 ReLU。
• 完整注意力和滑动窗口掩码保留 valid-key 掩码和 padded-query 规则。局部半径从 local_attention // 2 读取。
• 最终的标记 gather 用一个外部准备好的 one-hot selector 和一次 4D einsum 表达。图保留 32 个标记槽，包括非活动的槽；主机把非活动 logits 替换为原始的 -1e4 值。

在 PyTorch 布局检查中，一整个原始 FP32 编码器层与它的 BC1S 对应物最多相差 2.29e-5。FP16 Core ML 的层输出误差更大，这对这种精度/后端改动来说是预期的。主体探针最初包含一次自我比较；那份无效证据已被移除，它的整主体 PyTorch 布局检查被显式标为 not_measured。真正的完整模型验证改为对照存储的原始 FP32 logits 和决策。

test_ane_layout.py 中独立的 CPU 回归另外把一个完整的微型 ConvBody 与原始 DecisionModel 做比较，使用显式和 SDPA 注意力 oracle、三种问题类型、改变后的 padding 值、非零 norm bias、多个 RoPE 基底和一个非默认的 norm epsilon。这些测试覆盖布局和掩码语义，且不会把它们与完整 checkpoint 的 FP16 硬件精度混为一谈。全部五个布局测试都在本地通过。

注意力原型增加了一个有限的 -1e4 掩码偏置。这在有限的已验证激活上具有预期的掩码行为，但对于任意的极端输入，它与把被掩码的分数替换为 -1e4 或 -infinity 并非逐位相同。同样，也不声称 Core ML FP16 执行与原始 FP32 模型逐位相同。

runtime.py 为整个 transformer 提供一个 CPU→ANE→CPU 边界，而不是每一层都做一次设备切换：

1. CPU 对每个问题分词，只 gather 被请求的 embedding 行，并构造固定形状的加性掩码、类型向量和标记 selector。它不在问题之间复用上下文隐藏状态或 K/V。
2. 一次 Core ML 调用执行 embedding 归一化、全部 22 个编码器层、两个 decision-head 层和 scoring 卷积。
3. CPU 从 未校准的原始 logit softmax 推导 action 特征，并以 FP32 配 erf GELU 运行小的原始 action head。随后公开的校准和输出格式化使用现有实现。

导出的主体使用 FP16 计算，而小的主机端 action head 是 FP32。embedding 查找使用原始存储的权重。源 safetensors 文件包含 169 个 FP16 张量和一个 FP32 张量：“FP32 参考”描述的是原始 PyTorch 执行，并不是声称原始 checkpoint 全部以 FP32 存储。这个混合精度边界是原型数值契约的一部分。在饱和夹具上动作概率差为零，并不能证明 action-logit 相同。

adapter 会拒绝不兼容的包维度、超出范围的 ID/标记、无效的掩码值、空的 attention-key 行，以及超出其固定容量的输入长度。它的默认输入上限是 96 token，批大小为一；多个问题顺序执行。它绝不会为了适应较短的导出而静默截断请求。

复现、来源与质量门禁

在仓库根目录、钉定的 .venv 中运行。生成的包被 Git 忽略；不需要重复的 embedding NPZ 或大的权重产物。probe.py 拒绝已存在的输出目录。新的导出会在一个 manifest 中记录原始权重/配置 SHA256 值、包内容哈希、形状、工具版本和实验源指纹。

复现命令写入 artifacts/ane-repro/ 下，因为提交进仓库的实验目录已经包含报告和 manifest。重复导出时请另选一个新目录；两个导出器都不会静默复用已存在的包。

用 pip install -e '.[convert,dev,research]'（或对应的 uv sync extras）安装转换、开发和压缩研究依赖。research extra 钉住 kmeans1d==0.4.0；这个可选依赖用于 FP16 分组 K-means 实验，并记录在新的 manifest 中。

默认情况下，转换会把 laya-multilingual 解析为本仓库钉定的 Hugging Face checkpoint，并在必要时下载它。传入 --source /path/to/checkpoint 以使用已有的本地文件。验证默认使用包 manifest 中记录的源目录。ANEAgent 运行时本身只接受本地文件；它在加载前验证原始权重/配置哈希、固定形状和包内容哈希。验证还会拒绝源权重哈希不同的 golden reference。实测的多语言源权重 SHA256 是 9d628fd971b700382ac6f65920a86f149777b2e748e0c955fb3b19695aa8f204。

# Small placement probes, then the complete model.
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind mlp --length 96 --output artifacts/ane-repro/mlp96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind layer --length 96 --output artifacts/ane-repro/layer96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 96 --output artifacts/ane-repro/body96
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --length 96 --repeats 100 \
  --output artifacts/ane-repro/validation96.json

验证器要求所有被评估的 argmax 决策一致、校准概率和动作概率误差 <=0.02、输出有限、token 使用不变，以及重复的公开输出完全相同。失败的候选会写入 passed: false 并以非成功状态退出。即使固定形状子集通过，被跳过的用例仍未验证。最初的 L96 报告是在测量之后才显式套用这些门禁的；它已相应标记，且不改变其记录的计时。

更长的固定导出及其完整的 golden-reference 验证命令：

.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 192 --output artifacts/ane-repro/body192
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body192/model.mlpackage --length 192 \
  --output artifacts/ane-repro/validation192.json
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 1024 --output artifacts/ane-repro/body1024
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 \
  --output artifacts/ane-repro/validation1024.json
.venv/bin/python -m experiments.ane_engineering.benchmark \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 --long \
  --output artifacts/ane-repro/long1024-performance.json

这些命令复现上面列出的、经独立检查的更长导出。它们的放置和数值保真度是与 L96 图分开检查的；把短请求补齐到 1024 token 不是所提议的生产策略。

压缩筛选与 10× 目标

palettize.py 准备独立的 weight-only 调色板变体，用 8 位均匀查找表作为廉价的首个筛选候选。只选择大于 2048 个元素的卷积权重；RoPE 常量、归一化、激活数学和主机端 action 权重保持不变。分组输出通道使用独立的查找表。K-means 是更昂贵的模式。它对这些 FP16 分组使用已安装的 kmeans1d 实现。当 num_kmeans_workers > 1 时，Core ML Tools 9.0 通过进程池的 starmap 并行处理独立分组；实验对离线 K-means 使用八个 worker，每个 worker 一个数学库线程。worker 数量改变导出的吞吐量，而不改变所意图的码本目标。更低位的 6/4 位候选是单独的近似产物，不是精确实现。

压缩后的体积指的是 导出的 transformer 主体包。原始的 1.96608 亿条目 embedding 表留在主机上，每个请求只查找被请求的行，而小的主机端 action 权重保持不变。主体包缩小约一半，并不等于整个 checkpoint、运行时内存或每请求能耗也减半。

.venv/bin/python -m experiments.ane_engineering.palettize \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --bits 8 --mode uniform --group-size 32 \
  --output artifacts/ane-repro/body96-w8
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96-w8/model.mlpackage --length 96 \
  --output artifacts/ane-repro/validation96-w8.json

# Independent K-means candidates; inspect each validation exit status.
for bits in 8 6 4; do
  VECLIB_MAXIMUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 \
    .venv/bin/python -m experiments.ane_engineering.palettize \
    --package artifacts/ane-repro/body96/model.mlpackage \
    --bits "$bits" --mode kmeans --group-size 32 --workers 8 \
    --output "artifacts/ane-repro/body96-w${bits}km"
  .venv/bin/python -m experiments.ane_engineering.validate \
    --package "artifacts/ane-repro/body96-w${bits}km/model.mlpackage" --length 96 \
    --output "artifacts/ane-repro/validation96-w${bits}km.json"
done

两个均匀 W8 筛选都保留了全部 59 个夹具 argmax 决策并通过 100 次重复调用，但未通过概率门禁：分组大小 32 达到 0.023612 误差，分组大小 4 达到 0.033858。W8 K-means/group32 筛选以最大误差 0.014393 和 59/59 个决策通过。对于那个 K-means 候选，饱和的动作概率掩盖了高达 16.24 的 action-logit 差异；通过这个小回归夹具，不能证明保持了通用校准或任务准确率。压缩候选始终是单独标识的近似模型。

固定的 W6 K-means/group32 候选也保持 59/59 个 argmax 决策和稳定的重复输出，但以最大概率误差 0.052243 失败。它的最大 action-logit 误差是 76.62。因此，把权重降到六位并不能满足不变的验收门禁，尽管它的短请求筛选延迟仍接近 FP16。

W4 K-means/group32 也保持 59/59 个决策，但最大概率误差升到 0.200221，最大 action-logit 误差升到 605.30。它未通过同一门禁。全部五个压缩筛选中都没有 argmax 变化，这说明了为什么单凭该夹具的饱和决策作为验收测试是不够的。

全部压缩变体都保留 6,390 个 NE 首选的已分配算子、59/59 个夹具 argmax 一致和 100 次稳定的重复调用。其余计划条目包含常量和权重-LUT 重建表达式；单凭放置元数据不能证明一次请求中有多少压缩数据从 DRAM 传输。三个 K-means 导出在八个离线 worker 下分别耗时 186.50、56.13 和 23.90 秒。设置和 worker 进程都在每次推理测量之前结束。

这些 50 次调用的串行筛选并不能确立加速比。FP16 筛选早于最后的输入校验检查，而且这些筛选没有交错进行。W8 K-means 是 ANE_BENCHMARKS.md 中更强的同会话对比的唯一压缩入围方案。被拒绝的变体作为精度边界的证据保留，不是推荐的部署。压缩只在这个多语言 L96 子集上验证过；上面 63 问题的 L1024 结果针对的是单独的 FP16 导出。

Core ML 调色板压缩从带索引的查找表重建浮点权重；更小的存储张量本身不能确立更快的推理或更低的能耗。每个变体都需要相同的精度门禁、一份新的计算计划和一次配对的端到端速度/能耗对比。Core ML 调色板化文档。

最初成功的 ANE 放置确立了一条可信的优化路径，而不是一个 10× 结果。对比必须使用 MLX FP16，在更快时包含其编译变体，并报告相同的任务范围。功耗必须在完整的请求上积分。总的系统能耗和任何扣除空闲的估计都应当连同其测量限制一并报告。独立的数学审查见 ANE_MATH.md。

手工实现确立了什么

这里有用的手工工作，是把计算图完整地、经独立验证地重写为一个 Core ML 能映射到 ANE 的布局。这在保留训练参数的同时改变了执行放置。它比在原图上改 compute_units 有效得多。它并没有去掉那 24 个顺序的 attention/MLP 块或它们的密集投影工作。

长请求的下一个精确候选是真正只访问局部窗口的注意力，用固定的 query/key tile 和原始的 padding 与 RoPE 规则实现。当前的图仍然计算一个密集的分数矩阵并施加局部掩码。tile 化的图可以减少那部分工作，但更多的 slice、边界和小收缩可能损害 ANE 调度；它的放置、精度和端到端收益仍未测量。在上面这些失败之后，进一步的权重压缩需要校准或质量恢复。蒸馏或更少的层会引入一个新模型，并需要更广的任务质量评估。这些尚未实现的方向，没有一个能提供今天就有 10× 收益的证据。