[Feature Request] Support xDiT / Diffusion workloads in trace analysis & roofline

[haishuok0525]

[Feature Request] Support xDiT / Diffusion workloads in trace analysis & roofline

1. Background

We profile a class of diffusion (text-to-image / DiT models, e.g. FLUX, running on xDiT / xfuser) workloads and feed the profiler's torch traces (*.pt.trace.json.gz) into TraceLens profile + trace_analyze to produce roofline / per-operator bottleneck analysis.

For LLM-style workloads (vllm / sglang) this pipeline works fine. For diffusion workloads, however, trace_analyze always fails, so these sessions get no TraceLens analysis at all.

2. Symptoms

When running trace_analyze on a diffusion trace, the steady-state splitting stage errors out and aborts:

• steady_state_chunk_empty / steady_state_chunk_missing
• Final failure: trace_split_no_steady_state (no steady-state chunk can be extracted)

Auto-retry (re-issuing once with an alternate mode from non_empty_modes) does not recover, because all existing modes are designed around autoregressive decode steady state.

3. Root Cause

The current steady-state splitter assumes the steady-state structure of LLM autoregressive inference (prefill / decode / mixed): it expects to find a region of stable, repeated decode steps.

Diffusion workloads have a fundamentally different execution structure:

• There is no prefill / decode distinction and no growing KV cache.
• One inference is a loop of N nearly-identical "denoising steps" (N = num_inference_steps, e.g. 20 / 28 / 50).
• Each step runs a full forward pass over the entire DiT; the structure is highly repetitive and regular.

As a result, the existing splitter cannot find the "decode steady state" it is built for and fails as empty. In reality, a diffusion trace has an even more regular periodic structure (the denoising-step period) that can be used for splitting — the current logic just does not recognize it.

4. Requirements

Requirement 1 (primary): steady-state splitting for a diffusion mode

Add a diffusion-specific steady-state detection / splitting mode that chunks by the denoising-step period rather than by decode steady state.

Possible implementations (either one, or both — TraceLens team's choice):

• Option A (recommended, better for cross-workload standardization): Provide a generic "split by externally-provided annotation markers" mode. The caller emits standard paired step-boundary markers in the trace (e.g. record_function ranges named denoise_step), and TraceLens splits on those markers without needing to understand the specific model.
• Option B (TraceLens-side only, zero instrumentation): Detect the repeated denoising-step period via periodicity detection over kernel signatures and split accordingly.

Option A is model/framework-agnostic and provides a stable integration contract; Option B requires zero changes from the caller. Either is acceptable, but at least one is needed.

Requirement 2: PerfModel coverage for DiT operators

The current PerfModel per-operator modeling is oriented toward transformer-decoder + MoE (GEMM, SDPA-with-KV, FusedMoE, etc.). Diffusion / DiT has a different operator structure. Please support arithmetic-intensity / efficiency / bound attribution for:

• Full-sequence self-attention (latent token count N = H×W / patch², non-causal, no KV growth);
• Cross-attention to text embeddings;
• adaLN / modulation operators;
• VAE conv2d-class operators.

Requirement 3: emit workload type and throughput unit in the output

The analysis.md Executive Summary is currently token-oriented (tok/s, etc.). Please explicitly emit:

• workload_type = diffusion (or an equivalent field);
• throughput_unit = images/sec (rather than tokens/sec).

so that downstream parsers do not have to guess the workload type and unit.

5. Acceptance Criteria

1. For the provided diffusion trace samples, trace_analyze completes successfully (no more trace_split_no_steady_state / steady_state_chunk_empty).
2. The extracted steady-state chunks correspond to the denoising-step period (step count / period count can be cross-checked).
3. The per-operator analysis includes the DiT operators above (attention / cross-attention / modulation / conv), each with a bound type and efficiency.
4. The output carries workload_type and throughput_unit=images/sec.

6. Reproduction Materials (available on request)

• One or more real diffusion profiler traces (*.pt.trace.json.gz);
• The corresponding trace_analyze failure logs (with steady_state_chunk_empty / trace_split_no_steady_state errors and non_empty_modes contents);
• Model info (DiT structure, num_inference_steps, resolution, images-per-batch, etc.);
• If Option A is chosen: we can provide trace samples with standard denoise_step markers for integration.

[kyle-hoffmeyer]

PR addressing error: #786. Messaged @haishuok0525 on Teams.

@haishuok0525, please follow up here whether the fix works. Thanks!