Add TP2 prefill serving-path smoke experiment

This commit is contained in:
2026-07-23 16:09:52 +08:00
parent 9c1175a434
commit 2c3220c2be
12 changed files with 729 additions and 0 deletions

View File

@@ -0,0 +1,3 @@
fleet-artifacts/
fleet-state/
remote-outputs/

View File

@@ -0,0 +1,89 @@
#!/usr/bin/env python3
"""Freeze the EXP-TP2-PREFILL-SERVING entry audit from simulator ledgers."""
from __future__ import annotations
import json
from pathlib import Path
ROOT = Path(__file__).resolve().parent
REPO = ROOT.parents[1]
REPLAY = REPO / "runs/frontier-attn-structured-v0/replay"
REAL_CHUNK_MS = {2: 410.0, 4: 231.0}
def find_one(root: Path, name: str) -> Path:
matches = list(root.rglob(name))
if len(matches) != 1:
raise ValueError(f"expected one {name} below {root}: {matches}")
return matches[0]
def chunk_zero(root: Path) -> dict[str, float | int]:
ledger = find_one(root, "frontier_stage_batch_ledger.jsonl")
chunks = []
processed: dict[str, int] = {}
with ledger.open() as stream:
for line in stream:
row = json.loads(line)
request_ids = row.get("request_ids") or []
request_tokens = row.get("request_num_tokens") or []
if len(request_ids) != 1 or request_tokens != [8192]:
continue
request_id = request_ids[0]
before = processed.get(request_id, 0)
processed[request_id] = before + 8192
if before == 0:
components = row["execution_time"]["component_ledger_ms"]
chunks.append(
{
"total_ms": row["execution_time"]["total_time_ms"],
"attention_prefill_ms": components[
"attention_prefill_execution_time"
],
"moe_grouped_gemm_ms": components[
"moe_grouped_gemm_time"
],
}
)
if len(chunks) == 9:
break
if len(chunks) != 9:
raise ValueError(f"expected 9 initial q8k chunks, got {len(chunks)}")
return {
name: sum(float(row[name]) for row in chunks) / len(chunks)
for name in chunks[0]
} | {"samples": len(chunks)}
def main() -> None:
cells = {}
for tp, label in ((2, "tp2_rho0p0025"), (4, "tp4_rho0p0025")):
measured = chunk_zero(REPLAY / label)
real = REAL_CHUNK_MS[tp]
measured["real_total_ms"] = real
measured["total_bias"] = (measured["total_ms"] - real) / real
cells[f"tp{tp}"] = measured
tp2 = cells["tp2"]
required_moe = (
tp2["moe_grouped_gemm_ms"]
+ tp2["real_total_ms"]
- tp2["total_ms"]
)
payload = {
"schema": "frontier-tp2-prefill-serving-entry-v1",
"cells": cells,
"tp2_required_moe_if_residual_is_all_moe_ms": required_moe,
"tp2_required_moe_shift": (
required_moe / tp2["moe_grouped_gemm_ms"] - 1
),
"entry_gate": abs(tp2["total_bias"]) >= 0.10,
}
results = ROOT / "results"
results.mkdir(exist_ok=True)
(results / "entry-audit.json").write_text(json.dumps(payload, indent=2))
print(json.dumps(payload, indent=2))
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,146 @@
#!/usr/bin/env python3
"""Summarize the longest graph-on execute window in each TP rank."""
from __future__ import annotations
import argparse
import gzip
import json
from collections import Counter, defaultdict
from pathlib import Path
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--trace-root", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
return parser.parse_args()
def load_events(path: Path) -> list[dict]:
opener = gzip.open if path.suffix == ".gz" else open
with opener(path, "rt") as source:
return json.load(source)["traceEvents"]
def classify(name: str) -> str:
lower = name.lower()
if any(token in lower for token in ("nccl", "allreduce", "all_reduce")):
return "collective"
if "topkgating" in lower or "fused_topk" in lower:
return "router"
if any(
token in lower
for token in (
"fused_moe",
"moefcgemm",
"tensorrt_llm::kernels::cutlass_kernels",
"groupproblemshape",
"memcpy32_post",
)
):
return "moe"
if any(
token in lower
for token in (
"flashattn",
"flashattnfwd",
"reshape_and_cache",
"prepare_varlen_num_blocks",
)
):
return "attention"
if any(
token in lower
for token in ("nvjet", "cublaslt", "rms_norm", "rsqrt", "rope")
):
return "linear_norm_rope"
return "other"
def analyze_rank(path: Path) -> dict:
events = load_events(path)
kernels = [event for event in events if event.get("cat") == "kernel"]
windows = [
event
for event in events
if event.get("cat") == "gpu_user_annotation"
and str(event.get("name", "")).startswith("execute_")
]
if not windows:
raise ValueError(f"{path}: no execute annotation")
selected = max(windows, key=lambda event: float(event["dur"]))
start = float(selected["ts"])
end = start + float(selected["dur"])
current = [
kernel for kernel in kernels if start <= float(kernel["ts"]) < end
]
components: dict[str, float] = defaultdict(float)
kernel_totals: dict[str, float] = defaultdict(float)
for kernel in current:
duration_ms = float(kernel["dur"]) / 1000
name = str(kernel["name"])
components[classify(name)] += duration_ms
kernel_totals[name] += duration_ms
kernel_rows = [
{"name": name, "duration_ms": duration}
for name, duration in sorted(
kernel_totals.items(), key=lambda item: -item[1]
)
]
wall_ms = float(selected["dur"]) / 1000
busy_ms = sum(components.values())
return {
"trace": str(path),
"selected_execute_annotation": str(selected["name"]),
"execute_annotation_histogram": dict(
sorted(Counter(str(window["name"]) for window in windows).items())
),
"all_execute_windows": [
{
"name": str(window["name"]),
"duration_ms": float(window["dur"]) / 1000,
}
for window in sorted(windows, key=lambda event: float(event["ts"]))
],
"execute_wall_ms": wall_ms,
"gpu_kernel_busy_ms": busy_ms,
"non_kernel_gap_ms": wall_ms - busy_ms,
"components_ms": dict(sorted(components.items())),
"kernel_rows": kernel_rows,
}
def main() -> None:
args = parse_args()
traces = sorted(args.trace_root.rglob("*.pt.trace.json*"))
if not traces:
raise ValueError(f"no traces below {args.trace_root}")
ranks = [analyze_rank(path) for path in traces]
critical = max(ranks, key=lambda rank: rank["execute_wall_ms"])
payload = {
"schema": "frontier-tp2-prefill-serving-smoke.v1",
"contract": {
"selection": "longest execute annotation per TP rank",
"critical_path": "rank with largest selected execute wall",
"component_time": "sum of CUDA kernel duration within selected window",
},
"ranks": ranks,
"critical_rank": critical,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
print(
json.dumps(
{
"ranks": len(ranks),
"execute_wall_ms": critical["execute_wall_ms"],
"components_ms": critical["components_ms"],
},
sort_keys=True,
)
)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,79 @@
# 实验 EXP-TP2-PREFILL-SERVINGTP2 base-prefill residual 是否来自 serving-path MoE
> **状态:** GPU smoke harness 已冻结,待远端执行
>
> Parent campaign[`../frontier-simulator-gap-campaign-v0/README.md`](../frontier-simulator-gap-campaign-v0/README.md)
## Claim 与决策
- **Parent claim** structured predictor 后,单请求 q8k/ctx0 的 TP2 sim chunk
仍为 356 ms而 real anchor 为 410 ms13.1%TP4 为 231.3 vs
231 ms。剩余量是 TP2 独有的 base serving-path residual。
- **目的:** 判断约 54 ms/chunk 是否可由 TP2 prefill MoE 的
tactic/warmup/routing/serving composition 工程修复。
- **Competing hypotheses**
- H1standalone grouped-GEMM 在 TP2 的 q8k expert shape 上过于乐观;
warm real-routing 或 serving-path trace 会比当前 171.1 ms/chunk 高 ≥10%。
- H2MoE 三臂接近,残差位于其它 whole-layer/host/event path。
- **事前预测:** 若 H1 成立serving-path MoE 增量接近 54 ms并且只注入
TP2 row 后 chunk0 residual 降至 ≤5%;若 H2 成立,三臂差异 <10%。
- **判定规则**
- 先用一个 TP2 q8k smoke 验证 prefill execute annotation kernel
component 可提取
- B/C 相对 A 都无稳定 10% shift 停止 profile 注入 whole-layer
- C 解释 70% 54 ms residual 注入 TP2重放两个 TP2 trace cell
## Setup
- **自变量**
- A当前 Frontier standalone profile/predictorTP2 q8k/ctx0
`moe_grouped_gemm=171.118 ms/chunk`
- B相同 expert GEMM shapes充分 warmup/autotune使用 serving trace
抽取的 routing allocation
- CvLLM 0.20 graph-aligned serving path q8k prefill execute按最慢
TP rank 汇总 MoE kernels
- **控制变量** Qwen3-30B-A3B BF16H20TP2/EP1q8k/ctx0同模型
runtimeFlashInfer workspaceCUDA/driverTP4 仅作 control anchor
- **Hardware** 只用 `dash1`--`dash4`运行前要求目标主机 8×H20 全部
idle/healthy每个 arm fresh process避免 profiler one-shot tactic cache
交叉污染
- **Smoke contract** q81922TP28192 max batched tokens同一 fresh
process 先执行 2 profiler-off warmup profile 1 次请求每个 rank
选择 duration 最大的 `execute_*` annotation 作为 prefill windowcritical
path wall 最大的 rank
- **Metrics** chunk execute wallMoE/attention/collective/other kernel ms
最慢 rank独立 process repeat注入后的 chunk residual TTFT/E2E
## 预期产物与 review
- **预期数据** `results/entry-audit.json`三臂 trace/provenancecomponent
对照TP2 counterfactual replay
- **Figure prototype** `figure-prototype.png`左图为 TP2/TP4 q8k chunk
real vs sim右图为三臂 MoE 事前预测
- **人工 review** campaign 已批准entry gate 已通过
- **Review 意见** 先单 TP2 smoke只有 annotation component contract
通过才扩为 repeat/control不先铺满 GPU 网格
## 复现信息
- **Code** 当前 aituner worktreeserving profiler 复用
`runs/frontier-component-residual-v0` vLLM/Kineto harness
- **Environment** vLLM 0.20.0H20remote repo
`/home/admin/cpfs/wjh/aituner/aituner`
- **产物路径** 本目录
- **已知 deviation** real 410/231 ms 是已冻结的 pure-prefill chunk anchor
本实验不重新声称它来自当前 1h trace 的在线 batch
## 结果
- **Entry audit** PASSTP2 q8k/ctx0 structured sim=`356.25 ms`
real=`410 ms`bias=`13.11%`TP4=`231.33 vs 231 ms`。TP2
component ledger MoE=`171.12 ms`,若单独解释 residual 需增至约
`224.9 ms`+31.4%)。
- **观察事实** GPU
- **Fleet preflight2026-07-23** dash1--dash4 均为 8×H2032 张卡
`memory.used=0 MiB``utilization=0%` compute processdry-run 选择
`dash1:[0,1]`正式 job pin dash1
- **含义** GPU
- **Claim update** unchanged
- **下一步** probe fleet运行 TP2 serving-path prefill smoke

