集成 Metal FlashAttention:加速 Apple 生态系统中图像生成的核心

Draw Things是第一个在“边缘”直接在手机上运行成熟图像生成模型的实用应用程序。自推出以来,人们对本地运行的开源大型模型越来越感兴趣。LLaMA.cpp将大型语言模型引入笔记本电脑;MLC LLM在网络浏览器中执行语言和图像生成模型。最初的学术练习已经演变成一场运动:让我运行我的模型,就像云模型一样强大,本地且免费(“就像自由一样”)!

到目前为止,大多数算法创新和改进都发生在 NVIDIA CUDA 硬件上。当大多数人工智能相关计算发生在服务器端时,这是有意义的。然而,当我们将计算移近边缘时,如何为最常用的硬件平台之一推动相同的算法创新和改进激发了我们的想象力。

Metal FlashAttention

Metal FlashAttention 包含针对大型图像生成和语言模型中常见操作进行优化的 Metal 计算着色器。这包括薄矩阵乘法(例如 [4096, 320] x [320, 320])、缩放点积注意力(多头注意力或转换器的核心)和层归一化。它是Metal Performance Shaders ( MPS )的开源替代品。

GEMM

GEMM 计算通常出现在模型的稳定扩散变体(v1、v2、XL)中,无法达到 Apple 的 Metal Performance Shaders 或 MPSGraph 实现的最佳效果。Metal FlashAttention 利用simdgroup_async_copyAPI(自 A14 起),这是一种未记录的硬件功能,与计算和加载指令重叠。

猜你喜欢

转载自blog.csdn.net/iCloudEnd/article/details/132223452