SSM Selective Scan

agicy2026/6/6大约 2 分钟

SSM Selective Scan

原始题目：LeetGPU - SSM Selective Scan

题目描述

实现状态空间模型（SSM）选择性扫描的前向传播，这是 Mamba 风格序列模型的核心操作。给定输入序列 $u$ 、时间步参数 $\Delta$ 、状态转移矩阵 $A$ 、输入投影 $B$ 、输出投影 $C$ 和跳跃连接权重 $skip$ ，以 float32 计算输出序列 $y$ 。

对每个 batch $b$ 、位置 $t$ 和通道 $d$ ：

\begin{aligned} \bar{A}_{b,t,d,n} &= \exp(\Delta_{b,t,d} \cdot A_{d,n}) \\ \bar{B}_{b,t,d,n} &= \Delta_{b,t,d} \cdot B_{b,t,n} \\ h_{b,t,d,n} &= \bar{A}_{b,t,d,n} \cdot h_{b,t-1,d,n} + \bar{B}_{b,t,d,n} \cdot u_{b,t,d} \\ y_{b,t,d} &= \sum_n C_{b,t,n} \cdot h_{b,t,d,n} + skip_d \cdot u_{b,t,d} \end{aligned}

初始隐藏状态 $h_{b,-1,d,n} = 0$ 。所有通道 $d$ 相互独立——它们共享相同的 $B$ 、 $C$ 投影，但在 $A$ 中有各自独立的状态转移行。

实现要求

实现 solve(u, delta, A, B, C, skip, y, batch, seq_len, d_model, d_state)，签名保持不变。
不允许使用外部库。

示例

batch=1, seq_len=4, d_model=2, d_state=2
u = [[[1,0],[0,1],[1,1],[0,0]]], delta 全1
A = [[-0.5,-1.0],[-0.5,-1.0]] → A_bar ≈ [[0.607,0.368],[0.607,0.368]]
t=0: y=[1,0], t=1: y=[0,1], t=2: y=[2.368,2.368], t=3: y=[0.599,0.555]

约束条件

$1 \le batch \le 16$ ， $1 \le seq\_len \le 8{,}192$ ， $1 \le d_{model} \le 2{,}048$ ， $1 \le d\_state \le 64$ 。
$\Delta > 0$ ， $A < 0$ （确保 $\bar{A} \in (0, 1)$ ）。
性能测试在 $batch=4, seq\_len=4{,}096, d_{model}=512, d\_state=16$ 下进行。

解题思路

SSM 选择性扫描是 Mamba 区别于传统 SSM（S4）的关键—— $\Delta$ 、 $B$ 、 $C$ 是输入相关的（选择性），而非固定的。这破坏了卷积表示，使得必须逐时间步串行计算。与线性递推类似，可以使用并行扫描来处理选择性 SSM：将串行递推转化为关联操作，然后在 $O(\log L)$ 并行步中求解。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void ssm_scan(const float* u, const float* delta, const float* A,
    const float* B, const float* C, const float* skip, float* y, int batch, int L, int D, int N) {
    int b=blockIdx.x; if(b>=batch)return;
    // Simplified: per-batch sequential scan
    for(int d=0;d<D;d++){
        float h[64]={0}; float Ad[64];
        for(int n=0;n<N;n++)Ad[n]=expf(delta[b*L*D+0*D+d]*A[d*N+n]);
        for(int t=0;t<L;t++){
            float d_t=delta[(b*L+t)*D+d];
            for(int n=0;n<N;n++){
                float Abar=expf(d_t*A[d*N+n]);
                float Bbar=d_t*B[(b*L+t)*N+n];
                h[n]=Abar*h[n]+Bbar*u[(b*L+t)*D+d];
            }
            float yt=0;for(int n=0;n<N;n++)yt+=C[(b*L+t)*N+n]*h[n];
            y[(b*L+t)*D+d]=yt+skip[d]*u[(b*L+t)*D+d];
        }
    }
}
extern "C" void solve(const float* u,const float* delta,const float* A,const float* B,
    const float* C,const float* skip,float* y,int batch,int L,int D,int N) {
    ssm_scan<<<batch,1>>>(u,delta,A,B,C,skip,y,batch,L,D,N); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def ssm_scan(u_ptr,delta_ptr,A_ptr,B_ptr,C_ptr,skip_ptr,y_ptr, batch,L,D,N, BLOCK:tl.constexpr):
    b=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=b<batch
    d_r=tl.arange(0,D); n_r=tl.arange(0,N)
    h=tl.zeros((D,N),tl.float32)
    for t in range(L):
        dt=tl.load(delta_ptr+(b[:,None]*L+t)*D+d_r[None,:],mask=mask[:,None])
        ut=tl.load(u_ptr+(b[:,None]*L+t)*D+d_r[None,:],mask=mask[:,None])
        At=dt[:,:,None]*tl.load(A_ptr+d_r[:,None]*N+n_r[None,:])
        Bt=dt[:,:,None]*tl.load(B_ptr+(b[:,None]*L+t)*N+n_r[None,:],mask=mask[:,None])
        h=tl.exp(At)*h+Bt*ut[:,:,None]
        Ct=tl.load(C_ptr+(b[:,None]*L+t)*N+n_r[None,:],mask=mask[:,None])
        yt=tl.sum(Ct*h,axis=2)+tl.load(skip_ptr+d_r[None,:])*ut
        tl.store(y_ptr+(b[:,None]*L+t)*D+d_r[None,:],yt,mask=mask[:,None])