实际情况恰好相反，decode 是推理框架使用 CUDA Graph 的典型场景。

本文先从一次 kernel launch 的提交过程出发，解释 CUDA Graph保存了什么，再讨论其动态性，以及 PyTorch 如何管理 capture期间产生的临时显存。

一次 kernel launch提交了什么

CPU 启动 kernel 时，提交的是“执行这段 GPU程序所需的工作描述”。理解这个过程，需要区分三类内容：

例如：

1

scale<<<4, 256, 0, stream>>>(x_ptr, y_ptr, 1024, 2.0f);

这里需要指定运行哪个 kernel、启动多少线程，以及四个实参。对于x_ptr，传入的是地址值；它指向的整个 tensor 不会随这次launch 自动复制。kernel代码通常也可以在加载后重复使用，无须每次重新传输。参见 CUDADriver API。

下面以 NVIDIA GPU 上普通的 host launch为例，分别讨论代码加载、参数准备和驱动提交。

代码加载与 kernel 启动

首先需要区分 kernel 的代码与某一次执行它的请求。

CUDA C++ 或 Triton 产生的 device 代码，最终需要转换为目标 GPU能执行的机器指令。在 NVIDIA 平台上，PTX 是虚拟指令集，cubin中包含面向具体架构的机器码，反汇编时通常称为 SASS。如果加载的是PTX，驱动可以在加载过程中将其 JIT 编译为机器码；如果已经有适合当前架构的cubin，则可以加载该二进制。

Driver API 用 module 和 function handle 表达这些对象。例如：

1
2
3
4
5

// 假设 CUDA context 已初始化，scale.cubin 中包含名为 scale 的 kernel。
CUmodule module;
CUfunction function;
cuModuleLoad(&module, "scale.cubin");
cuModuleGetFunction(&function, module, "scale");

这一步使驱动能够找到 kernel 的代码和相关 metadata。后续可以使用同一个function handle 多次 launch，并不需要每次重新编译和传输整份 kernel代码。lazy loading可能将部分准备工作推迟到首次使用，但不改变代码加载与重复执行的区别。参见CUDADriver API 的 module 与 kernel execution。

参数如何传给 GPU

假设 scale 的签名是：

1

extern "C" __global__ void scale(const float* x, float* y, int n, float alpha);

在 Driver API 层，一次启动可以写成：

1
2
3
4
5
6
7
8
9
10
11
12
13

// d_x、d_y 是 CUdeviceptr 类型的 host 变量，保存已分配的 device buffer 地址。
// stream 是有效的 CUstream。
int n = 1024;
float alpha = 2.0f;
void* args[] = {&d_x, &d_y, &n, &alpha};

cuLaunchKernel(
    function,
    4, 1, 1,       // grid dimensions
    256, 1, 1,     // block dimensions
    0, stream,     // dynamic shared memory、stream
    args, nullptr
);

这里的 extern "C" 保证导出的 kernel 符号名为scale，与前面的 cuModuleGetFunction调用一致。cuLaunchKernel 在 CPU 上执行。args中的 &d_x 告诉驱动从哪个 host 地址读取实参；读取到的d_x 才是传给 kernel 的 device 地址值。参见 cuLaunchKernel参数说明。

1
2
3
4
5
6
7

CPU 内存                           GPU

&d_x ──→ d_x = 0xA000
                  │
                  └──复制参数值──→ x = 0xA000
                                         │
                                         └──读取──→ device buffer 中的数据

图中的地址仅作示意。n 和 alpha 同样由 host按值传入，因此 kernel 参数不必都是 device指针。需要区分的是：传递指针值与搬运指针指向的数组，是两件不同的事。

从提交到执行

驱动将函数、参数和执行配置组织为 GPU 可以处理的工作描述，再通知 GPU执行。这里需要区分启动计算的命令与计算本身的指令：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

CPU / host
    框架调用或 CUDA C++ launch
          ↓
    Runtime / Driver：准备参数与工作描述
          ↓
    提交工作，通知 GPU
          │
          ▼
GPU
    命令前端取得工作，处理相关依赖
          ↓
    将 thread blocks 分派给具备资源的 SM
          ↓
    SM 内的 warp scheduler 发射就绪指令
          ↓
    执行 kernel 的 load / store / compute

host 指定 grid 和 block，具体的 block 分派和 warp 调度由 GPU完成。因此，CPU 提交工作所需的时间，与 GPU 执行 kernel的时间需要分别考察。

CUDA Graph 记录了什么

在普通的 eager 执行中，CPU 依次调用算子，经过框架的 dispatcher 和CUDA runtime，将 kernel 提交到 GPU。对于计算时间很短的小算子，host端的开销可能成为显著的性能瓶颈。

CUDA Graph 将一组 GPU 工作及其依赖关系预先定义为图，实例化后通过一次graph launch 提交执行。stream capture是构建这张图的一种方式，另外也可以通过显式的 Graph API 添加节点。CUDA将这个过程划分为 definition、instantiation 和 execution三个阶段，PyTorch 通常将最后一个阶段称为 replay。参见 CUDAGraphs 文档。

对于一段简单的逐元素计算，可以将其理解为：

1
2
3
4
5
6
7
8

eager：
    CPU 提交 add → CPU 提交 relu → CPU 提交 mul

capture 后：
    graph = [add → relu → mul]

replay：
    CPU 提交 graph

这里的三个 kernel 仍然是三个 kernel。CUDA Graph本身不会将它们融合成一个kernel，也不会自动消除中间结果的显存读写。kernel fusion 与 CUDA Graph可以配合使用，但解决的是不同层面的问题。

API 层的图与执行时的表示

kernel 节点的参数结构 cudaKernelNodeParams 包含func、gridDim、blockDim、sharedMemBytes，以及kernelParams / extra所描述的实参。节点间的依赖另外表达。参见 cudaKernelNodeParams。

以两个 kernel 为例，图在语义上可以表示为：

1
2
3
4
5
6
7

node A：scale(x_ptr, tmp_ptr, 1024, 2.0)
        grid=(4,1,1), block=(256,1,1)

node B：relu(tmp_ptr, y_ptr, 1024)
        grid=(4,1,1), block=(256,1,1)

dependency：A → B

