# Kimi-K3 PP + DFlash 适配进度

2026-08-31：PP worker、PD 输入衔接、异构 TP 传输和 Kimi FlashInfer 接入已完成首版；SM120 编译及 Kimi SiTU GPU 集成回归通过。两项启动兼容修复及 24 项 CPU 回归已完成。Run `pd-dflash-kvbounds1-20260831-1740` 的 P 组已通过预热、Mooncake 初始化和健康检查（HTTP 200），D 组正在启动；尚无跨组请求或 GSM8K 结果。

## 源码和环境

- 基于官方 PR #33863 的 `6465a6f3d3b6c8b7fee40fba0fdc09cf5e9ca1c5`。
- 上游 Python 目录的 4619 个文件均与 Git blob SHA1 校验一致。
- `pp_dflash_integration.patch` 包含完整 SGLang 增量；`pp_dflash_files.json` 记录文件 SHA256。`mixed_kv_transfer.patch` 保留传输层首版记录。
- 601 独立源码路径：`/data/hzy/src/sglang-kimi-pp-dflash-33863`。
- 601 实验路径：`/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_pd_dflash_validation`。
- 不修改原 EAGLE3/DSpark 部署；不创建新分支。

## 本轮实现

| 部分 | 实现 |
|---|---|
| Kimi hidden 采集 | 接入 DFlash checkpoint 的 capture 层；PP 边界由下一 stage 使用实际 residual 聚合权重采集 |
| 分段投影 | 各 stage 计算对应权重列的局部投影，通过 PP proxy 累加；末 stage 统一 RMSNorm 并写 6 层 draft KV |
| P 侧 draft 执行 | 非末 stage 只保留最小 KV pool；P 不执行 draft decode CUDA Graph，D 的原有初始化路径保留 |
| Scheduler 与 PD | 接通 DFlash proxy、next draft input、FutureMap 发布与 idle 生命周期；仅放开 P 侧 PP |
| Mooncake 传输 | 注册每 entry 的 dtype/head/page stride/容量；target MLA 整页传输，draft GQA 按 head 交集逐 token 切片；写入前完成全部边界检查 |
| FlashInfer | 迁移已验证的 Kimi SM120 布局与 SiTU 参数适配，同时保留新上游 SwigluStep 行为 |

目标范围：Kimi hybrid MLA + 普通 NHD DFlash draft KV、Mooncake、CP=DCP=1，无 staging/unified KV/HiSparse。P TP4/PP8/EP4；D TP32/PP1；BF16 KV；chunk=8192。

## FlashInfer 依赖

官方 `#4460` 已于 2026-08-21 合并，提交为 `b460bc00cb373541102d2155aec35bd626e522ce`。但 `v0.6.18` 属于另一条发布分支，实际下载的官方 wheel 没有 CUTLASS `situ_beta/situ_linear_beta` API。

因此本镜像使用 **官方 0.6.18 wheel + 未改写的 #4460 合并补丁**。Dockerfile 对 2 个 Python 和 5 个 C++ 文件先执行 `git apply --check`，再应用上游补丁。它不是未经修改的官方 0.6.18，也不是旧的私人 SiTU kernel。

