第四章 Dao 接口与源码地图

按调用链阅读源码

这里固定阅读 v1.5.0。先弄清输入输出，再追踪模板与线程布局，比直接从 CUDA 文件第一行开始更容易建立整体理解。

文件	重点入口	要回答的问题
causal_conv1d_interface.py	causal_conv1d_fn、CausalConv1dFn、causal_conv1d_ref	autograd 保存什么，参考语义是什么
causal_conv1d.cpp	causal_conv1d_fwd/bwd/update	哪些输入合法，如何分派与分配输出
causal_conv1d_fwd.cu	普通路径与 channellast 路径	tile、halo、共享内存、激活融合
causal_conv1d_bwd.cu	两种布局的反向 kernel	重计算、dx 所有权、参数梯度归约
causal_conv1d_update.cu	update kernel	如何读取历史并原地推进状态
test_causal_conv1d.py	各类参数化测试	哪些布局、长度、状态和梯度被覆盖

Python 包装处理用户参数和 autograd，C++ 处理合法性、内存与设备流，CUDA 完成计算。移植时三层都属于算子行为，单独把 CUDA 数学循环翻译正确并不够。

完整序列接口的契约

参数	逻辑 shape	含义
x	[B,C,L]	当前输入序列
weight	[C,K]	按最旧到当前排序的每通道卷积核
bias	[C] 或 None	激活前偏置
seq_idx	[B,L] 或 None	packed 序列的所属标识
initial_states	[B,C,K-1] 或 None	当前块左侧的历史输入
return_final_states	bool	是否同时输出历史状态
final_states_out	[B,C,K-1] 或 None	可选的状态输出缓冲区
activation	None、silu、swish	后两者在这里走相同激活

返回 out 的逻辑 shape 始终为 [B,C,L]；请求最终状态时返回二元组。final_states_out 是输出存储，不应当作额外的可微历史输入。seq_idx 是离散元数据，没有对它的连续梯度。Python 接口

shape 相同为什么仍然有两种布局

普通连续 [B,C,L] 的 stride 为：

$$ (s_B,s_C,s_L)=(CL,L,1). $$

时间相邻元素连续。如果先生成连续 [B,L,C]，再 transpose(1,2) 得到逻辑 [B,C,L]，则 stride 为：

$$ (s_B,s_C,s_L)=(LC,1,C). $$

这才是该库所谓的 channel-last 路径：通道在物理内存中连续，但调用时仍然传 [B,C,L]。不能直接把逻辑 [B,L,C] 当作参数传入。

一般地址为：

$$ \operatorname{offset}(b,c,t)=b\,s_B+c\,s_C+t\,s_L. $$

源码参数里的 stride 以元素为单位，字节地址还要乘 dtype 大小。投影后切出的视图可能在相邻时间之间包含其他分支，因此 $s_L$ 不一定等于当前视图的 $C$。实现应该使用真实 stride，不能看到 channel-last 就自行重建为紧凑布局。

Python 包装在时间与通道两个维度都不连续时将输入转成 contiguous；这可能有额外拷贝成本。MLU 的端到端基准需要将此类转换计入。

v1.5.0 的具体限制

C++ 完整序列接口支持 $K=2,3,4$，输入和权重分别分派 FP32、FP16、BF16，bias dtype 要与权重一致。channel-last 分支检查 $C$ 为 8 的倍数，并要求输入 batch stride 和时间 stride 为 8 的倍数。seq_idx 为连续 CUDA int32 张量，且只走 channel-last。

初始状态与最终状态同样限于 channel-last 分支，状态 dtype 与输入一致，通道 stride 必须为 1。Python 层也禁止 seq_idx 与 initial_states 同时使用，并禁止使用 seq_idx 时请求最终状态。这些限制属于该版本实现，不是数学上无法组合这些功能。C++ 参数检查

不能从包装层的宽松检查推断底层一定接受某布局。例如某状态在 Python 看起来有一个连续维度，仍可能不满足 C++ 的 channel stride 条件。做兼容实现时，应沿调用链检查最终限制。

autograd 如何连接前反向

CausalConv1dFn.forward 保存 $x,w,\beta$、seq_idx 和 initial_states；保存激活开关，并记录是否需要最终状态梯度及初始状态梯度。反向收到 dout，如果前向返回状态，还会收到 dfinal_states。

C++ 反向可接受预分配的 dx；Python 公共包装传 None，由底层分配。这个入口便于上层融合时让梯度直接写入指定缓冲区。dweight 和存在时的 dbias 先分配 FP32 零张量，kernel 累加完成后转换为参数 dtype 返回。反向分配与返回

返回值的 dtype 与中间累加 dtype 是两件事。用 FP16 返回梯度，不代表所有乘加都用 FP16；用 FP32 累加，也不代表最终输出能保存 FP32 的全部有效位。

参考函数如何使用

causal_conv1d_ref 用框架 Conv1d 表达语义，是理解 padding、state 和切片的首选。它会把输入转换到 weight dtype，最终输出再转回输入 dtype。因此混合 dtype 时，参考路径和专用 kernel 的舍入过程可能不同，不能要求逐 bit 一致。

建议维护两层参考：一层用高精度数学公式检查正确性，另一层用目标版本参考函数检查兼容性。两者的误差阈值与用途不同。

下一章：Dao 前向 CUDA 实现。