Skip to content

Causal Convolution 从数学到 Dao 与 MLU 实现

Causal convolution 将当前位置和有限个历史位置组合成新的表示。理解这个算子,需要同时回答三个问题:语言模型为什么需要这种局部混合,反向传播究竟沿哪些位置传递梯度,以及硬件怎样复用重叠窗口而不反复搬运数据。

本教程围绕 causal depthwise Conv1d 展开,覆盖训练前向、训练反向、带状态的分块执行和增量推理。Dao 指 Dao-AILab/causal-conv1d,源码讲解固定采用 v1.5.0;Mamba 模型调用采用 state-spaces/mamba v2.2.2。这些版本用于建立可复查的学习基线,不表示它们是最新版本,也不表示你的 MLU 实现采用相同版本。

阅读顺序

章节 学习目标
第一章 大模型业务与模型位置 理解自回归、Mamba、短卷积与状态空间递推的分工
第二章 前向数学与传统卷积对比 从标准卷积得到 Dao 的精确下标、边界和计算量
第三章 反向传播的完整推导 独立推导 dx、dw、db 和 SiLU 梯度,用数值例子检查方向
第四章 Dao 接口与源码地图 理解 Python、C++、CUDA 分工及 shape、stride、dtype 约束
第五章 Dao 前向 CUDA 实现 分别理解时间连续和通道连续两条路径的搬运与计算
第六章 Dao 反向 CUDA 实现 理解重计算、右侧 halo、局部归约和全局累加
第七章 状态 分块与增量推理 理解状态梯度、packed sequence、循环缓存和请求槽位
第八章 MLU 迁移与验证实践 建立迁移设计、正确性实验、性能分析和排错方法

建议先读前三章并手算第三章的例子,再对照第四至六章阅读 Dao 源码。第七章说明实际训练与推理中最容易漏掉的边界,第八章将这些语义转化为 MLU 实现的核对步骤。

统一符号

符号 含义
$B,C,L,K$ batch、通道数、序列长度、卷积核宽度
$x\in\mathbb R^{B\times C\times L}$ 卷积输入,通常是投影后的特征而非 token ID
$w\in\mathbb R^{C\times K},\ \beta\in\mathbb R^C$ depthwise 权重与可选偏置
$z,y$ 激活前结果与最终输出;$y=\phi(z)$
$g=\partial\mathcal L/\partial y$ 上游传入的输出梯度,即接口中的 dout
$\delta=\partial\mathcal L/\partial z$ 经过激活反向后的梯度
$S,F$ 分块输入的历史状态与分块输出的最终状态
$N$ SSM 的状态维度;与序列长度 $L$ 区分

除特别说明外,主线取 stride=1、dilation=1、depthwise multiplier=1、零初始状态。矩阵公式先忽略有限精度舍入,精度问题在实现章节单独讨论。一次乘加按 2 FLOPs 计。

MLU 部分是依据算子语义建立的迁移设计,不绑定某一代 MLU 的片上容量、指令吞吐或任务配置。具体代码需要结合你的设备、SDK 和算子接口核对。

返回 AI 与系统工程