用 Triton + Symmetric Memory 写一个 AllReduce + RMSNorm 融合算子

本文最后更新于 2026年7月19日

背景：为什么要融合

TP 推理里，每个 Transformer block 的 attention 和 MLP 出来都会跟一次 all-reduce，紧接着是下一个 sub-layer 的 RMSNorm（外加一个残差相加）。decode 阶段 batch 很小，这几个算子全是 latency-bound：

all_reduce 走 NCCL，一次 launch + ring/tree 的多跳；
RMSNorm 又是一个独立 kernel launch；
中间还要把 all-reduce 的结果写回 HBM、再读进来给 norm。

这些固定开销有多夸张？算一笔账：M=1, N=2048, bf16 时数据一共 4KB，GB200 的 HBM 带宽约 8TB/s，纯搬运是纳秒级；但实测 NCCL all_reduce + eager RMSNorm 一次要 ~40µs——几乎全是 launch 开销、多跳延迟和 HBM 往返。

很自然的想法就是：把它们塞进一个 kernel，all-reduce 的结果还热在寄存器里的时候，顺手就把残差和 RMSNorm 算掉。

融合前后对比（一次 TP 输出）

NCCL + eager RMSNorm：

  ┌─────────────┐   HBM   ┌─────────────┐   HBM   ┌─────────────┐
  │ NCCL all-   │ ──写回─► │ load +      │ ──写回─► │ RMSNorm     │
  │ reduce      │         │ residual_add│         │ kernel      │
  └─────────────┘         └─────────────┘         └─────────────┘
      一次 launch             一次 launch             一次 launch
      多跳 ring/tree

Fused kernel：

  ┌─────────────────────────────────────────────────────────────┐
  │  all-reduce sum  ──►  + residual  ──►  RMSNorm  ──► 写回 HBM  │
  │         一个 Triton program，一次 launch，数据始终热在寄存器里     │
  └─────────────────────────────────────────────────────────────┘

这个算子在算什么

对一个 rank-local 的张量 x（shape [M, N]，也就是要被 all-reduce 的东西）：

ar   = sum over ranks of x                 # one-shot all-reduce
h    = ar + residual                       # residual_out
rms  = rsqrt(mean(h^2, dim=-1) + eps)
out  = (h * rms) * weight                  # norm_out

单个 token 的 RMSNorm 计算流程（hidden dim = N）

   x[0]   x[1]   x[2]        x[N-1]
    │      │      │             │
    └──────┴──────┴─────────────┘
                 ▼
          sum(x[i]^2) / N      ←── 均方
                 │
                 ▼
          rsqrt(均方 + eps)    ←── rms（每行一个标量）
                 │
        ┌────────┼────────┐
        ▼        ▼        ▼            ▼
     x[0]     x[1]     x[2]  ...   x[N-1]
      │        │        │              │
      *        *        *              *
      │        │        │              │
    w[0]     w[1]     w[2]  ...    w[N-1]
      │        │        │              │
      ▼        ▼        ▼              ▼
   out[0]   out[1]   out[2]  ...   out[N-1]

关键原语：symmetric memory

要在一个 kernel 里做 all-reduce，就需要各卡能直接读到对方的数据。这正是 PyTorch symmetric memory 干的事情：

用 symm_mem.empty + rendezvous 分配的一块 buffer，会被映射进每一张卡的地址空间（走 NVLink / NVSwitch）。于是在 kernel 里，我从 hdl.buffer_ptrs_dev 拿到所有 peer 的 buffer 指针，就能像读本地显存一样读别人的那份数据。

有了这个原语，one-shot all-reduce 就是：每张卡把自己的贡献放进对称显存，同步一下，然后每张卡各自把所有 peer 的对应行加起来。

