120 lines
10 KiB
Markdown
120 lines
10 KiB
Markdown
# 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 接受长度异常的根因。
|