Laya MLX 效能研究

研究日期：2026-09-19。目標：Apple M3 Max，40 個 GPU 核心，128 GiB 統一記憶體，MLX/MLX Metal 0.32.2。這是對原生執行時、已安裝的 MLX 實現、官方文件和現有基準 JSON 的靜態審閱。本次研究沒有執行任何 GPU 基準或模型推理。本報告中下面提出的任何最佳化都沒有實測加速。

最初的實驗應該是整模型編譯和有代表性的批排程，然後是選擇性的量化矩陣乘法。它們處理的是佔主導的重複工作。針對區域性注意力的專用 kernel 對長輸入來說是一個可信的長期專案。精確裁剪最後決策頭層是可行的，但它對整模型的算術節省只有幾個百分點。在不改變 checkpoint 的前提下取得大幅改進，需要改進稠密主幹、消除真正冗餘的請求，或找出一個經實測的實現瓶頸；僅僅替換一個啟用函式或再開啟一個 attention 標誌不太可能夠用。

現有測量確立了哪些結論

下面這些是現有的端到端中位延遲，包含提示準備和結果格式化，帶同步的 GPU 完成、五次預熱和 50 次測量迭代。模型載入和下載不計入。基準允許 64 個問題的批；公開執行時的預設值是 16，所以 50 問題的結果不是預設 API 配置。

來源：Laya FP16、Laya FP32、Laya MPS、multilingual FP16、multilingual FP32，以及 multilingual MPS。長輸入對 Laya 含 512 個 token，對 multilingual 含 1024 個；因此比較它們的長輸入延遲並不是在相同序列長度下比較。短輸入填充長度分別是 93 和 91。Torch 的比較必須保留 FP32 標籤：與 MLX FP16 相比時，它們把後端變化和精度變化合在了一起。

執行之間存在明顯的波動。例如，multilingual FP16 長 10 問題執行的 p50 是 389.487 ms、p95 是 462.319 ms、最大值是 619.663 ms。它短單問題的前向中位數是 8.023 ms，而獨立測得的端到端中位數是 7.390 ms。把這些中位數相減會得到一個荒謬的負預處理時間。當前檔案沒有隔離 tokenizer、Python 排程、單個 GPU kernel 或同步的開銷。它們確立了有用的基線，而不是 kernel 級的瓶頸診斷。

驗證報告 記錄了三個 checkpoint 在 FP32 和 FP16 下各 63/63 的 argmax 一致，以及每個變體 100 次有限、確定性的重複呼叫：共 378/378 個答案一致和 600 次重複呼叫。這是在一個小型固定樣本語料上的迴歸檢查，其中包含重複的問題。它們不是「未來的量化或架構改動會保持通用任務準確率」的證據。

為什麼稠密矩陣乘法值得優先

模型實現 對每一個填充過的 token 施加 QKV 和輸出投影、一個帶門控的編碼器 MLP，以及兩個常規的 Transformer 決策頭層。設 D 為隱藏維度，I 為編碼器中間維度，N 為編碼器層數，H 為決策頭層數。這些塊中每個 token 用到的矩陣權重數量是：

A = N * (4 * D^2 + 3 * D * I) + H * 12 * D^2
dense FLOPs per batch ~= 2 * B * L * A
dense attention FLOPs ~= 4 * B * (N + H) * L^2 * D

這些估計把乘法和加法分開計數，並排除歸一化、啟用、嵌入、評分、掩碼、記憶體搬運和 kernel 開銷。它們是一個算術模型，不是執行時剖面。

multilingual checkpoint 約 3.22 億的總參數裡有 1.966 億是嵌入參數。推理時只收集選中的嵌入行；這不是一次完整的詞表投影。它主要的逐 token 矩陣工作量大約是英文模型的三分之一，儘管兩者的總參數量看起來接近得多。這與實測的短批延遲差距一致，儘管它並不能證明某個特定的硬體瓶頸。只量化嵌入主要會減小常駐權重的體積，對 multilingual 尤其如此；它不必改善推理延遲。

在當前稠密注意力路徑下，attention 乘積在 Laya 的 L=93 時約佔建模 FLOPs 的 1.5%，Laya 的 L=512 時佔 7.9%，multilingual 的 L=1024 時佔 23.3%。它們佔真即時間的比例可能大不相同。在投入自定義 Metal 工作之前，剖析器應當區分矩陣 kernel、注意力 kernel、逐元素 kernel、CPU 圖構造和空閒間隙。

按優先順序排列的實驗

編譯完整的、被評估的推理路徑

