如果你最近在跟进 Mamba、线性注意力这类模型，大概率会看到 Gated DeltaNet 这个新名字。我第一次读那篇论文的时候，第一反应是“公式也不难嘛，不就是多了一个 delta 校正项外加一个门控”。但一打开官方实现，发现里面全是分块、状态扫描、重计算的细节，才意识到真正难的地方不是公式本身，而是“怎么在 GPU 上把它跑起来”。

这个问题的本质是：普通的线性注意力可以借助矩阵乘积的结合律，把整个序列上的递归直接展开成矩阵乘法，天然适合 GPU；但 Gated DeltaNet 因为引入了 delta 规则，状态更新里带了一个“上一时刻状态对当前 token 的响应”，导致它在序列方向上不像线性注意力那么好拆。你没法简单地把整条序列丢给一个并行 scan 然后万事大吉。

这篇文章我主要想聊两件事：第一，线性注意力家族的 GPU 并行化到底卡在哪里；第二，一个可行的、能落地的分块并行实现思路，以及我在实际调内核时踩到的一些坑。内容更适合已经会写 CUDA/Triton、或者至少对 FlashAttention 这类 IO-aware kernel 不陌生的读者。如果你只是刚接触线性注意力，后面的三段式分块流程也值得先读一遍，它能帮你建立对序列状态模型并行方式的直觉。

1. 怎么理解“GPU并行”这件事

1.1 自注意力与线性注意力的并行差异

先回忆一下标准的自注意力。给定 Q、K、V，某个位置的输出是：

o_t = sum_{i<=t} attention_weight(i,t) * v_i

由于因果 mask 的存在，t 位置的输出依赖前面所有位置，但在 GPU 上它其实很好并行。因为注意力权重是通过 QK^T 一次性算出来的，我们可以把序列切成块，每个块只需要拿到前面块的“全局信息”（比如 FlashAttention 里的 running max 和 running sum），就能并行计算。这种并行来自“softmax 的局部值可以被后期归并”。

而线性注意力把注意力矩阵简化成线性核，状态 S_t 是一个 d×d 的累加矩阵：

S_t = S_{t-1} + k_t v_t^T o_t = S_t^T q_t

看起来很朴素，但它有一个非常好的性质：状态更新操作是“可结合的”。展开之后：

S_t = sum_{i<=t} k_i v_i^T

这意味着状态只是一个历史 token 的外积累加。于是训练或者 prefill 阶段，整个序列上的运算可以重写成一个大矩阵乘法：

S = K^T V (同时做 causal mask)

GPU 上对这种矩阵乘法非常友好。即使你想保持子序列级别的并行性，也可以用“前缀扫描”把组合操作拆开。也就是说，线性注意力在序列维度上是“可并行扫描”的，这也是它和普通 softmax attention 在实现上一个很大的不同。

1.2 Gated DeltaNet 的公式与难点

Gated DeltaNet 属于 Delta Rule 这一支。常见的 Gated DeltaNet 更新形式可以写成：

S_t = alpha_t * S_{t-1} + beta_t * k_t * (v_t - S_{t-1}^T k_t)^T

多出来的部分就是 delta 校正项：模型先去读一下当前状态 S_{t-1} 对 k_t 的“预测结果” S_{t-1}^T k_t，如果这个结果和 v_t 不一致，就往状态里做一次低秩修正。这有点像在线学习里的 delta rule，状态不再只是简单累加，而是带反馈的更新。同时 alpha_t 控制遗忘，beta_t 控制写入强度。

问题就在这个反馈项上。 S_t 的计算依赖 S_{t-1}，而 S_{t-1} 又依赖更早的状态，整条链必须一步步推，不能直接展开成 K^T V 这种一次性矩阵乘法。也许你可以把每个 token 的更新看成一个小矩阵操作，但这些操作之间不满足普通线性注意力的那种简单结合律。如果你直接写：

for t in range(T): S = alpha[t] * S + beta[t] * outer(k[t], v[t] - S.T @ k[t])

