隻需 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 卷積加速。
它的原理簡單來說就是使用更多的加法計算來減少乘法計算，從而降低計算量。
但弊端也很明顯，将會帶來更多的顯存消耗和數值錯誤，尤其是在 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 秒以内。
論文地址：
https://arxiv.org/abs/2304.11267
