Generic phase splitting and steady state identification

[kyle-hoffmeyer]

Summary

Steady state identification and phase division previously required LLM inference serving annotations — execute_* from vLLM, step[DECODE/EXTEND] from SGLang, or prefill[]/decode[] from ATOM. Steady state identification for LLM inference traces without these annotations as well as general traces was flaky.

This PR introduces a tiered system where phase splitting and steady state identification adapt to the information available in the trace.

Implementation

Phase splitting (2 tiers):
Phase splitting is only relevant for LLM inference traces with distinct phases. Since there isn't a simple, reliable way of identifying LLM inference traces, we rely on the user to pass a new --llm-inference flag when splitting a LLM inference trace.

Tier	Condition	Method
Annotation-based	Trace has vLLM/SGLang/ATOM annotations	Parse context_requests / generation_requests from annotation names (unchanged)
Shape-based	--llm-inference flag, no serving annotations	Derive batch size per iteration from the most common first dimension of cpu_op Input Dims. Extracting batch size this way has been found to be the most reliable. Iterations with batch size > 2x median are classified as prefill-bearing; the rest as decode.

Tier 1 — Annotation-based phase splitting:

flowchart LR
    A[Iteration root] --> B[Parse annotation name]
    B --> C{context_requests > 0?}
    C -->|Yes| D[prefill_bearing]
    C -->|No| E{generation_requests > 0?}
    E -->|Yes| F[decode]
    E -->|No| G[skip / idle]

Loading

Tier 2 — Shape-based phase splitting (--llm-inference):

flowchart LR
    A[All iteration roots] --> B[infer_batch_sizes_from_shapes]
    B --> C[batch_sizes per iteration]
    C --> D[Compute median batch size]
    D --> E{batch_size > 2x median?}
    E -->|Yes| F[prefill_bearing]
    E -->|No| G[decode]

Loading

Steady state identification (3 tiers):

Tier	Condition	Method
Concurrency-based	Trace has serving annotations	Find regions where num_requests is near peak (unchanged). Select windows by prefill/decode mode within the largest region.
Decode-baseline	--llm-inference flag, no serving annotations	Use decode-iteration batch sizes as a concurrency proxy (in decode, batch size = number of sequences ≈ concurrency). Filter out prefill spikes, find the region where decode batch sizes are near peak, map back to full iteration indices (re-including prefill iterations). Within that region, use shape-derived phase labels for mode-based window selection (mixed / decode-only / max-prefilldecode).
Generic duration CV	No annotations, no --llm-inference flag	Find the most duration-consistent region via sliding-window CV. Extract a single window. No phase awareness.

Decode-baseline steady state (Tier 2):

batch_sizes:    [8, 8, 32, 32, 32, 946, 32, 32, 854, 32, 32, 32, 8, 8]
phase_labels:   [D, D,  D,  D,  D,  P,   D,  D,  P,   D,  D,  D, D, D]

Step 1: Filter to decode-only batch sizes
decode_bs:      [8, 8, 32, 32, 32, __, 32, 32, __, 32, 32, 32, 8, 8]

Step 2: Peak-proximity scan on decode batch sizes
global_max=32, threshold=3.2
→ decode steady state: indices where decode_bs ≈ 32

Step 3: Map back to full iteration indices (re-include prefill)
→ region [2, 12) — includes both decode (bs=32) and prefill (bs=946, 854)

Step 4: Mode-based window selection within region
  mixed:            pick num_steps iters with representative prefill/decode ratio
  decode_only:      longest contiguous decode run
  max_prefilldecode: longest contiguous prefill-bearing run

Duration CV sliding window (Tier 3):

Iterations:    [  0  ][  1  ][  2  ][  3  ][  4  ][  5  ][  6  ][  7  ][  8  ][  9  ]
Durations:      3100   3050   5200   3080   3070   3090   3060   3075   3085   4900

