BCG（breakable CUDA Graph）
Breakable CUDA Graph（BCG）将模型前向组织为 CUDA Graph 图段与 Eager 函数交替执行的序列：稳定计算通过图重放降低 Kernel 启动开销，动态操作在图段之间正常执行。SGLang 在 Prefill 中利用这一机制，将常规 Attention 留在图外，并捕获周围的投影、归一化和 MLP。
本文以 SGLang 提交 0b3bb0cbe3 的 CUDA 常规 Llama Prefill 路径为例。BCG 是框架层对已有 CUDA Graph 能力的组织方式；每个图段仍是独立的 torch.cuda.CUDAGraph。
1. Prefill 捕获约束
1.1 Kernel 启动开销
一次模型前向包含大量 GPU Kernel。Eager 模式下，CPU 执行框架代码并逐次提交设备操作；当 Kernel 较短时，提交开销可能使 GPU 在相邻计算之间等待。CUDA Graph 预先捕获操作及其依赖，后续通过重放提交这些工作，减少重复经过 Python、框架和驱动的开销。
可重放区域需要保持捕获时的操作结构和设备地址等约束。输入值可以变化，但通常需要先写入固定地址的输入缓冲区；仅替换 Python 变量指向的新张量，不会自动修改图中已经记录的地址。
1.2 请求结构的动态性
Prefill 通常把多个请求本轮新增的 token 拼成一个扁平 Batch。设请求数为
Norm、线性投影和普通 MLP 主要沿 token 行执行计算。固定 positions 张量更新。
Attention 还要读取请求边界、序列长度、缓存前缀长度和 KV Cache 索引。相同的
BCG 将需要运行时处理的部分设为显式断点，使其周围的计算仍能使用 CUDA Graph。支持静态元数据准备的 Attention 后端也可以采用 Full Graph；Attention 位于图外是本文分析路径的设计选择，并非 CUDA 的普遍限制。
2. 分段原理
2.1 Attention 边界
Llama 的一层计算可以简化为：
Norm → QKV 投影 → RoPE → Attention → O 投影 → 残差/Norm → MLP这里的 Attention 指消费 Q/K/V 的注意力后端调用。QKV 投影、RoPE 和 O 投影虽然属于模型的 Self-Attention 模块，却位于 BCG 断点之外。
在没有其他断点的情况下，跨层捕获结构为：
Graph 0
Embedding → 第 0 层 Norm → QKV 投影 → RoPE
↓
Eager：第 0 层 Attention
↓
Graph 1
第 0 层 O 投影 → 残差/Norm → MLP
第 1 层 残差/Norm → QKV 投影 → RoPE
↓
Eager：第 1 层 Attention
↓
……
↓
Graph L
最后一层 O 投影 → 残差/Norm → MLP → Final Norm
↓
Eager：LM Head 与 Logits Processor图段由断点之间的连续执行区域决定，可以跨越 Decoder Layer 边界。 若有
上述顺序依据 Llama 的 Attention 与 Decoder Layer 实现及 RadixAttention 断点整理。残差与 Norm 可能融合执行，图示表达数据流，不代表每个节点都对应一个独立 Kernel。
2.2 捕获与重放
eager_on_graph 标记需要留在图外的函数。捕获执行到标记处时，框架结束当前图段，执行该函数并记录重放所需的调用状态，然后开启下一张图的捕获。
捕获：capture_begin → 稳定计算 → capture_end
↓
执行并记录 Eager 函数
↓
capture_begin → 后续计算
重放：Graph 0 → Eager 0 → Graph 1 → Eager 1 → ……分段结构在捕获时确定，在线执行直接使用记录好的序列。这里的 Eager 仍可以提交 GPU Kernel；它表示该区域每次运行普通框架代码。图段之间的顺序也不等于每段结束都进行 CPU–GPU 同步，设备执行仍由 Stream 顺序和必要的依赖约束。
2.3 捕获策略对比
| 策略 | 图段的形成方式 | 图外区域 | 主要约束 |
|---|---|---|---|
| Full CUDA Graph | 对选定捕获区域整体建图 | 捕获区域之外的准备与尾部处理 | 整个捕获区域需要兼容 |
| TC Piecewise CUDA Graph | torch.compile 跟踪后在注册的算子边界切分 | 注册的切分操作 | 需要满足编译器跟踪要求 |
| BCG | 捕获过程中遇到显式标记就结束并新建图段 | 标记的普通函数 | 图段本身及跨段数据需要满足重放约束 |
BCG 与 TC Piecewise 都能形成分段执行结构，区别在于构建方式。BCG 的分段不依赖 torch.compile 理解整个前向，但仍需要开发者选择正确断点；它不会自动把任意动态 Python 流程转成可重放的图。
3. Token 分桶与数据传递
3.1 Token 分桶
算子分段解决捕获兼容性，Token 分桶解决输入形状变化。设预先捕获的 token 数集合为
图内隐藏状态按 [\hat{T}, D] 的固定形状执行，真实 token 数
| 请求组成 | 新增 token 总数 | 图内行数 | Attention 使用的信息 |
|---|---|---|---|
| 一个请求，200 tokens | 200 | 256 | 一个请求的长度与 KV 索引 |
| 两个请求，80 和 120 tokens | 200 | 256 | 两个请求各自的长度与 KV 索引 |
| 一个请求，240 tokens | 240 | 256 | 当前请求的真实长度与 KV 索引 |
前两种情况可以复用同一桶的图段，但 Attention 必须按各自的请求边界执行，防止跨请求访问。投影会为 padding 行付出额外计算；常规 Attention wrapper 则截取真实 token 行，再调用后端。
分桶也与 Chunked Prefill 分属不同层次：Chunked Prefill 决定本轮处理哪些 token，分桶决定本轮计算适配哪组图，BCG 断点决定这次前向中的哪些操作留在图外。
3.2 固定输出缓冲区
图 A 生成的 Q/K/V 和图 B 消费的 Attention 输出必须通过稳定地址衔接。SGLang 的常规路径在断点之前分配 output，将它作为参数交给图外 Attention：
Graph A：向固定地址写入本轮 Q/K/V
↓
Eager Attention：读取 Q/K/V 与当前 Batch 元数据
↓
向固定 output 地址写入 Attention 结果
↓
Graph B：从 output 读取结果，执行 O 投影等计算后端支持指定输出地址时，直接写入 output[:T]；否则返回临时结果，由 wrapper 复制到该区域。以常规 GQA 为例，Q 和输出可表示为 [\hat{T}, H_q, D_h]，K/V 为 [\hat{T}, H_{kv}, D_h]；
这条 Attention 包装函数返回 None，通过原地修改 output 传递结果。BCG 的通用 _copy_output() 还负责支持的返回值拷贝场景，但普通 Attention 并不依赖它传递主输出。
3.3 内存与执行顺序
各图段共享图内存池，BCG 对部分已捕获张量保存弱引用视图，避免 Python 引用让每层中间结果都长期占用独立空间。本文版本还在不同捕获尺寸之间复用主体输出缓冲区。
内存复用依赖既定的重放顺序：前一段写入的数据必须在被后续工作覆盖前消费完毕。Eager 函数能够动态执行，并不意味着它可以任意改变下一图段要求的输出形状、布局或地址。相关实现见 BCG 内存与捕获状态和 BCG Backend。
4. SGLang 实现
4.1 Runner 与 Backend
PrefillCudaGraphRunner 负责捕获尺寸、静态输入、Batch 元数据和前向入口；BreakableCudaGraphBackend 负责把一次主体前向捕获成图段序列，并保存各尺寸的图与输出。
准备和执行流程如下：
初始化
遍历捕获 token 桶
→ 构造对应形状的 ForwardBatch
→ 准备 Attention 元数据与预热
→ 捕获 Transformer 主体，记录图段与断点函数
在线 Prefill
检查重放条件 → 选择 token 桶
→ 填充固定输入缓冲区
→ 更新当前 Batch 的 Attention 元数据
→ 交替重放图段和执行 Attention
→ 执行 LM Head 与 Logits Processor本文版本只对 BCG 的 Transformer 主体建图，输出尾部保持 Eager，以适应请求数量与 logits 处理需求。普通 Attention 在重放前调用 init_forward_metadata()；选择专用协议的后端则更新按桶保存的元数据对象，不能将普通路径推广到所有模型。
入口见 capture_one_shape() 与 load_batch()、execute()。
4.2 Attention 断点
RadixAttention.forward() 在 Extend 模式且存在分段执行 context 时分配输出，再根据 BCG 状态选择包装函数。虽然 context 名称仍包含 tc_piecewise，BCG 也复用它来读取当前 Batch 信息，这不表示 BCG 经过了 torch.compile。
📄 radix_attention.py：断点与输出写入
# RadixAttention.forward：省略非 BCG 分支及输出维度差异处理
...
output = torch.empty_like(q)
if is_in_breakable_cuda_graph():
breakable_unified_attention_with_output(
q, k, v, output, save_kv_cache, self.layer_id, **kwargs
)
...
# unified_attention_with_output：省略参数整理与平台分支
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
...
real_num_tokens = forward_batch.num_token_non_padded_cpu
query = query[:real_num_tokens]
...
forward_batch._attn_output = output[:real_num_tokens]
ret = get_attn_backend().forward(
query, key, value, attention_layer, forward_batch, save_kv_cache, **kwargs
)
...
if ret.data_ptr() != output.data_ptr():
output[:real_num_tokens].view(ret.shape).copy_(ret)
...
breakable_unified_attention_with_output = eager_on_graph(True)(
unified_attention_with_output
)4.3 重放序列
BreakableCUDAGraph 保存 _segments 和 _break_fns。重放图段之后立即调用对应断点函数，断点函数再从当前 context 获取服务本轮请求所需的状态。
📄 breakable_cuda_graph.py：顺序重放
def replay(self) -> None:
...
try:
for i, seg in enumerate(self._segments):
seg.replay()
if i < len(self._break_fns):
self._break_fns[i]()
finally:
...图内普通 Python 不会在重放时重新执行。因此，需要每轮更新的 Python 状态必须由 Runner 或断点函数处理；图内 Kernel 则通过固定地址读取更新后的张量值。
5. 性能与代价
BCG 的收益来自减少图内区域的重复提交开销，并跳过 TC Piecewise 建图所需的编译器跟踪流程。它本身不减少模型 FLOPs，也不自动融合 Kernel。
| 方面 | 作用与代价 |
|---|---|
| 稳态延迟 | 减少图内 Kernel 的逐次提交；图段启动和 Eager 调用仍有开销 |
| 启动时间 | 不经过 torch.compile 切图，但仍需预热并捕获各 token 桶 |
| 兼容性 | 动态操作可留在断点中，图内区域仍需满足捕获约束 |
| Padding | 桶覆盖更多输入长度，同时为投影等计算增加无效行 |
| 输出衔接 | 固定缓冲区维持地址稳定，部分后端需要额外设备内复制 |
| 显存 | 图及其缓冲区常驻，内存池复用可降低重复分配，但不能视为零成本 |
当大量小 Kernel 的提交间隙占据明显时间时，BCG 更有机会降低 Prefill 延迟；长输入的大型 GEMM 或 Attention 已占据主要时间时，启动开销的相对占比会下降。断点越多、图外区域越大，能够消除的提交开销也越有限。
验证时应固定模型、Attention 后端、并行配置、请求长度分布和并发度，分别比较 Eager、BCG 与支持的其他捕获方式。预热后同时观察 Prefill 延迟、TTFT、GPU Kernel 间隙、padding 比例和图内存，另行记录启动建图时间。本文分析源码与执行机制，未执行 GPU 性能测试。