h3-blackwell-runtime/tools/summarize_ncu_profile.py
2026-08-26 16:14:12 +07:00

127 lines
5.2 KiB
Python

"""Extract targeted H3 block metrics from an Nsight Compute raw CSV export."""
from __future__ import annotations
import argparse
import csv
import json
from pathlib import Path
ROLES = ("qkv", "sage2", "attention_output", "fc1", "fc2")
EXPECTED_GRIDS = (
"(296, 168, 1)",
"(296, 56, 1)",
"(296, 42, 1)",
"(296, 224, 1)",
"(296, 42, 1)",
)
def number(row: dict[str, str], name: str) -> float | None:
value = row.get(name, "")
if value in {"", "no data", "n/a"}:
return None
return float(value)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--traffic", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
with args.input.open(newline="", encoding="utf-8") as handle:
reader = csv.DictReader(handle)
units = next(reader)
rows = list(reader)
if len(rows) != len(ROLES):
raise RuntimeError(f"expected five targeted launches, found {len(rows)}")
with args.traffic.open(newline="", encoding="utf-8") as handle:
traffic_reader = csv.DictReader(handle)
next(traffic_reader)
traffic_rows = list(traffic_reader)
if len(traffic_rows) != len(ROLES):
raise RuntimeError(f"expected five traffic launches, found {len(traffic_rows)}")
for index, (role, row, traffic_row, expected_grid) in enumerate(
zip(ROLES, rows, traffic_rows, EXPECTED_GRIDS, strict=True)
):
if int(row["ID"]) != index or int(traffic_row["ID"]) != index:
raise RuntimeError(f"{role} launch ID/order contract failed")
if row["Kernel Name"] != traffic_row["Kernel Name"]:
raise RuntimeError(f"{role} kernel differs between metric and traffic passes")
if row["Grid Size"] != expected_grid or traffic_row["Grid Size"] != expected_grid:
raise RuntimeError(f"{role} grid/order contract failed")
if role == "sage2" and "qk_int_sv_f8_attn_kernel" not in row["Kernel Name"]:
raise RuntimeError("Sage2 launch contract failed")
if role != "sage2" and "block_scaled" not in row["Kernel Name"]:
raise RuntimeError(f"{role} NVFP4 GEMM launch contract failed")
stall_prefix = "smsp__average_warps_issue_stalled_"
stall_suffix = "_per_issue_active.ratio"
output = {}
for role, row, traffic_row in zip(ROLES, rows, traffic_rows, strict=True):
stalls = []
for name in row:
if name.startswith(stall_prefix) and name.endswith(stall_suffix):
value = number(row, name)
if value is not None:
stalls.append({
"reason": name[len(stall_prefix):-len(stall_suffix)],
"warps_per_issue_active": value,
})
stalls.sort(key=lambda item: item["warps_per_issue_active"], reverse=True)
output[role] = {
"launch_id": int(row["ID"]),
"kernel_name": row["Kernel Name"],
"grid_size": row["Grid Size"],
"block_size": row["Block Size"],
"duration_ns": number(row, "gpu__time_duration.sum"),
"registers_per_thread": number(row, "launch__registers_per_thread"),
"achieved_occupancy_percent": number(
row, "sm__warps_active.avg.pct_of_peak_sustained_active"
),
"eligible_warps_per_scheduler": number(
row, "smsp__warps_eligible.avg.per_cycle_active"
),
"issue_active_percent": number(
row, "smsp__issue_active.avg.pct_of_peak_sustained_active"
),
"sm_throughput_percent": number(
row, "sm__throughput.avg.pct_of_peak_sustained_elapsed"
),
"tensor_pipe_active_percent": number(
row, "sm__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_elapsed"
),
"l2_requested_bytes": number(traffic_row, "lts__t_bytes.sum"),
"l2_hit_rate_percent": number(traffic_row, "lts__t_sector_hit_rate.pct"),
"l2_throughput_percent": number(
row, "lts__throughput.avg.pct_of_peak_sustained_elapsed"
),
"memory_throughput_percent": number(
traffic_row, "gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed"
),
"local_spilling_requests": number(row, "derived__local_spilling_requests"),
"top_scheduler_stalls": stalls[:5],
}
report = {
"source_csv": str(args.input),
"source_traffic_csv": str(args.traffic),
"launch_order_contract": list(ROLES),
"cache_control": "none (warmed/uncontrolled cache, as reported by NCU)",
"metrics": output,
"units": {
"duration_ns": "ns",
"l2_requested_bytes": "lts__t_bytes.sum",
"throughput_and_hit_rate": "%",
},
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()