工程調查：這個 MLX 埠還能再快 10× 嗎？

日期：2026-09-19。機器：Apple M3 Max，40 個 GPU 核心，128 GiB 統一記憶體， macOS 27.2，MLX / MLX Metal 0.32.2，FP16 推理。本報告包含真實的 本地實驗，其中包括一個手寫的 Metal kernel。它不改變 生產執行時，也不釋出量化權重。

這些經過測試的工程改動沒有帶來 10×。 交錯測量支援來自編譯和 裁剪最後決策頭層未使用輸出的、適度的、依賴形狀的改進。用成對的每輪中位數 衡量，選定用例改進了約 3–8%。一些更大的批區間不包含改進。一個自定義的 精確 erf GELU/gate kernel 在數值上成功，但相比 MLX 編譯沒有提供 一致的額外端到端收益。樸素的 8-bit 和 4-bit 主幹量化減小了儲存， 卻沒能加速更大的試點工作負載，並改變了預測或校準機率。

數學極限和近似權衡另行在 MATH_10X_RESEARCH.md 中考察。最初的實現審閱 在 PERFORMANCE_RESEARCH.md 中；已釋出 checkpoint 的基準仍然是 BENCHMARKS.md。

實驗控制與侷限

所有研究 GPU 工作都序列執行。其他 agent 工作只用 CPU/檔案系統/網路。 機器使用交流電，實驗期間沒有記錄到 pmset 散熱/效能警告，也沒有報告 使用交換區。正常的桌面活動繼續。這不是一個受控的恆溫箱，也不是一臺在 其他方面空閒的專用基準機器。

最初的篩選執行讓每個候選在全新程序裡跑，帶 4–5 次預熱和 12–16 個樣本。 它們揭示了明顯的執行間漂移。例如，英文單問題試點暗示有 1.24× 的編譯改進， 而隨後的交錯實驗只發現約 1.03×。因此順序試點的延遲是篩選證據，不是主要 的因果加速主張。

確認指令碼 paired.py：

• 在每一輪內輪換候選順序，併為該輪中的每個候選使用相同的輸入。
• 在各輪之間改變實際的 state 文本。它生成至多 16 個 state 變體，並保留 張量形狀相同的變體；multilingual 短用例有 10 個這樣的變體，而其他報告的 用例有 16 個。
• 使用不同的自然語言問題，最大的短工作負載包含 50 條不同指令。它檢查輸入 雜湊，不快取答案、不去重問題，也不復用上下文編碼器狀態。
• 在停止每個計時器之前求值結果並同步 GPU。它同時測量已準備的前向呼叫 和公開的預測路徑，包含分詞和輸出格式化。模型載入不計入。
• 在預熱之後，英文頭/編譯實驗跑 32 個測量輪，multilingual 和自定義 Metal 實驗跑 16 個。每個候選看到相同的輪數和輸入序列。

這些輸入不同於已釋出的基線固定樣本。下面的比較是研究實驗內部的， 而不是通過相除不相關的表得到的 before/after 比較。研究的批上限是 64， 而已釋出 API 預設是 16。重複的 state 變體是有意的重複測量；沒有結果快取。

analyze.py 計算每輪的 eager_time / candidate_time 比值及其中位數的探索性百分位 bootstrap 區間， 使用 2,000 次輪索引重取樣。這些區間沒有考慮作業系統噪聲或序列相關的 每一個來源，也不能替代多會話復現。把獨立計算的 p50 值相除得到的比值 可能與成對中位數比值不同。

原始 JSON 包含所有計時、輸入雜湊、環境後設資料、保真指標，以及測量時記錄的 源指紋。實驗指令碼隨後被格式化，並擴充套件了不重疊的可選候選；更早的指紋描述的 是那些更早的指令碼版本。

編譯與精確的末頭裁剪

比較了四條路徑：