那在 GPU 上就是一个顺序循环。 T 一长，全部时间都花在等待上一次更新完成上。这跟让 GPU 去跑一个串行程序差不多，浪费了它最擅长的并行能力。

所以核心矛盾是：

• 线性注意力：更新是“可结合”的策略，所以可以用 scan 或矩阵乘把序列维度并行化。
• Gated DeltaNet：更新是“带反馈的仿射变换”，需要在分块后把这种仿射变换本身变成可组合的对象，才能做块间并行。

2. 解决思路：分块（Chunking）与状态传递

2.1 线性注意力的分块并行

在进入 Gated DeltaNet 之前，先看一个更通用的分块策略，理解它会很有帮助。

假设序列被切成若干块，每个块长度为 C。对线性注意力来说，我们可以定义块内累积矩阵：

B = sum_{i in block} k_i v_i^T

再定义一个块结束时的状态：

S_end = S_start + B

块内和块间都被分开了。GPU 可以先并行计算所有块的 B，然后再单独处理块之间的状态传递。因为 S_{t} = S_{t-1} + 外积，所以这个“传递”就是加法，做前缀和即可。整个算法只需要两次遍历序列：

1. 第一次遍历：并行计算每个块的局部累积 B。
2. 第二次遍历：把块之间的前缀状态扫出来，再拿 state 去算输出。

这个思路等价于把长序列的 scan 拆成两层：块内顺序扫，块间并行扫。你可以根据 GPU 的并行能力选 C 的大小，比如 64、128、256。

2.2 把 Delta 更新改写成仿射变换

Gated DeltaNet 不能直接用“加法式”的状态传递，但我们可以把单 token 更新重新整理成更一般的形式：

S_t = (alpha_t I - beta_t k_t k_t^T) S_{t-1} + beta_t k_t v_t^T

注意看，这里右边第一项是 alpha_t S_{t-1} 加上一个低秩校正，第二项是一个与 S_{t-1} 无关的输入项。如果我们记：

A_t = alpha_t I - beta_t k_t k_t^T B_t = beta_t k_t v_t^T

那么单 token 更新就是：

S_t = A_t S_{t-1} + B_t

这是一个非常标准的仿射变换。只要 S_t 是仿射变换，我们就能做“变换的组合”。假设一个块内从初始状态 S_start 出发，连续经过若干 token 的更新，最终得到的 S_end 一定可以写成：

S_end = A_chunk S_start + B_chunk

然后问题就变成了怎么高效算出 A_chunk 和 B_chunk。

2.3 块间的关联扫描

既然每个块都可以被浓缩成一个仿射变换 (A_chunk, B_chunk)，那么块与块之间就有了明确的组合关系。给定一个更早的变换 (A_prev, B_prev) 和一个更晚的变换 (A_next, B_next)，它们的组合是：

S' = A_next (A_prev S + B_prev) + B_next = (A_next A_prev) S + (A_next B_prev + B_next)

于是：

compose((A_next, B_next), (A_prev, B_prev)) = (A_next A_prev, A_next B_prev + B_next)

这个组合运算满足结合律，因此可以在块级别上做并行前缀扫描。也就是说，Gated DeltaNet 的状态传递不是完全不能并行，而是需要先把每个块“压缩”成一个仿射变换，然后在块的数量级别上做一次并行 scan。块的数量通常远小于序列长度，所以这个 scan 的成本也能控制在可接受范围。

3. 实操：一个完整的 GPU 并行实现流程

3.1 三阶段流水

我在自己复现时，把整个流程切成了三个清晰的阶段，这样既方便调试，也方便替换成更激进的融合内核：

1. 投影阶段：从输入 x 计算 q、k、v、alpha、beta。
2. 块变换阶段：每个块独立计算自己的 A_chunk 和 B_chunk。
3. 块间扫描与重计算阶段：对块状态做并行 scan，再用块起始状态重新计算输出。

第一阶段基本就是普通矩阵乘法，没什么需要特别解释的。第二和第三阶段是关键，后面详细展开。

3.2 阶段1：投影与局部参数

