Linear Recurrence

原始题目：LeetGPU - Linear Recurrence

题目描述

给定两个矩阵 $a$ 和 $x$，形状均为 $[B, L]$（批量大小 $\times$ 序列长度），计算形状为 $[B, L]$ 的线性递推 $h$：

$$\begin{aligned} h[b, 0] &= x[b, 0] \\ h[b, t] &= a[b, t] \times h[b, t-1] + x[b, t], \quad t \ge 1 \end{aligned}$$

所有值为 float32。此操作是状态空间模型（SSM）（如 Mamba、S4 和 H3）的核心计算原语。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 结果必须存储在输出张量 $h$ 中。

示例

示例 1 — 指数衰减 (a=0.5, 单位脉冲):
  a = [0.5, 0.5, 0.5, 0.5], x = [1, 0, 0, 0]
  h = [1, 0.5, 0.25, 0.125]

示例 2 — 前缀和 (a=1, 全1输入):
  a = [1, 1, 1, 1], x = [1, 1, 1, 1]
  h = [1, 2, 3, 4]

示例 3 — B=2, L=4:
  a = [[0.5,0.5,0.5,0.5],[1,1,1,1]], x = [[1,0,0,0],[1,1,1,1]]
  h = [[1,0.5,0.25,0.125],[1,2,3,4]]

约束条件

• $1 \le B \le 256$，$1 \le L \le 65{,}536$。
• $a$ 和 $x$ 均为 float32。
• 性能测试在 $B = 64,\ L = 16{,}384$ 下进行。

解题思路

线性递推的串行依赖是 GPU 并行化的核心难点——$h[t]$ 依赖 $h[t-1]$ 使得不能直接全并行。对于小 $L$，可以使用并行扫描（parallel scan / prefix sum 的关联性推广）在 $O(\log L)$ 步骤中求解。对于大 $L$ 和不同的 $a, x$，分段并行+串行边界传递是常用策略。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void lr_kernel(const float* a, const float* x, float* h, int B, int L) {
    int b=blockIdx.x; if(b>=B)return;
    float acc=x[b*L];
    h[b*L]=acc;
    for(int t=1;t<L;t++){
        acc=a[b*L+t]*acc + x[b*L+t];
        h[b*L+t]=acc;
    }
}
extern "C" void solve(const float* a, const float* x, float* h, int B, int L) {
    lr_kernel<<<B,1>>>(a,x,h,B,L); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def lr_kernel(a_ptr,x_ptr,h_ptr,B,L,BLOCK:tl.constexpr):
    b=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=b<B
    acc=tl.load(x_ptr+b*L); tl.store(h_ptr+b*L,acc,mask=mask)
    for t in range(1,L):
        acc=tl.load(a_ptr+b*L+t)*acc + tl.load(x_ptr+b*L+t)
        tl.store(h_ptr+b*L+t,acc,mask=mask)