From da6c75813071a53f3763272002b9090f93d11037 Mon Sep 17 00:00:00 2001 From: Kao Date: Sat, 11 Apr 2026 06:48:33 +0800 Subject: [PATCH 1/4] [Feature] Add idle (gap) and wall time (span) rows to layer detail reports --- trace_module_analyzer.py | 47 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/trace_module_analyzer.py b/trace_module_analyzer.py index 3ce860e..6a95230 100755 --- a/trace_module_analyzer.py +++ b/trace_module_analyzer.py @@ -1306,6 +1306,21 @@ def print_layer_detail(self, stats_list: List[ModuleStats], mode: str, else: wall_time = 0 + # Compute active_time via sorted interval union; idle = wall - active + active_time = 0.0 + if all_details: + srt = sorted(all_details, key=lambda d: d.ts) + ms, me = srt[0].ts, srt[0].ts + srt[0].duration + for d in srt[1:]: + s, e = d.ts, d.ts + d.duration + if s <= me: + me = max(me, e) + else: + active_time += me - ms + ms, me = s, e + active_time += me - ms + idle_time = max(0.0, wall_time - active_time) + print(f"\n{'='*100}") print(f" Layer Detail: {selected.name} [{len(all_details)} items]") print(f" Kernel sum: {sum_dur:,.0f} us | Wall time: {wall_time:,.0f} us " @@ -1334,6 +1349,9 @@ def print_layer_detail(self, stats_list: List[ModuleStats], mode: str, print(f" {i:4d} {d.duration:13,.1f} {pct:5.1f}% {d.category:>15s} {leaf:30s} {kname:50s} {dims}") else: print(f" {i:4d} {d.duration:13,.1f} {pct:5.1f}% {d.category:>15s} {leaf:30s} {kname}") + idle_pct = idle_time / wall_time * 100 if wall_time > 0 else 0 + print(f" {'':4s} {idle_time:13,.1f} {idle_pct:5.1f}% {'':>15s} {'':30s} idle (gap)") + print(f" {'':4s} {wall_time:13,.1f} {'':6s} {'':>15s} {'':30s} wall time (span)") def _pick_median_instance(self, matches: List[ModuleStats], mode: str) -> Tuple[ModuleStats, str]: @@ -1755,6 +1773,22 @@ def export_excel(self, stats_list: List[ModuleStats], mode: str, cell.font = header_font cell.fill = header_fill cur_row += 1 + # Compute active_time via sorted interval union (all_details already + # sorted by ts above). idle_time = wall_time - active_time. + active_time = 0.0 + if all_details: + merge_start = all_details[0].ts + merge_end = all_details[0].ts + all_details[0].duration + for d in all_details[1:]: + s, e = d.ts, d.ts + d.duration + if s <= merge_end: + merge_end = max(merge_end, e) + else: + active_time += merge_end - merge_start + merge_start, merge_end = s, e + active_time += merge_end - merge_start + idle_time = max(0.0, wall_time - active_time) + detail_truncated = len(all_details) > MAX_ROWS_PER_TAB for i, d in enumerate(all_details[:MAX_ROWS_PER_TAB], 1): pct = d.duration / wall_time * 100 if wall_time > 0 else 0 @@ -1777,6 +1811,19 @@ def export_excel(self, stats_list: List[ModuleStats], mode: str, if detail_truncated: ws_det.cell(row=cur_row, column=1, value=f"... truncated at {MAX_ROWS_PER_TAB} rows") + cur_row += 1 + # Idle-gap synthetic row + idle_pct = idle_time / wall_time * 100 if wall_time > 0 else 0 + idle_font = Font(italic=True, color="888888") + idle_cell_name = ws_det.cell(row=cur_row, column=3, value="idle (gap)") + idle_cell_name.font = idle_font + ws_det.cell(row=cur_row, column=4, value=round(idle_time, 1)).font = idle_font + ws_det.cell(row=cur_row, column=5, value=round(idle_pct, 1)).font = idle_font + cur_row += 1 + # Wall-time-span footer row + span_font = Font(bold=True) + ws_det.cell(row=cur_row, column=3, value="wall time (span)").font = span_font + ws_det.cell(row=cur_row, column=4, value=round(wall_time, 1)).font = span_font ws_det.column_dimensions["A"].width = 35 ws_det.column_dimensions["B"].width = 60 ws_det.column_dimensions["C"].width = 80 From c4f8c44a12828995f2b32d57bc44191cbd6dec6b Mon Sep 17 00:00:00 2001 From: Alan Kao Date: Fri, 5 Jun 2026 14:59:39 +0800 Subject: [PATCH 2/4] Add utility for agent skill --- compare_merged.py | 289 +++++++++++++++ merge_graph_nograph.py | 737 +++++++++++++++++++++++++++++++++++++++ report_html.py | 406 +++++++++++++++++++++ trace_module_analyzer.py | 87 +++++ 4 files changed, 1519 insertions(+) create mode 100644 compare_merged.py create mode 100644 merge_graph_nograph.py create mode 100644 report_html.py diff --git a/compare_merged.py b/compare_merged.py new file mode 100644 index 0000000..bbdfe3d --- /dev/null +++ b/compare_merged.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +"""Side-by-side comparison of two merged-layer xlsx files. + +Reads two outputs of merge_graph_nograph.py and produces a single xlsx where +each Layer block lists the kernels from both runs in adjacent column groups, +aligned by leaf-module name (best-effort). + +Example: + python compare_merged.py b200.xlsx mi355.xlsx -o compare.xlsx \ + --labels B200 MI355 +""" + +import argparse +import re +from collections import defaultdict +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import openpyxl +from openpyxl.styles import Alignment, Font, PatternFill +from openpyxl.utils import get_column_letter + + +_INSTANCE_RE = re.compile(r"_\d+$") + + +def _strip(name: str) -> str: + return _INSTANCE_RE.sub("", name) if name else name + + +@dataclass +class Row: + layer: str # parent group, e.g. DeepseekV2AttentionMLA + module: str # leaf, e.g. RadixAttention (may be "(unmapped …)") + shape: str + kernel_name: str + duration_us: float + pct: str + properties: str + + +def load_merged(path: str) -> Tuple[List[Row], float]: + """Return (rows, total_wall_us) preserving file order.""" + wb = openpyxl.load_workbook(path, read_only=True) + ws = wb["Merged Layer"] + rows: List[Row] = [] + cur_layer = "" + total = 0.0 + for r in ws.iter_rows(values_only=True): + if r[0] == "TOTAL": + total = r[5] or 0 + continue + if r[0] == "Layer": # header row + continue + if r[0]: + cur_layer = r[0] + # Subtotal rows have no kernel name in col 4 + if r[3] is None: + continue + # Defensive: skip any other non-numeric duration row + try: + float(r[5] or 0) + except (TypeError, ValueError): + continue + rows.append(Row( + layer=cur_layer, + module=str(r[1] or ""), + shape=str(r[2] or ""), + kernel_name=str(r[3] or ""), + duration_us=float(r[5] or 0), + pct=str(r[6] or ""), + properties=str(r[7] or ""), + )) + wb.close() + return rows, total + + +def group_by_layer(rows: List[Row]) -> "dict[str, List[Row]]": + out: Dict[str, List[Row]] = defaultdict(list) + for r in rows: + out[r.layer].append(r) + return out + + +# --------------------------------------------------------------------------- +# Per-Layer side-by-side alignment. +# --------------------------------------------------------------------------- + +def align_layer_block(a_rows: List[Row], b_rows: List[Row] + ) -> List[Tuple[Optional[Row], Optional[Row]]]: + """Pair rows from two runs inside one Layer block. + + Strategy: walk both lists in order. At each step pair if the leaf-module + name matches (stripped); otherwise advance whichever side has a module + that the other side has somewhere downstream — falling back to greedy + insertion when leaf modules diverge entirely. This is intentionally + simple, since within a Layer block the leaf modules usually appear in + the same conceptual order on both runs. + """ + pairs: List[Tuple[Optional[Row], Optional[Row]]] = [] + i = j = 0 + # Precompute remaining leaf-module sets for lookahead + while i < len(a_rows) or j < len(b_rows): + if i >= len(a_rows): + pairs.append((None, b_rows[j])); j += 1; continue + if j >= len(b_rows): + pairs.append((a_rows[i], None)); i += 1; continue + a, b = a_rows[i], b_rows[j] + if a.module == b.module and a.module: + pairs.append((a, b)); i += 1; j += 1 + continue + # Look ahead: does B have a's module later, or does A have b's later? + a_later_in_b = any(b_rows[k].module == a.module + for k in range(j + 1, len(b_rows))) + b_later_in_a = any(a_rows[k].module == b.module + for k in range(i + 1, len(a_rows))) + if a_later_in_b and not b_later_in_a: + # B's current row has no match in A → emit B alone, advance j + pairs.append((None, b)); j += 1 + elif b_later_in_a and not a_later_in_b: + pairs.append((a, None)); i += 1 + else: + # Both or neither have downstream matches — emit side-by-side + # (treating them as the "same slot" even though leaf differs). + pairs.append((a, b)); i += 1; j += 1 + return pairs + + +def write_compare(a_path: str, b_path: str, output: str, + label_a: str, label_b: str): + a_rows, a_total = load_merged(a_path) + b_rows, b_total = load_merged(b_path) + a_by_layer = group_by_layer(a_rows) + b_by_layer = group_by_layer(b_rows) + + # Canonical Layer order — fall back to file-order for unknown groups. + CANONICAL = [ + "prepare_attn", + "DeepseekV2AttentionMLA", + "prepare_mlp", + "DeepseekV2MoE", + "DeepseekV2MLP", + "post_layer", + ] + seen = set(a_by_layer) | set(b_by_layer) + merged_layers: List[str] = [L for L in CANONICAL if L in seen] + # Append any unknown groups in file-order + for L in [r.layer for r in a_rows] + [r.layer for r in b_rows]: + if L not in merged_layers and L in seen: + merged_layers.append(L) + + wb = openpyxl.Workbook() + ws = wb.active + ws.title = "Compare" + + bold = Font(bold=True) + title_font = Font(bold=True, size=12) + hdr_fill = PatternFill(start_color="CCE5FF", end_color="CCE5FF", + fill_type="solid") + grp_fill = PatternFill(start_color="EFF6FF", end_color="EFF6FF", + fill_type="solid") + sub_fill = PatternFill(start_color="F0F0F0", end_color="F0F0F0", + fill_type="solid") + delta_pos_fill = PatternFill(start_color="FFE5E5", end_color="FFE5E5", + fill_type="solid") # slower = pink + delta_neg_fill = PatternFill(start_color="E5FFE5", end_color="E5FFE5", + fill_type="solid") # faster = green + + # ── Section 1: Summary table at top ── + ws.cell(row=1, column=1, + value=f"Side-by-side: {label_a} vs {label_b}").font = title_font + + ws.cell(row=3, column=1, value="Layer group").font = bold + ws.cell(row=3, column=2, value=f"{label_a} (us)").font = bold + ws.cell(row=3, column=3, value=f"{label_b} (us)").font = bold + ws.cell(row=3, column=4, value="Δ (us)").font = bold + ws.cell(row=3, column=5, value="Δ %").font = bold + for c in range(1, 6): + ws.cell(row=3, column=c).fill = hdr_fill + + summary_row = 4 + for L in merged_layers: + at = sum(r.duration_us for r in a_by_layer.get(L, [])) + bt = sum(r.duration_us for r in b_by_layer.get(L, [])) + d = bt - at + pct = (d / at * 100) if at else float("nan") + ws.cell(row=summary_row, column=1, value=L) + ws.cell(row=summary_row, column=2, + value=round(at, 1) if at else None) + ws.cell(row=summary_row, column=3, + value=round(bt, 1) if bt else None) + ws.cell(row=summary_row, column=4, value=round(d, 1)) + ws.cell(row=summary_row, column=5, + value=f"{pct:+.1f}%" if at else "") + if d > 0.5: + for c in range(1, 6): + ws.cell(row=summary_row, column=c).fill = delta_pos_fill + elif d < -0.5: + for c in range(1, 6): + ws.cell(row=summary_row, column=c).fill = delta_neg_fill + summary_row += 1 + + # TOTAL row + d = b_total - a_total + pct = (d / a_total * 100) if a_total else 0 + ws.cell(row=summary_row, column=1, value="TOTAL").font = bold + ws.cell(row=summary_row, column=2, value=round(a_total, 1)).font = bold + ws.cell(row=summary_row, column=3, value=round(b_total, 1)).font = bold + ws.cell(row=summary_row, column=4, value=round(d, 1)).font = bold + ws.cell(row=summary_row, column=5, + value=f"{pct:+.1f}%").font = bold + for c in range(1, 6): + ws.cell(row=summary_row, column=c).fill = sub_fill + + # ── Section 2: per-Layer side-by-side ── + row_idx = summary_row + 3 + # Header for kernel block + headers = [ + "Layer", + f"{label_a}: Module", f"{label_a}: Kernel_name", + f"{label_a}: t (us)", f"{label_a}: %", + "", # spacer + f"{label_b}: Module", f"{label_b}: Kernel_name", + f"{label_b}: t (us)", f"{label_b}: %", + ] + for c, h in enumerate(headers, 1): + cell = ws.cell(row=row_idx, column=c, value=h) + cell.font = bold + cell.fill = hdr_fill + cell.alignment = Alignment(horizontal="center") + row_idx += 1 + + for L in merged_layers: + a_grp = a_by_layer.get(L, []) + b_grp = b_by_layer.get(L, []) + pairs = align_layer_block(a_grp, b_grp) + + a_total_L = sum(r.duration_us for r in a_grp) + b_total_L = sum(r.duration_us for r in b_grp) + first = True + + for (a, b) in pairs: + ws.cell(row=row_idx, column=1, value=(L if first else "")) + if first: + ws.cell(row=row_idx, column=1).fill = grp_fill + ws.cell(row=row_idx, column=1).font = bold + first = False + + if a is not None: + ws.cell(row=row_idx, column=2, value=a.module) + ws.cell(row=row_idx, column=3, value=a.kernel_name) + ws.cell(row=row_idx, column=4, value=round(a.duration_us, 2)) + ws.cell(row=row_idx, column=5, value=a.pct) + if b is not None: + ws.cell(row=row_idx, column=7, value=b.module) + ws.cell(row=row_idx, column=8, value=b.kernel_name) + ws.cell(row=row_idx, column=9, value=round(b.duration_us, 2)) + ws.cell(row=row_idx, column=10, value=b.pct) + + row_idx += 1 + + # Subtotal (no per-row delta column) + ws.cell(row=row_idx, column=4, value=round(a_total_L, 2)).font = bold + ws.cell(row=row_idx, column=9, value=round(b_total_L, 2)).font = bold + for c in range(1, 11): + ws.cell(row=row_idx, column=c).fill = sub_fill + row_idx += 1 + + widths = [28, 26, 55, 10, 8, 2, 26, 55, 10, 8] + for i, w in enumerate(widths, 1): + ws.column_dimensions[get_column_letter(i)].width = w + + wb.save(output) + print(f"Wrote {output}") + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("a", help="First merged xlsx (baseline)") + p.add_argument("b", help="Second merged xlsx (target)") + p.add_argument("-o", "--output", required=True) + p.add_argument("--labels", nargs=2, default=("A", "B")) + args = p.parse_args() + write_compare(args.a, args.b, args.output, args.labels[0], args.labels[1]) + + +if __name__ == "__main__": + main() diff --git a/merge_graph_nograph.py b/merge_graph_nograph.py new file mode 100644 index 0000000..e0500f1 --- /dev/null +++ b/merge_graph_nograph.py @@ -0,0 +1,737 @@ +#!/usr/bin/env python3 +"""Merge a CUDA-graph trace with its no-graph counterpart for one decode layer. + +When SGLang runs with --disable-cuda-graph, each kernel is launched inside its +real nn.Module context, so the trace carries the full module hierarchy +(DeepseekV2AttentionMLA → RadixAttention, etc.). When CUDA graph is enabled, +the same kernels are replayed inside a captured graph and the module info is +lost — the analyzer can only label them as the synthetic "Layer_N". + +This tool produces a per-layer view that combines: + • module hierarchy + Input Dims ← from the no-graph trace + • per-kernel duration ← from the graph (perf) trace +aligned via longest-common-subsequence on kernel name within one decoder layer. + +The output xlsx has the layout illustrated in the screenshot: + Layer | Module | shape | Kernel_name | call-times | time (us) | percentage (%) | properties + +Where "Layer" is the parent module group (e.g. DeepseekV2AttentionMLA, +DeepseekV2MoE), "Module" is the leaf nn.Module (RMSNorm, RadixAttention, ...), +and rows are grouped by Layer with a subtotal time row at the end of each group. +""" + +import argparse +import os +import sys +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +if THIS_DIR not in sys.path: + sys.path.insert(0, THIS_DIR) + +import gzip +import json + +from trace_module_analyzer import ( # noqa: E402 + CpuOpShapeIndex, + CudaGraphCorrelator, + KernelCorrelator, + ModuleAggregator, + ModuleStats, + ModuleTreeBuilder, + PhaseDetector, + PythonSourceIndex, + _categorize_kernel, +) + +try: + from fix_rocm_trace_flow import fix_trace as _rocm_fix_trace + _HAS_ROCM_FIX = True +except ImportError: + _HAS_ROCM_FIX = False + + +@dataclass +class KernelRow: + """One kernel as it will appear in the merged report.""" + layer_group: str # parent module (e.g. DeepseekV2AttentionMLA) + module: str # leaf module (e.g. RadixAttention) + shape: str # Input Dims (from nograph cpu_op) + kernel_name: str + duration_us: float + category: str # 'attention', 'gemm', 'moe', 'communication', ... + source: str = "" # source-path (nograph) + + +# --------------------------------------------------------------------------- +# Step 1 — reuse the analyzer to produce a fully-correlated module tree. +# --------------------------------------------------------------------------- + +def _load_trace_json(path: str) -> dict: + opener = gzip.open if path.endswith(".gz") else open + with opener(path, "rt", encoding="utf-8") as f: + return json.load(f) + + +def analyze_trace(trace_path: str): + """Load trace, build module tree, correlate kernels, aggregate. Returns stats.""" + print(f" Loading {trace_path}") + data = _load_trace_json(trace_path) + if _HAS_ROCM_FIX: + data, _, _ = _rocm_fix_trace(data) + events = data.get("traceEvents", []) + + runtime_events, driver_events, kernel_events = [], [], [] + cpu_ops, module_events, py_events = [], [], [] + gpu_memcpy, gpu_memset = [], [] + phase_markers = [] + MODULE_PREFIX = "nn.Module: " + for e in events: + cat = e.get("cat", "") + if cat == "kernel": + kernel_events.append(e) + elif cat == "cuda_runtime": + runtime_events.append(e) + elif cat == "cuda_driver": + driver_events.append(e) + elif cat == "cpu_op": + cpu_ops.append(e) + elif cat == "gpu_memcpy": + gpu_memcpy.append(e) + elif cat == "gpu_memset": + gpu_memset.append(e) + elif cat == "python_function": + name = e.get("name", "") + if name.startswith(MODULE_PREFIX) and e.get("dur") is not None: + module_events.append(e) + elif e.get("dur") is not None: + py_events.append(e) + if "model_runner" in name: + if ": forward_extend" in name: + phase_markers.append((e["ts"], e["ts"] + e["dur"], + "prefill", e["tid"], e.get("pid"))) + elif ": forward_decode" in name: + phase_markers.append((e["ts"], e["ts"] + e["dur"], + "decode", e["tid"], e.get("pid"))) + + print(f" events: kernel={len(kernel_events):,} cpu_op={len(cpu_ops):,} " + f"nnModule={len(module_events):,}") + roots = ModuleTreeBuilder().build_from_module_events(module_events) + + shape_index = CpuOpShapeIndex(cpu_ops, runtime_events, driver_events) + source_index = PythonSourceIndex(py_events, runtime_events, driver_events) + all_gpu = kernel_events + gpu_memcpy + gpu_memset + correlator = KernelCorrelator(runtime_events, roots, driver_events) + correlator.correlate(all_gpu, roots, + shape_index=shape_index, source_index=source_index) + + graph_corr = CudaGraphCorrelator(runtime_events) + if graph_corr.has_graph_replays: + unmatched = [e for e in all_gpu if not e.get("_matched")] + new_roots, _ = graph_corr.correlate(unmatched, roots) + roots.extend(new_roots) + + PhaseDetector().detect_from_markers(roots, phase_markers) + stats = ModuleAggregator().aggregate(roots, mode="full") + + # Propagate phase from nodes onto stats (mirrors TraceModuleAnalyzer) + def _copy_phase(node, stat): + ph = getattr(node, "_phase", "") + if ph: + stat.phase = ph + for cn, cs in zip(node.children, stat.children_stats): + _copy_phase(cn, cs) + for r, s in zip(roots, stats): + _copy_phase(r, s) + return stats + + +# --------------------------------------------------------------------------- +# Step 2 — locate one instance of a module type by instance_id. +# --------------------------------------------------------------------------- + +def find_instance(stats_list: List[ModuleStats], + module_type: str, + instance_id: int) -> Optional[ModuleStats]: + for s in stats_list: + if s.module_type == module_type and s.instance_id == instance_id: + return s + found = find_instance(s.children_stats, module_type, instance_id) + if found is not None: + return found + return None + + +def list_instances(stats_list: List[ModuleStats], + module_type: str) -> List[ModuleStats]: + out: List[ModuleStats] = [] + + def _walk(slist): + for s in slist: + if s.module_type == module_type: + out.append(s) + _walk(s.children_stats) + _walk(stats_list) + return out + + +def _find_best_prof_halves(pf_stats: List[ModuleStats], module_type: str, + layer_index: int, + ng_kernels: List["FlatKernel"]) -> List[int]: + """Scan prof Layer instances near 2*N+1 and pick contiguous halves whose + combined kernel-name set best matches the nograph layer's kernel set.""" + target = {k.name[:80] for k in ng_kernels} + all_inst = list_instances(pf_stats, module_type) + by_id = {s.instance_id: s for s in all_inst} + n_inst = len(all_inst) + n_target = len(ng_kernels) + + # Decide overall mode by ratio of prof instances per nograph layer. + # KimiK2.5 has 61 decoder layers; if prof has ~61, it's full-layer mode; + # if ~122 per iter, it's halves mode. Use total/61 as a hint. + layers_per_iter = max(1, n_inst // max(1, _estimate_iter_count(all_inst))) + halves_mode = layers_per_iter > 80 # > ~61 * 1.5 + + if halves_mode: + candidate_bases = [layer_index * 2 + 1, layer_index * 2] + length_choices = (2,) # require attn+mlp pair + else: + candidate_bases = [layer_index] + length_choices = (1,) + + best_score = -1.0 + best_halves: List[int] = [] + for base in candidate_bases: + for start in range(max(0, base - 2), base + 3): + for length in length_choices: + ids = [start + k for k in range(length)] + insts = [by_id.get(i) for i in ids] + if any(s is None for s in insts): + continue + kernels = [] + for s in insts: + kernels.extend(flatten_decoder_layer(s)) + cand = {k.name[:80] for k in kernels} + if not cand: + continue + inter = len(target & cand) + union = len(target | cand) + jacc = inter / union + dist_penalty = abs(start - base) * 0.01 + score = jacc - dist_penalty + if score > best_score: + best_score = score + best_halves = ids + if not best_halves: + best_halves = [layer_index] + return best_halves + + +def _estimate_iter_count(all_inst: List[ModuleStats]) -> int: + """Heuristic: a decode trace usually runs N iterations. Return N by + looking at how many times instance_id resets (or, simplistically, the + max instance_id over the list, divided into the total count).""" + if not all_inst: + return 1 + # If instance_ids restart per iteration, max+1 is one iter's count. + # Otherwise total/max gives iter count. + max_id = max(s.instance_id for s in all_inst) + total = len(all_inst) + # If max_id == total-1, instance ids are unique → 1 iter (unusual) + # Usually there are 5 iterations and ids repeat 0..N-1. + return max(1, total // (max_id + 1)) + + +def find_decode_instance(stats_list: List[ModuleStats], + module_type: str, + layer_index: int) -> ModuleStats: + """Return the decode-phase instance with module_type whose ordinal among + decode instances equals layer_index (0-based in the decode pool). + + For nograph, decode pool of DeepseekV2DecoderLayer typically holds one full + sequence of 0..N-1 per decode iteration; we just take the layer_index-th + instance of the first decode iteration (instance_id == layer_index in + practice). For prof, "Layer" indices count both halves and accumulate + across iterations. + """ + all_inst = list_instances(stats_list, module_type) + decode = [s for s in all_inst if getattr(s, "phase", "") == "decode"] + if not decode: + decode = all_inst + # Prefer exact instance_id match within the first decode iteration + by_id = [s for s in decode if s.instance_id == layer_index] + if by_id: + return by_id[0] + if layer_index >= len(decode): + raise SystemExit( + f"layer_index={layer_index} out of range for {module_type} " + f"(decode instances: {len(decode)})") + return decode[layer_index] + + +# --------------------------------------------------------------------------- +# Step 3 — collect the in-order kernel list for one module instance +# (descend into children, sorted by timestamp). +# --------------------------------------------------------------------------- + +@dataclass +class FlatKernel: + name: str + duration: float + category: str + leaf_module: str # e.g. RadixAttention_40 + parent_module: str # the depth-1-under-decoderlayer ancestor name + shape: str + source: str + ts: float + + +def flatten_kernels(node_stats: ModuleStats, + parent_at_layer_level: str = "") -> List[FlatKernel]: + """Recursively collect kernels under node_stats in trace-time order. + + parent_at_layer_level: the module-name to record as the 'layer_group' + for direct-and-descendant kernels. When None, the immediate child of the + top-level DecoderLayer is used; otherwise inherited from caller. + """ + out: List[FlatKernel] = [] + own_parent = parent_at_layer_level or node_stats.name + + for kd in node_stats.kernel_details: + out.append(FlatKernel( + name=kd.name, + duration=kd.duration, + category=kd.category, + leaf_module=node_stats.name, + parent_module=own_parent, + shape=getattr(kd, "input_dims", "") or "", + source=getattr(kd, "source_path", "") or "", + ts=getattr(kd, "ts", 0.0), + )) + + for child in node_stats.children_stats: + # When walking out of the DecoderLayer down, set parent_at_layer_level + # to the child's name on first descent. + sub_parent = (own_parent + if parent_at_layer_level else child.name) + out.extend(flatten_kernels(child, sub_parent)) + + out.sort(key=lambda k: k.ts) + return out + + +_INSTANCE_SUFFIX_RE = __import__("re").compile(r"_\d+$") + + +def _strip_instance(name: str) -> str: + return _INSTANCE_SUFFIX_RE.sub("", name) + + +def flatten_decoder_layer(layer_stats: ModuleStats) -> List[FlatKernel]: + """Flatten kernels under one DecoderLayer instance with proper layer groups. + + The 'layer_group' (e.g. DeepseekV2AttentionMLA, DeepseekV2MoE) is the + name of the immediate child sub-module type. Layer-direct kernels get + synthetic group labels (prepare_attn, prepare_mlp, post_layer) based on + their time position relative to the first/last sub-module kernel. + """ + out: List[FlatKernel] = [] + layer_name = layer_stats.name + + for kd in layer_stats.kernel_details: + out.append(FlatKernel( + name=kd.name, + duration=kd.duration, + category=kd.category, + leaf_module=layer_name, + parent_module=layer_name, # placeholder; renamed below + shape=getattr(kd, "input_dims", "") or "", + source=getattr(kd, "source_path", "") or "", + ts=getattr(kd, "ts", 0.0), + )) + for child in layer_stats.children_stats: + # parent_at_layer_level is the type name (no instance suffix) + out.extend(flatten_kernels(child, + parent_at_layer_level=_strip_instance(child.name))) + + out.sort(key=lambda k: k.ts) + + sub_module_present = [i for i, k in enumerate(out) + if k.parent_module != layer_name] + if sub_module_present: + first_sub = sub_module_present[0] + last_sub = sub_module_present[-1] + mid_layer_ks = [i for i in range(first_sub + 1, last_sub) + if out[i].parent_module == layer_name] + prepare_mlp_idx = mid_layer_ks[0] if mid_layer_ks else None + for i, k in enumerate(out): + if k.parent_module != layer_name: + continue + if i < first_sub: + k.parent_module = "prepare_attn" + elif prepare_mlp_idx is not None and i == prepare_mlp_idx: + k.parent_module = "prepare_mlp" + elif i > last_sub: + k.parent_module = "post_layer" + else: + k.parent_module = "prepare_mlp" + return out + + +# --------------------------------------------------------------------------- +# Step 4 — LCS alignment by kernel name. +# --------------------------------------------------------------------------- + +def lcs_align(a: List[FlatKernel], b: List[FlatKernel] + ) -> List[Tuple[Optional[int], Optional[int]]]: + """Return list of (idx_a, idx_b) pairs covering both sequences. + + Matched pairs have both indices; insertions have (None, j); deletions + have (i, None). Match key is the (truncated) kernel name. + """ + def key(k: FlatKernel) -> str: + # Use first 80 chars to tolerate trivial template differences. + return k.name[:80] + + n, m = len(a), len(b) + # Standard LCS DP + dp = [[0] * (m + 1) for _ in range(n + 1)] + for i in range(n - 1, -1, -1): + for j in range(m - 1, -1, -1): + if key(a[i]) == key(b[j]): + dp[i][j] = dp[i + 1][j + 1] + 1 + else: + dp[i][j] = max(dp[i + 1][j], dp[i][j + 1]) + + out: List[Tuple[Optional[int], Optional[int]]] = [] + i = j = 0 + while i < n and j < m: + if key(a[i]) == key(b[j]): + out.append((i, j)) + i += 1 + j += 1 + elif dp[i + 1][j] >= dp[i][j + 1]: + out.append((i, None)) + i += 1 + else: + out.append((None, j)) + j += 1 + while i < n: + out.append((i, None)) + i += 1 + while j < m: + out.append((None, j)) + j += 1 + return out + + +# --------------------------------------------------------------------------- +# Step 5 — write the merged xlsx. +# --------------------------------------------------------------------------- + +CATEGORY_PROPERTY = { + "attention": "MLA", + "moe": "MoE", + "gemm": "GEMM", + "communication": "COMM", + "elementwise": "EW", + "embedding": "EMB", +} + + +def build_merged_rows(nograph_kernels: List[FlatKernel], + prof_kernels: List[FlatKernel] + ) -> List[KernelRow]: + """LCS-align prof against nograph; carry module info from nograph. + + For prof kernels not matched by LCS, infer the module by looking at the + nograph anchors on either side: any unmatched prof kernels between + anchors ng[a..b] get labeled with nograph kernels' module info from + that gap region (positional fallback). + """ + pairs = lcs_align(nograph_kernels, prof_kernels) + + # Build mapping: prof_idx -> (ng_idx or None) + prof_to_ng: Dict[int, Optional[int]] = {} + last_anchor_ng = -1 + next_anchor_ng_for: Dict[int, int] = {} + # First pass: matches + for (ig, ip) in pairs: + if ip is None: + continue + prof_to_ng[ip] = ig + + # Walk through pairs in order and assign nograph fallback labels to + # unmatched prof kernels. For each "gap" of prof kernels between two + # nograph anchors, distribute the nograph kernels in that gap evenly. + rows: List[KernelRow] = [] + # Convert pairs into (ng_idx or None, prof_idx or None) ordered list + # and walk grouped by prof kernel. + prof_seq: List[Tuple[int, Optional[int]]] = [] # (prof_idx, ng_idx) + # Collect mapping while preserving order from LCS + ip_seen = set() + ng_gap: List[int] = [] + for (ig, ip) in pairs: + if ip is None: + if ig is not None: + ng_gap.append(ig) + continue + # When we hit a prof kernel: emit any queued unmatched ng kernels + # by associating them positionally with prior unmatched prof rows. + if ig is not None: + prof_seq.append((ip, ig)) + else: + prof_seq.append((ip, None)) + + # Now produce rows in prof order, with fallback nograph labels for + # unmatched prof kernels from neighbouring anchors. + for k, (ip, ig) in enumerate(prof_seq): + pk = prof_kernels[ip] + cat = pk.category or _categorize_kernel(pk.name) + if ig is not None: + ng = nograph_kernels[ig] + layer_group = ng.parent_module + module = ng.leaf_module + shape = ng.shape + source = ng.source + else: + # Find nearest prior and next matched anchor in prof_seq + prev_ng = next( + (prof_seq[j][1] for j in range(k - 1, -1, -1) + if prof_seq[j][1] is not None), + None) + next_ng = next( + (prof_seq[j][1] for j in range(k + 1, len(prof_seq)) + if prof_seq[j][1] is not None), + None) + # Candidate nograph kernels in the gap (prev_ng, next_ng) + lo = (prev_ng + 1) if prev_ng is not None else 0 + hi = next_ng if next_ng is not None else len(nograph_kernels) + gap_ng = nograph_kernels[lo:hi] + # Count unmatched prof kernels in this gap to index into + gap_prof_indices = [ + j for j in range( + (next( + (jj + 1 for jj in range(k - 1, -1, -1) + if prof_seq[jj][1] is not None), 0)), + (next( + (jj for jj in range(k + 1, len(prof_seq)) + if prof_seq[jj][1] is not None), len(prof_seq)))) + if prof_seq[j][1] is None] + if gap_ng: + try: + rel = gap_prof_indices.index(k) + except ValueError: + rel = 0 + ng_idx = min(int(rel * len(gap_ng) / max(1, len(gap_prof_indices))), + len(gap_ng) - 1) + ng = gap_ng[ng_idx] + layer_group = ng.parent_module + # Use the gap-region nograph kernel's leaf module, tagged as + # unmapped (kernel name differs but spatially in same module). + module = "(unmapped " + _strip_instance(ng.leaf_module) + ")" + shape = "" + source = ng.source + elif prev_ng is not None: + # Trailing-tail kernels (no next anchor in nograph). + # post_layer only when it's a *large* comm kernel that + # is clearly the layer-output allreduce (>10 us). + # Otherwise inherit prev anchor's group with (unmapped) tag. + ng = nograph_kernels[prev_ng] + if cat == "communication" and pk.duration > 10: + layer_group = "post_layer" + module = "" + else: + layer_group = ng.parent_module + module = "(unmapped " + _strip_instance(ng.leaf_module) + ")" + shape = "" + source = "" + elif next_ng is not None: + ng = nograph_kernels[next_ng] + layer_group = ng.parent_module + module = "(unmapped " + _strip_instance(ng.leaf_module) + ")" + shape = "" + source = "" + else: + layer_group = "(unmatched)" + module = "(unknown)" + shape = "" + source = "" + rows.append(KernelRow( + layer_group=layer_group, + module=module, + shape=shape, + kernel_name=pk.name, + duration_us=pk.duration, + category=cat, + source=source, + )) + return rows + + +def write_xlsx(rows: List[KernelRow], output_path: str, + title: str = "Merged decode layer"): + import openpyxl + from openpyxl.styles import Alignment, Font, PatternFill + + wb = openpyxl.Workbook() + ws = wb.active + ws.title = "Merged Layer" + + bold = Font(bold=True) + header_fill = PatternFill(start_color="CCE5FF", end_color="CCE5FF", + fill_type="solid") + group_fill = PatternFill(start_color="EFF6FF", end_color="EFF6FF", + fill_type="solid") + subtotal_fill = PatternFill(start_color="F0F0F0", end_color="F0F0F0", + fill_type="solid") + + # Title row + ws.cell(row=1, column=1, value=title).font = Font(bold=True, size=12) + + # Total wall time = sum (treat as approximate; no overlap data here) + total = sum(r.duration_us for r in rows) or 1.0 + + headers = ["Layer", "Module", "shape", "Kernel_name", + "call-times", "time (us)", "percentage (%)", "properties", + "source"] + for col, h in enumerate(headers, 1): + c = ws.cell(row=3, column=col, value=h) + c.font = bold + c.fill = header_fill + c.alignment = Alignment(horizontal="center") + + row_idx = 4 + # Group rows by layer_group; emit in canonical order, unknowns after. + CANONICAL = [ + "prepare_attn", + "DeepseekV2AttentionMLA", + "prepare_mlp", + "DeepseekV2MoE", + "DeepseekV2MLP", + "post_layer", + ] + group_to_rows: Dict[str, List[KernelRow]] = {} + file_order: List[str] = [] + for r in rows: + if r.layer_group not in group_to_rows: + group_to_rows[r.layer_group] = [] + file_order.append(r.layer_group) + group_to_rows[r.layer_group].append(r) + seen_groups: List[str] = [g for g in CANONICAL if g in group_to_rows] + for g in file_order: + if g not in seen_groups: + seen_groups.append(g) + + for g in seen_groups: + group_rows = group_to_rows[g] + group_total = sum(r.duration_us for r in group_rows) + + for i, r in enumerate(group_rows): + ws.cell(row=row_idx, column=1, value=(g if i == 0 else "")) + # Strip _NN instance suffix from leaf module name for readability + mod_disp = r.module + if mod_disp and not mod_disp.startswith("("): + mod_disp = _strip_instance(mod_disp) + ws.cell(row=row_idx, column=2, value=mod_disp) + ws.cell(row=row_idx, column=3, value=r.shape) + ws.cell(row=row_idx, column=4, value=r.kernel_name) + ws.cell(row=row_idx, column=5, value=1) + ws.cell(row=row_idx, column=6, value=round(r.duration_us, 3)) + pct = r.duration_us / total * 100 if total else 0 + ws.cell(row=row_idx, column=7, value=f"{pct:.1f}%") + ws.cell(row=row_idx, column=8, + value=CATEGORY_PROPERTY.get(r.category, r.category)) + ws.cell(row=row_idx, column=9, value=r.source) + if i == 0: + ws.cell(row=row_idx, column=1).fill = group_fill + ws.cell(row=row_idx, column=1).font = bold + row_idx += 1 + # subtotal + ws.cell(row=row_idx, column=6, value=round(group_total, 3)).font = bold + for c in range(1, 10): + ws.cell(row=row_idx, column=c).fill = subtotal_fill + row_idx += 1 + + # Grand total + ws.cell(row=row_idx, column=1, value="TOTAL").font = bold + ws.cell(row=row_idx, column=6, value=round(total, 3)).font = bold + + widths = [28, 28, 30, 70, 10, 12, 12, 12, 60] + from openpyxl.utils import get_column_letter + for i, w in enumerate(widths, 1): + ws.column_dimensions[get_column_letter(i)].width = w + + wb.save(output_path) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def main(): + p = argparse.ArgumentParser( + description="Merge a CUDA-graph trace with its no-graph counterpart " + "for one decode layer.") + p.add_argument("--nograph", required=True, + help="Trace from a --disable-cuda-graph run (carries module info)") + p.add_argument("--prof", required=True, + help="Trace from the normal run with CUDA graph (real perf)") + p.add_argument("--layer-index", type=int, default=40, + help="Decoder layer index in the nograph trace (default 40)") + p.add_argument("--nograph-module", default="DeepseekV2DecoderLayer", + help="Module type name for a full decoder layer in nograph") + p.add_argument("--prof-module", default="Layer", + help="Synthetic graph layer module name in prof") + p.add_argument("--prof-halves", type=int, nargs="+", default=None, + help="Specific prof Layer indices to take (default: " + "2*layer_index+1 and +2, the attn+moe halves)") + p.add_argument("-o", "--output", required=True, help="Output xlsx path") + args = p.parse_args() + + print(f"[1/4] Analyzing nograph trace: {args.nograph}") + ng_stats = analyze_trace(args.nograph) + ng_layer = find_decode_instance(ng_stats, args.nograph_module, + args.layer_index) + print(f" Picked nograph instance: {ng_layer.name}") + ng_kernels = flatten_decoder_layer(ng_layer) + print(f" {len(ng_kernels)} kernels in nograph layer") + + print(f"[2/4] Analyzing prof trace: {args.prof}") + pf_stats = analyze_trace(args.prof) + + # Determine which prof Layer halves correspond to this decoder layer. + # If --prof-halves given, use it. Otherwise scan a window around 2N+1 + # and pick the contiguous run whose combined kernel set best matches + # nograph's kernel set (by short-name Jaccard similarity). + if args.prof_halves: + halves = args.prof_halves + print(f" Using user-specified prof Layer halves: {halves}") + else: + halves = _find_best_prof_halves( + pf_stats, args.prof_module, args.layer_index, ng_kernels) + print(f" Auto-picked prof Layer halves: {halves}") + pf_kernels: List[FlatKernel] = [] + for h in halves: + inst = find_instance(pf_stats, args.prof_module, h) + if inst is None: + print(f" WARNING: prof {args.prof_module}_{h} not found, skipping") + continue + pf_kernels.extend(flatten_decoder_layer(inst)) + print(f" {len(pf_kernels)} kernels combined from prof halves") + + print("[3/4] Aligning by kernel name (LCS) ...") + rows = build_merged_rows(ng_kernels, pf_kernels) + print(f" {len(rows)} merged rows") + + print(f"[4/4] Writing {args.output}") + title = (f"Merged decode layer {args.layer_index} | " + f"nograph instance {ng_layer.name} | " + f"prof halves {halves}") + write_xlsx(rows, args.output, title=title) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/report_html.py b/report_html.py new file mode 100644 index 0000000..4f8c73a --- /dev/null +++ b/report_html.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 +"""Generate a self-contained HTML report from two merged-layer xlsx files. + +The report has three sections: + 1. Category bar chart (like categorize.png) — kernels grouped by their + 'properties' tag (GEMM, MoE, MLA, COMM, EW, EMB, quantize, other) + with side-by-side bars per platform. + 2. Layer-group summary table — totals per Layer block + Δ. + 3. Per-Layer side-by-side detail — every kernel pair from the merged xlsx. + +Output is a single .html with no external dependencies (inline SVG + CSS). +""" + +import argparse +import re +from collections import defaultdict, OrderedDict +from dataclasses import dataclass +from html import escape +from typing import Dict, List, Optional, Tuple + +import openpyxl + +_INSTANCE_RE = re.compile(r"_\d+$") + + +def _strip(name: str) -> str: + return _INSTANCE_RE.sub("", name) if name else name + + +# ── Property/category canonicalization for the bar chart ──────────────────── + +PROPERTY_DISPLAY = { + "GEMM": "GEMM", + "MoE": "MoE", + "MLA": "MLA", + "COMM": "COMM", + "EW": "elementwise", + "EMB": "embedding", + "quantization": "quantize", + "quant": "quantize", + "normalization": "norm", + "norm": "norm", + "other": "others", + "": "others", +} + +CATEGORY_ORDER = ["GEMM", "MoE", "MLA", "COMM", "elementwise", + "embedding", "norm", "quantize", "others"] + +CATEGORY_COLORS = { + "GEMM": "#3b82f6", + "MoE": "#10b981", + "MLA": "#f59e0b", + "COMM": "#ef4444", + "elementwise": "#a855f7", + "embedding": "#06b6d4", + "norm": "#0ea5e9", + "quantize": "#84cc16", + "others": "#94a3b8", +} + + +@dataclass +class Row: + layer: str + module: str + shape: str + kernel_name: str + duration_us: float + pct: str + properties: str + + +def load_merged(path: str) -> Tuple[List[Row], float]: + wb = openpyxl.load_workbook(path, read_only=True) + ws = wb["Merged Layer"] + rows: List[Row] = [] + total = 0.0 + cur_layer = "" + for r in ws.iter_rows(values_only=True): + if r[0] == "TOTAL": + total = float(r[5] or 0) + continue + if r[0] == "Layer": + continue + if r[0]: + cur_layer = r[0] + if r[3] is None: + continue + try: + d = float(r[5] or 0) + except (TypeError, ValueError): + continue + rows.append(Row( + layer=cur_layer, module=str(r[1] or ""), shape=str(r[2] or ""), + kernel_name=str(r[3] or ""), duration_us=d, + pct=str(r[6] or ""), properties=str(r[7] or ""), + )) + wb.close() + return rows, total + + +def categorize_rows(rows: List[Row]) -> Dict[str, float]: + out: Dict[str, float] = defaultdict(float) + for r in rows: + cat = PROPERTY_DISPLAY.get(r.properties, r.properties or "others") + out[cat] += r.duration_us + return dict(out) + + +# ── Layer-block alignment (mirrors compare_merged.py) ─────────────────────── + +def align_layer_block(a_rows: List[Row], b_rows: List[Row] + ) -> List[Tuple[Optional[Row], Optional[Row]]]: + pairs: List[Tuple[Optional[Row], Optional[Row]]] = [] + i = j = 0 + while i < len(a_rows) or j < len(b_rows): + if i >= len(a_rows): + pairs.append((None, b_rows[j])); j += 1; continue + if j >= len(b_rows): + pairs.append((a_rows[i], None)); i += 1; continue + a, b = a_rows[i], b_rows[j] + if a.module == b.module and a.module: + pairs.append((a, b)); i += 1; j += 1 + continue + a_later_in_b = any(b_rows[k].module == a.module + for k in range(j + 1, len(b_rows))) + b_later_in_a = any(a_rows[k].module == b.module + for k in range(i + 1, len(a_rows))) + if a_later_in_b and not b_later_in_a: + pairs.append((None, b)); j += 1 + elif b_later_in_a and not a_later_in_b: + pairs.append((a, None)); i += 1 + else: + pairs.append((a, b)); i += 1; j += 1 + return pairs + + +# ── Inline SVG bar chart ──────────────────────────────────────────────────── + +def svg_grouped_bars(cats: List[str], a_vals: List[float], b_vals: List[float], + label_a: str, label_b: str, title: str) -> str: + W, H = 900, 380 + pad_left, pad_right, pad_top, pad_bot = 60, 30, 60, 130 + plot_w = W - pad_left - pad_right + plot_h = H - pad_top - pad_bot + max_v = max(max(a_vals + b_vals), 1) + # Round up max for grid + grid_step = 5.0 if max_v <= 50 else (10.0 if max_v <= 100 else 20.0) + max_grid = (int(max_v / grid_step) + 1) * grid_step + + n = len(cats) + group_w = plot_w / n + bar_w = group_w * 0.32 + gap = group_w * 0.06 + + parts = [f''] + # Title + parts.append(f'{escape(title)}') + + # Y axis grid + labels + n_ticks = int(max_grid / grid_step) + for t in range(n_ticks + 1): + val = t * grid_step + y = pad_top + plot_h - (val / max_grid * plot_h) + parts.append(f'') + parts.append(f'{val:g}') + parts.append(f'TIME(us)') + + # Bars + value labels + for i, cat in enumerate(cats): + cx = pad_left + i * group_w + group_w / 2 + a_h = (a_vals[i] / max_grid) * plot_h + b_h = (b_vals[i] / max_grid) * plot_h + a_x = cx - bar_w - gap / 2 + b_x = cx + gap / 2 + a_y = pad_top + plot_h - a_h + b_y = pad_top + plot_h - b_h + color = CATEGORY_COLORS.get(cat, "#94a3b8") + # A bar (red-ish for first) + parts.append(f'') + parts.append(f'{a_vals[i]:.1f}') + # B bar (green for second) + parts.append(f'') + parts.append(f'{b_vals[i]:.1f}') + # X-axis category label + parts.append(f'{escape(cat)}') + + # Embedded data table below the chart + row1_y = pad_top + plot_h + 36 + row2_y = pad_top + plot_h + 58 + parts.append(f'') + parts.append(f'{escape(label_a)}') + parts.append(f'') + parts.append(f'{escape(label_b)}') + for i, cat in enumerate(cats): + cx = pad_left + i * group_w + group_w / 2 + parts.append(f'{a_vals[i]:.2f}') + parts.append(f'{b_vals[i]:.2f}') + + # Legend at the very bottom + legend_y = H - 18 + parts.append(f'') + parts.append(f'{escape(label_a)}') + parts.append(f'') + parts.append(f'{escape(label_b)}') + + parts.append('') + return "".join(parts) + + +# ── HTML generation ───────────────────────────────────────────────────────── + +CSS = """ +body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + margin: 24px auto; max-width: 1400px; color: #111827; } +h1 { margin: 0 0 8px 0; font-size: 22px; } +h2 { margin: 28px 0 12px 0; font-size: 18px; color: #1f2937; + border-bottom: 1px solid #e5e7eb; padding-bottom: 4px; } +.subtitle { color: #6b7280; margin-bottom: 18px; } +table { border-collapse: collapse; font-size: 13px; margin-bottom: 8px; } +th, td { padding: 4px 10px; border: 1px solid #e5e7eb; text-align: right; } +th { background: #f3f4f6; font-weight: 600; } +td.l, th.l { text-align: left; } +tr.layer-row td { background: #eff6ff; font-weight: 600; } +tr.subtotal td { background: #f3f4f6; font-weight: 600; } +tr.total td { background: #d1d5db; font-weight: 700; } +.slower { background: #fee2e2 !important; } +.faster { background: #d1fae5 !important; } +.dim { color: #9ca3af; } +.kname { font-family: ui-monospace, Consolas, monospace; font-size: 11.5px; + max-width: 320px; overflow: hidden; text-overflow: ellipsis; + white-space: nowrap; } +.spacer { background: #fafafa; border-top: none; border-bottom: none; width: 8px; } +""" + + +def fmt_pct(base: float, target: float) -> str: + if base <= 0: + return "" + return f"{(target-base)/base*100:+.1f}%" + + +def render_summary_table(merged_layers: List[str], + a_by_layer: Dict[str, List[Row]], + b_by_layer: Dict[str, List[Row]], + a_total: float, b_total: float, + label_a: str, label_b: str) -> str: + parts = ['', + f'' + f'' + f'' + f''] + for L in merged_layers: + at = sum(r.duration_us for r in a_by_layer.get(L, [])) + bt = sum(r.duration_us for r in b_by_layer.get(L, [])) + d = bt - at + cls = "slower" if d > 0.5 else ("faster" if d < -0.5 else "") + parts.append( + f'' + f'' + f'') + td = b_total - a_total + parts.append( + f'' + f'' + f'') + parts.append('
Layer group{escape(label_a)} (us){escape(label_b)} (us)Δ (us)Δ %
{escape(L)}{at:.1f}{bt:.1f}{d:+.1f}{fmt_pct(at, bt)}
TOTAL{a_total:.1f}{b_total:.1f}{td:+.1f}{fmt_pct(a_total, b_total)}
') + return "".join(parts) + + +def render_detail_table(merged_layers: List[str], + a_by_layer: Dict[str, List[Row]], + b_by_layer: Dict[str, List[Row]], + label_a: str, label_b: str) -> str: + parts = ['', + f'' + f'' + f'' + f'' + f'' + f'' + f'' + f''] + for L in merged_layers: + a_grp = a_by_layer.get(L, []) + b_grp = b_by_layer.get(L, []) + pairs = align_layer_block(a_grp, b_grp) + first = True + for (a, b) in pairs: + cells = [ + f'', + f'', + f'', + f'' if a else '', + f'', + '', + f'', + f'', + f'' if b else '', + f'', + ] + cls = "layer-row" if first else "" + parts.append(f'{"".join(cells)}') + first = False + # Subtotal + at = sum(r.duration_us for r in a_grp) + bt = sum(r.duration_us for r in b_grp) + parts.append( + '' + '' + f'' + '' + f'') + parts.append('
Layer{escape(label_a)}: Module{escape(label_a)}: Kernelt (us)%{escape(label_b)}: Module{escape(label_b)}: Kernelt (us)%
{escape(L) if first else ""}{escape(a.module) if a else ""}{escape(a.kernel_name) if a else ""}{a.duration_us:.2f}{escape(a.pct) if a else ""}{escape(b.module) if b else ""}{escape(b.kernel_name) if b else ""}{b.duration_us:.2f}{escape(b.pct) if b else ""}
subtotal{at:.2f}subtotal{bt:.2f}
') + return "".join(parts) + + +def build_html(a_path: str, b_path: str, output: str, + label_a: str, label_b: str, title: str): + a_rows, a_total = load_merged(a_path) + b_rows, b_total = load_merged(b_path) + a_by_layer = defaultdict(list) + for r in a_rows: a_by_layer[r.layer].append(r) + b_by_layer = defaultdict(list) + for r in b_rows: b_by_layer[r.layer].append(r) + + CANONICAL = ["prepare_attn", "DeepseekV2AttentionMLA", "prepare_mlp", + "DeepseekV2MoE", "DeepseekV2MLP", "post_layer"] + seen = set(a_by_layer) | set(b_by_layer) + merged_layers = [L for L in CANONICAL if L in seen] + for L in [r.layer for r in a_rows] + [r.layer for r in b_rows]: + if L not in merged_layers and L in seen: + merged_layers.append(L) + + a_cat = categorize_rows(a_rows) + b_cat = categorize_rows(b_rows) + all_cats = [c for c in CATEGORY_ORDER if c in a_cat or c in b_cat] + for c in list(a_cat) + list(b_cat): + if c not in all_cats: + all_cats.append(c) + a_vals = [a_cat.get(c, 0.0) for c in all_cats] + b_vals = [b_cat.get(c, 0.0) for c in all_cats] + + chart_title = (f"Profiling DeepseekV2DecoderLayer in decode — " + f"{label_a} vs {label_b}") + svg = svg_grouped_bars(all_cats, a_vals, b_vals, + label_a, label_b, chart_title) + + summary = render_summary_table(merged_layers, a_by_layer, b_by_layer, + a_total, b_total, label_a, label_b) + detail = render_detail_table(merged_layers, a_by_layer, b_by_layer, + label_a, label_b) + + html = f""" + +{escape(title)} + +

{escape(title)}

+

{escape(label_a)} vs {escape(label_b)} · one DECODE layer (median instance)

+ +

Kernel time by category

+{svg} + +

Layer-group totals

+{summary} + +

Per-Layer side-by-side detail

+{detail} + +""" + with open(output, "w", encoding="utf-8") as f: + f.write(html) + print(f"Wrote {output}") + + +def main(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("a", help="Merged xlsx for run A (baseline)") + p.add_argument("b", help="Merged xlsx for run B (target)") + p.add_argument("-o", "--output", required=True) + p.add_argument("--labels", nargs=2, default=("A", "B")) + p.add_argument("--title", default="Profile compare report") + args = p.parse_args() + build_html(args.a, args.b, args.output, args.labels[0], args.labels[1], + args.title) + + +if __name__ == "__main__": + main() diff --git a/trace_module_analyzer.py b/trace_module_analyzer.py index 6a95230..52844ef 100755 --- a/trace_module_analyzer.py +++ b/trace_module_analyzer.py @@ -2725,6 +2725,93 @@ def _fallback_kernel_only(self, kernel_events: List[Dict]): pct = dur / total_dur * 100 if total_dur > 0 else 0 print(f" {cat:<20s} {dur:>14,.0f} {cnt:>8,d} {pct:>6.1f}%") + # Generate Excel if an output path was requested + xlsx_path = self.output_path + if not xlsx_path and self.model_info: + trace_dir = os.path.dirname(os.path.abspath(self.trace_path)) + trace_base = os.path.splitext(os.path.basename(self.trace_path))[0] + if trace_base.endswith(".json"): + trace_base = trace_base[:-5] + xlsx_path = os.path.join(trace_dir, f"{trace_base}_analysis.xlsx") + + if xlsx_path: + try: + from openpyxl import Workbook + from openpyxl.styles import Font, PatternFill + wb = Workbook() + ws = wb.active + ws.title = "Kernel-Only Summary" + header_font = Font(bold=True) + header_fill = PatternFill(start_color="CCE5FF", end_color="CCE5FF", + fill_type="solid") + headers = ["Category", "Duration (us)", "Count", "% of Total"] + for col, h in enumerate(headers, 1): + cell = ws.cell(row=1, column=col, value=h) + cell.font = header_font + cell.fill = header_fill + + row = 2 + for cat, (dur, cnt) in sorted(breakdown.items(), key=lambda x: -x[1][0]): + pct = dur / total_dur * 100 if total_dur > 0 else 0 + ws.cell(row=row, column=1, value=cat) + ws.cell(row=row, column=2, value=round(dur, 1)) + ws.cell(row=row, column=3, value=cnt) + ws.cell(row=row, column=4, value=round(pct, 2)) + row += 1 + + # Totals row + ws.cell(row=row, column=1, value="TOTAL").font = header_font + ws.cell(row=row, column=2, value=round(total_dur, 1)).font = header_font + ws.cell(row=row, column=3, + value=len(kernel_events)).font = header_font + ws.cell(row=row, column=4, value=100.0).font = header_font + + ws.column_dimensions["A"].width = 24 + + # GPU Kernels tab: per-kernel-name aggregation + ws_kn = wb.create_sheet(title="GPU Kernels") + kn_headers = ["Kernel Name", "Category", "Total Duration (us)", + "Count", "Avg (us)", "% of Total"] + for col, h in enumerate(kn_headers, 1): + cell = ws_kn.cell(row=1, column=col, value=h) + cell.font = header_font + cell.fill = header_fill + + kn_agg: Dict[str, List] = {} # name -> [dur, count, category] + for k in kernel_events: + kname = k.get("name", "") + dur = k.get("dur", 0) + cat = _categorize_kernel(kname) + entry = kn_agg.get(kname) + if entry: + entry[0] += dur + entry[1] += 1 + else: + kn_agg[kname] = [dur, 1, cat] + + sorted_kn = sorted(kn_agg.items(), key=lambda x: -x[1][0]) + truncated = len(sorted_kn) > MAX_ROWS_PER_TAB + kn_row = 2 + for kname, (dur, cnt, cat) in sorted_kn[:MAX_ROWS_PER_TAB]: + pct = dur / total_dur * 100 if total_dur > 0 else 0 + ws_kn.cell(row=kn_row, column=1, value=kname) + ws_kn.cell(row=kn_row, column=2, value=cat) + ws_kn.cell(row=kn_row, column=3, value=round(dur, 1)) + ws_kn.cell(row=kn_row, column=4, value=cnt) + ws_kn.cell(row=kn_row, column=5, + value=round(dur / cnt, 1) if cnt else 0) + ws_kn.cell(row=kn_row, column=6, value=round(pct, 1)) + kn_row += 1 + if truncated: + ws_kn.cell(row=kn_row, column=1, + value=f"... truncated at {MAX_ROWS_PER_TAB} rows") + ws_kn.column_dimensions["A"].width = 80 + + wb.save(xlsx_path) + print(f"\n Kernel-only Excel report saved to: {xlsx_path}") + except ImportError: + print(" (openpyxl not installed — skipping Excel export)") + @staticmethod def _load_trace(path: str) -> Dict[str, Any]: if path.endswith(".gz"): From 60fcc29c740d4194e689e5a6ee7ab73137e37dc0 Mon Sep 17 00:00:00 2001 From: Kao Date: Thu, 9 Jul 2026 13:11:18 +0800 Subject: [PATCH 3/4] =?UTF-8?q?Add=20--callstack-md:=20per-kernel=20CPU?= =?UTF-8?q?=E2=86=92GPU=20call=20stack=20markdown=20report?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - PythonSourceIndex gains get_callstack() / _find_callstack(): collects all python_function frames enclosing a kernel launch timestamp, sorted outermost-first (widest span first), stripping boilerplate wrappers. - KernelDetail gains a callstack slot populated during correlation. - KernelCorrelator.correlate() stores _callstack on each kernel event. - ReportGenerator.write_callstack_markdown() writes a Markdown file covering the same representative module instances as the Excel detail tabs, with per-kernel metadata tables and full call stack blocks. - New --callstack-md PATH argument; works standalone without -o. - Auto-selection uses self_kernel_time (not total) so container modules don't outrank the leaf layer types that own the actual work. Co-Authored-By: Claude --- trace_module_analyzer.py | 260 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 256 insertions(+), 4 deletions(-) diff --git a/trace_module_analyzer.py b/trace_module_analyzer.py index 6181bad..ae321ae 100755 --- a/trace_module_analyzer.py +++ b/trace_module_analyzer.py @@ -103,11 +103,12 @@ class KernelDetail: of instances for large traces). """ __slots__ = ("name", "duration", "category", "module_path", - "ts", "phase", "input_dims", "source_path") + "ts", "phase", "input_dims", "source_path", "callstack") def __init__(self, name: str, duration: float, category: str, module_path: str, ts: float = 0.0, phase: str = "", - input_dims: str = "", source_path: str = ""): + input_dims: str = "", source_path: str = "", + callstack: Optional[List[str]] = None): self.name = name self.duration = duration self.category = category @@ -116,6 +117,7 @@ def __init__(self, name: str, duration: float, category: str, self.phase = phase self.input_dims = input_dims self.source_path = source_path + self.callstack: List[str] = callstack if callstack is not None else [] @dataclass @@ -382,6 +384,7 @@ def __init__(self, pyfunc_events: List[Dict], runtime_events: List[Dict], } self._corr_to_source: Dict[int, str] = {} + self._corr_to_stack: Dict[int, List[str]] = {} all_launches = list(runtime_events) if driver_events: all_launches.extend(driver_events) @@ -395,10 +398,22 @@ def __init__(self, pyfunc_events: List[Dict], runtime_events: List[Dict], src = self._find_source(ts, pid, tid) if src: self._corr_to_source[corr] = src + stack = self._find_callstack(ts, pid, tid) + if stack: + self._corr_to_stack[corr] = stack def get_source(self, correlation_id: int) -> str: return self._corr_to_source.get(correlation_id, "") + def get_callstack(self, correlation_id: int) -> List[str]: + """Return the full CPU call stack for a kernel, outermost → innermost. + + Collects every python_function frame enclosing the kernel's launch + timestamp and returns them sorted widest-span-first (outermost caller + first). Skips pure boilerplate wrappers to keep the stack readable. + """ + return self._corr_to_stack.get(correlation_id, []) + def _find_source(self, ts: float, pid: int, tid: int) -> str: """Find the best Python source location enclosing *ts*. @@ -465,6 +480,42 @@ def _find_source(self, ts: float, pid: int, tid: int) -> str: return best_fallback return "" + # Trivial boilerplate to strip from call stacks (improves readability). + _BOILERPLATE_RE = re.compile( + r"threading\.py|multiprocessing/||tqdm/|importlib/" + r"|contextlib\.py|||", + re.IGNORECASE, + ) + + def _find_callstack(self, ts: float, pid: int, tid: int) -> List[str]: + """Collect all python_function frames enclosing *ts*, outermost first. + + Returns frames sorted by span descending (widest = outermost caller + first). Boilerplate-only frames are dropped to keep stacks concise. + """ + key = (pid, tid) + intervals = self._intervals.get(key) + if not intervals: + return [] + starts = self._starts_cache[key] + idx = bisect.bisect_right(starts, ts) - 1 + + # Collect every enclosing frame with its span for ordering. + frames: List[Tuple[float, str]] = [] # (span, name) — wider = earlier + lo = max(0, idx - 80) + hi = min(len(intervals), idx + 10) + for i in range(lo, hi): + s, e, name = intervals[i] + if s <= ts <= e: + if not self._BOILERPLATE_RE.search(name): + frames.append((e - s, name)) + elif s > ts + 100: + break + + # Sort widest-span first so callers appear before callees. + frames.sort(key=lambda x: -x[0]) + return [name for _, name in frames] + class KernelCorrelator: """Map GPU kernels to modules via correlation ID chain.""" @@ -509,6 +560,7 @@ def correlate(self, kernel_events: List[Dict], roots: List[ModuleNode], k["_input_dims"] = shape_index.get_shape(corr) if source_index is not None: k["_source_path"] = source_index.get_source(corr) + k["_callstack"] = source_index.get_callstack(corr) k["_matched"] = True module.kernels.append(k) matched += 1 @@ -1106,7 +1158,8 @@ def _aggregate_node(self, node: ModuleNode, depth: int, mode: str, name=kname, duration=dur, category=cat, module_path=path, ts=k.get("ts", 0), phase=node_phase, input_dims=k.get("_input_dims", ""), - source_path=k.get("_source_path", ""))) + source_path=k.get("_source_path", ""), + callstack=k.get("_callstack", []))) # Direct cpu_op stats for op in node.cpu_ops: @@ -2852,6 +2905,182 @@ def _collect_global_kernel_agg(self, stats_list: List[ModuleStats], agg[d.name] = [d.duration, 1, d.category, {s.module_type}] self._collect_global_kernel_agg(s.children_stats, agg) + # ------------------------------------------------------------------ + # Callstack markdown export + # ------------------------------------------------------------------ + + def write_callstack_markdown( + self, + stats_list: List[ModuleStats], + md_path: str, + mode: str, + max_detail_modules: int = 3, + detail_modules: Optional[List[str]] = None, + detail_instances: Optional[List[int]] = None, + ) -> None: + """Write a Markdown file with per-kernel CPU→GPU call stacks. + + For every kernel that appears in a "Layer (des)" detail sheet the + markdown shows: + + * The module path (nn.Module hierarchy) that owns the kernel + * The full Python call stack from the outermost caller down to the + kernel launch site + * Basic kernel metadata (category, duration, input dims) + + The selection logic for which module instances are covered mirrors + ``export_excel`` — the same representative instances used for detail + tabs are documented here. + """ + detail_modules = detail_modules or [] + + # Collect the representative instances the same way export_excel does. + # We reuse _select_detail_instances to get a consistent set. + selected = self._select_detail_instances( + stats_list, mode, + max_detail_modules=max_detail_modules, + detail_modules=detail_modules, + detail_instances=detail_instances, + ) + + lines: List[str] = [ + "# Kernel Call Stacks", + "", + "Per-kernel CPU → GPU call stacks for each representative module instance.", + "Frames are ordered **outermost caller → innermost (kernel launch site)**.", + "", + ] + + for module_stats, _variant in selected: + all_details = self._collect_all_details(module_stats) + all_details.sort(key=lambda d: d.ts) + + phase_tag = f" [{module_stats.phase}]" if module_stats.phase else "" + lines.append(f"## {module_stats.name}{phase_tag}") + lines.append("") + + has_stacks = any(getattr(d, "callstack", []) for d in all_details) + if not has_stacks: + lines.append("_No call stack data available for this module " + "(CUDA graph replays or cpu_only mode)._") + lines.append("") + continue + + for i, d in enumerate(all_details, 1): + stack = getattr(d, "callstack", []) + dims_str = f" | `{d.input_dims}`" if d.input_dims else "" + lines.append( + f"### Kernel {i}: `{d.name}`" + ) + lines.append("") + lines.append( + f"| Field | Value |" + ) + lines.append("| --- | --- |") + lines.append(f"| Module | `{d.module_path}` |") + lines.append(f"| Category | {d.category} |") + lines.append(f"| Duration | {d.duration:.1f} us |") + if d.input_dims: + lines.append(f"| Input Dims | `{d.input_dims}` |") + if d.source_path: + lines.append(f"| Launch site | `{d.source_path}` |") + lines.append("") + + if stack: + lines.append("**Call stack (CPU → kernel launch):**") + lines.append("") + lines.append("```") + for frame in stack: + lines.append(frame) + lines.append(f" → [GPU] {d.name}") + lines.append("```") + else: + lines.append("_Call stack unavailable for this kernel._") + lines.append("") + + with open(md_path, "w", encoding="utf-8") as fh: + fh.write("\n".join(lines)) + fh.write("\n") + + print(f"\n Call stack markdown saved to: {md_path}") + + def _select_detail_instances( + self, + stats_list: List[ModuleStats], + mode: str, + max_detail_modules: int = 3, + detail_modules: Optional[List[str]] = None, + detail_instances: Optional[List[int]] = None, + ) -> List[Tuple["ModuleStats", str]]: + """Return (stats, variant_label) pairs matching the detail-tab selection. + + Delegates to the same helper used by export_excel so the markdown + covers exactly the same instances as the spreadsheet detail tabs. + """ + detail_modules = detail_modules or [] + # Reuse the existing export_excel selection helpers. + # _find_detail_instances_for_export returns List[(ModuleStats, variant)] + return self._find_detail_instances_for_export( + stats_list, mode, + max_detail_modules=max_detail_modules, + detail_modules=detail_modules, + detail_instances=detail_instances, + ) + + def _find_detail_instances_for_export( + self, + stats_list: List[ModuleStats], + mode: str, + max_detail_modules: int = 3, + detail_modules: Optional[List[str]] = None, + detail_instances: Optional[List[int]] = None, + ) -> List[Tuple["ModuleStats", str]]: + """Shared logic: collect representative (stats, variant) for detail output.""" + detail_modules = detail_modules or [] + results: List[Tuple[ModuleStats, str]] = [] + + def _add_for_type(type_name: str): + all_of_type: List[ModuleStats] = [] + self._find_modules(stats_list, type_name, all_of_type) + if not all_of_type: + return + if detail_instances: + chosen = [m for m in all_of_type if m.instance_id in detail_instances] + else: + median_inst, _ = self._pick_median_instance(all_of_type, mode) + chosen = [median_inst] + for m in chosen: + results.append((m, "")) + + if detail_modules: + for dm in detail_modules: + _add_for_type(dm) + else: + # Mirror the top-N auto-selection from export_excel + type_times: Dict[str, float] = {} + self._collect_type_times(stats_list, type_times, mode) + top_types = sorted(type_times, key=lambda t: -type_times[t]) + seen = set() + for t in top_types: + if len(seen) >= max_detail_modules: + break + if t in seen: + continue + seen.add(t) + _add_for_type(t) + + return results + + def _collect_type_times(self, stats_list: List[ModuleStats], + acc: Dict[str, float], mode: str): + for s in stats_list: + # Use self_kernel_time (kernels directly owned, excluding children) + # so container modules like GptOssModel don't outrank the layer + # types that actually contain the work. + t = s.self_kernel_time if mode == "full" else s.self_cpu_op_time + acc[s.module_type] = acc.get(s.module_type, 0.0) + t + self._collect_type_times(s.children_stats, acc, mode) + # --------------------------------------------------------------------------- # Main orchestrator @@ -2869,7 +3098,8 @@ def __init__(self, trace_path: str, max_detail_modules: int = 3, auto_fix_rocm: bool = True, model_info: bool = False, - port: int = 8765): + port: int = 8765, + callstack_md: Optional[str] = None): self.trace_path = trace_path self.output_path = output_path self.detail_modules = detail_modules or [] @@ -2880,6 +3110,7 @@ def __init__(self, trace_path: str, self.auto_fix_rocm = auto_fix_rocm self.model_info = model_info self.port = port + self.callstack_md = callstack_md def run(self): import time as _time @@ -3124,6 +3355,15 @@ def _elapsed(): detail_instances=self.detail_instances) print(f" {_elapsed()} excel export done") + if self.callstack_md: + reporter.write_callstack_markdown( + stats_list, self.callstack_md, mode, + max_detail_modules=self.max_detail_modules, + detail_modules=self.detail_modules, + detail_instances=self.detail_instances, + ) + print(f" {_elapsed()} callstack markdown done") + if self.model_info and xlsx_path: html_path = os.path.splitext(xlsx_path)[0] + "_module_tree.html" result = _generate_model_info_html(xlsx_path, html_path, @@ -3659,6 +3899,11 @@ def main(): "(e.g. Prefill_0 + Decode_0)") parser.add_argument("--no-rocm-fix", action="store_true", help="Disable automatic ROCm trace fix (hipGraphLaunch flow events)") + parser.add_argument("--callstack-md", default=None, metavar="PATH", + help="Write a Markdown file with per-kernel CPU→GPU call stacks " + "for each representative module instance (same instances as " + "the 'Layer (des)' detail sheets in the Excel report). " + "Example: --callstack-md callstacks.md") parser.add_argument("-v", "--verbose", action="store_true", help="Enable debug logging") @@ -3692,6 +3937,12 @@ def main(): except ValueError: pass + # Resolve callstack-md path the same way as output + callstack_md = args.callstack_md + if callstack_md and not os.path.isabs(callstack_md): + trace_dir = os.path.dirname(os.path.abspath(args.trace_file)) + callstack_md = os.path.join(trace_dir, callstack_md) + try: analyzer = TraceModuleAnalyzer( trace_path=args.trace_file, @@ -3704,6 +3955,7 @@ def main(): auto_fix_rocm=not args.no_rocm_fix, model_info=args.model_info, port=args.port, + callstack_md=callstack_md, ) analyzer.run() except FileNotFoundError as e: From 10a1a86618c40f57782a30f5f8e090ce050564f3 Mon Sep 17 00:00:00 2001 From: Kao Date: Thu, 9 Jul 2026 13:26:59 +0800 Subject: [PATCH 4/4] Add --layer-html: interactive D3 kernel layer report MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generates a self-contained HTML file (no server needed) with: - Tab per selected module type showing instance count and kernel count - Kernel timeline chart: bar per kernel (execution order on X), Y = avg duration (µs) across all layer instances, error bars show min/max variance - Category bar chart: horizontal bars summing avg duration per category; updates live as kernels are reassigned via drag-and-drop - Category buckets: one chip per unique kernel type (not one per occurrence); dragging a chip reassigns all occurrences of that kernel in the layer; chip shows repeat count (e.g. "×2") when the kernel appears multiple times - Hover tooltips on all bars and chips with full kernel name, module path, category, avg/min/max, input dims, and launch source - Collapsible table view with all kernels in execution order - Light/dark mode via CSS custom properties and prefers-color-scheme - Categorical palette from the validated dataviz reference (8-slot fixed order, CVD-safe adjacent ordering) Works standalone without -o. Use --detail-module to target a specific layer type (e.g. --detail-module GptOssDecoderLayer). Auto-selection uses self_kernel_time ranking so container wrappers don't outrank leaf modules. Co-Authored-By: Claude --- trace_module_analyzer.py | 773 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 770 insertions(+), 3 deletions(-) diff --git a/trace_module_analyzer.py b/trace_module_analyzer.py index ae321ae..06d7405 100755 --- a/trace_module_analyzer.py +++ b/trace_module_analyzer.py @@ -2905,6 +2905,750 @@ def _collect_global_kernel_agg(self, stats_list: List[ModuleStats], agg[d.name] = [d.duration, 1, d.category, {s.module_type}] self._collect_global_kernel_agg(s.children_stats, agg) + # ------------------------------------------------------------------ + # Interactive HTML layer report + # ------------------------------------------------------------------ + + def write_layer_html( + self, + stats_list: List[ModuleStats], + html_path: str, + mode: str, + max_detail_modules: int = 3, + detail_modules: Optional[List[str]] = None, + detail_instances: Optional[List[int]] = None, + ) -> None: + """Write a self-contained interactive HTML report for per-layer kernel analysis. + + For each selected module type, collects ALL instances to compute per-kernel + avg / min / max duration across layers. The page shows: + - A kernel timeline bar+error chart (X = kernel slot, Y = avg us, error = min/max) + - Drag-and-drop category buckets for re-assigning kernels to categories + - A live category bar chart that updates as kernels are reassigned + All rendering is done with D3.js loaded from CDN; the page opens offline once cached. + """ + import json as _json + + detail_modules = detail_modules or [] + + # --- 1. Collect all instances per selected module type --- + type_instances: Dict[str, List[ModuleStats]] = {} + + if detail_modules: + for dm in detail_modules: + matches: List[ModuleStats] = [] + self._find_modules(stats_list, dm, matches) + if matches: + type_instances[dm] = matches + else: + type_times: Dict[str, float] = {} + self._collect_type_times(stats_list, type_times, mode) + top_types = sorted(type_times, key=lambda t: -type_times[t]) + seen: set = set() + for t in top_types: + if len(seen) >= max_detail_modules: + break + seen.add(t) + matches = [] + self._find_modules(stats_list, t, matches) + if matches: + type_instances[t] = matches + + # --- 2. Build per-module-type kernel slot data --- + # Align by kernel position within each instance's sorted kernel list. + # Instances may have different counts (e.g., prefill vs decode) — we + # group by phase first, then align within phase groups. + + def _build_layer_data(instances: List[ModuleStats]) -> List[Dict]: + """Return list of kernel slot dicts with avg/min/max across instances.""" + # Group instances by kernel count (proxy for same structural variant) + from collections import defaultdict as _dd + groups: Dict[int, List[List]] = _dd(list) + for inst in instances: + details = self._collect_all_details(inst) + details.sort(key=lambda d: d.ts) + groups[len(details)].append(details) + + # Pick the largest group (most common kernel count) + if not groups: + return [] + canonical_count = max(groups, key=lambda k: len(groups[k])) + group = groups[canonical_count] + + slots = [] + for slot_i in range(canonical_count): + durs = [inst_details[slot_i].duration for inst_details in group] + sample = group[0][slot_i] + slots.append({ + "slot": slot_i, + "name": sample.name, + "short_name": (sample.name[:20] + "…") if len(sample.name) > 20 else sample.name, + "category": sample.category, + "module": sample.module_path.split("/")[-1] if sample.module_path else "", + "avg": round(sum(durs) / len(durs), 2), + "min": round(min(durs), 2), + "max": round(max(durs), 2), + "input_dims": sample.input_dims, + "source": sample.source_path, + "instance_count": len(group), + }) + return slots + + layers_data = [] + for type_name, instances in type_instances.items(): + slots = _build_layer_data(instances) + if not slots: + continue + # Collect all unique categories present + cats = sorted(set(s["category"] for s in slots)) + layers_data.append({ + "type_name": type_name, + "instance_count": len(instances), + "slots": slots, + "categories": cats, + }) + + payload = _json.dumps(layers_data, ensure_ascii=False) + + # --- 3. Render HTML --- + html = self._layer_html_template(payload) + with open(html_path, "w", encoding="utf-8") as fh: + fh.write(html) + print(f"\n Layer HTML report saved to: {html_path}") + + @staticmethod + def _layer_html_template(payload_json: str) -> str: + return r""" + + + + +Kernel Layer Report + + + + +
+
+

Kernel Layer Report

+ +
+
+
+
+
+ + + +""" + # ------------------------------------------------------------------ # Callstack markdown export # ------------------------------------------------------------------ @@ -3099,7 +3843,8 @@ def __init__(self, trace_path: str, auto_fix_rocm: bool = True, model_info: bool = False, port: int = 8765, - callstack_md: Optional[str] = None): + callstack_md: Optional[str] = None, + layer_html: Optional[str] = None): self.trace_path = trace_path self.output_path = output_path self.detail_modules = detail_modules or [] @@ -3111,6 +3856,7 @@ def __init__(self, trace_path: str, self.model_info = model_info self.port = port self.callstack_md = callstack_md + self.layer_html = layer_html def run(self): import time as _time @@ -3364,6 +4110,15 @@ def _elapsed(): ) print(f" {_elapsed()} callstack markdown done") + if self.layer_html: + reporter.write_layer_html( + stats_list, self.layer_html, mode, + max_detail_modules=self.max_detail_modules, + detail_modules=self.detail_modules, + detail_instances=self.detail_instances, + ) + print(f" {_elapsed()} layer HTML done") + if self.model_info and xlsx_path: html_path = os.path.splitext(xlsx_path)[0] + "_module_tree.html" result = _generate_model_info_html(xlsx_path, html_path, @@ -3899,6 +4654,12 @@ def main(): "(e.g. Prefill_0 + Decode_0)") parser.add_argument("--no-rocm-fix", action="store_true", help="Disable automatic ROCm trace fix (hipGraphLaunch flow events)") + parser.add_argument("--layer-html", default=None, metavar="PATH", + help="Write a self-contained interactive HTML report with D3.js " + "charts: per-kernel avg/min/max bar chart across layer " + "instances, drag-and-drop category buckets, and a live " + "category bar chart. Opens directly in any browser. " + "Example: --layer-html report.html") parser.add_argument("--callstack-md", default=None, metavar="PATH", help="Write a Markdown file with per-kernel CPU→GPU call stacks " "for each representative module instance (same instances as " @@ -3937,12 +4698,17 @@ def main(): except ValueError: pass - # Resolve callstack-md path the same way as output + # Resolve relative output paths to same folder as trace file + trace_dir = os.path.dirname(os.path.abspath(args.trace_file)) + callstack_md = args.callstack_md if callstack_md and not os.path.isabs(callstack_md): - trace_dir = os.path.dirname(os.path.abspath(args.trace_file)) callstack_md = os.path.join(trace_dir, callstack_md) + layer_html = args.layer_html + if layer_html and not os.path.isabs(layer_html): + layer_html = os.path.join(trace_dir, layer_html) + try: analyzer = TraceModuleAnalyzer( trace_path=args.trace_file, @@ -3956,6 +4722,7 @@ def main(): model_info=args.model_info, port=args.port, callstack_md=callstack_md, + layer_html=layer_html, ) analyzer.run() except FileNotFoundError as e: