自定义卷积注意力算子的CUDA实现

发布时间：2023-07-18 16:00

背景

注意力机制的一种卷积替代方式

我们在上面这篇文章讲到了设计适当的卷积来替代注意力机制

最终设计的卷积如下

$Out_{i,c} = \sum^k_{j=1} X_{ i+j-\left \lceil \frac{k+1}{2} \right \rceil, c} \cdot W_{\left \lceil \frac{cH}{d} \right \rceil, j} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (1)$

其中 $X \in R^{B,n,d}, W \in R^{H, k}$ ，且有 $W_{h,:} = Softmax(W_{h, :})$ ；c代表d这个维度；

要实现这样的卷积，有一种折中的方式是，令

$\\X => X' \in R^{BH,n,\frac{d}{H}} \\W => W' \in R^{BH, n, n} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (2)$

的具体设计请参考上述文章，此处不再赘述；

为了更高效的实现这样的卷积，考虑对该自定义卷积算子进行CUDA实现；

方案

CUDA 实现自定义卷积一般需要几个固定的步骤，在 CUDA实现focal_loss 这篇文章中详细说明了；这里引用一张图，如下

图1：CUDA实现算子的一般步骤介绍

这里CUDA实现自定义卷积算子的步骤和上面差不多，最大的区别是 step4/5，尤其是step5的kernel设计；

以下按照上述流程讲述实现方法；

---------

待续。。。

自定义卷积注意力算子的CUDA实现

背景

方案

相关推荐