gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Align Frontier piecewise graph profiles

Browse Source
This commit is contained in:
Gahow Wang
2026-07-17 23:22:42 +08:00
parent 47355a9411
commit bdc357dc6c
10 changed files with 804 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

176
runs/frontier-fidelity-envelope-v1/run_frontier_qwen30_exact_trace_surface.py
View File

@@ -19,6 +19,27 @@ from typing import Any
TARGET_PASS_RATE = 0.95
TPOT_SLOS_MS = (50.0, 100.0, 150.0, 180.0)
WINDOW_SECONDS = 600.0
GRAPH_CAPTURE_SIZES_BY_MNS = {
8: (1, 2, 4, 8, 16),
16: (1, 2, 4, 8, 16, 24, 32),
32: (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64),
64: (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128),
}
REAL_NUM_BLOCKS_BY_CONFIG = {
(1, 8): 20137,
(1, 16): 20128,
(1, 32): 20108,
(1, 64): 20069,
(2, 8): 76639,
(2, 16): 76620,
(2, 32): 76583,
(2, 64): 76505,
(4, 8): 191930,
(4, 16): 191882,
(4, 32): 191786,
(4, 64): 191589,
}
KERNEL_DECODE_KV_CONTEXTS = (128, 1024, 2048, 4096, 8192, 16384, 32768, 40960)
BASE_RUNNER = (
Path(__file__).resolve().parents[1]
/ "frontier-phase-factorial-v0/run_frontier_qwen30_prefill_surface.py"
@@ -43,6 +64,7 @@ def parse_args() -> argparse.Namespace:
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("--kernel-profile-root", type=Path)
parser.add_argument("--python-deps", type=Path, required=True)
parser.add_argument("--output-root", type=Path, required=True)
parser.add_argument(
@@ -68,6 +90,21 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--allreduce-csv", type=Path)
parser.add_argument("--timeout-seconds", type=float, default=1800.0)
parser.add_argument("--predictor-training-job-threads", type=int, default=1)
parser.add_argument(
"--decode-cuda-graph-mode",
choices=("none", "full_decode_only", "piecewise"),
default="none",
)
parser.add_argument(
"--align-real-graph-runtime",
action="store_true",
help="Use real observed capture lists and per-(TP,MNS) KV blocks.",
)
parser.add_argument(
"--fresh-predictor-cache",
action="store_true",
help="Disable Frontier predictor cache reuse for this profile family.",
)
parser.add_argument("--resume", action="store_true")
parser.add_argument("--continue-on-failure", action="store_true")
return parser.parse_args()
@@ -241,15 +278,99 @@ def score(path: Path, expected_shapes: list[tuple[int, int]]) -> dict[str, Any]:
}
ttfts = [float(row["ttft_ms"]) for row in request_metrics]
tpots = [float(row["tpot_ms"]) for row in request_metrics if row["tpot_ms"] is not None]
e2es = [float(row["e2e_ms"]) for row in request_metrics]
return {
"ttft_mean_ms": sum(ttfts) / len(ttfts),
"ttft_p50_ms": percentile(ttfts, 0.50),
"ttft_p90_ms": percentile(ttfts, 0.90),
"ttft_p95_ms": percentile(ttfts, 0.95),
"tpot_mean_ms": sum(tpots) / len(tpots),
"tpot_p50_ms": percentile(tpots, 0.50),
"tpot_p90_ms": percentile(tpots, 0.90),
"tpot_p95_ms": percentile(tpots, 0.95),
"e2e_mean_ms": sum(e2es) / len(e2es),
"e2e_p50_ms": percentile(e2es, 0.50),
"e2e_p90_ms": percentile(e2es, 0.90),
"e2e_p95_ms": percentile(e2es, 0.95),
"slos": slos,
}
def kernel_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_kernel_profile(paths: dict[str, Path]) -> dict[str, Any]:
manifest = json.loads(paths["manifest"].read_text())
outputs = manifest.get("outputs", {})
for filename, name in (
("linear_op.csv", "linear"),
("attention.csv", "attention"),
("moe.csv", "moe"),
):
if outputs.get(filename) != BASE.sha256(paths[name]):
raise ValueError(f"kernel-only profile hash mismatch for {filename}")
with paths["linear"].open(newline="") as source:
linear_rows = list(csv.DictReader(source))
with paths["attention"].open(newline="") as source:
attention_rows = list(csv.DictReader(source))
with paths["moe"].open(newline="") as source:
moe_rows = list(csv.DictReader(source))
for label, rows in (("linear", linear_rows), ("attention", attention_rows), ("moe", moe_rows)):
if not rows or {row.get("measurement_type") for row in rows} != {"KERNEL_ONLY"}:
raise ValueError(f"{label} lacks an exclusive KERNEL_ONLY measurement family")
required_buckets = set(GRAPH_CAPTURE_SIZES_BY_MNS[64])
coverage: dict[str, Any] = {}
for tp in (1, 2, 4):
linear_tokens = {
int(float(row["num_tokens"]))
for row in linear_rows
if int(float(row["num_tensor_parallel_workers"])) == tp
}
moe_tokens = {
int(float(row["num_tokens"]))
for row in moe_rows
if int(float(row["num_tensor_parallel_workers"])) == tp
}
attention_pairs = {
(int(float(row["batch_size"])), int(float(row["kv_cache_size"])))
for row in attention_rows
if int(float(row["num_tensor_parallel_workers"])) == tp
and row["is_prefill"].lower() == "false"
and row.get("is_true_mixed_batch", "").lower() != "true"
}
missing_linear = required_buckets - linear_tokens
missing_moe = required_buckets - moe_tokens
missing_attention = {
(bucket, kv)
for bucket in required_buckets
for kv in KERNEL_DECODE_KV_CONTEXTS
if (bucket, kv) not in attention_pairs
}
if missing_linear or missing_moe or missing_attention:
raise ValueError(
f"kernel-only profile coverage TP{tp}: linear={sorted(missing_linear)}, "
f"moe={sorted(missing_moe)}, attention={sorted(missing_attention)}"
)
coverage[str(tp)] = {
"linear_tokens": sorted(linear_tokens),
"moe_tokens": sorted(moe_tokens),
"attention_decode_pairs": len(attention_pairs),
}
return {"manifest": manifest, "coverage": coverage}
def main() -> None:
args = parse_args()
if args.predictor_training_job_threads <= 0:
@@ -264,6 +385,10 @@ def main() -> None:
setattr(args, name, getattr(args, name).resolve())
if args.allreduce_csv is not None:
args.allreduce_csv = args.allreduce_csv.resolve()
if args.kernel_profile_root is not None:
args.kernel_profile_root = args.kernel_profile_root.resolve()
if args.decode_cuda_graph_mode == "none":
raise ValueError("--kernel-profile-root requires a non-none graph mode")
traces = [
parse_trace(
specification,
@@ -288,6 +413,11 @@ def main() -> None:
raise ValueError(f"unknown configs: {wanted - {config.name for config in selected}}")
paths = BASE.profile_paths(args.profile_root)
coverage = BASE.validate_profile(paths)
kernel_paths = None
kernel_coverage = None
if args.kernel_profile_root is not None:
kernel_paths = kernel_profile_paths(args.kernel_profile_root)
kernel_coverage = validate_kernel_profile(kernel_paths)
builder = BASE.load_module(
"qwen30_exact_trace_frontier_builder",
args.replayserve_root / "tools/run_frontier_sweep.py",
@@ -320,6 +450,18 @@ def main() -> None:
config_knobs = BASE.knobs(config, paths, args.output_root / "cache")
config_knobs["enable_prefix_caching"] = args.prefix_caching
config_knobs["prediction_max_tokens_per_request"] = 40960
config_knobs["decode_cuda_graph_mode"] = args.decode_cuda_graph_mode
config_knobs["no_cache"] = args.fresh_predictor_cache
if args.align_real_graph_runtime:
config_knobs["num_blocks"] = REAL_NUM_BLOCKS_BY_CONFIG[(config.tp, config.mns)]
if kernel_paths is not None:
config_knobs.update(
{
"linear_op_kernel_only_input_file": str(kernel_paths["linear"]),
"atten_kernel_only_input_file": str(kernel_paths["attention"]),
"moe_kernel_only_input_file": str(kernel_paths["moe"]),
}
)
for trace in traces:
run_dir = args.output_root / "runs" / config.name / trace["label"]
result_path = run_dir / "result.json"
@@ -340,6 +482,13 @@ def main() -> None:
str(args.predictor_training_job_threads),
]
)
if args.align_real_graph_runtime:
command.extend(
[
"--cudagraph_capture_sizes",
*(str(size) for size in GRAPH_CAPTURE_SIZES_BY_MNS[config.mns]),
]
)
command = BASE.configure_cc_command(
command,
backend=args.cc_backend,
@@ -514,6 +663,9 @@ def main() -> None:
"primary_tpot_slo_ms": 150.0,
"target_pass_rate": TARGET_PASS_RATE,
"predictor_training_job_threads": args.predictor_training_job_threads,
"decode_cuda_graph_mode": args.decode_cuda_graph_mode,
"align_real_graph_runtime": args.align_real_graph_runtime,
"fresh_predictor_cache": args.fresh_predictor_cache,
},
"frontier": {
"source": str(args.frontier_source),
@@ -530,6 +682,30 @@ def main() -> None:
"coverage": coverage,
"sha256": {name: BASE.sha256(path) for name, path in paths.items()},
},
"kernel_only_profiles": (
None
if kernel_paths is None
else {
"root": str(args.kernel_profile_root),
"coverage": kernel_coverage,
"sha256": {
name: BASE.sha256(path) for name, path in kernel_paths.items()
},
}
),
"runtime_alignment": {
"capture_sizes_by_mns": (
GRAPH_CAPTURE_SIZES_BY_MNS if args.align_real_graph_runtime else None
),
"num_blocks_by_config": (
{
f"tp{tp}_mns{mns}": blocks
for (tp, mns), blocks in REAL_NUM_BLOCKS_BY_CONFIG.items()
}
if args.align_real_graph_runtime
else None
),
},
"collective": {
"backend": args.cc_backend,
"allreduce_csv": str(args.allreduce_csv) if args.allreduce_csv else None,