企业官网建设流程全解析

TokenSpeed CUDA Graph深度指南：Prefill图、KDA子图与Breakable Graph原理

【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeed

🚀 TokenSpeed 是一款光速 LLM 推理引擎，其中CUDA Graph是它压榨 GPU 性能的核心武器之一。本文将用通俗易懂的方式，带你搞懂 TokenSpeed 中三张关键的技术牌：Breakable Graph（可打断图）、Prefill Graph（预填充图）和KDA 预填充子图，并说明它们各自解决什么问题、如何配置。

TokenSpeed 与其他推理引擎在 Kimi K2.5 上的吞吐对比，CUDA Graph 正是其领先的关键因素之一

为什么 LLM 推理需要 CUDA Graph

想象一下：模型每一层都有几十个 kernel（卷积、归一化、GEMM……），每个 kernel 都要由 CPU 上的 Python 代码"发射"一次。CPU 发射指令的速度跟不上 GPU 执行的速度时，GPU 就会出现一段段等待空隙——就像流水线上的工人等料，机器空转。

CUDA Graph的解法很直接：把一系列 kernel 调用在启动时"录制"成一张图，之后每次只需要一次replay（重放）命令，就能让整个序列在 GPU 上背靠背跑完，几乎消除 CPU 发射开销。

• ✅Decode 阶段：每步形状固定，天然适合录制整图重放
• ❌Prefill 阶段：请求长度千变万化，直接整图录制不可行——这就引出了 TokenSpeed 的 Breakable Graph

Breakable Graph：给 Prefill 图装上"打断键"

核心代码位于 breakable_cuda_graph.py。它的设计思想是：

把一次前向传播拆成若干图段（segment），在特殊位置插入打断点（break）。

工作原理可以理解为"三段式接力"：

1. 图段 A：录制一段纯计算 kernel（投影、归一化、MoE 等）
2. 打断点：注意力 / KV-Cache 操作的元数据依赖真实请求数据，无法录制，就在这里结束录制，让注意力以 eager 模式（即时模式）执行
3. 图段 B：从打断点之后继续录制，直到下一个打断点

重放时按顺序调用每个图段即可。整个机制有两条"承重不变量"：

• 所有图段共享同一个 CUDA 内存池，保证图中间结果的显存地址在每次重放间稳定
• 打断点输出的张量必须落回固定地址——为此图段内会预分配 handoff 缓冲区，把 eager 计算结果拷回固定地址，下一个图段再从这里读取

这样 TokenSpeed 无需torch.compile，就实现了"分段 CUDA Graph"：注意力留在图外处理数据依赖，其余部分全部图化。

Prefill 内核延迟对比：内核层面的每一微秒优化，都会被 CUDA Graph 的调度放大

Prefill Graph：分桶捕获，一次录制多次复用

MoE 等被图段捕获的 kernel，正是 Prefill Graph 消除启动开销的对象

Prefill 的 token 数不可能固定，怎么办？TokenSpeed 的答案是容量分桶（bucket）。核心实现在 prefill_graph.py：

• 图按(token 桶, 请求数桶)两个维度在启动时捕获
• 服务时，请求只选择一个现成的桶重放，绝不在服务中新增图
• 桶比实际大就填充 padding（用带掩码的 dummy token 占位），不会变成真实调度请求

配置示例（完整说明见 server.md）：

桶的粒度是权衡：桶越疏，捕获越快但 padding 浪费越多；桶越密，padding 越少但启动捕获时间和显存占用越高。

KDA 预填充子图：把线性注意力也"关进图里"

对于 KDA（Kimi 系列的线性注意力）模型，普通的"注意力打断点"反而成了瓶颈：KDA 层由一连串短 kernel（状态准备、卷积、门控投影、递归扫描、状态写回）组成，逐层从 Python 发射会留下大量 GPU 空隙。

TokenSpeed 的做法是内联捕获（inline capture）：符合条件的 KDA 批次不再打断，而是把 KDA 与相邻的投影、gated RMSNorm、输出投影一起录进同一个外层图段，连续 KDA 层可以共享一个段。完整设计文档见 kda-prefill-subgraphs.md。

对于带内部检查点的 KDA 层，逻辑流程为：

输入投影 → KDA main：处理到检查点边界 → 保留边界检查点 → KDA tail：从边界继续扫到请求终点 → 恢复 token 顺序、归一化并投影输出

main 和 tail 是同一次外层捕获的一部分，不是单独选择的图，从而省掉了逐层 KDA 图发射和打断处的输出拷贝。

桶复用一个精巧的细节：每个捕获为每个请求预留一个 checkpoint/tail 槽位，因此"哪些请求带检查点、带几个"可以变化而不需要重新捕获。文档中举了个真实例子：两个请求分别新增 868 / 869 个 token（granularity 128），main 段各扫 768 个、tail 段各扫 100 / 101 个，两者都落进(2048 tokens, BS 2)这一个捕获里。

兜底策略：图不是万能的，但一定有人兜底

工程上最见功力的是降级路径清晰而保守：

• 混合 prefill/decode 批次、请求数超出捕获容量 → 走普通外层图（保留注意力打断点）
• token 数超过最大桶→ 该 forward 走 eager；但长提示词仍可让"分块调度后恰好落进桶里"的 chunk 用上图
• 数据并行、layerwise PD 缓存传输→ 保留普通路径
• DeepSeek-V4.1 这类行收窄模型：某一层行数按"哪些请求完成"而变，无法用一张 token 形图表达。模型按NarrowingPrefillModel契约声明拆分，录制为 encoder 图 → eager 收窄 → decoder 图两段家族

设计文档 unified_path.md 中特别强调：捕获失败是致命错误，而不是静默降级——不能在启动时捕获的架构必须在类属性里静态声明，避免"看起来能用实则没上图"。

关键文件导航 📚

总结

TokenSpeed 的 CUDA Graph 体系是一套层层递进的方案：

1. Breakable Graph用"图段 + 打断点"解决 Prefill 变长问题
2. Prefill Graph用 token/请求双维度分桶，让一张图服务一类形状
3. KDA 子图进一步消除线性注意力的逐层发射空隙
4. 清晰的兜底与错误策略保证系统可预测性

对使用者的建议：开启图功能时围绕真实流量选择 token 桶（默认--prefill-graph-max-tokens即可起步），避免过疏的桶造成大量 padding 浪费。🎯

【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeed