Generic trace splitting support

[kyle-hoffmeyer]

Summary

This PR adds robustness to trace splitting logic, allowing it to gracefully handle traces from a variety of different frameworks and workload types.

Motivation

The trace splitter was built to split inference traces via the annotations that TraceLens recognizes. There was some support to handle generic traces of any workload and without any of these annotations, but it was faulty and could not handle traces that differed from the norm:

• Speculative decode traces
• Traces with limited call stack information

Solution

The detection pipeline is now a flat, three-step cascade in a single function (find_iteration_roots in execution_roots.py):

1. Annotation-based detection. If the trace has recognized iteration markers (vLLM's execute_*, SGLang's step[*], ProfilerStep#*, etc.), use them directly. These annotations are known to explictly mark each iteration. This mirrors the original trace splitting implementation. When a recognized pattern is found, the detector also looks for an enclosing family to widen the split (so a 512-iteration run isn't split into only the 15 iterations a regex happened to know). Coverage is audited by correlating GPU kernels back to the annotation spans — if the annotations explain enough GPU work, accept the split. If the known annotations don't provide good coverage try using unrecognized annotations.

After the annotation-based detection, the entire trace tree is used to determine what events can be used as splitting points. TraceToTree is ran and a new function called _reattach_worker_threads is ran. This function allows training traces to be traversable by reattaching the call stacks in the worker threads to the call stack in the main thread.
The following steps are then ran:

2. Branch descent on the call tree. BFS from the tree's roots looking for the first frame whose children form a repeating pattern with high GPU coverage. diffusion denoise steps, and any workload with a regular call-tree structure. The coverage gate (95% of GPU time) prevents the descent from stopping on a shallow sub-loop (like a grad-norm poll that repeats but does almost no work). If the best candidate found has poor coverage, the detector still returns it rather than None, letting the cascade's fallback logic decide what to do.

3. Sibling-root periodicity. For some traces (especially those with missing python call stacks) contain many roots with short call stacks instead of one root with a sprawling call stack. In this case, the repeating pattern exists across roots and the branch descent step can not find the pattern. Sibling-root detection finds the repeating period across the roots.

Each step returns a RootSet dataclass with information about the splitting result. Each RootSet object carries a coverage metric, indicating how much of the GPU time is covered by the splits. If the coverage is too (<95%) then the subsequent step is ran. If nothing passes this coverage gate, the best failed candidate is returned with its honest status so the caller can decide whether to use it anyway (with --allow-degraded).

Detection Flow

flowchart TD
  START["find_iteration_roots<br/>(entry point)"] --> PREP["Collect annotations"]

  %% Step 1: Annotations
  PREP --> ANN["Step 1: _detect_from_annotations<br/><br/> Try known annotations then unknown annotations"]
  ANN --> ANN_AUDIT["Audit GPU coverage:<br/>correlate kernels to annotation spans"]
  ANN_AUDIT --> ANN_GATE{"coverage ≥ 95%?"}
  ANN_GATE -->|yes| OK_ANN["SPLITTABLE<br/>method = annotation:tier / widened / family:unknown_only"]
  ANN_GATE -->|no| TREE["Build TraceToTree <br/>Reattach cross-thread worker roots<br/>"] --> BD

  %% Step 2: Branch descent
  BD["Step 2: detect_from_branch_descent<br/><br/> Search for repeating pattern of events"] --> BD_GATE{"SPLITTABLE<br/>(coverage ≥ 95%)?"}
  BD_GATE -->|yes| OK_BD["SPLITTABLE<br/>method = generic:branch_descent"]
  BD_GATE -->|no| SR

  %% Step 3: Sibling roots
  SR["Step 3: detect_from_sibling_roots<br/><br/>Check if the top-level entry roots themselves form a repeating pattern"] --> SR_GATE{"SPLITTABLE<br/>(coverage ≥ 95%)?"}
  SR_GATE -->|yes| OK_SR["SPLITTABLE<br/>method = generic:sibling_roots"]
  SR_GATE -->|no| BEST

  %% Pick best across all detectors
  BEST["Pick the best result across all detectors"]

  %% Converge to extraction
  OK_ANN --> EXT
  OK_BD --> EXT
  OK_SR --> EXT
  BEST --> ABRT{"NOT_SPLITTABLE and<br/>no --allow-degraded?"}
  ABRT -->|yes| STOP["Write manifest, ABORT"]
  ABRT -->|no| EXT

  EXT["Tiling + extraction:<br/>build gap-free per-thread tiles,<br/>assign events by start timestamp,<br/>follow correlation IDs to GPU kernels"] --> OUT["Write per-iteration slices +<br/>split_manifest.json +<br/>execution_details.csv"]

Loading

When does each detector fire?

Detector	Fires when	Example workloads
Annotation-based	Trace has recognized user_annotation events with enough GPU coverage	vLLM, SGLang, Megatron (ProfilerStep), any framework emitting profiler.step()
Branch descent	No annotations (or poor annotation coverage), but the call tree has a frame whose children repeat	Training loops (forward/backward), diffusion denoise steps, torch.compile workloads
Sibling roots	Branch descent fails (no repeating children), but the top-level frames themselves repeat	Workloads where call stack information is sparse

After splitting completes, split_inference_trace_annotation.py writes a split_manifest.json alongside the split traces. This is a machine-readable record of the detection and extraction result:

{
    "status": 0,
    "method": "family:unknown_only",
    "phase_confidence": "unknown",
    "n_roots": 11,
    "attribution_strategy": "projection",
    "coverage_any_annotation": 1.0,
    "coverage_selected_roots": 1.0,
    "root_family_skeleton": "ProfilerStep##",
    "gap_fill": true,
    "n_gpu_events_in_trace": 24069,
    "n_gpu_events_extracted": 24069,
    "gpu_events_duplicated": false,
    "gpu_event_retention": 1.0
}

Key fields:

• status: 0 = splittable, 1 = not splittable, 2 = degraded (needs --allow-degraded)
• method: which detection path produced the roots
• attribution_strategy: projection (GPU-side annotation spans) or correlation (CPU launch correlation IDs) — how GPU kernels were mapped to annotations for coverage auditing
• coverage_selected_roots: fraction of GPU time explained by the selected roots' tile windows — the primary quality metric
• gpu_event_retention: fraction of GPU kernels that survived extraction — should be 1.0 (all kernels accounted for across all splits)
• gpu_events_duplicated: whether any kernel was claimed by more than one split (indicates a tiling bug)

Next steps

• Implement generic steady state identification
• Allow for recognition of multiple patterns in a list of events (ex. prefill -> prefill/decode -> decode)

[gabeweisz]

Looks good to me - fix linting and I'll approve

[codecov-commenter]

Codecov Report

❌ Patch coverage is 85.49075% with 102 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
TraceLens/Trace2Tree/inference_iteration_roots.py	75.42%	43 Missing ⚠️
...eLens/TraceUtils/split_inference/root_detection.py	78.37%	24 Missing ⚠️
...Lens/TraceUtils/split_inference/execution_roots.py	83.56%	12 Missing ⚠️
...ens/TraceUtils/split_inference_trace_annotation.py	73.33%	12 Missing ⚠️
...aceLens/TraceUtils/split_inference/detect_utils.py	95.43%	10 Missing ⚠️
...ens/TraceUtils/split_inference/trace_extraction.py	97.22%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!