diff --git a/docs/agent_checkpoint_cache_design.md b/docs/agent_checkpoint_cache_design.md new file mode 100644 index 0000000000..be1f848aa2 --- /dev/null +++ b/docs/agent_checkpoint_cache_design.md @@ -0,0 +1,515 @@ +# 面向 Agent 的精确前缀检查点缓存设计 + +状态：已实现第一版，默认关闭。本文第 1–10 节保留目标架构与后续设计，**不表示每项建议均已落地**。以下实现范围及限制优先于后文。原代码基线：`cache-optim @ d7c6ef10`（PR #1558）。 + +## 本次实现 + +```bash +--enable_exact_prefix_cache \ +--exact_prefix_cache_mb 1024 \ +--exact_prefix_cache_entries 128 \ +--exact_prefix_cache_page_size 8192 \ +--exact_prefix_cache_capture_slots 4 +``` + +新模式使用普通 GPU token radix 和独立 CPU 检查点目录，不与旧 `--enable_cpu_cache` / `--enable_disk_cache` 同开。旧小页/大页参数不再决定新检查点的长度，也不再额外切分 prefill。关闭新开关恢复旧路径。 + +| 已实现 | 代码入口 | +|---|---| +| 压缩 token 前缀目录、任意精确终点、完整 CPU KV 覆盖、独立 conv/SSM/seed | `dynamic_prompt/checkpoint_cache.py` | +| 固定容量纯 KV 页、同计算来源整页共享、不可变尾页 COW、LRU、lease、flush generation | 同上 | +| batch 长度绑定、GPU 带 mask 的有界捕获、CPU stop 最终判定、全 TP admission | `model_infer/exact_prefix_cache.py` | +| CPU 恢复仅加载 GPU radix 缺失区间，GPU KV 淘汰不删除 CPU 状态目录 | 同上 | +| 双 infer 线程各自 ticket；microbatch overlap 使用四组 staging；伙伴计算时跳过过时的空闲休眠 | chunked prefill / DP backend | +| 独立 hidden 副本、批量 scheduler HEAD_ONLY、当前请求重新采样、独立 pinned 输出缓冲 | `basemodel.py` / `base_backend.py` | +| 精确 MTP conv 窗口与 SSM row、canonical 恢复、Qwen draft 私有尾槽重建 | `linear_att.py` / `proposers/exact_resume.py` | +| D→P CPU 缺页传输、全 TP 发布、首 token owner、D 完整命中零 KV 控制任务 | `pd/checkpoint_transport.py` / PD master | +| 原生 HTTP 边界、Agent、分支、stop、并发回归脚本 | `test/benchmark/agent_checkpoint_cache.py` | + +`exact_prefix_cache_mb` 是**每 TP rank 的缓存数据预算**，包括 KV 页物理容量、状态/seed、未发布数据及尚有 lease 的已淘汰数据。Python 元数据、HTTP 序列化、TP 转换和网络队列的临时内存另计；条目数、捕获槽及传输队列另有上限。GPU 自动容量 profiling 预留 staging 和一页 gather 临时空间。 + +启用时在启动阶段一次性预留该预算的 pinned CPU arena，请求处理中只划分和回收窗口，避免 `cudaHostAlloc` 停顿。每个窗口具有独立 tensor storage，序列化一页不会携带整座 arena；最后一个 tensor/view 释放后才归还窗口。finalizer 只向队列投递回收记录，分配时在锁内合并空闲区，避免 GC 重入锁或修改正在遍历的空闲表。已关闭 lease 和已消费的 pending 不再持有缓存实体。 + +捕获所需 CPU flags 按 staging slot 使用独立 pinned buffer，一次异步上传；已知输出长度且忽略 EOS、没有 token stop 时，直接使用 prompt/终止位置 mask。完整命中的 LM head、当前采样和 MTP 私有尾槽修复按批次执行，再统一等待。普通模式固定 MTP 的首次 decode 只验证有效 target 行，随后由正常 proposer 产生下一轮候选，避免验证尚未初始化的 draft 位置。启用该模式时，目标模型的 GDN、FA3 FP/MLA 始终使用可变行布局，保持 CUDA Graph 捕获和重放一致，SSM 的每请求物理槽跨度不变。捕获选择和请求起始位置 kernel 将实际行数作为非特化运行时参数，仅保留 block 大小的编译版本，避免新旧请求混合后按每个精确行数重新编译。启动时用空 mask 预热本 rank 容量范围内的有限 block 版本；不读写请求状态和 KV。offload 的最后一页 KV 与 state/seed 共用完成等待，多页之间仍释放上一页 gather 临时空间；异常路径同样等待已提交的读写，再归还窗口。 + +普通模式在 CPU post 决定接受/停止位置后冻结 tokens、origins 和状态视图，交给独立 CUDA stream 的后台线程写入 CPU 缓存。任务数量由 staging 槽约束；在全 TP 都完成复制之前，请求保留 GPU KV 和索引，不能释放或暂停，但可以继续计算后缀。worker 使用独立 Gloo 组，只准备尚不可见的数据；调度阶段取完成 epoch 的全 TP 交集后统一发布，再释放引用和 staging，并允许 worker 处理下一任务。目录发布与恢复因此保持串行，避免相同 tokens/长度在不同 TP rank 命中不同版本。clear 取消旧 generation 的发布，不提前释放在途资源。两个 infer 线程仍遵守原有全部握手，保证 post 发布登记与下一次分类轮询的顺序。若全 TP 的目录已经有同一精确前缀及所需 seed，worker 持 lease 共同确认后沿用既有 KV/state/seed/origins，跳过重复搬运；任一 rank 缺失仍执行原有准备流程。复用也经过完成队列和引用释放，不修改当前请求的计算来源。PD 保留原有同步发布和传输协议。 + +普通模式的 `_pre_post_handle` 只推进实际接受行；当输出长度已达 `max_new_tokens` 时，下一次分类不再安排多余 forward，仍保留原有握手和正常 post/finish/free 顺序。 + +捕获自然 prefill chunk 终点，以及 GPU token/EOS/长度提示选出的停止候选；CPU post 再核验实际接受与 stop 边界。普通 decode 生成 N 个 token，默认最多保存 `prompt_len + N - 1` 的状态：最后刚采样的 token 尚未计算。迟到的外部 abort 或字符串停止不保证留下精确末态，只保留此前合法检查点。 + +检查点额外保存每个 KV token 的计算来源 `origins`。同 token 前缀分别通过 decode/prefill 计算，可能留下不同的浮点 KV；不能把一份状态和另一条计算历史的 KV 混用。CPU 目录先验证实际 tokens，GPU radix 再按 origins 查找；CPU 页键同时包含 tokens 与 origins，COW 只复制两者都一致的前缀。每批新计算区间产生新的跨 TP 一致来源，恢复/PD 传输继承已有来源；普通 decode 不逐步复制整条前缀来源向量。 + +MTP CPU 页键还记录 successor 或终端标记，因为 draft slot i 可能依赖 token[i+1]。恢复末槽到请求私有 KV，再用 H@L 和本次 token[L] 重建 draft 尾部，并为这个 packed 槽赋新来源；捕获还冻结 packed 尾槽，防止下一批 proposal 提前改写它。DP 的单批次 draft 修复不能重放双微批 CUDA Graph，采用已有的普通执行路径；正常 DP 双微批仍使用图重放。 + +普通模式完整命中的恢复把缺失 KV、请求索引、conv/SSM 和 output seed 排在同一 CUDA stream 上，最后统一等待一次，然后发布请求恢复状态并关闭 CPU lease。`load_kv(wait=False)` 的调用方必须保持 lease 和源索引有效，直到这一等待完成；复制异常也必须等待后才能释放。逐页临时 GPU tensor 在该 stream 上循环复用，避免后一页分配时仍持有前一页。部分命中和 PD 保留各自的显式等待；该改动减少恢复屏障，不改变请求的 READY 协议，也不代表 onload 已能与活动 decode 并行。 + +### 当前限制 + +- **CPU onload 和 PD 发布仍有同步屏障，CPU cache 锁会跨页拷贝等待。** 普通模式 offload 在后台完成，活动 decode 不等待该复制；新请求的目录查找仍可能等待 worker 持有的锁。尚未实现 `LOAD_WAIT`、HEAD_ONLY 与 onload 并行，不能称为全异步 CPU cache。 +- 尾页采用 COW，没有原地追加、fragment arena；许多短分支可能浪费 CPU 容量和复制带宽。checkpoint 选择为 LRU/有界候选，没有成本模型或每会话优先级。 +- namespace 覆盖模型配置、实际加载的 safetensors（没有时为 `.bin`）文件名/大小/mtime、dtype、解析后的量化配置、expert dtype 及 draft 配置，不包含模型或量化文件路径。部署期间权重必须不可变，跨节点复制需保留权重元信息；这里不扫描权重内容计算摘要，若替换权重却保留大小和 mtime，必须更新 `weight_version` 并重启服务。启动拒绝在线 RL 更新。多模态、请求 prompt logprobs/routed experts 回退重算。 +- 完整命中要求模型具有末位置 hidden adapter。MTP 仅支持带 adapter 的单层 Qwen3.5 `vanilla_with_att` / `eagle_with_att`；其他 proposer 不使用新检查点。启动拒绝 MTP+expert parallelism、diverse mode、legacy DP cache fetch 组合。 +- 没有 seal-only forward、严格字符串 stop 历史环或后台末态重建；也没有 disk 持久化。 +- D→P 使用 CPU HTTP 传输与 canonical head 布局转换，尚非 RDMA/共享内存。相同 MTP layout 可传输；D MTP→P target-only 可剥离 draft 并重建页键，反方向拒绝。P→D 同时传递每段 KV 的计算来源，D→P 缺页协商保留这些来源。D 暂只复用完整输入或差一个 token 的检查点；其他部分命中放弃，接收 P 的全部输入 KV 和状态，避免混合两种计算历史。跨节点回流和 P 基页去重要求两端开启新模式；仅 D 开启时为收到的 KV 分配私有来源，仍可在 D 本地缓存。下一轮路由不保证返回原 P；owner 亲和及 D 短后缀 prefill 留待后续。 +- 浮点 decode 状态续 prefill 与完整冷 prefill 的计算路径可能不同，甚至使接近的 greedy 候选改变顺序。来源隔离保证 KV/state 配套，不保证不同算子/分块/批次形状之间逐位相同。验证分别检查保存时运行态的恢复等价性与冷算差异，不将后者直接计为数据传输损坏。 + +性能实验必须用 `exp -m`，记录实际 attention backend、代码/补丁、GPU、TP/DP 和容量参数。自动后端降级、首次 JIT 和客户端输出长度都可能混淆对比；命中率不是吞吐结论。 + +目标是在 LightLLM 的 token KV 寻址、CPU 页搬运、PD 分离和 CPU/GPU 折叠流水线上，支持任意已计算长度的混合模型检查点。正确性必须无条件成立；命中率与额外开销有明确策略，不承诺所有结束条件都能零代价保存精确状态。 + +## 1. 核心决策 + +采用 **精确前缀目录 + 独立状态存储 + 固定容量 KV 页 + 异步捕获事务 + 显式恢复计划**。 + +- KV 搬运页大小与状态检查点长度解耦；8192 是物理容量/传输组织参数，不是状态对齐条件。 +- 只在选定位置保存检查点，不保存每个 token 的历史状态；保存位置可以是任意整数长度。 +- 检查点绑定同一实际 token 前缀的 KV、线性状态、可选输出信息和辅助模型恢复信息。 +- 捕获、CPU 停止判定、跨 rank 就绪、发布、淘汰是不同阶段。 +- GPU/CPU 状态缓存不再由 GPU radix 节点的生死直接决定。 +- 完整命中使用显式 `HEAD_ONLY` 执行；不伪装成 0-token prefill。 +- D 输出必须显式发布到下一轮可访问的 cache owner；会话亲和不能代替正确的前缀验证。 + +“任意长度”表示：在长度 L 实际得到的合法状态可以存储、索引和恢复。它不表示能从 S@12000 倒推出未保存过的 S@11000。 + +## 2. 当前流程中不可忽略的约束 + +| 当前实现 | 对新设计的约束 | +|---|---| +| `InferReq.linear_att_cache_len` 来自输入 hash 数；chunk 在大页及输入尾部边界停下 | 改为独立捕获策略，默认不为哈希尾边界额外拆 chunk | +| `_pre_post_handle()` 先推进请求长度，`notify_forward()` 后才执行 CPU `_post_handle()` | 不能在 finished 回调里读取“当前 req 长度和当前 state”当作上一批状态 | +| CUDA graph 输出缓冲会重用 | Python tensor 引用不能冻结 hidden/state | +| MTP SSM 有多个候选槽，conv 有接受位置对应的滑动窗口 | 要按精确接受行导出，不可直接套用 prefill 的 canonical 槽 | +| Qwen3.5 MTP 会原地归一化 target collector hidden | 输出信息需要独立捕获，且必须声明 hidden 的具体格式 | +| 部分 proposer 的 draft KV 使用左移输入，尾部包含本次采样 token | target 与 draft 的可复用边界可能不同 | +| CPU onload 目前有逐请求 synchronize、rank barrier，FA3 有额外保护 | 全异步是后续工程目标，不能简单删同步 | +| D backend 禁用 CPU cache | D→cache-owner 是新增路径，不是打开一个现成开关 | +| PD 首 token 协议依赖 `input_len - 1` 和最后一个 KV task | 精确完整命中需要独立控制消息和明确的首 token 产生方 | +| tool arguments、reasoning 和 chat template 会重新渲染 | 新请求必须按实际 token 前缀验证，不能仅凭 session ID 恢复 | + +## 3. 数据模型与正确性不变量 + +### 3.1 前缀身份 + +逻辑身份为： + +```text +PrefixKey = namespace + execution_fingerprint + token_count + prefix_digest +``` + +`execution_fingerprint` 需要覆盖会改变被复用计算结果的模型版本、adapter、位置/模型配置、KV/state 格式，以及多模态 embedding 身份等。采样 temperature 不属于 target KV 的身份；draft 对采样 token 的依赖由辅助恢复描述单独表达。 + +使用服务端确定的 namespace；客户端 continuation/session ID 只能作为查找提示。不可根据一个不经验证的客户端 handle 直接加载其他请求的内容。 + +### 3.2 检查点记录 + +```text +CheckpointDescriptor + id, generation, namespace, execution_fingerprint + prefix_handle, exact_len, prefix_digest + state_handle + state_layout_version + kv_manifest_handle + output_seed_handle? + seed_format + auxiliary_resume_descriptor? + availability_by_location + retention_metadata +``` + +`prefix_handle` 引用不可变 token 序列/压缩前缀树；manifest 使用结构共享的范围链，而不是为每个检查点复制整条 token 索引表。 + +物理句柄必须带 generation，例如 `(pool_id, slot_id, generation)`，防止 slot 复用后旧异步回调或目录指向新数据。进程内指针不能作为跨进程/跨节点句柄。 + +### 3.3 五条不变量 + +1. `S@L`、`KV[0,L)`、`OutputSeed@L` 必须属于同一个实际执行前缀；不能用较晚状态冒充较早截断位置。 +2. 候选命中 L 必须存在完整 KV 覆盖 `[0,L)`。各层/各 TP 分片都要满足；检查点目录存在不等于数据仍完整。 +3. 完整输入命中 `L == input_len`，除 KV/state 外还需要足够的输出恢复信息；否则选更早检查点。 +4. READY 只表示相应位置的读路径已满足依赖。某副本正在 offload，不得被宣称为 CPU READY；可独立保留已可用的 GPU 路径。 +5. 发布后的可见字节不可覆盖；所有异步读写结束前，源/目标资源不得被回收。 + +### 3.4 三个不同的长度 + +显式记录 `computed_len`、`verified_len`、`visible_end`。普通采样刚生成的最后一个 token 还没有 KV/state；MTP 已验证的一串 token 又可能在中间因 EOS/stop 截断。 + +检查点选择来自 batch 的精确行与状态选择器，不能仅依赖已预推进的 `req.cur_kv_len`。恢复下一轮前，再用新请求实际 token LCP 确认该检查点是否可用。 + +用 0-based token 下标说明 off-by-one：本批开始已计算 B 个 token，请求内第 r 个 verify 输入行处理 `token[B+r]`，产生 `S@(B+r+1)` 和该位置 hidden；该行采样的是下一位置 `token[B+r+1]`。刚采样出来的 token 不因“已输出”而自动拥有 KV/state。 + +普通 decode 生成 N 个 token 时，若没有额外超前计算，最新可保存状态是 `S@(prompt_len+N-1)`。下一轮从这里补算最后一个 token 即可。若业务确实要求包含它的状态，可在后台预算内做一次不向用户继续采样的 seal-only forward；默认不为省下一轮一个 token 强制增加本轮计算。 + +## 4. 索引、KV 页和状态存储 + +### 4.1 逻辑前缀索引独立存在 + +保留现有普通 token radix 的按 token 比较/拆边能力，另建不由 GPU KV 淘汰销毁的 checkpoint 前缀目录。不能仅在现有 GPU `TreeNode` 上挂字段，因为它的生命周期、ref 和 split 仍由 GPU KV 管理。 + +目录使用压缩边，不为每个 token 建 Python 节点；检查点可挂在任意边界。查找是一次前缀遍历加沿途候选检查，不扫描整个状态池、也不为全池每一种长度重复 hash 整个输入。 + +现有固定粒度 hash 可以作为前缀块加速索引。精确检查点的 key 可以通过增量 prefix digest 或复用分块 hash 加尾段计算；hash 页不再是可保存长度的限制。插入/匹配需遵守同一格式，不能混用不兼容 hash 方案。 + +元数据目录按 cache owner 分片；GPU 本地只做自己的快速匹配和租约获取。不要在各推理进程各自维护无法一致回收的“全局” Python 树。 + +### 4.2 固定容量 CPU KV 页与任意有效范围 + +CPU 页只存 KV，暂以 8192 token 为物理容量。状态不嵌在 KV 页内。每个 manifest 的页视图保存： + +```text +page_handle, logical_token_start, physical_offset, valid_token_count +``` + +完整页可以继续使用现有前缀 hash 进行去重；尾页通过 checkpoint 的精确前缀身份关联，不能仅用更长完整页的 hash 来寻找较短历史尾部。 + +同一分支由 12000 延伸到 12317 时： + +- 若获得尾页唯一追加写租约，且新增 token 与已存在内容无冲突，可只写未提交区域。旧检查点仍使用固定的旧 `valid_token_count`。 +- 所有 rank 新增区域写完后，再发布更长页视图；旧读者不观察未提交区域。 +- 若另一个分支写入不同 token，使用新尾页/写时复制，禁止覆盖旧视图可见字节。 +- 空页初始化、padding 和复用也不能触碰正在被引用的有效区域。 + +首版允许尾页 COW，记录其字节开销；不要一开始引入复杂子页压缩。若 Agent 短分支造成大量尾页复制或内部空洞，再以测量结果决定是否增加 fragment arena。 + +页大小独立调节，不绑 snapshot 数量，也不绑 `chunked_prefill_size`。不同大小影响元数据数量、尾页浪费和搬运效率，不能先验保证 8192 对所有负载最佳。 + +### 4.3 状态池统一管理 + +使用一个逻辑 StateStore，取消大小页两个配额。其物理实现包含长期 CPU 锁页槽，以及有界 GPU capture staging；二者预算、用途和指标必须显式区分。 + +CPU 状态已在相同可长期持有的状态存储时，提交只做所有权转移，不再拷贝进 CPU KV 页。不同 owner 的状态仍需显式搬运。 + +CPU 共享内存需要使用现有 `CpuCacheCreator` 一类的显式映射/注册机制，不能把普通进程私有 pinned tensor 的 Python 引用当作跨服务共享。TP layout 要有版本，跨 TP 配置采用 canonical 全局 head 布局或显式转换。 + +### 4.4 查找输出是一个恢复计划 + +```text +ResumePlan + checkpoint_id + exact_len + target_kv_gpu_refs + missing_kv_spans_by_owner + state_source + output_seed_source? + auxiliary_restore_plan + leases + execution_kind = PREFILL_SUFFIX | HEAD_ONLY +``` + +先用 token LCP 限制候选，再验证 KV 覆盖和实际 capability。候选需获取租约，generation 和 READY 二次验证失败时释放并回退。 + +最长命中不一定最快。冷 CPU 的更长检查点与 GPU 上较短检查点，比较“加载缺口 + 状态恢复 + 剩余 prefill”的预计成本；无可靠测量时使用保守阈值，不引入无法解释的复杂评分。 + +## 5. 捕获事务与流水线折叠 + +### 5.1 新增本批不可变 CaptureTicket + +```text +CaptureTicket + request_id + request_slot_generation + batch_epoch + microbatch_id + token_prefix_handle + exact_candidate_lengths + state_selectors / accepted_rows + owned_state_staging + owned_output_seed + KV source leases + capture_event, copy_events, rank_completion + decision = PENDING | RETAIN | DISCARD +``` + +生命周期：`RESERVED → FROZEN → TRANSFERRING → READY`，任何中间阶段可进入取消流程。取消是“不再发布”，不是立即释放仍被 DMA/kernel 使用的内存。 + +### 5.2 时序 + +```text +计算流，批次 t + prepare：固定本批 token 范围和状态行 + target forward + 在会被 graph replay / draft norm 覆盖之前捕获所需 output seed + sample / verify + 选定状态行，条件 gather 到 ticket 独占 GPU staging + record capture_event + 允许后续批次写运行态 + +搬运流 + wait_event(capture_event) + 批量 state D2H、KV gather/offload + record copy_done + +CPU post / cache owner + 根据接受和停止结果决定 RETAIN / DISCARD + event.query + 全 TP 分片完成确认 + 短元数据事务发布对应位置 READY + 释放 ticket 的临时引用 +``` + +CPU 不必阻塞等待 freeze 完成来放行下一批，但 GPU 的执行依赖必须保证下一次状态写入晚于 freeze。跨流只等待 forward_done 不能保护源状态；下一批仍可能同时改写它。独占 staging 把慢 D2H 移出主计算依赖链，但稀疏 gather 本身仍有成本。 + +两个 infer 线程/microbatch 必须各自携带 ticket，不能用一个全局 last_hidden。同一 TP/模型分片组使用同一个逻辑 capture plan；不同 DP 组可以处理不同请求，但须正确参与 collective 协议，包括空 batch 和 padding。CUDA graph 路径使用预分配固定容量 staging/描述数组和有效 mask，不在 replay 中动态分配。 + +### 5.3 不每个 decode token 都做快照 + +默认捕获策略： + +| 位置 | 策略 | +|---|---| +| 精确 prompt 终点 | 预知，优先保存 state + output seed | +| 自然 chunk 终点 | 根据间隔预算保留一部分，不额外按页拆 chunk | +| EOS、max_new_tokens、可设备判定的 token stop | sample/verify 后条件捕获，CPU 后处理最终批准 | +| MTP 接受段中间遇 stop | 根据 token/state 的实际位置关系选择已冻结候选行；adapter 不具备该能力时回退 | +| CPU/外部晚到的字符串 stop、abort | 默认回退较早有效检查点，可在后台预算内重建结束检查点 | + +设备 stop mask 必须与 ignore_eos、min_tokens、token stop 序列等真实规则等价，不能仅判断 `token_id == eos`。CPU 仍是可见输出边界的最终判定者。 + +若 stop token 是第 r 行刚采样出来的，row r 的状态仍在它之前；只有下一行确实处理了该 token 且对应前缀有效，才有包含它的状态。若删除多 token stop 序列后，目标边界早于本批保存的版本窗口，也只能回退，不能用当前 MTP bank 冒充历史状态。 + +独立 detokenizer 的字符串停止可能延迟多批；两个状态版本不能覆盖任意延迟。若产品要求严格保留每个此类终点，必须选择： + +- 有界状态版本环 + 有界确认窗口，窗口满时回压；或 +- 从较早检查点重算尾段，支付重建代价。 + +不承诺同时获得“所有停止点精确保存、无额外显存、无额外拷贝、从不阻塞”。默认优先不破坏推理流水线和正确性，缓存 admission 失败只损失本次缓存。 + +### 5.4 释放与取消 + +state/seed 已冻结到独立槽后，原请求运行态不必一直等 CPU/network 传完；满足原执行依赖即可回收请求槽。KV transfer 任务持有独立 KV lease、不可变 token prefix/manifest，不能继续依赖已被重用的 Req 对象。 + +网络失败、超时或取消仅使缓存事务失败。真正的 buffer 回收要等相关设备操作结束；租约 deadline 不能直接释放仍在传输的内存。 + +## 6. MTP、完整命中与输出信息 + +### 6.1 状态导出接口必须按模型实现 + +定义 `ModelResumeAdapter`，职责为： + +```text +select_capture_candidates(batch_metadata, verify_result, device_stop_mask) +freeze_state(exact_len, state_selector, destination) +finalize_capture(ticket, cpu_stop_result) +restore_state(checkpoint, request) +capture_output_seed(final_hidden, exact_len, destination) +plan_auxiliary_resume(checkpoint, newly_sampled_tokens) +``` + +选择候选与 freeze 位于允许下一批覆盖之前；CPU `finalize_capture` 只能批准已冻结候选、丢弃或回退，不能再读取当前运行态补造较早状态。 + +Qwen3.5 MTP 的 SSM 槽是 `req_idx * (mtp_step + 1) + accepted_row`；这里的 row 是请求内 `b_mtp_index`，不是压缩 batch 的全局行号。conv 使用该 row 对应的窗口。恢复归一化到 canonical 槽，并初始化 MTP 状态索引。其他线性模型不得假定相同布局。 + +不能把所有已计算候选当作已提交 token，也不能把最大 accepted row 当作被 stop 截断后的终点。若所需状态行已经覆盖，回退，不伪造状态。 + +### 6.2 OutputSeed + +首版定义明确格式：`target_final_hidden_before_final_norm`，并记录模型、dtype、TP/DP layout 和精确长度。捕获发生在 target 最终层数据就绪、请求行顺序恢复后，且早于后续复用/原地修改。 + +`mtp_collector.spec_hidden` 不是通用 LM-head 输入：普通模式没有它，一些模式是中间层拼接，Qwen MTP 还会原地 norm。因此 seed 必须独立拥有内存，不能保存 collector 别名。 + +完整命中 `L == input_len`： + +```text +seed 到 GPU → HEAD_ONLY（final norm + LM head + vocab gather） + → 本次请求的 sampling/约束/计数更新 → 当前请求首 token + +缺失 KV/state onload ───────────────────┐ +首 token + auxiliary bootstrap ─────────┴→ DECODE_READY → decode +``` + +不重放旧采样结果，也不缓存经过旧请求 temperature/penalty 处理的 logits。不同 sampling seed/约束可以共享模型输出信息，但应按本次请求重新执行处理。 + +HEAD_ONLY 本身不读取 KV/state，因而 seed 与采样上下文就绪后，可以和大块 KV onload 折叠；下一步 decode 才等待完整恢复。工程上仍由正常调度器安排 LM-head 的 TP collective，不能在任意后台线程发起并打乱通信次序。若首 token 提前返回而 KV 尚未加载完，必须同时记录首个 decode 的等待，不能只用变好看的 TTFT 掩盖后续停顿。 + +`prompt_logprobs` 需要整个输入的分布，一个终点 hidden 不足够；缺少对应缓存时保留重算路径。其他需要逐 token 输出的接口按 capability 同样处理。 + +### 6.3 Draft 的有效边界单独表达 + +当前 Vanilla/EAGLE draft 填充可能把输入左移，并在尾部使用本次采样的首 token。新请求重新采样后，target KV/state 可以相同，但旧 draft 尾槽可能不再有效。 + +辅助描述必须含模式、有效长度、依赖 token/特征和恢复版本。adapter 根据新采样 token 重建受影响 draft 尾部，写请求独占槽；多级 draft 逐级传播依赖，不能拿 target 的 L 当作全部辅助层的 L。 + +没有 adapter 的模式不启用该检查点恢复能力，走原正确路径。可以以后增加显式 target-only bootstrap，但不能假设只跑一个 target token 就自动重建任意 proposer 的完整历史上下文。 + +## 7. Agent 与 PD 的完整流程 + +### 7.1 标准 Agent API + +标准 chat/tool 请求继续正常 render→tokenize。匹配验证的是本次实际 token 前缀：工具 JSON 规范化、reasoning 是否回传、模板和 stop token 都可能让它不同于上次 raw output。 + +保留精确 prompt 终点和输出终点，并按预算保留少量中间检查点。只有末态时，Agent 模板在更早位置分叉就无法恢复；需要回退到沿途已有状态。 + +可选 continuation hint 只用于定位 owner 和 prefix；当前 Responses API 的 `previous_response_id` 并没有现成状态化能力。原生 token continuation 应作为独立 API 能力设计，不能把 hint 解释为跳过验证的授权。 + +工具等待期间可以异步导出检查点、预取下一轮需要的副本。并行工具分支和 retry 保留不可变共同前缀，各自生成独立后缀。 + +### 7.2 P 和 D 之间的目录与数据 + +当前独立服务有独立 CPU shared-memory id 和目录，即便同机也不会自动共享。建议显式设置 cache owner： + +- 首版仍由 P 侧管理 CPU cache；D 增加受预算控制的 checkpoint exporter。 +- P 完成输入后发布其 KV 基础 manifest，P→D 控制信息携带可验证的 base handle。 +- D 结束时导出 state + 新增 KV + 可选 seed/auxiliary 描述，目的 owner 检查已有基础页。 +- 基础页仍在，只传增量；基础页已淘汰则补传缺失范围，或放弃本次 export。不能永远假设 P 的输入页还在。 +- 完成发布后，再将 location/generation 告知路由目录。下一轮优先可直接使用的 owner，同时考虑排队和传输成本。 + +长期 Agent 会话不持有全前缀硬租约。lease 用于实际查找/传输事务；工具等待阶段主要依靠普通保留策略和短期预算，防止挂起会话锁死缓存容量。 + +P TP2、D TP4 等布局不能直接交换 rank-local state 槽；复用已有 PD 全局 head 打包/拆分概念，但 decode 导出必须先按 committed state selector 归一化，不能照搬 prefill 固定槽导出。 + +### 7.3 PD 控制协议显式化 + +新增/扩展协商信息： + +```text +target_prefix_len +kv_ready_len / state_ready_len +output_seed_capability / auxiliary_capability +missing_data_plan +first_token_owner = P | D +transfer_epoch +``` + +`RESUME_READY` 和 `FIRST_TOKEN` 独立于数据任务，允许 0 KV 字节、仅 state、仅 seed 等计划。按 request generation + epoch 保证首 token exactly-once，收到重复通知不得再次递增输出计数。 + +这替换当前 `ready_kv_len == input_len-1` 的隐式判断，也消除“必须有最后一个 KV task 才能附带首 token”的耦合。 + +### 7.4 短后缀在 D 续算属于另一个调度扩展 + +这里的“流水线折叠”首先指现有 CPU/GPU、双 infer 线程和 microbatch overlap。若还希望 Agent 短后缀直接在持有状态的 D 续算，可作为后续 locality-aware 路由优化。 + +不能立即把所有 Agent prefill 发到 D。D 需要显式接收 suffix-prefill 的入口、token 预算和 admission；长工具结果仍可走 P。比较 D 短 prefill 对其他请求 TPOT 的影响与 P/D 搬运代价，再定阈值。 + +### 7.5 一个完整的 12000-token 例子 + +先取普通 decode，假设下一轮模板完整保留上轮 token 前缀，且停止后没有额外超前计算： + +```text +第一轮 P：输入 12000，chunk 预算 8192 + forward [0,8192) → 可按策略捕获 S@8192 + forward [8192,12000) → 捕获 S@12000 + H@12000 + KV 与 state 独立保存；CPU KV 视图为 8192 + 3808 + 不为 hash 对齐额外计算一个 224-token chunk + +第一轮 D：生成 317 个 token + 最后一枚刚采样，冻结的已计算前缀为 12316 + 导出 S@12316 + 新增 KV[12000,12316) + owner 验证基础 KV 后，发布 checkpoint@12316 + +工具运行后：新输入 = 原 12000 + 输出 317 + 新增 400 + 实际 token LCP 验证通过，命中 checkpoint@12316 + 恢复 KV[0,12316) + S@12316 + 只 prefill 剩余 401 个 token，再继续生成 + +另一个请求：输入恰好还是原来的 12000 token + 命中 checkpoint@12000 + 读取 H@12000，HEAD_ONLY → 新请求重新采样首 token +``` + +`H@L` 表示处理前 L 个 token 后，最后位置的 target hidden。它解决完整输入命中后的首 token 生成；对于后面还要追加工具结果的请求，正常 suffix prefill 就会产生新的 logits，并不需要先对旧 H 再采样。 + +CPU 尾页可在满足唯一追加写租约时由有效 3808 延伸到 4124；发生分支则使用独立尾页。若模板重渲染使前缀在更早位置变化，以实际 LCP 回退，不保证这个示例一定能命中 12316。 + +## 8. CPU onload/offload 的效率约束 + +### 8.1 搬运描述符与批处理 + +以 `(source handles/token indices, destination page, offset, valid_count, layout)` 描述范围，多请求、多页合并提交。KV 根据现有 token 索引 gather，状态从独立池按 batch 搬运。 + +- onload 只加载 GPU 缺失的 KV 范围和选中终点的一份 state，缺 seed 时才加载 seed；不沿途加载所有 state。 +- offload 只写新缺失范围，不重复发送已经 READY 的共同前缀；尾部 COW 单独计量。 +- 一次传输可以跨页，也可以只涉及尾部有效区；不为逻辑 token 逐个启动 kernel。 +- CPU lock 只保护元数据预留、租约和发布，不持锁跨 CUDA、网络或磁盘等待。 +- 全 TP 完成信号必须属于同一个 ticket/epoch，不可某一 rank 完成就发布整份状态。 + +### 8.2 折叠加载与调度 + +新请求先得到 `ResumePlan`，进入 `LOAD_WAIT`，异步发起缺口 onload；只有消费该结果的请求等待 load event，其他已 READY 请求可继续计算。 + +依赖应细化为 `HEAD_READY` 与 `DECODE_READY`：完整命中时前者只需 output seed、请求采样上下文及调度通信条件，后者还需 KV/state 和 auxiliary 恢复完成。suffix prefill 则必须先满足该计算真正需要的 KV/state 依赖。首版可以统一等待全部完成保正确，再单独启用 HEAD_ONLY 与 onload 的折叠。 + +不能只删掉当前 `synchronize()`：需要新增请求 ready 状态、调度过滤、源/目标 lifetime 和跨 rank 完成协议。FA3 等现有特殊同步保护先保留，资源级事件依赖验证通过后再替换。 + +增加独立 staging、pending-copy、pending-export 队列的字节上限。工具等待期间可以 offload，但不能让长 D2H 队列饿死新请求的 H2D onload；优先级/限流按目标硬件实际 copy engine 和链路竞争测试决定。 + +### 8.3 所有内存预算可见 + +启动日志至少报告： + +```text +GPU KV capacity +GPU live state bytes +GPU checkpoint staging bytes +CPU KV pinned bytes +CPU checkpoint state bytes +CPU output-seed bytes +pending transfer pinned bytes +``` + +GPU staging 和可选版本环必须在 KV 容量 profiling 前预留。取消“只配置小页槽，另有隐式大页池”的容量表达；按字节预算/实际单槽大小显示可用数量。 + +跨 TP 总内存和每 rank 内存分别展示，避免把单 rank 状态槽或 CPU 整份状态统计混为一谈。 + +## 9. 淘汰、过载和故障 + +- GPU KV 淘汰不直接释放 state。检查点是否仍可用取决于另一层是否有完整 KV。 +- state 淘汰不直接释放共享 KV。KV 可以继续支持其他检查点或纯全注意力缓存。 +- 普通目录/manifest 引用不等于永久 pin 所有物理页；generation 验证及无完整覆盖时的回退是必要的。 +- 页覆盖变化更新目录可用性；惰性校验允许短期陈旧提示，但取得使用租约前必须再次验证。 +- 只优先保留 Agent 最近检查点也不够：共享系统前缀、中间稳定前缀可能更有价值。首版沿用易解释的 LRU，加受限的最近回合优先级，避免每会话无限保留。 +- staging 满、CPU 池满、export 失败：放弃新缓存或回退旧检查点，不能让普通请求因为缓存优化而永久等待。 +- cache flush 使用 epoch/generation 失效新查找，再排空或取消旧事务；不得立即释放仍被设备读取的 slot。 +- 若启用 disk cache，新增纯 KV/state 格式必须版本化；首版未实现 state 持久化时不承诺任意 checkpoint 的跨重启恢复。旧打包页不得按新格式解释。 + +## 10. 工程拆分与验收 + +### 10.1 建议模块职责 + +| 模块/入口 | 改造职责 | +|---|---| +| `infer_batch.py` | 删除固定尾边界作为唯一恢复点的假设；引入 checkpoint/plan handles；分离请求释放与 transfer leases | +| 新 `CheckpointDirectory` | 精确前缀索引、版本句柄、租约、可用性，不访问 CUDA | +| 新 `StateStore` / `CaptureManager` | 有界 state/seed/staging 池、ticket、event 轮询、跨 rank 发布 | +| `req_manager/linear_att.py` | committed state export/restore adapter；明确 MTP row/window | +| `post_layer_infer` / `ModelOutput` | 专用 output seed 捕获，不复用 spec_hidden 别名 | +| `chunked_prefill/impl.py` / `InferReqUpdatePack` | batch epoch、精确长度、capture ticket 和停止结果的绑定 | +| `multi_level_kv_cache.py` / CPU client | 纯 KV 有效范围搬运、LOAD_WAIT、manifest、独立 state 加载 | +| PD master + P/D backend | 显式恢复/首 token 协议、D export、owner 路由提示 | +| 各 proposer | draft frontier、依赖和尾槽重建 adapter | + +接口可按项目风格合并到现有类，以上是职责边界，不要求机械地新增同样数量的类。 + +### 10.2 分阶段落地 + +1. **先做生命周期和模型 adapter**：精确长度、ticket、state selector、取消/回收。普通 decode 与 Qwen MTP 分别验收，再启用对应能力。 +2. **单服务、GPU KV 常驻的精确检查点**：token 前缀索引独立；state 仍可放 CPU 独立池，保留自然 chunk 末尾和精确输入尾部；增加 OutputSeed/HEAD_ONLY，删除仅为 hash 尾边界引起的切分。 +3. **CPU 解耦与范围搬运**：纯 KV 页、独立 state、尾页视图和 lease。初版可保留现有同步安全屏障，随后独立验证 LOAD_WAIT/event pipeline。 +4. **PD 精确恢复协议**：0-byte 控制任务、首 token owner、跨 TP state/auxiliary 转换。 +5. **D 输出导出与 Agent 路由**：基础 manifest 验证、增量传输、工具等待期 export，标准 chat 仍验证 token LCP。 +6. **依据瓶颈扩展**：严格字符串 stop 版本环、D 短后缀 prefill、CPU 尾片压缩。每项单独证明收益和成本。 + +新旧缓存模式通过配置与协议版本隔离，启动时验证模型 adapter 和服务间能力。不能让某个 unsupported proposer 在运行时悄悄使用不完整的恢复信息。回滚时释放新模式事务并冷启动对应缓存，不复用不兼容的共享内存布局。 + +### 10.3 必须证明的行为 + +正确性覆盖：任意长度/大页前后 1 token、完整命中重新采样、不同约束、prompt_logprobs 回退、MTP 部分拒绝和接受段中间 stop、工具序列化改变前缀、并行分支、GPU-only/CPU-only/混合覆盖、尾页追加与 COW、不同 TP、0-byte PD、重复 FIRST_TOKEN、下一批覆盖 state、graph replay 覆盖 hidden、slot ABA、拷贝时淘汰、flush/abort、队列满和延迟 rank。 + +验收不仅看 token 命中率： + +- 冷请求 TTFT：12000 输入在 chunk=8192 且无其他调度约束时应为 2 次 target prefill，而非为 11776 再拆一次。 +- 真正少算的 token，以及来自 D 输出的复用 token 数。 +- Agent 下一轮 token LCP、可用 checkpoint 长度、模板变化造成的回退长度。 +- 完整命中 HEAD_ONLY 次数和 MTP auxiliary 重建成本。 +- TTFT/TPOT p50/p95/p99，其他并发请求的尾延迟。 +- checkpoint 捕获次数、D2D/D2H 字节、对计算流的阻塞时间、状态/staging 峰值。 +- KV 增量传输、基础页补传、尾页 COW、CPU 内部空洞。 +- 因状态缺失、KV 缺口、未 READY、capability 不足而回退的次数。 + +任何 benchmark/eval 都按仓库要求使用 `exp -m ...` 记录；本设计稿没有执行新的性能实验，也不声称上述方案已有收益数据。 + +## 11. 代码依据 + +- `lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py`：forward、预更新、notify_forward、post_handle 的次序。 +- `lightllm/server/router/model_infer/mode_backend/overlap_events.py`：双 infer 线程握手。 +- `lightllm/server/router/model_infer/infer_batch.py`：固定 `linear_att_cache_len`、当前快照复制、请求释放、`InferReqUpdatePack`。 +- `lightllm/server/router/model_infer/mtp_speculative/utils.py`、`lightllm/common/basemodel/triton_kernel/mtp_utils.py`：接受状态索引。 +- `lightllm/common/basemodel/attention/linear/gdn.py`、`lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py`：SSM 候选行和 conv 窗口。 +- `lightllm/models/llama/layer_infer/post_layer_infer.py`：最后位置、final norm 和 LM head。 +- `lightllm/models/qwen3_5_mtp/layer_infer/pre_layer_infer.py`：target hidden 原地 norm。 +- `lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py`、`eagle_with_att.py`：draft 输入左移和采样 token 依赖。 +- `lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py`：现有同步、prompt-only offload 和 READY 发布。 +- `lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py`：禁用 CPU cache、最后一个输入 token 假设。 +- `lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py`：KV/state 传输和首 token 附着。 +- `lightllm/server/httpserver_for_pd_master/manager.py`、`pd_selector/pd_selector.py`：PD 握手、token 编码和路由。 +- `lightllm/server/build_prompt.py`、`api_openai.py`、`api_responses.py`：Agent 模板/工具重渲染与无状态 Responses 接口。 diff --git a/docs/agent_checkpoint_cache_validation.md b/docs/agent_checkpoint_cache_validation.md new file mode 100644 index 0000000000..d6a258f1f0 --- /dev/null +++ b/docs/agent_checkpoint_cache_validation.md @@ -0,0 +1,210 @@ +# 精确前缀缓存：H100 验证记录 + +日期：2026-09-10。基线为 upstream `1eb4810c78c6ea908b69c20c3c037b79979e1cd6`，候选按独立源码快照记录。最新R15合并了完整命中的恢复等待，降低搬运临时显存；**服务TPOT仍约2.54ms，尚未追平基线**。历史R12与最新R15的计时分别列出，不把不同轮次的baseline波动算作代码收益。所有 benchmark/eval 经 `exp -m` 执行，保留失败运行、源码、启动日志、原始请求及逐 token 概率。**本功能默认关闭；恢复正确性、冷算差异和性能分别报告，不宣称全面提升或整体全绿。** + +## R15：本轮恢复优化（2026-09-10） + +只保留两项相互配合的改动：normal完整命中时，KV、索引、conv/SSM和hidden异步排入同一stream，最后统一等待；逐页GPU临时tensor立即释放，让同一stream复用一页空间。`load_kv`默认仍等待，部分命中/PD保留原有屏障；异常路径在lease关闭前等待。发布请求长度和HEAD seed仍发生在完成屏障之后。 + +候选基于R12 commit `213155bccc7e1fb761e35f6c38369b5228667e02`，部署tar SHA256 `3924999ad1712246ee50abdf5ae5b31ea55c6a5fb31eeab9cfec7c4bd8546df2`。最终docstring只澄清异步调用方的lease/索引生命周期要求，执行代码与该快照一致。H100 GPU3依次运行R12/R14/R15，原基线位于GPU6；计时窗口无其他测试请求或GPU专项。 + +### 搬运及生命周期 + +实际CPU cache和模型的state恢复方法、真实0.8B张量形状、合成数据的GPU专项中，8请求顺序恢复结果如下（3次预热、10次重复）： + +| 场景 | R12 | R15 | +|---|---:|---:| +| 仅缺末尾1个KV token，另恢复state/seed | 4.857ms | 4.687ms | +| 8193-token完整CPU恢复，跨两页 | 23.416ms | 22.855ms | +| 1/2/4页搬运的最大临时GPU空间 | 112 / 224 / 224MiB | 约112MiB / 112MiB / 112MiB | + +每请求conv/SSM约18.63MiB；一整页KV约112MiB。此专项不加载模型权重、不走HTTP调度，**不把微测改善等同于TPOT改善**。独立随机数据复验覆盖8请求×8193、start=127、非连续目标槽位、KV/conv/SSM/seed逐字节、未使用state bank不变、clear后的lease及第二次copy异常。CPU边界替代专项42场景/310检查通过；真实CPU缓存759检查通过，覆盖默认/显式wait及跨页恢复。它们不代替GPU和服务数值验证。 + +### R15真实服务功能 + +三套独立功能测试共85请求、344项checks、18组恢复/replay比较通过，HTTP错误0。Agent/EOS/stop/分支为28请求、124checks/15比较；Triton非greedy uniform/mixed为32请求、144checks，16个warm全部完整命中；GPU容量压力为25请求、76checks/3比较。 + +压力请求使用正常缓存的fresh prefix，每轮10×8192，超过GPU容量65536；避免`disable_prompt_cache=True`的请求释放后不留radix、无法造成真实缓存淘汰。两轮日志均证明`GPU=0 CPU=8193`，分别验证8193完整命中和追加17-token后的8210输入恢复。临时测试脚本改动和原版同时归档，没有修改服务参数。 + +严格cold比较独立报告：R15三组中两组失败。Agent立即续接的532-token输入命中S@512，与完整冷算的输出IDs相同，第2token标量logprob最大差0.0648808479；stop后270-token输入命中S@259，恢复首token为16（logprob −2.0264773），冷算首token为760（−1.9266164），输出序列不同。branch组逐位一致。恢复与完整冷算的等价性仍未成立，不能用普通功能suite的exit0遮盖这两项失败。功能测试首次并发长输出seed还记录约12.3秒TTFT，未删除、不纳入性能均值。 + +随后在同一GPU3重启未改同步的R12，用相同run-id重跑28请求：输入摘要、cache命中长度、全部输出IDs及返回标量logprob与R15逐项相同；124checks/15恢复比较通过，严格cold的两例失败也完整复现。R12部署的1818个源码文件与其归档tar逐一核对。该对照证明这两例差异在R15之前已存在，不证明任意输入都与冷算等价。 + +### 服务性能和未解决的问题 + +原负载仍是8并发、257输入、MTP3、8输出。R15连续两次各20轮，每侧合计320个warm；704请求、1312比较通过，无样本剔除。两次candidate TPOT为2.5510/2.5227ms，合并结果： + +| 均值 | 同轮baseline | R15 | +|---|---:|---:| +| TTFT | 75.707ms | 50.320ms | +| TPOT | 2.0995ms | 2.5369ms | +| 总延迟 | 90.672ms | 68.350ms | + +TPOT比同轮baseline高20.8%，总延迟低24.6%。R15的TPOT与历史R12合并值2.5366ms几乎相同，不能据此称剩余TPOT回退得到改善。baseline历史值2.2454ms和本轮2.0995ms也说明必须区分计时轮次。 + +两侧输出改为128的完整20轮中，baseline/R15为TTFT70.747/48.887ms、TPOT1.3941/1.5781ms、总延迟248.122/249.758ms。**该组退出1：656比较中4项失败**，来自2个warm请求的seed与cold比较；IDs一致，但第16个token标量logprob差0.033751875，超过原定0.03。R12独立长输出对照也出现同一位置、同一概率值和差值，说明问题在本次恢复合并之前已存在；不能因此忽略失败，也不能宣称与冷算完全等价。 + +此前尝试的R13完整MTP预提案、R14单候选预提案均已撤回。8输出仍需3轮target验证，未得到稳定TPOT收益：R13两次2.6855/2.5948ms，R14两次2.6700/2.6238ms；同GPU的R12对照2.6181ms。既未用延迟发布首token美化TPOT，也未保留无收益的提案路径。原始失败和探索性运行一并保留。 + +## 环境与验证边界 + +- H100 80GB HBM3，driver 580.167.08，PyTorch 2.11.0+cu130。 +- Qwen3.5-0.8B：普通模式、DP2双微批、P/D分离，以及 `eagle_with_att` MTP step=3。DP组CLI `--tp 2 --dp 2`表示两个总rank、每个DP副本TP1，不能解释为每副本TP2；各组配置见启动工件。 +- Qwen3.5-27B：同一对 GPU 依次运行 baseline/candidate，TP2/DP1、MTP0、关闭 CUDA graph。两侧实际后端相同：Fa3 full-attention prefill、FlashInfer decode、FlashQLA linear prefill、Triton linear decode。 +- GPU KV 容量 65536 tokens，running requests 32，请求上限32768，chunk/batch预算8192；小页256、大页256×32、旧小页状态槽128。27B CPU预算8192 MiB/TP rank；MTP并发8和DP/CPU专项4096 MiB/rank；PD专项1024 MiB/rank。 +- R8上传tar SHA256：`fc329aa845f507fea9249707771c75276051aea8d200f84742b9e91e328f8da1`。R10：`46e489e9e9b85cf15a518708c27bc1eeced583e16c7b3cbf18fdfc389cc52466`。 +- 历史R12上传tar SHA256：`0c900118b9cbd83399f761cc8fb1052d4ef32c320b404b8728dce7d1cd31b6d9`。R12交付时的14个变更Python文件与该快照逐文件SHA256一致；R15改变了其中两个文件。 + +早期0.8B对比曾出现baseline的FlashQLA降级到Triton，该组不用于性能归因。首次shape编译和拆批抖动保留在原始记录中，不通过删除异常波使结果变好。下述短输出负载不是饱和吞吐测试。 + +## R3–R12修复与证据 + +| 问题 | 修复及验证 | +|---|---| +| 同tokens混入不同计算历史的KV | origins同时约束CPU页去重、GPU radix、P→D和D→P。此前S@284恢复混入95776个不同target KV元素、最大差0.185547；真实保存/恢复的state及KV字节断言验证修复。 | +| 任意终点与MTP状态 | 8193最后一个输入token进入decode时仍捕获prompt endpoint；保存实际accepted conv窗口/SSM row，私有packed尾槽重建后更新来源；HEAD使用独立输出缓冲命名空间。 | +| 输出发布约49.56ms | 首次pinned页分配35.53ms、state分配9.03ms是主要热点。启动预留有界arena，消除请求内大块`cudaHostAlloc`；窗口拥有独立storage边界，lease/COW/flush/序列化及OOM生命周期已验证。 | +| 折叠流水线额外等待/计算 | HEAD和aux修复批量执行；shared forward generation阻止伙伴已开始forward时仍睡20ms；max-output门控保留PASS握手和延迟释放，但不再白算达到输出上限后的整轮。 | +| capture与搬运同步 | pinned元数据异步上传；offload的最后一页KV与state/seed共用完成等待；R15另修复onload逐页临时空间复用。normal模式后台准备CPU快照，主调度统一提交；onload和PD发布仍同步。 | +| R8发布可见性竞态（P1） | R7后台commit可让两个TP rank在同tokens/length处取得不同历史：真实CPU cache＋TP2 Gloo复现rank0 origins222/state22、rank1 origins111/state11。R8 worker仅prepare，主调度交集完成epoch后commit/discard、释放hold并ACK。另修复candidate查询后被evict导致`acquire=None`的异常。 | +| R9首轮MTP布局错误 | 只移除无效draft行，却遗漏GDN/FA3的fixed4布局假设。真实模型20轮有320项比较失败；该运行不是收益证据。R10让attention builder、page table及graph元数据按实际请求/候选边界构造。 | +| 可变行数触发重复JIT | R11把实际行数改为runtime参数；H100改前/改后各384项检查通过，capture选择kernel变体32→4，MTP start-location变体89→6。R12启动预热有限BLOCK变体：独立空Triton cache的H100专项348项通过，预热后行数0..32及normal/prefill/DP局部边界无新变体，conv/SSM/KV及请求索引字节不变。 | +| 同一checkpoint重复发布 | R12仅在所有TP rank已有相同来源、满足seed要求的条目时复用，避免重建CPU页/state。真实CPU cache＋TP2 Gloo的13场景、26个rank结果全部通过，覆盖单rank缺失、clear、OOM、异常、完成时序差异和ACK；该专项不证明GPU性能。 | +| DP MTP辅助计算误用graph | R10真实服务8个seed成功，第9个请求首次HEAD超时：单批`resume_auxiliary`误replay双微批graph，复制缺失infer state时触发`vars(None)`。R12让该单批调用走现有eager路径，正常成对DP调用保留graph；132项真实dispatch方法的CPU契约先红后绿，随后真实DP MTP复验通过。 | + +R8 P1回归先红后绿：原R7两项缺陷均复现；R8通过版本一致性、eviction回退、两任务ACK、clear/discard、单rank第二槽OOM、fatal清理及generation分歧检查。这些使用生产方法、真实CPU cache、线程和TP2 Gloo，CUDA上下文被替代，**不能单独证明GPU数值或服务性能**；实际服务证据见下文。 + +R10在H100执行真实GDN、FA3 FP、FA3 MLA state builder和Triton元数据/page-table kernel，覆盖一行HEAD、完整候选、混合/补齐/空批、CUDA graph replay及DP workspace独立性。48个gate组合均通过；原summary误写96，纠正工件与原始记录同时保留。现有相关测试91项通过。该builder专项不加载模型权重，不把graph元数据验证等同于完整模型输出验证。 + +## R12真实MTP功能 + +R12 probe完成28请求、124项checks、15项恢复/replay比较，全部通过且无HTTP错误，覆盖Agent立即续接、256-token长输出、分支、MTP token stop、自然EOS和并发长输出。 + +R12长输出HTTP结束后的立即下一轮，输入532 tokens命中输出checkpoint S@512；token-stop后的下一轮命中S@259，随后重放分别完整命中532/270。较早R10相同Agent立即续接只命中旧prompt257，该次较短前缀回退记录保留。异步发布仍**不保证HTTP结束时输出checkpoint已可见**；尚未发布时允许安全回退计算。 + +probe的exit 0只代表上述恢复比较。R12独立strict-cold三组中，`stop_immediate_next_turn`输出IDs相同，但最大logprob差0.08048176765，超过0.03，比较失败；另两组Agent最大差0.0095384、branch为0。R10也出现同一stop比较失败。没有放宽阈值，也不声称恢复路径与完整冷算逐位等价。 + +R12 FlashInfer非greedy专项32请求、144项checks全部通过，16次warm全命中，无HTTP错误。本轮未再出现R10首组约55.7秒采样JIT；R10完整计时仍保留。R12非greedy与Agent功能请求并行执行，其延迟不用于性能结论。 + +R12 DP2、每副本TP1、双微批、MTP3真实服务完成26请求、84项checks、14项比较，全部通过且无HTTP错误；14对返回IDs和标量logprob逐位相同。实际后端为full attention Triton prefill/FA3 decode、linear attention FlashQLA prefill/Triton decode，target和draft均启用overlap graph。两个DP副本各自的隔离HEAD均只验证1 token、1 step；还覆盖旧128-token流运行中插入8个完整HEAD，以及输出S@264续接、276-token完整重放。该结果验证了R10失败路径的修复，但不声称逐forward追踪了所有1/4行组合或证明全部vocabulary logits相等。 + +## R8：27B、DP/CPU与R3 PD + +| 路径 | 实际结果及限制 | +|---|---| +| 27B R8完整对比＋追加20轮 | 55＋176＝231请求，HTTP错误0，193次warm全部完整命中，harness 77＋328项比较通过。seed→warm IDs全相同；返回logprob154/193对逐位相同，最大差0.0130941，不能写成全部bitwise通过。 | +| R8 DP2/双微批/空rank | MTP0，60请求无HTTP错误或功能断言失败；12/12组seed→full-hit的IDs及返回标量logprob逐位一致。原suite仍exit 1：8个cold-reference比较失败，其中1个greedy首token不同。 | +| R8 GPU淘汰后CPU恢复 | 26请求，checks/comparisons全通过。两轮各10×8192正常缓存请求施压，日志两次证明`GPU=0 CPU=6000`。锚点恢复及追加17-token后的6017-token完整重放，2/2组IDs/logprob逐位一致。 | +| R3 PD输出回流 | TP1/MTP0，31/12000输入两组共14个生成请求通过。D输出checkpoint在P命中46/12015；P完整HEAD命中60/12029。12015长度回流`pages_sent=1 pages_reused=1`，只传新尾页。 | +| R3 PD首token/权限 | P HEAD、D完整HEAD及D缓存S@L后输入L+1的重复输出IDs/logprob一致、首token不重复。零数据任务明确为KV控制消息、`first_token_owner=decode`；非空P→D origins完整且为正。未认证registry访问均403。 | + +27B非逐位一致的记录集中于拆批边界，主要影响首token及过渡处的第2/最后token，与batch3/5/7的计算形状相关；尚未用固定hidden seed纯head实验完成因果确认。返回token的标量logprob一致，也不等于全部vocabulary logits一致。 + +PD证据明确来自R3。其20个PD/`pd_io_struct.py`文件与R8哈希相同，但共用ExactPrefixCache/CPU cache已有变化，不能据此冒充R8或历史R12整条PD路径重验。该TP1/MTP0组也不证明异构TP、PD＋MTP或跨机带宽性能。 + +## 性能：历史回退与尚存问题 + +原始0.8B、MTP step3、8并发257输入/8输出的历史回退是baseline TPOT **2.32ms→13.01ms**，总延迟99.96ms→192.09ms。它是本轮调查的起点，不是最终候选性能。 + +**历史R12**使用同一原负载连续两次20轮，保留全部40波，每侧320个warm；共704个HTTP请求、1312项harness比较全部通过，HTTP错误0，候选320次warm全命中。seed→warm的320组IDs均相同；294组返回标量logprob逐位相同，其余最大绝对差0.0034699291，不能称为全部bitwise通过。其他服务无请求时执行的完整合并计时如下： + +| 指标 | baseline | R12 | 均值变化 | +|---|---:|---:|---:| +| TTFT mean | 73.21255ms | 52.53901ms | −28.24% | +| TPOT mean | 2.245395ms | 2.536606ms | **+12.97%** | +| e2e mean | 89.19329ms | 70.56587ms | −20.88% | +| TPOT p95 | 2.45881ms | 3.44297ms | — | +| TPOT max | 2.64958ms | 3.78341ms | — | + +两次candidate的TPOT全量均值分别2.4692/2.6040ms。历史13.01ms回退中的大停顿在本轮未再出现，但**约13%的TPOT损失仍在，不能宣布性能回退已全部修复**。短输出的TTFT收益与TPOT代价需同时评估；本结果不外推到饱和吞吐或27B。 + +下列均为原负载每侧20轮、160个warm的**完整中间结果**，各运行656项harness比较通过，候选warm全命中。TTFT/TPOT/e2e均为全量均值，单位ms： + +| 快照/运行 | baseline TTFT / TPOT / e2e | candidate TTFT / TPOT / e2e | 仍存问题 | +|---|---|---|---| +| R8 | 73.969 / 2.255 / 90.001 | 51.974 / 2.756 / 71.544 | TPOT仍约+20%，按全量均值为+22.2%。 | +| R11首次 | 73.942 / 2.26816 / 90.053 | 52.964 / 3.65916 / 78.843 | 有限BLOCK首次JIT仍发生在请求内。 | +| R11完整重跑 | 73.321 / 2.29844 / 89.680 | 50.877 / 2.65470 / 69.743 | TPOT仍+15.5%；不是历史R12结果。 | + +R10完整运行另保留在工件中：candidate TPOT mean3.3826ms、median2.5008ms、p95为10.2915ms，baseline mean2.2164ms；round9/18约70ms停顿。R11首次运行median2.440ms、p95为10.122ms；缓存全量扫描确认请求窗口内新编译四个变体：start-location的BLOCK8/32及capture选择的BLOCK8/32，源文件到cubin间隔24/24/32/40ms，与round1/19的三个长gap重叠。这证明停顿内发生首次编译，不能把每一毫秒都归因于编译。`r11-tail-jit-audit`中的剔除波次分析仅用于诊断，**不替换任何完整性能结果**。 + +27B R8与同GPU、同后端基线的对比如下。steady定义固定为warm repeat index 1/2，所有index 0和异常记录仍保留： + +| 负载 | baseline TTFT / TPOT / e2e | R8 TTFT / TPOT / e2e | +|---|---|---| +| 单请求257 | 137.909 / 41.610 / 429.214ms | 43.357 / 42.123 / 338.301ms | +| 单请求8193 | 126.096 / 40.876 / 412.268ms | 37.921 / 42.514 / 335.555ms | +| 单请求12000 | 157.912 / 40.780 / 443.423ms | 47.025 / 41.469 / 337.343ms | +| 8并发257 | 166.557 / 40.948 / 453.515ms | 86.369 / 45.254 / 403.437ms | + +27B并发TTFT改善48.1%、e2e改善11.0%，但**TPOT仍回退10.5%**；单请求TPOT回退1.2%–4.0%。追加160个warm保留全部20轮，其尾延迟为： + +| 指标 | mean | p50 | p95 | p99 | max | +|---|---:|---:|---:|---:|---:| +| TTFT | 89.390ms | 85.284ms | 114.817ms | 330.894ms | 331.753ms | +| TPOT | 47.026ms | 42.519ms | 78.536ms | 89.258ms | 89.327ms | +| e2e | 419.034ms | 385.917ms | 654.165ms | 674.077ms | 675.098ms | + +分位数采用`(n−1)×q`线性插值，无异常删除。round1、round6的平均TPOT分别69.14/79.74ms；已记录的TileLang编译在这些波之前结束，不能把它们统称为该次JIT。该20轮只提供候选尾延迟证据，未配套同期20轮baseline。 + +R6 trace解释了首gap的一部分：HEAD即时发布首token，之后首decode event约42ms已ready，但CPU等下一轮host forward返回，到约84ms才post第2个token；真正event同步仅0.015ms。基线首token也经过普通流水线，该等待包含在TTFT。单独把HEAD改为普通stage主要会把约40ms从TPOT移入TTFT，不据此宣称提速。R8确实删除了达到max-output后的整轮计算，单请求末gap从约40ms降到1.4–3.6ms；并发拆批仍留下损失与长尾。 + +## 使用限制 + +- 默认关闭。normal输出快照后台prepare、主调度统一commit；**CPU onload仍同步，PD发布仍同步**。新请求恢复可能阻塞已有请求的post，未宣称全异步或保证尾延迟。 +- checkpoint表示同一计算来源的KV/state。完整cold prefill、decode形成的状态、不同chunk和batch形状不保证浮点等价；实际cold比较出现过不同greedy token。要求与完整冷算逐位相同的使用方不能据本验证获得保证。 +- 1GiB不足以同时保留本负载8路MTP的16个prompt/output checkpoint：每个约130.64MiB，合计约2.04GiB；相关验证使用4GiB。容量拒绝和较短前缀回退是允许行为，不保证命中。 +- arena窗口、GPU槽、staging和CPU lease必须保留到DMA完成；clear/OOM/abort不得省略完成等待。部署须统一更新使用新共享请求/PD结构的进程。 + +## 复现与工件索引 + +原MTP性能负载（服务端均使用同一模型、MTP step3、4GiB CPU预算与已归档启动参数）： + +```bash +exp -m "MTP8 exact checkpoint original workload twenty waves" \ + python test/benchmark/agent_checkpoint_cache.py \ + --baseline-url http://127.0.0.1:BASELINE_PORT \ + --candidate-url http://127.0.0.1:CANDIDATE_PORT \ + --baseline-revision 1eb4810c --candidate-revision DEPLOYED_SNAPSHOT \ + --model-dir /models/Qwen3.5-0.8B \ + --suites concurrency --concurrency 8 --concurrency-input-len 257 \ + --max-new-tokens 8 --seed 1558 --settle-ms 200 --repeats 20 \ + --run-id mtp-concurrent8-4gb-20260910 --require-exact-hits \ + --output /path/to/new-immutable-run +``` + +27B完整对比使用`--suites boundaries,concurrency --lengths 257,8193,12000 --max-new-tokens 8 --repeats 3 --concurrency 8 --run-id tpot-27b-20260910`，追加组仅改为concurrency/repeats20并使用独立输出目录。脚本同时保留cold→seed/warm与seed→warm，不放宽阈值。 + +本机工件根：`~/experiments/artifacts/checkpoint-tpot-fix-20260910/`。目录内保留源码指纹、参数、日志、原始请求及审计；远端根为`/dev/shm/lightllm-cache-optim-20260910/artifacts/`。 + +| 内容 | 工件目录 | exp ID前缀 | +|---|---|---| +| R15原负载40轮及128输出20轮 | `r15-service-analysis/`（含原始请求、失败及汇总） | 本地`260910-123942/124116/124156`；最后一组exit1 | +| R15 Agent/非greedy/CPU压力及严格cold | `r15-http-functional/` | 本地`260910-124918/125022/125104`；strict-cold单列在`diagnostic.json` | +| R12同输入复核R15的严格cold失败 | `r15-http-functional/r12-same-input-async-control/` | 本地`260910-125547`；逐请求跨版本核对`root-r12-r15-comparison.json` | +| R15恢复边界与真实CPU cache | `r15-restore-fence-contract/` | `260910-123055`红、`123147`绿；矩阵`123858`、cache `124129`；临时runner语法失败`123833`保留 | +| R12/R15 GPU恢复、临时显存及随机槽位 | `r15-onload-gpu-contract/` | 远端`260910-123458/123638/123758/124545/124604` | +| R13/R14未保留的MTP尝试 | `tpot-priming-analysis/`、`r14-one-candidate-contract/`、`r14-two-row-gpu-contract/` | 完整逐运行记录见目录内meta/analysis；源码在`source-snapshots-r13-r15/` | +| R8 P1先红后绿 | `r8-review-regressions/` | `260910-093103`红；`093115/093243`绿 | +| R8 MTP原负载20轮 | `tpot-candidate-twenty-r8/` | `260910-093658`，exit0 | +| R9真实模型失败 | `tpot-candidate-twenty-r9/` | `260910-095145`，exit1 | +| R10 GDN/FA3 builder与graph | `variable-verify-layout-gpu-r10/` | `260910-101300/101301`；91项测试`101308` | +| R10原负载中间结果，长尾未解决 | `tpot-candidate-twenty-r10/` | `260910-101714/101715` | +| R10真实MTP功能/strict-cold | `tpot-async-publication-http-r10/` | `260910-101324/101325`；见`diagnostic.json` | +| R10非greedy功能 | `tpot-flashinfer-nongreedy-r10/` | `260910-101442/101443` | +| R10 DP MTP真实请求失败 | `dp-mtp-r10-failures/` | `260910-102313/102314`，exit1 | +| R11 runtime行数kernel | `variable-row-jit-r11/` | `260910-102227/102228`；相关测试`102310` | +| R11原20轮/完整重跑 | `tpot-candidate-twenty-r11/`、`tpot-candidate-twenty-r11-repeat/` | `260910-102757`、`103017` | +| R11首次BLOCK编译关联诊断 | `r11-tail-jit-audit/` | `260910-102911/102923/103132` | +| R12发布复用/DP辅助graph契约 | `publication-dedup-tp2-r12b/`、`dp-mtp-aux-graph-contract/` | `260910-103124`；`102722`红、`102747`绿 | +| R12空cache启动预热GPU专项 | `warmup-kernels-r12/` | `260910-103513`；早期临时目录满失败`103442`保留 | +| R12最终原负载40轮 | `tpot-candidate-twenty-r12/`、`tpot-candidate-twenty-r12-repeat/`；汇总`r12-final-performance.json` | `260910-103813/103814`、`103935/103936` | +| R12 Agent/stop/分支/严格冷算 | `tpot-async-publication-http-r12/` | 本地`260910-103702-….cmaYtv`；远端`103703`（pid932434） | +| R12 FlashInfer非greedy | `tpot-flashinfer-nongreedy-r12/` | 本地`260910-103702`；远端`103703`（pid932419） | +| R12 DP MTP真实服务 | `dp-mtp-r12-fa3-p72/` | `260910-103654/103655`；严格审计`103833`；审计输出权限失败`103744`保留 | +| 27B R8完整对比/20轮尾延迟 | `27b-r8/` | `260910-094142/094143`、`094231/094232` | +| R6首gap实际trace | `27b-profile-r6/` | `260910-092052/092053` | +| R8 DP/CPU及严格重放审计 | `dp-r8/` | `260910-094012`原DP exit1；`094051`CPU；`094130`审计 | +| R3 PD短/长输入 | `pd-r3/` | `260910-083441/083442`、`083520/083521` | + +历史R12的性能、Agent恢复、非greedy和DP MTP各自保留独立原始工件；较早R3/R8的PD、CPU压力和27B证据保留其原快照范围，不冒充最终快照重验。 diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index a7e2d8122a..73d4e844e1 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -66,8 +66,12 @@ def uses_dynamic_spec_verify_layout(self) -> bool: draft_step = self.model.mtp_manager.get_decode_draft_step(self.model.is_mtp_draft_model) is_main_model = not self.model.is_mtp_draft_model has_decode_draft_step = draft_step > 0 - dynamic_verify_enabled = args.mtp_dynamic_verify - return is_main_model and has_decode_draft_step and dynamic_verify_enabled + # Exact HEAD replay starts without proposals, so fixed planning can + # also mix one-row and full-width requests. Keep the attention layout + # service-wide to make CUDA Graph capture and replay use the same shape. + exact_head_enabled = getattr(args, "enable_exact_prefix_cache", False) and args.run_mode == "normal" + variable_layout_enabled = args.mtp_dynamic_verify or exact_head_enabled + return is_main_model and has_decode_draft_step and variable_layout_enabled def uses_causal_attention(self) -> bool: args = get_env_start_args() diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f1247e0ef4..119b48a49c 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -310,6 +310,55 @@ def _init_custom(self): def _init_hidden_collector(self): self.hidden_collector_prototype = self.mtp_manager.create_hidden_collector(model=self) + def supports_exact_output_seed(self) -> bool: + """Whether the target's output head has the implemented replay format. + + This only describes the target head. Speculative modes additionally + require the proposer's auxiliary-resume capability before a full hit + can be admitted. A matching seed alone is not sufficient for MTP. + """ + from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + + return not self.is_mtp_draft_model and type(self.post_infer) is LlamaPostLayerInfer + + def _capture_output_seed(self, hidden: torch.Tensor, infer_state: InferStateInfo): + if not getattr(self.args, "enable_exact_prefix_cache", False) or not self.supports_exact_output_seed(): + return None + if infer_state.is_prefill: + last_rows = torch.cumsum(infer_state.b_seq_len - infer_state.b_ready_cache_len, dim=0).long() - 1 + return hidden.index_select(0, last_rows) + return hidden[-infer_state.batch_size :].clone() + + @torch.no_grad() + def forward_output_seed(self, output_seed: torch.Tensor, microbatch_index: int = 0) -> ModelOutput: + """Run only the normal final norm/head/gather for a batch of exact hits. + + The scheduler must call this on its usual compute stream, in the same + collective order on all ranks of this TP group. This method neither + touches recurrent/KV state nor samples, updates request lengths, or + initializes a drafter. Those remain normal backend responsibilities. + """ + if not self.supports_exact_output_seed(): + raise NotImplementedError("this model does not implement exact output-seed replay") + if output_seed.ndim != 2 or output_seed.shape[1] != self.config["hidden_size"]: + raise ValueError("output seed must contain raw final hidden rows for this model") + if not output_seed.is_cuda or output_seed.dtype != self.data_type: + raise ValueError("output seed must use the model CUDA device and hidden dtype") + infer_state = self.infer_state_class() + infer_state.dist_group = dist_group_manager.get_group(microbatch_index) + batch_size = output_seed.shape[0] + if batch_size == 0: + vocab_size = self.pre_post_weight.lm_head_weight_.vocab_size + return ModelOutput(logits=torch.empty((0, vocab_size), dtype=torch.float32, device=output_seed.device)) + g_cache_manager.cache_env_in() + try: + logits = self.post_infer._lm_head_and_gather( + output_seed, batch_size, self.pre_post_weight, infer_state + ).clone() + finally: + g_cache_manager.cache_env_out() + return ModelOutput(logits=logits) + @torch.no_grad() def forward(self, model_input: ModelInput): model_input.to_cuda() @@ -472,7 +521,7 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): padded_batch_size = model_output.logits.shape[0] - if padded_batch_size == origin_batch_size: + if padded_batch_size == origin_batch_size and model_output.output_seed is None: return model_output new_model_output = copy.copy(model_output) new_model_output.logits = new_model_output.logits[0:origin_batch_size] @@ -480,6 +529,10 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, ) + if model_output.output_seed is not None: + # Graph output addresses are overwritten by the next replay even + # without padding. Detach them before returning to the backend. + new_model_output.output_seed = model_output.output_seed[:origin_batch_size].clone() return new_model_output def _create_unpad_prefill_model_output( @@ -488,6 +541,8 @@ def _create_unpad_prefill_model_output( new_model_output = copy.copy(padded_model_output) # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。 new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if padded_model_output.output_seed is not None: + new_model_output.output_seed = padded_model_output.output_seed[:origin_batch_size].clone() new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -585,9 +640,15 @@ def _decode( # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = self.graph is not None and self.graph.can_run( - batch_size=infer_batch_size, - max_len_in_batch=infer_max_kv_seq_len, + # Auxiliary single-batch calls cannot replay a graph captured with two + # microbatches; normal DP decode uses _microbatch_overlap_decode_cuda. + use_cuda_graph = ( + self.graph is not None + and not self.graph.enable_decode_microbatch_overlap + and self.graph.can_run( + batch_size=infer_batch_size, + max_len_in_batch=infer_max_kv_seq_len, + ) ) need_capture = False if use_cuda_graph: @@ -620,7 +681,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -674,6 +734,7 @@ def prefill_func(input_tensors, _infer_state): if infer_state.need_dp_prefill_balance: last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs) + output_seed = self._capture_output_seed(last_input_embs, infer_state) predict_logits = self.post_infer.token_forward(last_input_embs, infer_state, self.pre_post_weight) hidden_collector = infer_state.hidden_collector hidden_collector.add_final_hidden(last_input_embs) @@ -681,6 +742,7 @@ def prefill_func(input_tensors, _infer_state): logits=predict_logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, + output_seed=output_seed, ) # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候 @@ -701,6 +763,7 @@ def _token_forward(self, infer_state: InferStateInfo): hidden_collector.add(layer_index=i, hidden=input_embs) last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) + output_seed = self._capture_output_seed(last_input_embs, infer_state) predict_logits: torch.Tensor = self.post_infer.token_forward( last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight ) @@ -709,6 +772,7 @@ def _token_forward(self, infer_state: InferStateInfo): model_output = ModelOutput( logits=predict_logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), + output_seed=output_seed, ) # 在 cuda graph 模式下，输出需要转为 no ref tensor, 加强mem pool 的复用，降低显存的使用。 @@ -953,6 +1017,8 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs) last_input_embs1 = infer_state1._all_to_all_unbalance_get(data=last_input_embs1) + output_seed = self._capture_output_seed(last_input_embs, infer_state) + output_seed1 = self._capture_output_seed(last_input_embs1, infer_state1) predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward( last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight ) @@ -964,11 +1030,13 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state logits=predict_logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, + output_seed=output_seed, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, + output_seed=output_seed1, ) return model_output, model_output1 @@ -1003,6 +1071,8 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) last_input_embs1 = self.post_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1) + output_seed = self._capture_output_seed(last_input_embs, infer_state) + output_seed1 = self._capture_output_seed(last_input_embs1, infer_state1) predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward( last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight ) @@ -1012,10 +1082,12 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: model_output = ModelOutput( logits=predict_logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), + output_seed=output_seed, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), + output_seed=output_seed1, ) if infer_state.is_cuda_graph: diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index ae645d4b7b..075371873c 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -200,6 +200,13 @@ class ModelOutput: # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None + # Exact-prefix replay seed: one raw final hidden per logits row, before + # final norm. Public model.forward() returns independently owned storage; + # graph-internal outputs are cloned when leaving the graph replay wrapper. + # This is distinct from spec_hidden, which may contain intermediate layers + # and may be normalized in-place by an MTP draft model. + output_seed: Optional[torch.Tensor] = None + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() @@ -207,3 +214,5 @@ def __post_init__(self) -> None: def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) self.mtp_collector.to_no_ref_tensor() + if self.output_seed is not None: + self.output_seed = tensor_to_no_ref_tensor(self.output_seed) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/capture_state.py b/lightllm/common/basemodel/triton_kernel/linear_att/capture_state.py new file mode 100644 index 0000000000..a463b4f2b9 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/capture_state.py @@ -0,0 +1,164 @@ +"""Freeze a bounded number of recurrent states without a GPU-to-CPU decision. + +The output buffers are owned by the caller and must not be reused while a +checkpoint transfer still reads them. Candidate selection is stable across TP +ranks: logical input order, rather than the order of GPU atomics, chooses slots. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["N"]) +def _select_capture_candidates( + req_indices, + state_rows, + exact_lengths, + capture_mask, + selected_reqs, + selected_rows, + selected_lengths, + source_rows, + N, + CAPACITY: tl.constexpr, + MAX_REQS: tl.constexpr, + MTP_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + index = tl.arange(0, BLOCK) + tl.store(selected_reqs + index, -1, index < CAPACITY) + tl.store(selected_rows + index, 0, index < CAPACITY) + tl.store(selected_lengths + index, 0, index < CAPACITY) + tl.store(source_rows + index, -1, index < CAPACITY) + req = tl.load(req_indices + index, index < N, other=-1) + row = tl.load(state_rows + index, index < N, other=-1) + length = tl.load(exact_lengths + index, index < N, other=0) + enabled = tl.load(capture_mask + index, index < N, other=0) + valid = enabled & (req >= 0) & (req < MAX_REQS) & (row >= 0) & (row < MTP_SIZE) & (length > 0) + slot = tl.cumsum(valid.to(tl.int32)) - 1 + keep = valid & (slot < CAPACITY) + tl.debug_barrier() + tl.store(selected_reqs + slot, req, keep) + tl.store(selected_rows + slot, row, keep) + tl.store(selected_lengths + slot, length, keep) + tl.store(source_rows + slot, index, keep) + + +@triton.jit +def _freeze_linear_states( + conv, + ssm, + selected_reqs, + selected_rows, + out_conv, + out_ssm, + CONV_LAYER_STRIDE: tl.constexpr, + CONV_REQ_STRIDE: tl.constexpr, + CONV_DIM_STRIDE: tl.constexpr, + CONV_WIDTH_STRIDE: tl.constexpr, + SSM_LAYER_STRIDE: tl.constexpr, + SSM_REQ_STRIDE: tl.constexpr, + LAYERS: tl.constexpr, + CONV_WIDTH: tl.constexpr, + CONV_ELEMENTS: tl.constexpr, + SSM_ELEMENTS: tl.constexpr, + MTP_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + slot, layer, block = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(selected_reqs + slot) + if req < 0: + return + row = tl.load(selected_rows + slot) + index = block * BLOCK + tl.arange(0, BLOCK) + conv_source = ( + layer * CONV_LAYER_STRIDE + + req * CONV_REQ_STRIDE + + (index // CONV_WIDTH) * CONV_DIM_STRIDE + + (row + index % CONV_WIDTH) * CONV_WIDTH_STRIDE + ) + conv_value = tl.load(conv + conv_source, index < CONV_ELEMENTS, other=0) + tl.store(out_conv + (slot * LAYERS + layer) * CONV_ELEMENTS + index, conv_value, index < CONV_ELEMENTS) + ssm_source = layer * SSM_LAYER_STRIDE + (req * MTP_SIZE + row) * SSM_REQ_STRIDE + index + ssm_value = tl.load(ssm + ssm_source, index < SSM_ELEMENTS, other=0) + tl.store(out_ssm + (slot * LAYERS + layer) * SSM_ELEMENTS + index, ssm_value, index < SSM_ELEMENTS) + + +def freeze_linear_states( + conv: torch.Tensor, + ssm: torch.Tensor, + req_indices: torch.Tensor, + state_rows: torch.Tensor, + exact_lengths: torch.Tensor, + capture_mask: torch.Tensor, + selected_reqs: torch.Tensor, + selected_rows: torch.Tensor, + selected_lengths: torch.Tensor, + source_rows: torch.Tensor, + out_conv: torch.Tensor, + out_ssm: torch.Tensor, + mtp_size: int, +) -> None: + """Gather at most ``out_conv.shape[0]`` masked candidates into owned slots. + + ``state_rows`` contains request-local MTP row numbers, not flattened model + output rows. A row r is the state *after processing input row r*, before the + token sampled from that row. The caller supplies that exact prefix length. + No candidate metadata is read back to the CPU by this function. + """ + count = req_indices.numel() + capacity, layers, conv_dim, conv_width = out_conv.shape + assert capacity > 0 and mtp_size > 0 + for tensor in (req_indices, state_rows, exact_lengths, capture_mask): + assert tensor.is_cuda and tensor.ndim == 1 and tensor.numel() == count and tensor.is_contiguous() + assert tensor.device == conv.device + for tensor in (selected_reqs, selected_rows, selected_lengths, source_rows): + assert tensor.is_cuda and tensor.shape == (capacity,) and tensor.is_contiguous() + assert tensor.device == conv.device and tensor.dtype in (torch.int32, torch.int64) + assert req_indices.dtype in (torch.int32, torch.int64) + assert state_rows.dtype in (torch.int32, torch.int64) + assert exact_lengths.dtype in (torch.int32, torch.int64) + assert capture_mask.dtype == torch.bool + assert conv.is_cuda and ssm.is_cuda and conv.device == ssm.device + assert out_conv.device == conv.device and out_ssm.device == conv.device + assert conv.ndim == 4 and conv.shape[0] == layers and conv.shape[2] == conv_dim + assert conv.shape[-1] == conv_width + mtp_size - 1 + assert ssm.shape[0] == layers and ssm.shape[1] == conv.shape[1] * mtp_size + assert ssm.is_contiguous() and out_conv.is_contiguous() and out_ssm.is_contiguous() + assert out_ssm.shape == (capacity, layers, *ssm.shape[2:]) + assert conv.dtype == out_conv.dtype and ssm.dtype == out_ssm.dtype + _select_capture_candidates[(1,)]( + req_indices, + state_rows, + exact_lengths, + capture_mask, + selected_reqs, + selected_rows, + selected_lengths, + source_rows, + N=count, + CAPACITY=capacity, + MAX_REQS=conv.shape[1] - 1, # The final request slot is graph padding. + MTP_SIZE=mtp_size, + BLOCK=triton.next_power_of_2(max(count, capacity)), + ) + ssm_elements = ssm[0, 0].numel() + conv_elements = conv_dim * conv_width + _freeze_linear_states[(capacity, layers, triton.cdiv(max(conv_elements, ssm_elements), 256))]( + conv, + ssm, + selected_reqs, + selected_rows, + out_conv, + out_ssm, + *conv.stride(), + ssm.stride(0), + ssm.stride(1), + LAYERS=layers, + CONV_WIDTH=conv_width, + CONV_ELEMENTS=conv_elements, + SSM_ELEMENTS=ssm_elements, + MTP_SIZE=mtp_size, + BLOCK=256, + ) diff --git a/lightllm/common/basemodel/triton_kernel/mtp_utils.py b/lightllm/common/basemodel/triton_kernel/mtp_utils.py index 7e943f2925..9491b1b69c 100644 --- a/lightllm/common/basemodel/triton_kernel/mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/mtp_utils.py @@ -309,12 +309,11 @@ def mtp_scatter_next_token_ids( ) -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _fwd_kernel_gen_b_req_mtp_start_loc( b_mtp_index, b_req_mtp_start_loc, - num_reqs: tl.constexpr, - batch_size: tl.constexpr, + batch_size, BLOCK_SIZE: tl.constexpr, ): offset = tl.arange(0, BLOCK_SIZE) @@ -333,7 +332,6 @@ def gen_b_req_mtp_start_loc(b_mtp_index: torch.Tensor, num_reqs: int): _fwd_kernel_gen_b_req_mtp_start_loc[grid]( b_mtp_index=b_mtp_index, b_req_mtp_start_loc=b_req_mtp_start_loc, - num_reqs=num_reqs, batch_size=batch_size, BLOCK_SIZE=BLOCK_SIZE, num_warps=8, diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index d217e05c78..f5ae9040db 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -27,7 +27,6 @@ class MemoryManager: - operator_class = NormalMemOperator def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): @@ -91,7 +90,8 @@ def profile_size(self, mem_fraction): available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size + checkpoint_staging_size = self.get_checkpoint_staging_size() + available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size - checkpoint_staging_size self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") @@ -100,11 +100,26 @@ def profile_size(self, mem_fraction): logger.info( f"{str(available_memory)} GB space is available after load the model weight\n" f"{str(pd_kv_move_buffer_size / 1024 ** 2)} MB is reserved for PD KV transfer buffer\n" + f"{str(checkpoint_staging_size / 1024 ** 2)} MB is reserved for exact checkpoint capture/transfer\n" f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n" f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n" ) return + def get_checkpoint_staging_size(self): + args = get_env_start_args() + config = getattr(self, "linear_config", None) + if not getattr(args, "enable_exact_prefix_cache", False) or config is None: + return 0 + state_bytes = config.linear_layer_num * ( + math.prod(config.get_conv_state_shape()) * torch._utils._element_size(config.conv_state_dtype) + + math.prod(config.get_ssm_state_shape()) * torch._utils._element_size(config.ssm_state_dtype) + ) + batch_count = 4 if (args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap) else 2 + # Include frozen packed tail KV plus one bounded page gather temporary. + capture_bytes = batch_count * args.exact_prefix_cache_capture_slots * (state_bytes + self.get_cell_size()) + return capture_bytes + args.exact_prefix_cache_page_size * self.get_cell_size() + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): # 在初始化 kv_buffer 的时候，每层多初始化了一个 token，这个 token 永远不会被真的被对外 # 分配，内部实际也没有管理，这个token是预留来对一些特殊的运行模式，如多dp下，overlap microbatch diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index 907cc494a6..4ec3a2d65b 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -39,6 +39,11 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): return def _init_linear_att_buffers(self): + if getattr(get_env_start_args(), "enable_exact_prefix_cache", False): + self.linear_att_big_page_buffers = None + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = None + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = None + return big_page_token_num = ( get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size ) diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 967bc9bf7a..c40829a981 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -1,4 +1,5 @@ -from typing import TYPE_CHECKING +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Optional import torch @@ -14,6 +15,82 @@ from lightllm.server.router.model_infer.infer_batch import InferReq +@dataclass +class LinearStateSnapshot: + """Canonical state for exactly ``exact_len`` processed input tokens. + + Both tensors have independent storage, or are protected by the owning + capture ticket. A CPU snapshot retains its GPU source until the DMA event + completes. Merely recording an event does not preserve source lifetime. + """ + + conv_state: torch.Tensor + ssm_state: torch.Tensor + exact_len: int + ready_event: torch.cuda.Event + _sources: tuple = field(default_factory=tuple, repr=False) + + def wait(self, stream: Optional[torch.cuda.Stream] = None) -> None: + (stream or torch.cuda.current_stream()).wait_event(self.ready_event) + + def is_ready(self) -> bool: + ready = self.ready_event.query() + if ready: + self._sources = () + return ready + + def to_cpu(self, stream: Optional[torch.cuda.Stream] = None) -> "LinearStateSnapshot": + """Copy into pinned storage; callers must retain the result until ready.""" + if self.conv_state.device.type == "cpu": + return self + stream = stream or torch.cuda.current_stream(self.conv_state.device) + conv = torch.empty(self.conv_state.shape, dtype=self.conv_state.dtype, device="cpu", pin_memory=True) + ssm = torch.empty(self.ssm_state.shape, dtype=self.ssm_state.dtype, device="cpu", pin_memory=True) + with torch.cuda.stream(stream): + self.wait(stream) + conv.copy_(self.conv_state, non_blocking=True) + ssm.copy_(self.ssm_state, non_blocking=True) + self.conv_state.record_stream(stream) + self.ssm_state.record_stream(stream) + event = torch.cuda.Event() + event.record(stream) + return LinearStateSnapshot(conv, ssm, self.exact_len, event, (self,)) + + +@dataclass +class LinearStateCaptureStaging: + """Bounded, reusable GPU storage. The caller owns its exclusive lease. + + Reuse is allowed only once every consumer/transfer of the previous capture + has completed, not merely when ``ready_event`` is signaled. Slot metadata + with req_indices == -1 denotes an unused slot and its state must not be read. + """ + + conv_state: torch.Tensor + ssm_state: torch.Tensor + req_indices: torch.Tensor + mtp_rows: torch.Tensor + exact_lengths: torch.Tensor + source_rows: torch.Tensor + ready_event: Optional[torch.cuda.Event] = None + + def clone_snapshot(self, slot: int, exact_len: int) -> LinearStateSnapshot: + """Detach a selected slot before reusing this staging allocation. + + The caller must have validated the selected slot's metadata and length. + This method does not perform a hidden device-to-host metadata read. + """ + assert 0 <= slot < self.conv_state.shape[0] and exact_len > 0 + assert self.ready_event is not None + stream = torch.cuda.current_stream(self.conv_state.device) + stream.wait_event(self.ready_event) + conv = self.conv_state[slot].clone() + ssm = self.ssm_state[slot].clone() + event = torch.cuda.Event() + event.record(stream) + return LinearStateSnapshot(conv, ssm, exact_len, event) + + class ReqManagerForMamba(ReqManager): def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): super().__init__(max_request_num, max_sequence_length, mem_manager) @@ -61,6 +138,100 @@ def init_linear_att_state(self, req: "InferReq"): self.req_to_mtp_state_index[req.req_idx] = 0 return + def allocate_linear_state_staging(self, capacity: int) -> LinearStateCaptureStaging: + """Allocate once outside CUDA graph replay, subject to a caller budget.""" + if capacity <= 0: + raise ValueError("linear state staging capacity must be positive") + config = self.linear_config + device = self.req_to_conv_state.buffer.device + conv = torch.empty( + (capacity, config.linear_layer_num, *config.get_conv_state_shape()), + dtype=config.conv_state_dtype, + device=device, + ) + ssm = torch.empty( + (capacity, config.linear_layer_num, *config.get_ssm_state_shape()), + dtype=config.ssm_state_dtype, + device=device, + ) + metadata = [torch.empty(capacity, dtype=torch.int32, device=device) for _ in range(4)] + return LinearStateCaptureStaging(conv, ssm, *metadata) + + def freeze_linear_states( + self, + req_indices: torch.Tensor, + mtp_rows: torch.Tensor, + exact_lengths: torch.Tensor, + capture_mask: torch.Tensor, + staging: LinearStateCaptureStaging, + ) -> LinearStateCaptureStaging: + """Freeze masked candidates on the producing stream before the next write. + + Candidates are retained in logical row order until staging is full; + excess candidates are skipped. ``mtp_rows`` is request-local. In a + prefill batch pass zeros; after verification select the exact accepted + input row, accounting for stop truncation and the unprocessed sample. + This method performs no host synchronization or per-request allocation. + """ + from lightllm.common.basemodel.triton_kernel.linear_att.capture_state import freeze_linear_states + + freeze_linear_states( + self.req_to_conv_state.buffer, + self.req_to_ssm_state.buffer, + req_indices, + mtp_rows, + exact_lengths, + capture_mask, + staging.req_indices, + staging.mtp_rows, + staging.exact_lengths, + staging.source_rows, + staging.conv_state, + staging.ssm_state, + self.mtp_step + 1, + ) + staging.ready_event = torch.cuda.Event() + staging.ready_event.record() + return staging + + def freeze_linear_state(self, req_idx: int, exact_len: int, mtp_row: int = 0) -> LinearStateSnapshot: + """Freeze one known boundary; use the bounded batch API on decode paths.""" + if not 0 <= req_idx < self.HOLD_REQUEST_ID: + raise ValueError("cannot capture an invalid or padding request slot") + if not 0 <= mtp_row <= self.mtp_step or exact_len <= 0: + raise ValueError("invalid exact prefix length or request-local MTP row") + conv_width = self.linear_config.get_conv_state_shape()[-1] + conv = self.req_to_conv_state.buffer[:, req_idx, ..., mtp_row : mtp_row + conv_width].clone() + ssm = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1) + mtp_row, ...].clone() + event = torch.cuda.Event() + event.record() + return LinearStateSnapshot(conv, ssm, exact_len, event) + + def restore_linear_state(self, snapshot: LinearStateSnapshot, req_idx: int) -> torch.cuda.Event: + """Restore a canonical snapshot without carrying speculative state rows.""" + if not 0 <= req_idx < self.HOLD_REQUEST_ID: + raise ValueError("cannot restore an invalid or padding request slot") + config = self.linear_config + if snapshot.conv_state.shape != (config.linear_layer_num, *config.get_conv_state_shape()): + raise ValueError("incompatible convolution state layout") + if snapshot.ssm_state.shape != (config.linear_layer_num, *config.get_ssm_state_shape()): + raise ValueError("incompatible SSM state layout") + if snapshot.conv_state.dtype != config.conv_state_dtype or snapshot.ssm_state.dtype != config.ssm_state_dtype: + raise ValueError("incompatible linear state dtype") + snapshot.wait() + conv_width = config.get_conv_state_shape()[-1] + self.req_to_conv_state.buffer[:, req_idx, ...].zero_() + self.req_to_conv_state.buffer[:, req_idx, ..., :conv_width].copy_(snapshot.conv_state, non_blocking=True) + ssm_start = req_idx * (self.mtp_step + 1) + self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + self.mtp_step + 1, ...].zero_() + self.req_to_ssm_state.buffer[:, ssm_start, ...].copy_(snapshot.ssm_state, non_blocking=True) + if self.req_to_mtp_state_index is not None: + self.req_to_mtp_state_index[req_idx] = 0 + event = torch.cuda.Event() + event.record() + # The caller retains snapshot until this H2D/read event has completed. + return event + def get_mamba_cache(self, layer_idx_in_all: int): assert ( 0 <= layer_idx_in_all < self.linear_config.all_layer_num diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e6c077dbe4..f45475d1e5 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -891,6 +891,35 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=False, help="""Enable prefix prompt cache fetch for data parallel inference, disabled by default.""", ) + parser.add_argument( + "--enable_exact_prefix_cache", + action="store_true", + help="Cache hybrid attention checkpoints at exact token lengths, including generated prefixes.", + ) + parser.add_argument( + "--exact_prefix_cache_mb", + type=int, + default=1024, + help="CPU checkpoint KV, state and hidden budget in MiB per TP rank.", + ) + parser.add_argument( + "--exact_prefix_cache_entries", + type=int, + default=128, + help="Maximum number of retained exact checkpoints per TP rank.", + ) + parser.add_argument( + "--exact_prefix_cache_page_size", + type=int, + default=8192, + help="CPU KV page capacity in tokens; checkpoint lengths do not need page alignment.", + ) + parser.add_argument( + "--exact_prefix_cache_capture_slots", + type=int, + default=4, + help="Bounded checkpoint capture slots per overlapping batch; excess candidates are skipped.", + ) parser.add_argument( "--linear_att_hash_page_size", type=int, diff --git a/lightllm/server/api_http_pd.py b/lightllm/server/api_http_pd.py index 1d8b2112fc..b97f93a50c 100644 --- a/lightllm/server/api_http_pd.py +++ b/lightllm/server/api_http_pd.py @@ -12,7 +12,7 @@ import pickle import ujson as json -from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from fastapi import APIRouter, HTTPException, Request, WebSocket, WebSocketDisconnect from lightllm.server.pd_io_struct import ObjType from lightllm.utils.envs_utils import get_lightllm_websocket_max_message_size @@ -23,6 +23,22 @@ router = APIRouter() +@router.get("/pd_checkpoint/registry") +async def checkpoint_registry(request: Request): + """Internal PD discovery; bulk checkpoint data bypasses PD Master.""" + from .api_http import g_objs + from lightllm.server.router.model_infer.mode_backend.pd.checkpoint_transport import ( + checkpoint_registry_token, + read_checkpoint_registry, + ) + + if g_objs.args.run_mode not in ("prefill", "decode"): + return {"ranks": []} + if request.headers.get("Authorization") != f"Bearer {checkpoint_registry_token()}": + raise HTTPException(status_code=403, detail="Internal PD credentials required") + return {"ranks": read_checkpoint_registry()} + + @router.websocket("/pd_register") async def register_and_keep_alive(websocket: WebSocket): from .api_http import g_objs @@ -32,7 +48,8 @@ async def register_and_keep_alive(websocket: WebSocket): client_ip, client_port = websocket.client logger.info(f"Client connected from IP: {client_ip}, Port: {client_port}") regist_json = json.loads(await websocket.receive_text()) - logger.info(f"received regist_json {regist_json}") + log_registration = dict(regist_json, checkpoint_registry_token="") + logger.info(f"received regist_json {log_registration}") await g_objs.httpserver_manager.register_pd(regist_json, websocket) try: @@ -45,18 +62,18 @@ async def register_and_keep_alive(websocket: WebSocket): await g_objs.httpserver_manager.put_to_handle_queue(obj) except asyncio.TimeoutError: - logger.warning(f"client {regist_json} heartbeat timed out after {heartbeat_timeout_seconds} seconds") + logger.warning(f"client {log_registration} heartbeat timed out after {heartbeat_timeout_seconds} seconds") try: await websocket.close(code=1011, reason="PD heartbeat timed out") except BaseException: - logger.debug(f"failed to close timed-out client {regist_json}", exc_info=True) + logger.debug(f"failed to close timed-out client {log_registration}", exc_info=True) except WebSocketDisconnect as e: - logger.info(f"client {regist_json} disconnected: {str(e)}") + logger.info(f"client {log_registration} disconnected: {str(e)}") except BaseException as e: - logger.error(f"client {regist_json} has error {str(e)}") + logger.error(f"client {log_registration} has error {str(e)}") logger.exception(str(e)) finally: - logger.error(f"client {regist_json} removed") + logger.error(f"client {log_registration} removed") await g_objs.httpserver_manager.remove_pd(regist_json) return diff --git a/lightllm/server/api_lightllm.py b/lightllm/server/api_lightllm.py index fe6547d63e..14797bd79f 100644 --- a/lightllm/server/api_lightllm.py +++ b/lightllm/server/api_lightllm.py @@ -31,7 +31,6 @@ async def lightllm_get_score(request: Request, httpserver_manager: HttpServerMan async def lightllm_generate(request: Request, httpserver_manager: HttpServerManager) -> Response: - request_dict = await request.json() prompt = request_dict.pop("inputs") sample_params_dict = request_dict["parameters"] @@ -116,7 +115,6 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana async def lightllm_generate_stream(request: Request, httpserver_manager: HttpServerManager) -> Response: - request_dict = await request.json() prompt = request_dict.pop("inputs") sample_params_dict = request_dict["parameters"] @@ -156,6 +154,16 @@ async def stream_results() -> AsyncGenerator[bytes, None]: "input_usage": input_usage, } ret["token"]["logprobs"] = metadata["logprobs"] + # Keep native streaming observability aligned with /generate + # return_details. Omitted fields stay omitted on older backends. + for key in ( + "prompt_cache_len", + "mtp_accepted_token_num", + "mtp_verify_token_num", + "mtp_verify_step_num", + ): + if key in metadata: + ret["token"][key] = metadata[key] if "prompt_logprobs" in metadata: ret["prompt_logprobs"] = metadata["prompt_logprobs"] ret["prompt_token_ids"] = metadata.get("prompt_token_ids") diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 7c7ac9fe48..209aa99dec 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -256,6 +256,32 @@ def _launch_subprocesses(args: StartArgs): f"but got {args.batch_max_tokens}, {args.chunked_prefill_size}" ) + if args.enable_exact_prefix_cache: + if not is_linear_att_mixed_model(args.model_dir): + raise ValueError("--enable_exact_prefix_cache requires a hybrid linear-attention model") + for name in ( + "exact_prefix_cache_mb", + "exact_prefix_cache_entries", + "exact_prefix_cache_page_size", + "exact_prefix_cache_capture_slots", + ): + if getattr(args, name) <= 0: + raise ValueError(f"--{name} must be positive") + if args.disable_dynamic_prompt_cache: + raise ValueError("exact prefix cache requires dynamic prompt cache") + if args.enable_cpu_cache or args.enable_disk_cache: + raise ValueError( + "exact prefix cache owns its CPU store; disable legacy --enable_cpu_cache/--enable_disk_cache" + ) + if args.diverse_mode or args.enable_dp_prompt_cache_fetch: + raise ValueError("exact prefix cache does not support diverse mode or legacy DP cache fetch") + if args.mtp_step and args.enable_ep_moe: + raise ValueError("exact MTP resume with expert parallelism requires coordinated auxiliary scheduling") + if args.enable_rl: + raise ValueError( + "exact prefix cache requires immutable deployment weights; online RL updates are unsupported" + ) + # linear att cache 参数自动设置 if args.linear_att_cache_size is None: # linear_att_cache_size 只会在 qwen3.5 等混合线性层模型中生效。 diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 9729a8205c..e30c55b554 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -153,6 +153,8 @@ class Req(ctypes.Structure): ("token_hash_page_len_list", TokenPageLenList), # 用于保存查找匹配到的可以被复用的cpu cache 页面信息。 ("cpu_cache_match_page_indexes", CpuCachePageList), + # 历史碎页的实际前缀终点；0 表示沿用原始页边界。与 offload 的 hash/长度列表分开保存。 + ("cpu_cache_match_tail_len", ctypes.c_int), ] def get_str(self): @@ -189,6 +191,7 @@ def init( self.candetoken_out_len = 0 self.prompt_cache_len = 0 self.cpu_prompt_cache_len = 0 + self.cpu_cache_match_tail_len = 0 self.disk_prompt_cache_len = 0 self.finish_token_index = -1 self.can_released_mark = False @@ -533,5 +536,4 @@ def get_decode_need_tokens(self): return need_tokens def get_first_router_need_tokens(self): - return min(self.input_len + self.shm_cur_output_len, self.chunked_prefill_size) diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py index 20bd6c61ea..05b783d38c 100644 --- a/lightllm/server/core/objs/sampling_params.py +++ b/lightllm/server/core/objs/sampling_params.py @@ -290,6 +290,9 @@ class SamplingParams(ctypes.Structure): # P/D 节点的资源等待超时，由 PD Master 下发。非负值用于控制 shm_req 申请和 # Router 等待进入推理系统的时限；负数表示永久等待。 ("pd_node_resource_wait_timeout_seconds", ctypes.c_int), + # Set by PD Master from the selected, registered P node. Never accepted from an API caller. + ("pd_checkpoint_owner_url", ctypes.c_char * 512), + ("pd_checkpoint_owner_auth", ctypes.c_char * 128), ("suggested_dp_index", ctypes.c_int), # suggest dp index, deepseekv2 dp mode, use to suggest used dp_index # in pd split mode, use to keep the id of pd master ("pd_master_node_id", NodeUUId), @@ -336,6 +339,8 @@ def init(self, tokenizer, **kwargs): # 这两个字段是 PD Master 的内部调度信息，不能由外部请求参数开启或修改。 self.pd_high_priority_request = False self.pd_node_resource_wait_timeout_seconds = -1 + self.pd_checkpoint_owner_url = b"" + self.pd_checkpoint_owner_auth = b"" self.suggested_dp_index = kwargs.get("suggested_dp_index", -1) self.skip_special_tokens = kwargs.get("skip_special_tokens", SKIP_SPECIAL_TOKENS) diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9c89975de7..da866d575a 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -236,6 +236,11 @@ class StartArgs: # hybrid attention model (Qwen3Next) linear_att_hash_page_size: int = field(default=512) + enable_exact_prefix_cache: bool = field(default=False) + exact_prefix_cache_mb: int = field(default=1024) + exact_prefix_cache_entries: int = field(default=128) + exact_prefix_cache_page_size: int = field(default=8192) + exact_prefix_cache_capture_slots: int = field(default=4) linear_att_page_block_num: int = field(default=10000000) disable_linear_att_small_page_cpu_cache: bool = field(default=False) linear_att_cache_size: Optional[int] = field(default=None) diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index cea3cb6fc9..694584470e 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -328,7 +328,6 @@ async def generate( # 用于等待 pd_master 下发的交换信息 pd_event: asyncio.Event = None, ) -> AsyncGenerator[Tuple[int, str, dict, FinishStatus], None]: - start_time = time.time() request_headers = request.headers if request is not None else {} group_request_id = self.alloc_req_id(sampling_params) @@ -413,7 +412,10 @@ async def generate( decode_node_info: PDDecodeNodeInfo = pd_event.decode_node_info sampling_params.pd_kv_trans_params.set(pickle.dumps(decode_node_info)) - if decode_node_info.ready_kv_len == len(prompt_ids) - 1: + first_token_owner = getattr(decode_node_info, "first_token_owner", None) + if first_token_owner == "decode" or ( + first_token_owner is None and decode_node_info.ready_kv_len == len(prompt_ids) - 1 + ): # 如果 decode 节点的 ready_kv_len 和 prefill encode 的 len(prompt ids) -1 相等，说明不需要进行 prefill # 直接 raise PDPrefillNodeStopGenToken raise PDPrefillNodeStopGenToken(group_request_id=group_request_id) @@ -729,7 +731,6 @@ async def transfer_to_next_module( self, group_req_objs: Optional[GroupReqObjs] = None, ): - if self.pd_mode.is_P_or_NORMAL(): if not self.args.disable_vision: self.send_to_visual.send_pyobj(group_req_objs.to_group_req_index(), protocol=pickle.HIGHEST_PROTOCOL) @@ -772,7 +773,6 @@ async def _wait_to_token_package( req_status: "ReqStatus", request: Request, ): - event = req_status.event unfinished_count = sampling_params.best_of out_token_counter = 0 diff --git a/lightllm/server/httpserver/pd_loop.py b/lightllm/server/httpserver/pd_loop.py index 51b7eea33b..96d89b3eb0 100644 --- a/lightllm/server/httpserver/pd_loop.py +++ b/lightllm/server/httpserver/pd_loop.py @@ -94,7 +94,6 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O # 下方应用层心跳已负责存活检测，禁用协议层 keepalive，避免繁忙连接被误断。 ping_interval=None, ) as websocket: - sock = websocket.transport.get_extra_info("socket") sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) @@ -107,9 +106,15 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O "mode": manager.pd_mode.value, "start_args": args_dict, } + if getattr(manager.args, "enable_exact_prefix_cache", False): + from lightllm.server.router.model_infer.mode_backend.pd.checkpoint_transport import ( + checkpoint_registry_token, + ) + + regist_json["checkpoint_registry_token"] = checkpoint_registry_token() await websocket.send(json.dumps(regist_json)) - logger.info(f"Sent registration JSON: {regist_json}") + logger.info(f"Sent registration JSON: {dict(regist_json, checkpoint_registry_token='')}") # 转发任务 forwarding_tokens_task = asyncio.create_task(_up_tokens_to_pd_master(forwarding_queue, websocket)) @@ -289,7 +294,6 @@ async def _send_heartbeat_to_pd_master(websocket: ClientConnection): # 获取节点负载信息 def _get_load_info() -> dict: - from lightllm.server.api_http import g_objs assert g_objs.shared_token_load is not None, "shared_token_load is not initialized" diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index 96f0d7203e..817d6b7cd9 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -108,6 +108,8 @@ async def update_req_status(self, upkv_status: PDUpKVStatus): return def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwargs=None): + if isinstance(prompt, list): + return len(prompt) kwargs = {} if kwargs is None else kwargs prompt_ids = self.tokenizer.encode(prompt, None, **kwargs) image_tokens = 0 @@ -169,7 +171,12 @@ async def _generate( multimodal_params: MultimodalParams, request: Request, ): - assert isinstance(prompt, str), "prompt must be str" + if not isinstance(prompt, (str, list)): + raise ValueError("prompt must be a string or a list of token IDs") + if isinstance(prompt, list) and ( + not prompt or any(not isinstance(token, int) or isinstance(token, bool) or token < 0 for token in prompt) + ): + raise ValueError("prompt token IDs must be a nonempty list of nonnegative integers") start_time = time.time() await multimodal_params.verify_and_preload(request) # 计算输入的 input_token_num, 进行校验，如果输入+输出参数设置太长，则将 @@ -296,8 +303,11 @@ async def _generate_one_attempt( pending_prefill_load_chars = None try: + # Cache-aware selection remains an approximate text affinity hint. + # Token-ID requests retain their original IDs all the way to P/D. + routing_prompt = self.tokenizer.decode(prompt) if isinstance(prompt, list) else prompt p_node, d_node, selection_extra_info = await self.select_p_d_node( - prompt, origin_sampling_params, multimodal_params + routing_prompt, origin_sampling_params, multimodal_params ) if not p_node or not d_node: logger.error(f"{origin_request_id}: No p_node or d_node found") @@ -323,6 +333,7 @@ async def _generate_one_attempt( ) history_gen_token_strs = [] + history_gen_token_ids = [] origin_prompt_cache_len = None remaining_max_new_tokens = origin_sampling_params.max_new_tokens segment_index = 0 @@ -353,8 +364,12 @@ async def _generate_one_attempt( # 分段请求始终复用循环外选定的 P 节点；这里只按每段实际发送的 # prompt 更新该节点的在途 prefill 负载，不会重新选点。 - block_prompt = prompt + "".join(history_gen_token_strs) - pending_prefill_load_chars = len(block_prompt) + block_prompt = ( + prompt + history_gen_token_ids + if isinstance(prompt, list) + else prompt + "".join(history_gen_token_strs) + ) + pending_prefill_load_chars = len(routing_prompt) + sum(map(len, history_gen_token_strs)) p_node.dispatched_prompt_chars += pending_prefill_load_chars p_node.dispatched_req_num += 1 results_generator = self._wait_to_token_package( @@ -387,6 +402,8 @@ async def _generate_one_attempt( # 容量 marker 已在上方过滤，能走到这里的每个 token 都立即扣减全局剩余输出额度。 remaining_max_new_tokens -= 1 history_gen_token_strs.append(request_output) + if isinstance(prompt, list): + history_gen_token_ids.append(int(metadata["id"])) prompt_tokens = min(prompt_tokens, metadata["prompt_tokens"]) metadata["prompt_tokens"] = prompt_tokens if origin_prompt_cache_len is None: @@ -396,7 +413,7 @@ async def _generate_one_attempt( if not raw_finish_status.is_error_finished(): # 只有收到成功的推理结果后才将 prompt 写入前缀树，避免尚未进入 # 推理或已失败的请求被后续请求误判为可复用 cache。 - self.pd_manager.selector.insert_prompt_cache(prompt, p_node) + self.pd_manager.selector.insert_prompt_cache(routing_prompt, p_node) metadata["prompt_cache_len"] = origin_prompt_cache_len or 0 yield origin_request_id, request_output, metadata, raw_finish_status @@ -511,6 +528,10 @@ async def fetch_pd_stream( ): group_request_id = sampling_params.group_request_id sampling_params.pd_master_node_id.initialize(self.args.pd_node_id) + # This address comes from the registered node selected by the server, + # rather than an untrusted sampling parameter supplied by the caller. + sampling_params.pd_checkpoint_owner_url = f"http://{p_node.client_ip_port}".encode("utf-8") + sampling_params.pd_checkpoint_owner_auth = (p_node.checkpoint_registry_token or "").encode("utf-8") req_status = ReqStatus(group_request_id, p_node, d_node) self.req_id_to_out_inf[group_request_id] = req_status @@ -565,7 +586,12 @@ async def fetch_pd_stream( ) first_token_gen = False - needs_prefill_first_token = decode_node_info.ready_kv_len != len(prompt_ids) - 1 + first_token_owner = getattr(decode_node_info, "first_token_owner", None) + needs_prefill_first_token = ( + first_token_owner == "prefill" + if first_token_owner is not None + else decode_node_info.ready_kv_len != len(prompt_ids) - 1 + ) prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( req_status=req_status, request=request, diff --git a/lightllm/server/multi_level_kv_cache/manager.py b/lightllm/server/multi_level_kv_cache/manager.py index ef5b7369c9..5b9ed7905f 100644 --- a/lightllm/server/multi_level_kv_cache/manager.py +++ b/lightllm/server/multi_level_kv_cache/manager.py @@ -19,6 +19,7 @@ from lightllm.utils.process_check import start_parent_check_thread from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args +from lightllm.utils.config_utils import is_linear_att_mixed_model logger = init_logger(__name__) @@ -29,6 +30,7 @@ def __init__( args: StartArgs, ): self.args: StartArgs = args + self.is_linear_att_mixed_model = is_linear_att_mixed_model(args.model_dir) ports = get_shm_port_args() context = zmq.Context(2) self.zmq_recv_socket = context.socket(zmq.PULL) @@ -137,6 +139,28 @@ def _disk_cache_match(self, token_hash_list: List[int], all_pages: List[int]) -> self.cpu_cache_client.lock.release() return all_pages, len(new_page_indexes) + def _match_linear_att_tail(self, req: Req, pages: List[int]): + """在首个缺失页内回退到最长的历史碎页，命中后不再向后拼接页面。""" + if not self.is_linear_att_mixed_model: + return + page_lens = req.token_hash_page_len_list.get_all() + if len(pages) == len(page_lens): + return + page_start = len(pages) * self.args.cpu_cache_token_page_size + page_end = page_lens[len(pages)] + hash_size = self.args.linear_att_hash_page_size + hashes = req.linear_att_token_hash_list.get_all() + self.cpu_cache_client.lock.acquire_sleep1ms() + try: + for end in range(page_end - hash_size, page_start, -hash_size): + page_index, _ = self.cpu_cache_client.query_one_page(hashes[end // hash_size - 1]) + if page_index is not None: + pages.append(page_index) + req.cpu_cache_match_tail_len = end + return + finally: + self.cpu_cache_client.lock.release() + def _handle_group_req_multi_cache_match(self, group_req_indexes: GroupReqIndexes, start_time: float): """ match cpu cache and disk cache pages @@ -203,6 +227,9 @@ def _handle_group_req_multi_cache_match(self, group_req_indexes: GroupReqIndexes logger.exception(f"calculate disk prompt cache len has exception {str(e)}") raise e + # 优先保留完整 CPU/disk 页的命中机会，再查首个缺失页内的历史碎页。 + self._match_linear_att_tail(req, finded_page_indexes) + while not self.cpu_cache_client.check_allpages_ready(finded_page_indexes): time.sleep(0.01) diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 78f5fedc93..6155c6f62b 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -62,6 +62,7 @@ class PD_Client_Obj: dispatched_prompt_chars: int = 0 # 当前派发到该节点且尚未产出首 token 的请求数。 dispatched_req_num: int = 0 + checkpoint_registry_token: Optional[str] = field(default=None, repr=False) def __post_init__(self): if self.mode not in ["prefill", "decode"]: @@ -93,7 +94,6 @@ class PDUpKVStatus: pd_kv_trans_params: bytes # pd kv 传输建立连接所使用的元数据对象 def __post_init__(self): - if not isinstance(self.group_request_id, int): error_info = "group_request_id only can be int" logger.error(error_info) @@ -126,6 +126,8 @@ class PDDecodeNodeInfo: request_id: int ready_kv_len: int # decode 节点上已经准备好的kv长度 + # None preserves compatibility with older peers using the input_len - 1 convention. + first_token_owner: Optional[str] = None @dataclass @@ -182,6 +184,9 @@ class PDChunckedTransTask: page_kind: str = "kv" # Only valid for the local task owner; remote notify copies may carry the sender-local req_idx. req_idx: Optional[int] = None + first_token_owner: Optional[str] = None + # Logical KV producer IDs for exactly this transfer range, when enabled. + kv_origins: Optional[List[int]] = None def __post_init__(self): if self.start_kv_index < 0 or self.end_kv_index < self.start_kv_index: @@ -195,6 +200,8 @@ def __post_init__(self): assert len(self.mem_indexes) == 0 else: raise ValueError(f"unknown PD trans page kind {self.page_kind}") + if self.kv_origins is not None and len(self.kv_origins) != self.end_kv_index - self.start_kv_index: + raise ValueError("PD KV provenance must cover the transferred range") self.create_time = time.time() return @@ -221,6 +228,7 @@ def get_key(self) -> str: def to_str(self): obj: PDChunckedTransTask = copy.copy(self) obj.mem_indexes = None + obj.kv_origins = None if obj.decode_agent_metadata is not None: obj.decode_agent_metadata = b"xxx" if obj.prefill_agent_metadata is not None: @@ -248,6 +256,8 @@ def createRetObj(self) -> "PDChunckedTransTaskRet": error_info=self.error_info, first_gen_token_id=self.first_gen_token_id, first_gen_token_logprob=self.first_gen_token_logprob, + prefill_dp_index=self.prefill_dp_index, + kv_origins=self.kv_origins, ) return ret @@ -277,6 +287,8 @@ class PDChunckedTransTaskRet: error_info: str = None first_gen_token_id: Optional[int] = None first_gen_token_logprob: Optional[float] = None + prefill_dp_index: Optional[int] = None + kv_origins: Optional[List[int]] = None def get_key(self) -> str: return f"{self.request_id}_{self.start_kv_index}_{self.end_kv_index}" diff --git a/lightllm/server/router/dynamic_prompt/checkpoint_cache.py b/lightllm/server/router/dynamic_prompt/checkpoint_cache.py new file mode 100644 index 0000000000..6f41f1d596 --- /dev/null +++ b/lightllm/server/router/dynamic_prompt/checkpoint_cache.py @@ -0,0 +1,797 @@ +"""Bounded, rank-local CPU storage for exact hybrid-model checkpoints. + +The caller coordinates candidate selection and commit across its TP group. This +module never performs distributed collectives. In particular, a successful +``prepare`` is private until every participating rank can commit the checkpoint. + +Transfers use the caller's current CUDA stream. KV onload waits by default; +callers may defer that wait to a later fence covering KV and state together. +Sources must already describe a frozen, committed model position; reading a live +request's state after another forward has started is not safe. Returned leases +own the CPU source until onload/export finishes, including across ``clear``. +""" + +import hashlib +import threading +from collections import OrderedDict +from dataclasses import dataclass, field +from numbers import Integral +from typing import Dict, Optional, Sequence, Tuple + +import numpy as np +import torch + + +@dataclass(eq=False) +class _Page: + key: str + tokens: Tuple[int, ...] + origins: Tuple[int, ...] + tensor: torch.Tensor + references: int = 0 + + @property + def nbytes(self): + return self.tensor.numel() * self.tensor.element_size() + + +@dataclass(eq=False) +class _Entry: + serial: int + epoch: int + namespace: str + length: int + pages: Tuple[_Page, ...] + conv_state: torch.Tensor + ssm_state: torch.Tensor + output_seed: Optional[torch.Tensor] + leases: int = 0 + ready: bool = False + retired: bool = False + + @property + def state_bytes(self): + return sum( + tensor.numel() * tensor.element_size() + for tensor in (self.conv_state, self.ssm_state, self.output_seed) + if tensor is not None + ) + + def tokens(self): + return tuple(token for page in self.pages for token in page.tokens) + + def origins(self): + return tuple(origin for page in self.pages for origin in page.origins) + + +@dataclass +class _Node: + edge: Tuple[int, ...] = () + children: Dict[int, "_Node"] = field(default_factory=dict) + entry: Optional[_Entry] = None + parent: Optional["_Node"] = None + + +@dataclass +class PendingCheckpoint: + """Private, fully copied data. Only its originating cache may publish it.""" + + cache: "CpuCheckpointCache" + entry: _Entry + consumed: bool = False + _length: int = field(init=False) + + def __post_init__(self): + self._length = self.entry.length + + @property + def length(self): + return self._length + + +class CheckpointLease: + """Read-only source ownership; close after all consumers have completed.""" + + def __init__(self, cache, entry): + self.cache = cache + self.entry = entry + self._length = entry.length + self.closed = False + + @property + def length(self): + return self._length + + @property + def conv_state(self): + self._check_open() + return self.entry.conv_state + + @property + def ssm_state(self): + self._check_open() + return self.entry.ssm_state + + @property + def output_seed(self): + self._check_open() + return self.entry.output_seed + + @property + def page_keys(self): + self._check_open() + return [page.key for page in self.entry.pages] + + @property + def origins(self): + self._check_open() + return torch.tensor(self.entry.origins(), dtype=torch.int64, device="cpu") + + def _check_open(self): + if self.closed: + raise RuntimeError("checkpoint lease has been released") + + def close(self): + self.cache.release(self) + + def __enter__(self): + self._check_open() + return self + + def __exit__(self, *_): + self.close() + + +class CpuCheckpointCache: + """Compressed token-prefix directory with shared, immutable CPU KV pages. + + ``max_bytes`` includes physical page capacity, independent state/seed tensors, + unpublished preparations, and retired data still held by leases. Metadata is + bounded separately by ``max_entries``. A tail page uses the same allocation + size as a full page, but only its valid token range is copied or exported. + """ + + def __init__( + self, + max_bytes: int, + max_entries: int, + page_size: int = 8192, + pin_memory: bool = True, + draft_tail_dependency: bool = False, + preallocate: bool = False, + ): + if max_bytes < 0 or max_entries < 0 or page_size <= 0: + raise ValueError("invalid checkpoint cache capacity") + self.max_bytes = int(max_bytes) + self.max_entries = int(max_entries) + self.page_size = int(page_size) + self.pin_memory = pin_memory + self.draft_tail_dependency = bool(draft_tail_dependency) + self._lock = threading.RLock() + self._epoch = 0 + self._serial = 0 + self._bytes = 0 + self._roots: Dict[str, _Node] = {} + self._pages: Dict[str, _Page] = {} + self._entries: Dict[int, _Entry] = {} + self._nodes: Dict[int, _Node] = {} + self._lru: OrderedDict[int, _Entry] = OrderedDict() + self._arena = None + if preallocate and self.pin_memory: + from .checkpoint_memory import PinnedCheckpointArena + + self._arena = PinnedCheckpointArena(self.max_bytes) + + @staticmethod + def _tokens(tokens: Sequence[int]): + if isinstance(tokens, torch.Tensor): + if tokens.device.type != "cpu": + raise ValueError("checkpoint token IDs must be on CPU") + tokens = tokens.tolist() + return tuple(int(token) for token in tokens) + + @staticmethod + def _origins(origins, length): + if origins is None: + return (0,) * length + if isinstance(origins, torch.Tensor): + if origins.device.type != "cpu" or origins.dtype != torch.int64 or origins.ndim != 1: + raise ValueError("checkpoint origins must be a CPU int64 vector") + origins = origins.tolist() + origins = tuple(origins) + if len(origins) != length or any( + isinstance(origin, bool) or not isinstance(origin, Integral) or not 0 <= origin < 2 ** 63 + for origin in origins + ): + raise ValueError("checkpoint origins must contain one nonnegative int64 per token") + return tuple(int(origin) for origin in origins) + + def _page_specs(self, tokens, namespace, origins): + return self._derive_page_specs(tokens, namespace, self.page_size, self.draft_tail_dependency, origins) + + @staticmethod + def derive_page_keys( + tokens, + namespace="default", + page_size=8192, + draft_tail_dependency=False, + origins=None, + ): + """Rekey a validated transport payload after an explicit layout conversion.""" + if page_size <= 0: + raise ValueError("page_size must be positive") + tokens = CpuCheckpointCache._tokens(tokens) + origins = CpuCheckpointCache._origins(origins, len(tokens)) + return [ + key + for key, _, _ in CpuCheckpointCache._derive_page_specs( + tokens, namespace, page_size, draft_tail_dependency, origins + ) + ] + + @staticmethod + def _derive_page_specs(tokens, namespace, page_size, draft_tail_dependency, origins): + digest = hashlib.blake2b(digest_size=20, person=b"llm-checkpoint-2") + encoded_namespace = namespace.encode("utf-8") + digest.update(len(encoded_namespace).to_bytes(8, "little")) + digest.update(encoded_namespace) + specs = [] + for start in range(0, len(tokens), page_size): + part = tokens[start : start + page_size] + digest.update(np.asarray(part, dtype="= min(len(tokens), max_length): + break + child = node.children.get(tokens[offset]) + if child is None or offset + len(child.edge) > max_length: + break + if tuple(tokens[offset : offset + len(child.edge)]) != child.edge: + break + offset += len(child.edge) + node = child + return found + + def candidate_lengths(self, tokens, max_length=None, require_output_seed=False, namespace="default"): + """Return deepest-first candidates; callers intersect these across TP.""" + tokens = self._tokens(tokens) + limit = len(tokens) if max_length is None else min(len(tokens), int(max_length)) + with self._lock: + return [entry.length for entry in reversed(self._candidates(tokens, namespace, limit, require_output_seed))] + + def acquire(self, tokens, length: int, namespace="default"): + """Acquire exactly ``length``; never silently select a shorter prefix.""" + tokens = self._tokens(tokens) + with self._lock: + candidates = self._candidates(tokens, namespace, min(length, len(tokens)), False) + if not candidates or candidates[-1].length != length: + return None + entry = candidates[-1] + entry.leases += 1 + self._lru.move_to_end(entry.serial) + return CheckpointLease(self, entry) + + def match(self, tokens, max_length=None, require_output_seed=False, namespace="default"): + tokens = self._tokens(tokens) + with self._lock: + lengths = self.candidate_lengths(tokens, max_length, require_output_seed, namespace) + return self.acquire(tokens, lengths[0], namespace) if lengths else None + + def _insert_node(self, tokens, entry): + node = self._roots.setdefault(entry.namespace, _Node()) + offset = 0 + while offset < len(tokens): + child = node.children.get(tokens[offset]) + if child is None: + child = _Node(edge=tokens[offset:], parent=node) + node.children[child.edge[0]] = child + node = child + break + common = self._common(tokens[offset:], child.edge) + if common < len(child.edge): + middle = _Node(edge=child.edge[:common], parent=node) + node.children[middle.edge[0]] = middle + child.edge = child.edge[common:] + child.parent = middle + middle.children[child.edge[0]] = child + node = middle + else: + node = child + offset += common + old_entry = node.entry + if old_entry is not None: + self._retire(old_entry, prune=False) + node.entry = entry + self._nodes[entry.serial] = node + + def _remove_node(self, entry, prune=True): + node = self._nodes.pop(entry.serial, None) + if node is None: + return + node.entry = None + while prune and node.parent is not None and node.entry is None: + parent = node.parent + if not node.children: + del parent.children[node.edge[0]] + elif len(node.children) == 1: + child = next(iter(node.children.values())) + child.edge = node.edge + child.edge + child.parent = parent + parent.children[node.edge[0]] = child + else: + break + node = parent + + def _drop_entry(self, entry): + self._entries.pop(entry.serial) + self._bytes -= entry.state_bytes + for page in entry.pages: + self._unref_page(page) + + def _unref_page(self, page): + page.references -= 1 + if page.references == 0: + if self._pages.get(page.key) is page: + del self._pages[page.key] + self._bytes -= page.nbytes + + def _retire(self, entry, prune=True): + entry.ready = False + entry.retired = True + self._lru.pop(entry.serial, None) + self._remove_node(entry, prune) + if entry.leases == 0: + self._drop_entry(entry) + + def _make_room(self, needed_bytes): + if needed_bytes > self.max_bytes or self.max_entries == 0: + return False + # Retain IDs, not entries: retired tensors must be released immediately + # so their arena windows can satisfy the next preparation. + for serial in list(self._lru): + if self._bytes + needed_bytes <= self.max_bytes and len(self._entries) < self.max_entries: + break + entry = self._entries[serial] + if entry.leases == 0: + self._retire(entry) + del entry + return self._bytes + needed_bytes <= self.max_bytes and len(self._entries) < self.max_entries + + def _empty_cpu(self, shape, dtype): + if self._arena is not None: + with self._lock: + tensor = self._arena.allocate(shape, dtype) + if tensor is not None: + return tensor + # Physical fragmentation or outstanding exported views can + # require more space than logical byte admission alone predicts. + for serial in list(self._lru): + entry = self._entries[serial] + if entry.leases == 0: + self._retire(entry) + del entry + tensor = self._arena.allocate(shape, dtype) + if tensor is not None: + return tensor + raise MemoryError("checkpoint pinned arena has no available contiguous range") + return torch.empty(shape, dtype=dtype, device="cpu", pin_memory=self.pin_memory) + + @staticmethod + def _finish_stream(tensors): + devices = {tensor.device for tensor in tensors if tensor is not None and tensor.is_cuda} + for device in devices: + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(device)) + event.synchronize() + + def _prepare( + self, + tokens, + origins, + namespace, + kv_shape, + kv_dtype, + conv_state, + ssm_state, + output_seed, + copy_page, + copy_producers=(), + ): + if not tokens: + return None + if conv_state is None or ssm_state is None: + raise ValueError("both conv and SSM state are required") + specs = self._page_specs(tokens, namespace, origins) + # Hold matching sources while admission evicts other entries. A COW + # tail copies its existing CPU prefix; only the new suffix comes from GPU. + candidates = self._candidates(tokens, namespace, len(tokens), False) + base_entry = candidates[-1] if candidates else None + base_pages = base_entry.pages if base_entry is not None else () + common_origins = self._common(origins, base_entry.origins()) if base_entry is not None else 0 + del base_entry + del candidates + page_bytes = int(np.prod(kv_shape)) * torch.empty((), dtype=kv_dtype).element_size() + state_bytes = sum( + tensor.numel() * tensor.element_size() + for tensor in (conv_state, ssm_state, output_seed) + if tensor is not None + ) + existing = [] + for key, part, start in specs: + page = self._pages.get(key) + # Verify token content as well as its digest, and reject layout reuse. + if page is not None and ( + page.tokens != part + or page.origins != origins[start : start + len(part)] + or tuple(page.tensor.shape) != tuple(kv_shape) + or page.tensor.dtype != kv_dtype + ): + raise ValueError("checkpoint namespace reused with incompatible content or KV layout") + existing.append(page) + needed = state_bytes + sum(page is None for page in existing) * page_bytes + # Exact shared pages can otherwise disappear when their owning checkpoint + # is evicted by admission. Temporary refs make the plan stable. + for page in existing: + if page is not None: + page.references += 1 + for page in base_pages: + page.references += 1 + created = [] + state_tensors = [] + try: + if not self._make_room(needed): + return None + for index, ((key, part, start), page) in enumerate(zip(specs, existing)): + if page is None: + tensor = self._empty_cpu(kv_shape, kv_dtype) + copied = 0 + if index < len(base_pages): + old = base_pages[index] + copied = min(len(old.tokens), len(part), max(0, common_origins - start)) + if old.tokens[:copied] != part[:copied]: + copied = 0 + if self.draft_tail_dependency and copied == len(old.tokens): + # A previous checkpoint's terminal draft slot was + # not reusable. Refresh it as it becomes interior. + copied = max(0, copied - 1) + if copied: + tensor[:, :copied].copy_(old.tensor[:, :copied]) + copy_page(tensor, start, len(part), copied, key) + page = _Page(key, part, origins[start : start + len(part)], tensor) + created.append(page) + for source in (conv_state, ssm_state, output_seed): + destination = None + if source is not None: + destination = self._empty_cpu(source.shape, source.dtype) + destination.copy_(source, non_blocking=source.is_cuda) + state_tensors.append(destination) + # The producer owns all sources through this fence. No directory or + # page becomes visible while a device copy is still outstanding. + self._finish_stream((*copy_producers, conv_state, ssm_state, output_seed)) + pages = iter(created) + entry_pages = tuple(page if page is not None else next(pages) for page in existing) + self._serial += 1 + entry = _Entry(self._serial, self._epoch, namespace, len(tokens), entry_pages, *state_tensors) + for page in entry_pages: + page.references += 1 + for page in created: + self._pages[page.key] = page + self._bytes += page.nbytes + self._bytes += entry.state_bytes + self._entries[entry.serial] = entry + return PendingCheckpoint(self, entry) + except BaseException: + # A later state allocation may fail after an earlier asynchronous + # D2H copy. Fence the producer before releasing temporary buffers + # or returning its staging slot to the next capture. + try: + self._finish_stream((*copy_producers, conv_state, ssm_state, output_seed)) + finally: + created.clear() + state_tensors.clear() + tensor = destination = page = None + raise + finally: + for page in existing: + if page is not None: + self._unref_page(page) + for page in base_pages: + self._unref_page(page) + existing.clear() + base_pages = () + old = page = None + + def prepare( + self, + tokens, + kv_buffer, + mem_indexes, + conv_state, + ssm_state, + output_seed=None, + namespace="default", + frozen_tail_kv=None, + origins=None, + ): + """Copy a frozen position; call commit only after every TP rank succeeds. + + ``kv_buffer`` has shape [layers, token_slots, 2 * heads, head_dim]. + mem_indexes contains the complete logical prefix, in token order. + origins identifies the actual forward that produced each token's KV; + equal token IDs alone do not make numerically different histories safe + to combine with the captured state. Production callers supply origins. + """ + tokens = self._tokens(tokens) + origins = self._origins(origins, len(tokens)) + if kv_buffer.ndim != 4 or mem_indexes.ndim != 1 or len(mem_indexes) != len(tokens): + raise ValueError("invalid KV shape or checkpoint token indexes") + if frozen_tail_kv is not None and tuple(frozen_tail_kv.shape) != ( + kv_buffer.shape[0], + *kv_buffer.shape[2:], + ): + raise ValueError("frozen tail must contain one complete KV token") + kv_shape = (kv_buffer.shape[0], self.page_size, *kv_buffer.shape[2:]) + copy_sources = [] + + def copy_page(destination, start, valid, copied, _key): + try: + if copied < valid: + if copy_sources: + # Keep GPU gather storage bounded to one page. The last + # page shares its completion fence with state and seed. + self._finish_stream((kv_buffer,)) + copy_sources.clear() + indexes = mem_indexes[start + copied : start + valid].to(device=kv_buffer.device, dtype=torch.long) + copy_sources.append(indexes) + source = kv_buffer.index_select(1, indexes) + copy_sources.append(source) + if frozen_tail_kv is not None and start + valid == len(tokens): + source[:, -1].copy_(frozen_tail_kv) + destination[:, copied:valid].copy_(source, non_blocking=source.is_cuda) + finally: + source = indexes = destination = None + + with self._lock: + try: + return self._prepare( + tokens, + origins, + namespace, + kv_shape, + kv_buffer.dtype, + conv_state, + ssm_state, + output_seed, + copy_page, + copy_producers=(kv_buffer,), + ) + finally: + # _prepare fences this producer on success and failure, even + # when recurrent state was already on CPU. Clear retained + # temporaries so an exception traceback cannot keep them alive. + copy_sources.clear() + + def commit(self, pending): + """Publish one prepared rank shard after coordinated TP admission.""" + with self._lock: + if pending.cache is not self or pending.consumed: + raise ValueError("invalid or already consumed checkpoint preparation") + pending.consumed = True + entry = pending.entry + pending.entry = None + if entry.epoch != self._epoch: + self._retire(entry) + return False + self._insert_node(entry.tokens(), entry) + entry.ready = True + self._lru[entry.serial] = entry + return True + + def discard(self, pending): + with self._lock: + if pending.cache is not self or pending.consumed: + raise ValueError("invalid or already consumed checkpoint preparation") + pending.consumed = True + entry = pending.entry + pending.entry = None + self._retire(entry) + + def release(self, lease): + with self._lock: + if lease.cache is not self: + raise ValueError("lease belongs to another cache") + if lease.closed: + return + lease.closed = True + entry = lease.entry + lease.entry = None + entry.leases -= 1 + if entry.retired and entry.leases == 0: + self._drop_entry(entry) + + def load_kv(self, lease, kv_buffer, mem_indexes, start=0, *, wait=True): + """Load [start, lease.length) into the corresponding destination indexes. + + ``mem_indexes`` contains only that missing range, not the reused prefix. + With wait=False the caller must retain the lease and keep mem_indexes + unchanged until the current stream finishes all readers. Copy errors + fence before returning. + Destination state is restored separately by the model-specific adapter. + """ + with self._lock: + lease._check_open() + if lease.cache is not self or not 0 <= start <= lease.length: + raise ValueError("invalid checkpoint lease or load range") + if len(mem_indexes) != lease.length - start: + raise ValueError("destination indexes must cover precisely the missing prefix range") + try: + for index, page in enumerate(lease.entry.pages): + page_start = index * self.page_size + begin = max(start, page_start) + end = page_start + len(page.tokens) + if begin >= end: + continue + source = page.tensor[:, begin - page_start : end - page_start].to( + device=kv_buffer.device, non_blocking=kv_buffer.is_cuda + ) + # Upload pinned allocator indexes before converting their + # dtype; a CPU cast would discard the pinned allocation. + indexes = ( + mem_indexes[begin - start : end - start] + .to(device=kv_buffer.device, non_blocking=kv_buffer.is_cuda) + .long() + ) + kv_buffer.index_copy_(1, indexes, source) + # Same-stream allocator reuse bounds scratch space to one + # page without keeping a list of pending GPU page tensors. + source = indexes = None + except BaseException: + self._finish_stream((kv_buffer,)) + raise + else: + if wait: + self._finish_stream((kv_buffer,)) + finally: + source = indexes = page = None + + def export(self, lease, known_page_keys=()): + """Return torch.save-compatible CPU payload; retain lease until serialized.""" + with self._lock: + lease._check_open() + if lease.cache is not self: + raise ValueError("lease belongs to another cache") + known = set(known_page_keys) + entry = lease.entry + return { + "version": 1, + "namespace": entry.namespace, + "page_size": self.page_size, + "draft_tail_dependency": self.draft_tail_dependency, + "length": entry.length, + "tokens": list(entry.tokens()), + "origins": list(entry.origins()), + "page_keys": lease.page_keys, + # torch.save serializes an entire backing storage for views. + # Own the valid tail bytes so neither page padding nor unrelated + # old allocation contents enter the wire payload. + "pages": { + page.key: ( + page.tensor + if len(page.tokens) == self.page_size + else page.tensor[:, : len(page.tokens)].clone(memory_format=torch.contiguous_format) + ) + for page in entry.pages + if page.key not in known + }, + "conv_state": entry.conv_state, + "ssm_state": entry.ssm_state, + "output_seed": entry.output_seed, + } + + def missing_page_keys(self, page_keys): + with self._lock: + return [key for key in page_keys if key not in self._pages] + + def prepare_import(self, payload): + """Prepare a CPU payload, returning None if omitted base pages disappeared.""" + if payload["version"] != 1 or payload["page_size"] != self.page_size: + raise ValueError("incompatible checkpoint format") + if bool(payload.get("draft_tail_dependency", False)) != self.draft_tail_dependency: + raise ValueError("incompatible checkpoint draft tail dependency") + tokens = self._tokens(payload["tokens"]) + if "origins" not in payload: + raise ValueError("checkpoint manifest lacks KV computation origins") + origins = self._origins(payload["origins"], len(tokens)) + namespace = payload["namespace"] + specs = self._page_specs(tokens, namespace, origins) + if len(tokens) != payload["length"] or [key for key, _, _ in specs] != payload["page_keys"]: + raise ValueError("checkpoint manifest does not match its token prefix and origins") + provided = payload["pages"] + with self._lock: + if any(key not in provided and key not in self._pages for key, _, _ in specs): + return None + if not specs: + return None + first_key = specs[0][0] + sample = provided[first_key] if first_key in provided else self._pages[first_key].tensor + if sample.ndim != 4 or sample.device.type != "cpu": + raise ValueError("checkpoint imports require CPU KV tensors") + kv_shape = (sample.shape[0], self.page_size, *sample.shape[2:]) + for key, part, _ in specs: + tensor = provided.get(key) + if tensor is not None and ( + tensor.device.type != "cpu" + or tuple(tensor.shape) != (sample.shape[0], len(part), *sample.shape[2:]) + or tensor.dtype != sample.dtype + ): + raise ValueError("invalid imported KV page layout") + for name in ("conv_state", "ssm_state", "output_seed"): + tensor = payload[name] + if tensor is not None and tensor.device.type != "cpu": + raise ValueError("checkpoint imports require CPU state tensors") + + def copy_page(destination, _start, valid, copied, key): + if copied < valid: + destination[:, copied:valid].copy_(provided[key][:, copied:valid]) + + return self._prepare( + tokens, + origins, + namespace, + kv_shape, + sample.dtype, + payload["conv_state"], + payload["ssm_state"], + payload["output_seed"], + copy_page, + ) + + def clear(self): + """Invalidate lookup immediately; leased and pending buffers remain owned.""" + with self._lock: + self._epoch += 1 + for entry in list(self._lru.values()): + self._retire(entry) + self._roots.clear() + self._nodes.clear() + self._pages.clear() + + def stats(self): + with self._lock: + return { + "bytes": self._bytes, + "max_bytes": self.max_bytes, + "entries": len(self._entries), + "ready_entries": len(self._lru), + "pages": len(self._pages), + "epoch": self._epoch, + } diff --git a/lightllm/server/router/dynamic_prompt/checkpoint_memory.py b/lightllm/server/router/dynamic_prompt/checkpoint_memory.py new file mode 100644 index 0000000000..c2fc55e28d --- /dev/null +++ b/lightllm/server/router/dynamic_prompt/checkpoint_memory.py @@ -0,0 +1,87 @@ +"""Pinned checkpoint storage reserved before requests enter the inference loop.""" + +import bisect +import math +import queue +import threading +import weakref + +import torch + + +class PinnedCheckpointArena: + """Split one bounded allocation into independently owned tensor storages. + + A tensor view of the arena would make torch.save serialize the entire cache. + frombuffer instead gives each window a storage of precisely its own size. + The storage retains the memoryview; its finalizer returns the range only + after the last tensor/view has released it. This object never owns allocated + windows, entries, or the cache, so that finalizer cannot create an owner cycle. + + Callers must still fence asynchronous CUDA users before releasing storage: + these window storages are not owned by PyTorch's pinned caching allocator. + """ + + def __init__(self, capacity): + self.capacity = int(capacity) + self._buffer = memoryview(torch.empty(self.capacity, dtype=torch.uint8, pin_memory=True).numpy()) + self._free = [(0, self.capacity)] if self.capacity else [] + self._lock = threading.Lock() + self._released = queue.SimpleQueue() + + def allocate(self, shape, dtype): + shape = tuple(int(dim) for dim in shape) + element_size = torch.empty((), dtype=dtype).element_size() + size = math.prod(shape) * element_size + if size == 0: + return torch.empty(shape, dtype=dtype, device="cpu") + if size < 0 or any(dim < 0 for dim in shape): + raise ValueError("invalid checkpoint allocation shape") + with self._lock: + self._reclaim() + best = None + for index, (offset, available) in enumerate(self._free): + start = (offset + element_size - 1) // element_size * element_size + if start + size <= offset + available and (best is None or available < best[0]): + best = (available, index, offset, start) + if best is None: + return None + available, index, offset, start = best + remainder = [] + if start > offset: + remainder.append((offset, start - offset)) + end = start + size + if end < offset + available: + remainder.append((end, offset + available - end)) + self._free[index : index + 1] = remainder + try: + window = self._buffer[start:end] + # GC can run during free-list mutation. Its callback must neither + # acquire our lock nor mutate the list being searched. SimpleQueue + # supports reentrant put from finalizers; allocation drains it. + weakref.finalize(window, self._released.put, (start, size)) + except BaseException: + self._released.put((start, size)) + raise + try: + return torch.frombuffer(window, dtype=dtype).reshape(shape) + finally: + # An allocator error's traceback can outlive admission handling. + # Successful storage owns the window; failed frames must not own it. + window = None + + def _reclaim(self): + while True: + try: + start, size = self._released.get_nowait() + except queue.Empty: + return + index = bisect.bisect_left(self._free, (start,)) + if index and self._free[index - 1][0] + self._free[index - 1][1] == start: + previous, previous_size = self._free.pop(index - 1) + index -= 1 + start, size = previous, previous_size + size + if index < len(self._free) and start + size == self._free[index][0]: + _, next_size = self._free.pop(index) + size += next_size + self._free.insert(index, (start, size)) diff --git a/lightllm/server/router/model_infer/exact_prefix_cache.py b/lightllm/server/router/model_infer/exact_prefix_cache.py new file mode 100644 index 0000000000..531c22958b --- /dev/null +++ b/lightllm/server/router/model_infer/exact_prefix_cache.py @@ -0,0 +1,840 @@ +"""Exact hybrid checkpoints at the inference/cache ownership boundary. + +Capture runs on the producing stream before the next batch can overwrite a +request's recurrent state. Normal-mode publication uses a bounded worker; +the scheduler retains source requests until every TP rank finishes copying. +PD publication and CPU onload retain their explicit transfer fences. +""" + +import hashlib +import queue +import threading +from collections import Counter +from dataclasses import dataclass, field +from typing import Optional + +import torch +import torch.distributed as dist + +from lightllm.common.basemodel.batch_objs import ModelInput +from lightllm.common.req_manager.linear_att import LinearStateSnapshot +from lightllm.server.router.dynamic_prompt.checkpoint_cache import CpuCheckpointCache +from lightllm.utils.dist_utils import create_new_group_for_current_dp +from lightllm.utils.log_utils import init_logger +from lightllm.utils.checkpoint_identity import get_checkpoint_identity +from lightllm.utils.envs_utils import get_unique_server_name + +logger = init_logger(__name__) + + +@dataclass +class CaptureBatch: + reqs: list + output_lengths: list + is_prefill: bool + epoch: int + cache_generation: int + kv_origins: dict = field(default_factory=dict) + staging: object = None + output_seed: Optional[torch.Tensor] = None + tail_kv: Optional[torch.Tensor] = None + metadata: Optional[torch.Tensor] = None + + +@dataclass(frozen=True) +class FrozenCheckpoint: + req_id: int + length: int + output_len: int + tokens: object + origins: tuple + mem_indexes: torch.Tensor + conv_state: torch.Tensor + ssm_state: torch.Tensor + output_seed: Optional[torch.Tensor] + tail_kv: torch.Tensor + + +class ExactPrefixCache: + def __init__(self, backend): + self.backend = backend + self.args = backend.args + self.req_manager = backend.model.req_manager + self.mem_manager = backend.model.mem_manager + self.cache = CpuCheckpointCache( + max_bytes=self.args.exact_prefix_cache_mb * 1024 ** 2, + max_entries=self.args.exact_prefix_cache_entries, + page_size=self.args.exact_prefix_cache_page_size, + draft_tail_dependency=bool(self.args.mtp_step), + preallocate=True, + ) + self.group = create_new_group_for_current_dp("gloo") + self.world_size = dist.get_world_size(self.group) + self._epoch = 0 + self._aux_epoch = 0 + self._producer = f"{get_unique_server_name()}:{backend.global_dp_rank}" + self._cache_generation = 0 + self._staging = [ + self.req_manager.allocate_linear_state_staging(self.args.exact_prefix_cache_capture_slots) + for _ in range( + 4 if (self.args.enable_decode_microbatch_overlap or self.args.enable_prefill_microbatch_overlap) else 2 + ) + ] + self._busy = [False] * len(self._staging) + self._warmup_kernels() + self.transport = None + self.target_fingerprint, self.draft_fingerprint, self.namespace = get_checkpoint_identity(self.args) + if self.args.run_mode in ("prefill", "decode"): + from lightllm.server.router.model_infer.mode_backend.pd.checkpoint_transport import PDCheckpointTransport + + self.transport = PDCheckpointTransport( + backend, + self.cache, + namespace=self.namespace, + target_namespace=f"{self.target_fingerprint}:target-only", + ) + self.stats = dict(captured=0, skipped=0, hits=0, hit_tokens=0, head_only=0, loaded_tokens=0) + self._async_publication = self.args.run_mode == "normal" + self._pending_publications = {} + self._publication_holds = Counter() + self._completed_publications = {} + self._publication_error = None + if self._async_publication: + # Worker collectives must never share ordering with scheduler + # admission, restore, or the two CPU inference threads. + self._publication_group = create_new_group_for_current_dp("gloo") + self._publication_stream = torch.cuda.Stream(device=self.mem_manager.kv_buffer.device) + self._publication_queue = queue.SimpleQueue() + self._publication_done = queue.SimpleQueue() + self._publication_ack = threading.Event() + self._publication_thread = threading.Thread( + target=self._publication_loop, name="exact-checkpoint-copy", daemon=True + ) + self._publication_thread.start() + logger.info( + "exact prefix cache: CPU budget=%s MiB/rank, entries=%s, KV page=%s tokens, " + "capture slots=%s x %s; recurrent states are independent of KV pages", + self.args.exact_prefix_cache_mb, + self.args.exact_prefix_cache_entries, + self.args.exact_prefix_cache_page_size, + self.args.exact_prefix_cache_capture_slots, + len(self._staging), + ) + + def _warmup_kernels(self): + """Compile bounded row-block variants before accepting requests.""" + max_rows = self.req_manager.max_request_num * (self.args.mtp_step + 1) + row_capacity = 1 << (max_rows - 1).bit_length() + device = self.mem_manager.kv_buffer.device + zeros = torch.zeros(row_capacity, dtype=torch.int32, device=device) + lengths = torch.ones_like(zeros) + mask = torch.zeros(row_capacity, dtype=torch.bool, device=device) + if self.args.mtp_step: + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc + + count = 1 + while count <= row_capacity: + # An empty mask only initializes staging metadata. It cannot read + # or overwrite a live recurrent state or request's KV mapping. + self.req_manager.freeze_linear_states( + zeros[:count], zeros[:count], lengths[:count], mask[:count], self._staging[0] + ) + if self.args.mtp_step: + gen_b_req_mtp_start_loc(zeros[:count], num_reqs=count) + count *= 2 + torch.cuda.current_stream(device).synchronize() + + def _all(self, value): + if self.world_size == 1: + return bool(value) + flag = torch.tensor(int(bool(value)), dtype=torch.int32, device="cpu") + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=self.group) + return bool(flag.item()) + + def _intersection(self, values): + if self.world_size == 1: + return sorted(values, reverse=True) + lists = [None] * self.world_size + dist.all_gather_object(lists, list(values), group=self.group) + return sorted(set.intersection(*(set(v) for v in lists)), reverse=True) + + @staticmethod + def _is_allocation_failure(error): + if isinstance(error, (MemoryError, torch.OutOfMemoryError)): + return True + if not isinstance(error, RuntimeError): + return False + message = str(error).lower() + allocator = any( + name in message + for name in ("defaultcpuallocator", "cachinghostallocator", "cudahostalloc", "pinned memory") + ) + exhausted = any( + text in message + for text in ("can't allocate memory", "out of memory", "not enough memory", "allocation failed") + ) + return allocator and exhausted + + def _spec_engine(self): + engine = self.backend.spec_engine + return getattr(engine, "common_engine", engine) + + def gpu_radix_key(self, origins, length): + # CPU directory lookup has already checked the actual token prefix. + # GPU sharing additionally requires the *same computation history*: + # equal tokens computed by decode/prefill can have different KV bytes. + if len(origins) < length or any(origin <= 0 for origin in origins[:length]): + raise ValueError("GPU checkpoint references require complete KV provenance") + return torch.tensor(origins[:length], dtype=torch.int64, device="cpu") + + def _origin(self, req, epoch, kind="forward"): + value = f"{self._producer}:{kind}:{req.req_id}:{epoch}".encode() + return (int.from_bytes(hashlib.sha256(value).digest()[:8], "little") & ((1 << 63) - 1)) or 1 + + def _eligible(self, req): + if req.sampling_param.disable_prompt_cache or req.infer_aborted: + return False + if req.sampling_param.shm_param.prompt_logprobs >= 0 or self.args.enable_return_routed_experts: + # Endpoint hidden/state cannot reproduce per-token prompt outputs. + return False + # Image/audio embeddings and positional deltas require their own + # execution identity; token IDs alone do not establish an equivalent run. + if any(req.multimodal_params.get(key) for key in ("images", "audios", "videos")): + return False + if req.multimodal_params.get("mrope_position_delta", 0): + return False + engine = self._spec_engine() + return not self.args.mtp_step or (engine is not None and engine.supports_exact_prefix_resume()) + + def prepare_batch(self, model_input, run_reqs): + """Bind immutable logical lengths before ModelInput moves to CUDA.""" + self._epoch += 1 + lengths = model_input.b_seq_len.tolist() + ends = {} + for req, length in zip(run_reqs, lengths): + ends[req.req_idx] = (req, max(length, ends.get(req.req_idx, (None, 0))[1])) + origins = {} + for req_index, (req, length) in ends.items(): + prefix = req.exact_kv_origins + if len(prefix) < req.cur_kv_len: + raise ValueError("computed KV prefix is missing its provenance") + prefix[req.cur_kv_len :] = [self._origin(req, self._epoch)] * (length - req.cur_kv_len) + # Later forwards may only change the unaccepted/new suffix. Every + # candidate that CPU post can publish lies within this batch's + # accepted prefix, whose origins remain immutable. Keep its list + # reference instead of copying a million-token prefix each decode. + origins[req_index] = prefix + return CaptureBatch( + reqs=list(run_reqs), + output_lengths=[length - req.shm_req.input_len + 1 for req, length in zip(run_reqs, lengths)], + is_prefill=model_input.is_prefill, + epoch=self._epoch, + cache_generation=self._cache_generation, + kv_origins=origins, + ) + + def capture(self, ticket, model_input, model_output, next_token_ids, accepted_index=None): + """Freeze only selected rows, on the same stream as sample/verify.""" + if not ticket.reqs: + return ticket + slot = (ticket.epoch - 1) % len(self._staging) + if self._busy[slot]: + self.stats["skipped"] += 1 + return ticket + candidate = [self._eligible(req) for req in ticket.reqs] + if not ticket.is_prefill: + # Known-length generation with EOS disabled needs no device capture + # work between its prompt endpoint and its final output. Ineligible + # requests likewise must not pay for empty staging/readback each step. + candidate = [ + eligible + and ( + out == 1 + or out == req.sampling_param.shm_param.max_new_tokens + or not req.sampling_param.shm_param.ignore_eos + or bool(req.stop_sequences) + ) + for eligible, req, out in zip(candidate, ticket.reqs, ticket.output_lengths) + ] + if not any(candidate): + return ticket + from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager + + flags = candidate + eos_allowed = [] + if not ticket.is_prefill: + # A one-token remainder of chunked prefill is scheduled as decode. + # It still produces the exact prompt endpoint and the first sample. + flags = [ + eligible and (out == 1 or out == req.sampling_param.shm_param.max_new_tokens) + for eligible, req, out in zip(candidate, ticket.reqs, ticket.output_lengths) + ] + eos_allowed = [ + eligible and not req.sampling_param.shm_param.ignore_eos + for eligible, req in zip(candidate, ticket.reqs) + ] + if any(eos_allowed): + flags += eos_allowed + # Reuse pinned metadata for this staging slot. Its publication fence + # completes the H2D read before the slot can be reused by either infer + # thread. One asynchronous copy preserves the producing stream's order + # without waiting for the preceding forward/propose work on the CPU. + flags_gpu = g_pin_mem_manager.gen_from_list( + key=f"exact_capture_flags_{slot}", data=flags, dtype=torch.bool + ).cuda(non_blocking=True) + mask = flags_gpu[: len(candidate)] + if not ticket.is_prefill: + token_ids = next_token_ids.reshape(-1) + if any(eos_allowed): + eos_mask = flags_gpu[len(candidate) :] + for token_id in self.backend.eos_id: + mask |= (token_ids == token_id) & eos_mask + # A last-token match is an inexpensive candidate hint. CPU post + # checks the entire stop sequence and may discard a false positive. + for row, (eligible, req) in enumerate(zip(candidate, ticket.reqs)): + if eligible: + for sequence in req.stop_sequences: + if sequence: + mask[row] |= token_ids[row] == sequence[-1] + if accepted_index is not None: + mask &= accepted_index.to(dtype=torch.bool) + staging = self._staging[slot] + self._busy[slot] = True + ticket.staging = staging + self.req_manager.freeze_linear_states( + model_input.b_req_idx, + model_input.b_mtp_index, + model_input.b_seq_len, + mask, + staging, + ) + ticket.output_seed = model_output.output_seed + # Speculative draft fill may rewrite its tail KV before CPU publication. + # Freeze that packed token along with the recurrent state. + source_indexes = self.req_manager.req_to_token_indexs[ + staging.req_indices.clamp_min(0).long(), (staging.exact_lengths - 1).clamp_min(0).long() + ] + source_indexes = torch.where(staging.req_indices >= 0, source_indexes, self.mem_manager.HOLD_TOKEN_MEMINDEX) + ticket.tail_kv = self.mem_manager.kv_buffer.index_select(1, source_indexes.long()) + # The caller records its compute event after capture. Include metadata + # readback in that event so CPU post only reads completed pinned data. + ticket.metadata = g_pin_mem_manager.async_copy_from_gpu_tensor( + key=f"exact_capture_metadata_{slot}", + gpu_tensor=torch.stack((staging.req_indices, staging.exact_lengths, staging.source_rows)), + ) + return ticket + + def finalize(self, ticket): + if self._async_publication: + self._enqueue_publication(ticket) + return + self._finalize_sync(ticket) + + def _enqueue_publication(self, ticket): + if ticket is None: + return + # Capture admission may differ locally; queue epochs and resource + # holds must nevertheless remain identical within each TP group. + if not self._all(ticket.staging is not None): + if ticket.staging is not None: + self._busy[(ticket.epoch - 1) % len(self._staging)] = False + return + staging = ticket.staging + metadata = ticket.metadata.tolist() + descriptors = [] + for slot, (req_index, length, row) in enumerate(zip(*metadata)): + descriptor = None + if req_index >= 0: + req = ticket.reqs[row] + output_len = ticket.output_lengths[row] + visible = getattr(req, "exact_visible_end", -1) + if ( + req.req_idx == req_index + and self._eligible(req) + and ( + ticket.is_prefill or output_len == 1 or (req.finish_status.is_finished() and length <= visible) + ) + ): + # CPU post has resolved accepted tokens and stop sequences. + # Freeze all metadata here: a later batch can mutate req. + descriptor = FrozenCheckpoint( + req_id=req.req_id, + length=length, + output_len=output_len, + tokens=req.shm_req.shm_prompt_ids.arr[:length].copy(), + origins=tuple(ticket.kv_origins[req_index][:length]), + mem_indexes=self.req_manager.req_to_token_indexs[req_index, :length], + conv_state=staging.conv_state[slot], + ssm_state=staging.ssm_state[slot], + output_seed=None if ticket.output_seed is None else ticket.output_seed[row], + tail_kv=ticket.tail_kv[:, slot], + ) + descriptors.append(descriptor) + # Hold every request in the ticket, including unselected slots. The + # fixed ticket membership makes scheduler holds independent of local + # eligibility and worker timing. Staging slots bound queue occupancy. + self._pending_publications[ticket.epoch] = ticket + self._publication_holds.update({req.req_id for req in ticket.reqs}) + self._publication_queue.put((ticket.epoch, ticket.cache_generation, descriptors)) + + def _publication_status(self, value): + if self.world_size == 1: + return value + status = torch.tensor(value, dtype=torch.int32, device="cpu") + dist.all_reduce(status, op=dist.ReduceOp.MIN, group=self._publication_group) + return int(status.item()) + + def _publish_frozen(self, generation, descriptors): + published = [] + skipped = 0 + # Fixed slot count ensures every TP rank executes the same collectives + # even if a local abort or allocation failure rejects its descriptor. + pending = None + try: + for descriptor in descriptors: + if not self._publication_status(int(descriptor is not None and generation == self._cache_generation)): + continue + pending = None + error = None + existing = None + reusable = False + try: + try: + existing = self.cache.acquire(descriptor.tokens, descriptor.length, namespace=self.namespace) + reusable = existing is not None and ( + descriptor.output_seed is None or existing.output_seed is not None + ) + except BaseException as exc: + if self._is_allocation_failure(exc): + logger.warning("checkpoint lookup skipped req=%s: %s", descriptor.req_id, exc) + else: + error = exc + status = self._publication_status(-1 if error is not None else int(reusable)) + if status < 0: + raise RuntimeError("checkpoint publication failed on a TP rank") from error + if status == 1: + # Retain the existing KV/state/seed history as one unit. + # A lease keeps it alive across the TP presence check; + # clear() may still invalidate its directory generation. + if self._publication_status(int(generation == self._cache_generation)): + skipped += 1 + continue + finally: + if existing is not None: + existing.close() + try: + pending = self.cache.prepare( + descriptor.tokens, + self.mem_manager.kv_buffer, + descriptor.mem_indexes, + descriptor.conv_state, + descriptor.ssm_state, + output_seed=descriptor.output_seed, + namespace=self.namespace, + frozen_tail_kv=descriptor.tail_kv, + origins=descriptor.origins, + ) + except BaseException as exc: + if self._is_allocation_failure(exc): + logger.warning("checkpoint allocation skipped req=%s: %s", descriptor.req_id, exc) + else: + error = exc + status = self._publication_status( + -1 if error is not None else int(pending is not None and generation == self._cache_generation) + ) + if status == 1: + # Only the scheduler may change the visible directory. A + # worker commit could let concurrent TP restores acquire + # different histories for the same tokens and length. + published.append((descriptor.req_id, descriptor.length, descriptor.output_len, pending)) + else: + if pending is not None: + self.cache.discard(pending) + if status < 0: + raise RuntimeError("checkpoint publication failed on a TP rank") from error + skipped += 1 + except BaseException: + # A fatal rank error must not strand earlier prepared CPU pages. + for _, _, _, prepared in published: + if not prepared.consumed: + self.cache.discard(prepared) + if pending is not None and not pending.consumed: + self.cache.discard(pending) + raise + return published, skipped + + def _publication_loop(self): + try: + with torch.cuda.device(self.mem_manager.kv_buffer.device), torch.cuda.stream(self._publication_stream): + while True: + epoch, generation, descriptors = self._publication_queue.get() + published, skipped = self._publish_frozen(generation, descriptors) + # prepare fences every GPU reader before reporting done. + # Do not retain the previous job while blocking on get(). + descriptors = None + self._publication_done.put((epoch, published, skipped)) + # Let the scheduler commit/discard before the next prepare + # can take the cache lock across another device transfer. + self._publication_ack.wait() + self._publication_ack.clear() + published = None + except BaseException as error: + self._publication_error = error + logger.exception("exact checkpoint publication worker failed") + + def has_pending(self, req): + return bool(self._publication_holds.get(req.req_id, 0)) + + def poll_publications(self): + """Retire globally completed jobs without waiting for unfinished copies.""" + if not self._pending_publications: + return + if not self._all(self._publication_error is None): + raise RuntimeError("exact checkpoint publication worker failed") from self._publication_error + while True: + try: + epoch, published, skipped = self._publication_done.get_nowait() + self._completed_publications[epoch] = (published, skipped) + except queue.Empty: + break + for epoch in reversed(self._intersection(self._completed_publications)): + ticket = self._pending_publications.pop(epoch) + published, skipped = self._completed_publications.pop(epoch) + reqs = {req.req_id: req for req in ticket.reqs} + current_generation = ticket.cache_generation == self._cache_generation + for req_id, length, output_len, pending in published: + if not current_generation: + self.cache.discard(pending) + elif self.cache.commit(pending): + self.stats["captured"] += 1 + reqs[req_id].exact_checkpoint_length = length + logger.debug("checkpoint published req=%s length=%s output_len=%s", req_id, length, output_len) + for req_id in reqs: + self._publication_holds[req_id] -= 1 + if not self._publication_holds[req_id]: + del self._publication_holds[req_id] + if current_generation: + self.stats["skipped"] += skipped + self._busy[(epoch - 1) % len(self._staging)] = False + self._publication_ack.set() + + def _finalize_sync(self, ticket): + if ticket is None or ticket.staging is None: + return + staging = ticket.staging + slot_id = (ticket.epoch - 1) % len(self._staging) + try: + if ticket.cache_generation != self._cache_generation: + return + # Called after the batch's normal post handler and compute event. + metadata = ticket.metadata.tolist() + for slot, (req_index, length, row) in enumerate(zip(*metadata)): + if req_index < 0: + continue + req = ticket.reqs[row] + if req.req_idx != req_index or not self._eligible(req): + continue + # A later stop decision cannot relabel this frozen state as an + # earlier version. MTP rows after a stop are never admitted. + output_len = ticket.output_lengths[row] + visible_limit = getattr(req, "exact_visible_end", -1) + is_prompt_end = output_len == 1 + if ( + not ticket.is_prefill + and not is_prompt_end + and (not req.finish_status.is_finished() or length > visible_limit) + ): + continue + tokens = req.shm_req.shm_prompt_ids.arr[:length].copy() + seed = None if ticket.output_seed is None else ticket.output_seed[row] + pending = None + try: + pending = self.cache.prepare( + tokens, + self.mem_manager.kv_buffer, + self.req_manager.req_to_token_indexs[req_index, :length], + staging.conv_state[slot], + staging.ssm_state[slot], + output_seed=seed, + namespace=self.namespace, + frozen_tail_kv=ticket.tail_kv[:, slot], + origins=ticket.kv_origins[req_index][:length], + ) + except (MemoryError, RuntimeError) as error: + if not self._is_allocation_failure(error): + raise + # A local allocation failure is an admission rejection on + # every TP rank, not an early exit before the collective. + logger.warning("checkpoint allocation skipped req=%s: %s", req.req_id, error) + if self._all(pending is not None): + self.cache.commit(pending) + self.stats["captured"] += 1 + req.exact_checkpoint_length = length + logger.debug("checkpoint published req=%s length=%s output_len=%s", req.req_id, length, output_len) + if self.transport is not None and self.backend.is_pd_decode_mode: + owner = bytes(req.sampling_param.shm_param.pd_checkpoint_owner_url).decode() + if owner: + self.transport.publish_checkpoint( + tokens, + self.namespace, + owner, + f"{req.req_id}:{ticket.epoch}:{length}", + owner_dp_index=getattr(req, "pd_checkpoint_owner_dp_index", 0), + owner_auth=bytes(req.sampling_param.shm_param.pd_checkpoint_owner_auth).decode(), + ) + else: + if pending is not None: + self.cache.discard(pending) + self.stats["skipped"] += 1 + finally: + # prepare fences all readers of staging before it is reused. + self._busy[slot_id] = False + + def restore(self, req): + if req.cur_kv_len: + return + if not self._eligible(req): + self.req_manager.init_linear_att_state(req) + return + tokens = req.get_input_token_ids().copy() + full_hit = req.cur_output_len == 0 and self.backend.model.supports_exact_output_seed() + lengths = self.cache.candidate_lengths( + tokens, + max_length=len(tokens) if full_hit else len(tokens) - 1, + require_output_seed=bool(self.args.mtp_step), + namespace=self.namespace, + ) + if self.backend.is_pd_decode_mode: + # Partial D KV and the state later produced by P may come from + # different numerical histories. Until P/D negotiate a common + # origin frontier, only D-local completion avoids mixing them. + lengths = [length for length in lengths if length >= len(tokens) - 1] + if full_hit and len(tokens) in lengths: + lease = self.cache.acquire(tokens, len(tokens), namespace=self.namespace) + # A background preparation may evict the candidate between lookup + # and lease acquisition. Missing local candidates are intersected + # away before any TP rank begins a restore. + if lease is None or lease.output_seed is None: + lengths.remove(len(tokens)) + if lease is not None: + lease.close() + lengths = self._intersection(lengths) + if not lengths: + self.req_manager.init_linear_att_state(req) + return + length = lengths[0] + lease = self.cache.acquire(tokens, length, namespace=self.namespace) + if not self._all(lease is not None): + if lease is not None: + lease.close() + self.req_manager.init_linear_att_state(req) + return + try: + origins = lease.origins.tolist() + # A normal complete hit needs no draft repair before HEAD_ONLY. + # Queue its KV, canonical state and seed on one stream, then fence + # once while the CPU lease still protects every source window. + join_restore = self.args.run_mode == "normal" and length == len(tokens) + # The last packed KV slot must be private before the draft adapter + # rebuilds its sampling-dependent tail. + shared_limit = length - 1 if self.args.mtp_step else length + node, gpu_length, values = (None, 0, None) + if shared_limit: + node, gpu_length, values = self.backend.radix_cache.match_prefix( + self.gpu_radix_key(origins, shared_limit), + update_refs=True, + ) + needed = length - gpu_length + available = self.req_manager.mem_manager.allocator.can_use_mem_size + available += ( + self.backend.radix_cache.get_tree_total_tokens_num() - self.backend.radix_cache.get_refed_tokens_num() + ) + if not self._all(needed <= available): + if node is not None: + self.backend.radix_cache.dec_node_ref_counter(node) + self.req_manager.init_linear_att_state(req) + return + self.backend.radix_cache.free_radix_cache_to_get_enough_token(needed) + indexes = self.mem_manager.alloc(needed) + if gpu_length: + self.req_manager.req_to_token_indexs[req.req_idx, :gpu_length] = values.to( + device="cuda", dtype=torch.int32, non_blocking=join_restore + ) + if needed: + self.cache.load_kv(lease, self.mem_manager.kv_buffer, indexes, start=gpu_length, wait=not join_restore) + self.req_manager.req_to_token_indexs[req.req_idx, gpu_length:length] = indexes.to( + device="cuda", non_blocking=join_restore + ) + req.shared_kv_node = node + event = torch.cuda.Event() + event.record() + snapshot = LinearStateSnapshot(lease.conv_state, lease.ssm_state, length, event) + restored = self.req_manager.restore_linear_state(snapshot, req.req_idx) + seed = lease.output_seed + if join_restore: + output_seed = seed.unsqueeze(0).to(device="cuda", non_blocking=True) + torch.cuda.current_stream().synchronize() + else: + restored.synchronize() + req.cur_kv_len = length + req.exact_kv_origins = origins + req.shm_req.shm_cur_kv_len = length + req.shm_req.prompt_cache_len = length + if length == len(tokens): + req.exact_output_seed = output_seed if join_restore else seed.unsqueeze(0).to(device="cuda") + elif self.args.mtp_step: + self.resume_auxiliary(req, seed.unsqueeze(0).to(device="cuda"), int(tokens[length])) + self.stats["hits"] += 1 + self.stats["hit_tokens"] += length + self.stats["loaded_tokens"] += needed + logger.info("exact checkpoint hit req=%s length=%s GPU=%s CPU=%s", req.req_id, length, gpu_length, needed) + except BaseException: + # A failed state/seed allocation can leave an earlier H2D copy in + # flight. Arena windows are reclaimed when this lease is closed. + torch.cuda.current_stream().synchronize() + raise + finally: + lease.close() + + def resume_auxiliary(self, req, output_seed, next_token): + if not self.args.mtp_step: + return + from lightllm.server.router.model_infer.infer_batch import g_infer_context + from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager + + torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) + next_ids = g_pin_mem_manager.gen_from_list(key="exact_aux_next_ids", data=[next_token], dtype=torch.int64).cuda( + non_blocking=True + ) + self._resume_auxiliary_batch([req], output_seed, next_ids) + torch.cuda.current_stream().synchronize() + self._aux_epoch += 1 + req.exact_kv_origins[req.cur_kv_len - 1] = self._origin(req, self._aux_epoch, "draft-tail") + + def _resume_auxiliary_batch(self, reqs, output_seed, next_token_ids): + """Enqueue private draft-tail repairs; the caller owns the completion fence.""" + if not self.args.mtp_step or not reqs: + return [] + from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager + + lengths = [req.cur_kv_len for req in reqs] + req_indexes = g_pin_mem_manager.gen_from_list( + key="exact_aux_req_indexes", data=[req.req_idx for req in reqs], dtype=torch.int32 + ).cuda(non_blocking=True) + seq_lengths = g_pin_mem_manager.gen_from_list( + key="exact_aux_seq_lengths", data=lengths, dtype=torch.int32 + ).cuda(non_blocking=True) + position_deltas = g_pin_mem_manager.gen_from_list( + key="exact_aux_position_deltas", + data=[req.multimodal_params.get("mrope_position_delta", 0) for req in reqs], + dtype=torch.int32, + ).cuda(non_blocking=True) + # Each restored request already owns its last packed target/draft slot. + # Gathering these slots keeps repairs independent for different lengths. + mem_indexes = self.req_manager.req_to_token_indexs[req_indexes.long(), seq_lengths.long() - 1] + batch_size = len(reqs) + device = output_seed.device + # FlashInfer filtered sampling may return int32 IDs. ModelInput and the + # embedding path require int64; keep this conversion entirely on GPU. + next_token_ids = next_token_ids.to(dtype=torch.int64) + model_input = ModelInput( + batch_size=batch_size, + total_token_num=sum(lengths), + max_q_seq_len=1, + max_kv_seq_len=max(lengths), + input_ids=next_token_ids, + mem_indexes=mem_indexes, + b_req_idx=req_indexes, + b_seq_len=seq_lengths, + b_mtp_index=torch.zeros(batch_size, dtype=torch.int32, device=device), + b_position_delta=position_deltas, + b_shared_seq_len=torch.zeros(batch_size, dtype=torch.int32, device=device), + b_shared_radix_node_id=torch.full((batch_size,), -1, dtype=torch.int64, device=device), + is_prefill=False, + multimodal_params=[req.multimodal_params for req in reqs], + ) + self._spec_engine().resume_auxiliary(model_input, output_seed, next_token_ids) + + def process_head_only(self, reqs): + """Consume full-hit seeds on the scheduler stream before classification.""" + from lightllm.server.router.model_infer.infer_batch import InferReqUpdatePack + from lightllm.server.router.model_infer.infer_batch import g_infer_context + from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager + + head_reqs = [] + for req in reqs: + seed = getattr(req, "exact_output_seed", None) + if seed is None or req.cur_output_len or req.infer_aborted or req.finish_status.is_finished(): + continue + if self.backend.is_pd_decode_mode and ( + req.pd_task_failed_num or req.pd_task_num != req.pd_task_success_num + ): + continue + head_reqs.append(req) + if not head_reqs: + return + + torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) + seeds = torch.cat([req.exact_output_seed for req in head_reqs], dim=0) + model_output = self.backend.model.forward_output_seed(seeds) + req_indexes = g_pin_mem_manager.gen_from_list( + key="exact_head_req_indexes", data=[req.req_idx for req in head_reqs], dtype=torch.int32 + ).cuda(non_blocking=True) + mtp_indexes = torch.zeros(len(head_reqs), dtype=torch.int32, device=seeds.device) + next_ids, ids_cpu, logprobs_cpu, ranks_cpu = self.backend._sample_and_scatter_token( + logits=model_output.logits, + b_req_idx=req_indexes, + b_mtp_index=mtp_indexes, + run_reqs=head_reqs, + is_prefill=True, + b_prefill_has_output_cpu=[True] * len(head_reqs), + mask_func=self.backend.prefill_mask_func, + pin_memory_namespace="exact_head_", + ) + self._resume_auxiliary_batch(head_reqs, seeds, next_ids) + # One fence covers sampled CPU outputs and every draft repair. It also + # keeps all seeds and private slots alive before the next decode starts. + torch.cuda.current_stream().synchronize() + for row, req in enumerate(head_reqs): + if self.args.mtp_step: + # Target KV is unchanged; the repaired draft half has new provenance. + self._aux_epoch += 1 + req.exact_kv_origins[req.cur_kv_len - 1] = self._origin(req, self._aux_epoch, "draft-tail") + req.cur_output_len = 1 + req.exact_output_seed = None + if self.args.mtp_step and self.args.run_mode == "normal" and not self.args.mtp_dynamic_verify: + req.exact_mtp_needs_proposal = True + self.backend._post_handle( + run_reqs=[req], + next_token_ids=ids_cpu[row : row + 1], + next_token_logprobs=logprobs_cpu[row : row + 1], + next_token_ranks=ranks_cpu[row : row + 1], + run_reqs_update_packs=[InferReqUpdatePack(req, 1)], + extra_post_req_handle_func=self.backend.extra_post_req_handle_func, + pd_prefill_chunked_handle_func=self.backend.pd_prefill_chunked_handle_func, + ) + self.stats["head_only"] += 1 + + def clear(self): + self._cache_generation += 1 + if self.transport is not None: + self.transport.clear() + self.cache.clear() + + def poll_imports(self): + if self.transport is None: + return + imports = {item.import_id: item for item in self.transport.drain_imports()} + for import_id in self._intersection(imports): + item = imports[import_id] + pending = None + try: + pending = self.cache.prepare_import(item.payload) + except (ValueError, MemoryError, RuntimeError) as exc: + logger.warning("checkpoint import rejected: %s", exc) + success = self._all(pending is not None) + if success: + self.cache.commit(pending) + elif pending is not None: + self.cache.discard(pending) + self.transport.finish_import(import_id, success) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 3585c223e2..3e67f6ffca 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -46,6 +46,7 @@ class InferenceContext: overlap_stream: torch.cuda.Stream = None # 一些情况下推理进程进行异步折叠操作的异步流对象。 cpu_kv_cache_stream: torch.cuda.Stream = None # 用 cpu kv cache 操作的 stream is_linear_att_mixed_model: bool = False # 标记模型是否是full att 混合 linear att 的混合模型。 + exact_prefix_cache: object = None def register( self, @@ -130,10 +131,12 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: def free_a_req_mem(self, free_token_index: List, req: "InferReq"): if self.radix_cache is None: free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][0 : req.cur_kv_len]) - elif CacheTier.GPU not in req.cache_tiers: + elif CacheTier.GPU not in req.cache_tiers or ( + self.exact_prefix_cache is not None and req.sampling_param.disable_prompt_cache + ): self._free_req_mem_without_radix_insert(free_token_index=free_token_index, req=req) else: - if not self.is_linear_att_mixed_model: + if not self.is_linear_att_mixed_model or self.exact_prefix_cache is not None: self._full_att_free_req(free_token_index=free_token_index, req=req) else: self._linear_att_free_req(free_token_index=free_token_index, req=req) @@ -146,7 +149,7 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer shared_kv_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][shared_kv_len : req.cur_kv_len]) - if self.is_linear_att_mixed_model: + if self.is_linear_att_mixed_model and self.exact_prefix_cache is None: # 释放请求尾部尚未移交给 radix cache 的 linear attention 小页状态。 if req.tail_linear_att_small_page_buffer_id is not None: self.radix_cache.linear_att_small_page_buffers.free_state_cache( @@ -168,9 +171,17 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer def _full_att_free_req(self, free_token_index: List, req: "InferReq"): input_token_ids = req.get_input_token_ids() - key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu") + insert_len = req.cur_kv_len + if self.exact_prefix_cache is not None: + if self.args.mtp_step: + insert_len = min(insert_len, len(input_token_ids) - 1) + key = self.exact_prefix_cache.gpu_radix_key(req.exact_kv_origins, insert_len) + else: + key = torch.tensor(input_token_ids[0:insert_len], dtype=torch.int64, device="cpu") # .cpu() 是 流内阻塞操作 - value = self.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].detach().cpu() + value = self.req_manager.req_to_token_indexs[req.req_idx][:insert_len].detach().cpu() + if insert_len < req.cur_kv_len: + free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][insert_len : req.cur_kv_len]) prefix_len, _ = self.radix_cache.insert(key, value) old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len @@ -348,7 +359,6 @@ def filter_reqs(self, finished_reqs: List["InferReq"]): @torch.no_grad() def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if pause_reqs: - free_token_index = [] for req in pause_reqs: if self.args.diverse_mode: @@ -369,7 +379,6 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool, can_alloc_token_num: int): if paused_reqs: - for req in paused_reqs: prefill_need_token_num = req.get_cur_total_len() if prefill_need_token_num > can_alloc_token_num: @@ -402,6 +411,8 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L """ if not self.is_linear_att_mixed_model: return + if self.exact_prefix_cache is not None: + return # 大页对应的 linear att 的拷贝 big_page_token_num = self.args.linear_att_hash_page_size * self.args.linear_att_page_block_num @@ -602,6 +613,7 @@ def __init__( # Disk 需要 CPU 中转，因此表示为 (CPU, Disk)。 # 兼容策略可以同时包含 GPU、CPU 和 Disk 多层。 self.cache_tiers: Tuple[CacheTier, ...] = (CacheTier.GPU,) + self.exact_kv_origins: List[int] = [] # mtp_step 用来记录一个请求 draft模型每步需要生成的token数量 # 正常模式下，这个值为0，在 mtp 模式下，这个值为 draft 模型每步需要生成的token数量 @@ -611,7 +623,7 @@ def __init__( else: self.decode_need_token_num = self._normal_decode_need_token_num - if g_infer_context.is_linear_att_mixed_model: + if g_infer_context.is_linear_att_mixed_model and not self.args.enable_exact_prefix_cache: self.get_chuncked_input_token_len = self.get_chuncked_input_token_len_for_linear_att self.get_chuncked_input_token_ids = self.get_chuncked_input_token_ids_for_linear_att @@ -682,6 +694,9 @@ def _match_radix_cache(self): return def _linear_match_radix_cache(self): + if g_infer_context.exact_prefix_cache is not None: + g_infer_context.exact_prefix_cache.restore(self) + return assert ( g_infer_context.is_linear_att_mixed_model is True ), "current _linear_match_radix_cache only support linear att hybrid model, to do..." diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 560c3b6de4..e0613f07cf 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -59,6 +59,9 @@ def __init__(self) -> None: self.shm_req_manager = ShmReqManager() self.overlap_event_manager = OverlapEventManager() + # Shared by the two infer threads; idle handshakes must not delay a + # partner that has already started a real model forward. + self._forward_generation = 0 # 标识是否支持 overlap 功能，很多子类模式如 xgrammar 和 outlines 当前不支持 overlap 高性能模式 self.support_overlap = True @@ -153,7 +156,7 @@ def init_model(self, kvargs): set_random_seed(2147483647) self.is_linear_att_mixed_model = isinstance(self.model.req_manager, ReqManagerForMamba) - if self.is_linear_att_mixed_model: + if self.is_linear_att_mixed_model and not self.args.enable_exact_prefix_cache: self.linear_att_cache_manager = LinearAttCacheManager( size=self.args.linear_att_cache_size, linear_config=self.model.req_manager.linear_config, @@ -164,7 +167,7 @@ def init_model(self, kvargs): if not self.use_dynamic_prompt_cache: self.radix_cache = None else: - if self.is_linear_att_mixed_model: + if self.is_linear_att_mixed_model and not self.args.enable_exact_prefix_cache: self.radix_cache = LinearAttPagedRadixCache( unique_name=get_unique_server_name(), total_token_num=self.model.mem_manager.size, @@ -258,6 +261,13 @@ def init_model(self, kvargs): if self.args.enable_cpu_cache: self.multi_level_cache_module = MultiLevelKvCacheModule(self) + self.exact_prefix_cache = None + if self.args.enable_exact_prefix_cache: + from lightllm.server.router.model_infer.exact_prefix_cache import ExactPrefixCache + + self.exact_prefix_cache = ExactPrefixCache(self) + g_infer_context.exact_prefix_cache = self.exact_prefix_cache + prof_name = f"lightllm-model_backend-node{self.node_rank}_dev{get_current_device_id()}" prof_mode = self.args.enable_profiling self.profiler = ProcessProfiler(mode=prof_mode, name=prof_name, use_multi_thread=True) if prof_mode else None @@ -357,28 +367,29 @@ def _async_copy_next_token_infos_to_pin_mem( next_token_ids: torch.Tensor, next_token_logprobs: torch.Tensor, next_token_ranks: torch.Tensor, + pin_memory_namespace: str = "", ): """ 把 next token id / logprobs / ranks 异步拷到 pinned memory， 供后续 post_handle 读取。ranks 始终有值（不需要时为常量 -1）。 """ next_token_ids_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( - key="next_token_ids", + key=pin_memory_namespace + "next_token_ids", gpu_tensor=next_token_ids, ) next_token_logprobs_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( - key="next_token_logprobs", + key=pin_memory_namespace + "next_token_logprobs", gpu_tensor=next_token_logprobs, ) # 仅 enable_rl 需要真实 rank；否则跳过 D2H，返回常量 -1。 if self.args.enable_rl: next_token_ranks_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( - key="next_token_ranks", + key=pin_memory_namespace + "next_token_ranks", gpu_tensor=next_token_ranks, ) else: next_token_ranks_cpu = g_pin_mem_manager.get_const_cpu_tensor( - key="next_token_ranks", + key=pin_memory_namespace + "next_token_ranks", shape=next_token_ids_cpu.shape, fill_value=-1, dtype=torch.int32, @@ -551,8 +562,26 @@ def _read_pd_trans_io_buffer_and_update_req_status(self): # pd decode 节点需要预填充 prefill 节点发送过来的产生的首token信息，以使 # 推理过程可以继续。 if self.is_pd_decode_mode: + start, end = obj.start_kv_index, obj.end_kv_index + origins = obj.kv_origins + exact_cache = g_infer_context.exact_prefix_cache + if origins is None and exact_cache is not None and end > start: + # P may run with exact caching disabled. Its KV + # and state still arrive together; assign this + # received range a stable, private D identity. + origin = exact_cache._origin(req, 0, kind=f"pd-import:{start}:{end}") + origins = [origin] * (end - start) + if origins is not None: + if len(origins) != end - start: + raise ValueError("PD KV provenance does not cover the received range") + if len(req.exact_kv_origins) < end: + req.exact_kv_origins.extend([0] * (end - len(req.exact_kv_origins))) + req.exact_kv_origins[start:end] = origins + if obj.prefill_dp_index is not None: + req.pd_checkpoint_owner_dp_index = obj.prefill_dp_index if obj.first_gen_token_id is not None: - assert req.cur_output_len == 0 + if req.cur_output_len != 0: + continue req.cur_output_len += 1 req_to_next_token_ids = ( self.model.req_manager.req_sampling_params_manager.req_to_next_token_ids @@ -671,8 +700,14 @@ def _get_classed_reqs( 4. prefill_reqs 需要进行prefill操作的请求 5. decode_reqs 需要进行decode操作的请求 """ + exact_cache = self.exact_prefix_cache + if exact_cache is not None: + exact_cache.poll_publications() + # 定期对 radix cache 进行 merge，防止查询插入的操作效率下降 self._timer_merge_radix_tree() + if exact_cache is not None: + exact_cache.poll_imports() if self.args.enable_cpu_cache and len(g_infer_context.infer_req_ids) > 0: self.multi_level_cache_module.update_cpu_cache_task_states() @@ -684,6 +719,8 @@ def _get_classed_reqs( return [], [] ready_reqs = self._filter_not_ready_reqs(req_ids) + if exact_cache is not None: + exact_cache.process_head_only(ready_reqs) support_overlap = self.support_overlap ready_reqs = self._reorder_pd_high_priority_reqs(ready_reqs) ready_reqs = self._reorder_long_prefill_reqs(ready_reqs) @@ -703,13 +740,18 @@ def _get_classed_reqs( can_alloc_token_num = g_infer_context.get_can_alloc_token_num() for req_obj in ready_reqs: - + # Publication reads the request's KV prefix on another stream. + # Active requests can extend it, but its slots must remain owned + # until scheduler polling retires every outstanding publication. + pending_publication = exact_cache is not None and exact_cache.has_pending(req_obj) if req_obj.filter_mark: - finished_reqs.append(req_obj) + if not pending_publication: + finished_reqs.append(req_obj) continue if req_obj.wait_pause: - wait_pause_reqs.append(req_obj) + if not pending_publication: + wait_pause_reqs.append(req_obj) continue if req_obj.paused: @@ -721,9 +763,18 @@ def _get_classed_reqs( # 延迟处理 req_obj.filter_mark = True continue - else: + elif not pending_publication: finished_reqs.append(req_obj) - continue + continue + + if ( + self.args.run_mode == "normal" + and req_obj.cur_output_len >= req_obj.sampling_param.shm_param.max_new_tokens + ): + # Accepted rows already reached the output limit in pre-post. + # Keep ownership until the preceding batch's ordered post + # finishes; another forward only delays its final tokens. + continue if no_decode: is_decode = False @@ -738,7 +789,7 @@ def _get_classed_reqs( decode_reqs.append(req_obj) can_alloc_token_num -= token_num else: - if wait_pause_count < pause_max_req_num: + if wait_pause_count < pause_max_req_num and not pending_publication: if self.args.run_mode == "decode": # PD Decode 节点的 token 容量不足时，强制当前请求提前结束以释放资源。 # 单轮只处理 pause_max_req_num 个请求，避免所有资源不足的请求同时退出。 @@ -772,7 +823,7 @@ def _get_classed_reqs( prefill_reqs.append(req_obj) can_alloc_token_num -= token_num else: - if wait_pause_count < pause_max_req_num: + if wait_pause_count < pause_max_req_num and not pending_publication: req_obj.wait_pause = True wait_pause_count += 1 @@ -872,6 +923,7 @@ def _post_handle( ): req_obj: InferReq = req_obj pack: InferReqUpdatePack = pack + was_finished = req_obj.finish_status.is_finished() pack.handle( next_token_id=next_token_id, next_token_logprob=next_token_logprob, @@ -881,6 +933,8 @@ def _post_handle( extra_post_req_handle_func=extra_post_req_handle_func, pd_prefill_chunked_handle_func=pd_prefill_chunked_handle_func, ) + if not was_finished and req_obj.finish_status.is_finished(): + req_obj.exact_visible_end = req_obj.shm_req.input_len + pack.output_len - 1 g_infer_context.req_manager.req_sampling_params_manager.update_reqs_token_counter( req_objs=run_reqs, next_token_ids=next_token_ids @@ -889,6 +943,16 @@ def _post_handle( # 一些可以复用的通用功能函数 def _filter_reqs(self, reqs: List[InferReq]): + if reqs and self.exact_prefix_cache is not None: + ready = [] + for req in reqs: + if self.exact_prefix_cache.has_pending(req): + # Preserve explicit removal intent until the next scheduler + # poll observes that every reader has completed. + req.filter_mark = True + else: + ready.append(req) + reqs = ready if reqs: g_infer_context.filter_reqs(reqs) return @@ -916,8 +980,8 @@ def _sample_and_scatter_token( is_prefill: bool, b_prefill_has_output_cpu: torch.Tensor = None, mask_func: Optional[Callable] = None, + pin_memory_namespace: str = "", ): - if mask_func is not None: assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) @@ -927,7 +991,7 @@ def _sample_and_scatter_token( b_has_out = None if is_prefill: b_has_out = g_pin_mem_manager.gen_from_list( - key="b_has_out", data=b_prefill_has_output_cpu, dtype=torch.bool + key=pin_memory_namespace + "b_has_out", data=b_prefill_has_output_cpu, dtype=torch.bool ).cuda(non_blocking=True) scatter_token( @@ -950,6 +1014,7 @@ def _sample_and_scatter_token( next_token_ids, next_token_logprobs, next_token_ranks, + pin_memory_namespace=pin_memory_namespace, ) return next_token_ids, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..fb5c4bddf7 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -72,6 +72,7 @@ def infer_loop(self): run_way = self.control_state_machine.select_run_way(prefill_reqs=prefill_reqs, decode_reqs=decode_reqs) if run_way.is_prefill(): + self._forward_generation += 1 # 进行一次流同步，保证 _try_read_new_reqs 中的一些算子操作，必然已经完成。 # 防止后续的推理流程读取到显存中可能存在错误的数据。 g_infer_context.get_overlap_stream().wait_stream(torch.cuda.current_stream()) @@ -81,6 +82,7 @@ def infer_loop(self): ) continue elif run_way.is_decode(): + self._forward_generation += 1 # 进行一次流同步，保证 _try_read_new_reqs 中的一些算子操作，必然已经完成。 # 防止后续的推理流程读取到显存中可能存在错误的数据。 g_infer_context.get_overlap_stream().wait_stream(torch.cuda.current_stream()) @@ -90,10 +92,14 @@ def infer_loop(self): ) continue elif run_way.is_pass(): + idle_generation = self._forward_generation event_pack.notify_post_handle_and_wait_pre_post_handle() event_pack.notify_forward_and_wait_post_handle() event_pack.notify_pre_post_handle() - time.sleep(0.02) + # The partner may have started work during these handshakes. + # Its next post step needs us; only back off if both stayed idle. + if self._forward_generation == idle_generation: + time.sleep(0.02) continue except BaseException as e: @@ -107,10 +113,17 @@ def prefill_normal( ): # 第一阶段: 模型推理 model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + cache = self.exact_prefix_cache + ticket = cache.prepare_batch(model_input, run_reqs) if cache is not None else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + next_token_ids, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -123,6 +136,8 @@ def prefill_normal( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) + if cache is not None: + cache.capture(ticket, model_input, model_output, next_token_ids) sync_event = torch.cuda.Event() sync_event.record() @@ -142,6 +157,9 @@ def prefill_normal( extra_post_req_handle_func=self.extra_post_req_handle_func, pd_prefill_chunked_handle_func=self.pd_prefill_chunked_handle_func, ) + if cache is not None: + cache.finalize(ticket) + # 第四阶段 event_pack.notify_pre_post_handle() return @@ -152,9 +170,16 @@ def decode_normal( decode_reqs: List[InferReq], ): model_input, run_reqs = prepare_decode_inputs(decode_reqs) + cache = self.exact_prefix_cache + ticket = cache.prepare_batch(model_input, run_reqs) if cache is not None else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + next_token_ids, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -162,6 +187,8 @@ def decode_normal( is_prefill=False, mask_func=self.decode_mask_func, ) + if cache is not None: + cache.capture(ticket, model_input, model_output, next_token_ids) sync_event = torch.cuda.Event() sync_event.record() @@ -181,6 +208,9 @@ def decode_normal( extra_post_req_handle_func=self.extra_post_req_handle_func, ) + if cache is not None: + cache.finalize(ticket) + # 第四阶段 event_pack.notify_pre_post_handle() return @@ -191,6 +221,8 @@ def prefill_mtp( prefill_reqs: List[InferReq], ): model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + cache = self.exact_prefix_cache + ticket = cache.prepare_batch(model_input, run_reqs) if cache is not None else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) @@ -219,6 +251,8 @@ def prefill_mtp( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) + if cache is not None: + cache.capture(ticket, model_input, model_output, next_token_ids) sync_event = torch.cuda.Event() sync_event.record() @@ -240,6 +274,9 @@ def prefill_mtp( pd_prefill_chunked_handle_func=self.pd_prefill_chunked_handle_func, ) + if cache is not None: + cache.finalize(ticket) + # 第四阶段 event_pack.notify_pre_post_handle() return @@ -251,6 +288,8 @@ def decode_mtp( ): """Run the speculative draft-and-verify decode flow.""" model_input, run_reqs = prepare_decode_inputs(decode_reqs) + cache = self.exact_prefix_cache + ticket = cache.prepare_batch(model_input, run_reqs) if cache is not None else None spec_engine = self.spec_engine req_num = len(decode_reqs) @@ -271,6 +310,9 @@ def decode_mtp( async_selected_row_mask_cpu.wait() selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] + if ticket is not None: + ticket.reqs = run_reqs + ticket.output_lengths = [n for n, keep in zip(ticket.output_lengths, selected_rows) if keep] next_token_ids, next_token_logprobs = sample( model_output.logits, run_reqs, @@ -331,6 +373,8 @@ def decode_mtp( next_token_ranks=next_token_ranks, ) + if cache is not None: + cache.capture(ticket, model_input, model_output, next_token_ids, accepted_index=accepted_index) sync_event = torch.cuda.Event() sync_event.record() @@ -388,6 +432,9 @@ def decode_mtp( extra_mem_indexes_cpu=proposal.extra_mem_indexes_cpu, ) + if cache is not None: + cache.finalize(ticket) + # 第四阶段 event_pack.notify_pre_post_handle() return diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..6ec93521e8 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -142,6 +142,7 @@ def infer_loop(self): ) if run_way.is_prefill(): + self._forward_generation += 1 # 进行一次流同步，保证 _try_read_new_reqs 中的一些算子操作，必然已经完成。 # 防止后续的推理流程读取到显存中可能存在错误的数据。 g_infer_context.get_overlap_stream().wait_stream(torch.cuda.current_stream()) @@ -151,6 +152,7 @@ def infer_loop(self): ) continue elif run_way.is_decode(): + self._forward_generation += 1 # 进行一次流同步，保证 _try_read_new_reqs 中的一些算子操作，必然已经完成。 # 防止后续的推理流程读取到显存中可能存在错误的数据。 g_infer_context.get_overlap_stream().wait_stream(torch.cuda.current_stream()) @@ -160,10 +162,12 @@ def infer_loop(self): ) continue elif run_way.is_pass(): + idle_generation = self._forward_generation event_pack.notify_post_handle_and_wait_pre_post_handle() event_pack.notify_forward_and_wait_post_handle() event_pack.notify_pre_post_handle() - time.sleep(0.02) + if self._forward_generation == idle_generation: + time.sleep(0.02) continue except BaseException as e: @@ -177,12 +181,14 @@ def prefill_normal( ): model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) run_reqs_num = len(run_reqs) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket = checkpoint_cache.prepare_batch(model_input, run_reqs) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) if run_reqs_num > 0: ( - _, + next_token_ids, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu, @@ -199,6 +205,8 @@ def prefill_normal( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) + if checkpoint_cache is not None: + checkpoint_cache.capture(checkpoint_ticket, model_input, model_output, next_token_ids) sync_event = torch.cuda.Event() sync_event.record() @@ -219,6 +227,8 @@ def prefill_normal( extra_post_req_handle_func=self.extra_post_req_handle_func, pd_prefill_chunked_handle_func=self.pd_prefill_chunked_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket) # 第四阶段 event_pack.notify_pre_post_handle() else: @@ -231,11 +241,13 @@ def decode_normal(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq model_input, run_reqs = prepare_decode_inputs(req_objs=decode_reqs) model_input: ModelInput = model_input run_reqs_num = len(run_reqs) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket = checkpoint_cache.prepare_batch(model_input, run_reqs) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) if run_reqs_num > 0: ( - _, + next_token_ids, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu, @@ -247,6 +259,8 @@ def decode_normal(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq is_prefill=False, mask_func=None, ) + if checkpoint_cache is not None: + checkpoint_cache.capture(checkpoint_ticket, model_input, model_output, next_token_ids) sync_event = torch.cuda.Event() sync_event.record() @@ -266,6 +280,8 @@ def decode_normal(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq run_reqs_update_packs=update_packs, extra_post_req_handle_func=self.extra_post_req_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket) # 第四阶段 event_pack.notify_pre_post_handle() @@ -283,6 +299,9 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer run_reqs1, ) = overlap_prepare_prefill_inputs(prefill_reqs) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket0 = checkpoint_cache.prepare_batch(model_input0, run_reqs0) if checkpoint_cache else None + checkpoint_ticket1 = checkpoint_cache.prepare_batch(model_input1, run_reqs1) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) @@ -302,7 +321,7 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer if req_num0 + req_num1 > 0: ( - _, + next_token_ids, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu, @@ -319,6 +338,9 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer if g_infer_context.is_linear_att_mixed_model: g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + if checkpoint_cache is not None: + checkpoint_cache.capture(checkpoint_ticket0, model_input0, model_output0, next_token_ids[:req_num0]) + checkpoint_cache.capture(checkpoint_ticket1, model_input1, model_output1, next_token_ids[req_num0:]) sync_event = torch.cuda.Event() sync_event.record() @@ -340,6 +362,9 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer extra_post_req_handle_func=self.extra_post_req_handle_func, pd_prefill_chunked_handle_func=self.pd_prefill_chunked_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket0) + checkpoint_cache.finalize(checkpoint_ticket1) # 第四阶段 event_pack.notify_pre_post_handle() else: @@ -353,6 +378,9 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe run_reqs = run_reqs0 + run_reqs1 req_num0, req_num1 = len(run_reqs0), len(run_reqs1) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket0 = checkpoint_cache.prepare_batch(model_input0, run_reqs0) if checkpoint_cache else None + checkpoint_ticket1 = checkpoint_cache.prepare_batch(model_input1, run_reqs1) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output0, model_output1 = self.model.microbatch_overlap_decode(model_input0, model_input1) if req_num0 + req_num1 > 0: @@ -360,7 +388,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe b_req_idx = torch.cat((model_input0.b_req_idx, model_input1.b_req_idx), dim=0) b_mtp_index = torch.cat((model_input0.b_mtp_index, model_input1.b_mtp_index), dim=0) ( - _, + next_token_ids, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu, @@ -372,6 +400,9 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe is_prefill=False, mask_func=None, ) + if checkpoint_cache is not None: + checkpoint_cache.capture(checkpoint_ticket0, model_input0, model_output0, next_token_ids[:req_num0]) + checkpoint_cache.capture(checkpoint_ticket1, model_input1, model_output1, next_token_ids[req_num0:]) sync_event = torch.cuda.Event() sync_event.record() @@ -391,6 +422,9 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe run_reqs_update_packs=update_packs, extra_post_req_handle_func=self.extra_post_req_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket0) + checkpoint_cache.finalize(checkpoint_ticket1) # 第四阶段 event_pack.notify_pre_post_handle() @@ -407,6 +441,8 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] is_chuncked_mode=not self.disable_chunked_prefill, ) req_num = len(run_reqs) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket = checkpoint_cache.prepare_batch(model_input, run_reqs) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output: ModelOutput = self.model.forward(model_input) b_has_out_cpu = model_input.b_prefill_has_output_cpu @@ -442,11 +478,12 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] if req_num > 0: g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + if checkpoint_cache is not None: + checkpoint_cache.capture(checkpoint_ticket, model_input, model_output, next_token_ids) sync_event = torch.cuda.Event() sync_event.record() if req_num > 0: - # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) @@ -464,6 +501,8 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] extra_post_req_handle_func=self.extra_post_req_handle_func, pd_prefill_chunked_handle_func=self.pd_prefill_chunked_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket) # 第四阶段 event_pack.notify_pre_post_handle() @@ -480,6 +519,8 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): spec_engine = self.spec_engine req_num = len(decode_reqs) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket = checkpoint_cache.prepare_batch(model_input, run_reqs) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): spec_plan = spec_engine.plan_decode( model_input=model_input, @@ -496,6 +537,11 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): async_selected_row_mask_cpu.wait() selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] + if checkpoint_ticket is not None: + checkpoint_ticket.reqs = list(run_reqs) + checkpoint_ticket.output_lengths = [ + length for length, selected in zip(checkpoint_ticket.output_lengths, selected_rows) if selected + ] if req_num > 0: next_token_ids, next_token_logprobs = sample( @@ -573,6 +619,10 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): next_token_ranks=next_token_ranks, ) + if checkpoint_cache is not None and req_num > 0: + checkpoint_cache.capture( + checkpoint_ticket, model_input, model_output, next_token_ids, accepted_index=accepted_index + ) sync_event = torch.cuda.Event() sync_event.record() @@ -620,6 +670,8 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): run_reqs_update_packs=update_packs, extra_post_req_handle_func=self.extra_post_req_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket) mtp_utils.free_mem_indexes( backend=self, extra_mem_indexes_cpu=proposal.extra_mem_indexes_cpu, @@ -645,6 +697,9 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I model_input1, run_reqs1, ) = overlap_prepare_prefill_inputs(prefill_reqs) + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket0 = checkpoint_cache.prepare_batch(model_input0, run_reqs0) if checkpoint_cache else None + checkpoint_ticket1 = checkpoint_cache.prepare_batch(model_input1, run_reqs1) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) @@ -698,6 +753,9 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I if req_num > 0 and g_infer_context.is_linear_att_mixed_model: g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + if checkpoint_cache is not None: + checkpoint_cache.capture(checkpoint_ticket0, model_input0, model_output0, next_token_ids[:req_num0]) + checkpoint_cache.capture(checkpoint_ticket1, model_input1, model_output1, next_token_ids[req_num0:]) sync_event = torch.cuda.Event() sync_event.record() @@ -717,6 +775,9 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I extra_post_req_handle_func=self.extra_post_req_handle_func, pd_prefill_chunked_handle_func=self.pd_prefill_chunked_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket0) + checkpoint_cache.finalize(checkpoint_ticket1) event_pack.notify_pre_post_handle() else: event_pack.notify_post_handle_and_wait_pre_post_handle() @@ -737,6 +798,9 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf real_request_num1 = len(decode_reqs1) req_num = real_request_num0 + real_request_num1 spec_engine = self.decode_draft_engine + checkpoint_cache = self.exact_prefix_cache + checkpoint_ticket0 = checkpoint_cache.prepare_batch(model_input0, run_reqs0) if checkpoint_cache else None + checkpoint_ticket1 = checkpoint_cache.prepare_batch(model_input1, run_reqs1) if checkpoint_cache else None with torch.cuda.stream(g_infer_context.get_overlap_stream()): spec_plan = spec_engine.plan_decode( model_input0=model_input0, @@ -761,10 +825,24 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf selected_row_mask_cpu0.wait() selected_rows0 = selected_row_mask_cpu0.tensor.tolist() run_reqs0 = [req for req, selected in zip(run_reqs0, selected_rows0) if selected] + if checkpoint_ticket0 is not None: + checkpoint_ticket0.reqs = list(run_reqs0) + checkpoint_ticket0.output_lengths = [ + length + for length, selected in zip(checkpoint_ticket0.output_lengths, selected_rows0) + if selected + ] if selected_row_mask_cpu1 is not None: selected_row_mask_cpu1.wait() selected_rows1 = selected_row_mask_cpu1.tensor.tolist() run_reqs1 = [req for req, selected in zip(run_reqs1, selected_rows1) if selected] + if checkpoint_ticket1 is not None: + checkpoint_ticket1.reqs = list(run_reqs1) + checkpoint_ticket1.output_lengths = [ + length + for length, selected in zip(checkpoint_ticket1.output_lengths, selected_rows1) + if selected + ] verify_row_num0 = model_input0.batch_size verify_row_num1 = model_input1.batch_size @@ -858,6 +936,21 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf next_token_ids=next_token_ids, mask=accepted_index == 1, ) + if checkpoint_cache is not None and req_num > 0: + checkpoint_cache.capture( + checkpoint_ticket0, + model_input0, + model_output0, + next_token_ids[:verify_row_num0], + accepted_index=accepted_index[:verify_row_num0], + ) + checkpoint_cache.capture( + checkpoint_ticket1, + model_input1, + model_output1, + next_token_ids[verify_row_num0:], + accepted_index=accepted_index[verify_row_num0:], + ) sync_event = torch.cuda.Event() sync_event.record() @@ -911,6 +1004,9 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf run_reqs_update_packs=update_packs, extra_post_req_handle_func=self.extra_post_req_handle_func, ) + if checkpoint_cache is not None: + checkpoint_cache.finalize(checkpoint_ticket0) + checkpoint_cache.finalize(checkpoint_ticket1) mtp_utils.free_mem_indexes( backend=self, extra_mem_indexes_cpu=proposal.extra_mem_indexes_cpu, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 5b29ea0510..4f253299f5 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -139,7 +139,7 @@ def _top_p_top_k_sample( int64_batch_next_token_ids = torch.empty_like(batch_next_token_ids, dtype=torch.int64) int64_batch_next_token_ids[:] = batch_next_token_ids batch_next_token_probs = torch.gather(probs, dim=1, index=int64_batch_next_token_ids.view(-1, 1)) - return batch_next_token_ids.view(-1), torch.log(batch_next_token_probs).view(-1) + return int64_batch_next_token_ids.view(-1), torch.log(batch_next_token_probs).view(-1) else: assert False, "Unsupported sampling backend for top_p_top_k_sample" diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 22731439c4..64ff538889 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -94,6 +94,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[InferReq]]: run_reqs: List[InferReq] = [] + initialized_mtp_reqs = [] total_token_num = 0 b_req_idx = [] b_mtp_index = [] @@ -110,8 +111,14 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In total_token_num += seq_len b_mtp_index.append(0) multimodal_params.append(req.multimodal_params) - # process the draft tokens. - for step in range(req.mtp_step): + # HEAD_ONLY repaired the draft tail, but has not generated a proposal. + # The first target row seeds it through the normal proposer below. + draft_steps = req.mtp_step + if getattr(req, "exact_mtp_needs_proposal", False): + initialized_mtp_reqs.append(req) + if req.cur_output_len == 1: + draft_steps = 0 + for step in range(draft_steps): run_reqs.append(req) b_req_idx.append(req.req_idx) seq_len += 1 @@ -161,6 +168,10 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In is_prefill=False, multimodal_params=multimodal_params, ) + # Keep the marker if allocation fails. A pause followed by prefill may + # already advance output_len; consume that stale marker without narrowing. + for req in initialized_mtp_reqs: + del req.exact_mtp_needs_proposal return model_input, run_reqs diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 00489b9c27..9e8b328d29 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -75,8 +75,12 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): continue page_len_list = req.shm_req.token_hash_page_len_list.get_all() - page_len_start_list = [0] + page_len_list assert len(page_list) <= len(page_len_list) + # 只调整加载视图，不能把后续 offload 的新尾页边界替换成历史边界。 + page_len_list = page_len_list[: len(page_list)] + if page_list and req.shm_req.cpu_cache_match_tail_len: + page_len_list[-1] = req.shm_req.cpu_cache_match_tail_len + page_len_start_list = [0] + page_len_list if page_list: match_tokens = page_len_list[len(page_list) - 1] @@ -227,7 +231,6 @@ def _start_kv_cache_offload_task( assert len(token_hash_list) == len(page_len_list) if self.backend.is_master_in_dp: - find_index = bisect.bisect_right(page_len_list, req.cur_kv_len) move_block_size = find_index diff --git a/lightllm/server/router/model_infer/mode_backend/pd/checkpoint_transport.py b/lightllm/server/router/model_infer/mode_backend/pd/checkpoint_transport.py new file mode 100644 index 0000000000..d31854ead1 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/pd/checkpoint_transport.py @@ -0,0 +1,622 @@ +"""Best-effort D -> P checkpoint export, outside the inference critical path. + +The existing NIXL/NCCL mover is a unidirectional GPU P -> D pipeline. This +transport moves already-frozen CPU cache pages directly between the two nodes; +PD Master carries only the owner address. Missing base pages are negotiated +before transfer, and receiving threads only queue imports. The inference +coordinator must prepare/commit each import with TP consensus. +""" + +import atexit +import hashlib +import io +import json +import os +import queue +import secrets +import threading +import time +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import urlsplit + +import httpx +import torch + +from lightllm.server.router.dynamic_prompt.checkpoint_cache import CpuCheckpointCache +from lightllm.utils.envs_utils import get_unique_server_name +from lightllm.utils.log_utils import init_logger +from lightllm.utils.net_utils import get_hostname_ip + +logger = init_logger(__name__) +_PROTOCOL_VERSION = 2 + + +def checkpoint_registry_token(): + # The service ID is public in get_server_info's IPC address. Keep discovery + # credentials independent, while sharing one secret across HTTP workers. + directory = _registry_dir() + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + path = directory / ".registry_token" + if not path.exists(): + temporary = directory / f".registry_token.{os.getpid()}.{secrets.token_hex(8)}" + descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + with os.fdopen(descriptor, "w") as output: + output.write(secrets.token_urlsafe(32)) + try: + # Publish a complete file without replacing another worker's + # secret; O_EXCL creation above and link are both atomic. + os.link(temporary, path) + except FileExistsError: + pass + finally: + temporary.unlink(missing_ok=True) + token = path.read_text() + if len(token) < 32: + raise ValueError("invalid internal checkpoint registry credential") + return token + + +def _registry_dir(): + name = hashlib.sha256(get_unique_server_name().encode()).hexdigest()[:24] + return Path("/dev/shm") / f"lightllm-checkpoint-{name}" + + +def read_checkpoint_registry(): + """Called by the local HTTP process; entries contain only CPU transport metadata.""" + entries = [] + for path in _registry_dir().glob("*.json"): + try: + entry = json.loads(path.read_text()) + os.kill(entry["pid"], 0) + entries.append(entry) + except (OSError, ValueError, KeyError): + continue + return sorted(entries, key=lambda x: (x["dp_index"], x["tp_rank"])) + + +def _save_payload(payload): + output = io.BytesIO() + torch.save(payload, output) + return output.getvalue() + + +def _load_payload(data): + # No remote Python class or callable is deserialized. + return torch.load(io.BytesIO(data), map_location="cpu", weights_only=True) + + +def _validate_origins(payload): + origins = payload.get("origins") + if not isinstance(origins, (list, tuple)) or len(origins) != payload["length"]: + raise ValueError("checkpoint KV provenance must cover the token prefix") + if any(type(origin) is not int or not 0 < origin < (1 << 63) for origin in origins): + raise ValueError("checkpoint KV provenance must use positive integer IDs") + + +def _head_slice(rank, world_size, global_heads): + if global_heads >= world_size: + if global_heads % world_size: + raise ValueError("checkpoint heads do not divide TP size") + width = global_heads // world_size + return slice(rank * width, (rank + 1) * width) + if world_size % global_heads: + raise ValueError("checkpoint replicated heads do not divide TP size") + head = rank // (world_size // global_heads) + return slice(head, head + 1) + + +def _merge_heads(shards, dim, global_heads): + world_size = len(shards) + if global_heads >= world_size: + if global_heads % world_size: + raise ValueError("checkpoint heads do not divide source TP size") + return torch.cat(shards, dim=dim) + if world_size % global_heads: + raise ValueError("checkpoint replicated heads do not divide source TP size") + copies = world_size // global_heads + # Replicated KV/key heads carry the same values. Keep a single copy per head. + return torch.cat([shards[i * copies] for i in range(global_heads)], dim=dim) + + +def _check_layouts(source, destination): + ignored = {"kv_layers", "draft_layer_num"} + if {k: v for k, v in source.items() if k not in ignored} != { + k: v for k, v in destination.items() if k not in ignored + }: + raise ValueError("incompatible checkpoint model/state layout") + source_draft, destination_draft = source.get("draft_layer_num", 0), destination.get("draft_layer_num", 0) + if destination_draft and destination_draft != source_draft: + raise ValueError("destination requires unavailable draft KV") + if source["kv_layers"] - source_draft != destination["kv_layers"] - destination_draft: + raise ValueError("incompatible target KV layer layout") + return bool(source_draft and not destination_draft) + + +def reshard_checkpoint_payloads( + payloads, source_layout, destination_layout, destination_rank, destination_world, destination_namespace=None +): + """Convert local TP fragments through the existing Q/K/V global-head layout.""" + strip_draft = _check_layouts(source_layout, destination_layout) + first = payloads[0] + _validate_origins(first) + for part in payloads[1:]: + for key in ("version", "namespace", "page_size", "length", "tokens", "origins", "page_keys"): + if part[key] != first[key]: + raise ValueError(f"TP checkpoint manifest mismatch: {key}") + if set(part["pages"]) != set(first["pages"]): + raise ValueError("TP checkpoint page coverage mismatch") + if part.get("draft_tail_dependency", False) != first.get("draft_tail_dependency", False): + raise ValueError("TP checkpoint draft dependencies mismatch") + + kv_heads = source_layout["kv_heads"] + k_heads = source_layout["linear_k_heads"] + v_heads = source_layout["linear_v_heads"] + k_dim = source_layout["linear_k_dim"] + v_dim = source_layout["linear_v_dim"] + source_world = len(payloads) + result = { + key: first[key] for key in ("version", "namespace", "page_size", "length", "tokens", "origins", "page_keys") + } + result["draft_tail_dependency"] = False if strip_draft else first.get("draft_tail_dependency", False) + if destination_namespace is not None: + result["namespace"] = destination_namespace + if result["namespace"] != first["namespace"] or strip_draft: + result["page_keys"] = CpuCheckpointCache.derive_page_keys( + first["tokens"], + namespace=result["namespace"], + page_size=first["page_size"], + draft_tail_dependency=result["draft_tail_dependency"], + origins=first["origins"], + ) + renamed = dict(zip(first["page_keys"], result["page_keys"])) + result["pages"] = {} + for key in first["pages"]: + keys, values = [], [] + for part in payloads: + page = part["pages"][key] + if strip_draft: + page = page[: destination_layout["kv_layers"]] + local_heads = page.shape[2] // 2 + keys.append(page[:, :, :local_heads]) + values.append(page[:, :, local_heads:]) + head_range = _head_slice(destination_rank, destination_world, kv_heads) + result["pages"][renamed[key]] = torch.cat( + (_merge_heads(keys, 2, kv_heads)[:, :, head_range], _merge_heads(values, 2, kv_heads)[:, :, head_range]), + dim=2, + ).contiguous() + + q_parts, k_parts, v_parts = [], [], [] + for part in payloads: + conv = part["conv_state"] + local_k = max(1, k_heads // source_world) + local_v = max(1, v_heads // source_world) + q, k, v = conv.split((local_k * k_dim, local_k * k_dim, local_v * v_dim), dim=1) + q_parts.append(q.reshape(q.shape[0], local_k, k_dim, q.shape[-1])) + k_parts.append(k.reshape(k.shape[0], local_k, k_dim, k.shape[-1])) + v_parts.append(v.reshape(v.shape[0], local_v, v_dim, v.shape[-1])) + k_range = _head_slice(destination_rank, destination_world, k_heads) + v_range = _head_slice(destination_rank, destination_world, v_heads) + result["conv_state"] = torch.cat( + ( + _merge_heads(q_parts, 1, k_heads)[:, k_range].flatten(1, 2), + _merge_heads(k_parts, 1, k_heads)[:, k_range].flatten(1, 2), + _merge_heads(v_parts, 1, v_heads)[:, v_range].flatten(1, 2), + ), + dim=1, + ).contiguous() + result["ssm_state"] = _merge_heads([p["ssm_state"] for p in payloads], 1, v_heads)[:, v_range].clone( + memory_format=torch.contiguous_format + ) + # The LM-head input is replicated after target TP reduction. A transport + # must not interpret a vocabulary-sharded logits tensor as such a seed. + result["output_seed"] = first["output_seed"] + return result + + +@dataclass +class CheckpointImport: + import_id: str + payload: dict + created_at: float + + +class PDCheckpointTransport: + """One CPU-only endpoint per inference rank, with bounded background work.""" + + def __init__(self, backend, cache, namespace="default", target_namespace=None): + self.backend = backend + self.cache = cache + self.args = backend.args + self.tp_rank = backend.rank_in_dp + self.tp_world = backend.dp_world_size + self.dp_index = backend.dp_rank_in_node + self.namespace = namespace + self.target_namespace = namespace if target_namespace is None else target_namespace + self.timeout = 60.0 + self._lock = threading.Lock() + self._imports = {} + self._import_status = {} + self._exports = {} + self._jobs = queue.Queue(maxsize=4) + self._closed = False + self._generation = 0 + self._auth = secrets.token_urlsafe(32) + cfg = backend.model.mem_manager.linear_config + self.layout = { + "kind": "qwen-linear-v1", + "kv_heads": cfg.full_att_all_num_kv_heads, + "kv_head_dim": cfg.full_att_head_dim, + "kv_layers": cfg.get_full_att_kv_layer_num_with_draft_model(), + "target_layer_num": cfg.get_main_model_full_att_layer_num(), + "draft_layer_num": cfg.draft_full_att_kv_layer_num, + "linear_k_heads": cfg.global_linear_k_heads, + "linear_v_heads": cfg.global_linear_v_heads, + "linear_k_dim": cfg.head_linear_k_dim, + "linear_v_dim": cfg.head_linear_v_dim, + "linear_layers": cfg.linear_layer_num, + "conv_width": cfg.conv_kernel_size - 1, + "kv_dtype": str(cfg.full_att_dtype), + "conv_dtype": str(cfg.conv_state_dtype), + "ssm_dtype": str(cfg.ssm_state_dtype), + } + host = self.args.host + if host in ("0.0.0.0", "127.0.0.1", "localhost"): + host = get_hostname_ip() + transport = self + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + with transport._lock: + authorized = self.headers.get("Authorization") == f"Bearer {transport._auth}" + generation = transport._generation + if not authorized: + self.send_error(403) + return + try: + length = int(self.headers.get("Content-Length", "0")) + # A valid payload contains no more than one cache-sized + # checkpoint plus serialization metadata. + limit = max(16 << 20, int(getattr(cache, "max_bytes", 1 << 30)) * 2) + if not 0 < length <= limit: + self.send_error(413) + return + self.connection.settimeout(transport.timeout) + data = self.rfile.read(length) + if len(data) != length: + raise ValueError("incomplete checkpoint request") + value = _load_payload(data) + response = transport._handle(self.path, value, generation=generation) + body = response if isinstance(response, bytes) else _save_payload(response) + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except Exception as exc: + logger.warning(f"checkpoint transport request rejected: {exc}") + self.send_error(409, "checkpoint unavailable") + + self._server = ThreadingHTTPServer(("0.0.0.0", 0), Handler) + self._server.daemon_threads = True + self.descriptor = { + "protocol_version": _PROTOCOL_VERSION, + "pid": os.getpid(), + "dp_index": self.dp_index, + "tp_rank": self.tp_rank, + "tp_world": self.tp_world, + "url": f"http://{host}:{self._server.server_port}", + "auth": self._auth, + "layout": self.layout, + "namespace": self.namespace, + "target_namespace": self.target_namespace, + "draft_tail_dependency": bool(cache.draft_tail_dependency), + "page_size": cache.page_size, + } + directory = _registry_dir() + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + self._registry_path = directory / f"{backend.rank_in_node}.json" + self._write_registry() + threading.Thread(target=self._server.serve_forever, daemon=True).start() + if self.tp_rank == 0: + threading.Thread(target=self._export_loop, daemon=True).start() + atexit.register(self.close) + + def _write_registry(self): + temporary = self._registry_path.with_suffix(f".{os.getpid()}.tmp") + temporary.write_text(json.dumps(self.descriptor)) + os.replace(temporary, self._registry_path) + + def clear(self): + """Invalidate queued/in-flight imports without a device or network wait.""" + with self._lock: + self._generation += 1 + self._auth = secrets.token_urlsafe(32) + self.descriptor = dict(self.descriptor, auth=self._auth) + self._imports.clear() + self._import_status.clear() + leases = list(self._exports.values()) + self._exports.clear() + for _, lease in leases: + lease.close() + while True: + try: + self._jobs.get_nowait() + self._jobs.task_done() + except queue.Empty: + break + self._write_registry() + + def close(self): + if self._closed: + return + self._closed = True + self._registry_path.unlink(missing_ok=True) + self._server.shutdown() + self._server.server_close() + with self._lock: + leases = list(self._exports.values()) + self._exports.clear() + for _, lease in leases: + lease.close() + + def drain_imports(self): + """Peek pending imports. Call finish_import after a TP-consistent decision.""" + now = time.monotonic() + with self._lock: + expired = [key for key, item in self._imports.items() if now - item.created_at > self.timeout] + for key in expired: + self._imports.pop(key) + self._import_status[key] = (now, "expired") + return list(self._imports.values()) + + def finish_import(self, import_id, success): + with self._lock: + self._imports.pop(import_id, None) + self._import_status[import_id] = (time.monotonic(), "ready" if success else "rejected") + + def publish_checkpoint(self, tokens, namespace, owner_url, export_id, owner_dp_index=0, owner_auth=None): + """Enqueue only on the DP leader, after local TP checkpoint commit.""" + if self.tp_rank != 0 or not owner_url or self._closed: + return False + if isinstance(owner_url, bytes): + owner_url = owner_url.decode("utf-8") + if isinstance(owner_auth, bytes): + owner_auth = owner_auth.decode("utf-8") + if not owner_auth: + # A P node without exact caching did not register this capability. + return False + parsed = urlsplit(owner_url) + if parsed.scheme != "http" or not parsed.hostname or parsed.username or parsed.password or parsed.path: + return False + try: + # Request token views contain NumPy scalars; the wire format deliberately + # permits only built-in values and tensors for weights_only decoding. + self._jobs.put_nowait( + ( + [int(token) for token in tokens], + namespace, + owner_url, + str(export_id), + int(owner_dp_index), + owner_auth, + ) + ) + return True + except queue.Full: + logger.info("checkpoint export admission skipped: queue full") + return False + + def _handle(self, path, obj, generation=None): + if not isinstance(obj, dict): + raise ValueError("expected checkpoint request object") + with self._lock: + if generation is not None and generation != self._generation: + raise ValueError("checkpoint cache was cleared during transfer") + self._expire_exports() + if path == "/pin": + if obj["namespace"] != self.namespace: + raise ValueError("checkpoint namespace incompatible") + lease = self.cache.match(obj["tokens"], max_length=len(obj["tokens"]), namespace=obj["namespace"]) + if lease is None or lease.length != len(obj["tokens"]): + if lease is not None: + lease.close() + raise ValueError("checkpoint no longer present") + export_id = secrets.token_hex(16) + with self._lock: + if (generation is not None and generation != self._generation) or len(self._exports) >= 8: + lease.close() + raise ValueError("checkpoint export leases full") + self._exports[export_id] = (time.monotonic(), lease) + return { + "lease_id": export_id, + "page_keys": lease.page_keys, + "origins": lease.origins.tolist(), + "layout": self.layout, + } + if path == "/export": + with self._lock: + _, lease = self._exports[obj["lease_id"]] + # Serialize while this lease is protected from release/expiry. + # No additional full-prefix tensor copy or GPU read is needed. + payload = self.cache.export(lease, known_page_keys=obj.get("known_page_keys", ())) + return _save_payload(payload) + if path == "/release": + with self._lock: + entry = self._exports.pop(obj["lease_id"], None) + if entry is not None: + entry[1].close() + return {"ok": True} + if path == "/missing": + return {"missing": self.cache.missing_page_keys(obj["page_keys"])} + if path == "/import": + import_id = obj["import_id"] + _validate_origins(obj["payload"]) + if obj["layout"] != self.layout: + raise ValueError("checkpoint layout mismatch") + if obj["payload"]["namespace"] != self.namespace: + raise ValueError("checkpoint namespace mismatch") + with self._lock: + if generation is not None and generation != self._generation: + raise ValueError("checkpoint cache was cleared during transfer") + prior = self._import_status.get(import_id) + if prior is not None: + return {"status": prior[1]} + if import_id not in self._imports: + if len(self._imports) >= 4: + raise ValueError("checkpoint import queue full") + self._imports[import_id] = CheckpointImport(import_id, obj["payload"], time.monotonic()) + return {"status": "pending"} + if path == "/status": + with self._lock: + status = self._import_status.get(obj["import_id"]) + return {"status": status[1] if status is not None else "pending"} + if path == "/cancel": + self.finish_import(obj["import_id"], False) + return {"ok": True} + raise ValueError("unknown checkpoint operation") + + def _expire_exports(self): + now = time.monotonic() + with self._lock: + expired = [key for key, (created, _) in self._exports.items() if now - created > self.timeout * 2] + leases = [self._exports.pop(key)[1] for key in expired] + self._import_status = {key: value for key, value in self._import_status.items() if now - value[0] < 300} + for lease in leases: + lease.close() + + @staticmethod + def _call(client, endpoint, path, payload): + response = client.post( + endpoint["url"] + path, + content=_save_payload(payload), + headers={"Authorization": f"Bearer {endpoint['auth']}", "Content-Type": "application/octet-stream"}, + ) + response.raise_for_status() + return _load_payload(response.content) + + def _export_loop(self): + while not self._closed: + try: + job = self._jobs.get(timeout=1) + except queue.Empty: + self._expire_exports() + continue + try: + self._export_one(*job) + except Exception as exc: + logger.warning(f"checkpoint D -> P export skipped: {exc}") + finally: + self._jobs.task_done() + + def _export_one(self, tokens, namespace, owner_url, export_id, owner_dp_index, owner_auth): + sources = [entry for entry in read_checkpoint_registry() if entry["dp_index"] == self.dp_index] + if len(sources) != self.tp_world: + raise ValueError("source TP checkpoint endpoints incomplete") + pins, destinations, import_ids = [], [], [] + with httpx.Client(timeout=self.timeout, trust_env=False) as client: + try: + response = client.get( + owner_url + "/pd_checkpoint/registry", + headers={"Authorization": f"Bearer {owner_auth}"} if owner_auth else {}, + ) + response.raise_for_status() + destinations = [e for e in response.json()["ranks"] if e["dp_index"] == owner_dp_index] + if not destinations or len(destinations) != destinations[0]["tp_world"]: + raise ValueError("destination TP checkpoint endpoints incomplete") + origins, page_keys = None, None + for source in sources: + if source.get("protocol_version") != _PROTOCOL_VERSION: + raise ValueError("source checkpoint protocol incompatible") + pin = self._call(client, source, "/pin", {"tokens": tokens, "namespace": namespace}) + pins.append((source, pin["lease_id"])) + if pin["layout"] != self.layout: + raise ValueError("source checkpoint layouts differ") + if origins is not None and (pin["origins"] != origins or pin["page_keys"] != page_keys): + raise ValueError("source TP checkpoint provenance differs") + origins = pin["origins"] + page_keys = pin["page_keys"] + target = destinations[0] + if target["page_size"] != self.cache.page_size: + raise ValueError("checkpoint KV page sizes differ") + strip_draft = _check_layouts(self.layout, target["layout"]) + expected_namespace = self.target_namespace if strip_draft else namespace + if target["namespace"] != expected_namespace: + raise ValueError("destination checkpoint execution identity incompatible") + destination_page_keys = CpuCheckpointCache.derive_page_keys( + tokens, + namespace=expected_namespace, + page_size=self.cache.page_size, + draft_tail_dependency=target["draft_tail_dependency"], + origins=origins, + ) + missing = set() + for destination in destinations: + if destination.get("protocol_version") != _PROTOCOL_VERSION: + raise ValueError("destination checkpoint protocol incompatible") + if destination["layout"] != target["layout"] or destination["namespace"] != expected_namespace: + raise ValueError("destination checkpoint layout incompatible") + missing.update( + self._call(client, destination, "/missing", {"page_keys": destination_page_keys})["missing"] + ) + known = [ + source_key + for source_key, dest_key in zip(page_keys, destination_page_keys) + if dest_key not in missing + ] + payloads = [ + self._call(client, source, "/export", {"lease_id": lease_id, "known_page_keys": known}) + for source, lease_id in pins + ] + for destination in destinations: + payload = reshard_checkpoint_payloads( + payloads, + self.layout, + destination["layout"], + destination["tp_rank"], + len(destinations), + destination_namespace=expected_namespace, + ) + self._call( + client, + destination, + "/import", + { + "import_id": export_id, + "layout": destination["layout"], + "payload": payload, + }, + ) + import_ids.append((destination, export_id)) + deadline = time.monotonic() + self.timeout + while time.monotonic() < deadline: + states = [ + self._call(client, dst, "/status", {"import_id": key})["status"] for dst, key in import_ids + ] + if all(state == "ready" for state in states): + logger.info( + f"checkpoint D -> P ready length={len(tokens)} pages_sent={len(missing)} " + f"pages_reused={len(known)} source_tp={len(sources)} destination_tp={len(destinations)}" + ) + import_ids.clear() + return + if any(state not in ("pending", "ready") for state in states): + raise ValueError(f"checkpoint import rejected: {states}") + time.sleep(0.02) + raise TimeoutError("checkpoint import commit timed out") + finally: + for destination, key in import_ids: + try: + self._call(client, destination, "/cancel", {"import_id": key}) + except Exception: + pass + for source, lease_id in pins: + try: + self._call(client, source, "/release", {"lease_id": lease_id}) + except Exception: + pass diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py index 472049442f..0e1a2bd369 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py @@ -169,7 +169,9 @@ def _decode_node_gen_trans_tasks(self, req_obj: InferReq): page_kind="linear_att_state", ) else: - assert req_obj.cur_kv_len == input_len - 1 + # Exact checkpoints also preserve the output seed, so a full hit can + # produce the first token on D without transferring any KV from P. + assert req_obj.cur_kv_len in (input_len - 1, input_len) if not group.task_list: # 需要上报一个包含 0 长度的trans task，触发 kv move manager 给 pd master 上报 @@ -183,6 +185,13 @@ def _decode_node_gen_trans_tasks(self, req_obj: InferReq): ) if self.is_master_in_dp: + group.task_list[0].first_token_owner = ( + "decode" + if input_len - req_obj.cur_kv_len <= 1 + and len(group.task_list) == 1 + and group.task_list[0].transfer_kv_num() == 0 + else "prefill" + ) self.info_queue.put(group) return diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index b406405e8a..8c12e6a1a4 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -229,6 +229,7 @@ def dispatch_task_loop(self): page_reg_desc=self.transporter.local_page_mem_desc, request_id=task.request_id, ready_kv_len=task.start_kv_index, + first_token_owner=task.first_token_owner, ) up_status = PDUpKVStatus( @@ -293,6 +294,7 @@ def accept_peer_task_loop( local_trans_task.prefill_agent_name = remote_trans_task.prefill_agent_name local_trans_task.prefill_agent_metadata = remote_trans_task.prefill_agent_metadata local_trans_task.prefill_num_pages = remote_trans_task.prefill_num_pages + local_trans_task.prefill_dp_index = remote_trans_task.prefill_dp_index local_trans_task.prefill_page_reg_desc = remote_trans_task.prefill_page_reg_desc self.request_page_task_queue.put(local_trans_task) logger.info(f"recv WRITE request from prefill: {remote_trans_task.to_str()}") @@ -321,6 +323,7 @@ def accept_peer_task_loop( if local_trans_task is not None: local_trans_task.first_gen_token_id = remote_trans_task.first_gen_token_id local_trans_task.first_gen_token_logprob = remote_trans_task.first_gen_token_logprob + local_trans_task.kv_origins = remote_trans_task.kv_origins self.ready_page_task_queue.put(local_trans_task) logger.info(f"recv WRITE done from prefill: {remote_trans_task.to_str()}") else: diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py index 084e51e1e5..7e237e76aa 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py @@ -155,6 +155,11 @@ def _create_pd_trans_task( first_gen_token_logprob=None, page_kind=page_kind, req_idx=req_idx, + kv_origins=( + list(req_obj.exact_kv_origins[kv_start_index:kv_end_index]) + if page_kind == "kv" and getattr(self.args, "enable_exact_prefix_cache", False) + else None + ), ) req_obj.pd_task_num += 1 return trans_task diff --git a/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py b/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py index 2649c879fc..1e23726f88 100644 --- a/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py +++ b/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py @@ -50,6 +50,9 @@ def dispatch(self, op_name: str, op_args): return getattr(self, op_name)(op_args) def flush_cache(self, request: FlushCacheReq): + checkpoint_cache = getattr(self.backend, "exact_prefix_cache", None) + if checkpoint_cache is not None: + checkpoint_cache.clear() if self.backend.radix_cache is not None: self.backend.radix_cache.flush_cache() return True, "Succeeded to flush cache." diff --git a/lightllm/server/router/model_infer/mtp_speculative/engine.py b/lightllm/server/router/model_infer/mtp_speculative/engine.py index 9d59afd8a7..d684c96de8 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/engine.py +++ b/lightllm/server/router/model_infer/mtp_speculative/engine.py @@ -46,6 +46,15 @@ def __init__( # Prefill draft-state initialization. + def supports_exact_prefix_resume(self) -> bool: + return self.proposer.supports_exact_prefix_resume() + + def resume_auxiliary( + self, resume_input: ModelInput, output_seed: torch.Tensor, next_token_ids: torch.Tensor + ) -> None: + """Repair mode-specific draft state after an exact target checkpoint hit.""" + self.proposer.resume_auxiliary(resume_input, output_seed, next_token_ids) + def fill_draft_model_kv_state( self, target_model_input: ModelInput, diff --git a/lightllm/server/router/model_infer/mtp_speculative/planner/fixed.py b/lightllm/server/router/model_infer/mtp_speculative/planner/fixed.py index b46fdd1ceb..d2554c8c17 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/planner/fixed.py +++ b/lightllm/server/router/model_infer/mtp_speculative/planner/fixed.py @@ -15,11 +15,13 @@ def __init__(self, max_draft_step: int) -> None: self.max_draft_step = int(max_draft_step) def plan(self, decode_reqs: List, origin_batch_size: int) -> SpecDecodePlan: + target_only = bool(decode_reqs) and origin_batch_size == len(decode_reqs) return SpecDecodePlan( origin_batch_size=origin_batch_size, dynamic_batch_size=origin_batch_size, draft_step=self.max_draft_step, - pre_draft_step=self.max_draft_step, + pre_draft_step=0 if target_only else self.max_draft_step, + all_reqs_have_proposals=origin_batch_size == len(decode_reqs) * (self.max_draft_step + 1), ) def update_statics( diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py index 4b3d0ed6c6..edaa046386 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py @@ -59,6 +59,14 @@ def __init__(self, *, backend: "ModeBackend", enable_dynmaic_mtp: bool) -> None: self.backend = backend self.enable_dynmaic_mtp = bool(enable_dynmaic_mtp) + def supports_exact_prefix_resume(self) -> bool: + return False + + def resume_auxiliary( + self, resume_input: ModelInput, output_seed: torch.Tensor, next_token_ids: torch.Tensor + ) -> None: + raise NotImplementedError("the configured proposer has no exact-prefix resume adapter") + @abstractmethod def fill_draft_model_kv_state( self, diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 3d2c0a0e86..7ea5d7d08f 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -11,10 +11,11 @@ MtpMemIndexesToFree, ) from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import EagleSpecProposal +from lightllm.server.router.model_infer.mtp_speculative.proposers.exact_resume import Qwen35ExactResumeMixin from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -class EagleWithAttProposer(BaseSpecProposer): +class EagleWithAttProposer(Qwen35ExactResumeMixin, BaseSpecProposer): """使用 attention KV cache 的 EAGLE proposer。""" def fill_draft_model_kv_state( diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/exact_resume.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/exact_resume.py new file mode 100644 index 0000000000..638f0bd45a --- /dev/null +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/exact_resume.py @@ -0,0 +1,58 @@ +"""Auxiliary-cache repair for a single Qwen3.5 attention draft module.""" + +import copy + +import torch + +from lightllm.common.basemodel.batch_objs import ModelInput + + +class Qwen35ExactResumeMixin: + def supports_exact_prefix_resume(self) -> bool: + """Do not apply a final-hidden recipe to multi-layer-hidden drafters.""" + from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel + + return ( + self.backend.args.mtp_mode in ("vanilla_with_att", "eagle_with_att") + and len(self.backend.draft_models) == 1 + and isinstance(self.backend.draft_models[0], Qwen3_5MTPModel) + and self.backend.model.supports_exact_output_seed() + ) + + def resume_auxiliary( + self, + resume_input: ModelInput, + output_seed: torch.Tensor, + next_token_ids: torch.Tensor, + ) -> None: + """Rebuild draft[L-1] using H@L and the new token at logical index L. + + ``resume_input`` contains one decode row per request, b_seq_len=L and + an independently owned target/draft packed KV slot at token L-1. The + backend must copy the target layers to that slot before calling us; + this forward changes only the draft layer. Draft KV for [0,L-1) must + already be restored. Rebuilding a shared slot would corrupt another + checkpoint and is forbidden by the caller's restore contract. + + The new token is either a fresh HEAD_ONLY sample or the first input + suffix token, so the same repair applies to partial and complete hits. + This does not advance the target recurrent state or generate an extra + user-visible token. The normal first decode iteration seeds proposals. + """ + if not self.supports_exact_prefix_resume(): + raise NotImplementedError("the configured proposer has no exact-prefix resume adapter") + if resume_input.is_prefill: + raise ValueError("auxiliary resume requires one decode row per restored request") + batch_size = resume_input.batch_size + if output_seed.ndim != 2 or output_seed.shape[0] != batch_size: + raise ValueError("output seed rows must match resumed requests") + if next_token_ids.shape != (batch_size,) or not next_token_ids.is_cuda: + raise ValueError("resume tokens must be a CUDA vector with one token per request") + if not output_seed.is_cuda or output_seed.device != next_token_ids.device: + raise ValueError("resume hidden and tokens must use the same CUDA device") + draft_input = copy.copy(resume_input) + draft_input.input_ids = next_token_ids + draft_input.b_mtp_index = torch.zeros_like(resume_input.b_req_idx) + # The Qwen3.5 draft pre-layer normalizes this argument in-place. + draft_input.mtp_draft_input_hiddens = output_seed.clone() + self.backend.draft_models[0].forward(draft_input) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py index ac844cc8d6..fa9d9ccdd8 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/vanilla_with_att.py @@ -8,11 +8,12 @@ build_chained_mtp_decode_input_inplace, ) from lightllm.server.router.model_infer.mtp_speculative.proposers.base import BaseSpecProposer +from lightllm.server.router.model_infer.mtp_speculative.proposers.exact_resume import Qwen35ExactResumeMixin from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import VanillaSpecProposal from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -class VanillaWithAttProposer(BaseSpecProposer): +class VanillaWithAttProposer(Qwen35ExactResumeMixin, BaseSpecProposer): """使用 attention KV cache 的 Vanilla chained MTP proposer。""" def fill_draft_model_kv_state( diff --git a/lightllm/utils/checkpoint_identity.py b/lightllm/utils/checkpoint_identity.py new file mode 100644 index 0000000000..21662d515c --- /dev/null +++ b/lightllm/utils/checkpoint_identity.py @@ -0,0 +1,68 @@ +"""Deployment identities shared by checkpoint storage and PD transports.""" + +import hashlib +import json +from pathlib import Path + +import yaml + + +def _model_identity(directory): + path = Path(directory) + with (path / "config.json").open() as source: + config = json.load(source) + # Follow load_hf_weights: safetensors take precedence, otherwise load .bin. + # Immutable deployments preserve this metadata across P/D copies. This is + # not a content checksum: replacing weights while preserving size/mtime + # requires a new weight_version and restarting the affected services. + files = sorted(path.glob("*.safetensors")) or sorted(path.glob("*.bin")) + weights = [] + for file in files: + metadata = file.stat() + weights.append((file.name, metadata.st_size, metadata.st_mtime_ns)) + return dict(config=config, weights=weights) + + +def get_checkpoint_identity(args): + """Return target fingerprint, auxiliary fingerprint and store namespace.""" + + def digest(value): + return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + quant_config = None + if getattr(args, "quant_cfg", None) is not None: + # Quantcfg reads YAML, including its JSON subset. Include parsed content + # so equivalent configuration files can live at different paths. + with Path(args.quant_cfg).open() as source: + quant_config = yaml.safe_load(source) + execution = dict( + dtype=args.data_type, + kv_type=args.llm_kv_type, + quant_type=args.quant_type, + quant_config=quant_config, + expert_dtype=getattr(args, "expert_dtype", None), + ssm_dtype=args.linear_att_ssm_data_type, + ) + target = digest( + dict( + model=_model_identity(args.model_dir), + version=args.weight_version, + execution=execution, + ) + ) + draft = "" + if args.mtp_step: + directories = args.mtp_draft_model_dir or [args.model_dir] + if isinstance(directories, str): + directories = [directories] + draft = digest( + dict( + models=[_model_identity(directory) for directory in directories], + mode=args.mtp_mode, + version=args.weight_version, + # init_mtp_draft_model inherits the target's dtype, quant_cfg, + # quant_type and expert_dtype; they also identify draft state. + execution=execution, + ) + ) + return target, draft, f"{target}:{draft or 'target-only'}" diff --git a/test/benchmark/agent_checkpoint_cache.py b/test/benchmark/agent_checkpoint_cache.py new file mode 100644 index 0000000000..4f2c1100a5 --- /dev/null +++ b/test/benchmark/agent_checkpoint_cache.py @@ -0,0 +1,602 @@ +"""Native HTTP correctness and timing checks for exact Agent checkpoints. + +Run through the experiment ledger, for example:: + + exp -m "Agent checkpoint comparison" python test/benchmark/agent_checkpoint_cache.py \ + --baseline-url http://127.0.0.1:17880 --candidate-url http://127.0.0.1:17881 \ + --model-dir /models/Qwen3.5-0.8B --output /dev/shm/checkpoint-results + +Each request records the exact input/output token IDs and parameters. Streaming +timings are client-observed; MTP can deliver several tokens in one burst. Older +native streams omit cache metadata: their hit length remains null, and a +separately labeled nonstream probe records cache hits without inventing TTFT. +This is a functional workload with short outputs, not a saturation benchmark. +""" + +import argparse +import concurrent.futures +import hashlib +import json +import math +from pathlib import Path +import statistics +import subprocess +import threading +import time +import urllib.error +import urllib.request +import uuid + + +CACHE_FIELDS = ("prompt_cache_len", "mtp_accepted_token_num", "mtp_verify_token_num", "mtp_verify_step_num") +DEFAULT_LENGTHS = (255, 256, 257, 8191, 8192, 8193, 12000) + + +def fetch_json(url, timeout): + try: + with urllib.request.urlopen(url, timeout=timeout) as response: + return json.load(response) + except (OSError, ValueError) as error: + return {"unavailable": str(error)} + + +def git_metadata(): + result = {} + for key, command in ( + ("commit", ["git", "rev-parse", "HEAD"]), + ("branch", ["git", "branch", "--show-current"]), + ("status", ["git", "status", "--short"]), + ): + completed = subprocess.run(command, capture_output=True, text=True, check=False) + result[key] = completed.stdout.strip() if completed.returncode == 0 else None + return result + + +def token_logprob(token): + value = token.get("logprob") + if value is None: + values = token.get("logprobs", {}) + entry = values.get(str(token.get("id")), values.get(token.get("id"), {})) + value = entry.get("logprob") if isinstance(entry, dict) else entry + return float(value) if value is not None and math.isfinite(float(value)) else None + + +def request_once(url, case, parameters, phase, stream, timeout): + started = time.perf_counter() + record = { + "case": case["name"], + "kind": case["kind"], + "phase": phase, + "stream": stream, + "url": url, + "started_unix": time.time(), + "input_len": len(case["tokens"]), + "input_token_ids": case["tokens"], + "input_sha256": hashlib.sha256(json.dumps(case["tokens"], separators=(",", ":")).encode()).hexdigest(), + "parameters": parameters, + "events": [], + "output_token_ids": [], + "output_logprobs": [], + "cache_hit_len": None, + "cache_hit_source": None, + "ttft_ms": None, + "tpot_ms": None, + "error": None, + } + request = urllib.request.Request( + url.rstrip("/") + ("/generate_stream" if stream else "/generate"), + data=json.dumps({"inputs": case["tokens"], "parameters": parameters}).encode(), + headers={"Content-Type": "application/json"}, + ) + tokens = [] + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + record["http_status"] = response.status + if stream: + for line in response: + received = (time.perf_counter() - started) * 1000 + line = line.decode("utf-8").strip() + if not line or line.startswith(":"): + continue + if not line.startswith("data:"): + continue + body = line[5:].strip() + if body == "[DONE]": + continue + event = json.loads(body) + if "error" in event: + raise RuntimeError(f"SSE error: {event['error']}") + token = event.get("token") + if not isinstance(token, dict) or token.get("id") is None: + continue + tokens.append(token) + record["events"].append({"received_ms": received, **event}) + if record["ttft_ms"] is None: + record["ttft_ms"] = received + if event.get("finished"): + record["finish_reason"] = event.get("finish_reason") + if tokens and "finish_reason" not in record: + raise RuntimeError("stream ended without a finished event") + else: + body = json.load(response) + record["response"] = body + tokens = body.get("tokens", []) + if tokens and isinstance(tokens[0], list): + raise RuntimeError("expected one sequence, received multiple outputs") + record["finish_reason"] = body.get("finish_reason") + if not tokens: + raise RuntimeError("response contained no token IDs") + except (OSError, ValueError, RuntimeError) as error: + record["error"] = str(error) + if isinstance(error, urllib.error.HTTPError): + record["http_status"] = error.code + record["error_body"] = error.read().decode("utf-8", errors="replace")[:4000] + record["latency_ms"] = (time.perf_counter() - started) * 1000 + record["output_token_ids"] = [int(token["id"]) for token in tokens] + record["output_logprobs"] = [token_logprob(token) for token in tokens] + record["output_len"] = len(tokens) + record["prompt_tokens"] = tokens[0].get("prompt_tokens") if tokens else None + for field in CACHE_FIELDS: + values = [token[field] for token in tokens if field in token] + record[field] = max(values) if values else None + record["cache_hit_len"] = record["prompt_cache_len"] + if record["cache_hit_len"] is not None: + record["cache_hit_source"] = "stream.token.prompt_cache_len" if stream else "generate.tokens.prompt_cache_len" + if stream and len(tokens) > 1: + record["tpot_ms"] = (record["events"][-1]["received_ms"] - record["ttft_ms"]) / (len(tokens) - 1) + return record + + +def compare(reference, observed, tolerance): + expected = reference["output_token_ids"] + actual = observed["output_token_ids"] + same_ids = expected == actual + pairs = list(zip(reference["output_logprobs"], observed["output_logprobs"])) + # Chosen-token probabilities are comparable only for the same sequence. + # After a greedy divergence, later positions also have different prefixes. + errors = [abs(a - b) for a, b in pairs if a is not None and b is not None] if same_ids else [] + complete_logprobs = bool(expected) and same_ids and len(errors) == len(expected) + return { + "case": observed["case"], + "reference": reference["record_id"], + "observed": observed["record_id"], + "token_ids_equal": same_ids, + "first_token_difference": next( + (i for i, (a, b) in enumerate(zip(expected, actual)) if a != b), + min(len(expected), len(actual)) if not same_ids else None, + ), + "max_logprob_abs_error": max(errors) if errors else None, + "logprobs_complete": complete_logprobs, + "logprob_atol": tolerance, + "passed": ( + reference["error"] is None + and observed["error"] is None + and same_ids + and complete_logprobs + and max(errors) <= tolerance + ), + } + + +class Workload: + def __init__(self, args): + from transformers import AutoTokenizer + + self.args = args + self.output = Path(args.output) + self.output.mkdir(parents=True, exist_ok=True) + if (self.output / "requests.jsonl").exists(): + raise ValueError("output already contains requests.jsonl; choose a new experiment directory") + self.tokenizer = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=args.trust_remote_code) + self.run_id = args.run_id or uuid.uuid4().hex[:12] + self.endpoints = { + name: url for name, url in (("baseline", args.baseline_url), ("candidate", args.candidate_url)) if url + } + self.server_info = { + name: fetch_json(url.rstrip("/") + "/get_server_info", args.timeout) for name, url in self.endpoints.items() + } + self.records = [] + self.comparisons = [] + self.checks = [] + self.lock = threading.Lock() + self.ordinal = 0 + self.manifest = { + "run_id": self.run_id, + "created_unix": time.time(), + "script_version": 1, + "git": git_metadata(), + "arguments": vars(args), + "server_info": self.server_info, + "cases": [], + "notes": [ + "TTFT and TPOT are observed by this HTTP client; MTP tokens may arrive in bursts.", + "Missing streaming cache metadata is null; detail probes are different requests.", + "Cold correctness controls use disable_prompt_cache=true; a rejected control is a failure.", + "HTTP MTP counters prove activity, not the per-verify stopping row; kernel tests are also required.", + ], + } + self.write_json("manifest.json", self.manifest) + + def write_json(self, name, value): + (self.output / name).write_text(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n") + + def append(self, name, value): + with (self.output / name).open("a") as file: + file.write(json.dumps(value, ensure_ascii=False, allow_nan=False) + "\n") + + def make_prompt(self, length, tag): + prefix = self.tokenizer.encode( + f"Checkpoint run {self.run_id}, case {tag}. Read this technical context carefully.\n", + add_special_tokens=False, + ) + filler = self.tokenizer.encode( + "A request reads a shared prefix and then processes a new suffix. " + "The cache stores attention keys and values, while recurrent state summarizes earlier tokens. ", + add_special_tokens=False, + ) + tail = self.tokenizer.encode("\nContinue the numbered list: 1, 2, 3, 4, 5,", add_special_tokens=False) + remaining = length - len(prefix) - len(tail) + if remaining < 1: + raise ValueError(f"length {length} is too short for the reproducible prompt scaffold") + return prefix + (filler * math.ceil(remaining / len(filler)))[:remaining] + tail + + def case(self, name, kind, tokens, parameters=None, **metadata): + case = {"name": name, "kind": kind, "tokens": list(tokens), "parameters": parameters or {}, **metadata} + self.manifest["cases"].append(case) + self.write_json("manifest.json", self.manifest) + return case + + def run_request(self, server, case, phase, *, stream=True, cold=False, concurrency_round=None): + parameters = { + "do_sample": False, + "seed": self.args.seed, + "max_new_tokens": self.args.max_new_tokens, + "ignore_eos": True, + "add_special_tokens": False, + "skip_special_tokens": False, + "return_details": True, + **case["parameters"], + "disable_prompt_cache": cold, + } + record = request_once(self.endpoints[server], case, parameters, phase, stream, self.args.timeout) + record["server"] = server + record["concurrency_round"] = concurrency_round + with self.lock: + self.ordinal += 1 + record["record_id"] = f"{self.ordinal:05d}-{server}-{case['name']}-{phase}" + self.records.append(record) + self.append("requests.jsonl", record) + print( + f"{server:9s} {case['name']:28s} {phase:18s} input={record['input_len']} " + f"output={record['output_len']} cache={record['cache_hit_len']} " + f"ttft_ms={record['ttft_ms']} error={record['error']}", + flush=True, + ) + self.add_check(record, "input_length", record["prompt_tokens"] == len(case["tokens"])) + if cold and record["cache_hit_len"] is not None: + self.add_check(record, "cold_has_no_cache_hit", record["cache_hit_len"] == 0) + return record + + def add_check(self, record, name, passed, detail=None): + check = {"record_id": record["record_id"], "check": name, "passed": bool(passed), "detail": detail} + with self.lock: + self.checks.append(check) + self.append("checks.jsonl", check) + + def compare(self, reference, observed): + comparison = compare(reference, observed, self.args.logprob_atol) + self.comparisons.append(comparison) + self.append("comparisons.jsonl", comparison) + if not comparison["passed"]: + print("COMPARISON FAILED " + json.dumps(comparison), flush=True) + + def evaluate(self, case): + references = {} + for server in self.endpoints: + # Legacy servers may still insert at request teardown when reads + # are disabled. Probe a new branch before the cold control can + # populate that branch and hide its initial matching behavior. + seed = self.run_request(server, case, "cache_seed", stream=False) + cold = self.run_request(server, case, "cold_reference", cold=True) + references[server] = cold + self.compare(cold, seed) + if "max_initial_hit" in case and seed["cache_hit_len"] is not None: + self.add_check( + seed, "branch_hit_does_not_cross_divergence", seed["cache_hit_len"] <= case["max_initial_hit"] + ) + if self.args.require_exact_hits and server == "candidate" and "min_initial_hit" in case: + self.add_check( + seed, + "output_prefix_reused", + seed["cache_hit_len"] is not None and seed["cache_hit_len"] >= case["min_initial_hit"], + ) + time.sleep(self.args.settle_ms / 1000) + for repeat in range(self.args.repeats): + warm = self.run_request(server, case, f"warm_stream_{repeat}") + self.compare(cold, warm) + self.compare(seed, warm) + if self.args.require_exact_hits and server == "candidate": + self.add_check(warm, "exact_full_hit", warm["cache_hit_len"] == len(case["tokens"])) + if warm["cache_hit_len"] is None and not self.args.no_detail_probe: + detail = self.run_request(server, case, f"cache_detail_probe_{repeat}", stream=False) + self.compare(cold, detail) + if case["kind"] == "eos": + self.add_check( + seed, + "natural_eos_reached", + seed.get("finish_reason") == "stop" + and bool(seed["output_token_ids"]) + and seed["output_token_ids"][-1] in case["eos_token_ids"], + ) + if "stop_sequences" in case["parameters"]: + self.add_check(seed, "token_stop_reached", seed.get("finish_reason") == "stop") + if self.args.require_mtp_activity and server == "candidate": + self.add_check( + seed, "mtp_acceptance_observed_on_stop_request", (seed["mtp_accepted_token_num"] or 0) > 0 + ) + if len(references) == 2: + self.compare(references["baseline"], references["candidate"]) + return references + + def run(self): + suites = set(self.args.suites.split(",")) + if "boundaries" in suites: + for length in self.args.lengths: + self.evaluate(self.case(f"boundary_{length}", "boundary", self.make_prompt(length, str(length)))) + anchor = None + if suites.intersection({"agent", "branch", "stops"}): + anchor = self.case("agent_anchor", "anchor", self.make_prompt(self.args.agent_input_len, "anchor")) + refs = self.evaluate(anchor) + reference = refs.get("baseline", next(iter(refs.values()))) + if reference["error"]: + raise RuntimeError("anchor cold request failed; cannot construct trustworthy continuation cases") + outputs = reference["output_token_ids"] + if "agent" in suites: + tool = self.tokenizer.encode( + '\nTool result: {"status":"ok","value":42}. Continue the answer.\n', add_special_tokens=False + ) + self.evaluate( + self.case( + "agent_tool_suffix", + "agent", + anchor["tokens"] + outputs + tool, + source_case=anchor["name"], + source_output_token_ids=outputs, + min_initial_hit=len(anchor["tokens"]) + len(outputs) - 1, + ) + ) + if "branch" in suites: + split = max(1, len(anchor["tokens"]) - 17) + replacement = self.tokenizer.encode("Different branch. ", add_special_tokens=False) + old = anchor["tokens"][split] + different = next(token for token in replacement if token != old) + branch = anchor["tokens"][:split] + [different] + anchor["tokens"][split + 1 :] + self.evaluate(self.case("branch_before_tail", "branch", branch, max_initial_hit=split)) + if "stops" in suites: + stop = outputs[: min(3, len(outputs))] + stop_refs = self.evaluate( + self.case( + "accepted_prefix_token_stop", + "token_stop", + anchor["tokens"], + {"stop_sequences": [stop]}, + stop_output_position=len(stop), + ) + ) + stopped = stop_refs.get("baseline", next(iter(stop_refs.values())))["output_token_ids"] + stop_tool = self.tokenizer.encode("\nTool result: 7. Continue.\n", add_special_tokens=False) + self.evaluate( + self.case( + "token_stop_tool_suffix", + "stopped_agent", + anchor["tokens"] + stopped + stop_tool, + source_case="accepted_prefix_token_stop", + source_output_token_ids=stopped, + max_initial_hit=len(anchor["tokens"]) + len(stopped), + ) + ) + eos_ids = next(iter(self.server_info.values())).get("eos_id") or [] + if isinstance(eos_ids, int): + eos_ids = [eos_ids] + if eos_ids: + # allowed_token_ids is ignored unless the server uses outlines. + # A completed assistant answer exercises real EOS sampling. + eos_prompt = self.tokenizer.apply_chat_template( + [ + {"role": "user", "content": "Reply with exactly OK."}, + {"role": "assistant", "content": "OK"}, + ], + tokenize=False, + continue_final_message=True, + enable_thinking=False, + ) + self.evaluate( + self.case( + "natural_eos", + "eos", + self.tokenizer.encode(eos_prompt, add_special_tokens=False), + {"ignore_eos": False, "max_new_tokens": max(16, self.args.max_new_tokens)}, + eos_token_ids=eos_ids, + ) + ) + else: + self.manifest["notes"].append("EOS case unavailable: server did not expose an eos_id.") + if "concurrency" in suites: + cases = [ + self.case( + f"concurrent_{i}", + "concurrency", + self.make_prompt(self.args.concurrency_input_len, f"parallel-{i}"), + ) + for i in range(self.args.concurrency) + ] + for server in self.endpoints: + refs = {} + seeds = {} + for case in cases: + refs[case["name"]] = self.run_request(server, case, "cold_reference", cold=True) + seeds[case["name"]] = self.run_request(server, case, "cache_seed", stream=False) + self.compare(refs[case["name"]], seeds[case["name"]]) + time.sleep(self.args.settle_ms / 1000) + with concurrent.futures.ThreadPoolExecutor(max_workers=self.args.concurrency) as executor: + for round_index in range(self.args.repeats): + pending = [ + ( + case, + executor.submit( + self.run_request, + server, + case, + "concurrent_warm", + concurrency_round=round_index, + ), + ) + for case in cases + ] + for case, future in pending: + warm = future.result() + self.compare(refs[case["name"]], warm) + self.compare(seeds[case["name"]], warm) + if self.args.require_exact_hits and server == "candidate": + self.add_check( + warm, "concurrent_exact_full_hit", warm["cache_hit_len"] == len(case["tokens"]) + ) + return self.finish() + + def finish(self, fatal_error=None): + summary = { + "run_id": self.run_id, + "fatal_error": fatal_error, + "servers": {}, + "comparison_count": len(self.comparisons), + } + for server in self.endpoints: + records = [record for record in self.records if record["server"] == server] + warm = [ + record + for record in records + if record["phase"].startswith("warm_stream") or record["phase"] == "concurrent_warm" + ] + timings = {} + for key in ("ttft_ms", "tpot_ms", "latency_ms"): + values = [record[key] for record in warm if record[key] is not None and record["error"] is None] + timings[key] = { + "n": len(values), + "mean": statistics.mean(values) if values else None, + "median": statistics.median(values) if values else None, + } + info = self.server_info[server] + accepted = [ + record["mtp_accepted_token_num"] for record in records if record["mtp_accepted_token_num"] is not None + ] + summary["servers"][server] = { + "requests": len(records), + "errors": sum(record["error"] is not None for record in records), + "warm_timings": timings, + "warm_case_timings": [ + { + key: record[key] + for key in ( + "case", + "phase", + "concurrency_round", + "input_len", + "output_len", + "ttft_ms", + "tpot_ms", + "latency_ms", + "cache_hit_len", + ) + } + for record in warm + ], + "stream_cache_metadata_observed": sum(record["cache_hit_len"] is not None for record in warm), + "mtp_enabled": bool(info.get("mtp_step", 0)), + "mtp_accepted_tokens_observed": max(accepted) if accepted else None, + "mtp_interior_stop_coverage": "requires per-verify trace; HTTP counters alone are insufficient" + if info.get("mtp_step", 0) + else "not applicable: MTP disabled", + } + failed = [comparison for comparison in self.comparisons if not comparison["passed"]] + failed_checks = [check for check in self.checks if not check["passed"]] + summary["failed_comparisons"] = failed + summary["failed_checks"] = failed_checks + summary["passed"] = not ( + fatal_error or failed or failed_checks or any(record["error"] for record in self.records) + ) + self.write_json("manifest.json", self.manifest) + self.write_json("summary.json", summary) + print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) + return 0 if summary["passed"] else 1 + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--baseline-url") + parser.add_argument("--candidate-url") + parser.add_argument( + "--baseline-revision", help="Revision of the deployed baseline; independent of the client checkout" + ) + parser.add_argument("--candidate-revision", help="Revision/diff label of the deployed candidate") + parser.add_argument("--model-dir", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--run-id", help="Stable case identity; default is unique to avoid previous-run cache hits") + parser.add_argument("--seed", type=int, default=1558) + parser.add_argument( + "--lengths", type=lambda value: [int(item) for item in value.split(",")], default=list(DEFAULT_LENGTHS) + ) + parser.add_argument("--suites", default="boundaries,agent,branch,stops,concurrency") + parser.add_argument("--agent-input-len", type=int, default=12000) + parser.add_argument("--concurrency-input-len", type=int, default=257) + parser.add_argument("--max-new-tokens", type=int, default=8) + parser.add_argument("--repeats", type=int, default=1) + parser.add_argument("--concurrency", type=int, default=4) + parser.add_argument("--settle-ms", type=float, default=200) + parser.add_argument("--timeout", type=float, default=120) + parser.add_argument("--logprob-atol", type=float, default=0.03) + parser.add_argument("--no-detail-probe", action="store_true") + parser.add_argument( + "--require-exact-hits", + action="store_true", + help="Assert candidate full and Agent continuation hits; requires streaming cache metadata", + ) + parser.add_argument( + "--require-mtp-activity", + action="store_true", + help="Require accepted MTP tokens on the candidate token-stop request; does not prove the exact stop row", + ) + parser.add_argument("--trust-remote-code", action="store_true") + args = parser.parse_args() + if not args.baseline_url and not args.candidate_url: + parser.error("at least one of --baseline-url and --candidate-url is required") + if set(args.suites.split(",")) - {"boundaries", "agent", "branch", "stops", "concurrency"}: + parser.error("unknown suite") + if ( + min( + [ + args.max_new_tokens, + args.repeats, + args.concurrency, + args.agent_input_len, + args.concurrency_input_len, + *args.lengths, + ] + ) + <= 0 + ): + parser.error("lengths, repeats, and concurrency must be positive") + if args.settle_ms < 0 or args.logprob_atol < 0 or args.timeout <= 0: + parser.error("invalid timing or tolerance parameter") + return args + + +def main(): + workload = Workload(parse_args()) + try: + return workload.run() + except Exception as error: + return workload.finish(fatal_error=f"{type(error).__name__}: {error}") + + +if __name__ == "__main__": + raise SystemExit(main())