#!/usr/bin/env python3 """Create an immutable Frontier profile root with measured decode rows.""" from __future__ import annotations import argparse import csv import hashlib import json import shutil import tempfile from pathlib import Path from typing import Any MODEL = "Qwen3-235B-A22B-FP8" ATTENTION_RELATIVE_PATH = Path("compute/h20") / MODEL / "attention.csv" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--base-profile-root", type=Path, required=True) parser.add_argument("--decode-attention-csv", type=Path, required=True) parser.add_argument("--true-mixed-attention-csv", type=Path) parser.add_argument("--output-root", type=Path, required=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 read_csv(path: Path) -> tuple[list[str], list[dict[str, str]]]: with path.open(newline="") as source: reader = csv.DictReader(source) if reader.fieldnames is None: raise ValueError(f"missing CSV header: {path}") return reader.fieldnames, list(reader) def is_true(value: str) -> bool: return value.strip().lower() == "true" def profile_hashes(root: Path) -> dict[str, str]: return { str(path.relative_to(root)): sha256(path) for path in sorted(root.rglob("*")) if path.is_file() and path.name != "profile_closure_manifest.json" } def write_json(path: Path, payload: Any) -> None: path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") def main() -> None: args = parse_args() base_root = args.base_profile_root.resolve() decode_csv = args.decode_attention_csv.resolve() output_root = args.output_root.resolve() base_attention = base_root / ATTENTION_RELATIVE_PATH if output_root.exists(): raise FileExistsError(f"refusing to overwrite profile root: {output_root}") for path in (base_attention, decode_csv): if not path.is_file(): raise FileNotFoundError(path) base_fields, base_rows = read_csv(base_attention) decode_fields, decode_source_rows = read_csv(decode_csv) if base_fields != decode_fields: raise ValueError("base and decode attention CSV schemas differ") if not base_rows or any(not is_true(row["is_prefill"]) for row in base_rows): raise ValueError("base attention profile must contain only prefill rows") decode_rows = [row for row in decode_source_rows if not is_true(row["is_prefill"])] if not decode_rows: raise ValueError("decode attention profile contains no decode rows") if any(not row["time_stats.attn_decode.median"] for row in decode_rows): raise ValueError("decode attention profile has an empty median") dimensions = { ( int(row["num_tensor_parallel_workers"]), int(row["batch_size"]), int(row["kv_cache_size"]), row["attention_backend"], ) for row in decode_rows } if len(dimensions) != len(decode_rows): raise ValueError("decode attention profile has duplicate coverage coordinates") if any(dimension[-1] != "FLASHINFER" for dimension in dimensions): raise ValueError("decode attention profile is not entirely FlashInfer") true_mixed_csv = ( args.true_mixed_attention_csv.resolve() if args.true_mixed_attention_csv is not None else None ) true_mixed_rows: list[dict[str, str]] = [] output_fields = list(base_fields) if true_mixed_csv is not None: if not true_mixed_csv.is_file(): raise FileNotFoundError(true_mixed_csv) true_mixed_fields, true_mixed_rows = read_csv(true_mixed_csv) required = { "is_true_mixed_batch", "decode_batch_size", "decode_avg_kv_cache_size", "num_prefill_seqs", "time_stats.attn_decode.median", "time_stats.attn_prefill.median", } missing = required - set(true_mixed_fields) if missing: raise ValueError(f"true-mixed attention CSV lacks columns: {sorted(missing)}") if not true_mixed_rows or any( not is_true(row["is_true_mixed_batch"]) for row in true_mixed_rows ): raise ValueError("true-mixed attention CSV has non-mixed rows") if any( not row["time_stats.attn_decode.median"] or not row["time_stats.attn_prefill.median"] for row in true_mixed_rows ): raise ValueError("true-mixed attention profile has an empty median") if { int(row["num_tensor_parallel_workers"]) for row in true_mixed_rows } != {4, 8}: raise ValueError("true-mixed attention profile must cover TP4 and TP8") output_fields.extend( field for field in true_mixed_fields if field not in output_fields ) for row in [*base_rows, *decode_rows]: if "is_true_mixed_batch" in output_fields: row["is_true_mixed_batch"] = "False" output_root.parent.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory( prefix=f".{output_root.name}.", dir=output_root.parent ) as temporary: temporary_root = Path(temporary) / output_root.name shutil.copytree(base_root, temporary_root) merged_attention = temporary_root / ATTENTION_RELATIVE_PATH with merged_attention.open("w", newline="") as output: writer = csv.DictWriter(output, fieldnames=output_fields, lineterminator="\n") writer.writeheader() writer.writerows([*base_rows, *decode_rows, *true_mixed_rows]) payload = { "schema": "frontier-profile-closure-v1", "model": MODEL, "base_profile_root": str(base_root), "base_attention_sha256": sha256(base_attention), "decode_attention_csv": str(decode_csv), "decode_attention_sha256": sha256(decode_csv), "base_prefill_rows": len(base_rows), "decode_source_rows": len(decode_source_rows), "added_decode_rows": len(decode_rows), "true_mixed_attention_csv": ( str(true_mixed_csv) if true_mixed_csv is not None else None ), "true_mixed_attention_sha256": ( sha256(true_mixed_csv) if true_mixed_csv is not None else None ), "added_true_mixed_rows": len(true_mixed_rows), "merged_attention_rows": ( len(base_rows) + len(decode_rows) + len(true_mixed_rows) ), "decode_dimensions": [list(values) for values in sorted(dimensions)], "output_files_sha256": profile_hashes(temporary_root), } write_json(temporary_root / "profile_closure_manifest.json", payload) temporary_root.rename(output_root) print(json.dumps(payload, indent=2, sort_keys=True)) if __name__ == "__main__": main()