cudaGraph_t是图定义的句柄，对应上述节点、参数和依赖。调用cudaGraphInstantiate 后，CUDA校验图并完成执行所需的准备，返回一个可执行图实例的句柄cudaGraphExec_t。后续 cudaGraphLaunch接收的是这个实例，可以反复启动它：

1
2
3
4
5
6

// graph 是已经建好的 cudaGraph_t，stream 是执行所用的 CUDA stream。
cudaGraphExec_t executable;
cudaGraphInstantiate(&executable, graph, 0);

cudaGraphLaunch(executable, stream);
cudaGraphLaunch(executable, stream);  // 复用同一个实例，再执行一次

1
2
3
4
5
6
7
8
9

capture / 显式建图
    ↓
cudaGraph_t：图定义（节点、参数、依赖）
    ↓ instantiate
cudaGraphExec_t：可执行图实例，供后续 launch 使用
    ↓ upload（可以合并到首次 launch）
device 端可用的工作描述
    ↓ launch / replay
按依赖执行各节点

“不可变”的具体含义

指针值与指针指向的内容

考虑下面的 CUDA kernel：

1
2
3
4
5
6

extern "C" __global__ void scale(const float* x, float* y, int n, float alpha) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        y[i] = x[i] * alpha;
    }
}

假设 capture 时的调用是：

1

scale<<<4, 256, 0, stream>>>(x_ptr, y_ptr, 1024, 2.0f);

那么 graph 记录的是这次 launch：使用哪个 kernel、grid 和 block的大小，以及x_ptr、y_ptr、1024、2.0f这些参数值。

每次 replay 时，kernel 都会从同一个 x_ptr读取数据，向同一个 y_ptr 写入结果。但 x_ptr指向的 1024 个浮点数可以改变。graph保存了指针值，并没有保存这块显存内容的快照。

因此，修改输入的正确方式是在 replay 之前更新原有 buffer。下面假设new_x 已准备完成，拷贝与 replay 在同一 stream 上提交：

1
2

static_x.copy_(new_x)
g.replay()

而下面的赋值只改变 Python 变量的绑定关系：

1
2

static_x = new_x
g.replay()

graph 不会重新查找 static_x 这个名字，也不会因此改用new_x 的地址。如果原来的输入 tensor失去引用，其内存还可能被 allocator 回收复用。因此，capture之外分配的输入、权重等 buffer，需要由调用方保证生命周期。

标量也可以变成运行时数据

上面的 alpha 是一个按值传入的标量，capture 时已经固定为2.0f。即使随后修改 host 端的 alpha 变量，graph中的参数也不会随之变化。

如果需要每次 replay 使用不同的缩放系数，可以改为：

1
2
3
4
5
6
7
8

__global__ void scale_dynamic(
    const float* x, float* y, int n, const float* alpha_ptr
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        y[i] = x[i] * alpha_ptr[0];
    }
}

此时固定的是 alpha_ptr，缩放系数则由 GPU在执行时从显存读取。只要在 replay 之前更新该地址的内容，并保证更新与replay 之间的执行顺序，系数就可以变化。

同一个数值，作为 host 标量传入，或者作为 device buffer 的内容传入，对CUDA Graph 来说具有不同的含义。后面讨论的 sequence length、block table和条件判断，都与这个区别有关。

固定 shape 与可变的有效长度

继续使用前面的 scale。假设输入 tensor 的 shape 是[1024]，host 根据元素数量计算 grid，再启动 kernel：

1
2
3

int n = 1024;                    // 来自输入 tensor 的元素数量
int blocks = (n + 255) / 256;     // 每个线程处理一个元素
scale<<<blocks, 256, 0, stream>>>(x_ptr, y_ptr, n, 2.0f);

capture 得到的是 grid=4、n=1024 的这次launch。之后 replay 不会重新执行上面计算 n 和blocks 的 host 代码。

如果输入改为 [2048]，普通执行会重新计算，得到grid=8、n=2048。但原 graph 仍然启动 4 个block，并向 kernel 传入 1024：它只处理前 1024个元素。即使原来的输入输出 buffer足够大，这张图也不会自动扩大计算范围。

这就是 shape 限制的具体来源：由 shape 决定的 launch配置和参数，在 capture 时已经确定。 对普通 PyTorch CUDAGraph，输入需要保持 capture 时的 size 和 layout，参见 PyTorch的约束说明。

padding 与 kernel 内部的分支

CUDA Graph 并不要求 kernel内部每次执行完全相同的指令路径。

通过 padding，可以将不同长度的输入放入相同 shape 的buffer。为便于展示，假设 buffer 容量固定为 8，两次输入的有效长度分别为 5和 7：

1
2
3
4
5
6
7

下标：          0  1  2  3  4  5  6  7

第一次输入：   [1, 2, 3, 4, 5, 0, 0, 0]   有效长度 = 5
               └──有效数据──┘ └padding┘

第二次输入：   [1, 2, 3, 4, 5, 6, 7, 0]   有效长度 = 7
               └─────有效数据─────┘  padding

本例约定 0 表示 padding，有效数据均非零。kernel读取输入值，遇到 padding 时直接返回：

1
2
3
4
5
6
7
8
9
10
11
12
13

__global__ void scale_padded(
    const int* x, int* y, int alpha
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int value = x[i];
    if (value == 0) {
        return;  // padding 位置不参与计算
    }
    y[i] = value * alpha;
}

// 仅为演示：固定启动 8 个线程，每个线程对应一个元素。
scale_padded<<<1, 8, 0, stream>>>(x_ptr, y_ptr, 2);

第一次 replay 中，i = 5 的线程读到 0后直接返回；第二次 replay 中，i = 5 的线程读到6 并执行乘法。分支条件由 GPU执行时读取的输入值决定，无须修改 graph。

调用方只使用有效数据对应的输出，padding 位置不会被此 kernel更新。

为什么 decode 可以使用 CUDAGraph

首先需要明确 LLM 如何组织计算。用 T 表示本轮输入的总token 数，H 表示 hidden size，hidden states 通常表示为[T, H]：每一行对应一个 token 的隐藏向量。

常见 Transformer 中的 Linear、RMSNorm 和 MLP 都是 token level的运算：对每个 token 的隐藏向量独立应用相同的变换。例如，Linear 的矩阵乘 Y = X @ W 可以按行理解为：

