Skip to content

feat(cache): add exact hybrid checkpoints for agent continuations - #1558

Open
sufubao wants to merge 5 commits into
ModelTC:mainfrom
sufubao:cache-optim
Open

feat(cache): add exact hybrid checkpoints for agent continuations#1558
sufubao wants to merge 5 commits into
ModelTC:mainfrom
sufubao:cache-optim

Conversation

@sufubao

@sufubao sufubao commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

混合注意力模型的旧缓存把线性状态终点绑定在 hash 页边界,难以复用 Agent 的非对齐输入和已生成前缀。本 PR 增加默认关闭的精确检查点缓存:CPU KV 仍按固定容量页搬运,线性状态按实际 token 长度保存,GPU KV 淘汰后仍保留 CPU 检查点目录。

例如,12000-token 输入生成16个 token 后可以保存 S@12015;下一轮复用输入和已计算的输出,只补算最后一个尚未生成 KV/state 的 token 及新增内容。完整输入命中使用保存的 hidden 执行 HEAD_ONLY,按当前请求参数重新采样。

实现

  • 独立 CPU 前缀目录保存配套的 KV、conv/SSM、hidden;整页共享、不可变尾页 COW、LRU、lease 和 generation 管理。origins 贯穿 CPU、GPU radix 和 PD,避免混用不同计算历史。
  • GPU stream 冻结有界候选,CPU post 核验接受/停止边界。普通模式由后台线程准备 CPU 数据,调度线程在全 TP 完成后发布;请求和 staging 持有到复制结束。已有完整检查点经全 TP 确认后复用,跳过重复复制。
  • 启动预留 pinned arena;批量执行 HEAD 与 draft 尾部修复;移除伙伴计算后的过时空闲休眠和输出达到上限后的多余 forward。normal 完整命中的 KV/索引/conv/SSM/hidden 恢复只在结尾等待一次,lease 持有到复制完成;逐页复用 GPU 临时空间,部分命中/PD 保留显式等待。
  • 固定 MTP 完整命中的首轮只处理有效 target 行;GDN、FA3 FP/MLA 使用一致的可变布局。辅助 kernel 的行数不再逐值特化,启动预热有限 block 版本。DP 单批 draft 修复走普通执行路径,正常双微批仍使用 CUDA Graph。
  • Qwen3.5 单层 vanilla_with_att / eagle_with_att adapter、精确 accepted 状态、私有 draft 尾槽,以及 P→D 来源传递、D→P 缺页回流和首 token owner 协议。
  • 保留旧 CPU cache 历史尾页匹配修复及实际匹配终点传递。
--enable_exact_prefix_cache --exact_prefix_cache_mb 4096 \
--exact_prefix_cache_entries 128 --exact_prefix_cache_page_size 8192 \
--exact_prefix_cache_capture_slots 4

CPU 数据预算按 TP rank 计算;页容量不限制状态长度。新模式不能与旧 --enable_cpu_cache / --enable_disk_cache 同开。详见设计文档

H100 验证与性能

原 MTP 负载保持不变:Qwen3.5-0.8B、MTP step3、8并发、257输入/8输出、4GiB CPU预算、小页256、大页8192(256×32)、GPU KV容量65536。对照 upstream 1eb4810c;最新 R15 两次各20轮,每侧320次warm,704个HTTP请求、1312项比较通过,未剔除样本。

均值 同轮 upstream 基线 最新候选 R15
TTFT 75.71 ms 50.32 ms
TPOT 2.10 ms 2.54 ms
总延迟 90.67 ms 68.35 ms

剩余 TPOT 回退未解决。 R15 的2.5369ms与历史R12的2.5366ms基本相同;同轮baseline更快,因此当前差距20.8%,不能把历史13%的比例继续当作最新结果。总延迟降低24.6%,该短输出测试不能代表饱和吞吐。

本次确认的收益限于恢复搬运:实际CPU cache/state方法的H100微测,8请求只补末尾KV的恢复4.857→4.687ms,8193-token多页恢复23.416→22.855ms;多页GPU临时空间峰值224→约112MiB。使用真实张量形状、合成数据,未加载模型权重。随机数据、非连续目标槽位、跨页、clear/lease及copy异常通过逐字节验证;CPU专项42场景/310检查和真实CPU cache 759检查通过。不能用此微测宣称服务TPOT改善。

  • R13完整MTP预提案和R14单候选预提案未减少target验证轮数、未取得稳定收益,均已撤回。保留其全部测量和失败记录。
  • 128输出的配套20轮:baseline/R15 TPOT1.394/1.578ms,总延迟248.12/249.76ms。656比较中4项失败,涉及2个warm请求;IDs相同,第16个token标量logprob差0.033751875超过0.03。独立R12对照也出现相同位置/概率差,失败仍如实保留。
  • R15 Agent/EOS/stop/分支28请求、15组恢复比较通过;Triton非greedy32请求、16次warm完整命中通过;两轮真实GPU容量压力后25请求检查及3比较通过,日志确认完整8193与追加后8210均从CPU8193/GPU0恢复。严格冷算3组中2组失败:Agent续接logprob差0.06488;stop续接恢复首token16、冷算760。在同GPU用相同输入重跑R12,28请求的IDs/logprob/命中长度与R15全部一致,两项失败均复现;详见验证记录。恢复重放通过不等于完整冷算等价。
  • 历史R12 FlashInfer非greedy、DP2双微批MTP、较早27B和PD结果保留各自快照范围,本次不冒充重新验证。完整PR的45个文件pre-commit通过;未新增单元测试文件。

所有实验通过 exp -m 记录,保留不可变源码、启动参数、逐token原始请求及失败结果。本次部署tar SHA256 3924999ad1712246ee50abdf5ae5b31ea55c6a5fb31eeab9cfec7c4bd8546df2;最终仅增加docstring澄清,已核对执行AST一致。详见完整记录

当前限制

CPU onload 和 PD 发布仍有同步等待,新请求查找可能等待复制持有的缓存锁;尚非全异步 CPU cache。27B R8 的8并发 TPOT 也仍有10.5%回退,后续改动未重新跑该模型,不能宣称全面优化。

只缓存已计算、可核验的状态;不额外 seal 最后采样 token,不保证迟到 abort/字符串 stop 的精确末态。D 暂仅复用完整输入或差一个 token 的检查点;回流使用 CPU HTTP,路由不保证 owner 亲和。多模态、prompt logprobs/routed experts 回退,其他不支持组合和权重身份约束详见设计文档。

共享请求和PD结构有变更,部署需统一更新并重启相关进程。

@sufubao sufubao changed the title fix(cache): reuse historical CPU tail pages for hybrid attention feat(cache): add exact hybrid checkpoints for agent continuations Sep 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant