适用于 Apple Metal 的 Automatic1111，sd1.5 速度提升 40%

# 教 Automatic1111 在 M1 上說 Metal 的語言 来源：https://therad.ninja/from-8-10-seconds-to-3-7-teaching-automatic1111-to-speak-metal-on-an-m3-pro/ 我在 Apple 硬體上很常用 Draw Things，但 Automatic1111 一直有一件事讓我在意：*它感覺比應有的速度還慢。* 不是慢到不能用，就只是慢到你會注意到。 在我的 M3 Pro 上，Automatic1111 跑一次短短的五步 DPM++ SDE 生成，通常大約要 8–10 秒。Draw Things 早已讓我看見 Stable Diffusion 在 Apple Silicon 上可以比這即時得多。 所以我想知道，這中間的落差到底有多少是必要的。 GitHub - dmikey/stable-diffusion-webui-metal：為 Apple Silicon 微調的 automatic1111。為 Apple Silicon 微調的 automatic1111。前往 GitHub 建立帳號以參與 dmikey/stable-diffusion-webui-metal 的開發。dmikey (https://github.com/dmikey/stable-diffusion-webui-metal) 有一個重要的限制：*我不想取代 Automatic1111。* 我想要同樣的 WebUI、checkpoints、LoRAs、samplers、extensions、API、提示詞語法，以及整體工作流程。我沒有興趣把一切都轉成 Core ML，然後在它外面再建一個推論引擎。目標比那窄得多： **如果我們讓關鍵的部分表現得更像原生 Apple 軟體，Automatic1111 可以快到什麼程度？** 答案，至少就我正在跑的工作負載而言，是快了不少。 同一類生成在我的 M3 Pro 上原本要花大概 8–10 秒，現在通常落在 3 到 7 秒之間。而 13–20，**如今在我的 M1 Mac Mini 上落在 8–10。** 這些是我目前工作負載的觀察範圍，不是宣稱普遍有 2 倍提升的控制基準。另外，runtime 的改進和 NGMS 之間有一個重要區別：NGMS 實際上是減少了要做的引導（guidance）計算量。 儘管如此，實際使用上的差異是很大的。 不過，比最終數字更有趣的，是達到這個結果需要做什麼。 不是單一的一項優化。 ## 從工作量開始，而不是從基準開始 我在意的工作量相當特定： - Stable Diffusion 1.x - DPM++ SDE - Karras - 5 步 - CFG 約 1.15 - 384×640 和 512×512 - FP16 UNet 跑在 MPS - 預設 FP32 VAE 這種特定性很重要。 早期，DPM++ 2M 看起來像是省時間的好方法。它比較快，但在短排程上沒有產生我要的結果。 那不是優化，那是不同的工作量。 這基本上成了之後所有事情的規則：如果某個優化單獨看起來很棒，但並不能讓實際生成變快，同時又保留我想產生的結果，那就不算數。 注意力機制（Attention）是明顯的起點。 PyTorch 的 MPS 後端已經進步很多，但仍有一些 Stable Diffusion 的 attention shape，直接用 Metal 是合理的。 錯誤的作法，會是把自訂 Metal 實作視為「全部都比較快」。 它不是。 所以，我新增了一條 Metal Flash Attention 路徑，專門針對 SD 1.x 中那些在測試中真的勝出的 shape。 路由器的概念大概像這樣： `` if inference and fp16_mps and query_tokens >= 192 and head_dim in (40, 80, 160): return metal_flash_attention(q, k, v) return pytorch_sdpa(q, k, v) `` 還有一些額外檢查是關於 mask、training、dropout、tensor layout、grouped-query attention、支援的型別等等，但基本概念就是這樣。 Metal 不是因為聽起來比較快就當預設。它是被量測過、而且那個 shape 值得用 Metal，才拿到這個運算。 其他一切都回到 PyTorch。 這個 fallback 很重要。Automatic1111 支援的設定遠多於我的五步 SD 1.x 工作流程。我不要一個「只要有人動過設定就壞掉」的加速 fork。 ## Kernel 不是全部的問題 把 attention 放進 Metal 有幫助，但它揭露了更有趣的事。 原生 extension 在每次 attention 呼叫後都提交 MPS command buffer。 Stable Diffusion 在每次 UNet 評估內會一再呼叫 attention。以短短五步的生成而言，重複提交小塊工作開始變成總執行時間中相當可觀的一部分。 所以，我沒有把 Metal kernel 當成一個獨立的小應用程式，而是把它整合進 PyTorch 目前的 MPS stream。 extension 會結束 PyTorch 目前的 kernel 合併（coalescing），把 Metal Flash Attention 運算編碼進目前的 command buffer，然後讓 PyTorch MPS 後續的工作繼續從那裡接下去。 每次 attention 後的那個顯式 commit 消失了。 這最終成為整個專案中最重要的教訓之一。 **如果你每一次呼叫後都提交 command buffer，就算 kernel 最快也還是會輸。** 在這種生成時間尺度下，overhead 很要緊。你不只是在優化 GPU 可以多快做矩陣乘法；你是在優化 Python、PyTorch、MPSGraph 和 Metal 之間需要互相協調的頻率。 還有一個提醒——用一種很明顯的方式提醒你不要相信計時器：早期某個版本會產生一張綠色的圖。 它很快。 它也是綠色的。 現在 Metal 路徑在 WebUI 啟用它之前，會先跑一個隔離的 attention-plus-projection 正確性測試。 ## 統一記憶體改變了遊戲規則 下一個問題是記憶體。 Apple Silicon 沒有一塊獨立的 VRAM 放在系統 RAM 旁邊。GPU 和機器其他部分競爭的是同一塊實體記憶體。 這讓一些傳統 GPU 假設變得很糟。 一個 attention matrix 在技術上可以放進記憶體，但如果 macOS 有記憶體壓力、allocator 開始 thrashing、或機器開始 swap，它仍然是個糟糕的主意。 所以，這個 fork 不是用固定的 VRAM 門檻，而是同時根據總記憶體和目前可用記憶體來估計原生 attention 的成本。 概念上： `` attention_bytes = batch × heads × query_tokens × key_tokens × element_size estimated_peak = attention_bytes × 2.5 budget = min( 10% of total memory, 20% of currently available memory, 1.5 GiB ) `` 如果估計的峰值落在這個預算內，就可以跑原生 SDPA。 如果沒有，請求就會改走 memory-bounded 的 sub-quadratic 路徑。 那個 fallback 的 chunk size 也是動態的。8 GB Mac 不該跟 32 GB Mac 做一樣的決定，而且兩者都不該假裝 Chrome、Xcode 或其他正在跑的東西不存在。 我不把這算成全面的速度提升。這主要是讓效能可預測，並避免那些「表面上看起來很快的運算」造成足夠的記憶體壓力、讓整個生成反而變慢的情況。 ## 不要再保留每一個 attention chunk 我也改了 sub-quadratic fallback 處理 K/V chunk 的方式。 原本的做法是計算部分 attention 結果，把每個 chunk 的分子、正規化權重、最大值都留著，最後全部堆疊起來。 這沒有必要。 取而代之，這個 fork 維持一個 running maximum、正規化總和，以及加權輸出。每個新的 K/V chunk 被合併進這個 running state 之後就可以丟棄。 遞迴基本上長這樣： `` new_max = max(running_max, chunk_max) running_scale = exp(running_max - new_max) chunk_scale = exp(chunk_max - new_max) running_values = running_values × running_scale + chunk_values × chunk_scale running_weights = running_weights × running_scale + chunk_weights × chunk_scale `` 這和讓 Flash Attention 節省記憶體的是同一個 online-softmax 概念。 記憶體現在是依目前 chunk 的規模來擴展，而不是一路累積所有部分結果到最後。 我用 PyTorch SDPA 比對過 forward 結果，也用 float64 測過梯度。再說一次，目標不只是做出聰明的東西，它必須是一個安全的 fallback。 ## 有些 MPS 的 workaround 已經比 bug 活得還久 還有一類優化沒那麼光鮮：刪除舊的 workaround。 Apple 的 PyTorch 後端變了很多。 Automatic1111 為了較舊的 MPS 實作累積了許多防禦性行為，包括 clone `torch.narrow()` 的結果，以及把 LayerNorm 推到 FP32。 當時底層 MPS bug 存在時，那些修正有意義。在比較新的 PyTorch 版本上，它們只是變成複製、配置、轉換和記憶體流量。 所以這些行為現在改由 runtime 版本來決定是否啟用，而不是不分青紅皂白一律套用。 如果有人需要舊行為，仍然有 `A1111_MPS_FORCE_LEGACY_OPS=1` 這個逃生門。 我也啟用了 `PYTORCH_MPS_PREFER_METAL=1`，因為直接使用 Metal 矩陣乘法在我鎖定的 SD 1.x projection 尺寸上測試結果較好；同時也移除了預設的 sampling upcast，讓更多短 sampling 路徑維持在 FP16。 最後這個改動是真正的取捨。FP16 的 reduction 順序，以及移除 upcast，會影響同一個 seed 的輸出。 對這個工作流程來說，我可以接受。但這不應該被當成免費的效能。 ## 融合 GroupNorm 和 SiLU 在減少不必要的運算之後，我去找那些既必要又不斷重複的運算。 GroupNorm 接 SiLU 在 SD 1.x UNet 裡無所不在。 正常來說，它們是兩個獨立的 PyTorch 運算。那代表各自 dispatch，而且中間的 activation 會被寫出去、然後立刻讀回來。 所以，我寫了一個融合的 Metal kernel。 對於相容的 FP16 推論 tensor，每個 batch/group pair 由一個 256-thread 的 Metal threadgroup 處理。kernel 用 FP32 累加 sum 和 squared sum，把它們歸約成 mean 和 variance，套用正規化及 affine 參數，套用 SiLU，然後寫出 FP16 結果。 一次 dispatch。沒有中間 activation。 如果 tensor 不相容、我們在訓練、gradient 被啟用、dtype 不對，或是原生路徑失敗，就直接回到： `` F.silu(norm(input_tensor)) `` 我刻意停在那裡。 把整個 residual block 融合起來確實很誘人，但 GroupNorm 加 SiLU 是那一組我可以單獨隔離、測試、證明的小對。 結果證明，那個克制很重要。 ## NGMS 是另一回事 最終加速中有一部分，需要跟引擎本身的工作分開來看。 NGMS，也就是 Negative Guidance minimum sigma，可以在 sampling 的可符合期間跳過 unconditional guidance。 在使用 classifier-free guidance 時，UNet 通常會同時做 conditional 和 unconditional 的工作。在低 CFG（像是 1.15）搭配五步排程時，跳過可省略的 unconditional 工作可以移除相當可觀的運算量。 那當然很快，因為 GPU 根本沒做其中一些工作。 這個 fork 把 NGMS 預設設為 `1.0`，並為這個調校過的工作流程啟用 all-steps 行為。 但這跟讓 attention 或 GroupNorm 變快不是同一類別。 NGMS 會改變 denoising 的計算。它可能改變構圖和細節，而且在啟用時會記錄在 PNG metadata 中。 所以這裡其實有兩個效能故事。 第一個是讓既有引擎更便宜：Metal attention、更少的 command-buffer 提交、更好的記憶體行為、移除過時的轉換、融合運算。 第二個是透過 NGMS 讓引擎做更少的工作。 任何關於這個 fork 的控制基準測量，都必須同時呈現兩者。 ## 最有用的優化，是我刪掉的那些 這個專案有很大一部分是嘗試一些聽起來該有效的方法，然後把它們拿掉。 Packed QKV projections 就是一個例子。 我實作了它們。測試中產生的圖片是 byte-identical。 效能從 8.988 秒變成 9.011 秒。 大約慢了 0.