File size: 6,432 Bytes
3706f1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | #!/usr/bin/env python3
"""Resident MTP forward probe; temporarily lock clocks and restore on exit."""
import argparse
import datetime
import glob
import json
import os
from pathlib import Path
import signal
import subprocess
import time
def read(path):
try:
return Path(path).read_text().strip()
except OSError:
return None
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("binary")
parser.add_argument("model")
parser.add_argument("--output", required=True)
parser.add_argument("--runtime", required=True)
parser.add_argument("--seqlen", type=int, default=128)
parser.add_argument("--tokens", type=int, default=64)
parser.add_argument("--rounds", type=int, default=3)
parser.add_argument("--context", type=int, default=32768)
parser.add_argument("--cases")
args = parser.parse_args()
if os.geteuid() != 0:
parser.error("run with sudo for temporary clock locking and NPU debugfs sampling")
help_env = os.environ.copy()
help_env["LD_LIBRARY_PATH"] = args.runtime + ":" + help_env.get("LD_LIBRARY_PATH", "")
probe = subprocess.run([args.binary], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
text=True, env=help_env, timeout=15)
if "Usage:" not in probe.stdout:
parser.error("cannot query engine CLI: " + probe.stdout.strip())
if "mtp-benchmark" not in probe.stdout:
parser.error("MTP is not compiled; rebuild with -DLING3_EXPERIMENTAL_MTP=ON")
output = Path(args.output)
output.mkdir(parents=True, exist_ok=False)
original = {}
applied = {}
child = None
def write_clock(path, value):
original[path] = read(path)
if original[path] is None:
raise RuntimeError("cannot snapshot " + path)
Path(path).write_text(str(value))
applied[path] = read(path)
def stop(signum, frame):
raise KeyboardInterrupt
signal.signal(signal.SIGTERM, stop)
signal.signal(signal.SIGINT, stop)
try:
for policy in sorted(glob.glob("/sys/devices/system/cpu/cpufreq/policy*")):
write_clock(policy + "/scaling_governor", "performance")
write_clock(policy + "/scaling_min_freq", read(policy + "/scaling_max_freq"))
for device in ("fdab0000.npu", "dmc"):
base = "/sys/class/devfreq/" + device
frequencies = read(base + "/available_frequencies")
maximum = max(map(int, frequencies.split())) if frequencies else int(read(base + "/max_freq"))
write_clock(base + "/governor", "performance")
write_clock(base + "/min_freq", maximum)
env = os.environ.copy()
for key in ("LING3_GDN_PREFILL_DIR", "LING3_PREFILL_W4A4", "LING3_EXPERT_ALL_CORES",
"LING3_MLA_DECODE_NPU", "LING3_DISABLE_BATCH_INPUT_REUSE",
"LING3_DISABLE_PARALLEL_GATHER"):
env.pop(key, None)
for key in ("LING3_PREWARM_EXPERTS", "LING3_EXPERT_BALANCED", "LING3_EXPERT_ZERO_COPY",
"LING3_GDN_CPU_PREFILL", "LING3_GDN_CPU_FP32_STATE", "LING3_GDN_CPU_DECODE",
"LING3_GDN_FULL_FP32", "LING3_MLA_SIMD", "LING3_VECTOR_MATH"):
env[key] = "1"
env["LING3_MLA_BACKEND"] = "cpu"
env["LD_LIBRARY_PATH"] = args.runtime + ":" + env.get("LD_LIBRARY_PATH", "")
command = [args.binary, "mtp-benchmark", args.model, str(args.seqlen), str(args.tokens), str(args.rounds), str(args.context)]
if args.cases:
command.append(args.cases)
metadata = {"started_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(),
"command": command, "original_clocks": original, "meminfo_before": read("/proc/meminfo"), "applied_clocks": applied,
"profile": {k: v for k, v in env.items() if k.startswith("LING3_")}}
(output / "metadata.json").write_text(json.dumps(metadata, indent=2))
with (output / "benchmark.log").open("w") as log, (output / "telemetry.jsonl").open("w") as telemetry:
child = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=env)
print("pid=" + str(child.pid) + " output=" + str(output), flush=True)
start = time.monotonic()
while child.poll() is None:
row = {"elapsed_s": time.monotonic() - start,
"process_stat": read(f"/proc/{child.pid}/stat"),
"process_status": read(f"/proc/{child.pid}/status"),
"npu_load": read("/sys/kernel/debug/rknpu/load"),
"meminfo": read("/proc/meminfo")}
row["frequencies"] = {p: read(p) for p in glob.glob("/sys/devices/system/cpu/cpufreq/policy*/scaling_cur_freq")
+ ["/sys/class/devfreq/fdab0000.npu/cur_freq", "/sys/class/devfreq/dmc/cur_freq"]}
row["temperatures"] = {p: read(p) for p in glob.glob("/sys/class/thermal/thermal_zone*/temp")}
telemetry.write(json.dumps(row) + "\n")
telemetry.flush()
time.sleep(0.5)
metadata["exit_code"] = child.returncode
metadata["elapsed_s"] = time.monotonic() - start
print("exit_code=" + str(child.returncode), flush=True)
finally:
if child is not None and child.poll() is None:
child.terminate()
try:
child.wait(timeout=15)
except subprocess.TimeoutExpired:
child.kill()
child.wait()
restored = {}
for path, value in reversed(list(original.items())):
try:
Path(path).write_text(value)
restored[path] = read(path)
except OSError as error:
restored[path] = str(error)
final = locals().get("metadata", {"original_clocks": original, "applied_clocks": applied})
final["restored_clocks"] = restored
(output / "metadata.json").write_text(json.dumps(final, indent=2))
if "SUDO_UID" in os.environ:
for path in [output] + list(output.iterdir()):
os.chown(path, int(os.environ["SUDO_UID"]), int(os.environ["SUDO_GID"]))
print("clocks_restored", flush=True)
return child.returncode if child else 1
if __name__ == "__main__":
raise SystemExit(main())
|