Piecewise CUDA Graph
Piecewise CUDA Graph（PCG）将模型计算拆成多个区域，分别捕获稳定部分，并在图段之间执行动态操作。SGLang 的 TC Piecewise 路径先通过 torch.compile 跟踪模型，再按注册的算子边界切分 FX 图，使 Prefill 中的投影、Norm 和 MLP 能够复用 CUDA Graph。
本文将 PCG 的分段执行原理与 SGLang 的编译器实现分开说明，源码固定到提交 0b3bb0cbe3。底层捕获与重放约束见 CUDA Graph。
1. 部分捕获需求
Prefill 可以将多个请求的新增 token 拼为 [T, D] 的隐藏状态，其中
部分后端需要根据这些运行时信息规划 Attention。如果整个前向都要求可捕获，一处动态操作就可能使周围的大量稳定计算失去图重放机会。PCG 将该操作单独留在图外：
Graph A：Norm → QKV 投影 → RoPE
↓
Eager：Attention 后端调用
↓
Graph B：O 投影 → 残差/Norm → MLP → 下一层 Norm/QKV/RoPE
↓
Eager：下一层 Attention图段可以跨层，实际边界取决于注册的操作与模型执行路径。Attention 留在图外也不意味着使用 CPU 计算；其 Python 包装函数仍可提交 GPU Kernel。
2. 编译器切图原理
2.1 FX 图与 CUDA Graph
两种图处于不同层次：
| 图 | 表达对象 | 在 TC Piecewise 中的作用 |
|---|---|---|
| FX Graph | 张量操作及其调用关系 | 支持识别边界、切分与编译 |
| CUDA Graph | GPU 工作与设备依赖 | 支持低开销重复提交 |
SGLang 先让 torch.compile 得到 FX 图，再由自定义 Backend 切分。fullgraph=True 要求目标函数形成完整的编译器图，随后仍可由 Backend 将它拆成多张 CUDA Graph。它与“整个模型只捕获一张 CUDA Graph”是不同含义。编译入口
2.2 显式切分算子
SGLang 用 register_split_op() 收集切分点。普通 Attention 同时注册成自定义算子，使编译器能够保留调用边界，并利用 Fake 实现传播张量信息；切图 Backend 再识别这个节点。
完整 FX 图
投影/Norm → Attention 节点 → 投影/MLP
↓ 按注册节点切分
子图 0 子图 1 子图 2
可捕获 保持图外 可捕获split_graph() 遇到注册节点时，为它分配独立子图，并从下一个节点开始使用新的子图编号。调用 split_module(..., keep_original_order=True) 保持原始操作顺序，尤其避免带原地写入的操作被错误重排。之后只有非切分点子图会交给分段捕获 Backend。切图实现
2.3 编译与捕获
SGLang 的 tc_compiler 支持 eager 和 inductor。选择 eager 仍然经过 torch.compile 跟踪和 FX 切分，只是不使用 Inductor 对片段进行代码生成优化；CUDA Graph 捕获仍然存在。
每个可捕获片段由 CUDAPiecewiseBackend 管理，并按具体尺寸维护 ConcreteSizeEntry。其中保存可调用计算、预热状态、CUDA Graph 与输出。运行流程可归纳为：
跟踪与切分 → 准备片段计算 → 按尺寸预热 → 捕获片段
↓
在线执行：调用切分后的模型 → 片段选中对应尺寸 → replay动态形状 FX 图与固定尺寸 CUDA Graph 可以共存：前者表达多个尺寸的计算规则，后者为选定尺寸保存具体的设备执行计划。编译缓存与图捕获缓存因此不能混为一谈。分段 Backend
3. 独立分段示例
3.1 计算与边界
下面的教学示例手工指定两段图，展示 PCG 最终的执行结构；它不调用 SGLang，也不演示自动 FX 切图。
输入形状固定为 [8, 32]，最多容纳 8 个 token，每个 token 有 32 个特征。中间使用简化的单头因果自注意力：令 Q/K/V 都等于前一段的输出，按传入的请求长度分别计算；两个请求之间不共享 Attention。
Graph A：h = tanh(0.5 × x)
↓
Eager：按请求长度计算 Attention，写入固定 bridge
↓
Graph B：y = bridge + sigmoid(bridge)请求长度可在每次重放时变化，但桥接缓冲区始终为 [8, 32]。无效 token 行清零后，后续逐行计算仍会处理 padding，消费输出时只取真实 token 行。
3.2 完整代码
示例需要 NVIDIA GPU 与支持 CUDA 的 PyTorch，无需模型权重。每次改变请求长度后都与完全 Eager 的参考结果比较。
import math
import torch
CAPACITY, DIM = 8, 32
def before_attention(x):
return torch.tanh(0.5 * x)
def after_attention(x):
return x + torch.sigmoid(x)
def eager_attention(h, lengths, out):
# lengths 是 CPU 上的请求长度列表，动态控制流保留在图外。
if not lengths or any(n <= 0 for n in lengths):
raise ValueError("Request lengths must be positive.")
if sum(lengths) > CAPACITY:
raise ValueError("Batch exceeds the captured token capacity.")
out.zero_()
start = 0
for length in lengths:
qkv = h[start:start + length]
scores = (qkv @ qkv.T) / math.sqrt(DIM)
future = torch.ones(
length, length, device=h.device, dtype=torch.bool
).triu(diagonal=1)
scores.masked_fill_(future, float("-inf"))
out[start:start + length].copy_(scores.softmax(dim=-1) @ qkv)
start += length
@torch.no_grad()
def main():
if not torch.cuda.is_available():
raise RuntimeError("This example requires CUDA-enabled PyTorch and a GPU.")
static_x = torch.zeros(CAPACITY, DIM, device="cuda")
bridge = torch.zeros_like(static_x)
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
for _ in range(3):
warm_h = before_attention(static_x)
eager_attention(warm_h, [4, 4], bridge)
warm_y = after_attention(bridge)
stream.synchronize()
del warm_h, warm_y
graph_a = torch.cuda.CUDAGraph()
graph_b = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph_a, stream=stream):
static_h = before_attention(static_x)
with torch.cuda.graph(graph_b, stream=stream):
static_y = after_attention(bridge)
for lengths in ([3, 2], [4, 4], [1, 5]):
tokens = sum(lengths)
x = torch.randn_like(static_x)
reference_bridge = torch.zeros_like(bridge)
eager_attention(before_attention(x), lengths, reference_bridge)
expected = after_attention(reference_bridge)
# 图与中间操作按当前 Stream 的顺序执行。
static_x.copy_(x)
graph_a.replay()
eager_attention(static_h, lengths, bridge)
graph_b.replay()
actual = static_y[:tokens].clone()
torch.testing.assert_close(actual, expected[:tokens])
print(f"Passed: request lengths = {lengths}")
torch.cuda.synchronize()
if __name__ == "__main__":
main()graph_b 始终读取 bridge 的原地址，因此图外函数必须原地写入。把 bridge 重新赋值为另一个张量，无法改变已捕获图的输入地址。static_h 同理：它是前一图段反复写入的输出，也是 Eager Attention 每轮读取的输入。
示例保留了独立缓冲区和图各自的内存池，便于说明生命周期。生产实现还需要共享内存池、减少复制、管理 KV Cache 和多设备通信；这些优化不能破坏跨段的数据依赖。
4. SGLang 源码示例
4.1 Attention 注册
普通 Attention 使用以下两个装饰器声明边界。output 是预分配的输出张量，通过原地写入交给下一图段。
📄 radix_attention.py：切分算子
@register_custom_op(mutates_args=["output"])
@register_split_op()
def unified_attention_with_output(
query: torch.Tensor,
key: Optional[torch.Tensor],
value: Optional[torch.Tensor],
output: torch.Tensor,
save_kv_cache: bool,
layer_id: int,
...
) -> None:
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
...源码摘录中的省略号表示省略参数与实现，不是独立运行脚本。完整函数从当前 context 获取请求信息，将 Q/K/V 截到真实 token 数，再调用 Attention 后端并写入 output。
4.2 FX 切图
下面保留了为注册节点分配独立子图的核心代码：
📄 backend.py：split_graph()
subgraph_id = 0
node_to_subgraph_id = {}
split_op_graphs = []
for node in graph.graph.nodes:
if node.op in ("output", "placeholder"):
continue
if node.op == "call_function" and str(node.target) in ops:
subgraph_id += 1
node_to_subgraph_id[node] = subgraph_id
split_op_graphs.append(subgraph_id)
subgraph_id += 1
else:
node_to_subgraph_id[node] = subgraph_id
split_gm = torch.fx.passes.split_module.split_module(
graph, None, lambda node: node_to_subgraph_id[node], keep_original_order=True
)ops 来自注册表和配置追加的切分点。本文版本还会在使用 DeepEP 或 Mooncake MoE All-to-All 后端时加入相应 MoE 操作，因此切分边界并不限于 Attention。TC Piecewise 配置
4.3 形状选择与重放
Prefill Runner 将输入补齐到捕获 token 桶，并准备当前 Attention 元数据。切分后的模型进入某个片段时，该片段根据运行尺寸选择 ConcreteSizeEntry；已捕获的分支调用 entry.cudagraph.replay() 并返回保存的输出。
编译器的动态形状支持不会让一张 CUDA Graph 自由改变尺寸。实际仍需捕获多个桶，或对未覆盖的情况采用其他执行路径。调试模式还会比较捕获与重放时的输入地址，检查静态地址约束是否被破坏。具体尺寸状态与重放
5. 性能与边界
| 方案 | 构建方式 | 主要收益 | 主要代价 |
|---|---|---|---|
| Full Graph | 整体捕获选定区域 | 图启动次数少 | 捕获兼容性要求高 |
| TC Piecewise | 编译器跟踪、切图、分段捕获 | 动态操作周围的计算可重放，可选编译优化 | 跟踪与编译准备成本、边界注册和多段启动 |
| BCG | 捕获过程中遇标记直接分段 | 无需编译器完成整段跟踪 | 仍需维护断点、地址与图外执行 |
PCG 对启动开销明显的计算区域更有效。若多数时间消耗在图外 Attention、通信或大型 GEMM 上，收益会受限。较多图段和 token 桶还会增加建图时间、缓存状态及显存管理成本。
torch.compile 的普通 Graph Break 也不等于本文的切分点：TC Piecewise 先要求可跟踪的完整图，再由自定义 Backend 切分；无法跟踪的 Python 行为可能在切图之前就导致失败。BCG 避开了这项整体跟踪要求，但图内区域仍需兼容 CUDA 捕获。
本文独立示例包含结果断言，用于检查分段前后计算是否一致；当前环境未执行 CUDA 测试，不将示例视为性能基准。