移除镜像内旧版 `flashinfer-cubin`、`flashinfer-jit-cache`，保留版本校验；SM120 CUTLASS 通过官方 JIT 路径编译。[安装说明](https://docs.flashinfer.ai/installation.html)、[#4460](https://github.com/flashinfer-ai/flashinfer/pull/4460)。

## CPU 验证

```bash
# 独立传输测试，无需 GPU 或 torch
cd /data/hzy/src/sglang-kimi-pp-dflash-33863
python3 test/registered/disaggregation/test_mixed_kv_entry_layout.py

# 全部 CPU 回归需要带 torch 的容器；镜像构建时自动执行
python3 -m unittest discover -s test/registered/disaggregation -v
```

| 回归 | 数量 | 主要覆盖 |
|---|---:|---|
| 传输 | 7 | 144 组 KV-head/TP 组合、实际 P-TP4→D-TP32 全 32 rank、逐字节复制、越界/重复写拒绝、注册往返 |
| PP 与 PD 输入 | 13 | PP1/2/4/8/16 capture 归属、边界 residual、空 capture stage、BF16/FP32 分段投影、单次 norm、末 stage 写 KV、输入生命周期、P 普通预热不创建 verify metadata、混合池/普通池传输层范围 |
| Kimi MoE 合并 | 4 | 两类 gate/up 布局、SiTU 参数、activation 白名单、非连续输入、保留 SwigluStep、API 能力检查 |

24 项在本地与 601 `kvbounds1` 镜像中通过。此前 `pwarm1` 镜像的 23 项也在 602 导入容器中通过，并完成真实 DFlashWorkerV2、MooncakeKVManager 导入。8 条启动命令的 CLI 解析及两类服务的参数后处理已检查。CPU 测试对实际方法作 AST 提取，使用 CPU tensor 或记录型 engine，核验数学与调用契约，不替代服务级验证。

601 GPU6 的 `test_kimi_k3_sm120_situ_layout_and_noncontiguous_input` 已通过：检查 Kimi gate/up 及 scale 布局、SiTU(4,25)、非连续输入，以及 SGLang adapter 与直接 FlashInfer 调用的输出一致性。这是小 shape 的集成回归；端到端正确性由后续 PD/GSM8K 检验。

修复前镜像 `local/sglang:kimi-k3-pp-dflash-33863-fi0618-situ4460` 的独立 GPU 复测通过，Mooncake CUDA engine 导入成功，GPU 测试 1 passed、25.28 秒（复用 JIT 缓存）。

当前镜像为 `local/sglang:kimi-k3-pp-dflash-33863-fi0618-situ4460-kvbounds1`，在 `pwarm1` 上只替换 Python 层范围处理及回归测试，不改变 FlashInfer 层或 GPU kernel。Linux/amd64 manifest 为 `sha256:e74c9b1e7fb6a9509db32576d2c9196b42ef329118909eeebd620cb9db2f91e9`，image config 为 `sha256:26584138eb1e097183ce167cbffe4198be0611e0cfde7422cbe692100d865982`；OCI index 为 `sha256:7aa2be2d7af6cb11fea6d0c00d9bcd323ec986095dc9c4d3c997cff15ccb3b3f`。区分这三个摘要，不将 Docker 不同模式显示的 ID 当成代码不一致。

601 证据均位于实验目录 `results/`：

- `pp_and_transport_cpu_601_20260831.log`：18 项 CPU 回归。
- `patched_worker_import_20260831.log`、`patched_launch_help_20260831.log`：真实模块导入与 CLI。
- `image_build_situ4460_cpucheck_20260831.log`：集成镜像的完整 API/回归检查。
- `command_parse_20260831.log`：P/D 全部 8 条命令通过 `ServerArgs` argparse 检查；该检查不执行 ServerArgs 后处理或服务初始化。
- `args_resolve_20260831.log`、`resolved_p_20260831.json`、`resolved_d_20260831.json`：P/D 真实模型配置通过 `resolve_once()` 后处理；保持 BF16 KV、8K chunk 和 FlashInfer。Kimi 投机验证自动选择 `nv_cutedsl`，P 侧 PP 自动关闭 overlap scheduler；该检查不加载模型权重。
- `sm120_cutlass_compile_20260831.log`：SM120 预编译进度。首次构建被外部 Docker SIGKILL 终止，退出 137；Docker 事件为显式 kill，无 OOM 事件。已编译的对象文件保留在独立缓存。
- `sm120_cutlass_compile_resume_20260831.log`：恢复后编译成功，`SM120_CUTLASS_BUILD_OK`，退出 0。
- `gpu_situ_smoke_detached_20260831.log`：首次完整单测发现 SiTU activation 白名单遗漏；已修复并增加 CPU 回归。
- `gpu_situ_smoke_fix_20260831.log`：修复后 GPU 回归通过，1 passed；首次 JIT 在内的总用时 771.10 秒。
- `image_build_activation_fix_20260831.log`：固化修复后的镜像构建及 22 项 CPU 回归。
- `final_image_gpu_verify_20260831.log`、`final_image_gpu_container_20260831.json`：最终镜像独立 GPU 复测、Mooncake 导入、镜像 ID 和挂载证据。
- `pd-dflash-20260831-1617/logs/p_0.log` 至 `p_3.log`：首轮 P 预热报错原始证据；该 Run 未进入请求测试，退出码 1。
- `image_build_pwarm1_20260831.log`：预热修复镜像构建及 23 项 CPU 回归。
- `pwarm1_602_cpu.log`、`pwarm1_602_inspect.json`：约 49 MiB 增量包导入后的容器回归与平台摘要。
- `pd-dflash-pwarm1-20260831-1652/logs/p_0.log`：P 预热完成后，在初始化传输管理器时发现 `HybridLinearKVPool.end_layer` 缺失；未进入 D 启动和请求测试。
- `image_build_kvbounds1_20260831.log`、`kvbounds_image_context/`：复用原 FlashInfer 层、只修改 Python 的构建配方及 24 项回归。

已通过 compileall、Black 和 Ruff 的未定义变量/语法检查。

## 实验入口

代码位于 601 的 `/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_pd_dflash_validation`。只运行 `deploy_pd_dflash.sh`，它负责 P/D 启动、健康检查、Router、smoke 和评测。`bench_gsm8k_acceptance.py` 沿用原客户端的 prompt 与请求设置，将结果标签改为 DFlash，并增加原始 `meta_info`、verify 次数和输出结束原因归档。

```bash
# 601：命令检查，不需要 GPU 或 sudo
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_pd_dflash_validation
DRY_RUN=1 RUN_ID=command-audit-20260831 bash deploy_pd_dflash.sh all
```

此前两轮均在 P 侧启动阶段退出，日志已保存、失败容器已清理。修复镜像同步完成后使用新 Run ID 重试。正式顺序为 `start` → `smoke` → `bench` → `logs` → `stop`，也可用 `all` 串行执行。所有操作使用同一个 `RUN_ID`。失败时保留本任务容器和日志，入口不会自动杀其他实验。

配置选择：P 为 TP4/PP8/EP4，D 为 TP32/PP1/EP4，两侧 FlashInfer；KV 为 BF16，chunk=8192，page=64。为 C1/C8 评测将活跃请求与 Decode Graph 上限设为 8；未额外改变模型 context 上限。

GSM8K 为原数据集前 64 题，沿用历史客户端的 5-shot、temperature=0、输出上限 512，分别运行 C1/C8。输入不裁剪。历史 5-shot 也取自测试集前 5 题，和本次 64 题有重叠：结果适用于与历史流程的部署回归，不作为独立无泄漏的模型准确率评估。逐题输出和截断情况保留。

## 后续

确认 8 节点模型、镜像与网络一致，然后运行 C1 smoke 和固定 GSM8K 64 题 C1/C8。保留逐题答案、截断、接受直方图、总耗时、吞吐、原始 `meta_info` 与服务日志；不重跑 EAGLE3/no-spec 基线。当前客户端为非流式请求，不把客户端总耗时当作 TTFT 或 ITL。

运行时设置 `SGLANG_CACHE_DIR=/cache`，FlashInfer、Triton、PyTorch 扩展和 CUDA 编译缓存也指向该挂载目录，`TMPDIR=/cache/tmp`，宿主路径见 `config.env` 的 `JIT_CACHE`，避免容器可写层占满根盘；预检会创建所需临时目录。

## 首轮启动修复

`base_runner.py::_dummy_run` 已将 PD Prefill target 设为普通 Decode 预热，但随后仍创建 `DFlashVerifyInput`。普通 Triton Attention 读取 `kv_indptr` 时因此报 `AttributeError`。修复让该分支的 `spec_info=None`，D 侧真正的 TARGET_VERIFY 路径不变。没有强行添加字段，也没有修改 DFlash 接受算法。

第二轮通过预热后，`PrefillBootstrapQueue._init_kv_manager` 直接访问混合池不存在的 `end_layer`。修复以现有的完整注意力层偏移加 `full_layer_nums` 得到传输范围，KDA 状态和 draft KV 仍分别使用明确的全局层 ID。新增回归覆盖 MLA/MHA 混合池、空完整注意力 stage、普通池和 layer-shard 分支，不将 K/V 两组条目误计为两倍层数。

镜像同步须检查 Docker 所在根盘，而不只检查模型盘 `/data`。本镜像层展开约 34.4GB，压缩内容约 15GB；共享层会减少增量占用。607 已按用户授权删除两个无容器引用的 vLLM 镜像，根盘恢复到约 107GB，原镜像元数据保存在 `results/607_vllm_images_before_delete_20260831.json`。用户清理 606 后其根盘恢复到约 91GB。601 使用 `ctr images export` 将既有 OCI 压缩层直接导出到 `/data`，其旧 vLLM 镜像尚未删除。