Decaying Causal Attention

原始题目：LeetGPU - Decaying Causal Attention

题目描述

实现衰减因果注意力。给定查询矩阵 $Q$、键矩阵 $K$ 和值矩阵 $V$（形状均为 $seq\_len \times d_{model}$）以及标量衰减因子 $\gamma \in (0, 1]$，计算未归一化的因果注意力输出。位置 $n$ 以权重 $\gamma^{n-m}$ 关注所有过去位置 $m \le n$：

$$output[n] = \sum_{m=0}^{n} \gamma^{n-m} \cdot \frac{Q[n] \cdot K[m]}{\sqrt{d_{model}}} \cdot V[m]$$

与标准 softmax 注意力不同，这里没有归一化——权重从当前位置向后呈几何衰减。这是 RetNet（Retention Network）的并行形式，用作序列模型中注意力机制的一种递推友好替代方案。

实现要求

• 实现 solve 函数，签名保持不变。
• 不允许使用外部库。

示例

seq_len=2, d_model=4, gamma=0.5
Q=[[1,1,0,0],[1,1,0,0]], K=[[1,0,0,0],[0,1,0,0]], V=[[4,8,12,16],[4,8,12,16]]
Scores QK^T/√4: [[0.5,0.5],[0.5,0.5]]
Decay mask D[n,m]=0.5^(n-m): [[1,0],[0.5,1]]
Weighted A⊙D: [[0.5,0],[0.25,0.5]]
Output (A⊙D)V: [[2,4,6,8],[3,6,9,12]]

约束条件

• $1 \le seq\_len \le 8{,}192$，$1 \le d_{model} \le 256$。
• $0 < \gamma \le 1$。
• 性能测试在 $seq\_len = 4{,}096,\ d_{model} = 64$ 下进行。

解题思路

衰减因果注意力与标准 softmax attention 的区别在于用几何衰减替代了 softmax 归一化。计算流程仍是 $QK^T$ + 逐元素衰减掩码 + $@V$。衰减掩码是一个下三角矩阵，第 $(n,m)$ 个元素为 $\gamma^{n-m}$，对角线为 1。可以用递推方式避免 $O(seq\_len^2)$ 的计算成本——每个新位置 $n$ 的加权和可以从 $n-1$ 递推而来。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void dca(const float* Q,const float* K,const float* V,float* O,int S,int d,float gamma) {
    int i=blockIdx.x; if(i>=S)return;
    float scale=1.0f/sqrtf(d);
    for(int k=0;k<d;k++){float sum=0;for(int j=0;j<=i;j++){
        float dot=0;for(int kk=0;kk<d;kk++)dot+=Q[i*d+kk]*K[j*d+kk];
        sum+=powf(gamma,(float)(i-j))*dot*scale*V[j*d+k];
    }O[i*d+k]=sum;}
}
extern "C" void solve(const float* Q,const float* K,const float* V,float* O,int S,int d,float gamma) {
    dca<<<S,1>>>(Q,K,V,O,S,d,gamma); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def dca(Q_ptr,K_ptr,V_ptr,O_ptr, S,d,gamma, BLOCK:tl.constexpr):
    i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=i<S
    d_range=tl.arange(0,d)
    qi=tl.load(Q_ptr+i[:,None]*d+d_range[None,:],mask=mask[:,None])
    scale=1.0/tl.sqrt(d.to(tl.float32)); acc=tl.zeros((BLOCK,d),tl.float32)
    for j in range(S):
        kj=tl.load(K_ptr+j*d+d_range); vj=tl.load(V_ptr+j*d+d_range)
        score=tl.sum(qi*kj[None,:],axis=1)*scale
        decay=tl.where(i[:,None]>=j, gamma**(i[:,None].to(tl.float32)-j), 0.0)
        acc+=decay[:,None]*score[:,None]*vj[None,:]
    tl.store(O_ptr+i[:,None]*d+d_range[None,:],acc,mask=mask[:,None])