做长序列建模这行的，最近两年大概都绕不开一个词：Mamba。很多人第一次看到它，注意力都集中在“线性复杂度”“选择性状态空间”这些概念上，但真正决定它能不能在真实硬件上跑起来的，是标题里这半句话——硬件感知算法，并行关联扫描。我最初也只把它当成一个数学技巧，直到自己动手把一个朴素递归版本和并行扫描版本放在同一块卡上对比，才发现差距不是两三倍，而是几十倍，而且这个差距会随着序列长度、batch 大小、状态维度一起放大。并行关联扫描解决的核心问题很具体：状态空间模型在时间维度上天然是串行递归，第 t 步必须等第 t-1 步算完，这种依赖关系直接掐死了 GPU 最擅长的并行性。而关联扫描把“串行递归”重新表达成“满足结合律的区间合并”，从而把一个长度为 L 的串行链条，拆成可以在多个线程、多个 block 上同时推进的分段，再用一次小规模的跨段合并收口。这篇内容我会从问题来源、数学底座、硬件映射、动手实现、线上排查五个层面把它拆开讲，适合已经写过一两个自定义 kernel、想搞清楚“为什么这么设计”的中级读者，也适合只调过现成算子库、想理解内部到底发生了什么的应用侧同学。所有代码和参数都是我在实际调试中验证过的思路，你可以直接拿去改。
1. 先把问题摆上桌：SSM 的串行递归为什么在现代加速器上跑不动
1.1 递归式本身长什么样，为什么会卡住
先把公式摆清楚，后面所有讨论都建立在它上面。连续时间状态空间模型经过零阶保持离散化之后，变成这样一条递推：
h_t = Ā_t · h_{t-1} + B̄_t · x_t
y_t = C_t · h_t + D · x_t
其中 h_t 是维度为 N 的状态向量（常见取值 16、32、64），x_t 是维度为 D 的输入，Ā_t 是 N×N 的状态转移矩阵。为了让它具备“选择性”，也就是让记忆的保留和丢弃依赖当前输入，Ā_t、B̄_t、C_t 都是由输入 x_t 经过线性投影再经过离散化得到的，它们是随 t 变化的。
问题就出在第一行。h_t 依赖 h_{t-1}，h_{t-1} 依赖 h_{t-2}，一条长度为 L 的序列就是一条长度为 L 的依赖链。GPU 的执行模型是 SIMT，一块卡上有成千上万个线程，但只要有严格的先后依赖，同一时刻能真正干活的线程数就只有一个。你可能会想，那我把 D 个通道摊开，每个通道算自己的递归，不就有 D 路并行了？对，这确实能拿到一批并行度，但这是“通道并行”，不是“时间并行”。序列长度从 2K 涨到 32K，通道数不变，总计算量线性增长，墙钟时间也线性增长，并行度一点没变。这就是我们说的“串行诅咒”：序列维度上的工作量完全无法被并行资源摊薄。
我在第一版原型里就是老老实实按这个公式写了个循环，每个时间步做一次 N×N 的矩阵向量乘。跑 L=4096、D=1024、N=16、batch=8，一层要十几毫秒，叠上几十层，一个前向就上秒级。这时候你才会真正意识到，Mamba 的“线性复杂度”是算法复杂度上的说法，和硬件上的实际吞吐是两回事。
1.2 为什么不能像卷积模型那样预先算好核
有人会问：状态空间模型不是和卷积等价吗？时不变的情况下，h_t 可以展开成对历史上所有输入的加权和，权重是一个可以预先算好的卷积核，长度 L 的序列就变成了一次 FFT 卷积，复杂度 O(L log L)，而且高度可并行。这条路在早期的状态空间模型里确实走得通，因为那时候 Ā、B̄、C 是不随输入变化的，卷积核是固定的。
但“选择性”这个东西一旦引入，这条路就被堵死了。因为 Ā_t、B̄_t 依赖 x_t，卷积核不再是固定的，而是变成了一整个 L×L 的下三角矩阵，每一项都由输入决定。你没法预先算核，因为核本身是数据的函数；你也没法用 FFT，因为 FFT 要求卷积核平移不变。剩下能做的，只有沿着时间轴老老实实算。这是 Mamba 在设计上付出的代价，也是并行关联扫描存在的全部理由——既然躲不开时间轴上的递推，那就把这个递推本身改造成可并行的形式。
提示：如果你在复现时发现某个实现可以直接调卷积算子，先确认它的 A、B、C 是不是时不变的。时不变版本可以用卷积加速，选择性版本不行，这是两条完全不同的实现路径，混用会得到数值上完全错误的输出。
1.3 硬件视角：真正的瓶颈往往不是算力，是搬运
把递归写出来之后，第二个坑紧跟着来了：中间状态落不落盘。假设 batch=8，序列长度 L=2048，隐藏维度 D=2048，状态维度 N=16，如果把每个时间步的 h_t 都存下来，大小是 8 × 2048 × 2048 × 16 × 4 字节，算一下大概是 2.1 GB。这还只是一层。几十层叠起来，显存直接爆掉。就算显存放得下，写一遍读一遍意味着 4 GB 以上的显存往返流量，按一块带宽 2 TB/s 的数据中心级加速卡估算，光这一项就是 2 毫秒左右的纯搬运时间，而这一层真正的浮点计算量其实很小。
这就引出了“硬件感知”的第一层含义：算法设计必须围绕内存层次来组织，而不是围绕公式来组织。理想的实现是，把离散化、B̄x 的计算、扫描、C 的乘加、D 的残差、输出投影全部融合进同一个 kernel，中间状态只在寄存器或共享内存里流转，一次都不写回显存。融合之后，这个 kernel 的访存需求退化成读 x、读若干投影权重、写 y，流量从 O(B·L·D·N) 降到 O(B·L·D)，差了整整 N 倍。N=16 时就是 16 倍，这是能不能跑得动的分水岭。
第二层含义是：并行粒度必须和执行单元的物理尺寸对齐。GPU 上一个 warp 是 32 个线程，一个 block 通常是 128 到 1024 个线程，跨 block 通信只能走全局内存或者分布式共享内存，代价高得多。所以扫描算法的分层方式，必须和“warp 内 → block 内 → block 间”这个物理层次一一对应，否则你写出来的“并行算法”在硬件上会被通信开销吃干净。
2. 关联扫描的数学底子：把递归翻译成结合律
2.1 把一步递推抽象成一个“作用”
要并行化，第一步是把“计算一个值”变成“计算一个可组合的对象”。观察 h_t = a_t · h_{t-1} + b_t 这个形式（为了书写方便，这里把矩阵乘先当成标量乘来看，矩阵情形是逐元素的推广），它本质上是作用在 h 上的一个仿射变换：输入 h，输出 a·h + b。一个仿射变换由一对参数 (a, b) 唯一确定。
那么从 h_0 推到 h_t，就是把 t 个这样的仿射变换依次复合起来。关键在于，两个仿射变换的复合，仍然是仿射变换：
设先作用 (a₁, b₁)，再作用 (a₂, b₂)，即
h' = a₁·h + b₁
h'' = a₂·h' + b₂ = a₂·a₁·h + (a₂·b₁ + b₂)
所以复合结果是 (a₁·a₂, a₂·b₁ + b₂)。这个复合运算就是整个并行化的地基。注意参数的顺序：a 分量是直接相乘，与顺序无关；b 分量里出现了 a₂ 乘以 b₁，顺序敏感。写代码时这里是最容易搞反的地方，一旦反了，短序列可能看不出来，长序列会稳定地产生错误输出，而且错误是“平滑”的，不会报 NaN，特别难查。
复合运算满足结合律：((a₁,b₁) ∘ (a₂,b₂)) ∘ (a₃,b₃) 和 (a₁,b₁) ∘ ((a₂,b₂) ∘ (a₃,b₃)) 结果相同。这一点可以直接展开验证，也可以从“仿射变换的复合本来就满足结合律”这个角度理解——矩阵乘法满足结合律，仿射变换是矩阵乘加平移，同样满足。结合律意味着括号可以随便挪，而括号的位置决定了计算的先后顺序，也就决定了哪些部分可以并行。
2.2 从结合律到并行前缀扫描
有了结合律，长度为 L 的序列前缀计算就变成了一个标准的并行前缀扫描问题。最直观的版本是 Hillis-Steele 算法：第 d 轮，每个位置 i 把自己和位置 i-d 的结果合并，d 依次取 1、2、4、8……经过 ⌈log₂L⌉ 轮，每个位置都拿到了从开头到自己的完整复合结果。
这个算法在 GPU 上特别讨喜，因为它每一轮的操作对所有元素完全一致，没有分支，没有依赖分歧，天然适合映射到一个线程处理一个元素、用 shuffle 在 warp 内交换数据。代价是总工作量从 O(L) 涨到了 O(L log L)，多做了 log L 倍的无用功。放在 CPU 上这是亏的，放在 GPU 上却常常是赚的，因为 GPU 有的是闲置算力，缺的是并行度。这就是典型的“用算力换并行度”的取舍，也是硬件感知设计里最基本的一条判断：在算力过剩、并行度不足的场景下，宁可多算，也不要串行。
还有一类 work-efficient 的扫描算法，总工作量是 O(L)，但需要 up-sweep 和 down-sweep 两个阶段，同步次数更多，访存模式更绕。在 CPU 上它更优，在 GPU 上它常常跑不过 Hillis-Steele，因为每一轮都要过一遍共享内存，而共享内存的带宽和延迟都是有限的。我自己实测下来，在 warp 这个尺度（32 个元素）上，Hillis-Steele 用 shuffle 实现只要 5 轮，几乎是最优解；在 block 尺度上，则更适合“每个线程先串行处理若干个元素，再在 block 内做一次小规模扫描”，把线程内的串行当作免费的并行度压缩手段。
2.3 分段、重置与选择性带来的额外麻烦
真实的 batch 里序列长度不一致，padding 的部分不能被算进状态里。这时候需要“分段扫描”（segmented scan）：在 padding 的边界处，把复合算子替换成单位元 (1, 0)，这样前缀结果在边界处会被重置，等价于从零开始。这个技巧很实用，但它要求你在扫描前先构造好一个“段 ID”或者“有效掩码”，并且在 combine 函数里根据掩码决定是正常合并还是重置。掩码判断会带来分支分歧，所以更好的做法是把掩码预先乘进 a 和 b 里：无效位置令 a=1、b=0，这样合并时自然就是恒等，不需要分支。这个改写看起来不起眼，但在长序列、小 batch、padding 比例高的场景下能省掉可观的开销。
另一个麻烦是方向。前向是沿时间正序扫描，反向传播里对状态的梯度递推同样是线性递推，但方向相反。很多人第一次写反向 kernel 时会想当然地复用前向的 combine 函数，结果梯度全错。正确做法是：反向的递推形式是 dh_t = a_{t+1} · dh_{t+1} + grad_t，参数顺序和正向前缀扫描不同，需要单独写一个 combine，或者干脆把序列翻转后复用同一个函数。我倾向后者，代码量少一半，出错概率也低，代价是多一次翻转的访存——如果翻转能在读入时用索引直接倒着读，这个代价就接近零。
注意：判断 combine 顺序对不对，有个很便宜的验证方法。构造一条 a 全是 0.5、b 全是 1 的短序列（长度 8 就够），手算出理论值，再和你的 kernel 输出逐位对比。这个测试能在 30 秒内暴露 90% 的顺序错误，比拿真实数据跑 loss 看曲线有效得多。
2.4 数值稳定性：那些不会报错但会让模型悄悄变差的地方
关联扫描在数值上有一个特点：a 分量是连乘的。在状态空间模型里，Ā = exp(Δ·A)，而 A 被参数化成负值（通常写成 -exp(A_log)，保证恒为负），Δ 经过 softplus 之后为正，所以 Δ·A 恒为负，Ā 落在 (0, 1) 区间。这意味着一路连乘下去，a 会指数衰减到很小的值。这本身是设计意图——状态空间模型的“记忆”本来就是衰减的，远的输入影响小。但在浮点表示下，连乘几十上百次之后，a 会 underflow 到 0，此时 b 分量还在正常累加，结果就是“早期信息完全丢失，后期信息正常保留”。
这个行为在数值上是稳定的（不会爆 NaN），但在语义上可能不是你想要的。如果你的任务需要更长的有效记忆，需要调的不是扫描算法，而是 A 的参数化范围和 Δ 的取值区间（常见做法是给 Δ 加一个上界和一个下界，避免它被压到极小或者顶到极大）。反过来，如果 a 偶尔大于 1（比如 A 的参数化没约束好），连乘会指数发散，很快就溢出成 inf，这时候输出会先出现大数值再变 NaN，比较容易发现。
另一处容易被忽略的是累积精度。中间的前缀结果应该用 fp32 累积，哪怕输入输出是 bf16。原因是：a 的连乘在数值范围上跨度极大，bf16 只有 8 位指数和 7 位尾数，连乘几次就丢精度了；而 b 的累加在长序列上是典型的“大数吃小数”场景。把 h 的累积路径固定为 fp32，只在最后输出时转回 bf16，代价是寄存器占用翻倍，收益是长序列下的 loss 曲线明显更平。这个取舍在短序列上几乎看不出差别，所以很容易被漏掉——等到训练 32K 长度的任务时才突然发现收敛不了，回头查又是一轮。
3. 硬件感知落地：把扫描映射到真实的并行执行模型
3.1 三层分块：把逻辑并行度对齐到物理层次
一个能在真实硬件上跑得动的扫描实现，通常是三层结构，这套结构几乎是所有高性能扫描实现的标准套路：
第一层是线程内串行。每个线程负责一段连续的时间步，比如 8 个或 16 个，在这段内老老实实串行递推，结果留在寄存器里。这一步没有任何通信开销，而且因为时间连续，读内存天然是合并访问的。它把有效并行度需求压缩了 8 到 16 倍。
第二层是warp 内扫描。一个 warp 的 32 个线程各自算完自己那一段的局部复合结果（也就是这一段的 (a_prod, b_total)），然后用 shuffle 指令做 5 轮 Hillis-Steele，把前缀复合结果分发回每个线程。shuffle 走的是寄存器交换网络，延迟只有几十个周期，是 warp 内通信的最优手段。这里不需要共享内存，也不需要__syncthreads()。
第三层是block 间或 block 内扫描。如果一次性处理的序列长度超过一个 warp 能覆盖的范围（32 × 16 = 512 个时间步），就需要跨 warp 合并。两种做法：一是在 block 内用共享内存做一次小规模扫描，二是把每个 block 的局部结果写到一个小的全局缓冲区，启第二个 kernel 做跨 block 扫描，再启第三个 kernel 把结果回填。后者多两次 kernel 启动，但能让单个 block 的寄存器压力降下来，占用率更高。
选择哪种，取决于你的序列长度和状态维度。我的经验是：序列长度在 4K 以内、状态维度不大时，单 block 内搞定最省事；序列长度上万、batch 又小的时候，分 kernel 做跨 block 合并反而更快，因为单 block 那点并行度填不满整块卡。
3.2 内存布局：那个决定你是快 3 倍还是慢 3 倍的细节
扫描 kernel 的输入张量通常有几个维度：batch、序列长度、通道（隐藏维度）、状态维度。它们在内存里的排列顺序，直接决定了访存是不是合并的。GPU 的全局内存访问以 32 字节或 128 字节的事务为单位，同一个 warp 的 32 个线程如果访问的是地址连续的一段，一次事务就能覆盖，效率最高；如果访问的是跨步的地址，就会触发多次事务，有效带宽掉到几分之一。
对扫描来说，最自然的并行方向是“通道 + 状态”这两个维度，因为它们在时间维度上是独立的。所以内存布局应该让这两个维度的数据在地址上连续，让“时间”成为变化最慢的那一维。如果你的张量是按 (batch, length, channel) 排的，那么在某个固定的时间步上取不同 channel，地址是连续的，这是好的；如果你按 (batch, channel, length) 排，那么在固定时间步上取不同 channel，地址间隔是 length，就变成了最坏情况。
还有一个更隐蔽的问题：状态维度 N 常常是 16 或者 32，正好卡在“不够一个 warp”的位置。如果让每个线程负责一个状态元素，那么只有 16 或 32 个线程是活跃的，剩下的一半 warp 在空转。更好的做法是让每个线程负责多个 (channel, state) 组合，把线程填满；或者干脆把时间步也塞进线程维度，让一个 warp 里的 32 个线程覆盖“几个 channel × 几个时间步”的组合。
3.3 融合边界怎么划：多融一步，少一次往返
融合的收益是明确的，但融合到什么程度需要权衡。一个完整的融合 kernel 大致要完成这些事：读入 x，做三组线性投影得到 Δ、B、C 的原始值；对 Δ 做 softplus 和上下界裁剪；对 A 做参数化（取负指数）；计算离散化的 Ā 和 B̄；算出 b 分量 B̄·x；执行扫描；用 C 乘上状态得到输出；加上 D 的残差项；最后写出。
全融在一起的好处是中间张量一个都不用落盘，坏处是寄存器压力大、kernel 复杂度高、一旦某一步出错很难定位。我通常的做法是先写一个“半融合”版本：投影在外面做，把 Δ、B、C 作为输入传进来，kernel 内部只做离散化 + 扫描 + 输出。这个版本已经能拿到大部分带宽收益，而且好调试。等它稳定了，再把投影融进去，用 profiler 对比前后两版的耗时，确认融合确实带来了收益再保留。
这里有个反直觉的点值得提一下：融合不一定总是更快。当状态维度较大（比如 N=64 甚至 128）时，融合 kernel 的寄存器需求会暴涨，占用率掉下来，反而跑不过分离版本。判断标准很朴素：用占用率工具看一眼每个 SM 上能同时驻留多少个 block，如果掉到个位数，基本就该考虑拆分了。
3.4 反向传播：重计算换带宽这笔账怎么算
训练时反向传播需要前向的中间状态。如果全存下来，显存占用是 O(B·L·D·N)，前面算过，一层就是 GB 级，不可接受。所以实际实现都用重计算：前向只保存每个 chunk 的边界状态（也就是每个分段的 (a_prod, b_total)），反向时根据这些边界状态重新算 chunk 内的中间值。
这里有个粒度选择。chunk 太小，边界状态多，保存和读取的开销上来了；chunk 太大，重计算时串行的部分变长，反向的并行度下降。以我调试过的配置为例，序列长度 4096、每个线程串行处理 16 个时间步、一个 block 覆盖 512 个时间步时，chunk 取 512 是比较自然的：边界状态数量是 8 个每 (batch, channel)，占用可以忽略，重计算时每个 block 内部还能保持完整的并行度。
算一笔具体的账：如果不做重计算而是全存，每层多出 2.1 GB 的读写，按 2 TB/s 带宽算往返约 2 毫秒；如果做重计算，多出来的是一次前向的算力开销，而这次前向本身是访存受限的，代价通常小于 2 毫秒。所以重计算在这个场景下是划算的。但如果你的状态维度很小（比如 N=8），中间张量本身就小，全存的代价可能反而更低，这时候就不必强上重计算。
4. 动手实现：从朴素递推到分块并行 kernel
4.1 先写一个能对拍的朴素版本
任何并行实现之前，先写一个最笨的串行版本作为基准。它的作用不是跑得快，是给你一个数值上的真值参照。这一步千万别省，我见过太多人直接上手写 kernel，然后花两天时间在一个拼写错误上。
import torch def naive_scan(a, b, h0=None): # a, b: (batch, length, channel)，逐元素递推 # 返回所有时间步的状态 (batch, length, channel) batch, length, channel = a.shape h = torch.zeros(batch, channel, dtype=a.dtype, device=a.device) if h0 is None else h0 outs = [] for t in range(length): h = a[:, t] * h + b[:, t] outs.append(h) return torch.stack(outs, dim=1)写完立刻做两件事。第一，用 a 全 0.5、b 全 1 的构造数据手算对拍；第二，用真实分布的数据（a 在 0 到 1 之间、b 是标准正态）跑一遍，检查有没有 inf 或 nan。这两步加起来不到十行代码，能挡掉后面一大半的调试时间。
4.2 用高层 DSL 快速验证分块思路
在动手写底层 kernel 之前，可以用高层 kernel DSL 快速验证分块逻辑。这类 DSL 通常自带关联扫描原语，你只需要把 combine 函数定义对，分块、warp 内扫描、跨 warp 合并都由它代劳。这能让你在半小时内确认“算法思路是对的”，把剩下的时间花在性能调优上。
import triton import triton.language as tl @triton.jit def _combine(a1, b1, a2, b2): # 先作用 (a1,b1)，再作用 (a2,b2) return a1 * a2, a2 * b1 + b2 @triton.jit def chunk_scan_kernel(a_ptr, b_ptr, h0_ptr, out_ptr, L, CHUNK: tl.constexpr): pid = tl.program_id(0) offs = pid * CHUNK + tl.arange(0, CHUNK) mask = offs < L a = tl.load(a_ptr + offs, mask=mask, other=1.0).to(tl.float32) b = tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32) h0 = tl.load(h0_ptr).to(tl.float32) # 块内关联扫描，得到从块首到当前位置的复合结果 a_pref, b_pref = tl.associative_scan((a, b), axis=0, combine_fn=_combine) # 把块首的初始状态代进去 h = a_pref * h0 + b_pref tl.store(out_ptr + offs, h.to(out_ptr.dtype.element_ty), mask=mask)注意三个细节。第一，other=1.0和other=0.0对应单位元，这是 padding 位置自动重置的关键，写错了 padding 会污染真实位置的状态。第二，a和b都显式转成 fp32，扫描全程用 fp32 累积。第三，这个 kernel 只处理块内扫描，块首的初始状态h0是外部传进来的，跨块的递推需要额外一层。
跨块那层最简单粗暴的做法是：先用一个 kernel 算出每个块的 (a_prod, b_total)，这些值只有L / CHUNK个，用一个单 block 的小扫描串起来得到每个块的起始状态，再启第二个 kernel 做块内扫描并代入。块数通常只有几十个，这层开销可以忽略，但它把整个 kernel 的寄存器压力降到了最低。
如果你打算手写底层实现，warp 内那部分的核心就是一个 shuffle 循环，写出来大概是这样：
// 每个线程持有一段的局部复合结果 (a, b)，先做 warp 内前缀扫描 for (int d = 1; d < 32; d <<= 1) { float a_up = __shfl_up_sync(0xffffffff, a, d); float b_up = __shfl_up_sync(0xffffffff, b, d); if (lane >= d) { b = a * b_up + b; // 必须先算 b，因为它用的是旧的 a a = a * a_up; } }这两行的顺序是硬性要求：b的更新用到了当前的a，如果先更新a，b就用到了已经被覆盖的值，结果全错。而且这个错误很隐蔽——因为 a 和 b 的关系在数值上仍然“合理”，输出不会变成 NaN，只会和真值有稳定的偏差。
4.3 分块大小怎么定：一次带数字的推演
CHUNK这个参数没有万能值，但推演过程是可以复用的。假设序列长度 L=4096，隐藏维度 D=2048，状态维度 N=16，batch=4，用一块带宽大约 2 TB/s 的数据中心加速卡。
先看边界状态的存储代价。每个块的边界状态大小是 batch × D × N 个 fp32，也就是 4 × 2048 × 16 × 4 字节 = 512 KB。这是每个块一份，块数等于 L / CHUNK。
- CHUNK=256 时，块数 16，边界状态总计 8 MB；
- CHUNK=512 时，块数 8，总计 4 MB；
- CHUNK=1024 时，块数 4，总计 2 MB。
单看这一项，块越大越省。但块内扫描的并行度是 CHUNK 除以每线程串行步数。如果每线程串行 16 步，CHUNK=256 时需要 16 个线程，连半个 warp 都填不满；CHUNK=512 时需要 32 个线程，正好一个 warp；CHUNK=1024 时需要 64 个线程，两个 warp，需要额外的跨 warp 合并。
所以在这个配置下，CHUNK=512 是一个合理的落点：正好一个 warp 满负荷，不需要跨 warp 同步，边界状态总量 4 MB 完全可接受。如果序列长度拉到 16384，块数变成 32，边界状态涨到 16 MB，仍然不算大；但如果同时把 batch 拉到 32，就是 128 MB，这时候就该考虑把边界状态也做成分层（块内再分块），避免一次性物化。
还有一个隐含约束：CHUNK最好是 32 的整数倍，因为 warp 是 32 个线程。设成 48 或 96 这种数，会让最后一个 warp 有半个空转，白白浪费。这个约束在纸面上看不出来，但在 profiler 里会明确显示成低占用率。
提示：调
CHUNK的时候别只看单层耗时，要看端到端。块太小会让 kernel 启动次数变多，块太大会让占用率下降，两个效应方向相反，最优点往往在中间。用一个脚本扫 128、256、512、1024 四档，跑三次取中位数，通常二十几分钟就能定下来。
4.4 怎么确认自己写对了，以及怎么读性能数据
验证分三步走，一步都不能省。第一步是短序列逐位对拍：长度取 8 到 64，和朴素版本比较最大绝对误差，fp32 下应该在 1e-5 量级。第二步是长序列趋势检查：长度从 256 扫到 8192，看误差是否随长度增大而单调增长，如果增长得特别快（比如到 1e-1 量级），八成是累积精度的问题，回去检查中间状态是不是真的用了 fp32。第三步是端到端对齐：把扫描层替换进完整模型，用同一份权重、同一份输入，和参考实现比较输出的相对误差，同时比较 loss 在前几十步的下降曲线——如果曲线形状一致但整体偏移，通常是初始化或者某个残差项的细节没对上。
性能数据这块，我一般只看三个数：有效带宽占峰值的比例、每个 SM 上驻留的 block 数、以及 kernel 的实际耗时和纯访存理论耗时的比值。第一个数告诉你是不是访存受限，第二个数告诉你是不是寄存器压力过大，第三个数告诉你还有多少优化空间。如果第三个比值已经接近 1，说明这个 kernel 已经贴着带宽跑了，再怎么调算法都没用，该去调数据布局或者减少数据量了。我踩过最坑的一次是花了两天优化计算部分，最后发现带宽利用率已经 92%，根本不是我优化函数的瓶颈。
5. 线上排查：那些只在长序列和大 batch 下才冒出来的问题
5.1 短序列全部正常，长序列开始漂移
这是最常见也最折磨人的一类问题。训练一开始 loss 正常下降，跑到几千步之后突然开始抖动，最后发散。把输入长度从 8192 砍到 512 就又正常了。这种情况九成不是算法错，是数值范围失控。
排查顺序我固定成三件事。第一，dump 出 a 的实际分布，看最大值和最小值。如果最大值超过 1，说明 A 的参数化或者 Δ 的上界出了问题，连乘会发散；如果最小值小到 1e-30 以下，说明衰减过强，早期信息全丢了。第二，检查中间状态是不是真的用了 fp32，很多人写的时候想的是 fp32，实际上因为输入张量是 bf16，在某一步被隐式转换回去了。第三，检查 Δ 的裁剪区间，softplus 之后如果没有上下界，遇到异常大的输入会直接把 Δ 顶到几百，Ā 变成精确的 0，整个状态被强制清零。
这里补一个容易忽略的点：梯度路径上同样有连乘。反向递推里的 a 乘的是 ∏ a_t 的雅可比，如果所有 a 都远小于 1，梯度会指数消失，表现为“前面几层几乎不更新”。这不是 bug，是这类模型的固有特性，解决办法是在 A 的参数化和 Δ 的初始化上做文章，而不是去改扫描算法。
5.2 一张表把常见症状和根因对上号
下面这张表是我自己整理的问题速查表，遇到问题时先按症状对号，能省掉大量盲猜时间。
| 症状 | 最可能的根因 | 排查动作 |
|---|---|---|
| 输出和前向逐位对比时误差极大但无 NaN | combine 函数的参数顺序反了，或other单位元填错 | 用 a=0.5、b=1 的短序列手算对拍 |
| 长序列尾部出现 inf 或 NaN | a 连乘发散，A 未约束为负或 Δ 无上界 | dump a 的最大值，检查 A 的参数化方式 |
| 长序列头部信息完全丢失 | a 连乘 underflow，衰减过强 | 检查 a 的最小值，放宽 Δ 的下界 |
| 训练收敛但明显不如参考实现 | 中间状态累积精度不足（用了 bf16） | 强制扫描路径 fp32，只在输出转回 |
| padding 位置污染了真实位置的状态 | 单位元没设成 (1, 0)，或掩码没生效 | 构造含 padding 的 batch 对比单条推理 |
| 反向梯度整体偏移一个常数 | 反向 combine 顺序与前向不一致 | 翻转序列复用前向 combine 验证 |
| 短序列快、长序列慢得不成比例 | 边界状态没有分块，退化成全局串行 | 检查 chunk 划分是否真的生效 |
| 单层快、整模型慢 | 每个 kernel 都在物化中间张量 | 用 profiler 看 kernel 数量和访存量 |
| 占用率极低、耗时远高于理论值 | 寄存器压力过大，融合过度 | 拆分投影部分，或减小 chunk |
| 多流并发时结果偶发错误 | 共享内存缺同步，或存在写后读竞争 | 加__syncthreads()后复测 |
| 同一份输入两次运行结果不同 | 扫描的并行归约顺序不确定，浮点非结合 | 固定归约顺序，或接受误差并验证量级 |
表格里最后一条值得单独说两句。浮点加法不满足结合律，所以不同并行顺序得到的和会有微小差异。这不是 bug，是浮点运算的本性。但如果这个差异大到影响复现（比如同一个 seed 两次跑出来的 loss 差 0.01 以上），说明你的累积路径精度不够，需要提高中间精度或者固定归约树的结构。
5.3 几个只有实际调过才会知道的坑
第一个坑是 warp 内的串行步数和寄存器数量的关系。每个线程串行处理 k 个时间步，就要在寄存器里暂存 k 个中间结果，再加上 a、b、h 以及各种临时变量，很容易顶到寄存器上限。一旦超出，编译器会把寄存器溢出到本地内存，而本地内存在物理上是显存，延迟比寄存器高两个数量级。表现就是“代码看着没问题，但慢了五倍”。这时候要么减少 k，要么减少融合的范围，让编译器喘口气。
第二个坑是 shuffle 的活跃掩码。用__shfl_up_sync时如果传的掩码和实际活跃线程不一致，在某些执行路径下会得到未定义的值。这类问题在单元测试里通常不出现，因为测试用的 block 大小刚好都是满 warp；一旦真实数据让最后一个 block 有线程退出，问题就冒出来了。稳妥的写法是统一传全 1 掩码，并用条件判断控制后续合并，而不是依赖部分掩码。
第三个坑和序列长度对齐有关。当 L 不是 CHUNK 的整数倍时，最后一个块是不满的。如果掩码处理只是简单地跳过越界位置，那块的边界状态会算错，进而污染后面所有块。正确做法是在扫描之前就把越界位置的 a 设成 1、b 设成 0，让它们在数学上成为恒等元，这样边界状态永远正确。
第四个坑是关于 profiler 的解读。看 kernel 耗时的时候要区分“这个 kernel 慢”和“这个 kernel 在等数据”。如果前一个 kernel 写出的数据还没落盘，后一个 kernel 读的时候就会等。用同一块卡上跑的多个 kernel 之间，这种隐式依赖非常常见。判断方法是看时间线上有没有空隙，有空隙说明是调度或者依赖等待，不是计算慢。
我在实际项目里踩得最深的一次，是把扫描和输出投影融进同一个 kernel 之后，单层耗时反而涨了 40%。profiler 显示寄存器用量从前一版的 40 多个涨到了 90 多个，占用率从 50% 掉到 12%。拆回去之后立刻恢复。这件事让我改了一个习惯：每次融合之前先估算寄存器预算，融合之后第一件事就是看占用率，而不是看总耗时。因为总耗时受波动影响大，占用率是个稳定信号，占用率掉了，优化方向基本就是错的。
还有一个体会上，关于 chunk 大小和 batch 大小的联动。小 batch 的时候，块内并行度填不满，你会倾向于把 CHUNK 调大来榨并行度；但 CHUNK 一大，边界状态就多，又反过来压占用率。真正稳的做法是让 CHUNK 和 batch 的乘积大致等于“SM 数量 × 每 SM 驻留 block 数”这个容量，实在算不准就写个脚本扫一遍。我一般会扫三档：偏小（256）、居中（512）、偏大（1024），在真实数据上跑，不看单层看端到端，选那个端到端最稳的，而不是单层最快的。毕竟最终要交付的是整个模型，不是一个漂亮的 kernel 耗时数字。
