DIFFLET：面向 AWS Trainium 上 Diffusion Transformer 模型的快速推理引擎

arXiv:2610.23047v1 发布类型：cross 摘要：在 AWS Trainium2 上，扩散模型的每个组件都作为针对固定形状预先编译的程序运行，而不同模型的流水线、不同请求的提示词长度以及不同去噪步骤中的缓存决策各不相同。我们提出了 Difflet，这是一种服务引擎：五种模型流水线共享一个编译后端，联合注意力按注意力头分配到不同核心上，并由主机控制器在每个去噪步骤决定运行去噪器还是复用其先前的输出。一个小型探测程序在设备上计算缓存信号，并在每个步骤返回一个数值；常驻模型进程则在不同请求之间保持所有程序处于加载状态。在五类图像和视频模型上，与使用 PyTorch/XLA 注意力的同一框架相比，融合注意力使去噪器步骤加速 1.2–4.3 倍；当视觉 token 数达到 32.8k 时，它是唯一能够完成编译的注意力实现。在五个模型中的四个上，常驻式服务相比每次生成都从温文件缓存重新加载权重的独立运行快 1.5–3.4 倍。