[Profile] Attribute Kimi-K3 Prefill communication

This commit is contained in:
Zhiyi Hong 2026-08-20 15:56:58 +08:00
parent 6fac5ad567
commit 60f77cd4ef
3 changed files with 549 additions and 3 deletions

View File

@ -0,0 +1,241 @@
#!/usr/bin/env python3
"""Summarize CUDA/NCCL time and overlap from stage-scoped Nsight SQLite files."""
from __future__ import annotations
import argparse
import csv
import json
import math
import sqlite3
from collections import defaultdict
from pathlib import Path
def merge_intervals(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]:
merged: list[list[int]] = []
for start, end in sorted(intervals):
if end <= start:
continue
if not merged or start > merged[-1][1]:
merged.append([start, end])
else:
merged[-1][1] = max(merged[-1][1], end)
return [(start, end) for start, end in merged]
def interval_length(intervals: list[tuple[int, int]]) -> int:
return sum(end - start for start, end in intervals)
def intersection_length(
left: list[tuple[int, int]], right: list[tuple[int, int]]
) -> int:
i = j = total = 0
while i < len(left) and j < len(right):
start = max(left[i][0], right[j][0])
end = min(left[i][1], right[j][1])
if end > start:
total += end - start
if left[i][1] <= right[j][1]:
i += 1
else:
j += 1
return total
def percentile(values: list[float], pct: float) -> float:
if not values:
return math.nan
ordered = sorted(values)
pos = (len(ordered) - 1) * pct
low = math.floor(pos)
high = math.ceil(pos)
if low == high:
return ordered[low]
return ordered[low] * (high - pos) + ordered[high] * (pos - low)
def analyze_sqlite(path: Path) -> tuple[list[dict], dict[str, dict]]:
connection = sqlite3.connect(path)
rows = connection.execute(
"""
SELECT k.deviceId, k.start, k.end, s.value
FROM CUPTI_ACTIVITY_KIND_KERNEL AS k
JOIN StringIds AS s ON k.shortName = s.id
ORDER BY k.deviceId, k.start
"""
).fetchall()
connection.close()
by_device: dict[int, list[tuple[int, int, str]]] = defaultdict(list)
top_names: dict[str, dict[str, float]] = defaultdict(
lambda: {"count": 0, "duration_ns": 0}
)
for device, start, end, name in rows:
by_device[device].append((start, end, name))
top_names[name]["count"] += 1
top_names[name]["duration_ns"] += end - start
output = []
node = path.parent.name
for device, kernels in sorted(by_device.items()):
all_intervals = merge_intervals([(start, end) for start, end, _ in kernels])
nccl_rows = [row for row in kernels if row[2].lower().startswith("nccl")]
compute_rows = [row for row in kernels if not row[2].lower().startswith("nccl")]
nccl_intervals = merge_intervals([(start, end) for start, end, _ in nccl_rows])
compute_intervals = merge_intervals(
[(start, end) for start, end, _ in compute_rows]
)
wall_ns = max(end for _, end, _ in kernels) - min(start for start, _, _ in kernels)
kernel_union_ns = interval_length(all_intervals)
nccl_union_ns = interval_length(nccl_intervals)
compute_union_ns = interval_length(compute_intervals)
overlap_ns = intersection_length(nccl_intervals, compute_intervals)
allreduce_ms = [
(end - start) / 1e6
for start, end, name in nccl_rows
if "allreduce" in name.lower()
]
normal_allreduce_ms = [value for value in allreduce_ms if value < 50.0]
# The 8K Prefill trace has a clean gap between the two known payloads:
# [M, 7168] ~= 112 MiB and M * (3584 + 7168) ~= 168 MiB.
hidden_allreduce_ms = [value for value in normal_allreduce_ms if value < 7.5]
moe_allreduce_ms = [value for value in normal_allreduce_ms if value >= 7.5]
normal_allreduce_total_ms = sum(normal_allreduce_ms)
output.append(
{
"node": node,
"local_gpu": device,
"wall_ms": wall_ns / 1e6,
"gpu_busy_union_ms": kernel_union_ns / 1e6,
"gpu_busy_pct": 100.0 * kernel_union_ns / wall_ns,
"nccl_union_ms": nccl_union_ns / 1e6,
"nccl_wall_pct": 100.0 * nccl_union_ns / wall_ns,
"compute_union_ms": compute_union_ns / 1e6,
"compute_wall_pct": 100.0 * compute_union_ns / wall_ns,
"nccl_compute_overlap_ms": overlap_ns / 1e6,
"nccl_overlap_pct": (
100.0 * overlap_ns / nccl_union_ns if nccl_union_ns else 0.0
),
"nccl_exposed_ms": (nccl_union_ns - overlap_ns) / 1e6,
"nccl_exposed_wall_pct": 100.0
* (nccl_union_ns - overlap_ns)
/ wall_ns,
"allreduce_count": len(allreduce_ms),
"allreduce_normal_count": len(normal_allreduce_ms),
"allreduce_mean_ms": (
sum(normal_allreduce_ms) / len(normal_allreduce_ms)
if normal_allreduce_ms
else math.nan
),
"allreduce_p50_ms": percentile(normal_allreduce_ms, 0.50),
"allreduce_p95_ms": percentile(normal_allreduce_ms, 0.95),
"allreduce_max_ms": max(allreduce_ms, default=math.nan),
"hidden_allreduce_count": len(hidden_allreduce_ms),
"hidden_allreduce_sum_ms": sum(hidden_allreduce_ms),
"hidden_allreduce_p50_ms": percentile(hidden_allreduce_ms, 0.50),
"moe_allreduce_count": len(moe_allreduce_ms),
"moe_allreduce_sum_ms": sum(moe_allreduce_ms),
"moe_allreduce_p50_ms": percentile(moe_allreduce_ms, 0.50),
"moe_allreduce_time_share_pct": (
100.0 * sum(moe_allreduce_ms) / normal_allreduce_total_ms
if normal_allreduce_total_ms
else 0.0
),
}
)
return output, top_names
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("result_root", type=Path)
args = parser.parse_args()
sqlite_paths = sorted((args.result_root / "nsys").glob("node*/*.sqlite"))
if len(sqlite_paths) != 4:
raise SystemExit(f"expected four SQLite files, found {len(sqlite_paths)}")
devices: list[dict] = []
global_names: dict[str, dict[str, float]] = defaultdict(
lambda: {"count": 0, "duration_ns": 0}
)
for path in sqlite_paths:
device_rows, names = analyze_sqlite(path)
devices.extend(device_rows)
for name, values in names.items():
global_names[name]["count"] += values["count"]
global_names[name]["duration_ns"] += values["duration_ns"]
stable_devices = [
row for row in devices if row["allreduce_max_ms"] < 50.0
]
metric_names = [
"wall_ms",
"gpu_busy_pct",
"nccl_union_ms",
"nccl_wall_pct",
"compute_union_ms",
"compute_wall_pct",
"nccl_compute_overlap_ms",
"nccl_overlap_pct",
"nccl_exposed_ms",
"nccl_exposed_wall_pct",
"allreduce_count",
"allreduce_mean_ms",
"allreduce_p50_ms",
"allreduce_p95_ms",
"hidden_allreduce_count",
"hidden_allreduce_sum_ms",
"hidden_allreduce_p50_ms",
"moe_allreduce_count",
"moe_allreduce_sum_ms",
"moe_allreduce_p50_ms",
"moe_allreduce_time_share_pct",
]
stable_medians = {
name: percentile([float(row[name]) for row in stable_devices], 0.50)
for name in metric_names
}
top_kernels = [
{
"name": name,
"count": int(values["count"]),
"duration_ms": values["duration_ns"] / 1e6,
}
for name, values in sorted(
global_names.items(),
key=lambda item: item[1]["duration_ns"],
reverse=True,
)[:40]
]
payload = {
"result_root": str(args.result_root),
"sqlite_files": [str(path) for path in sqlite_paths],
"devices": devices,
"stable_device_rule": "allreduce_max_ms < 50; excludes profiler trigger lanes",
"stable_device_count": len(stable_devices),
"stable_medians": stable_medians,
"top_kernels": top_kernels,
}
output_json = args.result_root / "nsys_analysis.json"
output_csv = args.result_root / "nsys_device_metrics.csv"
output_json.write_text(json.dumps(payload, indent=2), encoding="utf-8")
with output_csv.open("w", newline="", encoding="utf-8") as handle:
writer = csv.DictWriter(handle, fieldnames=list(devices[0]))
writer.writeheader()
writer.writerows(devices)
print(json.dumps(payload["stable_medians"], indent=2))
print(f"stable devices: {len(stable_devices)}/{len(devices)}")
print(f"wrote {output_json}")
print(f"wrote {output_csv}")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,254 @@
#!/usr/bin/env bash
# Capture a short, stage-scoped Nsight Systems trace for Kimi-K3 prefill.
set -Eeuo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RUN_ID="${RUN_ID:-kimi3-prefill-comm-$(date '+%Y%m%d-%H%M%S')}"
RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/results}"
RESULT_ROOT="${RESULT_BASE}/${RUN_ID}"
MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}"
IMAGE="${IMAGE:-local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5}"
HEAD_HOST="${HEAD_HOST:-174.1.60.1}"
NODE_USER="${NODE_USER:-user}"
NODE_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
PORT="${PORT:-30000}"
DIST_PORT="${DIST_PORT:-20000}"
PROFILE_STEPS="${PROFILE_STEPS:-3}"
CONCURRENCY="${CONCURRENCY:-8}"
NUM_PROMPTS="${NUM_PROMPTS:-40}"
CONTAINER_PREFIX="kimi3_prefill_comm_profile"
SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10)
mkdir -p "${RESULT_ROOT}"/{bench,nsys,service,gpu}
log() {
printf '[%(%F %T)T] %s\n' -1 "$*" | tee -a "${RESULT_ROOT}/orchestrator.log"
}
require_password() {
if [[ -z "${SUDO_PASSWORD:-}" && -n "${SUDO_PASSWORD_FILE:-}" ]]; then
IFS= read -r SUDO_PASSWORD <"${SUDO_PASSWORD_FILE}"
fi
[[ -n "${SUDO_PASSWORD:-}" ]] || {
echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE" >&2
exit 2
}
}
is_head() {
[[ "$1" == "$HEAD_HOST" ]]
}
sudo_host() {
local host="$1"
shift
require_password
if is_head "$host"; then
printf '%s\n' "$SUDO_PASSWORD" | sudo -S -p '' -- "$@"
else
local command
printf -v command '%q ' "$@"
printf '%s\n' "$SUDO_PASSWORD" | ssh "${SSH_OPTS[@]}" "${NODE_USER}@${host}" \
"sudo -S -p '' -- ${command}"
fi
}
run_host() {
local host="$1"
shift
if is_head "$host"; then
"$@"
else
ssh "${SSH_OPTS[@]}" "${NODE_USER}@${host}" "$@"
fi
}
container_name() {
printf '%s_node%s' "$CONTAINER_PREFIX" "$1"
}
prepare_nodes() {
local host rank=0
for host in "${NODE_HOSTS[@]}"; do
run_host "$host" mkdir -p "${RESULT_ROOT}/nsys/node${rank}"
rank=$((rank + 1))
done
}
stop_service() {
local host rank=0
for host in "${NODE_HOSTS[@]}"; do
sudo_host "$host" docker rm -f "$(container_name "$rank")" >/dev/null 2>&1 || true
rank=$((rank + 1))
done
}
start_node() {
local rank="$1" host name output launch
host="${NODE_HOSTS[$rank]}"
name="$(container_name "$rank")"
output="${RESULT_ROOT}/nsys/node${rank}/prefill_node${rank}"
launch="export SGLANG_HOST_IP=174.1.60.$((rank + 1)); exec nsys profile --trace=cuda,nvtx,nccl --nccl-trace=api,group,gpu,coll,kernel-launch --sample=none --cpuctxsw=none --capture-range=cudaProfilerApi --capture-range-end=stop --cuda-graph-trace=node --force-overwrite=true --stats=true --output ${output} python3 -m sglang.launch_server --model-path ${MODEL_PATH} --served-model-name kimi-k3 --tp-size 32 --ep-size 4 --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend flashinfer_mxfp4 --chunked-prefill-size 8192 --mem-fraction-static 0.88 --cuda-graph-max-bs-decode 16 --mamba-radix-cache-strategy extra_buffer_lazy --disable-radix-cache --dist-timeout 3600 --mamba-full-memory-ratio 0.36 --host 0.0.0.0 --port ${PORT}"
local -a command=(
docker run -d --name "$name"
--gpus all --network host --ipc=host --ulimit memlock=-1
--device /dev/infiniband --shm-size 32g --entrypoint bash
-v "${MODEL_PATH}:${MODEL_PATH}:ro"
-v "${RESULT_ROOT}:${RESULT_ROOT}"
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-e NCCL_SOCKET_IFNAME=bond0 -e GLOO_SOCKET_IFNAME=bond0
-e NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
-e NCCL_IB_GID_INDEX=3 -e NCCL_IB_TIMEOUT=22 -e NCCL_IB_RETRY_CNT=7
-e NCCL_CUMEM_ENABLE=1 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
-e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0 -e SGLANG_MOE_FUSED_GATE_RADIX=1
-e FLASHINFER_DISABLE_JIT=1 -e FLASHINFER_DISABLE_VERSION_CHECK=1
"$IMAGE" -lc "$launch"
)
printf '%q ' "${command[@]}" >"${RESULT_ROOT}/service/node${rank}.cmd.txt"
printf '\n' >>"${RESULT_ROOT}/service/node${rank}.cmd.txt"
sudo_host "$host" "${command[@]}" >/dev/null
}
wait_health() {
local elapsed
for ((elapsed = 1; elapsed <= 2400; elapsed++)); do
if curl --fail --silent --max-time 5 "http://${HEAD_HOST}:${PORT}/health" >/dev/null 2>&1; then
log "service healthy after ${elapsed}s"
return 0
fi
if (( elapsed % 30 == 0 )); then
log "waiting for service: ${elapsed}s"
collect_logs starting
if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|RuntimeError' \
"${RESULT_ROOT}/service/starting_node"*.log; then
log "ERROR: startup log contains a fatal error"
return 1
fi
fi
sleep 1
done
return 1
}
collect_logs() {
local label="$1" host rank=0
for host in "${NODE_HOSTS[@]}"; do
sudo_host "$host" docker logs "$(container_name "$rank")" \
>"${RESULT_ROOT}/service/${label}_node${rank}.log" 2>&1 || true
rank=$((rank + 1))
done
}
collect_gpu() {
local label="$1" host rank=0
for host in "${NODE_HOSTS[@]}"; do
sudo_host "$host" nvidia-smi \
--query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu,power.draw \
--format=csv,noheader,nounits \
>"${RESULT_ROOT}/gpu/${label}_node${rank}.csv" 2>&1 || true
rank=$((rank + 1))
done
}
run_bench() {
local label="$1" prompts="$2" concurrency="$3"
local output="${RESULT_ROOT}/bench/${label}.jsonl"
sudo_host "$HEAD_HOST" docker run --rm --network host \
-v "${MODEL_PATH}:${MODEL_PATH}:ro" -v "${RESULT_ROOT}:${RESULT_ROOT}" \
-e PYTHONUNBUFFERED=1 --entrypoint python3 "$IMAGE" \
-m sglang.benchmark.serving --backend sglang --host "$HEAD_HOST" --port "$PORT" \
--tokenizer "$MODEL_PATH" --dataset-name random-ids \
--random-input-len 16384 --random-output-len 1 --random-range-ratio 1.0 \
--num-prompts "$prompts" --max-concurrency "$concurrency" --request-rate 10000 \
--warmup-requests 0 --output-file "$output" --output-details --disable-tqdm \
>"${RESULT_ROOT}/bench/${label}.log" 2>&1
}
trigger_profile() {
local payload response
payload="{\"activities\":[\"CUDA_PROFILER\"],\"num_steps\":${PROFILE_STEPS},\"profile_by_stage\":true,\"profile_id\":\"${RUN_ID}\",\"profile_prefix\":\"prefill\"}"
printf '%s\n' "$payload" >"${RESULT_ROOT}/service/start_profile_request.json"
response="$(curl --fail --silent --show-error -X POST \
"http://${HEAD_HOST}:${PORT}/start_profile" \
-H 'Content-Type: application/json' -d "$payload")"
printf '%s\n' "$response" | tee "${RESULT_ROOT}/service/start_profile_response.json"
}
copy_worker_artifacts() {
local host rank
for rank in 1 2 3; do
host="${NODE_HOSTS[$rank]}"
printf '%s\n' "$SUDO_PASSWORD" | ssh "${SSH_OPTS[@]}" "${NODE_USER}@${host}" \
"sudo -S -p '' -- tar -C '${RESULT_ROOT}' -cf - 'nsys/node${rank}'" | \
tar -C "${RESULT_ROOT}" -xf -
done
}
write_manifest() {
{
printf 'run_id=%s\n' "$RUN_ID"
printf 'image=%s\n' "$IMAGE"
printf 'model=%s\n' "$MODEL_PATH"
printf 'tp=32\nep=4\ndp=1\n'
printf 'moe_runner_backend=flashinfer_mxfp4\nmoe_a2a_backend=none\n'
printf 'input_len=16384\noutput_len=1\nconcurrency=%s\n' "$CONCURRENCY"
printf 'chunked_prefill_size=8192\nprofile_steps=%s\n' "$PROFILE_STEPS"
} >"${RESULT_ROOT}/manifest.env"
}
run_all() {
require_password
prepare_nodes
write_manifest
stop_service
collect_gpu before
log "starting current TP32/EP4/FlashInfer service under Nsight"
start_node 1
start_node 2
start_node 3
sleep 5
start_node 0
wait_health
collect_logs healthy
collect_gpu healthy
log "warming kernels without profiling"
run_bench warmup 2 2
log "arming ${PROFILE_STEPS} Prefill steps"
trigger_profile
run_bench profile_c8 "$NUM_PROMPTS" "$CONCURRENCY"
collect_gpu profiled
collect_logs profiled
log "stopping services to finalize reports"
stop_service
sleep 10
copy_worker_artifacts
collect_gpu after
find "${RESULT_ROOT}/nsys" -type f -printf '%p %s bytes\n' | sort \
>"${RESULT_ROOT}/nsys/files.txt"
log "capture complete: ${RESULT_ROOT}"
}
cleanup() {
collect_logs cleanup 2>/dev/null || true
stop_service 2>/dev/null || true
}
case "${1:-run}" in
run)
trap cleanup EXIT INT TERM
run_all
trap - EXIT INT TERM
;;
stop)
require_password
stop_service
;;
*)
echo "Usage: SUDO_PASSWORD=... bash $0 [run|stop]" >&2
exit 2
;;
esac

