Kimi-K3 Prefill PP baseline search

Goal

Select the pipeline-parallel configuration for the next MoE A2A experiment on four RTX 6000D nodes. The comparison fixes every workload and backend variable except PP/TP:

SGLang image: local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
Model: Kimi-K3
GPUs: 32 GPUs on 601-604
MoE: FlashInfer MXFP4, EP4, A2A disabled
Workload: 16K input, one output token, C=8/16, 40 requests
Chunked prefill: 8K
Repeats: three per point

run_pp_baseline_search.sh is the only entry point. It uses the upstream SGLang PP implementation and does not apply a runtime PP patch.

Result

Values are the medians of three runs. PP1 is the directly comparable existing run; PP2/4/8 come from the new search.

PP / TP / EP	C	Input TPS	TTFT p50	TTFT p95	Input TPS vs PP1	TTFT p50 vs PP1
1 / 32 / 4	8	3,257.96	39.19 s	41.62 s	baseline	baseline
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%
1 / 32 / 4	16	3,260.14	78.39 s	80.88 s	baseline	baseline
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%

All 18 new samples completed 40/40 requests. No OOM, traceback, NCCL failure, request error, or engine failure was found. The three repeats are stable: PP8 Input TPS varies by less than 0.7% at both concurrency points.

Memory

The table reports the highest per-GPU memory observed in the post-benchmark snapshot across all four nodes.

PP / TP / EP	Peak GPU memory	Remaining from 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

PP stages do not have identical model/state allocations, so memory is less balanced than PP1. PP8 nevertheless has the best worst-rank headroom in this run, which matters because A2A backends allocate additional communication buffers.

Why PP helps here

The existing PP1 Nsight profile measured a median 57.27% of the Prefill window inside exposed NCCL work, with about 557 AllReduce calls and effectively no compute/communication overlap. Its TP32 collectives span all four nodes.

Changing PP also changes the TP communication domain:

PP2/TP16: each TP group spans two nodes.
PP4/TP8: each TP group fits on one node; only pipeline activation transfers cross node boundaries.
PP8/TP4: two pipeline stages fit per node and each TP/EP4 group is node-local.

This replaces frequent cross-node TP AllReduce with a smaller number of pipeline activation transfers. At C=8/16 there is enough request-level work to fill the deeper pipeline, so the communication saving is larger than the pipeline bubble cost.

The current run verifies the rank topology, successful NCCL P2P communicators, and absence of communication failures. It does not contain per-PP Nsight or HCA bandwidth samples; the 57.27% NCCL attribution belongs specifically to the directly comparable PP1 profile and is used to explain, not fabricate, the PP8 timeline.

Why older PP tests could be slower

The old configurations were not an apples-to-apples PP-only comparison. In particular, an old PP4 profile set:

--pp-max-micro-batch-size 1

Current SGLang leaves this value unset and computes:

pp_max_micro_batch_size = max(max_running_requests // pp_size, 1)

The scheduler then limits admission with pp_max_micro_batch_size - running_bs. Setting it to one therefore permits only one running request per PP scheduler, serializing much of a C=8/16 load and exposing pipeline bubbles. Older tests also mixed Marlin, 16K chunks, EP1/EP8, PP16, older images, and a runtime PP-group patch. They cannot be used to conclude that PP itself is slower.

Decision

Use PP8 / TP4 / EP4 as the primary configuration for the next A2A compatibility and memory smoke test. Keep PP4 / TP8 / EP4 as the fallback: it is slower and has less memory headroom here, but provides a useful check if an A2A backend imposes TP/PP constraints not exercised by the baseline.

Do not carry forward --pp-max-micro-batch-size 1. Leave it unset unless a separate controlled experiment establishes a reason to cap it.

Reproduce and evidence

Run from 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

New PP2/4/8 evidence:

results/kimi3-pp-upstream-20260821-120208/
  summary.csv
  results.csv
  metadata/
  bench/
  raw/
  gpu/
  service/

Existing PP1 evidence:

../kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/
  results/kimi3-ep4-moe-full-20260818-151349/

Existing PP1 communication profile:

../kimi3_pro6000_sglang_prefill_communication_profile/
  results/kimi3-prefill-comm-20260820-143749/nsys_analysis.json

5.2 KiB Raw Blame History