187 lines
7.0 KiB
Python
187 lines
7.0 KiB
Python
#!/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()
|