Binary file not shown.

After

Width:  |  Height:  |  Size: 70 KiB

View File

@@ -0,0 +1,48 @@
version = 1
[paths]
state_dir = "runs/frontier-tp2-prefill-serving-v0/fleet-state"
artifacts_dir = "runs/frontier-tp2-prefill-serving-v0/fleet-artifacts"
[ssh]
connect_timeout_sec = 10
[scheduler]
gpu_free_memory_mb = 1024
gpu_free_utilization_pct = 10
prefer_pack = true
# Code is synchronized by Git. The fleet's mandatory rsync pass is intentionally
# empty so it only schedules and harvests the committed remote checkout.
[sync]
mode = "rsync"
local_path = "runs/frontier-tp2-prefill-serving-v0"
exclude = ["*"]
[[hosts]]
name = "dash1"
ssh_alias = "dash1"
enabled = true
sync_remote_path = "/home/admin/cpfs/wjh/aituner/aituner"
fleet_root = "/home/admin/cpfs/wjh/aituner/gpu-fleet-tp2-prefill-serving-v0/dash1"
[[hosts]]
name = "dash2"
ssh_alias = "dash2"
enabled = true
sync_remote_path = "/home/admin/cpfs/wjh/aituner/aituner"
fleet_root = "/home/admin/cpfs/wjh/aituner/gpu-fleet-tp2-prefill-serving-v0/dash2"
[[hosts]]
name = "dash3"
ssh_alias = "dash3"
enabled = true
sync_remote_path = "/home/admin/cpfs/wjh/aituner/aituner"
fleet_root = "/home/admin/cpfs/wjh/aituner/gpu-fleet-tp2-prefill-serving-v0/dash3"
[[hosts]]
name = "dash4"
ssh_alias = "dash4"
enabled = true
sync_remote_path = "/home/admin/cpfs/wjh/aituner/aituner"
fleet_root = "/home/admin/cpfs/wjh/aituner/gpu-fleet-tp2-prefill-serving-v0/dash4"

