Linear Self-Attention

原始题目：LeetGPU - Linear Self-Attention

题目描述

实现论文 "Transformers are RNNs" 中提出的线性注意力。给定 $Q$、$K$、$V$（均为 $M \times d$），计算：

$$\text{LinearAttention}(Q, K, V) = \frac{\phi(Q)(\phi(K)^T V)}{\phi(Q)\left(\sum_j \phi(K_j)\right)}$$

其中 $\phi(\cdot) = \text{elu}(\cdot) + 1$（elu 加 1 确保非负）。与标准 softmax attention（$O(M^2)$）不同，线性注意力利用 $\phi(Q)(\phi(K)^T V)$ 的计算顺序，将复杂度降为 $O(M \cdot d^2)$。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

约束条件

• $1 \le M \le 10{,}000$，$1 \le d \le 128$。
• 性能测试在 $M = 5{,}000$ 下进行。

解题思路

线性注意力的核心技巧是改变计算顺序：先算 $\phi(K)^T V$（$d \times d$），再与 $\phi(Q)$ 相乘——当 $d \ll M$ 时，这比 $O(M^2)$ 的 standard attention 快得多。$\phi$ 函数（elu+1）保证输出非负。但线性注意力缺少 softmax 的归一化效应，在建模长程依赖方面不如标准注意力。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void linear_attn(const float* Q, const float* K, const float* V, float* O, int M, int d) {
    int i = blockIdx.x;
    // phi(x) = elu(x) + 1
    for(int k = 0; k < d; k++) {
        float num = 0.0f, denom_kv = 0.0f, denom_k = 0.0f;
        for(int j = 0; j <= i; j++) {
            float qj = Q[i*d+k] > 0 ? Q[i*d+k] : expf(Q[i*d+k]) - 1.0f;
            float kj = K[j*d+k] > 0 ? K[j*d+k] : expf(K[j*d+k]) - 1.0f;
            qj += 1.0f; kj += 1.0f;
            num += kj * V[j*d+k];
            denom_k += kj;
        }
        float qi = Q[i*d+k] > 0 ? Q[i*d+k] : expf(Q[i*d+k]) - 1.0f; qi += 1.0f;
        O[i*d+k] = qi * num / (qi * denom_k + 1e-8f);
    }
}
extern "C" void solve(const float* Q, const float* K, const float* V, float* O, int M, int d) {
    linear_attn<<<M, 1>>>(Q, K, V, O, M, d);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def linear_attn(Q_ptr,K_ptr,V_ptr,O_ptr, M,d, BLOCK:tl.constexpr):
    i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=i<M
    d_range=tl.arange(0,d)
    qi=tl.load(Q_ptr+i[:,None]*d+d_range[None,:],mask=mask[:,None])
    phi_q = tl.where(qi>0, qi, tl.exp(qi)-1) + 1.0
    kv=tl.zeros((d,d),tl.float32); k_sum=tl.zeros((d,),tl.float32)
    for j in range(M):
        kj=tl.load(K_ptr+j*d+d_range); vj=tl.load(V_ptr+j*d+d_range)
        phi_k = tl.where(kj>0, kj, tl.exp(kj)-1) + 1.0
        kv += phi_k[:,None] * vj[None,:]
        k_sum += phi_k
    num = phi_q @ kv
    denom = tl.sum(phi_q * k_sum[None,:], axis=1, keepdim=True) + 1e-8
    tl.store(O_ptr+i[:,None]*d+d_range[None,:], num/denom, mask=mask[:,None])