1
2
3
4
5
6

X：[T, H]           W：[H, D]          Y：[T, D]

token 0 的向量  ────── @ W ──────→ token 0 的输出
token 1 的向量  ────── @ W ──────→ token 1 的输出
token 2 的向量  ────── @ W ──────→ token 2 的输出
      ...                              ...

这些 token 共享权重W，但每一行的结果不依赖其他行。实现时通常将所有行合并为一次矩阵乘，以提高GPU 利用率。RMSNorm 则沿每行的 hidden dimension 归一化，MLP也分别变换每个 token 的向量。因此，对这些算子而言，T 个token 来自一条请求还是多条请求，不影响逐 token 的计算方式。

attention 负责 token 之间的信息交互：每个当前 token 的query，需要读取该请求中允许访问的 token 的 K/V。因此，attention除了当前输入，还需要知道请求边界和历史 KV 长度。

在使用 KV cache 的普通单 token decode 中，每条请求每轮只输入一个新token。Linear、RMSNorm、MLP 等算子只处理本轮的 T 个token；attention 通过 KV cache 读取历史信息，无须重新计算历史 token的这些逐 token 变换。只要本轮 token 数 T 固定，当前 hiddenstates 就保持为 [T, H]，即使历史 KV 长度仍在增长。

例如，连续两轮 decode 的状态可以是：

1
2
3

                    第 t 轮              第 t+1 轮
hidden states       [4, H]               [4, H]
sequence lengths    [128, 256, 80, 512]  [129, 257, 81, 513]

sequence lengths tensor 的 shape 始终为[4]，变化的是其中的四个整数。

对于采用 paged KV cache 的实现，KV cache 可以来自预先分配的物理 block池。逻辑序列访问哪些 block，由 block table 描述。跨越一个 block边界时，框架可以分配新的逻辑映射，更新表项，而不移动整个 KV cache。

attention kernel 启动时，传入的是 sequence lengths buffer的指针，而不是其中每条请求的长度值。capture 固定了这个指针；每轮 decode只需更新 buffer 中的长度，例如从 [128, 256, 80, 512] 更新为[129, 257, 81, 513]。kernel 在 GPU上执行时，再从同一个地址读取本轮的长度，确定需要访问多少历史 KV。

sequence lengths 的内容改变，不需要更新 kernel的启动参数。 在 kernel、grid、block等其余启动配置也保持不变的前提下，不同 sequence length 就可以复用同一张CUDA Graph。

block table 和 slot mapping 也可以采用相同的方式：保持 buffer地址不变，通过更新表项指定读取哪些 KV block，以及将新 token 的 K/V写入哪里。

下面是这一类实现的数据流示意，具体字段和布局由 backend 决定：

1
2
3
4
5

固定地址的 hidden states ───────────────┐
固定地址的 KV cache block 池 ───────────┤
固定地址、内容可更新的 sequence lengths ─┼─→ attention kernel
固定地址、内容可更新的 block table ─────┤
固定地址、内容可更新的 slot mapping ────┘

这并不意味着所有 attention kernel都天然支持这种用法。如果某个实现根据 host 端的最大序列长度选择 split数量、分配 workspace，或者改变 launch grid，那么仅更新 device metadata就不够了。backend 仍需要提供与 CUDA Graph 兼容的调度方式。

再考虑 T 的变化。continuous batching中，请求会加入或退出，每轮输入的 token数也会随之变化。常见做法是预先为若干 T capturegraph，例如：

1
2
3

capture sizes：1, 2, 4, 8, 16, 32
实际 T：     13
选择 graph： T = 16

框架将输入填充到 16，并通过有效长度、mask 或安全的占位 metadata处理多出的条目。尤其是 KV cache 写入，padding请求必须具有正确的屏蔽语义。

因此，decode 的动态性通常由两部分共同承担：graph 外部选择适合当前T 的图，graph 内部通过 device metadata处理变化的序列状态。比如vLLM 的 batch descriptor 和 CUDA Graphdispatcher 就是在这一层管理 capture 与 replay，参见其 CUDAGraphs 设计文档。

prefill attention的限制来自哪里

普通 decode 中，每条请求本轮只有一个 query token；prefill中，每条请求本轮的 query token 数也会变化。即使总 token 数T 相同，attention 的 launch grid 仍可能不同。

例如，两次 prefill 都有两条请求、总共 256 个 token，但请求长度分别是[128, 128] 和 [192, 64]。Linear、RMSNorm、MLP的输入都可以保持为 [256, H]；attention则需要按请求处理各自的 query。

以 FlashAttention 2 的普通 forward 路径为例，host 按最长 query长度计算分块数，再为每条请求、每个 head 启动对应数量的 block。假设选用的kernel 每个 block 处理 128 个 query token，则有：

这里的分块数按 ceil(最长 query 长度 / 128) 计算；128是本例选取的 query 分块大小。对应逻辑见 FlashAttention的 grid 计算，最长 query 长度由 hostAPI 传入。

如果按第一组输入 capture，图中为每条请求记录的只有一个block。第二组输入中，第一条请求需要两个 block；直接 replay原图，就缺少处理后 64 个 query token 的 block。更新 devicebuffer 中的请求长度和边界，不会改变已经捕获的 grid。

这就是仅按 T 复用 prefill graph可能失败的具体原因。要让这两组输入复用同一张图，可以提前按足够大的 query长度上限建图，让 kernel 从 device metadata读取实际长度并跳过多余分块，同时保证其他启动参数和 buffer地址也兼容。

如果 attention 路径仍需要 CPU 每轮重新计算 launch，框架可以采用piecewise CUDA Graph：只捕获 attention 前后的计算，attention每轮正常启动。

动态控制流与 torch.cond

前面讨论的是同一个 kernel 内部的动态性。如果数据决定的是执行哪一组kernel，就涉及 graph 层面的控制流。

例如：

1
2
3
4

if pred.item():
    y = x.sin()
else:
    y = x.cos()

假设 pred 是 CUDA tensor，这里的 .item()需要将结果交给 CPU，Python 才能决定下一步调用哪个算子。这种依赖 GPU结果的 host 分支不能直接放入普通 stream capture。

