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