071、芯片级优化：扩散模型专用加速器设计手记

上周调一个Stable Diffusion推理管线，在Jetson AGX上跑出12秒的生成耗时。客户要求压到3秒内，常规的CUDA核优化、算子融合手段都用尽了，离目标还差一大截。盯着nsys性能分析报告里那些稀疏注意力矩阵和残差块的内存搬运开销，突然意识到：通用GPU的架构红利，在扩散模型这种特殊计算图面前已经见底了。

扩散模型的硬件不匹配困境

扩散推理的本质是迭代去噪——20到50轮的前向传播，每轮都是相似的U-Net计算图。通用GPU的算力很强，但大量功耗花在了不必要的地方。比如调度器在每一步都要重新加载权重，尽管这些权重在迭代中完全不变；再比如噪声预测中的分组归一化，在通用ALU上需要多次访存和同步。

更头疼的是那些条件控制模块。做文生图时，CLIP文本编码器的输出要反复和U-Net的交叉注意力层交互，这个过程中大量的张量转置和重塑操作，在GPU的SIMD架构上会产生大量线程束分化。你会在性能分析器里看到，某些核函数的执行效率只有理论峰值的30%。

专用加速器的设计切入点

去年开始接触几家初创公司的AI加速芯片，发现他们不约而同地在做扩散模型优化。拆开看，核心思路都是针对计算特征做硬化。

内存子系统必须重构。扩散模型权重通常在1-4GB，传统架构每轮迭代都要从DRAM读取。专用设计会在片上集成大容量SRAM作为权重缓存，50轮迭代只需加载一次。某款芯片甚至做了权重压缩，在加载时实时解压，带宽需求降了40%。

归一化层要硬件化。GroupNorm在扩散模型里出现频率极高，通用GPU上需