如果 pred 是 Python bool，capture时可以完成分支判断，但只会记录当时实际调用的分支。之后修改这个bool，replay 也不会重新执行 Python 的 if。

CUDA 的 conditional nodes提供了另一种表达方式：预先构建分支或循环体，由 device端的条件决定执行哪部分图。条件可以通过 kernel 中的cudaGraphSetConditional 设置，判断过程不需要将控制权交回CPU。参见 NVIDIA 的 ConditionalNodes 介绍。

PyTorch 中对应的高层表达是 torch.cond。这里使用torch.compile，是为了把 torch.cond表达的条件分支转换成 CUDA Graph 的条件节点。

下面将分支前的 x * 2、条件分支和分支后的z * 3 写入同一个 forward，再对整个forward 使用 cudagraphsbackend。capture 的范围是整个forward，包括分支前后的计算。 backend在首次几次调用中自动完成 warmup 和 capture，后续调用则 replay整张图：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

import torch


def positive_branch(h):
    return h + 1


def negative_branch(h):
    return h - 1


def forward(pred, x):
    h = x * 2
    z = torch.cond(pred, positive_branch, negative_branch, (h,))
    return z * 3


compiled_forward = torch.compile(forward, backend="cudagraphs", fullgraph=True)

with torch.inference_mode():
    x = torch.ones(8, device="cuda")
    pred = torch.tensor(True, device="cuda")

    # backend 在这些调用中完成整个 forward 的 warmup、capture 和首次 replay。
    for _ in range(3):
        compiled_forward(pred, x)

    pred.fill_(True)
    y = compiled_forward(pred, x)   # replay，结果为 8 个 9.0

    pred.fill_(False)
    y = compiled_forward(pred, x)   # replay 同一张图，结果为 8 个 3.0

对应的整张 CUDA Graph 可以表示为：

1
2
3
4
5
6

                    forward 的 CUDA Graph
┌────────────────────────────────────────────────────┐
│                  ┌─ true：z = h + 1 ──┐             │
│ h = x * 2 → 条件 ┤                    ├→ y = z * 3  │
│                  └─ false：z = h - 1 ─┘             │
└────────────────────────────────────────────────────┘

两次调用更新的是同一个 device bool tensor pred的内容。compiled_forward(pred, x) 内部通过一次 graph launch执行分支前的计算、选中的分支和分支后的计算。示例没有显式的with torch.cuda.graph(g):，是因为整张图的 capture 和 replay都由 backend 管理。

普通 CUDA Graph capture 本身不要求torch.compile。在本文实测的 PyTorch 2.13.0+cu130 中，仅将torch.cond 放入with torch.cuda.graph(g):，仍会在读取条件时触发 capture错误；提前 warmup 两个分支也不能改变这条执行路径。上面的 backend用法参见 DataDependent Control Flow。

torch.cond 保留了两个分支的结构，使运行时选择可以在device 端完成。示例中的两个分支都产生相同 shape、dtype 和 layout的输出，这也使后续计算能够使用一致的 buffer 约定。

它并不意味着任意 Python 控制流或任意动态输出形状都能进入 CUDAGraph。更准确的理解是：分支结构提前建立，每次 replay可以选择不同的执行路径。

capture 期间的执行语义

“执行”在这里需要区分 host 与 device。

在 with torch.cuda.graph(g): 内，Python代码仍然执行，算子的 host 端实现也会被调用。框架需要计算 metadata、选择kernel、分配输出 tensor，并调用 CUDAAPI，这样才能知道应当向图中添加哪些节点。

但提交到正在 capture 的 stream 中的 GPU工作不会在此时执行，而是被记录为 graph 节点。capture_end得到图，后续 graph launch 才会执行其中的工作。这是 CUDAstream capture 的定义。

下面的例子可以同时观察这两个层面的行为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

import torch


static_x = torch.zeros(4, device="cuda")
static_y = torch.full_like(static_x, -1)

# 在 side stream 上 warmup，使算子完成必要的初始化。
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
    for _ in range(3):
        torch.add(static_x, 1, out=static_y)
torch.cuda.current_stream().wait_stream(warmup_stream)

# 恢复一个便于观察的初值。
static_y.fill_(-1)
torch.cuda.synchronize()

g = torch.cuda.CUDAGraph()
host_calls = 0

with torch.cuda.graph(g):
    host_calls += 1
    torch.add(static_x, 1, out=static_y)

# 以下读取均位于 capture 之外。
print(host_calls, static_y.tolist())
# 预期：1 [-1.0, -1.0, -1.0, -1.0]

static_x.fill_(3)
g.replay()
print(host_calls, static_y.tolist())
# 预期：1 [4.0, 4.0, 4.0, 4.0]

static_x.fill_(5)
g.replay()
print(host_calls, static_y.tolist())
# 预期：1 [6.0, 6.0, 6.0, 6.0]

capture 之后，host_calls 已经增加，但static_y 仍然是事先写入的 -1。replay 之后，GPU输出发生变化，Python 计数器却不会继续增加。

这里使用原地写入，是为了保留一个确定的输出初值。如果在 capture 内使用static_y = static_x + 1，PyTorch 同样会创建输出tensor，但在首次 replay 前，不应把它的内容当作这次加法的结果。

一些封装会在 capture 前运行 warmup，在 capture 后立即replay，因此观察外层函数时，可能会看到输出已经完成计算。那是封装额外执行的结果，不能据此认为stream capture 本身执行了被记录的 kernel。

为什么同步和读取结果会失败

capture 期间，GPU 工作尚未被提交执行，自然不能要求 CPU等待其中某个结果完成。常见的不兼容操作包括：

• 对正在 capture 的 stream 调用 cudaStreamSynchronize或查询执行状态。
• 对 captured event 调用 cudaEventSynchronize或查询完成状态。
• 调用覆盖当前 capture 的 device 级同步，例如torch.cuda.synchronize()。
• 通过 CUDA tensor 的 .item()或需要同步的拷贝，将数据读回 CPU 并参与后续计算。

这些限制针对的是 capture 中的执行状态及其依赖，并不是所有与 stream 或event 有关的 API 都被禁止。例如，符合 capture 规则的cudaEventRecord 和 cudaStreamWaitEvent可以用来建立图内依赖。CUDA 文档对不允许的操作和跨stream 依赖作了分别说明。