Window 0-3:    [3100, 3050, 5200, 3080]  → CV = 0.25  ✗ (> 0.15)
Window 1-4:    [3050, 5200, 3080, 3070]  → CV = 0.24  ✗
Window 2-5:    [5200, 3080, 3070, 3090]  → CV = 0.24  ✗
Window 3-6:    [3080, 3070, 3090, 3060]  → CV = 0.004 ✓ mean=3075
Window 4-7:    [3070, 3090, 3060, 3075]  → CV = 0.004 ✓ mean=3074  ← fastest passing
Window 5-8:    [3090, 3060, 3075, 3085]  → CV = 0.004 ✓ mean=3078
Window 6-9:    [3060, 3075, 3085, 4900]  → CV = 0.22  ✗

Result: region [4, 8) — iterations with consistent durations, skipping
        warmup (iter 2 spike) and cooldown (iter 9 spike).

Other changes

Trace splitter improvements:

• Combined mode execution: --store-single-iteration, --find-steady-state, and --divide-phases can now all run in a single invocation.
• Window extraction with ancestors: When -i is specified without --store-single-iteration, the selected iterations are extracted as a single trace with ancestor context — parent frames from the split root up to the process entry are included, but not their other children.
• Bookend enhancement: When the generic detector (branch_descent or sibling_roots) finds iterations covering ≥50% of GPU time but below the splittable threshold, it checks for warmup (before first iteration) and wrapup (after last iteration) blocks among the parent's GPU-bearing children. If adding them improves the coverage grade, they're included as warmup / wrapup roots. This enables splitting of diffusion traces (e.g., HunyuanVideo: 4 denoising steps = 58% GPU → with bookends = 100%).
• GPU-bearing child filter: The branch_descent pattern finder now only considers GPU-bearing children when searching for repeating patterns, filtering out non-GPU setup/scheduling calls that obscure the iteration structure.

Tests

Verification of LLM inference splitting

For each of 26 LLM inference traces, we stripped the serving annotations (execute_*, step[*]) and ran the full pipeline on both versions: the annotated trace uses the annotation-based path, the stripped trace uses the generic detection cascade + shape-based phase classification.

Root detection (26 traces):

Result	Count	Notes
Root count matches	23/26	Stripped trace finds the same number of iterations as annotated.
Root count mismatch	3/26	One trace contains an annotation CompiledFxGraph annotation below every model layer in addition to the TraceLens recognized annotations. The annotation approach splits by iteration whereas stripped approach splits by the model layer. Another trace follows issue 3 described here: #1018.

Phase classification (23 traces where roots matched):

Result	Count	Notes
100% match	21/23	Shape-derived decode/prefill_bearing labels match annotation ground truth for every iteration.
97% (31/32)	1/23	Once trace contains an iteration which is pure decode, but the window extension (extending the end of one iteration to the start of the next) extends its window into the adjacent prefill iteration's territory, picking up cpu_ops with batch size 1024. The shape classifier sees the large batch size and labels it prefill_bearing.
0% (0/5)	1/23	One pure prefill trace. The median batch size is large, so no iteration exceeds the 2x median threshold required for iterations to be labeled prefill_bearing. All iterations are labeled "decode". This is rare in practice and can be ignored: an E2E trace won't contain pure prefill.

Steady state window overlap (23 root-matched traces):

Restricted to the 23 traces where both paths found the same number of iteration roots, so the comparison is apples-to-apples. Overlap = percentage of the stripped path's selected window that falls within the annotation path's selected window.

Metric	≥90%	<90%
Mixed window	23/23	0/23
Decode-only window	22/23	1/23

All 23 traces achieve 100% mixed window overlap — the shape-based path selects iterations that are entirely within the annotation path's steady state window.

Decode-only window overlap is 100% on 22 traces. The 1 exception is gemma_prefill: a pure prefill trace where the annotation path correctly finds no decode-only run (returns empty), but the shape-based path misclassifies all iterations as decode (the pure-prefill limitation) and returns a window.

The shape-based path uses decode-baseline region finding: it filters out prefill-bearing iterations, runs peak-proximity scan on the decode-only batch sizes (which approximate concurrency — 1 token per sequence in decode), then maps the region back to full iteration indices. This produces a region similar to the annotation path's concurrency-based region, since decode batch size ≈ num_requests. Within the region, mixed mode prefers windows containing at least one prefill-bearing iteration (matching the annotation path's behavior).