one-shot all-reduce  across world_size GPUs（一行数据）

  GPU 0          GPU 1          GPU 2          GPU 3
 ┌─────┐        ┌─────┐        ┌─────┐        ┌─────┐
 │ x_0 │        │ x_1 │        │ x_2 │        │ x_3 │   各卡本地贡献
 │row r│        │row r│        │row r│        │row r│
 └──┬──┘        └──┬──┘        └──┬──┘        └──┬──┘
    │              │              │              │
    └──────────────┴──────────────┴──────────────┘
                   ▼
        Symmetric Memory（NVLink / NVSwitch）
        同一块 buffer 被映射到每张卡的地址空间
                   │
    ┌──────────────┼──────────────┐
    ▼              ▼              ▼              ▼
 读 x_0         读 x_1         读 x_2        读 x_3
    │              │              │              │
    └──────────────┴──────────────┴──────────────┘
                   ▼
            sum = x_0 + x_1 + x_2 + x_3
                   │
                   ▼
              residual + RMSNorm

笔者的实现跑在 4×GB200 上，torch 2.12 + triton 3.7。

总体设计

Kernel 的结构很简单：一个 Triton program 负责一个 token（一行），hidden 维度放进一个 block 里，fp32 累加。每个 program 的数据流是：

1
2
3

读自己 + 所有 peer 的第 r 行  →  fp32 求和（all-reduce）
  →  加 residual，写出 residual_out
  →  求均方、rsqrt、乘 weight，写出 norm_out

计算本身没有任何难度，真正的难点在各卡怎么同步——必须保证「我读 peer 的第 r 行时，peer 已经把它写好了」。同步方式直接决定了延迟、支持的 dtype 和实现风险。笔者实现了两种：

	方案一：signal-pad barrier	方案二：Lamport
同步信号	显式屏障：signal pad 上写 1 / 等 1	数据本身：哨兵值被覆盖即「到了」
用的原语	`atomic_xchg` / `atomic_cas`（sys 作用域）	普通 store + `volatile` load 自旋
dtype	bf16 / fp16 / fp32 均可	仅 bf16（两个 bf16 打包进一个 u32）
跨 launch 状态	无，pad 自清	三缓冲 + 轮换 flag
稳态延迟	高（多一趟屏障 round-trip）	低（数据到了就能算）
实现风险	低，同步状态与数据解耦	高，哨兵 scrub / volatile / ABA 都是坑

读者可能会问：一个 program 负责一行，那 decode 只来一个 token 时，整张卡不就只剩一个 CTA 在干活，上百个 SM 全闲着？闲着就闲着——这个区间的瓶颈根本不是吞吐。M=1 时每个 program 要读的不过 world_size × N × 2B（N=2048 时 16KB），一个 CTA 拉这点数据绰绰有余；瓶颈在同步和 NVLink 小消息传输这些固定延迟上。这一点不用争辩，后面 benchmark 表的同一列就能看出来：N=2048 时 M 从 1 涨到 128（128 倍的 token），lamport 只从 ~7.3µs 涨到 ~9.7µs——多花的 ~2.3µs 才是「多干 127 行活」的真实边际成本，其余全是固定开销。

那能不能把一行沿 hidden 维切给多个 program？all-reduce 部分好切，但 RMSNorm 要整行的 sum of squares 才能算 rms——切了 N，每个 CTA 只有部分和，必须跨 CTA 再归约一次，等于把融合省掉的同步又加了回来（TRT-LLM/flashinfer 的 two-shot 就是这么干的，小 batch 下反而更慢）。所以整行放进一个 block 是刻意的取舍：宁可单 CTA 低利用率，也要让行统计量留在本地寄存器里。

下面分别说。

方案一：signal-pad barrier

第一版用显式屏障。先说清楚要同步什么：kernel 一启动，每张卡的每个 program 立刻就要去读其他卡对称 buffer 里的对应行，但对方那张卡可能还没把数据写进去——所以必须有一扇门：确认「所有卡都写完了」之后才放行去读。