这也解释了为什么“异步 API”不一定就能 capture：如果 host后续逻辑必须读取拷贝结果才能确定launch，仍然没有消除数据依赖。是否兼容，需要检查完整调用路径。

CUDA Graph 与 stream

stream 用来组织异步提交的 GPU 工作。在通常的执行语义下，同一 stream中的操作按提交顺序执行；不同 stream 之间需要通过 event等机制建立依赖。CUDA Graph则将这些操作及其依赖保存下来，供后续重复执行。

stream 在 capture 和 replay 中承担不同的作用：

多个 stream 如何进入同一张图

如果在一个 stream 中依次 capture 四个kernel，得到的依赖就是一条链。要表达可以并行的分支，可以在 capture中使用多个 stream。

下面用四个操作构造一个例子。所有 buffer 都提前分配，并在warmup、capture 和 replay 期间保留引用：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

import torch


static_x = torch.zeros(8, device="cuda")
z = torch.empty_like(static_x)
a = torch.empty_like(static_x)
b = torch.empty_like(static_x)
static_y = torch.empty_like(static_x)
s0 = torch.cuda.Stream()
s1 = torch.cuda.Stream()


def workload():  # 在 s0 为当前 stream 时调用。
    torch.mul(static_x, 2, out=z)       # K0

    s1.wait_stream(s0)
    torch.add(z, 1, out=a)             # K1，在 s0 上
    with torch.cuda.stream(s1):
        torch.mul(z, 3, out=b)         # K2，在 s1 上

    s0.wait_stream(s1)
    torch.add(a, b, out=static_y)       # K3，在 s0 上


s0.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s0):
    for _ in range(3):
        workload()
torch.cuda.current_stream().wait_stream(s0)

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g, stream=s0):
    workload()

第一次 wait_stream 使 K2 依赖K0，第二次使 K3 依赖 K2。结合s0 自身的顺序，最终得到：

1
2
3

              ┌── K1：a = z + 1 ──┐
K0：z = x * 2 ┤                    ├── K3：y = a + b
              └── K2：b = z * 3 ──┘

K1 和 K2 都读取z，但写入不同的输出，因此可以并行；K3必须等待两者完成。是否实际重叠执行，还取决于 GPU 资源和调度。

这里的 wait_stream 通过 event 表达依赖，capture时会将依赖加入图中。它不要求 CPU 等待 kernel 执行完成。参与同一次capture 的其他 stream，需要在结束前通过依赖汇合回起始 stream；仅切换当前stream，并不会自动建立这种关系。参见 PyTorch多 stream capture 示例。

PyTorch 如何处理 allocator

当我们计算一个巨大的神经网络时，除了权重是固定的、输入是给定的，很多中间隐藏层的输出都是临时创建、计算、销毁的。以不记录梯度的forward 为例：

1
2
3
4
5
6

@torch.no_grad()
def forward(x, w1, w2):
    h = x @ w1
    h = torch.relu(h)
    y = h @ w2
    return y

在这个函数返回时，局部变量 h 不再被 Python 持有。但 CUDAGraph 中的后续 replay仍然会向中间结果对应的地址写入数据，再从这些地址读取。显然，不能直接将这些地址交给无关tensor 使用。

另一方面，如果每个临时 tensor都必须永久保留独占显存，深层网络的内存开销又会非常大。解决这个问题，需要区分tensor 对象的生命周期、allocator 中 block的复用，以及底层显存的释放。

缓存分配与 stream 顺序

先看普通 eager 执行。PyTorch 的 native caching allocator会缓存已经申请的显存块，tensor 释放时，通常只是将对应 block归还缓存，并不立即调用 cudaFree。

对于同一 stream 上的使用，allocator可以利用顺序关系复用内存，不必每次都等 GPU 完成。

假设某段代码的内存使用关系如下：

1
2
3
4

kernel A：写入临时 buffer p
kernel B：读取 p，写入另一个 buffer
          p 对应的临时 tensor 释放
kernel C：将 p 复用为自己的输出

即使 CPU 提交 kernel C 时，GPU 尚未执行完 kernelA，也不妨碍这种复用。只要 A、B、C 按上述依赖顺序执行，C 对p 的覆盖就发生在 B 读取之后。

跨 stream 则不能仅凭 host 调用顺序判断安全性。allocator 通常需要结合record_stream 和 event 等机制，确认其他 stream 不再使用该block。相关实现可以参考 CUDACachingAllocator.cpp开头关于 stream 与内存复用的说明。

graph-private memory pool

capture 期间，PyTorch 的 native caching allocator 会将相应分配导向graph-private memory pool。临时 tensor仍然可以创建和释放，满足依赖条件的 block也仍然可以在池内复用；需要额外保证的是，这些地址在后续 replay 时仍由该graph 使用。

可以将前面的例子延伸为：

1
2
3
4
5
6
7
8

capture 期间记录：
    A 写 p → B 读 p → C 写 p

第一次 replay：
    A 写 p → B 读 p → C 写 p

第二次 replay：
    A 写 p → B 读 p → C 写 p

p 可以先后服务于不同临时结果。每次 replay都保留相同的依赖关系，因此也保留了复用所需的先后顺序。

replay 时不会重新进入 Pythonforward，也不会重新执行其中的 tensor 分配逻辑。kernel直接使用 capture 时确定的地址。这意味着 allocator 不需要在每次 replay中重新构造上述分配与释放过程。

PyTorch 源码中的Note [Interaction with CUDA graph capture]专门说明了这一点：private pool 内仍然允许复用，但 graph存活期间会保留该池所需的内存高水位。参见 allocator源码。

由此也可以解释，为什么 capture 后显存占用可能高于普通eager：普通路径中的空闲 block 可以被其他工作复用，而 graph 的 privatepool 需要为未来 replay 保留地址。这个占用并不等于“所有层的临时 tensor大小相加”，还取决于生命周期重叠、workspace、池内复用和碎片。

输出 tensor也具有同样的地址约束。如果需要保留某次输出用于之后的计算，应在它被后续replay 覆盖之前复制：

1
2
3

