| """Summarize a torch-profiler chrome trace: GPU kernel time by category, busy vs wall, per-graph-replay breakdown.""" |
| import gzip, json, sys, re, collections |
| d = json.load(gzip.open(sys.argv[1])) |
| ev = [e for e in d["traceEvents"] if e.get("ph") == "X"] |
| k = [e for e in ev if e.get("cat") == "kernel"] |
| cpu = [e for e in ev if e.get("cat") in ("cpu_op", "user_annotation", "python_function")] |
| def cat(n): |
| n2 = n.lower() |
| if "block_scaled" in n2 or "fp4" in n2 or "nvfp4" in n2 or "cudnn" in n2 and "gemm" in n2 or "sm120" in n2 or "quantize_with_block" in n2 or "cvt_fp16_to_fp4" in n2 or "e2m1" in n2: return "FP4 GEMM/quant" |
| if "gemv" in n2 or "gemm" in n2 or "sgemm" in n2 or "cublas" in n2 or "dot_kernel" in n2 or "wmma" in n2 or "splitk" in n2 or "nvjet" in n2: return "other GEMM/GEMV" |
| if "delta_rule" in n2 or "causal_conv" in n2 or "gdn" in n2 or "fla" in n2 or "chunk_" in n2: return "GDN/recurrent" |
| if "attn" in n2 or "attention" in n2 or "flash" in n2 and "norm" not in n2 or "decode" in n2 and "rope" not in n2: return "attention" |
| if any(s in n2 for s in ("exp_kernel", "log_kernel", "masked_fill", "absfunctor", "compareeq", "maxnan", "sum_functor", "add<float>", "sigmoid", "functor_add", "logsumexp")): return "elementwise/reduce (fp32, mHC-like)" |
| if "norm" in n2: return "norm" |
| return "other" |
| tot = collections.Counter(); cnt = collections.Counter(); byname = collections.Counter(); bync = collections.Counter() |
| for e in k: |
| c = cat(e["name"]); tot[c] += e["dur"]; cnt[c] += 1; byname[e["name"][:110]] += e["dur"]; bync[e["name"][:110]] += 1 |
| t0 = min(e["ts"] for e in k); t1 = max(e["ts"] + e["dur"] for e in k) |
| busy = sum(e["dur"] for e in k) |
| print(f"kernels={len(k)} window={((t1-t0)/1e3):.1f}ms gpu_kernel_sum={busy/1e3:.1f}ms ({100*busy/(t1-t0):.0f}% of window)") |
| for c, v in tot.most_common(): print(f" {c:40s} {v/1e3:9.1f} ms {100*v/busy:5.1f}% n={cnt[c]}") |
| print("top kernels:") |
| for n, v in byname.most_common(22): print(f" {v/1e3:8.1f} ms {bync[n]:7d}x {n}") |
| |
| ann = collections.Counter(); annd = collections.Counter() |
| for e in cpu: |
| n = e["name"] |
| if any(s in n for s in ("replay", "forward_", "run_batch", "verify", "draft", "cudaGraphLaunch")): |
| ann[n[:80]] += 1; annd[n[:80]] += e["dur"] |
| print("cpu annotations (count, total ms):") |
| for n, c in ann.most_common(30): print(f" {c:5d} {annd[n]/1e3:9.1f} {n}") |
| rt = [e for e in ev if e.get("cat") == "cuda_runtime"] |
| rc = collections.Counter(); rd = collections.Counter() |
| for e in rt: rc[e["name"]] += 1; rd[e["name"]] += e["dur"] |
| print("cuda runtime:", [(n, rc[n], round(rd[n]/1e3, 1)) for n, _ in rd.most_common(8)]) |
|
|
