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 和算子接口核对。