这扇门就是 signal pad。symmetric memory 除了数据 buffer，每个 rank 还带一块小的对称内存，专门用来做同步——可以把它理解成一块所有卡都能读写的记分板，布局和约定如下：

每个 rank 一块 pad：uint32[num_blocks, world_size]
（num_blocks 就是这次 launch 的 program 数，一个 program 一行，即 M）

约定：我 pad 上的槽位 [b, r] == 1   ⟺   「rank r 的 block b 已经报到」

按 block 分行的原因是：每个 program 只读写自己那一行，槽位互不冲突，也不用让全 grid 的 program 齐步走。

屏障协议只有两步：

第 1 步 arrive： rank r 的 block b 向全部 world_size 块 pad 写入 1，写的位置是每块 pad 上的槽位 [b, r]——即每个 rank 只写「自己那一列」，但写遍所有 pad（atomic_xchg，release）。

第 2 步 wait： rank r 的 block b 自旋读自己 pad 的第 b 行（atomic_cas，acquire），直到 world_size 个槽位全为 1——即所有 rank 的 block b 都完成了第 1 步，放行。

world_size = 4，rank 0 视角（pad 初始全 0）

  arrive（远端写）：rank r 把 1 写进全部 4 块 pad 的槽位 [b, r]

    rank 0 ──► 各 pad 的 [b,0]      rank 2 ──► 各 pad 的 [b,2]
    rank 1 ──► 各 pad 的 [b,1]      rank 3 ──► 各 pad 的 [b,3]

  wait（本地读）：rank 0 自旋读自己 pad 的第 b 行

    rank 0 的 pad[b, :] = [ 1  1  0  1 ]   ← [b,2] 还是 0：rank 2 没到，继续自旋
                          [ 1  1  1  1 ]   ← 全为 1：放行

写远端、读本地的原因：写是每个 block 只做一次的 world_size 次远端写；而等待时的自旋要反复读，放在本地内存上就不必跨 NVLink 空转。

为什么看到 flag 就等于数据写好了？两条 memory order 缺一不可：写 1 用 release，保证它对别人可见时、排在它前面的数据写入也可见；读 flag 用 acquire，保证看到 1 之后再发起的读不会拿到旧值；scope="sys" 让这个排序关系跨 GPU 成立。合起来就是跨卡的 happens-before：看到 flag=1 ⟹ 对方的数据已经写完。barrier 之后，每个 program 把自己这一行从 world_size 个 peer buffer 里读出来、fp32 求和，接着算 residual + RMSNorm。

落到代码：

@triton.jit
def _blockwise_barrier(signal_pad_ptrs, rank, world_size: tl.constexpr, sem: tl.constexpr):
    block_id = tl.program_id(0)   # 省略多维展开
    remote = tl.arange(0, world_size)
    ...
    send_addr = remote_pad + block_id * world_size + rank   # 写到每个 peer 的 [block, my_rank]
    wait_addr = local_pad  + block_id * world_size + remote  # 等自己的 [block, every_peer]

    ones  = tl.full((world_size,), 1, tl.int32)
    zeros = tl.full((world_size,), 0, tl.int32)
    tl.atomic_xchg(send_addr, ones, sem=sem, scope="sys")

    arrived = zeros
    while tl.sum(arrived) != world_size:
        got = tl.atomic_cas(wait_addr, ones, zeros, sem=sem, scope="sys")
        arrived = tl.maximum(arrived, got)

这里藏了个坑，笔者第一版就死锁了。最初等待循环是直接对每一轮 CAS 的返回值求和：

1 2	`while tl.sum(tl.atomic_cas(wait_addr, ones, zeros, ...)) != world_size: pass`

问题出在 atomic_cas(ptr, 1, 0) 会消费 flag——看到 1 就顺手写回 0。消费是故意的：pad 清干净了，下一次 launch 才能直接复用。但「每轮求和」只在所有 peer 恰好同一轮内凑齐时才对；只要到达有先后，先到的 flag 在某一轮被吃掉，后面的轮次里它就永远缺席：

