Causal Depthwise Conv1d

agicy2026/6/6大约 2 分钟

Causal Depthwise Conv1d

原始题目：LeetGPU - Causal Depthwise Conv1d

题目描述

对一个形状为 $(B, L, D)$ 的批量序列张量 $x$ 实现因果深度可分离一维卷积，产生相同形状的输出。在深度可分离卷积中，每个通道 $d$ 使用自己的核权重 $weight[d, :]$ 独立卷积——通道之间没有混合。卷积是因果的：输出位置 $l$ 只能依赖于输入位置 $0, 1, \dots, l$ （过去和现在），绝不能依赖未来位置。

这是 Mamba 等状态空间模型的关键组件，在选择性扫描之前用于混合每个特征通道内的局部上下文。

output[b, l, d] = bias[d] + \sum_{k=0}^{K-1} weight[d, k] \cdot x[b, l-k, d]

其中 $l - k < 0$ 的位置视为零（左侧零填充）。张量布局为 channels-last： $x[b, l, d]$ 存储在偏移 $b \times L \times D + l \times D + d$ 处。

实现要求

solve 函数签名必须保持不变。
不允许使用外部库。
序列开始之前的输入位置（ $l - k < 0$ ）必须视为零。

示例

B=1, L=4, D=2, K=3
x = [[[1,2],[3,4],[5,6],[7,8]]]  shape (1,4,2)
weight = [[1,0,-1],[1,1,1]]       shape (2,3), bias = [0,0]
output[l=0]: d0:1*1=1         d1:1*2=2        → [1,2]
output[l=1]: d0:3*1+1*0=3     d1:4*1+2*1=6    → [3,6]
output[l=2]: d0:5*1+3*0+1*(-1)=4  d1:6+4+2=12 → [4,12]
output[l=3]: d0:7*1+5*0+3*(-1)=4  d1:8+6+4=18 → [4,18]

约束条件

$1 \le B \le 16$ ， $1 \le L \le 8{,}192$ ， $1 \le D \le 8{,}192$ ， $1 \le K \le 8$ 。
所有张量使用 32 位浮点数，channels-last 布局。
性能测试在 $B=8, L=2{,}048, D=4{,}096, K=4$ 下进行。

#include <cuda_runtime.h>
__global__ void cdconv(const float* x, const float* w, const float* bias, float* y, int B, int L, int D, int K) {
    int b=blockIdx.x,d=blockIdx.y-blockDim.y+threadIdx.y,l=blockIdx.z*blockDim.x+threadIdx.x;
    if(b<B&&d<D&&l<L){float sum=bias[d];for(int k=0;k<K&&l-k>=0;k++)sum+=w[d*K+k]*x[(b*L+(l-k))*D+d];y[(b*L+l)*D+d]=sum;}
}
extern "C" void solve(const float* x,const float* w,const float* bias,float* y,int B,int L,int D,int K) {
    dim3 t(32,min(D,8)), b(B,(D+min(D,8)-1)/min(D,8),(L+31)/32);
    cdconv<<<b,t>>>(x,w,bias,y,B,L,D,K); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def cdconv(x_ptr,w_ptr,bias_ptr,y_ptr, B,L,D,K, BLOCK_L:tl.constexpr,BLOCK_D:tl.constexpr):
    b=tl.program_id(0); l=tl.program_id(1)*BLOCK_L+tl.arange(0,BLOCK_L)
    d=tl.program_id(2)*BLOCK_D+tl.arange(0,BLOCK_D)
    mask=(l[:,None]<L)&(d[None,:]<D)
    acc=tl.load(bias_ptr+d[None,:],mask=d[None,:]<D)
    for k in range(K):
        valid=l[:,None]-k>=0
        v=tl.load(x_ptr+(b*L+(l[:,None]-k))*D+d[None,:],mask=mask&valid,other=0.0)
        acc+=v*tl.load(w_ptr+d[None,:]*K+k,mask=d[None,:]<D)
    tl.store(y_ptr+(b*L+l[:,None])*D+d[None,:],acc,mask=mask)

Causal Depthwise Conv1d

Causal Depthwise Conv1d

题目描述

实现要求

示例

约束条件

解题思路

代码实现

CUDA

Triton