View File

@@ -0,0 +1,16 @@
version = 1
[[jobs]]
name = "tp2-prefill-serving-smoke-20260723-v0"
gpus = 2
gpu_model = "H20"
hosts = ["dash1"]
command = "timeout --signal=TERM --kill-after=60s 3600 bash runs/frontier-tp2-prefill-serving-v0/run_prefill_profile.sh"
artifacts = ["runs/frontier-tp2-prefill-serving-v0/remote-outputs/tp2-smoke"]
[jobs.env]
TP = "2"
SERVER_PORT = "9827"
OUTPUT_ROOT = "runs/frontier-tp2-prefill-serving-v0/remote-outputs/tp2-smoke"
WARMUP_REQUESTS = "2"
FLASHINFER_WORKSPACE_BASE = "/tmp/frontier-tp2-prefill-serving-v0-flashinfer"

View File

@@ -0,0 +1,50 @@
#!/usr/bin/env python3
"""Create one exact-length, prefix-disjoint prefill request."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=Path, required=True)
parser.add_argument("--input-tokens", type=int, default=8192)
parser.add_argument("--output-tokens", type=int, default=2)
parser.add_argument("--output", type=Path, required=True)
return parser.parse_args()
def main() -> None:
args = parse_args()
if min(args.input_tokens, args.output_tokens) <= 0:
raise ValueError("token counts must be positive")
if args.input_tokens + args.output_tokens > 40960:
raise ValueError("request exceeds the server max model length")
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
special = set(tokenizer.all_special_ids)
candidates = [
token for token in range(tokenizer.vocab_size) if token not in special
]
if len(candidates) < 2:
raise ValueError("tokenizer has too few non-special token IDs")
body = {
"model": "qwen30-prefill-profile",
"prompt": [candidates[1], *([candidates[0]] * (args.input_tokens - 1))],
"min_tokens": args.output_tokens,
"max_tokens": args.output_tokens,
"ignore_eos": True,
"temperature": 0,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(body, separators=(",", ":")) + "\n")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,50 @@
#!/usr/bin/env python3
"""Schematic figure frozen before EXP-TP2-PREFILL-SERVING GPU execution."""
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
ROOT = Path(__file__).resolve().parent
fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.0), dpi=160)
colors = {"real": "#111111", "sim": "#d95f02", "candidate": "#1b75bc"}
labels = ["TP2", "TP4"]
x = np.arange(2)
width = 0.34
axes[0].bar(x - width / 2, [410, 231], width, color=colors["real"], label="real anchor")
axes[0].bar(x + width / 2, [356.25, 231.33], width, color=colors["sim"], label="structured sim")
axes[0].set_xticks(x, labels)
axes[0].set_ylabel("q8k / ctx0 chunk time (ms)")
axes[0].set_title("(a) Entry residual is TP2-only", loc="left", fontsize=10)
axes[0].legend(frameon=False, fontsize=8)
arms = ["A\nstandalone", "B\nwarm+route", "C\nserving"]
expected = [171.1, 205, 225]
axes[1].bar(np.arange(3), expected, color=[colors["sim"], colors["candidate"], colors["real"]])
axes[1].axhline(224.9, color="#777777", linestyle="--", linewidth=1, label="needed to close 54 ms")
axes[1].set_xticks(np.arange(3), arms)
axes[1].set_ylabel("TP2 MoE per chunk (ms)")
axes[1].set_title("(b) Three-arm decision (schematic)", loc="left", fontsize=10)
axes[1].legend(frameon=False, fontsize=8)
for ax in axes:
ax.spines[["top", "right"]].set_visible(False)
ax.grid(axis="y", color="#dedbd2", linewidth=0.8)
ax.set_axisbelow(True)
fig.suptitle(
"MOCK / schematic — EXP-TP2-PREFILL-SERVING (B/C not measured)",
x=0.01,
ha="left",
color=colors["sim"],
fontsize=9,
)
fig.tight_layout(rect=(0, 0, 1, 0.95))
fig.savefig(ROOT / "figure-prototype.png")
fig.savefig(ROOT / "figure-prototype.svg")
print(ROOT / "figure-prototype.png")

View File

@@ -0,0 +1,24 @@
{
"schema": "frontier-tp2-prefill-serving-entry-v1",
"cells": {
"tp2": {
"total_ms": 356.2487769916667,
"attention_prefill_ms": 100.244853113,
"moe_grouped_gemm_ms": 171.118426598,
"samples": 9,
"real_total_ms": 410.0,
"total_bias": -0.1311005439227641
},
"tp4": {
"total_ms": 231.3315694846667,
"attention_prefill_ms": 51.95132805766667,
"moe_grouped_gemm_ms": 116.27978482099999,
"samples": 9,
"real_total_ms": 231.0,
"total_bias": 0.0014353657344878235
}
},
"tp2_required_moe_if_residual_is_all_moe_ms": 224.86964960633333,
"tp2_required_moe_shift": 0.314117094675072,
"entry_gate": true
}

View File

@@ -0,0 +1,163 @@
#!/usr/bin/env bash
set -euo pipefail
TP="${TP:?TP is required}"
OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
SERVER_PORT="${SERVER_PORT:?SERVER_PORT is required}"
VENV_ROOT="${VENV_ROOT:-/home/admin/cpfs/wjh/venvs/vllm-0.20.0-cu129-workload-regime-v2}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.92}"
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-900}"
WARMUP_REQUESTS="${WARMUP_REQUESTS:-2}"
SCRIPT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MATERIALIZER="${SCRIPT_ROOT}/materialize_prefill_request.py"
CLIENT="${SCRIPT_ROOT}/run_request.py"
SERVER_PID=""
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?GPU allocation is required}"
if [[ "${#GPU_IDS[@]}" -ne "${TP}" ]]; then
echo "ERROR: TP=${TP}, but CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}" >&2
exit 1
fi
if [[ ! "${WARMUP_REQUESTS}" =~ ^[1-9][0-9]*$ ]]; then
echo "ERROR: WARMUP_REQUESTS must be a positive integer" >&2
exit 1
fi
mkdir -p "${OUTPUT_ROOT}"
OUTPUT_ROOT="$(cd "${OUTPUT_ROOT}" && pwd)"
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" \
"${OUTPUT_ROOT}/requests" "${OUTPUT_ROOT}/results" \
"${OUTPUT_ROOT}/trace-staging" "${OUTPUT_ROOT}/traces/profile"
exec > >(tee -a "${OUTPUT_ROOT}/logs/controller.log") 2>&1
cleanup() {
if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
kill -TERM -- "-${SERVER_PID}" 2>/dev/null || true
for _ in $(seq 1 30); do
kill -0 "${SERVER_PID}" 2>/dev/null || break
sleep 1
done
kill -KILL -- "-${SERVER_PID}" 2>/dev/null || true
fi
SERVER_PID=""
}
trap cleanup EXIT INT TERM
sha256sum "${BASH_SOURCE[0]}" "${MATERIALIZER}" "${CLIENT}" \
"${MODEL_ROOT}/config.json" > "${OUTPUT_ROOT}/provenance/inputs.sha256"
"${VENV_ROOT}/bin/python" -c \
'import torch, transformers, vllm; print(f"torch={torch.__version__}"); print(f"transformers={transformers.__version__}"); print(f"vllm={vllm.__version__}")' \
> "${OUTPUT_ROOT}/provenance/runtime.versions"
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total \
--format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.before.csv"
ps -eo user,pid,ppid,etimes,pcpu,pmem,args --sort=pid \
> "${OUTPUT_ROOT}/provenance/processes.before.txt"
uptime > "${OUTPUT_ROOT}/provenance/uptime.before.txt"
env | sort > "${OUTPUT_ROOT}/provenance/environment.txt"
PROFILE_CONFIG="$("${VENV_ROOT}/bin/python" - "${OUTPUT_ROOT}/trace-staging" <<'PY'
import json
import sys
print(json.dumps({
"profiler": "torch",
"torch_profiler_dir": sys.argv[1],
"torch_profiler_with_stack": False,
"torch_profiler_record_shapes": True,
"torch_profiler_use_gzip": True,
"ignore_frontend": True,
"wait_iterations": 0,
"warmup_iterations": 0,
"active_iterations": 2,
}, separators=(",", ":")))
PY
)"
printf '%s\n' "${PROFILE_CONFIG}" > "${OUTPUT_ROOT}/provenance/profiler-config.json"
export TOKENIZERS_PARALLELISM=false
export VLLM_USE_V1=1
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${OUTPUT_ROOT}/flashinfer-workspace}"
mkdir -p "${FLASHINFER_WORKSPACE_BASE}"
ulimit -n 65536
"${VENV_ROOT}/bin/python" "${MATERIALIZER}" --model "${MODEL_ROOT}" \
--input-tokens 8192 --output-tokens 2 \
--output "${OUTPUT_ROOT}/requests/q8192-o2.json"
printf 'LAUNCH host=%s tp=%s gpus=%s output=%s warmups=%s\n' \
"$(hostname)" "${TP}" "${CUDA_VISIBLE_DEVICES}" "${OUTPUT_ROOT}" \
"${WARMUP_REQUESTS}"
setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
--host 127.0.0.1 --port "${SERVER_PORT}" \
--served-model-name qwen30-prefill-profile \
--tensor-parallel-size "${TP}" \
--gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \
--max-model-len 40960 --max-num-batched-tokens 8192 --max-num-seqs 16 \
--no-enable-prefix-caching --enable-chunked-prefill --no-enable-log-requests \
--enable-logging-iteration-details --profiler-config "${PROFILE_CONFIG}" \
> "${OUTPUT_ROOT}/logs/server.log" 2>&1 &
SERVER_PID=$!
READY=0
for _ in $(seq 1 "${SERVER_READY_ATTEMPTS}"); do
if curl -fsS --max-time 2 \
"http://127.0.0.1:${SERVER_PORT}/v1/models" \
> "${OUTPUT_ROOT}/results/models.json" 2>/dev/null; then
READY=1
break
fi
if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
tail -200 "${OUTPUT_ROOT}/logs/server.log"
exit 1
fi
sleep 3
done
if [[ "${READY}" -ne 1 ]]; then
echo "ERROR: server readiness timeout" >&2
tail -200 "${OUTPUT_ROOT}/logs/server.log"
exit 1
fi
for index in $(seq 1 "${WARMUP_REQUESTS}"); do
"${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" \
--request "${OUTPUT_ROOT}/requests/q8192-o2.json" \
--output "${OUTPUT_ROOT}/results/warmup-${index}.json"
done
curl -fsS -X POST "http://127.0.0.1:${SERVER_PORT}/start_profile" \
> "${OUTPUT_ROOT}/logs/start-profile.txt"
"${VENV_ROOT}/bin/python" "${CLIENT}" --port "${SERVER_PORT}" \
--request "${OUTPUT_ROOT}/requests/q8192-o2.json" \
--output "${OUTPUT_ROOT}/results/profile.json"
curl -fsS -X POST "http://127.0.0.1:${SERVER_PORT}/stop_profile" \
> "${OUTPUT_ROOT}/logs/stop-profile.txt"
deadline=$((SECONDS + 120))
while (( SECONDS < deadline )); do
trace_count="$(find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
-name '*.pt.trace.json*' | wc -l)"
if (( trace_count >= TP )); then
break
fi
sleep 2
done
trace_count="$(find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
-name '*.pt.trace.json*' | wc -l)"
if (( trace_count < TP )); then
echo "ERROR: expected ${TP} rank traces, found ${trace_count}" >&2
exit 1
fi
find "${OUTPUT_ROOT}/trace-staging" -maxdepth 1 -type f \
-name '*.pt.trace.json*' -exec mv -t "${OUTPUT_ROOT}/traces/profile" {} +
cleanup
nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total \
--format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.after.csv"
find "${OUTPUT_ROOT}" -type f \
! -path '*/provenance/artifacts.sha256' -print0 \
| sort -z | xargs -0 sha256sum \
> "${OUTPUT_ROOT}/provenance/artifacts.sha256"
echo PREFILL_PROFILE_COMPLETE

View File

@@ -0,0 +1,61 @@
#!/usr/bin/env python3
"""Issue one non-streaming completion request and record wall/usage."""
from __future__ import annotations
import argparse
import json
import time
import urllib.request
from pathlib import Path
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, required=True)
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--timeout-seconds", type=float, default=1800)
return parser.parse_args()
def main() -> None:
args = parse_args()
body = json.loads(args.request.read_text())
expected_input = len(body["prompt"])
expected_output = int(body["max_tokens"])
request = urllib.request.Request(
f"http://127.0.0.1:{args.port}/v1/completions",
data=json.dumps(body, separators=(",", ":")).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
started = time.monotonic()
with urllib.request.urlopen(
request, timeout=args.timeout_seconds
) as response:
payload = json.load(response)
wall_ms = (time.monotonic() - started) * 1000
usage = payload.get("usage") or {}
observed = (
int(usage.get("prompt_tokens", -1)),
int(usage.get("completion_tokens", -1)),
)
if observed != (expected_input, expected_output):
raise ValueError(
f"usage mismatch: expected {expected_input}+{expected_output}, "
f"observed {observed[0]}+{observed[1]}"
)
result = {
"wall_ms": wall_ms,
"input_tokens": observed[0],
"output_tokens": observed[1],
"usage": usage,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
print(json.dumps(result, sort_keys=True), flush=True)
if __name__ == "__main__":
main()