用你的手机/电脑运行文生图方案

大淘宝技术 · 公众号 · 程序员 · 2024-10-14 18:51

正文

请到「今天看啥」查看全文

一个共有节点，分别经过三个Linear层，得到Query/Key/Value，Query/Key经过形状变换进行BatchMatMul操作，再进行Scale，取Softmax操作；该结果和Value经过形状变换做BatchMatMul；之后把结果进行形状变换，得到最终的输出。可以看到上述总共有19个算子，包括12个形状变化算子，7个计算型算子。

大量的形状变化会带来很多的访存耗时，对于GPU高算力的硬件来说，访存耗时往往容易成为热点。因此，将上述结构，融合成2个算子，第一个是将三个Linear层权重融合在一起，只做一个Linear，这样形成更大的矩阵乘尺寸，更容易打满GPU算力，带来性能收益；第二个算子是将Attention算子融合成一个算子Fused-MultiHead-Attention，融合之后在该新算子内部仅需5个Kernel就可以实现整个Attention功能。消除了大量额外的形状变换算子，降低了访存压力，同时可以更容易基于Attention算子特性做进一步优化工作。

▐ GroupNorm/SplitGeLU融合

在Stable Diffusion中，有一个通用的结构ResnetBlock，其中包含了BroadCast Binary + GroupNorm + SiLU结构，在onnx模型图结构中包含了如下13个算子：

可以看到GroupNorm采用InstanceNorm+形变算子实现，gamma/beta被单独拆解为mul/add算子，细碎的算子会增加全局内存的访存次数、以及Kernel launch的压力。因此将上述通用结构合并成一个GroupNorm算子，该算子把前面的BroadCast Binary和后续的SiLU激活函数，融合在一起。高效的只需一个Kernel就可以实现上述计算需求。

同样的图融合原理，在Transformer激活函数中，Stable Diffusion Feed-Forward模块中采用GEGLU结构，对应onnx图结构如下。将该8个onnx图算子，融合为通用的SplitGeLU算子。