背景#

【图式数据表示与语言模型训练_交流说明_v5.docx】

图#

【这里没写好，假设规范了一个图的规则】

图是 DAG。

每个节点代表一个 token 序列，节点内部是一条链。

定义 $V(u)$ 为 $u$ 的可见上下文集合，集合元素为

定义两种边（$u\to v$）：

• 实边（inherit 行为）：$v$ 继承 $u$ 以及 $u$ 的可见上下文，后续引用 $v$ 的节点 $w$，若能看到 $v$，则能看到 $u$ 以及 $u$ 能看见的上下文。
• 虚边（artifact 行为）：$v$ 作为 $u$ 以及 $u$ 的可见上下文的总结，后续引用 $v$ 的节点 $w$，不能通过这条边看到 $u$ 以及 $u$ 能看见的上下文。

实边和虚边的区别是，该边是否具备上下文传递的性质。

一个节点的前驱不允许同时出现实边和虚边，但允许出现多条实边或多条虚边，多条实边很好理解，支持多条虚边

已有工作#

本项目工作#

测试#

Baseline 怎么选取？

同一套 vLLM + 同样的 DAG 语义上下文可见规则，和现成的 Attention 实现

• 比较：对于一个新节点的推理，暴力向上遍历

【以下为 draft】

正确性关注：

• 每个节点的可见上下文范围是否正确。
• 一段 token 序列区间被多次引用是否正确去重，包括共享祖先和重叠区间。
• 可见 token 之间的距离计算是否正确。

性能测试关注：

• 【没写】

功能#

在 SGLang 现有调度和分页 KV 管理的基础上，拟新增以下功能：

• DAG 上下文引用：将请求之间的上下文依赖表示为图，支持多个父状态和多种引用方式。

  • inherit 引入父状态的完整有效上下文，artifact 只引入指定输出区间，两者可以混用。
  • 支持公共历史分叉、多个分支合流，以及合流结果的后续引用。图结构不要求所有父状态具有公共前缀，具体可用的引用组合由模型适配层校验。

• 图上下文编译：将请求依赖转换为 Attention 可以直接读取的元数据。

  • 展开有效上下文，按 token 实例合并重复引用，处理共享祖先和重叠输出区间。
  • 构造分段页表、页内有效范围和可见性元数据，复用未变化的依赖展开结果，增量维护当前请求追加的 KV。

• 共享状态管理：使已完成请求的 KV 可以作为独立状态被后续请求引用。

  • 通过状态句柄和引用计数管理共享 KV，多个消费者只读引用父状态，新增 KV 写入各自的空间。
  • 将状态保留、请求结束和物理 KV 回收分开处理，覆盖正常完成、取消和异常退出，避免提前释放仍被使用的状态。

• 上下文位置与模型适配：将图上的可见关系、位置编码和 KV 的物理布局分开描述。

  • 根据模型规则生成位置元数据，不直接用合流后的数组下标代替 token 的位置关系。
  • 对接模型所需的 RoPE 位移、cache mapper 和 adapter 配置，读取时完成适配，保持共享父状态不变；不兼容的状态引用在执行前拒绝。

• 多段分页 Attention 后端：使用 CUDA 和 TileLang 优化图上下文的 KV 读取与合流计算。

  • 支持 prefill 和 decode 读取多个不连续的 KV 区间，处理页内切片和不满页边界，避免为了合流而完整复制父状态。
  • 融合适用的位置变换与 KV 读取，针对分支长度不均衡划分任务，通过局部 Attention 输出和 LSE 完成归并，减少中间读写与 kernel 启动开销。

• 对照执行与性能统计：为同一上下文提供可切换的执行后端，验证图语义和算子优化的效果。

  • 保留显式变换／打包后调用成熟 Attention、分段 Attention 加 LSE 归并两条对照路径，使用相同模型、输入和可见性规则比较。
  • 记录上下文编译、KV 准备、合流首 token 和后续 decode 的耗时，以及共享 KV 占用和临时缓冲区大小；同时检查数值误差与任务质量。

后续扩展：让共享状态支持 GPU／CPU 间的卸载与重载。卸载只改变 KV 的驻留位置，不改变状态身份和依赖关系；正在参与 GPU 计算的页面需要等读取完成后才能迁移或回收。

接口#

• Multiverse-Engine 已经在 SGLang（的旧版本）上实现分支调度和 KV 合流，但 Multiverse-Engine 有以下问题：

  • 受分支结构限制：要求合流时必须有公共前缀。
  • 位置编码采用特定规则：各分支从共同的位置起点展开，合流后的结论接在最长分支之后。该规则不适用于通用 DAG 上下文结构。
  • KV 存储能力有工程缺口：不支持 KV cache 的卸载与重载。

附录#

https://github.com/vllm-project/vllm