g.replay()
saved_output = static_output.clone()
g.replay()  # static_output 的内容会被覆盖

这里假定 clone 与 replay 具有正确的 stream 顺序。仅写saved_output = static_output 不会保存结果的副本。

private pool 的生命周期还与 graph 和其中仍然存活的 tensor 有关。销毁graph 并不意味着所有输出引用同时失效，因此也不保证显存立刻全部归还给CUDA。具体生命周期见 Graphmemory management。

capture 期间的底层显存分配

private pool 并不意味着所有内存都必须在 capture前精确分配完毕。需要区分 PyTorch 的 tensor 分配、底层cudaMalloc，以及 CUDA Graph 自身的内存节点。

对于 native allocator，已有可复用 block 时，tensor 分配主要是 host端的缓存管理。如果确实需要新的底层分配，当前源码中的cudaMallocMaybeCapturing 会在 capture 状态下临时使用cudaStreamCaptureModeRelaxed，再进行分配。

这种分配在 capture 阶段取得真实地址，随后由 private pool保留，并不会在每次 replay 时重新申请一份显存。对应实现同样位于 CUDACachingAllocator.cpp。

这与 cudaMallocAsync 被捕获为 graph 的 allocation node是不同的机制，后者由 CUDA 提供图内分配与释放语义，参见 GraphMemory Nodes。不能将 native caching allocator的工作方式直接套用到所有 allocator backend。

这里的 relaxed mode 也不表示 kernel 会在 capture 时执行，更不会让依赖GPU 结果的 host同步变得有效。它只是允许框架在自行保证地址生命周期的前提下完成这类分配。

无法查询 event 时如何回收

剩下的问题是，allocator 在 eager 下可能需要查询 event，capture时却不能照常依赖这些执行状态。

PyTorch 对此有专门的 capture 路径。native allocator 在 capture期间跳过常规的 process_events，对需要跨 stream跟踪的部分回收操作作延后处理。同一 stream中可以根据依赖顺序证明安全的复用，仍然能够进行。

较新的实现还提供 graph_capture_record_stream_reuse选项，尝试利用图的拓扑关系判断部分跨 stream内存何时可以安全复用。这里使用的是图中已经建立的依赖，而不是查询尚未执行的kernel 是否完成。对应逻辑可以从 allocator源码中的 process_events、deferred_blocks和该配置项继续阅读。

因此，capture中的内存管理并不是暂停所有分配和释放，而是将“可以根据依赖关系确定的复用”与“需要运行时完成状态才能确定的回收”分别处理。

多张 graph 能否共享内存池

如果每个 T 对应的 graph 都使用完全独立的 privatepool，推理服务会为许多不会同时执行的 graph 重复保留内存。PyTorch允许通过 pool 参数让不同 capture 共享池，例如：

1
2
3
4
5
6
7
8
9

# 假设输入已经准备好，两个 workload 都已完成 warmup。
g1 = torch.cuda.CUDAGraph()
g2 = torch.cuda.CUDAGraph()

with torch.cuda.graph(g1):
    out1 = workload1(static_input1)

with torch.cuda.graph(g2, pool=g1.pool()):
    out2 = workload2(static_input2)

共享池只是允许复用，不会自动判断调用方的数据依赖是否安全。对于前后相接的图段，需要维持与capture 兼容的执行顺序和 tensor生命周期，不能让后一个图段覆盖前一个图段尚未消费的数据。

对于相互独立、只会择一执行的不同 T 对应的graph，可以在不并发执行且旧输出不再需要时复用池内空间。PyTorch文档也列举了推理服务中多个 graph 共享池的用法，参见 Sharingmemory across captures。

从这个角度看，capture为临时内存确定了一套可以重复使用的地址安排。Python局部变量可以结束生命周期，但图所依赖的地址、执行顺序和内存所有权必须继续成立。

要解决的问题

长上下文推理有两个主要瓶颈：

• KV cache 显存随长度线性增长。MLA 虽然比 MHA/MQA省得多，但到 1M 长度依然扛不住。
• attention 计算量。即使有 DSA（DeepSeek SparseAttention，V3.2 引入），长上下文下的计算仍然是大头。

核心设计

V4 的 attention 由3个组件组合而成：

1. 跨 token 压缩 KV cache（省 4x ~ 128x）：
   • c4a：约 4 倍压缩。一个压缩 token 是 8 个原始 token的加权和，步长 4（相邻压缩 token 之间有重叠）。
   • c128a：约 128 倍压缩。一个压缩 token 是 128 个原始token 的加权和，步长 128（不重叠）。
2. DSA 稀疏注意力（把计算量卡住）：c4a 压缩后 1M上下文仍有 250k 条压缩 token，所以用 lightning indexer 选出top-k（k=512，Pro 为 1024）条来算注意力。c128a 压缩后 1M 上下文最多只有8192 条，直接对全部条目做 full attention 即可（等价于top-k=8192）。
3. 128 的短滑窗（保住局部性）：作用在未压缩的 token上。特别是 c128a 下，query 在压缩边界之前可能一条压缩 token都看不了（causality），滑窗保证它至少能看到本地信息。

CSA（c4a 路线）

CSA 层的完整算子图如下。输入 hidden 为 ，query 先下投影到 1024维潜向量 （RMSNorm 后按位置加RoPE），再上投影出 64 个头、head dim 512 的 q；KV 侧共享一份 512 维的latent（K=V），一路进 128 大小的滑窗 cache（未压缩），另一路经compressor 4 合 1 压缩后进压缩 KV cache。indexer 用 FP4 打分选出 top-512条压缩条目，与滑窗条目拼接后做注意力，输出后 64 维做 Inverse RoPE。

C4A Compressor

compressor 的结构：当前块和上一块各自做投影（ 与 ，各 512维），然后把两块的权重拼起来加偏置，得到每个压缩位置 8 个权重，在 512个通道上各自做逐通道 softmax归一；内容同样拼接后按权重加和，得到一个压缩条目，最后过 RMSNorm +RoPE（RoPE 的锚点位置是 ）。第0 块时上一块内容补零、权重置 ；decode时”上一块”是最近的完整块，“当前块”是滚动缓冲。

Lightning Indexer