View File

@ -167,6 +167,56 @@ CustomAllReduceV2 multicast 可用时自动开启;跨节点日志也明确显
`CustomAllreduce is disabled because this process group spans across nodes` `CustomAllreduce is disabled because this process group spans across nodes`
所以以上标准 NCCL 调用仍是当前 EP4/FlashInfer 配置的真实结构。 所以以上标准 NCCL 调用仍是当前 EP4/FlashInfer 配置的真实结构。
### 当前 EP4 + FlashInfer 真机时间线
为消除旧 Trace 使用 EP32/Marlin 的不确定性,又对当前正式配置做了三步
Prefill Nsight 捕获:
```text
TP32 / EP4 / DP1
MoE runner: flashinfer_mxfp4
MoE A2A: none
16K -> 1, C=8, Chunk=8K
```
原始结果位于 601
```text
/data/hzy/sskj/experiments/pro6000/
kimi3_pro6000_sglang_prefill_communication_profile/results/
kimi3-prefill-comm-20260820-143749/
```
四个节点均生成 `.nsys-rep``.sqlite`,每份报告覆盖本节点 8 张 GPU。
`nsys_analysis.json``nsys_device_metrics.csv` 是离线分析结果。为避免
`cudaProfilerStart` 边界开销污染结论,统计排除了出现单次 AllReduce
大于 50 ms 的 5 条触发 lane保留 27/32 条稳定 lane 取中位数。
| 指标 | 稳定 rank 中位数 | 解释 |
|---|---:|---|
| 捕获窗口 | 7274.23 ms | 三个 8K Prefill step |
| GPU busy | 99.37% | 几乎没有 GPU 空洞 |
| NCCL 时间 | 4161.36 ms57.27% | 当前 Prefill 的第一瓶颈 |
| 非 NCCL 计算 | 3062.48 ms42.10% | Attention、MoE GEMM、KDA 等 |
| NCCL 与计算重叠 | 0 ms0% | Collective 完整暴露在关键路径上 |
| AllReduce 次数 | 557 | 捕获边界略少于理论 561 次 |
| AllReduce P50 / P95 | 7.10 / 9.14 ms | 全部选择 `RING_LL` |
AllReduce 时长呈现无重叠的双峰:约 `6.1 ms``8.7 ms`,与源码 Trace
中的 112 MiB hidden 消息和 168 MiB MoE 消息相符。按两个峰的累计时间:
| 通信类别 | 三步累计时间 | 占 AllReduce 时间 | 占完整窗口 |
|---|---:|---:|---:|
| hidden / Attention 路径 | 1719.78 ms | 41.27% | 23.64% |
| MoE 尾部归约 | 2440.52 ms | 58.73% | 33.55% |
稳定 rank 的 AllReduce 平均值约为 `7.347.56 ms`,未发现固定慢节点。
这说明问题不是某张卡掉速,而是 TP32 每层同步通信本身占据关键路径。
带 Nsight 的完整 benchmark 为 40/40 请求成功,但 Input TPS 为 2573低于
无 Profiler 基线 3257.96;该值只用于证明 workload 完整成功,不能作为性能
回归结果。性能比较仍使用无 Profiler 的正式基线。
### 优化优先级 ### 优化优先级
1. **MoE A2A / SP-MoE 优先验证。** 它瞄准占输入字节 59.2% 的 MoE 尾部 1. **MoE A2A / SP-MoE 优先验证。** 它瞄准占输入字节 59.2% 的 MoE 尾部
@ -180,9 +230,10 @@ CustomAllReduceV2 multicast 可用时自动开启;跨节点日志也明确显
4. **不优先做 NCCL 算法或 launch fusion。** Prefill 单消息为 112/168 MiB 4. **不优先做 NCCL 算法或 launch fusion。** Prefill 单消息为 112/168 MiB
首要矛盾是字节量和 32-rank 跨节点通信域,不是小消息启动延迟。 首要矛盾是字节量和 32-rank 跨节点通信域,不是小消息启动延迟。
因此,通用 TP Reduce Scatter 即使成功,也只触及约 40.8% 的 hidden 路径, 因此,通用 TP Reduce Scatter 即使成功,也只触及字节口径约 40.8%、时间
且 K3 gate/KDA 又迫使完整 hidden 存在;它不如先处理占比更大的 MoE 尾部 口径约 41.3% 的 hidden 通信;它无法处理占完整 Prefill 时间约 33.6% 的
通信。 MoE 尾部归约,而且 K3 gate/KDA 又迫使完整 hidden 存在。它不如先处理
MoE 通信,再用 PP 缩小剩余 TP collective 的通信域。
## FlashInfer MoE 的关系 ## FlashInfer MoE 的关系