可以把 q、k、v 都投影到 d 维（通常 d 是 16、64、128 这类值），beta 和 alpha 一般通过 sigmoid 得到。在训练时，我会用标准的 einsum 一次性投影整条序列，然后在 kernel 里直接读取这些参数。

这一步里有一个重要决定：参数精度。推荐保留 fp32 的 q、k、v 用于状态更新，尤其在第二阶段算 A_t 和 B_t 时。GPU 在 fp16 下做外积和矩阵乘法很快，但 delta 规则里的“S^T k”会反复叠加误差，后面你会看到 fp16 状态很容易发散。

3.3 阶段2：块状态变换

假设块大小是 C。这一步输入是当前块的 k、v、alpha、beta，输出是 A_chunk 和 B_chunk。

我用一个偏教学向的 PyTorch 代码来描述一下逻辑，实际部署时肯定要写成 CUDA/Triton kernel，但数据流是一样的：

def chunk_transform(k, v, alpha, beta): # k, v: [C, D] C, D = k.shape A = torch.eye(D, device=k.device, dtype=torch.float32) B = torch.zeros(D, D, device=k.device, dtype=torch.float32) for i in range(C): # alpha[i] 可以是标量，也可以是 D 维向量 Ai = alpha[i].diag() - beta[i] * torch.outer(k[i], k[i]) Bi = beta[i] * torch.outer(k[i], v[i]) # 把当前 token 的仿射变换组合进块变换 B = Ai @ B + Bi A = Ai @ A return A, B

这个循环是块内的串行部分，理论上没法完全消掉，但因为块大小 C 通常只有 64 到 256，而且每个块是独立计算的，GPU 上可以同时跑“序列块数 × batch × head”这么多个循环实例。并行度主要来自块间，不来自块内。

这里有两个可以优化的点：

• 如果 alpha[i] 是标量，A_i 是“标量对角阵减一个秩 1 矩阵”，那么组合出来的 A_chunk 其实也带低秩结构。存储 A_chunk 时不要总是用 D×D 完整矩阵，能用低秩表示就用低秩表示。
• 计算 B_chunk 时，可以把循环里的Ai @ B改成(alpha*B - beta * k (k^T B))，用左乘一个向量或小矩阵的方式来避免显式构造 D×D 的 A_i。公式上一致，效率高很多。

3.4 阶段3：扫描与重计算输出

块变换算完后，我们对所有 chunk 的 (A_chunk, B_chunk) 做一次关联扫描。这里可以用自定义 kernel，也可以用 JAX 的lax.associative_scan来验证正确性。PyTorch 目前没有一个特别方便的公开 associative_scan API，所以你如果要在 PyTorch 里快速验证，可以自己写一个简单的递归二分 scan：

def compose(next_t, prev_t): A_next, B_next = next_t A_prev, B_prev = prev_t return (A_next @ A_prev, A_next @ B_prev + B_next) def scan_chunks(chunks): # chunks: list of (A_chunk, B_chunk) # 先用简单的顺序版本验证，再替换成并行 scan acc = None results = [] for chunk in chunks: acc = chunk if acc is None else compose(chunk, acc) results.append(acc) return results

跑通以后，你用 Hillis-Steele 扫描或者 Blelloch 扫描把它并行化。每个 block 的最终状态是：

S_start_chunk = A_prefix @ S_initial + B_prefix

拿到 S_start_chunk 之后，最后一个阶段是对每个 chunk 重新执行顺序 forward。这一步可以复用已经算好的 k、v、alpha、beta，相当于在块内再做一次带初始状态的小循环：

def block_forward(k, v, alpha, beta, S_start): S = S_start.clone() outs = [] for i in range(k.shape[0]): out_i = S.T @ k[i] # 这是 key 与状态的交互 outs.append(out_i) # 也可以直接写成 out_i = q[i]^T S 的形式，取决于你的线性注意力公式 S = alpha[i] * S + beta[i] * torch.outer(k[i], v[i] - S.T @ k[i]) return torch.stack(outs), S

