# 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 结果。详见 [实现进度与命令](evidence/kimi_k3_pd_dflash/README.md)。

暂停 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](https://github.com/sgl-project/sglang/pull/33863)

## 2. 为什么可以迁移

DSparkDraftModel 继承 DFlashDraftModel。普通 Kimi draft 的 prompt 上下文处理都包含：采集多个 target 层的 hidden，拼接后线性投影，再归一化，生成 draft KV。

#33863 将这一步改为：

```text
各 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](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/models/dflash.py#L683)。最后阶段的归一化及 KV 写入实现位于 [write_projected_context_kv](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/models/dspark.py#L751)。

这一路线保留 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 局部层采集](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/models/kimi_k3.py#L2912)。
- [DSpark PP 上下文处理](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py#L655)。
- [PP scheduler 的 DSpark 专用入口](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/managers/scheduler_pp_mixin.py#L1391)。
- [main 的 PD 输入分派](https://github.com/sgl-project/sglang/blob/3139ceaeec50868a441d68dc231663f3777e0d93/python/sglang/srt/speculative/spec_info.py#L174)：DFlash 仍返回 None。仓库里的 `dflash_disaggregation.py` 辅助函数尚未被这个入口调用，不能仅凭文件存在认定 PD 已支持。

相关 [#36140](https://github.com/sgl-project/sglang/issues/36140) 记录了输入缺失和 draft 状态缺失问题；[#36277](https://github.com/sgl-project/sglang/pull/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](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/models/dflash.py#L142) 使用 `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 传输分支](https://github.com/sgl-project/sglang/blob/6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5/python/sglang/srt/disaggregation/mooncake/conn.py#L686) 按 layer ID 配对后，对源偏移、目标偏移和复制字节数都使用源端 `item_len`。Kimi hybrid MLA 会进入这条路径。层号正确只能确认“哪层到哪层”，没有解决“该层哪些 head 到哪个 rank”。

### CPU 验证结果

从该 PR 源码直接用 AST 提取 `_send_kvcache_generic` 和 `build_transfer_entry_pairs`，使用 synthetic 地址、单页索引和记录型传输函数执行。未调用 GPU、Mooncake 或 RDMA。

[地址规划结果与源文件 SHA256](evidence/kimi_k3_pp_dflash_pr33863_transfer_plan.json)。

| 用例 | 地址/长度是否符合目标布局 |
|---|---|
| 收发双方均为 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. 实现与验收顺序

1. **局部数学与协议验证。** 检查各 stage capture 层覆盖、feature 顺序、空 capture stage、单次 RMSNorm；比较拼接投影与分段投影误差。测试 TP4→TP32 的全部 8 个 draft KV head、K/V、6 层、页索引以及失败清理。源/目标布局不兼容时在传输前报错。
2. **最小 PD 请求。** 保留 P TP4/PP8、D TP32/PP1，先 C1。核对真实 buffer metadata、draft KV 到达情况、首轮 draft logits，以及成功/失败后的资源释放。
3. **GSM8K 固定 64 题。** 相同 prompt/template、采样设置，依次测试 C1/C8；保留准确率、完成率、接受长度、接受直方图、Output TPS、总耗时与显存状态。当前客户端非流式，不把总耗时换算成 TTFT/ITL。输入不裁剪，输出上限 512，报告截断，不新增 16K synthetic 测试。历史 5-shot 与前 64 题有重叠，本轮沿用以便部署回归，不能作为独立无泄漏的模型准确率评估。
4. **验收后决定继续投入。** 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 接受长度异常的根因。