1. Eager： 已釋出的 FP16 DecisionModel。
2. 編譯： 在已載入、已求值、凍結的模型外包裹 mx.compile，使用 常規的形狀特化。
3. 選中的 Q + 編譯： 在最後一個頭層中保留全長的 QKV 投影和 K/V，但只 發出 CLS/選項標記的注意力 query。只對這些選中的輸出執行輸出投影和 FFN。
4. 全注意力 + 選中輸出 + 編譯： 保留原始的全長 QKV 和 SDPA 呼叫，然後 在輸出投影和 FFN 之前收集 CLS/選項輸出。這保留了原始的注意力 kernel 形狀， 同時移除大部分未使用的末頭稠密工作。

兩個裁剪原型都保留了模型的數學依賴關係。它們仍然計算所有 QKV 投影； 它們沒有實現數學上界計算中那部分額外的僅 Q 投影節省。改變 GEMM 和 SDPA 形狀可能改變浮點舍入。兩個原型都不是解碼器快取、提前退出，或丟棄 更早 Transformer 層的近似。

端到端 p50 延遲，毫秒：

來源：英文成對資料和 multilingual 成對資料。

對全注意力/選中輸出路徑，成對中位加速和探索性 95% 區間包括：

英文 50 問題和 multilingual 長批的區間包含 1。它們不能確立一個可重複的 改進。選中的 Q 路徑在 multilingual 短 16 和長 1 用例上略好，但沒有哪條 裁剪路徑在每個形狀上都佔優。所有候選區間、前向測量和原始每輪比值都在 paired_analysis.json。

在這個頭/編譯實驗中，編譯在兩個模型系列的 1,530 個變化輸入問題比較上， 與 eager 的 logits、動作 logits 和校準機率完全一致。兩條裁剪路徑在所有 1,530 個 argmax 決策上一致，最大的校準機率差異為 0.0001883。全注意力 裁剪路徑還通過了每個模型單獨的 63 問題固定樣本套件：126/126 一致，英文 最大機率差異 4.31e-5，multilingual 為 6.48e-6。這些是迴歸檢查，不是對 1,530 個獨立帶標籤樣本的任務準確率主張。

整模型編譯和逐塊編譯都做了篩選。塊實驗也保留了全部 63 個英文固定樣本 輸出，但沒有確立相對於整模型編譯的實質優勢。形狀特化在一個服務裡必須有界。 模型使用依賴形狀的 Python reshape 和掩碼，所以不加區別地應用 shapeless=True 是不安全的。官方編譯指南 記錄了形狀特化和狀態捕獲。

第一次英文整模型候選呼叫耗時 2,166.7 ms，隨後在那個試點中是約 12.75 ms 的熱態前向 p50；一個新的 B16 形狀首次呼叫耗時 272.4 ms。該 JSON 欄位名叫 cold_forward，但它指的是在 eager 參考推理之後的第一次候選呼叫，不是 完全冷啟動的應用或全新初始化的 Metal driver。後續候選複用了此前編譯的 Metal kernel，所以它們的首次呼叫時間不是冷啟動成本的有控排名。試點中已編譯 英文短 1 的 MLX 活動/峰值記憶體約為 803.6/918.6 MiB；multilingual 約為 614.1/676.9 MiB。這些分配器測量不包含每一處宿主側編譯器分配，也不能確定 在無界形狀變動下的記憶體上限。見 英文編譯試點和 multilingual 編譯試點。

選擇性量化：儲存節省有用，但不適合作為加速主張

這個原型在載入稠密 FP16 模型之後呼叫 nn.quantize。它只選擇 encoder.layers.* 線性模組，仿射組大小 64，然後編譯得到的模型。嵌入、 歸一化、決策頭、評分器和動作頭保持 FP16。這避免了讓打包的整數權重經過 當前的稠密載入器，也避開了動作頭不可整除的 1028/772 輸入寬度。沒有釋出 任何量化 checkpoint 格式或載入契約。官方 MLX 量化層實現 提供了這個選擇機制。

