只需12秒，只憑手機(jī)自己的算力，就能拿Stable Diffusion生成一張圖像。

而且是完成了20次迭代的那種。

要知道，現(xiàn)在的擴(kuò)散模型基本都超過(guò)了10億參數(shù)，想要快速生成一張圖片，要么基于云計(jì)算，要么就是要本地硬件夠強(qiáng)大了。

而隨著大模型應(yīng)用逐漸普及開(kāi)來(lái)，在個(gè)人電腦、手機(jī)上跑大模型很可能是未來(lái)的新趨勢(shì)。

由此，谷歌的研究員們帶來(lái)了這篇新成果，名字就叫Speed is all you need：通過(guò)GPU優(yōu)化加速大規(guī)模擴(kuò)散模型在設(shè)備上的推理速度。

三步走優(yōu)化加速

該方法是針對(duì)Stable Diffusion來(lái)做的優(yōu)化，但同時(shí)也能適應(yīng)其他擴(kuò)散模型。面向的任務(wù)是從文本生成圖像。

具體優(yōu)化可以分成三個(gè)部分：

• 設(shè)計(jì)專門的內(nèi)核
• 提升Attention模型效率
• Winograd卷積加速

首先來(lái)看專門設(shè)計(jì)的內(nèi)核，它包括了組歸一化和GELU激活函數(shù)。

組歸一化是在整個(gè)UNet體系結(jié)構(gòu)中實(shí)現(xiàn)，這種歸一化的工作原理是將特征映射的通道劃分為更小的組，并對(duì)每個(gè)組獨(dú)立歸一，使組歸一化較少依賴于批大小，并且能適應(yīng)更大范圍的批處理大小和網(wǎng)絡(luò)架構(gòu)。

研究人員以GPU著色器（shader）的形式設(shè)計(jì)了一個(gè)獨(dú)特的核，能在沒(méi)有任何中間張量的情況下，在單個(gè)GPU命令中執(zhí)行所有內(nèi)核。

GELU激活函數(shù)中，包含大量的數(shù)值計(jì)算，如懲罰、高斯誤差函數(shù)等。

通過(guò)一個(gè)專用著色器來(lái)整合這些數(shù)值計(jì)算以及伴隨的分割和乘法操作，使得這些計(jì)算能放在一個(gè)簡(jiǎn)單的draw call里。

Draw call是CPU調(diào)用圖像編程接口，命令GPU進(jìn)行渲染的操作。

接下來(lái)，到了提升Attention模型效率方面，論文介紹了兩種優(yōu)化方法。

其一是部分融合softmax函數(shù)。

為了避免在大矩陣A上執(zhí)行整個(gè)softmax計(jì)算，該研究設(shè)計(jì)了一個(gè)GPU著色器來(lái)計(jì)算L和S向量以減少計(jì)算，最終得到一個(gè)大小為N×2的張量。然后將softmax計(jì)算和矩陣V的矩陣乘法融合。

這種方法大幅減少了中間程序的內(nèi)存占用張量和總體延遲。

需要強(qiáng)調(diào)的是從A到L、S的計(jì)算映射的并行是有限的，因?yàn)榻Y(jié)果張量中的元素比輸入張量A中的元素?cái)?shù)量要少得多。

為了增加并行、進(jìn)一步降低延遲，該研究將A中的元素組成block，將歸約操作（reduction operations）切分為多個(gè)部分進(jìn)行。

然后在每個(gè)block上執(zhí)行計(jì)算，然后將其簡(jiǎn)化為最終結(jié)果。

利用精心設(shè)計(jì)的線程和內(nèi)存緩存管理，可以在多個(gè)部分實(shí)現(xiàn)使用單個(gè)GPU命令降低延遲。

另一種優(yōu)化方法是FlashAttention。

這是去年火起來(lái)的IO感知精確注意力算法，具體有兩種加速技術(shù)：按塊遞增計(jì)算即平鋪、并在后向傳遞中重新計(jì)算注意力，將所有注意力操作融合到CUDA內(nèi)核中。

相較于標(biāo)準(zhǔn)Attention，這種方法能減少HBM（高帶寬內(nèi)存）訪問(wèn)，提高整體效率。

不過(guò)FlashAttention內(nèi)核的緩存器密集度非常高（register-intensive），所以該團(tuán)隊(duì)是有選擇性地使用這一優(yōu)化方法。

他們?cè)谧⒁饬仃嘾=40的Adreno GPU和Apple GPU上使用FlashAttention，其他情況下使用部分融合softmax函數(shù)。

第三部分是Winograd卷積加速。

它的原理簡(jiǎn)單來(lái)說(shuō)就是使用更多的加法計(jì)算來(lái)減少乘法計(jì)算，從而降低計(jì)算量。

但弊端也很明顯，將會(huì)帶來(lái)更多的顯存消耗和數(shù)值錯(cuò)誤，尤其是在tile比較大的情況時(shí)。

Stable Diffusion的主干非常依賴3×3卷積層，尤其是在圖像解碼器方面，這里90%的層都是由3×3卷積層構(gòu)成的。

研究人員分析后發(fā)現(xiàn)，在使用4×4大小的tile時(shí)，是模型計(jì)算效率和顯存利用率的最佳平衡點(diǎn)。

實(shí)驗(yàn)結(jié)果

為了評(píng)估提升效果，研究人員先在手機(jī)上進(jìn)行了基準(zhǔn)測(cè)試。

結(jié)果表明，兩部手機(jī)在使用了加速算法后，生成圖片的速度都明顯提升。

其中三星S23 Ultra的延遲降低了52.2%，iPhone 14 Pro Max上的延遲降低了32.9%。

在三星S23 Ultra上端到端從文本生成一張512×512像素的圖片，迭代20次，耗時(shí)在12秒以內(nèi)。

論文地址：https://arxiv.org/abs/2304.11267