“重计算”这个步骤听起来浪费，但它是有意为之：如果我们把每个 token 时刻的 S 都全部存下来，显存会爆得很快。状态 S 是 D×D 矩阵，存下每个时刻的状态等于额外消耗 T×D×D 个浮点数。不如只存每块的 q/k/v/alpha/beta，然后在输出阶段重算一遍。时间换显存，在长序列场景下非常划算。

3.5 使用 Triton 时的内核注意点

如果你想写成 Triton kernel，一个比较自然的划分是：

• 每个 program 处理一个 (batch, head, chunk)。
• k、v、alpha、beta 加载到寄存器或者 shared memory。
• 用一个小的 D×D 状态矩阵 S 在块内做顺序更新。
• 把这个 chunk 的 (A_chunk, B_chunk) 写出去，供后续 scan 使用。

在块内循环里，你大概率不必对“序列长度”做并行，而是依靠 batch、head、num_chunks 三个维度填充 SM。当 batch × head 不够大时，可以再把 chunk 维度再切细，比如一个 GPU block 只处理半个 chunk，提高占用率。

写 Triton 的时候我建议这样组织：

@triton.jit def chunk_transform_kernel( k_ptr, v_ptr, alpha_ptr, beta_ptr, A_ptr, B_ptr, D, C, stride_k, stride_a, ... ): pid = tl.program_id(0) # 每个 pid 负责一个 chunk # 在循环里维护 A[B,D,D] 和 B[B,D,D] pass

Triton 里直接操作二维稠密矩阵比较顺手，但要注意tl.dot的尺寸限制。如果 D=64 或 128，直接用块矩阵没问题；如果 D 很大，可能要拆成多个 tile 来更新状态矩阵。

还有一个容易忽略的点：A_chunk 和 B_chunk 的输出尺寸是 “batch × head × num_chunks × D × D”，对长序列来说仍然不小。所以更工程化的做法是在扫描阶段开始时，不一次性把所有 chunk 的 A_chunk、B_chunk 写到全局内存，而是生成一个 chunk 后立刻送入扫描流程的对应流水线。但要验证“块间并行扫描的正确性”时，先把所有 chunk 的变换存下来更稳妥。

4. 不同方案的取舍与调参

4.1 线性注意力还是 GDN？

我画过一张很小的对照表，用来判断一个项目到底该用哪种实现路径：

说实话，如果你的序列长度只有几百，batch 又大，那直接用普通线性注意力的结合扫描就行，没必要上 GDN。GDN 的收益主要体现在你希望用更少参数和更高压缩比去模拟更复杂依赖的时候，这时候才值得为它的状态更新规则付出实现成本。

4.2 块大小怎么定

块大小 C 的选择会直接影响性能，我的经验是：

• C 太小（比如 16）：块数量太多，块间 scan 和状态组合的开销变大，同时每个 chunk 的 A_chunk、B_chunk 数量也变多，存储不划算。
• C 太大（比如 1024）：块内顺序循环太长，单个 block 的计算延迟变大，占用率会掉下来，而且重计算阶段的访存压力也增大。
• 折中区间通常在 64 到 256。具体选多少，要结合 D 的大小和 GPU 型号。如果 D=64，C=128 是个很稳的起点。

另一个判断方式：让一个 chunk 的块内循环次数和 D 的规模匹配。状态更新每次都要读一遍 D×D 的矩阵，C × D^2 级别的计算量可以很好地用矩阵乘指令掩盖，C 和 D 数值越接近越好。

4.3 数值精度和参考实现

我强烈建议把状态 S 和 A_chunk、B_chunk 都用 fp32 来算。很多线性注意力模型在 fp16 下训练还凑合，是因为累加项比较简单；但 delta 规则等于在状态上不断做减法修正，fp16 的尾数精度不够时，微小误差会累积，最后出现对角线上的“假注意力”。

如果因为性能问题必须用 bf16，那至少要做到以下几点：

• k、v 进入状态更新前，用 fp32 做 outer product，再转回 bf16。
• 对特征维度做某种 normalization，避免 k 的范数过大。
• 对 alpha、beta 保持 fp32，门控值对稳定性的影响比 k/v 更大。

