sskj/experiments/pro6000/kimi3_pro6000_pd_dflash_validation

Kimi-K3 PP + DFlash 适配进度

2026-08-31PP worker、PD 输入衔接、异构 TP 传输和 Kimi FlashInfer 接入已完成首版SM120 编译及 Kimi SiTU GPU 集成回归通过。两项启动兼容修复及 24 项 CPU 回归已完成。Run pd-dflash-kvbounds1-20260831-1740 的 P 组已通过预热、Mooncake 初始化和健康检查HTTP 200D 组正在启动;尚无跨组请求或 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 poolP 不执行 draft decode CUDA GraphD 的原有初始化路径保留
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/EP4D TP32/PP1BF16 KVchunk=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-cubinflashinfer-jit-cache保留版本校验SM120 CUTLASS 通过官方 JIT 路径编译。安装说明#4460

CPU 验证

# 独立传输测试,无需 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:e74c9b1e7fb6a9509db32576d2c9196b42ef329118909eeebd620cb9db2f91e9image config 为 sha256:26584138eb1e097183ce167cbffe4198be0611e0cfde7422cbe692100d865982OCI index 为 sha256:7aa2be2d7af6cb11fea6d0c00d9bcd323ec986095dc9c4d3c997cff15ccb3b3f。区分这三个摘要,不将 Docker 不同模式显示的 ID 当成代码不一致。

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

  • pp_and_transport_cpu_601_20260831.log18 项 CPU 回归。
  • patched_worker_import_20260831.logpatched_launch_help_20260831.log:真实模块导入与 CLI。
  • image_build_situ4460_cpucheck_20260831.log:集成镜像的完整 API/回归检查。
  • command_parse_20260831.logP/D 全部 8 条命令通过 ServerArgs argparse 检查;该检查不执行 ServerArgs 后处理或服务初始化。
  • args_resolve_20260831.logresolved_p_20260831.jsonresolved_d_20260831.jsonP/D 真实模型配置通过 resolve_once() 后处理;保持 BF16 KV、8K chunk 和 FlashInfer。Kimi 投机验证自动选择 nv_cutedslP 侧 PP 自动关闭 overlap scheduler该检查不加载模型权重。
  • sm120_cutlass_compile_20260831.logSM120 预编译进度。首次构建被外部 Docker SIGKILL 终止,退出 137Docker 事件为显式 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.logfinal_image_gpu_container_20260831.json:最终镜像独立 GPU 复测、Mooncake 导入、镜像 ID 和挂载证据。
  • pd-dflash-20260831-1617/logs/p_0.logp_3.log:首轮 P 预热报错原始证据;该 Run 未进入请求测试,退出码 1。
  • image_build_pwarm1_20260831.log:预热修复镜像构建及 23 项 CPU 回归。
  • pwarm1_602_cpu.logpwarm1_602_inspect.json:约 49 MiB 增量包导入后的容器回归与平台摘要。
  • pd-dflash-pwarm1-20260831-1652/logs/p_0.logP 预热完成后,在初始化传输管理器时发现 HybridLinearKVPool.end_layer 缺失;未进入 D 启动和请求测试。
  • image_build_kvbounds1_20260831.logkvbounds_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 次数和输出结束原因归档。

# 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 重试。正式顺序为 startsmokebenchlogsstop,也可用 all 串行执行。所有操作使用同一个 RUN_ID。失败时保留本任务容器和日志,入口不会自动杀其他实验。

配置选择P 为 TP4/PP8/EP4D 为 TP32/PP1/EP4两侧 FlashInferKV 为 BF16chunk=8192page=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=/cacheFlashInfer、Triton、PyTorch 扩展和 CUDA 编译缓存也指向该挂载目录,TMPDIR=/cache/tmp,宿主路径见 config.envJIT_CACHE,避免容器可写层占满根盘;预检会创建所需临时目录。

首轮启动修复

base_runner.py::_dummy_run 已将 PD Prefill target 设为普通 Decode 预热,但随后仍创建 DFlashVerifyInput。普通 Triton Attention 读取 kv_indptr 时因此报 AttributeError。修复让该分支的 spec_info=NoneD 侧真正的 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 镜像尚未删除。