Agent.forward() 目前構造陣列、呼叫 self.model 並求值結果。在模型或塊層級沒有外層的 mx.compile。固定形狀編譯可以減少 Python 圖構造並融合受支援的操作。MLX 記錄了形狀特化和顯式狀態捕獲；它也警告說，無形狀編譯無法安全地保留任意依賴形狀的 Python 操作。見官方編譯指南。

從載入、轉換並求值權重之後建立的已編譯可呼叫物件開始，使用常規的形狀特化。保留現有的未編譯可呼叫物件，用於一致性比較和 CPU 相容。對於一個凍結的推理模型，權重可以保持為該模型例項所捕獲；如果權重或模組結構改變，重建該可呼叫物件或顯式捕獲相關狀態。不要在替換 checkpoint 時複用一個已編譯閉包。

測試一個整模型包裝器；如果追蹤限制或編譯成本讓它不划算，就分別編譯編碼器塊和頭。當前程式碼把 x.shape 讀進 Python 整數、用顯式的批/長度值做 reshape、建立 arange(length) 掩碼，並用由形狀推匯出的行範圍索引標記。在不重新設計的情況下對這個完整圖應用 shapeless=True 是不安全的。把動態掩碼構造移到已編譯塊之外，並使用與形狀無關的展平/還原操作，可能使之後的無形狀變體成為可能；要在變化的 B、L 和標記數量上驗證它。

形狀桶可以限制重追蹤，但填充有計算成本。把 L=93 填充到 96 會增加約 3.2% 的逐 token 工作；填充到 128 會增加約 37.6%。比較精確形狀編譯與小的長度倍數，以及一組有界、來自工作負載的桶。把 (batch size, padded length, marker slots, dtype, device/model instance) 納入快取策略決策，並測量形狀頻繁變動下的冷編譯延遲和保留記憶體。

worker.py 裡獨立的 forward 基準直接呼叫 agent.model。如果只在 Agent.forward 里加入編譯，當前的 forward 基準會繞過它，而端到端基準會使用它。為了讓比較有意義，兩條路徑都必須顯式選擇同一個候選實現。在計時流程中保留 mx.eval 和 GPU 同步：只測量圖構造並不會測量推理。

按有用的 token 批處理，然後檢查矩陣排程

collate_items 把每個塊右側填充到它最長的序列；執行時按插入順序對問題分組。對異構流量，按準備後的長度排序或分桶，在問題數量上限之外再加一個 token 預算，並恢復原始的問題 ID 和輸出順序。只在能代表該服務的地方比較 1、2、4、8、16、32 和 64 的批。對線上請求，要把等待一個批的時間算進去；僅看離線的問題/秒會掩蓋不可接受的延遲。

當前的短 50 問題工作負載對 Laya 浪費約 8.9% 的填充 token，對 multilingual 浪費 5.8%。長基準行沒有填充浪費。因此，僅靠去填充或排序在這些固定樣本上的算術收益有限。要揭示生產收益，需要一種混合長度的分佈，比如許多短問題中夾一個長問題。移除填充必須保留逐樣本的 RoPE 位置、標記位置和注意力邊界；在沒有隔離掩碼的情況下把樣本拼進一個序列會改變模型。

對稠密 kernel，檢查實際的形狀和步長。QKV 已經是一次投影，編碼器 MLP 的兩個輸入分支已經共享一次投影。不加區分地拆分會增加 launch。只有在剖析器或 MLX 排程追蹤顯示有不理想的批 GEMM 時，才比較把連續的 [B,L,D] 輸入顯式展平成 [B*L,D]；框架可能已經高效地展平了。僅憑原始碼審查不能成為「漏掉了一個 GEMM 最佳化」的理由。

不要在生產實現裡每一層之後插入同步。當前的執行時每個塊只求值一次。額外的等待可能消除 CPU/GPU 的重疊並掩蓋排程改進；層級的剖析應當是一次單獨的診斷執行。

量化：瞄準主幹並實現其儲存契約

已安裝的 MLX 0.32.2 量化層實現 提供 nn.quantize(..., class_predicate=...) 和 QuantizedLinear，經由 mx.quantized_matmul 做僅權重的矩陣乘法。分組仿射量化支援 8-bit 和 4-bit 實驗。從組大小 64 的編碼器線性層開始，把啟用、歸一化、型別嵌入、評分器和動作頭保留為 FP16。然後獨立地加入決策頭線性層，並可選地加入嵌入量化。測量每個變體；在高 token 數量下，僅權重 kernel 可能不如 FP16 GEMM。

當前的載入器和模型裡有具體的整合隱患：

