10 KiB
Kimi-K3：从 PP + DSpark 迁移到 PP + DFlash
更新：2026-08-31；PP、传输与 MoE 集成首版。
1. 决策与范围
优先推进 DFlash 的 PP + PD 适配。 #33863 的分段投影设计能复用于 DFlash，但需要补齐 DFlash worker、调度入口和异构 TP 下的 draft KV 传输。
当前已实现 DFlash PP worker、Kimi 跨 stage capture、PD 输入衔接、异构 TP 的 draft KV 传输，以及 Kimi FlashInfer 布局/SiTU 接入。首轮 P 启动暴露普通预热误带 verify metadata 的问题，已修复并增加回归，23 项 CPU 测试通过。SM120 编译及 Kimi SiTU GPU 集成回归已通过。修复镜像已同步全部八节点，Run pd-dflash-pwarm1-20260831-1652 正在重新验证服务启动；尚无 PD 请求和 GSM8K 结果。详见 实现进度与命令。
暂停 EAGLE3 baseline。KV cache 保持 BF16，chunk 保持 8192。FlashInfer 使用官方 0.6.18 加已合并 #4460 的显式 backport；原版 0.6.18 wheel 尚未包含所需 CUTLASS SiTU 接口。
目标拓扑：P 使用 601–604 的 TP4/PP8/EP4 与已适配的 FlashInfer MXFP4；D 使用 605–608 的 TP32/PP1，接入 DFlash。这里适配的是 P 侧 PP 生成 draft 上下文，D 侧投机执行仍为 PP1。
审计固定版本：
| 对象 | 版本 |
|---|---|
| SGLang main | 3139ceaeec50868a441d68dc231663f3777e0d93 |
| PR #33863 head | 6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5 |
| PR 状态 | Open；维护者要求拆分 DSv4 PP/PD、DSpark、Kimi/线性注意力工作 |
PR 公布的准确性数据来自 DSv4 Flash / H20，不能代替本次 Kimi-K3 / 6000D / P-TP4 与 D-TP32 验证。PR #33863
2. 为什么可以迁移
DSparkDraftModel 继承 DFlashDraftModel。普通 Kimi draft 的 prompt 上下文处理都包含：采集多个 target 层的 hidden，拼接后线性投影，再归一化，生成 draft KV。
#33863 将这一步改为：
各 P/PP stage：采集自己负责的 hidden → 使用对应权重列做局部投影
PP stage 间：传递并累加投影结果 [token 数, 7168]
最后一个 P/PP stage：统一 RMSNorm → 各 draft 层 K/V 投影、K norm、RoPE → 写 draft KV
P → D：传 target KV、KDA 状态及 draft KV
D：使用已有 DFlash proposer / verify / commit 流程
依据是线性运算 concat(h_i) × W^T = sum(h_i × W_i^T)。归一化必须在求和后做。BF16 分段累加会改变舍入顺序，因此需要比较中间张量、logits 和 token 接受行为。
关键接口已经加在共用 DFlash 模型中：project_target_hidden_partial。最后阶段的归一化及 KV 写入实现位于 write_projected_context_kv。
这一路线保留 P 侧 PP8 的收益，不要求把各层原始 hidden 通过 PD 网络传给 D。
3. 需要补哪些代码
以下路径均相对于 SGLang 的 python/sglang/srt/。
| 模块 | 已有基础 | 本次增量 |
|---|---|---|
models/kimi_k3.py |
PR 已支持各 PP stage 局部采集 DSpark 特征 | 接通 DFlash capture setter，按 DFlash checkpoint 的层号及 Kimi residual 语义采集 |
models/dflash.py |
PR 已有局部线性投影 | 复用投影；最后阶段归一化后复用 DFlash 的逐层 KV 写入路径 |
speculative/dflash_worker_v2.py |
非 PP prefill 的 hidden → draft KV | 接收、转发 PP proxy；仅最后阶段写完整 draft KV；处理非末 stage 没有 logits/next token 的情况 |
managers/scheduler_pp_mixin.py |
PR 的成功交集、失败并集及一致释放机制 | 将目前 DSpark 专用的 proxy/input 衔接扩展到 DFlash，保持请求顺序和完成条件一致 |
speculative/spec_info.py |
当前 PD 输入分派只有 EAGLE、DSpark | 接入 DFlash 输入构造，并验证 overlap FutureMap 与 idle batch 生命周期 |
disaggregation/prefill.py、utils.py、mooncake/conn.py |
最后 PP stage 传 draft KV、按层号配对 | 加入 draft GQA 的 head 分片/复制与收发端独立 stride，保留 target MLA/KDA 的已有路径 |
arg_groups/speculative_hook.py、validation_hook.py |
当前 main 拒绝 PP + DFlash | 仅放开已实现的 PD-prefill + PP 组合；D 端 PP1 范围不变 |
关键证据：
- Kimi 局部层采集。
- DSpark PP 上下文处理。
- PP scheduler 的 DSpark 专用入口。
- main 的 PD 输入分派：DFlash 仍返回 None。仓库里的
dflash_disaggregation.py辅助函数尚未被这个入口调用，不能仅凭文件存在认定 PD 已支持。
相关 #36140 记录了输入缺失和 draft 状态缺失问题；#36277 是提前拒绝不支持配置的保护性改动，不是 DFlash PD 的实现。
4. 异构 TP 的明确风险
605 的 /data/hf_models/Kimi-K3-DFlash/config.json 为 6 层 draft、8 个 KV head、head_dim=128、BF16、4096 sliding window；target capture 层是 [19,37,54,66,78,90]。
DFlashAttention 使用 max(1, total_kv_heads // tp_size)：
| draft 布局 | 每 rank KV heads | BF16、page=64 时单层 K 的页大小 |
|---|---|---|
| P TP4 | 2 | 64 × 2 × 128 × 2 = 32768 bytes |
| D TP32 | 1，按 head 复制到多个 rank | 64 × 1 × 128 × 2 = 16384 bytes |
PR 当前 Mooncake flat 传输分支 按 layer ID 配对后，对源偏移、目标偏移和复制字节数都使用源端 item_len。Kimi hybrid MLA 会进入这条路径。层号正确只能确认“哪层到哪层”，没有解决“该层哪些 head 到哪个 rank”。
CPU 验证结果
从该 PR 源码直接用 AST 提取 _send_kvcache_generic 和 build_transfer_entry_pairs，使用 synthetic 地址、单页索引和记录型传输函数执行。未调用 GPU、Mooncake 或 RDMA。
| 用例 | 地址/长度是否符合目标布局 |
|---|---|
| 收发双方均为 1 KV head | 符合 |
| P 为 2 head、D 为 1 head | 不符合：目标偏移和复制长度均按源端大条目计算 |
具体用例：D buffer 起始地址设为 2000000，目标 page ID=2。正确目标地址为 2000000 + 2×16384 = 2032768，当前函数生成 2065536，复制 32768 bytes 而不是单 head 页的 16384 bytes。
这是特定布局的地址规划复现，尚未启动完整服务复现。迁移时必须加入每条目的布局校验，并实现 GQA head-aware 传输；不能直接复用 flat copy。还需确认 D 接收到全部 target/draft 组件后才进入首轮 draft。
5. 实现与验收顺序
- 局部数学与协议验证。 检查各 stage capture 层覆盖、feature 顺序、空 capture stage、单次 RMSNorm；比较拼接投影与分段投影误差。测试 TP4→TP32 的全部 8 个 draft KV head、K/V、6 层、页索引以及失败清理。源/目标布局不兼容时在传输前报错。
- 最小 PD 请求。 保留 P TP4/PP8、D TP32/PP1，先 C1。核对真实 buffer metadata、draft KV 到达情况、首轮 draft logits，以及成功/失败后的资源释放。
- GSM8K 固定 64 题。 相同 prompt/template、采样设置，依次测试 C1/C8；保留准确率、完成率、接受长度、接受直方图、Output TPS、总耗时与显存状态。当前客户端非流式，不把总耗时换算成 TTFT/ITL。输入不裁剪，输出上限 512，报告截断，不新增 16K synthetic 测试。历史 5-shot 与前 64 题有重叠，本轮沿用以便部署回归，不能作为独立无泄漏的模型准确率评估。
- 验收后决定继续投入。 DFlash 正确性与部署通过且有实际 Decode 收益，继续 PD 调优；若关键适配无法通过，或接受行为同样异常，切换到下面的 D-only 诊断。
真实实验代码继续维护在 601 的 /data/hzy/sskj 工作区，不改同事部署目录，不再创建分支。本轮核实当前检出名为 hzy-kimi3-pd-pp8-standard，保持现状；独立 SGLang 源码位于 /data/hzy/src/sglang-kimi-pp-dflash-33863。
6. DSpark 接受长度诊断的备用路径
DSpark 低接受长度的问题先独立于 PD 排查。PP + PD 适配本身不会自动修复 D-only 已存在的问题。
| 检查项 | 目的 |
|---|---|
| checkpoint、tokenizer、mask token、target_layer_ids、RoPE | 排除 draft/target 配置失配及错误 hidden 来源 |
| 相同 token 前缀下，逐位置 draft token 与 target logits | 找到拒绝从第几个位置开始，区分首 token 错位与持续预测质量问题 |
| verify 的 KDA/SSM 状态更新、回滚与 token 位置 | 判断是否首次可用、随后状态漂移 |
| fused/replay 路径与参考执行比较 | 将 kernel/状态管理错误与 drafter 本身质量分开 |
| 接受统计定义、实际 proposal 数、bonus token | 统一接受长度口径，防止指标解释错误 |
605 的 DSpark checkpoint 为 block=7、5 层、capture [7,23,51,67,83]；DFlash 为 block=16、6 层、另一组 capture 与 RoPE。二者的原始平均接受长度不能单独判断实现异常，需同时看每次实际提议数、逐位置接受比例和实际吞吐。
本轮不重跑 no-spec 性能基线，保留此前结果作为参考。当前尚未确定 DSpark 接受长度异常的根因。