第五章 Dao 前向 CUDA 实现

从重叠窗口看到可复用数据

假设一个任务负责时间区间 $[t_0,t_0+T-1]$。前向输入范围是：

$$ [t_0-(K-1),\ t_0+T-1]. $$

其中左侧多出来的 $K-1$ 个值称为 halo。若每个输出独立从全局内存读取 $K$ 个元素，共需约 $TK$ 次读取；先搬入 $T+K-1$ 个元素，再在片上复用，可以显著减少重复访问。

halo 重复加载的相对开销约为 $(K-1)/T$。增大 $T$ 会减少该比例，但也增加片上容量、寄存器和任务粒度。块越大不一定越快：并行任务数不足或资源占用过高，都可能抵消复用收益。

时间连续路径的任务划分

causal_conv1d_fwd_kernel 的 grid 为 (B,C)，一个 block 负责一个 batch 的一个通道，并在 block 内循环处理该通道的时间 chunks。该版本启动 128 个线程；每线程处理 FP32 的 4 个元素或 FP16/BF16 的 8 个元素，对应每线程 16 字节。

所以一个 chunk 分别覆盖 512 或 1024 个时间位置。这里时间 chunk 不对应独立 block；它是同一 block 内的循环迭代。若 $BC$ 很小而 $L$ 很长，block 总数仍由 $BC$ 决定，这是理解该路径并行度的关键。

权重很短，先加载为每线程的 FP32 weight_vals，随后循环复用。输入通过 CUB 的 block load/store 组织；能走向量化时使用向量路径，否则使用带有效长度的路径处理尾部。

相邻线程与相邻 chunk 怎样交换历史

一个线程持有自己的若干连续时间元素，但它负责的第一个输出还需要前一线程末尾的数据。源码使用 smem_exchange 交换这些值。第一个线程需要的历史来自上一 chunk 的末尾；最开始将其初始化为零。

关键时序是：其他线程先发布当前数据，线程 0 读取上一 chunk 遗留的尾部，然后最后一个线程才能覆盖共享槽位，把当前 chunk 的尾部留给下一轮。相关同步不是装饰；过早覆盖会让 chunk 边界读取错误时间的数据。

获得“历史片段 + 当前片段”后，每个输出在寄存器中执行：

$$ z_i=\beta+\sum_{k=0}^{K-1}w_k x_{i-K+1+k}. $$

因为 $K$ 是模板常量，短循环可展开。FP32 累加完成后可直接计算 SiLU，最后转为输出 dtype 写回。融合省掉了单独的预激活输出与激活输入之间的一次大张量往返。

通道连续路径为什么要换组织

对 channel-last，连续地址对应不同通道。如果仍让相邻线程都沿某一个通道的时间前进，访存会跨越较大的时间 stride。Dao 因此先按适合合并访存的方式搬运时间×通道 tile，再在共享内存中以适合计算的方式读取。

该版本前向取 $T=64$。通道 tile 宽度为 $128/\operatorname{sizeof}(input)$，FP32 是 32 个通道，FP16/BF16 是 64 个通道。grid 为：

$$ \left(B,\left\lceil L/64\right\rceil, \left\lceil C/C_{tile}\right\rceil\right). $$

这条路径沿时间也产生独立 block。每个 block 读取自己的左 halo，因此不需要等待前一个 block。

先搬运再改变计算访问方向

加载阶段每线程取 16 字节，一组 8 个线程覆盖一行 128 字节。数据写入共享内存：

$$ x_{smem}[K-1+T][C_{tile}+kNElts]. $$

前 $K-1$ 行保存 halo，随后是本块输入。额外的通道 padding 改变共享内存行跨度，是布局设计的一部分；具体冲突和收益应结合硬件分析，而不能把“多加若干列”机械套到 MLU 上。

加载完成后，线程重新按通道与时间小段分工：每个线程拿某通道的一段时间及其历史，在寄存器中连续计算多个输出。算完再写回共享内存，最后恢复适合连续通道访问的方式写到全局内存。

可将其理解为四步：

全局内存按通道连续读
  → 共享内存保存 tile 和左 halo
  → 按单通道时间小段计算
  → 共享内存重排后按通道连续写

“搬运时线程怎样分工”和“计算时线程怎样分工”可以不同。这是比某个 CUDA 线程数更值得迁移的思路。

三类边界分别处理

序列开头的 halo 从零或 initial_states 获取；中间 tile 的 halo 读取真实历史输入；最后不足一个 tile 的时间或通道区域通过边界条件处理。它们不是同一种 padding。

使用 seq_idx 时，每个窗口项比较源位置与目标位置的序列标识，不同标识的贡献取零。它阻止 packed 数据中前一条样本的结尾污染后一条样本开头。必须确保标识编码确实区分窗口内的不同片段。

最终状态由最后一个时间 tile 写出。该 tile 已经包含自己的左 halo，所以即便当前序列很短，也有机会从历史与当前输入的组合中提取尾部 $K-1$ 个值。最终状态保存的是卷积输入，不是 SiLU 输出。

两条路径的对照

设计点	时间连续	通道连续
block 所有权	一个 batch、一个通道	一个 batch、一个时间 tile、一个通道 tile
时间推进	block 内循环	时间 tile 之间独立
halo 来源	线程交换及前一 chunk 遗留	显式读左侧输入或初始状态
连续搬运方向	时间	通道
片上主要用途	交换邻接片段及 CUB IO	tile 存储与访问方向转换

不要只看两条路径的乘加表达式相同。真正影响性能的是数据地址、任务数量、共享内存访问和同步方式。

下一章：Dao 反向 CUDA 实现。