1. Agent.__init__ 把每一個儲存的權重轉換為浮點 dtype，並在嚴格載入之前只例項化稠密模組。量化 checkpoint 需要描述所選模組、組大小、位寬和模式的後設資料；在載入前例項化匹配的量化模組並保留打包的整數權重。把打包權重轉換成浮點並不是合法的載入。
2. 動作頭的第一個線性層輸入寬度是 D+4，即 1028 或 772，不能被 32、64 或 128 的仿射組大小整除。因此一刀切的量化呼叫不合適。轉換前檢查每一個所選層的輸入寬度。
3. DecisionModel.__call__ 從 self.act_head.layers[0].weight.dtype 選擇動作頭的輸入 dtype。對量化層來說，那個權重會是打包的整數儲存，而不是所期望的啟用 dtype。最初排除動作頭可以避開這條路徑；之後要支援它需要一個顯式的啟用 dtype 契約。
4. 量化會改變 logits 和校準機率。現有小型 FP16 固定樣本的一致性不足以證明 4-bit 質量。使用留出的帶標籤 choice、score 和 noul 任務、多語言輸入、接近的決策、不同選項數量和升級例子。跟蹤 argmax 一致、任務準確率、評分誤差、機率漂移、校準和動作機率。飽和的動作輸出可能掩蓋巨大的動作 logit 變化。

對組大小 64、FP16 仿射 scale 和 offset 的情況，近似的矩陣儲存是每個參數 bits/8 + 4/64 位元組：8-bit 時 1.0625 位元組，4-bit 時 0.5625 位元組，對比 FP16 的 2 位元組。這些是量化矩陣的儲存估計，不含其他張量和打包開銷；它們不是加速估計。官方 quantize API 描述了組整除性和格式。

更新的低位格式應當針對實際的 M3 Max 後端來評估，而不是假定它會使用來自更晚 Apple 晶片的硬體。MLX 的 NAX 可用性檢查要求比所記錄的 applegpu_g15s 裝置更新的架構代。見 MLX 0.32.2 裝置檢查。

在輸入確實相同時複用 CPU 準備

build_sequence 為每一個問題分別把同一個 state 序列化、清洗和分詞。它還分別對每條指令和每個選項分詞。Rust tokenizer 已經被直接使用；替換 Transformers 分詞不是一個尚未做的最佳化。

每次 prepare 呼叫把 state 只序列化並清洗一次、編碼一次，再把它的 token ID 切片到每個問題可用的空間。當同一個評分細則跨 state 使用時，快取不可變的已準備問題字首，鍵要包含 tokenizer 身份/revision、問題型別、有序 criteria、指令序列化、特殊 token 清洗和 token 預算。有界快取在 tokenizer 或配置變化後不得複用結果。批次的 tokenizer 編碼是另一個實驗，前提是它的輸出與當前那一串獨立編碼完全一致。

不要把對新拼接的提示分詞當作「拼接獨立編碼片段」的替代：子詞邊界可能改變。逐位元組校驗輸入 ID、注意力掩碼、標記位置、qtypes、截斷行為、結構化 criteria、mask 字面量、空輸入和輸出對映。

對長基準來說，state 在截斷前把一句話重複了 200 次。避免 N 次重複的 state 編碼可能有助於 CPU 準備，但現有的端到端/前向計時差並沒有測量這項節省。分別對 prepare、collate/陣列構造、forward 和後處理做基準，然後用一個不重複的留出語料確認端到端結果。

解碼器 KV 快取不適用於這個編碼器。 它的第一層是全域性雙向注意力；state token 表示取決於問題、選項及其位置。在不同的問題之間複用 state 隱藏狀態或 K/V 會改變結果。分詞和完全相同的整輸入結果可以快取；任意的上下文編碼器狀態不能。

精確裁剪最終決策頭的輸出

在最後一個 HeadLayer 之後，只有 [CLS] token 和選項標記 token 會被用到。更早的頭層仍必須產生所有 token，因為最後一層讀取它們的 K/V。只在最後一層中：

1. 歸一化所有輸入 token 並計算所有 K/V。
2. 在 [CLS] 和合法選項位置收集 Q，並把這些 query 跑在整個帶掩碼的 K/V 序列上。
3. 只對這些選定位置施加輸出投影、殘差、第二次歸一化和前饋網路。
4. 用選中的 [CLS] 輸出給動作頭，用選中的標記輸出做評分；保留標記填充和原始順序。

第一版實現可以保留融合的完整 QKV 投影，之後再收集 Q。更激進的變體把它的權重拆成一個全長 KV 投影和一個選定 token 的 Q 投影。那會節省更多算術，但可能讓 GEMM 排程效率更低。重複的填充標記索引只有在被掩碼的結果保持不可見時才無害。1 選項、多選項和可變標記的情形需要顯式的保真檢查。由於更少的 query 可能選中不同的 SDPA kernel，數學等價並不意味著浮點結果逐位相同。

