序列建模做到一定规模,谁都绕不开一个老问题:想让模型记住更长的上下文,又不想让显存和时延跟着一起爆炸。Mamba 这套状态空间模型之所以在工程圈里被反复讨论,很大一部分原因不在于它的公式多优雅,而在于它配套的那套硬件感知算法真正把理论上的并行性落到了硅片上。并行关联扫描就是其中最核心的一环——它把一个天生串行的递推过程,拆成了可以同时开工的块,再拼回正确结果。我第一次读这段实现的时候,最直观的感受是:这份代码不是在写数学,是在写"数据放在哪、什么时候搬、搬多少"。
这篇内容面向的是想真正读懂、甚至自己动手复现这套扫描逻辑的工程师。我会从动机讲到原理,从 CUDA 层面讲到一个能跑的最小实现,再把调试过程中踩过的坑摊开说。不需要你有状态空间模型的背景,但最好知道 GPU 的共享内存和全局内存大概是个什么速度量级。
1. 硬件感知的由来:这个算法到底在解决什么麻烦
1.1 序列建模的两难:并行的宽度与记忆的长度
循环神经网络的问题大家都清楚:时间步之间存在硬依赖,第 t 步必须等第 t-1 步算完。哪怕一个时间步的计算量很小,整条时间线也只能一个挨一个地走。GPU 上几千个核心闲着,真正干活的只有那么几个,利用率惨不忍睹。这就是所谓的串行深度问题——不是算力不够,是依赖链把并行度锁死了。
注意力机制走的是另一条路:它把任意两个位置直接连起来,训练时可以一次性算完整个矩阵,并行度拉满。代价是计算量和显存对序列长度呈平方增长,序列翻一倍,开销翻四倍。处理几千个 token 还能接受,再长就有点顶不住了。
状态空间模型想两头讨好。它的递推形式和 RNN 一样简洁:
$$h_t = \bar{A}t , h{t-1} + \bar{B}_t , x_t, \qquad y_t = C_t , h_t$$
形式上看是串行的,但关键在于——这个递推是线性的,而且系数结构规整。线性这个性质意味着它可以被"重新结合",而"可重新结合"就意味着可以并行。这一句话就是并行关联扫描存在的全部理由。
我个人的理解是,Mamba 这套设计真正的巧思不在于发明了新算子,而在于它认清了"哪些依赖是真的、哪些是假象"。时间步之间的依赖是真的,但这个依赖的形式足够简单,可以被折叠、被分块、被批处理。认清这一点之后,剩下的工作就变成了纯工程:怎么在有限的片上存储里把折叠做干净。
1.2 关联扫描的数学地基:结合律
扫描(scan)这个词在并行计算里是老面孔了,前缀和就是最经典的例子。给定一个序列 x₁…x_n,想要得到每个位置的前缀和 s_t = x₁ + … + x_t。串行做法是扫一遍,O(n) 时间、O(n) 深度。但加法满足结合律,所以可以两两相加、四四相加,用一棵树把深度压到 O(log n)。
我们的递推长得不像求和,它每一项还要乘以一个系数。把单步状态转移写成一个二元组 $(a_t, b_t)$,含义是"输入状态 h,输出 $a_t h + b_t$"。那么连续两步的复合是:
$$(a_2, b_2) \circ (a_1, b_1) = (a_2 a_1,; a_2 b_1 + b_2)$$
自己验算一下:先做第一步得 $h' = a_1 h + b_1$,再做第二步得 $h'' = a_2 h' + b_2 = a_2 a_1 h + a_2 b_1 + b_2$。对得上。
这个复合运算满足结合律:$((c \circ b) \circ a) = (c \circ (b \circ a))$。而且它有单位元 $(1, 0)$——作用在任意状态上不改变它。有结合律、有单位元,这就构成了一个幺半群,而幺半群上的归约就可以做成任意形状的并行树。这是并行关联扫描能成立的充要条件。
注意:结合律在数学上成立,不代表在浮点数上完全成立。加法不满足严格的结合律,不同的归约顺序会产生不同的舍入误差。这在短序列上无所谓,序列一长、乘子一堆,误差会累积。后面第 4 章会专门讲这个问题怎么处理。
对应到具体的状态空间模型里,单步系数是 $a_t = \exp(\Delta_t A)$,$b_t = \Delta_t B_t x_t$。$\Delta_t$ 是离散步长,一般用 softplus 从输入投影出来,保证为正。$A$ 通常初始化为负实数(或者负对角矩阵),这样 $a_t \in (0,1)$,递推天然稳定。
1.3 硬件感知感知的是什么:带宽、层级与并行粒度
"硬件感知"这个词容易被讲得玄乎。说白了它就是三个具体的东西:
第一是存储层级的速度差。寄存器最快,共享内存次之,L2 缓存再次,全局显存最慢,差出两三个数量级。同样一份数据,放在共享内存里反复用和每次回全局显存取,性能能差十倍以上。所以算法的结构必须围绕"这块数据在生命周期内待在哪一级"来设计。
第二是算术强度。这个指标衡量每搬运一个字节能换来多少次浮点运算。GPU 的峰值算力和峰值带宽之间有个拐点,算术强度低于拐点的算子必然是带宽受限的,再怎么优化计算也没用,只能想办法少搬数据。状态空间扫描恰好是一个典型的低算术强度算子——每个元素就两三次乘加,数据量却很大。
第三是并行粒度。GPU 的调度单位是线程束(warp,32 个线程)和线程块。块内的线程可以通过共享内存和同步原语协作,块之间只能通过全局内存和多次 kernel 启动来通信。这个约束直接决定了分块策略:块内做扫描,块间做一轮小规模的层级合并。
把这三条摆在一起,你就明白为什么论文里要花那么大篇幅讲 kernel 怎么写了。数学上"可以并行"是一回事,硬件上"并行得起来且不浪费带宽"完全是另一回事。
2. 并行关联扫描的原理拆解
2.1 从递推到扫描算子:把状态转移做成可组合的二元组
我们先把抽象层搭好。整个扫描的输入是一串二元组 $(a_1,b_1), (a_2,b_2), \ldots, (a_L,b_L)$,输出是每个位置的前缀复合结果 $(A_t, B_t)$,满足 $h_t = A_t h_0 + B_t$。其中 $h_0$ 是整个序列的初始状态。
有两种经典的并行扫描形态:
Hillis-Steele 扫描,也叫包含式扫描。每一轮把步长翻倍,用前一位置的结果更新当前位置:
for d in [1, 2, 4, 8, ...]: 并行地对每个 i >= d: (a[i], b[i]) = (a[i], b[i]) ∘ (a[i-d], b[i-d])$L$ 个元素只需要 $\lceil \log_2 L \rceil$ 轮,深度极小,实现也简单。缺点是总工作量是 $O(L \log L)$,比串行多做了不少冗余计算。
Blelloch 扫描,也就是工作高效的树形扫描。先做上扫(reduction)建树,再做下扫(down-sweep)分发前缀。总工作量 $O(L)$,深度 $O(\log L)$。理论最优,但访存模式在树的不同层之间来回跳,缓存不友好,实现起来也绕。
实际工程里很少直接用这两种中的一个,而是取个折中:块内用 Hillis-Steele 或者串行累加,块间用一轮树形合并。块的大小通常在几百到几千之间,块内的 $\log L$ 因子被压得很小,冗余计算的代价可以接受;块间只有几十到几百个元素,用一次小规模扫描就解决了。
这里有个容易被忽略的点:因为算子有单位元 $(1,0)$,我们可以给序列前面补上若干个单位元,把长度补到 2 的幂,这样扫描树的形状规整,不用处理边界。补位成本很低,因为单位元的复合是退化操作。
2.2 分块与两趟传递:GPU 上的落地形态
把上面的思路映射到 GPU,自然会得到这样的结构:
序列沿长度方向切成 $K$ 个块,每块长度 $C$(chunk size)。所有块同时启动,各自独立地做块内扫描。这时候每个块算出来的前缀都是相对于块首的局部量,缺一个正确的初始状态。于是需要第二趟:把所有块的总转移 $(A^{(k)}, B^{(k)})$ 收集起来,做一次长度为 $K$ 的扫描,得到每块的正确初始状态 $h^{(k)}_{init}$,再回填修正。
这就是标准的两趟方案。它的优点是逻辑清晰、块内并行度拉满;缺点是中间状态要多存一份,而且需要两次 kernel 启动或者一次 kernel 内的网格级同步。
另一条路是单趟流水线:第 $k$ 个块在等待前一块传状态的同时,先把自己的局部扫描算完,状态一到就立刻做一次修正。这种方案把等待和计算重叠起来,但实现复杂度陡增,而且块之间的传递仍然是串行的——第 $K$ 块必须等前面所有块都完成。序列特别长的时候,尾部延迟会很显眼。
我在实际项目里的取舍是:序列长度在几千以内,两趟方案足够好,实现清爽、调试方便;序列上到几万甚至更长,就得考虑把块间扫描也做成树形,或者干脆用分层的两级分块。别小看这个决策,它直接决定了你的 kernel 是"跑得动"还是"跑得快"。
还有第三种做法,也是很多高效实现采用的:沿通道维并行,时间步串行。因为状态空间模型的隐藏状态维度是 $d_{state} \times d_{inner}$,每个通道的递推是独立的。一个线程块负责一批通道,在块内用共享内存缓存 $h$,然后把时间步循环放在块内串行执行。这种做法的并行度来自通道数(通常是几百到几千),不需要块间通信,实现最简单,对小模型小批量场景非常友好。缺点是序列方向的并行度是零,序列一长就顶不住。
两种思路各有适用面,很多实现会根据序列长度动态切换。这不是投机,是权衡。
2.3 选择性带来的变化:参数随输入走,但结合律不变
早期状态空间模型的 $A$、$B$、$C$ 是时不变的,全序列共用一套参数。这意味着可以把整个递推解析地展开成一个长卷积,用 FFT 加速,训练时完全不需要扫描。这也是为什么老一代模型在效率上有优势。
Mamba 的"选择性"(selective)改掉了这一点:$\Delta_t$、$B_t$、$C_t$ 都由当前输入 $x_t$ 线性投影得到,变成时变的。好处是模型能根据内容动态决定"记住什么、忘掉什么",表达能力上去了。坏处是长卷积那套技巧失效了——系数每个位置都不一样,没法预先算好卷积核。
但注意,变的是系数,不变的是结构。单步转移依然是一阶线性的,复合运算依然满足结合律。所以并行扫描照用不误,只是每一步都得先把 $a_t$、$b_t$ 现算出来。这也是为什么现代实现里扫描 kernel 和前面的投影、离散化是融合在一起的——分开算的话,中间那堆 $(a_t, b_t)$ 要写回显存再读出来,带宽全浪费在搬运上了。
我的经验是:判断一个序列模型能不能用并行扫描,不要看它的参数是不是时变的,要看它的状态更新是不是"一阶线性 + 可结合"。二阶、非线性、带门控交叉的,统统不行。
3. 手写内核:从数据布局到共享内存的实操
3.1 数据布局与线程映射
先说布局,这个决定了后面所有访存能不能合并。
输入张量通常是(batch, seq_len, d_inner)的连续布局。扫描需要按(batch, d_inner)分组,沿seq_len方向推进。如果我们把线程映射成"每个线程负责一个(batch, d_inner)通道",那么同一时刻所有线程访问的是x[b, t, :]这一行,地址连续,访存完美合并。
但这样有个问题:并行度只有batch × d_inner。假设 batch=16、d_inner=1536,也就两万多个并发通道,够填满 GPU 了,但如果 batch 小、模型窄,并行度就会不足。
所以更常见的映射是块内混合:一个线程块负责(一个 batch, 一段 d_inner),块内线程进一步细分到d_state维度。这样并发单元数变成batch × d_inner × d_state,但要注意d_state这一维在内存里往往不是最内层连续的(隐藏状态h是单独的一块内存),访问它会产生跨步,需要靠共享内存来兜底。
我在实际调试里最常犯的错误就是这个:想当然地认为所有维度都是连续的,结果h的访问模式是一堆散点,带宽利用率只有理论值的两三成。定位方法很简单,用性能分析工具看一眼全局内存的加载效率,如果远低于 100%,八成就是合并没做好。
3.2 块内并行扫描的实现细节
下面是我常用的一个块内 Hillis-Steele 扫描的骨架,双缓冲避免读写冲突:
// 块内扫描示意,CHUNK 为块长度,编译期常量 // sa[], sb[] 为共享内存中的 (a, b) 二元组 extern __shared__ float smem[]; float* sa = smem; float* sb = smem + CHUNK; float* sa2 = smem + 2 * CHUNK; float* sb2 = smem + 3 * CHUNK; // 载入本块数据 for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) { sa[i] = load_a(base + i); sb[i] = load_b(base + i); } __syncthreads(); for (int offset = 1; offset < CHUNK; offset <<= 1) { for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) { if (i >= offset) { // (a[i], b[i]) ∘ (a[i-offset], b[i-offset]) sa2[i] = sa[i] * sa[i - offset]; sb2[i] = sa[i] * sb[i - offset] + sb[i]; } else { sa2[i] = sa[i]; sb2[i] = sb[i]; } } __syncthreads(); // 交换缓冲 float* t; t = sa; sa = sa2; sa2 = t; t = sb; sb = sb2; sb2 = t; __syncthreads(); }几个细节值得单独拎出来说。
双缓冲不是可选项。如果用原地更新,第i个线程写的时候,第i+offset个线程可能还在读它,产生读写冒险。加一次__syncthreads()可以解决顺序问题,但会引入一次额外的全局同步,性能损失明显。双缓冲把同步点减到每次迭代一次,代价是共享内存翻倍。共享内存是很紧的资源,块大小要根据它反推。
共享内存用量需要算清楚。上面的骨架用了 4 个长度为CHUNK的 float 数组。如果CHUNK = 512,就是 8KB。加上编译器可能分配的其他开销,一个块占 12KB 左右的话,一块 48KB 共享内存的 GPU 能同时驻留 4 个块,占用率还不错。如果CHUNK加到 2048,共享内存 32KB,只能驻留 1 个块,占用率骤降,延迟隐藏能力就没了。这个取舍我在第 3.4 节会给出具体的计算。
分支预测。if (i >= offset)在每轮迭代里会让一部分线程空转。第一轮只浪费一半,后面几轮浪费的比例越来越小,整体开销可以接受。但如果块长度很小(比如 32),浪费比例会很显眼,这种时候不如直接用 warp shuffle 指令做扫描——一个 warp 内的扫描完全可以用__shfl_up_sync完成,不需要共享内存,速度飞快。
累加的数值顺序。Hillis-Steele 每轮都在做乘法链,$a$ 的连乘长度会指数增长。如果 $a$ 的平均值明显小于 1(状态衰减快),连乘几十步之后就会下溢到 0。虽然结果上没啥大问题(表示"早就忘光了"),但如果你的模型需要捕捉极长程的弱依赖,精度损失就不可忽略。缓解办法是把块切小,或者在对数域里累加。
3.3 块间状态传递的两种做法
块内扫描算完之后,每个块手里有一个总转移 $(A^{(k)}, B^{(k)})$,含义是"从块首到块尾的整体变换"。块数量一般是几十到几百,量级很小,处理方式有两种。
做法一:串行传递。起一个 kernel,让块 0 把自己的状态写出去,块 1 读进来修正自己,再写出去。这个链条是串行的,$K$ 个块就有 $K$ 步延迟。$K$ 小的时候无所谓,$K$ 大了会拖慢整体。
做法二:块间再做一次扫描。把那 $K$ 个二元组当成新序列,再做一次扫描,得到每个块的初始状态,然后并行回填。这一次扫描规模小($K$ 通常 ≤ 1024),一个线程块就能搞定,或者干脆在同一个 kernel 里做网格级同步。
我更倾向于做法二,因为它把串行深度从 $O(K)$ 压到了 $O(\log K)$。实测下来,序列长度 16384、块大小 256、块数 64 时,做法二的尾部延迟比做法一低了将近四成。
回填那一步的数学要写对:
$$h_{init}^{(k)} = A^{(k-1)}{total} \cdot h{init}^{(k-1)} + B^{(k-1)}_{total}$$
也就是用前一块的总转移去递推前一块的初始状态,得到本块的初始状态。这里最容易错的是拿成前一块的块内最后一个元素的状态,或者把下标错开一位。差一位带来的结果是整个序列平移一格,表现为输出看起来"差不多但就是不对",极难靠肉眼发现。我建议写完之后先用seq_len=8, chunk=4这种极小规模跑一遍,把每一步的中间值打印出来对,确认无误再放大规模。
3.4 参数选择与显存账本
现在来算一笔具体的账。假设一个中等规模的配置:
| 参数 | 取值 |
|---|---|
| batch | 16 |
| seq_len | 2048 |
| d_model | 768 |
| expand | 2 |
| d_inner | 1536 |
| d_state | 16 |
| chunk_size | 256 |
块数:$K = 2048 / 256 = 8$。
块状态显存占用:每个块要存一份 $(A^{(k)}, B^{(k)})$,形状是(batch, d_inner, d_state),即 $16 \times 1536 \times 16 = 393216$ 个元素。用 fp32 存,约 1.5 MB。8 个块合计约12.6 MB。这个量级完全可控。
如果保存全部中间状态呢?形状变成(batch, seq_len, d_inner, d_state),即 $393216 \times 2048 \approx 8.05 \times 10^8$ 个元素,fp32 下约3.2 GB。这就从"可控"变成"要命"了。而这只算了一层,真实模型有几十层。
这个对比就是"硬件感知"最直白的体现:存块状态不存全状态,把开销从 $O(L)$ 压到 $O(L/C)$,代价是反向传播时要重算一遍前向。省下来的显存能换来更大的 batch 或者更长的序列,这笔账怎么算都划算。
线程块配置:假设一个块负责d_inner_chunk = 64个通道,共享内存需要缓存64 × 16 = 1024个 float 的 $h$,约 4 KB。再加扫描用的 8KB,总共 12KB 左右。取 256 个线程,每个线程负责 4 个(通道, 状态)组合,寄存器压力不大。
占用率估算:一块 48KB 共享内存的 GPU,12KB 的块配置下最多驻留 4 个块。256 线程 × 4 = 1024 线程,如果 GPU 单个 SM 最多支持 2048 线程,占用率就是 50%。这个水平对带宽受限的算子够用了,因为主要还是靠访存合并和预取来隐藏延迟,不需要靠超高的 warp 数量来掩盖计算延迟。
一个反直觉的点:这类算子不是占用率越高越好。我试过把共享内存压到 6KB 换来 8 个块的驻留,结果性能反而降了,因为块变小之后块数变多,块间扫描的开销翻倍了。共享内存用量和块数之间存在一个最优区间,得实测。
4. 融合与重计算:硬件感知最见功力的地方
4.1 算术强度决定了优化方向
先算一下这个扫描算子的算术强度,这个数字会直接告诉你优化该往哪使劲。
处理一个元素(一次t迭代、一个通道、一个状态维),大致要做:一次乘加更新h = a*h + b,一次乘加累加y += c*h。加起来 4 次浮点运算。访存方面:读a、b、c各一次(可以复用),写y一次,$h$ 在寄存器里不落地。折算下来往少了说也要读写 4 个 float,16 个字节。
算术强度 ≈ 4 / 16 =0.25 FLOP/Byte。
这是一个非常低的值。现代 GPU 的算力带宽比(也就是 roofline 模型的拐点)通常在 10 到 20 FLOP/Byte 之间。0.25 远低于拐点,意味着这个算子百分之百是带宽受限的,算力利用率能到 5% 就算不错了。
既然是带宽受限,优化方向就只有一条:少搬数据。具体三条路:
融合。把投影、离散化、扫描、输出投影全都塞进一个 kernel,中间结果全部留在寄存器和共享内存里,一个字节都不往全局显存写。这一步的收益最大,实测下来能省掉 60% 到 80% 的访存量。
重计算。前向不存中间激活,反向重算。用计算换显存,用显存换更大的 batch。前面算过,一层能省 3 GB 级别的空间。
精度压缩。中间激活用 bf16 存,累加用 fp32。带宽直接减半,精度损失可控。
这三条是配套的,缺一不可。只做融合不做重计算,显存还是爆;只做重计算不做融合,重算的时候又从全局显存读一遍,时间全花在搬运上了。
4.2 训练时的重计算与显存占用实测
重计算的逻辑不复杂:前向只保留每块的边界状态(就是第 3.4 节算的那 12.6 MB),反向传播走到这一层时,从边界状态出发,把每块内部的 $h$ 重新算一遍,边算边做反向。
这里有个细节值得说:重算的时候,扫描的顺序是从块尾往块头(反向传播的方向),但状态的递推方向是从块头往块尾。所以实际上需要先把整块正向扫一遍拿到所有 $h_t$,再做反向。这一块的额外计算量等于整个前向的扫描部分,而扫描本身计算量很小,所以总开销增加得有限。
我在一个 24 层的模型上做过对比,配置就是第 3.4 节那张表:
| 方案 | 单层激活显存 | 24 层合计 | 单步耗时 |
|---|---|---|---|
| 保存全部中间态 | 3.2 GB | 76.8 GB | 基准 |
| 只存块边界态 | 12.6 MB | 302 MB | +18% |
| 块边界态 + bf16 激活 | 6.3 MB | 151 MB | +21% |
显存从 76.8 GB 降到 302 MB,两个数量级的差距,而时间只多了不到两成。这个交换比在任何场景下都是值得的。bf16 那一档进一步把显存砍半,代价是多 3 个百分点的耗时,如果你的卡本来就紧张,选它。
需要提醒的是,重计算的收益随 batch 增大而摊薄。batch 很小的时候,省下来的显存用不上,白白多花了两成时间。所以这个开关应该做成可配置的,而不是写死。
4.3 数值稳定性与精度策略
扫描的数值问题主要有两个来源。
第一个是乘子的连乘下溢。$a_t = \exp(\Delta_t A)$,$A$ 是负的,所以 $a_t < 1$。连乘 $C$ 次之后大约是 $\exp(C \cdot \Delta A)$。假设平均 $\Delta A \approx -0.05$,块长 256,连乘结果是 $e^{-12.8} \approx 2.8 \times 10^{-6}$,还在 fp32 的正常范围内(最小正规数约 $1.2 \times 10^{-38}$)。但如果 $\Delta A$ 再大一点,比如 -0.3,256 步之后就是 $e^{-76.8} \approx 4 \times 10^{-34}$,已经逼近非正规数区间了,精度会断崖式下跌。
缓解的办法很直接:控制块长。块长和 $\Delta A$ 的量级成反比,实践中块长取 128 到 512 之间比较稳妥。另外可以在块内做分段重归一化,把极小的乘子单独记录一个指数偏移,用的时候再加回去。
第二个是累加的顺序误差。fp32 只有 24 位有效数字,累加几千项之后,相对误差能到 $10^{-5}$ 量级。对训练来说通常无所谓(梯度噪声本来就在这个量级),但如果做推理或者要算高精度的对比指标,就得用 fp64 累加,或者用 Kahan 求和这类补偿算法。
精度策略上的实践:
- 参数和激活用 bf16 存储,累加和归约用 fp32。这是目前最主流的组合,精度损失基本感知不到。
- 指数运算
exp在 bf16 上误差偏大,建议先转 fp32 算完再转回去。 - 块间扫描那一步规模很小(几十到几百个元素),无条件用 fp32,几乎没有性能影响。
- 序列特别长的场景,考虑把 $A$ 参数化成对数域的形式,用加减代替乘除,能显著降低下溢风险。
5. 常见问题与排查实录
5.1 结果错位类问题的定位路径
这类问题最典型的表现是:输出看起来"形状对、量级对、就是数值不对",或者只在序列开头几个位置对、后面全错。我的排查路径是固定的三步。
第一步,缩到最小复现。batch=1, seq_len=8, chunk=4,只有两个块。手算一遍正确答案,跟 kernel 输出逐个对比。这时候能立刻区分是"块内扫描错"还是"块间传递错"——如果前 4 个位置对、后 4 个错,那肯定是块间传递的问题。
第二步,单独验证块内扫描。把块间传递代码注释掉,让每个块都从零初始状态开始扫,跟串行参考实现对拍。这一步能排除掉 90% 的实现错误。
第三步,检查块间传递的下标。前面说过,最常见的 bug 就是差一位或者用了错误的量。写成伪代码就是:回填第 $k$ 块的初始状态时,必须用第 $k-1$ 块的总转移去作用在第 $k-1$ 块的初始状态上。写成公式是 $h^{(k)}{init} = A^{(k-1)}{total} h^{(k-1)}{init} + B^{(k-1)}{total}$。用最后一个元素的状态是错的,用第 $k$ 块自己的总转移也是错的。
还有一个隐蔽的坑:同步点缺失。块内扫描用共享内存,如果漏了一次__syncthreads(),在某些执行时序下结果是对的,换个 GPU 或者换个占用率就跑飞了。这种"时对时错"的 bug 最折磨人。解决办法是在开发阶段打开计算卫生检查工具(比如 compute-sanitizer 那类),它会直接告诉你哪两个线程存在竞争。
5.2 性能不达预期时的排查顺序
调优的时候别乱试,按影响从大到小依次排查。
先看访存效率。用 profiler 看全局内存的加载/存储效率,理想值接近 100%。如果只有三四十,说明有严重的跨步访问,大概率是张量布局和线程映射没对齐。这个是最常见、也是收益最大的优化点。
再看共享内存的 bank 冲突。共享内存被分成若干个 bank,同一 warp 内的线程如果访问同一个 bank 的不同地址,就会串行化。我们的 $h$ 数组第二维是d_state=16,如果线程按(通道, 状态)排列且步长为 16,恰好是 32 的约数,冲突会很严重。解决办法是给数组的第二维补一个 padding,把步长改成 17。
然后看占用率。前面说过,占用率不是越高越好,但太低肯定不行。低于 25% 的话,延迟隐藏能力就不够了。
最后看块大小的取舍。块大,块间开销小,但共享内存占用高、占用率低;块小则反过来。这个要扫一遍参数,我一般在 [64, 128, 256, 512, 1024] 里挑,配合不同的块内通道数做二维扫描,通常能找到一个明显的甜点区。
5.3 常见问题速查表
| 现象 | 可能原因 | 排查动作 | 处理办法 |
|---|---|---|---|
| 前若干位置对、后面全错 | 块间状态传递错误 | 缩小到两块规模手算对比 | 核对方程下标,确认用的是前一块的总转移 |
| 结果每轮跑都不一样 | 共享内存竞争 | 打开计算卫生检查工具 | 补上缺失的同步点 |
| 带宽利用率低 | 访存未合并、跨步太大 | 看 profiler 的访存效率指标 | 调整张量布局,让最内层维连续 |
| 性能明显低于理论值 | 共享内存 bank 冲突 | 用 padding 前后对比 | 数组第二维补 1,把步长改成质数 |
| 长序列精度骤降 | 乘子连乘下溢 | 打印块内总乘子的量级 | 减小块长,或用对数域累加 |
| 反向传播显存爆掉 | 未开重计算 | 看激活显存占比 | 只存块边界态,反向重算块内状态 |
| 小 batch 时提速不明显 | 重计算开销占比过高 | 对比开/关重计算的耗时 | 做成可配置开关,按 batch 动态选 |
| 块数很多时尾部延迟大 | 块间串行传递链条太长 | 看 kernel 末尾的空转时间 | 块间改用树形扫描,深度降到对数级 |
最后分享一个我自己的习惯:每次实现这类扫描 kernel,我都会先写一个纯 Python 的串行版本作为参考,用同样的随机种子生成输入,然后拿 GPU 版本的输出跟它逐元素对齐。容忍度从1e-6开始,对不上就往回找是哪一个块、哪一个位置开始偏的。这套流程看着笨,但它把我从"改一行跑一次、全靠猜"的循环里彻底解放了出来。扫描这东西,一旦下标错了一位,肉眼是绝对看不出来的,唯一可靠的办法就是让机器替你比。
另外,如果你的场景里序列长度变化很大,别指望一套块大小打天下。我做过的项目里,短序列(几千)和长序列(几万)的最优配置能差出三倍以上,老老实实做一组模板实例化,编译期把块长定死,比运行时动态判断省事得多,也不会有分支开销。