工程探究：把 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× 收益的證據。