专栏名称: 极市平台
极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台,为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯,行业动态,在线分享信息,线下活动等。 网站: http://cvmart.net/
今天看啥  ›  专栏  ›  极市平台

Linear Attention的cuda kernel实现解析

极市平台  · 公众号  · 算法 科技媒体  · 2024-10-21 22:00

主要观点总结

文章主要介绍了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 ………………………………

原文地址:访问原文地址
快照地址: 访问文章快照
总结与预览地址:访问总结与预览