5. 我踩过的坑

5.1 状态更新顺序写反

这个坑非常低级但很容易犯。Delta 规则里输出 o_t 用的是“更新前的状态 S_t”还是“更新后的状态”，不同论文习惯不一样。如果你先更新 S，再算输出，数值和官方实现会差很多，但 loss 又不会完全崩，只是模型收敛变慢。后来我把单 token 的 PyTorch 参考实现和官方的结果逐位置对了一下，发现差就差在这。

建议最开始时先把参考实现写成最简单、最不容易出错的顺序版本：

S = S_start for i in range(C): out.append(S.T @ k[i]) S = alpha * S + beta * outer(k[i], v[i] - S.T @ k[i])

先让这个版本对齐，再去做分块重计算。

5.2 fp16 状态下状态爆掉

我一开始为了提速，把 S 直接存在 fp16 里。结果训练 loss 很怪，先下降一小段，然后开始剧烈震荡。把状态打印出来发现 S 的对角线元素已经变成 1e4 级别的值。

原因是 delta 规则里每次更新都会有“读旧状态、做预测、修正”的过程，如果 S 的数值范围没有约束，误差修正项会不断放大。Gated DeltaNet 里门控 alpha 和 beta 的作用就是控制这个范围，但 fp16 的动态范围太小，状态很容易溢出。

解决方法是：

state = state.to(torch.float32)

状态更新用 fp32，只在最终输出 o_t 时转回低精度。这个改动带来的加速损失远小于它避免的不稳定问题。

5.3 扫描内核的同步问题

做 chunk 间 scan 时，我最初天真地以为可以直接在多个 CUDA stream 上并行算 prefix，然后合并。结果因为在 stream 之间没有做合适的事件同步，某些 chunk 读到了还没写完的前序 chunk 状态，导致结果时对时错。

不是算法问题，是工程问题。标准做法是：

• chunk transform 阶段结束后，用一个全局 memory fence。
• 扫描阶段和重计算阶段放在同一个 stream，或用 event 串起来。
• 调试时可以先回到单 stream 顺序执行，确认结果正确后再并发。

5.4 用 CUB / warp scan 时注意 scan 对象维度

在线性注意力里，scan 的“元素”是一个 d×d 矩阵或外积，不是标量。如果你在 CUDA 里直接用 CUB 的 BlockScan 扫标量数组，会发现状态是错的。scan 元素的 typesize 越大，寄存器压力越高，block 数量反而下降。

所以在块间 scan 阶段，我后来更倾向用 Blelloch 扫描，并且把每个 chunk 的变换拆成两个矩阵来扫，而不是放在同一个结构体里。这样能避免记忆体 padding 和 alignment 带来的额外开销。

6. 最后的实操经验

Gated DeltaNet 这类带反馈的线性注意力，核心难点不在公式，而在你能不能把“顺序状态更新”重新组织成适合 GPU 的分块计算。我的经验是，不要一上来就追求最融合、最高性能的 kernel，而是先做一个“顺序参考版 + 块输出重计算版 + 块间并行 scan 版”的三角验证：

• 顺序参考版用来保证数学正确。
• 块输出重计算版用来验证分块逻辑。
• 块间并行 scan 版用来验证 GPU 并行改造。

三个版本结果一致以后，再逐步把循环写进 Triton 或 CUDA。性能优化时可以大胆尝试不同块大小，但正确性基线一定要先锁住。

我自己实际跑下来，最值得优先优化的其实是阶段 2 里 A_chunk 和 B_chunk 的计算。很多实现会把注意力放在最后的重计算上，但一旦序列长了，每个 chunk 都要做一块内部仿射组合，这部分算力一点不少。换成低秩更新表示以后，耗时能明显降下来。

如果你手头项目也需要实现类似的模型，建议把状态维度 D 控制在 64 以内。D 一超过 128，D×D 状态矩阵的读写就会变成瓶颈，这时你就需要更强的 kernel 融合能力，比如把多个 chunk 的变换直接放在一个 CTA 里计算，而不是写完再读。希望这篇能帮你少走点弯路。