multilingual 的 4-bit 結果說明了為什麼僅靠小型固定樣本套件不夠：它的 63 個 固定樣本 argmax 保持不變，但 26 個不同工作負載決策中有 7 個改變。這些是 與 FP16 的一致性測量，不是真值準確率測量。0.8008 的絕對機率變化就是 80.08 個百分點。

在英文短 16 的試點輸入上，FP16 eager/編譯的端到端 p50 是 91.26/87.94 ms； 8-bit/4-bit 編譯是 96.66/93.20 ms。短 1 的量化在那次篩選執行中看起來略快， 而更大的形狀沒有。multilingual 大形狀篩選也沒能顯示速度優勢，但它的順序 執行有明顯的漂移。這些觀察為拒絕一個無條件的加速或釋出主張提供了理由， 而不是在沒有交錯量化復現的情況下指定精確的減速因子。進一步的量化工作需要 啟用感知的校準或微調，以及一個有代表性的帶標籤質量套件。

原始來源：英文 8-bit、 英文 4-bit、 multilingual 8-bit、 multilingual 4-bit。

手寫 Metal：精確的 GELU/gate 融合已實現並測試

kernels.py 實現了一個真正的自定義 Metal kernel，它讀取兩個拼接的 MLP 分支，計算相同的基於 erf 的 GELU，乘以門控， 並寫出單個輸出。它不替換成 tanh-GELU 或 sigmoid 近似。該 kernel 使用 MLX v0.32.2 自帶的 erf 和 expm1 輔助函式，並在 vendor/README.md 中保留它們的許可和宣告。它顯式 只支援 FP16，並使用安全的 Metal 數學模式。官方自定義 kernel 指南 描述了這套 API 及其數學模式控制。

在八個有代表性的啟用形狀上，27,958,016 個隨機生成的 FP16 輸出元素與原始 操作的值完全相等。這個微基準比較的是數值相等，不是零的符號位。全模型 變化輸入測試也完全匹配：跨兩個模型系列的 474/474 個問題比較，外加兩個 63 問題固定樣本套件，logit、動作 logit 和校準機率差異均為零。

這個正確性結果並沒有轉化為相對 MLX 融合編譯表示式的一致速度優勢。例如， 在 1,312 token、中間寬度 2,624 時，每次呼叫同步的啟用計時對 eager 的 GELU-then-gate 是 0.378 ms，對 mx.compile 是 0.268 ms，對自定義 kernel 是 0.280 ms。在 8,192 token、寬度 1,152 時，對應值是 0.846/0.764/0.714 ms。 這些微基準包含排程和同步開銷，是篩選探針；它們不是對隔離裝置執行時間的 測量。完整輸入、原始計時和相等性檢查都在 microbench.json。

隨後該自定義 kernel 被裝入每一個編碼器 MLP，並在輪換候選順序和變化 輸入下於完整模型中測量：

完整的自定義 kernel 成對執行使用一個權重相同的第二個模型例項，這樣未修改 和自定義實現可以共存而不發生變更或陳舊的編譯捕獲。它們的絕對計時不得與 更早的末頭裁剪執行比較。這些適度的混合結果不支援把自定義 kernel 作為通用 效能改進來發布。來源： 英文 Metal 成對資料和 multilingual Metal 成對資料。

自定義工程在哪些地方值得進一步調查

模型已經呼叫 mx.fast.scaled_dot_product_attention、mx.fast.rope 和最佳化過的層歸一化。它的 D64 布林掩碼 SDPA 路徑是融合的；並不存在一個能 解釋 10× 差距的缺失 Flash Attention 開關。區域性注意力仍然遍歷稠密的 key/value tile。一個真正的雙向視窗 kernel 可以跳過那些 tile，同時保留 包含端點的距離 <=64 和填充語義，但它的整模型算術機會在短輸入上很小， 在已釋出的長形狀上也是有界的。現有原始碼審閱 和數學報告量化了這一區別。

有用的後續專案及各自的證據要求是：

