Save inference parity implementation and evaluation harness

This commit is contained in:
Georg Bauer
2026-09-10 22:45:59 +02:00
parent b99ce2aa10
commit 02db0968ae
198 changed files with 111205 additions and 586 deletions
+202
View File
@@ -0,0 +1,202 @@
"""Summarize exported xctrace intervals for the attached target only.
Reads XML incrementally; never prints the TOC environment or other processes.
GPU intervals are not shader timings. Nested intervals are not added to level 0.
Per-buffer covered time merges overlaps; group sums are not GPU utilization.
"""
import argparse
import io
import json
import math
import re
import statistics
import sys
import xml.etree.ElementTree as ET
def summarize(source, schemas, pid):
ids, totals, buffers = {}, {}, {}
scalar_tags = {"start-time", "duration", "process", "gpu-channel-name",
"gpu-state", "metal-nesting-level", "metal-command-buffer-id"}
schema = None
def value(element):
if element is None:
return None
if "ref" in element.attrib:
return ids[element.attrib["ref"]]
return element.get("fmt") if element.tag == "process" else element.text
for event, element in ET.iterparse(source, events=("start", "end")):
if event == "start":
if element.tag == "node":
index = int(re.search(r"table\[(\d+)\]$", element.attrib["xpath"])[1])
schema = schemas[index - 1]
continue
if element.tag in scalar_tags and "id" in element.attrib:
ids[element.attrib["id"]] = value(element)
if element.tag == "row":
process = value(element.find("process"))
if process and process.endswith(f"({pid})"):
duration = value(element.find("duration"))
if duration is not None:
ms = int(duration) / 1e6
channel = value(element.find("gpu-channel-name"))
state = value(element.find("gpu-state"))
depth = value(element.find("metal-nesting-level"))
key = "/".join(str(x) for x in (schema, channel, state, depth))
total = totals.setdefault(key, {"rows": 0, "sum_ms": 0, "max_ms": 0})
total["rows"] += 1
total["sum_ms"] += ms
total["max_ms"] = max(total["max_ms"], ms)
if (schema == "metal-gpu-intervals" and channel == "Compute"
and depth == "0" and state == "Active"):
command = value(element.find("metal-command-buffer-id"))
start = int(value(element.find("start-time"))) / 1e6
record = buffers.setdefault(command, {"command_buffer": command,
"start_ms": start, "end_ms": start, "gpu_interval_sum_ms": 0,
"max_interval_ms": 0, "intervals": 0, "_spans": []})
record["start_ms"] = min(record["start_ms"], start)
record["end_ms"] = max(record["end_ms"], start + ms)
record["gpu_interval_sum_ms"] += ms
record["max_interval_ms"] = max(record["max_interval_ms"], ms)
record["intervals"] += 1
record["_spans"].append((start, start + ms))
element.clear()
elif element.tag == "node":
element.clear()
for record in buffers.values():
covered, end = 0, float("-inf")
for start, stop in sorted(record.pop("_spans")):
covered += max(0, stop - max(start, end))
end = max(end, stop)
record["gpu_covered_ms"] = covered
return {"pid": pid, "interval_groups": totals,
"largest_compute_buffers": sorted(buffers.values(),
key=lambda x: x["gpu_interval_sum_ms"], reverse=True)[:12]}
def self_test():
xml = """<trace-query-result><node xpath="//table[1]">
<row><start-time id="1">1000000</start-time><duration id="2">2000000</duration>
<process id="3" fmt="ds4-server (42)"/><gpu-channel-name id="4">Compute</gpu-channel-name>
<gpu-state id="5">Active</gpu-state><metal-nesting-level id="6">0</metal-nesting-level>
<metal-command-buffer-id id="7">100</metal-command-buffer-id></row>
<row><start-time>4000000</start-time><duration ref="2"/><process ref="3"/>
<gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level ref="6"/>
<metal-command-buffer-id ref="7"/></row>
<row><start-time>2000000</start-time><duration ref="2"/><process ref="3"/>
<gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level ref="6"/>
<metal-command-buffer-id ref="7"/></row>
<row><start-time>1000000</start-time><duration ref="2"/><process ref="3"/>
<gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level>1</metal-nesting-level>
<metal-command-buffer-id ref="7"/></row>
<row><duration ref="2"/><process fmt="other (43)"/></row>
</node></trace-query-result>"""
result = summarize(io.StringIO(xml), ["metal-gpu-intervals"], 42)
assert list(result["interval_groups"].values()) == [
{"rows": 3, "sum_ms": 6, "max_ms": 2}, {"rows": 1, "sum_ms": 2, "max_ms": 2}]
assert result["largest_compute_buffers"] == [{"command_buffer": "100", "start_ms": 1,
"end_ms": 6, "gpu_interval_sum_ms": 6, "gpu_covered_ms": 5,
"max_interval_ms": 2, "intervals": 3}]
print("PASS: XML references, PID filtering, depth filtering and overlap union")
def summarize_stages(source):
"""Validate native counter receipts; group encoder spans, never shaders."""
headers, stages, groups = {}, {}, {}
for line in source:
# Supervisor JSON can precede a worker's complete native stderr record.
# Keep strict field/completeness checks; never discard a damaged record.
_, marker, record = line.partition("ds4: stage-counter ")
if not marker:
continue
values = dict(word.split("=", 1) for word in record.split())
cb = int(values["cb"])
if "stages" in values:
assert cb not in headers, f"duplicate CB {cb}"
assert values["resolved"] == "1" and values["dropped"] == "0", values
headers[cb] = values
stages[cb] = {}
continue
assert cb in headers and values["valid"] == "1", values
index = int(values["stage"])
assert index not in stages[cb], (cb, index)
start, end = int(values["start"]), int(values["end"])
duration = float(values["ms"])
scale = float(headers[cb]["ns_per_tick"])
assert end >= start and math.isfinite(duration) and duration >= 0
assert abs(duration - (end - start) * scale / 1e6) <= 0.000002
stages[cb][index] = (start, end)
key = (values["first"], values["last"], int(values["qwen_dispatches"]))
groups.setdefault(key, []).append(duration)
assert headers, "no counter receipts"
covered_total, cb_total = 0.0, 0.0
for cb, header in headers.items():
assert set(stages[cb]) == set(range(int(header["stages"]))), cb
covered, prior_end = 0, 0
for start, end in sorted(stages[cb].values()):
covered += max(0, end - max(start, prior_end))
prior_end = max(prior_end, end)
covered_ms = covered * float(header["ns_per_tick"]) / 1e6
cb_ms = float(header["cb_ms"])
assert math.isfinite(cb_ms) and cb_ms > 0
assert covered_ms <= cb_ms * 1.05 + 0.002, (cb, covered_ms, cb_ms)
covered_total += covered_ms
cb_total += cb_ms
return {"scope": "Qwen four-row command buffers, existing compute encoder boundaries",
"command_buffers": len(headers), "encoder_spans": sum(map(len, stages.values())),
"counter_covered_ms": covered_total, "command_buffer_gpu_ms": cb_total,
"groups_are_encoder_spans_not_individual_kernels": True,
"encoder_groups": sorted([
{"first_qwen_kernel": first, "last_qwen_kernel": last,
"qwen_dispatches": dispatches, "samples": len(values),
"sum_ms": sum(values), "median_ms": statistics.median(values), "max_ms": max(values)}
for (first, last, dispatches), values in groups.items()
], key=lambda row: row["sum_ms"], reverse=True)}
def stage_self_test():
log = ("ds4: stage-counter cb=1 rows=4 stages=2 dropped=0 resolved=1 ns_per_tick=1 cb_ms=0.05\n"
"ds4: stage-counter cb=1 stage=0 valid=1 start=100000 end=130000 ms=0.03 qwen_dispatches=2 first=a last=b\n"
"ds4: stage-counter cb=1 stage=1 valid=1 start=120000 end=140000 ms=0.02 qwen_dispatches=1 first=c last=c\n")
result = summarize_stages(io.StringIO(log))
assert result["encoder_spans"] == 2 and result["counter_covered_ms"] == 0.04
assert summarize_stages(io.StringIO(log.replace("ds4:", '{"event":ds4:'))) == result
for bad in (log.replace("resolved=1", "resolved=0"), log.replace("dropped=0", "dropped=1"),
log.replace("ms=0.03", "ms=1.25"), "\n".join(log.splitlines()[:-1])):
try:
summarize_stages(io.StringIO(bad))
except AssertionError:
continue
raise AssertionError("invalid stage trace accepted")
print("PASS: stage completeness, failure reporting, clock scale and overlap coverage")
if __name__ == "__main__":
if sys.argv[1:] == ["--self-test"]:
self_test()
stage_self_test()
else:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("toc", nargs="?")
parser.add_argument("intervals", nargs="?")
parser.add_argument("--stage-log", help="Native DS4_METAL_STAGE_PROFILE stderr receipt")
args = parser.parse_args()
if args.stage_log:
if args.toc or args.intervals:
parser.error("do not combine --stage-log with XML inputs")
with open(args.stage_log) as source:
print(json.dumps(summarize_stages(source), indent=2))
sys.exit(0)
if not args.toc or not args.intervals:
parser.error("provide both TOC and intervals, or --stage-log")
toc = ET.parse(args.toc)
if len(toc.findall("./run")) != 1:
parser.error("export a trace containing exactly one run")
pid = int(toc.find("./run/info/target/process").attrib["pid"])
schemas = [table.attrib["schema"] for table in toc.findall("./run/data/table")]
result = summarize(args.intervals, schemas, pid)
result["trace_duration_seconds"] = float(toc.findtext("./run/info/summary/duration"))
print(json.dumps(result, indent=2))