優化 Apple Silicon 裝置端推論
可改善預填與解碼吞吐量的自訂本機引擎。
Apple Silicon 上的混合運算協調雲端中的前沿智慧與 Mac 上的本機模型之間的任務。雲端模型處理研究與推理,而本機模型則處理 Mac 上的私密檔案與應用程式。
為了讓這種分工感覺順暢,本機推論必須跟上其餘任務的步伐。這需要一個能夠快速處理提示並維持高 Token 生成率的引擎。
Lily 係我哋專門為 Apple 晶片同 Qwen3.6-35B-A3B 打造嘅輕量級本機推論引擎,針對預先填充(prefill)同解碼(decode)作出咗獨立優化。獨立示範版本已經喺 GitHub 公開。
簡介
在 Mac 上執行 LLM 的常見方式是使用 MLX,這是 Apple 針對 Apple Silicon 推出的開源機器學習框架。其附屬函式庫 MLX-LM 新增了載入與使用各種語言模型生成文字所需的元件。MLX 與 MLX-LM 結合,為本機 LLM 推論提供了一個現成的通用堆疊。
Qwen3.6-35B-A3B 是一個稀疏混合模型:它使用混合專家(MoE)路由,並將固定大小的遞迴狀態與完整注意力機制相結合。這些架構選擇減少了所需的運算量,但也造成了不規則的工作負載。Token 會路由至不同的專家權重,且遞迴狀態本質上具有循序性。
MLX-LM 已經會針對推論階段和常見的工作負載形狀選擇最佳化核心,但其可重複使用的運算必須支援許多模型架構。專為 Qwen 設計的引擎能夠在模型與執行階段層級進行專門化,圍繞模型的固定結構來協調核心、資料移動與排程。
Lily 在單一程序中端對端地實作了此專門化。Rust 執行階段會載入模型檢查點並管理工作階段狀態與生成迴圈,相容於 OpenAI 的聊天完成 API 接受請求並串流傳輸 Token,而自訂 Metal 核心則執行 Qwen 專屬的運算。PyTorch 與 MLX 皆不在執行路徑中。
我們分開測量預填與解碼效能。預填吞吐量捕捉引擎處理提示的速度;解碼吞吐量則捕捉引擎生成輸出 Token 的速度。
我們在一部配備 M5 Max、40 核心 GPU 與 128 GB 統一記憶體的單一 MacBook Pro 上對 Qwen3.6-35B-A3B 進行基準測試。在預填寫與解碼的十個提示長度和十個上下文長度中,從 256 到 128K token(K = 1,024),該引擎的預填寫吞吐量平均為 MLX-LM 的 1.22 倍,解碼吞吐量平均為其 1.32 倍。當提示為 4K token、解碼上下文為 4K token 時,自訂引擎的預填寫速度達到每秒 5,749.9 個 token,解碼速度達到每秒 186.6 個 token;相較之下,MLX-LM 分別為 4,737.5 和 140.9。於多輪工作階段中,這些時間節省會隨著每次額外的模型呼叫而累積。
接下來,我們會解釋 Qwen 的架構如何在 Apple Silicon 上創造模型專屬的最佳化機會。接著,我們將逐步解說由此產生的預填與解碼變更。我們也會探討額外最佳化在何處會停止帶來效益,最後以與 MLX-LM 的端對端比較作為總結。
Apple Silicon 上的 Qwen 專屬最佳化機會
Qwen 建立了三種不同的工作負載形狀
Qwen3.6-35B-A3B 包含 35 百億個參數,但每個 Token 僅會啟動大約 3 百億個。路由器會對 256 個專家子網路進行評分並選擇 8 個,同時搭配一個處理每個 Token 的共用專家。這種稀疏 MoE 設計減少了運算量,但產生了不均勻的工作負載:專家接收到不同數量的 Token,且每個 Token 都需要來自不同專家組合的權重。
Qwen 同時結合了 10 個完整注意力層與 30 個 Gated DeltaNet 層。這兩種層類型以不同的方式保留先前的資訊。
注意力層使用群組查詢注意力(GQA)。Qwen 擁有 16 個查詢頭和 2 個鍵值(KV)頭,每 8 個查詢頭共用一個 KV 頭。共用可使 KV 快取更小,並允許在查詢頭之間重複使用快取資料。快取仍然會為每個 Token 儲存新的鍵與值,因此隨著內容增長,每個解碼步驟會讀取更多資料。
相對地,Gated DeltaNet 會將先前的資訊壓縮成固定大小的遞迴狀態。學習到的閘門會控制要保留多少現有狀態,而增量更新則會納入來自目前 Token 的資訊。模型以遞迴方式定義了這些更新,因此每個 Token 都取決於前一個 Token 所產生的狀態。然而,在預填期間,引擎可以用兩種方式評估相同的運算。它可以直接掃描 Token 同時向前傳遞狀態,或者將更新重新組織成公開更多矩陣運算與 Token 層級平行處理的區塊。哪種方法較快取決於模型維度、工作負載與硬體。
這些結構共同造就了三種運算模式:不均勻的專家群組、針對不斷增長的快取的注意力機制,以及可以直接評估或以區塊評估的固定大小遞迴。
Apple Silicon 針對不同的工作負載提供不同的路徑
預填會一次處理許多提示 Token 啟動列。此處考量的本機工作負載通常一次解碼一個請求(批次 1),並且每個步驟處理一個新列。這種差異改變了相同模型權重的使用方式。預填可以在數百或數千個列之間重複使用每個權重區塊。解碼則大抵無法如此,因為每個新 Token 都需要再次通過權重。
Apple Silicon 將 CPU 和 GPU 置於統一記憶體(兩者皆可存取的單一實體記憶體池)背後。這允許模型保持常駐狀態而無需維護獨立的 GPU 複本,但它並不會使資料移動變得免費。讀取權重和中間值仍然會消耗記憶體頻寬,而暫存器和其他晶片上儲存空間雖較快但小得多。
M5 GPU 也提供不同的運算路徑。預填的線性層使用一般矩陣乘法(GEMM),將權重矩陣同時套用到許多列。相容的 GEMM 可以透過 Metal 4 張量運算使用每個 GPU 核心中的神經加速器。相對地,批次 1 解碼使用一般矩陣向量乘法(GEMV),將相同的權重套用到單一列。由於權重重複使用率低,GEMV 主要受限於記憶體頻寬,且比起專為具有較高資料重複使用率的矩陣運算所設計的神經加速器,它更適合 GPU 的向量算術邏輯單元(ALU)。
這些執行路徑並非 Lily 獨有。MLX 在相同的統一記憶體上運作,並根據工作負載形狀選擇最佳化的矩陣與向量核心。MLX-LM 的 Qwen 實作已經將專家工作分組、使用融合的遞迴 Metal 核心來評估 Gated DeltaNet,並使用感知 GQA 的注意力機制。這些功能是在 Apple Silicon 上進行高效 Qwen 推論的共同起點。
最佳化策略
Lily 較窄的範疇讓它能夠圍繞 Qwen 精確的架構與維度來協調這些共用的執行路徑。它使用階段專屬的 GPU 路徑,對應 Qwen 的專家、遞迴與注意力工作負載以將資料移動降至最低,並從測量到的工作負載形狀中選擇核心與配置。此策略包含三個部分:
- 使 GPU 路徑符合推論階段。當預填可以在許多列之間重複使用權重時,使用矩陣導向的執行方式;當批次 1 解碼一次處理一列時,使用向量導向的執行方式。
- 將 Qwen 的結構對應至 GPU 上,同時將資料移動降至最低。在權重使用前保持壓縮狀態,組織路由的專家工作而不需返回 CPU,透過其遞迴掃描將 Gated DeltaNet 狀態保留在晶片上,並重複使用由群組查詢注意力共用的 KV 資料。
- 使核心適應工作負載形狀。在每個階段中,從可用的列數、列在專家間的分佈、運算的維度以及目前的內容長度中,選擇磚塊大小、執行配置與注意力路徑。
以下章節將說明這些選擇。對於在 M5 Max 上的配對消融實驗中所評估的最佳化,我們透過比較除了正在研究的最佳化之外、其他方面完全相同的引擎設定來估算其影響。由於這些實驗是將我們引擎的不同版本進行自我比較,因此它們解釋的是機制,而非將最終結果相對於 MLX-LM 進行分解。
預填:重複使用權重並將路由保留在 GPU 上
預填一次公開了許多 Token 列,但 Qwen 會將這些列不均勻地路由至各個專家,並透過序列更新遞迴狀態。其最佳化分為三組:圍繞路由列組織稀疏專家工作、將 Gated DeltaNet 掃描保留在晶片上,以及將長提示分割成受限的區塊。
最佳化稀疏專家運算
在矩陣乘法期間進行反量化權重
Qwen3.6-35B-A3B 檢查點使用群組式仿射 4 位元量化。每個權重皆儲存為 4 位元整數代碼,而每組 64 個權重則共用一個用於重建其值的 bfloat16 比例與偏差。這將這個擁有 35 百億參數的模型從大約 70 GB 的 bfloat16 權重縮減至 19.4 GB 的檢查點,使得將模型常駐於 Mac 上變得切實可行。
用於矩陣乘法的 Metal 4 張量運算消耗的是 bfloat16運算元,而非打包的 4 位元表示法。在乘法之前,GPU 必須以 bfloat16 重建權重。Lily 中經過最佳化的群組 GEMM 一次執行一個小型權重磚塊的此種轉換,並將結果保留在晶片上的 threadgroup 記憶體中,其時間僅夠將其乘以路由的啟動列。累積運算使用 32 位元浮點數,輸出則以 bfloat16 寫入。完整的展開權重陣列絕對不會在統一記憶體中建立。
在消融實驗中,反量化是作為獨立運算執行的:它將 4 位元權重展開至統一記憶體的 bfloat16 陣列中,在此之後矩陣核心會將該陣列讀回。在 512 Token 的提示下,將反量化移至群組 GEMM 內部透過消除此中間寫入與讀取,使端對端預填吞吐量提升了 77.4%。
將專家路由保留在 GPU 上
群組 GEMM 需要將指派給每個專家的啟動列儲存在一起。在每個 Token 選擇 8 個專家後,直方圖會計算有多少指派分配給每個專家。前綴掃描會將這些計數轉化為起始偏移量,散布(Scatter)步驟會將列置入其專家群組中,而區塊對應表則會列出群組 GEMM 必須處理的固定大小矩陣區塊。
經過最佳化的路徑會將每個提示區塊的整個序列保留在單一命令緩衝區(一組有序的 GPU 運算批次)中。相對地,消融實驗會暫停,以便 CPU 檢查路由中間值並提交下一個運算。將直方圖與前綴掃描保留在 GPU 上會增加兩個核心,但消除了每個 MoE 層內部所需的 CPU 與 GPU 同步。
在 512 Token 的提示下,啟用 GPU 常駐路由使端對端預填提升了 89%。這也顯示了為什麼單純的核心數量可能會具有誤導性:較快的路徑會啟動更多核心,但從不在層內部等待 CPU。
使磚塊大小符合專家負載
在 2K Token 的提示下,將每個 Token 路由至 256 個專家中的 8 個會產生 16,384 個 Token 與專家的配對,平均每個專家有 64 個啟動列。實際分佈並不均勻:有些專家接收了許多列,而其他專家則接收很少。
群組 GEMM 將每個專家的輸出劃分為磚塊,這些是矩陣乘法輸出的小型矩陣長方形區塊。每個磚塊皆被指派給一個 GPU threadgroup。在 Apple Silicon 的 GPU 上,一個 threadgroup 包含一或多個 simdgroups,每個 simdgroups 由 32 個以鎖步(Lockstep)執行指令的執行緒組成。
較大的磚塊(Tile)可將設定成本分散到更多列中,並公開更多平行工作,但當專家只接收到幾列時,大型磚塊的一部分仍會保持閒置。因此,磚塊大小與 simdgroup 數量是互相搭配的。
消融實驗將磚塊固定為 16 列。相較於該控制組,啟用包含四個 simdgroups 的 32 列磚塊使端對端預填在 2K Token 下提升了 13.2%。
將遞迴狀態保留在晶片上
在預填期間,每個 Gated DeltaNet 層會依序掃描提示,同時向前傳遞其遞迴狀態。若停用暫存器常駐,消融實驗將會使用區塊式掃描。在 2K Token 的提示下,該路徑每層會移動 256 MiB(百萬位元組)的狀態,並且會在屏障(所有參與執行緒都必須彼此等待的同步點)處重複停止合作執行緒。
遞迴狀態是一個矩陣。經過最佳化的核心會將每一欄指派給一個 simdgroup。simdgroup 會在其執行緒之間分割該欄、將該欄載入其暫存器一次,並在整個掃描過程中持續傳遞狀態。執行緒透過 simdgroup 運算來交換中間結果,而不是透過在 threadgroup 之間共用的晶片上儲存空間(即 threadgroup 記憶體)。完成的狀態僅會在掃描結束後才寫回。
狀態及其閘門(Gate)使用 32 位元浮點格式,因為微小的捨入誤差會在循序更新中不斷累積。查詢與鍵的啟動值則維持在 bfloat16。
在 2K Token 的提示下,啟用暫存器常駐掃描使端對端預填效能提升了 5.6%。專家 GEMM 佔了預填時間的大約 90%。循序掃描無法公開足夠的可重複使用矩陣工作,因此無法受惠於神經加速器(Neural Accelerators)。
透過提示分塊來限制暫存記憶體
執行階段將長提示處理為一系列受限的區塊,而不是將每個提示 Token 的暫存資料全部同時保留在記憶體中。模型權重常駐於統一記憶體中,而遞迴狀態與 KV 快取則將內容從一個區塊傳遞到下一個區塊。不會捨棄任何先前的內容。
若無分塊機制,暫時啟動陣列會隨著完整提示而增長,並與模型權重、遞迴狀態以及 KV 快取競爭統一記憶體。分塊機制一次只讓一個區段的暫存值保持活躍,然後在處理下一個區段之前釋放或重複使用該儲存空間。這限制了尖峰工作記憶體,並允許引擎在不改變模型輸出的情況下處理更長的提示。
分塊預填在許多引擎中很常見,對於在這些記憶體受限環境中服務長多輪軌跡至關重要。注意力層的總預填時間相對於提示長度維持二次方關係,並且由於先前區塊的重複 KV 載入而帶有一些額外負荷。
解碼:將每個 Token 移動的位元組降至最低
批次 1 解碼一次處理一個新列。由於權重重複使用率低,其吞吐量主要取決於引擎為每個 Token 移動了多少位元組。解碼變更分為四組:最佳化單列權重路徑、將每個步驟保留在 GPU 上、減少中間與狀態流量,以及有效率地讀取注意力快取。
最佳化單列權重路徑
MLX 已經將單列工作分派給專門的矩陣向量核心。由於 Lily 不使用 MLX,因此自訂執行階段必須提供相同的基本策略。我們的列平行 GEMV 是專為單一啟動列所設計。simdgroup 會在輸出上進行合作,同時平行讀取權重矩陣的不同部分。
將每個解碼步驟保留在 GPU 上
將 Token 交接保留在 GPU 上
每個解碼步驟皆以選擇下一個 Token 作為結尾;下一個步驟則以該 Token 作為輸入開始。將選擇傳送至 CPU 然後再傳回 GPU 會為每個 Token 增加一個同步點。相對地,我們的執行階段會在兩個命令緩衝區與兩個 GPU 常駐 Token 槽之間交替。GPU 會選擇最高得分的 Token,並將其 Token ID 直接寫入下一個解碼步驟的輸入槽中,同時 CPU 則準備後續的工作。
重疊執行獨立的 GPU 工作
在一個記錄的批次 1 解碼步驟中,生成一個 Token 啟動了 795 個 GPU 核心。它們的相依性形成了 555 個循序階段,使某些核心能夠同時執行。儘管如此,Metal 的循序執行模式仍依序執行了每一個核心。
經過最佳化的解碼路徑會在並行的 Metal 通道中記錄實際的資料相依性。當 GPU 資源允許時,獨立的核心啟動可以同時執行。只有在後續工作需要先前結果時才會插入屏障。
減少中間與狀態流量
獨立的核心通常會具現化一個中間結果:一個核心將暫時結果寫入記憶體,下一個核心則將該結果讀回。經過最佳化的解碼路徑融合了四個鏈結:兩個帶有其閘控啟動的專家輸入投影;帶有其路由分數與共用專家結果的專家輸出投影;注意力之前的查詢與鍵準備;以及帶有其正規化的遞迴更新。每個融合的核心都會將暫存值保留在暫存器中,而不是透過記憶體傳送它們。
融合也縮短了相依性圖表:當中間寫入消失時,保護其消費者的屏障(Barrier)也會隨之消失。
有效率地讀取注意力快取
合併注意力快取讀取
注意力機制會在每個解碼步驟期間從 KV 快取中讀取鍵與值。在消融實驗中,相鄰的 GPU 執行緒並不總是請求相鄰的位元組,這迫使記憶體系統服務更多獨立的交易。啟用合併載入會使得相鄰執行緒請求相鄰的位元組,以便硬體可以組合成批讀取其讀取要求。
在 bfloat16 設定下,合併(Coalescing)將鍵(Key)頻寬從 33.8 提升至 47.9 GB/s,將值(Value)頻寬從 42.0 提升至 61.8 GB/s,並在 3,840 Token 的內容長度下將端對端解碼效能提升了 2.1%。
打包查詢頭以重複使用 KV 列
群組查詢注意力機制讓八個查詢頭共用一個 KV 頭。在消融實驗中,每個查詢頭都在獨立的 simdgroup 中執行,因此所有八個查詢頭都會各自獨立請求相同的快取 KV 列。經過最佳化的核心將四個查詢頭打包到一個 threadgroup 中,該 threadgroup 會載入每個 KV 列一次,並在四個注意力計算中共用它。第二個 threadgroup 則負責處理其餘的四個頭。
此技術通常稱為 GQA 打包,它執行相同的算術並產生相同的輸出位元組,同時將八個獨立的 KV 請求減少為兩個共用載入。相較於未打包的消融實驗,它在 32K Token 的內容長度下將端對端解碼吞吐量提升了 23.8%。
在長內容下切換注意力配置
完整注意力層中的每一個解碼步驟都會掃描現有的 KV 快取。固定區塊配置會將該快取分割成相等的片段,供 GPU 進行平行處理。當快取較小時,其額外的排程並不划算,但隨著內容增長,固定區塊配置能更平均地平衡工作負載。
針對此模型,執行階段在 32K Token 以下時會保持一般注意力路徑,並在 32K 或以上時使用固定區塊路徑。當每個頭擁有 256 個值且八個查詢頭共用一個 KV 頭時,就會套用此切換;其他形狀則維持在一般路徑上。消融實驗會停用此切換並一律使用一般路徑。啟用固定區塊路由使端對端解碼在 32K 時提升了 7.7%,在 64K 時提升了 27.4%,在 128K 時則提升了 40.2%。
進一步最佳化的限制
某些變更改善了獨立的運算,但並未改善端對端推論。
推測式解碼(Speculative decoding)使用較小的模型來提出 Token 供完整模型驗證,這使批次 1 解碼慢了 18%。驗證處理了兩到五列的群組,對於這項硬體而言這是不具效率的形狀,且這些列通常會選擇不同的專家,從而增加了所讀取的專家權重資料量。減少草稿器(Drafter)的輸出詞彙表使草稿器吞吐量提升了 4.7–5.1%,但並未使完整的推測迴圈更快。此結果具有工作負載專屬性:我們在 Blackwell 上的批次化 Qwen 部署會在不同的條件下使用推測式解碼。
其他實驗包括減少 GPU 啟動次數、重疊執行整個階段、使用更大的預填磚塊、套用更廣泛的融合、加速路由器,以及將輸出投影與 Token 選擇相結合。沒有任何一項實驗改善了完整的推論迴圈。
硬體限制的測量也顯示,主要的預填與解碼運算幾乎沒有剩餘的運作空間。MoE GEMM 和 GEMV 達到了其存取模式下最快持續權重讀取速率的 97.9% 和 90.3%。移除稀疏 GEMV 中的算術運算僅使吞吐量改變了 0.2%,這證實了權重讀取而非運算才是受限資源。預填的矩陣乘法在獨立情況下同樣達到了理論矩陣極限的 93%,而在測試模型內部則達到 80–86%。
端對端效能
端對端比較會在兩個引擎中載入相同的 4 位元檢查點位元組,並在一台具備 40 核心、128 GB 記憶體的 M5 Max 上一次執行一個請求。在每一輪中,兩個引擎會以交替順序執行,以減少背景負載與晶片溫度變化所帶來的偏差。我們是與 MLX-LM 最快的直接生成路徑進行比較,而非其伺服器,因此測量重點在於模型執行而非服務額外負擔。
此掃描涵蓋從 256 到 128K 個 Token 的十個預填提示長度和十個解碼內容長度。隨著引擎將固定的設定成本分散到更多 Token 上,預填吞吐量最初會上升。預填會在大約 4K Token 的提示附近達到峰值,然後下降,因為隨著提示增長,十個完整注意力層會執行更多工作。解碼在短內容時維持幾乎平坦,一旦讀取不斷增長的 KV 快取變得顯著時便會下降。自訂引擎在每一個記錄的長度上都更快。
由於專門化的執行可能會改變浮點運算的順序,我們也針對 MLX-LM 檢查了數值一致性。在教師強制(Teacher-forced)比較中,兩個引擎在 192 個位置中的每一個位置皆從相同的參考前綴預測下一個 Token,以防止先前的差異影響後續的輸入。Lily 的困惑度(Perplexity)僅高出 0.04%,並且在 96.35% 的受測位置中選擇了相同的最高排名 Token。
為本機平台打造
Apple Silicon 並非較小的資料中心 GPU。它是一個完整的本機推論平台,擁有自身的硬體與軟體特性。統一記憶體賦予單一節點極高的上限,使其能容納大量的模型與狀態。M5 神經加速器負責處理預填中的密集矩陣工作。向量 ALU 則處理解碼中受限於頻寬且低重複使用的剩餘部分。
Qwen 為專門化提供了進一步的機會:將專家路由與遞迴狀態保留在 GPU 上、消除不必要的中斷、重疊執行獨立工作、重複使用共用的 KV 資料,以及使核心適應工作負載形狀。
透過模型與平台專屬的最佳化,單一台 Mac 就能高效執行大型稀疏模型。未來的工作將擴大涵蓋各種模型、晶片與服務工作負載,並將此處在單一設定上驗證的機制轉化為更通用的執行階段策略。
更廣泛的原則是讓引擎同時配合模型的架構以及硬體特定的運算與記憶體路徑。隨著前沿開源權重模型與硬體的演進,高效能本機推論將日益取決於專門針對兩者量身打造的引擎,而非抽象化其差異的引擎。
引用
@online{perplexity2026OptimizingOnDeviceInferenceForAppleSilicon,
author = {{Perplexity} and Yibo Wu and Lequn Chen and Kevin Hu and Phi Hoang and Claire Beamer and Jeremy Yang and Denis Yarats},
title = {{Optimizing on-device inference for Apple Silicon}},
date = {2026-09-01},
url = {https://www.perplexity.ai/hub/blog/optimizing-on-device-inference-for-apple-silicon}
}