【轻量化复刻大模型 / KIMI K3】KDA Serving 基础设施 — 设计 / 实现

【轻量化复刻大模型 / KIMI K3】KDA Serving 基础设施 --- 设计 / 实现

目录

• [【轻量化复刻大模型 / KIMI K3】KDA Serving 基础设施 --- 设计 / 实现](#【轻量化复刻大模型 / KIMI K3】KDA Serving 基础设施 — 设计 / 实现)
  • [0x00 概要](#0x00 概要)
    • [0.1 主要内容](#0.1 主要内容)
    • [0.2 注意说明](#0.2 注意说明)
  • [0x01 先懂 KDA：为什么它会颠覆 serving](#0x01 先懂 KDA：为什么它会颠覆 serving)
    • [1.1 从最朴素的「KV 缓存」说起](#1.1 从最朴素的「KV 缓存」说起)
    • [1.2 KDA 把这张清单换成了一个「状态矩阵」](#1.2 KDA 把这张清单换成了一个「状态矩阵」)
    • [1.3 关键：M 是「逐 token 相关」的矩阵](#1.3 关键：M 是「逐 token 相关」的矩阵)
    • [1.4 释疑](#1.4 释疑)
  • [0x02 总纲：三个天然属性如何变成三层 Infra](#0x02 总纲：三个天然属性如何变成三层 Infra)
    • [2.1 为什么要做这一套](#2.1 为什么要做这一套)
    • [2.2 范围与边界（本项目做了什么 / 刻意没做什么）](#2.2 范围与边界（本项目做了什么 / 刻意没做什么）)
    • [2.3 术语](#2.3 术语)
  • [0x03 实现全貌](#0x03 实现全貌)
    • [3.1 总览图](#3.1 总览图)
    • [3.2 组件数据流](#3.2 组件数据流)
    • [3.3 组件快速浏览](#3.3 组件快速浏览)
  • [0x04 组件一：统一页池](#0x04 组件一：统一页池)
    • [4.1 它要解决什么问题](#4.1 它要解决什么问题)
    • [4.2 难在哪](#4.2 难在哪)
    • [4.3 总体思路](#4.3 总体思路)
    • [4.4 落地架构（类 / 函数 / 数据结构）](#4.4 落地架构（类 / 函数 / 数据结构）)
    • [4.5 数据流 / 执行流](#4.5 数据流 / 执行流)
    • [4.6 反方辩手（自我质疑）](#4.6 反方辩手（自我质疑）)
    • [4.7 何时用 / 何时不用](#4.7 何时用 / 何时不用)
    • [4.8 可视化：](#4.8 可视化：)
  • [0x05 组件二：双粒度前缀缓存](#0x05 组件二：双粒度前缀缓存)
    • [5.1 它要解决什么问题](#5.1 它要解决什么问题)
    • [5.2 难在哪](#5.2 难在哪)
    • [5.3 总体思路](#5.3 总体思路)
    • [5.4 背后的原理](#5.4 背后的原理)
    • [5.5 关键：group 到底是用来做啥的？](#5.5 关键：group 到底是用来做啥的？)
      • [5.5.1 先看清 K3 的层排布（真实 layer_schedule()）](#5.5.1 先看清 K3 的层排布（真实 layer_schedule()）)
      • [5.5.2 group 是"层"吗？------ 不是，它是"对状态的分组维度"](#5.5.2 group 是"层"吗？—— 不是，它是"对状态的分组维度")
      • [5.5.3 为什么"层数"和"组数"分开？（设计动机）](#5.5.3 为什么"层数"和"组数"分开？（设计动机）)
      • [5.5.4 图：group vs 层示意](#5.5.4 图：group vs 层示意)
    • [5.6 checkpoint 的原子性与「all()」的含义](#5.6 checkpoint 的原子性与「all()」的含义)
      • [5.6.1 图：checkpoint 原子性示意](#5.6.1 图：checkpoint 原子性示意)
    • [5.7 落地架构与数据结构](#5.7 落地架构与数据结构)
    • [5.8 数据流 / 执行流（lookup 的两段式）](#5.8 数据流 / 执行流（lookup 的两段式）)
    • [5.9 反方辩手](#5.9 反方辩手)
    • [5.10 何时用 / 何时不用](#5.10 何时用 / 何时不用)
    • [5.11 对比：SGLang 是如何处理 KDA prefix cache 的？](#5.11 对比：SGLang 是如何处理 KDA prefix cache 的？)
      • [5.11.1 图：SGLang KDA 前缀缓存机制](#5.11.1 图：SGLang KDA 前缀缓存机制)
      • [5.11.2 证据链：SGLang 的具体机制](#5.11.2 证据链：SGLang 的具体机制)
      • [5.11.3 上右图实测（chunk 对齐量化，CPU 可复现）](#5.11.3 上右图实测（chunk 对齐量化，CPU 可复现）)
      • [5.11.4 与我们的设计对比](#5.11.4 与我们的设计对比)
    • [5.12 可视化](#5.12 可视化)
  • [0x06 组件三：投机解码 · 投影重放（decode.py）](#0x06 组件三：投机解码 · 投影重放（decode.py）)
    • [6.1 要解决什么问题](#6.1 要解决什么问题)
    • [6.2 难在哪](#6.2 难在哪)
    • [6.3 一个推测解码例子](#6.3 一个推测解码例子)
    • [6.4 总体思路](#6.4 总体思路)
    • [6.5 背后的原理](#6.5 背后的原理)
    • [6.6 落地架构与数据结构](#6.6 落地架构与数据结构)
    • [6.6 数据流 / 执行流](#6.6 数据流 / 执行流)
    • [6.7 反方辩手](#6.7 反方辩手)
    • [6.8 何时用 / 何时不用](#6.8 何时用 / 何时不用)
    • [6.9 可视化](#6.9 可视化)
  • [0x07 组件四：投机解码 · RawInputRing 环形缓冲（replayssm.py）](#0x07 组件四：投机解码 · RawInputRing 环形缓冲（replayssm.py）)
    • [7.1 问题所在](#7.1 问题所在)
    • [7.2 SGLang ReplaySSM](#7.2 SGLang ReplaySSM)
      • 定义
      • [ReplaySSM（= 存 raw input + 重演）的优势](#ReplaySSM（= 存 raw input + 重演）的优势)
      • 适用场景
      • 关键结论
    • [7.3 总体思路](#7.3 总体思路)
    • [7.4 背后的原理](#7.4 背后的原理)
      • [7.4.1 第一性原理：RawInputRing 到底解决什么](#7.4.1 第一性原理：RawInputRing 到底解决什么)
      • [7.4.1 什么场景下有效](#7.4.1 什么场景下有效)
      • [7.4.2 带宽账](#7.4.2 带宽账)
        • 量化对比
        • [raw input 解读](#raw input 解读)
    • [7.5 落地架构与数据结构](#7.5 落地架构与数据结构)
      • [7.5.1 RawInputRing 类设计](#7.5.1 RawInputRing 类设计)
      • [7.5.2 内部数据结构](#7.5.2 内部数据结构)
      • [7.5.3 RawInputRing vs 假想快照](#7.5.3 RawInputRing vs 假想快照)
      • [7.5.4 与 decode.py 的协作](#7.5.4 与 decode.py 的协作)
    • [7.6 何时调用](#7.6 何时调用)
      • [7.6.1 写入路径（draft 期间）](#7.6.1 写入路径（draft 期间）)
      • [7.6.2 读取路径（commit 时）](#7.6.2 读取路径（commit 时）)
      • [7.6.3 容量约束](#7.6.3 容量约束)
    • [7.7 关键对比：投影重放 vs RawInputRing ------ 两条路线各在何时适用](#7.7 关键对比：投影重放 vs RawInputRing —— 两条路线各在何时适用)
    • [7.8 示例](#7.8 示例)
    • [7.9 可视化](#7.9 可视化)
  • [0x08 组件五：KDA 上下文并行（KCP）](#0x08 组件五：KDA 上下文并行（KCP）)
    • [8.1 它要解决什么问题](#8.1 它要解决什么问题)
    • [8.2 难在哪](#8.2 难在哪)
    • [8.3 总体思路](#8.3 总体思路)
    • [8.4 背后的原理](#8.4 背后的原理)
      • [8.4.1 图：KCP 前缀扫描](#8.4.1 图：KCP 前缀扫描)
      • [8.4.2 常见疑问：KCP 用的 kda_chunkwise 到底是什么并行？](#8.4.2 常见疑问：KCP 用的 kda_chunkwise 到底是什么并行？)
    • [8.5 落地架构与数据结构（parallel/kcp.py）](#8.5 落地架构与数据结构（parallel/kcp.py）)
      • [8.5.1 local_transition_and_state（计算每段的 M̄ 和 S̃）](#8.5.1 local_transition_and_state（计算每段的 M̄ 和 S̃）)
      • [8.5.2 prefix_scan_states（前缀扫描重组）](#8.5.2 prefix_scan_states（前缀扫描重组）)
      • [8.5.3 kcp_forward（三阶段总入口）](#8.5.3 kcp_forward（三阶段总入口）)
      • [8.5.4 通信量 & 并行](#8.5.4 通信量 & 并行)
    • [8.6 数据流 / 执行流](#8.6 数据流 / 执行流)
    • [8.7 反方辩手](#8.7 反方辩手)
    • [8.8 何时用 / 何时不用](#8.8 何时用 / 何时不用)
    • [8.9 实测数值（CPU 可复现）](#8.9 实测数值（CPU 可复现）)
    • [8.10 验证](#8.10 验证)
  • [0x09 组件六：Fleet 层调度（serve/fleet.py）](#0x09 组件六：Fleet 层调度（serve/fleet.py）)
    • [9.1 它要解决什么问题](#9.1 它要解决什么问题)
    • [9.2 难在哪](#9.2 难在哪)
    • [9.3 总体思路（两个策略）](#9.3 总体思路（两个策略）)
    • [9.4 背后的原理](#9.4 背后的原理)
    • [9.5 落地架构与数据结构](#9.5 落地架构与数据结构)
    • [9.6 数据流 / 执行流](#9.6 数据流 / 执行流)
    • [9.7 反方辩手](#9.7 反方辩手)
    • [9.8 何时用 / 何时不用](#9.8 何时用 / 何时不用)
    • [9.10 可视化](#9.10 可视化)
  • [0x10 收官引擎：把它们咬合成一个「第二个请求到底省了什么」](#0x10 收官引擎：把它们咬合成一个「第二个请求到底省了什么」)
    • [10.1 它要解决什么问题](#10.1 它要解决什么问题)
    • [10.2 总体思路与原理](#10.2 总体思路与原理)
    • [10.3 落地与实测](#10.3 落地与实测)
• 附录
  • [附录 A：跨仓库核对 --- group 语义与每层状态的归属](#附录 A：跨仓库核对 — group 语义与每层状态的归属)
    • [A.1 kimi3-master（本项目镜像的来源）](#A.1 kimi3-master（本项目镜像的来源）)
    • [A.2 kimi-k3-pytorch-main（模块化参考）](#A.2 kimi-k3-pytorch-main（模块化参考）)
    • [A.3 SGLang K3 分支（生产实现）------ 前缀复用的事实真相](#A.3 SGLang K3 分支（生产实现）—— 前缀复用的事实真相)
    • [A.4 合成的结论](#A.4 合成的结论)
  • [附录B: 与 vLLM KV Cache Group 的对比](#附录B: 与 vLLM KV Cache Group 的对比)
    • [B.1. 项目中的 group 是什么](#B.1. 项目中的 group 是什么)
    • [B.2. vLLM 中的 KVCache Group 是什么](#B.2. vLLM 中的 KVCache Group 是什么)
    • [B.3. 逐项对比](#B.3. 逐项对比)
    • [B.4. 两者为什么容易被认为相同](#B.4. 两者为什么容易被认为相同)
    • [B.5. 最接近的对应关系](#B.5. 最接近的对应关系)
    • [B.6. 关键差异：checkpoint 原子性 vs. 物理 cache 管理](#B.6. 关键差异：checkpoint 原子性 vs. 物理 cache 管理)
      • [6.1 项目 group 解决的是恢复正确性](#6.1 项目 group 解决的是恢复正确性)
      • [6.2 vLLM group 解决的是存储和调度一致性](#6.2 vLLM group 解决的是存储和调度一致性)
      • [6.3 两者可以同时存在](#6.3 两者可以同时存在)
    • [B.7. 最终判断](#B.7. 最终判断)
  • [0xFF 参考](#0xFF 参考)

---

0x00 概要

0.1 主要内容

这份文档写给谁 ：给想理解「Kimi Delta Attention（KDA）这种新注意力，到底把整个 serving（推理服务）体系颠覆成了什么样」的人。假定你知道一些 Transformer，但完全不懂 KDA 的服务端细节------我们从最朴素的问题出发，一路讲到每一层基础设施为什么要那么设计。

serving 相关的主要组件如下：

• 页池把「两类异构缓存」归一成同一种「页 + kind 标签」；
• 前缀缓存只信「hash 块的整数边界 + 每个 group 都有 checkpoint」（组 = "必须一起存在"的抽象分区，默认整模型 1 组）；
• 投机不是「删 KV」而是「不存状态存投影、失败时重演」（或「存快照、失败时 O(1) 取最旧」）；
• KCP 用「每段分解 (M̄,S̃) + 一次固定 all-gather + 前缀扫描」击穿「状态不可加」；
• Fleet 层用「亲和路由 + 一致哈希双绑定 + 预算隔离」把命中率和故障、长短流都管起来；
• 收官引擎把它们串成一个「第二个请求到底省了什么」的 before/after。

这些"因固定状态而生的 Infra 挑战"在 KCP、投机解码、前缀缓存等相关章节展开。

0.2 注意说明

• 项目 https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
• 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
• 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
• 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
• 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
• 每个主要组件都有可直接运行的 Gradio demo------无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
• 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
• 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。

0x01 先懂 KDA：为什么它会颠覆 serving

1.1 从最朴素的「KV 缓存」说起

标准的 Transformer 注意力（softmax attention）在服务阶段有一个叫 KV cache 的东西。它的作用是：记住已经看过的 token，这样生成下一个字时，不用把整个句子重新算一遍，只要用「记住的」K 和 V 去算就行。

KV cache 有一个朴素的特性：它是一张「逐 token 的清单」。每个位置一个条目，随句子变长而变长。正因为它是一张清单，你可以：

• 按 token 索引它（第 3 个位置是什么）；
• 把它截断（不要后面几个）；
• 把它拼接（两份清单接起来）；
• 把它哈希复用（两条请求共享同一段前缀，直接抄）。

这一整套 serving 基础设施（前缀缓存、内存池、上下文并行、投机解码的回滚）------全都是建立在「KV 是一张可以随便索引、截断、拼接的清单」这个前提上的。

1.2 KDA 把这张清单换成了一个「状态矩阵」

Kimi Delta Attention（KDA）是 Kimi K3 混合架构（3 份 KDA + 1 份 MLA）里的主力。它不存逐 token 的 KV 清单，而是维护一个固定尺寸的状态矩阵 S，每读一个 token 就原地更新一次：

S ← M · S + β · k · vᵀ

其中：

M = (I − β k kᵀ) · Diag(α)

这里 k, v 是当前 token 的 key 和 value，α 是逐通道衰减，β 是门控。这就是论文 Eq.1 的 delta rule（δ 规则）。

为什么是固定状态而不是 KV cache ：因为核心卖点是"1M 上下文可负担"。KV cache 随 T 线性增长，KDA 状态固定 d_k×d_v/head。但固定状态带来三方面代价，也是我们整个文档反复出现的主题：

1. 检索是压缩的（近似）→ 需要 MLA 混合。
2. 状态跨段不可加（转移矩阵非可交换）→ 需要 KCP 前缀扫描。
3. 状态无法简单截断回滚（整体快照）→ 需要 clone()/ReplaySSM。

State 与 KV Cache 的硬件类比

• Linear state 像 register file: 固定大小；访问快；反复更新；容量有限。
• Full/Sparse KV 像 cache: 地址明确；容量更大；读写更依赖带宽和缓存命中；可以共享 prefix。
• 模型架构和芯片内存系统因此是共同设计问题。

1.3 关键：M 是「逐 token 相关」的矩阵

这是整个故事的源头 。注意 M 里有一个 − β k kᵀ，它依赖当前的 k。也就是说：

每次状态更新，不是简单「往清单末尾加一行」，而是先拿一个跟当前 token 有关的矩阵，把旧状态整体缩放一遍，再加一条新记录。

这带来几个让 serving 工程师头疼的「天然属性」，正是它们，把上面那套「建立在清单前提上」的基础设施全部作废：

#	状态的天然属性	它把哪个「清单前提」击穿了	需要的新基础设施
A	固定尺寸 （S 不随序列变长，每请求一份）	内存池假设「KV 随 T 增长、逐 token 分页」	统一页池：让两种缓存共用一套内存
B	整体不可分割（状态是一团，不能按 token 截断/拼接/哈希）	前缀缓存假设「前缀可逐 token 哈希复用」	双粒度前缀缓存 + KDA checkpoint
C	原地更新不可回滚（没有「删最后几个」这种操作）	投机解码假设「验证失败可以回滚 KV」	投影重放 / ReplaySSM 快照回滚
C'	跨段不可加（每段都先被各自的 M 缩放）	上下文并行假设「各段自状态直接求和」	KCP 前缀扫描

这份文档，就是依次解决 A、B、C、C'，再加上把这些组件调度起来的 fleet 层。

通俗解释：普通注意力让服务器「记账本」，KDA 让服务器「算一个不断被改写的滚动数字」。以前所有「翻到第几页」的技巧都不能用了，你得发明新的。

---

1.4 释疑

Q1：KDA 的"状态"到底包含哪几部分？decode 时每一步都在更新什么？

KDA 每一层的解码状态（在本项目里由 KDACache 承载）其实不只一块状态矩阵，而是由四块东西拼起来的：

1. state ：(B, H, dk, dv) 的递推矩阵 S------这是 KDA 与 MLA 最本质的差别，delta rule 的核心，随每步 S ← M·S + β·k·vᵀ 原地更新。
2. conv_q / conv_k / conv_v ：三条 ShortConv 的滑动窗口------每条形状 (B, k_dim/v_dim, kernel-1)（kernel 通常 = 4），存着"前 kernel-1 个 token 的 q/k/v 投影输入"，每步把当前 token 推进窗口、最老那个滑出去。ShortConv 给无位置编码的 KDA 注入 token 级"位移"信息，是位置编码的替代。

所以完整地说，KDA 的"cache"是一份状态矩阵 + 三个卷积窗口。它们每 token 都在推进，都是 decode 不能丢的状态。

Q2：为什么"KDA 状态会随着 decode 继续更新"？

这话听起来是废话，但它恰好是 KDA 颠覆整个 serving 栈的根因：

• 对 MLA（Gated MLA / softmax 注意力），decode 时 KV cache 的动作是追加 ------每生成一个 token，往末尾追加一行新的 K/V。前面已有的 KV 原封不动。所以"回滚到第 n 步" = 把末尾若干行删掉，"截断/拼接/哈希复用" = 对清单做列表操作。
• 对 KDA，decode 时状态的动作是原地改写 ------S ← M·S + β·k·vᵀ，每走一步，整张矩阵的每个元素都被缩放 + 叠加新记录 。前面"曾经的"状态没有任何痕迹残留。所以不存在"删最后几行"这种操作，不存在"按 token 索引旧状态"这种能力，也不存在"直接拼两段状态"这种操作。

这条性质往下传导：前缀缓存不能做逐 token 哈希（只能在稀疏边界快照整块状态）、投机解码失败不能直接回滚（只能快照或重演）、上下文并行不能直接相加（必须前缀扫描重组）、内存池不能再按 token 分页（得把状态当整块来管）------这篇文档后面所有组件，全都是从这个"状态是原地更新的整体"这一条性质长出来的。

0x02 总纲：三个天然属性如何变成三层 Infra

这一节把全文的因果链先立起来。下面的表 ------每个基础设施层 = 状态的一个「不适配」 + 一个「工程解法」：

基础设施层	softmax/KV 的假设	KDA 状态的不适配	工程解法
内存池	KV 逐 token 分页、随 T 增长	S 固定大小、每请求单份；且与 MLA KV 异构	统一页池：byte 对齐，KDA 页 + MLA 页共用一套 alloc/evict
前缀缓存	前缀可逐 token 哈希复用	状态是整体，只有稀疏边界可快照	双粒度解耦：512-token hash 块（MLA 细粒度）+ KDA checkpoint
上下文并行	各段 KV 可直接拼接	状态跨段不可加（M 逐 token 相关）	KCP：每段分解 (M̄, S̃)，一次固定 all-gather + 前缀扫描
decode/投机	KV 可逐 token 截断回滚	S 原地更新无法回溯	缓存投影而非状态；失败时重放
集群调度	传统负载均衡	前缀缓存 miss 比 hit 贵几个数量级	Fleet 层：亲和路由 + 预算隔离

一句话总纲 ：KDA 的 serving 设计，本质是把状态 S 的三种天然属性（固定 / 整体 / 原地更新）系统性转换成对应的基础设施机制 。没有哪一层是「顺带就能支持」的------每一层都要针对性改造。这就是 K3 报告把 算法-系统协同设计（algorithm--system co-design） 列为关键基础设施创新的原因。

2.1 为什么要做这一套

K3 = 3 份 KDA + 1 份 MLA 的混合架构（每个 block 内 3 层 KDA + 1 层 Gated MLA）。任何 serving 引擎都必须同时管理两种根本不同的缓存 ：MLA 的 KV（逐 token、随 T 增长）和 KDA 的状态 S（固定、每请求一份）。它们物理性质完全不同，却要共享同一台机器的显存/内存、走同一套分配/淘汰/传输机制。这是 K3 与所有「纯 softmax」模型最大的工程差异。

2.2 范围与边界（本项目做了什么 / 刻意没做什么）

做了（6 个组件，全部 CPU 可运行、可测试、有 before/after 可视化）：

1. 统一页池（models/llm/kimi_k3/serve/paged_cache.py）
2. 双粒度前缀缓存（models/llm/kimi_k3/serve/prefix_cache.py）
3. 投机解码·投影重放（models/llm/kimi_k3/serve/decode.py）
4. 投机解码·ReplaySSM 环形缓冲（models/llm/kimi_k3/serve/replayssm.py）
5. KDA 上下文并行 KCP（models/llm/kimi_k3/parallel/kcp.py）
6. Fleet 层调度（models/llm/kimi_k3/serve/fleet.py）

刻意没做：

• 不做真实多机通信 / TP / DP 调度（KCP 用单进程分段模拟 world 个 rank）；
• 不做 Triton/GPU 内核分发；
• 不做真实 1M token 计算（代价用公式折算成时间）；
• 不做 WebSocket 实时前端。

为什么这些「没做」仍值得写？因为原理正确性 与工程形态是两回事。我们要证明的是「这套设计在数学上成立、在代码上可复现」，而不是「能跑到多快」。CPU 上能逐位对拍出正确性，就已经把设计讲明白了。

2.3 术语

术语	含义
S（状态）	KDA 的递归状态矩阵 (B,H,dk,dv)，固定尺寸，每 token 原地更新
M（过渡矩阵）	(I − β k kᵀ)Diag(α)，逐 token 相关，先缩放旧状态再加新记录
KDA group	一组 KDA 层的集合（组件二详解：为什么是「组」而不是「层」）
checkpoint	在某个稀疏边界上把当时的 S 快照下来，用于跳过递推
投影（projection）	投机的每 draft token 的 (q,k,v,g,β) 输入（远小于整个状态）
replay / 重演	用已接受的投影，从已提交状态重新跑 kda_step，得到正确状态
ReplaySSM	SGLang 的快照路线：环形缓冲存最近若干状态快照，回滚时直接取
KCP	KDA 上下文并行：把序列分段，用前缀扫描重组各段入态
fleet	一组 serving 集群（默认 8 个）
P95 TTFT	time-to-first-token（首个 token 的延迟）的 95 分位

---

0x03 实现全貌

3.1 总览图

先给一张总览图。

3.2 组件数据流

我们再看看「这六个东西为什么放在一起、数据怎么流」：

依赖方向（单向，不反向 import）：

paged_cache.py（基础，无依赖）
    ↑
prefix_cache.py（依赖 paged_cache 的物理块 id）
    ↑
decode.py（投影重放）  ← 与 replayssm.py（快照）互为对照的两种策略
    ↑
kcp.py（上下文并行，独立，依赖 kda_ref 的 kda_chunkwise）
    ↑
fleet.py（调度，复用 prefix_cache + paged_cache）
    ↑
engine.py / demo_engine.py（收官汇总，依赖全部）

3.3 组件快速浏览

Q1：这个组件在解决什么「真问题」？（不是「实现了什么功能」，而是「不解决会怎样」）

• 页池：两种缓存怎么共用一块内存？（不解决 → 内存割裂浪费）
• 前缀缓存：状态不可分割，怎么复用前缀？（不解决 → 每个请求重算 400K）
• 投影重放 / ReplaySSM：投机失败怎么恢复状态？（不解决 → 投机不成立，或快照成本吃掉收益）
• KCP：状态不可加，怎么并行？（不解决 → 长序列没法跨设备并行）
• Fleet：命中率与故障怎么两全？（不解决 → miss 贵几个数量级 / 集群一挂全断）

Q2：这个组件的「核心机制」一句话是什么？

• 页池：字节层归一化 + 一个 free list + kind 标签
• 前缀缓存：双粒度（MLA hash 块 + KDA checkpoint）+ 两段式 lookup + group 原子性
• 投影重放：不存状态存投影，失败时重演
• ReplaySSM：环形缓冲存快照，失败时 O(1) 取最旧
• KCP：每段分解 (M̄,S̃) + 一次固定 all-gather + 前缀扫描
• Fleet：亲和路由 + 一致哈希双绑定 + 预算隔离

Q3：这个组件的「边界 / 何时不用」是什么？（识别滥用场景）

• 页池：纯 softmax / 纯 KDA 时退化成单种页
• 前缀缓存：无共享前缀时是纯开销
• 投影重放 vs ReplaySSM：状态宽/长 draft 选重演，状态窄/短 draft 选快照
• KCP：短序列时通信是浪费
• Fleet：单实例 / 无共享前缀时退化

下面从底层往上，一章一个组件。每个组件都按同一套模板讲：① 它要解决什么问题 → ② 难在哪 → ③ 总体思路 → ④ 背后的原理 → ⑤ 落地架构（类/函数/数据结构）→ ⑥ 数据流/执行流 → ⑦ 反方辩手（自我质疑）→ ⑧ 何时用/何时不用。

---

0x04 组件一：统一页池

K3 把缓存看成 MLA KV + KDA recurrent/convolution state 的联合对象；以共同前缀边界保存 checkpoint；区分稳定共享副本和请求活动副本；在缓存预算内选择值得保存的恢复点。核心是把缓存键从 "token prefix" 升级为包含模型版本、tokenizer、前缀边界和多种状态的联合键。

4.1 它要解决什么问题

K3 一个请求要同时存两种完全不同的东西：

• MLA 的 KV：随序列变长，像一条会增长的清单；
• KDA 的 S：固定尺寸，每请求一份（每个 KDA 层一份）。

它们占同一块显存/内存，但「长得」完全不一样。传统 serving 引擎会为每种缓存各开一块内存（一份 KV 池、一份状态池）。

问题：两块池子边界固定，浪费（某时 KV 挤爆、状态闲置），而且每次 load/store/evict 要分别处理，代码重复。

4.2 难在哪

难在**「异构」**：一个按 token 增长、一个固定尺寸；一个内容是 latent_kv、一个是状态矩阵。你怎么让两个「性格完全不同」的东西，共用一套内存管理代码，还不互相踩踏？

4.3 总体思路

在字节层面把它们「归一化」成同一种东西------页（page） 。把所有缓存切成固定字节宽度的页，页上贴一个 kind 标签（是 MLA 页还是 KDA 页），然后用同一个 free list 管理所有页的分配/释放/引用计数。

背后的原理如下：虽然 MLA 页和 KDA 页内容结构完全不同，但在「分配/释放/引用计数」这个操作层面，它们是一样的 ------都是一块固定宽度的内存。只要在页上标清 kind，上层只需要对着一个 PagePool 对象操作，永远不用关心「这块是 MLA 还是 KDA」。这就是「统一」的含义。

这个「标签」还带来了一个免费的 sanity check ：两种页内容结构完全不同，任何「类型混淆」的访问都会产生明显的垃圾数据，而不是「看起来正常的数据」------零成本的布局正确性自检。

4.4 落地架构（类 / 函数 / 数据结构）

models/llm/kimi_k3/serve/paged_cache.py 里三个关键角色：

PageKind (str Enum)          # 页的类型标签：MLA_KV | KDA_STATE
    │
Page (@dataclass)            # 一页物理内存
    │   ident     全局唯一 id（分配时从 free list 弹，释放时还回去）
    │   kind      页类型（决定内容是 KV 还是状态）
    │   data      固定字节宽的扁平 Tensor
    │   refcount  引用计数（共享页 / 前缀复用）；归 0 的那一刻回到 free list
    │   filled    已用的元素数（load 按它切「真数据」和「补齐空位」）
    │
PagePool (@dataclass)        # 管理所有页：分配 / 释放 / 保留 / 利用率
        free   ------ 一个共享的 id 列表，两种页共用同一个分配池
        allocate(kind)  → 从 free 弹 id，refcount=1
        release(ident)  → refcount 减到 0 则删页、id 还回 free
        retain(ident)   → 显式 refcount+1（共享页）
        utilization()   → 1 - 空闲页/总页

配套的两个「布局转换」函数：

• pack_heads(state)：把 (B,H,dk,dv) 的状态压平成一维字节流 ，且每个 head 的字节保持连续。
• unpack_heads(flat, heads, dk, dv)：反向还原；只切 batch*H*dk*dv 个元素，把 padding 的零丢弃，保证补位值永远不进重建的状态。

PagedCache 是「请求级」的视图，不独立持内存，而是指向共享的 PagePool：

• mla_pages: list[int] ------ 本请求按 token 顺序的 MLA 页 id；
• kda_pages: dict[int,int] ------ 层号 → KDA 页 id（每层一页）；
• append_mla(n_tokens)：PagedAttention 语义，当前尾页满了就新开一页，逐 token 填充；
• store_kda(layer, state)：一层一页，重复写同层复用同一页（覆盖）；
• load_kda(layer, ...)：按 filled 切片读回，unpack_heads 还原；
• release()：把所有页还回池。

常见疑问：ShortConv 的卷积窗口（conv_q/conv_k/conv_v）也存在页池里吗？

目前没有 。本项目的 PageKind 枚举只有两值：MLA_KV 和 KDA_STATE。store_kda / load_kda 只存/取递推矩阵 S，conv_q / conv_k / conv_v 三条滑动窗口不在 paged_cache 里管理，也不在 prefix_cache 的 checkpoint 里。

这是因为目前的 serving 实现偏"设计验证"层面，重点放在状态矩阵 S 的管理（这是 KDA 与 MLA 最本质的差别、也是最颠覆 serving 的地方）。卷积窗口只是「宽度 kernel-1 的小尾巴」，字节开销远小于 S（3 × k_dim × (kernel-1) vs H × dk × dv），工程上可以先简化处理。如果要做完整的生产级 serving，需要把三条 conv 状态也纳入页池 / checkpoint 原子性------它和 S 一样，是"必须一起恢复"的状态，只是小很多。

4.5 数据流 / 执行流

一次 prefill：
  请求进来 → view1 = PagedCache(pool, tokens_per_page=512)
           → view1.append_mla(512)               # MLA：512 token 铺进页（满则新页）
           → 每层 KDA: view1.store_kda(layer, S)  # KDA：一层一页，固定尺寸
  请求结束 → view1.release()                       # 页还回池，id 回 free list

append_mla 的细节（为什么 is_full 是关键）：PagedAttention 的核心是「KV 按 token 分页、只在页满时分配新页」。Page.is_full(tokens_per_page) 检查当前尾页是否装满，装满才 allocate 下一页------这是 MLA 侧「随 T 增长」的落地。

4.6 反方辩手（自我质疑）

「MHA 也能做成池，分页不稀罕。」 对，分页本身不稀罕。但关键是：这个池必须同时装得下「随 T 增长的 token」和「尺寸对 T 不敏感的整块状态」 ------这是 MLA+KDA 的独特混合，普通的 KV 桶做不到。PageKind 的双值枚举正是「两种物理性质共存于一个分配池」的体现。

「为什么 pack_heads 要保持 head 连续，直接 flatten 不行吗？」 因为 head 是跨节点传输的最小语义单元。未来的 KCP（上下文并行）按 head 拆分通信；如果一个 head 的字节被切到不同页/不同段，传输时就要额外拼装。保持 head 连续 = 每个 head 的字节流「自包含」，可独立搬运。

4.7 何时用 / 何时不用

• 何时用：任何要同时管理「随 T 增长的 KV」和「固定尺寸状态」的 K3 混合 serving。这是底层，上面所有组件都建立在它之上。
• 何时不用 ：纯 softmax 模型（没有 KDA 状态）不需要 KDA_STATE 页；纯 KDA（没有 MLA）不需要 MLA_KV 页。PageKind 的设计让这两种情况退化成「只用一种页」，依然能用同一个池。

4.8 可视化：

plot_paged_cache.py --- 统一页池 (serve/paged_cache.py)

• 表达：MLA 与 KDA 共享同一个 free list 页池，但占页规律完全不同
• 测试：PagePool + PagedCache.append_mla(L)，扫描 L=512...65536，统计 MLA 页数（log-log 斜率为 1 的直线）；对照 KDA 一层递归状态 S(B,H,dk,dv) 恒定 1 页（水平参考线）
• 产出：plot_paged_cache.png

plot_paged_cache.py 产出 plot_paged_cache.png（对比 MLA 页数随 seq_len 增长 vs KDA 页数固定）。

---

0x05 组件二：双粒度前缀缓存

5.1 它要解决什么问题

前缀缓存 （prompt/prefix cache）解决的问题是：很多请求共享同一段 prompt 前缀 （比如系统提示词、few-shot 示例、长文档开头）。如果第一个请求算过这段前缀，第二个请求能不能直接复用，而不从头重算？

对普通 KV：能。因为 KV 是「逐 token 的清单」，只要哈希比对前缀一致，就能把整段 KV 抄过来。

对 KDA 呢？

5.2 难在哪

KDA 的状态 S 是一个递归算出来的整体，不是一张可以逐 token 索引的清单。它有两条硬约束：

1. S 不是按 token 可取的 ：要「跳过硬算」一段前缀，唯一的办法是在某个边界 把当时的中间状态快照 下来（checkpoint）。但 checkpoint 只能在稀疏边界存------不是每个位置都适合存。
2. checkpoint 必须「原子」存在 ：状态是逐层递推 的。一个 K3 block 有 3 层 KDA + 1 层 MLA。要在边界 B 处恢复状态，每一层都必须在那存了 checkpoint。少一层，就无法正确恢复。

这里就有了「组（group）」的概念。

5.3 总体思路

双粒度解耦 ：把「能安全复用的前缀」定义成两个条件的交集：

MLA 侧逐块匹配 （细粒度：512-token hash 块）+ KDA 侧每个 group 都有 checkpoint（在某个稀疏边界）。

于是 lookup 分两阶段：

• 阶段一：按 chained hash 逐块匹配 MLA 侧，能走多远走多远；
• 阶段二 ：从最深匹配边界往回退，找到一个「每个 KDA group 都存了 checkpoint」的边界，才真正命中。

命中永远是 hash block 的端点 ，不要求是 physical block 的整数倍。

5.4 背后的原理

（1）chained hash 而非每块独立 hash ：每一块的 digest 认证「直到该端点的整个前缀 」。两个请求只要在某端点 hash 相同，就保证其前面全部 token 完全一致------不存在「前缀撞车」的歧义。命中 = 内容完全一致，无需再逐 token 比一遍。

（2）为什么 KDA 只能复用「边界上的整块状态」 ：状态更新 S ← M·S + β·k·vᵀ 是串行递推。要恢复 S，要么从头一路递推到目标位置，要么从一个先前存好的 checkpoint 接上继续递推。中途任何位置都没有「现成的 S」可用------所以 checkpoint 只能在「我们主动存的那些边界」存在。

（3）ready 标志：为什么注册是异步的 ：先登记 digest（ready=False）、后填充内容（mark_ready 置 True）。ready=False 的块参与索引但不参与 lookup------防止命中「登记过但还没填好」的假前缀。

（4）KDA Prefix Cache 为什么常在请求结尾保存？

KDA state 是递推结果，只有处理完某个前缀位置后才得到该位置的完整状态。保存请求结尾快照最简单：快照自治；恢复后可继续递推；不必保存每个 token 的完整中间状态。如果要支持任意中间分支，就需要更多快照或从最近快照重放。

5.5 关键：group 到底是用来做啥的？

先澄清一个容易误解的地方 ：group 不是「一层 KDA」的别名 ，它也不由代码强制绑定到任何具体层。在 kimi3-master（本项目对照的参考实现）里，n_kda_groups 是一个抽象的"必须共存的分区数" ------意思是"一个 checkpoint 必须同时在几个 group 里都存在，这个边界才算可用"。默认 n_kda_groups=1（整个模型当 1 个 group）。

参考实现里可以这样用：

make_cache(n_kda_groups=3)                        # 3 个抽象 group
cache.persist_kda_checkpoint(512, groups=[0, 1])  # 只给 group 0、1 存 → 缺 group 2，暂不可用
cache.persist_kda_checkpoint(512, groups=[2])     # 补上 group 2 → 才可命中

也就是说：group 是"同生共死"的那一组对象的抽象容器。它可以是"一层的状态"、可以是"一个 block 里 3 层 KDA 的状态"、也可以是"任意一组你必须一起恢复的状态"。代码不管它具体指啥------只负责"每个 group 都存了 checkpoint 才算数"。

为什么需要"必须共存"这个约束 ：K3 的递归状态是逐层 的（每层有自己的 S_l）。要恢复"整个 block"在边界 B 的状态以继续算，每一层 在 B 的 checkpoint 都得在------少一层，算出的下一 token 就是错的。所以"一起存在才算数"是硬约束。group 就是把"哪几个东西必须一起存在"打包成一个抽象单元的机制。

5.5.1 先看清 K3 的层排布（真实 layer_schedule()）

def layer_schedule(self):
    body = self.n_layers - 1
    period = 4
    schedule = []
    for i in range(body):
        schedule.append("mla" if (i + 1) % period == 0 else "kda")
    schedule.append("mla")   # 尾部再挂一个全局 MLA
    return schedule

实测（脚本输出）：

• 16 层示例排布：[kda,kda,kda,mla] × 4 → 12 KDA + 4 MLA（3:1）；
• 真实 K3 规模：93 层 = 69 KDA + 24 MLA（3:1，含尾部全局 MLA）。

所以"层"是明确的：每个 block 3 层 KDA + 1 层 MLA，每个 KDA 层自己持有状态 S(B,H,dk,dv)。

5.5.2 group 是"层"吗？------ 不是，它是"对状态的分组维度"

• group ≠ Transformer 层 。group 是 prefix_cache 里对 KDA 状态如何打包成原子组 的分组维度，用 n_kda_groups 配置。
• 每一层不一定有一个 group 。两种极端：
  • n_kda_groups=1（默认）：全部 KDA 状态当一个整体 blob ，kda_checkpoints = {0: set}------原子判定只查这一组。这正是 SGLang 的做法；
  • n_kda_groups = n_kda_layers：每个 KDA 层单独一组，kda_checkpoints = {0:set, 1:set, 2:set, ...}------原子判定要求每个 KDA 层在该边界都有快照。
• 无论哪种粒度，原子性的语义不变：恢复边界必须所有组齐；只是"组"的粒度可由服务方选择。

5.5.3 为什么"层数"和"组数"分开？（设计动机）

把"层"和"组"解耦是刻意的：

1. 真实 K3 是 69 个 KDA 层 ，若强制每层一组，原子判定要遍历 69 个集合，且任何一层缺快照都会让边界作废------恢复约束过强；
2. 但 KDA 状态在语义上是"逐层递推的整体"，服务方常选择整组原子 （n_kda_groups=1），既能正确恢复又不会因为"某一层丢了"而全盘失效；
3. 中间粒度（如每 block 一组）也允许------这正是"分组维度"而不是"层数"被抽象成参数的原因。

一句话：group 不是「层」的别名，而是「同生共死」的那一组对象的抽象容器------它把"哪几个东西必须一起存在"打包成编号，checkpoint 必须"每个 group 都有"才算数，少一个就整个边界作废。默认整个模型是 1 个 group；只有当你显式把"必须一起恢复的层集合"设成一个 group 时，才和"层"发生联系。它把「状态是一个整体」这个数学事实，翻译成了「checkpoint 必须原子存在」这个工程约束。

5.5.4 图：group vs 层示意

5.6 checkpoint 的原子性与「all()」的含义

checkpoint 的原子性，用代码表达就是：

def checkpoint_available(self, boundary: int) -> bool:
    """边界被每个 group 都有 checkpoint 才算可用。

    这是「原子性」的查询面：KDA 递推状态必须整组一起恢复，
    只要有一个 group 缺快照，这个边界就不可复用。
    """
    return all(boundary in group for group in self.kda_checkpoints.values())

即：边界 B 被每个 group 都有 checkpoint，才算可用；只要有一个 group 缺，这个边界就不可复用。

上面代码逐字拆解：

• self.kda_checkpoints 是一个 dict：{KDA group 编号 → 该 group 已存 checkpoint 的边界位置集合}。
• self.kda_checkpoints.values() 取出所有 group 的集合（n_kda_groups 个 set）。
• boundary in group：判断某个边界位置 boundary 是否在这个 group 的 checkpoint 集合里。
• all(...)：要求每个 group 的集合里都包含 boundary，任何一个不包含就返回 False。

为啥要 all（原子性） ： KDA 的递归状态 S(B,H,dk,dv) 是逐层递推 的整体------它按 group（KDA 混合架构里多个并行的状态组）一起递推，各层的 checkpoint 必须同时存在才能恢复。少任何一层（group）都无法正确恢复前缀状态。所以：

• 注册时 persist_kda_checkpoint 往每个 group 都登记；
• 查询时 checkpoint_available 要求所有 group 都有（少一个就算"半原子状态"，不可复用）。

all 是"与"语义------它是这个原子性约束的查询面。如果只用 any（任一 group 有），就可能拿到缺层的 checkpoint，前缀恢复出来是错的。正是 all 保证了命中前缀的 KDA 状态完整可恢复。

5.6.1 图：checkpoint 原子性示意

左图用真实 PrefixCache(n_kda_groups=3) 构造了一个故意不齐的 checkpoint 矩阵：

• 512 ：三个 group 都有快照 → checkpoint_available=True（可用）；
• 1024 ：只有 group 0、1 有，缺 group 2 → False；
• 1536 ：只有 group 0、2 有，缺 group 1 → False。

绿色"可用"字样标出了 all() 通过的两个边界，红色"不可用"标出了缺一组的边界。

右图是两场景的 lookup 结果：

• 场景 A （2048 全组齐）：阶段一 MLA 逐 hash 块匹配到 2048，阶段二 all() 校验通过 → 命中 2048；
• 场景 B （2048 抽走 group 2）：阶段一 MLA 依然能匹配到 2048（MLA 侧不受影响），但阶段二 all() 在 2048 失败 → 回退到 512。

实测数值（CPU 可复现）：

boundary=  512  checkpoint_available(all groups) = True
boundary= 1024  checkpoint_available(all groups) = False
boundary= 1536  checkpoint_available(all groups) = False
boundary= 2048  checkpoint_available(all groups) = True
[场景A] lookup matched_tokens=2048 kda_checkpoint=2048
[场景B] lookup matched_tokens= 512 kda_checkpoint= 512

一句话结论 ：all() 锁的是恢复的原子性：KDA 状态必须"每个 group 都有 checkpoint"才能在某个边界恢复，缺一个组该边界就作废、lookup 会回退到更浅的可用边界（实测场景 B 从 2048 回退到 512）。

5.7 落地架构与数据结构

HashBlock (@dataclass)         # 一个 hash 块的索引条目
    digest            链式 hash，索引 key（也是 dict 的 key）
    end_token         本块覆盖到的 token（含，块尾端点，必为 512 的整数倍）
    physical_block    该端点落在页池的哪个物理块
    ready             内容是否已填充（异步登记的 False→True）

LookupResult (@dataclass)      # 一次 lookup 的返回值
    matched_tokens     MLA 侧匹配到的端点 token 数
    physical_blocks    命中的物理块 id（已去重，供上层 pin/共享）
    kda_checkpoint     命中的 KDA checkpoint 边界（None 表示没命中）
    @property hit      matched_tokens > 0

PrefixCache (@dataclass)       # 双粒度缓存主体
    hash_block_tokens=512
    physical_block_tokens=6144 (=12 个 hash block)
    n_kda_groups                KDA group 数（默认 1 = 整个模型 1 个 group）
    index: dict[str,HashBlock]  digest → 块 的倒排索引
    kda_checkpoints: dict[int,set[int]]  group → 边界集合
    free_list / _next_block     物理块分配

关键方法：

• chained_hash(previous, tokens)：上块 digest + 本块 token → sha256，链式认证整个前缀。
• register(tokens)：把请求前缀分成 hash 块登记进索引，返回各块 digest。
• mark_ready(digest)：异步登记后内容填充完，置 ready。
• persist_kda_checkpoint(boundary, groups)：在某 hash 块边界为指定（或全部）group 存 checkpoint；非 512 整数倍直接 raise，防止错误对齐的 checkpoint 污染索引。
• evict_kda_checkpoint(boundary, groups)：从指定 group 移除（discard 幂等），用于部分缺少时的状态补齐或逐出。
• checkpoint_available(boundary)：all(group 有 boundary)，原子性查询面。
• lookup(tokens)：两段式。
• recompute_saved(result, total_tokens)：报告「本次真正免于重算的 token 数」，上限取 min(matched, total)。

5.8 数据流 / 执行流（lookup 的两段式）

lookup(tokens):
   阶段一 (MLA 逐块)：
    previous = ""
    for 每个 512-token 块:
        previous = chained_hash(previous, chunk)
        entry = index.get(previous)
        if entry is None or not entry.ready:  break   # 不连续 → 停
        candidates.append((entry.end_token, entry.physical_block))
        # ↑ 前缀命中必须是「连续」的，中断的后继即使 hash 在索引里也不算

  阶段二 (KDA 原子性回退)：
    for (end, block) in reversed(candidates):          # 从最深往回退
        if checkpoint_available(end):                  # 每个 group 都有 checkpoint?
            return LookupResult(end, 去重后的物理块, end)
    return LookupResult(0, [], None)                   # 一个都没命中

为什么「命中必为 hash block 端点，却不必是 physical block 倍数」：MLA 能逐 512 块匹配，所以命中在 512 的整数倍上；而 KDA checkpoint 只要求「在每个 group 都有」这个条件满足的边界------它落在某个 hash block 端点即可，不强制是 6144（physical block）的整数倍。这就是「双粒度」：MLA 的细粒度（512）与 KDA 的稀疏边界（checkpoint）在 lookup 时取「兼容的交点」。

具体 lookup 完整代码如下：

def lookup(self, tokens: list[int]) -> LookupResult:
    """两段式 lookup：MLA 一键逐块 + KDA 回退到有 checkpoint 的边界。

    阶段一（MLA）：
        从 t=0 开始按 hash_block_tokens 连续计算 chained hash，
        每个 digest 去 index 查。只要块在且 ready，就继续；
        一旦 miss 或未就绪就立刻停止------前缀命中必须是**连续**的。
        不连续的后继 token 不属于该前缀，即使它的 hash 恰好在索引里。

    阶段二（KDA 原子性约束）：
        MLA 匹配得多深并不直接决定复用有多长------还必须「最深匹配边界
        vs 每个 group 是否有 checkpoint」。从最深匹配边界往回退，
        找到 endpoint 处 checkpoint_available 而又能覆盖前面的。
        这样命中的 matched_tokens 一定落在 hash block 端点上，
        不必是 physical block 的整数倍（6144 的倍率线）。

        同时收集「该前缀实际用到的物理块」------阶段 1 产出的候选块
        集去重，供上层执行「共享这些物理块」的动作。
    """
    previous = ""
    candidates: list[tuple[int, int]] = []  # (end_token, physical_block)

    # 阶段 1：连续哈希块（MLA 细粒度匹配）
    for i in range(0, len(tokens) - self.hash_block_tokens + 1, self.hash_block_tokens):
        chunk = tokens[i : i + self.hash_block_tokens]
        previous = chained_hash(previous, chunk)
        entry = self.index.get(previous)
        if entry is None or not entry.ready:
            break
        candidates.append((entry.end_token, entry.physical_block))

    # 阶段 2 ：从最深边界回退，找每个 KDA group 都有 checkpoint 的边界
    for end, _block in reversed(candidates):
        if self.checkpoint_available(end):
            # 去重收集命中物理块：多个 hash block 共享同一物理块时只看一次
            seen: list[int] = []
            for boundary, block in candidates:
                if boundary <= end and block not in seen:
                    seen.append(block)
            return LookupResult(end, seen, end)
    return LookupResult(0, [], None)

5.9 反方辩手

「Q1：physical_block / 6144 是做什么的？」 它是页池共享单元。上层（页池）需要知道「这个前缀命中了哪些物理块，我要 pin 住、让它们能被共享」。6144 = 12 个 hash block 拼一个大物理块，是页池层 share/pin 的粒度。

「Q2：一个 group 就够了，为什么要支持多个 group？」 因为"必须一起存在"的对象不一定是整层模型 。真实 K3 是 3:1 混合，一个 block 里 3 层 KDA 天然构成「必须一起存在」的约束------但那是语义上 的事，代码不强制。n_kda_groups 只是"把多少组东西当作必须共存的抽象分区"：默认 1（整个模型一组）；如果你想让"每个必须一起恢复的层集合"各自算一组，就设成那个数，并在 persist_kda_checkpoint(..., groups=[...]) 时逐个写上。注意：代码不检查 group 对应哪层，这层映射留给调用方。

「Q3：chained hash 会不会有碰撞？」 sha256 的碰撞概率在实际 token 规模下可以忽略；真正的保证是「链式」------每块 digest 依赖全部前缀，所以「命中=内容一致」不需要额外比对 token，这是正确性关键。

5.10 何时用 / 何时不用

• 何时用：任何「请求共享前缀」且「需要复用 KDA 状态」的 serving。命中一次，KDA 免去整段递推。
• 何时不用：没有共享前缀的独立请求（每次都 miss，缓存是纯开销）；纯 MLA 模型（没有 KDA checkpoint 约束，退化成纯 hash 块前缀缓存）。

5.11 对比：SGLang 是如何处理 KDA prefix cache 的？

工业界（SGLang K3 分支）怎么给 KDA 做前缀缓存？它跟我们的 prefix_cache.py 设计有什么关系？这是一个常被问到的问题，我们把整条证据链摆出来。

SGLang 用的是**"每请求一整份状态 + 按 chunk 对齐的稀疏快照"**，与我们的"chained-hash 逐块 + 稀疏 checkpoint"是同一思路的两套实现。

5.11.1 图：SGLang KDA 前缀缓存机制

右图橙色阴影就是"MLA 命中但 KDA 用不上的前缀"（对齐损失），阶梯线是 KDA 真正可续算的边界。

5.11.2 证据链：SGLang 的具体机制

① 状态存放：每请求一个槽的 temporal 状态

layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_states = layer_cache.conv[0]
ssm_states = layer_cache.temporal          # (num_slots, heads, dk, dv)
...
cache_indices = self.forward_metadata.mamba_cache_indices

KDA 状态存在 mamba_cache.temporal，每请求占一个 slot ，num_slots = pool_size + 1（slot 0 保留）。一个 slot 同时装下所有层------等价于我们默认的 n_kda_groups=1。

② radix 树节点存两样东西（mamba_value）

TreeNode 同时持有：

• value：该前缀端点的 MLA KV 页索引列表（逐 token、细粒度）；
• mamba_value：该端点的 KDA/SSM 状态快照（整块、粗粒度，可能是 slot 索引、clone 或 int8 checkpoint）。

③ 复用边界的 chunk 对齐

aligned_seqlen = (result.full_kv_hit_length // self.mamba_cache_chunk_size) \
                 * self.mamba_cache_chunk_size

mamba_cache_chunk_size = max(chunk_size, page_size)（典型 chunk=64 / page=256 → 256）。KDA 只能从 chunk 整数倍边界续算------MLA 命中得再深，也要对齐截断。

④ 续算前提

has_initial_state = forward_batch.extend_prefix_lens > 0

命中后 extend kernel 直接从已存状态起步，不再重算前缀段。

⑤ 写入时机 ：decode 时在 mamba_track_interval 边界快照（_track_mamba_state_decode），extend 时在 chunk 边界抓中间态 h（_track_mamba_state_extend），保证每次到达 radix 端点时状态恰好落在该节点。

5.11.3 上右图实测（chunk 对齐量化，CPU 可复现）

mamba_cache_chunk_size=256
full_kv_hit= 512 -> aligned= 512 (wasted=  0)
full_kv_hit= 700 -> aligned= 512 (wasted=188)
full_kv_hit=1024 -> aligned=1024 (wasted=  0)
full_kv_hit=1100 -> aligned=1024 (wasted= 76)
full_kv_hit=1536 -> aligned=1536 (wasted=  0)
full_kv_hit=2048 -> aligned=2048 (wasted=  0)

5.11.4 与我们的设计对比

维度	我们的 prefix_cache.py	SGLang
前缀身份	chained-hash 逐块认证（512-token 块）	radix 树按 token 前缀精确匹配
KDA 状态粒度	可配 n_kda_groups（默认 1）	固定整块（≈n_kda_groups=1）
恢复边界约束	必须是 hash block 端点且每组有 checkpoint	必须是 mamba_cache_chunk_size 整数倍
原子性	checkpoint_available = all()	单块快照天然原子（无需 all()）
对齐损失	回退到更浅的可用边界	截断到 chunk 整数倍

一句话总结 ：SGLang 的 MLA 走逐 token 页缓存、KDA 走"每请求一槽 + chunk 对齐稀疏快照"，两者共用一棵 radix 树，但 KDA 侧不做逐 token 索引；而我们的设计用 chained-hash 认证前缀身份 + 每 group 原子 checkpoint，更偏"内容寻址"风格。本质是同一套思路的两种工程落地：MLA 细粒度、KDA 粗粒度，两者一起决定"前缀能复用多长"。

5.12 可视化

plot_prefix_cache.py --- 前缀缓存 (serve/prefix_cache.py)

• 表达：前缀缓存的收益不来自「hash 索引」，而来自 KDA checkpoint 的原子性约束
• 测试：同一个 PrefixCache 注册 512-token 前缀 + persist_kda_checkpoint(512) 后，用不同前缀查询（miss）vs 完全一致查询（hit），对比 recompute_saved 的可复用 token 数
• 产出：plot_prefix_cache.png（柱状图，0 vs 512）

plot_prefix_cache.py → plot_prefix_cache.png（耗时：no-cache vs hash vs kda 三级）。关键结论：如果 hit 带 checkpoint，0ms；否则 hash 命中仍是 MISCHANGE=$ timeline。

---

0x06 组件三：投机解码 · 投影重放（decode.py）

6.1 要解决什么问题

投机解码（speculative decoding）：用「小/快的草稿模型」先一口气生成 k 个 token，再用「大/慢的真模型」一次性验证------对的就全部接受，错的就回退到第一个错的地方重来。它能大幅加速，同时保证输出与不投机时完全一致。

投机解码有一个隐含前提：验证失败时要能「回滚」------回到「draft 还没开始」的状态，从正确位置重新生成。

对普通 KV cache：能回滚。KV 是逐 token 的清单，验证失败就「删掉最后几个 KV 条目」即可回滚。

短卷积状态比较容易回滚：线性注意力中的短卷积状态 conv_state 也保存的是一组独立的向量。对于推测解码，只需要把滑动窗口拓宽到：conv_kernel - 1 + num_spec。再把候选 Token 依次追加到窗口末尾。下一步根据 num_accepted_tokens - 1 做偏移读取即可。如果某些 draft Token 被拒绝，因为它们是独立保存的，所以不会污染前面已经确认的历史。

对 KDA 呢？ KDA 的状态 S 每 token 原地更新 （S ← M·S + β·k·vᵀ），M 逐 token 相关------没有任何「删最后几个」的操作 。draft 验证失败时，状态已经「前进」到错误位置之后，无法简单回滚。

即，KDA（Kimi Delta Attention）的递归状态 S 每步就地在 cache 里更新：

S_r = (I - beta_r k_r k_r^T) Diag(alpha_r) S_{r-1} + beta_r k_r v_r^T

所有 Token 都会原地混合进同一个矩阵：num_heads, head_dim, head_dim。一旦把候选 draft Token 更新进最终状态，之后如果发现这些 Token 被拒绝，之前的状态已经被它们污染，不能像独立 KV 那样简单删除某一条记录。因此，默认方案需要为每个请求额外分配：num_speculative_blocks 个物理块，把每个候选 Token 更新后的状态分别保存到不同槽位。验证完成后，根据接受的 Token 数量，选择对应的干净状态作为下一轮起点。

这与 Transformer 的 KV cache 有本质区别：

维度	KV cache	KDA 状态 S
结构	逐 token 独立条目	整体矩阵，每步覆盖更新
回滚到第 k 步	截断前 k 条即可	无法截断，必须知道当时的完整 S
存储	O(T) 随序列增长	O(1) 固定大小，但每步覆盖

既然KDA 的状态无法像 KV cache 那样截断，必须有一种机制来恢复任意中间点的状态。

6.2 难在哪

Delta rule 的状态递推不可加、不可截断。

Vanilla 线性注意力的状态是 S = sum k_i v_i^T，天然可加，可以拆段计算再拼起来。但 KDA 的 delta rule 每步有一个逐 token 依赖的过渡矩阵 (I - beta k k^T)，它缩放的是整个既有状态：

vanilla: S_{[i+1]} = S_{[i]} + 段内累加量        （可加，可前缀和）
KDA:     S_{[i+1]} = Stilde_{[i+1]} + M_{[i+1]} S_{[i]}  （仿射，逐段不同 M）

这意味着：

• 不能按 token 索引状态：S_5 不是 S_3 + 某个增量，而是 S_3 经过两步仿射变换后的结果
• 不能提前计算任意中间点：必须按顺序递推
• 快照成本高 ：每步保存完整 S 要 H·dk·dv·itemsize 字节，draft 越长越昂贵

两难：

• 若投机失败无法回滚 → 投机解码根本不成立（这是必须解决的）；
• 若每 draft token 都存一个状态快照 （用于回滚）→ 额外状态写入很重 ：每个快照占H·dk·dv·itemsize 字节，draft 越长、状态越宽，成本越高。具体是否抵消投机收益，仍需结合 batch、dtype 和实现 benchmark。

关键：解决「能回滚」的问题不能靠「存快照」，否则引入的成本抵消了投机本身的收益。

投机解码三连挑战：

1. 状态回滚：verify 拒绝部分 token → S 须回滚到拒绝点。KV 可截断，S 不能。
2. 分支状态：beam/多候选需独立 S → KDACache.clone()。
3. draft/verify 一致性：draft 推进了 S，verify 须从正确旧状态重算 → RawInputRing 重放轨迹。

6.3 一个推测解码例子

假设目标模型已经处理到前缀 p。上一轮采样得到了一个确定 Token x0，但目标模型还没有把 x0 作为输入执行前向。

草稿模型给出了三个候选 Token：

d1, d2, d3

因此本轮目标模型真正处理的输入是：

[x0, d1, d2, d3]

目标模型的输出用途依次是：

• 处理 x0 的输出，用来验证 d1；
• 处理 d1 的输出，用来验证 d2；
• 处理 d2 的输出，用来验证 d3；
• 最后一个位置的输出，可以用于全部候选都接受时的 bonus Token。

验证 Kernel 为每个输入位置保存一个状态快照：

状态列索引	本轮处理到的输入	保存的递推状态
0	x0	S_{p+1}
1	x0, d1	S_{p+2}
2	x0, d1, d2	S_{p+3}
3	x0, d1, d2, d3	S_{p+4}

假设验证结果是：

• d1 被接受；
• d2 被接受；
• d3 被拒绝；
• 重新采样得到修正 Token r。

本轮确认生成的 Token 是：

[d1, d2, r]

但是，目标模型真正计算并确认写入状态的输入只有：

[x0, d1, d2]

也就是三个 Token，对应的有效状态是：

S_{p+3}

r 虽然已经被采样出来，但目标模型还没有把 r 作为输入执行前向，因此不能说 r 的状态已经被缓存。

下一轮应该从 S_{p+3} 开始，并把尚未输入目标模型的 r 作为下一轮的第一个确定输入 x0'。

关键结论是：采样出一个 Token，不等于目标模型已经计算并缓存了这个 Token 的递归状态。

6.4 总体思路

不缓存状态，只缓存「投影」（projection）。

关键观察：任何「已接受前缀之后」的状态，完全由这些 draft token 的投影输入决定------而投影输入远小于状态本身。

不可行：cache 每个 draft token 的 S 快照   （每步 H·dk·dv 字节）
可行：  cache 每个 draft token 的 (q,k,v,g,β)（每步 H·(3·dk+dv+1) 字节）

投机期间只缓存这些小投影；验证时从「已接受的最长前缀」位置，用 kda_step 精确重演一遍，得到正确的状态。

6.5 背后的原理

为什么「投影决定状态」成立？因为 kda_step 是确定性的单步更新 ：给定当前状态和这一 token 的 (k,v,g,β)，下一状态唯一确定。所以：

只要我有「已提交状态」+ 「被接受的那几个 draft token 的投影」，我就能一步步精确重算出「应该的当前状态」------不需要存任何中间快照。拒绝的 token 不用算，只是「少重演几步」。

这就是「重演（replay）」而非「回滚（rollback）」：不需要显式 undo （KDA 下不存在 undo），只需带着投影从头走到「已接受长度」。这是「算法-系统协同设计」给系统设计者的第一课：把数学里不存在的操作，翻译成数学里天然存在的操作。

量化对比（snapshot_cost_ratio）：

每步快照字节  = H · dk · dv
每步投影字节  = H · (3·dk + dv + 1)     # q/k/g 各 dk 宽 + v 的 dv + β 的 1
比值          = (H·dk·dv) / (H·(3dk+dv+1))   # draft_len 分子分母抵消

比值必须代入实际 H, dk, dv 计算；dtype 和 draft_len 会在比值中抵消，但不能脱离 shape 宣称固定倍数。以本文 H=2, dk=dv=8 为例，比值为128/66 ≈ 1.94。draft 越长，绝对字节差随 draft 线性增长。

6.6 落地架构与数据结构

DraftProjections (@dataclass)      # draft 窗口内逐 token 的投影
    q,k,v,g,beta                    # 每个都是 (draft, B, H, D) 张量
    __len__()   → draft 长度
    nbytes()    → 已缓存投影总字节（与快照对比）
    slice(n)    → 只取前 n 个位置（用于「只重放被接受前缀」）

SpeculativeKDAState (@dataclass)   # 投机前的已接受状态
    state          最新已接受的 KDA 状态
    projections    未确认的 draft 投影
    accepted       已通过验证的步数
    history        历次接受数（调优用）
    state_bytes()  → 完整状态字节（快照策略要持久化的量）
    cache_bytes()  → 投影缓存字节

verify_and_replay(committed, projections, n_accepted)
    state = committed.clone()      # 拷贝，避免污染共享的 committed 状态
    for i in range(n_accepted):    # 只重演被接受的步
        out, state = kda_step(proj.q[i], proj.k[i], proj.v[i],
                              proj.g[i], proj.beta[i], state)
    return state, stacked_outputs   # 新状态 + 各步输出

具体代码如下：

"""投机解码用于 KDA。

设计动机（为什么要"重演"而不是"快照"）
---------------------------------------
KDA 的递归状态 S 每步就地在 cache 里更新，
无法像 KV cache 那样随意越程截断或按 token 指定位拷贝一份"旧状态"。
若投机 draft 中某一位置被验证拒绝，状态已经前移，无法直接回滚。

解法：不缓存完整状态（快照），投机期间只缓存 draft 的一小撮逐 step 投影
(q, k, v, g, beta)，commit 时从"已接受的最长前缀"位置用 kda_step
精确重演一遍 ------ 零快照、零状态深拷贝，只有一次短重放（O(n_accepted) 步）。
这与 SGLang 的 ReplaySSM（raw-input 重演）是**同一路线**，
见 replayssm.py 的 RawInputRing（流式写入版，与本模块的批量构造互补）。
"""

@dataclass
class DraftProjections:
    """draft 窗口内逐 step 的投影，每个都是 (draft, B, H, Dk/Dv/Beta) 张量。

    只缓存重放所需的输入（q/k/v/g/beta），不缓存每步中间状态，
    因此整个 draft 窗口附带的缓存远比持久化的状态快照小。
    """

    q: Tensor
    k: Tensor
    v: Tensor
    g: Tensor
    beta: Tensor

    def __len__(self) -> int:
        return int(self.q.shape[0])

    def nbytes(self) -> int:
        return sum(t.numel() * t.element_size() for t in (self.q, self.k, self.v, self.g, self.beta))

    def slice(self, n: int) -> "DraftProjections":
        return DraftProjections(*(t[:n] for t in (self.q, self.k, self.v, self.g, self.beta)))

@dataclass
class SpeculativeKDAState:
    """投机前的已接受状态 + 可选缓存投影，提交时用于精确重演。"""

    state: Tensor
    projections: DraftProjections | None = None
    accepted: int = 0
    history: list[int] = field(default_factory=list)

    def state_bytes(self) -> int:
        return self.state.numel() * self.state.element_size()

    def cache_bytes(self) -> int:
        return self.projections.nbytes() if self.projections is not None else 0

def verify_and_replay(
    committed: Tensor,
    projections: DraftProjections,
    n_accepted: int,
) -> tuple[Tensor, Tensor]:
    """从已接受位置用 draft 投影精确重演，返回 (新状态, 各步输出)。

    投机被拒绝时，已接受即短于目标数，只需在可接受的前缀上重放：
    每一步的 q/k/v/g 取 (B, H, Dk/Dv)、beta 取 (B, H)，
    直接喂给 kda_step（与解码路径一致）。被拒的草稿位丢弃。
    """
    state = committed.clone()           # 避免修改共享的 committed 状态
    outs: list[Tensor] = []
    for i in range(n_accepted):
        out, state = kda_step(
            projections.q[i], projections.k[i], projections.v[i],
            projections.g[i], projections.beta[i], state,
        )
        outs.append(out)
    if outs:
        stacked = torch.stack(outs, dim=2)
    else:
        stacked = committed.new_zeros(
            committed.shape[0], committed.shape[1], 0, committed.shape[-1]
        )
    return state, stacked

def snapshot_cost_ratio(heads: int, dk: int, dv: int, draft_len: int) -> float:
    """完整快照 vs 投影缓存的每步字节比（draft 长度在分子分母抵消）。"""
    state_per_step = heads * dk * dv
    proj_per_step = heads * (3 * dk + dv + 1)
    return (state_per_step * draft_len) / (proj_per_step * draft_len)

6.6 数据流 / 执行流

投机窗口（draft_len 步）：
  1. draft 模型一口气生成 k 个 token，同时对每步缓存 (q,k,v,g,β) 到 DraftProjections
  2. 真模型验证：只重演「被接受」的前 n_accepted 步：
        state = committed.clone()
        for i in range(n_accepted):
            state = kda_step(proj[i], state)
        # 拒绝的 token 零成本：只是少重演几步
  3. 把重演得到的 state 写回 cache，作为下一步的 committed

6.7 反方辩手

「验证也要跑 kda_step，那投机到底省了什么？」 省的是主模型的前向 。投机时 draft 模型做 k 次便宜的前向，主模型只做一次验证前向 （对 k 个 draft token 批量验证）。如果全对，一次主模型前向 = 生成 k 个 token。重演 kda_step 是很轻的操作（相对主模型前向），所以投影重放的额外开销可忽略。

「为什么不直接存快照，反正有 4 倍差距？」 因为差距是每步的。draft 越长（比如 8 步），绝对差是 8×4 倍；而在线 serving 大 batch 时，这个「每 token 深拷贝状态」的成本占主导------投机省下的算力全被快照拷贝吃掉。投影方案让「省下的」真正省下来。

6.8 何时用 / 何时不用

• 何时用 ：状态宽（dk·dv 大）、draft 长、batch 大------快照存储成本占主导的场景，投影重放是唯一划算的。
• 何时不用 ：状态很窄（dk·dv 接近投影大小）时，快照和投影成本相当，此时可以选快照；不投机（直接解码）时两者都不用。

6.9 可视化

plot_decode_replay.py --- 投机回滚策略成本 (serve/decode.py)

• 表达：投机被拒时 KDA 状态回滚的两条路线成本比较
• 测试：调 snapshot_cost_ratio(H=16, dk=16, dv=16, draft)，扫描 draft=1...32，画「投影缓存 vs 完整快照」的每步字节比。结论是比值恒定 > 1（投影永远更省，且与 draft 长度无关）
• 产出：plot_decode_replay.png

---

0x07 组件四：投机解码 · RawInputRing 环形缓冲（replayssm.py）

RawInputRing：SGLang ReplaySSM 的真实语义与实现。

7.1 问题所在

当前项目中，DraftProjections 借鉴了 K3-master。DraftProjections 的教学演示存在问题。因为当前实现是批量构造，即 draft 全部跑完后，一次性切片包装。

# draft 结束后，一次性把 5 步数据 permute 成 DraftProjections
draft = DraftProjections(
    q=q[:, :, N_COMMIT:].permute(2, 0, 1, 3),  # (5, B, H, D)
    ...
)

为什么这样设计？ 因为 decode.py 的 DraftProjections 定位是离线分析/教学演示 的载体------假设你已经有了完整的 draft 数据，需要包装成 verify_and_replay 能消费的格式。

批量构造并不适合实际在线场景，应该在 draft token 生成过程中逐步存储。因此，我们来看看SGLang 如何处理。

7.2 SGLang ReplaySSM

ReplaySSM 是 SGLang 的方案，我们来学习下。

定义

看 SGLang 的真实定义（server_args.py:2688-2696）：

ReplaySSM spec-verify: fold-every-commit --- a per-slot raw-input window replaces the recurrent verify's per-draft full-state snapshots. The verify stores each draft step's raw inputs into a per-slot window and the commit replays the accepted prefix into the fp32 checkpoint.

ReplaySSM-KDA: the accepted drafts live in the per-slot ring (written during verify); no intermediate_ssm is allocated. Replay the accepted prefix into temporal instead of scattering an intermediate state.

SGLang 的 ReplaySSM 不是"存快照"，而是"存原始输入 + 重演"！ 它和 decode.py 的投影重放是同一路线 ，只是实现细节不同（GDN 存 raw input 窗口，KDA 存 (v, k, gate, beta) 环）。

另外，业界还有 RecoverSSM 方式。

• ReplaySSM：验证后重新回放接受的部分；
• RecoverSSM：只提交接受部分的状态更新。

ReplaySSM（= 存 raw input + 重演）的优势

ReplaySSM 不是"回滚到任意点更快"，而是省 HBM 带宽 。看 server_args.py:2670-2685 的两个独立开关：

开关	场景	核心收益
enable_linear_replayssm	纯 decode（非投机）	每 L 步才把完整状态 flush 到 HBM，期间只在 ring 里累积 raw input。省的是每步写状态矩阵的带宽
enable_linear_replayssm_spec	投机 verify	verify 期间不存每步的完整状态快照，只存 raw input 窗口；commit 时重演被接受前缀到 fp32 checkpoint

带宽账（这是真正的卖点）：

• 传统 decode：每步把 H·dk·dv 的状态矩阵写回 HBM
• ReplaySSM decode：每步只把 (v, k, gate, beta)（H·(3dk+dv+1)）写进 ring，每 L 步才 flush 一次完整状态

注释里写得很清楚："Primarily a GDN (scalar-gate) decode-bandwidth optimization (~1.2-1.5x at batch >= 64)"。batch 大时状态矩阵的 HBM 往返是瓶颈，raw input 环能把带宽砍一个数量级。

KDA 的坑 ：注释也说了------KDA uses its selected Triton or Helion implementation, but its per-K gate ring is larger and ReplaySSM is typically slower than packed KDA decode。KDA 的 gate 是 per-channel 向量（H·dk），比 GDN 的标量 gate 大得多，所以 ring 反而更占带宽，通常更慢，要 benchmark 后再开。

适用场景

ReplaySSM 有效的场景非常具体：

1. GDN 类标量 gate 模型 + 大 batch decode（>= 64）：带宽收益 1.2-1.5x
2. 投机 verify 的内存优化：不想为每步 draft 存完整状态快照时
3. 线性链（topk ≤ 1）：树状投机不支持

对 KDA 来说，SGLang 自己都建议默认别开------per-K gate ring 的带宽开销可能超过省下的状态写带宽。这也解释了为什么 decode.py 的投影重放对 KDA 是更自然的选择：两者都是"存 raw input + 重演"，但投影重放的表述更直接，没有 ReplaySSM 这个名字带来的"快照回放"误导。

关键结论

关键结论 ：SGLang 的 ReplaySSM 与 decode.py 的投影重放是同一路线------都是存 raw input + commit 时重演。差异只在实现形式（流式窗口 vs 批量构造）。

因此，我们构建了 RawInputRing，即 replayssm 方案。

7.3 总体思路

存 raw input（小），不存状态（大）；commit 时重演，不依赖快照。

核心洞察：回滚到 S_k 不需要"当时存下 S_k"，只需要"从某个已知安全点出发，用 raw input 重新递推 k 步"。raw input (q, k, v, g, beta) 每步只占 H·(3dk+dv+1) 字节，远小于完整状态 H·dk·dv。

两条实现路线（同宗同源，共享重演引擎）：

路线	写入时机	适用场景
投影重放（decode.py）	draft 结束后一次性批量构造 DraftProjections	离线分析、教学演示、draft 长度固定
RawInputRing（replayssm.py）	draft 每步实时流式 push()，满则覆盖最旧	在线投机、draft 长度可变、内存受限

7.4 背后的原理

7.4.1 第一性原理：RawInputRing 到底解决什么

KDA 的状态更新可以抽象成：

S_next = F(S_current, q, k, v, g, beta)

其中 S_current 是一个已经被前面 Token 混合过的整体矩阵。投机验证失败时，我们想得到某个较早的 S_k。

但 KDA 没有通用的 undo：

• 状态不是"每个 Token 一格"的可截断列表；
• 被拒 Token 的影响已经写进后续状态；
• 仅凭当前污染后的状态，通常无法反推出被拒 Token 之前的状态。

所以只有两条原则上可靠的路：

1. 保存每个中间状态 ：失败后直接取 S_k；
2. 保存一个已知安全状态和之后的输入轨迹 ：从安全状态重新执行F，重建 S_k。

RawInputRing 选择第二条路。它本质上不是"回滚器"，而是：

为状态重建保留一段可重放的输入轨迹。

在当前项目中，这段轨迹不是原始 token ID，而是已经完成模型投影、足以驱动kda_step 的 (q, k, v, g, beta)。其中：

• (k, v, g, beta) 决定下一状态；
• q 不参与状态更新，但重放时需要它产生对应输出，因此一起保存；
• RawInputRing 只负责流式保存和按时间顺序读出；
• verify_and_replay 才负责从 committed state 逐步重演。

因此，RawInputRing 的作用可以压缩成一句话：

用较小的输入轨迹，替代每一步昂贵的完整状态快照。

7.4.1 什么场景下有效

raw-input 重演路线（SGLang ReplaySSM / decode.py / replayssm.py）：

场景	是否有效	原因
KDA 类 per-channel gate 模型	⚠️ 算法上可用，收益需 benchmark	gate 是向量 H·dk，但 raw-input ring 也更大；SGLang 注释指出 KDA ReplaySSM 通常更慢
任意接受模式（全拒/部分/全接受）	✅ 通用	重演长度 = n_accepted，精确到任意前缀
大 batch decode（>= 64）	✅ 带宽收益明显	每步写 raw input 比写状态矩阵省 HBM 带宽
GDN 类标量 gate 模型	✅ SGLang 注释提到 1.2-1.5x 加速	gate 是标量，ring 更小
KDA + 假想快照路线	⚠️ 可作为基线	快照访问快但写带宽和显存成本高；不能简单断言"不可用"

假想快照路线（用于成本和正确性对照）：

场景	是否有效	原因
整轮全拒	⚠️ 理论高效	pop() 一次到位，零重演
部分接受	❌ 退化	仍需从安全点重演，且每步多付了快照存储
KDA	⚠️ 成本高	gate 向量让快照成本上升，但并非算法上不可用
任意中间点查询	⚠️ 取决于保存范围	环形缓冲覆盖后，超出窗口的中间点无法直接取得

假想快照路线就是存储全部的 KDA state。

7.4.2 带宽账

量化对比

以 H=2, dk=8, dv=8, fp64 为例：

方案	每步缓存字节	draft=5 总字节	相对比值
假想快照	H·dk·dv·8 = 1024	5 × 1024 = 5120	1.94x
raw input	H·(3dk+dv+1)·8 = 528	5 × 528 = 2640	1.00x

• dk 是 key/query/state 的地址或记忆通道维度 。代码中 q, k, g 的最后一维都是 Dk；KDA 状态为
  S.shape = (B, H, dk, dv)。
• dv 是 value/output 的内容维度 。代码中 v 的最后一维是 Dv，单步输出 out 的最后一维也为 Dv。

• 状态矩阵可以按一个 head 写成：S ∈ R^(dk × dv)

所以保存完整状态需要 dk·dv 个元素；而保存 raw input 只需要3dk+dv+1 个元素。两者的差别正是"存矩阵"与"存驱动下一步计算所需的投影"的差别。

另外，这里的 128 和 66 是 fp64 元素数量，不是字节数；fp64 必须再乘sizeof(fp64)=8。相对比值不受 dtype 和 draft 长度影响：

ratio = (H·dk·dv) / (H·(3dk+dv+1))
      = 128 / 66
      ≈ 1.94

dtype 只改变绝对字节数；draft 长度只把两种方案同时线性放大。当 dk·dv 相对 3dk+dv+1 更大时，raw-input 的相对优势才会更明显。

raw input 解读

H·(3dk+dv+1) 是每个 Token、每个缓存时间步的元素数量，不是脱离 dtype 后的字节数。实际缓存字节数应为：

每步字节数 = H · (3dk + dv + 1) · itemsize

这里的每个符号都能在 RawInputRing 中找到对应变量：

数量	具体数据	代码变量	单个 head 的形状	为什么要存
dk	query 向量 q	_q	(dk,)	重放时计算当前 Token 的输出 o = S^T q
dk	key 向量 k	_k	(dk,)	检索旧状态，并决定新信息写入状态的方向
dk	per-channel log gate g	_g	(dk,)	生成逐通道保留因子 exp(g)，先衰减旧状态
dv	value 向量 v	_v	(dv,)	提供当前 Token 要写入状态的目标内容
1	delta-rule 强度 beta	_beta	标量 ()	控制本次写入幅度

因此：

3dk = dk(q) + dk(k) + dk(g)
dv  = dv(v)
1   = 1(beta)

乘上 H 是因为每个 Token 有 H 个 head，每个 head 都各自保存一组q, k, g, v, beta。RawInputRing 的五个 list 正好对应上表：_q, _k, _g, _v, _beta。它们按相同的时间顺序追加和淘汰，所以第 i 个位置的五个张量必须视为一个不可拆分的投影组。

7.5 落地架构与数据结构

7.5.1 RawInputRing 类设计

class RawInputRing:
    """per-slot raw-input 环形窗口：流式缓存 draft 每步的 (q,k,v,g,beta)。

    环形缓冲存的是每步的原始输入投影，**不是**状态快照。满则覆盖最旧
    （draft 长度封顶）。commit 时调用方取 :meth:`window` 得到
    ``DraftProjections``，喂给 ``decode.verify_and_replay`` 重演到目标状态。

    .. note:: 为什么不做回滚？
       回滚逻辑（取前 n_accepted 项、逐步重演）已经在 ``decode.verify_and_replay``
       里实现并对拍验证过。本类只做"流式写入 + 窗口读出"，保持单一职责。
    """

    def __init__(self, capacity: int) -> None:
        """capacity 必须为正整数，通常取 draft_len 上限。

        环形窗口的逻辑语义：``_q[0]`` 永远是当前窗口里**最旧**的一组投影，
        ``_q[-1]`` 是最新。满员时 push 会先淘汰最旧再追加，保证窗口永远
        保留最近 capacity 步。
        """
        if capacity <= 0:
            raise ValueError(f"capacity must be positive, got {capacity}")
        self.capacity = capacity
        # 5 个并列的环形缓冲，每个元素是 (B, H, D) 张量（beta 是 (B, H)）。
        # 用 list 而非定长 tensor：写入次数不固定（draft 长度可变），
        # list 的 append/pop(0) 语义更直白，capacity 是小常数（draft_len 级）。
        self._q: list[Tensor] = []
        self._k: list[Tensor] = []
        self._v: list[Tensor] = []
        self._g: list[Tensor] = []
        self._beta: list[Tensor] = []

    def __len__(self) -> int:
        """当前窗口中的投影组数（<= capacity）。"""
        return len(self._q)

    def push(self, q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor) -> None:
        """每步 draft 推入一组投影。

        Args:
            q, k, g: ``(B, H, Dk)``
            v: ``(B, H, Dv)``
            beta: ``(B, H)``

        满员时先淘汰最旧（各 buffer 的 ``[0]``），再追加到队尾------等价于
        "新投影覆盖最旧槽位"，但保留列表有序性，``window()`` 读出的
        ``DraftProjections.q[i]`` 永远对应"窗口内第 i 旧的投影"。
        """
        if len(self._q) == self.capacity:
            self._q.pop(0)
            self._k.pop(0)
            self._v.pop(0)
            self._g.pop(0)
            self._beta.pop(0)
        self._q.append(q)
        self._k.append(k)
        self._v.append(v)
        self._g.append(g)
        self._beta.append(beta)

    def window(self) -> DraftProjections:
        """返回当前窗口为 ``DraftProjections``，布局 ``(draft, B, H, D)``。

        空窗口时抛 ``IndexError``------调用方应先检查 ``len(ring) > 0``。
        ``verify_and_replay`` 对 n_accepted=0 有专门处理，但空窗口无法
        构造有效 DraftProjections（没有样本决定形状），故选择显式报错。
        """
        if not self._q:
            raise IndexError("RawInputRing empty: call len() before window()")
        return DraftProjections(
            q=torch.stack(self._q, dim=0),
            k=torch.stack(self._k, dim=0),
            v=torch.stack(self._v, dim=0),
            g=torch.stack(self._g, dim=0),
            beta=torch.stack(self._beta, dim=0),
        )

    def reset(self) -> None:
        """清空窗口，恢复到刚构造的状态。"""
        self._q.clear()
        self._k.clear()
        self._v.clear()
        self._g.clear()
        self._beta.clear()

7.5.2 内部数据结构

# 5 个并列的环形缓冲，每个元素是 (B, H, D) 张量（beta 是 (B, H)）
self._q: list[Tensor]   # (B, H, dk)
self._k: list[Tensor]   # (B, H, dk)
self._v: list[Tensor]   # (B, H, dv)
self._g: list[Tensor]   # (B, H, dk)
self._beta: list[Tensor] # (B, H)

逐项解释：

1. g：衰减旧记忆。 g 是逐 channel 的 log decay，alpha = exp(g)。alpha[d] 会乘到状态矩阵的第 d 行，因此不同 key channel可以有不同的记忆保留率。没有 _g，就无法重建每一步对旧状态的衰减。
2. k：检索并定位写入位置。
   sk = S^T k 从旧状态中检索出当前 key 已经记住的 value 内容；外积 k ⊗ delta 再沿 key 方向把新信息写回状态。没有 _k，既不能计算检索结果，也不能确定写入状态的哪些行。
3. v：提供目标内容。
   v - sk 是当前 value 与状态已有预测之间的差值，即 delta rule 要补写的新信息。没有 _v，只能知道已有记忆，无法知道要写入什么。
4. beta：控制写入强度。
   delta = beta · (v - sk)。beta 越大，本 Token 对状态的修正越强；beta 越小，更新越保守。没有 _beta，重放得到的状态会与原递推不同。
5. q：生成对应输出。
   q 不参与状态更新，因此如果目标只是重建最终 state，理论上可以不保存 q；但 verify_and_replay 同时返回每个接受 Token 的输出，需要用 out = S^T q 计算这些输出，所以项目将 _q 与其余投影一起保存。

_g 和 _beta 不是附加日志，而是 KDA 状态转移的一部分。

一个 head 的具体数字

以 dk=8, dv=8 为例，一个 head、一个 Token 的缓存是：

q:    8 个元素
k:    8 个元素
g:    8 个元素
v:    8 个元素
beta: 1 个元素
----------------
合计: 33 个元素 = 3·8 + 8 + 1

若 H=2：

2 个 head × 33 = 66 个元素

若 dtype 是 fp64，每个元素 8 bytes：

66 × 8 = 528 bytes/step
draft=5 时：5 × 528 = 2640 bytes

为什么用 list 而非定长 tensor？

• draft 长度可变，写入次数不固定
• append/pop(0) 语义直白，capacity 是小常数（draft_len 级）
• window() 时 torch.stack(list, dim=0) 一步转 DraftProjections

RawInputRing 是逐步调用的：

for i in range(DRAFT_LEN):
    ring.push(q[:, :, i], k[:, :, i], ...)  # 每步实时写入

这正是 SGLang enable_linear_replayssm_spec 的在线语义------draft 每产出一步就立即缓存，不等全部跑完。

如何重放

对某个 Token，kda_step 的重建过程等价于：

# state: (B, H, dk, dv)
# q, k, g: (B, H, dk)
# v:       (B, H, dv)
# beta:    (B, H)

state = g.exp().unsqueeze(-1) * state
sk = torch.einsum("bhd,bhdv->bhv", k, state)
delta = beta.unsqueeze(-1) * (v - sk)
state = state + k.unsqueeze(-1) * delta.unsqueeze(-2)
out = torch.einsum("bhd,bhdv->bhv", q, state)

重放时，window() 将五个 list 分别 torch.stack(..., dim=0)，得到 (draft, B, H, D) 的 DraftProjections；随后verify_and_replay 从 committed state 出发，只遍历被接受的前缀，

依次调用 kda_step(q[i], k[i], v[i], g[i], beta[i], state)。因此，RawInputRing 本身不"恢复"状态；它保存重建所需的五类数据，真正的状态重建发生在 verify_and_replay 的逐 Token 递推中。

DraftProjections 和 RawInputRing 的区别

存储内容完全相同 ------都是 (q, k, v, g, beta) 五元组。区别只在写入接口：

	DraftProjections	RawInputRing
本质	一个 dataclass（5 个 tensor 的容器）	一个环形缓冲管理器
写入方式	构造时一次性传入所有数据	逐 step push()，满则覆盖最旧
容量控制	无（draft 多长就存多长）	有 capacity 上限，超出覆盖最旧
使用场景	已知全部 draft 数据后批量包装	在线流式写入，内存受限

RawInputRing 不比 DraftProjections 存得多------存的是同一批数据，只是写入方式不同 。但当前实现不会自动推进已安全提交的 checkpoint；若一次投机窗口最多可能接受draft_len 个 Token，capacity 必须覆盖这段完整前缀。capacity < draft_len只有在调用方另有安全 checkpoint、且能证明不会重放被覆盖前缀时才成立。

7.5.3 RawInputRing vs 假想快照

一句话：假想快照在"全拒且 capacity 足够"时理论最优，但 KDA 的状态矩阵太大、覆盖语义让中间点不可靠、SGLang 工程验证更慢------所以被淘汰。

假想快照理论上可行，工程上代价不可接受。

维度	假想快照	raw-input 重演
每步存储	H·dk·dv·itemsize 字节（完整状态矩阵）	H·(3dk+dv+1)·itemsize 字节
draft=5, H=2, dk=dv=8, fp64	5 × 1024 = 5120 B	5 × 528 = 2640 B
draft=10, H=16, dk=dv=64	10 × 65536 = 640 KB	10 × 1040 = 10 KB
回滚到任意中间点	直接取对应快照（若未被覆盖）	从安全点重演 n 步
计算代价	零（若快照在）或重演（若被覆盖）	固定 n 步重演

关键问题不是"能不能"，而是"值不值"：

• KDA 的 per-channel gate 让状态矩阵膨胀 ：dk·dv 是乘积关系，而 raw input 的 3dk+dv+1 是加和关系。dk=dv=64, H=16 时比值约 10.7x，不是 2x
• SGLang 明确说 KDA 上更慢 ：server_args.py 注释 "KDA ... its per-K gate ring is larger and ReplaySSM is typically slower than packed KDA decode"
• 环形覆盖语义下，中间点可能已被顶掉：想取 S_5 但 capacity=3 时只剩 S_6,S_7,S_8，反而还得从 S_3 重演------快照白存了

7.5.4 与 decode.py 的协作

RawInputRing 必须和 decode.py 配合使用吗？

不是必须，但目前设计是强耦合的。

层面	是否必须	说明
数据格式	是	window() 返回 DraftProjections，这是 verify_and_replay 的唯一输入类型
重演引擎	否	你可以不用 verify_and_replay，自己写循环调 kda_step
依赖关系	是	replayssm.py import 了 decode.py 的 DraftProjections

解耦的可能性 ：RawInputRing 内部的 _q/_k/_v/_g/_beta 是独立存储的，你可以：

# 不用 window()，直接访问内部 buffer（不推荐，破坏封装）
for i in range(len(ring)):
    q, k, v, g, beta = ring._q[i], ring._k[i], ...
    out, state = kda_step(q, k, v, g, beta, state)

但这绕过了 DraftProjections 的类型安全，目前没有这个需求。

当前定位 ：RawInputRing 是 decode.py 投影重放路线的流式写入前端 ，两者是同一方案的不同封装，不是独立方案。

关键设计 ：RawInputRing 不做回滚 ，只做"流式写入 + 窗口读出"。回滚逻辑复用 decode.verify_and_replay，保持单一职责。

7.6 何时调用

7.6.1 写入路径（draft 期间）

from models.llm.kimi_k3.serve.replayssm import RawInputRing

ring = RawInputRing(capacity=draft_max)  # capacity >= draft_len 上限
for step in range(draft_len):
    q, k, v, g, beta = model.project(token[step])  # 模型前向
    ring.push(q, k, v, g, beta)                     # 流式写入

7.6.2 读取路径（commit 时）

from models.llm.kimi_k3.serve.decode import verify_and_replay

n_accepted = verify(draft_tokens)                   # 验证接受数
if len(ring) > 0:
    projections = ring.window()                     # 取出 raw input
    new_state, outs = verify_and_replay(
        committed_state, projections, n_accepted    # 重演到目标状态
    )
ring.reset()                                        # 清空窗口，准备下一轮

7.6.3 容量约束

capacity 必须 >= draft_len 的最大可能值。

若 capacity < draft_len，早期投影会被覆盖，导致 window() 返回的 DraftProjections 不包含重演所需的前缀。demo 中刻意演示了这种溢出场景作为警示。

7.7 关键对比：投影重放 vs RawInputRing ------ 两条路线各在何时适用

一句话 ：RawInputRing 是 DraftProjections 的流式写入封装 ，解决"在线投机时内存受限"的问题，两者是协作关系而非替代关系。

不是升级，是互补------就像"流式写入"和"批量读取"的关系。

维度	DraftProjections	RawInputRing
定位	数据载体（DTO）	写入管理器
职责	包装已有数据，传给 verify_and_replay	控制写入时机、容量、覆盖策略
生命周期	不可变（构造后只读）	可变（push/reset 动态变化）
内存控制	无（数据多大就多大）	有（capacity 上限）

类比：

• DraftProjections = torch.Tensor（数据本身）
• RawInputRing = torch.utils.data.DataLoader（控制怎么喂数据）

不是升级版的原因：

1. RawInputRing 依赖 DraftProjections ：window() 的产物就是 DraftProjections
2. DraftProjections 可以独立存在：demo 方案 A 直接构造，不经过 RawInputRing
3. 场景不同，无优劣：批量场景用 DraftProjections 更直接，流式场景用 RawInputRing 更自然

两者都是 raw-input 路线，共享 verify_and_replay，差异只在写入时机：

维度	投影重放（decode.py）	RawInputRing（replayssm.py）
写入方式	批量构造 DraftProjections（draft 结束后一次性）	流式 push()（draft 每步实时）
内存峰值	draft_len × 66 B	capacity × 66 B（capacity 可 < draft_len）
适用场景	离线分析、教学演示、draft 长度固定	在线投机、draft 长度可变、内存受限
与 SGLang 对应	无直接对应	enable_linear_replayssm_spec
回滚引擎	verify_and_replay	verify_and_replay（共享）
结果一致性	与参考 S_k 逐位对齐	与参考 S_k 逐位对齐

假想快照路线（已被淘汰）：

维度	假想快照
每步缓存	完整状态 S（H·dk·dv 字节）
回滚方式	pop 最旧快照 + 可能重演兜底
KDA 适用性	❌ 不适用（per-channel gate 让快照成本失控）
SGLang 对应	❌ 无（SGLang 从未用快照）

7.8 示例

见 experiments/models/llm/kimi_k3/serve/demo_replay_vs_snapshot.py，场景：

• 已 commit 3 个 token（S_3）
• draft 提出 5 个 token，验证后接受前 2 个（d3 起被拒）
• 目标回滚到 S_5

输出：

========================================================================
场景：已 commit 3 个 token（S_3），draft 5 个，验证后接受前 2 个（d3 起被拒），目标回滚到 S_5
========================================================================

[A] 投影重放：draft 结束后一次性批量构造 DraftProjections
    缓存投影字节：2640 B  (5 步 x H*(3*dk+dv+1)=66 B/步)
    内存峰值：draft 结束时一次性占 2640 B
    重演 2 步 -> S_5
    与参考 S_5 的最大误差：0.00e+00

[B] RawInputRing：draft 每步实时流式 push，容量=3
    （capacity < draft_len，演示环形覆盖语义）
    push 第 0 步后，窗口大小 = 1
    push 第 1 步后，窗口大小 = 2
    push 第 2 步后，窗口大小 = 3
    push 第 3 步后，窗口大小 = 3（最旧被覆盖）
    push 第 4 步后，窗口大小 = 3（最旧被覆盖）

    警告：capacity=3 < draft_len=5，第 0,1 步投影已被覆盖，无法重演前 2 步！
    真实部署中 capacity 必须 >= draft_len 上限，本例仅演示覆盖语义，下面用 capacity=5 重跑
    capacity=5 时，窗口完整保留 5 步投影
    缓存投影字节：2640 B  （与方案 A 相同）
    内存峰值：全程只占 5 步 x 66 B/步（流式写入的优势）
    重演 2 步 -> S_5
    与参考 S_5 的最大误差：0.00e+00

========================================================================
汇总
========================================================================
  最终状态一致性：A == B == 参考 S_5  (max err 0.00e+00)
  缓存内容：完全相同（raw input (q,k,v,g,beta)，2640 B）
  写入时机：A 批量（draft 结束后）；B 流式（draft 每步）
  内存峰值：A = draft_len x 66 B；B = capacity x 66 B（可 < draft_len）
  重演引擎：共享 verify_and_replay（SGLang ReplaySSM 的真实语义）

  假想快照路线（已被淘汰）：每步存完整状态 H*dk*dv=128 B，是 raw input 的 1.94 倍
  KDA 的 per-channel gate 让快照成本雪上加霜，SGLang 注释明确指出 KDA 上 ReplaySSM 通常更慢

7.9 可视化

见 experiments/models/llm/kimi_k3/serve/gradio_serve_demo.py Tab 4，四个子图：

子图	内容
左上	RawInputRing push/窗口溢出时序（capacity=4，演示覆盖语义）
右上	window() 转 DraftProjections 的布局可视化（draft × H × D）
左下	RawInputRing vs 投影重放：共享 verify_and_replay，误差均为 0
右下	带宽账：raw input vs 假想快照（draft 越长，快照劣势越明显）

---

0x08 组件五：KDA 上下文并行（KCP）

8.1 它要解决什么问题

上下文并行（context parallelism, CP） ：当序列极长（比如 1M token）时，单块 GPU 放不下整个序列，需要把序列切成多段，分到多块设备（rank）上并行计算。每段各算各的，最后把结果拼起来。

对普通 KV 的上下文并行，做法很简单：各段先算好局部的 KV，一次 all-gather 拼起来，然后各段自由地重算后续 。因为 vanilla 线性注意力的递推是加性的 （S = Σ_t k_t v_tᵀ），每段可以从 S=0 出发算自己的「自生成状态」，整体入态 = 前面各段自状态的直接求和。

对 KDA 呢？ KDA 的状态递推不是加性的：

S_t = M_t S_{t-1} + β_t k_t v_tᵀ ,   M_t = (I − β_t k_t k_tᵀ) Diag(α_t)

M_t 是逐 token 相关的矩阵 ，先作用在入状态上，再加本次写入 。所以局部段对结果的贡献，取决于「进入这一段时的状态」------不能只靠「从 S=0 出发的自生成状态」来恢复。

8.2 难在哪

状态跨段「不可加」。直观地说：

• vanilla 线性注意力：段 A 和段 B 是「并行算好、直接相加」的关系------因为每段都是纯「累加」；
• KDA：段 B 的输入状态，依赖于段 A 末尾的状态；而段 A 末尾的状态，又取决于更前面的段。直接把它们加起来是错的。

8.3 总体思路

把每段的效果分解成两个「本地可算」的量，再用一次「仿射链的前缀扫描」重组。

每个段 [t0, tN) 可以看作一个仿射变换 T(S) = M̄·S + S̃：

• M̄（段内累积过渡矩阵） ：把段内每个 token 的 M_t 依次复合------它负责「把上一段末尾传进来的状态，缩放成本段末尾的状态」；
• S̃（段从零状态自生成的状态） ：从 S=0 出发，由本段 token 生成的「自有状态」。

于是段 i+1 的入态可以写成（K3 技术报告 Eq.17）：

S_{[i+1]} = S̃_{[i+1]} + M̄_{[i+1]} · S_{[i]}

这两个量 M̄ 和 S̃ 只依赖本段的本地 token ，在「上一段的入态可用之前」就能算好。而且 rank 级更新可结合 ------因此每个 rank 的入态可以用一个前缀扫描 恢复。通信量只随模型维度（与序列长度无关），算力线性扩展------这就是百万 token 上下文训练可行的基础。

8.4 背后的原理

为什么「分解成 M̄ 和 S̃」就能绕过「不可加」？因为把一个「逐 token 的、非线性的递推」变成了「段级别的、仿射的、可结合的链」。

看图。把序列切成 4 段，分别放到 rank 0-3。三段式流程：

阶段 1：每段本地算 (M̄, S̃) ------ 各 rank 并行，互不等待
         [段0]      [段1]      [段2]      [段3]
         (M̄0,S̃0)   (M̄1,S̃1)    (M̄2,S̃2)    (M̄3,S̃3)      ← 只依赖本地 token

阶段 2：一次固定大小 all-gather ------ 交换的只有 (M̄, S̃)，与序列长度无关
         ┌────── all-gather (M̄, S̃) ──────┐
         ▼       ▼     ▼     ▼
         rank0 rank1 rank2 rank3       每个 rank 都拿到完整列表：
                                       [(M̄0,S̃0),(M̄1,S̃1),(M̄2,S̃2),(M̄3,S̃3)]
         （all-gather = "每个 rank 把自己的 (M̄,S̃) 广播给所有人，最终每人手上有全部"
           为什么非要它？因为 rank i 的入态依赖所有更前面段的 (M̄,S̃)，
           所以每个 rank 必须先拿到全部，才能独立做下一步前缀扫描）

阶段 3：前缀扫描恢复各段入态
         incoming[0] = 0
         incoming[1] = M̄0·incoming[0] + S̃0
         incoming[2] = M̄1·incoming[1] + S̃1
         incoming[3] = M̄2·incoming[2] + S̃2
         每个 rank 用 incoming[i] 作为本段 initial_state，重跑 chunkwise 得输出

补充说明 ：KCP 严格说是 4 个阶段。代码 docstring 里的"三阶段"是把阶段 1（分段）和阶段 2（本地算）合并了。按真实流水拆开：

1. 分段/调度 ：沿 T 均分，发到各 rank（kcp.py:150）
2. 并行本地计算 （M̄, S̃)：各 rank 互不等待（kcp.py:152-165）
3. 一次 all-gather + 前缀扫描 （串行瓶颈）：交换的只有 (M̄, S̃)，通信量与序列长度无关；扫描本身是 n_segments 步的串行递推（kcp.py:166）
4. 并行重算输出 ：各 rank 用 incoming[i] 独立跑 kda_chunkwise 产出真实输出（kcp.py:171-180）

关键点：阶段 3 的"前缀扫描"是串行部分（仿射链不可并行，这正是 delta rule 与 vanilla 线性注意力的根本差异），但它的输入是常数大小的 (M̄, S̃)，不随 T 增长，所以瓶颈可控。阶段 2 和 4 都是完全并行的。

关键洞察 ：虽然每段的状态「不可直接相加」，但把每段写成 入态 → 出态 的仿射变换 后，整条链就是「一个接着一个的仿射变换」------这正是前缀扫描 能处理的形态（把「段级转移」当作「标量转移」，只是转移量从标量变成了矩阵）。这跟 vanilla「直接相加」的根本差别在于：vanilla 的转移矩阵是恒等（I），KDA 的转移矩阵逐段不同（M̄），所以必须扫描、不能直接求和。

8.4.1 图：KCP 前缀扫描

左图：跨段状态重组两种方式对精确递推的误差（log 轴）------"直接相加"误差 9.78e-01 （基本全错），"KCP 前缀扫描"误差 1.11e-16（机器精度，与整段递推逐位一致）。

右图：KCP 通信量与序列长度无关（常数），而 softmax attention 的 KV 交换量随长度线性/超线性增长。

8.4.2 常见疑问：KCP 用的 kda_chunkwise 到底是什么并行？

KCP 在阶段 1 和阶段 3 都调用了 kda_chunkwise。要理解 KCP，就得先懂这个 kernel 的并行结构。

kda_chunkwise 是"块内并行、块间串行" （K3 技术报告 Eq.3-4）。它把序列切成 chunk_size=16 的块（chunk）：

• 每个 chunk 内部：token 级并行 ------块内所有 token 的相互作用一次算完（q_bar @ k_hatᵀ、A 矩阵、UT 三角求解都是矩阵乘法），这正是"chunk 内并行"；
• chunk 之间：串行递推 ------每个 chunk 的入态来自上一个 chunk 的尾态，一个 for i in range(n_chunks) 循环串行推进。

序列 T = 64，chunk_size = 16 → 4 个 chunk

  chunk0        chunk1        chunk2        chunk3
 [0..15]       [16..31]      [32..47]      [48..63]
    │             │             │             │
    │  token 并行  │             │             │
    │  (矩阵乘)    │             │             │
    ▼             ▼             ▼             ▼
   S₀ ──►  S₁  ──►  S₂  ──►  S₃  ──►  S₄     （状态串行在 chunk 间传递）
        chunk内并行算完所有 token 输出，再把尾态 S 传给下一个 chunk

对比 kda_recurrent：recurrent 是纯逐 token 串行（一个 token 一个循环）；chunkwise 把"一个 chunk 内"重排成并行矩阵乘，只保留"chunk 间"的串行状态传递。两者数学等价，chunkwise 是 GPU 友好的并行形式。

KCP 为什么用 chunkwise 算 (M̄,S̃)？ 因为 local_transition_and_state 里：

• S̃ = 用 kda_chunkwise(k, v=真实值, S₀=0) 跑出"段自生成状态"；
• M̄ = 用 kda_chunkwise(k, v=0, S₀=I) 跑出"纯过渡矩阵"（v=0 时写入项消失，只剩过渡项）。

两个量都是"从零/恒等出发跑一段 chunkwise"得到的，所以共用同一个 kernel，数学上严格吻合 delta rule。

8.5 落地架构与数据结构（parallel/kcp.py）

local_transition_and_state(k, v, g, beta)     计算单段 (M̄, S̃)
    S̃:  用 kda_chunkwise(k, v=真实值, S0=0)   → 段自生成状态
    M̄:  用 kda_chunkwise(k, v=0, S0=I)         → 纯过渡矩阵（homogeneous 响应）
         # 为什么 v=0, S0=I 得到 M̄？v=0 时写入项 β·k·vᵀ 消失，
         # 只剩过渡项，从 S0=I 出发跑完段得到的就是累积转移矩阵

prefix_scan_states(transitions, zero_states)   前缀扫描
    incoming[0]   = 0
    incoming[i+1] = transitions[i] @ incoming[i] + zero_states[i]

kcp_forward(q, k, v, g, beta, n_segments)      三段式总入口
    # 阶段1: 每段算 (M̄, S̃)
    # 阶段2: prefix_scan 重组入态
    # 阶段3: 每段用 incoming[i] 重跑 chunkwise

KDAContextParallel                       封装 + 通信量
    communication_volume(batch, heads, dk, dv, world, itemsize)
        = batch·heads·(dk² + dk·dv)·itemsize · world
        # 注意：公式里没有序列长度 T！这就是 KCP 的核心卖点

具体代码如下：

"""KDA Context Parallelism（KCP：Kimi Delta Attention 上下文并行，论文 §5.1.2 / Eq.17）。

设计动机（why）
--------------
KDA 的状态更新是 **delta rule**：每个 token 不仅"写入"一条 k v 外积，还会通过
(I - beta_k k k^T) 这个 **逐 token 依赖的过渡矩阵** 去缩放左右两侧既有的状态。

    S_r = (I - beta_r k_r k_r^T) Diag(alpha_r) S_{r-1} + beta_r k_r v_r^T

因此，把序列切成多段分别计算时，一段的效果**不能**像 vanilla 线性注意力
（S = sum k v^T，天然可加）那样简单相加或拼出一个闭式前缀和：

    vanilla : S_{[i+1]} = S_{[i]} + 段内累加量          （可加，可前缀和）
    KDA     : S_{[i+1]} = Stilde_{[i+1]} + M_{[i+1]} S_{[i]}  （仿射，逐段不同 M → 需扫描重组）

KCP 的解法是把每段内部的效果拆成**两个局部量**，再做一次仿射链的"前缀扫描"：

  * M_{[i]}    ------ 段内累积过渡矩阵（homogeneous / 齐次响应）：
                    令 v = 0 且初态 S_0 = I，delta rule 完全退化为纯过渡
                    S -> (I - beta k k^T) Diag(alpha) S，所有 token 的过渡矩阵
                    逐段复合得到 M_{[i]}；
  * Stilde_{[i]} ------ 段从零初态"自生成"的状态：S_0 = 0，直接跑 chunkwise。

读出口按 S_{[i+1]} = Stilde_{[i+1]} + M_{[i+1]} S_{[i]} 前缀扫描重组各段的**入向
状态**（incoming），再用 kda_chunkwise(initial_state=incoming[i]) 得到段内输出。

本实现 CPU 单进程：以 n_segments 分段模拟真实的 world 个 rank，不引入
torch.distributed。rank 间只交换 (M, Stilde)，通信量（字节）与序列长度无关。
"""

（下接函数定义）

8.5.1 local_transition_and_state（计算每段的 M̄ 和 S̃）

def local_transition_and_state(
    k: Tensor, v: Tensor, g: Tensor, beta: Tensor,
    *, chunk_size: int = 16, dtype: torch.dtype | None = None,
) -> tuple[Tensor, Tensor]:
    """计算单段的两个局部量 (M_{[i]}, Stilde_{[i]})。

    Returns:
        (m_cum, s_tilde)：
        - m_cum    (B, H, Dk, Dk)：段内累积过渡矩阵 M_{[i]}。
        - s_tilde  (B, H, Dk, Dv)：段从零状态自生成的状态 Stilde_{[i]}。

    数值设定：
        用 eye 作 initial_state + zero_v 的调用，从 S_0 = I、v = 0
        出发，状态递推退化为纯过渡矩阵连乘，得到的最终状态正是 M_{[i]}。
        这样的好处是 M 与 S 在同一次 chunk 化实现里用**相同代码路径**算出，
        数学上严格吻合 delta rule 的复合结构。
    """
    dtype = dtype or (torch.float64 if k.dtype == torch.float64 else torch.float32)
    b, h, _t, dk = k.shape

    # Stilde_{[i]}：段从 S=0 开始，自身激活（v 非零）产生的状态。
    # q 用 k 占位 ------ 我们只需要返回值中的 state，与序号无关，q 只影响输出 out（丢弃）。
    _, s_tilde = kda_chunkwise(k, k, v, g, beta, chunk_size=chunk_size, dtype=dtype)
    # M_{[i]}：homogeneous 响应。v=0 之后 KDA 的写入项 (beta k v^T) 消失，
    # 只剩过渡项，给定 S0=I，最终 state 就等于段内的累积过渡矩阵。
    eye = torch.eye(dk, dtype=dtype, device=k.device).expand(b, h, dk, dk).contiguous()
    zero_v = torch.zeros_like(k)
    _, m_cum = kda_chunkwise(k, k, zero_v, g, beta, chunk_size=chunk_size,
                            initial_state=eye, dtype=dtype)
    return m_cum, s_tilde

要点：M 与 Stilde 用同一套 kda_chunkwise 代码路径 算出来，只是换初态（S0=I vs S0=0）和 v（v=0 vs v 原值）------数学上严格吻合 delta rule 的复合结构。

8.5.2 prefix_scan_states（前缀扫描重组）

def prefix_scan_states(transitions: Tensor, zero_states: Tensor) -> Tensor:
    """基于每段 (M, Stilde) 的前缀扫描，恢复各段的**入向**状态。

    Args:
        transitions: (n, B, H, Dk, Dk) 每段的过渡矩阵 M。
        zero_states: (n, B, H, Dk, Dv) 每段零态激活产生的 Stilde。

    Returns:
        incoming (n, B, H, Dk, Dv)：incoming[i] 是段 i **带着**
        进入执行的状态（incoming[0] = 0）。

    核心递推（对扫描 self-consistency）:
        incoming[0]     = 0
        incoming[i+1]   = transitions[i] @ incoming[i] + zero_states[i]

    即：段 i+1 的入态 = 「段 i 过渡到它自身」+ 「段 i 从零状态自主生成的状态」。
    这正是 delta rule 与 vanilla 线性注意力的根本差异所在 ------ 状态不直接可加，
    必须用一个逐段不同的仿射变换来传递。
    """
    n = transitions.shape[0]
    incoming = torch.zeros_like(zero_states)
    state = torch.zeros_like(zero_states[0])
    for i in range(n):
        incoming[i] = state
        state = transitions[i] @ state + zero_states[i]
    return incoming

8.5.3 kcp_forward（三阶段总入口）

def kcp_forward(
    q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor,
    *, chunk_size: int = 16, n_segments: int = 4,
    dtype: torch.dtype | None = None,
) -> tuple[Tensor, Tensor]:
    """把序列切成 n_segments 段，按 KCP 三段式顺序计算。

    三阶段：
      1. 分片：把序列沿 T 均分成 n 段（要求 T % n_segments == 0）。
      2. 并行计算各段的局部量 (M, Stilde)（真正的分布式中这段在 rank 之间）。
      3. 前缀扫描重组出各段入状态，再用 kda_chunkwise(initial_state=incoming[i])
         逐段产出输出并把状态推到末尾。
    """
    dtype = dtype or (torch.float64 if q.dtype == torch.float64 else torch.float32)
    b, h, t, dk = q.shape
    dv = v.shape[-1]
    if t % n_segments != 0:
        raise ValueError(f"length {t} not divisible by n_segments {n_segments}")

    step = t // n_segments
    segs = [slice(i * step, (i + 1) * step) for i in range(n_segments)]

    # --- 阶段 1：每段局部 (M̄, S̃)（分布式下这段在各自 rank 本地并行执行）---
    m_all, s_all = [], []
    for sl in segs:
        m_, st_ = local_transition_and_state(
            k[:, :, sl], v[:, :, sl], g[:, :, sl], beta[:, :, sl],
            chunk_size=chunk_size, dtype=dtype,
        )
        m_all.append(m_)
        s_all.append(st_)

    Ms = torch.stack(m_all, dim=0)   # (n, b, h, dk, dk)
    Ss = torch.stack(s_all, dim=0)   # (n, b, h, dk, dv)
    incoming = prefix_scan_states(Ms, Ss)   # (n, b, h, dk, dv)

    # --- 阶段 2：逐段用 incoming 重新计算真实输出 ---
    outs = []
    state = incoming[0]
    for i in range(n_segments):
        out, state = kda_chunkwise(
            q[:, :, segs[i]], k[:, :, segs[i]], v[:, :, segs[i]],
            g[:, :, segs[i]], beta[:, :, segs[i]],
            chunk_size=chunk_size, initial_state=incoming[i], dtype=dtype,
        )
        outs.append(out)
    return torch.cat(outs, dim=2), state

8.5.4 通信量 & 并行

class KDAContextParallel:
    """携带 chunk_size 的 KCP 封装器，供报告使用。"""

    @staticmethod
    def communication_volume(
        batch: int, heads: int, dk: int, dv: int, world: int, itemsize: int = 8
    ) -> int:
        """每 rank 每 KDA 层交换 (M, Stilde) 的字节数 ------ 与序列长度无关。

        M:    Dk x Dk 过渡矩阵
        Stilde: Dk x Dv 零状态生发
        每个 rank 存外传这两块给其余 world-1 个 rank。

        返回整个集群一次 KDA 层的交换总量（字节）。动态大小时保持不变；
        world 翻倍时总量翻倍（一条链路多收发一份）。
        """
        per_rank = batch * heads * (dk * dk + dk * dv) * itemsize
        return per_rank * world

通信量 （kcp.py:200-214）：每 rank 每 KDA 层交换的是 (M, S̃ / Stilde) 两个定长张量，大小与序列长度无关------这是 KCP 对长上下文的核心卖点：

• M_{[i]}：(B, H, Dk, Dk) --- 段内累积过渡矩阵（d×d）
• S̃_{[i]}：(B, H, Dk, Dv) --- 段从零初态自生成的状态（d×d）

每 rank 总字节 = B · H · (Dk² + Dk·Dv) · itemsize，与序列长度 T 无关（这就是 KCP 的核心卖点）。

并行 vs 串行 （kcp.py:152-180）：

量	计算方式	并行性
M_{[i]}	每段独立调用 kda_chunkwise(q=k, v=0, S₀=I)	完全并行（段与段之间无依赖，分布式下各 rank 本地算自己的）
S̃_{[i]}	每段独立调用 kda_chunkwise(q=k, v, S₀=0)	完全并行（同上）
incoming[i]（前缀扫描）	incoming[i+1] = M[i] @ incoming[i] + S̃[i]	本质串行（递推依赖上一段结果；但只是 n 次 d×d 矩阵乘，开销远小于段内计算）
S_out[i]（最终输出）	每段用 kda_chunkwise(initial_state=incoming[i]) 重算	完全并行（拿到入态后段间又独立了）

即**"两并行 + 一短串行 + 再并行"** ：第一段并行求 (M, S̃)，然后一次 O(n) 前缀扫描（n = rank 数，通常 ≤ 8），最后再并行计算输出。串行段的代价是 O(n · d³)，与序列长度无关，可忽略。

8.6 数据流 / 执行流

kcp_forward（n_segments=4, chunk_size=16）：
  1. 分片：把序列沿 T 均分成 4 段（要求 T % 4 == 0）
  2. 每段本地算 (M̄, S̃)：local_transition_and_state()
  3. 前缀扫描：prefix_scan_states() → 得到每段入态 incoming[i]
  4. 每段用 incoming[i] 作为 initial_state，重跑 kda_chunkwise() → 输出 + 推到末尾
  5. 拼接各段输出，返回 (out, s_final)

8.7 反方辩手

「现在大上下文直接整段塞上去不就行了，何必卷 KCP？」 seq_len → ∞ 时，整段传输成本与 T 相关（随序列翻倍）；KCP 的通信量与 T 无关 （只随模型维度），是唯一能随上下文翻倍的方案。而且 KCP 依赖「状态廉价搬移」（统一页池的 pack_heads head-contiguous 布局），两个组件在这里汇合------这不是巧合，是设计使然。

「为什么说通信量与序列长度无关？」 因为跨 rank 交换的只有 (M̄, S̃)------它们都是固定尺寸 （dk×dk 和 dk×dv），不随段内 token 数变化。communication_volume 公式里没有 T，这正是要断言的属性。

8.8 何时用 / 何时不用

• 何时用：长序列（上下文 > 单设备容量）需要跨设备并行；此时「每段独立算 (M̄,S̃) + 一次固定 all-gather + 前缀扫描」是唯一通信量不随 T 增长的方案。
• 何时不用：短序列（单设备放得下）------KCP 的通信开销反而是浪费；纯 MLA 模型没有 KDA 状态，用普通 KV 拼装即可。

8.9 实测数值（CPU 可复现）

• 直接相加误差（错误做法）：9.782e-01
• KCP 前缀扫描误差（正确做法）：1.110e-16
• 通信量对比：KDA 恒定 2048 字节/rank/层，softmax 在 1M token 时达 1.07e9 字节

一句话结论 ：KCP 把"不可加的仿射状态链"拆成每段两个局部量（M 齐次响应 + Stilde 零态自生成），再用 S ← M S + Stilde 前缀扫描重组各段入状态，最后每段用 initial_state=incoming[i] 重算输出。误差到机器精度，通信量 O(1)。

8.10 验证

tests/parallel/test_kcp.py（3 项）：

• test_kcp_forward_matches_recurrent（并行分段与全局递推数值一致 ✅）
• test_communication_volume_scales_with_world（固定 rel 值线性）
• test_local_transition_state（M̄ 用 identity/S₀ 构造 + 覆盖）

---

0x09 组件六：Fleet 层调度（serve/fleet.py）

9.1 它要解决什么问题

前 5 个组件解决「单个实例内 」的请求效率（缓存怎么用、状态怎么复用、怎么回滚、怎么并行）。Fleet 层把视角提升一层：真实线上不是一台机器，而是几十上百台组成的 fleet（集群组）。Fleet 层回答两个单实例不存在的问题：

1. 缓存命不命中，不由单实例决定 ：一个请求可能被路由到任何一个集群。如果落在「没缓存这个前缀」的集群，就得整个重新 prefill。K3 报告的关键数字：1M 上下文下，典型编码请求的 prefix 达 400K token ，而 prefill increment 只有 4K ------命中一次避免重算 400K 前缀，比 miss 便宜几个数量级。
2. 一个超长请求能拖垮整个 fleet：短请求 <2K token、超长请求可达 1M token，单请求长度横跨 3 个数量级。长请求突发一旦占满总容量，会把所有短请求的 TTFT 全线拖垮。

9.2 难在哪

• 亲和性 vs 故障的矛盾：要让命中率高，就得把 session「绑定」到持有其前缀缓存的集群（cache-aware affinity）。但绑定 = 绑死，主集群一旦故障，所有绑定它的 session 全部中断。
• 长度方差大：3 个数量级的长度差，固定总容量下，长突发必然挤死短请求。

9.3 总体思路（两个策略）

策略 1：cache-aware affinity scheduling + 一致哈希双绑定

把请求路由到「持有其 prefix cache」的集群。但为了避免「绑死单集群」的故障风险，用一致哈希 把每个 session 绑定到 primary + secondary 两个集群：

• primary 正常服务（持有该 session 的前缀缓存，命中率高）；
• secondary 在 primary 失效时接管，重新 prefill；
• 各 session 的 secondary 由一致哈希均匀分布 到整个 fleet------re-fill 的工作量被分散，而不是集中到故障点一处。

策略 2：budget-based admission control（预算准入）

为不同请求类分配独立资源预算 ，短类、长类各查各的配额，互不占用。突发长上下文只能消耗自己那份 capacity，不拖累其他类的 TTFT。

9.4 背后的原理

为什么一致哈希能「分散 re-fill 工作量」？ 一致哈希的最小扰动性质 （hopscotch）：移除/新增一个集群时，只有 O(1/N) 的 session 需要改路由。因为每个集群在环上放多个虚拟节点 （vnodes），session 落在哪个节点上是「均匀随机」的------去掉一个集群，只影响「原来指向它」的那部分 session，它们各自换到新的 primary/secondary，于是 re-prefill 的负载被均匀地摊到全 fleet，而不是某个集群被压垮。

为什么预算隔离能保住短请求 TTFT？ 因为短类和长类各占各的预算 。长请求突发时，它只能在长类的预算内接；超出的被拒绝/降级。短类配额不受影响，所以短请求的 TTFT 不会因为「后面排了一堆长请求」而飙升。这是「predictability（可预期性）」------报告用这个词表达：长突发牺牲自己的延迟，换取全系统短请求的 SLO 可预期。

9.5 落地架构与数据结构

ConsistentHashRing             # 一致哈希环
    _ring: [(pos, cluster_id)]   按位置升序
    add_cluster(id, vnodes=8)   每个集群放 8 个虚拟节点
    remove_cluster(id)          摘除该集群全部虚拟节点（最小扰动）
    route(session) -> (primary, secondary)   md5(session) 顺时针 → primary，
                                             再顺移取下一个不同集群 → secondary

Request (@dataclass)           # 一条请求
    session_id / request_class(short|long) / prefix / prompt_len / arrive_time

ClusterInstance                # 一个 serving 集群
    cache: PrefixCache         复用现有模块（真实命中判定）
    pool: PagePool             复用现有模块（页占用）
    register(tokens)           登记前缀 + persist KDA checkpoint
    reused_tokens(...)         PrefixCache.lookup + recompute_saved（真实命中）
    serve_cost(prompt_len, reused)  命中只算增量，miss 整段

BudgetController               # 预算准入控制
    budgets: {short, long}     各自独立配额
    weights: {short:1, long:3} 长请求占 3 份容量
    admit(class, usage)        该类再加一请求（占 weight 份）是否在预算内

FleetSimulator                 # 顶层事件驱动仿真
    run(requests)              按到达时间排序 → 预算 → 路由 → 命检 → 定价 → 统计
    fail_cluster(id)           令集群失效（摘环 + alive=False）
    results()                  hit_rate / total_cost / P95 TTFT / rejected / failover

Fleet 调度模块纯 CPU、纯标准库依赖（md5 + bisect），与 prefix_cache.py 同级独立。

一致哈希环：把 session 稳定映射到 (primary, secondary) 两个集群。虚拟节点保证各集群负载均匀，同时保证「移除/新增一个集群只扰动 O(1/N) 的 session」------这正是 failover 后 re-prefill 工作量被均匀分散的基础。

ClusterInstance 的关键点：命中判定不是 拍一个概率参数，而是委托现有 PrefixCache 的 lookup / recompute_saved ------ 因此这里的命中率与 serve 层真实模块一致（chained-hash 匹配 + KDA checkpoint 原子性），这是本仿真「可信」的来源：代价数字不是凭空编的，而是真实缓存行为折算出来的。

代价模型 ：work / throughput。命中时只重算 prompt_len - reused 的增量；未命中则整段重新 prefill。当 reused 接近 prompt_len 时，miss/hit 的代价比是数量级。

BudgetController：短请求 <2K 与超长 1M 混流，长度横跨 3 个数量级------若总容量固定且没有隔离，长请求突发会把短请求的 TTFT 全线拖垮。给每类独立预算后：长突发只能吃掉自己的配额，短请求配额不受影响。

权重：每类请求占用不同份容量（短=1、长=3），预算数字是「容量单位」而非「请求条数」。

时间模型（CPU 无法真跑 1M token，用公式折算）：

• prefill 代价 = serve_cost；TTFT = 集群当前累计负载 / throughput（队列近似）+ 本次 prefill 代价。
• 路由：session 绑定的 primary 存活走 primary；失效则走预分配的 secondary 并累加一个 re-prefill 成本。

run() 产出每请求明细；results() 汇总命中率 / 总代价 / 分请求类 P95 TTFT / 拒绝数 / failover / 负载标准差。

9.6 数据流 / 执行流

FleetSimulator.run(requests)：
  对每个请求（按 arrive_time 排序）：
    1. budget.admit(request_class, usage)   → 拒绝则 rejected[class]+1，跳过
    2. primary, secondary = ring.route(session_id)
       若 primary 失效 → 走 secondary，failover_count+1，re_prefill_cost += prompt_len
    3. reused = cluster.reused_tokens(prefix, prompt_len)   # 真实命中
       cost = cluster.serve_cost(prompt_len, reused)        # 命中=增量，miss=整段
    4. 先查后登记：cluster.register(prefix)   # 前缀在"这次请求之后"才对未来请求可见
       ttft = 集群累计负载/throughput + cost
    5. 更新统计

为什么「先查后登记」？ 这保证：首次请求必然 miss，第二次同前缀请求才命中------这正是「命中一次避免重算整个前缀」的语义来源。如果先登记再查，第一次请求就会命中「自己刚塞进去的前缀」，那是假命中。

9.7 反方辩手

「budget 会把合法请求拒绝掉？」 是的，长类超预算会被拒。但这是设计的权衡 ：牺牲长类突发延迟，换取全系统短请求 TTFT（SLO）可预期------这正是报告「predictability」一词的意思。代价不是 bug，是特性。

「一致哈希的 secondary 平时闲着不浪费吗？」 不浪费------它正是「primary 故障时的既定容器」。各 session 的 secondary 由一致哈希均匀分布，所以 re-fill 负载被均匀分散。这是「接受 secondary 闲置」换取「故障时负载分散」的权衡。

「CPU 上没法真跑 1M token，代价数字可信吗？」 用 cost = work / throughput 折算时间，命中/未命中的数量级差异由 token 数决定（400K vs 4K），与推理设备无关------比值仍是数量级（实测 ~781×）。数字是「真实缓存行为折算出来的」，不是拍脑袋。

9.8 何时用 / 何时不用

• 何时用：多实例 serving，请求共享前缀（亲和路由能命中）、流量长短混排（预算隔离能保 SLO）。
• 何时不用：单实例（没有「路由到哪个集群」的问题）；请求无共享前缀（亲和路由无处命中，退化成普通负载均衡）。

9.10 可视化

plot_fleet.py --- 集群级调度 (serve/fleet.py，报告 §5.4.3 / V3 附录 D)

• fig_affinity()：ConsistentHashRing 的 cache-aware 亲和路由 ------ 200 个 session 各发两次同前缀请求（有路由）vs 400 个不同 session（无路由），对比命中率；再用 serve_cost 对比 miss/hit 的 prefill 代价倍数（对数轴）
• fig_budget()：预算隔离 ------ 20 个长请求（1M token）与 60 个短请求混流，short_budget=long_budget=8 vs 无预算（1e9），对比短/长请求 P95 TTFT 以及各类被拒数量（长请求被限流、短请求零拒绝）
• 产出：plot_fleet_affinity.png、plot_fleet_budget.png

---

0x10 收官引擎：把它们咬合成一个「第二个请求到底省了什么」

10.1 它要解决什么问题

前 5 个组件各是独立子系统。收官引擎（serve/engine.py + experiments/serve/demo_engine.py）回答工程上一个很朴素的问题：第二个请求复用了同一个前缀时，到底省了什么？ 它把「统一页池 + 前缀缓存 + KCP」串成一个可以对照的 before/after。

10.2 总体思路与原理

• before（无前缀缓存）：第二个请求从头 re-compute------MLA 侧重新 prefill 全部 512 token（再占一页），KDA 侧重新走完整递推（重算 checkpoint）；
• after（有前缀缓存）：MLA 的 KV 页逐块命中直接复用，KDA 从 checkpoint 快照处恢复 O(1) 状态，免去全部重算。

省下的正是 recompute_saved 报告的 token 数（本例 512 vs 0）。run_once(hash_present) 是可导入的纯函数（无 IO、无随机），实验脚本与集成测试共用。

10.3 落地与实测

run_once 返回四项指标：reused_tokens（512 vs 0，本实验核心）、pages（首请求在统一页池占的页数，验证「一请求一份状态」）、comm_bytes（KCP 每层跨 8 rank 交换 (M̄,S̃) 的总字节，验证「与序列长度无关」）、pool_utilization（页池利用率）。

demo_engine.py (serve/engine.py + parallel/kcp.py)是命令行可视化入口，产出 before/after 对比图 + 通信字节对比。

• 表达：把前面各模块串成一次端到端引擎调用的三合一汇总图
• 测试：run_once(hash_present=True/False) 对比第二个请求的 reused_tokens（前缀复用）；MLA KV 字节随长度线性 vs KDAContextParallel.communication_volume 恒定；第三子图单独强调 KCP 通信量对序列长度是 O(1)
• 产出：demo_engine.png

---

附录

---

附录 A：跨仓库核对 --- group 语义与每层状态的归属

本附录回答一个反复出现的疑问：kda_checkpoints 到底是每层一份、每 token 一份，还是每个 group 一份？一个 PrefixCache 对应一个前缀吗？ 结论来自三处独立实现的交叉核对：kimi3-master、kimi-k3-pytorch-main、SGLang K3 分支。

A.1 kimi3-master（本项目镜像的来源）

kda_checkpoints: dict[int, set[int]] 是"记账本"，不是"储物架"。

# 原文注释：
#: Number of independent KDA cache groups a checkpoint must exist in.
n_kda_groups: int = 1
#: ``group -> set of token boundaries with a persisted KDA checkpoint``.
kda_checkpoints: dict[int, set[int]]

• n_kda_groups 默认 1；kda_checkpoints 的 key = 抽象 group 编号 ，value = 一组 token 边界位置（512 对齐的整数）。
• 一个 PrefixCache 存的是"大量 token 边界"，不是"一个前缀" ------value 是 set[int]，可以装上千个边界。"一个 ClusterInstance 只有一个 PrefixCache → 只有一个前缀"是误读。
• 真正的"每层一份状态"在 paged_cache.py ：kda_pages: dict[int, int]（layer -> page）与 store_kda(layer, state) / load_kda(layer, ...)。但 PrefixCache 与 PagedCache 互相之间没有任何关联代码------一个管"哪个边界有没有 checkpoint"，一个管"每层状态存到哪里"。

A.2 kimi-k3-pytorch-main（模块化参考）

状态确实每层独立一份，但仓库里完全没有前缀缓存机制。

• HybridLayerCache = (attention_type, state)，KDA 层配 KDAState，MLA 层配 MLACache；HybridBackboneCache 是每层一个 cache 的有序元组。
• 没有跨请求复用：推理流程是单次 prefill + 逐 token decode。
• README 明确："不包含 paged attention、continuous batching、quantized caches、distributed serving"。因此该仓库只提供状态对象与数学 参考（KDAState、recurrent_kda/chunkwise_kda），前缀缓存机制需要自行设计。
• "group"在此仓库只是架构拓扑的命名 （(kda, kda, kda, gated_mla) 模式重复 num_hybrid_groups 次），不进入任何向量数学。

A.3 SGLang K3 分支（生产实现）------ 前缀复用的事实真相

一条请求的 KDA 状态，是一个"把所有 KDA 层叠在一起"的大张量，"存放到某个 token 边界"。

# memory_pool.py：
temporal_state = torch.zeros(
    size=(num_mamba_layers, size + 1) + temporal_state_shape, ...)

• 注意第一维就是 num_mamba_layers ------一个 slot 装下"token 边界处所有 KDA 层的递归状态"。每层在自己维度，但整体作为一块分配与搬运。

前缀树（radix）在"对齐的 token 边界"存整包状态，而不是按层存：

• 在 len(prefix_indices)+extend_len 到达 mamba_cache_chunk_size 对齐点时，把当前状态写入 radix 缓存节点。
• radix 树每个节点持有 mamba_value（device 上的状态快照），命中前缀时 copy-on-write 把它装到新的活跃 slot 再继续解码。
• 缓存侧状态 int8 压缩（按 per-head/k-channel 对称量化）存放在独立 checkpoint 池；命中一次反量化一次，"从没重入一次量化-反量化循环"。

"group"字面量在 SGLang K3 分支全部 grep 为零 ------SGLang 的做法是把"所有 KDA 层的整包状态"当成一个不可分割的缓存值。

A.4 合成的结论

疑问	答案
kda_checkpoints 是每层一份吗？	不是。key=抽象 group 编号（默认 1），value=token 边界集合
每个 PrefixCache 只有一个前缀？	不是。value 是 set，可存无数 token 边界
每层的状态到底存哪？	kimi3-master：paged_cache.py 的 kda_pages（未接线）；kimi-k3-pytorch：每层 HybridLayerCache；SGLang：单 slot 内 dim=layers 的整包张量
生产里"一个边界"是什么？	一个 token 位置，对应"所有 KDA 层那一份整包状态快照"（radix 节点）
"group"字面量在生产中存在吗？	在 SGLang 不存在；在 kimi3-master 是设计抽象（默认 1），对应"整包被当作一个对象"

核心洞见 ：所有实现最终都收敛到同一件事------cache 的原子性单位 = "某一 token 边界的整包多层状态" 。kimi3 把这一体化思想用 group（抽象编号，默认 1）表达，SGLang 直接用一个"叠了所有层的 slot"。group 不是层的别名，而是------把"某一 token 处必须一起存在的东西"表达为一个原子单位的编号。

附录B: 与 vLLM KV Cache Group 的对比

当前项目中的group，与 vLLM 的 KVCache Group 有相似的"分组后共同管理/共同满足约束"的思想，但不是同一个抽象，也不存在一一对应关系。

更准确地说：

• 项目中的 group 是 KDA checkpoint 的逻辑原子性分区；
• vLLM 中的 KVCache Group 是 具体 cache layer/spec 的物理与调度管理分组。

因此，项目的 group 更接近"多个缓存状态必须同时可恢复"的约束；vLLM 的KVCache Group 则是"哪些层共享一个 block table、如何共同分配和布局缓存"的实现对象。

B.1. 项目中的 group 是什么

核心数据结构是：

kda_checkpoints: dict[int, set[int]]

它表示：

group_id -> 该 group 已保存 KDA checkpoint 的 token 边界集合

n_kda_groups 只决定需要维护多少个逻辑 group。代码没有记录 group 对应哪些 Transformer 层，也没有记录 group 的 page size、block table 或物理内存布局。

项目通过以下逻辑判断某个边界是否可恢复：

def checkpoint_available(self, boundary: int) -> bool:
    return all(
        boundary in group
        for group in self.kda_checkpoints.values()
    )

也就是说，边界 B 只有在所有 group 都有 checkpoint 时才可用。缺少任意一个group，该边界就不能用于恢复 KDA 状态，lookup() 会退回到更浅的可用边界。

因此，这里的 group 主要回答：

在 token 边界 B，恢复所需的所有状态分区是否都已经准备好？

它不直接回答：

这些状态应该放在哪个物理 block？
哪些层共用一个 block table？
每个 block 需要多少显存？

B.2. vLLM 中的 KVCache Group 是什么

vLLM 首先为每一层建立 KVCacheSpec，然后将适合共同管理的 cache layer 放进 KVCacheGroupSpec。

一个 vLLM KVCache Group 通常包含：

• 明确的 layer_names；
• 这些层各自的 KVCacheSpec；
• 共同使用的一张 block_table；
• 由组内各层 page size 共同决定的 block 大小；
• 与共享 BlockPool 配合的分配、回收和物理布局。

例如，MLA 层和 Indexer 层可以放进同一个 UniformTypeKVCacheSpecs，使一个 block 同时包含多个 MLA page 和多个 Indexer page。该 group 的 block 不是抽象的"checkpoint 存在性标记"，而是真实的缓存存储与寻址单位。

vLLM 的 group 划分还会受到以下因素影响：

• 层的 cache spec 是否兼容；
• block size 是否相同；
• page size 是否适合放在同一个物理 block；
• 不同层的生命周期是否一致；
• block table 和底层 kernel 的寻址、对齐要求。

因此，即使两层的 cache spec 看起来相同，也可能因为物理布局或 page size 约束被拆成不同的 vLLM group。

B.3. 逐项对比

维度	当前项目 PrefixCache 的 group	vLLM KVCache Group
主要对象	KDA checkpoint 集合	具体 cache layer 与 cache spec
主要目的	保证某边界的 KDA 状态可整体恢复	共同管理 cache block、block table 和物理布局
组成员是否显式记录	否，只记录整数 group ID	是，通过 layer_names 和 spec 记录
是否绑定 Transformer 层	不绑定	直接绑定
是否描述 page/block 大小	不描述；physical_block 只是抽象 ID	描述并参与显存分配和布局
是否有 block table	没有	有，每个 group 对应一张 block table
lookup 约束	所有逻辑 group 在同一边界都有 checkpoint	各 cache group 的命中结果需要协调到共同可执行长度
KDA 状态处理	用 checkpoint 边界集合抽象表示	通过 Mamba/KDA cache group 管理状态 block
物理内存语义	PrefixCache 中没有真实内存布局	直接关联共享 BlockPool 和物理 block
当前默认粒度	n_kda_groups=1，整个 KDA 状态视为一个逻辑组	根据层、spec、block size 和布局规则划分多个组

B.4. 两者为什么容易被认为相同

两者都体现了一个共同思想：某些缓存对象不能完全独立处理，需要把它们作为一个整体协调。

在当前项目中，这体现为：

所有 group 都有 checkpoint => 该边界可恢复

在 vLLM 中，这体现为：

相关 cache layer 放进同一个 KVCache Group
        => 共享 block table，并以统一方式分配/访问

两者还都服务于 hybrid attention 的 prefix cache：

• MLA/Full Attention 可以按 token block 进行前缀匹配；
• KDA/Mamba/Linear Attention 需要在某些边界保存可恢复状态；
• 最终可执行的 prefix hit 需要满足不同缓存部分的共同约束。

因此，概念上可以说它们都在表达"缓存复用的协调单元"，但协调的对象不同。

B.5. 最接近的对应关系

如果只从"必须共同满足条件"这个角度看，项目中的多个 group 可以类比为：

多个需要共同命中的缓存状态分区

而不是 vLLM 中的某一个 KVCache Group。

更接近的关系是：

项目的 group 集合
    ≈ 一组需要共同满足 checkpoint 条件的逻辑缓存分区

vLLM 的 KVCache Group
    ≈ 一组明确映射到层和 cache spec、共享物理管理结构的缓存对象

如果未来项目将每个 KDA 状态分区显式映射到具体层，并为每个分区实现独立的 cache spec、block table、物理 block 分配和回收，那么这些逻辑 group 才可能进一步演化为类似 vLLM KVCache Group 的对象。

当前代码尚未完成这一步，因此不能把：

n_kda_groups

直接解释为：

vLLM 中的 KVCache Group 数量

B.6. 关键差异：checkpoint 原子性 vs. 物理 cache 管理

6.1 项目 group 解决的是恢复正确性

项目的核心不变量是：

KDA 在边界 B 的完整状态必须可恢复

所以使用 all() 检查所有 group 是否都有 checkpoint。group 的边界是"缺少任意一份状态就不能恢复"的逻辑边界。

6.2 vLLM group 解决的是存储和调度一致性

vLLM 的核心问题是：

不同 cache layer 如何共享 block table、BlockPool 和物理布局

因此 vLLM group 的划分会考虑 page size、block size、内存碎片、kernel 寻址

和生命周期等工程因素。它不是简单地保存一个"该边界是否存在"的集合。

6.3 两者可以同时存在

在一个完整的 hybrid serving 实现中，二者可能同时出现：

1. vLLM KVCache Group 负责实际的 layer/block/table/显存管理；
2. 更高层的 checkpoint 原子性逻辑负责判断多个 group 的状态是否能共同恢复；
3. prefix-cache coordinator 再将 MLA 命中长度与 KDA checkpoint 边界对齐。

因此，二者不是互斥的替代关系，而可能是不同层次的抽象。

B.7. 最终判断

判断问题	结论
项目 group 是否就是 vLLM KVCache Group？	不是
是否存在设计思想上的对应？	存在，都用于协调一组不能完全独立处理的缓存对象
项目 group 是否对应具体 Transformer 层？	当前代码没有表达
项目 n_kda_groups 是否等于 vLLM group 数？	不能直接等同
当前项目默认 group 的含义	全部 KDA checkpoint 被视为一个逻辑原子整体
vLLM group 的核心含义	一组显式 cache layer/spec 共享 block table 和物理缓存管理

一句话总结：当前项目的 group 是"checkpoint 是否齐全"的逻辑分组；vLLM 的 KVCache Group 是"哪些 cache layer 共享 block 管理"的物理/调度分组。两者有抽象层面的关联，但当前项目中的 group 不能直接视为 vLLM 的KVCache Group。

---

---

0xFF 参考

AI Infra进阶：3万字深入理解线性注意力机制