只需12秒，只憑手機自己的算力，就能拿Stable Diffusion生成一張圖像，而且完成了20次反覆運算。
現在的擴散模型基本都超過了10億參數，想要快速產生一張圖片，要嘛基於雲端運算，要嘛本地端硬體要夠強大。 但隨著大模型應用逐漸普及開來，在個人電腦、手機上跑大模型很可能是未來的新趨勢。
由此，谷歌的研究員們帶來了這篇新成果，名字就叫Speed is all you need：透過GPU最佳化加速大規模擴散模型在設備上的推理速度。
三步驟最佳化加速
該方法是針對Stable Diffusion來做的最佳化，但同時也能適應其他擴散模型。面向的任務是從文本生成圖像。
具體最佳化可以分成三個部分：
- 設計專門的核心
- 提升Attention模型效率
- Winograd卷積加速
首先來看專門設計的核心，它包括了組歸一化和GELU啟動函數。
組歸一化是在整個UNet體系結構中實現，這種歸一化的工作原理是將特徵映射的通道劃分為更小的組，並對每個組獨立歸一，使組歸一化較少依賴於批大小，並且能適應更大範圍的批次處理大小和網路架構。
研究人員以GPU著色器（shader）的形式設計了一個獨特的核，能在沒有任何中間張量的情況下，在單個GPU命令中執行所有核心。
GELU啟動函數中，包含大量的數值運算，如懲罰、高斯誤差函數等。
通過一個專用著色器來整合這些數值運算以及伴隨的分割和乘法操作，使得這些運算能放在一個簡單的draw call裡。
Draw call是CPU調用圖像程式設計介面，命令GPU進行算繪的操作。
接下來，到了提升Attention模型效率方面，論文介紹了兩種最佳化方法。
其一是部分融合softmax函數。
為了避免在大矩陣A上執行整個softmax運算，該研究設計了一個GPU著色器來運算L和S向量以減少運算，最終得到一個大小為N×2的張量。然後將softmax運算和矩陣V的矩陣乘法融合。
這種方法大幅減少了中間程式的佔用記憶體張量和總體延遲。
需要強調的是從A到L、S的運算映射的並行是有限的，因為結果張量中的元素比輸入張量A中的元素數量要少得多。
為了增加並行、進一步降低延遲，該研究將A中的元素組成block，將歸約操作（reduction operations）切分為多個部分進行。
然後在每個block上執行運算，然後將其簡化為最終結果。
利用精心設計的執行緒和記憶體緩存管理，可以在多個部分實現使用單個GPU命令降低延遲。
另一種最佳化方法是FlashAttention。
這是去年熱門起來的IO感知精確注意力演算法，具體有兩種加速技術：按塊遞增運算即平鋪、並在後向傳遞中重新運算注意力，將所有注意力操作融合到CUDA核心中。
相較於標準Attention，這種方法能減少HBM（高頻寬記憶體）存取，提高整體效率。
不過FlashAttention核心的暫存器密集度非常高（register-intensive），所以該團隊是有選擇性地使用這一最佳化方法。
他們在注意力矩陣d=40的Adreno GPU和Apple GPU上使用FlashAttention，其他情況下使用部分融合softmax函數。
第三部分是Winograd卷積加速。
它的原理簡單來說就是使用更多的加法運算來減少乘法運算，進而降低運算量。
但弊端也很明顯，將會帶來更多的VRAM消耗和數值錯誤，尤其是在tile比較大的情況時。
Stable Diffusion的主幹非常依賴3×3卷積層，尤其是在圖像解碼器方面，這裡90%的層都是由3×3卷積層構成的。
研究人員分析後發現，在使用4×4大小的tile時，是模型運算效率和顯存利用率的最佳平衡點。
實驗結果
為了評估提升效果，研究人員先在手機上進行了基準測試。
結果表明，兩部手機在使用了加速演算法後，產生圖片的速度都明顯提升。
其中三星S23 Ultra的延遲降低了52.2%，iPhone 14 Pro Max上的延遲降低了32.9%。
在三星S23 Ultra上端到端從文本生成一張512×512像素的圖片，反覆運算20次，耗時在12秒以內。
論文連結：
加入電腦王Facebook粉絲團