轮次   pad 槽位 [r0 r1 r2 r3]   本轮 CAS 看到     本轮 sum
 1     [ 1  0  0  0 ]          [ 1  0  0  0 ]      1   ← r0 的 flag 被吃掉（槽回 0）
 2     [ 0  1  0  0 ]          [ 0  1  0  0 ]      1   ← r0 的到达已被「遗忘」
 3     [ 0  0  1  0 ]          [ 0  0  1  0 ]      1
 4     [ 0  0  0  1 ]          [ 0  0  0  1 ]      1   ← 其实四卡都到过了
 5     [ 0  0  0  0 ]          [ 0  0  0  0 ]      0   ← 永远凑不齐 4：死锁

（这个 bug 阴险在它有时能跑过：各 peer 到达得足够集中时，一轮就能 sum 出 4。）

修复办法就是代码里的 arrived 数组：每个 lane 把历轮观测 OR 累积——槽位一旦被看到过 1 就永久记住，CAS 照常消费清 pad。既不丢到达记录，pad 也还是干净的。

这一版简单、鲁棒——同步状态（signal pad）和数据完全解耦，所以每次 launch 都是无状态的，bf16 / fp16 / fp32 都能直接跑。

方案二：Lamport（负零轮询）

方案一里的屏障，本质上是「全员到齐再开工」：每个 block 要等最慢的那个 peer 把 flag 写过来，才能开始读数据。对 latency-bound 的小消息，这一趟 release/acquire round-trip 就是延迟的大头。flashinfer 的 one-shot 走的是另一条路：完全不要屏障，同步信息由数据本身携带——读者看到数据的值变了，就知道它到了。笔者也照着实现了一版，几个关键点：

1. 两个 bf16 打包进一个 u32。 这是整个方案能成立的地基。4 字节对齐的 load/store 是原子的，所以读者要么看到整包旧值、要么整包新值，绝不会读到「一半旧一半新」。

2. 哨兵 + scrub。 comm buffer 预先填满哨兵字 0x80008000（两个 bf16 的 -0.0）。写数据前把真实数据里的 -0.0（0x8000）抹成 +0.0（0x0000），这样真实数据永远不可能等于哨兵。

