主要观点总结
文章主要介绍了Linear Transformer的核心组件Linear Attention的原理,并对实现Linear Attention组件的forward cuda kernel进行了详细解析。文章通过解析三个cuda kernel,分别在不同的数据规模和优化选项下触发,展示了Linear Attention的计算过程,包括加载和存储Q、K、V,计算K*V^T乘积,加载Q值并计算输出,汇总计算结果,存储输出,更新指针和缓冲区等操作。文章还介绍了Linear Attention的目的,即将Self Attention的平方根序列长度级别复杂度降低为线性的复杂度,并给出了Linear Attention的公式解释。
关键观点总结
关键观点1: Linear Attention原理
Linear Attention是为了将Self Attention的复杂度从平方根序列长度级别降低为线性复杂度。
关键观点2: Linear Attention实现
通过解析三个cuda kernel,展示了Linear Attention的计算过程,包括加载和存储Q、K、V,计算K*V^T乘积,加载Q值并计算输出,汇总计算结果,存储输出,更新指针和缓冲区等操作。
关键观点3: cuda kernel调用
根据不同的数据规模和优化选项,文章使用了三个不同的cuda kernel进行Linear Attention的计算,其中包括lmha_kernel和lmha_low_occupancy_kernel。
文章预览
↑ 点击 蓝字 关注极市平台 作 者丨BBuf 来源丨GiantPandaCV 编辑丨极市平台 极市导读 本文主要是对Linear Transformer的核心组件Linear Attention进行了原理讲解,并对实现Linear Attention组件的forward cuda kernel进行了详细解析。 >> 加入极市CV技术交流群,走在计算机视觉的最前沿 欢迎来 https://github.com/BBuf/how-to-optim-algorithm-in-cuda 踩一踩。 0x0. 问题引入 Linear Attention的论文如下: Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention: https://arxiv.org/pdf/2006.16236.pdf 。官方给出实现代码地址: https://github.com/idiap/fast-transformers 。虽然这个仓库是Linear Attention的原始实现,但基于这个codebase也引出了后续的一系列线性Attention的工作比如:Efficient Attention: Attention with Linear Complexities ( https://arxiv.org/abs/1812.01243 ), Linformer: SelfAttention with Linear Complexity( https://arxiv.org/abs/2006
………………………………