94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Assemble immutable flat profiles for the Qwen235 Frontier runner."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
|
|
MODEL = "Qwen3-235B-A22B"
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
for name in ("cuda_common", "cuda_moe_tp4", "cuda_moe_ep8", "kernel_common", "kernel_moe_tp4", "kernel_moe_ep8"):
|
|
parser.add_argument(f"--{name.replace('_', '-')}", type=Path, required=True)
|
|
parser.add_argument("--output-root", type=Path, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def digest(path: Path) -> str:
|
|
value = hashlib.sha256()
|
|
value.update(path.read_bytes())
|
|
return value.hexdigest()
|
|
|
|
|
|
def model_file(root: Path, name: str) -> Path:
|
|
path = root / "compute" / "h20" / MODEL / name
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
return path
|
|
|
|
|
|
def merge_csv(inputs: list[Path], output: Path) -> None:
|
|
fields = None
|
|
rows = []
|
|
for path in inputs:
|
|
with path.open(newline="") as source:
|
|
reader = csv.DictReader(source)
|
|
if fields is None:
|
|
fields = reader.fieldnames
|
|
elif reader.fieldnames != fields:
|
|
raise ValueError(f"CSV schema mismatch: {path}")
|
|
rows.extend(reader)
|
|
if not fields or not rows:
|
|
raise ValueError("cannot merge empty profile CSV")
|
|
with output.open("w", newline="") as target:
|
|
writer = csv.DictWriter(target, fieldnames=fields, lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
output = args.output_root.resolve()
|
|
if output.exists():
|
|
raise FileExistsError(output)
|
|
output.mkdir(parents=True)
|
|
sources = {
|
|
"linear_op.csv": model_file(args.cuda_common, "linear_op.csv"),
|
|
"attention.csv": model_file(args.cuda_common, "attention.csv"),
|
|
"linear_op_kernel_only.csv": model_file(args.kernel_common, "linear_op_kernel_only.csv"),
|
|
"attention_kernel_only.csv": model_file(args.kernel_common, "attention_kernel_only.csv"),
|
|
}
|
|
for name, source in sources.items():
|
|
shutil.copy2(source, output / name)
|
|
merge_csv(
|
|
[model_file(args.cuda_moe_tp4, "moe.csv"), model_file(args.cuda_moe_ep8, "moe.csv")],
|
|
output / "moe.csv",
|
|
)
|
|
merge_csv(
|
|
[model_file(args.kernel_moe_tp4, "moe_kernel_only.csv"), model_file(args.kernel_moe_ep8, "moe_kernel_only.csv")],
|
|
output / "moe_kernel_only.csv",
|
|
)
|
|
outputs = {path.name: digest(path) for path in sorted(output.glob("*.csv"))}
|
|
manifest = {
|
|
"schema": "qwen235-v020-frontier-profile-v1",
|
|
"model": MODEL,
|
|
"measurement_families": ["CUDA_EVENT", "KERNEL_ONLY"],
|
|
"moe_runtime_paths": {"tp4_ep1": "TRITON", "tp8_ep8": "FLASHINFER_CUTLASS"},
|
|
"inputs": {name: str(getattr(args, name).resolve()) for name in ("cuda_common", "cuda_moe_tp4", "cuda_moe_ep8", "kernel_common", "kernel_moe_tp4", "kernel_moe_ep8")},
|
|
"outputs": outputs,
|
|
}
|
|
(output / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps(manifest, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|