101 lines
3.6 KiB
Python
Executable File
101 lines
3.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Append measured true-mixed attention rows to an immutable Qwen235 profile."""
|
|
|
|
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()
|
|
parser.add_argument("--base-frozen", type=Path, required=True)
|
|
parser.add_argument("--cuda-tp4", type=Path, required=True)
|
|
parser.add_argument("--cuda-tp8", type=Path, required=True)
|
|
parser.add_argument("--kernel-tp4", type=Path, required=True)
|
|
parser.add_argument("--kernel-tp8", 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 read_rows(path: Path) -> tuple[list[str], list[dict[str, str]]]:
|
|
with path.open(newline="") as source:
|
|
reader = csv.DictReader(source)
|
|
return list(reader.fieldnames or []), list(reader)
|
|
|
|
|
|
def merge_csv(inputs: list[Path], output: Path, measurement_type: str) -> int:
|
|
fields: list[str] = []
|
|
rows: list[dict[str, str]] = []
|
|
true_mixed_count = 0
|
|
for index, path in enumerate(inputs):
|
|
input_fields, input_rows = read_rows(path)
|
|
for field in input_fields:
|
|
if field not in fields:
|
|
fields.append(field)
|
|
if index:
|
|
for row in input_rows:
|
|
if row.get("is_true_mixed_batch", "").lower() != "true":
|
|
raise ValueError(f"non-true-mixed row in {path}")
|
|
if row.get("measurement_type") != measurement_type:
|
|
raise ValueError(f"measurement type mismatch in {path}")
|
|
true_mixed_count += len(input_rows)
|
|
rows.extend(input_rows)
|
|
if not fields or not rows or not true_mixed_count:
|
|
raise ValueError("cannot assemble an empty true-mixed augmentation")
|
|
with output.open("w", newline="") as target:
|
|
writer = csv.DictWriter(target, fieldnames=fields, lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
return true_mixed_count
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
output = args.output_root.resolve()
|
|
if output.exists():
|
|
raise FileExistsError(output)
|
|
output.mkdir(parents=True)
|
|
for name in ("linear_op.csv", "linear_op_kernel_only.csv", "moe.csv", "moe_kernel_only.csv"):
|
|
shutil.copy2(args.base_frozen / name, output / name)
|
|
cuda_count = merge_csv(
|
|
[args.base_frozen / "attention.csv", args.cuda_tp4, args.cuda_tp8],
|
|
output / "attention.csv",
|
|
"CUDA_EVENT",
|
|
)
|
|
kernel_count = merge_csv(
|
|
[args.base_frozen / "attention_kernel_only.csv", args.kernel_tp4, args.kernel_tp8],
|
|
output / "attention_kernel_only.csv",
|
|
"KERNEL_ONLY",
|
|
)
|
|
base_manifest = json.loads((args.base_frozen / "manifest.json").read_text())
|
|
manifest = {
|
|
**base_manifest,
|
|
"schema": "qwen235-v020-frontier-profile-v2-true-mixed",
|
|
"true_mixed_inputs": {
|
|
name: str(getattr(args, name).resolve())
|
|
for name in ("cuda_tp4", "cuda_tp8", "kernel_tp4", "kernel_tp8")
|
|
},
|
|
"true_mixed_rows": {"CUDA_EVENT": cuda_count, "KERNEL_ONLY": kernel_count},
|
|
"outputs": {path.name: digest(path) for path in sorted(output.glob("*.csv"))},
|
|
}
|
|
(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()
|