Verification of generic splitting

Ran evaluation script on 8 multi-iteration traces spanning training (OWLv2, DLRM) and LLM inference (DeepSeek, GPT-OSS) workloads. The evaluation checks root detection, GPU coverage, kernel count consistency (CV across iterations), duration consistency, and kernel shape similarity.

Trace	Method	Roots	Coverage	Kernel CV	Duration CV	Shape Sim	Verdict
sglang_deepseek_mi300	annotation:tier	6	97.9%	1.867	0.687	0.392	PASS
vllm_gptoss_mi300	annotation:tier	7	100%	0.000	0.026	1.000	PASS
dlrm-rank5	generic:branch_descent	6	100%	0.000	0.011	1.000	PASS
google_owlv2_h100	generic:branch_descent	5	100%	0.000	0.009	1.000	PASS
owlv2-h100	generic:branch_descent	5	100%	0.000	0.009	1.000	PASS
google_owlv2_mi300	generic:branch_descent	5	100%	0.000	0.008	1.000	PASS
owlv2-mi300	generic:branch_descent	5	100%	0.000	0.008	1.000	PASS
owlv2-vision-mi300	generic:branch_descent	5	100%	0.000	0.008	1.000	PASS

All pass.

Regression tests (tests/test_trace_split.py)

New test file that runs the trace splitting tool on real traces and compares output against checked-in reference splits. Covers inference fixtures (tests/traces/inference/) and multi-iteration fixtures (tests/traces/trace_splitter_traces/) spanning training, LLM inference, and diffusion workloads.

5 test functions:

Test	What it does
test_trace_split	Runs --store-single-iteration for first, middle, and last iterations. Compares each output .json.gz against split_traces/ references.
test_trace_split_steady_state	Runs --find-steady-state. Compares output against steady_state_traces/ references.
test_trace_split_divide_phases	Runs --divide-phases. Compares output against phase_split_traces/ references (LLM traces only).
test_trace_split_no_annotations	Strips user_annotation events from multi-iteration traces, re-runs splitting, and asserts the output matches the annotated references (minus annotations). Passes --llm-inference for sglang/vllm traces. Validates that generic detection produces equivalent splits.
test_trace_not_splittable	Verifies that small/simple traces in small_traces/ return NOT_SPLITTABLE.

Trace fixtures:

Location	Traces	Purpose
tests/traces/trace_splitter_traces/	8 multi-iteration traces	Training (OWLv2, DLRM), LLM inference (DeepSeek, GPT-OSS)
tests/traces/trace_splitter_traces/small_traces/	3 traces (BERT, NSFW, ResNet)	Expected NOT_SPLITTABLE

Multi-iteration trace breakdown (see tests/traces/trace_splitter_traces/README.md):

Directory	Workload	Iterations
dlrm-rank5	Training — DLRM (MI300)	6
google_owlv2_h100	Training — OWLv2 (H100)	5
google_owlv2_mi300	Training — OWLv2 (MI300)	5
owlv2-h100	Training — OWLv2 (H100)	5
owlv2-mi300	Training — OWLv2 (MI300)	5
owlv2-vision-mi300	Inference — OWLv2 (MI300)	5
sglang_deepseek_mi300	LLM inference — DeepSeek (MI300)	6
vllm_gptoss_mi300	LLM inference — GPT-OSS (MI300)	7

All reference directories support --update-references to regenerate golden outputs.

Notes

Known limitations of the shape-based path:

• Pure prefill traces (0% phase match): When every iteration is prefill, there's no decode baseline for the batch-size spike heuristic. All iterations get labeled "decode". This is rare in practice — pure prefill only occurs during the initial warmup of a serving trace, which would be excluded from steady state anyway.
• Fully-graphed traces with zero cpu_op shapes: Some traces (e.g., amd_vllm_gptoss_20b) have no Input Dims on any cpu_op — the entire model runs inside CUDA graphs. The --llm-inference flag detects this (all batch sizes are None) and falls through to the generic duration-CV path gracefully.