• 長輸入視窗注意力： 針對 D64、實際的雙向視窗和填充批特化 tile 邊界。 在 512/1024 token 上與融合的稠密 SDPA 比較，然後在完整模型中比較。 本報告沒有構建或基準測試這個 kernel。
• 稠密 kernel 的尾聲與排程： 調查把帶門控 MLP 的尾聲融合進 GEMM，或改進 短 M 的矩陣排程。MLX 已經使用專用的 Metal GEMM 實現，所以替換它們需要 真實的排程/kernel 剖析，以及在精確 M/N/K 形狀上實測的收益。單獨的啟用 結果表明，為什麼僅再寫一個逐元素 kernel 並不夠。
• 長度感知的批處理與共享 CPU 準備： 在構造每個問題序列之前，把共享 state 文本只分詞一次，同時保留精確的輸入 ID，並避免把小項填充到不相關的 長項。multilingual 長 8 試點花約 13.1 ms 準備輸入， 而端到端是數百毫秒。即使完全消除那部分準備，在這個工作負載上也不會產生 10×。排隊延遲和唯一推理數必須成為任何批處理主張的一部分。
• 一個更小的聯合作答學生： 如果 10× 是產品需求，就蒸餾或重設計模型， 以移除大部分稠密工作，或用一次上下文編碼回答許多固定問題。這會改變學習到 的模型，需要代表性的帶標籤訓練/評估；它不是一次精確的埠最佳化。在當前的 雙向編碼器裡跨問題複用任意的上下文 state/KV 是無效的。

八個獨立的 FP16 編碼器輸入投影 GEMM 探針達到了 0.66–11.55 TFLOP/s（含每次呼叫的同步）。英文大型的 M=4096, N=5248, K=1024 探針達到 11.55 TFLOP/s；multilingual 的 M=8192, N=2304, K=768 探針達到 7.96 TFLOP/s。這些是觀測到的吞吐 值，不是硬體峰值指標，也不是全圖吞吐的上界。小 M 的測量尤其被提交和 同步成本主導；流式的圖會以不同方式攤薄它們。它們顯示哪些形狀值得剖析， 而不是證明不存在更好的 kernel。數學報告中的 10× 同工作量吞吐預算仍是理論 要求，而不是實測的裝置能力。

復現與釋出決定

指令碼使用現有的 .venv 和本地鎖定版本的 checkpoint。按順序執行 GPU 命令， 絕不要與正式基準並行：

# Screening: repeat for eager, compiled, blocks, q8, q4, selected-compiled.
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant compiled --iterations 12 --warmup 4 --quality \
  --output experiments/engineering/reproduced-compiled.json

# Primary confirmation, including 50 genuinely different questions.
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 32 \
  --output experiments/engineering/reproduced-laya-paired.json
.venv/bin/python -m experiments.engineering.paired \
  --model laya-multilingual --iterations 16 --cases short1,short16,long1,long8 \
  --output experiments/engineering/reproduced-multilingual-paired.json

# Hand-written kernel microbench and complete-model comparison.
.venv/bin/python -m experiments.engineering.microbench
.venv/bin/python -m experiments.engineering.paired \
  --model laya --iterations 16 --cases short1,short16,long1 --metal \
  --output experiments/engineering/reproduced-metal-paired.json
.venv/bin/python -m experiments.engineering.run_variants \
  --model laya --variant metal-compiled --iterations 5 --warmup 3 \
  --cases short1 --quality --output experiments/engineering/reproduced-metal-quality.json

# CPU-only paired analysis.
.venv/bin/python -m experiments.engineering.analyze

所有實驗性 Python 檔案都通過 Ruff 格式化和 lint 檢查。穩定的執行時、原始 基準結果和已釋出的 FP16 checkpoint 仍是釋出工件。在冷形狀/快取策略和更 廣泛的質量驗證之後，編譯和精確的末頭裁剪是可信的可選的未來最佳化；實測 收益不足以證明悄悄把編譯延遲或自定義 kernel 加入預設路徑是合理的。沒有 聲稱 10× 加速、可用於生產的量化 checkpoint，或實測的區域性視窗 kernel 收益。