From 19892d4ac9c02786616f196bfb73c679d01c8539 Mon Sep 17 00:00:00 2001 From: Deval Shah Date: Tue, 18 Aug 2026 18:50:43 -0500 Subject: [PATCH 1/4] refactoring to add clear flow --- .../TraceUtils/split_inference/__init__.py | 37 + .../split_inference/execution_roots.py | 42 + .../split_inference/steady_state_window.py | 365 +++++++ .../split_inference/trace_extraction.py | 487 +++++++++ .../split_inference_trace_annotation.py | 943 +----------------- .../test_split_inference_trace_annotation.py | 80 -- 6 files changed, 951 insertions(+), 1003 deletions(-) create mode 100644 TraceLens/TraceUtils/split_inference/__init__.py create mode 100644 TraceLens/TraceUtils/split_inference/execution_roots.py create mode 100644 TraceLens/TraceUtils/split_inference/steady_state_window.py create mode 100644 TraceLens/TraceUtils/split_inference/trace_extraction.py diff --git a/TraceLens/TraceUtils/split_inference/__init__.py b/TraceLens/TraceUtils/split_inference/__init__.py new file mode 100644 index 000000000..f25444d9d --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/__init__.py @@ -0,0 +1,37 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Inference trace splitting: execution roots, steady-state windows, extraction.""" + +from .execution_roots import find_iteration_roots +from .steady_state_window import ( + compute_reference_pd_ratio, + find_steady_state_window, + identify_steady_state_regions, +) +from .trace_extraction import ( + divide_phases_and_save, + extract_and_save, + extract_iteration, + extract_phases_and_save, + get_filename, + parse_range, + preprocess_trace, +) + +__all__ = [ + "compute_reference_pd_ratio", + "divide_phases_and_save", + "extract_and_save", + "extract_iteration", + "extract_phases_and_save", + "find_iteration_roots", + "find_steady_state_window", + "get_filename", + "identify_steady_state_regions", + "parse_range", + "preprocess_trace", +] diff --git a/TraceLens/TraceUtils/split_inference/execution_roots.py b/TraceLens/TraceUtils/split_inference/execution_roots.py new file mode 100644 index 000000000..c8ba5daa5 --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/execution_roots.py @@ -0,0 +1,42 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Stage 1: find iteration execution roots in an inference trace.""" + +from typing import List, Optional + +from ..annotation_utils import ( + ITERATION_BACKUP_PATTERNS, + ITERATION_PATTERNS, + find_events_by_patterns, +) + + +def find_iteration_roots(events: List[dict]) -> Optional[List[dict]]: + """Return iteration-root events. + + Tries the primary annotation pattern first, then backup patterns, then + falls back to generic call-tree traversal via Trace2Tree. + """ + roots = find_events_by_patterns( + events, ITERATION_PATTERNS, label="execution steps (iteration)", verbose=True + ) + if len(roots) == 0: + print("No primary annotations found; falling back to backup patterns...") + roots = find_events_by_patterns( + events, + ITERATION_BACKUP_PATTERNS, + label="execution steps (iteration, backup)", + verbose=True, + ) + if len(roots) == 0: + print("No annotation patterns found; trying generic call-tree traversal...") + from ...Trace2Tree.inference_iteration_roots import ( + find_iteration_roots_generic, + ) + + roots = find_iteration_roots_generic(events) + return roots diff --git a/TraceLens/TraceUtils/split_inference/steady_state_window.py b/TraceLens/TraceUtils/split_inference/steady_state_window.py new file mode 100644 index 000000000..5d4c9c651 --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/steady_state_window.py @@ -0,0 +1,365 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Stage 2: steady-state region detection and window selection.""" + +import math +from statistics import mean +from typing import List, Optional, Tuple + +from ..annotation_utils import ( + has_context, + is_decode_only, + is_mixed, + iteration_details, +) + + +def identify_steady_state_regions( + iter_details: List[dict], num_steps: int +) -> Tuple[List[Tuple[int, int]], int]: + """Detect contiguous steady-state regions based on num_requests proximity to global max. + + Returns ``(regions, global_max)`` where ``regions`` is a list of + ``(start, end)`` index pairs and ``global_max`` is the peak concurrency + observed across all iterations. + """ + n = len(iter_details) + thresh = 0.1 if n >= num_steps else 0.2 + global_max = max(t["num_requests"] for t in iter_details) + + steady_state_started = False + steady_state_ended = False + prev_events_in_steady = 0 + start_index = 0 + regions = [] + + for i, t in enumerate(iter_details): + if abs(t["num_requests"] - global_max) <= max(1, thresh * global_max): + if not steady_state_started: + prev_events_in_steady += 1 + else: + if steady_state_started: + prev_events_in_steady -= 1 + + if prev_events_in_steady > 5 and not steady_state_started: + print(f"Steady state started at index {i - 5}") + steady_state_started = True + start_index = i - prev_events_in_steady + 1 + + if ( + prev_events_in_steady <= 0 + and steady_state_started + and not steady_state_ended + ): + print(f"Steady state ended at index {i}") + steady_state_ended = True + regions.append((start_index, i)) + steady_state_started = False + steady_state_ended = False + prev_events_in_steady = 0 + + if steady_state_started and not steady_state_ended: + regions.append((start_index, i)) + + print(f"Steady state regions: {regions}") + + if len(regions) == 0: + delta = min(n, max(8, num_steps - n)) + start = max(0, delta // 2) + end = max(start + 1, min(n, n - delta // 2)) + regions = [(start, end)] + print( + "Warning: no steady state region found; discarding initial/final iterations " + "and selecting middle region" + ) + + return regions, global_max + + +def compute_reference_pd_ratio( + regions: List[Tuple[int, int]], iter_details: List[dict] +) -> Tuple[Tuple[int, int], float, float]: + """ + Return the largest steady-state region, a reference PD ratio, and the + median PD ratio across all regions. + + The reference ratio starts as the PD/total ratio of the largest region. A + sanity check compares it to the median ratio across ALL regions. If the + largest region deviates by more than 50 % relative from the median, the + median is used instead and a warning is printed. + """ + region_stats = [] + total_steps = 0 + total_pd_steps = 0 + for s, e in regions: + window = iter_details[s:e] + total = len(window) + total_steps += total + pd_count = sum(1 for t in window if has_context(t)) + total_pd_steps += pd_count + ratio = pd_count / total if total > 0 else 0.0 + region_stats.append({"start": s, "end": e, "size": total, "pd_ratio": ratio}) + print( + f" Region [{s}, {e}): size={total}, " + f"prefilldecodemix_steps={pd_count}, prefilldecodemix_to_totalsteps_ratio={ratio:.3f}" + ) + + largest = max(region_stats, key=lambda x: x["size"]) + average_ratio = total_pd_steps / total_steps if total_steps > 0 else 0.0 + largest_window_ratio = largest["pd_ratio"] + print( + f"Reference prefilldecodemix_to_totalsteps_ratio={largest_window_ratio:.3f} (largest region [{largest['start']}, {largest['end']}), Average across all regions={average_ratio:.3f})" + ) + + return (largest["start"], largest["end"]), average_ratio, largest_window_ratio + + +def find_steady_state_window( + iteration_roots: List[dict], + num_steps: int, + steady_state_regions: List[Tuple[int, int]], + mode: str = "mixed", + CONC: Optional[int] = None, + OSL: Optional[float] = None, + R: Optional[float] = None, +) -> List[dict]: + """ + Find the best contiguous window of up to ``num_steps`` iterations. + + Parameters + ---------- + iteration_roots : list of iteration-root events + num_steps : requested window size + steady_state_regions : pre-computed steady-state region list as ``(start, end)`` + index pairs. Pass ``[(0, len(iteration_roots))]`` to treat the entire + slice as steady state. + mode : one of ``"mixed"``, ``"decode_only"``, ``"max_prefilldecode"`` + CONC : expected peak concurrency (number of concurrent requests). + If provided, a warning is printed when the observed peak in the trace + differs from this value. + OSL : average output sequence length (decode tokens per request). + Combined with ``R`` to derive the ideal PD ratio. + R : OSL window ratio in [0, 1]. The actual OSL per request is sampled from + ``[R * OSL, OSL]``, giving mean OSL = OSL * (1 + R) / 2. + + When ``CONC``, ``OSL``, and ``R`` are all provided the ideal PD ratio is + + ideal_pd_ratio = (CONC * 2) / (OSL * (1 + R)) + + and ``num_steps`` is automatically raised to ``ceil(1 / ideal_pd_ratio)`` + if it is too small to capture the true DO/PD distribution. + + Modes + ----- + ``"mixed"`` + Pick the sub-window whose pd_ratio is closest to the reference ratio + (ideal when available, otherwise largest-region / median sanity-checked). + Ties broken by highest average num_requests. + ``"decode_only"`` + Fewest-PD window: sub-window with lowest pd_ratio. + ``"max_prefilldecode"`` + Most-PD window: sub-window with highest pd_ratio. + """ + iter_details = iteration_details(iteration_roots) + regions = steady_state_regions + global_max = max(t["num_requests"] for t in iter_details) + + (largest_start, largest_end), reference_ratio, largest_window_ratio = ( + compute_reference_pd_ratio(regions, iter_details) + ) + + # --- Optional: CONC / OSL / R validation and ideal ratio override ---------- + ideal_pd_ratio: Optional[float] = None + + if CONC is not None and global_max != CONC: + print( + f"Warning: expected peak concurrency CONC={CONC} but the trace peak is " + f"global_max={global_max}. The trace may not contain requests at the " + f"intended concurrency level." + ) + + if CONC is not None and OSL is not None and R is not None: + if not (0.0 <= R <= 1.0): + print(f"Warning: R={R} is outside [0, 1]; clamping to valid range.") + R = max(0.0, min(1.0, R)) + mean_osl = OSL * (1.0 + R) / 2.0 + ideal_pd_ratio = (CONC * 2.0) / (OSL * (1.0 + R)) + print( + f"Ideal prefilldecodemix_to_totalsteps_ratio = (CONC={CONC} * 2) / (OSL={OSL} * (1 + R={R})) " + f"= {ideal_pd_ratio:.4f} [mean OSL = {mean_osl:.1f}]" + ) + + min_steps_for_ratio = math.ceil(1.0 / ideal_pd_ratio) + if num_steps < min_steps_for_ratio: + print( + f"Warning: --num-steps={num_steps} is too small to capture the true " + f"decode_only/prefilldecodemix distribution. At prefilldecodemix_to_totalsteps_ratio={ideal_pd_ratio:.4f} you need at " + f"least {min_steps_for_ratio} steps to see a representative mix. " + f"Raising num_steps to {min_steps_for_ratio}." + ) + num_steps = min_steps_for_ratio + else: + print(f"num_steps={num_steps} >= min required {min_steps_for_ratio} — OK.") + + # Ideal ratio overrides the empirical reference for the mixed mode + reference_ratio = ideal_pd_ratio + print( + f"Using ideal prefilldecodemix_to_totalsteps_ratio={ideal_pd_ratio:.4f} as reference (overrides empirical {reference_ratio:.4f})" + ) + print(f"\n --------------------------------") + # --------------------------------------------------------------------------- + + divider = max(1, min(int(num_steps / 2), 10)) + step = max(1, num_steps // divider) + + # Build candidate sub-windows from the largest region + candidates = [] + s, e = largest_start, largest_end + + def _count_mixed(window: List[dict]) -> int: + """Count truly-mixed steps (both context and generation requests > 0).""" + return sum(1 for t in window if is_mixed(t)) + + if (e - s) >= num_steps: + for s1 in range(s, e - num_steps + 1, step): + window = iter_details[s1 : s1 + num_steps] + pd_count = sum(1 for t in window if has_context(t)) + candidates.append( + { + "start": s1, + "end": s1 + num_steps, + "pd_count": pd_count, + "pd_ratio": pd_count / num_steps, + "mixed_count": _count_mixed(window), + "avg_requests": mean(t["num_requests"] for t in window), + } + ) + else: + # Region is smaller than num_steps — use the whole region + window = iter_details[s:e] + pd_count = sum(1 for t in window if has_context(t)) + candidates.append( + { + "start": s, + "end": e, + "pd_count": pd_count, + "pd_ratio": pd_count / len(window) if window else 0.0, + "mixed_count": _count_mixed(window), + "avg_requests": ( + mean(t["num_requests"] for t in window) if window else 0 + ), + } + ) + + if mode == "mixed": + # Prefer candidate windows that contain at least one prefill-bearing + # step (pure prefill OR truly mixed, i.e. context_requests > 0). Fall + # back to all candidates only when no window contains any prefill + # activity at all. + pd_candidates = [c for c in candidates if c["pd_count"] > 0] + if pd_candidates: + print( + f"[mixed] Filtering to {len(pd_candidates)}/{len(candidates)} " + f"candidate windows that contain at least one prefill or " + f"prefill-decode step." + ) + selection_pool = pd_candidates + else: + print( + "[mixed] No candidate window contains a prefill or prefill-decode " + "step; falling back to the full candidate set." + ) + selection_pool = candidates + + best = min( + selection_pool, + key=lambda c: (abs(c["pd_ratio"] - reference_ratio), -c["avg_requests"]), + ) + print( + f"[mixed] Selected window [{best['start']}, {best['end']}): " + f"prefilldecodemix_to_totalsteps_ratio={best['pd_ratio']:.3f} (target={reference_ratio:.3f}), " + f"avg_requests={best['avg_requests']:.1f}, " + f"pd_count={best['pd_count']}, mixed_count={best['mixed_count']}" + ) + + elif mode == "decode_only": + # Find the longest contiguous run of pure decode-only steps (active + # generation with no context requests) in the largest steady-state + # region, capped at num_steps. + do_runs: List[Tuple[int, int]] = [] # (start, end) in iter_details coords + run_start: Optional[int] = None + for idx in range(largest_start, largest_end): + if is_decode_only(iter_details[idx]): + if run_start is None: + run_start = idx + else: + if run_start is not None: + do_runs.append((run_start, idx)) + run_start = None + if run_start is not None: + do_runs.append((run_start, largest_end)) + + if do_runs: + longest = max(do_runs, key=lambda r: r[1] - r[0]) + run_s, run_e = longest + win_s = run_s + win_e = min(run_e, run_s + num_steps) + print( + f"[decode_only] Longest pure decode-only run: [{run_s}, {run_e}) " + f"({run_e - run_s} steps). " + f"Selected [{win_s}, {win_e}) ({win_e - win_s} steps, " + f"capped at num_steps={num_steps})." + ) + return iteration_roots[win_s:win_e] + else: + print( + "[decode_only] No pure decode-only run found in steady-state region; " + ) + return [] + + elif mode == "max_prefilldecode": + # Find the longest contiguous run of pure PD steps (no decode-only) in + # the largest steady-state region, capped at num_steps. + pd_runs: List[Tuple[int, int]] = [] # (start, end) in iter_details coords + run_start: Optional[int] = None + for idx in range(largest_start, largest_end): + if has_context(iter_details[idx]): + if run_start is None: + run_start = idx + else: + if run_start is not None: + pd_runs.append((run_start, idx)) + run_start = None + if run_start is not None: + pd_runs.append((run_start, largest_end)) + + if pd_runs: + # Pick the longest pure-PD run + longest = max(pd_runs, key=lambda r: r[1] - r[0]) + run_s, run_e = longest + # Cap to num_steps from the start of the run + win_s = run_s + win_e = min(run_e, run_s + num_steps) + print( + f"[max_prefilldecode] Longest pure prefilldecodemix run: [{run_s}, {run_e}) " + f"({run_e - run_s} steps). " + f"Selected [{win_s}, {win_e}) ({win_e - win_s} steps, " + f"capped at num_steps={num_steps})." + ) + return iteration_roots[win_s:win_e] + else: + print( + "[max_prefilldecode] No pure prefilldecodemix run found in steady-state " + ) + return [] + + else: + raise ValueError( + f"Unknown mode: {mode!r}. Use 'mixed', 'decode_only', or 'max_prefilldecode'." + ) + + return iteration_roots[best["start"] : best["end"]] diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py new file mode 100644 index 000000000..78f88dc44 --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -0,0 +1,487 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Stage 3: preprocess traces, extract iteration windows, and write output files.""" + +import gzip +import json +import os +import zipfile +from typing import List, Optional, Set, Tuple + +from tqdm import tqdm + +from ..annotation_utils import ( + ITERATION_BACKUP_PATTERNS, + ITERATION_PATTERNS, + find_phase_from_window, + has_context, + has_generation, + is_decode_only, + iteration_details, +) + +GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"] + + +def get_filename(filepath: str) -> dict: + """Load trace JSON from file (.json, .json.gz, or .zip).""" + print(f"Loading trace: {filepath}") + if filepath.endswith(".zip"): + with zipfile.ZipFile(filepath, "r") as zf: + # Find the JSON file inside the zip + json_files = [f for f in zf.namelist() if f.endswith(".json")] + if not json_files: + raise ValueError(f"No .json file found in {filepath}") + json_file = json_files[0] + print(f" Reading {json_file} from zip...") + return json_file + return filepath + + +def preprocess_trace(events: List[dict]): + gpu_corr_map = {} + flow_corr_map = {} + meta_events = [] + for e in tqdm(events): + ts = e.get("ts") + ph = e.get("ph") + cat = e.get("cat") + if ts is None: + meta_events.append(e) + continue + if ph in ("s", "f"): + corr = e.get("id") + if corr is not None: + flow_corr_map.setdefault(corr, []).append(e) + continue + if cat in GPU_EVENT_CATEGORIES: + corr = e.get("args", {}).get("correlation") + if corr is not None: + gpu_corr_map.setdefault(corr, []).append(e) + continue + return gpu_corr_map, flow_corr_map, meta_events + + +def extract_iteration( + iteration_roots: List[dict], + events: List[dict], + trace_json: dict, + gpu_corr_map: dict, + flow_corr_map: dict, + meta_events: List[dict], +) -> dict: + """Extract a single iteration trace.""" + + filtered_events = [] + gpu_dur = 0 + gpu_busy = 0 + num_gpu_events = 0 + batch_list = [] + + # Pre-index GPU and flow events by correlation id + + # Compute the global time window for all iteration roots + if not iteration_roots: + return trace_json.copy(), [], 0, 0, 0 + min_iter_ts = min(root.get("ts", 0) for root in iteration_roots) + max_iter_end = max( + root.get("ts", 0) + root.get("dur", 0) for root in iteration_roots + ) + # Collect all relevant tid/pid pairs + tid_pid_set = set((root.get("tid"), root.get("pid")) for root in iteration_roots) + + # Pre-filter all CPU events in the global window and by tid/pid + cpu_events = [] + for e in events: + ts = e.get("ts") + if ts is None: + continue + dur = e.get("dur") + if dur is None: + continue + e_end = ts + dur + e_tid = e.get("tid") + e_pid = e.get("pid") + if (e_tid, e_pid) in tid_pid_set and ( + min_iter_ts <= ts and e_end <= max_iter_end + ): + cpu_events.append(e) + + # For each iteration root, filter CPU events and collect correlation ids + for iteration_root in tqdm(iteration_roots): + start_time = [] + end_time = [] + iter_tid = iteration_root.get("tid") + iter_pid = iteration_root.get("pid") + iter_ts = iteration_root.get("ts", 0) + iter_end = iter_ts + iteration_root.get("dur", 0) + + correlation_ids: Set[int] = set() + + # CPU events: filter from pre-filtered list + for e in cpu_events: + ts = e.get("ts") + dur = e.get("dur") + e_end = ts + dur + e_tid = e.get("tid") + e_pid = e.get("pid") + if ( + e_tid == iter_tid + and e_pid == iter_pid + and (iter_ts <= ts and e_end <= iter_end) + ): + filtered_events.append(e) + corr = e.get("args", {}).get("correlation") + if corr is not None: + correlation_ids.add(corr) + + # Add matching flow events + for corr in correlation_ids: + for e in flow_corr_map.get(corr, []): + filtered_events.append(e) + # Add matching GPU events + for corr in correlation_ids: + for e in gpu_corr_map.get(corr, []): + filtered_events.append(e) + start_time.append(e.get("ts")) + end_time.append(e.get("ts") + e.get("dur")) + gpu_busy += e.get("dur") + num_gpu_events += 1 + gpu_dur += max(end_time) - min(start_time) if start_time else 0 + + # Add all meta events (no timestamp) + filtered_events.extend(meta_events) + + for e in tqdm(filtered_events): + if "vllm::unified_attention_with_output" in e.get( + "name", "" + ) or "sgl_kernel::sgl_per_token_group_quant_8bit" in e.get("name", ""): + dims = e.get("args", {}).get("Input Dims") + if dims and len(dims) > 0 and len(dims[0]) > 0: + batch_list.append(dims[0][0]) + # Create output trace + output = trace_json.copy() + output["traceEvents"] = filtered_events + return output, list(set(batch_list)), num_gpu_events, gpu_dur, gpu_busy + + +def parse_range(range_str: str, max_len: int) -> Tuple[int, int]: + """Parse a range string like '10:20' or 'all'.""" + if range_str == "all": + return 0, max_len + parts = range_str.split(":") + start = int(parts[0]) + end = int(parts[1]) if len(parts) > 1 else start + 1 + return start, min(end, max_len) + + +def extract_and_save( + roots: List[List[dict]], + events: List[dict], + trace_json: dict, + output_dir: str, + base_name: str, + prefix: str, + start: int, + end: int, + gpu_corr_map: dict, + flow_corr_map: dict, + meta_events: List[dict], + output_label: Optional[str] = None, +): + """Extract and save a range of iterations/dummy runs. + + If ``output_label`` is provided the output filename becomes + ``{output_label}_{name_append}_{base_name}.json.gz`` instead of the + default ``{base_name}_{prefix}_{idx}_{name_append}.json.gz``. + """ + extraction_summary = [] + # print(f"roots: {roots}") + if len(roots) == 0 or len(roots[0]) == 0: + print(f"No {prefix} events found in the specified range, skipping extraction") + return extraction_summary + selected = roots[start:end] + indices = range(start, end) + if len(selected) == 0: + print(f"No {prefix} events found in the specified range, skipping extraction") + return extraction_summary + for idx, root in zip(indices, selected): + iter_details = iteration_details(root) + iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = extract_iteration( + root, events, trace_json, gpu_corr_map, flow_corr_map, meta_events + ) + is_annotation = "annotation_iteration" in prefix + # Use the structured phase-aware name for any annotation extraction + # produced by the steady-state code paths (output_label is set), and + # for any multi-step annotation window. Single-step annotations from + # --store-single-iteration keep their literal step name. + is_structured = is_annotation and (output_label is not None or len(root) > 1) + + if is_structured or not is_annotation: + if len(batch_list) == len(iter_details): + for bs, iteration in zip(batch_list, iter_details): + iteration["batch_size"] = bs + + phase_details = find_phase_from_window(iter_details) + + if is_structured: + name_append = ( + f"prefill_{phase_details['num_prefill']}" + f"_prefilldecode_{phase_details['num_prefilldecode']}" + f"_decode_{phase_details['num_decode']}" + f"_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" + ) + elif is_annotation and len(root) == 1: + root_name = root[0]["name"] + is_known_annotation = any( + pat.match(root_name) + for pat in ITERATION_PATTERNS + ITERATION_BACKUP_PATTERNS + ) + if is_known_annotation: + name_append = ( + root_name.replace("/", "_") + .replace("(", "_") + .replace(")", "") + .replace(":", "") + .replace(" ", "_") + ) + else: + name_append = "" + else: + if len(batch_list) == len(iter_details): + name_append = f"batch{int(sum(batch_list)/len(batch_list))}_gpu{prefix}" + else: + name_append = f"batch_NA_gpu{prefix}" + + if output_label is not None: + out_path = os.path.join( + output_dir, f"{output_label}_{name_append}_{base_name}.json.gz" + ) + else: + suffix = f"_{name_append}" if name_append else "" + out_path = os.path.join( + output_dir, f"{base_name}_{prefix}_{idx}{suffix}.json.gz" + ) + with gzip.open(out_path, "wb") as f: + f.write(json.dumps(iter_trace).encode("utf-8")) + + print( + f" {prefix} {idx}: {len(iter_trace['traceEvents'])} events -> {out_path}" + ) + extraction_summary.append( + { + "idx": idx, + "output_path": out_path, + "event_count": len(iter_trace["traceEvents"]), + "num_gpu_events": num_gpu_events, + "gpu_duration": gpu_dur, + "gpu_busy_duration": gpu_busy, + "steps": iter_details, + "phase": phase_details, + } + ) + return extraction_summary + + +def extract_phases_and_save( + roots: List[List[dict]], + events: List[dict], + trace_json: dict, + output_dir: str, + base_name: str, + prefix: str, + start: int, + end: int, + gpu_corr_map: dict, + flow_corr_map: dict, + meta_events: List[dict], +): + """Extract and save a range of iterations/dummy runs.""" + extraction_summary = [] + + if "annotation_iteration" not in prefix: + print("phase extraction only supported for annotation iterations, skipping") + return extraction_summary + for root in roots: + iter_details = iteration_details(root) + prefilldecode_steps = [r for r, i in zip(root, iter_details) if has_context(i)] + decode_steps = [r for r, i in zip(root, iter_details) if is_decode_only(i)] + + if len(prefilldecode_steps) > 0: + iter_details = iteration_details(prefilldecode_steps) + phase_details = find_phase_from_window(iter_details) + + iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = ( + extract_iteration( + prefilldecode_steps, + events, + trace_json, + gpu_corr_map, + flow_corr_map, + meta_events, + ) + ) + name_append = f"prefilldecode_{phase_details['num_prefilldecode']}_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" + + out_path = os.path.join(output_dir, f"{name_append}_{base_name}.json.gz") + with gzip.open(out_path, "wb") as f: + f.write(json.dumps(iter_trace).encode("utf-8")) + + print(f" {prefix}: {len(iter_trace['traceEvents'])} events -> {out_path}") + extraction_summary.append( + { + "idx": 0, + "output_path": out_path, + "event_count": len(iter_trace["traceEvents"]), + "num_gpu_events": num_gpu_events, + "gpu_duration": gpu_dur, + "gpu_busy_duration": gpu_busy, + "steps": iter_details, + "phase": phase_details, + } + ) + if len(decode_steps) > 0: + iter_details = iteration_details(decode_steps) + phase_details = find_phase_from_window(iter_details) + iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = ( + extract_iteration( + decode_steps, + events, + trace_json, + gpu_corr_map, + flow_corr_map, + meta_events, + ) + ) + name_append = f"decode_{phase_details['num_decode']}_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" + + out_path = os.path.join(output_dir, f"{name_append}_{base_name}.json.gz") + with gzip.open(out_path, "wb") as f: + f.write(json.dumps(iter_trace).encode("utf-8")) + + print(f" {prefix}: {len(iter_trace['traceEvents'])} events -> {out_path}") + extraction_summary.append( + { + "idx": 0, + "output_path": out_path, + "event_count": len(iter_trace["traceEvents"]), + "num_gpu_events": num_gpu_events, + "gpu_duration": gpu_dur, + "gpu_busy_duration": gpu_busy, + "steps": iter_details, + "phase": phase_details, + } + ) + return extraction_summary + + +def divide_phases_and_save( + iteration_roots: List[dict], + events: List[dict], + trace_json: dict, + output_dir: str, + base_name: str, + gpu_corr_map: dict, + flow_corr_map: dict, + meta_events: List[dict], + steady_state_regions: List[Tuple[int, int]], +) -> List[dict]: + """ + Group contiguous steps of the same phase within steady-state regions and + save each contiguous run as a single trace file into one of two sub-folders: + + - ``{output_dir}/prefilldecodemix/`` — runs where every step has ``context_requests > 0`` + - ``{output_dir}/decode_only/`` — runs where every step has ``context_requests == 0`` + and ``generation_requests > 0`` + + A phase transition (PD → DO or DO → PD) always starts a new file. + + Parameters + ---------- + steady_state_regions + Pre-computed steady-state region list as ``(start, end)`` index pairs. + Pass ``[(0, len(iteration_roots))]`` to treat the entire slice as steady state. + """ + iter_details = iteration_details(iteration_roots) + regions = steady_state_regions + print(f"[divide-phases] Steady-state regions: {regions}") + + # Build an ordered list of (phase_label, root) for all steady-state steps + steady_steps: List[Tuple[str, dict]] = [] + for s, e in regions: + for idx in range(s, e): + detail = iter_details[idx] + root = iteration_roots[idx] + if has_context(detail): + steady_steps.append(("prefilldecodemix", root)) + elif has_generation(detail): + steady_steps.append(("decode_only", root)) + # steps that are neither (e.g. idle) are skipped + + # Group into contiguous runs of the same phase + runs: List[Tuple[str, List[dict]]] = [] # (phase, [roots]) + for phase, root in steady_steps: + if runs and runs[-1][0] == phase: + runs[-1][1].append(root) + else: + runs.append((phase, [root])) + + pd_count = sum(1 for p, _ in runs if p == "prefilldecodemix") + do_count = sum(1 for p, _ in runs if p == "decode_only") + total_pd_steps = sum(len(r) for p, r in runs if p == "prefilldecodemix") + total_do_steps = sum(len(r) for p, r in runs if p == "decode_only") + print( + f"\n[divide-phases] {pd_count} prefilldecodemix runs ({total_pd_steps} steps) and " + f"{do_count} decode_only runs ({total_do_steps} steps) across all steady-state regions." + ) + + pd_dir = os.path.join(output_dir, "prefilldecodemix") + do_dir = os.path.join(output_dir, "decode_only") + if pd_count: + os.makedirs(pd_dir, exist_ok=True) + if do_count: + os.makedirs(do_dir, exist_ok=True) + + extraction_summary = [] + pd_chunk_idx = 0 + do_chunk_idx = 0 + + for phase, chunk_roots in runs: + if phase == "prefilldecodemix": + out_dir = pd_dir + chunk_idx = pd_chunk_idx + pd_chunk_idx += 1 + else: + out_dir = do_dir + chunk_idx = do_chunk_idx + do_chunk_idx += 1 + + phase_details = find_phase_from_window(iteration_details(chunk_roots)) + name_append = ( + f"chunk{chunk_idx}_" + f"steps{len(chunk_roots)}_" + f"bs{phase_details['avg_bs']}_" + f"conc{phase_details['avg_conc']}" + ) + extraction_summary.extend( + extract_and_save( + [chunk_roots], + events, + trace_json, + out_dir, + base_name, + "annotation_iteration", + 0, + 1, + gpu_corr_map, + flow_corr_map, + meta_events, + output_label=f"{phase}_{name_append}", + ) + ) + + return extraction_summary diff --git a/TraceLens/TraceUtils/split_inference_trace_annotation.py b/TraceLens/TraceUtils/split_inference_trace_annotation.py index 10549b118..48706bcbe 100644 --- a/TraceLens/TraceUtils/split_inference_trace_annotation.py +++ b/TraceLens/TraceUtils/split_inference_trace_annotation.py @@ -11,7 +11,6 @@ - Individual execution iterations - Steady-state regions (representative execution windows) - Per-phase traces (prefill-decode vs decode-only) -- Dummy run traces (graph capture phases) This enables efficient performance analysis and comparison without processing massive tracefiles. @@ -33,12 +32,7 @@ 'N' - Single iteration N 'START:END' - Iterations START through END-1 - -d, --dummy Dummy run range (default: 'all'): - 'all' - All dummy runs - 'N' - Single dummy run N - 'START:END' - Dummy runs START through END-1 - - --store-single-iteration Store each iteration/dummy run as an individual file + --store-single-iteration Store each iteration as an individual file --find-steady-state Automatically detect steady-state and extract three representative contiguous windows (no idle gaps): @@ -115,15 +109,6 @@ ───────────────────────────────────────────────────────────────────────────── -5. EXTRACT DUMMY RUNS (Graph Capture Phases) - - $ python split_inference_trace_annotation.py trace.json.gz \\ - -o ./dummy_runs --store-single-iteration - - → One trace file per dummy run in ./dummy_runs/ - -───────────────────────────────────────────────────────────────────────────── - Generated outputs: ✓ Individual .json.gz trace files in output directory @@ -167,871 +152,36 @@ """ import argparse -import gzip import json -import math import os -import re -import zipfile -from typing import List, Set, Tuple, Optional -from statistics import mean +from typing import List, Tuple + +import pandas as pd + from ..util import DataLoader from .annotation_utils import ( - ITERATION_PATTERNS, ITERATION_BACKUP_PATTERNS, - find_events_by_patterns, - find_phase_from_window, - has_context, - has_generation, - is_decode_only, - is_mixed, + ITERATION_PATTERNS, iteration_details, ) -import pandas as pd - -# Try to use faster JSON parser (orjson is 2-10x faster than json) -from tqdm import tqdm - -GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"] - - -def get_filename(filepath: str) -> dict: - """Load trace JSON from file (.json, .json.gz, or .zip).""" - print(f"Loading trace: {filepath}") - if filepath.endswith(".zip"): - with zipfile.ZipFile(filepath, "r") as zf: - # Find the JSON file inside the zip - json_files = [f for f in zf.namelist() if f.endswith(".json")] - if not json_files: - raise ValueError(f"No .json file found in {filepath}") - json_file = json_files[0] - print(f" Reading {json_file} from zip...") - return json_file - return filepath - - -def find_iteration_roots(events: List[dict]) -> Optional[List[dict]]: - """Return iteration-root events. - - Tries the primary annotation pattern first, then backup patterns, then - falls back to generic call-tree traversal via Trace2Tree. - """ - roots = find_events_by_patterns( - events, ITERATION_PATTERNS, label="execution steps (iteration)", verbose=True - ) - if len(roots) == 0: - print("No primary annotations found; falling back to backup patterns...") - roots = find_events_by_patterns( - events, - ITERATION_BACKUP_PATTERNS, - label="execution steps (iteration, backup)", - verbose=True, - ) - if len(roots) == 0: - print("No annotation patterns found; trying generic call-tree traversal...") - from ..Trace2Tree.inference_iteration_roots import ( - find_iteration_roots_generic, - ) - - roots = find_iteration_roots_generic(events) - return roots - - -def preprocess_trace(events: List[dict]): - gpu_corr_map = {} - flow_corr_map = {} - meta_events = [] - for e in tqdm(events): - ts = e.get("ts") - ph = e.get("ph") - cat = e.get("cat") - if ts is None: - meta_events.append(e) - continue - if ph in ("s", "f"): - corr = e.get("id") - if corr is not None: - flow_corr_map.setdefault(corr, []).append(e) - continue - if cat in GPU_EVENT_CATEGORIES: - corr = e.get("args", {}).get("correlation") - if corr is not None: - gpu_corr_map.setdefault(corr, []).append(e) - continue - return gpu_corr_map, flow_corr_map, meta_events - - -def extract_iteration( - iteration_roots: List[dict], - events: List[dict], - trace_json: dict, - gpu_corr_map: dict, - flow_corr_map: dict, - meta_events: List[dict], -) -> dict: - """Extract a single iteration trace.""" - - filtered_events = [] - gpu_dur = 0 - gpu_busy = 0 - num_gpu_events = 0 - batch_list = [] - - # Pre-index GPU and flow events by correlation id - - # Compute the global time window for all iteration roots - if not iteration_roots: - return trace_json.copy(), [], 0, 0, 0 - min_iter_ts = min(root.get("ts", 0) for root in iteration_roots) - max_iter_end = max( - root.get("ts", 0) + root.get("dur", 0) for root in iteration_roots - ) - # Collect all relevant tid/pid pairs - tid_pid_set = set((root.get("tid"), root.get("pid")) for root in iteration_roots) - - # Pre-filter all CPU events in the global window and by tid/pid - cpu_events = [] - for e in events: - ts = e.get("ts") - if ts is None: - continue - dur = e.get("dur") - if dur is None: - continue - e_end = ts + dur - e_tid = e.get("tid") - e_pid = e.get("pid") - if (e_tid, e_pid) in tid_pid_set and ( - min_iter_ts <= ts and e_end <= max_iter_end - ): - cpu_events.append(e) - - # For each iteration root, filter CPU events and collect correlation ids - for iteration_root in tqdm(iteration_roots): - start_time = [] - end_time = [] - iter_tid = iteration_root.get("tid") - iter_pid = iteration_root.get("pid") - iter_ts = iteration_root.get("ts", 0) - iter_end = iter_ts + iteration_root.get("dur", 0) - - correlation_ids: Set[int] = set() - - # CPU events: filter from pre-filtered list - for e in cpu_events: - ts = e.get("ts") - dur = e.get("dur") - e_end = ts + dur - e_tid = e.get("tid") - e_pid = e.get("pid") - if ( - e_tid == iter_tid - and e_pid == iter_pid - and (iter_ts <= ts and e_end <= iter_end) - ): - filtered_events.append(e) - corr = e.get("args", {}).get("correlation") - if corr is not None: - correlation_ids.add(corr) - - # Add matching flow events - for corr in correlation_ids: - for e in flow_corr_map.get(corr, []): - filtered_events.append(e) - # Add matching GPU events - for corr in correlation_ids: - for e in gpu_corr_map.get(corr, []): - filtered_events.append(e) - start_time.append(e.get("ts")) - end_time.append(e.get("ts") + e.get("dur")) - gpu_busy += e.get("dur") - num_gpu_events += 1 - gpu_dur += max(end_time) - min(start_time) if start_time else 0 - - # Add all meta events (no timestamp) - filtered_events.extend(meta_events) - - for e in tqdm(filtered_events): - if "vllm::unified_attention_with_output" in e.get( - "name", "" - ) or "sgl_kernel::sgl_per_token_group_quant_8bit" in e.get("name", ""): - dims = e.get("args", {}).get("Input Dims") - if dims and len(dims) > 0 and len(dims[0]) > 0: - batch_list.append(dims[0][0]) - # Create output trace - output = trace_json.copy() - output["traceEvents"] = filtered_events - return output, list(set(batch_list)), num_gpu_events, gpu_dur, gpu_busy - - -def parse_range(range_str: str, max_len: int) -> Tuple[int, int]: - """Parse a range string like '10:20' or 'all'.""" - if range_str == "all": - return 0, max_len - parts = range_str.split(":") - start = int(parts[0]) - end = int(parts[1]) if len(parts) > 1 else start + 1 - return start, min(end, max_len) - - -def extract_and_save( - roots: List[List[dict]], - events: List[dict], - trace_json: dict, - output_dir: str, - base_name: str, - prefix: str, - start: int, - end: int, - gpu_corr_map: dict, - flow_corr_map: dict, - meta_events: List[dict], - output_label: Optional[str] = None, -): - """Extract and save a range of iterations/dummy runs. - - If ``output_label`` is provided the output filename becomes - ``{output_label}_{name_append}_{base_name}.json.gz`` instead of the - default ``{base_name}_{prefix}_{idx}_{name_append}.json.gz``. - """ - extraction_summary = [] - # print(f"roots: {roots}") - if len(roots) == 0 or len(roots[0]) == 0: - print(f"No {prefix} events found in the specified range, skipping extraction") - return extraction_summary - selected = roots[start:end] - indices = range(start, end) - if len(selected) == 0: - print(f"No {prefix} events found in the specified range, skipping extraction") - return extraction_summary - for idx, root in zip(indices, selected): - iter_details = iteration_details(root) - iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = extract_iteration( - root, events, trace_json, gpu_corr_map, flow_corr_map, meta_events - ) - is_annotation = "annotation_iteration" in prefix - # Use the structured phase-aware name for any annotation extraction - # produced by the steady-state code paths (output_label is set), and - # for any multi-step annotation window. Single-step annotations from - # --store-single-iteration keep their literal step name. - is_structured = is_annotation and (output_label is not None or len(root) > 1) - - if is_structured or not is_annotation: - if len(batch_list) == len(iter_details): - for bs, iteration in zip(batch_list, iter_details): - iteration["batch_size"] = bs - - phase_details = find_phase_from_window(iter_details) - - if is_structured: - name_append = ( - f"prefill_{phase_details['num_prefill']}" - f"_prefilldecode_{phase_details['num_prefilldecode']}" - f"_decode_{phase_details['num_decode']}" - f"_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" - ) - elif is_annotation and len(root) == 1: - root_name = root[0]["name"] - is_known_annotation = any( - pat.match(root_name) - for pat in ITERATION_PATTERNS + ITERATION_BACKUP_PATTERNS - ) - if is_known_annotation: - name_append = ( - root_name.replace("/", "_") - .replace("(", "_") - .replace(")", "") - .replace(":", "") - .replace(" ", "_") - ) - else: - name_append = "" - else: - if len(batch_list) == len(iter_details): - name_append = f"batch{int(sum(batch_list)/len(batch_list))}_gpu{prefix}" - else: - name_append = f"batch_NA_gpu{prefix}" - - if output_label is not None: - out_path = os.path.join( - output_dir, f"{output_label}_{name_append}_{base_name}.json.gz" - ) - else: - suffix = f"_{name_append}" if name_append else "" - out_path = os.path.join( - output_dir, f"{base_name}_{prefix}_{idx}{suffix}.json.gz" - ) - with gzip.open(out_path, "wb") as f: - f.write(json.dumps(iter_trace).encode("utf-8")) - - print( - f" {prefix} {idx}: {len(iter_trace['traceEvents'])} events -> {out_path}" - ) - extraction_summary.append( - { - "idx": idx, - "output_path": out_path, - "event_count": len(iter_trace["traceEvents"]), - "num_gpu_events": num_gpu_events, - "gpu_duration": gpu_dur, - "gpu_busy_duration": gpu_busy, - "steps": iter_details, - "phase": phase_details, - } - ) - return extraction_summary - - -def extract_phases_and_save( - roots: List[List[dict]], - events: List[dict], - trace_json: dict, - output_dir: str, - base_name: str, - prefix: str, - start: int, - end: int, - gpu_corr_map: dict, - flow_corr_map: dict, - meta_events: List[dict], -): - """Extract and save a range of iterations/dummy runs.""" - extraction_summary = [] - - if "annotation_iteration" not in prefix: - print("phase extraction only supported for annotation iterations, skipping") - return extraction_summary - for root in roots: - iter_details = iteration_details(root) - prefilldecode_steps = [r for r, i in zip(root, iter_details) if has_context(i)] - decode_steps = [r for r, i in zip(root, iter_details) if is_decode_only(i)] - - if len(prefilldecode_steps) > 0: - iter_details = iteration_details(prefilldecode_steps) - phase_details = find_phase_from_window(iter_details) - - iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = ( - extract_iteration( - prefilldecode_steps, - events, - trace_json, - gpu_corr_map, - flow_corr_map, - meta_events, - ) - ) - name_append = f"prefilldecode_{phase_details['num_prefilldecode']}_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" - - out_path = os.path.join(output_dir, f"{name_append}_{base_name}.json.gz") - with gzip.open(out_path, "wb") as f: - f.write(json.dumps(iter_trace).encode("utf-8")) - - print(f" {prefix}: {len(iter_trace['traceEvents'])} events -> {out_path}") - extraction_summary.append( - { - "idx": 0, - "output_path": out_path, - "event_count": len(iter_trace["traceEvents"]), - "num_gpu_events": num_gpu_events, - "gpu_duration": gpu_dur, - "gpu_busy_duration": gpu_busy, - "steps": iter_details, - "phase": phase_details, - } - ) - if len(decode_steps) > 0: - iter_details = iteration_details(decode_steps) - phase_details = find_phase_from_window(iter_details) - iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = ( - extract_iteration( - decode_steps, - events, - trace_json, - gpu_corr_map, - flow_corr_map, - meta_events, - ) - ) - name_append = f"decode_{phase_details['num_decode']}_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" - - out_path = os.path.join(output_dir, f"{name_append}_{base_name}.json.gz") - with gzip.open(out_path, "wb") as f: - f.write(json.dumps(iter_trace).encode("utf-8")) - - print(f" {prefix}: {len(iter_trace['traceEvents'])} events -> {out_path}") - extraction_summary.append( - { - "idx": 0, - "output_path": out_path, - "event_count": len(iter_trace["traceEvents"]), - "num_gpu_events": num_gpu_events, - "gpu_duration": gpu_dur, - "gpu_busy_duration": gpu_busy, - "steps": iter_details, - "phase": phase_details, - } - ) - return extraction_summary - - -def identify_steady_state_regions( - iter_details: List[dict], num_steps: int -) -> Tuple[List[Tuple[int, int]], int]: - """Detect contiguous steady-state regions based on num_requests proximity to global max. - - Returns ``(regions, global_max)`` where ``regions`` is a list of - ``(start, end)`` index pairs and ``global_max`` is the peak concurrency - observed across all iterations. - """ - n = len(iter_details) - thresh = 0.1 if n >= num_steps else 0.2 - global_max = max(t["num_requests"] for t in iter_details) - - steady_state_started = False - steady_state_ended = False - prev_events_in_steady = 0 - start_index = 0 - regions = [] - - for i, t in enumerate(iter_details): - if abs(t["num_requests"] - global_max) <= max(1, thresh * global_max): - if not steady_state_started: - prev_events_in_steady += 1 - else: - if steady_state_started: - prev_events_in_steady -= 1 - - if prev_events_in_steady > 5 and not steady_state_started: - print(f"Steady state started at index {i - 5}") - steady_state_started = True - start_index = i - prev_events_in_steady + 1 - - if ( - prev_events_in_steady <= 0 - and steady_state_started - and not steady_state_ended - ): - print(f"Steady state ended at index {i}") - steady_state_ended = True - regions.append((start_index, i)) - steady_state_started = False - steady_state_ended = False - prev_events_in_steady = 0 - - if steady_state_started and not steady_state_ended: - regions.append((start_index, i)) - - print(f"Steady state regions: {regions}") - - if len(regions) == 0: - delta = min(n, max(8, num_steps - n)) - start = max(0, delta // 2) - end = max(start + 1, min(n, n - delta // 2)) - regions = [(start, end)] - print( - "Warning: no steady state region found; discarding initial/final iterations " - "and selecting middle region" - ) - - return regions, global_max - - -def compute_reference_pd_ratio( - regions: List[Tuple[int, int]], iter_details: List[dict] -) -> Tuple[Tuple[int, int], float, float]: - """ - Return the largest steady-state region, a reference PD ratio, and the - median PD ratio across all regions. - - The reference ratio starts as the PD/total ratio of the largest region. A - sanity check compares it to the median ratio across ALL regions. If the - largest region deviates by more than 50 % relative from the median, the - median is used instead and a warning is printed. - """ - region_stats = [] - total_steps = 0 - total_pd_steps = 0 - for s, e in regions: - window = iter_details[s:e] - total = len(window) - total_steps += total - pd_count = sum(1 for t in window if has_context(t)) - total_pd_steps += pd_count - ratio = pd_count / total if total > 0 else 0.0 - region_stats.append({"start": s, "end": e, "size": total, "pd_ratio": ratio}) - print( - f" Region [{s}, {e}): size={total}, " - f"prefilldecodemix_steps={pd_count}, prefilldecodemix_to_totalsteps_ratio={ratio:.3f}" - ) - - largest = max(region_stats, key=lambda x: x["size"]) - average_ratio = total_pd_steps / total_steps if total_steps > 0 else 0.0 - largest_window_ratio = largest["pd_ratio"] - print( - f"Reference prefilldecodemix_to_totalsteps_ratio={largest_window_ratio:.3f} (largest region [{largest['start']}, {largest['end']}), Average across all regions={average_ratio:.3f})" - ) - - return (largest["start"], largest["end"]), average_ratio, largest_window_ratio - - -def divide_phases_and_save( - iteration_roots: List[dict], - events: List[dict], - trace_json: dict, - output_dir: str, - base_name: str, - gpu_corr_map: dict, - flow_corr_map: dict, - meta_events: List[dict], - steady_state_regions: List[Tuple[int, int]], -) -> List[dict]: - """ - Group contiguous steps of the same phase within steady-state regions and - save each contiguous run as a single trace file into one of two sub-folders: - - - ``{output_dir}/prefilldecodemix/`` — runs where every step has ``context_requests > 0`` - - ``{output_dir}/decode_only/`` — runs where every step has ``context_requests == 0`` - and ``generation_requests > 0`` - - A phase transition (PD → DO or DO → PD) always starts a new file. - - Parameters - ---------- - steady_state_regions - Pre-computed steady-state region list as ``(start, end)`` index pairs. - Pass ``[(0, len(iteration_roots))]`` to treat the entire slice as steady state. - """ - iter_details = iteration_details(iteration_roots) - regions = steady_state_regions - print(f"[divide-phases] Steady-state regions: {regions}") - - # Build an ordered list of (phase_label, root) for all steady-state steps - steady_steps: List[Tuple[str, dict]] = [] - for s, e in regions: - for idx in range(s, e): - detail = iter_details[idx] - root = iteration_roots[idx] - if has_context(detail): - steady_steps.append(("prefilldecodemix", root)) - elif has_generation(detail): - steady_steps.append(("decode_only", root)) - # steps that are neither (e.g. idle) are skipped - - # Group into contiguous runs of the same phase - runs: List[Tuple[str, List[dict]]] = [] # (phase, [roots]) - for phase, root in steady_steps: - if runs and runs[-1][0] == phase: - runs[-1][1].append(root) - else: - runs.append((phase, [root])) - - pd_count = sum(1 for p, _ in runs if p == "prefilldecodemix") - do_count = sum(1 for p, _ in runs if p == "decode_only") - total_pd_steps = sum(len(r) for p, r in runs if p == "prefilldecodemix") - total_do_steps = sum(len(r) for p, r in runs if p == "decode_only") - print( - f"\n[divide-phases] {pd_count} prefilldecodemix runs ({total_pd_steps} steps) and " - f"{do_count} decode_only runs ({total_do_steps} steps) across all steady-state regions." - ) - - pd_dir = os.path.join(output_dir, "prefilldecodemix") - do_dir = os.path.join(output_dir, "decode_only") - if pd_count: - os.makedirs(pd_dir, exist_ok=True) - if do_count: - os.makedirs(do_dir, exist_ok=True) - - extraction_summary = [] - pd_chunk_idx = 0 - do_chunk_idx = 0 - - for phase, chunk_roots in runs: - if phase == "prefilldecodemix": - out_dir = pd_dir - chunk_idx = pd_chunk_idx - pd_chunk_idx += 1 - else: - out_dir = do_dir - chunk_idx = do_chunk_idx - do_chunk_idx += 1 - - phase_details = find_phase_from_window(iteration_details(chunk_roots)) - name_append = ( - f"chunk{chunk_idx}_" - f"steps{len(chunk_roots)}_" - f"bs{phase_details['avg_bs']}_" - f"conc{phase_details['avg_conc']}" - ) - extraction_summary.extend( - extract_and_save( - [chunk_roots], - events, - trace_json, - out_dir, - base_name, - "annotation_iteration", - 0, - 1, - gpu_corr_map, - flow_corr_map, - meta_events, - output_label=f"{phase}_{name_append}", - ) - ) - - return extraction_summary - - -def find_steady_state_window( - iteration_roots: List[dict], - num_steps: int, - steady_state_regions: List[Tuple[int, int]], - mode: str = "mixed", - CONC: Optional[int] = None, - OSL: Optional[float] = None, - R: Optional[float] = None, -) -> List[dict]: - """ - Find the best contiguous window of up to ``num_steps`` iterations. - - Parameters - ---------- - iteration_roots : list of iteration-root events - num_steps : requested window size - steady_state_regions : pre-computed steady-state region list as ``(start, end)`` - index pairs. Pass ``[(0, len(iteration_roots))]`` to treat the entire - slice as steady state. - mode : one of ``"mixed"``, ``"decode_only"``, ``"max_prefilldecode"`` - CONC : expected peak concurrency (number of concurrent requests). - If provided, a warning is printed when the observed peak in the trace - differs from this value. - OSL : average output sequence length (decode tokens per request). - Combined with ``R`` to derive the ideal PD ratio. - R : OSL window ratio in [0, 1]. The actual OSL per request is sampled from - ``[R * OSL, OSL]``, giving mean OSL = OSL * (1 + R) / 2. - - When ``CONC``, ``OSL``, and ``R`` are all provided the ideal PD ratio is - - ideal_pd_ratio = (CONC * 2) / (OSL * (1 + R)) - - and ``num_steps`` is automatically raised to ``ceil(1 / ideal_pd_ratio)`` - if it is too small to capture the true DO/PD distribution. - - Modes - ----- - ``"mixed"`` - Pick the sub-window whose pd_ratio is closest to the reference ratio - (ideal when available, otherwise largest-region / median sanity-checked). - Ties broken by highest average num_requests. - ``"decode_only"`` - Fewest-PD window: sub-window with lowest pd_ratio. - ``"max_prefilldecode"`` - Most-PD window: sub-window with highest pd_ratio. - """ - iter_details = iteration_details(iteration_roots) - regions = steady_state_regions - global_max = max(t["num_requests"] for t in iter_details) - - (largest_start, largest_end), reference_ratio, largest_window_ratio = ( - compute_reference_pd_ratio(regions, iter_details) - ) - - # --- Optional: CONC / OSL / R validation and ideal ratio override ---------- - ideal_pd_ratio: Optional[float] = None - - if CONC is not None and global_max != CONC: - print( - f"Warning: expected peak concurrency CONC={CONC} but the trace peak is " - f"global_max={global_max}. The trace may not contain requests at the " - f"intended concurrency level." - ) - - if CONC is not None and OSL is not None and R is not None: - if not (0.0 <= R <= 1.0): - print(f"Warning: R={R} is outside [0, 1]; clamping to valid range.") - R = max(0.0, min(1.0, R)) - mean_osl = OSL * (1.0 + R) / 2.0 - ideal_pd_ratio = (CONC * 2.0) / (OSL * (1.0 + R)) - print( - f"Ideal prefilldecodemix_to_totalsteps_ratio = (CONC={CONC} * 2) / (OSL={OSL} * (1 + R={R})) " - f"= {ideal_pd_ratio:.4f} [mean OSL = {mean_osl:.1f}]" - ) - - min_steps_for_ratio = math.ceil(1.0 / ideal_pd_ratio) - if num_steps < min_steps_for_ratio: - print( - f"Warning: --num-steps={num_steps} is too small to capture the true " - f"decode_only/prefilldecodemix distribution. At prefilldecodemix_to_totalsteps_ratio={ideal_pd_ratio:.4f} you need at " - f"least {min_steps_for_ratio} steps to see a representative mix. " - f"Raising num_steps to {min_steps_for_ratio}." - ) - num_steps = min_steps_for_ratio - else: - print(f"num_steps={num_steps} >= min required {min_steps_for_ratio} — OK.") - - # Ideal ratio overrides the empirical reference for the mixed mode - reference_ratio = ideal_pd_ratio - print( - f"Using ideal prefilldecodemix_to_totalsteps_ratio={ideal_pd_ratio:.4f} as reference (overrides empirical {reference_ratio:.4f})" - ) - print(f"\n --------------------------------") - # --------------------------------------------------------------------------- - - divider = max(1, min(int(num_steps / 2), 10)) - step = max(1, num_steps // divider) - - # Build candidate sub-windows from the largest region - candidates = [] - s, e = largest_start, largest_end - - def _count_mixed(window: List[dict]) -> int: - """Count truly-mixed steps (both context and generation requests > 0).""" - return sum(1 for t in window if is_mixed(t)) - - if (e - s) >= num_steps: - for s1 in range(s, e - num_steps + 1, step): - window = iter_details[s1 : s1 + num_steps] - pd_count = sum(1 for t in window if has_context(t)) - candidates.append( - { - "start": s1, - "end": s1 + num_steps, - "pd_count": pd_count, - "pd_ratio": pd_count / num_steps, - "mixed_count": _count_mixed(window), - "avg_requests": mean(t["num_requests"] for t in window), - } - ) - else: - # Region is smaller than num_steps — use the whole region - window = iter_details[s:e] - pd_count = sum(1 for t in window if has_context(t)) - candidates.append( - { - "start": s, - "end": e, - "pd_count": pd_count, - "pd_ratio": pd_count / len(window) if window else 0.0, - "mixed_count": _count_mixed(window), - "avg_requests": ( - mean(t["num_requests"] for t in window) if window else 0 - ), - } - ) - - if mode == "mixed": - # Prefer candidate windows that contain at least one prefill-bearing - # step (pure prefill OR truly mixed, i.e. context_requests > 0). Fall - # back to all candidates only when no window contains any prefill - # activity at all. - pd_candidates = [c for c in candidates if c["pd_count"] > 0] - if pd_candidates: - print( - f"[mixed] Filtering to {len(pd_candidates)}/{len(candidates)} " - f"candidate windows that contain at least one prefill or " - f"prefill-decode step." - ) - selection_pool = pd_candidates - else: - print( - "[mixed] No candidate window contains a prefill or prefill-decode " - "step; falling back to the full candidate set." - ) - selection_pool = candidates - - best = min( - selection_pool, - key=lambda c: (abs(c["pd_ratio"] - reference_ratio), -c["avg_requests"]), - ) - print( - f"[mixed] Selected window [{best['start']}, {best['end']}): " - f"prefilldecodemix_to_totalsteps_ratio={best['pd_ratio']:.3f} (target={reference_ratio:.3f}), " - f"avg_requests={best['avg_requests']:.1f}, " - f"pd_count={best['pd_count']}, mixed_count={best['mixed_count']}" - ) - - elif mode == "decode_only": - # Find the longest contiguous run of pure decode-only steps (active - # generation with no context requests) in the largest steady-state - # region, capped at num_steps. - do_runs: List[Tuple[int, int]] = [] # (start, end) in iter_details coords - run_start: Optional[int] = None - for idx in range(largest_start, largest_end): - if is_decode_only(iter_details[idx]): - if run_start is None: - run_start = idx - else: - if run_start is not None: - do_runs.append((run_start, idx)) - run_start = None - if run_start is not None: - do_runs.append((run_start, largest_end)) - - if do_runs: - longest = max(do_runs, key=lambda r: r[1] - r[0]) - run_s, run_e = longest - win_s = run_s - win_e = min(run_e, run_s + num_steps) - print( - f"[decode_only] Longest pure decode-only run: [{run_s}, {run_e}) " - f"({run_e - run_s} steps). " - f"Selected [{win_s}, {win_e}) ({win_e - win_s} steps, " - f"capped at num_steps={num_steps})." - ) - return iteration_roots[win_s:win_e] - else: - print( - "[decode_only] No pure decode-only run found in steady-state region; " - ) - return [] - - elif mode == "max_prefilldecode": - # Find the longest contiguous run of pure PD steps (no decode-only) in - # the largest steady-state region, capped at num_steps. - pd_runs: List[Tuple[int, int]] = [] # (start, end) in iter_details coords - run_start: Optional[int] = None - for idx in range(largest_start, largest_end): - if has_context(iter_details[idx]): - if run_start is None: - run_start = idx - else: - if run_start is not None: - pd_runs.append((run_start, idx)) - run_start = None - if run_start is not None: - pd_runs.append((run_start, largest_end)) - - if pd_runs: - # Pick the longest pure-PD run - longest = max(pd_runs, key=lambda r: r[1] - r[0]) - run_s, run_e = longest - # Cap to num_steps from the start of the run - win_s = run_s - win_e = min(run_e, run_s + num_steps) - print( - f"[max_prefilldecode] Longest pure prefilldecodemix run: [{run_s}, {run_e}) " - f"({run_e - run_s} steps). " - f"Selected [{win_s}, {win_e}) ({win_e - win_s} steps, " - f"capped at num_steps={num_steps})." - ) - return iteration_roots[win_s:win_e] - else: - print( - "[max_prefilldecode] No pure prefilldecodemix run found in steady-state " - ) - return [] - - else: - raise ValueError( - f"Unknown mode: {mode!r}. Use 'mixed', 'decode_only', or 'max_prefilldecode'." - ) - - return iteration_roots[best["start"] : best["end"]] +from .split_inference import ( + compute_reference_pd_ratio, + divide_phases_and_save, + extract_and_save, + extract_iteration, + extract_phases_and_save, + find_iteration_roots, + find_steady_state_window, + get_filename, + identify_steady_state_regions, + parse_range, + preprocess_trace, +) def main(): parser = argparse.ArgumentParser( - description="Split vLLM trace into per-iteration or per-dummy-run traces" + description="Split vLLM trace into per-iteration traces" ) parser.add_argument("trace_path", help="Path to trace file (.json or .json.gz)") @@ -1042,12 +192,6 @@ def main(): default="all", help="Iteration range: 'all', single index '50', or range '10:20'", ) - parser.add_argument( - "--dummy", - "-d", - default="all", - help="Dummy run range: 'all', single index '0', or range '0:2'", - ) parser.add_argument( "--store-single-iteration", action="store_true", @@ -1107,12 +251,6 @@ def main(): ) args = parser.parse_args() execution_details = [] - RUNTIME_EVENT_PATTERN = [ - re.compile(r"vllm/v1/worker/gpu_model_runner\.py\(\d+\): _dummy_run"), - re.compile( - r"/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py\(\d+\): _capture_graph" - ), - ] # Load trace trace_json = DataLoader.load_data(get_filename(args.trace_path)) @@ -1120,11 +258,7 @@ def main(): gpu_corr_map, flow_corr_map, meta_events = preprocess_trace(events) print(f"Loaded {len(events)} events") - # Find iterations and dummy runs iteration_roots = find_iteration_roots(events) - dummy_roots = find_events_by_patterns( - events, RUNTIME_EVENT_PATTERN, cat=None, label="_dummy_run", verbose=True - ) # Create output directory os.makedirs(args.output_dir, exist_ok=True) @@ -1239,43 +373,6 @@ def main(): [pd_roots], *_extract_args, output_label="prefilldecode_steady_state" ) execution_details.extend(temp_execution_details) - # Extract dummy runs - if args.dummy and dummy_roots: - start, end = parse_range(args.dummy, len(dummy_roots)) - if args.store_single_iteration: - print(f"\nExtracting dummy runs {start} to {end-1}...") - temp_execution_details = extract_and_save( - [[root] for root in dummy_roots], - events, - trace_json, - args.output_dir, - base_name, - "run_iteration", - start, - end, - gpu_corr_map, - flow_corr_map, - meta_events, - ) - execution_details.extend(temp_execution_details) - if args.dummy != "all": - print(f"\nExtracting dummy runs {start} to {end-1}...") - temp_execution_details = extract_and_save( - [dummy_roots[start:end]], - events, - trace_json, - args.output_dir, - base_name, - "run_iteration", - 0, - 1, - gpu_corr_map, - flow_corr_map, - meta_events, - ) - execution_details.extend(temp_execution_details) - if args.find_steady_state: - print(f"finding steady state without annotations not supported") print(f"\nDone! Extracted {len(execution_details)} traces to {args.output_dir}") if len(execution_details) > 0: diff --git a/tests/test_split_inference_trace_annotation.py b/tests/test_split_inference_trace_annotation.py index a356cfe9e..26a12cd1a 100644 --- a/tests/test_split_inference_trace_annotation.py +++ b/tests/test_split_inference_trace_annotation.py @@ -963,48 +963,6 @@ def test_main_explicit_iteration_range(tmp_path): assert out_dir.exists() -def test_main_dummy_runs_and_steady_state_message(tmp_path, capsys): - dummy_name = "vllm/v1/worker/gpu_model_runner.py(99): _dummy_run" - trace = { - "traceEvents": [ - { - "name": dummy_name, - "cat": "cpu_op", - "ph": "X", - "ts": 0, - "dur": 100, - "tid": 1, - "pid": 1, - "args": {}, - } - ], - "schemaVersion": 1, - } - trace_path = tmp_path / "dummy_trace.json" - trace_path.write_text(json.dumps(trace)) - out_dir = tmp_path / "dummy" - - old_argv = sys.argv - sys.argv = [ - "split_inference_trace_annotation", - str(trace_path), - "-o", - str(out_dir), - "--dummy", - "0:1", - "--find-steady-state", - ] - try: - split.main() - finally: - sys.argv = old_argv - - assert ( - "finding steady state without annotations not supported" - in capsys.readouterr().out - ) - - class TestCaptureMergeAndMoe: def test_align_streams_multistream(self): graph = [ @@ -1040,44 +998,6 @@ def test_moe_sort_scatter_missing_kernel_details(self): assert model.bytes() is None or model.bytes() >= 0 -class TestSplitAnnotationDummyPhase10: - def test_main_dummy_store_single_iteration(self, tmp_path): - - dummy_name = "vllm/v1/worker/gpu_model_runner.py(99): _dummy_run" - trace = {"traceEvents": [], "schemaVersion": 1} - for i in range(5): - trace["traceEvents"].append( - { - "name": dummy_name, - "cat": "user_annotation", - "ph": "X", - "ts": 1000 + i * 1000, - "dur": 100, - "tid": 10, - "pid": 1, - "args": {}, - } - ) - trace_path = tmp_path / "trace.json" - trace_path.write_text(json.dumps(trace)) - out_dir = tmp_path / "out" - old_argv = sys.argv - sys.argv = [ - "split_inference_trace_annotation", - str(trace_path), - "--output-dir", - str(out_dir), - "--dummy", - "1:3", - "--store-single-iteration", - ] - try: - split.main() - finally: - sys.argv = old_argv - assert (out_dir / "execution_details.json").exists() - - class TestCaptureMergeDeep: def test_find_closest_batch_size_and_execution_details(self): assert find_closest_batch_size(128, [64, 256, 512]) == 256 From cc6ea1d9ec7b5305ff62ab125563e914e2f57adc Mon Sep 17 00:00:00 2001 From: Deval Shah Date: Tue, 18 Aug 2026 19:00:45 -0500 Subject: [PATCH 2/4] Reformatting --- .../split_inference/execution_roots.py | 4 +- .../split_inference/steady_state_window.py | 37 ++++++------ .../split_inference/trace_extraction.py | 59 +++++++++---------- .../split_inference_trace_annotation.py | 15 ++--- .../inference_analysis/build_docker_sglang.sh | 12 +++- 5 files changed, 67 insertions(+), 60 deletions(-) diff --git a/TraceLens/TraceUtils/split_inference/execution_roots.py b/TraceLens/TraceUtils/split_inference/execution_roots.py index c8ba5daa5..568ea7ac0 100644 --- a/TraceLens/TraceUtils/split_inference/execution_roots.py +++ b/TraceLens/TraceUtils/split_inference/execution_roots.py @@ -6,8 +6,6 @@ """Stage 1: find iteration execution roots in an inference trace.""" -from typing import List, Optional - from ..annotation_utils import ( ITERATION_BACKUP_PATTERNS, ITERATION_PATTERNS, @@ -15,7 +13,7 @@ ) -def find_iteration_roots(events: List[dict]) -> Optional[List[dict]]: +def find_iteration_roots(events: list[dict]) -> list[dict] | None: """Return iteration-root events. Tries the primary annotation pattern first, then backup patterns, then diff --git a/TraceLens/TraceUtils/split_inference/steady_state_window.py b/TraceLens/TraceUtils/split_inference/steady_state_window.py index 5d4c9c651..a792b6915 100644 --- a/TraceLens/TraceUtils/split_inference/steady_state_window.py +++ b/TraceLens/TraceUtils/split_inference/steady_state_window.py @@ -8,7 +8,6 @@ import math from statistics import mean -from typing import List, Optional, Tuple from ..annotation_utils import ( has_context, @@ -19,8 +18,8 @@ def identify_steady_state_regions( - iter_details: List[dict], num_steps: int -) -> Tuple[List[Tuple[int, int]], int]: + iter_details: list[dict], num_steps: int +) -> tuple[list[tuple[int, int]], int]: """Detect contiguous steady-state regions based on num_requests proximity to global max. Returns ``(regions, global_max)`` where ``regions`` is a list of @@ -81,8 +80,8 @@ def identify_steady_state_regions( def compute_reference_pd_ratio( - regions: List[Tuple[int, int]], iter_details: List[dict] -) -> Tuple[Tuple[int, int], float, float]: + regions: list[tuple[int, int]], iter_details: list[dict] +) -> tuple[tuple[int, int], float, float]: """ Return the largest steady-state region, a reference PD ratio, and the median PD ratio across all regions. @@ -119,14 +118,14 @@ def compute_reference_pd_ratio( def find_steady_state_window( - iteration_roots: List[dict], + iteration_roots: list[dict], num_steps: int, - steady_state_regions: List[Tuple[int, int]], + steady_state_regions: list[tuple[int, int]], mode: str = "mixed", - CONC: Optional[int] = None, - OSL: Optional[float] = None, - R: Optional[float] = None, -) -> List[dict]: + CONC: int | None = None, + OSL: float | None = None, + R: float | None = None, +) -> list[dict]: """ Find the best contiguous window of up to ``num_steps`` iterations. @@ -168,12 +167,12 @@ def find_steady_state_window( regions = steady_state_regions global_max = max(t["num_requests"] for t in iter_details) - (largest_start, largest_end), reference_ratio, largest_window_ratio = ( + (largest_start, largest_end), reference_ratio, _largest_window_ratio = ( compute_reference_pd_ratio(regions, iter_details) ) # --- Optional: CONC / OSL / R validation and ideal ratio override ---------- - ideal_pd_ratio: Optional[float] = None + ideal_pd_ratio: float | None = None if CONC is not None and global_max != CONC: print( @@ -210,7 +209,7 @@ def find_steady_state_window( print( f"Using ideal prefilldecodemix_to_totalsteps_ratio={ideal_pd_ratio:.4f} as reference (overrides empirical {reference_ratio:.4f})" ) - print(f"\n --------------------------------") + print("\n --------------------------------") # --------------------------------------------------------------------------- divider = max(1, min(int(num_steps / 2), 10)) @@ -220,7 +219,7 @@ def find_steady_state_window( candidates = [] s, e = largest_start, largest_end - def _count_mixed(window: List[dict]) -> int: + def _count_mixed(window: list[dict]) -> int: """Count truly-mixed steps (both context and generation requests > 0).""" return sum(1 for t in window if is_mixed(t)) @@ -290,8 +289,8 @@ def _count_mixed(window: List[dict]) -> int: # Find the longest contiguous run of pure decode-only steps (active # generation with no context requests) in the largest steady-state # region, capped at num_steps. - do_runs: List[Tuple[int, int]] = [] # (start, end) in iter_details coords - run_start: Optional[int] = None + do_runs: list[tuple[int, int]] = [] # (start, end) in iter_details coords + run_start: int | None = None for idx in range(largest_start, largest_end): if is_decode_only(iter_details[idx]): if run_start is None: @@ -324,8 +323,8 @@ def _count_mixed(window: List[dict]) -> int: elif mode == "max_prefilldecode": # Find the longest contiguous run of pure PD steps (no decode-only) in # the largest steady-state region, capped at num_steps. - pd_runs: List[Tuple[int, int]] = [] # (start, end) in iter_details coords - run_start: Optional[int] = None + pd_runs: list[tuple[int, int]] = [] # (start, end) in iter_details coords + run_start: int | None = None for idx in range(largest_start, largest_end): if has_context(iter_details[idx]): if run_start is None: diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index 78f88dc44..f8b101e6b 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -10,7 +10,6 @@ import json import os import zipfile -from typing import List, Optional, Set, Tuple from tqdm import tqdm @@ -42,7 +41,7 @@ def get_filename(filepath: str) -> dict: return filepath -def preprocess_trace(events: List[dict]): +def preprocess_trace(events: list[dict]): gpu_corr_map = {} flow_corr_map = {} meta_events = [] @@ -67,12 +66,12 @@ def preprocess_trace(events: List[dict]): def extract_iteration( - iteration_roots: List[dict], - events: List[dict], + iteration_roots: list[dict], + events: list[dict], trace_json: dict, gpu_corr_map: dict, flow_corr_map: dict, - meta_events: List[dict], + meta_events: list[dict], ) -> dict: """Extract a single iteration trace.""" @@ -92,7 +91,7 @@ def extract_iteration( root.get("ts", 0) + root.get("dur", 0) for root in iteration_roots ) # Collect all relevant tid/pid pairs - tid_pid_set = set((root.get("tid"), root.get("pid")) for root in iteration_roots) + tid_pid_set = {(root.get("tid"), root.get("pid")) for root in iteration_roots} # Pre-filter all CPU events in the global window and by tid/pid cpu_events = [] @@ -120,7 +119,7 @@ def extract_iteration( iter_ts = iteration_root.get("ts", 0) iter_end = iter_ts + iteration_root.get("dur", 0) - correlation_ids: Set[int] = set() + correlation_ids: set[int] = set() # CPU events: filter from pre-filtered list for e in cpu_events: @@ -141,8 +140,7 @@ def extract_iteration( # Add matching flow events for corr in correlation_ids: - for e in flow_corr_map.get(corr, []): - filtered_events.append(e) + filtered_events.extend(flow_corr_map.get(corr, [])) # Add matching GPU events for corr in correlation_ids: for e in gpu_corr_map.get(corr, []): @@ -169,7 +167,7 @@ def extract_iteration( return output, list(set(batch_list)), num_gpu_events, gpu_dur, gpu_busy -def parse_range(range_str: str, max_len: int) -> Tuple[int, int]: +def parse_range(range_str: str, max_len: int) -> tuple[int, int]: """Parse a range string like '10:20' or 'all'.""" if range_str == "all": return 0, max_len @@ -180,8 +178,8 @@ def parse_range(range_str: str, max_len: int) -> Tuple[int, int]: def extract_and_save( - roots: List[List[dict]], - events: List[dict], + roots: list[list[dict]], + events: list[dict], trace_json: dict, output_dir: str, base_name: str, @@ -190,8 +188,8 @@ def extract_and_save( end: int, gpu_corr_map: dict, flow_corr_map: dict, - meta_events: List[dict], - output_label: Optional[str] = None, + meta_events: list[dict], + output_label: str | None = None, ): """Extract and save a range of iterations/dummy runs. @@ -221,10 +219,11 @@ def extract_and_save( # --store-single-iteration keep their literal step name. is_structured = is_annotation and (output_label is not None or len(root) > 1) - if is_structured or not is_annotation: - if len(batch_list) == len(iter_details): - for bs, iteration in zip(batch_list, iter_details): - iteration["batch_size"] = bs + if (is_structured or not is_annotation) and len(batch_list) == len( + iter_details + ): + for bs, iteration in zip(batch_list, iter_details): + iteration["batch_size"] = bs phase_details = find_phase_from_window(iter_details) @@ -288,8 +287,8 @@ def extract_and_save( def extract_phases_and_save( - roots: List[List[dict]], - events: List[dict], + roots: list[list[dict]], + events: list[dict], trace_json: dict, output_dir: str, base_name: str, @@ -298,7 +297,7 @@ def extract_phases_and_save( end: int, gpu_corr_map: dict, flow_corr_map: dict, - meta_events: List[dict], + meta_events: list[dict], ): """Extract and save a range of iterations/dummy runs.""" extraction_summary = [] @@ -315,7 +314,7 @@ def extract_phases_and_save( iter_details = iteration_details(prefilldecode_steps) phase_details = find_phase_from_window(iter_details) - iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = ( + iter_trace, _batch_list, num_gpu_events, gpu_dur, gpu_busy = ( extract_iteration( prefilldecode_steps, events, @@ -347,7 +346,7 @@ def extract_phases_and_save( if len(decode_steps) > 0: iter_details = iteration_details(decode_steps) phase_details = find_phase_from_window(iter_details) - iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = ( + iter_trace, _batch_list, num_gpu_events, gpu_dur, gpu_busy = ( extract_iteration( decode_steps, events, @@ -380,16 +379,16 @@ def extract_phases_and_save( def divide_phases_and_save( - iteration_roots: List[dict], - events: List[dict], + iteration_roots: list[dict], + events: list[dict], trace_json: dict, output_dir: str, base_name: str, gpu_corr_map: dict, flow_corr_map: dict, - meta_events: List[dict], - steady_state_regions: List[Tuple[int, int]], -) -> List[dict]: + meta_events: list[dict], + steady_state_regions: list[tuple[int, int]], +) -> list[dict]: """ Group contiguous steps of the same phase within steady-state regions and save each contiguous run as a single trace file into one of two sub-folders: @@ -411,7 +410,7 @@ def divide_phases_and_save( print(f"[divide-phases] Steady-state regions: {regions}") # Build an ordered list of (phase_label, root) for all steady-state steps - steady_steps: List[Tuple[str, dict]] = [] + steady_steps: list[tuple[str, dict]] = [] for s, e in regions: for idx in range(s, e): detail = iter_details[idx] @@ -423,7 +422,7 @@ def divide_phases_and_save( # steps that are neither (e.g. idle) are skipped # Group into contiguous runs of the same phase - runs: List[Tuple[str, List[dict]]] = [] # (phase, [roots]) + runs: list[tuple[str, list[dict]]] = [] # (phase, [roots]) for phase, root in steady_steps: if runs and runs[-1][0] == phase: runs[-1][1].append(root) diff --git a/TraceLens/TraceUtils/split_inference_trace_annotation.py b/TraceLens/TraceUtils/split_inference_trace_annotation.py index 48706bcbe..e8a928955 100644 --- a/TraceLens/TraceUtils/split_inference_trace_annotation.py +++ b/TraceLens/TraceUtils/split_inference_trace_annotation.py @@ -154,17 +154,18 @@ import argparse import json import os -from typing import List, Tuple import pandas as pd from ..util import DataLoader from .annotation_utils import ( - ITERATION_BACKUP_PATTERNS, - ITERATION_PATTERNS, + ITERATION_BACKUP_PATTERNS, # noqa: F401 + ITERATION_PATTERNS, # noqa: F401 iteration_details, ) -from .split_inference import ( + +# Re-exports for tests and downstream callers. +from .split_inference import ( # noqa: F401 compute_reference_pd_ratio, divide_phases_and_save, extract_and_save, @@ -290,10 +291,10 @@ def main(): # Determine the working set and compute steady-state regions once, # shared across all downstream calls. - steady_state_regions: List[Tuple[int, int]] = [] + steady_state_regions: list[tuple[int, int]] = [] if args.iterations != "all": working_roots = iteration_roots[start:end] - steady_state_regions: List[Tuple[int, int]] = [(0, end - start)] + steady_state_regions: list[tuple[int, int]] = [(0, end - start)] print( f"\nUsing explicit iteration range [{start}, {end}) as the working region." ) @@ -384,7 +385,7 @@ def main(): rows = [] for entry in execution_details: row = {k: v for k, v in entry.items() if k not in ("steps", "phase")} - if "phase" in entry and entry["phase"]: + if entry.get("phase"): for pk, pv in entry["phase"].items(): row[f"phase_{pk}"] = pv row["num_steps"] = len(entry.get("steps", [])) diff --git a/examples/custom_workflows/inference_analysis/build_docker_sglang.sh b/examples/custom_workflows/inference_analysis/build_docker_sglang.sh index f5ababe19..bf1126f23 100755 --- a/examples/custom_workflows/inference_analysis/build_docker_sglang.sh +++ b/examples/custom_workflows/inference_analysis/build_docker_sglang.sh @@ -31,6 +31,11 @@ Options: - 0.5.17 / v0517 : sglang_roofline_patches/sglang_0_5_17/ --gpu-type mi300 | mi350 | mi355 (default: mi350) --base-image Override the default base image + --patch-dir Patch directory under sglang_roofline_patches/ + instead of the one derived from --sglang-version. + Use for bases that have drifted from a release tag, + e.g. --patch-dir sglang_0_5_17_sgldev for the + rocm/sgl-dev nightly. -h, --help Show this help Legacy positional (still supported): @@ -72,6 +77,7 @@ EOF SGLANG_VERSION="0.5.9" GPU_TYPE="mi350" CUSTOM_BASE_IMAGE="" +CUSTOM_PATCH_DIR="" TRACELENS_PATH="" DOCKER_ARGS=() @@ -103,6 +109,10 @@ while [[ $# -gt 0 ]]; do CUSTOM_BASE_IMAGE="$2" shift 2 ;; + --patch-dir) + CUSTOM_PATCH_DIR="$2" + shift 2 + ;; -h|--help) usage ;; @@ -246,7 +256,7 @@ if [ -n "${CUSTOM_BASE_IMAGE}" ]; then BASE_IMAGE="${CUSTOM_BASE_IMAGE}" fi -PATCH_DIR="sglang_$(echo "${SGLANG_VERSION}" | tr '.' '_')" +PATCH_DIR="${CUSTOM_PATCH_DIR:-sglang_$(echo "${SGLANG_VERSION}" | tr '.' '_')}" PATCH_DIR_PATH="${PATCHES_ROOT}/${PATCH_DIR}" if [ ! -d "${PATCH_DIR_PATH}" ]; then echo "Error: patch directory not found: ${PATCH_DIR_PATH}" From 80c3398e438e8d93979918c1649d9ee1bff89aa3 Mon Sep 17 00:00:00 2001 From: Deval Shah Date: Wed, 19 Aug 2026 15:30:40 -0500 Subject: [PATCH 3/4] fixing lint import and stale references --- .importlinter | 2 +- TraceLens/TraceUtils/split_inference/trace_extraction.py | 4 ++-- docs/reference/api-reference.md | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.importlinter b/.importlinter index 83490b050..1885c79db 100644 --- a/.importlinter +++ b/.importlinter @@ -10,4 +10,4 @@ ancestors = depth = 10 ignore_imports = TraceLens.Trace2Tree.trace_capture_merge_experimental -> TraceLens.TraceUtils.annotation_utils - TraceLens.TraceUtils.split_inference_trace_annotation -> TraceLens.Trace2Tree.inference_iteration_roots + TraceLens.TraceUtils.split_inference.execution_roots -> TraceLens.Trace2Tree.inference_iteration_roots diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index f8b101e6b..a035ea7b6 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -191,7 +191,7 @@ def extract_and_save( meta_events: list[dict], output_label: str | None = None, ): - """Extract and save a range of iterations/dummy runs. + """Extract and save a range of iterations. If ``output_label`` is provided the output filename becomes ``{output_label}_{name_append}_{base_name}.json.gz`` instead of the @@ -299,7 +299,7 @@ def extract_phases_and_save( flow_corr_map: dict, meta_events: list[dict], ): - """Extract and save a range of iterations/dummy runs.""" + """Extract and save a range of iterations.""" extraction_summary = [] if "annotation_iteration" not in prefix: diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 6078fb478..4ae3abcc8 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -172,7 +172,7 @@ Split an inference trace into per-iteration or per-phase sub-traces. | `trace_path` | required | Path to the trace (`.json` or `.json.gz`). | | `-o`, `--output-dir` | required | Output directory. | | `-i`, `--iterations` | `all` | Iteration range: `all`, a single index (`50`), or a range (`10:20`). | -| `-d`, `--dummy` | `all` | Dummy-run range. | +| `--store-single-iteration` | off | Write each iteration as a separate trace file. | | `--find-steady-state` / `--num-steps` | off / 32 | Extract a steady-state region of N iterations. | | `--divide-phases` | off | Store steady-state steps into `prefilldecodemix/` and `decode_only/` sub-folders. | | `--CONC`, `--OSL`, `--R` | None | Expected concurrency and output-sequence-length window parameters. | From 2088cb11a333cf9b9635246e7adb83ec5c127d2f Mon Sep 17 00:00:00 2001 From: Deval Shah Date: Tue, 25 Aug 2026 22:10:41 -0500 Subject: [PATCH 4/4] restructuring the splitter flow to integrate fallbacks for corner cases --- .importlinter | 3 + .../Trace2Tree/inference_iteration_roots.py | 347 ++++++-- TraceLens/TraceUtils/annotation_utils.py | 91 +- .../TraceUtils/split_inference/__init__.py | 20 +- .../split_inference/detect_utils.py | 562 +++++++++++++ .../split_inference/execution_roots.py | 199 ++++- .../split_inference/root_detection.py | 345 ++++++++ .../TraceUtils/split_inference/root_probes.py | 254 ++++++ .../split_inference/steady_state_window.py | 161 +++- .../split_inference/trace_extraction.py | 96 ++- .../split_inference_trace_annotation.py | 170 +++- tests/test_split_root_detection.py | 793 ++++++++++++++++++ 12 files changed, 2917 insertions(+), 124 deletions(-) create mode 100644 TraceLens/TraceUtils/split_inference/detect_utils.py create mode 100644 TraceLens/TraceUtils/split_inference/root_detection.py create mode 100644 TraceLens/TraceUtils/split_inference/root_probes.py create mode 100644 tests/test_split_root_detection.py diff --git a/.importlinter b/.importlinter index 1885c79db..d860257d8 100644 --- a/.importlinter +++ b/.importlinter @@ -11,3 +11,6 @@ depth = 10 ignore_imports = TraceLens.Trace2Tree.trace_capture_merge_experimental -> TraceLens.TraceUtils.annotation_utils TraceLens.TraceUtils.split_inference.execution_roots -> TraceLens.Trace2Tree.inference_iteration_roots + TraceLens.TraceUtils.split_inference.root_detection -> TraceLens.Trace2Tree.inference_iteration_roots + TraceLens.TraceUtils.split_inference.root_probes -> TraceLens.Trace2Tree.inference_iteration_roots + TraceLens.TraceUtils.split_inference.root_probes -> TraceLens.Trace2Tree.trace_to_tree diff --git a/TraceLens/Trace2Tree/inference_iteration_roots.py b/TraceLens/Trace2Tree/inference_iteration_roots.py index fcb5b8d6b..5df693dd4 100644 --- a/TraceLens/Trace2Tree/inference_iteration_roots.py +++ b/TraceLens/Trace2Tree/inference_iteration_roots.py @@ -6,124 +6,291 @@ """Generic iteration-root detection via TraceToTree call-tree traversal.""" -from typing import List, Optional, Tuple +from collections import Counter, deque +from dataclasses import dataclass +from statistics import mean, pstdev +from typing import Dict, List, Optional, Sequence, Tuple from .trace_to_tree import TraceToTree +# A period must explain more than half the sequence, matching the original rule. +MIN_PERIOD_COVERAGE = 0.5 +# Candidate periods come from gaps between recurrences of one label; cap the +# number verified so a pathological sequence cannot blow up the search. +MAX_PERIOD_CANDIDATES = 64 +# A longer period is only preferred over a shorter one it is a multiple of when +# it explains meaningfully more of the sequence. +DIVISOR_COVERAGE_TOLERANCE = 0.05 + +# Label sequences shorter than this are utility-function child lists, not loops. +MIN_LABEL_CHILDREN = 6 + +# Preferred sources of labels for period detection, best first. Python frames +# correspond to semantic loop bodies; the rest are fallbacks for traces captured +# without stack recording. +PYTHON_TIER = "python_function" +CPU_OP_TIER = "cpu_op" +ALL_CHILDREN_TIER = "all_children" + +PERIOD_EXACT = "exact" +PERIOD_INTEGER_RATIO = "integer_ratio" +PERIOD_CONFLICT = "conflict" + + +@dataclass +class PeriodCandidate: + """A verified repeating period, with the evidence for it.""" + + period: int + start: int + repeats: int + coverage: float + duration_cv: float + + @property + def rank(self) -> tuple: + """Sort key: explain the most sequence, most evenly, with the shortest unit. + + Coverage is rounded so float noise cannot outrank a steadier candidate, + and period breaks ties downward since every multiple explains as much. + """ + return (-round(self.coverage, 3), round(self.duration_cv, 3), self.period) + + +def _candidate_periods(codes: Sequence[int], min_repeats: int) -> List[int]: + """Plausible periods, taken from the gaps between one label's recurrences. + + If a sequence has period ``p`` every label recurs every ``p`` positions, so + any single label's gaps contain ``p``. Anchoring on the rarest eligible + label keeps the list short and makes a non-repeating sequence cost nothing. + """ + counts = Counter(codes) + eligible = [(n, code) for code, n in counts.items() if n >= min_repeats] + if not eligible: + return [] + _, anchor = min(eligible) + positions = [i for i, code in enumerate(codes) if code == anchor] + gaps = {b - a for a, b in zip(positions, positions[1:]) if b > a} + return sorted(gaps)[:MAX_PERIOD_CANDIDATES] + + +def _longest_periodic_run(codes: Sequence[int], period: int) -> Tuple[int, int]: + """Start index and block count of the longest ``period``-aligned run. + + Scanning for the longest run finds the loop wherever it sits, so a warmup + prefix is skipped without retrying the search from every offset. + """ + limit = len(codes) - period + best_start = best_len = 0 + i = 0 + while i < limit: + if codes[i] != codes[i + period]: + i += 1 + continue + j = i + while j < limit and codes[j] == codes[j + period]: + j += 1 + if j - i > best_len: + best_start, best_len = i, j - i + i = j + 1 + return best_start, (best_len + period) // period + + +def _duration_cv( + durations: Optional[Sequence[float]], start: int, period: int, repeats: int +) -> float: + """Coefficient of variation of per-occurrence duration. + + A real iteration takes about the same time every time, which separates a + genuine loop from a coincidental label match. + """ + if not durations: + return 0.0 + blocks = [ + sum(durations[start + i * period : start + (i + 1) * period]) + for i in range(repeats) + ] + blocks = [b for b in blocks if b > 0] + if len(blocks) < 2: + return 0.0 + average = mean(blocks) + return pstdev(blocks) / average if average else 0.0 + + +def _drop_multiples(candidates: List[PeriodCandidate]) -> List[PeriodCandidate]: + """Keep primitive periods; any multiple of one is valid and explains no more.""" + kept: List[PeriodCandidate] = [] + for cand in candidates: + if any( + cand.period % k.period == 0 + and cand.coverage <= k.coverage + DIVISOR_COVERAGE_TOLERANCE + for k in kept + ): + continue + kept.append(cand) + return kept + + +def find_period_candidates( + labels: Sequence[str], + durations: Optional[Sequence[float]] = None, + min_repeats: int = 3, +) -> List[PeriodCandidate]: + """Every qualifying repeating period in ``labels``, best first. + + Scored candidates rather than one answer let callers cross-check against an + independent detection instead of trusting a single verdict. + """ + # Labels as small ints, so comparisons are cheap in the verification loop. + table: Dict[str, int] = {} + codes = [table.setdefault(label, len(table)) for label in labels] + total = len(codes) + found: List[PeriodCandidate] = [] + for period in _candidate_periods(codes, min_repeats): + if period * min_repeats > total: + continue + start, repeats = _longest_periodic_run(codes, period) + if repeats < min_repeats: + continue + coverage = repeats * period / total + if coverage <= MIN_PERIOD_COVERAGE: + continue + found.append( + PeriodCandidate( + period, + start, + repeats, + coverage, + _duration_cv(durations, start, period, repeats), + ) + ) + return _drop_multiples(sorted(found, key=lambda c: c.rank)) + + +def compare_periods(a: Optional[int], b: Optional[int]) -> Tuple[str, Optional[int]]: + """Whether two independently detected periods agree. + + Differing by an exact integer factor means the same loop at different + granularities -- a confirmation, not a conflict. + """ + if not a or not b: + return PERIOD_CONFLICT, None + low, high = min(a, b), max(a, b) + if low == high: + return PERIOD_EXACT, 1 + if high % low == 0: + return PERIOD_INTEGER_RATIO, high // low + return PERIOD_CONFLICT, None + def _find_repeating_period( names: List[str], min_repeats: int = 3 ) -> Tuple[Optional[int], Optional[List[str]], Optional[int]]: - """Find the shortest repeating name sequence anywhere in ``names``. + """Best repeating name sequence in ``names`` as ``(period, pattern, start)``.""" + candidates = find_period_candidates(names, min_repeats=min_repeats) + if not candidates: + return None, None, None + best = candidates[0] + return best.period, list(names[best.start : best.start + best.period]), best.start + - Slides a start offset forward to skip any non-repeating prefix (setup - events before the loop body). Returns ``(period, pattern, start_offset)`` - where ``start_offset`` is the index in ``names`` where the first block - begins. Returns ``(None, None, None)`` if no qualifying period is found. +def _nearest_descendants(tree: TraceToTree, node: dict, cat: str) -> List[dict]: + """Nearest descendants of ``node`` in category ``cat``, in time order. - Requires at least ``min_repeats`` consecutive repetitions covering more - than half of the suffix starting at ``start_offset``. + Descent stops at each match, so the result is one abstraction layer. Direct + children are not enough: a python frame's children are often ATen ops whose + own children are the next python frames, so filtering them returns nothing. """ - n = len(names) - for start in range(n): - suffix = names[start:] - m = len(suffix) - for p in range(1, m // 2 + 1): - pattern = suffix[:p] - count = 0 - i = 0 - while i + p <= m and suffix[i : i + p] == pattern: - count += 1 - i += p - if count >= min_repeats and count * p > m * 0.5: - return p, pattern, start - return None, None, None - - -def _detect_iteration_roots_from_tree(tree: TraceToTree, roots) -> Optional[List[dict]]: - """BFS down the tree from one or more root nodes to find and return synthetic - iteration-root events. - - ``roots`` may be a single event dict or a list of event dicts — all are - seeded into the BFS at depth 0 so they are explored level-by-level together. - - Pattern detection uses all children (not just GPU-path ones) so that - leading CPU-only events (e.g. ``next`` in the OWL pipeline) are included - as part of the iteration anchor. A minimum child count guards against false - positives from short utility-function child lists. - - Returns a list of synthetic root events, one per detected iteration, where - each event's ``dur`` spans from the first to the last child of the block. + found: List[dict] = [] + queue = deque(tree.get_children_events(node)) + while queue: + child = queue.popleft() + if child.get("cat") == cat: + found.append(child) + else: + queue.extend(tree.get_children_events(child)) + found.sort(key=lambda e: e.get("ts", 0)) + return found + + +def _label_events(tree: TraceToTree, node: dict) -> Tuple[List[dict], str]: + """Events under ``node`` to run period detection over, and which tier they are. + + Launches recur many times per iteration and swamp the iteration-level + signal, so python frames -- the semantic loop bodies -- are preferred. The + ladder exists because a capture without stack recording has none at all. """ - from collections import deque + for cat in (PYTHON_TIER, CPU_OP_TIER): + found = _nearest_descendants(tree, node, cat) + if len(found) >= MIN_LABEL_CHILDREN: + return found, cat + return tree.get_children_events(node), ALL_CHILDREN_TIER + +def _detect_iteration_roots_from_tree( + tree: TraceToTree, roots, diagnostics: Optional[dict] = None +) -> Optional[List[dict]]: + """BFS down from ``roots`` for a repeating block, returned as synthetic roots. + + Each returned event spans one block, from the first child's start to the last + child's end, so CPU-only leading work stays inside the iteration. + """ if isinstance(roots, dict): roots = [roots] queue = deque((node, 0) for node in roots) while queue: current, depth = queue.popleft() - children = tree.get_children_events(current) - if not children: + labelled, tier = _label_events(tree, current) + if not labelled: continue - # Only recurse into GPU-bearing subtrees. - if not any(c.get("gpu_events") for c in children): + # Only recurse into GPU-bearing subtrees, tested on the events actually + # being used as labels rather than on the raw child list. + if not any(e.get("gpu_events") for e in labelled): continue - p, _, start = _find_repeating_period([c.get("name", "") for c in children]) - if p is None: - for child in children: - if child.get("gpu_events"): - queue.append((child, depth + 1)) + period, _, start = _find_repeating_period([e.get("name", "") for e in labelled]) + if period is None: + for event in labelled: + if event.get("gpu_events"): + queue.append((event, depth + 1)) continue - print( - f"Generic fallback: repeating pattern found under '{current.get('name')}' at depth {depth}" - ) - print(f"Generic fallback: period={p}") - - # Anchor each iteration between the Nth occurrence of the first and last - # events in the detected pattern. Using all-children anchors means - # CPU-only leading/trailing events are included naturally. - first_anchor_name = children[start]["name"] - last_anchor_name = children[start + p - 1]["name"] - - first_anchors = [ - i - for i, c in enumerate(children) - if i >= start and c.get("name") == first_anchor_name - ] - last_anchors = [ - i - for i, c in enumerate(children) - if i >= start and c.get("name") == last_anchor_name - ] - + blocks = (len(labelled) - start) // period iteration_roots = [] - for n in range(min(len(first_anchors), len(last_anchors))): - block_start = first_anchors[n] - block_end = last_anchors[n] - if block_end < block_start: - break - block = children[block_start : block_end + 1] + for index in range(blocks): + block = labelled[start + index * period : start + (index + 1) * period] first, last = block[0], block[-1] root_event = dict(first) root_event["dur"] = (last["ts"] + last.get("dur", 0)) - first["ts"] iteration_roots.append(root_event) + print( + f"Generic fallback: repeating pattern found under " + f"'{current.get('name')}' at depth {depth} (tier={tier}, period={period})" + ) print(f"Generic fallback: identified {len(iteration_roots)} iterations.") - return iteration_roots if iteration_roots else None + if diagnostics is not None: + diagnostics.update( + { + "period_label_tier": tier, + "period": period, + "period_depth": depth, + } + ) + return iteration_roots or None return None -def find_iteration_roots_generic(events: List[dict]) -> Optional[List[dict]]: - """Fallback: detect iteration roots by finding a repeating child pattern in - the call tree, using TraceToTree for parent/child relationships. +def find_iteration_roots_generic( + events: List[dict], diagnostics: Optional[dict] = None +) -> Optional[List[dict]]: + """Fallback: detect iteration roots from a repeating child pattern. Works for any workload (diffusion, training, etc.) where the iteration loop - body is a repeating sequence of top-level calls under a common parent. + body is a repeating sequence of calls under a common parent. """ try: tree = TraceToTree(events, prune_nongpu_paths=False) @@ -132,26 +299,26 @@ def find_iteration_roots_generic(events: List[dict]) -> Optional[List[dict]]: print(f"Generic fallback: TraceToTree build failed ({e}), skipping.") return None - # Walk every cpu_root_node upward through python_function parents until - # reaching a parentless node — these are the true per-thread entry points. + # Walk every cpu_root_node upward to a parentless node -- those are the true + # per-thread entry points. seen_roots: set = set() trace_roots = [] for uid in tree.cpu_root_nodes: - e = tree.get_UID2event(uid) + event = tree.get_UID2event(uid) while True: - parent = tree.get_parent_event(e) + parent = tree.get_parent_event(event) if parent is None: break - e = parent - if id(e) not in seen_roots: - seen_roots.add(id(e)) - trace_roots.append(e) + event = parent + if id(event) not in seen_roots: + seen_roots.add(id(event)) + trace_roots.append(event) if not trace_roots: print("Generic fallback: no root nodes found.") return None - roots = _detect_iteration_roots_from_tree(tree, trace_roots) + roots = _detect_iteration_roots_from_tree(tree, trace_roots, diagnostics) if roots is None: print("Generic fallback: no repeating child pattern found.") return roots diff --git a/TraceLens/TraceUtils/annotation_utils.py b/TraceLens/TraceUtils/annotation_utils.py index d8bcf7f1a..e6250f0ad 100644 --- a/TraceLens/TraceUtils/annotation_utils.py +++ b/TraceLens/TraceUtils/annotation_utils.py @@ -13,7 +13,8 @@ """ import re -from typing import Callable, List, Optional +from functools import lru_cache +from typing import Callable, Dict, Iterable, List, Optional, Tuple # --- patterns -------------------------------------------------------------- # Each block lists matching annotations, prefill/decode/mixed where the format @@ -248,6 +249,8 @@ def __init__(self, annotation: str): self.c_sq = self.c_sk = self.c_sqsq = self.c_sqsk = 0 self.g_sq = self.g_sk = self.g_sqsq = self.g_sqsk = 0 self.has_sqsk = False + # Only _fill_diffusion sets this; declared here so every instance has it. + self.resolution = None self.meta = {} for kind, pattern, parser in self.FORMATS: if pattern.match(annotation) and parser(self, annotation) is not False: @@ -413,6 +416,90 @@ def find_iteration_roots_by_priority( return [] +# --- family keys ------------------------------------------------------------ +_DIGIT_RUN = re.compile(r"\d+") +SKELETON_PLACEHOLDER = "#" + + +@lru_cache(maxsize=None) +def name_skeleton(name: str) -> str: + """Canonical family key for an event name. + + Collapses every run of digits so that instances of one logical operation + share a key regardless of the batch sizes and sequence lengths baked into + their names:: + + execute_1_context_3(sq8sk8) -> execute_#_context_#(sq#sk#) + execute_7_context_2(sq4sk4) -> execute_#_context_#(sq#sk#) + + Keying on structure rather than on a fixed-length prefix matters both ways: + a prefix merges families that diverge after the cutoff and splits families + that differ only in a numeric tail. + """ + return _DIGIT_RUN.sub(SKELETON_PLACEHOLDER, name) + + +def cluster_by_skeleton(names: Iterable[str]) -> Dict[str, List[str]]: + """Group names by skeleton, keeping first-seen order for stable output.""" + groups: Dict[str, List[str]] = {} + for name in names: + groups.setdefault(name_skeleton(name), []).append(name) + return groups + + +def dominant_cluster(groups: Dict[str, List[str]]) -> Tuple[Optional[str], float]: + """Largest group's skeleton and its share of every clustered name.""" + if not groups: + return None, 0.0 + total = sum(len(v) for v in groups.values()) + skeleton = max(groups, key=lambda k: len(groups[k])) + return skeleton, len(groups[skeleton]) / total + + +# --- cached identity and inheritance ---------------------------------------- +# Stamped onto roots whose window and identity come from different annotations. +PROVENANCE_KEY = "split_provenance" + + +@lru_cache(maxsize=None) +def parse_annotation(name: str) -> IterationAnnotation: + """Memoized parse. Treat the result as read-only; instances are shared. + + Construction runs up to seven regex matches, and detection parses the same + few hundred distinct names across hundreds of thousands of instances. + """ + return IterationAnnotation(name) + + +def is_parseable(name: str) -> bool: + """True when a parser recognized the name, so its metadata is real. + + The distinction matters because unparseable names still yield a full detail + dict, just a fabricated one (one decode-equivalent request), so callers that + test the numbers instead of this predicate cannot tell the two apart. + """ + return parse_annotation(name).matched + + +def inherit_identity(target: dict, source: dict) -> dict: + """Copy of ``target`` that parses as ``source``. + + Timestamps, process and thread stay with ``target``: the outer span decides + the extraction window while the inner annotation supplies phase and batch + size. Both names are recorded so the two-level choice stays auditable. + """ + prior = target.get(PROVENANCE_KEY) or {} + out = dict(target) + out["name"] = source.get("name", "") + out[PROVENANCE_KEY] = { + # Keep the original window owner when a root is relabelled more than + # once, otherwise the second pass erases where the span came from. + "window_from": prior.get("window_from") or target.get("name", ""), + "identity_from": source.get("name", ""), + } + return out + + def has_context(detail: dict) -> bool: """Step runs at least one prefill (context) request.""" return detail.get("context_requests", 0) > 0 @@ -450,7 +537,7 @@ def classify_phase(detail: dict) -> Optional[str]: # --- per-window aggregation ------------------------------------------------- def iteration_details(roots: List[dict], full: bool = False) -> List[dict]: """Parse iteration-root events into one detail dict per step.""" - annotations = (IterationAnnotation(r.get("name", "")) for r in roots) + annotations = (parse_annotation(r.get("name", "")) for r in roots) if full: return [a.full_details() for a in annotations] return [a.iter_details() for a in annotations] diff --git a/TraceLens/TraceUtils/split_inference/__init__.py b/TraceLens/TraceUtils/split_inference/__init__.py index f25444d9d..90a6c89d8 100644 --- a/TraceLens/TraceUtils/split_inference/__init__.py +++ b/TraceLens/TraceUtils/split_inference/__init__.py @@ -6,13 +6,23 @@ """Inference trace splitting: execution roots, steady-state windows, extraction.""" -from .execution_roots import find_iteration_roots +from .execution_roots import ( + DetectStatus, + PhaseConfidence, + RootSet, + find_iteration_roots, + find_iteration_roots_ex, +) from .steady_state_window import ( + classify_workload, compute_reference_pd_ratio, + find_max_pattern_window, find_steady_state_window, identify_steady_state_regions, + select_window, ) from .trace_extraction import ( + build_root_tiles, divide_phases_and_save, extract_and_save, extract_iteration, @@ -23,15 +33,23 @@ ) __all__ = [ + "DetectStatus", + "PhaseConfidence", + "RootSet", + "build_root_tiles", + "classify_workload", "compute_reference_pd_ratio", "divide_phases_and_save", "extract_and_save", "extract_iteration", "extract_phases_and_save", "find_iteration_roots", + "find_iteration_roots_ex", + "find_max_pattern_window", "find_steady_state_window", "get_filename", "identify_steady_state_regions", "parse_range", "preprocess_trace", + "select_window", ] diff --git a/TraceLens/TraceUtils/split_inference/detect_utils.py b/TraceLens/TraceUtils/split_inference/detect_utils.py new file mode 100644 index 000000000..7af435d39 --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/detect_utils.py @@ -0,0 +1,562 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared machinery for execution-root detection. + +Four concerns that several detection steps each need, kept together so their +semantics are defined once: + +- the result contract every step speaks (:class:`RootSet` and friends) +- containment queries over event spans (:class:`IntervalIndex`) +- attributing GPU kernels to annotations and measuring coverage + (:class:`GpuAttribution`) +- enumerating alternative root candidates, and running escalation probes +""" + +from bisect import bisect_left +from dataclasses import dataclass, field +from enum import Enum, IntEnum +from statistics import median +from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple + +from ..annotation_utils import PROVENANCE_KEY, name_skeleton + +# Projections *enclose* the kernels they describe, so summing GPU time over both +# double-counts. Kept apart here and recombined by consumers that want both. +GPU_KERNEL_CATEGORIES = ("kernel", "gpu_memcpy", "gpu_memset") +PROJECTION_CATEGORY = "gpu_user_annotation" + +# Coverage to accept roots outright, and the floor below which a trace is +# unsplittable rather than degraded. +COVERAGE_GATE = 0.95 +COVERAGE_FLOOR = 0.75 + +# Least share of captured GPU time the annotation spans must explain themselves. +# Below this, coverage comes from window extension, so the roots are too sparse. +MIN_SPAN_SHARE = 0.5 + +# Fewer roots than this usually means only a warmup loop matched. +MIN_ROOTS = 8 + + +# --- result contract -------------------------------------------------------- +class DetectStatus(IntEnum): + """Whether the trace can be split. Deliberately separate from phase trust.""" + + SPLITTABLE = 0 + NOT_SPLITTABLE = 1 + DEGRADED = 2 + + +class PhaseConfidence(str, Enum): + """How much to trust the phase and batch-size labels on the roots.""" + + HIGH = "high" # parsed from a recognized annotation + LOW = "low" # inherited onto a synthetic root + UNKNOWN = "unknown" # derived from kernel or python-frame periodicity + + +@dataclass +class CoverageReport: + """Result of a GPU-time coverage audit. + + ``covered_selected`` measures the roots' extraction windows, gaps included, + since that is what the output contains; ``covered_spans`` measures the bare + annotation spans. Judging on bare spans alone hunts for extra roots whenever + an iteration's sampling step runs just outside its annotation. + """ + + strategy: str + covered_any: float + covered_selected: float + covered_spans: float + gpu_busy: float + window: Optional[Tuple[float, float]] = None + + @property + def span_share(self) -> float: + """How much of the captured work the annotations themselves explain. + + Near 0 means a few annotations stretched over many iterations, which + would pass a coverage check while bundling iterations into each slice. + """ + if self.covered_selected <= 0: + return 0.0 + return min(1.0, self.covered_spans / self.covered_selected) + + @property + def passes(self) -> bool: + """Whether the roots explain enough GPU work, without being stretched. + + Gating on the roots rather than on all annotations is what makes this a + real check: a run whose annotations blanket the timeline while the roots + cover fifteen of five hundred iterations scores near-perfectly on the + permissive measure. + """ + return self.covered_selected >= COVERAGE_GATE and ( + self.span_share >= MIN_SPAN_SHARE + ) + + @property + def better_roots_exist(self) -> bool: + """Annotations cover work the roots miss, so escalation should help.""" + return self.covered_any - self.covered_selected > 1 - COVERAGE_GATE + + +@dataclass +class RootSet: + """Roots plus how they were found and how much we trust them.""" + + roots: List[dict] + method: str + phase_confidence: PhaseConfidence = PhaseConfidence.UNKNOWN + status: DetectStatus = DetectStatus.SPLITTABLE + coverage: Optional[CoverageReport] = None + diagnostics: Dict = field(default_factory=dict) + + def __len__(self) -> int: + return len(self.roots) + + def to_manifest(self) -> dict: + cov = self.coverage + manifest = { + "status": int(self.status), + "method": self.method, + "phase_confidence": self.phase_confidence.value, + "n_roots": len(self.roots), + "attribution_strategy": cov.strategy if cov else None, + "coverage_any_annotation": round(cov.covered_any, 4) if cov else None, + "coverage_selected_roots": round(cov.covered_selected, 4) if cov else None, + "coverage_root_spans_only": round(cov.covered_spans, 4) if cov else None, + "root_span_share": round(cov.span_share, 4) if cov else None, + } + manifest.update(self.diagnostics) + return manifest + + +# --- containment queries ---------------------------------------------------- +class IntervalIndex: + """Event spans grouped by ``(pid, tid)`` and sorted by start timestamp. + + Grouping by thread keeps containment from crossing threads and inventing a + false parent; sorting makes the query a bisect rather than a full scan. + """ + + def __init__(self, events: Iterable[dict]): + buckets: Dict[Tuple, List[dict]] = {} + for e in events: + ts, dur = e.get("ts"), e.get("dur") + if ts is None or dur is None: + continue + buckets.setdefault((e.get("pid"), e.get("tid")), []).append(e) + self._threads: Dict[Tuple, Tuple[List[float], List[float], List[dict]]] = {} + for key, evs in buckets.items(): + evs.sort(key=lambda x: x["ts"]) + starts = [x["ts"] for x in evs] + ends = [x["ts"] + x["dur"] for x in evs] + self._threads[key] = (starts, ends, evs) + + def contained_in(self, span: dict, exclude_self: bool = True) -> List[dict]: + """Events lying entirely within ``span``, on ``span``'s own thread.""" + entry = self._threads.get((span.get("pid"), span.get("tid"))) + if entry is None: + return [] + starts, ends, evs = entry + start = span.get("ts", 0) + end = start + span.get("dur", 0) + out = [] + i = bisect_left(starts, start) + while i < len(starts) and starts[i] < end: + if ends[i] <= end and not (exclude_self and evs[i] is span): + out.append(evs[i]) + i += 1 + return out + + +def gaps_between(roots: Sequence[dict]) -> List[dict]: + """Idle spans between consecutive roots, per thread, as span dicts.""" + gaps = [] + for (pid, tid), group in group_by_thread(roots).items(): + for prev, nxt in zip(group, group[1:]): + prev_end = prev.get("ts", 0) + prev.get("dur", 0) + if nxt.get("ts", 0) > prev_end: + gaps.append( + { + "ts": prev_end, + "dur": nxt["ts"] - prev_end, + "pid": pid, + "tid": tid, + } + ) + return gaps + + +def build_root_tiles(roots: Sequence[dict]) -> Tuple[dict, int]: + """Gap-free extraction windows, one per root, keyed by ``(pid, tid, ts)``. + + Each window reaches to the next root's start on the same thread, so work + between two roots belongs to somebody; the last gets the median length + rather than running to the end of the trace and swallowing teardown. + + Grouping by thread is required, not tidy: a global sort interleaves threads + and a window spanning two of them describes nothing. Overlapping roots keep + their own span and are counted instead. + """ + tiles: dict = {} + overlaps = 0 + for (pid, tid), group in group_by_thread(roots).items(): + spans = [] + for index, root in enumerate(group): + start = root.get("ts", 0) + own_end = start + root.get("dur", 0) + end = own_end + if index + 1 < len(group): + following = group[index + 1].get("ts", 0) + if following < own_end: + overlaps += 1 + else: + end = following + spans.append((start, end)) + if len(spans) > 1: + typical = median(end - start for start, end in spans[:-1]) + last_start, last_end = spans[-1] + spans[-1] = (last_start, max(last_end, last_start + typical)) + for start, end in spans: + tiles[(pid, tid, start)] = (start, end) + return tiles, overlaps + + +def group_by_thread(events: Iterable[dict]) -> Dict[Tuple, List[dict]]: + """Events grouped by ``(pid, tid)`` and sorted by timestamp within group.""" + groups: Dict[Tuple, List[dict]] = {} + for e in events: + groups.setdefault((e.get("pid"), e.get("tid")), []).append(e) + for group in groups.values(): + group.sort(key=lambda x: x.get("ts", 0)) + return groups + + +class SpanSet: + """A disjoint, sorted union of ``(start, end)`` time spans. + + Overlapping input is the norm -- annotations nest, projections repeat per + stream -- so merging on construction is what makes membership a bisect. + """ + + def __init__(self, spans: Iterable[Tuple[float, float]] = ()): + self.spans: List[Tuple[float, float]] = [] + for start, end in sorted(spans): + if self.spans and start <= self.spans[-1][1]: + last_start, last_end = self.spans[-1] + self.spans[-1] = (last_start, max(last_end, end)) + else: + self.spans.append((start, end)) + + @classmethod + def of_events(cls, events: Iterable[dict]) -> "SpanSet": + return cls((e["ts"], e["ts"] + e["dur"]) for e in events) + + def __or__(self, other: "SpanSet") -> "SpanSet": + return SpanSet(self.spans + other.spans) + + def __bool__(self) -> bool: + return bool(self.spans) + + def covers(self, point: float) -> bool: + i = bisect_left(self.spans, (point, float("inf"))) + if i > 0 and self.spans[i - 1][0] <= point <= self.spans[i - 1][1]: + return True + return i < len(self.spans) and self.spans[i][0] <= point <= self.spans[i][1] + + @property + def bounds(self) -> Optional[Tuple[float, float]]: + """Time frame enclosing every span, or ``None`` when there are none.""" + if not self.spans: + return None + return self.spans[0][0], max(end for _, end in self.spans) + + +# --- GPU attribution and coverage ------------------------------------------- +class GpuAttribution: + """Attributes GPU kernels to annotations and measures coverage. + + Two strategies, chosen per trace. ``projection`` (the kernel starts inside a + ``gpu_user_annotation`` span) is preferred: cheaper, and immune to graph + capture, where correlations cannot be walked at all. ``correlation`` (the + launch traces back to a CPU op inside an annotation) is a mandatory + fallback, since some traces have no projections -- notably any trace that is + itself a previous split output. + + Attribution to *any* annotation counts, not just the selected root: the + question is "are we missing whole regions of work", not fine accounting. + """ + + STRATEGY_PROJECTION = "projection" + STRATEGY_CORRELATION = "correlation" + + def __init__(self, events: Iterable[dict]): + self.kernels: List[dict] = [] + self.projections: List[dict] = [] + corr_cpu: List[dict] = [] + self._corr_kernels: Dict[int, List[dict]] = {} + self._launch_by_corr: Dict[int, dict] = {} + for e in events: + ts, dur, cat = e.get("ts"), e.get("dur"), e.get("cat") + if ts is None or dur is None: + continue + corr = (e.get("args") or {}).get("correlation") + if cat == PROJECTION_CATEGORY: + self.projections.append(e) + elif cat in GPU_KERNEL_CATEGORIES: + self.kernels.append(e) + if corr is not None: + self._corr_kernels.setdefault(corr, []).append(e) + elif corr is not None: + corr_cpu.append(e) + self._launch_by_corr.setdefault(corr, e) + + self.kernels.sort(key=lambda x: x["ts"]) + self._kernel_starts = [k["ts"] for k in self.kernels] + self.strategy = ( + self.STRATEGY_PROJECTION if self.projections else self.STRATEGY_CORRELATION + ) + # Built on demand: the projection path normally never needs it, but it + # remains reachable as a cross-check when projections look untrustworthy. + self._corr_cpu = corr_cpu + self._cpu_index_cache: Optional[IntervalIndex] = None + + @property + def _cpu_index(self) -> IntervalIndex: + if self._cpu_index_cache is None: + self._cpu_index_cache = IntervalIndex(self._corr_cpu) + return self._cpu_index_cache + + def _kernels_in(self, window: Optional[Tuple[float, float]]) -> List[dict]: + """Kernels whose *start* lies in ``window``. + + Selecting on start rather than clipping durations keeps coverage from + exceeding 1 through partially-overlapping kernels. + """ + if window is None: + return self.kernels + lo, hi = window + i = bisect_left(self._kernel_starts, lo) + out = [] + while i < len(self.kernels) and self._kernel_starts[i] <= hi: + out.append(self.kernels[i]) + i += 1 + return out + + def kernels_for(self, spans: Sequence[dict]) -> List[dict]: + """Kernels launched from CPU ops inside ``spans`` (correlation path).""" + seen, out = set(), [] + for span in spans: + for cpu in self._cpu_index.contained_in(span, exclude_self=False): + corr = (cpu.get("args") or {}).get("correlation") + for k in self._corr_kernels.get(corr, ()): + if id(k) not in seen: + seen.add(id(k)) + out.append(k) + return out + + def cpu_launches_for(self, kernels: Sequence[dict]) -> List[dict]: + """CPU launch sites of ``kernels``, empty under graph capture.""" + seen, out = set(), [] + for kernel in kernels: + corr = (kernel.get("args") or {}).get("correlation") + launch = self._launch_by_corr.get(corr) + if launch is not None and id(launch) not in seen: + seen.add(id(launch)) + out.append(launch) + return out + + def gpu_time_by_correlation(self, spans: Sequence[dict]) -> float: + """GPU time launched from CPU ops inside ``spans``.""" + return sum(k["dur"] for k in self.kernels_for(spans)) + + def gpu_time_for_family(self, skeleton: str, instances: Sequence[dict]) -> float: + """GPU time attributable to one annotation family.""" + if self.strategy != self.STRATEGY_PROJECTION: + return self.gpu_time_by_correlation(instances) + spans = self._projection_union({skeleton}) + if not spans: + return 0.0 + return sum( + k["dur"] for k in self._kernels_in(spans.bounds) if spans.covers(k["ts"]) + ) + + # -- coverage ------------------------------------------------------------ + def audit( + self, annotations: Sequence[dict], roots: Sequence[dict] + ) -> CoverageReport: + """Measure GPU coverage by all annotations, and by the roots alone. + + High ``covered_any`` next to low ``covered_selected`` means the roots + sit at the wrong nesting level, and widening should fix it. + """ + # Credit the roots two ways: projections survive graph capture, and + # launch correlations work for roots no annotation produced. Matching by + # name alone scores every synthetic root zero and stalls escalation. + if self.strategy == self.STRATEGY_PROJECTION: + by_name = self._projection_union(_window_names(roots)) + any_spans = self._projection_union(None) + window = SpanSet.of_events(self.projections).bounds + else: + by_name = SpanSet() + any_spans = SpanSet.of_events(self.kernels_for(annotations)) + window = any_spans.bounds + + in_window = self._kernels_in(window) + busy = sum(k["dur"] for k in in_window) + if busy <= 0: + return CoverageReport(self.strategy, 0.0, 0.0, 0.0, 0.0, window) + + tiles, _ = build_root_tiles(roots) + tile_spans = [ + {"pid": pid, "tid": tid, "ts": start, "dur": end - start} + for (pid, tid, _), (start, end) in tiles.items() + ] + covered = {} + for label, spans in ( + ("any", any_spans), + ("spans", by_name | SpanSet.of_events(self.kernels_for(roots))), + ("tiles", by_name | SpanSet.of_events(self.kernels_for(tile_spans))), + ): + covered[label] = sum(k["dur"] for k in in_window if spans.covers(k["ts"])) + return CoverageReport( + self.strategy, + covered["any"] / busy, + covered["tiles"] / busy, + covered["spans"] / busy, + busy, + window, + ) + + def uncovered_kernels(self, annotations: Sequence[dict]) -> List[dict]: + """Kernels in the audit window that no annotation accounts for.""" + if self.strategy == self.STRATEGY_PROJECTION: + spans = self._projection_union(None) + window = SpanSet.of_events(self.projections).bounds + else: + spans = SpanSet.of_events(self.kernels_for(annotations)) + window = spans.bounds + return [k for k in self._kernels_in(window) if not spans.covers(k["ts"])] + + def _projection_union(self, names: Optional[set]) -> SpanSet: + """Union of projection spans, optionally restricted by annotation name. + + Taken across every GPU thread rather than per stream: a kernel on a side + stream inside an annotated region is genuinely covered, and per-stream + matching would call it uncovered and escalate for nothing. + """ + return SpanSet.of_events( + p + for p in self.projections + if names is None or name_skeleton(p.get("name", "")) in names + ) + + +def _window_names(roots: Sequence[dict]) -> set: + """Skeletons to match projections against for the selected roots. + + A root enriched in step 1.5 carries the inner annotation's name, so its + projections are recorded under the outer span it actually came from. + """ + names = set() + for r in roots: + prov = r.get(PROVENANCE_KEY) or {} + names.add(name_skeleton(prov.get("window_from") or r.get("name", ""))) + return names + + +# --- candidate enumeration -------------------------------------------------- +def ancestors_of(tree, event: dict, max_depth: int = 8) -> List[dict]: + """Enclosing spans above ``event``, nearest first. + + The depth bound is load-bearing: walk far enough and you reach the thread + entry point, a single root spanning the trace that passes coverage uselessly. + """ + out, current = [], event + for _ in range(max_depth): + parent = tree.get_parent_event(current) + if parent is None: + break + out.append(parent) + current = parent + return out + + +def in_gap_candidates( + index: IntervalIndex, gaps: Sequence[dict], categories: Sequence[str] +) -> List[dict]: + """Events of the given categories sitting in inter-root gaps.""" + out = [] + for gap in gaps: + out.extend( + e + for e in index.contained_in(gap, exclude_self=False) + if e.get("cat") in categories + ) + return out + + +# --- escalation ------------------------------------------------------------- +@dataclass +class Probe: + """One escalation attempt: ``applies_to`` gates it, ``run`` proposes roots.""" + + name: str + applies_to: Callable[[RootSet], bool] + run: Callable[[RootSet], Optional[RootSet]] + + +def run_probes( + root_set: RootSet, + probes: Sequence[Probe], + audit: Callable[[RootSet], CoverageReport], +) -> RootSet: + """Try probes in declared order, re-measuring coverage after each. + + Stops at the first that clears the gate, recording what every attempt did to + coverage. If none clears it, the best result is kept and graded against the + floor as degraded or unsplittable. + """ + attempts: List[dict] = [] + best = root_set + for probe in probes: + if not probe.applies_to(best): + attempts.append({"probe": probe.name, "outcome": "skipped"}) + continue + candidate = probe.run(best) + if candidate is None or not candidate.roots: + attempts.append({"probe": probe.name, "outcome": "no_candidate"}) + continue + candidate.coverage = audit(candidate) + before = best.coverage.covered_selected if best.coverage else 0.0 + after = candidate.coverage.covered_selected + attempts.append( + { + "probe": probe.name, + "outcome": "adopted" if after > before else "rejected", + "coverage_before": round(before, 4), + "coverage_after": round(after, 4), + "n_roots": len(candidate.roots), + } + ) + if after > before: + best = candidate + if best.coverage and best.coverage.passes: + break + + best.diagnostics["probes_run"] = attempts + coverage = best.coverage + if coverage and coverage.passes: + best.status = DetectStatus.SPLITTABLE + elif coverage and coverage.covered_selected >= COVERAGE_FLOOR: + best.status = DetectStatus.DEGRADED + else: + best.status = DetectStatus.NOT_SPLITTABLE + return best diff --git a/TraceLens/TraceUtils/split_inference/execution_roots.py b/TraceLens/TraceUtils/split_inference/execution_roots.py index 568ea7ac0..0bb0f50be 100644 --- a/TraceLens/TraceUtils/split_inference/execution_roots.py +++ b/TraceLens/TraceUtils/split_inference/execution_roots.py @@ -4,13 +4,58 @@ # See LICENSE for license information. ############################################################################### -"""Stage 1: find iteration execution roots in an inference trace.""" +"""Stage 1: find iteration execution roots in an inference trace. + +Two entry points. :func:`find_iteration_roots` is the original first-tier-wins +lookup, kept exactly as it was for callers that just want a root list. +:func:`find_iteration_roots_ex` runs the coverage-gated flow and reports how much +of the GPU's work the roots actually account for, which is the only way to tell a +correct root set from one that locked onto a warmup loop. +""" + +from typing import List, Optional, Sequence, Tuple from ..annotation_utils import ( ITERATION_BACKUP_PATTERNS, ITERATION_PATTERNS, + PROVENANCE_KEY, find_events_by_patterns, + find_iteration_roots_by_priority, + inherit_identity, + is_parseable, + name_skeleton, +) +from .detect_utils import ( + COVERAGE_FLOOR, + COVERAGE_GATE, + MIN_ROOTS, + DetectStatus, + GpuAttribution, + IntervalIndex, + PhaseConfidence, + RootSet, + run_probes, +) +from .root_detection import ( + NESTING_MAJORITY, + AnnotationFamily, + build_families, + collect_annotations, + detect_from_unknown_family, + detect_generic, + resolve_nesting, ) +from .root_probes import build_probes + +__all__ = [ + "COVERAGE_FLOOR", + "COVERAGE_GATE", + "DetectStatus", + "PhaseConfidence", + "RootSet", + "find_iteration_roots", + "find_iteration_roots_ex", +] def find_iteration_roots(events: list[dict]) -> list[dict] | None: @@ -38,3 +83,155 @@ def find_iteration_roots(events: list[dict]) -> list[dict] | None: roots = find_iteration_roots_generic(events) return roots + + +# --- step 1.5: separate the extraction window from the metadata -------------- +def _widen_to_outer_family( + roots: Sequence[dict], + families: Sequence[AnnotationFamily], + index: IntervalIndex, +) -> Tuple[Optional[List[dict]], Optional[AnnotationFamily]]: + """Replace ``roots`` with an enclosing family's instances, where one exists. + + Enclosing the known roots is how the family is *found*, but the whole family + is adopted: instances enclosing no known root are iterations whose inner + annotation is merely unrecognized. Keeping only the matching ones is how a + 512-iteration run gets split into the 15 a regex happened to know. + + Ranking by GPU time separates the iteration boundary from the bookkeeping + spans that wrap it equally well but do almost no work. + """ + root_ids = {id(r) for r in roots} + own = {name_skeleton(r.get("name", "")) for r in roots} + threshold = NESTING_MAJORITY * len(roots) + best = None + for family in families: + if family.skeleton in own or not family.regular: + continue + # Count the *roots* that end up wrapped, not the instances doing the + # wrapping. One outer span often holds several known roots, and counting + # instances then reports half when every root is in fact covered. + wrapped = { + id(e) + for instance in family.instances + for e in index.contained_in(instance) + if id(e) in root_ids + } + if len(wrapped) > threshold and (best is None or family.rank < best.rank): + best = family + if best is None: + return None, None + return sorted(best.instances, key=lambda e: e.get("ts", 0)), best + + +def _relabel_from_inner(roots: Sequence[dict], index: IntervalIndex) -> List[dict]: + """Give each root the identity of the parseable annotation inside it. + + The longest *parseable* inner annotation wins: an unparseable one leaves the + metadata fabricated and silently collapses batch size to one. + """ + relabelled = [] + for root in roots: + own = name_skeleton(root.get("name", "")) + inner = [ + e + for e in index.contained_in(root) + if is_parseable(e.get("name", "")) + and name_skeleton(e.get("name", "")) != own + ] + if inner: + root = inherit_identity(root, max(inner, key=lambda e: e.get("dur", 0))) + relabelled.append(root) + return relabelled + + +def _detect_annotated( + events: Sequence[dict], attribution: GpuAttribution +) -> Optional[RootSet]: + """Steps 1 and 1.5 for a trace with at least one recognized annotation.""" + known = find_iteration_roots_by_priority(events) + if not known: + return None + + annotations = collect_annotations(events) + index = IntervalIndex(annotations) + families = build_families(annotations, attribution) + resolve_nesting(families, index) + + diagnostics = { + "n_families": len(families), + "n_known_roots": len(known), + } + widened, outer = _widen_to_outer_family(known, families, index) + roots = widened if widened else list(known) + if outer is not None: + diagnostics["root_family_skeleton"] = outer.skeleton + diagnostics["root_family_known"] = outer.parseable + + roots = _relabel_from_inner(roots, index) + inherited = { + r[PROVENANCE_KEY]["identity_from"] for r in roots if PROVENANCE_KEY in r + } + if inherited: + diagnostics["inherited_from_skeleton"] = sorted( + {name_skeleton(n) for n in inherited} + ) + diagnostics["suspiciously_few_roots"] = len(roots) < MIN_ROOTS + + # Trust the phase labels only as far as they were actually parsed. Adopting a + # whole family means some of its iterations may carry no recognizable + # annotation at all, and calling that "high" would launder a guess. + labelled = sum(1 for r in roots if is_parseable(r.get("name", ""))) + diagnostics["n_roots_with_phase"] = labelled + if labelled == len(roots): + confidence = PhaseConfidence.HIGH + elif labelled: + confidence = PhaseConfidence.LOW + else: + confidence = PhaseConfidence.UNKNOWN + + return RootSet( + roots=sorted(roots, key=lambda e: e.get("ts", 0)), + method="annotation:widened" if widened else "annotation:tier", + phase_confidence=confidence, + diagnostics=diagnostics, + ) + + +def find_iteration_roots_ex(events: Sequence[dict]) -> RootSet: + """Find iteration roots and report how much GPU work they account for. + + Escalates only when it has to: recognized annotations, the nesting level + around them, a coverage audit, probes, then call-tree periodicity. + """ + attribution = GpuAttribution(events) + annotations = collect_annotations(events) + + root_set = _detect_annotated(events, attribution) + if root_set is None: + root_set = detect_from_unknown_family(events, attribution) + if root_set is None: + generic = detect_generic(events, attribution) + if generic is not None: + return generic + return RootSet( + roots=[], + method="none", + status=DetectStatus.NOT_SPLITTABLE, + diagnostics={"reason": "no annotations and no repeating call pattern"}, + ) + + def audit(candidate: RootSet): + return attribution.audit(annotations, candidate.roots) + + root_set.coverage = audit(root_set) + # Coverage can look fine while only a warmup loop matched, so a suspiciously + # small root count also escalates. + if root_set.coverage.passes and not root_set.diagnostics.get( + "suspiciously_few_roots" + ): + root_set.status = DetectStatus.SPLITTABLE + root_set.diagnostics["probes_run"] = [] + return root_set + + return run_probes(root_set, build_probes(events, attribution), audit) diff --git a/TraceLens/TraceUtils/split_inference/root_detection.py b/TraceLens/TraceUtils/split_inference/root_detection.py new file mode 100644 index 000000000..4c5c0c27e --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/root_detection.py @@ -0,0 +1,345 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Annotation families, and the detection steps built on them. + +Step 1 groups every annotation into families, known and unknown alike. Step 1.5 +chooses which nesting level is the iteration and where its metadata comes from. +Step 3 handles traces with no recognizable annotation, and step 5 falls back to +call-tree periodicity. +""" + +from collections import Counter +from dataclasses import dataclass, field +from statistics import mean, pstdev +from typing import Dict, List, Optional, Sequence, Tuple + +from ...Trace2Tree.inference_iteration_roots import ( + PERIOD_CONFLICT, + compare_periods, + find_iteration_roots_generic, + find_period_candidates, +) +from ..annotation_utils import ( + ANNOTATION_CAT, + inherit_identity, + is_parseable, + name_skeleton, +) +from .detect_utils import ( + MIN_ROOTS, + DetectStatus, + GpuAttribution, + IntervalIndex, + PhaseConfidence, + RootSet, +) + +# A family must enclose more than this share of another family's instances for +# the nesting relation to hold. Majority rather than unanimity, because real +# traces have ragged edges: warmup instances outside the loop, a truncated final +# iteration. Requiring every instance rejects the correct relation nearly always. +NESTING_MAJORITY = 0.5 + + +@dataclass +class AnnotationFamily: + """All instances of one logical annotation, keyed by its skeleton. + + CPU-side only: projections duplicate one annotation across streams, so any + count taken from them is inflated. They give the "has GPU work" signal only. + """ + + skeleton: str + instances: List[dict] + gpu_time: float = 0.0 + parseable: bool = False + interarrival_cv: float = 0.0 + encloses: Counter = field(default_factory=Counter) + + @property + def count(self) -> int: + return len(self.instances) + + @property + def regular(self) -> bool: + """Enough instances to be a per-iteration event rather than a one-off.""" + return self.count >= MIN_ROOTS + + @property + def rank(self) -> tuple: + """Sort key for choosing between families: most GPU work, steadiest.""" + return (-self.gpu_time, round(self.interarrival_cv, 3), -self.count) + + def is_outer_to(self, inner: "AnnotationFamily") -> bool: + """Whether this family encloses a majority of ``inner``'s instances.""" + if self.skeleton == inner.skeleton or not inner.count: + return False + return self.encloses.get(inner.skeleton, 0) / inner.count > NESTING_MAJORITY + + +def _interarrival_cv(instances: Sequence[dict]) -> float: + """Variation in the spacing between consecutive instances. + + A once-per-iteration annotation arrives on a steady cadence. Used to rank + families, not reject them, so it needs no threshold. + """ + stamps = sorted(e.get("ts", 0) for e in instances) + gaps = [b - a for a, b in zip(stamps, stamps[1:]) if b > a] + if len(gaps) < 2: + return 0.0 + average = mean(gaps) + return pstdev(gaps) / average if average else 0.0 + + +def collect_annotations(events: Sequence[dict]) -> List[dict]: + """CPU-side annotation events, in time order.""" + annotations = [ + e + for e in events + if e.get("cat") == ANNOTATION_CAT + and e.get("ts") is not None + and e.get("dur") is not None + ] + annotations.sort(key=lambda e: e["ts"]) + return annotations + + +def build_families( + annotations: Sequence[dict], attribution: GpuAttribution +) -> List[AnnotationFamily]: + """Group annotations into families and drop the ones with no GPU work. + + Pruning first is the point of the ordering: roughly half the families are + pure CPU scheduling chatter that nesting analysis would otherwise carry. + + A family survives if *any* instance has GPU work, not all -- a scheduler + family may have thousands of instances where only hundreds wrap real work. + """ + grouped: Dict[str, List[dict]] = {} + for event in annotations: + grouped.setdefault(name_skeleton(event.get("name", "")), []).append(event) + + families = [ + AnnotationFamily( + skeleton=skeleton, + instances=instances, + gpu_time=attribution.gpu_time_for_family(skeleton, instances), + parseable=any(is_parseable(e.get("name", "")) for e in instances), + interarrival_cv=_interarrival_cv(instances), + ) + for skeleton, instances in grouped.items() + ] + + # "No projection means no GPU work" is an inference, not a guarantee. If it + # would delete every family on a trace that plainly has kernels, it is wrong + # here -- fall back to following launch correlations instead. + if families and attribution.kernels and not any(f.gpu_time for f in families): + for family in families: + family.gpu_time = attribution.gpu_time_by_correlation(family.instances) + + return [f for f in families if f.gpu_time > 0] + + +def resolve_nesting(families: Sequence[AnnotationFamily], index: IntervalIndex) -> None: + """How many instances of other families each family encloses. + + Computed per family, not per event, which would be quadratic. + """ + for family in families: + counts: Counter = Counter() + for instance in family.instances: + for inner in index.contained_in(instance): + counts[name_skeleton(inner.get("name", ""))] += 1 + family.encloses = counts + + +# --- step 1.5: choose the nesting level and where metadata comes from -------- +def select_root_family( + families: Sequence[AnnotationFamily], +) -> Tuple[Optional[AnnotationFamily], List[AnnotationFamily]]: + """Outermost regular family that wraps something parseable. + + Separating "which span" from "which label" is the point: a scheduler span + can be the right window while carrying no metadata, and the annotation that + carries it can be too narrow. Returns the family and what it wraps. + """ + parseable = [f for f in families if f.parseable] + if not parseable: + return None, [] + + candidates = [] + for family in families: + if not family.regular: + continue + wrapped = [p for p in parseable if family.is_outer_to(p)] + if wrapped: + candidates.append((family, wrapped)) + if not candidates: + return None, [] + + # Outermost means not enclosed by another candidate. Being outermost is not + # sufficient on its own -- a whole-run wrapper would win every time -- which + # is why only regular families were considered above. + outermost = [ + (family, wrapped) + for family, wrapped in candidates + if not any(other.is_outer_to(family) for other, _ in candidates) + ] + pool = outermost or candidates + return min(pool, key=lambda item: item[0].rank) + + +def enrich_roots( + family: AnnotationFamily, index: IntervalIndex +) -> Tuple[List[dict], int, float]: + """Roots from ``family``, each carrying its inner annotation's identity. + + Instances enclosing nothing parseable do no iteration work and are dropped. + Their count and span share are returned because a large dropped share next + to passing coverage means the wrong nesting level. + """ + roots: List[dict] = [] + dropped_dur = 0.0 + total_dur = 0.0 + for instance in family.instances: + total_dur += instance.get("dur", 0) + inner = [ + e for e in index.contained_in(instance) if is_parseable(e.get("name", "")) + ] + if not inner: + dropped_dur += instance.get("dur", 0) + continue + # Longest *parseable* child, not longest child: an unparseable winner + # leaves the metadata fabricated and silently collapses batch size to 1. + source = max(inner, key=lambda e: e.get("dur", 0)) + roots.append(inherit_identity(instance, source)) + roots.sort(key=lambda e: e.get("ts", 0)) + dropped = family.count - len(roots) + return roots, dropped, (dropped_dur / total_dur if total_dur else 0.0) + + +# --- steps ------------------------------------------------------------------ +def detect_from_families( + events: Sequence[dict], attribution: GpuAttribution +) -> Optional[RootSet]: + """Steps 1 and 1.5: families, then the chosen nesting level. + + ``None`` when no family wraps anything parseable, handing off to step 3. + """ + annotations = collect_annotations(events) + if not annotations: + return None + + families = build_families(annotations, attribution) + if not families: + return None + + index = IntervalIndex(annotations) + resolve_nesting(families, index) + family, wrapped = select_root_family(families) + + diagnostics = { + "n_families": len(families), + "n_annotations": len(annotations), + } + if family is None: + return None + + roots, dropped, dropped_share = enrich_roots(family, index) + if not roots: + return None + + inner = min(wrapped, key=lambda f: f.rank) + diagnostics.update( + { + "root_family_skeleton": family.skeleton, + "root_family_known": family.parseable, + "inherited_from_skeleton": inner.skeleton, + "n_root_instances_dropped": dropped, + "dropped_gpu_time_share": round(dropped_share, 4), + "suspiciously_few_roots": len(roots) < MIN_ROOTS, + } + ) + known = "known" if family.parseable else "unknown" + return RootSet( + roots=roots, + method=f"family:{known}_outer+parseable_inner", + phase_confidence=PhaseConfidence.HIGH, + diagnostics=diagnostics, + ) + + +def detect_from_unknown_family( + events: Sequence[dict], attribution: GpuAttribution +) -> Optional[RootSet]: + """Step 3: no parseable annotation anywhere, but a regular family exists. + + Adopts the family doing the most GPU work on the steadiest cadence. Phases + are unknowable from an unrecognized name; the trace is still splittable. + """ + annotations = collect_annotations(events) + if not annotations: + return None + families = [f for f in build_families(annotations, attribution) if f.regular] + if not families: + return None + + family = min(families, key=lambda f: f.rank) + return RootSet( + roots=sorted(family.instances, key=lambda e: e.get("ts", 0)), + method="family:unknown_only", + phase_confidence=PhaseConfidence.UNKNOWN, + diagnostics={ + "n_families": len(families), + "root_family_skeleton": family.skeleton, + "root_family_known": False, + }, + ) + + +def detect_generic( + events: Sequence[dict], attribution: GpuAttribution +) -> Optional[RootSet]: + """Step 5: call-tree periodicity, cross-checked against the kernel stream. + + Two independent detections make this trustworthy without a coverage gate. + They are compared on *iteration count*, not raw period: the periods are + measured in different units -- python frames versus kernels -- so a kernel + loop at an exact multiple is the same loop at finer grain, and confirms it. + """ + diagnostics: dict = {} + roots = find_iteration_roots_generic(list(events), diagnostics) + if not roots: + return None + + kernel_names = [k.get("name", "") for k in attribution.kernels] + kernel_candidates = find_period_candidates(kernel_names) + kernel_blocks = kernel_candidates[0].repeats if kernel_candidates else None + verdict, ratio = compare_periods(len(roots), kernel_blocks) + + diagnostics.update( + { + "kernel_loop_blocks": kernel_blocks, + "period_agreement": verdict, + "generic_ratio_k": ratio, + } + ) + if verdict == PERIOD_CONFLICT: + return RootSet( + roots=roots, + method="generic:python_function", + phase_confidence=PhaseConfidence.UNKNOWN, + status=DetectStatus.NOT_SPLITTABLE, + diagnostics=diagnostics, + ) + return RootSet( + roots=roots, + method="generic:python_function", + phase_confidence=PhaseConfidence.UNKNOWN, + status=DetectStatus.SPLITTABLE, + diagnostics=diagnostics, + ) diff --git a/TraceLens/TraceUtils/split_inference/root_probes.py b/TraceLens/TraceUtils/split_inference/root_probes.py new file mode 100644 index 000000000..3fdc48c26 --- /dev/null +++ b/TraceLens/TraceUtils/split_inference/root_probes.py @@ -0,0 +1,254 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Step 4: escalation probes, tried in declared order when coverage falls short. + +Each probe proposes a different root set. They are ordered cheapest and most +reliable first, and the runner re-measures coverage after each one, so the order +is a declared preference rather than an implicit one -- every attempt and its +effect on coverage is recorded either way. +""" + +from typing import Callable, Dict, List, Optional, Sequence + +from ...Trace2Tree.inference_iteration_roots import ( + PERIOD_EXACT, + compare_periods, + find_iteration_roots_generic, + find_period_candidates, +) +from ..annotation_utils import inherit_identity, name_skeleton +from .detect_utils import ( + GpuAttribution, + IntervalIndex, + PhaseConfidence, + Probe, + RootSet, + ancestors_of, + gaps_between, + in_gap_candidates, +) + +# An ancestor level is only a plausible iteration boundary if there is roughly +# one ancestor per root. All roots sharing a single ancestor is the whole-run +# wrapper, which passes a coverage check while being useless. +ANCESTOR_COUNT_TOLERANCE = 0.8 +# How far to walk upward. Unbounded walking always reaches the thread entry point. +MAX_ANCESTOR_DEPTH = 8 +# Categories that can plausibly mark an iteration in an inter-root gap. +GAP_CANDIDATE_CATEGORIES = ("python_function", "cpu_op") + +SYNTHETIC_KEY = "synthetic" + + +def _synthesize(template: dict, ts: float, dur: float, probe: str, name: str) -> dict: + """A root that no annotation produced, tagged with where it came from.""" + return { + **{k: v for k, v in template.items() if k in ("pid", "tid", "cat")}, + "name": name, + "ts": ts, + "dur": dur, + SYNTHETIC_KEY: True, + "probe": probe, + } + + +def enclosing_ancestor_probe(events: Sequence[dict], attribution: GpuAttribution): + """Widen each root to an enclosing span that covers more GPU work. + + Narrow scope: step 1 already handles parents that are annotations, so what + is left here are parents that were never annotated at all. + """ + cache: Dict[str, object] = {} + + def _tree(): + if "tree" not in cache: + from ...Trace2Tree.trace_to_tree import TraceToTree + + try: + tree = TraceToTree(list(events), prune_nongpu_paths=False) + tree.build_tree(add_python_func=True) + except Exception as exc: # a malformed trace must not abort detection + print(f"Probe 4a: tree build failed ({exc}); skipping.") + tree = None + cache["tree"] = tree + return cache["tree"] + + def run(root_set: RootSet) -> Optional[RootSet]: + tree = _tree() + if tree is None: + return None + baseline = attribution.gpu_time_by_correlation(root_set.roots) + chains = [ancestors_of(tree, r, MAX_ANCESTOR_DEPTH) for r in root_set.roots] + for depth in range(MAX_ANCESTOR_DEPTH): + level = {} + for root, chain in zip(root_set.roots, chains): + if depth < len(chain): + level.setdefault(id(chain[depth]), (chain[depth], root)) + if len(level) < ANCESTOR_COUNT_TOLERANCE * len(root_set.roots): + continue + ancestors = [a for a, _ in level.values()] + if attribution.gpu_time_by_correlation(ancestors) <= baseline: + continue + roots = sorted( + (inherit_identity(a, r) for a, r in level.values()), + key=lambda e: e.get("ts", 0), + ) + return RootSet( + roots=roots, + method="probe:enclosing_ancestor", + phase_confidence=PhaseConfidence.LOW, + diagnostics={**root_set.diagnostics, "ancestor_depth": depth + 1}, + ) + return None + + return Probe( + name="4a_enclosing_ancestor", + applies_to=lambda rs: bool(rs.roots), + run=run, + ) + + +def generic_ratio_probe(events: Sequence[dict], attribution: GpuAttribution): + """Cross-check against call-tree periodicity, adopting only on exact agreement. + + An integer multiple of the root count is the same loop at finer grain, which + confirms the roots rather than replacing them. The ratio is recorded anyway. + """ + + def run(root_set: RootSet) -> Optional[RootSet]: + diagnostics: dict = {} + roots = find_iteration_roots_generic(list(events), diagnostics) + if not roots: + return None + verdict, ratio = compare_periods(len(roots), len(root_set.roots)) + merged = {**root_set.diagnostics, "generic_ratio_k": ratio, **diagnostics} + if verdict != PERIOD_EXACT: + return None + return RootSet( + roots=roots, + method="probe:generic_ratio", + phase_confidence=PhaseConfidence.UNKNOWN, + diagnostics=merged, + ) + + return Probe( + name="4b_generic_ratio", + applies_to=lambda rs: True, + run=run, + ) + + +def gap_family_probe(events: Sequence[dict], attribution: GpuAttribution): + """Recover work sitting in the gaps between roots. + + A family firing about once per gap is an iteration marker the annotations + missed. Recovered events are added to the roots rather than replacing them, + but carry no parseable identity, so phase confidence drops. + """ + index = IntervalIndex(events) + + def run(root_set: RootSet) -> Optional[RootSet]: + gaps = gaps_between(root_set.roots) + if not gaps: + return None + candidates = in_gap_candidates(index, gaps, GAP_CANDIDATE_CATEGORIES) + if not candidates: + return None + grouped: Dict[str, List[dict]] = {} + for event in candidates: + grouped.setdefault(name_skeleton(event.get("name", "")), []).append(event) + best = max( + grouped.items(), + key=lambda kv: attribution.gpu_time_by_correlation(kv[1]), + ) + skeleton, found = best + if attribution.gpu_time_by_correlation(found) <= 0: + return None + roots = sorted([*root_set.roots, *found], key=lambda e: e.get("ts", 0)) + return RootSet( + roots=roots, + method="probe:gap_family", + phase_confidence=PhaseConfidence.LOW, + diagnostics={**root_set.diagnostics, "gap_family_skeleton": skeleton}, + ) + + return Probe( + name="4c_gap_family", + applies_to=lambda rs: len(rs.roots) > 1, + run=run, + ) + + +def kernel_series_probe(events: Sequence[dict], attribution: GpuAttribution): + """Build roots from a repeating period in the kernels nothing accounts for. + + Last resort, and the only probe that works when the uncovered work has no + annotation near it. Roots must be CPU-side spans since extraction windows by + thread, so launch sites are used; a block launched entirely inside a + captured graph has none and is skipped. + """ + + def run(root_set: RootSet) -> Optional[RootSet]: + from .root_detection import collect_annotations + + uncovered = attribution.uncovered_kernels(collect_annotations(events)) + if len(uncovered) < 2: + return None + candidates = find_period_candidates([k.get("name", "") for k in uncovered]) + if not candidates: + return None + best = candidates[0] + roots = [] + for index in range(best.repeats): + lo = best.start + index * best.period + block = uncovered[lo : lo + best.period] + launches = attribution.cpu_launches_for(block) + if not launches: + continue + start = min(e["ts"] for e in launches) + end = max(e["ts"] + e["dur"] for e in launches) + roots.append( + _synthesize( + launches[0], + start, + end - start, + "4d_kernel_series", + f"synthetic:{name_skeleton(block[0].get('name', ''))}", + ) + ) + if not roots: + return None + roots.sort(key=lambda e: e.get("ts", 0)) + return RootSet( + roots=roots, + method="probe:kernel_series", + phase_confidence=PhaseConfidence.UNKNOWN, + diagnostics={ + **root_set.diagnostics, + "kernel_series_period": best.period, + "n_uncovered_kernels": len(uncovered), + }, + ) + + return Probe( + name="4d_kernel_series", + applies_to=lambda rs: bool(attribution.kernels), + run=run, + ) + + +PROBE_FACTORIES: Sequence[Callable] = ( + enclosing_ancestor_probe, + generic_ratio_probe, + gap_family_probe, + kernel_series_probe, +) + + +def build_probes(events: Sequence[dict], attribution: GpuAttribution) -> List[Probe]: + """Every probe, in the order they should be attempted.""" + return [factory(events, attribution) for factory in PROBE_FACTORIES] diff --git a/TraceLens/TraceUtils/split_inference/steady_state_window.py b/TraceLens/TraceUtils/split_inference/steady_state_window.py index a792b6915..9f0e6c072 100644 --- a/TraceLens/TraceUtils/split_inference/steady_state_window.py +++ b/TraceLens/TraceUtils/split_inference/steady_state_window.py @@ -4,9 +4,15 @@ # See LICENSE for license information. ############################################################################### -"""Stage 2: steady-state region detection and window selection.""" +"""Stage 2: steady-state region detection and window selection. + +Which window is "best" depends on what the workload is. Request concurrency is +the right signal for a serving trace and meaningless for a diffusion or training +one, so the workload is classified first and the objective chosen to match. +""" import math +from collections import Counter from statistics import mean from ..annotation_utils import ( @@ -14,8 +20,19 @@ is_decode_only, is_mixed, iteration_details, + name_skeleton, + parse_annotation, ) +WORKLOAD_SERVING = "serving" +WORKLOAD_DIFFUSION = "diffusion" +WORKLOAD_GENERIC = "generic" + +DIFFUSION_KINDS = ("diffusion_native",) + +# Share of roots a workload's annotations must account for to pick its objective. +CLASSIFICATION_MAJORITY = 0.5 + def identify_steady_state_regions( iter_details: list[dict], num_steps: int @@ -362,3 +379,145 @@ def _count_mixed(window: list[dict]) -> int: ) return iteration_roots[best["start"] : best["end"]] + + +# --- workload classification ------------------------------------------------ +def classify_workload(iteration_roots: list[dict]) -> tuple[str, dict]: + """Decide which window objective applies, from what the parsers recognized. + + Reads ``kind``, never the detail dict: an unrecognized name still yields a + full but fabricated detail dict reporting one decode-equivalent request, so + a detail-based test sees a flawless serving trace on any workload. + """ + kinds = Counter(parse_annotation(r.get("name", "")).kind for r in iteration_roots) + diffusion = sum(n for kind, n in kinds.items() if kind in DIFFUSION_KINDS) + serving = sum( + n for kind, n in kinds.items() if kind and kind not in DIFFUSION_KINDS + ) + + # A majority, not merely one match. Where fifteen of five hundred roots parse, + # the request counts driving the serving heuristic are fabricated for the rest, + # and a concurrency curve built from them describes nothing. + majority = CLASSIFICATION_MAJORITY * len(iteration_roots) + if serving > majority: + workload = WORKLOAD_SERVING + elif diffusion > majority: + workload = WORKLOAD_DIFFUSION + else: + workload = WORKLOAD_GENERIC + return workload, { + "workload_class": workload, + "n_recognized_roots": serving + diffusion, + "annotation_kinds": {k: n for k, n in kinds.items() if k}, + } + + +def _pattern_key(root: dict) -> tuple: + """What makes two iterations "the same shape". + + Diffusion steps at different resolutions are different work under one name, + so resolution joins the key when known. + """ + annotation = parse_annotation(root.get("name", "")) + resolution = annotation.resolution if annotation.kind in DIFFUSION_KINDS else None + return (name_skeleton(root.get("name", "")), resolution) + + +def _prefix_sums(values: list[float]) -> tuple[list[float], list[float]]: + """Running sums of ``values`` and of their squares.""" + totals, squares = [0.0], [0.0] + for value in values: + totals.append(totals[-1] + value) + squares.append(squares[-1] + value * value) + return totals, squares + + +def find_max_pattern_window( + iteration_roots: list[dict], + num_steps: int, + steady_state_regions: list[tuple[int, int]] | None = None, +) -> list[dict]: + """Pick the window that best matches the run's dominant iteration shape. + + With no request concurrency to track, "steady state" means the stretch that + looks most like the repeating pattern and runs most evenly. Ties break toward + the steadiest durations, which is what skips warmup: the same shape, but + slower and more erratic while caches fill and autotuning settles. + + The serving heuristic cannot answer this. Given fabricated request counts it + sees one request per step, so every step is "at peak", the region never + closes, and the result collapses to the first ``num_steps`` iterations. + """ + total = len(iteration_roots) + if not total: + return [] + + keys = [_pattern_key(r) for r in iteration_roots] + dominant, dominant_count = Counter(keys).most_common(1)[0] + matches, squares = _prefix_sums([1.0 if k == dominant else 0.0 for k in keys]) + durations, duration_squares = _prefix_sums( + [float(r.get("dur", 0)) for r in iteration_roots] + ) + + size = min(num_steps, total) + windows = [] + for start, end in steady_state_regions or [(0, total)]: + end = min(end, total) + if end - start < size: + if end > start: + windows.append((start, end)) + continue + windows.extend((s, s + size) for s in range(start, end - size + 1)) + if not windows: + windows = [(0, size)] + + def score(window: tuple[int, int]) -> tuple: + start, end = window + count = end - start + coverage = (matches[end] - matches[start]) / count + mean_dur = (durations[end] - durations[start]) / count + variance = (duration_squares[end] - duration_squares[start]) / count - ( + mean_dur * mean_dur + ) + cv = (max(variance, 0.0) ** 0.5 / mean_dur) if mean_dur else 0.0 + return (coverage, -cv) + + best = max(windows, key=score) + coverage, negative_cv = score(best) + print( + f"[pattern] Dominant iteration shape {dominant[0]!r} covers " + f"{dominant_count}/{total} roots. Selected [{best[0]}, {best[1]}): " + f"pattern_coverage={coverage:.3f}, duration_cv={-negative_cv:.3f}" + ) + return iteration_roots[best[0] : best[1]] + + +def select_window( + iteration_roots: list[dict], + num_steps: int, + steady_state_regions: list[tuple[int, int]] | None = None, + mode: str = "mixed", + **kwargs, +) -> tuple[list[dict], dict]: + """Choose a window with whichever objective the workload supports. + + Returns the window and a record of how it was chosen. + """ + workload, info = classify_workload(iteration_roots) + if workload == WORKLOAD_SERVING: + regions = steady_state_regions + if regions is None: + regions, _ = identify_steady_state_regions( + iteration_details(iteration_roots), num_steps + ) + window = find_steady_state_window( + iteration_roots, num_steps, regions, mode=mode, **kwargs + ) + info["window_strategy"] = f"steady_state:{mode}" + else: + window = find_max_pattern_window( + iteration_roots, num_steps, steady_state_regions + ) + info["window_strategy"] = "max_pattern_coverage" + info["n_window_roots"] = len(window) + return window, info diff --git a/TraceLens/TraceUtils/split_inference/trace_extraction.py b/TraceLens/TraceUtils/split_inference/trace_extraction.py index a035ea7b6..d14e1ad48 100644 --- a/TraceLens/TraceUtils/split_inference/trace_extraction.py +++ b/TraceLens/TraceUtils/split_inference/trace_extraction.py @@ -23,7 +23,16 @@ iteration_details, ) -GPU_EVENT_CATEGORIES = ["kernel", "gpu_memcpy", "gpu_memset", "gpu_user_annotation"] +from .detect_utils import ( + GPU_KERNEL_CATEGORIES, + PROJECTION_CATEGORY, + build_root_tiles, +) + +# Kernels plus the annotation projections that describe them. Anything summing +# GPU *time* must use GPU_KERNEL_CATEGORIES instead, since a projection encloses +# the kernels it describes and counting both double-counts. +GPU_EVENT_CATEGORIES = [*GPU_KERNEL_CATEGORIES, PROJECTION_CATEGORY] def get_filename(filepath: str) -> dict: @@ -72,8 +81,25 @@ def extract_iteration( gpu_corr_map: dict, flow_corr_map: dict, meta_events: list[dict], + root_tiles: dict | None = None, + gap_fill: bool = True, ) -> dict: - """Extract a single iteration trace.""" + """Extract a single iteration trace. + + Events are assigned to the window containing their *start* timestamp, which + makes the windows a partition of the timeline. Testing for full containment + instead would drop any event straddling a boundary from both neighbours, so + closing the gaps alone would not stop kernels going missing. + + Events longer than their window are enclosing spans -- thread roots, outer + python frames -- which belong to no single iteration and are left out. They + carry no correlation id, so no kernel is lost with them. + + ``root_tiles`` should come from :func:`build_root_tiles` over the *whole* root + list: built from a selected window instead, the window's last root would lose + the boundary of the root that follows it. Pass ``gap_fill=False`` to score + each root by its own span. + """ filtered_events = [] gpu_dur = 0 @@ -81,15 +107,29 @@ def extract_iteration( num_gpu_events = 0 batch_list = [] - # Pre-index GPU and flow events by correlation id - - # Compute the global time window for all iteration roots if not iteration_roots: return trace_json.copy(), [], 0, 0, 0 - min_iter_ts = min(root.get("ts", 0) for root in iteration_roots) - max_iter_end = max( - root.get("ts", 0) + root.get("dur", 0) for root in iteration_roots - ) + + if not gap_fill: + windows = [ + (r.get("ts", 0), r.get("ts", 0) + r.get("dur", 0)) for r in iteration_roots + ] + else: + tiles = ( + root_tiles + if root_tiles is not None + else build_root_tiles(iteration_roots)[0] + ) + windows = [ + tiles.get( + (r.get("pid"), r.get("tid"), r.get("ts", 0)), + (r.get("ts", 0), r.get("ts", 0) + r.get("dur", 0)), + ) + for r in iteration_roots + ] + + min_iter_ts = min(start for start, _ in windows) + max_iter_end = max(end for _, end in windows) # Collect all relevant tid/pid pairs tid_pid_set = {(root.get("tid"), root.get("pid")) for root in iteration_roots} @@ -102,22 +142,19 @@ def extract_iteration( dur = e.get("dur") if dur is None: continue - e_end = ts + dur e_tid = e.get("tid") e_pid = e.get("pid") - if (e_tid, e_pid) in tid_pid_set and ( - min_iter_ts <= ts and e_end <= max_iter_end - ): + if (e_tid, e_pid) in tid_pid_set and min_iter_ts <= ts <= max_iter_end: cpu_events.append(e) # For each iteration root, filter CPU events and collect correlation ids - for iteration_root in tqdm(iteration_roots): + for iteration_root, (win_ts, win_end) in zip(tqdm(iteration_roots), windows): start_time = [] end_time = [] iter_tid = iteration_root.get("tid") iter_pid = iteration_root.get("pid") - iter_ts = iteration_root.get("ts", 0) - iter_end = iter_ts + iteration_root.get("dur", 0) + win_dur = win_end - win_ts + is_last = win_end == max_iter_end correlation_ids: set[int] = set() @@ -125,14 +162,12 @@ def extract_iteration( for e in cpu_events: ts = e.get("ts") dur = e.get("dur") - e_end = ts + dur e_tid = e.get("tid") e_pid = e.get("pid") - if ( - e_tid == iter_tid - and e_pid == iter_pid - and (iter_ts <= ts and e_end <= iter_end) - ): + # Half-open so neighbouring windows cannot both claim an event; the + # final window is closed so nothing at the very end is orphaned. + within = win_ts <= ts < win_end or (is_last and ts == win_end) + if e_tid == iter_tid and e_pid == iter_pid and within and dur <= win_dur: filtered_events.append(e) corr = e.get("args", {}).get("correlation") if corr is not None: @@ -190,12 +225,16 @@ def extract_and_save( flow_corr_map: dict, meta_events: list[dict], output_label: str | None = None, + root_tiles: dict | None = None, ): """Extract and save a range of iterations. If ``output_label`` is provided the output filename becomes ``{output_label}_{name_append}_{base_name}.json.gz`` instead of the default ``{base_name}_{prefix}_{idx}_{name_append}.json.gz``. + + ``root_tiles`` should be built over the whole root list so that a root at the + edge of a selected window still knows where its successor begins. """ extraction_summary = [] # print(f"roots: {roots}") @@ -210,7 +249,13 @@ def extract_and_save( for idx, root in zip(indices, selected): iter_details = iteration_details(root) iter_trace, batch_list, num_gpu_events, gpu_dur, gpu_busy = extract_iteration( - root, events, trace_json, gpu_corr_map, flow_corr_map, meta_events + root, + events, + trace_json, + gpu_corr_map, + flow_corr_map, + meta_events, + root_tiles=root_tiles, ) is_annotation = "annotation_iteration" in prefix # Use the structured phase-aware name for any annotation extraction @@ -298,6 +343,7 @@ def extract_phases_and_save( gpu_corr_map: dict, flow_corr_map: dict, meta_events: list[dict], + root_tiles: dict | None = None, ): """Extract and save a range of iterations.""" extraction_summary = [] @@ -322,6 +368,7 @@ def extract_phases_and_save( gpu_corr_map, flow_corr_map, meta_events, + root_tiles=root_tiles, ) ) name_append = f"prefilldecode_{phase_details['num_prefilldecode']}_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" @@ -354,6 +401,7 @@ def extract_phases_and_save( gpu_corr_map, flow_corr_map, meta_events, + root_tiles=root_tiles, ) ) name_append = f"decode_{phase_details['num_decode']}_bs{phase_details['avg_bs']}_conc{phase_details['avg_conc']}" @@ -388,6 +436,7 @@ def divide_phases_and_save( flow_corr_map: dict, meta_events: list[dict], steady_state_regions: list[tuple[int, int]], + root_tiles: dict | None = None, ) -> list[dict]: """ Group contiguous steps of the same phase within steady-state regions and @@ -480,6 +529,7 @@ def divide_phases_and_save( flow_corr_map, meta_events, output_label=f"{phase}_{name_append}", + root_tiles=root_tiles, ) ) diff --git a/TraceLens/TraceUtils/split_inference_trace_annotation.py b/TraceLens/TraceUtils/split_inference_trace_annotation.py index e8a928955..6744cf5bf 100644 --- a/TraceLens/TraceUtils/split_inference_trace_annotation.py +++ b/TraceLens/TraceUtils/split_inference_trace_annotation.py @@ -166,18 +166,76 @@ # Re-exports for tests and downstream callers. from .split_inference import ( # noqa: F401 + DetectStatus, + build_root_tiles, + classify_workload, compute_reference_pd_ratio, divide_phases_and_save, extract_and_save, extract_iteration, extract_phases_and_save, find_iteration_roots, + find_iteration_roots_ex, + find_max_pattern_window, find_steady_state_window, get_filename, identify_steady_state_regions, parse_range, preprocess_trace, + select_window, ) +from .split_inference.detect_utils import GPU_KERNEL_CATEGORIES +from .split_inference.steady_state_window import WORKLOAD_SERVING + +MANIFEST_NAME = "split_manifest.json" + + +def _write_manifest(output_dir: str, manifest: dict) -> None: + """Record how the split was decided, next to the slices it produced. + + A split is only trustworthy if its quality is written down, so this is + emitted even when nothing was extracted. + """ + os.makedirs(output_dir, exist_ok=True) + path = os.path.join(output_dir, MANIFEST_NAME) + with open(path, "w") as f: + json.dump(manifest, f, indent=2) + print(f"Wrote split manifest to {path}") + + +def _conservation(events: list, per_iteration_details: list | None, args) -> dict: + """Check that slicing did not lose or duplicate GPU events. + + Only the one-file-per-iteration pass can be checked this way, and only over + the whole trace: those windows partition the timeline, so every kernel should + land in exactly one slice. Steady-state windows deliberately re-extract the + same kernels, so counting them too would compare a total against itself plus + overlap. + + The tiles span the iterations, not the capture, so a healthy trace still + leaves kernels unclaimed: warmup launched before the first root and teardown + after the last one. Those are excluded on purpose, which is why the failure + this reports is duplication rather than a shortfall -- extracting more than + exists means some kernel was counted under two iterations, and that is a bug. + The shortfall is reported as a quantity instead, since only its size is + interesting. + """ + kernels_in_trace = sum(1 for e in events if e.get("cat") in GPU_KERNEL_CATEGORIES) + report = {"n_gpu_events_in_trace": kernels_in_trace} + partitioned = ( + per_iteration_details is not None + and args.iterations == "all" + and not args.no_gap_fill + ) + if not (partitioned and kernels_in_trace): + return report + + extracted = sum(entry.get("num_gpu_events", 0) for entry in per_iteration_details) + report["n_gpu_events_extracted"] = extracted + report["n_gpu_events_outside_iterations"] = kernels_in_trace - extracted + report["gpu_events_duplicated"] = extracted > kernels_in_trace + report["gpu_event_retention"] = round(extracted / kernels_in_trace, 4) + return report def main(): @@ -250,8 +308,29 @@ def main(): "output_dir/decode_only/. Each step is a separate trace file." ), ) + parser.add_argument( + "--no-gap-fill", + action="store_true", + default=False, + help=( + "Score each iteration by its own annotation span instead of extending " + "it to the next root. Work between two roots is then dropped, as it " + "was before gap-free extraction; use this only to reproduce old output." + ), + ) + parser.add_argument( + "--allow-degraded", + action="store_true", + default=False, + help=( + "Continue even when the detected roots do not account for enough GPU " + "time. The manifest records the shortfall either way." + ), + ) args = parser.parse_args() execution_details = [] + # Only the partitioning pass can be checked for kernel conservation. + per_iteration_details: list | None = None # Load trace trace_json = DataLoader.load_data(get_filename(args.trace_path)) @@ -259,10 +338,53 @@ def main(): gpu_corr_map, flow_corr_map, meta_events = preprocess_trace(events) print(f"Loaded {len(events)} events") - iteration_roots = find_iteration_roots(events) + detection = find_iteration_roots_ex(events) + iteration_roots = detection.roots + manifest = detection.to_manifest() + print( + f"\nDetection: {detection.method} -> {len(iteration_roots)} roots, " + f"status={detection.status.name}, phase_confidence=" + f"{detection.phase_confidence.value}" + ) + if detection.coverage: + print( + f"GPU coverage ({detection.coverage.strategy}): " + f"{detection.coverage.covered_any:.1%} by any annotation, " + f"{detection.coverage.covered_selected:.1%} by the selected roots" + ) # Create output directory os.makedirs(args.output_dir, exist_ok=True) + + if detection.status is DetectStatus.NOT_SPLITTABLE and not args.allow_degraded: + manifest["aborted"] = True + _write_manifest(args.output_dir, manifest) + print( + "\nRefusing to split: the detected roots do not account for enough of " + "the GPU's work, so per-iteration slices would be misleading. " + f"See {MANIFEST_NAME} for the coverage breakdown, or pass " + "--allow-degraded to continue anyway." + ) + return + + workload, window_info = classify_workload(iteration_roots) + manifest.update(window_info) + print(f"Workload class: {workload}") + + # Built over every root, not just a selected window: the last root of a + # window still needs to know where the next one starts. + root_tiles = None + if not args.no_gap_fill and iteration_roots: + root_tiles, overlaps = build_root_tiles(iteration_roots) + manifest["gap_fill"] = True + manifest["n_overlapping_roots"] = overlaps + if overlaps: + print( + f"Warning: {overlaps} roots overlap their successor and keep their " + "own span; their windows are not gap-free." + ) + else: + manifest["gap_fill"] = False base_name = os.path.basename(args.trace_path) base_name = ( base_name.replace(".pt.trace", "").replace(".json.gz", "").replace(".json", "") @@ -286,7 +408,9 @@ def main(): gpu_corr_map, flow_corr_map, meta_events, + root_tiles=root_tiles, ) + per_iteration_details = temp_execution_details execution_details.extend(temp_execution_details) # Determine the working set and compute steady-state regions once, @@ -300,7 +424,12 @@ def main(): ) else: working_roots = iteration_roots - if args.find_steady_state or args.divide_phases: + # Request concurrency only means something for a serving trace. Asked + # about anything else it sees one request per step, calls every step a + # peak, and returns the first num_steps iterations -- warmup included. + if (args.find_steady_state or args.divide_phases) and ( + workload == WORKLOAD_SERVING + ): _iter_details = iteration_details(working_roots) steady_state_regions, _ = identify_steady_state_regions( _iter_details, args.num_steps @@ -330,7 +459,25 @@ def main(): gpu_corr_map, flow_corr_map, meta_events, - steady_state_regions=steady_state_regions, + steady_state_regions=steady_state_regions or [(0, len(working_roots))], + root_tiles=root_tiles, + ) + execution_details.extend(temp_execution_details) + + elif args.find_steady_state and workload != WORKLOAD_SERVING: + # Prefill and decode do not exist here, so there is one window to + # find: the stretch that best matches the run's repeating shape. + print("\n--- Finding representative window by iteration pattern ---") + pattern_roots = find_max_pattern_window( + working_roots, + num_steps=args.num_steps, + steady_state_regions=steady_state_regions or None, + ) + temp_execution_details = extract_and_save( + [pattern_roots], + *_extract_args, + output_label="pattern_steady_state", + root_tiles=root_tiles, ) execution_details.extend(temp_execution_details) @@ -347,7 +494,10 @@ def main(): R=args.R, ) temp_execution_details = extract_and_save( - [mixed_roots], *_extract_args, output_label="mixed_steady_state" + [mixed_roots], + *_extract_args, + output_label="mixed_steady_state", + root_tiles=root_tiles, ) execution_details.extend(temp_execution_details) @@ -359,7 +509,10 @@ def main(): mode="decode_only", ) temp_execution_details = extract_and_save( - [do_roots], *_extract_args, output_label="decode_only_steady_state" + [do_roots], + *_extract_args, + output_label="decode_only_steady_state", + root_tiles=root_tiles, ) execution_details.extend(temp_execution_details) @@ -371,11 +524,16 @@ def main(): mode="max_prefilldecode", ) temp_execution_details = extract_and_save( - [pd_roots], *_extract_args, output_label="prefilldecode_steady_state" + [pd_roots], + *_extract_args, + output_label="prefilldecode_steady_state", + root_tiles=root_tiles, ) execution_details.extend(temp_execution_details) print(f"\nDone! Extracted {len(execution_details)} traces to {args.output_dir}") + manifest.update(_conservation(events, per_iteration_details, args)) + _write_manifest(args.output_dir, manifest) if len(execution_details) > 0: json_path = os.path.join(args.output_dir, "execution_details.json") with open(json_path, "w") as f: diff --git a/tests/test_split_root_detection.py b/tests/test_split_root_detection.py new file mode 100644 index 000000000..e4e09cb8d --- /dev/null +++ b/tests/test_split_root_detection.py @@ -0,0 +1,793 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for the coverage-gated splitter: components and the flow built on them.""" + +from TraceLens.Trace2Tree.inference_iteration_roots import ( + PERIOD_CONFLICT, + PERIOD_EXACT, + PERIOD_INTEGER_RATIO, + compare_periods, + find_period_candidates, +) +from TraceLens.TraceUtils.annotation_utils import ( + PROVENANCE_KEY, + cluster_by_skeleton, + dominant_cluster, + inherit_identity, + is_parseable, + name_skeleton, + parse_annotation, +) +from TraceLens.TraceUtils.split_inference import ( + DetectStatus, + PhaseConfidence, + build_root_tiles, + classify_workload, + extract_iteration, + find_iteration_roots_ex, + find_max_pattern_window, + preprocess_trace, + select_window, +) +from TraceLens.TraceUtils.split_inference.detect_utils import ( + COVERAGE_GATE, + GpuAttribution, + IntervalIndex, + gaps_between, + group_by_thread, +) +from TraceLens.TraceUtils.split_inference.root_detection import ( + build_families, + collect_annotations, + resolve_nesting, +) + +VLLM = "execute_{i}_context_3(sq128sk256sqsq1sqsk1)_generation_2(sq1sk300sqsq1sqsk1)" + + +# --------------------------------------------------------------------------- # +# Event builders +# --------------------------------------------------------------------------- # +def annotation(name, ts, dur, pid=1, tid=10): + return { + "name": name, + "cat": "user_annotation", + "ph": "X", + "ts": ts, + "dur": dur, + "pid": pid, + "tid": tid, + "args": {}, + } + + +def launch(ts, corr, pid=1, tid=10, dur=2): + return { + "name": "hipLaunchKernel", + "cat": "cuda_runtime", + "ph": "X", + "ts": ts, + "dur": dur, + "pid": pid, + "tid": tid, + "args": {"correlation": corr}, + } + + +def kernel(ts, dur, corr, name="gemm", pid=1, tid=99): + return { + "name": name, + "cat": "kernel", + "ph": "X", + "ts": ts, + "dur": dur, + "pid": pid, + "tid": tid, + "args": {"correlation": corr}, + } + + +def projection(name, ts, dur, pid=1, tid=99): + return { + "name": name, + "cat": "gpu_user_annotation", + "ph": "X", + "ts": ts, + "dur": dur, + "pid": pid, + "tid": tid, + "args": {}, + } + + +def serving_trace(count=16, name_template=VLLM, period=1000, with_projection=False): + """One annotation per iteration, each launching one kernel.""" + events, corr = [], 500 + for i in range(count): + base = 1000 + i * period + name = name_template.format(i=i) + events.append(annotation(name, base, 100)) + events.append(launch(base + 10, corr)) + events.append(kernel(base + 200, 40, corr)) + if with_projection: + events.append(projection(name, base + 200, 40)) + corr += 1 + return events + + +# --------------------------------------------------------------------------- # +# C1: family keys +# --------------------------------------------------------------------------- # +class TestNameSkeleton: + def test_collapses_digit_runs(self): + assert ( + name_skeleton("execute_1_context_3(sq8sk8)") + == "execute_#_context_#(sq#sk#)" + ) + + def test_instances_of_one_operation_share_a_key(self): + a = name_skeleton("execute_1_context_3(sq128sk256)") + b = name_skeleton("execute_77_context_9(sq4sk4)") + assert a == b + + def test_distinguishes_genuinely_different_operations(self): + assert name_skeleton("step[DECODE bs=4]") != name_skeleton("step[EXTEND bs=4]") + + def test_separates_families_differing_only_after_a_prefix(self): + """A fixed-length prefix key would merge these; the skeleton must not.""" + long_a = "scheduler.process_batch_result_decode" + long_b = "scheduler.process_batch_result_extend" + assert name_skeleton(long_a) != name_skeleton(long_b) + + def test_cluster_and_dominant(self): + names = ["step[DECODE bs=1]", "step[DECODE bs=2]", "step[EXTEND bs=1 toks=8]"] + groups = cluster_by_skeleton(names) + assert len(groups) == 2 + skeleton, share = dominant_cluster(groups) + assert skeleton == "step[DECODE bs=#]" + assert share == 2 / 3 + + def test_dominant_of_nothing(self): + assert dominant_cluster({}) == (None, 0.0) + + +# --------------------------------------------------------------------------- # +# C6: identity and inheritance +# --------------------------------------------------------------------------- # +class TestAnnotationIdentity: + def test_memoized_instance_is_shared(self): + assert parse_annotation("step[DECODE bs=4]") is parse_annotation( + "step[DECODE bs=4]" + ) + + def test_parseable_reflects_recognition_not_the_detail_dict(self): + assert is_parseable("step[DECODE bs=4]") + assert not is_parseable("scheduler.process_batch_result") + + def test_unparseable_name_still_yields_a_full_detail_dict(self): + """Why classification must not read the numbers: the stub looks real.""" + details = parse_annotation("scheduler.process_batch_result").iter_details() + assert details["num_requests"] == 1 + assert details["context_requests"] == 0 + + def test_resolution_exists_on_every_annotation(self): + """Stage 2 keys on resolution; a missing attribute would raise.""" + assert parse_annotation("step[DECODE bs=4]").resolution is None + + def test_inherit_keeps_window_and_takes_identity(self): + outer = annotation("scheduler.process_batch_result", 1000, 500) + inner = annotation("step[DECODE bs=7]", 1050, 100) + merged = inherit_identity(outer, inner) + + assert merged["ts"] == 1000 and merged["dur"] == 500 + assert merged["name"] == "step[DECODE bs=7]" + assert parse_annotation(merged["name"]).generation_requests == 7 + assert merged[PROVENANCE_KEY] == { + "window_from": "scheduler.process_batch_result", + "identity_from": "step[DECODE bs=7]", + } + assert outer["name"] == "scheduler.process_batch_result" + + def test_second_inheritance_keeps_the_original_window_owner(self): + outer = annotation("scheduler.process_batch_result", 1000, 500) + once = inherit_identity(outer, annotation("step[DECODE bs=1]", 1050, 100)) + twice = inherit_identity(once, annotation("step[EXTEND bs=2 toks=8]", 1060, 10)) + assert twice[PROVENANCE_KEY]["window_from"] == "scheduler.process_batch_result" + assert twice[PROVENANCE_KEY]["identity_from"] == "step[EXTEND bs=2 toks=8]" + + +# --------------------------------------------------------------------------- # +# C2: periodicity +# --------------------------------------------------------------------------- # +class TestPeriodicity: + def test_skips_a_warmup_prefix(self): + best = find_period_candidates(["setup", "a", "b", "a", "b", "a", "b"])[0] + assert (best.period, best.start, best.repeats) == (2, 1, 3) + + def test_no_repetition_yields_nothing(self): + assert find_period_candidates(["a", "b", "c", "d"]) == [] + + def test_too_few_repeats_rejected(self): + assert find_period_candidates(["a", "b", "a", "b"]) == [] + + def test_reports_a_primitive_period_not_a_multiple(self): + """Every multiple of a valid period is valid; only the unit is useful.""" + candidates = find_period_candidates(["a", "b"] * 12) + assert candidates[0].period == 2 + assert all(c.period % 2 or c.period == 2 for c in candidates) + + def test_sub_iteration_noise_does_not_win(self): + """A launch repeating within the iteration must not become the period.""" + labels = (["step"] + ["launch"] * 5) * 8 + best = find_period_candidates(labels)[0] + assert best.period == 6 + assert best.repeats == 8 + + def test_duration_variance_is_reported(self): + labels = ["a", "b"] * 5 + steady = find_period_candidates(labels, durations=[10, 20] * 5)[0] + erratic = find_period_candidates( + labels, durations=[10, 20, 900, 20, 10, 20, 10, 500, 10, 20] + )[0] + assert steady.duration_cv == 0.0 + assert erratic.duration_cv > steady.duration_cv + + def test_coverage_reported(self): + best = find_period_candidates(["x", "y"] * 10)[0] + assert best.coverage == 1.0 + + def test_compare_periods(self): + assert compare_periods(4, 4) == (PERIOD_EXACT, 1) + assert compare_periods(4, 12) == (PERIOD_INTEGER_RATIO, 3) + assert compare_periods(4, 7) == (PERIOD_CONFLICT, None) + assert compare_periods(4, None) == (PERIOD_CONFLICT, None) + assert compare_periods(0, 4) == (PERIOD_CONFLICT, None) + + +# --------------------------------------------------------------------------- # +# C3: containment queries +# --------------------------------------------------------------------------- # +class TestIntervalIndex: + def test_finds_events_inside_a_span_excluding_itself(self): + outer = annotation("outer", 100, 100) + inner = annotation("inner", 120, 20) + index = IntervalIndex([outer, inner]) + assert index.contained_in(outer) == [inner] + assert index.contained_in(outer, exclude_self=False) == [outer, inner] + + def test_partial_overlap_is_not_containment(self): + outer = annotation("outer", 100, 100) + straddling = annotation("straddling", 150, 100) + index = IntervalIndex([outer, straddling]) + assert index.contained_in(outer) == [] + + def test_containment_never_crosses_threads(self): + outer = annotation("outer", 100, 100, tid=10) + other_thread = annotation("inner", 120, 20, tid=11) + index = IntervalIndex([outer, other_thread]) + assert index.contained_in(outer) == [] + + def test_ignores_events_without_duration(self): + outer = annotation("outer", 100, 100) + flow = {"name": "ac2g", "ph": "s", "ts": 110, "pid": 1, "tid": 10} + assert IntervalIndex([outer, flow]).contained_in(outer) == [] + + def test_gaps_between_consecutive_roots(self): + roots = [annotation("r", 100, 10), annotation("r", 200, 10)] + (gap,) = gaps_between(roots) + assert (gap["ts"], gap["dur"]) == (110, 90) + + def test_no_gap_when_roots_touch(self): + roots = [annotation("r", 100, 100), annotation("r", 200, 10)] + assert gaps_between(roots) == [] + + def test_group_by_thread_sorts_within_group(self): + events = [annotation("b", 200, 1, tid=10), annotation("a", 100, 1, tid=10)] + groups = group_by_thread(events) + assert [e["name"] for e in groups[(1, 10)]] == ["a", "b"] + + +# --------------------------------------------------------------------------- # +# C4: GPU attribution +# --------------------------------------------------------------------------- # +class TestGpuAttribution: + def test_prefers_projections_when_present(self): + attribution = GpuAttribution(serving_trace(4, with_projection=True)) + assert attribution.strategy == GpuAttribution.STRATEGY_PROJECTION + + def test_falls_back_to_correlation_without_projections(self): + attribution = GpuAttribution(serving_trace(4)) + assert attribution.strategy == GpuAttribution.STRATEGY_CORRELATION + + def test_projections_are_excluded_from_gpu_busy_time(self): + """Counting a projection as GPU time double-counts the kernels inside it.""" + events = serving_trace(4, with_projection=True) + annotations = collect_annotations(events) + assert GpuAttribution(events).audit(annotations, annotations).gpu_busy == 4 * 40 + + def test_full_coverage_when_every_kernel_is_annotated(self): + events = serving_trace(8) + annotations = collect_annotations(events) + report = GpuAttribution(events).audit(annotations, annotations) + assert report.covered_any == 1.0 + assert report.covered_selected == 1.0 + assert report.passes + + def test_work_just_outside_a_root_counts_once_windows_extend(self): + """Extraction captures the tile, so the audit must judge the tile. + + Mirrors vLLM sampling: each iteration launches work just after its + annotation ends. Judging bare spans would report a tenth of the GPU + unaccounted for and send detection looking for extra roots, splitting + every iteration in two to find work already being captured. + """ + events, corr = [], 10 + for i in range(12): + base = 1000 + i * 1000 + events.append(annotation(f"step[DECODE bs={i + 1}]", base, 400)) + events.append(launch(base + 10, corr)) + events.append(kernel(base + 100, 90, corr)) + corr += 1 + events.append(launch(base + 500, corr)) # after the annotation ends + events.append(kernel(base + 600, 10, corr)) + corr += 1 + + roots = collect_annotations(events) + report = GpuAttribution(events).audit(roots, roots) + assert report.covered_spans < 1.0 + assert report.covered_selected == 1.0 + assert report.span_share > 0.9 + assert report.passes + + def test_sparse_roots_stretched_over_many_iterations_do_not_pass(self): + """Coverage from window extension alone is not root coverage.""" + events = serving_trace(40) + annotations = collect_annotations(events) + report = GpuAttribution(events).audit(annotations, annotations[::20]) + assert report.covered_selected > report.covered_spans + assert report.span_share < 0.5 + assert not report.passes + + def test_gate_measures_the_roots_not_every_annotation(self): + """The 0.5.17 lesson: blanket annotation coverage is not a root check. + + A run whose annotations cover the whole timeline while the chosen roots + cover a fraction of its iterations must not pass. + """ + events = serving_trace(40) + annotations = collect_annotations(events) + report = GpuAttribution(events).audit(annotations, annotations[:2]) + assert report.covered_any == 1.0 + assert not report.passes + assert report.better_roots_exist + + def test_unannotated_work_lowers_coverage(self): + events = serving_trace(8) + # A kernel with no launch site inside any annotation. + events.append(kernel(1500, 4000, 99999, name="orphan")) + attribution = GpuAttribution(events) + report = attribution.audit(collect_annotations(events), []) + assert report.covered_any < COVERAGE_GATE + assert [ + k["name"] + for k in attribution.uncovered_kernels(collect_annotations(events)) + ] == ["orphan"] + + def test_selected_roots_can_cover_less_than_all_annotations(self): + """The signature of roots sitting at the wrong nesting level.""" + events = serving_trace(8) + annotations = collect_annotations(events) + attribution = GpuAttribution(events) + report = attribution.audit(annotations, annotations[:2]) + assert report.covered_any == 1.0 + assert report.covered_selected < report.covered_any + + def test_family_gpu_time_and_launch_sites(self): + events = serving_trace(4) + attribution = GpuAttribution(events) + annotations = collect_annotations(events) + skeleton = name_skeleton(annotations[0]["name"]) + assert attribution.gpu_time_for_family(skeleton, annotations) == 4 * 40 + assert len(attribution.cpu_launches_for(attribution.kernels)) == 4 + + def test_graph_launched_kernels_have_no_launch_site(self): + attribution = GpuAttribution([kernel(100, 10, 4242)]) + assert attribution.cpu_launches_for(attribution.kernels) == [] + + +# --------------------------------------------------------------------------- # +# C9: families +# --------------------------------------------------------------------------- # +class TestFamilies: + def _events(self): + """A scheduler family wrapping a decode family, plus CPU-only chatter.""" + events = [] + corr = 800 + for i in range(12): + base = 1000 + i * 1000 + events.append(annotation("scheduler.process_batch_result", base, 500)) + events.append(annotation(f"step[DECODE bs={i + 1}]", base + 50, 200)) + events.append(annotation("scheduler.log_stats", base + 700, 10)) + events.append(launch(base + 60, corr)) + events.append(kernel(base + 600, 300, corr)) + corr += 1 + return events + + def test_prunes_families_with_no_gpu_work(self): + events = self._events() + families = build_families(collect_annotations(events), GpuAttribution(events)) + skeletons = {f.skeleton for f in families} + assert "scheduler.log_stats" not in skeletons + assert {"scheduler.process_batch_result", "step[DECODE bs=#]"} <= skeletons + + def test_nesting_is_directional(self): + events = self._events() + annotations = collect_annotations(events) + families = build_families(annotations, GpuAttribution(events)) + index = IntervalIndex(annotations) + resolve_nesting(families, index) + by_skeleton = {f.skeleton: f for f in families} + outer = by_skeleton["scheduler.process_batch_result"] + inner = by_skeleton["step[DECODE bs=#]"] + assert outer.encloses["step[DECODE bs=#]"] == 12 + assert inner.encloses["scheduler.process_batch_result"] == 0 + + def test_regularity_needs_enough_instances(self): + events = serving_trace(3) + families = build_families(collect_annotations(events), GpuAttribution(events)) + assert families and not families[0].regular + + def test_parseability_recorded_per_family(self): + events = self._events() + families = build_families(collect_annotations(events), GpuAttribution(events)) + by_skeleton = {f.skeleton: f for f in families} + assert by_skeleton["step[DECODE bs=#]"].parseable + assert not by_skeleton["scheduler.process_batch_result"].parseable + + +# --------------------------------------------------------------------------- # +# Stage 1 end to end +# --------------------------------------------------------------------------- # +class TestDetectionFlow: + def test_healthy_trace_resolves_without_probes(self): + result = find_iteration_roots_ex(serving_trace(16)) + assert result.status is DetectStatus.SPLITTABLE + assert result.phase_confidence is PhaseConfidence.HIGH + assert result.method == "annotation:tier" + assert len(result) == 16 + assert result.coverage.covered_any == 1.0 + assert result.diagnostics["probes_run"] == [] + + def test_whole_outer_family_is_adopted_not_just_matching_instances(self): + """The 0.5.17 shape: most iterations wrap an unrecognized annotation. + + Only the first three iterations carry a name a parser knows. Keeping just + those would split a twenty-iteration run into three. + """ + events, corr = [], 400 + for i in range(20): + base = 1000 + i * 1000 + events.append(annotation("scheduler.run_batch", base, 500)) + inner = ( + f"step[DECODE bs={i + 1}]" if i < 3 else f"step[TARGET_VERIFY bs={i}]" + ) + events.append(annotation(inner, base + 50, 200)) + events.append(launch(base + 60, corr)) + events.append(kernel(base + 600, 300, corr)) + corr += 1 + + result = find_iteration_roots_ex(events) + assert len(result) == 20 + assert result.method == "annotation:widened" + assert result.diagnostics["root_family_skeleton"] == "scheduler.run_batch" + # Three roots parsed, seventeen did not, so the phases are not all real. + assert result.diagnostics["n_roots_with_phase"] == 3 + assert result.phase_confidence is PhaseConfidence.LOW + assert result.status is DetectStatus.SPLITTABLE + + def test_outer_family_found_when_it_wraps_several_roots_each(self): + """One outer span per two known roots still counts as wrapping them. + + Counting enclosing instances rather than enclosed roots reports half here + and abandons the widening, keeping the two roots instead of all twelve. + """ + events, corr = [], 600 + for i in range(12): + base = 1000 + i * 1000 + events.append(annotation("scheduler.run_batch", base, 800)) + for step in range(2): + inner = ( + f"step[DECODE bs={step + 1}]" if i < 6 else f"step[UNKNOWN {step}]" + ) + events.append(annotation(inner, base + 50 + step * 300, 200)) + events.append(launch(base + 60 + step * 300, corr)) + events.append(kernel(base + 900 + step * 50, 100, corr)) + corr += 1 + + result = find_iteration_roots_ex(events) + assert result.method == "annotation:widened" + assert len(result) == 12 + assert result.diagnostics["root_family_skeleton"] == "scheduler.run_batch" + + def test_unknown_outer_family_becomes_the_window(self): + """The useful span has a name no regex knows.""" + events, corr = [], 900 + for i in range(20): + base = 1000 + i * 1000 + events.append(annotation("scheduler.process_batch_result", base, 500)) + events.append(annotation(f"step[DECODE bs={i + 1}]", base + 50, 200)) + events.append(launch(base + 60, corr)) + events.append(kernel(base + 600, 300, corr)) + corr += 1 + + result = find_iteration_roots_ex(events) + assert result.status is DetectStatus.SPLITTABLE + assert result.method == "annotation:widened" + assert len(result) == 20 + # Window from the scheduler span, identity from the decode annotation. + assert result.roots[0]["dur"] == 500 + assert result.roots[0][PROVENANCE_KEY] == { + "window_from": "scheduler.process_batch_result", + "identity_from": "step[DECODE bs=1]", + } + assert result.phase_confidence is PhaseConfidence.HIGH + assert result.diagnostics["root_family_known"] is False + + def test_inner_annotation_relabels_a_known_root(self): + """The 0.5.11 shape: the outer name parses, the inner one is truer.""" + events, corr = [], 700 + for i in range(12): + base = 1000 + i * 1000 + events.append(annotation(VLLM.format(i=i), base, 500)) + events.append(annotation(f"step[DECODE bs={i + 1}]", base + 50, 200)) + events.append(launch(base + 60, corr)) + events.append(kernel(base + 600, 300, corr)) + corr += 1 + + result = find_iteration_roots_ex(events) + assert len(result) == 12 + assert result.roots[0]["dur"] == 500 + assert result.roots[0][PROVENANCE_KEY]["identity_from"] == "step[DECODE bs=1]" + + def test_unrecognized_annotations_are_still_splittable(self): + events, corr = [], 300 + for i in range(10): + base = 1000 + i * 1000 + events.append(annotation("my_custom_step", base, 400)) + events.append(launch(base + 10, corr)) + events.append(kernel(base + 500, 200, corr)) + corr += 1 + + result = find_iteration_roots_ex(events) + assert len(result) == 10 + assert result.method == "family:unknown_only" + assert result.phase_confidence is PhaseConfidence.UNKNOWN + assert result.status is not DetectStatus.NOT_SPLITTABLE + + def test_empty_trace_reports_not_splittable(self): + result = find_iteration_roots_ex([]) + assert result.status is DetectStatus.NOT_SPLITTABLE + assert len(result) == 0 + + def test_uncovered_work_triggers_probes_and_is_recorded(self): + events = serving_trace(16) + events.append(kernel(1500, 500_000, 99999, name="unaccounted")) + result = find_iteration_roots_ex(events) + assert result.diagnostics["probes_run"], "escalation must be recorded" + assert result.coverage.covered_any < COVERAGE_GATE + assert result.status in (DetectStatus.DEGRADED, DetectStatus.NOT_SPLITTABLE) + + def test_manifest_reports_quality(self): + manifest = find_iteration_roots_ex(serving_trace(16)).to_manifest() + assert manifest["status"] == 0 + assert manifest["phase_confidence"] == "high" + assert manifest["n_roots"] == 16 + assert manifest["coverage_any_annotation"] == 1.0 + assert manifest["attribution_strategy"] == "correlation" + + +# --------------------------------------------------------------------------- # +# Stage 3: tiling +# --------------------------------------------------------------------------- # +class TestRootTiles: + def test_windows_touch_so_gaps_belong_to_somebody(self): + roots = [annotation("r", 1000, 100), annotation("r", 2000, 100)] + tiles, overlaps = build_root_tiles(roots) + assert overlaps == 0 + assert tiles[(1, 10, 1000)] == (1000, 2000) + + def test_last_window_uses_the_median_length(self): + roots = [annotation("r", 1000, 100), annotation("r", 2000, 100)] + tiles, _ = build_root_tiles(roots) + assert tiles[(1, 10, 2000)] == (2000, 3000) + + def test_single_root_keeps_its_own_span(self): + tiles, _ = build_root_tiles([annotation("r", 1000, 100)]) + assert tiles[(1, 10, 1000)] == (1000, 1100) + + def test_threads_are_tiled_independently(self): + roots = [ + annotation("r", 1000, 100, tid=10), + annotation("r", 1500, 100, tid=11), + annotation("r", 3000, 100, tid=10), + ] + tiles, _ = build_root_tiles(roots) + assert tiles[(1, 10, 1000)] == (1000, 3000) + assert tiles[(1, 11, 1500)] == (1500, 1600) + + def test_overlapping_roots_keep_their_own_span_and_are_counted(self): + roots = [annotation("r", 1000, 900), annotation("r", 1500, 100)] + tiles, overlaps = build_root_tiles(roots) + assert overlaps == 1 + assert tiles[(1, 10, 1000)] == (1000, 1900) + + +class TestGapFreeExtraction: + def _trace(self): + """Each iteration launches one kernel inside its root and one after it.""" + events, corr = [], 10 + for i in range(3): + base = 1000 + i * 1000 + events.append(annotation(f"step[DECODE bs={i + 1}]", base, 100)) + events.append(launch(base + 10, corr)) + events.append(kernel(base + 300, 20, corr, name="k_in_root")) + corr += 1 + events.append(launch(base + 400, corr)) # after the annotation ends + events.append(kernel(base + 500, 30, corr, name="k_in_gap")) + corr += 1 + return events + + def test_gap_kernels_are_recovered(self): + events = self._trace() + trace = {"traceEvents": events} + gpu_map, flow_map, meta = preprocess_trace(events) + roots = collect_annotations(events) + tiles, _ = build_root_tiles(roots) + + _, _, dropped, _, _ = extract_iteration( + roots, events, trace, gpu_map, flow_map, meta, gap_fill=False + ) + out, _, kept, _, busy = extract_iteration( + roots, events, trace, gpu_map, flow_map, meta, root_tiles=tiles + ) + assert dropped == 3 + assert kept == 6 + assert busy == 3 * 50 + names = {e["name"] for e in out["traceEvents"] if e.get("cat") == "kernel"} + assert names == {"k_in_root", "k_in_gap"} + + def test_every_kernel_lands_in_exactly_one_window(self): + events = self._trace() + trace = {"traceEvents": events} + gpu_map, flow_map, meta = preprocess_trace(events) + roots = collect_annotations(events) + tiles, _ = build_root_tiles(roots) + + per_root = [ + extract_iteration( + [r], events, trace, gpu_map, flow_map, meta, root_tiles=tiles + )[2] + for r in roots + ] + total_in_trace = sum(1 for e in events if e.get("cat") == "kernel") + assert sum(per_root) == total_in_trace + + def test_warmup_before_the_first_root_is_excluded(self): + """Tiles span the iterations, not the capture. + + Work launched before the first root belongs to no iteration, so the + per-iteration counts legitimately fall short of the trace total. Reading + that shortfall as lost kernels reports healthy traces as broken; the + failure worth detecting is a kernel claimed by two iterations. + """ + events = self._trace() + events.append(launch(500, 99)) + events.append(kernel(600, 40, 99, name="k_warmup")) + trace = {"traceEvents": events} + gpu_map, flow_map, meta = preprocess_trace(events) + roots = collect_annotations(events) + tiles, _ = build_root_tiles(roots) + + per_root = [ + extract_iteration( + [r], events, trace, gpu_map, flow_map, meta, root_tiles=tiles + )[2] + for r in roots + ] + total_in_trace = sum(1 for e in events if e.get("cat") == "kernel") + assert sum(per_root) == total_in_trace - 1 + assert sum(per_root) <= total_in_trace + + def test_enclosing_spans_are_left_out(self): + """An outer frame belongs to no single iteration.""" + events = self._trace() + events.append( + { + "name": "whole_run", + "cat": "python_function", + "ph": "X", + "ts": 900, + "dur": 5000, + "pid": 1, + "tid": 10, + "args": {}, + } + ) + trace = {"traceEvents": events} + gpu_map, flow_map, meta = preprocess_trace(events) + roots = collect_annotations(events) + tiles, _ = build_root_tiles(roots) + out, _, _, _, _ = extract_iteration( + [roots[0]], events, trace, gpu_map, flow_map, meta, root_tiles=tiles + ) + assert "whole_run" not in {e["name"] for e in out["traceEvents"]} + + +# --------------------------------------------------------------------------- # +# Stage 2: workload classification and window choice +# --------------------------------------------------------------------------- # +class TestWorkloadClassification: + def test_serving_recognized(self): + roots = collect_annotations(serving_trace(8)) + workload, info = classify_workload(roots) + assert workload == "serving" + assert info["n_recognized_roots"] == 8 + + def test_unrecognized_names_are_generic_not_decode_only(self): + """Classification must not be fooled by the fabricated detail dict.""" + roots = [annotation("denoise_step", 1000 + i * 100, 50) for i in range(8)] + workload, info = classify_workload(roots) + assert workload == "generic" + assert info["n_recognized_roots"] == 0 + assert info["annotation_kinds"] == {} + + def test_a_recognized_minority_does_not_make_it_serving(self): + """Most steps unparsed means most request counts are invented.""" + roots = [ + annotation("scheduler.run_batch", 1000 + i * 100, 50) for i in range(40) + ] + roots += [ + annotation(f"step[DECODE bs={i}]", 9000 + i * 100, 50) for i in range(3) + ] + workload, info = classify_workload(roots) + assert workload == "generic" + assert info["n_recognized_roots"] == 3 + + def test_pattern_window_skips_erratic_warmup(self): + """Same shape throughout, so duration steadiness decides.""" + roots = [] + for i in range(4): # warmup: same name, wildly uneven durations + roots.append(annotation("denoise_step", 1000 + i * 5000, 4000 - i * 900)) + for i in range(8): # steady state + roots.append(annotation("denoise_step", 30000 + i * 1000, 900)) + + window = find_max_pattern_window(roots, num_steps=4) + assert len(window) == 4 + assert all(r["ts"] >= 30000 for r in window) + + def test_pattern_window_prefers_the_dominant_shape(self): + roots = [annotation("odd_step", 1000 + i * 100, 50) for i in range(3)] + roots += [annotation("denoise_step", 2000 + i * 100, 50) for i in range(9)] + window = find_max_pattern_window(roots, num_steps=4) + assert {r["name"] for r in window} == {"denoise_step"} + + def test_pattern_window_handles_a_short_run(self): + roots = [annotation("denoise_step", 1000 + i * 100, 50) for i in range(2)] + assert len(find_max_pattern_window(roots, num_steps=8)) == 2 + + def test_pattern_window_of_nothing(self): + assert find_max_pattern_window([], num_steps=4) == [] + + def test_dispatch_records_the_strategy_used(self): + serving = collect_annotations(serving_trace(24)) + _, info = select_window(serving, num_steps=4, steady_state_regions=[(0, 24)]) + assert info["window_strategy"].startswith("steady_state:") + + generic = [annotation("denoise_step", 1000 + i * 100, 50) for i in range(12)] + window, info = select_window(generic, num_steps=4) + assert info["window_strategy"] == "max_pattern_coverage" + assert info["n_window_roots"] == len(window) == 4