設 R = 1 + number of option slots，保留完整 QKV 會從最後一個頭層移除約 18 * B * (L-R) * D^2 的稠密 FLOPs 和 4 * B * L * (L-R) * D 的注意力 FLOPs。拆分 Q/KV 會把稠密係數從 18 變成 20。相對於上面的整模型算術估計，取 R=5 得到：

這些是靜態的 FLOP 削減，不是預測的延遲削減。該技術從一個頭層移除了大部分工作，而不是從模型移除大部分工作。它有用，是因為它保留了依賴關係並能不改訓練地實現，而不是因為它承諾整模型速度的某個倍數。

只有在測量其貢獻之後，才構建真正的區域性注意力

所有編碼器注意力呼叫已經使用 mx.fast.scaled_dot_product_attention；RoPE 已經是 mx.fast.rope；nn.LayerNorm 呼叫快速歸一化原語。MLX 的注意力 API 接受布林掩碼並在 FP32 中做 softmax。當前的 head 維度是 64。MLX 0.32.2 Metal 排程 支援這種形狀配陣列掩碼，並且在推理時不會為它選擇未融合的兜底實現。沒有證據表明 Laya 的布林掩碼會停用融合注意力。 已安裝版本中可用的 force_fused=True 作為診斷斷言很有用，但不該在這裡被宣傳成一條新的快速路徑。

剩下的限制是結構化稀疏。該實現構造一個形狀為 [B,1,L,L] 的稠密區域性布林掩碼。在常規的 Metal 注意力 kernel 中，非因果迴圈遍歷完整的 KV tile 範圍；陣列掩碼在 QK 乘法之後作用於分數。它保留了區域性注意力語義，卻沒有利用區域性 tile 範圍。

一個精確的專用 kernel 可以把每個 query tile 限制在重疊的 K/V 視窗內、保持 FP32 softmax 累加，並避免一個稠密的 L×L 掩碼。正確的視窗是雙向且包含端點的：abs(query_position - key_position) <= 64。內部 query 能看到 129 個位置，儘管配置名是 local_attention=128。全注意力層和兩個決策頭層必須保持全域性。填充的 key 必須繼續被排除，未使用的填充 query 需要有定義的有限行為。

一個省力的原型可以把 query 塊與重疊的 K/V 切片分組，並用一個更小的精確掩碼呼叫現有的 SDPA。使用已經定位好的 RoPE Q/K，或顯式保留絕對偏移。優先使用成批的塊而不是大量 Python 呼叫，並考慮重複的 K/V 物化。這個原型在短長度下可能不如當前 kernel；在維護自定義 Metal 之前，它是一個正確性和盈虧平衡實驗。

移除所有被禁止的區域性注意力對所帶來的建模整模型 FLOPs 最大削減，對短輸入很小，對長輸入更有希望：

這些估計對 L>r 使用 local_pairs = L*(2*r+1) - r*(r+1)，取 r=64。它們包含所有稠密投影和兩個全注意力決策頭層。它們排除掩碼生成和記憶體流量。執行時收益可能高於或低於 FLOP 比例，因為注意力和 GEMM 的效率不同；只有剖析才能確定。在 8192 token 時權衡會不同，但隨附的 agent 把輸入限制在 512 或 1024，所以一個 8192 token 的說法需要一個單獨支援的工作負載。

編譯之外的融合

已安裝的精確 nn.gelu 已經用無形狀編譯裝飾，nn.Linear 已經在合適時使用帶偏置感知的 addmm。整塊編譯仍可能把 GELU 與它的門控乘法、殘差加法、型別轉換、掩碼和小型評分特徵操作融合。在實現等價的專用 kernel 之前，檢查已編譯的 kernel 圖。

如果啟用流量仍然顯著，原型化精確的 GELU-and-gate 或 residual-and-LayerNorm 融合。保留當前精確的基於 erf 的 GELU；tanh 或 sigmoid 近似會改變模型，需要單獨的質量測量。在添加布局轉換之前檢查實際的 Q/K/V 步長和複製：MLX 的全注意力實現接受一個連續 head 維度配其他步長，並寫出便於合併 head 的輸出佈局。無條件的連續複製可能增加工作。

標記 softmax、top-two 排序、熵、小型動作頭和 NumPy 結果格式化只有在實測過之後才是合理的後續目標。與數百次全寬 Transformer 操作相比，選項位置很少，所以先最佳化它們不太可能觸及佔主導的路徑。