沿用 V3.2 DSA 的 lightning indexer，QK 通路全程 FP4：query 侧从 投影出 64 个头、128 维的索引query，另有 产生逐头权重；key侧用与主路同款的 compressor（4 合 1）生成 128 维的索引键并以 FP4缓存。打分是 FP4 GEMM 后接 ReLU（负分直接置零），再按头加权求和（分数存BF16），最后在满足因果性的范围（）内取 top-512。

HCA（c128a 路线）

HCA 层没有 indexer：压缩比 128 使得 1M 上下文也只有至多 8192条压缩条目，直接对”滑窗 128 条 + 全部压缩条目”做稠密注意力（MQA形式，K=V，含 attention sink）。其余结构（潜向量维度、64 头、逐头RMSNorm、Inverse RoPE）与 CSA 一致，RoPE 锚点位置为 。

KV cache 账

• DeepSeek V3.2 风格的 KV cache：每层每 token 存 MLA latent B + indexer B = 1408 B，61 层共约83.9 GiB / 序列
• DeepSeek V4-Pro：每层 CSA 存压缩 KV 条 + indexer 条 = 320 MiB，每层 HCA 只存 条 = 8 MiB，共约9.62 GiB / 序列，省了约 8.7x

主要参考： - vLLM 的博客 DeepSeek V4 in vLLM:Efficient Long-context Attention

• DeepSeek V4技术报告

• https://zhuanlan.zhihu.com/p/2036707452293535593

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 算掉。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

融合前后对比（一次 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 的东西）：

1
2
3
4

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

单个 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 symmetricmemory 干的事情：

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

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 和实现风险。笔者实现了两种：

读者可能会问：一个 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还带一块小的对称内存，专门用来做同步——可以把它理解成一块所有卡都能读写的记分板，布局和约定如下：

1
2
3
4

每个 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 步，放行。

1
2
3
4
5
6
7
8
9
10
11

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。

落到代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

@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在某一轮被吃掉，后面的轮次里它就永远缺席：

1
2
3
4
5
6

轮次   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），这样真实数据永远不可能等于哨兵。

1
2
3
4
5
6
7
8

@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 的数据先到就先算。代价是写放大：每行要 pushworld_size 份，这个伏笔到 benchmark 再收。

1
2
3
4
5
6

@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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

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、逻辑时钟那位）关于无锁共享内存通信的工作（ConcurrentReading 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）：

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

• 和 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-padbarrier（简单鲁棒、支持任意 dtype）和 Lamport负零轮询（无屏障、更低延迟，但只支持 bf16，需要三缓冲 +scrub）。

参考资料：

• flashinferinclude/flashinfer/comm/trtllm_allreduce_fusion.cuh
• PyTorchSymmetric Memory
• Leslie Lamport, Concurrent Reading and Writing, CACM1977

attention 之前，每张卡拿的是一段 sequence、完整 hidden；attention的时候，每张卡换成完整 sequence、一部分 head。

也就是说，Ulysses Parallelism 做的事情，本质上就是在 Transformerblock 内部找了一个很自然的分界点：MLP、Linear 这类 token-wise的计算适合沿着 sequence 切；但是 Attention 需要看到完整上下文，所以在进Attention 之前，要把数据重新洗牌成「每张卡负责一部分 head」的形态。

背景：为什么要序列并行

长上下文训练或者多模态训练里，sequence length往往会变得很长。假设输入张量是 ，其中 是序列长度， 是 hidden size。

如果每张卡都拿完整的 ，那attention 里的激活、KV、中间 attention score都会跟着长序列一起膨胀。最直接的想法就是：既然 batch 可以切，sequence当然也可以切。把 按卡数 切开，每张卡只拿 个token，这样每张卡上的激活显存自然就降下来了。

问题在于，Attention 不像 MLP。

MLP 对每个 token 是独立计算的，local token 只需要 local hidden state就能算；但是 self-attention 里，每个 token 都要和完整序列上的 K/V交互。如果只是朴素地把 sequence切掉，那么每张卡只能看到局部上下文，结果肯定不对。

所以 Ulysses 要解决的问题可以概括成一句话：

如何让大部分计算保持 sequence parallel，同时在 Attention处又能算出数值等价的 full attention？

具体介绍

Ulysses 的做法是两个 All-To-All。

第一个 All-To-All 发生在 QKV projection 之后、Attention之前，把数据从：

1

每张卡持有：N/P 个 token，全部 attention heads

换成：

1

每张卡持有：完整 N 个 token，1/P 的 attention heads

这样每张卡就可以在自己负责的 head 上独立算 full attention。

第二个 All-To-All 发生在 Attention 输出之后，把数据再换回：

1

每张卡持有：N/P 个 token，全部 hidden

这样后面的 、MLP、下一层Transformer block 又可以继续按 sequence parallel 的方式执行。

注意这里有一个非常关键的点：Ulysses 不是把单个 head 内部的 headdimension 再切碎，而是按 head 切。也就是说，假设 attention head 数量是num_heads，并行度是 ，通常需要满足num_heads % P == 0。如果一个 head 自己的 hidden维度被切掉，那 的数值就不完整了。

为了讲起来简单，下面先假设 正好等于 head 数量。实际工程里只要 head 数能被 整除即可。

第一步：先沿 sequence 切输入

一开始，输入 沿着 sequence维度切分。每个 rank 上拿到一个 local block：

1

Local X: [N/P, d]

在 Ulysses 里，单独使用时参数通常没有像 Tensor Parallelism 那样按hidden 维度切开，所以每张卡上都有完整的 。于是每张卡可以直接用自己的 local token 去做 QKVprojection：

1

Local Q/K/V: [N/P, d]

这个时候的 Local Q/K/V 有一个特点：sequence 维度是不完整的，只有；但是 hidden/head维度是完整的，包含所有 attention heads。

到这里为止，MLP 类似的 token-wise 计算都很舒服，但 Attention还不行。因为每个 head 需要完整序列上的Q/K/V，当前每张卡只拿到了序列的一小段。

第二步：对 head 维度做All-To-All

现在要做第一次数据重排。我们把每个 rank 上的 Local Q/K/V 按 head维度切成 份：

1

[N/P, d] -> P 个 [N/P, d/P]

