401 lines
15 KiB
Python
401 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""Freeze the Qwen30 fixed-shape prefill-only Frontier surface."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
MODEL = "qwen3-a3b-30b-moe"
|
|
RATES = (4.0, 8.0, 16.0, 32.0, 64.0)
|
|
TTFT_SLO_MS = 1256.0
|
|
TARGET_PASS_RATE = 0.95
|
|
NUM_BLOCKS = {1: 20080, 2: 76537, 4: 191727}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Config:
|
|
tp: int
|
|
mns: int
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return f"tp{self.tp}_mns{self.mns}"
|
|
|
|
|
|
GRID = tuple(Config(tp, mns) for tp in (1, 2, 4) for mns in (8, 16, 32, 64))
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--frontier-source", type=Path, required=True)
|
|
parser.add_argument("--replayserve-root", type=Path, required=True)
|
|
parser.add_argument("--profile-root", type=Path, required=True)
|
|
parser.add_argument("--python-deps", type=Path, required=True)
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
parser.add_argument("--requests", type=int, default=64)
|
|
parser.add_argument("--rate", type=float, action="append")
|
|
parser.add_argument("--config", action="append")
|
|
parser.add_argument("--timeout-seconds", type=float, default=900.0)
|
|
parser.add_argument("--resume", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1 << 20), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def write_json(path: Path, payload: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def load_module(name: str, path: Path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
if spec is None or spec.loader is None:
|
|
raise ImportError(path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def write_trace(path: Path, *, requests: int, rate: float) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
fields = [
|
|
"arrived_at",
|
|
"num_prefill_tokens",
|
|
"num_decode_tokens",
|
|
"session_id",
|
|
"block_hash_ids",
|
|
]
|
|
with path.open("w", newline="") as output:
|
|
writer = csv.DictWriter(output, fieldnames=fields)
|
|
writer.writeheader()
|
|
for request_id in range(requests):
|
|
writer.writerow(
|
|
{
|
|
"arrived_at": f"{request_id / rate:.12f}",
|
|
"num_prefill_tokens": 2048,
|
|
"num_decode_tokens": 1,
|
|
"session_id": request_id,
|
|
"block_hash_ids": "|".join(
|
|
str(request_id * 128 + block + 1) for block in range(128)
|
|
),
|
|
}
|
|
)
|
|
|
|
|
|
def profile_paths(root: Path) -> dict[str, Path]:
|
|
paths = {
|
|
"linear": root / "linear_op.csv",
|
|
"attention": root / "attention.csv",
|
|
"moe": root / "moe.csv",
|
|
"manifest": root / "manifest.json",
|
|
}
|
|
missing = [str(path) for path in paths.values() if not path.is_file()]
|
|
if missing:
|
|
raise FileNotFoundError(missing)
|
|
return paths
|
|
|
|
|
|
def validate_profile(paths: dict[str, Path]) -> dict[str, Any]:
|
|
manifest = json.loads(paths["manifest"].read_text())
|
|
expected = manifest["outputs"]
|
|
for filename in ("linear_op.csv", "attention.csv", "moe.csv"):
|
|
actual = sha256(paths[{"linear_op.csv": "linear", "attention.csv": "attention", "moe.csv": "moe"}[filename]])
|
|
if actual != expected[filename]:
|
|
raise ValueError(f"profile hash mismatch for {filename}")
|
|
with paths["attention"].open(newline="") as source:
|
|
rows = list(csv.DictReader(source))
|
|
coverage = {}
|
|
for tp in (1, 2, 4):
|
|
exact = [
|
|
row
|
|
for row in rows
|
|
if int(row["num_tensor_parallel_workers"]) == tp
|
|
and row["is_prefill"].lower() == "true"
|
|
and row.get("is_true_mixed_batch", "").lower() != "true"
|
|
and int(float(row["total_tokens"])) == 2048
|
|
]
|
|
if len(exact) != 1:
|
|
raise ValueError(f"expected one exact TP{tp} 2048-token prefill row, got {len(exact)}")
|
|
coverage[str(tp)] = {"exact_prefill_2048_rows": len(exact), "profile_batch_size": int(exact[0]["batch_size"])}
|
|
return {"manifest": manifest, "attention": coverage}
|
|
|
|
|
|
def knobs(config: Config, paths: dict[str, Path], cache: Path) -> dict[str, Any]:
|
|
return {
|
|
"simulation_mode": "online",
|
|
"sys_arch": "co-location",
|
|
"num_replicas": 1,
|
|
"cluster_scheduler": "sticky_round_robin",
|
|
"model_name": MODEL,
|
|
"device": "h20",
|
|
"network_device": "h20_dgx",
|
|
"attn_tensor_parallel_size": config.tp,
|
|
"attn_data_parallel_size": 1,
|
|
"moe_tensor_parallel_size": config.tp,
|
|
"moe_expert_parallel_size": 1,
|
|
"num_pipeline_stages": 1,
|
|
"replica_scheduler": "vllm_v1",
|
|
"decode_cuda_graph_mode": "none",
|
|
"batch_size_cap": config.mns,
|
|
"max_tokens_in_batch": 8192,
|
|
"long_prefill_token_threshold": 0,
|
|
"block_size": 16,
|
|
"num_blocks_mode": "explicit",
|
|
"num_blocks": NUM_BLOCKS[config.tp],
|
|
"gpu_memory_utilization": 0.92,
|
|
"non_kv_cache_overhead_bytes": 0,
|
|
"trace_max_tokens": 40960,
|
|
"cache_dir": str(cache),
|
|
"enable_dummy_mode": False,
|
|
"linear_op_input_file": str(paths["linear"]),
|
|
"atten_input_file": str(paths["attention"]),
|
|
"moe_input_file": str(paths["moe"]),
|
|
"prediction_max_prefill_chunk_size": 18000,
|
|
"prediction_max_batch_size": 128,
|
|
"prediction_max_tokens_per_request": 32768,
|
|
"no_cache": False,
|
|
"skip_cpu_overhead_modeling": True,
|
|
"enable_prefix_caching": False,
|
|
"enable_chunked_prefill": True,
|
|
}
|
|
|
|
|
|
def find_metrics(run_dir: Path) -> tuple[Path, Path]:
|
|
systems = list((run_dir / "metrics").rglob("system_metrics.json"))
|
|
requests = list((run_dir / "metrics").rglob("request_metrics.csv"))
|
|
if len(systems) != 1 or len(requests) != 1:
|
|
raise RuntimeError(f"metric pair mismatch: {len(systems)}/{len(requests)}")
|
|
return systems[0], requests[0]
|
|
|
|
|
|
def score(system_path: Path, request_path: Path, expected_requests: int) -> dict[str, Any]:
|
|
system = json.loads(system_path.read_text())
|
|
metadata = system["simulation_metadata"]
|
|
if int(metadata["completed_requests"]) != expected_requests:
|
|
raise ValueError("Frontier completion count mismatch")
|
|
with request_path.open(newline="") as source:
|
|
rows = list(csv.DictReader(source))
|
|
if len(rows) != expected_requests:
|
|
raise ValueError("request metric count mismatch")
|
|
values = []
|
|
passed = 0
|
|
for row in rows:
|
|
if int(float(row["request_num_prefill_tokens"])) != 2048:
|
|
raise ValueError("prefill shape drift")
|
|
if int(float(row["request_num_decode_tokens"])) != 1:
|
|
raise ValueError("decode shape drift")
|
|
ttft = float(row["ttft"])
|
|
if not math.isfinite(ttft) or ttft < 0:
|
|
raise ValueError("invalid TTFT")
|
|
values.append(ttft)
|
|
passed += int(ttft <= TTFT_SLO_MS)
|
|
ordered = sorted(values)
|
|
pass_rate = passed / expected_requests
|
|
return {
|
|
"ttft_p50_ms": ordered[math.ceil(0.50 * len(ordered)) - 1],
|
|
"ttft_p95_ms": ordered[math.ceil(0.95 * len(ordered)) - 1],
|
|
"ttft_max_ms": max(ordered),
|
|
"passed": passed,
|
|
"pass_rate": pass_rate,
|
|
"feasible": pass_rate >= TARGET_PASS_RATE,
|
|
"throughput_requests_per_second": float(system["throughput_metrics"]["requests_per_second"]),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
args.frontier_source = args.frontier_source.resolve()
|
|
args.replayserve_root = args.replayserve_root.resolve()
|
|
args.profile_root = args.profile_root.resolve()
|
|
args.python_deps = args.python_deps.resolve()
|
|
args.output_root = args.output_root.resolve()
|
|
rates = tuple(args.rate or RATES)
|
|
selected = list(GRID)
|
|
if args.config:
|
|
wanted = set(args.config)
|
|
selected = [config for config in GRID if config.name in wanted]
|
|
if {config.name for config in selected} != wanted:
|
|
raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
|
|
paths = profile_paths(args.profile_root)
|
|
coverage = validate_profile(paths)
|
|
builder = load_module(
|
|
"qwen30_prefill_frontier_builder",
|
|
args.replayserve_root / "tools/run_frontier_sweep.py",
|
|
)
|
|
frontier_head = subprocess.run(
|
|
["git", "-C", str(args.frontier_source), "rev-parse", "HEAD"],
|
|
check=True,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
).stdout.strip()
|
|
traces = {}
|
|
for rate in rates:
|
|
trace = args.output_root / "traces" / f"r{rate:g}.csv"
|
|
write_trace(trace, requests=args.requests, rate=rate)
|
|
traces[rate] = trace
|
|
|
|
config_results = []
|
|
for config in selected:
|
|
loads = []
|
|
config_knobs = knobs(config, paths, args.output_root / "cache")
|
|
for rate in rates:
|
|
run_dir = args.output_root / "runs" / config.name / f"r{rate:g}"
|
|
result_path = run_dir / "result.json"
|
|
if args.resume and result_path.is_file():
|
|
loads.append(json.loads(result_path.read_text()))
|
|
continue
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
command = builder.build_frontier_command(
|
|
python_bin="/usr/bin/python3",
|
|
trace_file=traces[rate],
|
|
metrics_root=run_dir / "metrics",
|
|
run_id=f"qwen30_prefill_{config.name}_r{rate:g}",
|
|
knobs=config_knobs,
|
|
)
|
|
write_json(run_dir / "command.json", command)
|
|
environment = os.environ.copy()
|
|
pythonpath = [str(args.python_deps), str(args.frontier_source)]
|
|
if environment.get("PYTHONPATH"):
|
|
pythonpath.append(environment["PYTHONPATH"])
|
|
environment.update(
|
|
{
|
|
"PYTHONPATH": ":".join(pythonpath),
|
|
"CUDA_VISIBLE_DEVICES": "",
|
|
"NVIDIA_VISIBLE_DEVICES": "void",
|
|
"WANDB_DISABLED": "true",
|
|
"VIDUR_DISABLE_WANDB": "1",
|
|
"FRONTIER_LOG_LEVEL": "WARNING",
|
|
"PYTHONDONTWRITEBYTECODE": "1",
|
|
}
|
|
)
|
|
started = time.time()
|
|
with (run_dir / "stdout.log").open("w") as stdout, (
|
|
run_dir / "stderr.log"
|
|
).open("w") as stderr:
|
|
completed = subprocess.run(
|
|
command,
|
|
cwd=args.frontier_source,
|
|
env=environment,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
timeout=args.timeout_seconds,
|
|
check=False,
|
|
)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(
|
|
f"Frontier failed for {config.name} rate={rate}: {completed.returncode}"
|
|
)
|
|
system_path, request_path = find_metrics(run_dir)
|
|
result = {
|
|
"status": "completed",
|
|
"config": asdict(config) | {"name": config.name},
|
|
"offered_request_rate": rate,
|
|
"offered_request_rate_per_gpu": rate / config.tp,
|
|
"request_count": args.requests,
|
|
"elapsed_seconds": time.time() - started,
|
|
"trace_sha256": sha256(traces[rate]),
|
|
"request_metrics_sha256": sha256(request_path),
|
|
"score": score(system_path, request_path, args.requests),
|
|
}
|
|
write_json(result_path, result)
|
|
loads.append(result)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"config": config.name,
|
|
"rate": rate,
|
|
"pass_rate": result["score"]["pass_rate"],
|
|
"feasible": result["score"]["feasible"],
|
|
},
|
|
sort_keys=True,
|
|
),
|
|
flush=True,
|
|
)
|
|
config_results.append({"config": asdict(config) | {"name": config.name}, "loads": loads})
|
|
|
|
capacities = []
|
|
for item in config_results:
|
|
feasible = [
|
|
load["offered_request_rate"]
|
|
for load in item["loads"]
|
|
if load["score"]["feasible"]
|
|
]
|
|
capacity = max(feasible) if feasible else None
|
|
capacities.append(
|
|
{
|
|
"config": item["config"],
|
|
"maximum_tested_feasible_request_rate": capacity,
|
|
"maximum_tested_feasible_request_rate_per_gpu": (
|
|
capacity / item["config"]["tp"] if capacity is not None else None
|
|
),
|
|
"lower_censored": capacity is None,
|
|
"upper_censored": capacity == max(rates) if capacity is not None else False,
|
|
}
|
|
)
|
|
capacities.sort(
|
|
key=lambda row: (
|
|
-(row["maximum_tested_feasible_request_rate_per_gpu"] or -1),
|
|
row["config"]["name"],
|
|
)
|
|
)
|
|
full = selected == list(GRID) and rates == RATES and args.requests == 64
|
|
manifest = {
|
|
"schema": "frontier-qwen30-prefill-surface-v1",
|
|
"status": "frozen_before_real" if full else "partial_not_decision_bearing",
|
|
"contract": {
|
|
"rates": rates,
|
|
"requests_per_anchor": args.requests,
|
|
"input_tokens": 2048,
|
|
"output_tokens": 1,
|
|
"ttft_slo_ms": TTFT_SLO_MS,
|
|
"target_pass_rate": TARGET_PASS_RATE,
|
|
"prefix_caching": False,
|
|
"arrival": "open_loop_uniform",
|
|
},
|
|
"frontier": {
|
|
"source": str(args.frontier_source),
|
|
"git_head": frontier_head,
|
|
"git_status_short": subprocess.run(
|
|
["git", "-C", str(args.frontier_source), "status", "--short"],
|
|
check=True,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
).stdout,
|
|
},
|
|
"profiles": {
|
|
"root": str(args.profile_root),
|
|
"coverage": coverage,
|
|
"sha256": {name: sha256(path) for name, path in paths.items()},
|
|
},
|
|
"config_results": config_results,
|
|
"capacity": capacities,
|
|
}
|
|
write_json(args.output_root / "frontier_surface_frozen.json", manifest)
|
|
print(args.output_root / "frontier_surface_frozen.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|