第七章 状态 分块与增量推理

用扩展序列统一定义状态

令初始状态 $S$ 存放当前块之前的 $K-1$ 个卷积输入，按最旧到最新排列。把它们拼接成：

$$ u=[S_0,\ldots,S_{K-2},x_0,\ldots,x_{L-1}]. $$

于是当前块的前向和最终状态可以统一写为：

$$ z_t=\beta+\sum_{k=0}^{K-1}w_ku_{t+k},\qquad F=u[L:L+K-1]. $$

这里 $F$ 是长度 $K-1$ 的尾部切片。无初始状态时取 $S=0$。这一表达对 $L<K-1$ 也成立，避免错误地认为最终状态始终完全来自当前输入。Dao 状态参考实现

例如 $K=4,S=[a,b,c],x=[d]$，则当前窗口为 $[a,b,c,d]$，最终状态为 $[b,c,d]$。如果只从当前输入取尾部，不足部分补零，会丢失仍然有效的 $b,c$。

带状态反向的统一推导

假设损失依赖输出 $y$ 和最终状态 $F$。令 $q_r=\partial\mathcal L/\partial F_r$，先对拼接序列 $u$ 求导：

$$ \frac{\partial\mathcal L}{\partial u_j} =\sum_{t=0}^{L-1}\sum_{k=0}^{K-1} \delta_t w_k\mathbf1[j=t+k] +\mathbf1[L\le j<L+K-1]q_{j-L}. $$

第一项来自卷积，第二项来自尾部切片的直接路径。随后切开：

$$ dS=du[0:K-1],\qquad dx=du[K-1:K-1+L]. $$

权重与偏置梯度仍为：

$$ dw_k=\sum_t\delta_tu_{t+k},\qquad d\beta=\sum_t\delta_t, $$

batch 与通道维按前文补全。由于 $F$ 直接由卷积输入构成，$q$ 不会直接贡献到当前块的 $dw$；它通过输入继续传给更早的计算节点。

单 token 分块的检查例子

令 $K=3,S=[a,b],x=[c]$。则：

$$ z=\beta+w_0a+w_1b+w_2c,\qquad F=[b,c]. $$

给定 $\delta$ 和 $q=[q_0,q_1]$：

$$ dS=[w_0\delta,\ w_1\delta+q_0],\qquad dx=[w_2\delta+q_1]. $$

这个例子适合检测短块中状态直接梯度遗漏的问题。

分块训练与截断反向传播

将长序列切成 A、B 两块，令 $S_B=F_A$。要与整段训练等价，B 的 $dS_B$ 必须作为 A 的 $dF_A$ 继续反传。只传前向状态而将状态 detach，前向可以保持一致，梯度却成为截断反向传播。

对于短卷积，跨边界依赖只涉及有限历史；对于完整 Mamba，还要处理 SSM 状态的跨块传播。卷积状态传对，不代表完整模型的分块梯度已经等价。

用于推理的原地 update 没有承担上述 autograd 连接。可微分的分块训练应使用明确支持状态输入输出梯度的路径，不能直接依赖推理缓存原地修改来建立训练图。

单 token update 的数学步骤

维护最小状态 $S_t=[x_{t-K+1},\ldots,x_{t-1}]$，新 token 的特征是 $x_t$：

$$ y_t=\phi\left(\beta+\sum_{k=0}^{K-2}w_kS_{t,k}+w_{K-1}x_t\right). $$

随后删除最旧输入，追加当前输入。注意保存 $x_t$ 而非 $y_t$，否则下一个时间步计算的已不是原始卷积。

Dao update 接口也支持一次输入多个 token，逻辑 shape 为 [B,C] 或 [B,C,L_new]。kernel 对同一通道的这些新 token 顺序推进寄存器窗口，通道和 batch 之间并行。Dao update CUDA

物理缓存可以大于最小历史长度

conv_state 的 state_len 可以大于或等于 $K-1$。非循环模式保留最新的 state_len 个输入，但真正计算卷积时只需最近 $K-1$ 个历史值。有些 Mamba 调用直接分配长度 $K$ 的状态；这不改变算子的感受野。

循环模式通过 cache_seqlens 确定写指针。设缓存长度为 $S_{len}$，已处理长度为 $p$，新输入写入：

$$ \operatorname{slot}{new}=p\bmod S{len}. $$

历史第 $j$ 个位置读取 $(p-j)\bmod S_{len}$。例如 $K=3,S_{len}=4,p=5$，最近两个历史位置是逻辑 3、4，对应槽位 3、0；新输入写槽位 1。

必须先保留计算仍需的旧值，再覆盖缓存。Dao 将窗口保存在寄存器中，随后依次推进。该 kernel 使用传入长度计算索引，但不替调用者推进 cache_seqlens；服务层需要维护下一次调用的长度。update CUDA 状态访问

Continuous batching 的状态索引

请求在服务中不断加入和退出，当前 batch 第 $b$ 行不一定对应状态池第 $b$ 行。conv_state_indices 指定当前请求到缓存槽位的映射。

因此需要区分输入 batch 索引和状态池索引。若多个并行请求错误地映射到同一可写槽位，会产生相互覆盖；这是服务调度和算子调用契约的问题。请求结束后复用槽位，也必须明确清零或初始化所需历史。

Packed sequence 的边界

若将样本 A 和 B 沿时间拼在同一行，B 开头不能读取 A 结尾。seq_idx 的作用是逐窗口屏蔽不同序列标识的贡献。它不是 attention mask 的完整替代，也不会自动替上层决定哪些标签参与 loss。

源码采用标识相等判断。应用应让不同片段在相关窗口内拥有不同标识；若使用 A、B、A 这样复用 ID 的布局，仅比较 ID 不会自动识别中间发生过重置。

对推理状态切换、packed 边界、分块延续这三种情况，业务语义分别是“选择另一请求的历史”“阻止跨样本历史”“继续同一请求历史”。正确实现首先要把这三者区分开。

下一章：MLU 迁移与验证实践。