然后对这些小块做 All-To-All。直观理解就是：每张卡把自己手里不同 head的那一小段 sequence 分别发给对应的 rank；同时也从其他 rank 收到同一个head 在其他 sequence 分片上的数据。

以 rank 0 为例，All-To-All 之前，它只有自己那段 sequence 上所有 head的 Q/K/V；All-To-All 之后，它会收到所有 rank 发来的 head 0分片。把这些分片沿 sequence 维度 concat 起来，就得到了：

1

Qh/Kh/Vh: [N, d/P]

也就是某一组 attention heads 在完整 sequence 上的 Q/K/V。

这个地方就是 Ulysses 最妙的地方。第一次 All-To-All做完以后，数据从「按 sequence 分片」变成了「按 head 分片」。而 Attention恰好天然可以按 head 并行，因为不同 head 之间本来就是相互独立的。

第三步：每张卡独立算 fullattention

现在每个 rank 都有自己负责的 head 在完整 sequence 上的Q/K/V，所以可以直接执行标准的 attention：

1

Oh = softmax(Qh @ Kh.T) @ Vh

输出形状仍然是：

1

Oh: [N, d/P]

这一步和普通 MHA 没有本质区别，只是每张卡只算一部分head。换句话说，Ulysses 并没有近似 attention，也没有把 attention改成局部attention；它只是通过通信把数据换了个摆法，然后在每张卡上算数值等价的full attention。

第四步：再做一次 All-To-All换回来

Attention 算完后，每个 rank 手里是：

1

Oh: [N, d/P]

但后面的 和 MLP更希望看到：

1

Local O: [N/P, d]

也就是每张卡重新只负责一段 sequence，但是 hidden维度完整。于是这里需要第二次 All-To-All。

这一次的通信方向可以理解为第一次的反操作：每个 rank 把自己完整sequence 上的一部分 head，按 sequence 再切回 份，然后发给对应rank。通信结束后，同一个 sequence 分片上的不同 head会重新聚到一张卡上。

concat 之后，每个 rank 得到：

1

Local O: [N/P, d]

这就回到了 Transformer block 一开始的 local sequence 形态。后面的output projection 和 MLP 都可以直接在 local token 上继续算：

1
2

Local O @ Wo -> [N/P, d]
MLP(Local O) -> [N/P, d]

通信量为什么是 Ulysses的卖点

Ulysses 的核心卖点之一是通信量比较低。forward 里主要有两次通信：

• Attention 前，对 Q/K/V 做一次 All-To-All，消息量大约是
• Attention 后，对 O 做一次 All-To-All，消息量大约是

所以总消息量可以粗略看成 。但All-To-All 的一个重要性质是，当总消息被分散到 个 rank 上时，每条链路上的通信量大约是。

这就带来一个很有意思的结论：如果 sequence length 变长，同时并行度 也能跟着变大，那么单链路通信压力可以不随着 线性爆炸。

相比之下，Megatron-LM 里的 Sequence Parallelism 通常要和 TensorParallelism 绑定，通信模式也更多依赖 All-Gather /Reduce-Scatter。Ulysses 这个「两次 All-To-All完成形态转换」的设计，在长序列场景下就显得非常直接。

当然，漂亮归漂亮，限制也很明显。

第一个限制是 head 数。因为 Ulysses 是按 head 切的，所以并行度 会受 num_heads限制。尤其是 GQA/MQA 场景里，K/V head 数更少，这个限制会更明显。

第二个限制是 All-To-All 本身。All-To-All 对网络拓扑、NCCL实现、跨机带宽都比较敏感。如果机器之间互联不够好，理论上的通信量优势不一定能完全变成实际速度优势。

总结

Ulysses Parallelism 可以用一句话概括：

用两次 All-To-All，在 sequence parallel 和 head parallel之间来回切换。

进 Attention 前，把 [N/P, all heads] 变成[N, partial heads]，让每张卡都能在自己负责的 head 上算 fullattention；Attention 后，再把 [N, partial heads] 变回[N/P, all heads]，让后面的 projection 和 MLP 继续沿sequence 并行。

所以 Ulysses 的重点不是发明了一个新的 attention，而是非常巧妙地利用MHA 本身「head 之间独立」这个结构，把长序列训练里最麻烦的 full context需求变成了一次数据重排问题。

本文参考 https://zhuanlan.zhihu.com/p/5750410146

本文是笔者的理论力学（鞠国兴）课程论文，谨以此文纪念我的本科物理生涯，蓦然回首，感慨颇多

引言

《朗道力学》在有心力部分指出，只有两种类型的有心力场，其中一切有界运动的轨道是封闭的，这两种场的势能与 或者 成正比。实际上这就是 Bertrand定理，在 1873 年由法国数学家 Bertrand证明。本文将给出两种证明方法，并尝试探讨背后更深层次的物理意义。

微扰法证明

有心力场中的拉格朗日量可以写作

可以得到

为了讨论轨道的形状，我们更关心 随 的变化，作变换 ，得到 Binet 公式

圆轨道时 ，，显然有

上式其实等价于

接下来我们讨论，什么形式的力能够给出稳定闭合的轨道。

设偏差量 （ 不是很大），对于 在 附近作 Taylor 展开：

这样轨道方程可以被改写为

记 ，则上式变为

如果只保留到一阶，有：

显然我们需要 才能使轨道稳定，代入 的具体形式

因此，轨道稳定性的条件是

这个结论我们在课内已经得到。

选择合适的初值条件，简谐方程 (3) 的解为

如果要求轨道闭合，则 必须是一个有理数。根据稳定性条件 (4) 可以得到

为了进一步的计算，我们需要考虑更高阶项的贡献。将 (5) 式看作轨道方程的Fourier 展开，当然其主要贡献来自 项。可以尝试写出

我们可以大胆猜测

所以接下来我们只关注 的相关项。

开始计算：

代回 (2) 式，可以得到

比较对应余弦项的系数，有：

将 (6) 式代入 Binet 公式 (1) 得：

首先我们来处理式 (7a) 和 (7c)：

由式 (7d) 得：

式 (7b) 才是我们的核心内容：

化简可得：

该方程有三个解：。其中 是平凡的，因为它对应着圆轨道，显然是闭合的。另外两种分别对应着：