120 lines
10 KiB
Markdown
120 lines
10 KiB
Markdown
# Kimi-K3：从 PP + DSpark 迁移到 PP + DFlash
|
||
|
||
更新：2026-08-31；PP8 + PD + DFlash 集成与 GSM8K 64 题 C1/C8 验证完成。
|
||
|
||
## 1. 决策与范围
|
||
|
||
**优先推进 DFlash 的 PP + PD 适配。** #33863 的分段投影设计能复用于 DFlash，但需要补齐 DFlash worker、调度入口和异构 TP 下的 draft KV 传输。
|
||
|
||
已实现 DFlash PP worker、Kimi 跨 stage capture、PD 输入衔接、异构 TP 的 draft KV 传输及 Kimi FlashInfer 布局/SiTU 接入。24 项 CPU 回归及 SM120 GPU 集成回归通过。Run `pd-dflash-kvbounds1-20260831-1740` 已完成跨组 smoke 与 GSM8K 64 题 C1/C8：两轮各 64/64 请求和答案检查通过，均无输出截断；按验证步数加权的接受长度为 6.4211/6.5158，整体 Output TPS 为 33.38/117.56。八节点实验容器已清理、GPU 已释放。详见 [实现、结果与实际命令](../experiments/pro6000/kimi3_pro6000_pd_dflash_validation/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；D 使用 605–608 的 TP32/PP1/EP4，两组均为 FlashInfer MXFP4。这里适配的是 **P 侧 PP 生成 draft 上下文**，D 侧投机执行仍为 PP1。
|
||
|
||
前期源码审计快照（PR 状态为该次审计记录，不表示实时状态）：
|
||
|
||
| 对象 | 版本 |
|
||
|---|---|
|
||
| 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。
|
||
|
||
这是原 PR 的特定布局地址规划复现，未在 GPU 上故意执行越界复制。当前实现已补齐每条目布局校验及 GQA head-aware 传输，CPU 测试覆盖全部 P-TP4→D-TP32 head 映射；修复版本的完整服务已通过跨组请求和 C1/C8 评测。原 flat-copy 失败证据保留用于说明修复动机。
|
||
|
||
## 5. 验收与结果口径
|
||
|
||
1. **局部数学与协议验证已通过。** 覆盖各 stage capture、feature 顺序、空 capture stage、单次 RMSNorm、BF16/FP32 分段投影与 TP4→TP32 传输规划。CPU 方法级回归不等于完整模型 logits 等价验证。
|
||
2. **最小 PD 请求已通过。** P TP4/PP8、D TP32/PP1，经 Router 返回 128 token，服务日志有实际 RDMA ready ACK，D Decode 使用 CUDA Graph。
|
||
3. **GSM8K 固定 64 题 C1/C8 已完成。** 两轮各 64 个响应均成功，独立解析最终答案也全部通过；总耗时 195.42/55.55 秒，无 512-token 截断。完整逐题输出、直方图、服务日志和资源清理证据见结果目录。客户端非流式，不把总耗时换算成 TTFT/ITL。
|
||
4. **统计口径已复核。** 接口的每题等权接受长度为 6.9433/7.0070；直方图按 verify 步数加权为 6.4211/6.5158。两者权重和首尾 token 处理不同，不混用。5-shot 与前 5 题重叠，排除后剩余 59 题也均通过；这仍是历史口径部署回归，不是独立无泄漏的模型能力评测。本轮不重跑 no-spec，因此不宣称投机加速倍数。
|
||
|
||
真实实验代码继续维护在 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 接受长度异常的根因。
|