120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
"""Generate a markdown summary for representative H3 block profiles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
MAIN_TIMINGS = (
|
|
"norm1",
|
|
"modulate_msa",
|
|
"attn_qkv_proj",
|
|
"attn_qkv_split_view",
|
|
"attn_qk_rms_rope",
|
|
"attn_q_bshd_contiguous",
|
|
"attn_k_bshd_contiguous",
|
|
"attn_v_bshd_contiguous",
|
|
"attention_kernel",
|
|
"attn_output_reshape",
|
|
"attn_out_proj",
|
|
"gate_msa",
|
|
"norm2",
|
|
"modulate_mlp",
|
|
"mlp_fc1",
|
|
"mlp_swiglu",
|
|
"mlp_fc2",
|
|
"gate_mlp",
|
|
"block_total",
|
|
)
|
|
|
|
LINEARS = ("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2")
|
|
LINEAR_STAGES = (
|
|
"activation_scale",
|
|
"activation_quant_pack",
|
|
"gemm",
|
|
"pre_quant_scale",
|
|
"packed_weight_wrapper",
|
|
"bias_cast",
|
|
"scale_to_device",
|
|
"activation_quant_wrap",
|
|
"flatten_contiguous",
|
|
"slice_reshape",
|
|
)
|
|
|
|
|
|
def ms(value: float) -> float:
|
|
return value * 1000.0
|
|
|
|
|
|
def row(timings: dict, name: str) -> tuple[float, float, float]:
|
|
item = timings[name]
|
|
return ms(item["mean_s"]), ms(item["p50_s"]), ms(item["p95_s"])
|
|
|
|
|
|
def add_table(lines: list[str], timings: dict, names: tuple[str, ...]) -> None:
|
|
lines.extend(["| Stage | Mean ms | P50 ms | P95 ms |", "| --- | ---: | ---: | ---: |"])
|
|
for name in names:
|
|
if name not in timings:
|
|
continue
|
|
mean, p50, p95 = row(timings, name)
|
|
lines.append(f"| `{name}` | {mean:.3f} | {p50:.3f} | {p95:.3f} |")
|
|
|
|
|
|
def linear_table(lines: list[str], timings: dict) -> None:
|
|
lines.extend(["| Linear | Scale P50 ms | Pack P50 ms | GEMM P50 ms | Total P50 ms |", "| --- | ---: | ---: | ---: | ---: |"])
|
|
for linear in LINEARS:
|
|
prefix = f"linear.{linear}"
|
|
scale = ms(timings[f"{prefix}.activation_scale"]["p50_s"])
|
|
pack = ms(timings[f"{prefix}.activation_quant_pack"]["p50_s"])
|
|
gemm = ms(timings[f"{prefix}.gemm"]["p50_s"])
|
|
total = ms(timings[linear]["p50_s"])
|
|
lines.append(f"| `{linear}` | {scale:.3f} | {pack:.3f} | {gemm:.3f} | {total:.3f} |")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
|
|
data = json.loads(args.input.read_text(encoding="utf-8"))
|
|
timings = data["timings"]
|
|
block_p50 = ms(timings["block_total"]["p50_s"])
|
|
attention_layout_p50 = sum(ms(timings[name]["p50_s"]) for name in ("attn_q_bshd_contiguous", "attn_k_bshd_contiguous", "attn_v_bshd_contiguous", "attention_kernel", "attn_output_reshape"))
|
|
mlp_linear_p50 = sum(ms(timings[name]["p50_s"]) for name in ("mlp_fc1", "mlp_fc2"))
|
|
norm_mod_gate_p50 = sum(ms(timings[name]["p50_s"]) for name in ("norm1", "modulate_msa", "gate_msa", "norm2", "modulate_mlp", "gate_mlp"))
|
|
qk_rope_p50 = ms(timings["attn_qk_rms_rope"]["p50_s"])
|
|
|
|
lines = [
|
|
"# H3 Block Profile Baseline",
|
|
"",
|
|
f"Input: block `{data['block_index']}`, hidden shape `{data['hidden_shape']}`, segments `{data['segments']}`.",
|
|
f"Config: attention `{data['attention']}`, `{data['warmup']}` warmups, `{data['iterations']}` iterations.",
|
|
"",
|
|
"## Summary",
|
|
"",
|
|
f"- Block total p50: `{block_p50:.3f} ms`; p95: `{ms(timings['block_total']['p95_s']):.3f} ms`.",
|
|
f"- Sol attention plus BSHD layout p50: `{attention_layout_p50:.3f} ms`.",
|
|
f"- MLP linears p50: `{mlp_linear_p50:.3f} ms`; SwiGLU p50: `{ms(timings['mlp_swiglu']['p50_s']):.3f} ms`.",
|
|
f"- Norm/modulate/gate p50: `{norm_mod_gate_p50:.3f} ms`.",
|
|
f"- Q/K RMS+RoPE p50: `{qk_rope_p50:.3f} ms`.",
|
|
"- The next likely targets are Q/K RMS+RoPE, Q/K/V contiguous extraction, and the modulation/gating helpers; Sol output reshape remains negligible.",
|
|
"",
|
|
"## Main Stages",
|
|
"",
|
|
]
|
|
add_table(lines, timings, MAIN_TIMINGS)
|
|
lines.extend(["", "## NVFP4 Linear Breakdown", ""])
|
|
linear_table(lines, timings)
|
|
for linear in LINEARS:
|
|
lines.extend(["", f"## `{linear}` Stages", ""])
|
|
add_table(lines, timings, tuple(f"linear.{linear}.{stage}" for stage in LINEAR_STAGES))
|
|
lines.extend(["", "## Source", "", f"- JSON: `{args.input}`"])
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|