10 KiB
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 已释放。详见 实现、结果与实际命令。
暂停 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
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。
这是原 PR 的特定布局地址规划复现，未在 GPU 上故意执行越界复制。当前实现已补齐每条目布局校验及 GQA head-aware 传输，CPU 测试覆盖全部 P-TP4→D-TP32 head 映射；修复版本的完整服务已通过跨组请求和 C1/C8 评测。原 flat-copy 失败证据保留用于说明修复动机。
5. 验收与结果口径
- 局部数学与协议验证已通过。 覆盖各 stage capture、feature 顺序、空 capture stage、单次 RMSNorm、BF16/FP32 分段投影与 TP4→TP32 传输规划。CPU 方法级回归不等于完整模型 logits 等价验证。
- 最小 PD 请求已通过。 P TP4/PP8、D TP32/PP1，经 Router 返回 128 token，服务日志有实际 RDMA ready ACK，D Decode 使用 CUDA Graph。
- GSM8K 固定 64 题 C1/C8 已完成。 两轮各 64 个响应均成功，独立解析最终答案也全部通过；总耗时 195.42/55.55 秒，无 512-token 截断。完整逐题输出、直方图、服务日志和资源清理证据见结果目录。客户端非流式，不把总耗时换算成 TTFT/ITL。
- 统计口径已复核。 接口的每题等权接受长度为 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 接受长度异常的根因。