@triton.jit
def _pack_pairs_scrubbed(vals):
    v_lo, v_hi = tl.split(vals.reshape([vals.shape[0] // 2, 2]))
    lo16 = v_lo.to(tl.uint16, bitcast=True)
    hi16 = v_hi.to(tl.uint16, bitcast=True)
    lo16 = tl.where(lo16 == 0x8000, tl.full(lo16.shape, 0, tl.uint16), lo16)
    hi16 = tl.where(hi16 == 0x8000, tl.full(hi16.shape, 0, tl.uint16), hi16)
    return lo16.to(tl.uint32) | (hi16.to(tl.uint32) << 16)

3. push + poll。 每张卡把自己 scrub+pack 过的行写进每个 peer 的 buffer，然后 volatile-load 自己的 buffer，逐 slab 自旋，直到那个字不再是哨兵——这个「跳变」本身就是「数据到了」的信号。没有 atomic、没有屏障，哪个 peer 的数据先到就先算。代价是写放大：每行要 push world_size 份，这个伏笔到 benchmark 再收。

@triton.jit
def _poll_pairs(ptr_u32, offs, mask):
    v = tl.load(ptr_u32 + offs, mask=mask, other=0, volatile=True)
    while tl.max(tl.where(mask & (v == 0x80008000), 1, 0)) == 1:
        v = tl.load(ptr_u32 + offs, mask=mask, other=0, volatile=True)
    return v

Lamport push + poll 示意（rank 0 的某一行）

  rank 0 的 bf16 行          peer 0 的 comm buffer（预填哨兵 0x80008000）
       │                              │
       ▼                              ▼
  ┌──────────┐    scrub + pack     ┌──────────────┐
  │  h[0]    │ ─────u32[0]────────►│  0x80008000  │ ◄── volatile load 自旋
  │  h[1]    │ ─────u32[0]────────►│  0x80008000  │     直到不再是哨兵
  │  h[2]    │ ─────u32[1]────────►│  0x80008000  │
  │  h[3]    │ ─────u32[1]────────►│  0x80008000  │
  └──────────┘                     └──────────────┘
       │
       ├─────────u32[*]─────────► peer 1 buffer
       ├─────────u32[*]─────────► peer 2 buffer
       └─────────u32[*]─────────► peer 3 buffer

  哨兵值 0x80008000 = 两个 bf16 -0.0，真实数据 scrub 后不可能等于它
  读到非哨兵 → 数据已到，拆包求和

注意这里 volatile=True 是必须的，否则编译器会把这个 load 提到循环外，自旋就永远读的是同一个寄存器值。

4. 三缓冲防 ABA。 由于 buffer 要跨 launch 复用，如果这次的哨兵还没来得及被下一次的写覆盖……总之直接复用同一块会出竞态。flashinfer 的做法是三缓冲：region = flag % 3，本次用一块、清另一块，flag 每次 launch 递增。

关于 Lamport 这个名字：它来自 Leslie Lamport（对，就是 LaTeX、Paxos、逻辑时钟那位）关于无锁共享内存通信的工作（Concurrent Reading and Writing, 1977）——不用锁、不用特殊原子指令，纯靠读写和单字原子性来协调，数据的值本身就是信号。负零轮询正是这个思路的 GPU 版。NVIDIA 的 TensorRT-LLM 把这种一发式 all-reduce 命名为 Lamport，flashinfer 又 vendor 了这套 kernel，名字就沿用下来了。

Benchmark

这两个算子都是 launch-bound 的，eager 下测出来大头全是 Python/launch 开销，看不到真实的 device 延迟。所以笔者用 CUDA graph 把 host 端开销抹掉之后再测。4×GB200，bf16，每次调用的延迟（µs）：

M	N	lamport	barrier	nccl+norm	lam vs base
1	2048	7.32	8.77	39.31	5.37×
8	2048	7.52	12.23	43.41	5.77×
128	2048	9.67	17.62	47.80	4.94×
1	4096	8.66	9.16	39.86	4.60×
128	4096	12.78	19.21	55.38	4.33×
1	8192	10.74	11.61	41.73	3.89×
128	8192	18.74	27.10	57.93	3.09×

几个观察，正好回收前文的伏笔：

和 NCCL + eager 比快 3~6 倍，对应背景里算的那笔账：省掉的就是两次独立 launch、NCCL 的多跳和中间那趟 HBM 往返。batch 越小固定开销占比越大，加速比越高。
Lamport 处处快过 barrier，小尺寸差距最大——这就是方案二开头说的，省掉的那趟屏障 round-trip 在 latency-bound 区间最值钱。
大尺寸差距收窄——方案二 push + poll 的写放大开始还债：Lamport 要 push 给所有 peer，写 world_size× 的数据，变成 bandwidth-bound 之后带宽开销追上来了。

总结

用 symmetric memory 让各卡直接读对方显存，一个 Triton kernel 就能把 all-reduce + residual + RMSNorm 全做完，small batch 下比 NCCL + eager 快好几倍。
同步有两条路：signal-pad barrier（简单鲁棒、支持任意 dtype）和 Lamport 负零轮询（无屏障、更低延迟，但只支持 bf16，需要三缓冲 + scrub）。

参考资料：

flashinfer include/flashinfer/comm/trtllm_allreduce_fusion.cuh
PyTorch Symmetric Memory
Leslie Lamport, Concurrent Reading and Writing, CACM 1977

ulysses-parallel 下一篇