Kimi K3 推理优化 - 优化 Prefill - Deep PP
1. 工作背景
Kimi-K3 在 4 台 RTX 6000D、共 32 张 GPU 上部署时，最直接的方式是
PP1 / TP32 / EP4：完整模型只有一个 Pipeline Stage，Tensor Parallel
覆盖全部 32 张卡，MoE Expert Parallel 使用 4 张卡。
这套配置可以正确运行，但 16K 长输入的 Prefill 延迟很高。既有 Nsight
Systems Profile 显示，在 16K -> 1、C=8、chunk=8K 场景中：
- 稳定 GPU 的 Prefill 窗口中，NCCL 暴露时间占比中位数为 57.27%；
- 单个 Profile 窗口约有 557 次 AllReduce；
- AllReduce p50 约为 7.10 ms；
- 通信与计算重叠率接近 0%；
- 其中 MoE 相关 AllReduce 约占已分类 AllReduce 时间的 58.73%。
这说明 PP1/TP32 并不只是“32 张卡一起算得更快”。Kimi-K3 每层中频繁出现 的 TP Collective 会跨越 4 台服务器，在没有 NVLink 的 RTX 6000D 上通过 RoCE 执行；等待通信的时间已经超过实际计算时间。
2. 为什么尝试 Deep PP
本实验所说的 Deep PP，是把一个模型划分为更多 Pipeline Stage，同时相应 缩小每个 Stage 内的 TP Group：
总 GPU 数 = PP × TP = 32
在固定 32 张 GPU 的前提下，本实验考察：
| 配置 | 每个 Stage 使用的 GPU | Stage 与节点的关系 |
|---|---|---|
| PP1 / TP32 | 32 | 一个 Stage 跨 4 个节点 |
| PP2 / TP16 | 16 | 一个 Stage 跨 2 个节点 |
| PP4 / TP8 | 8 | 一个 Stage 完整落在 1 个节点 |
| PP8 / TP4 | 4 | 每个节点放置 2 个 Stage |
| PP16 / TP2 | 2 | 每个节点放置 4 个 Stage |
PP 加深后会新增 Pipeline Activation 的 Stage-to-Stage 传输，也会产生流水线 填充和排空造成的 bubble，因此 PP 并不是越大越快。它可能带来收益的原因是：
- 每层高频 TP AllReduce 的通信域从 32 卡逐步缩小到 16、8、4 卡；
- PP4 开始，TP Group 完整落在单机内，不再让每次 TP AllReduce 都跨节点；
- PP8 中 TP4/EP4 都位于单机，跨节点通信主要集中在少数 Pipeline 边界；
- 只要 C=8/16 提供了足够多的请求，Scheduler 就可能用不同请求交错填充 Pipeline，抵消一部分 bubble。
因此，这一阶段要回答的不是“PP 能不能启动”，而是：在 Kimi-K3、6000D 和固定 Prefill 压力下，缩小 TP 通信域的收益能否超过 Pipeline 开销，以及 拐点出现在 PP2、PP4、PP8 还是 PP16。
3. 实现方式
3.1 唯一实验入口
实验只保留一个入口脚本：
run_pp_baseline_search.sh
脚本负责：
- 检查 601-604 的模型、镜像与
/dev/infiniband； - 依次启动 PP2/TP16、PP4/TP8、PP8/TP4、PP16/TP2 服务；
- 通过
/get_server_info校验实际 PP、TP、EP 和 MoE backend； - 每个配置执行 C=8/16，各重复 3 次；
- 校验每轮是否完成 40/40 请求以及是否存在请求错误；
- 保存完整 Docker 命令、服务日志、benchmark 日志、原始 JSONL 和 GPU 快照；
- 生成逐轮
results.csv和中位数summary.csv； - 无论成功或失败，最终清理 4 台服务器上的实验容器。
3.2 SGLang PP 路径
本轮使用镜像：
local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
服务直接通过 SGLang 官方参数启动：
python3 -m sglang.launch_server \
--model-path /data/hf_models/Kimi-K3 \
--tp-size "${TP_SIZE}" \
--pp-size "${PP_SIZE}" \
--ep-size "${EP_SIZE}" \
--nnodes 4 \
--node-rank "${NODE_RANK}" \
--dist-init-addr 174.1.60.1:20000 \
--moe-runner-backend flashinfer_mxfp4 \
--chunked-prefill-size 8192
脚本默认取 EP_SIZE=min(4, TP_SIZE)。因此 PP2/4/8 使用 EP4，PP16 因 TP2
必须使用 EP2，满足 SGLang 的 TP % EP == 0 约束。
实验没有复制同事旧代码，也没有注入 patch_pp_group.py。准备阶段曾验证过
一版运行时 PP Group patch，但它会让 Python Object 通信和 Tensor/P2P 通信
复用错误的 Gloo Group，出现消息大小不匹配。回到当前 SGLang 上游实现后，
PP2、PP4 和 PP8 均能正常启动，因此该 patch 被完全移除，没有进入正式数据。
3.3 网络配置
4 个节点都向容器映射 /dev/infiniband，并使用：
NCCL_SOCKET_IFNAME=bond0
GLOO_SOCKET_IFNAME=bond0
NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
NCCL_IB_GID_INDEX=3
bond0 承担 Bootstrap/控制连接，mlx5_0~3 是 NCCL 的 RoCE HCA。正式运行
中没有出现 NCCL、Gloo、RDMA、connection refused 或 communicator failure。
4. 实验设计
为了让结果只反映 PP/TP 变化，其他主要变量保持固定：
| 项目 | 固定值 |
|---|---|
| 模型 | Kimi-K3 |
| GPU | 601-604，32 × RTX 6000D |
| MoE runner | FlashInfer MXFP4 |
| A2A backend | none |
| EP | PP1/2/4/8 为 4；PP16 为 2 |
| 输入/输出 | 16,384 / 1 token |
| Chunked Prefill | 8,192 token |
| 并发 | C=8、C=16 |
| 每轮请求数 | 40 |
| 重复次数 | 3 |
| 数据集 | random-ids，输入长度比例固定为 1.0 |
| Prefix Cache | --disable-radix-cache |
| 显存比例 | 0.88 |
输出长度固定为 1，是为了让 TTFT 和 Input TPS 主要反映 Prefill，而不是被长 Decode 混入。每个点重复 3 次，最终使用中位数，避免首轮 JIT、缓存状态或系统 瞬时抖动主导结果。
PP1 没有重复跑。它来自前一阶段已经完成的 FlashInfer MXFP4、EP4、chunk8K
实验。经逐项核对，PP1 与本轮使用相同镜像、模型、MoE backend、EP、输入输出、
并发、请求数、显存比例和网络参数，区别只有 PP1/TP32，因此可以直接并入。
5. 一个关键实现差异：PP Micro Batch
同事旧 PP4 配置显式设置过：
--pp-max-micro-batch-size 1
当前 SGLang 在该参数未设置时，会在 Scheduler 初始化后自动计算：
pp_max_micro_batch_size = max(max_running_requests // pp_size, 1)
接纳新请求时又会执行：
allocatable_requests = pp_max_micro_batch_size - running_batch_size
因此这里的 1 不只是一个普通的 Kernel batch 参数，它会把 PP Scheduler
允许同时运行的请求数压到 1。对于 C=8/16，这会让请求大量排队，Pipeline
难以通过请求交错保持各 Stage 忙碌，bubble、P2P 和调度固定成本都会被放大。
本实验不传这个参数，使用 SGLang 当前的自动值。旧实验还混合了 Marlin、 chunk16K、EP1/EP8、PP16、旧镜像和运行时 PP patch，因此旧数据不能作为 “PP 天生更慢”的单变量证据。
6. 实验过程与正确性
正式 Run ID：
kimi3-pp-upstream-20260821-120208
kimi3-pp16-20260821-161721
服务启动时间分别为：
| 配置 | 健康检查通过时间 |
|---|---|
| PP2 / TP16 / EP4 | 417 s |
| PP4 / TP8 / EP4 | 405 s |
| PP8 / TP4 / EP4 | 338 s |
| PP16 / TP2 / EP2 | 538 s |
PP2、PP4、PP8 共形成 18 个正式样本；补充的 PP16 形成 6 个正式样本。 每个样本均完成 40/40 请求，原始结果中的 errors 为空。服务和 benchmark 日志没有发现：
- CUDA OOM / OutOfMemory；
- Traceback / Exception；
- NCCL / Gloo / RDMA failure；
- Engine failure；
- 请求连接失败或提前终止。
PP8 在 C=8 和 C=16 下三次 Input TPS 的极差都小于 0.7%，说明结果不是某 一次运行的偶然峰值。
PP16 首次启动在 decode CUDA Graph 捕获时触发 custom all-reduce IPC buffer
注册错误；关闭 PP16 的 decode graph 后，16-stage warmup 又超过默认 300 秒
watchdog。正式 PP16 使用 --disable-decode-cuda-graph --watchdog-timeout 3000000，538 秒后健康。该差异不改变本轮 Prefill 路径：输出仅 1 token，
没有迭代 Decode；但 PP16 的 Decode 性能不能由本实验推断。PP16 三次 Input
TPS 极差不超过 0.7%，结果同样稳定。
7. 性能结果
下表均为 3 次正式运行的中位数：
| PP / TP / EP | C | Input TPS | TTFT p50 | TTFT p95 | Input TPS 相对 PP1 | TTFT p50 相对 PP1 |
|---|---|---|---|---|---|---|
| 1 / 32 / 4 | 8 | 3,257.96 | 39.19 s | 41.62 s | 基线 | 基线 |
| 2 / 16 / 4 | 8 | 5,015.95 | 25.22 s | 26.69 s | +53.96% | -35.65% |
| 4 / 8 / 4 | 8 | 7,155.20 | 17.28 s | 18.62 s | +119.62% | -55.90% |
| 8 / 4 / 4 | 8 | 8,809.50 | 13.47 s | 16.45 s | +170.40% | -65.64% |
| 16 / 2 / 2 | 8 | 6,701.04 | 17.94 s | 21.39 s | +105.68% | -54.22% |
| 1 / 32 / 4 | 16 | 3,260.14 | 78.39 s | 80.88 s | 基线 | 基线 |
| 2 / 16 / 4 | 16 | 5,018.38 | 50.70 s | 52.14 s | +53.93% | -35.32% |
| 4 / 8 / 4 | 16 | 7,152.58 | 34.84 s | 35.95 s | +119.39% | -55.56% |
| 8 / 4 / 4 | 16 | 8,824.12 | 27.31 s | 29.77 s | +170.67% | -65.16% |
| 16 / 2 / 2 | 16 | 6,700.96 | 35.88 s | 39.49 s | +105.54% | -54.23% |
从 PP1 到 PP8，Input TPS 沿着 PP2、PP4、PP8 单调上升，TTFT 也同步下降； PP16 则出现反转。相对 PP8，PP16 在 C=8/16 的 Input TPS 分别下降 23.93% 和 24.06%，TTFT p50 分别增加 33.18% 和 31.37%。这说明继续把 TP4 缩到 TP2 的收益已无法抵消 16-stage bubble、更多 P2P 边界和调度开销。
PP16 同时把 EP4 改为 EP2，因此它不是严格的纯 PP 单变量对照。但它准确代表
32 卡上可实际部署的 PP16/TP2 组合，足以判定该组合不应替代 PP8 基线。
C=16 的 TTFT 大致是 C=8 的两倍，而同一 PP 下 Input TPS 基本不变。这表明 这批 Prefill 工作已接近该配置的稳定吞吐区：增加并发主要增加排队等待，而不是 继续提高输入吞吐。PP 改善的是服务处理这批 token 的速率，因此两个并发点都 出现了接近一致的 TPS 增益。
8. 显存结果
下表取每个配置 benchmark 完成后，4 个节点 32 张卡中的最高显存：
| PP / TP / EP | 峰值显存 | 相对 85,651 MiB 的剩余空间 |
|---|---|---|
| 1 / 32 / 4 | 81,355 MiB | 4,296 MiB |
| 2 / 16 / 4 | 80,335 MiB | 5,316 MiB |
| 4 / 8 / 4 | 82,417 MiB | 3,234 MiB |
| 8 / 4 / 4 | 79,407 MiB | 6,244 MiB |
| 16 / 2 / 2 | 79,989 MiB | 5,662 MiB |
PP 的 Stage 不会获得完全相同的层、输出头和状态，因此各 PP 配置的显存不再 像 PP1 那样均匀。PP4 的最重 Stage 已使用约 82.4 GiB；PP8 的最重 Stage 约 79.4 GiB，在本轮反而保留了最多的最坏 Rank 余量。
这一点对下一阶段 A2A 很重要，因为 DeepEP 或 FlashInfer A2A 都可能申请额外 dispatch/combine buffer。只看 TPS 而忽略最重 Rank，可能出现基线可跑、打开 A2A 后启动或 Prefill OOM 的情况。
9. 通信结果如何解释
9.1 累计收益变大，但每次增加 PP 的边际收益在下降
第 7 节相对 PP1 的累计数字会让人感觉“PP 越大，收益增长越快”。逐级比较后， 实际规律相反。以 C=8 为例：
| 变化 | Input TPS 边际变化 | TTFT p50 边际变化 |
|---|---|---|
| PP1 -> PP2 | +53.96% | -35.65% |
| PP2 -> PP4 | +42.65% | -31.48% |
| PP4 -> PP8 | +23.12% | -22.05% |
| PP8 -> PP16 | -23.93% | +33.18% |
所以 PP1 到 PP8 的累计收益持续扩大，但每次把 PP 翻倍带来的新增收益已经逐步 减小；PP16 则越过最优点。C=16 也呈现相同趋势。
9.2 节点内 PCIe 与节点间 RoCE 的实测
在 601--604 上使用 NCCL 2.28.9 + CUDA 13 的 all_reduce_perf，分别测量与
PP Stage 对应的 TP4、TP8、TP16 和 TP32 通信域。所有测试均为 0 个错误值：
| 通信域 | 实际拓扑 | 128 MiB 延迟 | 128 MiB busbw | 128/256 MiB 平均 busbw |
|---|---|---|---|---|
| TP4 | 单节点 GPU0--3，同一 PCIe/NUMA 岛 | 4.93 ms | 40.85 GB/s | 41.00 GB/s |
| TP8 | 单节点 8 卡，跨两个 NUMA 岛 | 6.00 ms | 39.13 GB/s | 38.89 GB/s |
| TP16 | 两节点 | 6.24 ms | 40.31 GB/s | 40.60 GB/s |
| TP32 | 四节点 | 6.59 ms | 39.47 GB/s | 40.08 GB/s |
大消息带宽都在约 39--41 GB/s，说明 PP 收益不是因为节点间 400G RoCE 的 持续带宽远低于节点内 PCIe。真正变化更明显的是 collective 的参与 Rank 数、 固定延迟和同步范围：128 MiB AllReduce 从 TP32 的 6.59 ms 降到 TP4 的 4.93 ms。
两个拓扑边界尤其重要：
- PP2 -> PP4 把每个 Stage 的 TP 域从两节点 TP16 缩到单节点 TP8，逐层 AllReduce 不再跨节点；
- PP4 -> PP8 把 TP8 缩到 TP4，每个 Stage 可落在同一个 PCIe/NUMA 岛，避免 单节点内跨 CPU Root Complex，同时每卡得到更大的 GEMM 分片，减轻 TP32 过度切分造成的小矩阵低效率。
9.3 PP1 与 PP8 的 Nsight 对照
为避免把 Profiler 开销混入第 7 节性能结果，另起同场景短窗口 Nsight：
16K -> 1, C=8, chunk=8K。以下是稳定 Rank 的中位数：
| 每个活跃 Rank/Stage 的指标 | PP1 / TP32 | PP8 / TP4 | 变化 |
|---|---|---|---|
| 普通 AllReduce 调用数 | 557 | 66 | -88.2% |
| 普通 AllReduce 累计时间 | 4,160.6 ms | 392.3 ms | -90.6% |
| 单次 AllReduce p50 | 7.10 ms | 5.83 ms | -17.9% |
| 计算与 NCCL 时间线重叠 | 约 0 | 约 0 | 未形成有效重叠 |
这里的 557 -> 66 是单个 Rank/Stage的 Profile 窗口，不是一个请求在整个 模型上的 collective 总数。PP8 把模型层分给 8 个 Stage，一个请求仍要依次经过 全部 Stage；它减少的是每个活跃 Rank 在一个 Stage 内负责的层数和同步工作， 并让不同请求的 Stage 可以流水并行，而不是凭空删除模型全部通信。
PP8 Trace 还出现约 3.43 s 的非 AllReduce NCCL 区间，主要对应 Pipeline Send/Recv 及 Stage 等待。该值包含异步 Stage 等待，不能直接当作 3.43 s 的 网络传输时间，也不能把 PP8 的 NCCL wall 百分比与 PP1 做等口径比较。可靠的 对照是上表中的调用数、普通 AllReduce 累计时间和单次延迟。
9.4 为什么 PP8 最优，PP16 反而退化
PP1 -> PP8 的收益由四项共同产生：
- TP 通信域由 32 Rank 缩到 4 Rank，单次 collective 延迟降低；
- PP4/PP8 的 TP collective 保持在节点内，PP8 进一步保持在同一 PCIe/NUMA 岛；
- 每个 Rank 只执行所在 Stage 的层，C=8/16 的不同请求能在多个 Stage 并行， 将通信和计算流水化；
- TP32 过度切分的 GEMM 在 TP4 下变成更大的本地矩阵，Kernel 效率更高。
收益没有继续延伸到 PP16：C=8 时请求数甚至少于 16 个 Stage，流水线无法填满； C=16 也只有一波请求。更多 Stage 边界、P2P Send/Recv、调度和 bubble 超过了 TP4 -> TP2 的剩余收益，而且 PP16 还因整除约束从 EP4 改成 EP2。因此当前证据 支持的是“PP8 在这个 Prefill 压力点达到通信、计算粒度与流水线开销的平衡”， 而不是“PP 越大越快”。
10. 最终判断与后续使用
在本次固定场景中，PP8 / TP4 / EP4 同时得到：
- 最高的 Input TPS；
- 最低的 TTFT p50/p95；
- 三次重复中稳定的结果；
- 比 PP1 和 PP4 更充足的最坏 Rank 显存余量；
- 单节点内的 TP4/EP4 通信域。
因此，下一阶段 MoE A2A 的首个兼容性和显存 Smoke 应固定使用：
PP8 / TP4 / EP4
16K -> 1
C=8/16
chunked_prefill_size=8192
FlashInfer MXFP4
PP4 / TP8 / EP4 保留为回退对照：如果某个 A2A backend 对 PP8/TP4 有当前
代码未覆盖的约束，可以快速判断问题来自 backend 兼容性，还是 A2A 本身。
PP16 / TP2 / EP2 不进入下一阶段默认矩阵。它相对 PP1 仍有明显收益，但比
PP8 慢约 24%，并引入 decode graph 兼容性、超长启动 warmup 和 EP2 变量。
这项结果不能外推到低并发 C=1、长 Decode、在线混合流量或更多节点。当前 TTFT 即使在 PP8 下仍为 13.47 s（C8）和 27.31 s（C16），距离 5 s 目标仍有 明显差距。Deep PP 解决了一个主要通信瓶颈，但不是 Prefill 优化的终点。
11. 复现命令与证据
在 601 执行：
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search
RUN_ID=kimi3-pp-$(date +%Y%m%d-%H%M%S) \
SUDO_PASSWORD_FILE=/path/to/password-file \
bash run_pp_baseline_search.sh run
PP2/4/8 正式结果：
/data/hzy/sskj/experiments/pro6000/
kimi3_pro6000_sglang_pp_baseline_search/
results/kimi3-pp-upstream-20260821-120208/
其中：
summary.csv 6 个 PP×C 中位数
results.csv 18 个逐轮结果
metadata/ 实际服务配置与 Run manifest
service/ 四节点 Docker 命令和完整服务日志
bench/ benchmark 文本输出
raw/ benchmark JSONL
gpu/ 四节点启动前、健康、完成和清理后的显存快照
PP1 原始结果：
/data/hzy/sskj/experiments/pro6000/
kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/
results/kimi3-ep4-moe-full-20260818-151349/
PP16 正式结果：
/data/hzy/sskj/experiments/pro6000/
kimi3_pro6000_sglang_pp_baseline_search/
results/kimi3-pp16-20260821-161721/
PP16 启动兼容性证据（未产生 benchmark 数据）：
results/kimi3-pp16-20260821-155734/ # custom AR decode graph IPC 错误
results/kimi3-pp16-20260821-160555/ # 默认 300s watchdog 超时
PP1 Nsight 通信归因：
/data/hzy/sskj/experiments/pro6000/
kimi3_pro6000_sglang_prefill_communication_profile/
results/kimi3-prefill-comm-20260820-143749/nsys_analysis.json
PP8 Nsight 通信归因与四组 NCCL 微基准：
/data/hzy/sskj/experiments/pro6000/
kimi3_pro6000_sglang_prefill_communication_profile/
results/kimi3-prefill-comm-pp8-20260821-172717/
nsys_analysis.json # PP8 各 Rank CUDA/NCCL 汇总
nsys_device_metrics.csv # PP8 各设备时间线指标
nsys/node0/prefill_node0.nsys-rep # 601 原始 Nsight Trace
nsys/node1/prefill_node1.nsys-rep # 602 原始 Nsight Trace
nsys/node2/prefill_node2.nsys-rep # 603 原始 Nsight Trace
nsys/node3/prefill_node3.nsys-rep # 604 原始 Nsight Trace
nccl_tp4_intra_node.log # 单节点同 PCIe/NUMA 岛
nccl_tp8_intra_node.log # 单节点跨 NUMA
nccl_tp16_2node.log # 两节点
nccl_tp32_4node.log # 四节点