12 秒内 AI 在手机上完成作画！谷歌提出扩散模型推理加速新方法

本文源自：量子位

只需 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 

卡奥斯开源社区是为开发者提供便捷高效的开发服务和可持续分享、交流的 IT 前沿阵地，包含技术文章、群组、互动问答、在线学习、大赛活动、开发者平台、OpenAPI 平台、低代码平台、开源项目等服务，社区使命是让每一个知识工人成就不凡。

 官网链接：Openlab.cosmoplat—打造工业互联网顶级开源社区