在追求激進收益的同時保護基準的含義

當前的工作負載生成器 迴圈使用三個問題定義來構造 5、10 或 50 個問題。那些批裡至多隻有三個唯一的模型輸入。逐呼叫的精確去重可以在真實應用裡避免冗餘推理，但它會不成比例地改善這些固定樣本。把這項特性與 kernel 最佳化分開，並報告 questions、unique_forward_inputs、快取命中和實際求值的 token。用各自的標籤、有序 criteria 和校準後設資料重建每一個原始答案。保留一個含 50 個真正不同問題的套件和一個刻意重複輸入的套件。

不要在主推理基準裡使用跨呼叫結果快取：它會反覆呼叫完全相同的請求。任何快取實驗都要如實標註。tokenizer 快取基準應當同時包含重複評分細則的場景和全新輸入的場景。

對每個候選，使用這樣的實驗設計：

1. 保持目標不變。 記錄源 revision/雜湊、模型 revision、dtype、編譯標誌、所選量化模組、token ID 或其雜湊、形狀、標記數量、批策略、預熱、同步和裝置。現有的 input_sha256 雜湊的是 state/questions，而不是實際的 token 張量；它不能確立跨 tokenizer 的張量同一性。從同一個最終源 revision 重跑基線，因為歷史 JSON 的源雜湊不同。
2. 分離工作負載類別。 受控的 kernel 比較用固定形狀；排程和編譯用真實的可變長度；吞吐用不同的問題；合法的 CPU 快取用重複的評分細則；去重用刻意重複的問題。包含長長度尾部區間和不同的選項數量。在每個後端比較內部保持源文本和截斷一致。
3. 測量冷熱行為。 分別記錄模型載入和首次編譯。測量準備、圖構造、同步求值的前向、輸出轉換和端到端延遲，不要把獨立的中位數當作可加的。在單獨一次執行裡剖析 kernel，因為追蹤會擾動延遲。
4. 一次只改一個最佳化。 用現有的迭代次數篩選，然後在交替的基線/候選塊裡重複入圍者，並收集足夠樣本以獲得可信的 p95，例如每個工作負載至少 200 個計時請求。只跑一個活動的 GPU 基準，保持一致的功耗/散熱條件，並保留原始樣本。要求改進大於實測的執行波動。
5. 檢查正確性和穩定性。 在相同 dtype 下與 MLX 以及現有 FP32 參考比較；對排程和 CPU 準備的改動強制執行 token 同一性。測試 64、128 附近以及桶邊界處的長度；塊上限附近的批；單選項和多選項；多語言輸入；填充；以及編譯後的形狀變化。重複變化形狀的請求，在預熱後同時觀察活動記憶體和快取記憶體，並驗證每種配置內的結果有限且確定。
6. 對近似改動施加更強的門檻。 量化、啟用近似、token 裁剪、提前退出和蒸餾需要小型迴歸固定樣本之外的留出任務與校準結果。改變訓練行為時要保留單獨的模型身份和基準標籤。更快的 multilingual checkpoint 或蒸餾模型是另一種模型，而不是同一個 Laya checkpoint 的加速。

對任何佔據端到端時間比例 f 並以因子 s 加速的實測熱點，用 Amdahl 上界 1 / (1 - f + f/s) 評估預期的總影響。f 要用實測的時間比例；上面的算術比例不能替代它。下一個具體的工程決策應當跟隨編譯/批處理消融和短對長的 kernel 剖析，而不是一個未經驗證的倍數。

文件與可復現性說明

文件查詢使用了規定的 Context7 工作流：一次 library MLX 解析，然後分別對編譯和快速注意力做官方文件查詢，使用 /websites/ml-explore_github_io_mlx_build_html（共三條命令）。檢查了已安裝的 .pyi 和 Python 原始碼以核實 MLX 0.32.2 的行為，包括 force_fused、量化 API、已編譯的 GELU 和快速 LayerNorm。閱讀了鎖定版本的上下游 C++/Metal 原始碼以瞭解排程和 tile 迴圈細節。本次研究沒有升級任何庫。

用於詳細觀察的兩個 FP16 基線檔案在檢查時的 SHA-256 摘要如下：

laya-mlx-float16.json
63146dd664d039dde1a728b17aad896e491bd01ea36ea0786953691180e55b09

laya-multilingual-mlx-float16.json
9af74bd5a11e4edc15e6a8c9dc929a7b9fd2d19cb06f348cd0e04a076f912473

在本研究進行期間，主基準程序仍在產生額外的工件。表格刻意引用審閱時已經可用的完整基線檔案，並且不對未測量的候選實現作任何聲稱。