第八章 MLU 迁移与验证实践
先迁移语义 再选择任务划分
从 Dao 参考实现迁移到 MLU，首先固定核方向、输入输出 stride、bias 和 SiLU 次序、状态含义与 dtype。CUDA 的 block、warp 和共享内存技巧，是实现这些语义的一种方法；MLU 应根据自己的计算单元、向量指令和片上存储重新安排。
下面使用片外内存与片上 buffer 描述设计。实际落到 NRAM、SRAM 或其他存储层时，应以目标 MLU 与 SDK 为准；公式中的容量预算也应代入设备的真实限制。
建议分阶段建立能力：先实现无状态完整序列前向和反向，再添加 SiLU，随后支持非紧凑 stride、状态梯度和 packed sequence，最后优化 update。每阶段保留可独立比较的输出，避免多个新特性同时引入而无法定位错误。
前向的两种布局设计
时间连续时，可以让一个任务负责某个 batch、一个或多个通道的连续时间 tile。搬入 $T+K-1$ 个元素，在片上形成移位视图，通过 $K$ 轮向量乘加计算 $T$ 个输出。
通道连续时，更自然的 tile 是 $[T,C_t]$。搬入 $[T+K-1,C_t]$，每个 tap 对应时间方向上一个偏移视图，通道权重按时间广播。根据指令能力，可直接计算，也可先局部转置再计算；需要比较转置成本与后续访问收益。
概念伪代码如下，省略 dtype 与边界细节：
for each owned (batch, time_tile, channel_tile):
local_x = load(current tile and left halo)
local_z = broadcast(bias)
for k in 0 .. K-1:
local_z += local_x[k:k+T, :] * weight[:, k]
local_y = silu(local_z) if enabled else local_z
store(valid part of local_y)
这里广播只说明数学关系，不要求实际复制一个完整权重矩阵。$K$ 很小时，权重可复用，索引也可专门化。
用容量公式约束 tile
以输入低精度存储、FP32 累加为例，单缓冲前向的粗略容量为：
$$ M\approx (T+K-1)C_ts_x+TC_t\cdot4+KC_ts_w+M_{extra}. $$
$M_{extra}$ 包括输出转换、转置 scratch、对齐和激活临时数据。若输入与输出保留多份或使用双缓冲，应逐项增加，不能简单把整个公式乘 2 或假设所有 buffer 都可原地复用。
双缓冲可以尝试在计算当前 tile 时预取下一 tile，但需要验证搬运与计算能否真实重叠，以及 buffer 切换前的同步是否正确。只在代码里使用异步拷贝，不代表时间线已经重叠。
反向带 SiLU 时，输入 buffer 可能需 $(T+2K-2)C_t$，上游梯度需 $(T+K-1)C_t$，还需 FP32 的局部梯度与参数部分和。通常不能直接复用前向的最大 tile 配置。
反向如何分离输出所有权
让每个任务负责唯一的 $dx$ 区域，并读取右侧梯度，可以避免 $dx$ 的跨任务原子加。$dw$ 与 $d\beta$ 则显式做局部归约和全局归约。
若每个通道有 $P$ 个 partial groups，FP32 参数梯度 workspace 约为：
$$ M_{partial}=4PC(K+1)\ \text{bytes}, $$
其中 $+1$ 是有 bias 时的部分和。这个公式假设每个通道恰有 $P$ 组部分结果；若任务覆盖不同通道片段，要按真实布局计算。
两阶段归约的好处是所有权清晰、便于规定累加顺序；代价是额外 kernel 和 workspace IO。原子方案则需要确认支持的 dtype、累加冲突和数值要求。应比较完整反向时间，而不仅是主 kernel 的时间。
怎样构造可信的数学参考
附带的 reference.py 使用 Python 标量浮点实现单 batch、单通道的带状态前反向，不需要 PyTorch 或设备运行时。它采用直接窗口公式与显式梯度累加，便于独立核对优化 kernel 的下标。
python3 docs/ai-infra/causal-convolution/reference.py
脚本验证三件事：第三章手算结果；含 SiLU 和状态输入输出的有限差分；整段执行与逐 token 执行、分两块反传之间的一致性。由于 depthwise 通道独立，可对各 batch、各通道分别调用参考，再沿 batch 汇总参数梯度。
有限差分使用：
$$ \frac{\partial\mathcal L}{\partial\theta_i} \approx\frac{\mathcal L(\theta_i+\epsilon)-\mathcal L(\theta_i-\epsilon)}{2\epsilon}. $$
损失取 $\mathcal L=\langle g,y\rangle+\langle q,F\rangle$。随机 $g,q$ 比对所有输出简单求和更容易暴露索引错误。有限差分在高精度参考上进行，避免低精度舍入吞掉微小扰动。
该脚本检验的是数学关系，不覆盖真实设备的 stride、竞争写和同步，也不代替 Dao 或 MLU kernel 的运行测试。
设备正确性测试矩阵
| 维度 | 建议覆盖 | 主要发现的问题 |
|---|---|---|
| 核宽 | 2、3、4 | 模板分派、权重方向 |
| 长度 | 1、$K-1$、$K$、tile 前后相邻长度、多个 tile | padding、halo、尾块 |
| 通道 | 向量宽度附近、tile 边界、合法非整 tile 通道数 | 通道掩码与越界 |
| batch | 1 与多个 batch | 参数梯度全局归约 |
| 布局 | 时间连续、通道连续、合法带间隙视图 | 把 shape 当 stride 的错误 |
| 激活 | 无激活与 SiLU | 重计算与激活导数 |
| 参数 | 有无 bias、各受支持 dtype 组合 | dtype 分派、空指针分支 |
| 状态 | 零与非零初态、短块、非零 dF | 状态顺序与直接梯度 |
| packed | 边界落在 tile 内部和附近 | 跨样本污染 |
| update | 单 token、多 token、循环回绕、槽位置换 | 原地更新、索引与缓存一致性 |
不满足接口对齐或 dtype 限制的输入，应另外测试是否按契约报错。不能把本来就不支持的形状混入数值比较后误判为 kernel 数学错误。
另一个很有效的测试是因果性扰动：只修改未来输入 $x_{t+1:}$，检查输出 $y_{:t+1}$ 不变。对 packed 数据，修改前一条样本，检查后一条样本输出不变。两者分别检查时间依赖和样本隔离。
误差怎样报告
同时记录最大绝对误差与有稳定分母的相对误差，例如：
$$ e_{abs}=\max_i|a_i-r_i|,\qquad e_{rel}=\max_i\frac{|a_i-r_i|}{\max(|r_i|,\tau)}. $$
$\tau$ 防止参考值接近零时相对误差失真；其取值要与数据尺度一起报告。参数梯度有长归约，不能照搬前向逐元素结果的阈值。可以补充整体范数相对误差，观察误差是普遍舍入还是少数位置的严重偏差。
至少先比较 FP32，再比较 FP16/BF16；分开记录 $y,dx,dw,db,dS,F$。若 FP32 已出现 tile 边界的大误差，优先修复语义或搬运，不要靠放宽低精度阈值掩盖问题。
性能测量必须覆盖完整路径
设备执行通常异步。计时应使用可靠设备事件或在计时边界同步，进行预热，并避免把首次编译和内存分配混入稳定 kernel 延迟。对 update 多次计时时，要明确每次是重置相同状态还是连续推进同一请求。
同时报告三种范围：单个 kernel、完整算子调用、模型局部子图。布局转换、临时 workspace 清零、第二阶段归约和 dtype 转换，都可能让单 kernel 的优势在完整算子中缩小。
比较对象应包括框架参考路径、Dao 同语义路径和 MLU 实现；不同设备之间的绝对延迟同时受硬件影响，不能仅凭延迟比就判定移植算法质量。
用性能模型提出可验证的假设
简化下界可以写为：
$$ t\gtrsim\max\left(\frac{F}{P_{eff}},\frac{D}{BW_{eff}}\right), $$
再考虑不能被隐藏的启动、同步和归约成本。$F$ 是实际工作量，$D$ 是实际搬运字节数，$P_{eff},BW_{eff}$ 是该形状和实现能够达到的有效速率，不应直接等同于峰值规格。
若长序列前向随搬运字节近似线性增长，可检查带宽和重复加载；若很小的工作量延迟几乎不变，可检查启动成本；若只有反向随时间 tile 数明显变慢，可检查原子冲突与部分和开销。每个推断都应通过改变单一因素验证。
第二章示例的理想读写量不包括布局转换。若为了进入某条 kernel 路径额外复制一遍输入，新增的读写可能与原始算子的主 IO 同量级。因此应尽量理解上下游自然布局，再决定是否值得转换。
建议的学习实验
- 使用非对称权重手算一个 $K=3$ 的前向，写出对应带状矩阵，再转置推导 dx。
- 令 SiLU 输入在零附近，检查上游梯度约乘 $1/2$；再用非零输入验证完整导数。
- 将长度为 5 的序列切为 1 与 4 两块，分别比较保留状态梯度和 detach 状态后的结果。
- 用同样逻辑 shape 构造两种 stride，画出相邻地址属于哪个 token、哪个通道。
- 对多个 tile 大小估算 halo 比例与片上容量，再用设备测量验证最优点。
- 分开测 dx 与参数归约的瓶颈，比较原子方案和两阶段方案的完整延迟。
完成这些实验后，再回看 Dao 源码中的共享内存交换、布局转换和归约，能把每个优化对应到明确的数据依赖与成本，而不必记忆某个硬件的线程常数。