diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 833fa65d1..e6132f149 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -14,7 +14,7 @@ # TraceDiff module and docs require Spandan review /TraceLens/TraceDiff/ @spandoesai -/docs_original/TraceDiff.md @spandoesai +/docs/how-to/compare-traces.md @spandoesai # Extensions and experimental trace merge require Deval review /TraceLens/PerfModel/extensions/ @devalshahamd @@ -25,7 +25,8 @@ /TraceLens/Reporting/generate_perf_report_pytorch_inference.py @devalshahamd /TraceLens/TraceUtils/split_inference_trace_annotation.py @devalshahamd /examples/custom_workflows/inference_analysis/ @devalshahamd -/docs_original/Inference_analysis.md @devalshahamd +/docs/how-to/generate-perf-report-pytorch-inference.md @devalshahamd +/docs/conceptual/inference-analysis.md @devalshahamd /tests/test_inference_perf_report.py @devalshahamd /tests/traces/inference/ @devalshahamd diff --git a/README.md b/README.md index 3356c0e6c..6abde1e88 100755 --- a/README.md +++ b/README.md @@ -44,8 +44,8 @@ pip install git+https://github.com/AMD-AGI/TraceLens.git ### 2. Collect Traces TraceLens analyses profiler traces from PyTorch, JAX, and AMD rocprofv3; see [Supported Profile Formats](#supported-profile-formats) for the full list. The instructions below cover collecting a PyTorch trace: -- **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. The [PyTorch profiling guide](docs_original/conceptual/torch_profiling_guide.ipynb) walks through this end to end. -- **Inference Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Inference Analysis](docs_original/Inference_analysis.md). The [Profiling skill](TraceLens/Agent/Profiling/README.md) automates vLLM/SGLang benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. +- **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. The [PyTorch profiling walkthrough](docs/tutorials/torch-profiling.ipynb) walks through this end to end. +- **Inference Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Generate a PyTorch inference report](docs/how-to/generate-perf-report-pytorch-inference.md). The [Profiling skill](TraceLens/Agent/Profiling/README.md) automates vLLM/SGLang benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. To try out TraceLens without collecting your own trace, use the [demo traces](tests/traces) bundled in the repository. @@ -59,9 +59,9 @@ Generate a performance analysis report from an eager execution PyTorch trace wit TraceLens_generate_perf_report_pytorch --profile_json_path path/to/your/trace.json ``` -This produces an Excel workbook with GPU timeline breakdown, ops summary, roofline metrics and more. For additional details, see [Performance Report Generation](docs_original/generate_perf_report.md) and [Performance Report Column Definitions](docs_original/perf_report_columns.md). For other input formats, see [Supported Profile Formats](#supported-profile-formats). +This produces an Excel workbook with GPU timeline breakdown, ops summary, roofline metrics and more. For additional details, see [Generate a PyTorch performance report](docs/how-to/generate-perf-report-pytorch.md) and [Performance report column reference](docs/reference/perf-report-columns.md). For other input formats, see [Supported Profile Formats](#supported-profile-formats). -Compare two reports to quantify the impact of a change (see [compare_perf_reports_pytorch.md](docs_original/compare_perf_reports_pytorch.md)): +Compare two reports to quantify the impact of a change (see [Compare performance reports](docs/how-to/compare-perf-reports.md)): ```bash TraceLens_compare_perf_reports_pytorch \ @@ -71,7 +71,7 @@ TraceLens_compare_perf_reports_pytorch \ -o comparison.xlsx ``` -For multi-rank runs, generate a collective-communication report across ranks (see [generate_multi_rank_collective_report_pytorch.md](docs_original/generate_multi_rank_collective_report_pytorch.md)): +For multi-rank runs, generate a collective-communication report across ranks (see [Generate a collective-communication report](docs/how-to/collective-report.md)): ```bash TraceLens_generate_multi_rank_collective_report_pytorch \ @@ -105,10 +105,10 @@ Analyze a workload autonomously using an agentic system that automates performan | Format | Tool | Documentation | | --------------------- | ------------------------ | -------------------------------------------------------------------------------------------------------------- | -| **PyTorch** | `torch.profiler` | [docs_original/generate_perf_report.md](docs_original/generate_perf_report.md) | -| **JAX** | XPlane protobuf | [docs_original/jax_analyses.md](docs_original/jax_analyses.md) | -| **rocprofv3 JSON** | AMD ROCm rocprofiler-sdk | [docs_original/generate_perf_report_rocprof.md](docs_original/generate_perf_report_rocprof.md) | -| **rocprofv3 pftrace** | Perfetto-style | [docs_original/generate_perf_report_rocprof_pftrace.md](docs_original/generate_perf_report_rocprof_pftrace.md) | +| **PyTorch** | `torch.profiler` | [docs/how-to/generate-perf-report-pytorch.md](docs/how-to/generate-perf-report-pytorch.md) | +| **JAX** | XPlane protobuf | [docs/how-to/generate-perf-report-jax.md](docs/how-to/generate-perf-report-jax.md) | +| **rocprofv3 JSON** | AMD ROCm rocprofiler-sdk | [docs/how-to/generate-perf-report-rocprof.md](docs/how-to/generate-perf-report-rocprof.md) | +| **rocprofv3 pftrace** | Perfetto-style | [docs/how-to/generate-perf-report-rocprof.md](docs/how-to/generate-perf-report-rocprof.md) | Each format's linked doc covers its full CLI reference. For PyTorch report comparison and multi-rank collective analysis, see the corresponding docs in the [Documentation](#documentation) table. @@ -118,18 +118,19 @@ Each format's linked doc covers its full CLI reference. For PyTorch report compa | Module | Doc | | ---------------------------- | -------------------------------------------------------------------------------------------------------------------------------- | -| Trace2Tree | [docs_original/Trace2Tree.md](docs_original/Trace2Tree.md) | -| TreePerf | [docs_original/TreePerf.md](docs_original/TreePerf.md) | -| NCCL Analyser | [docs_original/NcclAnalyser.md](docs_original/NcclAnalyser.md) | -| TraceDiff | [docs_original/TraceDiff.md](docs_original/TraceDiff.md) | -| Event Replay | [docs_original/EventReplay.md](docs_original/EventReplay.md) | -| TraceFusion | [docs_original/TraceFusion.md](docs_original/TraceFusion.md) | -| GPU Event Analyser | [docs_original/gpu_event_analyser.md](docs_original/gpu_event_analyser.md) | -| JAX Analyses | [docs_original/jax_analyses.md](docs_original/jax_analyses.md) | -| pftrace Reports | [docs_original/generate_perf_report_rocprof_pftrace.md](docs_original/generate_perf_report_rocprof_pftrace.md) | -| Compare PyTorch Reports | [docs_original/compare_perf_reports_pytorch.md](docs_original/compare_perf_reports_pytorch.md) | -| Multi-Rank Collective Report | [docs_original/generate_multi_rank_collective_report_pytorch.md](docs_original/generate_multi_rank_collective_report_pytorch.md) | -| Performance Report Columns | [docs_original/perf_report_columns.md](docs_original/perf_report_columns.md) | +| Trace2Tree | [docs/conceptual/trace2tree.md](docs/conceptual/trace2tree.md) | +| TreePerf | [docs/how-to/tree-perf-analysis.md](docs/how-to/tree-perf-analysis.md) | +| NCCL Analyser | [docs/how-to/nccl-analysis.md](docs/how-to/nccl-analysis.md) | +| TraceDiff | [docs/how-to/compare-traces.md](docs/how-to/compare-traces.md) | +| Event Replay | [docs/how-to/event-replay.md](docs/how-to/event-replay.md) | +| TraceFusion | [docs/how-to/trace-fusion.md](docs/how-to/trace-fusion.md) | +| GPU Event Analyser | [docs/how-to/gpu-event-analysis.md](docs/how-to/gpu-event-analysis.md) | +| JAX Analyses | [docs/how-to/generate-perf-report-jax.md](docs/how-to/generate-perf-report-jax.md) | +| pftrace Reports | [docs/how-to/generate-perf-report-rocprof.md](docs/how-to/generate-perf-report-rocprof.md) | +| Compare PyTorch Reports | [docs/how-to/compare-perf-reports.md](docs/how-to/compare-perf-reports.md) | +| Multi-Rank Collective Report | [docs/how-to/collective-report.md](docs/how-to/collective-report.md) | +| Performance Report Columns | [docs/reference/perf-report-columns.md](docs/reference/perf-report-columns.md) | +| TraceLens Agent | [docs/how-to/agent.md](docs/how-to/agent.md) | --- @@ -153,9 +154,6 @@ Please see [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on branching, commi ## Additional Resources -- [PyTorch Conference 2025 Poster](docs_original/TraceLens%20-%20Democratizing%20AI%20Performance%20Analysis%20-%20Adeem%20Jassani%2C%20AMD.pdf) -- [GEMMs in AI Models: Conceptual Tutorial](docs_original/conceptual/aimodels_gemms.md) -- [Trace2Tree Motivation](docs_original/conceptual/trace2tree_motivation.md) -- [PyTorch Profiling Guide](docs_original/conceptual/torch_profiling_guide.ipynb) - -For more background and conceptual tutorials, browse `[docs_original/conceptual/](docs_original/conceptual/)`. \ No newline at end of file +- [GEMM analysis in TraceLens](docs/conceptual/gemm-analysis.md) +- [The Trace2Tree data model](docs/conceptual/trace2tree.md) +- [PyTorch profiling walkthrough](docs/tutorials/torch-profiling.ipynb) \ No newline at end of file diff --git a/TraceLens/Agent/Analysis/README.md b/TraceLens/Agent/Analysis/README.md index cdc096a01..6f8180fb3 100644 --- a/TraceLens/Agent/Analysis/README.md +++ b/TraceLens/Agent/Analysis/README.md @@ -60,8 +60,8 @@ pip install -e . The orchestrator runs against a single PyTorch profiler trace (`.json` or `.json.gz`). Collection is workload-specific: -- **Training and eager inference traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of training/inference steps, post-warmup) and dump the trace via `prof.export_chrome_trace(...)`. A single rank's trace is sufficient for per-rank analysis. -- **vLLM / SGLang inference traces**: Trace collection has framework-, version-, and execution-mode-specific requirements (custom Docker images or framework patches to add roofline annotations, profiler-config flags for graph-capture profiling, steady-state window selection, optional trace splitting). Follow the canonical guide in [Inference Analysis](../../../docs_original/Inference_analysis.md). For graph-mode workloads you will produce **two artifacts**: a graph-replay trace and a graph-capture folder; TraceLens (in inference mode with execution mode `graph replay + capture`) merges call-stack and shape information from the capture folder into the replay tree before analysis. +- **Generic Eager Traces**: Instrument your loop with `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state window (a handful of steps, post-warmup) and log the trace with `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank analysis. The [PyTorch profiling walkthrough](../../../docs/tutorials/torch-profiling.ipynb) walks through this end to end. +- **Inference Traces with Graph Capture**: Collection has framework-specific requirements. Follow guidelines in [Generate a PyTorch inference report](../../../docs/how-to/generate-perf-report-pytorch-inference.md). The [Profiling skill](../Profiling/README.md) automates vLLM/SGLang benchmarking and PyTorch profiler trace collection via [Magpie](https://github.com/AMD-AGI/Magpie), producing analysis-ready traces. For graph-mode workloads you produce two artifacts: a graph-replay trace and a graph-capture folder. In inference mode with execution mode `graph replay + capture`, TraceLens merges call-stack and shape information from the capture folder into the replay tree before analysis. ### 3. Establish a hardware performance baseline @@ -80,7 +80,11 @@ Roofline analysis compares each measured kernel against your GPU's max-achievabl ``` "Follow the analysis orchestrator installed with TraceLens and run the full agentic analysis workflow on " ``` - - Comparative (two traces): + - Standalone (single graph replay trace with capture trace directory): + ``` + "Follow the analysis orchestrator installed with TraceLens and run the full agentic analysis workflow on with capture folder " + ``` + - Comparative (two eager traces): ``` "Follow the analysis orchestrator installed with TraceLens and run the full agentic analysis workflow on and " ``` @@ -267,16 +271,17 @@ flowchart TD Sdpa --> CatFindings Others --> CatFindings - SysFindings --> Step8["Step 8: Validate"] - CatFindings --> Step8 - Step8 --> Step9["Step 9: Aggregate"] - Step9 --> Step11["Step 11: Final Report"] + SysFindings --> Agg["Step 7.5: Aggregate"] + CatFindings --> Agg + Agg --> Step8["Step 8: Validate"] + Step8 --> Step9["Step 9: Report Prep + Model ID"] + Step9 --> Step10["Step 10: Final Report"] ``` ### Orchestrator The **Analysis Orchestrator** skill coordinates the entire analysis workflow. -It queries user inputs, runs TraceLens to pre-compute trace data, and invokes system-level and compute kernel sub-agents in parallel. Finally, it validates outputs, aggregates findings, and generates a prioritized stakeholder report. +It queries user inputs, runs TraceLens to pre-compute trace data, and invokes system-level and compute kernel sub-agents in parallel. Finally, it aggregates and validates findings, identifies the model, and generates a prioritized stakeholder report. ### Workflow Steps @@ -284,12 +289,12 @@ It queries user inputs, runs TraceLens to pre-compute trace data, and invokes sy 0. Query User Inputs (Comparison scope, Trace path(s), Platform(s), Analysis Mode, Environment Setup) 1. Generate Performance Report (branches on analysis mode and comparison scope) 2-5. Prepare Category Data (GPU Util, Top Ops, Tree Data, Multi-Kernel Data, Category Filtering) + Fusion Candidate Extraction → category_data/fusion_candidates.json + kernel_fusion_metrics.json -5.5. Model Identification (subagent) → metadata/model_info.json 6. System-Level Analysis (CPU/Idle + Multi-Kernel + Kernel Fusion, PARALLEL) → system_findings/ 7. Compute Kernel Subagents (PARALLEL) → category_findings/ +7.5. Aggregate per-category Compute Kernel findings → priority_data.json (globally ranked) 8. Validate Subagent Outputs (time sanity, efficiency anomalies, coverage) -9. Aggregate Results: System-Level + Kernel Fusion + Compute Kernel Recommendations -10. Generate Final Report (analysis.md) +9. Prepare Report Data + Model Identification (subagent) → metadata/model_info.json +10. Generate Final Report (analysis.md): compose System-Level + Kernel Fusion + Compute Kernel sections ``` ### Sub-Agents @@ -320,7 +325,7 @@ It queries user inputs, runs TraceLens to pre-compute trace data, and invokes sy The orchestrator and all 13 sub-agents currently run on **`claude-opus-4-7-high`**, declared in each agent file's front matter under `.cursor/agents/`. The full set: `cpu-idle-analyzer`, `multi-kernel-analyzer`, `kernel-fusion-analyzer`, `model-identification-agent`, `gemm-analyzer`, `sdpa-analyzer`, `elementwise-analyzer`, `reduce-analyzer`, `triton-analyzer`, `moe-analyzer`, `norm-analyzer`, `convolution-analyzer`, `generic-op-analyzer`. -### Supported Standalone Analysis Modes +### Supported Analysis Modes The orchestrator supports two analysis modes, selected during Step 0: diff --git a/TraceLens/Agent/Profiling/README.md b/TraceLens/Agent/Profiling/README.md index 09d2c575c..8ace11fa4 100644 --- a/TraceLens/Agent/Profiling/README.md +++ b/TraceLens/Agent/Profiling/README.md @@ -151,7 +151,7 @@ The skill parses the `case` blocks of these scripts at runtime to discover curre ## Trace Splitting and Handoff to Analysis -Step 6 produces split traces in `torch_trace/trace_split/` via `TraceLens.TraceUtils.split_inference_trace_annotation`. The skill then prints (but does **not** run) a `generate_perf_report_pytorch_inference.py` command that the user can launch to feed the split traces into the [Analysis Orchestrator](../Analysis/README.md). See [`docs_original/Inference_analysis.md`](../../../docs_original/Inference_analysis.md) for splitting heuristics and prefill/decode mix selection. +Step 6 produces split traces in `torch_trace/trace_split/` via `TraceLens.TraceUtils.split_inference_trace_annotation`. The skill then prints (but does **not** run) a `generate_perf_report_pytorch_inference.py` command that the user can launch to feed the split traces into the [Analysis Orchestrator](../Analysis/README.md). See [Generate a PyTorch inference performance report](../../../docs/how-to/generate-perf-report-pytorch-inference.md) for splitting heuristics and prefill/decode mix selection. ## Bug Reporting diff --git a/docs/about/release-notes.md b/docs/about/release-notes.md index 6cce191ff..5c0c9e1dc 100644 --- a/docs/about/release-notes.md +++ b/docs/about/release-notes.md @@ -34,7 +34,7 @@ This release includes the following report generators. a unique-argument table, roofline metrics, and an optional short-kernel study. Compressed traces (`.zip`, `.gz`) are supported. - **PyTorch inference reports:** `TraceLens_generate_perf_report_pytorch_inference` - adds analysis for LLM-serving traces (vLLM/SGLang). It merges CUDA-graph + adds analysis for LLM-serving traces (vLLM/SGLang). It merges HIP/CUDA-graph capture traces into graph-replay traces (`--capture_folder`) to recover call-stack and shape information lost in graph mode, and uses per-step request annotations (prefill/decode counts and token statistics) to drive @@ -67,7 +67,7 @@ This release includes the following analysis features. dtypes, strides, and concrete inputs to isolate problematic input patterns. - **Roofline modeling:** Computes arithmetic intensity (FLOPs/byte) and classifies operations as compute- or memory-bound relative to a target - accelerator's roofline knee point. Optional Origami-based simulated GEMM/SDPA + GPU's roofline knee point. Optional Origami-based simulated GEMM/SDPA timings are available when a GPU architecture specification is provided. - **Activation-recompute detection** and **kernel-overlap** sheets for deeper PyTorch analysis. diff --git a/docs/conceptual/gemm-analysis.md b/docs/conceptual/gemm-analysis.md new file mode 100644 index 000000000..c340170dc --- /dev/null +++ b/docs/conceptual/gemm-analysis.md @@ -0,0 +1,470 @@ + + +# GEMM analysis in TraceLens +```{meta} +:description: Understand how model dimensions map to GEMM shapes and BLAS calls, and how TraceLens reports GEMM dimension-efficiency metrics such as tile and wave quantization efficiency. +:keywords: TraceLens, GEMM, general matrix multiply, BLAS, hipBLAS, PyTorch, roofline, tile efficiency, wave quantization, dimension efficiency, Tensile, ROCm, MI300X, LLM, prefill, decode +``` + +General matrix multiply (GEMM) operations are the *primary compute primitive* used in AI models. Efficient implementations of GEMMs are readily available through vendor-tuned libraries such as cuBLAS and hipBLAS, so whenever possible the goal is to *reduce computation to a matrix multiply*. + +This topic explains how *model-level parameters* like batch size, sequence length, and hidden dimension translate into GEMM shapes, how those shapes map to specific basic linear algebra subprograms (BLAS) kernel calls, and how TraceLens reports deeper GEMM dimension-efficiency metrics. By the end you should be equipped to understand the GEMM shapes, counts, and BLAS calls involved in *any AI model you encounter*. + +## From model dimensions to GEMM shapes + +### Linear layers in LLMs + +Consider how linear layers in a large language model (LLM), such as the multi-layer perceptron (MLP) *up projection*, correspond to GEMM calls. + +The input tensor shape for this operation is: + +``` +X: [B, L, d_model] +``` + +Here: + +- `B` represents the batch size. +- `L` denotes the sequence length. +- `d_model` is the input (or hidden) dimension. + +This operation outputs a tensor with the shape: + +``` +Y: [B, L, d_ff] +``` + +The projection for each token can be expressed individually as: + +``` +Y[b, l, :] = X[b, l, :] @ Wᵀ +``` + +Where `W` is the weight matrix with a shape of `[d_ff, d_model]`. + +### Flattening for GEMM + +To express this entire operation as a single GEMM, flatten the batch and sequence dimensions of the input tensor: + +``` +X_flat: [B·L, d_model] +Wᵀ: [d_model, d_ff] +Y = X_flat @ Wᵀ +``` + +This flattening yields the following GEMM shape parameters: + +- `param: M = B·L` +- `param: N = d_ff` +- `param: K = d_model` + +Here, `K` represents the *inner or shared dimension* between the input tensors involved in the multiplication. + +## Prefill versus decode + +Contrast the GEMM behavior during the *prefill* and *decode* phases of inference, focusing on how the sequence length (`L`) changes and affects the GEMM shapes. + +In both phases, the input tensor for an MLP computation within an LLM initially has a shape like: + +``` +X: [B, L, d_model] +``` + +Where `B` is the batch size, `L` is the sequence length, and `d_model` is the hidden size. This projects to an output shape of `[B, L, d_ff]`, where `d_ff` is the MLP's expansion size. + +As established earlier, to process this with a single GEMM the batch and sequence dimensions of the input are flattened. The input effectively becomes `[B·L, d_model]` for the GEMM `X_flat @ Wᵀ`. + +The key difference between prefill and decode lies in the value of `L`: + +- *Prefill phase*: `L` is the actual input sequence length (which can be large). +- *Decode phase*: `L` is always `1`, as the model processes one token at a time to generate the next. + +This difference in `L` directly impacts the `M` parameter of the GEMM `(M, N, K)`. + +#### GEMM shape summary + +| Mode | Input shape (conceptual) | Flattened input shape | GEMM shape `(M, N, K)` | Notes | +|----------|--------------------------|-----------------------|------------------------|--------------------------------------| +| Prefill | `[B, L, d_model]` | `[B·L, d_model]` | `(B·L, d_ff, d_model)` | `L` is the prompt length | +| Decode | `[B, 1, d_model]` | `[B·1, d_model]` | `(B, d_ff, d_model)` | `L=1` for generating one token at a time | + +Notice that in the decode phase, because `L=1`, the `M` parameter of the MLP GEMM becomes simply `B`. This means the computational cost of the MLP layers in decode remains constant per token regardless of the total sequence length generated so far. The dominant `O(L)` scaling cost during decode comes from the attention mechanism, not the MLPs. + +### Real-world example: LLaMA-2 7B + +The table below shows actual data from TraceLens profiling, filtered specifically for MLP *up* and *gate* projection GEMMs in a LLaMA-2 7B model inference trace. + +For this trace: + +- `d_model = 4096` +- `d_ff = 11008` +- Batch size: `1` (`B=1`) +- Input length (for prefill): `597` (`L=597`) +- The trace included 36 decode steps. + +| name | param: M | param: N | param: K | param: bias | counts | +|----------|----------|----------|----------|-------------|--------| +| aten::mm | 1 | 11008 | 4096 | FALSE | 2304 | +| aten::mm | 597 | 11008 | 4096 | FALSE | 64 | + +Interpreting these entries based on the `M` parameter: + +- The entry with `param: M = 597` corresponds to the *prefill* phase GEMM `(B·L = 1·597)`, which happens once per layer at the beginning of inference. Since there are 32 layers, this GEMM is called `64` times (32 up + 32 gate). +- The entry with `param: M = 1` corresponds to the *decode* phase GEMM `(B = 1)`, where `L=1`. These occur at each decode step for every layer. With 36 decode steps and 64 GEMMs per step (32 layers * 2), this GEMM is called `36 × 64 = 2304` times. + +## Backward pass GEMMs + +Next, consider the backward pass during training. A forward pass GEMM operation like `Y = X @ Wᵀ + b` necessitates *two corresponding backward GEMMs* to compute gradients: + +```python +dX = dY @ W # Gradient with respect to the input → resulting shape: [B·L, d_model] +dW = dYᵀ @ X # Gradient with respect to the weight → resulting shape: [d_ff, d_model] +db = dY.sum(dim=0) # Gradient with respect to the bias → resulting shape: [d_ff] +``` + +### GEMM shapes + +| Operation | GEMM shape `(param: M, param: N, param: K)` | Description | +|-------------|---------------------------------------------|--------------------------------------| +| Forward | `(B·L, d_ff, d_model)` | `X @ Wᵀ` | +| Backward dX | `(B·L, d_model, d_ff)` | `dY @ W` (result of `[B·L, d_ff] @ [d_ff, d_model]`) | +| Backward dW | `(d_ff, B·L, d_model)` | `dYᵀ @ X` (result of `[d_ff, B·L] @ [B·L, d_model]`) | + +A closer look at the backward GEMM shapes: + +- For `dX = dY @ W`, the operation is `[B·L, d_ff] @ [d_ff, d_model]`, which results in a shape of `[B·L, d_model]`. +- For `dW = dYᵀ @ X`, the operation is `[d_ff, B·L] @ [B·L, d_model]`, yielding a shape of `[d_ff, d_model]`. + +### Real-world example: GPT-3-XL + +This table presents data from a TraceLens analysis of a single training step for GPT-3-XL. + +For this example: + +- `d_model = 2048` +- `d_ff = 8192` +- Batch size: `5` +- Sequence length: `2048` +- Thus, `param: M` for the flattened dimension is `5 × 2048 = 10240`. + +| name | param: M | param: N | param: K | count | +|-------------|----------|----------|----------|-------| +| aten::addmm | 10240 | 8192 | 2048 | 24 | +| aten::mm | 10240 | 2048 | 8192 | 24 | +| aten::mm | 8192 | 2048 | 10240 | 24 | + +Each entry can be interpreted based on the GEMM shapes: + +- The `aten::addmm` call represents the forward pass GEMM (`X @ Wᵀ`). +- The first `aten::mm` call corresponds to the backward pass for `dX` (`dY @ W`). +- The second `aten::mm` call represents the backward pass for `dW` (`dYᵀ @ X`). + +Each of these operations appears once per layer in the network. Given that GPT-3-XL has 24 layers, each of these GEMMs is called 24 times per training step, aligning with the `count` column in the table. + +## How PyTorch calls BLAS + +To fully grasp how PyTorch leverages BLAS for operations like GEMM, first understand the fundamental concept of *memory layout* for tensors and how BLAS libraries interpret the data buffers they receive. + +### Memory layout and stride + +Despite tensors often being represented as multi-dimensional arrays, their elements are stored in linear memory. For a 2D matrix, the two primary storage conventions are: + +- *Row-major*: Elements of the same row are stored consecutively in memory. PyTorch adopts this as its default layout. +- *Column-major*: Elements of the same column are stored consecutively in memory. Many traditional BLAS libraries primarily optimize for this layout. + +PyTorch's `.stride()` method provides insight into a tensor's memory arrangement. It returns a tuple where each value indicates the byte (or element, depending on datatype size) distance in linear memory to move to the next element along that dimension. + +- For a 2D tensor `T[i][j]` in *row-major* layout, `.stride()` is typically `(num_cols, 1)`. Moving to `T[i][j+1]` requires stepping 1 element, while moving to `T[i+1][j]` requires stepping `num_cols` elements. +- For a 2D tensor `T[i][j]` in *column-major* layout, `.stride()` is typically `(1, num_rows)`. Moving to `T[i+1][j]` requires stepping 1 element, while moving to `T[i][j+1]` requires stepping `num_rows` elements. + +### BLAS transpose and row-major output + +The core BLAS GEMM routine typically computes $C = \alpha \cdot op(A) \cdot op(B) + \beta \cdot C$, where $op(X)$ is either $X$ or $X^T$ depending on the `transA` and `transB` flags (`'N'` for no transpose, `'T'` for transpose) passed to the function. By default, BLAS expects input matrices corresponding to the `'N'` flag to be in column-major layout. Crucially, the resulting matrix $C$ is written into the output buffer in *column-major* format by default. + +PyTorch, however, uses row-major layout internally and desires the result of a GEMM operation to also be in row-major layout *without* an extra copy or transpose step outside of the BLAS call. PyTorch achieves this by cleverly leveraging the `trans` flags and the relationship between row-major and column-major layouts. + +A matrix $M$ stored in row-major memory has the exact same element ordering as the matrix $M^T$ stored in column-major memory. PyTorch uses this identity. To get a row-major result $C$ from a BLAS call that outputs in column-major, PyTorch requests BLAS to compute $C^T$ and write it in column-major. Since $C^T$ in column-major is $C$ in row-major, the output buffer will contain the desired row-major $C$. + +Mathematically, the operation $C = A @ B$ (where $A, B, C$ are desired in row-major) is equivalent to computing $C^T = (A @ B)^T = B^T @ A^T$. PyTorch therefore configures the BLAS call to compute $B^T @ A^T$ using the row-major data of $B$ and $A$. + +Here's how the `transA` and `transB` flags work in this context when passing *row-major data* to BLAS via a wrapper like PyTorch's: + +- Passing row-major data for matrix $M$ with `trans = 'T'` tells BLAS to mathematically treat this data as $M$. (BLAS expects row-major data for `'T'` if it wants to use the matrix directly.) +- Passing row-major data for matrix $M$ with `trans = 'N'` tells BLAS to mathematically treat this data as $M^T$. (BLAS expects column-major data for `'N'`; giving it row-major data makes it see the transpose.) + +So, to compute $C^T = B^T @ A^T$ using row-major data for $B$ and $A$ and get $C$ row-major in the output buffer: + +- Pass $B$'s row-major data as the first operand data (`A_data` in BLAS call). To make BLAS see $B^T$, use `transA = 'N'`. +- Pass $A$'s row-major data as the second operand data (`B_data` in BLAS call). To make BLAS see $A^T$, use `transB = 'N'`. +- The BLAS call becomes `gemm(transA='N', transB='N', ..., B_data, ..., A_data, ...)`. This computes $B^T @ A^T = C^T$. The result $C^T$ is written in column-major into the output buffer, which is precisely the desired $C$ in row-major. + +This standard trick using `transA='N'` and `transB='N'` with swapped, row-major inputs is a common way PyTorch achieves row-major output for a general matrix multiply `C = A @ B` where A, B are row-major. + +### Linear layer: `Y = X @ Wᵀ` + +For a linear layer computation `Y = X @ Wᵀ`, where `X` (`[M, K]`) and `W` (`[N, K]`) are in row-major layout, PyTorch desires `Y` (`[M, N]`) also in row-major. To achieve this with a BLAS routine outputting column-major, PyTorch configures BLAS to compute $Y^T = W @ X^T$. + +This involves a BLAS call computing $op(A) @ op(B)$ where $op(A)$ is $W$ and $op(B)$ is $X^T$. Using the rule that row-major data with `trans='T'` yields the matrix ($M$) and `trans='N'` yields the transpose ($M^T$): + +- BLAS operand A uses $W$'s row-major data. To see $W$, `transA = 'T'`. +- BLAS operand B uses $X$'s row-major data. To see $X^T$, `transB = 'N'`. + +The BLAS call uses `(transA='T', transB='N')` with $W$'s data as the first operand and $X$'s data as the second. It computes $W @ X^T = Y^T$, writing the result in column-major, which PyTorch interprets as the desired row-major $Y$. + +### Backward pass operations + +The backward pass similarly uses GEMMs configured to produce row-major gradients: + +- `dX = dY @ W`: With `dY` (`[M, K]`) and `W` (`[K, N]`) row-major, `dX` (`[M, N]`) is needed row-major. BLAS computes $dX^T = W^T @ dY^T$. + - BLAS operand A uses $W$'s row-major data. Needs $W^T \implies$ `transA = 'N'`. + - BLAS operand B uses $dY$'s row-major data. Needs $dY^T \implies$ `transB = 'N'`. + - BLAS call uses `(transA='N', transB='N')` on $W$'s and $dY$'s data, computing $W^T @ dY^T$. +- `dW = dYᵀ @ X`: With `dY` (`[K, N]`) and `X` (`[K, M]`) row-major, `dW` (`[N, M]`) is needed row-major. BLAS computes $dW^T = X^T @ dY$. + - BLAS operand A uses $X$'s row-major data. Needs $X^T \implies$ `transA = 'N'`. + - BLAS operand B uses $dY$'s row-major data. Needs $dY \implies$ `transB = 'T'`. + - BLAS call uses `(transA='N', transB='T')` on $X$'s and $dY$'s data, computing $X^T @ dY$. + +In summary, for PyTorch's row-major operations: + +- Forward pass `Y = X @ Wᵀ` maps to BLAS calculating $W @ X^T$ using `(T, N)` flags on the row-major data of $W$ and $X$. +- Backward pass `dX = dY @ W` maps to BLAS calculating $W^T @ dY^T$ using `(N, N)` flags on the row-major data of $W$ and $dY$. +- Backward pass `dW = dYᵀ @ X` maps to BLAS calculating $X^T @ dY$ using `(N, T)` flags on the row-major data of $X$ and $dY$. + +Revisiting the GPT-3-XL model GEMM table from TraceLens: + +| name | param: M | param: N | param: K | param: bias | param: stride_A | param: stride_B | param: transpose | +|-------------|----------|----------|----------|-------------|-----------------|-----------------|------------------| +| aten::addmm | 10240 | 8192 | 2048 | TRUE | (2048, 1) | (1, 2048) | (True, False) | +| aten::mm | 10240 | 2048 | 8192 | FALSE | (8192, 1) | (2048, 1) | (False, False) | +| aten::mm | 8192 | 2048 | 10240 | FALSE | (1, 8192) | (2048, 1) | (False, True) | + +This table shows how `aten::addmm` (forward) and `aten::mm` (backward) calls map to underlying GEMM operations. The `param: M, N, K` values are likely the dimensions of the *PyTorch operation result* (`M x N` with inner dim `K`). The `param: transpose | (transA, transB)` are the BLAS flags used by the wrapper for the operands passed to the BLAS call. + +Interpret the trace entries based on the understanding that PyTorch uses row-major data and BLAS receives flags to mathematically interpret this data for computing the transpose of the desired result: + +1. `aten::addmm` (forward): Corresponds to `Y = X @ Wᵀ`. Trace flags: `(True, False)`. This matches the `(T, N)` needed for BLAS to compute $W @ X^T$. +2. First `aten::mm` (backward dX): Corresponds to `dX = dY @ W`. Trace flags: `(False, False)`. This matches the `(N, N)` needed for BLAS to compute $W^T @ dY^T$. +3. Second `aten::mm` (backward dW): Corresponds to `dW = dYᵀ @ X`. Trace flags: `(False, True)`. This matches the `(N, T)` needed for BLAS to compute $X^T @ dY$. + +This confirms how the trace flags correspond to the BLAS transpose configurations used with row-major input data to achieve row-major output via the $C^T = B^T A^T$ trick. + +The following pseudo code summarizes the transpose flag logic. + +```{note} +This Python code snippet provides a simplified view; PyTorch's actual implementation is more intricate, accounting for the specific GEMM variant and output requirements. +``` + +```python +def is_col_major(T): + return T.stride(0) == 1 and T.stride(1) >= T.shape[0] + +def get_blas_transpose_flags(A, B): + transA = 'N' if is_col_major(A) else 'T' # If A is col-major, BLAS sees it as is ('N') + transB = 'N' if is_col_major(B) else 'T' # If B is col-major, BLAS sees it as is ('N') + return transA, transB +``` + +### Edge cases + +```{warning} +One common assumption is that flattening a tensor shape like `[B, L, d_model]` to `[B⋅L, d_model]` is a cost-free metadata operation. This is only true if the last dimension (`d_model`) is contiguous. +``` + +- If the last dimension is not contiguous, PyTorch may be forced to insert a *copy* or *transpose* operation to create a physically contiguous tensor that BLAS can work with efficiently. +- Furthermore, even if the tensor layout *could* theoretically be used by BLAS (for example, certain striding patterns), some highly tuned BLAS libraries might lack kernels optimized for those specific layouts. In such instances, a *copy* or *transpose buffer* is inserted behind the scenes by PyTorch or the BLAS wrapper. Consequently, what the BLAS routine actually operates on might not be the original tensor directly, but rather a *temporary buffer* created for compatibility or performance. + +## GEMM dimension efficiency + +Beyond mapping shapes to BLAS calls, TraceLens provides deeper insights into the efficiency of GEMM operations, specifically focusing on Tensile kernels used on ROCm GPUs. This analysis complements the existing roofline metrics by breaking down performance limitations related to how GEMM computations are tiled and scheduled onto the GPU. + +By default, the tool provides roofline metrics summarizing overall performance for operations like `aten::mm`: + +*Original output example:* + +| name | M | N | K | bias | FLOPS/Byte_first | TFLOPS/s_mean | +|----------|------|------|-------|-------|------------------|---------------| +| aten::mm | 2048 | 2048 | 10240 | False | 930.90 | 521.57 | + +The enhanced analysis adds specific efficiency metrics for Tensile GEMM kernels, derived from the kernel's structure and the problem dimensions. See the [`examples/gemm_dim_eff.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/gemm_dim_eff.ipynb) notebook for example usage. + +*New output example (Tensile kernels):* + +| name | M | N | K | ... | mt_m | mt_n | num_tiles | tile_eff | wq_eff | dim_eff | ... | TFLOPS/s_mean | +|----------|------|------|-------|-----|------|------|-----------|----------|--------|---------|-----|---------------| +| aten::mm | 2048 | 2048 | 10240 | ... | 256 | 64 | 256 | 1.00 | 0.84 | 0.84 | ... | 521.57 | + +### Key metrics + +- `mt_m`, `mt_n`: Macro-tile dimensions extracted from the kernel name. +- `num_tiles`: Total number of tiles after padding. +- `tile_eff`: Tile quantization efficiency. Measures efficiency loss due to input matrix dimensions not being perfectly divisible by tile dimensions. +- `wq_eff`: Wave quantization efficiency. Measures efficiency loss due to the total number of tiles not perfectly filling all compute units in the final processing wave. +- `dim_eff`: Net dimension efficiency. The product of `tile_eff` and `wq_eff`, representing the combined efficiency impact of tiling and scheduling. + +### Understanding the concepts + +#### Tile quantization efficiency (`tile_eff`) + +Tiled computations divide matrices into smaller sub-blocks (tiles: `mt_m x mt_n`) for processing. If the matrix dimensions (M, N) are not exact multiples of these tile sizes, the implementation effectively pads the matrices to the nearest multiple of the tile size. + +Padded dimensions: + +``` +M_pad = ceil(M / mt_m) × mt_m +N_pad = ceil(N / mt_n) × mt_n +``` + +This padding introduces extra computations on regions that don't contribute to the final result. + +Tile efficiency: + +``` +tile_eff = (M × N) / (M_pad × N_pad) +``` + +Values `< 1` indicate wasted computation due to padding. + +#### Wave quantization efficiency (`wq_eff`) + +GPUs execute tiles across many compute units (CUs), also known as streaming multiprocessors (SMs). If the total number of tiles isn't a multiple of the number of available CUs, the final wave will leave some CUs idle. + +Total tiles: + +``` +B = (M_pad × N_pad) / (mt_m × mt_n) +``` + +Number of waves: + +``` +num_waves = ceil(B / num_cus) +``` + +Wave efficiency: + +``` +wq_eff = B / (num_waves × num_cus) +``` + +#### Net dimension efficiency (`dim_eff`) + +This is the combined impact: + +``` +dim_eff = tile_eff × wq_eff +``` + +### Why these metrics matter: diagnosing bottlenecks + +For compute-bound GEMMs, actual performance is often less than peak theoretical. These metrics help explain why: + +- Low `tile_eff`: Indicates high padding overhead. +- Low `wq_eff`: Indicates poor utilization of compute units in the final wave. + +GEMM tuning can help improve these metrics to a certain extent. + +However, these are just part of the picture. Even when tiling and wave quantization are optimal, there can still be performance degradation due to: + +- Shader clock (`sclk`) throttling: Some workloads may not run at peak clock speeds due to power or thermal constraints. +- Cache behavior: Cache misses can occur even for compute-bound GEMMs, affecting throughput and instruction efficiency. + +These factors should be analyzed alongside the dimension efficiency metrics to get a complete performance picture. + +### How it's calculated + +Step by step: + +1. Identify problem shape: Extract M, N, K from input tensors. + + ```{note} + In BLAS libraries, M and N are swapped relative to PyTorch's view, as BLAS libraries are column-major while PyTorch is row-major. This mapping is handled internally. See [PyTorch source - Blas.cpp](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/Blas.cpp#L102-L129) for more detail. This is accounted for in TraceLens as well when computing the tiles across the M and N dimensions. + ``` + +2. Extract tile size: Parse `mt_m`, `mt_n` from kernel name (for example, `MT256x144x32` → `mt_m = 256, mt_n = 144`). + +3. Calculate tiles: + + ``` + float_tiles_m = M / mt_m + float_tiles_n = N / mt_n + tiles_m = ceil(float_tiles_m) + tiles_n = ceil(float_tiles_n) + num_tiles = tiles_m * tiles_n + ``` + +4. `tile_eff`: + + ``` + tile_eff = (M * N) / (tiles_m * mt_m * tiles_n * mt_n) + ``` + +5. `wq_eff` (assume `num_cus` known, for example 304 for MI300X): + + ``` + float_rounds = num_tiles / num_cus + rounds = ceil(float_rounds) + wq_eff = num_tiles / (rounds * num_cus) + ``` + +6. `dim_eff`: + + ``` + dim_eff = tile_eff * wq_eff + ``` + +### Calculation examples + +Assuming `num_cus = 304` (for example, MI300X). + +#### Example 1: Perfect tiling, suboptimal wave quantization + +- Input: M = 10240, N = 2048, K = 2048 +- Kernel tile: `mt_m = 256`, `mt_n = 64` + +```text +float_tiles_m = 40 +float_tiles_n = 32 +tiles_m = 40, tiles_n = 32 +tile_eff = 1.0 +num_tiles = 1280 +float_rounds = 4.21 +rounds = 5 +wq_eff = 0.842 +dim_eff = 0.842 +``` + +#### Example 2: Slight padding, good wave quantization + +- Input: M = 2048, N = 10240, K = 2048 +- Kernel tile: `mt_m = 256`, `mt_n = 144` + +```text +float_tiles_m = 8 +float_tiles_n = 71.11 +tiles_m = 8, tiles_n = 72 +tile_eff ≈ 0.9877 +num_tiles = 576 +float_rounds = 1.895 +rounds = 2 +wq_eff ≈ 0.947 +dim_eff ≈ 0.935 +``` + +### Important considerations + +- Scope: This analysis assumes standard tiled GEMMs. Techniques like Stream-K or Split-K are not yet modeled. +- Relevance: These metrics are most useful for compute-bound GEMMs. Use FLOPS/Byte to determine whether a GEMM is compute- or memory-bound. + +For hands-on usage, work through the [`examples/gemm_dim_eff.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/gemm_dim_eff.ipynb) notebook. This is a new feature, so please report any issues or suggestions to the TraceLens team. + +## Related topics + +- [Trace2Tree](../conceptual/trace2tree.md) +- [Perf model walkthrough](../conceptual/triton-perf-model-walkthrough.md) +- [Torch profiling analysis](../conceptual/torch-profiling-analysis.md) +- [Performance report columns](../reference/perf-report-columns.md) +- [Generate a performance report from PyTorch](../how-to/generate-perf-report-pytorch.md) diff --git a/docs_original/conceptual/imgs/profiling_analysis_fig1.png b/docs/conceptual/imgs/profiling_analysis_fig1.png similarity index 100% rename from docs_original/conceptual/imgs/profiling_analysis_fig1.png rename to docs/conceptual/imgs/profiling_analysis_fig1.png diff --git a/docs_original/conceptual/imgs/profiling_analysis_fig2.png b/docs/conceptual/imgs/profiling_analysis_fig2.png similarity index 100% rename from docs_original/conceptual/imgs/profiling_analysis_fig2.png rename to docs/conceptual/imgs/profiling_analysis_fig2.png diff --git a/docs_original/conceptual/imgs/profiling_analysis_fig3.png b/docs/conceptual/imgs/profiling_analysis_fig3.png similarity index 100% rename from docs_original/conceptual/imgs/profiling_analysis_fig3.png rename to docs/conceptual/imgs/profiling_analysis_fig3.png diff --git a/docs_original/conceptual/imgs/profiling_analysis_fig4.png b/docs/conceptual/imgs/profiling_analysis_fig4.png similarity index 100% rename from docs_original/conceptual/imgs/profiling_analysis_fig4.png rename to docs/conceptual/imgs/profiling_analysis_fig4.png diff --git a/docs_original/conceptual/imgs/profiling_analysis_fig5.png b/docs/conceptual/imgs/profiling_analysis_fig5.png similarity index 100% rename from docs_original/conceptual/imgs/profiling_analysis_fig5.png rename to docs/conceptual/imgs/profiling_analysis_fig5.png diff --git a/docs_original/conceptual/imgs/profiling_analysis_fig6.png b/docs/conceptual/imgs/profiling_analysis_fig6.png similarity index 100% rename from docs_original/conceptual/imgs/profiling_analysis_fig6.png rename to docs/conceptual/imgs/profiling_analysis_fig6.png diff --git a/docs/conceptual/inference-analysis.md b/docs/conceptual/inference-analysis.md new file mode 100644 index 000000000..757d34c68 --- /dev/null +++ b/docs/conceptual/inference-analysis.md @@ -0,0 +1,284 @@ + + +# Inference performance analysis in TraceLens +```{meta} +:description: Understand the concepts behind TraceLens inference analysis - the roofline model for paged attention (prefill and decode), and the steady-state region that trace splitting targets. +:keywords: TraceLens, inference, roofline, paged attention, prefill, decode, chunked prefill, FLOPS, arithmetic intensity, steady state, vLLM, SGLang, LLM serving, ROCm, MI300X +``` + +This topic explains the concepts behind TraceLens's inference-serving analysis: +how the *roofline model* is derived for paged attention, and what the +*steady-state region* is and why it is the part of a serving trace worth +profiling. For the corresponding procedures, see +[Generate a PyTorch inference performance report](../how-to/generate-perf-report-pytorch-inference.md). + +In inference serving, multiple requests are batched together and each request +carries its own query and key/value sequence lengths. A single execution step +typically mixes *prefill* (context) and *decode* (generation) requests, which +makes per-request accounting the natural basis for a roofline model. + +## Roofline analysis for inference attention + +The following notation is used throughout this section. + +| Symbol | Description | +|--------|-------------| +| B | Batch size (1 per request in paged attention) | +| N_Q | Number of query tokens | +| N_KV | Number of key/value tokens (context length) | +| H_Q | Number of query heads | +| H_KV | Number of KV heads (H_KV ≤ H_Q; equal for MHA, smaller for GQA/MQA) | +| d_h_qk | Head dimension for queries and keys | +| d_h_v | Head dimension for values | +| R_C | Number of context (prefill) requests in the batch | +| R_G | Number of generation (decode) requests in the batch | +| R | Total number of requests (R = R_C + R_G) | +| N_Q(i), N_KV(i) | Query and KV token counts for the i-th request | + +### Standard SDPA attention (single request) + +Scaled dot-product attention consists of two matrix multiplications per head: + +1. *QK^T (score computation)*: `2 * B * N_Q * N_KV * H_Q * d_h_qk` +2. *Score × V (value aggregation)*: `2 * B * N_Q * N_KV * H_Q * d_h_v` + +For causal attention, roughly half the score matrix is masked out: + +``` +FLOPS = (2 * B * N_Q * N_KV * H_Q * d_h_qk + 2 * B * N_Q * N_KV * H_Q * d_h_v) / 2 +``` + +``` +Elements Moved = +Q: B * N_Q * H_Q * d_h_qk +K: B * N_KV * H_KV * d_h_qk +V: B * N_KV * H_KV * d_h_v +Output: B * N_Q * H_Q * d_h_v +``` + +### Paged attention (multiple requests) + +For total FLOPS and bytes moved in inference paged attention, *sum over the +compute requirement of every request individually* (B = 1 per request). Requests +fall into two categories: + +- *Context (prefill) requests* process input tokens; attention is causal within + the current chunk. +- *Generation (decode) requests* generate new tokens; attention is non-causal + (queries attend to all past KV tokens). Typically N_Q = 1, but approaches like + speculative decoding may produce multiple query tokens per request. + +#### FLOPS: prefill (context) requests + +When chunked prefill is not enabled, the first (and only) chunk has N_KV = N_Q +and attention is causal: + +``` + R_C +FLOPS_prefill = Σ (2 * N_Q(i) * N_KV(i) * H_Q * d_h_qk + 2 * N_Q(i) * N_KV(i) * H_Q * d_h_v) / 2 + i=1 +``` + +With chunked prefill, the nth chunk already has the KV tokens from all previous +chunks cached. The attention matrix for one such request splits into a +non-causal rectangle (current chunk attending to previous chunks) plus a causal +self-attention triangle (current chunk attending to itself): + +``` + Keys (N_KV) + |<---- N_KV - N_Q ---->|<---- N_Q ---->| + +----------------------+---------------+ ^ + | |\ | | + | Non-causal | \ (masked) | | + Queries | (full rectangle) | \ | N_Q + | attend to previous | Causal \ | | + | chunks' KV cache | (self) \ | | + +----------------------+---------------+ v + <----- previous ----->|<-- current --> + chunks chunk +``` + +For the first chunk (no chunking, or the first chunk of chunked prefill), +N_KV = N_Q, so the rectangle vanishes and the whole matrix is causal: + +``` + Keys (N_KV = N_Q) + |<---- N_Q ---->| + +---------------+ ^ + |\ | | + | \ (masked) | | + | \ | | + Queries | \ | N_Q + | Causal \ | | + | (entire) \ | | + +---------------+ v +``` + +Both cases (first chunk and nth chunk) simplify to a single unified formula: + +``` + R_C +FLOPS_context = Σ [ (2 * N_Q(i) * N_KV(i) * H_Q * d_h_qk + 2 * N_Q(i) * N_KV(i) * H_Q * d_h_v) + i=1 + - (2 * N_Q(i)^2 * H_Q * d_h_qk + 2 * N_Q(i)^2 * H_Q * d_h_v) / 2 ] +``` + +This works because: + +- When N_KV = N_Q (first chunk): `full - full/2 = full/2`, which is causal. +- When N_KV > N_Q (nth chunk): `full rectangle - self-triangle`, which is the + non-causal rectangle plus the causal self-attention. + +#### FLOPS: generation (decode) requests + +Generation requests attend to all cached KV tokens (N_KV = context length so +far). Typically N_Q = 1 (autoregressive decoding), but techniques like +speculative decoding may have N_Q > 1. The attention is non-causal: + +``` + R_G +FLOPS_generation = Σ (2 * N_Q(i) * N_KV(i) * H_Q * d_h_qk + 2 * N_Q(i) * N_KV(i) * H_Q * d_h_v) + i=1 +``` + +``` +FLOPS_total = FLOPS_context + FLOPS_generation +``` + +#### Elements moved + +Total memory traffic sums over all requests. Each request reads its Q, K, V +tensors and writes the output. With GQA/MQA, the KV cache uses H_KV heads (not +H_Q), reducing KV memory traffic. Ignoring cases with shared KV pages between +requests: + +``` + R +Elements_moved = Σ ( N_Q(i) * H_Q * d_h_qk // Q read + i=1 + + N_KV(i) * H_KV * d_h_qk // K read (from paged KV cache) + + N_KV(i) * H_KV * d_h_v // V read (from paged KV cache) + + N_Q(i) * H_Q * d_h_v ) // Output write +``` + +where R = R_C + R_G is the total number of requests. Because B = 1 per request in +paged attention, the batch dimension is absorbed into the summation. + +### Roofline analysis without per-request details + +Importantly, individual request details are *not* required to perform roofline +analysis. Inspecting the formulas above, the only per-request quantities that +appear are `N_Q(i)`, `N_KV(i)`, and their products. The full FLOPS and +memory-traffic expressions can be evaluated from just these aggregate +statistics, computed separately for context and generation requests: + +| Aggregate | Used in | +|-----------|---------| +| R_C, R_G | Request counts | +| Σ N_Q | Elements moved (Q read, Output write) | +| Σ N_KV | Elements moved (K read, V read) | +| Σ (N_Q * N_KV) | FLOPS (full-rectangle term) | +| Σ (N_Q^2) | FLOPS (causal self-attention correction for prefill) | + +TraceLens obtains these aggregates by applying `torch.record_function(annotation)` +to the framework's execution steps. Because a single execution step can contain a +mix of context and generation requests, the annotation encodes the aggregate +statistics *separately* for context and generation requests within that step (for +example R_C, R_G, Σ N_Q for context, Σ N_Q for generation, and so on). These +annotations are stored as `user_annotation` events in the PyTorch profiler trace, +making roofline analysis possible directly from the trace without any additional +instrumentation or runtime logging. + +The framework patches referenced in +[Generate a PyTorch inference performance report](../how-to/generate-perf-report-pytorch-inference.md) +are what add these annotations to the trace. + +## Steady-state region + +Inference-serving execution consists of three phases: + +1. *Ramp-up*: the initial few steps where requests are still being batched. +2. *Steady state*: the execution steps with the highest concurrency. +3. *Ramp-down*: the trailing steps where the final batch of requests finishes. + +Once steady state is reached, execution consists of: + +- *Decode-only steps*. +- *Prefill-decode steps*, typically containing one prefill request packed with + ~CONC−1 decode requests. + +For performance analysis, only the steady-state steps are of interest, +specifically prefill-decode steps and decode-only steps with large context sizes +(towards the end of a request). + +### Parameters relevant to inference serving + +- *NUM_PROMPTS*: typically `10 * CONC`. +- *CONC*: number of concurrent requests that can be batched together. +- *R*: random-range ratio used for sampling input and output sequence lengths. +- *OSL*: maximum output sequence length. The per-request output length is sampled + uniformly in `[R * OSL, OSL]`. +- *ISL*: assumed to be lower than the chunk size. + +Assuming the benchmark schedules requests at an effectively infinite rate, the +first CONC steps are conservatively treated as the ramp-up phase. The execution +step ranges where groups of CONC requests complete are: + +``` +1 * R * OSL to 1 * OSL e.g. 0.8 OSL - 1 OSL +2 * R * OSL to 2 * OSL e.g. 1.6 OSL - 2 OSL +3 * R * OSL to 3 * OSL e.g. 2.4 OSL - 3 OSL +4 * R * OSL to 4 * OSL e.g. 3.2 OSL - 4 OSL +... +N * R * OSL to N * OSL + +where N = NUM_PROMPTS / CONC +``` + +``` +TOTAL_STEPS = NUM_PROMPTS * Avg(OSL) / CONC +TOTAL_PrefillDecode_Steps = NUM_PROMPTS +TOTAL_DecodeOnly_Steps = NUM_PROMPTS * ((Avg(OSL) - CONC) / CONC) +``` + +### Recommended profiling window + +Since serving benchmarks commonly sample output sequence lengths from +`[R * OSL, OSL]`, the most useful steady-state profiling window lies in: + +``` +5 * R * OSL to 5 * OSL +``` + +Here, the probability of a request finishing at step t is non-uniform, peaking at +`((R+1)/2) * 5 * OSL`. This region exhibits fully saturated concurrency, a +representative mix of decode-only and prefill-decode steps, and minimal warm-up or +tail artifacts. That makes the recommended profiling window: + +``` +max_iters = 16 * OSL / CONC + # Number of execution steps to profile. The multiplier 16 targets + # ~16 execution steps with a prefill+decode mix. Clamp max_iters + # for extreme values of OSL/CONC. + +delay_iters = ((R+1)/2) * 5 * OSL - (max_iters / 2) + # Execution step where the profiler starts. +``` + +The trace-splitting workflow uses these definitions to detect the steady-state +region automatically and to separate prefill-decode from decode-only steps. When +the benchmark parameters are known, pass `--CONC`, `--OSL`, and `--R` to the +splitter to override the empirically estimated ratio with the analytically +derived one; see +[Split inference traces](../how-to/generate-perf-report-pytorch-inference.md#split-inference-traces-optional). + +## Related topics + +- [Generate a PyTorch inference performance report](../how-to/generate-perf-report-pytorch-inference.md) +- [Compare two traces in TraceLens](../how-to/compare-traces.md) +- [GEMM analysis in TraceLens](./gemm-analysis.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/conceptual/shape-metadata.md b/docs/conceptual/shape-metadata.md new file mode 100644 index 000000000..58af9ef62 --- /dev/null +++ b/docs/conceptual/shape-metadata.md @@ -0,0 +1,151 @@ + + +# Tensor shape metadata in PyTorch profiler traces +```{meta} +:description: A conceptual explanation of when tensor shapes appear in PyTorch profiler traces, why some operations lack them, and how to register operations so shapes are recorded. +:keywords: tensor shapes, PyTorch profiler, cpu_op, dispatcher, torch.library, custom op, Triton, FlashInfer, vLLM, SGLang, Input Dims, backward events, TraceLens +``` + +Tensor shapes are one of the most useful pieces of metadata in a PyTorch profiler trace: they let you match trace events to model architecture, analyze kernel efficiency, and spot unusual stride or dtype patterns. But shapes aren't always present. This topic explains when they're available, why they sometimes go missing, and how to get them back. + +## When are shapes available? + +The simple rule is that shapes are available when an operation goes through PyTorch's dispatcher as a `cpu_op`. + +In profiler traces, look for events with: + +- `"cat": "cpu_op"` +- `"args": { "Input Dims": [...], "Input type": [...] }` + +### Operations that have shapes + +| Operation type | Example | Why it works | +|---|---|---| +| Native ATen ops | `aten::mm`, `aten::linear` | Built into the PyTorch dispatcher | +| Registered custom ops | `torch.ops.mylib.my_op` | Registered via `torch.library` | +| TorchScript ops | JIT-compiled functions | Goes through the dispatcher | +| Custom `autograd.Function` | User-defined forward | Forward call is dispatched | +| Distributed collectives | `record_param_comms` | Instrumented by PyTorch | + +### Operations that don't have shapes + +| Operation type | Example | Why it fails | +|---|---|---| +| Plain Python functions | `def my_kernel(x, y): ...` | Bypasses the dispatcher | +| Triton kernels | `@triton.jit` | Called as a Python function | +| FlashInfer ops | GEMM, MoE, attention | Registration intentionally disabled | +| Backward engine events | `autograd::engine::evaluate_function:*Backward` | Empty inputs passed to the profiler | + +### Quick reference: event categories + +``` +cpu_op → Usually has shapes (exception: backward events) +python_function → No shapes +kernel → No shapes (GPU-side event) +cuda_runtime → No shapes (API-level event) +``` + +## How to get shapes when they're missing + +If an operation lacks shapes, the fix is to register it with the dispatcher through `torch.library`. There are two approaches. + +### Option 1: register as a custom op + +Wrap the operation with `torch.library.custom_op`: + +```python +@torch.library.custom_op("mylib::triton_matmul", mutates_args=()) +def triton_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return my_triton_kernel(a, b) + +# Now call via torch.ops +result = torch.ops.mylib.triton_matmul(a, b) +``` + +The event then appears as `cpu_op` with `Input Dims`. + +The `mutates_args` argument tells PyTorch which input arguments the function modifies in-place: + +- `mutates_args=()` — no inputs are modified (a pure function). +- `mutates_args=("output",)` — the `output` argument is modified in-place. +- Required for correctness with `torch.compile` and autograd. + +### Option 2: lighter-weight registration (vLLM approach) + +Use the lower-level `Library` API directly. For reference, see [`vllm/utils/torch_utils.py:742` — `direct_register_custom_op`](https://github.com/vllm-project/vllm/blob/0b225fb7b22f8ae1f5fc8ee618640ae0983c76de/vllm/utils/torch_utils.py#L742-L780). + +```python +from torch.library import Library +from torch._library.infer_schema import infer_schema + +my_lib = Library("mylib", "FRAGMENT") + +def register_op(op_name, op_func, mutates_args=None): + schema = infer_schema(op_func, mutates_args=mutates_args or []) + my_lib.define(op_name + schema) + my_lib.impl(op_name, op_func, dispatch_key="CUDA") + +# Register +register_op("triton_matmul", triton_matmul_impl) +``` + +### Choosing between the two options + +| Aspect | Option 1: `custom_op` | Option 2: `Library.define()` | +|---|---|---| +| Simplicity | Decorator, minimal code | More boilerplate | +| Overhead | Higher (full dispatcher) | Lower (direct to CUDA) | +| `torch.compile` | Full support | Works with `register_fake` | +| Use when | Prototyping, one-off ops | Performance-critical paths | + +## Framework-specific status + +| Framework | Current state | How to get shapes | +|---|---|---| +| PyTorch ATen | Has shapes | Already works | +| vLLM (standard) | Has shapes | Uses `direct_register_custom_op` | +| vLLM (OAI Triton) | Missing | Needs registration | +| FlashInfer | Disabled | Set `FLASHINFER_ENABLE_PROFILER_METADATA=1` (proposed) | +| SGLang (CUDA) | Has shapes | Uses `torch.ops.sgl_kernel.*` | +| SGLang (Triton) | Missing | Needs registration | + +## Summary + +- Shapes are tied to dispatcher registration. If it's a `cpu_op`, it has shapes. +- The fix is straightforward: register operations via `torch.library`. +- Start simple with `@torch.library.custom_op`. +- Optimize if needed by switching to `Library.define()` for lower overhead. +- FlashInfer disabled shape recording on purpose — a performance versus observability trade-off that can be re-enabled. + +## Appendix + +### Backward events and sequence number linking + +Backward engine events (`autograd::engine::evaluate_function:*Backward`) are `cpu_op` but don't have `Input Dims`. They have `Sequence number` instead, which links to the corresponding forward op. + +```python +def get_backward_shapes(trace_events, backward_event): + seq_num = backward_event["args"].get("Sequence number") + if seq_num is None: + return None + + # Find forward op with same sequence number + for event in trace_events: + if event.get("cat") == "cpu_op" and \ + event["args"].get("Sequence number") == seq_num and \ + "Backward" not in event.get("name", ""): + return event["args"].get("Input Dims") + return None +``` + +## Related topics + +- [Trace2Tree](../conceptual/trace2tree.md) +- [Understanding PyTorch traces](../conceptual/torch-profiling-analysis.md) +- [Generate a performance report from PyTorch traces](../how-to/generate-perf-report-pytorch.md) +- [Performance report columns](../reference/perf-report-columns.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/conceptual/torch-profiling-analysis.md b/docs/conceptual/torch-profiling-analysis.md new file mode 100644 index 000000000..165e46f68 --- /dev/null +++ b/docs/conceptual/torch-profiling-analysis.md @@ -0,0 +1,151 @@ + + +# Understanding PyTorch traces +```{meta} +:description: A conceptual walkthrough of how PyTorch traces capture host and GPU execution, how a single Python operation expands into GPU kernels, and how to read a single-GPU trace in Perfetto. +:keywords: PyTorch profiler, trace analysis, Perfetto, GPU kernels, ATen, MIOpen, cuDNN, host device execution, tensor shapes, autograd, memory copy, TraceLens +``` + +Once you've collected a PyTorch trace and opened it in [Perfetto](https://ui.perfetto.dev/), the next question is what the trace actually means and how you read insights from it. This topic unpacks that for a *single-GPU* trace. Multi-GPU profiling is covered separately. + +When you run a PyTorch model on a GPU, there's a hidden interplay between the CPU (host) and the GPU (device). A PyTorch trace lets you observe this choreography, revealing how your code executes, where time is spent, and what might be slowing you down. + +## The execution model: host, GPU, and asynchronous launches + +![Async execution example](./imgs/profiling_analysis_fig1.png) + +*Figure 1: Asynchronous execution example.* + +To understand profiling, you need a mental model of how PyTorch — or any application — runs code across the CPU and GPU. + +The host (CPU) is the conductor. It manages processes and threads, moving through the application call stack until it issues GPU runtime or driver application programming interface (API) calls such as `cudaLaunchKernel` or `hipLaunchKernel`. The GPU is the performer. It executes tasks with massive parallelism, often involving thousands of threads working on matrix multiplications, convolutions, or reductions. + +Most of these commands are asynchronous. The host queues the commands (kernels) and immediately continues, while the GPU independently drains its queue. This overlap is what enables high performance. Sometimes, however, the host introduces synchronization points such as `cudaDeviceSynchronize` or `hipDeviceSynchronize`. In those cases the CPU waits until the GPU finishes all queued work. Synchronization is necessary for correctness, but too many synchronizations can quickly turn into bottlenecks. + +One more distinction: host events come with detailed call stacks (Python to C++ to runtime), so you can see exactly how an operation was triggered. GPU events, on the other hand, are flat. A GPU kernel can't call another GPU kernel on its own — the launch must always come from the host. That's why GPU traces appear as a simple queue of kernels, memory copies, and memory sets with no stack information. + +A kernel itself is just a GPU function that runs across thousands of lightweight threads in parallel. In the trace, each kernel shows up as a single event: you can see when it started, how long it ran, and on which stream — but not what happened inside. To inspect that level of detail (thread execution, memory transactions, warp occupancy), you need dedicated kernel profilers such as [rocprof-compute](https://rocm.docs.amd.com/projects/rocprofiler-compute/en/latest/) for AMD GPUs or Nsight Compute for NVIDIA GPUs. These tools complement the trace view by showing what happens inside the GPU during a kernel event. + +## A simple example + +To dig deeper into what a trace shows, consider the snippet below. + +![Perfetto trace showing idle time between kernels](./imgs/profiling_analysis_fig2.png) + +*Figure 2: Perfetto trace highlighting idle time between convolution and batch normalization kernels.* + +The red dashed region marks idle time: the GPU is waiting for the host to issue the next command. This happens when the CPU frontend can't keep up with the GPU's execution speed. Reducing such idle gaps is an important optimization goal. To see why these gaps occur, follow the path from a high-level PyTorch call down to the GPU kernels that actually run. + +Take a simple example: + +```python +y = nn.Conv2d(3, 64, kernel_size=3)(x) +``` + +At the Python level, this looks like a single operation. In the trace, however, it expands into several layers: + +- *Python frontend*: `nn.Conv2d` in `torch.nn`. +- *ATen*: the call lowers into PyTorch's tensor library, [ATen](https://pytorch.org/cppdocs/notes/aten.html), and shows up as `aten::conv2d` in the trace. +- *Backend wrapper*: ATen provides wrappers that call into vendor libraries. On AMD GPUs, you'll see `aten::miopen_convolution`, which wraps [MIOpen](https://rocm.docs.amd.com/projects/MIOpen/en/latest/) commands. On NVIDIA GPUs, the equivalent is `aten::cudnn_convolution`, which wraps [cuDNN](https://developer.nvidia.com/cudnn) calls. +- *GPU kernels*: the backend library enqueues device kernels such as `igemm_fwd_gtcx2_nhwc` that perform the actual convolution on the GPU. + +This gives a clear understanding of how high-level code is translated into GPU execution. + +## Tensor metadata in the trace + +Another important detail is tensor metadata. Python-level operations in the trace don't carry shape information. But if you enable `record_shapes=True` in the profiler, the `cpu_op` category (such as `aten::convolution`) includes input shapes, strides, and dtypes. + +For example, in the figure below the first tensor is the activation (`[5, 64, 56, 56]`) and the second tensor is the convolution filter (`[64, 64, 3, 3]`). + +![Perfetto trace showing recorded input shapes](./imgs/profiling_analysis_fig3.png) + +*Figure 3: Shapes recorded on the backend op when `record_shapes=True`.* + +This metadata becomes very useful for deeper analysis and debugging — for example, when matching trace events to model architecture, analyzing kernel efficiency, or spotting unusual stride or dtype patterns. + +That's all you need for a clean first pass: know what the host and GPU are doing, read the timeline, and use recorded shapes to ground what you see. Perfetto gives you the raw signals; turning them into insight is a skill you build with practice, and TraceLens helps accelerate that process. + +The appendix below covers UI shortcuts, how memory copies show up, and the structure of the raw trace file. + +## Appendix + +### Perfetto UI tips + +Perfetto is a powerful trace viewer, but it takes some practice to navigate effectively. A few basics: + +- *Zoom and pan* with `Ctrl + scroll` to zoom in and out; click and drag to pan. +- *Event details*: click any event to open the *Current Selection* panel below. This shows start time, duration, and arguments. With `record_shapes=True`, backend ops (in the `cpu_op` category) also show tensor shapes, dtypes, and strides. + +Perfetto also links host launches and GPU execution with arrows called *flows*. These are crucial for connecting what you see on the CPU timeline with what actually runs on the GPU. + +When you select a runtime launch event such as `hipExtModuleLaunchKernel`, the *Following Flow* jumps you forward to the GPU kernel it triggered: + +![Following flow from host launch to GPU kernel](./imgs/profiling_analysis_fig4.png) + +*Figure 4: Following flow from a host launch (`hipExtModuleLaunchKernel`) to the corresponding GPU kernel event.* + +Conversely, when you select a GPU kernel event, the *Preceding Flow* takes you back to the runtime call on the host that launched it: + +![Preceding flow from GPU kernel back to host launch](./imgs/profiling_analysis_fig5.png) + +*Figure 5: Preceding flow from a GPU kernel (`SubTensorOpWithScalar1d`) back to its launch on the host.* + +These flows are the bridge between the Python-level trace and the GPU execution timeline. They let you answer both: + +- Which Python or ATen op caused this kernel to run? +- What GPU work did this runtime call produce? + +### Memory copy events + +Not all GPU activity is compute. Profiling traces also show memory transfers: + +- *H2D (host to device)*: copies data from CPU to GPU, usually synchronous and PCIe bandwidth-limited. +- *D2H (device to host)*: copies results back to CPU, also synchronous and PCIe bandwidth-limited. +- *D2D (device to device)*: moves data between GPU buffers, asynchronous and limited by HBM bandwidth. + +Recent versions even record measured bandwidth for these events in the trace arguments. Importantly, memory copy events use the GPU's DMA engines, not compute cores, so they don't directly compete with kernel execution. + +### JSON format + +Behind the Perfetto UI, the PyTorch profiler saves traces as JSON. Each entry is an event dictionary with: + +- *Timestamps*: `ts` (start) and `dur` (duration). +- *Process and thread IDs*: `pid` and `tid`. CPU events use real PIDs and TIDs; GPU events use pseudo-PIDs for devices and TIDs for streams. +- *Category*: for example, `python_function`, `cpu_op`, `cuda_runtime`, `kernel`, or `gpu_memcpy`. +- *Args*: extra information such as shapes, dtypes, strides, or bandwidth. + +The JSON format makes traces scriptable: you can parse them to build custom reports or run automated analysis outside Perfetto. This is exactly what TraceLens does. + +### Autograd in the trace + +Autograd introduces another dimension to the trace: the forward and backward passes are captured on different threads within the same process. In the timeline, you typically see: + +- `python3.x [main thread]` for the forward pass. +- `pt_autograd_0` (or similar) for the backward pass. + +These are linked at the `aten::convolution` layer of the call stack. Perfetto uses flows to connect the forward convolution op to its corresponding backward node. + +![Forward and backward convolution linked in the trace](./imgs/profiling_analysis_fig6.png) + +*Figure 6: Example of `aten::convolution` (forward) linked to `ConvolutionBackward0` (backward) via flows.* + +Key points: + +- The linkage happens at `aten::convolution`, not `aten::conv2d`. The higher-level `aten::conv2d` call funnels into `aten::convolution` before reaching the backend. +- Forward and backward both eventually call into the same GPU backend (MIOpen or cuDNN), so you see similar kernel launches in both directions. +- The backward pass typically runs more kernels, since it must compute gradients for both activations and weights, making it more expensive than the forward pass. +- Because autograd runs on its own thread, you can easily separate model execution (forward path) from gradient computation (backward path) when analyzing traces. + +Together with the flows discussed earlier, this lets you connect a forward convolution op to its exact backward counterpart and then follow the chain down to GPU kernels. + +## Related topics + +- [Trace2Tree](../conceptual/trace2tree.md) +- [GEMM analysis](../conceptual/gemm-analysis.md) +- [Generate a performance report from PyTorch traces](../how-to/generate-perf-report-pytorch.md) +- [Performance report columns](../reference/perf-report-columns.md) +- [What is TraceLens?](../what-is-tracelens.md) diff --git a/docs/conceptual/trace2tree.md b/docs/conceptual/trace2tree.md new file mode 100644 index 000000000..49263c24c --- /dev/null +++ b/docs/conceptual/trace2tree.md @@ -0,0 +1,60 @@ + + +# The Trace2Tree data model +```{meta} +:description: Understand why Trace2Tree exists and how its four-layer call tree links Python code and backend operations down to GPU kernels for portable, interpretable performance analysis. +:keywords: TraceLens, Trace2Tree, TreePerfAnalyzer, call tree, GPU kernel, PyTorch profiler, cpu_op, aten, HLO, JAX, ROCm, AMD, performance analysis +``` + +In GPU application performance analysis, understanding the relationship between host CPU operations and the corresponding GPU kernel executions is crucial for finding bottlenecks. The PyTorch profiler provides a JSON trace file containing events with timestamps and durations, but it lacks explicit call-stack dependency information. + +`Trace2Tree` is the underlying tree-structure component that TraceLens uses to parse trace files and build hierarchical dependency relationships, from host CPU operations down to GPU kernels. + +```{note} +It's recommended that you access this functionality through the `TreePerfAnalyzer` interface rather than using `Trace2Tree` directly. +``` + +## Why kernel names aren't enough + +Directly inspecting GPU kernel names has two fundamental limitations: + +* *Ambiguous semantics (and weak reproducibility)*: a single kernel name can map to many different computations depending on shape, dtype, strides or layout, and so on. Shape strongly affects performance — one shape might select a fast tiled path while another shape of the same op type falls onto a slower algorithm. Because the name omits this argument context, you can't reliably understand, compare, or reproduce the workload from the kernel string alone. Many kernel names are also cryptic and unreadable, for example `Cijk_Ailk_Bljk_*` or `void cutlass_*`. + +* *Platform-dependent, unstable naming (and weak cross-platform comparison)*: the same high-level operation appears under different kernel names across platforms. For example, a single GEMM shows up as `nvjet_*` or `cutlass_*` on NVIDIA H100, and as a Tensile kernel `Cijk_Ailk_Bljk_*` on AMD MI300. These names also shift across software versions, so raw kernel strings aren't a stable abstraction for comparison. + +## What Trace2Tree does + +`Trace2Tree` reconstructs a full call tree from the Python front end down to each GPU kernel. The tree has exactly four layers: + +1. *Python front end* — user code or `nn.Module`. +2. *Operation* — on PyTorch this is the dispatch operation (`cpu_op`), for example `aten::mm`, `aten::addmm`, and so on. On JAX the corresponding layer is the HLO operations. This is the layer that contains the argument information used to contextualize the kernel. +3. *HIP / CUDA runtime* — launch API calls. +4. *GPU kernel* — the executed kernel. + +## How this solves the problem + +* *Disambiguates semantics*: argument metadata at the backend op layer lets you group identical computations, attribute time, and deterministically reproduce slow cases. +* *Enables fair comparison*: operations such as `aten::mm` and HLO are stable across platforms. By anchoring analysis there, you can compare the same operation and arguments regardless of how the kernels are named underneath. +* *Flexible attribution*: GPU time can be viewed at any level — module (through its backend ops), dispatch op, runtime, or kernel — depending on the question. As an additional benefit, time can be attributed all the way up to the Python `nn.Module` level, making performance insights accessible even to users outside the performance-engineering field. This helps bridge the gap between model developers and low-level hardware execution. + +Kernel names are volatile and context-free. Trace2Tree anchors analysis at the stable backend operation, enriches it with arguments, and maps the full execution stack to deliver portable, interpretable performance insight. + +That said, kernel names are often useful — they can offer clues about the backend implementation, algorithm variant, or compiler choices. TraceLens intends to serve as a one-stop solution for extracting every bit of useful signal from a trace file, so it includes features to extract and parse relevant information from kernel names where applicable. But it treats them as supplementary, not foundational. + +## A powerful IR for deeper analysis + +This tree is a powerful intermediate representation (IR) that forms the foundation of many TraceLens features. It captures both the structural and performance semantics of a workload: the hierarchical dependency between CPU operations and GPU kernel launches, the argument metadata that disambiguates each operation, and the forward-versus-backward relationship between operations. Because it's a lightweight, framework-agnostic structure (built today for PyTorch profiler JSON, with JAX support and room for more), higher-level analyses — GPU-timeline breakdowns, roofline metrics, `nn.Module` attribution, event replay, and trace diffing — all build on top of it. + +You typically don't construct or traverse the tree by hand. It's accessed through the `TreePerfAnalyzer` interface, which builds the tree from a trace and exposes the navigation and metric APIs. To work with the tree in practice, see [Analyze traces with the TraceLens SDK](../how-to/sdk-analysis.md). + +## Related topics + +- [Analyze traces with the TraceLens SDK](../how-to/sdk-analysis.md) +- [Replay a single operation in TraceLens](../how-to/event-replay.md) +- [Understanding PyTorch traces](../conceptual/torch-profiling-analysis.md) +- [API reference](../reference/api-reference.md) +- [What is TraceLens?](../what-is-tracelens.md) diff --git a/docs_original/triton_perf_model_walkthrough.md b/docs/conceptual/triton-perf-model-walkthrough.md similarity index 76% rename from docs_original/triton_perf_model_walkthrough.md rename to docs/conceptual/triton-perf-model-walkthrough.md index 968b33908..2f5edc4e1 100644 --- a/docs_original/triton_perf_model_walkthrough.md +++ b/docs/conceptual/triton-perf-model-walkthrough.md @@ -1,37 +1,23 @@ -# Triton Kernel Performance Model +# Triton kernel performance model -How TraceLens computes GFLOPS, TB/s, and other performance metrics for -torch.compile-generated Triton kernels. - -**Source code:** `TraceLens/PerfModel/triton_compiled_perf_model.py` - -## Table of Contents +```{meta} +:description: How TraceLens computes GFLOPS, TB/s, and other performance metrics for torch.compile-generated Triton kernels, including the metric formulas and worked examples. +:keywords: TraceLens, Triton, torch.compile, TorchInductor, performance model, GFLOPS, TB/s, roofline, pointwise, reduction, PyTorch profiler, ROCm, AMD +``` -1. [Background](#1-background) -2. [Requirements and PyTorch Version Compatibility](#2-requirements-and-pytorch-version-compatibility) - - [Trace capture requirements](#21-trace-capture-requirements) - - [PyTorch version compatibility](#22-pytorch-version-compatibility) - - [`TORCHINDUCTOR_UNIQUE_KERNEL_NAMES`](#23-torchinductor_unique_kernel_names) - - [Summary: recommended setup](#24-summary-recommended-setup) -3. [How It Works](#3-how-it-works) - - [Data flow](#31-data-flow) - - [Two-tier metadata extraction](#32-two-tier-metadata-extraction) -4. [Inductor Triton Kernel Types](#4-inductor-triton-kernel-types) -5. [Metric Formulas](#5-metric-formulas) -6. [Worked Example: Pointwise Kernel (SwiGLU)](#6-worked-example-pointwise-kernel-swiglu) -7. [Worked Example: Reduction Kernel (RMSNorm)](#7-worked-example-reduction-kernel-rmsnorm) -8. [Results for the Example TransformerBlock](#8-results-for-the-example-transformerblock) -9. [References](#references) +This topic explains how TraceLens computes GFLOPS, TB/s, and other performance +metrics for `torch.compile`-generated Triton kernels. ---- +The performance model is implemented in +[`TraceLens/PerfModel/triton_compiled_perf_model.py`](https://github.com/AMD-AGI/TraceLens/blob/main/TraceLens/PerfModel/triton_compiled_perf_model.py). -## 1. Background +## Background When you run `torch.compile(model)`, [TorchDynamo](https://docs.pytorch.org/docs/stable/torch.compiler_deepdive.html) @@ -49,17 +35,15 @@ triton_poi_fused__unsafe_view_mul_silu_2 └── generated by Triton/Inductor ``` -**Pointwise** kernels (`triton_poi_*`) process each element independently. +*Pointwise* kernels (`triton_poi_*`) process each element independently. One loop dimension: `xnumel` (total elements). -**Reduction** kernels (`triton_red_*`, `triton_per_*`) reduce along one axis. +*Reduction* kernels (`triton_red_*`, `triton_per_*`) reduce along one axis. Two loop dimensions: `xnumel` (outer) and `rnumel` (reduction axis). ---- +## Requirements and PyTorch version compatibility -## 2. Requirements and PyTorch Version Compatibility - -### 2.1 Trace capture requirements +### Trace capture requirements TraceLens uses a two-tier approach (V2 primary, V1 fallback) to extract perf metrics. For V2 to work, the trace must be captured with @@ -76,10 +60,12 @@ with torch.profiler.profile( model(inputs) ``` -Without `record_shapes=True`, the trace will not contain `Concrete Inputs`, -`Input Dims`, or `Input type` — and V2 will fall back to V1. +```{note} +Without `record_shapes=True`, the trace won't contain `Concrete Inputs`, +`Input Dims`, or `Input type` — and V2 falls back to V1. +``` -### 2.2 PyTorch version compatibility +### PyTorch version compatibility V2 requires specific trace event fields. The following table shows verified availability (confirmed from actual traces captured on each version): @@ -94,33 +80,35 @@ availability (confirmed from actual traces captured on each version): | `num_warps` | No | Yes | Not used (informational) | | `num_stages` | No | Yes | Not used (informational) | -**V2 core fields are available since PyTorch 2.4.** The extra fields -(`kernel_file`, `num_warps`, etc.) were added later but are not required +V2 core fields are available since PyTorch 2.4. The extra fields +(`kernel_file`, `num_warps`, and so on) were added later but aren't required by TraceLens. -### 2.3 `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES` +### `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES` Inductor controls kernel naming via `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES`. -When **enabled**, kernels get descriptive names that encode the kernel -category and fused ops (e.g., `triton_poi_fused_add_mul_silu_0`). -When **disabled**, kernels are named generically (`triton_0`, `triton_1`). +When *enabled*, kernels get descriptive names that encode the kernel +category and fused ops (for example, `triton_poi_fused_add_mul_silu_0`). +When *disabled*, kernels are named generically (`triton_0`, `triton_1`). -This is a **compile-time** setting — it must be set when the trace is +```{note} +This is a *compile-time* setting — it must be set when the trace is captured, not when TraceLens analyzes it. +``` -**Default by PyTorch version:** +Default by PyTorch version: | PyTorch Version | Default | Notes | |-----------------|---------|-------| -| 2.4 – 2.5 | **Disabled** | Must opt in: `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1` | -| 2.6+ | **Enabled** | [Default changed to `"1"`](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/_inductor/config.py) | +| 2.4 – 2.5 | Disabled | Must opt in: `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1` | +| 2.6+ | Enabled | [Default changed to `"1"`](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/_inductor/config.py) | Source: [`torch/_inductor/config.py`](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/config.py) — verified at tags [v2.4.0](https://github.com/pytorch/pytorch/blob/v2.4.0/torch/_inductor/config.py) and [v2.6.0](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/_inductor/config.py). -**Impact on TraceLens:** +Impact on TraceLens: | Metric | Unique names enabled | Unique names disabled | |--------|---------------------|----------------------| @@ -136,19 +124,17 @@ capture to get FLOPs: TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 python your_training_script.py ``` -### 2.4 Summary: recommended setup +### Recommended setup Verified on PyTorch 2.4 and 2.11 (with `record_shapes=True`): -- **PyTorch 2.11** — recommended, everything works out of the box -- **PyTorch 2.4–2.5** — works with `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1` for FLOPs -- **PyTorch 2.6–2.10** — expected to work (unique names enabled by default per [source code](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/_inductor/config.py)) but not directly verified - ---- +- PyTorch 2.11 — recommended, everything works out of the box +- PyTorch 2.4–2.5 — works with `TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1` for FLOPs +- PyTorch 2.6–2.10 — expected to work (unique names enabled by default per [source code](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/_inductor/config.py)) but not directly verified -## 3. How It Works +## How it works -### 3.1 Data flow +### Data flow ``` Chrome Trace (.json.gz) @@ -193,13 +179,13 @@ Chrome Trace (.json.gz) │ (tree_perf.py:compute_perf_metrics) ``` -### 3.2 Two-tier metadata extraction +### Two-tier metadata extraction TraceLens needs three pieces of information to compute metrics: -1. **Element counts** — `xnumel` and `rnumel` -2. **Data types and shapes** — to calculate bytes moved -3. **Fused ATen ops** — to calculate FLOPs +1. Element counts — `xnumel` and `rnumel` +2. Data types and shapes — to calculate bytes moved +3. Fused ATen ops — to calculate FLOPs These are obtained through a two-tier approach. V2 is tried first; V1 is the fallback: @@ -211,10 +197,10 @@ if self._meta is None: self._meta = _lookup(self.name, kwargs.get("inductor_cache_dir")) # V1: cache fallback ``` -### V2: Trace-Intrinsic (primary) +### V2: trace-intrinsic (primary) Extracts metadata directly from the Chrome trace event's `args` dict. -Available in PyTorch 2.4+ traces (see [Section 2.2](#22-pytorch-version-compatibility)). +Available in PyTorch 2.4+ traces (see [PyTorch version compatibility](#pytorch-version-compatibility)). No disk I/O needed — the trace is self-contained. ```json @@ -234,13 +220,14 @@ No disk I/O needed — the trace is self-contained. | `Input Dims` | Per-tensor shapes for exact byte calculation | | `Input type` | C10 dtype strings (`c10::BFloat16` = 2 bytes, `c10::Float` = 4 bytes) | -### V1: Cache-Based (fallback) +### V1: cache-based (fallback) -Parses the Inductor wrapper `.py` files from the torch.compile cache directory. +Parses the Inductor wrapper `.py` files from the `torch.compile` cache directory. Used when trace args are missing (older PyTorch traces) or when the trace was collected without `record_shapes=True`. Cache dirs searched (first match wins): + 1. `--inductor_cache_dir` CLI argument (explicit) 2. `$TORCHINDUCTOR_CACHE_DIR` environment variable 3. `~/.cache/torchinductor` @@ -260,9 +247,9 @@ size_hints=[268435456] # older list format 'signature': {'in_out_ptr0': '*bf16', 'in_ptr0': '*bf16', 'xnumel': 'i32'} ``` -### V1 vs V2 comparison +### V1 compared with V2 -| | V2 (Trace-Intrinsic) | V1 (Cache-Based) | +| | V2 (trace-intrinsic) | V1 (cache-based) | |--|----------------------|-------------------| | Data source | Chrome trace `event["args"]` | Inductor cache `.py` files on disk | | PyTorch version | 2.4+ | Any | @@ -272,9 +259,7 @@ size_hints=[268435456] # older list format | External dependency | None — trace is self-contained | Requires cache directory on disk | | Works on shared traces | Yes | Only if cache is accessible | ---- - -## 4. Inductor Triton Kernel Types +## Inductor Triton kernel types Inductor generates Triton kernels in several categories, each with a distinct 3-character prefix derived from the category name (`category[:3]`). @@ -283,25 +268,25 @@ These 6 categories have been stable from PyTorch 2.3 through 2.11+ | Prefix | Category | Description | |--------|----------|-------------| -| `triton_poi_` | `pointwise` | Element-wise ops (add, mul, relu, silu, etc.). One loop dimension: `xnumel`. Fused ops stay in registers — no intermediate global memory writes. | -| `triton_red_` | `reduction` | Reduction ops (sum, mean, batch norm, etc.) using a looped strategy. Two dimensions: `xnumel` (outer), `rnumel` (reduction axis). Used when the reduction axis is too large for persistent strategy. | +| `triton_poi_` | `pointwise` | Element-wise ops (add, mul, relu, silu, and so on). One loop dimension: `xnumel`. Fused ops stay in registers — no intermediate global memory writes. | +| `triton_red_` | `reduction` | Reduction ops (sum, mean, batch norm, and so on) using a looped strategy. Two dimensions: `xnumel` (outer), `rnumel` (reduction axis). Used when the reduction axis is too large for persistent strategy. | | `triton_per_` | `persistent_reduction` | Same reduction ops, but keeps the entire reduction axis in registers/SRAM. Used when `rnumel` is small enough (~1024 elements). Inductor's [`should_use_persistent_reduction()`](https://karthick.ai/blog/2025/Learn-By-Doing-Torchinductor-Reduction/) heuristic decides. | | `triton_tem_` | `template` | Complex ops like matrix multiplication (mm, addmm, _scaled_mm) using pre-defined Triton templates with optional fused epilogues. | -| `triton_for_` | `foreach` | Fused operations across lists of tensors (e.g., optimizer step updates applied to all parameters at once). | -| `triton_spl_` | `split_scan` | Split scan operations (e.g., cumulative sums with decoupled lookback). | +| `triton_for_` | `foreach` | Fused operations across lists of tensors (for example, optimizer step updates applied to all parameters at once). | +| `triton_spl_` | `split_scan` | Split scan operations (for example, cumulative sums with decoupled lookback). | -**TraceLens coverage:** +TraceLens coverage: | Category | Modeled by TraceLens? | Notes | |----------|----------------------|-------| | `triton_poi_` | Yes | `TritonCompiledPerfModel` computes FLOPs and bytes | | `triton_red_` | Yes | Same perf model, two-dimensional (xnumel × rnumel) | | `triton_per_` | Yes | Treated identically to `triton_red_` | -| `triton_tem_` | No (not needed) | GEMM ops already modeled by `aten_mm`, `aten_scaled_mm`, etc. | +| `triton_tem_` | No (not needed) | General matrix multiplication (GEMM) ops already modeled by `aten_mm`, `aten_scaled_mm`, and so on | | `triton_for_` | No | Typically optimizer-step ops, not performance-critical | | `triton_spl_` | No | Scan operations, uncommon in typical workloads | -**Naming pattern:** +Naming pattern: ``` triton_{category}_fused_{op1}_{op2}_{...}_{index} @@ -311,9 +296,7 @@ The fused ops are listed alphabetically. The index is a sequential counter within the compiled graph. When `unique_kernel_names` is disabled, the entire name collapses to `triton_{index}`. ---- - -## 5. Metric Formulas +## Metric formulas ### FLOPs @@ -322,19 +305,19 @@ flops = sum(flops_per_elem[op] for op in fused_aten_ops) * xnumel * rnumel ``` The `flops_per_elem` table is in `_FLOPS_PER_ELEM`. Examples: `add`=1, `mul`=1, -`silu`=4, `gelu`=8, `exp`=4, `pow`=2, `rsqrt`=2. Memory-only ops (`clone`, +`silu`=4, `gelu`=8, `exp`=4, `pow`=2, `rsqrt`=2. Memory-only ops (`clone`, `_to_copy`, `_unsafe_view`) contribute 0. ### Bytes -**V2 (trace args):** +V2 (trace args): ``` bytes = sum(prod(shape_i) * bytes_per_elem_i for each tensor input) + ptr_bytes[0] * xnumel # only for pointwise (output write) ``` -**V1 (cache files):** +V1 (cache files): ``` # Reduction (rnumel > 1): @@ -349,7 +332,7 @@ read and written — they appear once in the signature but move data twice. ### Throughput -All throughput metrics use the **busy kernel time** from `GPUEventAnalyser` +All throughput metrics use the *busy kernel time* from `GPUEventAnalyser` (aggregated GPU execution time, not the CPU-side `event["dur"]`): ``` @@ -359,14 +342,12 @@ FLOPS/Byte = flops / bytes_moved Data Moved (MB) = bytes_moved / (1024 * 1024) ``` ---- - -## 6. Worked Example: Pointwise Kernel (SwiGLU) +## Worked example: pointwise kernel (SwiGLU) Kernel: `triton_poi_fused__unsafe_view_mul_silu_2` Operation: `silu(gate_proj(x)) * up_proj(x)` in the SwiGLU MLP -### 6.1 Trace event args (V2 input) +### Trace event args (V2 input) ``` Concrete Inputs: ['', '', '268435456'] @@ -374,18 +355,19 @@ Input Dims: [[8, 4096, 8192], [32768, 8192], []] Input type: ['c10::BFloat16', 'c10::BFloat16', 'Scalar'] ``` -### 6.2 Element count extraction +### Element count extraction Integer scalars from `Concrete Inputs` (empty strings and non-integer values like floats or booleans are skipped): `[268435456]` Pointwise kernel → last integer scalar is `xnumel`: + - `xnumel = 268,435,456` (= 8 x 4,096 x 8,192) - `rnumel = 1` -### 6.3 Bytes calculation +### Bytes calculation -**Step 1:** Calculate bytes for each tensor input: +Step 1: Calculate bytes for each tensor input: | Input | Dims | Type | Bytes per elem | Elements | Bytes | |-------|------|------|---------------|----------|-------| @@ -397,17 +379,17 @@ Pointwise kernel → last integer scalar is `xnumel`: input_bytes = (268,435,456 * 2) + (268,435,456 * 2) = 1,073,741,824 ``` -**Step 2:** Add output write for pointwise kernels. +Step 2: Add output write for pointwise kernels. In pointwise kernels, `in_out_ptr0` is both read and written — the output -overwrites the input buffer. The input bytes above already count the read. +overwrites the input buffer. The input bytes above already count the read. V2 adds one extra write of `ptr_bytes[0] * xnumel`: ``` output_write = 2 * 268,435,456 = 536,870,912 ``` -**Total:** +Total: ``` bytes_moved = 1,073,741,824 + 536,870,912 @@ -415,7 +397,7 @@ bytes_moved = 1,073,741,824 + 536,870,912 = 1,536 MB ``` -### 6.4 FLOPs calculation +### FLOPs calculation Fused ops parsed from the kernel name: `aten._unsafe_view`, `aten.mul`, `aten.silu` @@ -424,7 +406,7 @@ Fused ops parsed from the kernel name: `aten._unsafe_view`, `aten.mul`, `aten.si | `aten._unsafe_view` | 0 (memory-only) | | `aten.mul` | 1 | | `aten.silu` | 4 | -| **Total** | **5** | +| Total | 5 | ``` flops = 5 * xnumel * rnumel @@ -433,16 +415,16 @@ flops = 5 * xnumel * rnumel GFLOPS = 1,342,177,280 / 1e9 = 1.342 ``` -### 6.5 Throughput metrics +### Throughput metrics -GPU kernel time from the Chrome trace: **505.3 us** +GPU kernel time from the Chrome trace: 505.3 us ``` TB/s = (1,610,612,736 / 1e12) / (505.3 / 1e6) = 3.19 TB/s TFLOPS/s = (1.342 / 1e3) / (505.3 / 1e6) = 2.66 TFLOPS/s ``` -### 6.6 Summary +### Summary | Metric | Value | |--------|-------| @@ -452,14 +434,12 @@ TFLOPS/s = (1.342 / 1e3) / (505.3 / 1e6) = 2.66 TFLOPS/s | TB/s | 3.19 | | TFLOPS/s | 2.66 | ---- - -## 7. Worked Example: Reduction Kernel (RMSNorm) +## Worked example: reduction kernel (RMSNorm) Kernel: `triton_red_fused_add_mean_mul_pow_rsqrt_0` Operation: RMSNorm — `x * rsqrt(mean(x^2) + eps) * weight` -### 7.1 Trace event args (V2 input) +### Trace event args (V2 input) ``` Concrete Inputs: ['', '', '', '32768', '2048'] @@ -467,15 +447,16 @@ Input Dims: [[8, 4096, 2048], [2048], [8, 4096, 2048], [], []] Input type: ['c10::BFloat16', 'c10::BFloat16', 'c10::BFloat16', 'Scalar', 'Scalar'] ``` -### 7.2 Element count extraction +### Element count extraction Integer scalars from `Concrete Inputs` (non-integer values skipped): `[32768, 2048]` Reduction kernel → last two integer scalars are `xnumel` and `rnumel`: + - `xnumel = 32,768` (= 8 x 4,096 = batch x seq_len) - `rnumel = 2,048` (= hidden_dim, the reduction axis) -### 7.3 Bytes calculation (V2) +### Bytes calculation (V2) | Input | Dims | Type | Bytes per elem | Elements | Bytes | |-------|------|------|---------------|----------|-------| @@ -492,7 +473,7 @@ bytes_moved = 134,217,728 + 4,096 + 134,217,728 ≈ 256 MB ``` -### 7.4 Bytes calculation (V1 — for comparison) +### Bytes calculation (V1 — for comparison) V1 uses `xnumel`, `rnumel`, and per-pointer byte widths from the signature instead of per-tensor shapes: @@ -514,11 +495,13 @@ bytes = sum(ptr_bytes[:-1]) * xnumel * rnumel + ptr_bytes[-1] * xnumel ≈ 256 MB ``` -Note: V1 and V2 give slightly different values because V1 assumes all input +```{note} +V1 and V2 give slightly different values because V1 assumes all input pointers access the full `xnumel * rnumel` grid, while V2 uses exact tensor shapes (the weight tensor `in_ptr1` is only [2048], not [32768, 2048]). +``` -### 7.5 FLOPs calculation +### FLOPs calculation Fused ops: `aten.add`, `aten.mean`, `aten.mul`, `aten.pow`, `aten.rsqrt` @@ -529,28 +512,26 @@ Fused ops: `aten.add`, `aten.mean`, `aten.mul`, `aten.pow`, `aten.rsqrt` | `aten.mul` | 1 | | `aten.pow` | 2 | | `aten.rsqrt` | 2 | -| **Total** | **7** | +| Total | 7 | ``` flops = 7 * 32,768 * 2,048 = 469,762,048 GFLOPS = 0.470 ``` -### 7.6 Throughput metrics +### Throughput metrics -GPU kernel time: **79.4 us** +GPU kernel time: 79.4 us ``` TB/s = (268,439,552 / 1e12) / (79.4 / 1e6) = 3.38 TB/s TFLOPS/s = (0.470 / 1e3) / (79.4 / 1e6) = 5.92 TFLOPS/s ``` ---- - -## 8. Results for the Example TransformerBlock +## Results for the example transformer block Model: `TransformerBlock(dim=2048, n_heads=16, mlp_ratio=4)`, batch=8, -seq_len=4096, dtype=bf16. Traced with PyTorch 2.11+rocm7.2 on MI300X. +seq_len=4096, dtype=bf16. Traced with PyTorch 2.11+rocm7.2 on MI300X. | Kernel | GFLOPS | Data Moved (MB) | FLOPS/Byte | TB/s | TFLOPS/s | |--------|--------|-----------------|------------|------|----------| @@ -559,8 +540,6 @@ seq_len=4096, dtype=bf16. Traced with PyTorch 2.11+rocm7.2 on MI300X. | `triton_poi_fused__unsafe_view_mul_silu_2` (SwiGLU) | 1.342 | 1,536 | 0.83 | 3.19 | 2.66 | | `triton_poi_fused__unsafe_view_add_3` (residual add) | 0.067 | 512 | 0.12 | 3.81 | 0.48 | ---- - ## References - [PyTorch Blog: Why Is PyTorch Compile So Fast — Kernel Fusion](https://pytorch.org/blog/why-is-pytorch-compile-so-fast-kernel-fusion/) @@ -571,3 +550,11 @@ seq_len=4096, dtype=bf16. Traced with PyTorch 2.11+rocm7.2 on MI300X. - [PyTorch Inductor config.py](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/config.py) - [PyTorch Inductor codegen/triton.py](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/codegen/triton.py) - [PyTorch Inductor wrapper_benchmark.py](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/wrapper_benchmark.py) + +## Related topics + +- [GEMM analysis](../conceptual/gemm-analysis.md) +- [Trace2Tree](../conceptual/trace2tree.md) +- [Generate a performance report from a PyTorch trace](../how-to/generate-perf-report-pytorch.md) +- [Performance report columns](../reference/perf-report-columns.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/conf.py b/docs/conf.py index be3247276..94041d68f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -27,6 +27,11 @@ external_toc_path = "./sphinx/_toc.yml" +# rocm-docs-core (>=1.x) renders Jupyter notebooks through myst-nb. The tutorial +# notebooks ship with their captured outputs and need traces or GPUs to run, so +# don't re-execute them at build time. +nb_execution_mode = "off" + html_theme = "rocm_docs_theme" # fmt: off html_theme_options = { diff --git a/docs/how-to/agent.md b/docs/how-to/agent.md new file mode 100644 index 000000000..a00d77046 --- /dev/null +++ b/docs/how-to/agent.md @@ -0,0 +1,233 @@ + + + +# Generate optimization recommendations with the TraceLens Agent +```{meta} +:description: Learn how to run the TraceLens Agent, an agentic workflow that turns a GPU trace into a prioritized, stakeholder-facing performance optimization report. +:keywords: TraceLens, TraceLens Agent, agentic analysis, performance optimization, roofline, kernel fusion, GEMM, ROCm, AMD Instinct, optimization report +``` + +This topic shows how to run the *TraceLens Agent*, an agentic performance-analysis +workflow that reads a GPU trace and produces a single stakeholder-facing report, +`analysis.md`, organized as a prioritized bottleneck list. Findings are ranked and +grouped into three tiers: compute kernel optimizations, kernel fusion +opportunities, and system-level optimizations. Each finding carries the supporting +evidence, the reasoning behind the call-out, and a concrete resolution. + +The agent combines a structured, skill-driven workflow with codified TraceLens +analysis for repeatable, reliable results. + +## Analysis modes + +The agent runs in one of two modes: + +- *Standalone*: Single-trace roofline analysis. Use this when you have one trace + and want to find where performance falls short of hardware limits, find system bottlenecks or fusion opportunities. This is the recommended default. +- *Comparative*: Two-trace gap analysis. The agent compares a primary trace + against a reference trace, for example a different platform or a tuned config, + and identifies inefficiencies in the primary trace relative to the reference. + Comparative analysis works best when both traces come from the same framework. + Cross-framework comparisons can produce misleading gap estimates because of + structural differences in operation call stacks. + +Support depends on the execution mode of the traced workload: + +| Execution mode | Standalone | Comparative | +|---|---|---| +| Eager | Supported | Supported | +| Graph + capture | Supported | Not supported | +| Graph | Not supported | Not supported | + +## Before you begin + +### Install TraceLens + +Install locally, or into your container or virtual environment (see +[Install TraceLens](../install/install.md)): + +```bash +pip install git+https://github.com/AMD-AGI/TraceLens.git +``` + +### Collect a trace + +The orchestrator runs against a single `torch.profiler` trace (`.json` or +`.json.gz`). Collection is workload-specific: + +- *Generic Eager Traces*: Instrument your loop with + `torch.profiler.profile(...)`, enabling CPU-side call-stack and shape capture + (`with_stack=True`, `record_shapes=True`). Profile a representative steady-state + window of a handful of post-warmup steps, then log the trace with + `prof.export_chrome_trace(...)`. A single rank's trace is enough for per-rank + analysis. +- *Inference Traces with Graph Capture*: Collection has framework-specific + requirements. Follow + [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md). + The Profiling Skill automates + vLLM/SGLang benchmarking and PyTorch profiler trace collection via + Magpie, producing analysis-ready traces. For + graph-mode workloads you produce two artifacts: a graph-replay trace and a + graph-capture folder. In inference mode with execution mode + `graph replay + capture`, TraceLens merges call-stack and shape information from + the capture folder into the replay tree before analysis. + +### Establish a hardware baseline + +Roofline analysis compares each measured kernel against your GPU's max-achievable +TFLOPS and HBM bandwidth, so it needs a `.json` arch file for your +hardware. Bundled arch files ship with the package. If your platform isn't +included, or you want stack-specific measured values instead of published specs, +generate benchmark-derived peak TFLOPS and HBM bandwidth with the GPU +microbenchmarking suite. It writes the arch JSON in the shape the roofline +expects. + +## Run the agent from a chat + +```{note} +The examples use the Cursor IDE and CLI, but the orchestrator skills are portable +and also work with other agentic runners that support skill-file discovery. +``` + +In a chat with a capable model, invoke one of: + +- Standalone (single trace): + + ```text + Follow the analysis orchestrator installed with TraceLens and run the full + agentic analysis workflow on + ``` + +- Comparative (two traces) + + ```text + Follow the analysis orchestrator installed with TraceLens and run the full + agentic analysis workflow on and + ``` + +If prompted, provide the trace file path, the platform of the first trace, the +analysis mode (`default` for training and non-vLLM/SGLang eager inference, or +`inference` for vLLM/SGLang), the execution mode and capture-folder path for +inference, environment details (node, container, or virtual environment), and an +optional output directory. + +## Run the agent headless (CLI) + +Use the `agent` CLI to run the orchestrator non-interactively. Install it with: + +```bash +curl https://cursor.com/install -fsS | bash +``` + +Pass every parameter inline so no interactive prompts are needed. This is useful +for batch runs and continuous-integration pipelines. For example, for a default +standalone run on a remote node with a container: + +```bash +agent --model --print --force --trust \ + "Follow the analysis orchestrator installed with the TraceLens pip package + (look under TraceLens/Agent/Analysis/.cursor/skills/ in the package + installation directory) and run the full agentic analysis workflow on + with platform , analysis mode default, + node , container , output to " +``` + +For vLLM or SGLang inference, set `analysis mode inference` and add +`execution mode eager`, or `execution mode graph replay + capture` together with +`capture folder `. + +## Read the results + +Only `analysis.md` is intended for end-user review. Everything else under +`analysis_output/` is agent internal: intermediates the orchestrator and +sub-agents pass between steps. + +Every report has the same top-level structure: + +1. *Executive summary*: A one-paragraph workload characterization, a metrics + table (total time, compute percentage, idle percentage, exposed-communication + percentage, and top bottleneck category), and a representative chart. +2. *Compute kernel optimizations*: Top-operations table followed by per-category + priority items (`P1`, `P2`, and so on) sorted by `impact_score`. Each item + carries an Insight, Action, and Impact triplet. +3. *Kernel fusion opportunities* (experimental): Candidate modules to fuse. +4. *System-level optimizations* (experimental): Idle time, memory-copy overhead, + and compute/communication overlap. +5. *Detailed analysis*: Per-priority-item drill-down with identification + rationale, data tables, and reasoning. + +### Programmatic interface + +Every report embeds HTML comment markers so a downstream system can consume it +without parsing prose: + +- `` … + `` wraps each priority item's Impact line; `mid` is the + canonical `impact_score` and `` is the analyzer category (`gemm`, + `sdpa_fwd`, `elementwise`, `norm_fwd`, and so on). +- `` … `` wraps the Top + Operations table at the start of the Compute kernel optimizations section. +- Per-item anchors (`#detailed-analysis-compute-pN`, `#detailed-analysis-fusion-PN`, + `#detailed-analysis-system-pN`) link each summary card to its detailed reasoning + section. + +## Analysis workflow + +The workflow splits into three independently composable tiers: + +- *System-level optimizations*: Issues that affect the GPU pipeline as a whole, + such as idle time, memory-copy overhead, collective-communication blocking, and + compute/communication overlap. +- *Kernel fusion opportunities* (experimental): Multi-kernel modules that could be + fused, with estimated savings. +- *Compute kernel optimizations*: Per-category kernel analysis (GEMM, attention, + elementwise, etc.) focused on individual operation efficiency. + +The analysis orchestrator skill coordinates the workflow: it queries user +inputs, runs TraceLens to pre-compute trace data, invokes the system-level and +compute-kernel sub-agents in parallel, aggregates and validates their findings, +identifies the model, and generates the prioritized report. The orchestrator and +its sub-agents run on a capable reasoning model declared in each agent's front +matter. + +The high-level steps are: + +1. Query user inputs (comparison scope, trace paths, platforms, analysis mode, + environment setup). +2. Generate the performance report (branching on analysis mode and comparison + scope). +3. Prepare category data (GPU utilization, top operations, tree data, multi-kernel + data, category filtering) and extract fusion candidates. +4. Run system-level analysis (CPU/idle, multi-kernel, and kernel fusion) in + parallel. +5. Run the compute-kernel sub-agents in parallel. +6. Aggregate the per-category findings into a globally ranked list. +7. Validate sub-agent outputs for time sanity, efficiency anomalies, and coverage. +8. Prepare report data and identify the model. +9. Generate the final `analysis.md` report. + +### Sub-agents + +System-level sub-agents cover CPU and idle analysis, memory-copy and +collective-communication patterns, and kernel fusion. Compute-kernel sub-agents +cover GEMM, scaled dot-product attention, elementwise, +reduction, Triton-compiled kernels, Mixture-of-Experts, normalization, +convolution, and a generic analyzer for uncategorized operations. + +### Execution environments + +During the input step the orchestrator asks whether you're running locally or on a +cluster, and builds the appropriate command prefixes automatically — direct +execution, a virtual-environment activation prefix, an SSH wrapper for a remote +node, or an SSH plus container-exec wrapper for a containerized node. + +## Related topics + +- [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md) +- [Analyze traces with the TraceLens SDK](./sdk-analysis.md) +- [Trace2Tree data model](../conceptual/trace2tree.md) +- [GEMM analysis](../conceptual/gemm-analysis.md) \ No newline at end of file diff --git a/docs/how-to/collective-report.md b/docs/how-to/collective-report.md index b231d51b5..249f166f3 100644 --- a/docs/how-to/collective-report.md +++ b/docs/how-to/collective-report.md @@ -59,6 +59,11 @@ neither `--output_xlsx_path` nor `--output_csvs_dir` is given, the workbook is written to `--trace_dir` (or the lowest common directory of the resolved traces). +```{note} +Excel output requires `openpyxl`; TraceLens installs it automatically if it's +missing. Use `--output_csvs_dir` to write CSV files instead. +``` + ## Step 2: Read the report sheets Each collective is assigned a `collective_id` (its Process Group name plus an @@ -107,6 +112,7 @@ The all2allv sheets instead report: | `throughput (GB/s)` | Total data moved / wall-clock time. Compare against link bandwidth to gauge efficiency. | | `wall_time` | End-to-end time from the first rank entering to the last finishing. | | `size_imbalance` | Max rank's data ÷ mean. `1.0` = balanced; `>> 1.0` = some ranks carry disproportionate load (common when some MoE experts are hotter than others). | +| `max_rank_dur / min_rank_dur` | Spread in per-rank kernel time. A large spread together with high `size_imbalance` points to rebalancing expert routing or dropping tokens. | | `skew in start time` | How far apart ranks enter the collective — large skew means some ranks are blocked by upstream compute. | To diagnose low throughput, check `size_imbalance`: @@ -130,6 +136,19 @@ the others to wait in implicit synchronization. Open the `straggler_summary` sheet: it's sorted so the *straggler is the first row* (lowest total wait time — it arrives last, so it rarely waits itself). +Example from an 8-rank Llama 70B FSDP run (`straggler_summary`, truncated): + +| rank | total_wait_time_us | mean_wait_time_us | times_arrived_last | times_arrived_first | pct_arrived_last | num_collectives | total_nccl_dur_us | +|------|-------------------|-------------------|--------------------|---------------------|------------------|-----------------|-------------------| +| 4 | 27,195 | 55.7 | 420 | 6 | 86.1% | 488 | 3,910,753 | +| 5 | 4,660,353 | 9,549.9 | 39 | 10 | 8.0% | 488 | 8,463,896 | +| ... | ... | ... | ... | ... | ... | ... | ... | +| 0 | 13,358,753 | 27,374.5 | 3 | 320 | 0.6% | 488 | 17,203,432 | + +Rank 4 is the straggler: the lowest total wait time, last to arrive 86% of the +time, and the lowest NCCL kernel duration (other ranks' durations are inflated +by the time they spend waiting for it). + | Column | What it tells you | |--------|-------------------| | `total_wait_time_us` | Sum of this rank's wait across all collectives. The straggler has the *lowest* value. | diff --git a/docs/how-to/compare-perf-reports.md b/docs/how-to/compare-perf-reports.md new file mode 100644 index 000000000..db2885432 --- /dev/null +++ b/docs/how-to/compare-perf-reports.md @@ -0,0 +1,118 @@ + + +# Compare performance reports in TraceLens +```{meta} +:description: Learn how to compare two or more TraceLens performance reports with the compare_perf_reports_pytorch tool to diff per-op, kernel, and roofline metrics in a single workbook. +:keywords: TraceLens, compare performance reports, compare_perf_reports_pytorch, xlsx, diff, ops_summary, kernel_summary, roofline, ROCm, PyTorch profiler +``` + +This topic shows how to compare multiple performance reports that you've already +generated with `generate_perf_report_pytorch.py`, using the +`compare_perf_reports_pytorch.py` script +(`TraceLens/Reporting/compare_perf_reports_pytorch.py`). The result is a single +workbook with side-by-side metrics and per-metric diffs. + +```{note} +This topic covers comparing generated *reports* (the `.xlsx` or `.csv` outputs). +To compare raw traces by their morphological tree structure, see +[Compare two traces in TraceLens](./compare-traces.md), which also documents the +SDK-based `TraceDiff` workflow. +``` + +## Before you begin + +- TraceLens installed (see [Install TraceLens](../install/install.md)). +- At least two reports (`.xlsx`) generated with + [TraceLens_generate_perf_report_pytorch](./generate-perf-report-pytorch.md). + +## What it takes in + +| Input | Description | +|-------|-------------| +| `*.xlsx` files | TraceLens reports you want to compare. Provide at least two. | +| Optional `--names` | Human-readable tags for each report. If omitted, the script falls back to the base filenames. | + +## How to call it + +```bash +python compare_perf_reports.py \ + baseline.xlsx \ + candidate.xlsx \ + --names baseline candidate \ + --sheets all \ + -o comparison.xlsx +``` + +Alternatively, you can call the entry point directly after installing TraceLens, +which avoids needing to know the script's path: + +```bash +TraceLens_compare_perf_reports_pytorch \ + baseline.xlsx \ + candidate.xlsx \ + --names baseline candidate \ + --sheets all \ + -o comparison.xlsx +``` + +Common flags: + +| Flag | Default | Purpose | +|------|---------|---------| +| `-o, --output` | `comparison.xlsx` | Name of the merged workbook. | +| `--names` | `` | Custom tags (must match the number of reports). | +| `--sheets` | `all` | Limit processing to a subset: `gpu_timeline`, `ops_summary`, `kernel_summary`, `ops_all`, `roofline`, or `all`. | + +## What comes out + +The script writes one workbook (`comparison.xlsx` unless you override it) +containing multiple sheets: + +| Sheet | When you get it | What it shows | +|-------|-----------------|---------------| +| `gpu_timeline` | `--sheets gpu_timeline` or `all` | End-to-end GPU activity by type (`compute`, `memcpy`, and so on) with per-report timings plus `time ms___diff` and `time ms___pct`. | +| `ops_summary` | `--sheets ops_summary` or `all` | Per-op aggregates keyed on `name`. Sorted by the baseline's `total_direct_kernel_time_ms`. Unhelpful columns (for example cumulative %) are stripped from non-baseline reports. When comparing PyTorch traces, this sheet is generated. For rocprof traces without `ops_summary`, it automatically falls back to `kernel_summary`. | +| `kernel_summary` | `--sheets kernel_summary` or `all` | Kernel-level summary for rocprof reports, keyed on `name`. Similar to `ops_summary` but uses rocprof-specific columns like `Total Kernel Time (ms)` and `Count`. When `--sheets all` is used with rocprof reports, this is generated automatically if `ops_summary` is not available. | +| `ops_all_*` | `--sheets ops_all` or `all` | Three sheets per variant tag: `ops_all_intersect_` (op instances present in both baseline and variant), `ops_all_only_baseline_` (ops only the baseline ran), and `ops_all_only_variant_` (ops only the variant ran). Columns irrelevant to a given view are hidden, not deleted. | +| `_*` | `--sheets roofline` or `all` | Same intersect / only-baseline / only-variant breakdown for each roofline group: `GEMM`, `SDPA_fwd`, `SDPA_bwd`, `CONV_fwd`, `CONV_bwd`, `UnaryElementwise`, `BinaryElementwise`. | + +Hidden columns stay in the file (for power users) but are invisible in Excel by +default. + +## Diff math + +For every metric you ask it to track (`diff_cols` in the code), the script +computes: + +```text +metric___diff # variant - baseline +metric___pct # 100 * diff / baseline +``` + +## Design decisions to know + +- *Outer merge, never inner* — if an op vanished, you'll see it. +- *Baseline is the first report* — choose its order deliberately. +- *Column prefixing* — every metric becomes `::metric`, so you can safely concatenate arbitrary reports. +- *Sheet-specific pruning* — the script aggressively hides noise (for example `median`, `UID`) to keep the output readable. You can always unhide these in Excel if you need them. +- *Excel 31-char rule* — sheet names are truncated to fit; no data loss, just shorter labels. + +```{note} +A planned enhancement is morphology-aware diffing: understanding the call-stack +tree and comparing at the lowest common call-stack level. For example, if a +baseline leaf op is `cudnn_convolution` and the variant is `miopen_convolution`, +the diff would recognize that the lowest common level is `convolution` and +compare the two there. +``` + +## Related topics + +- [Compare two traces in TraceLens](./compare-traces.md) +- [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Generate a rocprof performance report](./generate-perf-report-rocprof.md) +- [Generate a report](./generate-reports.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/how-to/compare-traces.md b/docs/how-to/compare-traces.md index 63233e5a3..e2030806a 100644 --- a/docs/how-to/compare-traces.md +++ b/docs/how-to/compare-traces.md @@ -26,13 +26,13 @@ the matching to be: compares traces by their *morphological tree structure*, matching ops at the lowest common node. Use it when op names differ between traces (for example, across hardware, libraries, or framework versions) or when you need finer-grained, - programmatic analysis. + programmatic analysis. For an end-to-end analysis, the **[Comparative TraceLens Agent](#tracelens-agent-comparative-mode)** wraps TraceDiff in an agentic workflow that automates the comparison and provides optimization recommendations. ## Before you begin - TraceLens installed (see [Install TraceLens](../install/install.md)). -- Two generated reports (`.xlsx`), one per configuration. Generate them with - [TraceLens_generate_perf_report_pytorch](./generate-perf-report-pytorch.md). +- For perf-report comparison: two generated reports (`.xlsx`), one + per configuration. Generate them with [TraceLens_generate_perf_report_pytorch](./generate-perf-report-pytorch.md). ## Perf-report comparison @@ -186,9 +186,17 @@ See the [`trace_diff_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/trace_diff_example.ipynb) notebook for a worked example. +## TraceLens Agent (comparative mode) + +The Comparative TraceLens Agent is an agentic workflow that uses TraceDiff to analyze traces +and outputs a prioritized, stakeholder-facing optimization report. It's best used when you want +automated gap analysis with concrete optimization recommendations rather than +raw data tables. For more information, see [TraceLens Agent](./agent.md). + ## Related topics - [What is TraceLens?](../what-is-tracelens.md) - [Install TraceLens](../install/install.md) - [Generate a report](generate-reports.md) +- [Generate optimization recommendations with the TraceLens Agent](./agent.md) - [API reference](../reference/api-reference.md) diff --git a/docs/how-to/generate-perf-report-pytorch-inference.md b/docs/how-to/generate-perf-report-pytorch-inference.md index b838fd5b7..5c848932b 100644 --- a/docs/how-to/generate-perf-report-pytorch-inference.md +++ b/docs/how-to/generate-perf-report-pytorch-inference.md @@ -7,8 +7,8 @@ See LICENSE for license information. # Generate a PyTorch inference performance report ```{meta} -:description: Learn how to generate a TraceLens performance report from a PyTorch LLM-serving (vLLM/SGLang) trace captured in graph mode, including capture-trace merging. -:keywords: TraceLens, PyTorch profiler, inference, vLLM, SGLang, CUDA graph, HIP graph, graph capture, LLM serving, FusedMoE, roofline, ROCm +:description: Learn how to collect, split, and analyze PyTorch LLM-serving (vLLM/SGLang) traces with TraceLens, including graph-capture merging for graph-mode inference. +:keywords: TraceLens, PyTorch profiler, inference, vLLM, SGLang, HIP graph, CUDA graph, graph capture, LLM serving, FusedMoE, roofline, trace splitting, steady state, ROCm ``` `TraceLens_generate_perf_report_pytorch_inference` is the inference-oriented @@ -17,22 +17,280 @@ vLLM or SGLang) that run in CUDA/HIP graph mode, and can merge the graph-capture traces back into the graph-replay trace to recover the call-stack and input-shape metadata that graph execution drops. -For training or single-run PyTorch traces, use +This topic covers the end-to-end inference workflow: collecting traces, splitting +them into steady-state windows, and generating the report. For training or +single-run PyTorch traces, use [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) -instead. +instead. For the underlying concepts (the paged-attention roofline model and the +steady-state region), see +[Inference performance analysis in TraceLens](../conceptual/inference-analysis.md). -## Before you begin +## Supported frameworks and execution modes + +TraceLens inference features are primarily tested with vLLM and SGLang. Feature +coverage depends on how the model is executed: + +| Mode | Shapes / roofline analysis | Agent analysis | Limitations | +|------|----------------------------|----------------|-------------| +| Eager only | Yes | Supported; proposed patches are recommended to include roofline information for attention operations | Eager execution may use different compilation strategies, resulting in different kernels and fusions than graph execution. | +| Graph execution only | Non-graph kernels | Limited | Categorization, call stacks, and shapes are available only for attention kernels if `full_and_piecewise` mode is used. | +| Graph execution + eager-mode trace | Limited | Limited | Kernel categorization may be less accurate than eager or graph+capture. | +| Graph execution + graph capture [^1] | Yes | Yes (patches required) | | -Before generating a report, confirm you have the following: +[^1]: Graph-mode analysis using graph-capture and graph-replay traces is +supported for vLLM and SGLang (proposed patches required). + +## Before you begin - TraceLens installed (see [Install TraceLens](../install/install.md)). -- A graph-replay `torch.profiler` trace (`.json` or `.json.gz`). -- (Optional) The folder of graph-capture traces from the same run, if you want to - recover shapes and call stacks (see [Merge capture traces](#merge-capture-traces)). +- An LLM inference setup to profile (this guide uses vLLM or SGLang on MI300X). + +## Collect inference traces + +There are two ways to collect inference traces for TraceLens analysis. + +### Option A: Profiling agent (recommended) + +The Profiling agent is an agentic skill bundled with TraceLens that drives an LLM +inference benchmark, applies the Docker/framework patches for you, tunes the +profiler window, runs the benchmark, and splits the resulting trace into +analysis-ready windows. Follow +[`TraceLens/Agent/Profiling/README.md`](https://github.com/AMD-AGI/TraceLens/blob/main/TraceLens/Agent/Profiling/README.md). +The agent handles collection and splitting end-to-end; if you use it you can skip +the rest of this section and [Split inference traces](#split-inference-traces-optional). + +### Option B: Manual profiling + +Use manual profiling for full control over every step. Build (or patch) an image, +configure the profiler yourself, and run the benchmark by hand. + +#### Build a Docker image + +Build scripts for vLLM and SGLang are provided under +[`examples/custom_workflows/inference_analysis/`](https://github.com/AMD-AGI/TraceLens/tree/main/examples/custom_workflows/inference_analysis). + +**vLLM.** A unified build script supports multiple vLLM versions. It takes a +version tag as its first argument, followed by the path to your local TraceLens +clone and any standard `docker build` flags. The script selects the correct base +image and patch file automatically: + +| Version | Base image | vLLM version | Patch file | +|---------|-----------|--------------|------------| +| `v14` | `rocm/vllm-dev:preview_releases_rocm_v0.14.0_20260120` | v0.14.0 | `config_vllm_v0.14.0.patch` | +| `v15` | `rocm/vllm-dev:preview_releases_rocm_v0.15.0_20260130` | v0.15.0 | `config_vllm_v0.15.0.patch` | +| `v16` | `rocm/vllm-dev:preview_rocm70_releases_rocm_v0.16.0_20260223` | v0.16.0 | `config_vllm_v0.16.0.patch` | +| `v17` | `vllm/vllm-openai-rocm:v0.17.0` | v0.17.0 | `config_vllm_v0.17.0.patch` | +| `v18` | `vllm/vllm-openai-rocm:v0.18.0` | v0.18.0 | `config_vllm_v0.18.0.patch` | +| `v19` | `vllm/vllm-openai-rocm:v0.19.0` | v0.19.0 | `config_vllm_v0.19.0.patch` | +| `v20` | `rocm/vllm-dev:preview_v0.20.0_20260429` | v0.20.0 | `config_vllm_v0.20.0.patch` | +| `v21` | `vllm/vllm-openai-rocm:v0.21.0` | v0.21.0 | `config_vllm_v0.21.0.patch` | +| `v22` | `vllm/vllm-openai-rocm:v0.22.0` | v0.22.0 | `config_vllm_v0.22.0.patch` | +| `v23` | `vllm/vllm-openai-rocm:v0.23.0` | v0.23.0 | `config_vllm_v0.23.0.patch` | +| `v24` | `vllm/vllm-openai-rocm:v0.24.0` | v0.24.0 | `config_vllm_v0.24.0.patch` | + +```bash +bash examples/custom_workflows/inference_analysis/build_docker_vllm.sh \ + v16 \ + /path/to/TraceLens \ + -t tracelens-vllm +``` + +To use a custom base image instead of the default for the selected version, pass +`--base-image`: + +```bash +bash examples/custom_workflows/inference_analysis/build_docker_vllm.sh \ + v18 \ + /path/to/TraceLens \ + --base-image my-registry/vllm:nightly \ + -t tracelens-vllm:custom +``` + +**SGLang.** The SGLang build script takes the path to your local TraceLens clone, +the SGLang version (`--sglang-version`, default `0.5.9`), and the GPU type +(`--gpu-type`, default `mi350`). MI300 and MI350/MI355 are supported: + +| SGLang version | GPU type | Base image | +|----------------|----------|------------| +| `0.5.9` | MI300 | `lmsysorg/sglang:v0.5.9-rocm700-mi30x` | +| `0.5.9` | MI350/MI355 | `lmsysorg/sglang:v0.5.9-rocm700-mi35x` | +| `0.5.11` | MI300 | `lmsysorg/sglang:v0.5.11-rocm720-mi30x` | +| `0.5.11` | MI350/MI355 | `lmsysorg/sglang:v0.5.11-rocm720-mi35x` | +| `0.5.12` | MI300 | `lmsysorg/sglang:v0.5.12-rocm720-mi30x` | +| `0.5.12` | MI350/MI355 | `lmsysorg/sglang:v0.5.12-rocm720-mi35x` | + +```bash +bash examples/custom_workflows/inference_analysis/build_docker_sglang.sh \ + /path/to/TraceLens \ + --sglang-version 0.5.11 \ + --gpu-type mi350 \ + -t tracelens-sglang +``` + +Create a container from the resulting image. + +#### Apply framework patches manually + +If you prefer to patch an existing environment instead of building a new image, +apply patches to your inference framework to: + +- Add custom annotations with request-packing information (these annotations feed + the [roofline model](../conceptual/inference-analysis.md#roofline-analysis-for-inference-attention)). +- Capture graph-mode execution phases for augmentation by TraceLens. + +To apply them: + +1. Locate your inference engine: + + ```bash + # vLLM + python -c "import vllm; import os; print(os.path.dirname(vllm.__file__))" + # SGLang + python -c "import sglang; import os; print(os.path.dirname(sglang.__file__))" + ``` + +2. Select the patch for your framework and version and apply it: + + ```bash + cd /path/to/framework/../ && git apply /path/to/patchfile + ``` + + vLLM patches are in + [`vllm_patches`](https://github.com/AMD-AGI/TraceLens/tree/main/examples/custom_workflows/inference_analysis/vllm_patches); + SGLang patches are in + [`sglang_roofline_patches`](https://github.com/AMD-AGI/TraceLens/tree/main/examples/custom_workflows/inference_analysis/sglang_roofline_patches). + +#### Collection parameters + +- *Steady-state window*: Inference traces are large. Most serving benchmarks use + `NUM_PROMPTS = 10 * CONC` with OSL sampling ratio R. Trace + `((R+1)/2) * 5 * OSL ± (16 * OSL / CONC)` execution steps to capture peak + concurrency with a prefill-decode mix. See + [Steady-state region](../conceptual/inference-analysis.md#steady-state-region) + for the derivation. You may need to raise the trace-store timeout in some + frameworks to allow writing the trace mid-execution (for example + `VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200` for vLLM). +- *Graph-capture mode*: The recommended patch traces the graph-capture phase and + stores the corresponding trace files. +- *Profiler setup*: Enable CPU-side call-stack and shape capture (for example, + vLLM's `profiler-config.torch_profiler_record_shapes` and + `profiler-config.torch_profiler_with_stack`). + +#### Trace collection flags + +**vLLM.** The `config_vllm_v*.patch` patches (available for v0.14-v0.24) add two +`ProfilerConfig` flags. Pass them as server arguments: + +| Flag | Type | Default | Description | +|------|------|---------|-------------| +| `--profiler-config.capture_torch_profiler_dir DIR` | `str` | `""` | Directory where a profiler trace of the HIP/CUDA-graph capture phase is saved (rank 0 only). Requires `--profiler-config.profiler torch`. Leave empty to disable graph-capture profiling. | +| `--profiler-config.detailed_trace_annotation` | `bool` | `False` | When `True`, execution-step annotations include roofline metrics (`sk`, `sqsq`, `sqsk`) for context and generation requests. When `False`, annotations record only request and token counts. Enable for full roofline analysis. | + +Example - enable both flags alongside a steady-state window profile: + +```bash +--profiler-config.profiler torch \ +--profiler-config.torch_profiler_dir /workspace/torch_trace \ +--profiler-config.capture_torch_profiler_dir /workspace/torch_trace/capture_traces \ +--profiler-config.detailed_trace_annotation True \ +--profiler-config.delay_iterations 5402 \ +--profiler-config.max_iterations 256 \ +--profiler-config.ignore_frontend True +``` + +```{note} +`capture_torch_profiler_dir` is only available when `--profiler-config.profiler torch` +is set. The capture trace is written once at server startup during HIP/CUDA-graph +construction; the steady-state replay trace is written to `torch_profiler_dir` +during the benchmark. Pass both paths to the report generator via +`--capture_folder` and `--profile_json_path` respectively. +``` + +**SGLang.** In the profile request for the execution step: + +1. Pass `shape_discovery=True` to enable shape discovery and registration for + operations not covered by the default SGLang profile. +2. Pass `roofline_annotations=True` to annotate the trace with the detailed + information used for roofline analysis. +3. To profile the graph-capture phase, pass the `--enable-profile-cuda-graph` + server argument at startup. This saves one trace file per batch size but misses + shape information for some operations; add + `--enable-shape-discovery-for-cuda-graph-profile` for more diverse coverage. + +## Split inference traces (optional) + +Large traces can be split into steady-state windows or per-step files before +report generation, using `TraceLens.TraceUtils.split_inference_trace_annotation`. +Splitting assumes vLLM v0.14 or higher (tested through v0.24) or SGLang v0.5.9, +or use of the provided patches, so that annotations (batch size, request counts, +and so on) are present in the execution-step metadata. + +**Find the steady-state region** (highest concurrency) and separate +prefill-decode from decode-only steps. This is the recommended option for large +traces where you want a few representative steps extracted automatically: + +```bash +python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz \ + -o ./steady_state_analysis \ + --find-steady-state --num-steps 256 +``` + +The output is a trace of `--num-steps` contiguous steps near peak concurrency, +plus contiguous prefill-decode-mix and decode-only steady-state traces extracted +from that window with no idle gaps between steps. + +By default the mixed window is selected by matching the empirically observed +prefill-decode-to-total-steps ratio. If the benchmark parameters are known, +passing `--CONC`, `--OSL`, and `--R` computes the *ideal* ratio analytically and +uses that to drive window selection instead: + +| Argument | Type | Description | +|----------|------|-------------| +| `--CONC` | `int` | Expected peak concurrency (number of concurrent requests). A warning is printed if the observed trace peak differs. | +| `--OSL` | `float` | Maximum output sequence length (decode tokens per request). Each request's OSL is sampled from `[R * OSL, OSL]`. | +| `--R` | `float` | OSL sampling range ratio in `[0, 1]`. `R=0` means all requests use exactly `OSL` tokens; `R=1` means OSL is uniform in `[0, OSL]`. | + +When all three are provided, the tool derives: + +``` +ideal_prefilldecodemix_to_totalsteps_ratio = (CONC * 2) / (OSL * (1 + R)) +``` + +and uses it as the reference ratio, overriding the empirical estimate. +`--num-steps` is automatically raised to +`ceil(1 / ideal_prefilldecodemix_to_totalsteps_ratio)` if it is too small to +capture a representative mix: + +```bash +python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz \ + -o ./steady_state_analysis \ + --find-steady-state --num-steps 256 \ + --CONC 32 --OSL 1024 --R 0.8 +``` + +**One trace file per execution step** (vLLM v0.13+, SGLang v0.5.9, Atom 0.1.1) - +use when you want to analyze an isolated execution step: + +```bash +python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz \ + -o ./output --store-single-iteration +``` + +**Limit the steady-state search to a window** with `--iterations`: + +```bash +python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz \ + -o ./output --iterations 10:20 --find-steady-state --num-steps 256 \ + --CONC 32 --OSL 1024 --R 0.8 +``` + +See [Steady-state region](../conceptual/inference-analysis.md#steady-state-region) +for what these phases mean and why this window is the one worth profiling. ## Generate the report -Pass the graph-replay trace to generate the default Excel report: +Report generation is supported for both eager-mode and graph-mode (capture + +replay) traces. Pass the graph-replay trace to generate the default Excel report: ```bash TraceLens_generate_perf_report_pytorch_inference \ @@ -87,7 +345,8 @@ Run the tool with `--help` for the complete, version-specific argument list. ## Related topics +- [Inference performance analysis in TraceLens](../conceptual/inference-analysis.md) - [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) -- [Generate TraceLens reports](./generate-reports.md) - [Compare two traces](./compare-traces.md) +- [Agentically orchestrate and generate optimization recommendations](./agent.md) - [API reference](../reference/api-reference.md) diff --git a/docs/how-to/generate-perf-report-pytorch.md b/docs/how-to/generate-perf-report-pytorch.md index d57507813..c5888af32 100644 --- a/docs/how-to/generate-perf-report-pytorch.md +++ b/docs/how-to/generate-perf-report-pytorch.md @@ -23,7 +23,7 @@ Before generating a report, confirm you have the following: ```{note} If you don't have a trace yet, see the -[PyTorch profiling guide](https://github.com/AMD-AGI/TraceLens/blob/main/docs_original/conceptual/torch_profiling_guide.ipynb) +[PyTorch profiling walkthrough](../tutorials/torch-profiling.ipynb) for instructions on capturing one with `torch.profiler`. ``` @@ -77,7 +77,7 @@ For the GPU timeline, a low computation percentage with significant idle time indicates poor compute/communication overlap; use `--micro_idle_thresh_us` to split very short idle gaps into their own category. -See [Performance report column definitions](https://github.com/AMD-AGI/TraceLens/blob/main/docs_original/perf_report_columns.md) +See [Performance report column reference](../reference/perf-report-columns.md) for what each column means. ## Roofline classification diff --git a/docs/how-to/nccl-analysis.md b/docs/how-to/nccl-analysis.md new file mode 100644 index 000000000..16313c175 --- /dev/null +++ b/docs/how-to/nccl-analysis.md @@ -0,0 +1,100 @@ + + +# Analyze collectives with the NcclAnalyser SDK +```{meta} +:description: Learn how to use the TraceLens NcclAnalyser Python SDK to extract per-rank collective metrics like communication latency, bandwidth, and synchronization skew from PyTorch traces. +:keywords: TraceLens, NcclAnalyser, NCCL, RCCL, collective communication, algorithm bandwidth, bus bandwidth, skew, alltoallv, MoE, ROCm, PyTorch profiler +``` + +In distributed deep learning, analyzing the performance of collective +communication operations is crucial for diagnosing and optimizing performance at +scale. `NcclAnalyser` is a Python SDK that parses and analyzes NVIDIA Collective +Communications Library (NCCL) and ROCm Collective Communications Library (RCCL) +kernel events from PyTorch trace files (JSON). It computes key metrics like +communication latency, message sizes, algorithm bandwidth, bus bandwidth, and +synchronization metrics (for example skew in start and end times), providing +insights into communication patterns and potential bottlenecks in your +distributed training or inference workflows. + +```{note} +This topic covers the `NcclAnalyser` SDK and its API. To generate a multi-rank +collective-communication report from the command line instead, see +[Generate a collective-communication report](./collective-report.md). +``` + +## Key features + +- *Detailed NCCL event extraction*: Parses JSON trace files to extract per-rank NCCL kernel events (those containing `nccl` in the kernel name). +- *Correlated events via collective ID*: Each collective operation is assigned a unique `collective_id`, generated by combining the Process Group Name and an `index_in_group`. The `index_in_group` reflects the order in which that collective appears in the trace. +- *Latency computation for implicit-sync collectives*: For collectives like allreduce, allgather, reducescatter, and alltoall, the tool computes communication latency by taking the minimum duration across ranks to eliminate waiting time. +- *Bandwidth analysis*: Computes algorithm bandwidth based on input message size and communication latency, and finally bus bandwidth (link utilization). +- *Synchronization (wait-time) metrics*: Tracks skew in start times (how late certain ranks arrive) and skew in end times, providing insights for diagnosing load imbalance or synchronization overheads. +- *Alltoallv support*: Provides raw data for alltoallv events, which don't enforce an implicit sync pattern. Includes per-instance metrics (wall time, throughput, size imbalance) and a summary aggregation grouped by process group and dtype. An optional heatmap mode (`build_df_all2allv_heatmap`) produces per rank-pair send volumes, useful for diagnosing Mixture-of-Experts (MoE) / expert-parallel routing imbalance. +- *Summary and detailed dataframes*: Supports both high-level summaries for quick insights and detailed dataframes for advanced analysis, including per-rank metadata and performance metrics. Power users can build custom pipelines and analyses on top of the detailed dataframes. +- *PyTorch support*: Currently built for PyTorch trace files, with the potential to extend support to other frameworks. + +## Before you begin + +- TraceLens installed (see [Install TraceLens](../install/install.md)). +- Per-rank PyTorch profiler traces from a distributed run (one trace per rank). + +## Quick start + +`NcclAnalyser` can generate both a per-collective dataframe and a global summary. +More features are demonstrated in the +[`nccl_analyser_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/nccl_analyser_example.ipynb) +notebook. + +### Quick summary of implicit-sync collectives + +```python +from TraceLens import NcclAnalyser +import os + +# Define the root directory and create file paths for all ranks +root_dir = '/path/to/your/trace/files/directory' +world_size = 8 +list_profile_filepaths = [os.path.join(root_dir, f'rank{i}_trace.json') for i in range(world_size)] + +# Initialize the NCCL Analyser +my_nccl_analyser = NcclAnalyser(list_profile_filepaths, world_size) + +# Generate the summarized dataframe for implicit-sync collectives +df_summary = my_nccl_analyser.build_df_summary_nccl_implicit_sync_cat(agg_metrics=['mean']) +print(df_summary.head()) + +# Save the summary to CSV +df_summary.to_csv('nccl_summary.csv', index=False) +``` + +The summarized dataframe: + +- Groups events by collective type, message size, and data type to compute aggregated metrics. +- For communication performance, computes communication latency, algorithm bandwidth, and bus bandwidth. +- For synchronization delay, computes metrics like skew in start times, which shows the wait time. This is based on the difference between the earliest and latest arrival across the ranks for a given collective. + +### Example summarized dataframe + +| Collective name | In msg size (MB) | dtype | comm latency (µs)_mean | count | Total latency (ms) | algo bw (GB/s)_mean | bus bw (GB/s)_mean | skew in start time (µs)_mean | +|-----------------|------------------|-------|------------------------|-------|--------------------|---------------------|--------------------|------------------------------| +| _allgather_base | 204.00 | BFloat16 | 6041.88 | 318 | 1921.32 | 33.00 | 28.88 | 11779.36 | +| _reduce_scatter_base | 3264.06 | Float | 11662.77 | 160 | 1866.04 | 273.43 | 239.25 | 60238.77 | +| _reduce_scatter_base | 8016.03 | Float | 22988.50 | 2 | 45.98 | 340.53 | 297.96 | 146.48 | +| _allgather_base | 501.00 | BFloat16 | 11920.84 | 2 | 23.84 | 41.04 | 35.91 | 15405.14 | +| allreduce | 0.00 | Float | 18.58 | 6 | 0.11 | 0.00 | 0.00 | 936.63 | + +Note that the last row's message size is just rounded down to 0. + +Modify the file paths for your profiles in the `nccl_analyser_example.ipynb` +notebook to get analysis instantly. + +## Related topics + +- [Generate a collective-communication report](./collective-report.md) +- [Compare two traces in TraceLens](./compare-traces.md) +- [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/how-to/origami-integration.md b/docs/how-to/origami-integration.md new file mode 100644 index 000000000..c3eb9e79c --- /dev/null +++ b/docs/how-to/origami-integration.md @@ -0,0 +1,149 @@ + + +# Estimate simulated kernel times with Origami in TraceLens +```{meta} +:description: Learn how to enable Origami in TraceLens to estimate simulated GEMM and SDPA kernel times from a GPU architecture description. +:keywords: TraceLens, Origami, rocm-origami, GEMM, SDPA, performance model, simulation, gpu_arch_json_path, ROCm, roofline +``` + +TraceLens can estimate simulated GEMM and SDPA kernel times using *Origami* +(ROCm's performance modeling library) when you provide a GPU architecture +description and explicitly opt in. This is optional: `pip install TraceLens` +doesn't install Origami. + +## What Origami does in TraceLens + +- GEMM and SDPA perf models call Origami's Python bindings to predict a duration in microseconds for forward (and SDPA backward where applicable). +- Results show up in perf reports under columns such as `Origami Time (µs)`, `Origami TFLOPS/s`, `Origami TB/s`, and `Pct Origami` (relative to measured kernel busy time), when simulation uses Origami. +- Roofline metrics from `--gpu_arch_json_path` are separate; they don't require Origami. Origami adds *simulated* timing on top when enabled. + +## Installation + +### Python package + +Install the published package (PyPI name `rocm-origami`): + +```bash +pip install rocm-origami +``` + +Note that installing Origami currently (as of April 2026) requires that ROCm is +installed on the system, because Origami can use the HIP library to detect the +GPU currently installed in the system and use it as an architectural model. +TraceLens doesn't currently use this functionality, but it currently can't be +removed from Origami. + +### System environment + +Origami's wheels and bindings expect a ROCm runtime on the machine (a GPU isn't +required for pure Python simulation in many cases, but library loading may +depend on your setup). The project's continuous integration (CI) uses an AMD +ROCm container and installs `rocm-origami` alongside TraceLens (see +[`.github/workflows/regression-tests.yml`](https://github.com/AMD-AGI/TraceLens/blob/main/.github/workflows/regression-tests.yml)). + +If `import origami` fails after `pip install`, check: + +- Python version and wheel compatibility for `rocm-origami`. +- `LD_LIBRARY_PATH` and ROCm install paths required by the Origami wheel you use. + +## When TraceLens uses Origami + +Simulation is implemented in `TraceLens/PerfModel/perf_model.py` +(`GEMM.get_simulation_time_func`): + +- If `enable_origami` is true and the perf model has the needed architecture and parameters, TraceLens uses Origami. +- If `enable_origami` is false, TraceLens doesn't call Origami; simulated times are omitted for that path. + +So you need both: + +- A valid GPU architecture (see below), and +- `enable_origami=True` (CLI flag or Python API). + +## GPU architecture JSON + +Pass the same JSON file you use for roofline analysis with +`--gpu_arch_json_path`. It must include fields Origami expects (for example GPU +name, frequency, memory bandwidth, and CU count), as consumed by +`OrigamiHelper.get_hardware` in `TraceLens/PerfModel/origami_helper.py`. + +See [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +for the general format of the architecture JSON and its use in roofline +analysis. + +## Command-line usage + +### PyTorch perf report + +```bash +TraceLens_generate_perf_report_pytorch \ + --profile_json_path path/to/profile.json.gz \ + --gpu_arch_json_path path/to/gpu_arch.json \ + --enable-origami \ + --output_csvs_dir ./out_csvs +``` + +Or: + +```bash +python -m TraceLens.Reporting.generate_perf_report_pytorch \ + --profile_json_path path/to/profile.json.gz \ + --gpu_arch_json_path path/to/gpu_arch.json \ + --enable-origami \ + --output_csvs_dir ./out_csvs +``` + +### vLLM-oriented PyTorch report + +Same pattern; the entry point mirrors the PyTorch script (`--enable-origami`). + +### JAX perf report + +```bash +TraceLens_generate_perf_report_jax \ + --profile_path path/to/trace.xplane.pb \ + --gpu_arch_json_path path/to/gpu_arch.json \ + --enable-origami \ + --output_csvs_dir ./out_csvs +``` + +### Standalone GEMM/SDPA simulator helper + +`TraceLens/PerfModel/run_perf_model.py`: + +```bash +python -m TraceLens.PerfModel.run_perf_model --op gemm ... --enable_origami +``` + +## Python API + +When building a `TreePerfAnalyzer` (or `JaxTreePerfAnalyzer`) in code, pass: + +```python +TreePerfAnalyzer.from_file( + profile_filepath="profile.json.gz", + arch=gpu_arch_dict, # or load JSON + enable_origami=True, +) +``` + +Reporting helpers such as `generate_perf_report_pytorch` accept +`enable_origami=True` and forward it to the analyzer. + +## Troubleshooting + +| Symptom | What to check | +|---------|---------------| +| No Origami columns in CSVs | Confirm `--enable-origami` (or API `enable_origami=True`) and `--gpu_arch_json_path`. | +| Message on stderr about `origami` import | Install `rocm-origami` and fix ROCm/library paths. | +| Unsupported dtype warning | The Origami path supports a fixed set of dtypes (for example fp16, bf16, fp32, fp64, fp8); others skip simulation. | + +## Related topics + +- [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Generate a JAX performance report](./generate-perf-report-jax.md) +- [Generate a PyTorch inference performance report](./generate-perf-report-pytorch-inference.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/how-to/perf-model-without-trace.md b/docs/how-to/perf-model-without-trace.md new file mode 100644 index 000000000..8d543ab5e --- /dev/null +++ b/docs/how-to/perf-model-without-trace.md @@ -0,0 +1,105 @@ + + +# Model op performance without a trace +```{meta} +:description: Learn how to call the TraceLens perf model directly as an analytical library - compute FLOPs, bytes, arithmetic intensity, and roofline-bounded time from op shapes alone, with no trace or kernel launch. +:keywords: TraceLens, perf model, roofline, FLOPS, arithmetic intensity, GEMM, SDPA, analytical model, sizing study, ceiling, no trace, ROCm, MI300X +``` + +In normal use, TraceLens fuses two sources per kernel: the **perf model** supplies +analytical FLOPs and bytes from an op's shapes and dtypes, and the **trace** +supplies the measured wall-clock time. Dividing one by the other gives the +achieved `TFLOPS/s` and `TB/s` in every report. + +The perf model is just a Python library — you can call the analytical half on its +own, with **no trace and no kernel launch**. Hand it shapes, get back FLOPs and +bytes; arithmetic intensity and roofline-bounded time follow. This is useful for +sizing studies, end-to-end ceilings, and regression bounds before you run +anything. + +## Before you begin + +- TraceLens installed (see [Install TraceLens](../install/install.md)). +- The shapes and dtypes of the ops you want to model (no trace required). + +## Op coverage + +The perf model ships curated formulas for the fundamental op classes in modern +workloads — `GEMM`, `SDPA`, `CONV`, `Softmax`, `Normalization`, `GroupedGemm`, +`MoEComm`, `MambaSSD`, `FusedRoPE`, `CrossEntropy`, `CausalConv1d`, plus the +elementwise and reduction primitives — over a dozen base classes, all maintained +and tested against real traces. Every base class in +`TraceLens.PerfModel.perf_model` exposes `flops_func` and `bytes_func` (and +`*_bwd_func` variants where a backward pass applies). + +## Compute FLOPs, bytes, and arithmetic intensity + +Import the op class and call its `flops_func` / `bytes_func` with the op's shapes. +For a GEMM of shape `M × N × K`: + +```python +from TraceLens.PerfModel.perf_model import GEMM + +M, N, K = 40960, 6144, 1536 +bpe = 2 # bytes per element (bf16) + +flops = GEMM.flops_func(M, N, K, bias=False) # 2 * M * N * K +bytes_ = GEMM.bytes_func(M, N, K, False, bpe, bpe, bpe, bpe) # read A + read B + write C +arithmetic_intensity = flops / bytes_ # FLOPs per byte +``` + +Attention works the same way, with forward and backward variants: + +```python +from TraceLens.PerfModel.perf_model import SDPA + +# B, S_q, H_Q, S_kv, H_KV, d_h_qk, d_h_v, causal +flops_fwd = SDPA.flops_func(1, 4096, 64, 4096, 8, 128, 128, True) +bytes_fwd = SDPA.bytes_func(1, 4096, 64, 4096, 8, 128, 128, True, bpe) +flops_bwd = SDPA.flops_bwd_func(1, 4096, 64, 4096, 8, 128, 128, True, flash_impl=True) +``` + +## Bound the time with a roofline + +Arithmetic intensity tells you which side of the roofline an op sits on. Compare +it against the hardware's roofline knee (peak compute divided by peak bandwidth), +then take the larger of the compute-bound and bandwidth-bound times as the +achievable lower bound: + +```python +PEAK_TFLOPS_BF16 = 708.0 # MI300X matrix bf16 +HBM_BW_GBPS = 5300.0 # HBM3 + +ai_knee = (PEAK_TFLOPS_BF16 * 1e12) / (HBM_BW_GBPS * 1e9) # FLOPs/byte crossover + +def roofline_time_us(flops, bytes_): + compute_us = flops / (PEAK_TFLOPS_BF16 * 1e12) * 1e6 + bandwidth_us = bytes_ / (HBM_BW_GBPS * 1e9) * 1e6 + return max(compute_us, bandwidth_us) +``` + +This is the same analytical model TraceLens's built-in roofline analysis uses — +here applied to shapes directly instead of to a trace. The roofline is an upper +bound on achievable performance (a real kernel won't reach it), but the *ratios* +between layers and regimes are solid planning numbers that fall out of pure +arithmetic. + +## Worked example + +The +[`perf_model_without_trace.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/perf_model_without_trace.ipynb) +notebook works a full example end to end: every linear in a Llama-3-70B decoder +layer plotted as arithmetic intensity, SDPA forward and backward across sequence +length, and a per-token roofline rollup (predicted ms per token for prefill and +decode, bf16, MI300X) — all without a single trace. + +## Related topics + +- [Analyze traces with the TraceLens SDK](./sdk-analysis.md) +- [GEMM analysis in TraceLens](../conceptual/gemm-analysis.md) +- [Performance report columns](../reference/perf-report-columns.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/how-to/sdk-analysis.md b/docs/how-to/sdk-analysis.md new file mode 100644 index 000000000..6f64eb539 --- /dev/null +++ b/docs/how-to/sdk-analysis.md @@ -0,0 +1,282 @@ + + +# Analyze traces with the TraceLens SDK +```{meta} +:description: Learn how to use the TraceLens Python SDK (TreePerfAnalyzer) to build a call tree from a trace, navigate it, and compute GPU-timeline, per-operation, roofline, and nn.Module-level performance metrics. +:keywords: TraceLens, TreePerfAnalyzer, TreePerf, Trace2Tree, GPUEventAnalyser, SDK, GPU timeline, roofline, FLOPS, nn.Module, PyTorch profiler, JAX, ROCm +``` + +The TraceLens SDK builds a hierarchical call tree from a profiler trace and +computes performance metrics on top of it. `TreePerfAnalyzer` is the main entry +point: it wraps [`Trace2Tree`](../conceptual/trace2tree.md) (which parses the +trace into a CPU-to-GPU dependency tree) and exposes methods to navigate that tree +and derive metrics at the timeline, operation, roofline, and `nn.Module` levels. + +This topic is the narrated reference for the SDK. Each section links a runnable +notebook you can execute against your own trace. + +If you just need the standard metrics in a spreadsheet, the +[`generate_perf_report_pytorch.py`](https://github.com/AMD-AGI/TraceLens/blob/main/TraceLens/Reporting/generate_perf_report_pytorch.py) +script packages everything covered here into a ready-made Excel report. Reach for +the SDK when you need more than that report offers — custom filtering, ad-hoc +aggregations, support for a new operation, or programmatic access to the tree. + +## On this page + +- [Build and navigate the tree](#build-and-navigate-the-tree) +- [GPU timeline breakdown](#gpu-timeline-breakdown) +- [Per-operation GPU time](#per-operation-gpu-time) +- [Roofline metrics](#roofline-metrics) +- [nn.Module attribution](#nnmodule-attribution) + +## Before you begin + +- TraceLens installed (see [Install TraceLens](../install/install.md)). +- A PyTorch profiler trace JSON file for the model you want to analyze. +- A basic understanding of how `Trace2Tree` builds the tree (see + [The Trace2Tree data model](../conceptual/trace2tree.md)). + +## Build and navigate the tree + +Load a trace with `TreePerfAnalyzer.from_file`. Pass `add_python_func=True` to +include the Python call stack, so you can trace GPU kernels all the way back to +your Python code: + +```python +from TraceLens.TreePerf import TreePerfAnalyzer + +analyzer = TreePerfAnalyzer.from_file("/path/to/trace.json", add_python_func=True) +tree = analyzer.tree +``` + +Once you have an event of interest, the tree exposes navigation helpers: + +- `tree.traverse_subtree_and_print(event)` — visualize the subtree rooted at an + operation (its runtime calls and GPU kernels). +- `tree.traverse_parents_and_print(event, cpu_op_fields=(...))` — walk the parent + chain up to the Python frame that launched it; `cpu_op_fields` optionally shows + op details (`'Input Dims'`, `'Input type'`, `'Input Strides'`, `'Concrete Inputs'`). +- `tree.get_parent_event(event)`, `tree.get_children_events(event)`, and + `tree.get_gpu_events(event)` — direct parent/child/kernel access. + +For example, `traverse_subtree_and_print` on an `aten::convolution` op shows how +the dispatch op lowers to runtime launches and kernels: + +``` +└── UID: 41, Category: cpu_op, Name: aten::convolution + └── UID: 42, Category: cpu_op, Name: aten::_convolution + └── UID: 43, Category: cpu_op, Name: aten::miopen_convolution + ├── UID: 104314, Category: cuda_runtime, Name: hipExtModuleLaunchKernel + │ └── UID: 107846, Category: kernel, Name: Im2d2Col_v2, Duration: 45.063 + └── UID: 104318, Category: cuda_runtime, Name: hipExtModuleLaunchKernel + └── UID: 107848, Category: kernel, Name: Cijk_Ailk_Bljk_BBS_BH... +``` + +Walking the other direction, `traverse_parents_and_print` traces an op back +through the runtime and Python frames all the way to the model's `forward`, and +with `cpu_op_fields` it annotates each level with argument metadata: + +``` +Node: + UID: 41, Category: cpu_op, Name: aten::convolution + Input Dims: [[1, 768, 24, 24], [768, 768, 3, 3], []] + Input type: [float, float, float] +1-up: + UID: 40, Category: cpu_op, Name: aten::conv2d +2-up: + UID: 40139, Category: python_function, Name: +... +5-up: + UID: 40136, Category: python_function, Name: torch/nn/modules/conv.py(558): forward +``` + +For a full interactive walkthrough of tree navigation, see +[`trace2tree_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/trace2tree_example.ipynb). + +## GPU timeline breakdown + +`get_df_gpu_timeline()` returns a high-level view of GPU activity — computation, +communication, memcpy, exposed (non-overlapped) communication, busy, and idle +time: + +```python +analyzer.get_df_gpu_timeline() +``` + +Example output: + +| type | time ms | percent | +|------|---------|---------| +| busy_time | 6521.46 | 99.93 | +| computation_time | 6318.26 | 96.81 | +| exposed_communication_time | 203.06 | 3.11 | +| exposed_memcpy_time | 0.14 | 0.00 | +| idle_time | 4.72 | 0.07 | +| total_time | 6526.18 | 100.00 | + +The metrics are defined as: + +- **computation_time** — time in compute kernels (GEMMs, convolutions, and so on). +- **communication_time** — time in collective/communication kernels (for example, NCCL/RCCL). +- **memcpy_time** — time in host-device or device-device memory copies. +- **exposed_communication_time** / **exposed_memcpy_time** — the portion of + communication or memcpy that does *not* overlap computation (the part that + actually costs wall-clock time). +- **busy_time** / **idle_time** — time the GPU is doing anything vs. nothing. + +This breakdown is computed by `GPUEventAnalyser`, which works directly from a list +of GPU events. It operates on **any GPU timeline** (PyTorch, JAX, and others) and +uses **no host-side call stack**, so you can run it standalone without building +the CPU-side tree: + +```python +import json +from TraceLens import GPUEventAnalyser + +with open("/path/to/trace.json") as f: + events = json.load(f)["traceEvents"] + +df = GPUEventAnalyser(events).get_breakdown_df() +``` + +For JAX traces (which describe events differently and include all GPUs in one +trace), use `JaxGPUEventAnalyser`; `get_breakdown_df_multigpu()` returns a +`{gpu_index: DataFrame}` mapping. To adapt the analyzer to another profiling +format, subclass it and reimplement `get_gpu_event_lists()`. + +## Per-operation GPU time + +Break GPU time down by the lowest-level CPU operations (from the call-stack +perspective) and the time they *induce* on the GPU. This is a more stable, +interpretable abstraction than raw kernel durations: + +```python +df = analyzer.get_df_kernel_launchers(include_kernel_details=True) +analyzer.get_df_kernel_launchers_summary(df) # grouped by op name +analyzer.get_df_kernel_launchers_unique_args(df, include_pct=True) # by op + unique args +``` + +Example op-name summary (top rows by induced GPU time): + +| name | Count | total_direct_kernel_time_ms | Percentage (%) | +|------|-------|-----------------------------|----------------| +| aten::mm | 126 | 5576.76 | 88.73 | +| flash_attn::_flash_attn_backward | 8 | 213.62 | 3.40 | +| flash_attn::_flash_attn_forward | 8 | 119.65 | 1.90 | +| aten::copy_ | 4 | 69.96 | 1.11 | + +Pass `event_name="aten::mm"` to `get_df_kernel_launchers_unique_args` to drill +into a single op's argument combinations. + +## Roofline metrics + +For modeled op classes (GEMM, SDPA, CONV, and more), `TreePerf` computes +analytical FLOPs and bytes and combines them with measured time to yield achieved +efficiency (`TFLOPS/s`, `TB/s`, `FLOPS/Byte`): + +```python +from TraceLens.PerfModel.torch_op_mapping import build_sheet_category_to_op_names + +categories = build_sheet_category_to_op_names(analyzer.op_to_perf_model_class_map) +gemm_events = [e for e in analyzer.tree.events if e["name"] in categories["GEMM"]] + +df_gemm = analyzer.build_df_perf_metrics(gemm_events, include_kernel_details=True) +analyzer.summarize_df_perf_metrics(df_gemm, ["mean"]) # group by M, N, K, bias +``` + +Example `aten::mm` summary (top rows by induced GPU time, grouped by shape): + +| name | param: M | param: N | param: K | param: bias | FLOPS/Byte_first | TFLOPS/s_mean | +|---------|----------|----------|----------|-------------|------------------|---------------| +| aten::mm | 73728 | 28672 | 8192 | False | 5864.73 | 698.19 | +| aten::mm | 73728 | 8192 | 28672 | False | 5864.73 | 719.59 | +| aten::mm | 28672 | 8192 | 73728 | False | 5864.73 | 740.51 | +| aten::mm | 73728 | 128256 | 8192 | False | 6972.01 | 628.10 | + +`FLOPS/Byte` is the operation's arithmetic intensity; comparing achieved +`TFLOPS/s_mean` against the device peak tells you whether each shape is compute- +or memory-bound. + +For how these metrics are defined and modeled, see +[GEMM analysis](../conceptual/gemm-analysis.md) and the +[performance report columns](../reference/perf-report-columns.md) reference. To +compute the same metrics from shapes alone, without a trace, see +[Model op performance without a trace](./perf-model-without-trace.md). The +[`tree_perf_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/tree_perf_example.ipynb) +notebook walks through roofline metrics for every modeled category. + +### Add a new operation + +Roofline metrics are only available for modeled op classes, but extending the +perf model to a new op is straightforward, and contributions are welcome: + +- In [`perf_model.py`](https://github.com/AMD-AGI/TraceLens/blob/main/TraceLens/PerfModel/perf_model.py): + parse the operation's shapes from the trace and write (or reuse) a performance + model that returns FLOPs and bytes. +- In [`torch_op_mapping.py`](https://github.com/AMD-AGI/TraceLens/blob/main/TraceLens/PerfModel/torch_op_mapping.py): + map the operation name to that perf model. + +Once mapped, the op is picked up automatically by `build_df_perf_metrics` and by +the report generator. + +## nn.Module attribution + +GPU time can be attributed all the way up to the Python `nn.Module` hierarchy, +making performance insights accessible without call-stack expertise. Build a +latency tree rooted at a module and traverse it: + +```python +analyzer = TreePerfAnalyzer.from_file("/path/to/trace.json", add_python_func=True) +tree = analyzer.tree + +event = next(e for e in tree.events if e["name"] == "nn.Module: DeepseekV2DecoderLayer_2") +analyzer.build_nn_module_latency_tree(event) + +# navigate the module hierarchy +children = tree.get_nn_module_children(event) # list of module UIDs +parent = tree.get_nn_module_parent(event) # parent module UID +``` + +Each node carries its aggregated `GPU Time`, so a recursive walk prints a +module-level breakdown that mirrors your model architecture. The example below is +a single DeepSeek-V2 decoder layer, with each module annotated by its GPU time +(µs) and a `Non-nn.Module GPU Ops` entry for GPU work not owned by a child +module: + +``` + (GPU Time μs) +└── nn.Module: DeepseekV2DecoderLayer_2 ..................... (10233.51) + ├── nn.Module: RMSNorm_8 ................................ (148.46) + ├── nn.Module: DeepseekV2AttentionMLA_2 ................. (6775.96) + │ ├── nn.Module: ReplicatedLinear_4 ................... (817.78) + │ ├── nn.Module: RMSNorm_9 ............................ (19.67) + │ ├── nn.Module: ColumnParallelLinear_2 ............... (268.80) + │ ├── nn.Module: ReplicatedLinear_5 ................... (469.98) + │ ├── nn.Module: RMSNorm_10 ........................... (10.21) + │ ├── nn.Module: DeepseekScalingRotaryEmbedding_0 ..... (139.27) + │ ├── nn.Module: RadixAttention_2 ..................... (3204.62) + │ ├── nn.Module: RowParallelLinear_4 ................. (1446.40) + │ └── Non-nn.Module GPU Ops ........................... (399.22) + ├── nn.Module: RMSNorm_11 .............................. (154.43) + └── nn.Module: DeepseekV2MLP_2 ......................... (3154.66) + ├── nn.Module: MergedColumnParallelLinear_2 ........ (1592.94) + ├── nn.Module: SiluAndMul_2 ....................... (84.44) + └── nn.Module: RowParallelLinear_5 ................ (1477.28) +``` + +See +[`nn_module_view.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/nn_module_view.ipynb) +for the full traversal against your own trace, including the recursive +pretty-printer. + +## Related topics + +- [The Trace2Tree data model](../conceptual/trace2tree.md) +- [Model op performance without a trace](./perf-model-without-trace.md) +- [Generate a PyTorch performance report](./generate-perf-report-pytorch.md) +- [Replay a single operation in TraceLens](./event-replay.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/index.rst b/docs/index.rst index 708717e24..e61866ef2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -32,10 +32,26 @@ The TraceLens source code is hosted at `github.com/AMD-AGI/TraceLens ` * :doc:`Replay a single operation ` * :doc:`Fuse multi-rank traces ` + * :doc:`Analyze traces with the TraceLens SDK ` + * :doc:`Agentic performance analysis with the TraceLens Agent ` + + .. grid-item-card:: Concepts + + * :doc:`The Trace2Tree data model ` + * :doc:`Understanding PyTorch traces ` + * :doc:`GEMM analysis ` + * :doc:`Triton kernel performance model ` + + .. grid-item-card:: Tutorials + + * :doc:`PyTorch profiling walkthrough ` + * :doc:`Distributed profiling walkthrough ` + * :doc:`Anomaly detection walkthrough ` .. grid-item-card:: Reference * :doc:`API reference ` + * :doc:`Performance report columns ` For information on contributing to TraceLens, see the `Contributing guide `_. diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index ff32866fa..6078fb478 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -180,23 +180,23 @@ Split an inference trace into per-iteration or per-phase sub-traces. ## Python SDK The SDK modules live under the `TraceLens` package and can be imported to build -custom workflows. Each module has a dedicated guide in the repository `docs_original/` -directory and an example notebook under `examples/`. +custom workflows. Each module has a dedicated topic in the documentation and an +example notebook under `examples/`. | Module | Purpose | Reference | |--------|---------|-----------| -| `Trace2Tree` | Build and navigate the hierarchical event tree (Python ops → CPU dispatch → GPU kernels). | `docs_original/Trace2Tree.md`, `examples/trace2tree_example.ipynb` | -| `TreePerf` | GPU-timeline breakdown, per-op performance, and roofline metrics. | `docs_original/TreePerf.md`, `examples/tree_perf_example.ipynb` | -| `PerfModel` | Compute and roofline performance models for operators. | `docs_original/gemm_dim_eff.md`, `docs_original/triton_perf_model_walkthrough.md` | -| `NcclAnalyser` | Multi-rank collective latency/bandwidth/skew analysis. | `docs_original/NcclAnalyser.md`, `examples/nccl_analyser_example.ipynb` | -| `TraceDiff` | Morphological comparison of two trace trees. | `docs_original/TraceDiff.md`, `examples/trace_diff_example.ipynb` | -| `EventReplay` | Extract and replay isolated operations. | `docs_original/EventReplay.md`, `examples/event_replayer_example.ipynb` | -| `TraceFusion` | Merge multi-rank traces for Perfetto visualization. | `docs_original/TraceFusion.md`, `examples/trace_fusion_example.py` | -| `Reporting` | The report generators behind the CLI tools; importable to return pandas data frames. | `docs_original/generate_perf_report.md` | +| `Trace2Tree` | Build and navigate the hierarchical event tree (Python ops → CPU dispatch → GPU kernels). | [Trace2Tree data model](../conceptual/trace2tree.md), [`trace2tree_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/trace2tree_example.ipynb) | +| `TreePerf` | GPU-timeline breakdown, per-op performance, and roofline metrics. | [Analyze traces with the TraceLens SDK](../how-to/sdk-analysis.md), [`tree_perf_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/tree_perf_example.ipynb) | +| `PerfModel` | Compute and roofline performance models for operators. | [GEMM analysis](../conceptual/gemm-analysis.md), [Triton kernel performance model](../conceptual/triton-perf-model-walkthrough.md) | +| `NcclAnalyser` | Multi-rank collective latency/bandwidth/skew analysis. | [Analyze collectives with NcclAnalyser](../how-to/nccl-analysis.md), [`nccl_analyser_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/nccl_analyser_example.ipynb) | +| `TraceDiff` | Morphological comparison of two trace trees. | [Compare two traces](../how-to/compare-traces.md), [`trace_diff_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/trace_diff_example.ipynb) | +| `EventReplay` | Extract and replay isolated operations. | [Replay a single operation](../how-to/event-replay.md), [`event_replayer_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/event_replayer_example.ipynb) | +| `TraceFusion` | Merge multi-rank traces for Perfetto visualization. | [Fuse multi-rank traces](../how-to/trace-fusion.md), [`trace_fusion_example.py`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/trace_fusion_example.py) | +| `Reporting` | The report generators behind the CLI tools; importable to return pandas data frames. | [Generate a PyTorch performance report](../how-to/generate-perf-report-pytorch.md) | | `TraceUtils` | Trace utilities, including inference-trace splitting. | — | -For report-column definitions across all sheets, see -`docs_original/perf_report_columns.md` in the repository. +For report-column definitions across all sheets, see the +[Performance report column reference](./perf-report-columns.md). ```{note} For a class- and function-level SDK reference generated directly from diff --git a/docs/reference/perf-report-columns.md b/docs/reference/perf-report-columns.md new file mode 100644 index 000000000..94595863f --- /dev/null +++ b/docs/reference/perf-report-columns.md @@ -0,0 +1,922 @@ + + +# Performance report column reference +```{meta} +:description: Reference for every sheet and column in a TraceLens performance report, including the gpu_timeline breakdown, operation-analysis sheets, roofline metrics, and collective communication analysis. +:keywords: TraceLens, performance report, column reference, gpu_timeline, ops, ops_summary, ops_unique_args, roofline, GEMM, SDPA, convolution, NCCL, coll_analysis, GFLOPS, TFLOPS, arithmetic intensity, ROCm, AMD Instinct, MI300X +``` + +This topic explains the columns in each sheet of the TraceLens performance report. The report is an Excel file with multiple sheets, and each sheet analyzes a different aspect of GPU performance. Use it as a lookup while you read a report generated with one of the how-to topics listed under [Related topics](#related-topics). + +## Overview + +The performance report Excel file contains multiple sheets analyzing different aspects of GPU performance. The core sheets are: + +| Sheet | Description | +|-------|-------------| +| `gpu_timeline` | High-level GPU time breakdown. | +| `ops` | Detailed per-operation data (base data). | +| `ops_summary_by_category` | Operations summarized by category. | +| `ops_summary` | Operations summarized by name. | +| `ops_unique_args` | Operations summarized by unique argument combinations. | + +Additional sheets can include operation-specific analysis (GEMM, SDPA_fwd, CONV_fwd, and so on), kernel summary, short kernels, and collective analysis. + +### Unit conventions + +- *Time*: all times from the trace are in microseconds (µs) unless explicitly stated otherwise (for example, `time ms` in `gpu_timeline` is in milliseconds). +- *Compute*: GFLOPS (billions of FLOPs), TFLOPS/s (trillions of FLOPs per second). +- *Memory*: MB (megabytes), GB/s (gigabytes per second), TB/s (terabytes per second). + +## The gpu_timeline sheet + +The `gpu_timeline` sheet provides a high-level breakdown of GPU time into computation, communication, memory copy, and idle time, accounting for overlaps between different types of operations. + +### Example output + +Here's a typical `gpu_timeline` sheet from a distributed training workload: + +| type | time ms | percent | +| --------------------- | --------- | --------- | +| computation_time | 56305.19 | 99.30 | +| exposed_comm_time | 240.88 | 0.42 | +| exposed_memcpy_time | 14.44 | 0.03 | +| busy_time | 56560.52 | 99.75 | +| idle_time | 143.16 | 0.25 | +| total_time | 56703.68 | 100.00 | +| total_comm_time | 17203.43 | 30.34 | +| total_memcpy_time | 14.47 | 0.03 | + +In this example: + +- 99.30% computation time, so the GPU is very efficiently utilized. +- 0.42% exposed communication, which is excellent. Most communication (30.34% total - 0.42% exposed = 29.92%) overlaps with computation. +- 0.25% idle time, so there are minimal gaps and good kernel launch efficiency. +- This workload demonstrates effective computation/communication overlap. + +### Event classification + +GPU events are classified into three categories based on their `cat` field and kernel name: + +``` +GPU Event +│ +├─ cat = "gpu_memcpy"? +│ └─ YES → Memory Copy (H2D, D2H, D2D) +│ +├─ cat = "kernel"? +│ └─ YES → Does kernel name contain "nccl"? +│ ├─ YES → Communication (AllReduce, AllGather, ReduceScatter, etc.) +│ └─ NO → Computation (GEMM, Conv, Elementwise, etc.) +│ +└─ cat = "gpu_memset"? + └─ YES → Computation (grouped here for simplicity; typically very short) +``` + +```{note} +- "Computation" includes all kernels doing work (GEMM, convolution, elementwise, and so on), not just compute-bound operations. +- Communication kernels (NCCL) include both synchronization delay and actual data transfer time. In single-rank traces, these components can't be separated, so `total_comm_time` represents the total duration. See the NCCL Analyzer documentation for multi-rank analysis that can break down communication into sync and transfer phases. +``` + +### How time is calculated + +*Step 1: merge events by category.* Events of the same category are merged across all GPU streams into non-overlapping intervals. + +Example, merging computation events: + +``` +Stream 0: ──[Kernel A]────────[Kernel B]── +Stream 1: ─────[Kernel C]──[Kernel D]───── + ↓ +Merged: ──[──────────]──[──────────]── + (overlapping kernels merged into single intervals) +``` + +*Step 2: create interval sets.* After merging, there are four sets of non-overlapping intervals: + +- `COMP` = merged computation intervals +- `COMM` = merged communication intervals +- `MEMCPY` = merged memcpy intervals +- `ALL_GPU` = merged intervals of ALL GPU events (computation + communication + memcpy) + +*Step 3: apply set arithmetic.* Exposed metrics are calculated by subtracting overlaps: + +``` +exposed_comm intervals = COMM - COMP +exposed_memcpy intervals = MEMCPY - COMP - COMM +``` + +In simple terms: + +- Exposed communication is communication time that doesn't overlap with computation. +- Exposed memcpy is memcpy time that doesn't overlap with computation or communication. + +*Step 4: sum interval durations.* Time for each metric is the sum of durations of all intervals in that set: + +``` +computation_time = sum of durations in COMP +exposed_comm_time = sum of durations in (COMM - COMP) +exposed_memcpy_time = sum of durations in (MEMCPY - COMP - COMM) +busy_time = sum of durations in ALL_GPU +idle_time = total_time - busy_time +total_comm_time = sum of durations in COMM +total_memcpy_time = sum of durations in MEMCPY +total_time = end of last GPU event - start of first GPU event +``` + +The equation: + +``` +computation_time + exposed_comm_time + exposed_memcpy_time + idle_time = total_time +``` + +### Columns + +| Column | Description | Values | +|--------|-------------|--------| +| `type` | Category of GPU time | `computation_time`, `exposed_comm_time`, `exposed_memcpy_time`, `busy_time`, `idle_time` (or `micro_idle_time` and `macro_idle_time`), `total_comm_time`, `total_memcpy_time`, `total_time` | +| `time ms` | Time in milliseconds | Actual duration for each category | +| `percent` | Percentage of total time | `(time ms / total_time) * 100` | + +### Interpreting results + +- High `computation_time` (>80%): workload is efficiently using the GPU. +- High `exposed_comm_time`: communication isn't overlapped with computation, so optimize with computation/communication overlap. +- High `exposed_memcpy_time`: memory transfers are blocking work, so optimize data movement. +- High `idle_time`: the GPU is sitting idle, so check for CPU bottlenecks, kernel launch overhead, or synchronization. + +Generated by `TreePerfAnalyzer.get_df_gpu_timeline()` using `GPUEventAnalyser.compute_metrics()`. + +## Operation analysis + +The following sections analyze CPU operations that launch GPU kernels. Each operation is analyzed for its direct GPU time impact. + +### The ops sheet (base data) + +The `ops` sheet provides detailed per-operation data showing each CPU operation that launches GPU kernels. This is the base unsummarized data; the following sheets provide different ways of summarizing this information. Each row represents one CPU operation instance from your workload that launched GPU kernels, showing exactly which kernels it launched and how long they took. + +Generated by `TreePerfAnalyzer.get_df_kernel_launchers()`. + +#### Why analyze operations instead of kernels + +When you look at a PyTorch trace, you might be tempted to analyze GPU kernels directly. However, this has significant limitations: + +- Kernel names are cryptic and ambiguous. A kernel named `Cijk_Ailk_Bljk_BBS_BH_Bias_HAS_SAV...` tells you almost nothing about what computation it's doing. The same kernel name can also map to different computations depending on input shape, dtype, and memory layout. +- Kernel names vary across platforms. The same matrix multiply appears as `nvjet_*` or `cutlass_*` on NVIDIA GPUs, but as `Cijk_*` on AMD GPUs. This makes cross-platform comparison nearly impossible. +- Operations are stable and meaningful. Names like `aten::addmm` (matrix multiply with add) or `aten::conv2d` are platform-independent and immediately tell you what computation is happening. Combined with argument information (shapes, dtypes, strides), they fully define the computation in a reproducible way. + +By analyzing at the operation level, you get insights that are portable across platforms, interpretable without deep kernel knowledge, and reproducible. + +#### Which operations are analyzed + +TraceLens analyzes leaf operations, the lowest-level CPU operations in the call stack that directly launch GPU kernels. This gives you the most granular view without double-counting. + +Example call stack (showing all layers from the trace): + +``` +└── (python_function) nn.Module: Linear_75 + └── (python_function) forward @ linear.py:124 + └── (cpu_op) aten::linear + ├── (cpu_op) aten::to + │ └── (cpu_op) aten::_to_copy + │ └── (cpu_op) aten::copy_ ← Leaf op + │ └── (cuda_runtime) hipLaunchKernel ← Runtime layer + │ └── (kernel) elementwise_kernel... ← GPU kernel + ├── (cpu_op) aten::to + │ └── (cpu_op) aten::_to_copy + │ └── (cpu_op) aten::copy_ ← Leaf op + │ └── (cuda_runtime) hipLaunchKernel ← Runtime layer + │ └── (kernel) elementwise_kernel... ← GPU kernel + └── (cpu_op) aten::linear + └── (cpu_op) aten::addmm ← Leaf op + ├── (cuda_runtime) hipLaunchKernel ← Runtime layer + │ └── (kernel) elementwise_kernel... ← GPU kernel + └── (cuda_runtime) hipExtModuleLaunchKernel ← Runtime layer + └── (kernel) Cijk_Alik_Bljk_BBS_BH... ← GPU kernel +``` + +```{note} +The `cuda_runtime` label is used by the PyTorch profiler even on ROCm platforms. It's just a naming convention. +``` + +In this example: + +- Leaf operations analyzed: `aten::copy_` (2 instances) and `aten::addmm`. +- Not included: `aten::linear`, `aten::to`, `aten::_to_copy` (these are higher-level and don't directly launch kernels). +- Why this matters: if higher-level operations were included, the same GPU time would be counted multiple times. By focusing on leaf operations, each kernel's time is attributed to exactly one operation. + +#### What CPU operations (cpu_op) are + +When you write PyTorch code in Python (like `output = model(input)`), it goes through several layers before reaching the GPU. The trace captures all these layers: + +1. *Python frontend*: your Python code calls `nn.Module` methods (labeled `python_function` in the trace). +2. *Torch dispatcher*: operations are routed through PyTorch's dispatcher, which selects the appropriate implementation based on device, dtype, and other factors (labeled `cpu_op` in the trace). +3. *Runtime layer*: CUDA/HIP API calls to launch kernels (labeled `cuda_runtime` or `cuda_driver` in the trace). +4. *GPU kernels*: the actual computation on the GPU (labeled `kernel`, `gpu_memcpy`, or `gpu_memset` in the trace). + +Operations marked as `cpu_op` in the trace are operations registered in the PyTorch dispatcher. These can be registered in C++ or Python, and represent various levels of abstraction in the execution hierarchy. The key distinction is that TraceLens analyzes the ones that directly launch kernels (the leaf operations in the call stack). The dispatcher layer is where PyTorch attaches rich metadata about the computation. + +```{note} +In the trace hierarchy there's also a runtime layer (labeled `cuda_runtime` or `cuda_driver`) between CPU operations and GPU kernels. This represents the CUDA/HIP API calls that actually launch kernels. The PyTorch profiler uses the `cuda_runtime`/`cuda_driver` naming convention even on ROCm/HIP platforms. It's a naming convention inherited from CUDA. +``` + +#### What arguments operations contain + +Each operation in the trace contains argument information that fully characterizes the computation. These arguments come directly from the JSON event's `args` field that the PyTorch profiler captures. + +Example JSON event from a trace: + +```json +{ + "name": "aten::addmm", + "cat": "cpu_op", + "ts": 1234567890, + "dur": 150, + "args": { + "Input Dims": [[1024, 512], [512, 256], [256]], + "Input type": ["c10::BFloat16", "c10::BFloat16", "c10::BFloat16"], + "Input Strides": [[512, 1], [256, 1], [1]], + "Concrete Inputs": ["", "", "1.0"] + }, + ... +} +``` + +These arguments are extracted and presented in the `ops` sheet: + +| Argument type | Description | Example (from the `aten::addmm` above) | Why it matters | +|---------------|-------------|----------------------------------------|----------------| +| `Input Dims` | Shape of input tensors | `((1024, 512), (512, 256), (256,))` | Different shapes give different performance (tiling, memory access patterns). | +| `Input type` | Data types | `('c10::BFloat16', 'c10::BFloat16', 'c10::BFloat16')` | FP32 vs BF16 vs FP16 affects speed and memory. | +| `Input Strides` | Memory layout (row-major, column-major, and so on) | `((512, 1), (256, 1), (1,))` | Strided vs contiguous affects memory bandwidth. | +| `Concrete Inputs` | Scalar/string arguments | `('', '', '1.0')` | Additional parameters that can affect algorithm selection. | + +Why shapes matter: + +``` +aten::addmm with shape (1024, 512) × (512, 256): Fast (optimized tiling) +aten::addmm with shape (1023, 511) × (511, 255): Slower (odd dimensions, less efficient) +``` + +Why strides matter: + +``` +Same operation, same shape, different strides can have different performance: + stride=(512, 1) vs stride=(1, 512) +``` + +Strides affect memory access patterns. Some strides work better for certain operations depending on how kernels access data. TraceLens captures strides so you can identify stride-dependent performance variations. + +By capturing these arguments, you can group operations by their unique argument combinations. This enables targeted analysis: + +- Identify slow cases: "Which specific input shapes are causing slowdowns?" +- Reproduce issues: "I can reproduce the slow case with these exact inputs." +- Compare across platforms: "Does this shape perform better on NVIDIA or AMD?" + +This combination of operation name plus arguments forms a unique signature that can be aggregated in the summary sheets (`ops_summary_by_category`, `ops_summary`, `ops_unique_args`), enabling you to ask questions at different levels of granularity. + +#### How GPU time is calculated + +For each operation, TraceLens calculates `total_direct_kernel_time_us` using GPU Event Analyzer to account for kernel overlaps. When an operation launches multiple kernels, some can run concurrently on different GPU streams. GPU Event Analyzer computes the actual "busy time", the wall-clock time during which at least one of the operation's kernels is executing. + +Example: + +``` +Operation: aten::addmm + ├── Kernel A: 100 µs (stream 0, executes 0-100 µs) + └── Kernel B: 80 µs (stream 1, executes 20-100 µs, overlaps with Kernel A) + +Naive sum: 100 + 80 = 180 µs ✗ (incorrect) +Actual busy time: 100 µs ✓ (kernels overlapped from 20-100 µs) +``` + +This gives you the actual wall-clock time each operation's kernels occupied the GPU. + +#### Understanding kernel_details + +`kernel_details` captures detailed information about each GPU kernel launched by an operation: kernel name, duration (in microseconds), and stream ID. + +Example from `aten::addmm`: + +```python +[ + {'name': 'elementwise_kernel', 'dur': 2.3, 'stream': 7}, # dur in µs + {'name': 'Cijk_Alik_Bljk_BBS_BH_Bias_HAS_SAV...', 'dur': 45.2, 'stream': 7} # dur in µs +] + +# Analysis: addmm = elementwise (2.3 µs) + GEMM (45.2 µs) +# → GEMM is bottleneck (95% of time), both sequential (same stream) +# → Cijk_* = Tensile GEMM (AMD), cutlass_* would be NVIDIA +``` + +Key insights from `kernel_details`: + +- Backend implementation: kernel names reveal which library is used. + - `Cijk_*`: Tensile GEMM kernels (AMD ROCm) + - `cutlass_*`: CUTLASS kernels (NVIDIA) + - `ck_tile::kentry<...FmhaFwd...>`: Composable Kernel (CK) Flash Attention (AMD) + - `void at::native::*`: PyTorch native kernels (element-wise, and so on) +- Execution breakdown: shows the performance contribution of each kernel in multi-kernel operations. +- Concurrency: same stream = sequential, different streams = potentially concurrent. + +#### Columns + +| Column | Description | Details | +|--------|-------------|---------| +| `name` | Operation name | PyTorch operation name (for example, `aten::addmm`, `aten::index_select`) | +| `op category` | Operation category | Categorized type (for example, GEMM, CONV_fwd, SDPA_fwd, other) | +| `UID` | Unique event identifier | Unique ID for this specific operation instance in the trace | +| `total_direct_kernel_time` | GPU busy time in microseconds | Wall-clock time the GPU was busy executing this operation's kernels (accounts for overlaps using GPU Event Analyzer) | +| `direct_kernel_count` | Number of kernels launched | How many GPU kernels this operation launched | +| `Input Dims` | Input tensor shapes | Tuple of shapes for each input tensor, for example `((30522, 512), (), (141,))` | +| `Input type` | Input data types | Tuple of data types for each input, for example `('c10::BFloat16', 'Scalar', 'long int')` | +| `Input Strides` | Input tensor strides | Tuple of stride tuples for each input, for example `((512, 1), (), (1,))` | +| `Concrete Inputs` | Scalar/string arguments | Additional arguments passed to the operation, for example `('', '0', '')` | +| `kernel_details` | Kernel execution details | List of dicts with kernel name, duration (µs), and stream for each kernel launched, for example `[{'name': '...', 'dur': 3.136, 'stream': 7}]` | + +#### What this data represents + +The PyTorch trace has a hierarchical structure (Python → Operations → Runtime → Kernels). This sheet provides a MECE (mutually exclusive, collectively exhaustive) flat representation focusing on the leaf operations that launch GPU work. + +- Granularity: each row is a single CPU operation instance from the trace. +- Completeness: only operations that launch GPU kernels are included (pure CPU operations aren't shown). +- Raw data: this is unsummarized, so you see every individual operation occurrence. The following sheets (`ops_summary_by_category`, `ops_summary`, `ops_unique_args`) provide different ways to aggregate this data. + +The following three sheets provide different levels of aggregation of the base `ops` data, following a progressive disclosure of complexity. + +### The ops_summary_by_category sheet (most aggregated) + +This is the highest-level summary. It groups operations into broad computational categories (GEMM, convolution, attention, and so on). Use it to quickly identify which types of operations dominate your workload. + +#### Columns + +| Column | Description | Calculation | +|--------|-------------|-------------| +| `op category` | Operation category | Possible values: GEMM, CONV_fwd, CONV_bwd, SDPA_fwd, SDPA_bwd, BN_fwd, BN_bwd, triton, elementwise, reduce, multi_tensor_apply, other | +| `Count` | Number of operations in this category | Count of unique CPU operations | +| `total_direct_kernel_time_ms` | Total GPU time in milliseconds | Sum of all GPU kernel durations launched by operations in this category | +| `Percentage (%)` | Percentage of total GPU time | `(total_direct_kernel_time_ms / sum of all kernel time) * 100` | +| `Cumulative Percentage (%)` | Running cumulative percentage | Sum of percentages from top to current row | + +#### How operations are categorized + +Operations are automatically categorized based on name patterns and kernel characteristics. Some categories are detected from the operation name, while others are detected by inspecting the launched kernel names. + +| Category | Detection method | Example operations | +|----------|------------------|--------------------| +| GEMM | Operation name | `aten::addmm`, `aten::mm`, `aten::bmm`, `aten::baddbmm` | +| CONV_fwd | Operation name | `aten::convolution`, `aten::miopen_convolution`, `aten::cudnn_convolution` | +| CONV_bwd | Operation name | `aten::convolution_backward` | +| SDPA_fwd | Operation name | `aten::_scaled_dot_product_flash_attention`, `aten::_flash_attention_forward` | +| SDPA_bwd | Operation name | `aten::_scaled_dot_product_flash_attention_backward` | +| BN_fwd | Operation name | `aten::batch_norm`, `aten::native_batch_norm`, `aten::cudnn_batch_norm` | +| BN_bwd | Operation name | `aten::native_batch_norm_backward`, `aten::cudnn_batch_norm_backward` | +| triton | Operation name | Operations starting with `triton` | +| elementwise | Kernel name inspection | Operations launching `at::native` elementwise kernels (for example, `aten::relu`, `aten::add`, `aten::mul`) | +| reduce | Kernel name inspection | Operations launching `at::native` reduce kernels | +| multi_tensor_apply | Kernel name inspection | Operations launching `multi_tensor_apply` kernels | +| other | Default | Operations not matching any above patterns | + +Use this sheet when you want a high-level answer to "What types of operations are taking the most time?" + +Generated by `TreePerfAnalyzer.get_df_kernel_launchers_summary_by_category()`. + +### The ops_summary sheet (by operation name) + +This is a mid-level summary. It groups operations by their name (for example, all `aten::addmm` calls together, all `aten::conv2d` calls together). Use it when you know which category is expensive and want to see which specific operations within that category are the culprits. + +#### Columns + +| Column | Description | Calculation | +|--------|-------------|-------------| +| `name` | Operation name | PyTorch operation name (for example, `aten::addmm`, `aten::_flash_attention_forward`) | +| `total_direct_kernel_time_sum` | Total GPU time in microseconds | Sum of all GPU kernel durations launched by this operation type | +| `Count` | Number of operations | Count of instances of this operation | +| `total_direct_kernel_time_ms` | Total GPU time in milliseconds | `total_direct_kernel_time_sum / 1000` | +| `Percentage (%)` | Percentage of total GPU time | `(total_direct_kernel_time_ms / sum of all kernel time) * 100` | +| `Cumulative Percentage (%)` | Running cumulative percentage | Sum of percentages from top to current row | +| `is_recompute` | *(Optional)* Whether this operation is part of activation recomputation | `True` or `False`. Only present when `--detect_recompute` is enabled. The same op name can appear in two rows (one for recompute, one for non-recompute). | + +What this sheet shows: + +- Each unique operation name gets one row (for example, one row for all `aten::addmm` calls). +- All instances of the same operation are aggregated together, regardless of their input shapes, dtypes, or other arguments. +- More granular than the category view, but still hides variation due to different input arguments. +- When `--detect_recompute` is enabled, rows are further split by recompute status. + +Use this sheet when you're asking, "I see GEMM is expensive. Which specific GEMM operation is the problem: `aten::addmm`, `aten::mm`, or `aten::bmm`?" + +Generated by `TreePerfAnalyzer.get_df_kernel_launchers_summary()`. + +### The ops_unique_args sheet (most detailed) + +This is the most detailed summary. It groups operations by unique combinations of operation name AND input arguments (shapes, dtypes, strides, concrete inputs). Use it to identify which specific input patterns are causing performance issues. + +#### Columns + +| Column | Description | Details | +|--------|-------------|---------| +| `name` | Operation name | PyTorch operation name | +| `Input Dims` | Input tensor dimensions | Shape of input tensors (for example, `[[1, 512, 768], [768, 768]]`) | +| `Input type` | Input tensor data types | Data types of inputs (for example, `['c10::BFloat16', 'c10::BFloat16']`) | +| `Input Strides` | Input tensor strides | Memory layout strides for each input tensor | +| `Concrete Inputs` | Scalar input values | Non-tensor inputs like kernel size, stride, padding | +| `operation_count` | Number of occurrences | How many times this exact operation+args combination appeared | +| `total_direct_kernel_time_sum` | Total GPU time in microseconds | Sum of GPU kernel time for all occurrences | +| `total_direct_kernel_time_mean` | Mean GPU time in microseconds | Average GPU time per occurrence | +| `total_direct_kernel_time_median` | Median GPU time in microseconds | Median GPU time across occurrences | +| `total_direct_kernel_time_std` | Standard deviation | Standard deviation of GPU time across occurrences (microseconds) | +| `total_direct_kernel_time_min` | Minimum GPU time | Fastest occurrence (microseconds) | +| `total_direct_kernel_time_max` | Maximum GPU time | Slowest occurrence (microseconds) | +| `ex_UID` | Example event UID | UID of one example event (for further analysis) | +| `kernel_details_summary` | Kernel execution details | Summary of which GPU kernels were launched and their statistics | +| `trunc_kernel_details` | Truncated kernel details | Shortened version of `kernel_details_summary` for readability | +| `Percentage (%)` | Percentage of total GPU time | `(total_direct_kernel_time_sum / sum of all kernel time) * 100` | +| `Cumulative Percentage (%)` | Running cumulative percentage | Sum of percentages from top to current row | +| `is_recompute` | *(Optional)* Whether this operation is part of activation recomputation | `True` or `False`. Only present when `--detect_recompute` is enabled. | + +What makes arguments "unique": + +- Same operation with different input shapes gives different rows (for example, `aten::addmm` with (1024, 512) vs (2048, 1024)). +- Same operation with different data types gives different rows (for example, BF16 vs FP32). +- Same operation with different strides gives different rows (contiguous vs transposed). +- Each unique combination of (name + Input Dims + Input type + Input Strides + Concrete Inputs) gets one row. + +What this sheet shows: + +- Statistics for each unique operation+args combination (count, mean, median, std, min, max). +- Which specific input patterns are slow vs fast. +- Performance variation across different calls to the same operation. + +Use this sheet when you're asking, "I see `aten::addmm` is expensive. Is it slow for all input shapes, or just specific ones? Are there outliers?" + +#### Understanding kernel_details_summary + +`kernel_details_summary` provides aggregated kernel statistics across all occurrences of an operation+args combination. It shows which kernels are consistently launched and their performance distribution. All durations are in microseconds (µs). + +Example for `aten::addmm` (aggregated across 150 occurrences): + +```python +[ + {'kernel_name': 'elementwise_kernel', 'count': 150, 'mean_duration_us': 2.3, 'std_dev_duration_us': 0.1}, + {'kernel_name': 'Cijk_Alik_Bljk_BBS_BH...', 'count': 150, 'mean_duration_us': 45.2, 'std_dev_duration_us': 3.5} +] + +# Analysis: addmm = elementwise (2.3 µs) + GEMM (45.2 µs) across 150 calls +# → GEMM is bottleneck, with moderate variance (std_dev=3.5) → check for outliers with ex_UID +# → elementwise is consistent (std_dev=0.1) +``` + +Use this to identify: + +- Bottlenecks: which kernel in a multi-kernel operation dominates time. For example, in the GPT-3 XL analysis, `aten::addmm` consistently showed a Tensile GEMM kernel taking 97 µs (mean) while setup kernels took <5 µs each. +- Consistency: low `std_dev` = stable, high `std_dev` = investigate outliers with `ex_UID`. For example, if `std_dev_duration_us` is 15.3 when `mean_duration_us` is 120.5, that's 12.7% variance, worth investigating. +- Backend recipe: which kernels are always launched together (for example, elementwise + GEMM for addmm). +- Call count validation: verify all occurrences launch the same kernels (same `count`). + +```{note} +`trunc_kernel_details` provides a shortened version for spreadsheet readability. +``` + +#### Deep-dive with ex_UID + +The `ex_UID` column provides a UID of one example event with this operation+arguments combination. You can use it to access the actual event object for deeper analysis: + +```python +# Get the event object +event = tree.get_UID2event(ex_UID) + +# Analyze call stack and context: +parent = tree.get_parent_event(event) # Get parent operation +children = tree.get_children_events(event) # Get child operations +gpu_events = tree.get_gpu_events(event) # Get launched GPU kernels +tree.traverse_parents_and_print(event) # See full call stack above this +tree.traverse_subtree_and_print(event) # See full call stack below this + +# Replay the operation for benchmarking: +from TraceLens import EventReplayer +replayer = EventReplayer(event, device='cuda') +replayer.replay() # Replays the exact operation with same inputs/args +``` + +This is useful when you want to investigate a specific slow case in detail or benchmark it in isolation. See [Replay a traced operation](../how-to/event-replay.md) for more on replaying operations. + +Replaying from the perf report (no trace required): you can also replay operations directly from the Excel report without needing the original trace file: + +```python +import pandas as pd +import ast +from TraceLens import EventReplayer + +# Read ops from perf report +df = pd.read_excel('perf_report.xlsx', sheet_name='ops_unique_args') + +# Convert row to event format +def row_to_event(row): + return { + 'name': row['name'], + 'args': { + 'Input Dims': ast.literal_eval(row['Input Dims']), + 'Input Strides': ast.literal_eval(row['Input Strides']), + 'Input type': ast.literal_eval(row['Input type']), + 'Concrete Inputs': ast.literal_eval(row['Concrete Inputs']), + } + } + +# Replay an operation +row = df.iloc[0] # or filter by name/args +event = row_to_event(row) +replayer = EventReplayer(event, device='cuda') +replayer.replay() + +# Get standalone repro artifacts +repro_info = replayer.get_repro_info() # Returns serializable JSON +``` + +This is particularly useful for creating standalone reproducers or benchmarking specific operations without the full model or trace. See [`examples/event_replayer_example.ipynb`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/event_replayer_example.ipynb) for complete examples including batched replay. + +Generated by `TreePerfAnalyzer.get_df_kernel_launchers_unique_args()`. + +## Performance metrics sheets (roofline analysis) + +For certain operation categories (GEMM, CONV, SDPA, UnaryElementwise, BinaryElementwise), TraceLens generates additional sheets with roofline model metrics. These sheets help you understand how efficiently operations are using the GPU's computational and memory bandwidth capabilities. + +```{note} +While hardware counter profilers like `rocprof compute` and `nsight compute` reveal what the GPU actually executed, including effects of padding, redundant memory movement, and cache behavior, TraceLens focuses on the useful work dictated by operator semantics. Used together, these two perspectives provide a richer picture: hardware counters expose low-level execution characteristics, while TraceLens reveals the efficiency of the computation in context. +``` + +### Understanding the metrics pipeline + +The performance metrics are calculated through a 4-step pipeline: + +``` +Trace Event → Parameter Extraction → Static Metrics + Runtime → Performance Metrics + ↑ ↑ + (Compute Model) (GPU Time) +``` + +#### Step 1: event from trace + +Starting with an operation event from the trace JSON: + +```json +{ + 'name': 'aten::addmm', + 'args': { + 'Input Dims': [[6144], [40960, 1536], [1536, 6144], [], []], + 'Input type': ['c10::BFloat16', 'c10::BFloat16', 'c10::BFloat16', 'Scalar', 'Scalar'], + 'Input Strides': [...] + } +} +``` + +#### Step 2: parameter extraction + +TraceLens uses operation-specific performance models to extract relevant parameters from the event's `args` field: + +```python +# For aten::addmm: A @ B + C where A is (M, K), B is (K, N), C is (M, N) +params = { + 'M': 40960, # From Input Dims[1][0] + 'N': 6144, # From Input Dims[2][1] + 'K': 1536, # From Input Dims[1][1] + 'bias': True, # C tensor present + 'dtype': 'c10::BFloat16' # From Input type +} +``` + +Each operation type (GEMM, Conv, SDPA, and so on) has its own `get_param_details()` method that knows how to extract the relevant parameters from that operation's argument structure. See [Operation parameters reference](#operation-parameters-reference) for details on what gets extracted for each operation type. + +#### Step 3: parallel calculation + +Once parameters are extracted, two pieces of information are computed in parallel. + +Static operation metrics (via the compute model): the performance model calculates the theoretical work based on the operation type: + +```python +# GEMM compute model +GFLOPS = (2 * M * N * K) / 1e9 # 773.35 GFLOPS +Bytes_moved = (M*K + K*N + M*N) * bytes_per_elem # 618.01 MB (for BF16: 2 bytes/elem) +FLOPS_per_Byte = GFLOPS * 1e9 / Bytes_moved # 1193.38 (arithmetic intensity) +``` + +These metrics are static; they depend only on the operation parameters, not on actual execution. + +Kernel time (from Trace2Tree): the actual GPU execution time is extracted by: + +1. Finding all GPU kernels launched by this operation (using Trace2Tree's hierarchical analysis). +2. Using `GPUEventAnalyser` to compute total busy time, accounting for kernel overlaps. + +```python +T = 1884 µs # Measured kernel execution time +``` + +#### Step 4: runtime performance metrics + +Finally, static work metrics are combined with measured time to produce runtime performance: + +```python +TFLOPS/s = GFLOPS / (T / 1e6) = 773.35 / (1884 / 1e6) = 410.48 TFLOPS/s +TB/s = (Bytes_moved / 1e12) / (T / 1e6) = 0.34 TB/s +``` + +These metrics tell you: + +- TFLOPS/s: how many trillion floating-point operations per second were achieved. +- TB/s: how many terabytes per second of memory bandwidth were utilized. + +### Why this matters: roofline analysis + +The combination of TFLOPS/s, TB/s, and FLOPS/Byte (arithmetic intensity) allows you to perform roofline analysis: + +- High FLOPS/Byte (compute-intensive): performance is limited by compute throughput (TFLOPS/s). + - Example: large GEMM operations (for example, 6144×2048 × 2048×8192). + - Goal: maximize TFLOPS/s (approach the GPU's peak compute). +- Low FLOPS/Byte (memory-intensive): performance is limited by memory bandwidth (TB/s). + - Example: element-wise operations, small GEMMs (for example, 2048×2048). + - Goal: maximize TB/s (approach the GPU's peak bandwidth). + +The roofline "knee point": the boundary between memory-bound and compute-bound is determined by the GPU's hardware characteristics: + +``` +Arithmetic Intensity Threshold = Peak FLOPS / Peak Bandwidth +``` + +For example: + +- MI325X: 1300 TFLOPS / 6000 GB/s = ~217 FLOPs/Byte +- H100: 1000 TFLOPS / 3350 GB/s = ~298 FLOPs/Byte +- MI300X: 1300 TFLOPS / 5300 GB/s = ~245 FLOPs/Byte + +Operations with FLOPs/Byte below this threshold are memory-bound; operations above are compute-bound. This is why small GEMMs (FLOPs/Byte ~8-50) are memory-bound, while large GEMMs (FLOPs/Byte ~100-300) are compute-bound on these GPUs. By comparing your achieved TFLOPS/s and TB/s against the GPU's theoretical peaks, you can identify optimization opportunities. + +Interpreting performance numbers: understanding what "good" performance looks like requires comparing against theoretical peaks and max-achievable performance: + +| GPU | Peak compute (BF16) | Peak memory BW | Example utilization | +|-----|---------------------|----------------|---------------------| +| MI325X | ~1.3 PFLOPS | ~6 TB/s | 500-800 TFLOPS/s = 38-62% of peak (typical for medium GEMMs) | +| H100 | ~1.0 PFLOPS | ~3.35 TB/s | 400-700 TFLOPS/s = 40-70% of peak | +| MI300X | ~1.3 PFLOPS | ~5.3 TB/s | Similar to MI325X | + +Understanding theoretical vs. real-world performance: TraceLens uses idealized assumptions that represent upper bounds on performance. The actual roofline model has two key differences from the theoretical one: + +1. Peak FLOPS: the theoretical peak represents hardware limits, but real-world applications typically achieve lower performance due to realistic constraints. AMD's Max-Achievable FLOPS (MAF) methodology provides more realistic performance targets. For details, see: + - [Understanding Peak and Max-Achievable FLOPS](https://rocm.blogs.amd.com/software-tools-optimization/Understanding_Peak_and_Max-Achievable_FLOPS/README.html) + - [Measuring Max-Achievable FLOPS (Part 2)](https://rocm.blogs.amd.com/software-tools-optimization/measuring-max-achievable-flops-part2/README.html) +2. Arithmetic intensity: TraceLens assumes 100% cache hit rate, so each memory location is accessed once from global memory, then perfectly cached. For example, in a GEMM, matrices A and B are counted only once even if elements are reused. In reality: + - Cache misses, evictions, and redundant loads cause actual memory movement to be higher. + - This means actual arithmetic intensity is lower than TraceLens calculates. + - Impact: an operation that TraceLens shows as compute-bound (high FLOPs/Byte) might actually be memory-bound in practice. + - This is a common practice in performance modeling (see [NVIDIA's approach](https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#math-mem-bounds)). + - For typical compute-bound operations (large GEMMs, SDPA, convolutions), this limitation has minimal practical impact since they remain compute-bound even with lower arithmetic intensity. + +```{note} +Hardware profilers like `rocprof compute` or `nsight compute` measure actual memory transactions and are needed to determine true arithmetic intensity and memory-boundedness. +``` + +Real examples (from GPT-3 XL on MI325X, BF16): + +``` +Operation TFLOPS/s TB/s FLOPs/Byte Interpretation +────────────────────────────────────────────────────────────────────────────── +mm(256×256 × 256×256) 45 2.1 ~8 Memory-bound, 35% BW utilization +mm(2048×6144 × 6144×2048) 531 0.61 ~138 Compute-bound, 40% efficiency +mm(6144×2048 × 2048×2048) 624 0.71 ~117 Compute-bound, 48% efficiency +addmm(6144×2048 × 2048×8192) 762 0.59 ~203 Compute-bound, 58% efficiency +``` + +Understanding compute-bound vs. memory-bound: + +- Memory-bound (FLOPs/Byte < ~50): small GEMM with arithmetic intensity of ~8. Performance is limited by memory bandwidth (~2.1 TB/s), not compute. To optimize, focus on improving memory access patterns. +- Compute-bound (FLOPs/Byte > ~100): large GEMMs with high arithmetic intensity. Performance is limited by compute throughput (531-762 TFLOPS/s). The 40% vs 58% efficiency difference reflects kernel optimization quality (tile sizes, wave occupancy), not the compute vs. memory boundedness. + +```{note} +Arithmetic intensity (FLOPs/Byte) determines whether an operation is compute-bound or memory-bound. The percentage of peak achieved indicates optimization quality within that constraint. See [NVIDIA's GEMM Performance Guide](https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html) for more on this distinction. +``` + +### What these sheets contain + +Performance metrics sheets are generated for operations with available performance models: + +- GEMM: matrix multiply operations (addmm, mm, bmm, baddbmm, and so on). +- CONV_fwd / CONV_bwd: convolution operations. +- SDPA_fwd / SDPA_bwd: scaled dot-product attention. +- UnaryElementwise / BinaryElementwise: element-wise operations. + +Each sheet contains: + +- All columns from `ops_unique_args` (operation name, arguments, occurrences, and so on). +- Static metrics: GFLOPS, Data Moved (MB), FLOPS/Byte (calculated once from parameters). +- Compute Spec: combined compute type and precision (for example, `matrix_bf16`, `vector_fp32`). This indicates: + - `matrix_*`: operations using matrix compute units (GEMM, CONV, SDPA). + - `vector_*`: operations using vector compute units (elementwise). + - Precision: `fp8`, `fp16`, `bf16`, `fp32`, `fp64`. +- Roofline metrics (when `--gpu_arch_json_path` is provided): Roofline Time (µs), Pct Roofline; compares achieved time to the theoretical roofline bound. +- Runtime metrics: Kernel Time (µs), TFLOPS/s, TB/s (statistics across all occurrences). + - `mean`, `median`: central tendency; use median if variance is high. + - `std_dev`: variability; high `std_dev` (>10% of mean) suggests inconsistent performance. + - `min`, `max`: range; a large spread can indicate outliers worth investigating. + +```{note} +These sheets contain all the columns from `ops_unique_args`, so you can replay operations from these sheets using the same approach described in [Understanding kernel_details_summary](#understanding-kernel_details_summary). Simply read the desired performance metrics sheet (for example, `GEMM`, `SDPA_fwd`) instead of `ops_unique_args`. +``` + +Generated by `TreePerfAnalyzer.build_df_perf_metrics()` and `TreePerfAnalyzer.summarize_df_perf_metrics()`. + +### Operation parameters reference + +TraceLens extracts operation-specific parameters from trace events to calculate theoretical FLOPs and memory traffic. Each operation type has its own parameter extraction logic based on what information is needed for its performance model. + +Common sources: + +- `Input Dims`: tensor shapes. +- `Input type`: data types. +- `Concrete Inputs`: scalar arguments (stride, padding, and so on). + +Data type sizes: `fp64` = 8 bytes, `fp32` = 4 bytes, `bf16/fp16` = 2 bytes, `fp8` = 1 byte. + +#### GEMM (matrix multiply) + +| Parameter | Meaning | Used in | +|-----------|---------|---------| +| `M` | Number of rows in first matrix | All GEMM ops | +| `N` | Number of columns in second matrix | All GEMM ops | +| `K` | Inner dimension (A cols = B rows) | All GEMM ops | +| `B` | Batch size | bmm, baddbmm | +| `bias` | Whether bias is added | addmm, baddbmm | +| `dtype` | Data type (BFloat16, Float32, and so on) | All GEMM ops | + +#### Convolution + +| Parameter | Meaning | +|-----------|---------| +| `input_shape` | Input tensor shape (N, C_in, H, W, ...) | +| `filter_shape` | Filter/weight shape (C_out, C_in/groups, kH, kW, ...) | +| `stride` | Stride for convolution (for example, (1, 1)) | +| `padding` | Padding applied (for example, (0, 0)) | +| `dilation` | Dilation factor | +| `groups` | Number of groups for grouped convolution | +| `bias` | Whether bias is present | +| `dtype` | Data type | + +#### SDPA (scaled dot-product attention) + +| Parameter | Meaning | +|-----------|---------| +| `B` | Batch size | +| `N_Q` | Query sequence length | +| `N_KV` | Key/Value sequence length | +| `H_Q` | Number of query attention heads | +| `H_KV` | Number of key/value heads (for GQA/MQA) | +| `d_h_qk` | Head dimension for Q and K | +| `d_h_v` | Head dimension for V | +| `causal` | Whether causal masking is applied | +| `dropout` | Dropout probability | +| `dtype` | Data type | + +```{note} +Different attention implementations use different tensor layouts (BHND vs BNHD). +``` + +#### Element-wise operations + +| Parameter | Meaning | +|-----------|---------| +| `shape` | Tensor shape | +| `dtype` | Data type | + +For binary ops (add, mul, and so on), two shapes are extracted and broadcasting is handled automatically. + +For the specific formulas used to calculate FLOPs and memory traffic from these parameters, see the performance model implementations in [`TraceLens/PerfModel/perf_model.py`](https://github.com/AMD-AGI/TraceLens/blob/main/TraceLens/PerfModel/perf_model.py). + +Extending with custom operations: performance metrics sheets are only generated for operations that have a registered performance model. If your workload includes custom operations (for example, from Megatron, vLLM, or other libraries), you can extend TraceLens by: + +1. Creating a performance model class for your operation (inherit from `GEMM`, `CONV`, `SDPA`, and so on). +2. Implementing `get_param_details()`, `flops()`, and `bytes()` methods. +3. Passing an extension file via `--extension_file` when generating the report. + +See [`examples/example_megatron_extension.py`](https://github.com/AMD-AGI/TraceLens/blob/main/examples/example_megatron_extension.py) for a complete example of extending TraceLens with custom Megatron operations. + +```{note} +If your custom operation is frequently used across multiple projects, it can be added as a native operation in TraceLens. Open an issue or reach out to discuss integration. +``` + +## Collective communication analysis + +The `coll_analysis` sheet analyzes NCCL collective communication operations from a single rank's perspective. It shows which collectives are taking the most time on this rank. + +```{note} +This is single-rank analysis. For multi-rank synchronization analysis (skew, stragglers), see [Generate a multi-rank collective report](../how-to/collective-report.md). +``` + +This sheet provides aggregated statistics for collective operations seen by this rank, grouped by collective type, process group, data type, and message size. + +### Collective identification + +| Column | Description | +|--------|-------------| +| `rank` | Rank ID (single rank from this trace) | +| `Process Group Name` | Process group identifier | +| `Process Group Ranks` | List of ranks in this process group (for example, `[0, 1, 2, 3, 4, 5, 6, 7]`) | +| `Collective name` | Type of collective operation (for example, `allreduce`, `allgather`, `reduce_scatter`) | +| `Group size` | Number of ranks in the process group | +| `dtype` | Data type (for example, `Float`, `BFloat16`) | +| `In msg nelems` | Number of elements in input message | +| `Out msg nelems` | Number of elements in output message | +| `In split size` | Input split configuration (for split collectives) | +| `Out split size` | Output split configuration (for split collectives) | +| `stream` | GPU stream ID where the collective executes | + +### Message sizes + +| Column | Description | +|--------|-------------| +| `In msg size (MB)_first` | Input message size in megabytes | +| `Out msg size (MB)_first` | Output message size in megabytes | + +### Duration statistics + +All times are in microseconds. + +| Column | Description | +|--------|-------------| +| `dur_sum` | Total duration across all occurrences | +| `dur_mean` | Mean duration per occurrence | +| `dur_std` | Standard deviation of duration | +| `dur_min` | Minimum duration | +| `dur_max` | Maximum duration | +| `operation_count` | Number of times this collective appeared | + +### Understanding the metrics + +Duration is the time this rank spends in the collective operation (in microseconds). It includes: + +- Synchronization delay waiting for other ranks. +- Actual data transfer time. + +```{note} +From a single rank's perspective, you can't determine whether time is spent waiting for other ranks (sync) or in actual communication. For multi-rank analysis with skew detection and straggler identification, use NCCL Analyzer with traces from all ranks. See [Generate a multi-rank collective report](../how-to/collective-report.md). +``` + +### Interpreting results + +- High total duration: sort by `dur_sum` to find the most time-consuming collective operations on this rank. +- High variance (large `dur_std` or difference between `dur_min` and `dur_max`): + - Indicates inconsistent performance. + - Can suggest network contention or varying degrees of synchronization delay. + - Consider investigating with multi-rank traces to identify stragglers. +- Frequent operations (high `operation_count`): + - Even if individual operations are fast, high frequency can add up (`dur_sum` accounts for this). + - Check if collectives can be batched or overlapped with computation. + +Example analysis: + +``` +If you see: +- allreduce with dur_mean=3596 µs (3.6 ms), operation_count=2 +- But dur_std=35.5 µs (low variance) +→ Collective is consistent and moderately fast + +If you see: +- allreduce with dur_mean=192.6 µs, dur_std=313.7 µs, dur_max=1019.4 µs +→ High variance! One occurrence took 1ms while another took 60µs +→ Investigate with multi-rank trace to find if this rank or another is the straggler +``` + +Generated by `NcclAnalyser.build_df_summary_long()`. Enable with the `--disable_coll_analysis` flag (enabled by default). + +## Additional sheets + +Depending on the command-line arguments used when generating the report, additional analysis sheets can be present: + +- `kernel_summary`: per-kernel statistics aggregated by kernel name (enabled with `--kernel_summary`). +- `short_kernel_histogram`, `short_kernels_summary`: analysis of very short kernels (enabled with `--short_kernel_study`). +- `unified_perf_summary`: unified perf metrics for all ops with perf models or leaf ops that launch GPU kernels. Includes GFLOPS, TFLOPS/s, Data Moved, FLOPS/Byte, TB/s metrics aggregated by unique args. When `--detect_recompute` is enabled, an `is_recompute` column is added to split rows by recompute status. + +## Common analysis workflows + +A typical top-down analysis flow: + +1. Start with `gpu_timeline`: get a high-level time breakdown (computation, communication, idle). + - High idle time points to a CPU bottleneck or kernel launch issues. + - High exposed communication points to poor overlap with computation. +2. Identify bottleneck categories (`ops_summary_by_category`): which operation types dominate? GEMM, convolution, attention, elementwise, and so on. +3. Find expensive operations (`ops_summary`): which specific operations take the most time? +4. Analyze variants (`ops_unique_args`): which input shapes/arguments are slow? + - Compare performance across different dimensions. + - Check for stride or alignment issues. +5. Deep dive with roofline (GEMM, SDPA_fwd, and so on): assess efficiency vs. hardware peaks. +6. Check collectives (`coll_analysis`): analyze communication costs in distributed training. +7. Replay operations: reproduce and benchmark specific cases in isolation. + +## Related topics + +- [Generate a PyTorch performance report](../how-to/generate-perf-report-pytorch.md) +- [Generate a JAX performance report](../how-to/generate-perf-report-jax.md) +- [Generate a rocprofv3 performance report](../how-to/generate-perf-report-rocprof.md) +- [Replay a traced operation](../how-to/event-replay.md) +- [API reference](../reference/api-reference.md) diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index 2e99fcdf0..099ebb965 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -25,26 +25,64 @@ subtrees: title: Generate reports entries: - file: how-to/generate-perf-report-pytorch - title: PyTorch performance report + title: PyTorch performance report - file: how-to/generate-perf-report-pytorch-inference - title: PyTorch inference performance report + title: PyTorch inference performance report - file: how-to/generate-perf-report-jax - title: JAX performance report + title: JAX performance report - file: how-to/generate-perf-report-rocprof - title: rocprof performance report + title: rocprof performance report - file: how-to/collective-report - title: Collective-communication report + title: Collective-communication report + - file: how-to/compare-perf-reports + title: Compare performance reports - file: how-to/compare-traces title: Compare two traces - file: how-to/event-replay title: Replay a single operation - file: how-to/trace-fusion title: Fuse multi-rank traces + - file: how-to/sdk-analysis + title: Analyze traces with the SDK + - file: how-to/perf-model-without-trace + title: Model op performance without a trace + - file: how-to/nccl-analysis + title: Analyze collective communication + - file: how-to/origami-integration + title: Estimate kernel times with Origami + - file: how-to/agent + title: Agentic performance analysis with the TraceLens Agent + +- caption: Concepts + entries: + - file: conceptual/trace2tree + title: The Trace2Tree data model + - file: conceptual/torch-profiling-analysis + title: Understanding PyTorch traces + - file: conceptual/shape-metadata + title: Tensor shape metadata + - file: conceptual/gemm-analysis + title: GEMM analysis + - file: conceptual/inference-analysis + title: Inference performance analysis + - file: conceptual/triton-perf-model-walkthrough + title: Triton kernel performance model + +- caption: Tutorials + entries: + - file: tutorials/torch-profiling + title: PyTorch profiling walkthrough + - file: tutorials/distributed-profiling + title: Distributed profiling walkthrough + - file: tutorials/anomaly-detection + title: Anomaly detection walkthrough - caption: Reference entries: - file: reference/api-reference title: API reference + - file: reference/perf-report-columns + title: Performance report columns - caption: About entries: diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index f40e2d3af..c831419df 100644 --- a/docs/sphinx/requirements.txt +++ b/docs/sphinx/requirements.txt @@ -1 +1 @@ -rocm-docs-core==1.22.0 +rocm-docs-core==1.32.0 diff --git a/docs_original/conceptual/anomaly_detection_guide.ipynb b/docs/tutorials/anomaly-detection.ipynb similarity index 100% rename from docs_original/conceptual/anomaly_detection_guide.ipynb rename to docs/tutorials/anomaly-detection.ipynb diff --git a/docs_original/conceptual/distributed_profiling_guide.ipynb b/docs/tutorials/distributed-profiling.ipynb similarity index 100% rename from docs_original/conceptual/distributed_profiling_guide.ipynb rename to docs/tutorials/distributed-profiling.ipynb diff --git a/docs_original/conceptual/torch_profiling_guide.ipynb b/docs/tutorials/torch-profiling.ipynb similarity index 100% rename from docs_original/conceptual/torch_profiling_guide.ipynb rename to docs/tutorials/torch-profiling.ipynb diff --git a/docs_original/Agentic_Mode.md b/docs_original/Agentic_Mode.md deleted file mode 120000 index 9fa51d05a..000000000 --- a/docs_original/Agentic_Mode.md +++ /dev/null @@ -1 +0,0 @@ -../TraceLens/Agent/Analysis/README.md \ No newline at end of file diff --git a/docs_original/EventReplay.md b/docs_original/EventReplay.md deleted file mode 100644 index c8db3514e..000000000 --- a/docs_original/EventReplay.md +++ /dev/null @@ -1,152 +0,0 @@ - - -# Event Replay - -Optimizing GPU performance in deep learning requires isolating and benchmarking individual operations to identify bottlenecks. However, reproducing operations directly from complex model code or large profiles can be cumbersome. - -Event Replay is a Python-based tool within TraceLens that extracts and replays almost arbitrary PyTorch operations using minimal, portable Intermediate Representation (IR). It enables users to easily reproduce, analyze, and benchmark specific operators independently from the original model execution, streamlining performance optimization workflows. - ---- - -## Key Features - -- **Generic Operator Replay**: Reconstructs and benchmarks any PyTorch operator from profile data, including convolutions, GEMMs, reductions, element-wise operations, and more. -- **Minimalistic IR**: Extracts essential operator attributes (tensor shapes, strides, dtypes, and other arguments) into a lightweight, portable JSON-based IR. -- **Portable Artifacts**: Enables sharing standalone artifacts (JSON IR and scripts) with teammates or upstream repositories without requiring access to the model or TraceLens. - ---- - - -## Quick Start - -### Example: Replay a Single Event - -```python -from TraceLens import TreePerfAnalyzer, EventReplayer - -# Load profile and get event -perf_analyzer = TreePerfAnalyzer.from_file('/path/to/profile.json') -uid = 12345 # Replace with actual UID of interest -event = perf_analyzer.tree.get_UID2event(uid) - -# Initialize and replay -replayer = EventReplayer(event, device='cuda') -replayer.replay() -``` - ---- - -## Batch Replay and Benchmark - -### Extract Operator IR from TraceLens Profiles - -```python -import json - -# Extract replay IR for events of interest -repro_data = [EventReplayer(event, lazy=True).get_repro_info() for event in events_of_interest] - -with open('event_replay_ir.json', 'w') as f: - json.dump(repro_data, f, indent=4) -``` - -```bash -python batched_replay.py event_replay_ir.json -``` - -#### Example Output - -``` -[7/11] Replaying: aten::convolution - Reconstructing arguments for 'aten::convolution'... - Positional Args: - input Tensor: {'shape': [20, 128, 28, 28], 'dtype': 'c10::BFloat16', 'strides': [100352, 784, 28, 1]} - weight Tensor: {'shape': [256, 128, 3, 3], 'dtype': 'c10::BFloat16', 'strides': [1152, 9, 3, 1]} - bias Tensor?: None - stride SymInt[]: [2, 2] - padding SymInt[]: [1, 1] - dilation SymInt[]: [1, 1] - transposed bool: False - output_padding SymInt[]: [0, 0] - groups SymInt: 1 - Keyword Args: - Average time taken: 100.38 microseconds - Successfully executed aten::convolution. - Result: Tensor(shape=torch.Size([20, 256, 14, 14]), dtype=torch.bfloat16, device=cuda:0) - -[8/11] Replaying: aten::convolution - Reconstructing arguments for 'aten::convolution'... - Positional Args: - input Tensor: {'shape': [20, 256, 14, 14], 'dtype': 'c10::BFloat16', 'strides': [50176, 196, 14, 1]} - weight Tensor: {'shape': [512, 256, 3, 3], 'dtype': 'c10::BFloat16', 'strides': [2304, 9, 3, 1]} - bias Tensor?: None - stride SymInt[]: [2, 2] - padding SymInt[]: [1, 1] - dilation SymInt[]: [1, 1] - transposed bool: False - output_padding SymInt[]: [0, 0] - groups SymInt: 1 - Keyword Args: - Average time taken: 92.83 microseconds - Successfully executed aten::convolution. - Result: Tensor(shape=torch.Size([20, 512, 7, 7]), dtype=torch.bfloat16, device=cuda:0) -... ---- Replay Summary --- -Total operations in file: 11 -Attempted replays: 11 -Successful replays: 11 -Errors encountered: 0 - -``` -------------------- -### Creating Standalone Replay Artifacts - -You can optionally package the extracted replay IR and scripts into a standalone zip file for easy sharing and reproduction, independent of the original model code or TraceLens repository. - -Artifacts included: -- `event_replay_ir.json`: Serialized operator replay instructions. -- `utils.py`: Tensor creation and helper utilities. -- `batched_replay.py`: Script to batch replay and benchmark operations. -- `batched_replay_readme.md`: Instructions for running the replay. - -Example packaging code: - -```python -import zipfile -import os -from TraceLens.EventReplay import utils as tl_utils -from TraceLens.EventReplay import batched_replay - -files = [ - OUTPUT_REPRO_FILE, - tl_utils.__file__, - batched_replay.__file__, - batched_replay.__file__.replace('batched_replay.py', 'batched_replay_readme.md') -] - -zip_file_path = '/path/to/replay_code.zip' -with zipfile.ZipFile(zip_file_path, 'w') as zipf: - for file in files: - zipf.write(file, arcname=os.path.basename(file)) - -print(f"Created zip file: {zip_file_path}") -``` ---- - -## Use Cases - -- **Performance Debugging**: Quickly isolate and reproduce performance issues from large models. -- **Regression Testing**: Automate benchmarks to detect performance regressions at the operator level. -- **Kernel Development**: Extract minimal reproducers for GPU kernel optimization and debugging. -- **Numerical Validation**: Evaluate numerical correctness and stability of isolated operations across hardware. -- **Hardware Counter Profiling**: Use with hardware counters to analyze performance bottlenecks in specific operations. - ---- - -## Notes - -- Event Replay uses randomized data based on extracted tensor shapes; thus, replay timings approximate real-world performance. diff --git a/docs_original/Inference_analysis.md b/docs_original/Inference_analysis.md deleted file mode 100644 index 52d6c1db9..000000000 --- a/docs_original/Inference_analysis.md +++ /dev/null @@ -1,728 +0,0 @@ - - -# 🚀 TraceLens Inference Performance Analysis - -TraceLens now provides comprehensive support for inference use cases, with a focus on inference serving optimization. This documentation covers: - -- 📋 **Overview** - New features for inference trace analysis -- 🔧 **Trace Collection** - Methodologies and setup -- 📊 **Analysis Tools** - Available workflows and usage -- 🗺️ **Roadmap** - Upcoming improvements - -## ✨ Key Features - - -| Feature | Description | -| ---------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | -| **Agentic Analysis** | Agentic workflows for single-trace analysis for performance improvement recommendations | -| **Graph Execution Analysis** | Merge graph capture and graph reply traecfiles for augmenting graph execution tracefile with callstack and shape information | -| **TraceDiff** | Extended to support inference traces with Lowest Common Ancestor (LCA) analysis for kernel correlation across platforms | -| **Roofline Analysis** | Custom roofline models for key inference operations (fused MoE, unified attention) with prefill/decode request annotations. | -| **Trace Splitting** | Splitting of large tracefiles into steady-state regions, per-iteration traces, and phase-specific analyses | - - -## Supported Frameworks and Execution Modes - -TraceLens features for inference analysis have been primarily tested with vLLM and SGLang. Here is a summary of the different execution modes and supported features. - - -| Mode | Shapes/Roofline analysis | Agent Analysis | Limitations | -| ----------------------------------- | ------------------------ | ---------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Eager only | Yes | Supported; proposed patches are recommended to include roofline information for attention operations | Eager mode execution may employ different compilation strategies, which can result in differences in kernels and fusions compared to graph execution mode. | -| Graph execution only | Non‑graph kernels | Limited | Categorization, call stacks, and shapes are available only for attention kernels if full_and_piecewise mode is used | -| Graph execution + eager mode trace | Limited | Limited | Kernel categorization might not be as accurate as eager or graph+capture | -| Graph execution + Graph capture$^1$ | Yes | Yes (patches required) | | - - - $^1$ Graph mode analysis using graph capture and graph replay traces is supported for vLLM and SGLang (proposed patches required). - -## 📖 Quickstart Guide - -### Step 1: Installation - -Install TraceLens from GitHub (requires AMD-AGI organization access): - -```bash -pip install git+https://github.com/AMD-AGI/TraceLens.git -``` - -### Step 2: Trace Collection - -There are two ways to collect inference traces for TraceLens analysis: - -- **Option A — Use the Profiling agent** *(recommended)*: an agentic skill bundled with TraceLens that drives a Magpie LLM inference benchmark, applies the Docker / framework patches for you, tunes the profiler window, runs the benchmark, and splits the resulting trace into analysis-ready windows. -- **Option B — Manual profiling**: build (or patch) your own image, configure the profiler yourself, and run the benchmark by hand. Use this if you are not using Magpie or want full control over every step. - -#### Option A: Use the Profiling Agent (recommended) - -Follow [`TraceLens/Agent/Profiling/README.md`](../TraceLens/Agent/Profiling/README.md). The skill handles trace collection and splitting end-to-end; you can skip the rest of Step 2 and Step 3. - -#### Option B: Manual Profiling - -The remainder of Step 2 covers manual trace collection — building or patching the inference framework, configuring the profiler, and running the benchmark yourself. - -##### Build a Docker image using the [provided scripts](../examples/custom_workflows/inference_analysis/) (recommended for manual flow) - -###### vLLM Script - -A unified build script is provided that supports multiple vLLM versions. It takes a version tag (`v14`, `v15`, `v16`, `v17`, `v18`, `v19`, `v20`, `v21`, `v22`, `v23`, or `v24`) as the first argument, followed by the path to your local TraceLens clone and any standard `docker build` flags. The script selects the correct base image and patch file automatically. - - -| Version | Base Image | vLLM Version | Patch File | -| ------- | ------------------------------------------------------------- | ------------ | --------------------------- | -| `v14` | `rocm/vllm-dev:preview_releases_rocm_v0.14.0_20260120` | v0.14.0 | `config_vllm_v0.14.0.patch` | -| `v15` | `rocm/vllm-dev:preview_releases_rocm_v0.15.0_20260130` | v0.15.0 | `config_vllm_v0.15.0.patch` | -| `v16` | `rocm/vllm-dev:preview_rocm70_releases_rocm_v0.16.0_20260223` | v0.16.0 | `config_vllm_v0.16.0.patch` | -| `v17` | `vllm/vllm-openai-rocm:v0.17.0` | v0.17.0 | `config_vllm_v0.17.0.patch` | -| `v18` | `vllm/vllm-openai-rocm:v0.18.0` | v0.18.0 | `config_vllm_v0.18.0.patch` | -| `v19` | `vllm/vllm-openai-rocm:v0.19.0` | v0.19.0 | `config_vllm_v0.19.0.patch` | -| `v20` | `rocm/vllm-dev:preview_v0.20.0_20260429` | v0.20.0 | `config_vllm_v0.20.0.patch` | -| `v21` | `vllm/vllm-openai-rocm:v0.21.0` | v0.21.0 | `config_vllm_v0.21.0.patch` | -| `v22` | `vllm/vllm-openai-rocm:v0.22.0` | v0.22.0 | `config_vllm_v0.22.0.patch` | -| `v23` | `vllm/vllm-openai-rocm:v0.23.0` | v0.23.0 | `config_vllm_v0.23.0.patch` | -| `v24` | `vllm/vllm-openai-rocm:v0.24.0` | v0.24.0 | `config_vllm_v0.24.0.patch` | - - -```bash -bash examples/custom_workflows/inference_analysis/build_docker_vllm.sh \ - v16 \ - /path/to/TraceLens \ - -t tracelens-vllm -``` - -To use a custom base Docker image instead of the default for the selected version, pass `--base-image`: - -```bash -bash examples/custom_workflows/inference_analysis/build_docker_vllm.sh \ - v18 \ - /path/to/TraceLens \ - --base-image my-registry/vllm:nightly \ - -t tracelens-vllm:custom -``` - -Then create a container from the image. - -###### SGLang Script - -The build script for SGLang supports SGLang 0.5.9 and 0.5.11. It takes the path to the local TraceLens clone, the SGLang version (`--sglang-version`, default 0.5.9), and the GPU type (`--gpu-type`, default mi350). MI300 and MI350/MI355 are supported. - - -| SGLang Version | GPU Type | Base Image | -| -------------- | -------- | --------------------------------------- | -| `0.5.9` | MI300 | `lmsysorg/sglang:v0.5.9-rocm700-mi30x` | -| `0.5.9` | MI350/MI355 | `lmsysorg/sglang:v0.5.9-rocm700-mi35x` | -| `0.5.11` | MI300 | `lmsysorg/sglang:v0.5.11-rocm720-mi30x` | -| `0.5.11` | MI350/MI355 | `lmsysorg/sglang:v0.5.11-rocm720-mi35x` | -| `0.5.12` | MI300 | `lmsysorg/sglang:v0.5.12-rocm720-mi30x` | -| `0.5.12` | MI350/MI355 | `lmsysorg/sglang:v0.5.12-rocm720-mi35x` | - - -```bash -bash examples/custom_workflows/inference_analysis/build_docker_sglang.sh \ - /path/to/TraceLens \ - --sglang-version 0.5.11 \ - --gpu-type mi350 \ - -t tracelens-sglang -``` - -Then create a container from the image. - -##### Apply framework patches manually (alternative to building a new image) - -If you prefer to patch an existing environment instead of building a new image, apply patches to your inference framework to: - -- Add custom annotations with request packing information (See [roofline conceptual details](#roofline-analysis)) -- Capture graph mode execution phases for augmentation by TraceLens - -**Steps:** - -1. **Locate your inference engine:** - For vLLM: - ```bash - python -c "import vllm; import os; print(os.path.dirname(vllm.__file__))" - ``` - For SGLang: - ```bash - python -c "import sglang; import os; print(os.path.dirname(sglang.__file__))" - ``` - -2. **Find and apply the relevant patch:** - - Select by framework and version - - Apply: `cd /path/to/framework/../ && git apply /path/to/patchfile` - vLLM patches are in [vllm_roofline_patches](../examples/custom_workflows/inference_analysis/vllm_patches) - SGLang patches are in [sglang_roofline_patches](../examples/custom_workflows/inference_analysis/sglang_roofline_patches/) - -##### Collection Parameters - -- **Eager or Graph Execution Steady-State Window:** Large tracefiles are expected. Most inference serving benchmarks use `NUM_PROMPTS = 10 × CONC` with OSL sampling ratio R. We recommend tracing `(((R+1)/2) * 5 * OSL) ± (16 * OSL / CONC)` execution steps (which represents peak concurrency with prefill-decode mix). See [steady-state region identification](#steady-state-region-and-trace-splitting) for more details. User might need to increase the timeout limit in certain inference frameworks to allow storing the trace in the middle of the execution (e.g., VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200 for vLLM). -- **Graph Capture Mode:** The recommended patchfile will trace the graph capture phase and store corresponding tracefiles. -- **Profiler Setup:** Enable CPU-side call-stack and shape capture. For example, vLLM supports `profiler-config.torch_profiler_record_shapes` and `profiler-config.torch_profiler_with_stack`. - -##### Trace collection options - -###### vLLM - -The `config_vllm_v*.patch` patches add two `ProfilerConfig` flags that control graph-capture profiling and trace annotation detail. These patches are available for v0.14–v0.24. Pass the flags as server arguments: - - -| Flag | Type | Default | Description | -| -------------------------------------------------- | ------ | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `--profiler-config.capture_torch_profiler_dir DIR` | `str` | `""` | Absolute path to a directory where a PyTorch profiler trace of the CUDA graph capture phase will be saved (rank 0 only). Requires `--profiler-config.profiler torch`. Leave empty to disable graph capture profiling. | -| `--profiler-config.detailed_trace_annotation` | `bool` | `False` | When `True`, execution-step annotations include roofline metrics (`sk`, `sqsq`, `sqsk`) for both context and generation requests. When `False`, annotations record only request counts and token counts. Enable this for full roofline analysis; leave disabled for lighter-weight traces. | - - -**Example** — enable both flags alongside a standard steady-state window profile: - -```bash ---profiler-config.profiler torch \ ---profiler-config.torch_profiler_dir /workspace/torch_trace \ ---profiler-config.capture_torch_profiler_dir /workspace/torch_trace/capture_traces \ ---profiler-config.detailed_trace_annotation True \ ---profiler-config.delay_iterations 5402 \ ---profiler-config.max_iterations 256 \ ---profiler-config.ignore_frontend True -``` - -> **Note:** `capture_torch_profiler_dir` is only available when `--profiler-config.profiler torch` is set. The capture trace is written once at server startup during CUDA graph construction; the steady-state replay trace is written to `torch_profiler_dir` during the benchmark run. Pass both paths to `generate_perf_report_pytorch_inference` via `--capture_folder` and `--profile_json_path` respectively. - -###### SGLang - -1. While doing the profiling of the execution step, pass the parameter `shape_discovery=True` in the profile request to enable shape discovery and registration for operations which are not covered in default SGLang profile. -2. While doing the profiling of the execution step, pass the parameter `roofline_annotations=True` in the profile request to annotate trace with more detailed information useful for roofline annotations. -3. To profile the graph capture phase, while server startup provide the `--enable-profile-cuda-graph` server argument. This will save a trace file per batch size but it misses shape information for some operations, to ensure more diverse coverage, provide the `--enable-shape-discovery-for-cuda-graph-profile` server argument. - -### Step 3: Trace Preparation (Optional) - -This optional step reads the collected trace and splits it into smaller trace files or execution‑phase‑specific trace files. - -Option 1: Find steady-state region of execution (highest concurrency) and separate prefill-decode and decode-only execution steps (supports vLLM v0.14–v0.24 and SGLang v0.5.9; using the patchfile is recommended). This is recommended if the tracefile is large and the user wants to extract a few representative steps automatically. - -```python -python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz -o ./steady_state_analysis \ - --find-steady-state --num-steps 256 -``` - -Output: A tracefile containing {num-steps} contiguous execution steps where close to maximum concurrency is observed, plus contiguous prefill-decode mix and decode-only steady-state tracefiles extracted from this window with no idle gaps between execution steps. - -**Refining steady-state window selection with `--CONC`, `--OSL`, and `--R`** - -By default, the mixed steady-state window is selected by matching the empirically observed prefill-decode to total-steps ratio of the trace. If the benchmark parameters are known, passing `--CONC`, `--OSL`, and `--R` lets the tool compute an *ideal* perfilldecodemix_steps/total_steps ratio analytically and use that to drive window selection instead: - - -| Argument | Type | Description | -| -------- | ------- | ----------------------------------------------------------------------------------------------------------------------------------- | -| `--CONC` | `int` | Expected peak concurrency (number of concurrent requests). A warning is printed if the observed trace peak differs from this value. | -| `--OSL` | `float` | Maximum output sequence length (decode tokens per request). Each request's OSL is sampled from `[R × OSL, OSL]`. | -| `--R` | `float` | OSL sampling range ratio in `[0, 1]`. `R=0` means all requests use exactly `OSL` tokens; `R=1` means OSL is uniform in `[0, OSL]`. | - - -When all three are provided, the tool derives: - -``` -ideal_prefilldecodemix_to_totalsteps_ratio = (CONC × 2) / (OSL × (1 + R)) -``` - -and uses this as the reference ratio for mixed-window selection, overriding the empirical estimate. `--num-steps` is also automatically raised to `ceil(1 / ideal_prefilldecodemix_to_totalsteps_ratio)` if it is too small to capture a representative decode/prefill-decode mix. - -Example — benchmark with CONC=32, OSL=1024, R=0.8: - -```python -python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz \ - -o ./steady_state_analysis \ - --find-steady-state --num-steps 256 \ - --CONC 32 --OSL 1024 --R 0.8 -``` - -Option 2: One tracefile per eager/graph execution step (supports vLLM v0.13 or higher, SGLang v0.5.9, and Atom 0.1.1). This is recommended if the user wants to perform analysis on an isolated execution step. - -```python -python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz -o ./output --store-single-iteration -``` - -Output: Single trace file per execution step. - -Option 3: Limit the search of steady-state region to a limited window - -```python -python -m TraceLens.TraceUtils.split_inference_trace_annotation trace.json.gz -o ./output --iterations 10:20 --find-steady-state --num-steps 256 \ - --CONC 32 --OSL 1024 --R 0.8 -``` - -### Step 4: Generate Performance Report - -Performance report generation is supported for both eager-mode and graph-mode (capture + replay) traces. - -**Eager or graph replay traces (no graph capture folder):** - -```bash -python -m TraceLens.Reporting.generate_perf_report_pytorch_inference \ - --profile_json_path /path/to/trace.json \ - --output_xlsx_path perf_report.xlsx \ - --group_by_parent_module \ - --group_by_kernel_count \ - --enable_pseudo_ops -``` - -**Graph replay traces augmented with graph capture traces:** - -When a `--capture_folder` is provided, the script automatically classifies graph capture traces (batch sizes, full vs. piecewise mode) and merges their call-stack and shape information into the graph replay tree before generating the report. - -```bash -python -m TraceLens.Reporting.generate_perf_report_pytorch_inference \ - --profile_json_path /path/to/graph/replay/trace.json \ - --capture_folder /path/to/capture/traces/folder \ - --output_xlsx_path perf_report.xlsx \ - --group_by_parent_module \ - --group_by_kernel_count \ - --enable_pseudo_ops -``` - -**Additional options:** - - -| Flag | Description | -| --------------------------- | ------------------------------------------------------------ | -| `--output_csvs_dir DIR` | Write per-sheet CSV files instead of a single Excel workbook | -| `--gpu_arch_json_path PATH` | Provide a GPU architecture spec for roofline analysis | - - -### Step 5: Compare Traces with TraceDiff (Eager-mode only) - -Compare two tracefiles and analyze execution differences using Lowest Common Ancestor (LCA) analysis: - -```python -import sys -from TraceLens import TreePerfAnalyzer, TraceDiff - -file1, file2 = sys.argv[1], sys.argv[2] - -# Build performance trees -print(f"Creating tree1 from {file1}...") -perf_analyzer1 = TreePerfAnalyzer.from_file(file1, add_python_func=True) -tree1 = perf_analyzer1.tree - -print(f"Creating tree2 from {file2}...") -perf_analyzer2 = TreePerfAnalyzer.from_file(file2, add_python_func=True) -tree2 = perf_analyzer2.tree - -# Generate diff report -td = TraceDiff(tree1, tree2) -td.generate_tracediff_report() -td.print_tracediff_report_files("rprt_diff_pruned", prune_non_gpu=True) - -print("✅ Pruned TraceDiff reports (GPU only) written to rprt_diff_pruned/") -``` - -> **Recommendations:** Ensure both tracefiles use similar execution setup (profiled steps, OSL range, concurrency) and the same execution mode (eager/graph) for meaningful comparisons. - -### Step 6: Agentic Trace Analysis (Skip Step 4 and 5) - -Generate a performance analysis and comparison report (if comparing two traces), along with optimization opportunity analysis, automatically using an LLM agent. - -- Performance analysis: This is the recommended first step, and it leverages TraceLens roofline models for performance bridge gap analysis. Please follow [these instructions](../TraceLens/Agent/Analysis/README.md). - ---- - -## 🐞 TraceLens: Report a Bug or Feature Request - -Please include the following details when reporting an issue. Please use internal or direct channels to share sensitive data. - -1. 🖥️ Environment Details - - -| Item | Details | -| -------------------------------- | ----------------------------------- | -| **Inference Engine and Version** | (e.g., vLLM, SGLang) | -| **Execution Mode** | (e.g., Eager, Graph, Graph+Capture) | -| **Hardware** | (e.g., GPU model) | -| **Profiler Config** | (e.g. Torch profiler config) | - - -1. ▶️ Scripts/Commands Used - -The scripts and commands used to generate a performance analysis report using TraceLens to reproduce the issue. - -1. ❗ Error/Unexpected Behavior -2. 📂 Trace Files Used for Analysis -3. (Optional) 🧪 Expected Output Overview for Feature Request - ---- - -## 📚 Examples & Use Cases - -*Example notebooks and scripts coming soon* 🔄 - ---- - -## 🔬 Conceptual Details - -### [TraceDiff: Lowest Common Ancestor (LCA) Analysis](#tracediff-lca-analysis) - -`TraceDiff` is a trace comparison tool that analyzes execution differences between two inference traces (baseline and variant) by constructing a merged tree and identifying structural similarities and differences using **Lowest Common Ancestor (LCA) analysis**. - -#### The Problem - -When comparing two execution traces from different platforms, frameworks, or configurations: - -- Kernel names may differ (e.g., platform-specific optimizations) -- Execution paths may have insertions, deletions, or reorderings -- GPU operations may be fused differently -- We need to **correlate related operations** across traces - -#### The Solution: Lowest Common Ancestor - -The **Lowest Common Ancestor** is the nearest parent CPU operation or Python function that is **common to both traces** in the merged execution tree. A combination of **position**- and **name**-based comparison rules is used to match two operations or functions. It serves as an anchor point for correlating GPU kernels and operations that differ between traces. - -**Key Insight:** If two GPU kernels have the same LCA, they likely serve the same computational purpose, even if their implementations differ. - -#### How It Works - -1. Tree Alignment with Wagner-Fischer Algorithm - -- Uses dynamic programming to align execution trees from both traces -- Identifies three types of nodes: - - **Combined**: Operations present in both traces (potential LCA candidates) - - **Trace1-only**: Operations unique to baseline - - **Trace2-only**: Operations unique to variant -- Normalizes operation names by removing variable parts (memory addresses, line numbers) - -1. Merged Tree Construction - Creates a unified tree structure where: - -- Each node has a unique `merged_id` -- Nodes track UIDs from both original traces (`uid1`, `uid2`) -- Parent-child relationships are preserved from both traces -- The merged tree maintains execution hierarchy - -1. LCA Identification - For GPU operations that differ between traces: - -``` -Trace 1: Trace 2: -┌─────────────┐ ┌─────────────┐ -│ attention │ ◄── LCA ──►│ attention │ (Combined node) -└──────┬──────┘ └──────┬──────┘ - │ │ - ┌───┴───┐ ┌───┴───┐ - │ │ │ │ - GPU_K1 GPU_K2 GPU_K3 GPU_K4 (Different kernels) -(trace1-only) (trace2-only) -``` - -The `attention` operation is the **LCA** for all four GPU kernels, indicating they all serve the same high-level computation despite different implementations. - -1. Performance Correlation - The LCA enables meaningful comparisons: - -- **Kernel Grouping**: All GPU kernels under the same LCA are functionally related -- **Time Aggregation**: Sum kernel times under each LCA for apples-to-apples comparison -- **Shape Analysis**: Compare input dimensions at the LCA level -- **Optimization Identification**: Spot fusion opportunities or inefficiencies - -1. LCA example: - -Example snippet - -``` -├── nn.Module: Attention_0 -│ └── torch/nn/modules/module.py(1779): _call_impl -│ └── combined: vllm/attention/layer.py(310): forward | vllm/attention/layer.py(290): forward -│ ├── combined: torch/_ops.py(1243): __call__ | torch/_ops.py(1244): __call__ -│ │ └── combined: | -│ │ └── vllm::unified_attention_with_output -│ │ └── vllm/attention/utils/kv_transfer_utils.py(36): wrapper -│ │ └── combined: vllm/attention/layer.py(858): unified_attention_with_output | vllm/attention/layer.py(852): unified_attention_with_output -│ │ └── combined: vllm/v1/attention/backends/rocm_attn.py(256): forward | vllm/v1/attention/backends/flashinfer.py(1064): forward -│ │ ├── combined: vllm/attention/ops/paged_attn.py(31): write_to_paged_cache | torch/_ops.py(1244): __call__ -│ │ │ └── combined: vllm/_custom_ops.py(2156): reshape_and_cache | -│ │ │ └── combined: torch/_ops.py(1243): __call__ | _C_cache_ops::reshape_and_cache_flash -│ │ │ ├── >> trace1: -│ │ │ │ └── >> trace1: _C_cache_ops::reshape_and_cache -│ │ │ │ └── >> trace1: hipLaunchKernel -│ │ │ │ └── >> trace1: void vllm::reshape_and_cache_kernel<__hip_bfloat16, __hip_bfloat16, (vllm::Fp8KVCacheDataType)0>(__hip_bfloat16 const*, __hip_bfloat16 const*, __hip_bfloat16*, __hip_bfloat16*, long const*, int, int, int, int, int, int, float const*, float const*) -│ │ │ └── << trace2: cudaLaunchKernel -│ │ │ └── << trace2: void vllm::reshape_and_cache_flash_kernel<__nv_bfloat16, unsigned char, (vllm::Fp8KVCacheDataType)1>(__nv_bfloat16 const*, __nv_bfloat16 const*, unsigned char*, unsigned char*, long const*, long, long, long, long, long, int, int, int, float const*, float const*) -``` - -#### Output: TraceDiff Report - -The generated report includes: - - -| Column | Description | -| --------------------------------- | ------------------------------------ | -| `name` | GPU kernel name | -| `cpu_op_name` | Immediate parent CPU operation | -| `source` | `trace1` or `trace2` | -| `Input Dims` | Tensor shapes at CPU operation level | -| `kernel_time` | GPU kernel execution time (μs) | -| `**lowest_common_ancestor_name**` | **Name of the LCA operation** | -| `**lowest_common_ancestor_id`** | **Merged tree ID of the LCA** | -| `nn_module_stack` | PyTorch module hierarchy | - - -### [Roofline Analysis](#roofline-analysis) - -#### Inference Attention - -In inference serving, multiple requests are batched together. Each request has its own sequence lengths (N_Q, N_KV). -**Notation:** - - -| Symbol | Description | -| ----------------- | ------------------------------------------------------------------- | -| B | Batch size (1 per request in paged attention) | -| N_Q | Number of query tokens | -| N_KV | Number of key/value tokens (context length) | -| H_Q | Number of query heads | -| H_KV | Number of KV heads (H_KV ≤ H_Q; equal for MHA, smaller for GQA/MQA) | -| d_h_qk | Head dimension for queries and keys | -| d_h_v | Head dimension for values | -| R_C | Number of context (prefill) requests in the batch | -| R_G | Number of generation (decode) requests in the batch | -| R | Total number of requests (R = R_C + R_G) | -| N_Q^(i), N_KV^(i) | Query and KV token counts for the i-th request | - - -**Standard SDPA Attention (Single Request)** - -Attention consists of two matrix multiplications per head: - -1. **QK^T (score computation):** `2 * B * N_Q * N_KV * H_Q * d_h_qk` -2. **Score × V (value aggregation):** `2 * B * N_Q * N_KV * H_Q * d_h_v` - -For causal attention, roughly half the score matrix is masked out: - -``` -FLOPS = (2 * B * N_Q * N_KV * H_Q * d_h_qk + 2 * B * N_Q * N_KV * H_Q * d_h_v) / 2 -``` - -``` -Elements Moved = -Q: B * N_Q * H_Q * d_h_qk -K: B * N_KV * H_KV * d_h_qk -V: B * N_KV * H_KV * d_h_v -Output: B * N_Q * H_Q * d_h_v -``` - -**Inference Paged Attention** - -For calculating total Flops and bytes moved for inference paged attention, we **sum over the computation requirement of all requests individually** (B = 1 per request). - -Requests fall into two categories: - -- **Context (prefill) requests** — processing input tokens; attention is causal within the current chunk -- **Generation (decode) requests** — generating new tokens; attention is non-causal (queries attend to all past KV tokens). Typically N_Q = 1, but approaches like speculative decoding may produce multiple query tokens per request. - -**1. Flops Calculation** - -**Prefill Requests (First Chunk or Full Context)** - -When chunked prefill is not enabled, this is the first (and only) chunk, so N_KV = N_Q and attention is causal: - -``` - R_C -FLOPS_prefill = Σ (2 * N_Q(i) * N_KV(i) * H_Q * d_h_qk + 2 * N_Q(i) * N_KV(i) * H_Q * d_h_v) / 2 - i=1 -``` - -**Prefill Requests (Chunked Prefill, nth Chunk)** - -With chunked prefill, the nth chunk has KV tokens from all previous chunks already cached. The attention matrix for one such request looks like: - -``` - Keys (N_KV) - ◄──── N_KV - N_Q ────►◄──── N_Q ───► - ┌──────────────────────┬──────────────┐ ▲ - │ │╲ │ │ - │ │ ╲ (masked) │ │ - │ Non-causal │ ╲ │ │ - Queries │ (full rectangle) │ ╲ │ N_Q - │ │ ╲ │ │ - │ attend to previous │ Causal ╲ │ │ - │ chunks' KV cache │ (self) ╲ │ │ - └──────────────────────┴──────────────┘ ▼ - ◄────── previous ──────►◄── current ──► - chunks chunk -``` - -The attention computation splits into two regions: - -a. **Current chunk attending to previous chunks** — this region is a full (non-causal) rectangle of shape N_Q × (N_KV − N_Q) -b. **Current chunk attending to itself** — this region is causal (lower-triangular), so we halve - -For the **first chunk** (no chunking, or first chunk of chunked prefill), N_KV = N_Q, so the rectangle vanishes and the entire matrix is causal: - -``` - Keys (N_KV = N_Q) - ◄──── N_Q ─────► - ┌──────────────┐ ▲ - │╲ │ │ - │ ╲ (masked) │ │ - │ ╲ │ │ - │ ╲ │ N_Q Queries - │ ╲ │ │ - │ Causal ╲ │ │ - │ (entire) ╲ │ │ - └──────────────┘ ▼ -``` - -``` - R_C -FLOPS_chunked = Σ [ 2 * N_Q(i) * (N_KV(i) - N_Q(i)) * H_Q * d_h_qk - i=1 - + 2 * N_Q(i) * (N_KV(i) - N_Q(i)) * H_Q * d_h_v - + (2 * N_Q(i)² * H_Q * d_h_qk + 2 * N_Q(i)² * H_Q * d_h_v) / 2 ] -``` - -Both cases (first chunk and nth chunk) simplify to a **single unified formula**: - -``` - R_C -FLOPS_context = Σ [ (2 * N_Q(i) * N_KV(i) * H_Q * d_h_qk + 2 * N_Q(i) * N_KV(i) * H_Q * d_h_v) - i=1 - - (2 * N_Q(i)² * H_Q * d_h_qk + 2 * N_Q(i)² * H_Q * d_h_v) / 2 ] -``` - -This works because: - -- When N_KV = N_Q (first chunk): `full - full/2 = full/2`, which is causal -- When N_KV > N_Q (nth chunk): `full rectangle - self-triangle`, which is the non-causal rectangle plus the causal self-attention - -**Generation Requests** - -Generation requests attend to all cached KV tokens (N_KV = context length so far). Typically N_Q = 1 (autoregressive decoding), but techniques like speculative decoding may have N_Q > 1. The attention is non-causal: - -``` - R_G -FLOPS_generation = Σ (2 * N_Q(i) * N_KV(i) * H_Q * d_h_qk + 2 * N_Q(i) * N_KV(i) * H_Q * d_h_v) - i=1 -``` - -``` -FLOPS_total = FLOPS_context + FLOPS_generation -``` - -**2. Elements Moved** - -The total memory traffic sums over all requests. Each request reads its Q, K, V tensors and writes the output. With GQA/MQA, the KV cache uses H_KV heads (not H_Q), reducing KV memory traffic. - -Ignoring cases with shared KV pages between requests: - -``` - R -Elements_moved = Σ ( N_Q(i) * H_Q * d_h_qk // Q read - i=1 - + N_KV(i) * H_KV * d_h_qk // K read (from paged KV cache) - + N_KV(i) * H_KV * d_h_v // V read (from paged KV cache) - + N_Q(i) * H_Q * d_h_v ) // Output write -``` - -where R = R_C + R_G is the total number of requests. - -Note that B = 1 per request in paged attention, so the batch dimension is absorbed into the summation. - -**Practical Roofline Analysis Without Per-Request Details** - -Importantly, we do **not** need the details of individual requests to perform roofline analysis. Inspecting the formulas above, the only per-request quantities that appear are `N_Q(i)`, `N_KV(i)`, and their products. The full FLOPS and memory traffic expressions can be evaluated using just these **aggregate statistics**, computed separately for context and generation requests: - - -| Aggregate | Used in | -| -------------- | ---------------------------------------------------- | -| R_C, R_G | Request counts | -| Σ N_Q | Elements moved (Q read, Output write) | -| Σ N_KV | Elements moved (K read, V read) | -| Σ (N_Q * N_KV) | FLOPS (full rectangle term) | -| Σ (N_Q²) | FLOPS (causal self-attention correction for prefill) | - - -We obtain these aggregates by applying `torch.record_function(annotation)` to vLLM's execution steps. A single execution step can contain a mix of both context (prefill) and generation (decode) requests, so the annotation encodes the aggregate statistics **separately** for context and generation requests within that step (e.g., R_C, R_G, Σ N_Q for context, Σ N_Q for generation, etc.). These annotations are stored as `user_annotation` events in the PyTorch profiler trace, making roofline analysis possible directly from the trace without any additional instrumentation or runtime logging. - -### [Steady-State Region and Trace Splitting](#steady-state-region-and-trace-splitting) - -Inference serving execution consists of three phases: - -1. **Ramp-up:** The initial few steps where one or more requests are being batched. -2. **Ramp-down:** The last few trailing steps where the final batch of requests finishes. -3. **Steady state:** Defined as the execution steps with the highest concurrency. Once steady state is reached, execution consists of: - - Decode‑only steps - - Prefill‑decode steps, typically containing one prefill request packed with ~CONC−1 decode requests. - -For performance analysis, we are interested in profiling only the steady‑state steps: - -1. Prefill‑decode steps -2. Decode‑only steps with large context sizes (towards the end of a request) - -**Parameters Relevant to Inference Serving** - -- **NUM_PROMPTS**: typically `10 × CONC` -- **CONC**: number of concurrent requests that can be batched together -- **R**: Random‑range ratio used for sampling ISL and OSL -- **OSL**: Maximum output sequence length. Output sequence length per request is sampled uniformly in: -`[ R × OSL , OSL ]` -- **ISL**: assumed to be lower than the chunk size - -We assume inference serving benchmark schedules requests at an infinite rate, and conservatively treat the first **CONC** steps as the *ramp‑up* phase. - -With these parameters, execution step ranges where groups of CONC requests complete: - -``` -1 × R × OSL to 1 × OSL e.g. 0.8 OSL – 1 OSL -2 × R × OSL to 2 × OSL e.g. 1.6 OSL – 2 OSL -3 × R × OSL to 3 × OSL e.g. 2.4 OSL – 3 OSL -4 × R × OSL to 4 × OSL e.g. 3.2 OSL – 4 OSL -... -N × R × OSL to N × OSL - -Where, N = NUM_PROMPTS / CONC -``` - -``` -TOTAL_STEPS = NUM_PROMPTS × Avg(OSL) / CONC -TOTAL_PrefillDecode_Steps = NUM_PROMPTS -TOTAL_DecodeOnly_Steps= NUM_PROMPTS × ( (Avg(OSL) − CONC) / CONC ) -``` - -Since inference serving benchmarks commonly use `R = Random Range Ratio` to sample output sequence lengths from `[R * OSL, OSL]`, the most useful steady‑state profiling window lies in: - -``` -5 * R * OSL – 5 * OSL -``` - -Here, the probability distribution of a request finishing at step t is non-uniform, with the highest probability at ((R+1)/2) * 5 * OSL. - -This region exhibits: - -- Fully saturated concurrency -- Representative mix of decode-only and prefill-decode steps -- Minimal warm-up or tail artifacts - -This makes the recommended window for performance profiling: - -``` -max_iters = ( 16 * OSL/CONC ) # The number of execution steps to be profiled. The multiplier 16 is used to ensure that the profiling window has ~16 execution steps with prefill+decode mix requests. The max_iters value might need clamping for extreme values of OSL/CONC. - -delay_iters = ( ((R+1)/2) * 5 * OSL ) - (max_iters/2) # The execution step where profiler starts. -``` - -#### Trace Splitting - -The trace splitting workflow provides three key features. Note that trace splitting assumes vLLM v0.14 or higher (tested through v0.24), or use of our provided patches, to ensure that relevant annotations (batch size, request counts, etc.) are included in execution step metadata. - -1. **Split into individual execution steps:** Decompose the entire trace into per-step files, extracting batch size from annotations or kernels for shape-focused analysis and comparison. -2. **Identify steady-state region:** Detect execution steps with near-maximum concurrency. The algorithm identifies large windows with concurrency close to peak levels and selects a representative steady-state region based on prefill-decode and decode-only step composition. When benchmark parameters are known, pass `--CONC`, `--OSL`, and `--R` to `split_inference_trace_annotation` to override the empirical reference ratio with the analytically derived ideal PD ratio — see [Step 3](#step-3-trace-preparation-optional) for details. -3. **Separate phase analysis:** Further decompose steady-state into prefill-decode and decode-only traces. Since prefill and decode have different computational bottlenecks, separate analysis enables targeted performance optimization. - -### [Trace Availability-Analysis Trade-off](#trace-availability-analysis-trade-off) - -Balancing complete trace capture versus analysis complexity. - ---- - -**Last Updated:** May 2026 -**Maintainers:** AMD-AGI Performance and Optimization Team -**Repository:** [github.com/AMD-AGI/TraceLens](https://github.com/AMD-AGI/TraceLens) diff --git a/docs_original/NcclAnalyser.md b/docs_original/NcclAnalyser.md deleted file mode 100755 index acddc5e5b..000000000 --- a/docs_original/NcclAnalyser.md +++ /dev/null @@ -1,80 +0,0 @@ - - -# NCCL Analyser - -In distributed deep learning, analyzing the performance of collective communication operations is crucial for diagnosing and optimizing **performance at scale**. **NCCL Analyser** is a Python SDK designed to parse and analyze NCCL kernel events from PyTorch trace files (JSON). It computes key metrics like **communication latency**, **message sizes**, **algorithm bandwidth**, **bus bandwidth**, and **synchronization metrics** (e.g., skew in start/end times), providing insights into communication patterns and potential bottlenecks in your distributed training or inference workflows. - - ---- - -## Key Features - - -- **Detailed NCCL Event Extraction**: Parses JSON trace files to extract per-rank NCCL kernel events (those containing _nccl_ in the kernel name). - -- **Correlated Events via Collective ID**: Each collective operation is assigned a unique `collective_id`, generated by combining the **Process Group Name** and an `index_in_group`. The `index_in_group` reflects the order in which that collective appears in the trace. - -- **Latency Computation for Implicit-Sync Collectives**: For collectives like allreduce, allgather, reducescatter, and alltoall, the tool computes communication latency by taking the minimum duration across ranks to eliminate waiting time. - -- **Bandwidth Analysis**: Computes algorithm bandwidth based on input msg size and communication latency and finally bus bandwidth (link utilization). - -- **Synchronization (Wait-Time) Metrics**: Tracks **skew in start times** (how late certain ranks arrive) and **skew in end times**, providing insights for diagnosing load imbalance or synchronization overheads. - -- **Alltoallv Support**: Provides raw data for alltoallv events, which do not enforce an implicit sync pattern. Includes per-instance metrics (wall time, throughput, size imbalance) and a summary aggregation grouped by process group and dtype. An optional heatmap mode (`build_df_all2allv_heatmap`) produces per rank-pair send volumes, useful for diagnosing MoE/expert-parallel routing imbalance. - -- **Summary & Detailed Dataframes**: Supports both high-level summaries for quick insights and detailed dataframes for advanced analysis, including per-rank metadata and performance metrics. We suggest **power users** to build their custom pipelines and analyses on top of detailed dfs. - -- **PyTorch Support**: Currently built for PyTorch trace files, with the potential to extend support to other frameworks. - ---- - -## Quick Start - -The NcclAnalyser can generate both per collective data frame as well as a global summary. There are more features demonstrated in [example notebook](../examples/nccl_analyser_example.ipynb) - - -### Example: Quick Summary of Implicit Sync Collectives - -```python -from TraceLens import NcclAnalyser -import os - -# Define the root directory and create file paths for all ranks -root_dir = '/path/to/your/trace/files/directory' -world_size = 8 -list_profile_filepaths = [os.path.join(root_dir, f'rank{i}_trace.json') for i in range(world_size)] - -# Initialize the NCCL Analyser -my_nccl_analyser = NcclAnalyser(list_profile_filepaths, world_size) - -# Generate the summarized dataframe for implicit-sync collectives -df_summary = my_nccl_analyser.build_df_summary_nccl_implicit_sync_cat(agg_metrics=['mean']) -print(df_summary.head()) - -# Save the summary to CSV -df_summary.to_csv('nccl_summary.csv', index=False) - -``` -**Summarized Dataframe**: -- Groups events by collective type, message size, and data type to compute aggregated metrics. -- For communication performance we compute communication latency, algorithm bandwidth, and bus bandwidth. -- For synchronisation delay we compute metrics like skew in start times which shows the wait time. This is based on the difference of earliest and latest arrival across the ranks for a given collective. - -#### Example Summarized Dataframe - -| Collective name | In msg size (MB) | dtype | comm latency (µs)_mean | count | Total latency (ms) | algo bw (GB/s)_mean | bus bw (GB/s)_mean | skew in start time (µs)_mean | -|----------------------|------------------|----------|------------------------|-------|---------------------|---------------------|---------------------|-----------------------------| -| _allgather_base | 204.00 | BFloat16 | 6041.88 | 318 | 1921.32 | 33.00 | 28.88 | 11779.36 | -| _reduce_scatter_base | 3264.06 | Float | 11662.77 | 160 | 1866.04 | 273.43 | 239.25 | 60238.77 | -| _reduce_scatter_base | 8016.03 | Float | 22988.50 | 2 | 45.98 | 340.53 | 297.96 | 146.48 | -| _allgather_base | 501.00 | BFloat16 | 11920.84 | 2 | 23.84 | 41.04 | 35.91 | 15405.14 | -| allreduce | 0.00 | Float | 18.58 | 6 | 0.11 | 0.00 | 0.00 | 936.63 | - - -Note that the last row in msg size is just rounded down to 0. - -**Modify the filepaths for your profiles in nccl_analyser_example.ipynb notebook and get analysis instantly!** diff --git a/docs_original/Trace2Tree.md b/docs_original/Trace2Tree.md deleted file mode 100755 index 2e8dd9f28..000000000 --- a/docs_original/Trace2Tree.md +++ /dev/null @@ -1,109 +0,0 @@ - - -# Trace2Tree - -In GPU applications performance analysis, understanding the relationship between host CPU operations and corresponding GPU kernel executions is crucial for analyzing bottlenecks. The PyTorch profiler provides a JSON trace file containing events with timestamps and durations but lacks explicit call stack dependency information. - -Trace2Tree is the underlying tree structure component used by TraceLens to parse trace files and build hierarchical dependency relationships from host CPU operations to GPU kernels. **It is recommended to access this functionality through the `TreePerfAnalyzer` interface** rather than using Trace2Tree directly. - ---- - -## Key Features - -- **Hierarchical Dependency Tree**: Constructs a tree structure linking CPU operations to GPU kernel launches, enabling detailed analysis of the ops lowering and perfomance. -- **Extensible SDK**: Provides a framework for custom analyses, such as identifying GPU time for CPU operations or pinpointing bottlenecks. -- **Lightweight Design**: Minimal dependencies and a straightforward codebase for easy integration and use. -- **PyTorch Support**: Built for PyTorch profiler JSON traces, with potential for future support of other frameworks. - ---- - -## Quick Start - -> **💡 Tip:** See [`examples/trace2tree_example.ipynb`](../examples/trace2tree_example.ipynb) for a complete interactive tutorial. - -### Example: Build and traverse tree - -#### 1. Load the Trace data via TreePerfAnalyzer -```python -from TraceLens.TreePerf import TreePerfAnalyzer - -# Load trace data using TreePerfAnalyzer -# Set add_python_func=True to include Python function call stack in the tree -# This allows you to trace GPU kernels all the way back to your Python code -trace_file = '/path/to/trace.json' -analyzer = TreePerfAnalyzer.from_file(trace_file, add_python_func=True) - -# Access the underlying tree structure -tree = analyzer.tree -``` - -#### 2. Find an Operation to Analyze - -```python -# Find an operation of interest -event_interest = next( - evt for evt in tree.events - if evt.get('name') == 'aten::convolution' and evt.get('cat') == 'cpu_op' -) -``` - -#### 3. Traverse Subtree - -Visualize the entire subtree rooted at this operation: - -```python -tree.traverse_subtree_and_print(event_interest) -``` - -``` -└── UID: 41, Category: cpu_op, Name: aten::convolution - └── UID: 42, Category: cpu_op, Name: aten::_convolution - └── UID: 43, Category: cpu_op, Name: aten::miopen_convolution - ├── UID: 104314, Category: cuda_runtime, Name: hipExtModuleLaunchKernel - │ └── UID: 107846, Category: kernel, Name: Im2d2Col_v2, Duration: 45.063 - └── UID: 104318, Category: cuda_runtime, Name: hipExtModuleLaunchKernel - └── UID: 107848, Category: kernel, Name: Cijk_Ailk_Bljk_BBS_BH... -``` - -#### 4. Traverse Parent Chain - -Trace back through all parent events to see the full call stack. You can optionally include CPU operation details like input dimensions, types, and strides using the `cpu_op_fields` parameter: - -Available fields: `'Input Dims'`, `'Input type'`, `'Input Strides'`, `'Concrete Inputs'` - -```python -root = tree.traverse_parents_and_print( - event_interest, - cpu_op_fields=('Input Dims', 'Input type') -) -``` - -``` -Node: - UID: 41, Category: cpu_op, Name: aten::convolution - Input Dims: [[1, 768, 24, 24], [768, 768, 3, 3], []] - Input type: [float, float, float] -1-up: - UID: 40, Category: cpu_op, Name: aten::conv2d - Input Dims: [[1, 768, 24, 24], [768, 768, 3, 3], [1], [2, 2], [1, 1], [1, 1], [1]] - Input type: [float, float, int, int, int, int, int] -2-up: - UID: 40139, Category: python_function, Name: -3-up: - UID: 40138, Category: python_function, Name: torch/utils/_device.py(100): __torch_function__ -4-up: - UID: 40137, Category: python_function, Name: torch/nn/modules/conv.py(554): _conv_forward -5-up: - UID: 40136, Category: python_function, Name: torch/nn/modules/conv.py(558): forward -6-up: - UID: 40135, Category: python_function, Name: torch/nn/modules/module.py(1736): _wrapped_call_impl -7-up: - UID: 40134, Category: python_function, Name: torch/nn/modules/module.py(1747): _call_impl -8-up: - UID: 40133, Category: python_function, Name: transformers/models/owlv2/modeling_owlv2.py(395): forward -... -``` \ No newline at end of file diff --git a/docs_original/TraceDiff.md b/docs_original/TraceDiff.md deleted file mode 100644 index 8cb8c8355..000000000 --- a/docs_original/TraceDiff.md +++ /dev/null @@ -1,145 +0,0 @@ - - -# TraceDiff - -TraceDiff is a Python API and a component of TraceLens for comparing two PyTorch Kineto trace files. It allows users to identify, visualize, and analyze differences between execution traces at both the operation and kernel levels. - -Unlike a simple leaf-level operation comparison, TraceDiff considers the morphological structure of traces to automatically determine the lowest common node. This improves accuracy when dealing with differences in operator lowering at the host call stack — for example, `aten::convolution` lowering to `aten::miopen_convolution` on ROCm and `aten::cudnn_convolution` on CUDA. A leaf-level comparison alone would treat these as completely different operations, whereas TraceDiff can automatically match them at the appropriate level. - -TraceDiff is particularly useful for regression analysis, performance debugging, and assessing the impact of code or environment changes on GPU workloads. - - ---- - -## Key Features - -- **Automated Tree Comparison**: Builds hierarchical event trees from two traces and identifies points of difference (PODs) using a recursive diff algorithm. -- **Tree Diff Visualization**: Produces a diff output file that highlights matched and unmatched operations between traces. -- **Detailed and Summary Reports**: Generates CSV reports with kernel time statistics and aggregated summaries for each operation. -- **UID Mapping**: Provides a mapping between event UIDs in the two traces, enabling cross-referencing and deeper analysis. -- **Seamless Integration**: Designed to work with TraceLens's TraceToTree objects and PyTorch profiler JSON traces. - ---- - -## Quick Start - - -### Example: Compare Two Traces and Generate Reports - -```python -from TraceLens import TreePerfAnalyzer, TraceDiff -import json - -# Load two trace files into tree perf analyzer -trace_file1 = "/path/to/trace1.json" -trace_file2 = "/path/to/trace2.json" - -perf_analyzer1 = TreePerfAnalyzer.from_file(trace_file1) -perf_analyzer2 = TreePerfAnalyzer.from_file(trace_file2) -tree1 = perf_analyzer1.tree -tree2 = perf_analyzer2.tree - -# Compare and analyze the trees -td = TraceDiff(tree1, tree2) -td.generate_tracediff_report() # Generates DataFrames, does NOT write files -td.print_tracediff_report_files('rprt_diff') # Writes all reports to files in 'rprt_diff/' -``` - - - -**Output files:** -- `rprt_diff/merged_tree_output.txt`: Text visualization of the merged tree, showing matched and unmatched nodes. -- `rprt_diff/diff_stats.csv`: Detailed kernel and op statistics for each operation (see below for example and explanation). -- `rprt_diff/diff_stats_unique_args_summary.csv`: Aggregated summary statistics by op name and unique args. -- `rprt_diff/diff_stats_names_summary_df`: Aggregated summary stats by op name - giving top level summary. - ---- - - -## Output File Examples and Interpretation - -#### diff_stats_names_summary_df.csv - - - -**Example (first 5 rows):** - -| name | row_count | kernel_time_trace1_sum_ms | kernel_time_trace2_sum_ms | diff_sum_ms | abs_diff_sum_ms | -|-----------------------------|-----------|---------------------------|---------------------------|--------------|-----------------| -| aten::convolution_backward | 736 | 541.957809 | 366.136090 | -175.821719 | 198.297619 | -| aten::_convolution | 736 | 229.175081 | 157.807700 | -71.367381 | 85.731275 | -| aten::_batch_norm_impl_index| 448 | 129.995936 | 43.081600 | -86.914335 | 86.914335 | -| aten::mm | 300 | 78.684444 | 84.654093 | 5.969649 | 11.982847 | -| FlashAttnFuncBackward | 25 | 59.776381 | 54.930648 | -4.845733 | 4.845733 | - - - -#### diff_stats.csv - -This file contains detailed statistics for every op instance, including input shapes, types, kernel times, and kernel names. It is useful for fine-grained analysis and debugging. - -#### diff_stats_unique_args_summary.csv - -Midway between the detailed op wise view and the name summary this is a summary per argument of an operation. - -**How to use:** -- Drill down to individual op instances to investigate outliers or mismatches. -- Use the detailed input and kernel info to correlate with model code or trace events. - ---- - -## Accessing DataFrames and UID Mapping - -TraceDiff provides methods to access the detailed and summary DataFrames directly, as well as a `merged_uid_map` to cross-reference events between the two traces. This is useful for linking statistics or visualizations. - -### Accessing DataFrames - -```python -# After running td.generate_tracediff_report(): -df_stats = td.diff_stats_df # Detailed per-op DataFrame -df_unique_args_summary = td.diff_stats_unique_args_summary_df -df_name_summary = td.diff_stats_name_summary_df - -if df is not None: - print(df.head()) -if df_unique_args_summary is not None: - print(df_unique_args_summary.head()) -if df_name_summary is not None: - print(df_name_summary.head()) -``` - -### UID Mapping Example - -```python -# Get the corresponding UID in tree2 for a given UID in tree1 -uid1 = next(iter(td.baseline.cpu_root_nodes)) -uid2 = td.get_corresponding_uid(1, uid1) -if uid2 != -1: - print(f"Tree1 UID {uid1} corresponds to Tree2 UID {uid2}") -else: - print("No match found for this UID in tree2.") -``` - ---- - - -## Use Cases - -- **Performance Regression Analysis**: Quickly identify which operations or kernels have changed between two runs. -- **Debugging and Optimization**: Pinpoint new bottlenecks or regressions introduced by code or environment changes. -- **Cross-Trace Linking**: Map and compare specific events or kernels between two traces for deeper investigation. - ---- - - -## Notes - -- TraceDiff is designed for PyTorch profiler JSON traces and requires TraceToTree objects as input. -- For more advanced usage, see the example notebook: `examples/trace_diff_example.ipynb`. -- Output folder and file names can be customized via the API. -- The API now separates report generation (`generate_tracediff_report`) from file output (`print_tracediff_report_files`). -- DataFrames are only available after running `generate_tracediff_report()`. diff --git a/docs_original/TraceFusion.md b/docs_original/TraceFusion.md deleted file mode 100755 index 07ca3ffca..000000000 --- a/docs_original/TraceFusion.md +++ /dev/null @@ -1,102 +0,0 @@ - - -## TraceFusion - -In distributed deep learning, diagnosing issues like straggling ranks, load imbalance, or bottlenecks requires a **global view** of events across all ranks. TraceFusion is a Python SDK for merging trace files across ranks in distributed training and inference setups. With customization options for filtering events and defining file paths, TraceFusion simplifies the preparation of traces for seamless rendering in **Perfetto**. -Note that: TraceFusion is **only for visual analysis in PerfettoUI** and not for automated analysis. - ---- - -## Key Features - -- **Custom Filtering**: Easily define filtering logic to include or exclude specific events. For example, merge traces for a subset of ranks, focus only on GPU events, or narrow down further to NCCL kernel events only. -- **PyTorch Support**: Built primarily for PyTorch trace files, with potential support for other frameworks in the future. -- **Lightweight and Simple**: A dependency-free and straightforward codebase makes it easy to integrate and extend. - ---- - -## Quick Start - -Here's how to use TraceFusion to merge and process trace files for distributed training or inference: - -### Example 1: Basic Usage - -```python -from TraceLens import TraceFuse -import os - -# Define file paths for each rank -root_profiles = '/path/to/profiles/' -world_size = 8 -list_profile_files = [os.path.join(root_profiles, f'pytorch_profile_rank{i}_step120.json') for i in range(world_size)] - -# Initialize TraceFusion (sequential by default) -fuser = TraceFuse(list_profile_files) - -# Merge and Save traces -output_file = os.path.join(root_profiles, 'merged_trace_all_events.json') -fuser.merge_and_save(output_file) - -# By default, Python function category events are skipped to save memory. -# To include them, set include_pyfunc=True. -``` - -### Example 1b: Using Multiprocessing for Faster Processing - -```python -from TraceLens import TraceFuse -import os - -# Define file paths for each rank -root_profiles = '/path/to/profiles/' -world_size = 8 -list_profile_files = [os.path.join(root_profiles, f'pytorch_profile_rank{i}_step120.json') for i in range(world_size)] - -# Initialize TraceFusion with multiprocessing for faster processing -# Can provide significant speedup (system-dependent) -fuser = TraceFuse(list_profile_files, use_multiprocessing=True) - -# Optional: control number of workers (defaults to os.cpu_count()) -# fuser = TraceFuse(list_profile_files, use_multiprocessing=True, max_workers=32) - -# Merge and Save traces -output_file = os.path.join(root_profiles, 'merged_trace_all_events.json') -fuser.merge_and_save(output_file) -``` - -### Example 2: Advanced Usage - -```python -from TraceLens import TraceFuse -import os - -# Define file paths for rank 0 on each node -world_size = 64 -profile_files = {i: os.path.join('/path/to/profiles/', f'pytorch_profile_rank{i}_step120.json') for i in range(0, world_size, 8)} - -# Initialize TraceFusion -fuser = TraceFuse(profile_files) - -# Custom filter for NCCL kernels -def filter_nccl_kernels(event): - return ('cat' in event and 'args' in event and event['cat'] in ['kernel', 'gpu_user_annotation'] and 'nccl' in event['name']) - -# Merge and Save traces -output_file = '/path/to/profiles/merged_trace_nccl.json' -fuser.merge_and_save(output_file, filter_fn=filter_nccl_kernels) -``` - ---- - -### What's Inside? - -TraceFusion merges `traceEvents` across ranks by: -1. **Appending Events**: Combines all events from multiple ranks into a single list. -2. **Adjusting Process IDs**: Modifies `pid` so that traces for each rank render correctly in the UI. -3. **Correcting Flow Linking**: Updates `External id` for events and `id` for corresponding `ac2g` events to ensure accurate flow linking in the UI. - -These adjustments ensure seamless visualization in **Perfetto**, with clear rank separation and correct flow rendering. diff --git a/docs_original/TraceLens - Democratizing AI Performance Analysis - Adeem Jassani, AMD.pdf b/docs_original/TraceLens - Democratizing AI Performance Analysis - Adeem Jassani, AMD.pdf deleted file mode 100644 index 1ad86b289..000000000 Binary files a/docs_original/TraceLens - Democratizing AI Performance Analysis - Adeem Jassani, AMD.pdf and /dev/null differ diff --git a/docs_original/TreePerf.md b/docs_original/TreePerf.md deleted file mode 100755 index 2a918adcb..000000000 --- a/docs_original/TreePerf.md +++ /dev/null @@ -1,68 +0,0 @@ - - -# TreePerf - -TreePerf is a Python SDK that works in conjunction with the Trace2Tree project. PyTorch generates a trace JSON file during profiling, which Trace2Tree parses into a tree data structure representing hierarchical dependencies between CPU operations and GPU kernel executions. TreePerf builds on this tree structure to compute performance metrics at both the model and operation levels. It enables users to analyze, interpret, and optimize AI models by providing detailed performance insights essential for architectural design and performance optimization. - ---- -## Key Ideas - -Key Ideas -1. Tree Structure for GPU Execution Analysis: The hierarchical tree structure, generated by Trace2Tree, enables straightforward computation of GPU execution times. By linking CPU operations to their corresponding GPU kernel launches, it allows seamless aggregation of kernel execution times and performance metrics at various levels of granularity. -2. Performance Metrics from JSON Parsing: Metrics such as FLOPS and FLOPS/Byte are derived by extracting operation parameters from the JSON event data structure. These parameters are parsed and fed into performance models, enabling precise computation of performance metrics. - ---- - -## Key Featues - -1. **GPU timeline breakdown**: Provides a high-level view of activity of the GPU which includes busy time, idle time, communication time, etc -Example output: - -| type | time ms | percent | -|-------------------------------|--------------|------------| -| busy_time | 6521.458211 | 99.927717 | -| computation_time | 6318.257587 | 96.814092 | -| exposed_communication_time | 203.057890 | 3.111438 | -| exposed_memcpy_time | 0.142734 | 0.002187 | -| idle_time | 4.717306 | 0.072283 | -| total_time | 6526.175517 | 100.000000 | - - -2. **GPU compute time breakdown by CPU op**: Analyze the performance breakdown by examining the lowest-level CPU operations (from the call stack perspective) and the time they "induce" on the GPU. -Unlike traditional methods that directly look at CPU and GPU durations, this feature provides insight into how CPU operations translate into GPU kernel launches, offering a more stable and interpretable abstraction of performance. -Example output showing top 5 ops sorted by the total GPU time each op induces: - -| name | Count | total_direct_kernel_time_ms | Percentage (%) | Cumulative Percentage (%) | -|------------------------------------------------------------|-------|-----------------------------|----------------|---------------------------| -| aten::mm | 126 | 5576.76 | 88.73 | 88.73 | -| flash_attn::_flash_attn_backward | 8 | 213.62 | 3.40 | 92.13 | -| flash_attn::_flash_attn_forward | 8 | 119.65 | 1.90 | 94.04 | -| aten::copy_ | 4 | 69.96 | 1.11 | 95.15 | -| triton_poi_fused_add_fill_mul_sigmoid_silu_sub_0 | 8 | 43.19 | 0.69 | 95.84 | - - -3. **Roofline metrics** - Example output for aten::mm showing top 5 param combos sorted by the total GPU time each param combo induces: - -| name | param: M | param: N | param: K | param: bias | FLOPS/Byte_first | TFLOPS/s_mean | -|---------|----------|----------|----------|-------------|------------------|---------------| -| aten::mm | 73728 | 28672 | 8192 | False | 5864.73 | 698.19 | -| aten::mm | 73728 | 8192 | 28672 | False | 5864.73 | 719.59 | -| aten::mm | 28672 | 8192 | 73728 | False | 5864.73 | 740.51 | -| aten::mm | 73728 | 128256 | 8192 | False | 6972.01 | 628.10 | -| aten::mm | 8192 | 28672 | 73728 | False | 5864.73 | 599.95 | - -- Adding new operations is simple (Contributions are welcome!): - - perf_model.py - - Parse operation shapes from the JSON trace - - Write the performance model (FLOPS, bytes) or reuse an existing one - - torch_op_mapping.py: Map the operation to the perf model here - -For more details and walkthrough checkout the base_example.ipynb notebook -**Replace the profile path in base_example.ipynb by your profile file and get insights instantly!** - -Once you are familiar with the workflows you can directly use or modify the **generate_perf_report.py** script to generate a report excel sheet quickly. \ No newline at end of file diff --git a/docs_original/compare_perf_reports_pytorch.md b/docs_original/compare_perf_reports_pytorch.md deleted file mode 100644 index 062c046c2..000000000 --- a/docs_original/compare_perf_reports_pytorch.md +++ /dev/null @@ -1,92 +0,0 @@ - - -## Compare TraceLens Performance Reports - -This Python script (`TraceLens/Reporting/compare_perf_reports_pytorch.py`) compares multiple performance reports generated by `generate_perf_report_pytorch.py`. - - ---- - -### 1 . What It Takes In - -| Input | Description | -| ------------------ | -------------------------------------------------------------------------------------------------------------------------- | -| `*.xlsx` files | TraceLens reports you want to compare. Provide **at least two**. | -| Optional `--names` | Human-readable tags for each report. If omitted, the script falls back to the base filenames (handy, but sometimes messy). | - ---- - -### 2 . How to Call It - -```bash -python compare_perf_reports.py \ - baseline.xlsx \ - candidate.xlsx \ - --names baseline candidate \ - --sheets all \ - -o comparison.xlsx -``` - -Alternatively, you can call the entry point directly after installing TraceLens, which avoids needing to know the script's path: -```bash -TraceLens_compare_perf_reports_pytorch \ - baseline.xlsx \ - candidate.xlsx \ - --names baseline candidate \ - --sheets all \ - -o comparison.xlsx -``` - - -Common flags: - -| Flag | Default | Purpose | -| -------------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------- | -| `-o, --output` | `comparison.xlsx` | Name of the merged workbook. | -| `--names` | `` | Custom tags (must match the number of reports). | -| `--sheets` | `all` | Limit processing to a subset:
`gpu_timeline`, `ops_summary`, `kernel_summary`, `ops_all`, `roofline`, or `all`. | - ---- - -### 3 . What Comes Out - -The script writes **one workbook** (`comparison.xlsx` unless you override it) containing multiple sheets: - -| Sheet | When You Get It | What It Shows | -| ---------------- | ------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| `gpu_timeline` | `--sheets gpu_timeline` or `all` | End-to-end GPU activity by **type** (`compute`, `memcpy`, etc.) with per-report timings plus:
- `time ms___diff`
- `time ms___pct` | -| `ops_summary` | `--sheets ops_summary` or `all` | Per-op aggregates keyed on **`name`**. Sorted by the baseline's `total_direct_kernel_time_ms`. Unhelpful columns (e.g., cumulative %) are stripped from non-baseline reports. When comparing PyTorch traces, this sheet is generated. For rocprof traces without `ops_summary`, automatically falls back to `kernel_summary`. | -| `kernel_summary` | `--sheets kernel_summary` or `all` | Kernel-level summary for rocprof reports, keyed on **`name`**. Similar to `ops_summary` but uses rocprof-specific columns like `Total Kernel Time (ms)` and `Count`. When `--sheets all` is used with rocprof reports, this is generated automatically if `ops_summary` is not available. | -| `ops_all_*` | `--sheets ops_all` or `all` | Three sheets **per variant tag**:
• `ops_all_intersect_` – op instances present in both baseline and variant.
• `ops_all_only_baseline_` – ops only the baseline ran.
• `ops_all_only_variant_` – ops only the variant ran.
Columns irrelevant to a given view are hidden, not deleted. | -| `_*` | `--sheets roofline` or `all` | Same intersect / only\_\* breakdown for each roofline group:
`GEMM`, `SDPA_fwd`, `SDPA_bwd`, `CONV_fwd`, `CONV_bwd`, `UnaryElementwise`, `BinaryElementwise`. | - -Hidden columns stay in the file (for power users) but are invisible in Excel by default. - ---- - -### 4 . Diff Math - -For every metric you ask it to track (`diff_cols` in the code), the script computes: - -```text -metric___diff # variant - baseline -metric___pct # 100 * diff / baseline -``` ---- - -### 5 . Design Decisions You Should Know - -* **Outer merge, never inner** – if an op vanished, you’ll see it. -* **Baseline = first report** – choose wisely. -* **Column prefixing** – every metric becomes `::metric`, so you can safely concatenate arbitrary reports. -* **Sheet-specific pruning** – the script aggressively hides noise (e.g., median, UID) to keep the output readable. You can always unhide them in Excel if you need them. -* **Excel 31-char rule** – sheet names are truncated to fit; no data loss, just shorter labels. - ---- - -### 6 . Future Enhancements -1. Morphology-aware diffing – Understand the call stack tree and compare at lowest common call stack level. For example, if a baseline leaf op is 'cudnn_convolution' and the variant is 'miopen_convolution', the diff algorithm should recognize that lowest common level is 'convolution' and compare the two. \ No newline at end of file diff --git a/docs_original/conceptual/aimodels_gemms.md b/docs_original/conceptual/aimodels_gemms.md deleted file mode 100644 index 558457cb5..000000000 --- a/docs_original/conceptual/aimodels_gemms.md +++ /dev/null @@ -1,274 +0,0 @@ - - -## 📌 GEMMs in AI Workloads - -General Matrix Multiplications (GEMMs) are the **primary compute primitive** used in AI models. Efficient implementations of GEMMs are readily available through vendor-tuned libraries such as cuBLAS and hipBLAS. Therefore, whenever possible, the goal is to **reduce computation to a matrix multiply**. - -This document explains how **model-level parameters** like batch size, sequence length, and hidden dimension translate into GEMM shapes, and subsequently how these shapes map to specific **BLAS kernel calls**. - -By the end of this post, you should be equipped to understand the GEMM shapes, counts, and BLAS calls involved in **any AI model you encounter**. - ---- - -## From Model Dimensions to GEMM Shapes - -### Linear Layers in LLMs - -Let's begin by understanding how linear layers, such as the MLP **up projection**, correspond to GEMM calls. - -The input tensor shape for this operation is: -``` -X: [B, L, d_model] -``` -Here: -- `B` represents the batch size. -- `L` denotes the sequence length. -- `d_model` is the input (or hidden) dimension. - -This operation outputs a tensor with the shape: -``` -Y: [B, L, d_ff] -``` -The projection for each token can be expressed individually as: -``` -Y[b, l, :] = X[b, l, :] @ Wᵀ -``` -Where: -- `W` is the weight matrix with a shape of `[d_ff, d_model]`. - -### Flattening for GEMM - -To express this entire operation as a single GEMM, we flatten the batch and sequence dimensions of the input tensor: -``` -X_flat: [B·L, d_model] -Wᵀ:     [d_model, d_ff] -Y = X_flat @ Wᵀ -``` -This flattening yields the following GEMM shape parameters: -- `param: M = B·L` -- `param: N = d_ff` -- `param: K = d_model` - -Here, K represents the **inner or shared dimension** between the input tensors involved in the multiplication. - - -*** - -## Prefill vs Decode - -Let's now contrast the GEMM behavior during the **prefill** and **decode** phases of inference, focusing on how the sequence length (`L`) changes and affects the GEMM shapes. - -In both phases, the input tensor for an MLP computation within an LLM initially has a shape like: -``` -X: [B, L, d_model] -``` -Where `B` is the batch size, `L` is the sequence length, and `d_model` is the hidden size. This projects to an output shape of `[B, L, d_ff]`, where `d_ff` is the MLP's expansion size. - -As established earlier, to process this with a single GEMM, the batch and sequence dimensions of the input are flattened. The input effectively becomes `[B·L, d_model]` for the GEMM `X_flat @ Wᵀ`. - -The key difference between prefill and decode lies in the value of `L`: -- **Prefill Phase**: `L` is the actual input sequence length (which can be large). -- **Decode Phase**: `L` is always `1`, as the model processes one token at a time to generate the next. - -This difference in `L` directly impacts the `M` parameter of the GEMM `(M, N, K)`: - -### GEMM Shape Summary: -| Mode     | Input Shape (Conceptual) | Flattened Input Shape | GEMM Shape `(M, N, K)`              | Notes                                | -|----------|--------------------------|-----------------------|-------------------------------------|---------------------------------------| -| Prefill  | `[B, L, d_model]`        | `[B·L, d_model]`      | `(B·L, d_ff, d_model)`              | `L` is the prompt length            | -| Decode   | `[B, 1, d_model]`        | `[B·1, d_model]`      | `(B, d_ff, d_model)`                | `L=1` for generating one token at a time | - -Notice that in the decode phase, because `L=1`, the `M` parameter of the MLP GEMM becomes simply `B`. This means the computational cost of the MLP layers in decode remains constant per token regardless of the total sequence length generated so far. The dominant **O(L)** scaling cost during decode comes from the attention mechanism, not the MLPs. - -### Real-world Example: LLaMA-2 7B - -The table below shows actual data from **Tracelens** profiling, filtered specifically for MLP **up** and **gate** projection GEMMs in a LLaMA-2 7B model inference trace. - -For this trace: -- `d_model = 4096` -- `d_ff = 11008` -- Batch size: `1` (`B=1`) -- Input length (for prefill): `597` (`L=597`) -- The trace included 36 decode steps. - -| name     | param: M | param: N | param: K | param: bias | counts | -|----------|-----------|-----------|-----------|---------------|--------| -| aten::mm | 1         | 11008     | 4096      | FALSE         | 2304   | -| aten::mm | 597       | 11008     | 4096      | FALSE         | 64     | - -Interpreting these entries based on the `M` parameter: -- The entry with `param: M = 597` corresponds to the **prefill** phase GEMM `(B·L = 1·597)`, which happens once per layer at the beginning of inference. Since there are 32 layers, this GEMM is called `64` times (32 up + 32 gate). -- The entry with `param: M = 1` corresponds to the **decode** phase GEMM `(B = 1)`, where `L=1`. These occur at each decode step for every layer. With 36 decode steps and 64 GEMMs per step (32 layers * 2), this GEMM is called `36 × 64 = 2304` times. - ---- - -## 🔁 Backward Pass GEMMs - -Next, let's explore the backward pass during training. A forward pass GEMM operation like `Y = X @ Wᵀ + b` necessitates **two corresponding backward GEMMs** to compute gradients: - -```python -dX = dY @ W        # Gradient with respect to the input → resulting shape: [B·L, d_model] -dW = dYᵀ @ X       # Gradient with respect to the weight → resulting shape: [d_ff, d_model] -db = dY.sum(dim=0) # Gradient with respect to the bias   → resulting shape: [d_ff] -``` - -### GEMM Shapes: -| Operation         | GEMM Shape `(param: M, param: N, param: K)`       | Description                          | -|------------------|---------------------------------------------------|--------------------------------------| -| Forward          | `(B·L, d_ff, d_model)`                             | `X @ Wᵀ`                             | -| Backward dX      | `(B·L, d_model, d_ff)`                             | `dY @ W` (result of `[B·L, d_ff] @ [d_ff, d_model]`) | -| Backward dW      | `(d_ff, B·L, d_model)`                             | `dYᵀ @ X` (result of `[d_ff, B·L] @ [B·L, d_model]`) | - -Let's look closer at the backward GEMM shapes: -- For `dX = dY @ W`, the operation is `[B·L, d_ff] @ [d_ff, d_model]`, which results in a shape of `[B·L, d_model]`. -- For `dW = dYᵀ @ X`, the operation is `[d_ff, B·L] @ [B·L, d_model]`, yielding a shape of `[d_ff, d_model]`. - -### Real-world Example: GPT-3-XL - -This table presents data from a **Tracelens** of a single training step for **GPT-3-XL**. - -For this example: -- `d_model = 2048` -- `d_ff = 8192` -- Batch size: `5` -- Sequence length: `2048` -- Thus, `param: M` for the flattened dimension is `5 × 2048 = 10240`. - -| name        | param: M | param: N | param: K | count | -|-------------|-----------|-----------|-----------|--------| -| aten::addmm | 10240     | 8192      | 2048      | 24     | -| aten::mm    | 10240     | 2048      | 8192      | 24     | -| aten::mm    | 8192      | 2048      | 10240     | 24     | - -We can interpret each entry based on the GEMM shapes: -- The `aten::addmm` call represents the forward pass GEMM (`X @ Wᵀ`). -- The first `aten::mm` call corresponds to the backward pass for `dX` (`dY @ W`). -- The second `aten::mm` call represents the backward pass for `dW` (`dYᵀ @ X`). - -Each of these operations appears once per layer in the network. Given that GPT-3-XL has 24 layers, each of these GEMMs is called 24 times per training step, aligning with the 'count' column in the table. - ---- - - -## ⚙️ How PyTorch Calls BLAS - -To fully grasp how PyTorch leverages BLAS for operations like GEMM, we must first understand the fundamental concept of **memory layout** for tensors and how BLAS libraries interpret the data buffers they receive. - -### Memory Layout and Stride - -Despite tensors often being represented as multi-dimensional arrays, their elements are stored in linear memory. For a 2D matrix, the two primary storage conventions are: - -- **Row-major**: Elements of the same row are stored consecutively in memory. PyTorch adopts this as its default layout. -- **Column-major**: Elements of the same column are stored consecutively in memory. Many traditional BLAS libraries primarily optimize for this layout. - -PyTorch's `.stride()` method provides insight into a tensor's memory arrangement. It returns a tuple where each value indicates the byte (or element, depending on datatype size) distance in linear memory to move to the next element along that dimension. -- For a 2D tensor `T[i][j]` in **Row-major** layout, `.stride()` is typically `(num_cols, 1)`. Moving to `T[i][j+1]` requires stepping 1 element, while moving to `T[i+1][j]` requires stepping `num_cols` elements. -- For a 2D tensor `T[i][j]` in **Column-major** layout, `.stride()` is typically `(1, num_rows)`. Moving to `T[i+1][j]` requires stepping 1 element, while moving to `T[i][j+1]` requires stepping `num_rows` elements. - ---- - -### BLAS Transpose and Row-Major Output - -The core BLAS GEMM routine typically computes $C = \alpha \cdot op(A) \cdot op(B) + \beta \cdot C$, where $op(X)$ is either $X$ or $X^T$ depending on the `transA` and `transB` flags ('N' for No Transpose, 'T' for Transpose) passed to the function. By default, BLAS expects input matrices corresponding to the 'N' flag to be in column-major layout. Crucially, the resulting matrix $C$ is written into the output buffer in **column-major** format by default. - -PyTorch, however, uses row-major layout internally and desires the result of a GEMM operation to also be in row-major layout *without* an extra copy or transpose step outside of the BLAS call. PyTorch achieves this by cleverly leveraging the `trans` flags and the relationship between row-major and column-major layouts. - -A matrix $M$ stored in row-major memory has the exact same element ordering as the matrix $M^T$ stored in column-major memory. PyTorch uses this identity. To get a row-major result $C$ from a BLAS call that outputs in column-major, PyTorch requests BLAS to compute $C^T$ and write it in column-major. Since $C^T$ in column-major is $C$ in row-major, the output buffer will contain the desired row-major $C$. - -Mathematically, the operation $C = A @ B$ (where $A, B, C$ are desired in row-major) is equivalent to computing $C^T = (A @ B)^T = B^T @ A^T$. PyTorch therefore configures the BLAS call to compute $B^T @ A^T$ using the row-major data of $B$ and $A$. - -Here's how the `transA` and `transB` flags work in this context when passing **row-major data** to BLAS via a wrapper like PyTorch's: -- Passing row-major data for matrix $M$ with `trans = 'T'` tells BLAS to mathematically treat this data as $M$. (BLAS expects row-major data for 'T' if it wants to use the matrix directly). -- Passing row-major data for matrix $M$ with `trans = 'N'` tells BLAS to mathematically treat this data as $M^T$. (BLAS expects column-major data for 'N'; giving it row-major data makes it see the transpose). - -So, to compute $C^T = B^T @ A^T$ using row-major data for $B$ and $A$ and get $C$ row-major in the output buffer: -- Pass $B$'s row-major data as the first operand data (`A_data` in BLAS call). To make BLAS see $B^T$, use `transA = 'N'`. -- Pass $A$'s row-major data as the second operand data (`B_data` in BLAS call). To make BLAS see $A^T$, use `transB = 'N'`. -- The BLAS call becomes `gemm(transA='N', transB='N', ..., B_data, ..., A_data, ...)`. This computes $B^T @ A^T = C^T$. The result $C^T$ is written in column-major into the output buffer, which is precisely the desired $C$ in row-major. - -This standard trick using `transA='N'` and `transB='N'` with swapped, row-major inputs is a common way PyTorch achieves row-major output for a general matrix multiply `C = A @ B` where A, B are row-major. - ---- - - -### Linear Layer: `Y = X @ Wᵀ` - -For a linear layer computation `Y = X @ Wᵀ`, where `X` (`[M, K]`) and `W` (`[N, K]`) are in row-major layout, PyTorch desires `Y` (`[M, N]`) also in row-major. To achieve this with a BLAS routine outputting column-major, PyTorch configures BLAS to compute $Y^T = W @ X^T$. - -This involves a BLAS call computing $op(A) @ op(B)$ where $op(A)$ is $W$ and $op(B)$ is $X^T$. Using the rule that row-major data with `trans='T'` yields the matrix ($M$) and `trans='N'` yields the transpose ($M^T$): -- BLAS operand A uses $W$'s row-major data. To see $W$, `transA = 'T'`. -- BLAS operand B uses $X$'s row-major data. To see $X^T$, `transB = 'N'`. - -The BLAS call uses `(transA='T', transB='N')` with $W$'s data as the first operand and $X$'s data as the second. It computes $W @ X^T = Y^T$, writing the result in column-major, which PyTorch interprets as the desired row-major $Y$. - -### Backward Pass Operations: - -The backward pass similarly uses GEMMs configured to produce row-major gradients: - -- **`dX = dY @ W`**: With `dY` (`[M, K]`) and `W` (`[K, N]`) row-major, we need `dX` (`[M, N]`) row-major. BLAS computes $dX^T = W^T @ dY^T$. - - BLAS operand A uses $W$'s row-major data. Needs $W^T \implies$ `transA = 'N'`. - - BLAS operand B uses $dY$'s row-major data. Needs $dY^T \implies$ `transB = 'N'`. - - BLAS call uses `(transA='N', transB='N')` on $W$'s and $dY$'s data, computing $W^T @ dY^T$. - -- **`dW = dYᵀ @ X`**: With `dY` (`[K, N]`) and `X` (`[K, M]`) row-major, we need `dW` (`[N, M]`) row-major. BLAS computes $dW^T = X^T @ dY$. - - BLAS operand A uses $X$'s row-major data. Needs $X^T \implies$ `transA = 'N'`. - - BLAS operand B uses $dY$'s row-major data. Needs $dY \implies$ `transB = 'T'`. - - BLAS call uses `(transA='N', transB='T')` on $X$'s and $dY$'s data, computing $X^T @ dY$. - -In summary, for PyTorch's row-major operations: -- Forward pass `Y = X @ Wᵀ` maps to BLAS calculating $W @ X^T$ using `(T, N)` flags on the row-major data of $W$ and $X$. -- Backward pass `dX = dY @ W` maps to BLAS calculating $W^T @ dY^T$ using `(N, N)` flags on the row-major data of $W$ and $dY$. -- Backward pass `dW = dYᵀ @ X` maps to BLAS calculating $X^T @ dY$ using `(N, T)` flags on the row-major data of $X$ and $dY$. - -Let’s revisit the GPT-3-XL model gemm table from **Tracelens**: - -| name        | param: M | param: N | param: K | param: bias | param: stride_A | param: stride_B | param: transpose | -|--------------|-----------|-----------|-----------|--------------|------------------|------------------|--------------------| -| aten::addmm  | 10240     | 8192      | 2048      | TRUE         | (2048, 1)        | (1, 2048)        | (True, False)      | -| aten::mm     | 10240     | 2048      | 8192      | FALSE        | (8192, 1)        | (2048, 1)        | (False, False)     | -| aten::mm     | 8192      | 2048      | 10240     | FALSE        | (1, 8192)        | (2048, 1)        | (False, True)      | - -This table shows how `aten::addmm` (forward) and `aten::mm` (backward) calls map to underlying GEMM operations. The `param: M, N, K` values are likely the dimensions of the *PyTorch operation result* (`M x N` with inner dim `K`). The `param: transpose | (transA, transB)` are the BLAS flags used by the wrapper for the operands passed to the BLAS call. - -Let's interpret the trace entries based on our understanding that PyTorch uses row-major data and BLAS receives flags to mathematically interpret this data for computing the transpose of the desired result: - -1. **`aten::addmm` (Forward)**: Corresponds to `Y = X @ Wᵀ`. Trace flags: `(True, False)`. This matches the `(T, N)` needed for BLAS to compute $W @ X^T$. -2. **First `aten::mm` (Backward dX)**: Corresponds to `dX = dY @ W`. Trace flags: `(False, False)`. This matches the `(N, N)` needed for BLAS to compute $W^T @ dY^T$. - -3. **Second `aten::mm` (Backward dW)**: Corresponds to `dW = dYᵀ @ X`. Trace flags: `(False, True)`. This matches the `(N, T)` needed for BLAS to compute $X^T @ dY$. - - -This confirms how the trace flags correspond to the BLAS transpose configurations used with row-major input data to achieve row-major output via the $C^T = B^T A^T$ trick. - - -Lets summarize our understanding of the transpose flag by writing pseudo code for the logic: - -*(Note: This Python code snippet provides a simplified view; PyTorch's actual implementation is more intricate, accounting for the specific GEMM variant and output requirements.)* - - - -```python - -def is_col_major(T): -    return T.stride(0) == 1 and T.stride(1) >= T.shape[0] - -def get_blas_transpose_flags(A, B): -    transA = 'N' if is_col_major(A) else 'T' # If A is col-major, BLAS sees it as is ('N') -    transB = 'N' if is_col_major(B) else 'T' # If B is col-major, BLAS sees it as is ('N') -    return transA, transB - -``` - -## ⚠️ Edge Cases - - --   One common assumption is that flattening a tensor shape like `[B, L, d_model]` to `[B⋅L, d_model]` is a cost-free metadata operation. This is only true if the last dimension (`d_model`) is contiguous. - --   If the last dimension is not contiguous, PyTorch may be forced to insert a **copy** or **transpose** operation to create a physically contiguous tensor that BLAS can work with efficiently. - --   Furthermore, even if the tensor layout *could* theoretically be used by BLAS (e.g., certain striding patterns), some highly tuned BLAS libraries might lack kernels optimized for those specific layouts. In such instances, a **copy** or **transpose buffer** is inserted behind the scenes by PyTorch or the BLAS wrapper. Consequently, what the BLAS routine actually operates on might not be the original tensor directly, but rather a **temporary buffer** created for compatibility or performance. - diff --git a/docs_original/conceptual/shape_metadata_guide.md b/docs_original/conceptual/shape_metadata_guide.md deleted file mode 100644 index 55ccd80ae..000000000 --- a/docs_original/conceptual/shape_metadata_guide.md +++ /dev/null @@ -1,151 +0,0 @@ - - -# Tensor Shape Metadata in PyTorch Profiler Traces - -A practical guide to understanding when tensor shapes are available in profiler traces and how to get them when they're missing. - ---- - -## Part 1: When Do We Have Shapes? - -### The Simple Rule - -**Shapes are available when an operation goes through PyTorch's dispatcher as a `cpu_op`.** - -In profiler traces, look for events with: -- `"cat": "cpu_op"` -- `"args": { "Input Dims": [...], "Input type": [...] }` - -### What HAS Shapes - -| Operation Type | Example | Why It Works | -|---------------|---------|--------------| -| Native ATen ops | `aten::mm`, `aten::linear` | Built into PyTorch dispatcher | -| Registered custom ops | `torch.ops.mylib.my_op` | Registered via `torch.library` | -| TorchScript ops | JIT-compiled functions | Goes through dispatcher | -| Custom `autograd.Function` | User-defined forward | Forward call is dispatched | -| Distributed collectives | `record_param_comms` | Instrumented by PyTorch | - -### What DOESN'T Have Shapes - -| Operation Type | Example | Why It Fails | -|---------------|---------|--------------| -| Plain Python functions | `def my_kernel(x, y): ...` | Bypasses dispatcher | -| Triton kernels | `@triton.jit` | Called as Python function | -| FlashInfer ops | GEMM, MoE, Attention | Registration intentionally disabled | -| Backward engine events | `autograd::engine::evaluate_function:*Backward` | Empty inputs passed to profiler | - -### Quick Reference: Event Categories - -``` -cpu_op → Usually HAS shapes (exception: backward events) -python_function → NO shapes -kernel → NO shapes (GPU-side event) -cuda_runtime → NO shapes (API-level event) -``` - ---- - -## Part 2: How to Get Shapes When Missing - -### Option 1: Register as Custom Op - -Wrap the operation with `torch.library.custom_op`: - -```python -@torch.library.custom_op("mylib::triton_matmul", mutates_args=()) -def triton_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - return my_triton_kernel(a, b) - -# Now call via torch.ops -result = torch.ops.mylib.triton_matmul(a, b) -``` - -**Result:** Event appears as `cpu_op` with `Input Dims`. - -**What is `mutates_args`?** - -Tells PyTorch which input arguments the function modifies in-place: -- `mutates_args=()` — No inputs are modified (pure function) -- `mutates_args=("output",)` — The `output` argument is modified in-place -- Required for correctness with `torch.compile` and autograd - -### Option 2: Lighter-Weight Registration (vLLM approach) - -Use the lower-level `Library` API directly: - -**Reference:** [vllm/utils/torch_utils.py:742 - `direct_register_custom_op`](https://github.com/vllm-project/vllm/blob/0b225fb7b22f8ae1f5fc8ee618640ae0983c76de/vllm/utils/torch_utils.py#L742-L780) - -```python -from torch.library import Library -from torch._library.infer_schema import infer_schema - -my_lib = Library("mylib", "FRAGMENT") - -def register_op(op_name, op_func, mutates_args=None): - schema = infer_schema(op_func, mutates_args=mutates_args or []) - my_lib.define(op_name + schema) - my_lib.impl(op_name, op_func, dispatch_key="CUDA") - -# Register -register_op("triton_matmul", triton_matmul_impl) -``` - -### Option 1 vs Option 2: When to Use Which? - -| Aspect | Option 1: `custom_op` | Option 2: `Library.define()` | -|--------|----------------------|------------------------------| -| **Simplicity** | ✅ Decorator, minimal code | ❌ More boilerplate | -| **Overhead** | Higher (full dispatcher) | Lower (direct to CUDA) | -| **torch.compile** | ✅ Full support | ✅ Works with `register_fake` | -| **Use when...** | Prototyping, one-off ops | Performance-critical paths | - ---- - -## Framework-Specific Status - -| Framework | Current State | How to Get Shapes | -|-----------|--------------|-------------------| -| **PyTorch ATen** | ✅ Has shapes | Already works | -| **vLLM (standard)** | ✅ Has shapes | Uses `direct_register_custom_op` | -| **vLLM (OAI Triton)** | ❌ Missing | Needs registration | -| **FlashInfer** | ❌ Disabled | Set `FLASHINFER_ENABLE_PROFILER_METADATA=1` (proposed) | -| **SGLang (CUDA)** | ✅ Has shapes | Uses `torch.ops.sgl_kernel.*` | -| **SGLang (Triton)** | ❌ Missing | Needs registration | - ---- - -## Summary - -1. **Shapes are tied to dispatcher registration** — If it's a `cpu_op`, it has shapes -2. **The fix is straightforward** — Register operations via `torch.library` -3. **Start simple** — Use `@torch.library.custom_op` first -4. **Optimize if needed** — Switch to `Library.define()` for lower overhead -5. **FlashInfer disabled it on purpose** — Performance vs. observability trade-off; can be re-enabled - ---- - -## Appendix - -### A1: Backward Events and Sequence Number Linking - -Backward engine events (`autograd::engine::evaluate_function:*Backward`) are `cpu_op` but don't have `Input Dims`. They have `Sequence number` instead, which links to the corresponding forward op. - -```python -def get_backward_shapes(trace_events, backward_event): - seq_num = backward_event["args"].get("Sequence number") - if seq_num is None: - return None - - # Find forward op with same sequence number - for event in trace_events: - if event.get("cat") == "cpu_op" and \ - event["args"].get("Sequence number") == seq_num and \ - "Backward" not in event.get("name", ""): - return event["args"].get("Input Dims") - return None -``` diff --git a/docs_original/conceptual/torch_profiling_analysis.md b/docs_original/conceptual/torch_profiling_analysis.md deleted file mode 100644 index 86607f949..000000000 --- a/docs_original/conceptual/torch_profiling_analysis.md +++ /dev/null @@ -1,144 +0,0 @@ - - -## Understanding PyTorch Trace: A Deep Dive into Execution and Analysis - -You followed the [PyTorch profiling guide](../conceptual/torch_profiling_guide.ipynb), collected a trace, and even opened it in [Perfetto](https://ui.perfetto.dev/). But what does that trace actually mean, and how do you read insights from it? -That’s exactly what we’ll unpack here. - -When you run a PyTorch model on a GPU, there’s a hidden dance between the CPU (host) and the GPU (accelerator). PyTorch trace lets you peek into this choreography—revealing how your code executes, where time is spent, and what might be slowing you down. - -This blog is the first in a series—here we’ll focus on **understanding single-GPU traces**. In a follow-up, we’ll move on to **multi-GPU profiling**. - - -### The Execution Model: Host, GPU, and the Async Magic -![Async Execution Example](./imgs/profiling_analysis_fig1.png) -To understand profiling, you need a mental model of how PyTorch (or any application) runs code across the CPU and GPU. -**Figure 1**: Async Execution Example - -The **host (CPU)** is the conductor. It manages processes and threads, moving through the application call stack until it issues GPU runtime or driver API calls such as `cudaLaunchKernel` or `hipLaunchKernel`. The **GPU** is the performer. It executes tasks with massive parallelism, often involving thousands of threads working on matrix multiplications, convolutions, or reductions. - -Most of these commands are asynchronous. The host queues the "commands" (kernels) and immediately continues, while the GPU independently drains its queue. This overlap is what enables high performance. Sometimes, however, the host introduces synchronization points such as `cudaDeviceSynchronize` or `hipDeviceSynchronize`. In those cases the CPU waits until the GPU finishes all queued work. Synchronization is necessary for correctness, but too many syncs can quickly turn into bottlenecks. - -One more twist: host events come with detailed call stacks (Python → C++ → runtime), so you can see exactly how an operation was triggered. GPU events, on the other hand, are flat. A GPU kernel cannot call another GPU kernel on its own—the launch must always come from the host. That’s why GPU traces appear as a simple queue of kernels, memcopies, and memsets with no stack information. - -A **kernel** itself is just a GPU function that runs across thousands of lightweight threads in parallel. In the trace, each kernel shows up as a single event: you can see when it started, how long it ran, and on which stream—but not what happened inside. To peer into that level of detail (thread execution, memory transactions, warp occupancy), you need dedicated kernel profilers such as [rocprof-compute](https://rocm.docs.amd.com/projects/rocprofiler-compute/en/latest/) for AMD GPUs or Nsight Compute for NVIDIA GPUs. These tools complement the trace view by showing what happens inside the GPU during a kernel event. - - - - ---- - -### A Simple Example - -Let’s dig deeper into what a trace actually shows. Below is a snippet from our trace. - -![Perfetto trace showing idle time between kernels](./imgs/profiling_analysis_fig2.png) -**Figure 2**: Perfetto trace highlighting idle time between convolution and batch normalization kernels. - -The red dashed region marks idle time: the GPU is waiting for the host to issue the next command. This happens when the CPU frontend cannot keep up with the GPU’s execution speed. Reducing such idle gaps is an important optimization goal. To see why these gaps occur, let’s follow the path from a high-level PyTorch call down to the GPU kernels that actually run. - ---- - -Take a simple example: - - y = nn.Conv2d(3, 64, kernel_size=3)(x) - -At the Python level, this looks like a single operation. In the trace, however, it expands into several layers: -- **Python frontend**: `nn.Conv2d` in `torch.nn` -- **ATen**: the call lowers into PyTorch’s tensor library, [ATen](https://pytorch.org/cppdocs/notes/aten.html), and shows up as `aten::conv2d` in the trace -- **Backend wrapper**: ATen provides wrappers that call into vendor libraries. On AMD GPUs, you’ll see `aten::miopen_convolution`, which wraps [MIOpen](https://rocm.docs.amd.com/projects/MIOpen/en/latest/) commands. On NVIDIA GPUs, the equivalent is `aten::cudnn_convolution`, which wraps [cuDNN](https://developer.nvidia.com/cudnn) calls -- **GPU kernels**: the backend library enqueues device kernels such as `igemm_fwd_gtcx2_nhwc` that perform the actual convolution on the GPU - -This gives us a clear understanding of how high-level code is translated into GPU execution. - - ---- -Another important detail is **tensor metadata**. Python-level operations in the trace do not carry shape information. But if you enable `record_shapes=True` in the profiler, the "cpu_op" category (like `aten::convolution`) include **input shapes, strides, and dtypes**. - -For example, in the figure below the first tensor is the activation (`[5, 64, 56, 56]`) and the second tensor is the convolution filter (`[64, 64, 3, 3]`). -![Perfetto trace showing recorded input shapes](./imgs/profiling_analysis_fig3.png) -*Figure 3: Shapes recorded on the backend op when `record_shapes=True`.* - -This metadata becomes very useful for deeper analysis and debugging—for instance, when matching trace events to model architecture, analyzing kernel efficiency, or spotting unusual stride/dtype patterns. - ---- - -That’s all you need for a clean first pass: know what the host and GPU are doing, read the timeline, and use recorded shapes to ground what you see. If you want UI shortcuts, how memcpys show up, or the structure of the raw trace file, jump to the appendix below. - -**Your Profiling Journey Begins**: Perfetto view give you the raw signals; turning them into insight is a skill you’ll build with practice. And TraceLens is here to turbocharge your journey. - ---- -### Appendix - -#### Perfetto UI Tips - -Perfetto is a powerful trace viewer, but it takes some practice to navigate effectively. A few basics: - -- **Zoom and pan** with `Ctrl + scroll` to zoom in/out; click and drag to pan. -- **Event details**: click any event to open the *Current Selection* panel below. This shows start time, duration, and arguments. With `record_shapes=True`, backend ops (in the `cpu_op` category) also show tensor shapes, dtypes, and strides. - -Perfetto also links host launches and GPU execution with arrows called *flows*. These are crucial for connecting what you see on the CPU timeline with what actually runs on the GPU. - -When you select a runtime launch event such as `hipExtModuleLaunchKernel`, the **Following Flow** jumps you forward to the GPU kernel it triggered: - -![Following flow from host launch to GPU kernel](./imgs/profiling_analysis_fig4.png) -*Figure 4: Following flow from a host launch (`hipExtModuleLaunchKernel`) to the corresponding GPU kernel event.* - -Conversely, when you select a GPU kernel event, the **Preceding Flow** takes you back to the runtime call on the host that launched it: - -![Preceding flow from GPU kernel back to host launch](./imgs/profiling_analysis_fig5.png) -*Figure 5: Preceding flow from a GPU kernel (`SubTensorOpWithScalar1d`) back to its launch on the host.* - -These flows are the bridge between the Python-level trace and the GPU execution timeline. They let you answer both: -- *Which Python/ATen op caused this kernel to run?* -- *What GPU work did this runtime call produce?* - - - ---- - -#### Memory Copy Events -Not all GPU activity is compute. Profiling traces also show **memory transfers**: - -- **H2D (Host to Device)**: copies data from CPU to GPU, usually synchronous and PCIe bandwidth-limited. -- **D2H (Device to Host)**: copies results back to CPU, also synchronous and PCIe bandwidth-limited. -- **D2D (Device to Device)**: moves data between GPU buffers, asynchronous and limited by HBM bandwidth. - -Recent versions even record measured bandwidth for these events in the trace arguments. -Importantly, memcpy events use the GPU’s DMA engines, not compute cores—so they don’t directly compete with kernel execution. - ---- - -#### JSON Format -Behind the Perfetto UI, the PyTorch profiler saves traces as JSON. Each entry is an event dictionary with: - -- **Timestamps**: `ts` (start), `dur` (duration). -- **Process/thread IDs**: `pid`, `tid`. CPU events use real PIDs/TIDs; GPU events use pseudo-PIDs for devices and TIDs for streams. -- **Category**: e.g. `python_function`, `cpu_op`, `cuda_runtime`, `kernel`, `gpu_memcpy`. -- **Args**: extra info such as shapes, dtypes, strides, or bandwidth. - -The JSON format makes traces scriptable: you can parse them to build custom reports or run automated analysis outside Perfetto. This is what exactly TraceLens does. - - -#### Autograd in the Trace - -Autograd introduces another dimension to the trace: the forward and backward passes are captured on **different threads within the same process**. In the timeline, you’ll typically see: -- `python3.x [main thread]` for the forward pass, -- `pt_autograd_0` (or similar) for the backward pass. - -These are linked at the `aten::convolution` layer of the call stack. Perfetto uses flows to connect the forward convolution op to its corresponding backward node. - -![Forward and backward convolution linked in the trace](./imgs/profiling_analysis_fig6.png) -*Figure 6: Example of `aten::convolution` (forward) linked to `ConvolutionBackward0` (backward) via flows.* - -Key points: -- The linkage happens at **`aten::convolution`**, not `aten::conv2d`. The higher-level `aten::conv2d` call funnels into `aten::convolution` before reaching the backend. -- Forward and backward both eventually call into the same GPU backend (MIOpen/cuDNN), so you’ll see similar kernel launches in both directions. -- The backward pass typically runs more kernels, since it must compute gradients for both activations and weights, making it more expensive than forward. -- Because autograd runs on its own thread, you can easily separate model execution (forward path) from gradient computation (backward path) when analyzing traces. - -Together with the flows we discussed earlier, this lets you connect *a forward convolution op* to its *exact backward counterpart* and then follow the chain down to GPU kernels. diff --git a/docs_original/conceptual/trace2tree_motivation.md b/docs_original/conceptual/trace2tree_motivation.md deleted file mode 100644 index 0cbe260d1..000000000 --- a/docs_original/conceptual/trace2tree_motivation.md +++ /dev/null @@ -1,36 +0,0 @@ - - -# Trace2Tree Motivation - -## Limitations of Direct Kernel Analysis - -Directly inspecting GPU kernel names has **two fundamental limitations**: - -* **Ambiguous semantics (and weak reproducibility)**: A single kernel name can map to many different computations depending on shape, dtype, strides/layout, etc. Shape strongly affects performance—one shape may select a fast tiled path while another shape (same op type) falls onto a slower algorithm. Because the name omits this argument context, you cannot reliably understand, compare, or reproduce the workload from the kernel string alone. Further, many kernel names are cryptic and unreadable (e.g., `Cijk_Ailk_Bljk_*`, `void cutlass_*`). - -* **Platform‑dependent / unstable naming (and weak cross‑platform comparison)**: The same high‑level operation appears under different kernel names across pltforms. For example: a single GEMM shows up as `nvjet_*` or `cutlass_*` on NVIDIA H100, and as a Tensile kernel `Cijk_Ailk_Bljk_*` on AMD MI300. These names also shift across software versions. Raw kernel strings are therefore not a stable abstraction for comparison. - -## What Trace2Tree Does - -**Trace2Tree** reconstructs a full call tree from the Python front‑end down to each GPU kernel. The tree has exactly four layers: - -1. **Python front end** – user code / `nn.Module`. -2. **Operation** – On PyTorch it is the dispatch operation ("cpu_op") e.g. `aten::mm`, `aten::addmm`, etc. On JAX the corresponding layer is the HLO operations. This is the layer that contains the arg information to contextualize the kenrel. -3. **HIP / CUDA runtime** – launch API calls. -4. **GPU kernel** – executed kernel. - -## How This Solves the Problem - -* **Disambiguates semantics**: Argument metadata at the backend op layer lets us group identical computations, attribute time, and deterministically reproduce slow cases. -* **Enables fair comparison**: Operations (`aten::mm`, HLO) are stable across platforms; by anchoring analysis there, we can compare the same operation and arguments regardless of how kernels are named underneath. -* **Flexible attribution**: GPU time can be viewed at any level—module (via its backend ops), dispatch op, runtime, or kernel—depending on the question. As an additional benefit, we can also attribute time all the way up to the Python nn.Module level, making performance insights accessible even to users outside the performance engineering field. This helps bridge the gap between model developers and low-level hardware execution. - -## Takeaway - -Kernel names are volatile and context‑free. Trace2Tree anchors analysis at the stable backend operation, enriches it with arguments, and maps the full execution stack to deliver portable, interpretable performance insight. - -That said, **kernel names are often useful** — they can offer clues about the backend implementation, algorithm variant, or compiler choices. TraceLens intends to serve as a **one-stop solution** for extracting every bit of useful signal from a trace file. Therefore, it includes features to extract and parse relevant information from kernel names where applicable. But we treat them as supplementary, not foundational. diff --git a/docs_original/gemm_dim_eff.md b/docs_original/gemm_dim_eff.md deleted file mode 100644 index 332d3739b..000000000 --- a/docs_original/gemm_dim_eff.md +++ /dev/null @@ -1,190 +0,0 @@ - - -## What's New: Deeper Efficiency Metrics -This document details a new feature in TraceLens, providing deeper insights into the efficiency of GEMM (General Matrix Multiplication) operations, specifically focusing on Tensile kernels used on ROCm GPU. This analysis complements the existing Roofline metrics by breaking down performance limitations related to how gemm computations are tiled and scheduled onto the GPU. - -By default, the tool provides Roofline metrics summarizing overall performance for operations like `aten::mm`: - -**Original Output Example:** - -| name | M | N | K | bias | FLOPS/Byte_first | TFLOPS/s_mean | -|----------|----------|----------|----------|-------------|------------------|----------------| -| aten::mm | 2048 | 2048 | 10240 | False | 930.90 | 521.57 | - - -The enhanced analysis now adds specific efficiency metrics for Tensile GEMM kernels, derived from the kernel's structure and the problem dimensions: - -**New Output Example (Tensile Kernels):** -Checkout the examples/gemm_dim_eff.ipynb for example usage. - -| name | M | N | K | ... | mt_m | mt_n | num_tiles | tile_eff | wq_eff | dim_eff | ... | TFLOPS/s_mean | -|------------|-------|------|------|-----|------|-----|---|----------|--------|----------|--------|--------| -| aten::mm | 2048 | 2048 | 10240 | ... | 256 | 64 | 256 | 1.00 | 0.84 | 0.84 | ... | 521.57 | - - - ---- - -## Key New Metrics - -- **`mt_m`, `mt_n`**: Macro-tile dimensions extracted from the kernel name. -- **`num_tiles`**: Total number of tiles after padding. -- **`tile_eff`**: Tile Quantization Efficiency. Measures efficiency loss due to input matrix dimensions not being perfectly divisible by tile dimensions. -- **`wq_eff`**: Wave Quantization Efficiency. Measures efficiency loss due to the total number of tiles not perfectly filling all compute units in the final processing wave. -- **`dim_eff`**: Net Dimension Efficiency. The product of tile_eff and wq_eff, representing the combined efficiency impact of tiling and scheduling. - ---- - -## Understanding the Concepts - -### 1. Tile Quantization Efficiency (`tile_eff`) - -Tiled computations divide matrices into smaller sub-blocks (tiles: `mt_m x mt_n`) for processing. If the matrix dimensions (M, N) are not exact multiples of these tile sizes, the implementation effectively pads the matrices to the nearest multiple of the tile size. - -**Padded Dimensions:** - - M_pad = ceil(M / mt_m) × mt_m - N_pad = ceil(N / mt_n) × mt_n - -This padding introduces extra computations on regions that don’t contribute to the final result. - -**Tile Efficiency:** - - tile_eff = (M × N) / (M_pad × N_pad) - -Values < 1 indicate wasted computation due to padding. - ---- - -### 2. Wave Quantization Efficiency (`wq_eff`) - -GPUs execute tiles across many Compute Units (CUs) also known as Streaming multi-processors (SMs). If the total number of tiles isn’t a multiple of the number of available CUs, the final wave will leave some CUs idle. - -**Total Tiles:** - - B = (M_pad × N_pad) / (mt_m × mt_n) - -**Number of Waves:** - - num_waves = ceil(B / num_cus) - -**Wave Efficiency:** - - wq_eff = B / (num_waves × num_cus) - ---- - -### 3. Net Dimension Efficiency (`dim_eff`) - -This is the combined impact: - - dim_eff = tile_eff × wq_eff ---- - -## Why These Metrics Matter: Diagnosing Bottlenecks - -For compute-bound GEMMs, actual performance is often less than peak theoretical. These metrics help explain why: - -- **Low tile_eff**: Indicates high padding overhead. -- **Low wq_eff**: Indicates poor utilization of compute units in the final wave. - -GEMM tuning can help improving these metrics to a certain extent. - -However, these are just part of the picture. Even when tiling and wave quantization are optimal, there can still be performance degradation due to: - -- **Shader clock (`sclk`) throttling**: Some workloads may not run at peak clock speeds due to power or thermal constraints. -- **Cache behavior**: Cache misses can occur even for compute-bound GEMMs, affecting throughput and instruction efficiency. - -These factors should be analyzed alongside the dimension efficiency metrics to get a complete performance picture. - ---- - -## How It's Calculated - -### Step-by-Step: - -1. **Identify Problem Shape**: Extract M, N, K from input tensors. - - > Note: In BLAS libraries, M and N are swapped relative to PyTorch's view as BLAS libraries are column major while PyTorch is row major. This mapping is handled internally. See [PyTorch source - Blas.cpp](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/Blas.cpp#L102-L129) for more detail. This is accounted in TraceLens as well when computing the tiles across the M and N dimensions. - -2. **Extract Tile Size**: Parse `mt_m`, `mt_n` from kernel name (e.g., `MT256x144x32` → `mt_m = 256, mt_n = 144`). - -3. **Calculate Tiles:** - ``` - float_tiles_m = M / mt_m - float_tiles_n = N / mt_n - tiles_m = ceil(float_tiles_m) - tiles_n = ceil(float_tiles_n) - num_tiles = tiles_m * tiles_n - ``` - -4. **tile_eff**: - ``` - tile_eff = (M * N) / (tiles_m * mt_m * tiles_n * mt_n) - ``` - -5. **wq_eff** (assume `num_cus` known, e.g., 304 for MI300X): - ``` - float_rounds = num_tiles / num_cus - rounds = ceil(float_rounds) - wq_eff = num_tiles / (rounds * num_cus) - ``` - -6. **dim_eff**: - ``` - dim_eff = tile_eff * wq_eff - ``` - ---- - -## Calculation Examples - -**Assuming `num_cus = 304` (e.g., MI300X)** - -### Example 1: Perfect Tiling, Suboptimal Wave Quantization -- **Input**: M = 10240, N = 2048, K = 2048 -- **Kernel Tile**: `mt_m = 256`, `mt_n = 64` - -```text -float_tiles_m = 40 -float_tiles_n = 32 -tiles_m = 40, tiles_n = 32 -tile_eff = 1.0 -num_tiles = 1280 -float_rounds = 4.21 -rounds = 5 -wq_eff = 0.842 -dim_eff = 0.842 -``` - -### Example 2: Slight Padding, Good Wave Quantization -- **Input**: M = 2048, N = 10240, K = 2048 -- **Kernel Tile**: `mt_m = 256`, `mt_n = 144` - -```text -float_tiles_m = 8 -float_tiles_n = 71.11 -tiles_m = 8, tiles_n = 72 -tile_eff ≈ 0.9877 -num_tiles = 576 -float_rounds = 1.895 -rounds = 2 -wq_eff ≈ 0.947 -dim_eff ≈ 0.935 -``` - ---- - -## Important Considerations - -- **Scope**: This analysis assumes standard tiled GEMMs. Techniques like Stream-K or Split-K are not yet modeled. -- **Relevance**: These metrics are most useful for compute-bound GEMMs. Use FLOPS/Byte to determine whether a GEMM is compute- or memory-bound. - - -## Example Usage -Checkout the [Notebook](../examples/gemm_dim_eff.ipynb) for example usage. -This is a new feature so please report any issues or suggestions to the TraceLens team. diff --git a/docs_original/generate_multi_rank_collective_report_pytorch.md b/docs_original/generate_multi_rank_collective_report_pytorch.md deleted file mode 100644 index 2905f1358..000000000 --- a/docs_original/generate_multi_rank_collective_report_pytorch.md +++ /dev/null @@ -1,180 +0,0 @@ - - -# Generate Multi-Rank Collective Report - -This utility analyzes **PyTorch JSON profile traces** from **multiple ranks** and produces a comprehensive **NCCL communication report** (Excel workbook and/or CSVs). - ---- - -## 🚀 Quick Start - -Run the script with **one** of: a directory containing trace files, a trace file **pattern** with a single `*` placeholder for rank id, or a **glob** pattern with a regex to extract the rank. - -### Directory mode -```bash -python generate_multi_rank_collective_report_pytorch.py --trace_dir /path/to/traces --world_size 8 -``` - -### Pattern mode (strict rank substitution) -```bash -python generate_multi_rank_collective_report_pytorch.py --trace_pattern "/logs/job123/rank*/trace.json" --world_size 8 -``` - -### Glob mode (tensorboard-style / arbitrary filenames) -When traces are produced by `torch.profiler.tensorboard_trace_handler`, filenames are typically not `rank{i}_trace.json`. Use `--trace_glob` with an optional `--rank_regex` to extract the rank id: -```bash -python generate_multi_rank_collective_report_pytorch.py \ - --trace_glob "/path/to/tensorboard/**/*.pt.trace.json.gz" \ - --rank_regex "rank\[(?P\d+)\]" \ - --world_size 16 -``` - -### Parallel mode (for faster processing) -```bash -python generate_multi_rank_collective_report_pytorch.py --trace_dir /path/to/traces --world_size 8 --use_multiprocessing -``` - -> **Note:** Speedup varies based on available CPU cores, trace file sizes, and disk I/O. - -> The pattern must contain exactly one `*`, which is replaced with `0, 1, ..., world_size-1` to enumerate rank file paths. - -### Installed entry point (if packaged) -If you install this script as part of a package with a console entry point, you can call it directly, e.g.: -```bash -TraceLens_generate_multi_rank_collective_report_pytorch --trace_dir /path/to/traces --world_size 8 -``` ---- - -## 📥 Input - -- **PyTorch JSON traces** captured per-rank (e.g., via `torch.profiler` with JSON export). -- Naming patterns: - - Directory mode (files inside `--trace_dir`): `rank0_trace.json`, `rank1_trace.json`, ... - - Pattern mode: `--trace_pattern "/path/to/rank*_trace.json"` - - Glob mode: `--trace_glob "/path/to/**/*.pt.trace.json.gz"` with `--rank_regex` for rank extraction - ---- - -## ⚙️ Command-Line Options - -> Provide **exactly one** of `--trace_dir`, `--trace_pattern`, or `--trace_glob`. - -| Argument | Default | Description | -|---|---|---| -| `--trace_dir` | `None` | Directory containing trace files (expects names like `rank*_trace.json`). | -| `--trace_pattern` | `None` | Template path with a **single** `*` placeholder for the rank id (strict substitution: `* → 0..world_size-1`). | -| `--trace_glob` | `None` | Glob for trace files with arbitrary names (supports `**`). Requires `--world_size` and uses `--rank_regex` to map files to ranks. | -| `--rank_regex` | `rank[\[\-_/]?(?P\d+)` | Regex used with `--trace_glob` to extract rank id. Must contain a named group `rank` or a single capture group. | -| `--world_size` | `Required` | Number of ranks in the distributed run. | -| `--output_xlsx_path` | _auto-inferred_ | Path to save the Excel workbook. If omitted, a sensible default is created (see Output Behavior). | -| `--output_csvs_dir` | `None` | If provided, writes each sheet as an individual `.csv` under this directory. | -| `--detailed_analysis` | `False` | Include detailed per-rank sheets (see below). | -| `--agg_metrics` | `mean median min max` | Aggregations to compute in summaries. Allowed: `mean`, `median`, `min`, `max` (space-separated). | -| `--use_multiprocessing` | `False` | Enable parallel trace loading using multiprocessing. Can provide speedup (system-dependent) but uses more CPU resources. | -| `--max_workers` | `os.cpu_count()` | Maximum number of worker processes for parallel loading (requires `--use_multiprocessing`). Override to limit resource usage if needed. | -| `--all2allv_heatmap` | `False` | Add an `nccl_all2allv_heatmap` sheet with per rank-pair send volumes across all all2allv invocations. Useful for diagnosing MoE/expert-parallel routing imbalance. | -| `--gpus_per_node` | _auto-detected_ | Number of GPUs per node. When known, `node_id` and `node_span` columns are added to report sheets, labeling each collective's process group as `intra_node` or `inter_node`. Auto-detected from trace `deviceProperties` if omitted. | - - ---- - -## 📊 Excel Workbook Sheets - -| Sheet Name | What it contains | -|---|---| -| `nccl_summary_implicit_sync` | Aggregated view of collectives that incurred implicit sync (AllReduce, AllGather, ReduceScatter, AllToAll balanced). Provides latency, skew, and aggregated bandwidth metrics (algorithm and bus). | -| `nccl_summary_long` | Aggregated NCCL operation stats per (rank, process group, collective, dtype, size). Includes counts and duration aggregates (sum/mean/std/min/max). | -| `nccl_long` | Per-event, per-rank NCCL records with timestamps, durations, stream, message sizes, and other attributes. Useful for drilling into specific slow ops. | -| `nccl_implicit_sync` | Per-collective implicit synchronization breakdown. Includes per-rank wait_time columns indicating time lost to implicit sync, plus timing/skew and bandwidth metrics. | -| `nccl_summary_all2allv` | Aggregated all2allv metrics by (Process Group, dtype). Includes throughput, wall time, size imbalance, and timing skew. Unlike implicit-sync collectives, all2allv uses aggregate throughput rather than algo/bus bandwidth since per-rank data sizes vary. | -| `nccl_all2allv` | Detailed All2AllV analysis: variable send/recv sizes and splits per rank, with timing and skew columns per rank. | -| `nccl_all2allv_heatmap` | *(when `--all2allv_heatmap` is set)* Per rank-pair total send volumes across all all2allv invocations. Useful for identifying hot pairs in MoE/expert-parallel routing. | -| `straggler_summary` | One row per rank with aggregated straggler metrics: total/mean wait time, how often the rank arrived last or first, and total NCCL duration. Sorted ascending by wait time so the straggler (lowest wait time) is at the top. See [Identifying Straggler Ranks](#-identifying-straggler-ranks). | - -_Notes:_ -- Summary sheets (`nccl_summary_*`, `nccl_summary_all2allv`) are **always** produced when the corresponding collective type exists in the trace; detailed per-event sheets (`nccl_long`, `nccl_implicit_sync`, `nccl_all2allv`) appear when `--detailed_analysis` is set. -- When `gpus_per_node` is known (auto-detected or set via `--gpus_per_node`), sheets that contain `rank` and/or `Process Group Ranks` columns gain `node_id` and `node_span` columns (fully aggregated summary sheets without these columns remain unchanged). `node_span` is `intra_node` when all ranks in the process group reside on the same node, or `inter_node` otherwise. You can filter or pivot on these columns in Excel to compare intra- vs inter-node communication. -- Column sets can evolve; the above reflects the provided example workbook. ---- - -## 📤 Output Behavior - -- If **neither** `--output_xlsx_path` nor `--output_csvs_dir` is specified, the tool writes an Excel file to: - - the provided `--trace_dir`, or - - the **lowest common directory** of the resolved trace files (from `--trace_pattern` or `--trace_glob`). -- If `--output_csvs_dir` is set, all sheets are written as individual CSV files under that directory. -- If `--output_xlsx_path` is set, a single Excel workbook containing all sheets is created. -- `openpyxl` is required for Excel; the tool will attempt to install it when missing. - -**Lowest Common Directory (pattern mode)** -When using `--trace_pattern`, the tool expands the pattern by replacing `*` with `0..world_size-1`, then computes the **lowest common ancestor directory** of those paths to pick a default output location. - ---- - -## 🔍 Identifying Straggler Ranks - -A **straggler** is the rank that consistently arrives last at collectives, forcing all other ranks to wait in implicit synchronization. - -### Quick answer: open `straggler_summary` - -The `straggler_summary` sheet is sorted so **the straggler is in the first row** (lowest total wait time). Example from an 8-rank Llama 70B FSDP run: - -| rank | total_wait_time_us | mean_wait_time_us | times_arrived_last | times_arrived_first | pct_arrived_last | num_collectives | total_nccl_dur_us | -|------|-------------------|-------------------|-------------------|--------------------|-----------------:|----------------:|------------------:| -| 4 | 27,195 | 55.7 | 420 | 6 | 86.1% | 488 | 3,910,753 | -| 5 | 4,660,353 | 9,549.9 | 39 | 10 | 8.0% | 488 | 8,463,896 | -| ... | ... | ... | ... | ... | ... | ... | ... | -| 0 | 13,358,753 | 27,374.5 | 3 | 320 | 0.6% | 488 | 17,203,432 | - -**How to read it:** - -| Column | What it tells you | -|--------|-------------------| -| `total_wait_time_us` | Sum of this rank's wait time across all collectives. The straggler has the **lowest** value — it arrives last, so it rarely waits. | -| `times_arrived_last` | Number of collectives where this rank was the last to arrive. The straggler dominates this column. | -| `times_arrived_first` | Number of collectives where this rank arrived first. The rank that appears here most often pays the biggest sync penalty. | -| `pct_arrived_last` | `times_arrived_last / num_collectives`. A rank at 86% is a persistent straggler. | -| `total_nccl_dur_us` | Total time inside NCCL kernels. The straggler typically has the **lowest** value — other ranks' durations are inflated by implicit-sync wait time. | - -### Next step: visualize with TraceFusion - -Once you know the straggler rank, use [TraceFusion](TraceFusion.md) to merge the straggler's trace with a fast rank (e.g. the one with the highest `times_arrived_first`) into a single file and open it in Perfetto UI. Viewing both timelines side-by-side makes it straightforward to see what compute or memory work is delaying the straggler before each collective. - -```python -from TraceLens import TraceFuse - -fuser = TraceFuse( - trace_files=["rank4_trace.json.gz", "rank0_trace.json.gz"], - output_file="straggler_vs_fast.json.gz", -) -fuser.fuse() -``` - ---- - -## Interpreting all2allv Metrics - -**Why different metrics from other collectives?** - -All2allv sends variable amounts of data per rank — there is no single "message size" or uniform workload. The standard NCCL `algo bw` and `bus bw` formulas assume ring/tree algorithms with equal data on every rank, so they don't apply. - -| Metric | What it tells you | -|--------|-------------------| -| `throughput (GB/s)` | Total data moved / wall-clock time of the collective. Compare against theoretical link bandwidth to gauge efficiency. | -| `wall_time (us)` | End-to-end time from first rank entering to last rank finishing. High values + low throughput = bottleneck. | -| `size_imbalance` | Ratio of max rank's data to mean. 1.0 = balanced. Values >> 1.0 indicate some ranks carry disproportionate load — common in MoE when some experts are "hotter" than others. | -| `max_rank_dur / min_rank_dur` | Spread in per-rank kernel time. Large spread + high imbalance = consider rebalancing expert routing or token dropping. | -| `skew in start time` | How far apart ranks enter the collective. High skew means some ranks are blocked by prior compute. | - -**Actionable guidance:** -- If `throughput` is much lower than link bandwidth and `size_imbalance` ≈ 1.0 → likely a software/driver overhead issue. -- If `throughput` is low and `size_imbalance` >> 1.0 → expert routing imbalance; the busiest rank gates the collective. -- If `skew in start time` is large → ranks are entering the collective at very different times, indicating upstream compute imbalance. -- Use the heatmap sheet (`--all2allv_heatmap`) to identify which rank pairs carry the most traffic. - ---- - diff --git a/docs_original/generate_perf_report.md b/docs_original/generate_perf_report.md deleted file mode 100644 index d5ee1b97e..000000000 --- a/docs_original/generate_perf_report.md +++ /dev/null @@ -1,177 +0,0 @@ - - -# Generate Performance Report - -This Python script (`TraceLens/Reporting/generate_perf_report_pytorch.py`) processes a PyTorch JSON profile trace and outputs an Excel workbook or CSVs with relevant information. - -Similarly (`TraceLens/Reporting/generate_perf_report_jax.py`) processes JAX XPLANE.PB profile trace. - ---- - -## 🚀 Quick Start - -Run the script with a profile JSON to generate an Excel report: - -```bash -python generate_perf_report_pytorch.py --profile_json_path path/to/profile.json -``` - -Alternatively you can directly call the entry point with the same command line args, - without the need to copy the script out of the repo. - -```bash -TraceLens_generate_perf_report_pytorch --profile_json_path path/to/profile.json -``` - -Similarly for JAX profile trace: - -```bash -python generate_perf_report_jax.py --profile_path path/to/profile.xplane.pb -``` - ---- - -## 📋 Excel Workbook Sheets - -| Sheet Name | Description | -|----------------------------|------------------------------------------------------------------------------------------------------------------| -| `gpu_timeline` | End-to-end GPU activity summary, including compute, memory copies, communication, and idle time. | -| `ops_summary_by_category` | Summary of compute time grouped by operation category (e.g., GEMM, SDPA_fwd, elementwise). | -| `ops_summary` | Summary of compute time at the individual operation level; each row corresponds to a unique operation name. | -| `ops_unique_args` | Detailed operation-level summary; each row corresponds to a unique (operation name, argument) combination. | -| `unified_perf_summary` | Unified perf metrics for all ops with perf models OR leaf ops that launch GPU kernels. Includes GFLOPS, TFLOPS/s, Data Moved, FLOPS/Byte, TB/s metrics aggregated by unique args. | -| `short_kernels_histogram` | Histogram showing the distribution of kernel durations below the short-duration threshold. | -| `short_kernels_all_details`| Detailed list of short-duration kernels, including count, total/mean time, runtime percentage, and parent op. | -| Roofline Sheets | Roofline analysis for each operation category (GEMM, CONV, etc.), including TFLOPs, TB/s, and FLOPs/byte metrics. | - -Note: JAX outputs do not include `short_kernels_histogram` or `short_kernels_all_details`, these are for PyTorch only. - -### 🎯 GPU Architecture for Roofline Analysis - -When `--gpu_arch_json_path` is provided, the report includes additional roofline metrics: - -- **Compute Spec**: Combined compute type and precision (e.g., `matrix_bf16`, `vector_fp32`) -- **Roofline Time (µs)**: Theoretical minimum time based on GPU peak capabilities -- **Pct Roofline**: Percentage of roofline achieved (higher is better) - -The GPU arch file specifies Max Achievable FLOPS (MAF) for different compute types and precisions. See [GPU Architecture Example](../examples/gpu_arch_example.md) for format details and [AMD MAF measurements](https://rocm.blogs.amd.com/software-tools-optimization/measuring-max-achievable-flops-part2/README.html#amd-maf-results) for reference values. - -📖 **For detailed column definitions and usage guide**, see [Performance Report Column Definitions](perf_report_columns.md). - ---- - -## ⚙️ Optional Arguments - -The script supports several optional arguments to customize the output report. By default, it generates an Excel file (`.xlsx`) in the trace directory. If `--output_csvs_dir` is specified, individual CSV files are written instead. - -| Argument | Default | Description | -|-----------------------------------|-------------------|-----------------------------------------------------------------------------| -| `--gpu_arch_json_path PATH` | `None` | Path to GPU architecture JSON file for roofline analysis. See [GPU Architecture Example](../examples/gpu_arch_example.md) for format details. | -| `--topk_ops N` | `None` | Limit the number of rows in the unique-args launcher table. | -| `--topk_short_kernels N` | `None` | Limit the number of rows in the short-kernel table. | -| `--topk_roofline_ops N` | `None` | Limit the number of rows in the roofline sheet. | -| `--extension_file` | `None` | Path to extension python file | -| `--include_unlinked_kernels` | `False` | Include all kernels in the gpu timeline analysis - including kernels not linked to host call stack. By default these unlinked kernels are excluded in the analysis.| -| `--micro_idle_thresh_us X` | `None` | Threshold (in microseconds) to classify idle intervals as micro idle in GPU timeline analysis. If None, all idle times are included in one category. | -| `--short_kernel_study` | `False` | Include short-kernel analysis in the report. | -| `--short_kernel_threshold_us X` | `10` | Threshold (in microseconds) to classify a kernel as "short". | -| `--short_kernel_histogram_bins B` | `100` | Number of bins to use for the short-kernel duration histogram. | -| `--detect_recompute` | `False` | Detect activation recomputation (checkpointing) and add an `is_recompute` column to `ops_summary`, `ops_unique_args`, and `unified_perf_summary`. See [Activation Recompute Detection](#-activation-recompute-detection) below. | -| `--output_xlsx_path PATH` | `` | Path to save the Excel report. Auto-inferred if not provided. | -| `--output_csvs_dir DIR` | `None` | If set, saves each sheet as a CSV file in the specified directory. | - -Note: currently JAX supports only two optional arguments `--output_xlsx_path` and `--output_csvs_dir`. - -### 📦 Output Behavior - -- If `--output_csvs_dir` is set, all output sheets are saved as individual CSV files in that directory. -- Otherwise, the script saves a single Excel file: - - If `--output_xlsx_path` is not provided, it is inferred from the input JSON trace name (e.g., `profile.json` → `profile_perf_report.xlsx`). -- The `openpyxl` package is required only for the case when we write Excel files; it will be auto-installed if missing. - -#### 🧪 Example Usage to write CSVs - - -```bash -python generate_perf_report_pytorch.py \ - --profile_json_path traces/profile.json \ - --output_csvs_dir output_csvs/ \ - --topk_ops 50 \ -``` - -## 🔍 Activation Recompute Detection - -When training with activation checkpointing (`torch.utils.checkpoint`), some forward-pass ops are recomputed during the backward pass to save memory. The `--detect_recompute` flag identifies these ops and adds an `is_recompute` column to the report, so you can see exactly how much GPU time and compute is spent on recomputation. - -### How it works - -TraceLens walks the CPU call-stack tree and finds `python_function` events from `torch/utils/checkpoint.py` corresponding to `recompute_fn`. All ops in those subtrees are marked `is_recompute=True`. This requires `python_function` events in the trace, which TraceLens enables automatically when this flag is set. - -### Usage - -```bash -TraceLens_generate_perf_report_pytorch \ - --profile_json_path path/to/trace.json \ - --detect_recompute -``` - -### What changes in the report - -The following sheets gain an `is_recompute` column that splits rows into recompute vs non-recompute: - -| Sheet | Effect | -|-------|--------| -| `ops_summary` | Same op name appears in separate rows for `is_recompute=True` and `False` | -| `ops_unique_args` | Unique (op, shape, is_recompute) combinations, each with their own time/count | -| `unified_perf_summary` | Full perf metrics split by recompute status | - -This makes it straightforward to answer questions like: -- What percentage of GPU time is recomputation? -- Which layers are being recomputed and at what cost? -- Is the recompute overhead acceptable given the memory savings? - -### Python API - -```python -from TraceLens.TreePerf import TreePerfAnalyzer - -analyzer = TreePerfAnalyzer.from_file("trace.json", detect_recompute=True) -df = analyzer.get_df_kernel_launchers(include_kernel_details=True) -print(df["is_recompute"].value_counts()) -``` - -> **Note**: When `--detect_recompute` is not set (the default), there is zero overhead — no extra columns, no tree changes, and no `python_function` parsing. - ---- - -## 🧩 Extensions: Custom Hooks for Tree and PerfModel - -The `--extension_file` argument allows users to inject custom logic into the performance report generation pipeline. This is useful for experimenting with: - -- Tree post-processing (e.g., injecting pseudo ops) -- Custom performance models for new op types -- Additional operation category definitions - -### 🔧 How to Use - -Pass a Python file path via `--extension_file`. The file can define one or more of the following optional symbols: - -| Symbol Name | Type | Description | -|-----------------------------|-----------|-----------------------------------------------------------------------------| -| `tree_postprocess_extension`| `Callable`| Called with `perf_analyzer.tree`. Use to modify the tree structure post-construction. | -| `perf_model_extension` | `dict` | A mapping from op name (str) to a custom performance model class. These will override or extend existing models. | -| `op_category_extension` | `dict` | Mapping from category-only op names to final categories, used when an op should appear in unified reports without a perf model. | - -#### 📄 Example Extension File for MegatronLM in the examples dir - -### ✅ Example Usage - -```bash -python generate_perf_report_pytorch.py \ - --profile_json_path traces/profile.json \ - --extension_file my_extension.py -``` \ No newline at end of file diff --git a/docs_original/generate_perf_report_rocprof.md b/docs_original/generate_perf_report_rocprof.md deleted file mode 100644 index 92c5fcaad..000000000 --- a/docs_original/generate_perf_report_rocprof.md +++ /dev/null @@ -1,245 +0,0 @@ - - -# Generate Performance Report from rocprofv3 - -This Python script (`TraceLens/Reporting/generate_perf_report_rocprof.py`) processes a rocprofv3 JSON profile trace and outputs an Excel workbook or CSVs with relevant performance information. - ---- - -## 🚀 Quick Start - -Run the script with a rocprofv3 results JSON to generate an Excel report: - -```bash -TraceLens_generate_perf_report_rocprof --profile_json_path path/to/908_results.json -``` - -Or use the Python module directly: - -```bash -python -m TraceLens.Reporting.generate_perf_report_rocprof --profile_json_path path/to/908_results.json -``` - ---- - -## 📋 Excel Workbook Sheets - -| Sheet Name | Description | -|---------------------------------|-------------------------------------------------------------------------------------------------------| -| `gpu_timeline` | End-to-end GPU activity summary, including kernel execution, memory operations, and idle time. | -| `kernel_summary` | Summary of kernel execution time at the individual kernel level; each row is a unique kernel name. | -| `kernel_summary_by_category` | Summary of kernel time grouped by category (e.g., GEMM, Elementwise, Attention). | -| `kernel_details` (optional) | Detailed kernel information including grid/block dimensions for each dispatch. | -| `short_kernels_summary` (opt) | Summary of kernels with duration below the short-duration threshold. | -| `short_kernel_histogram` (opt) | Histogram showing the distribution of kernel durations below the short-duration threshold. | - ---- - -## 🎯 Command-Line Options - -### Required Arguments - -- `--profile_json_path`: Path to the rocprofv3 `*_results.json` file - -### Output Options - -- `--output_xlsx_path`: Path to the output Excel file (default: auto-generated from input filename) -- `--output_csvs_dir`: Directory to save output as CSV files instead of Excel - -### Analysis Options - -- `--disable_kernel_summary`: Disable kernel summary sheets (enabled by default) -- `--kernel_details`: Include detailed kernel information with grid/block dimensions -- `--short_kernel_study`: Include short kernel analysis in the report -- `--short_kernel_threshold_us`: Threshold in microseconds for "short" kernels (default: 10) -- `--short_kernel_histogram_bins`: Number of bins for short-kernel histogram (default: 100) -- `--topk_kernels`: Limit kernel details to top K kernels by time - ---- - -## 💡 Examples - -### Basic Report Generation - -Generate a default Excel report: - -```bash -TraceLens_generate_perf_report_rocprof \ - --profile_json_path run_20251209_070939/smci350-zts-gtu-c6-25/908_results.json -``` - -This creates `908_perf_report.xlsx` in the same directory. - -### Short Kernel Analysis - -Include analysis of short-duration kernels: - -```bash -TraceLens_generate_perf_report_rocprof \ - --profile_json_path 908_results.json \ - --short_kernel_study \ - --short_kernel_threshold_us 20 -``` - -### Detailed Kernel Information - -Include grid/block dimensions for top 100 kernels: - -```bash -TraceLens_generate_perf_report_rocprof \ - --profile_json_path 908_results.json \ - --kernel_details \ - --topk_kernels 100 -``` - -### Generate CSV Files - -Output as CSV files instead of Excel: - -```bash -TraceLens_generate_perf_report_rocprof \ - --profile_json_path 908_results.json \ - --output_csvs_dir ./rocprof_analysis -``` - -### Custom Output Path - -Specify a custom output filename: - -```bash -TraceLens_generate_perf_report_rocprof \ - --profile_json_path 908_results.json \ - --output_xlsx_path my_custom_report.xlsx \ - --kernel_details \ - --short_kernel_study -``` - ---- - -## 📊 Understanding the Output - -### GPU Timeline - -Shows the breakdown of GPU activity: -- **total_time**: Total profiling duration -- **kernel**: Time spent executing kernels -- **memory**: Time spent on memory operations -- **idle**: Time when GPU was idle - -### Kernel Summary - -Aggregated statistics for each unique kernel: -- **Count**: Number of times the kernel was dispatched -- **Total Kernel Time (µs)**: Sum of all dispatch durations -- **Mean/Median/Std/Min/Max**: Statistical distribution of durations -- **Percentage (%)**: Percentage of total kernel time -- **Cumulative Percentage (%)**: Running total percentage - -### Kernel Categories - -Kernels are automatically categorized based on name patterns: -- **GEMM**: Matrix multiplication kernels -- **Elementwise**: Element-wise operations -- **Reduction**: Reduction operations -- **Convolution**: Convolution kernels -- **Normalization**: Batch norm, layer norm, etc. -- **Attention**: Flash attention and related kernels -- **Memory**: Memory copy operations -- **Other**: Uncategorized kernels - ---- - -## 🔍 rocprofv3 Format Details - -rocprofv3 uses the `rocprofiler-sdk-tool` format, which is different from PyTorch's Chrome Trace Event format: - -### Input File Structure - -```json -{ - "rocprofiler-sdk-tool": [{ - "metadata": {"pid": ..., "init_time": ..., "fini_time": ...}, - "agents": [...], - "kernel_symbols": [...], - "buffer_records": { - "kernel_dispatch": [...], - "memory_copy": [...], - "hip_api": [...], - ... - } - }] -} -``` - -### Key Data Sources - -- **kernel_dispatch**: Kernel execution records with timestamps and dispatch info -- **kernel_symbols**: Kernel names and metadata indexed by `kernel_id` -- **memory_copy**: Memory transfer operations -- **hip_api**, **hsa_api**: API call traces (when available) - ---- - -## 🆚 Comparison with PyTorch Profiler - -| Feature | rocprofv3 | PyTorch Profiler | -|---------|-----------|------------------| -| Format | rocprofiler-sdk JSON | Chrome Trace Event | -| Kernel Names | Direct from ROCm | Via PyTorch ops | -| Grid/Block Dims | ✅ Always available | ✅ Available | -| CPU Operations | ❌ Limited | ✅ Full trace | -| Memory Ops | ✅ Basic support | ✅ Full support | -| API Calls | ✅ HIP/HSA | ✅ CUDA runtime | -| Tree View | ❌ N/A | ✅ CPU call stack | - ---- - -## ⚠️ Limitations - -1. **No CPU Call Stack**: rocprofv3 focuses on GPU operations; CPU-side profiling is limited -2. **No Operator Hierarchy**: Unlike PyTorch profiler, there's no tree of CPU operations -3. **Basic Categorization**: Kernel categories are inferred from names, not semantic information -4. **Limited Metadata**: Some PyTorch-specific metadata (input shapes, args) is not available - ---- - -## 🐛 Troubleshooting - -### "Not a valid rocprofv3 file" - -Ensure your file is a `*_results.json` file from rocprofv3, not a different format. - -### "No kernel events found" - -The trace may not have captured any GPU activity. Check that: -- GPU operations were actually executed during profiling -- rocprofv3 was configured to capture kernel dispatches - -### openpyxl not installed - -Install openpyxl for Excel output: - -```bash -pip install openpyxl -``` - -Or use CSV output instead: - -```bash -TraceLens_generate_perf_report_rocprof \ - --profile_json_path trace.json \ - --output_csvs_dir ./output -``` - ---- - -## 📚 Related topics - -- [PyTorch Performance Report](generate_perf_report.md) -- [JAX Performance Report](generate_perf_report_jax.md) -- [ROCm rocprofiler-sdk Documentation](https://rocm.docs.amd.com/projects/rocprofiler-sdk/en/latest/) - diff --git a/docs_original/generate_perf_report_rocprof_pftrace.md b/docs_original/generate_perf_report_rocprof_pftrace.md deleted file mode 100644 index 312b8c3c0..000000000 --- a/docs_original/generate_perf_report_rocprof_pftrace.md +++ /dev/null @@ -1,62 +0,0 @@ - - -# Generate Performance Reports from pftrace (rocprofv3 / Perfetto) - -For **Perfetto-style** traces (e.g. rocprofv3 with `--output-format pftrace`, or JSON with a `traceEvents` array), TraceLens provides two report generators that share the same trace loading and optional **traceconv** handling (for `.pftrace` → JSON conversion). - -## CLI Reports - -| CLI | Description | Output | -|-----|-------------|--------| -| `TraceLens_generate_perf_report_pftrace_hip_activity` | Category summary per GPU, kernel/HIP/XLA summaries (NSYS-style), optional Markdown | Excel, CSV, optional `.md` | -| `TraceLens_generate_perf_report_pftrace_hip_api` | HIP API ↔ kernel correlation; latency breakdown **T = A + Q + K** (API duration, queue delay, kernel duration) | Excel, CSV | - -## Input Formats - -- `.json` — Perfetto-style JSON with a `traceEvents` array -- `.json.gz` — Gzip-compressed JSON -- `.pftrace` — Perfetto binary format - -For `.pftrace` files, **traceconv** is resolved from `PATH` or downloaded into the trace file's directory if not provided via `--traceconv`. - -## Quick Start - -### 1. Record a pftrace - -```bash -rocprofv3 --hip-trace --kernel-trace --memory-copy-trace --rccl-trace \ - --output-format pftrace -d ./v3_traces -- python3 your_app.py -``` - -### 2. Generate reports - -```bash -# Activity report: category summary per GPU, kernel/HIP/XLA summaries (NSYS-style). Optional Markdown output. -TraceLens_generate_perf_report_pftrace_hip_activity --trace_path sample.pftrace --write_md - -# API↔Kernel report: latency breakdown T = A + Q + K (API duration, queue delay, kernel duration) -TraceLens_generate_perf_report_pftrace_hip_api --trace_path sample.pftrace - -# Memory copy report: count per copy_bytes with direction (h2d/d2h/d2d and which GPU(s)) -TraceLens_generate_perf_report_pftrace_memory_copy --trace_path sample.pftrace - -# Custom Excel path: -TraceLens_generate_perf_report_pftrace_memory_copy --trace_path sample.pftrace --output_xlsx_path report.xlsx - -# CSV directory: -TraceLens_generate_perf_report_pftrace_memory_copy --trace_path sample.pftrace --output_csvs_dir ./out -``` - -## Library Usage - -Both scripts can be imported and called with a trace path; they return a dictionary of pandas DataFrames (e.g. `api_kernel_summary`, `category_summary`, `kernel_summary`, `hip_summary`). - -## Testing - -```bash -python -m pytest tests/test_pftrace_hip_api_perf_report.py tests/test_pftrace_hip_activity_report.py -v -``` diff --git a/docs_original/gpu_event_analyser.md b/docs_original/gpu_event_analyser.md deleted file mode 100644 index 65034c8d6..000000000 --- a/docs_original/gpu_event_analyser.md +++ /dev/null @@ -1,95 +0,0 @@ - - -# GPUEventAnalyser - -GPUEventAnalyser is a reusable component designed to analyze GPU timeline and extract key performance metrics. While it is used within TreePerf, it can also be used independently. - ---- - -## Key Features - -1. **GPU Timeline Breakdown**: Computes key GPU activity metrics, including: - - **Computation Time**: Time spent in computation kernels (e.g., matrix multiplications, convolutions). - - **Communication Time**: Time spent in communication kernels (e.g., NCCL operations for distributed training). - - **Memcpy Time**: Time spent in memory copy operations between host and device or across devices. - - **Idle Time**: Periods where the GPU is not executing any computation, communication, or memcpy operations. - - **Exposed Communication**: Communication time that does not overlap with computation. - - **Exposed Memcpy**: Memcpy time that does not overlap with computation or communication. - -2. **Reusable Across Profiling Formats**: Although GPUEventAnalyser is designed for PyTorch's JSON trace format, it can be adapted to other profiling formats by inheriting the class and reimplementing `get_gpu_event_lists()`. - ---- - -## Usage Example - -```python -import json -import sys -from TraceLens import GPUEventAnalyser - -path = sys.argv[1] # this expects the JSON file (unzipped) - -with open(path, 'r') as f: - data = json.load(f) - -events = data['traceEvents'] -my_gpu_event_analyser = GPUEventAnalyser(events) -df = my_gpu_event_analyser.get_breakdown_df() -print(df) -``` - -Example output: - -| type | time ms | percent | -| --------------------- | --------- | --------- | -| computation_time | 4184.32 | 96.10 | -| exposed_comm_time | 160.85 | 3.69 | -| exposed_memcpy_time | 0.19 | 0.00 | -| busy_time | 4345.36 | 99.80 | -| idle_time | 8.53 | 0.20 | -| total_time | 4353.88 | 100.00 | -| total_comm_time | 292.92 | 6.73 | -| total_memcpy_time | 0.19 | 0.00 | - ---- - -## Jax support -Jax describes events slightly differently, and includes all GPUs in a single trace, using the JaxGPUEventAnalyser class. -When using GPUEventAnalyser for JAX traces (enabled by passing "jax=True" to the analysis functions), -the compute_metrics() and get_breakdown_df() methods will return events from GPU 0. -To access the traces for all devices, call get_gpu_event_lists_jax() to obtain a dictionary of -{pid: event_lists} where pid is the process id (1-n for n GPUs, a number >100 for CPU the CPU event list). -To create a Pandas dataframe for each GPU, call get_breakdown_df_multigpu(), which will return a dictionary of -{gpu_index: DataFrame} for gpuindex in [0, num_gpus). -The inherited function get_breakdown_df will return the results from GPU 0 - -### Jax Usage Example -```python -import gzip -import json -import sys -from TraceLens import JaxGPUEventAnalyser - -path = sys.argv[1] # this expects the zipped JSON file produced by the - -with gzip.open(path, 'r') as fin: - data = json.loads(fin.read().decode('utf-8')) - -events = data['traceEvents'] -my_gpu_event_analyser = JaxGPUEventAnalyser(events) -for gpu, df in my_gpu_event_analyser.get_breakdown_df_multigpu().items(): - print(f"GPU {gpu}") - print(df) -``` - -## Customizing for Other Profiling Formats - -To adapt GPUEventAnalyser for other profiling formats, subclass it and reimplement the `get_gpu_event_lists()` method to correctly extract GPU events. -See the class JaxGPUEventAnalyser for an example. - ---- - diff --git a/docs_original/jax_analyses.md b/docs_original/jax_analyses.md deleted file mode 100644 index 19726e038..000000000 --- a/docs_original/jax_analyses.md +++ /dev/null @@ -1,98 +0,0 @@ - - -22 October 2025: The following JAX analysis modules—`summarize_gpu_events`, `summarize_gpu_gemm_events`, and related utilities—have been merged into TraceLens and will no longer be maintained as standalone modules. Example usage is similar to PyTorch. See 'docs/generate_perf_report.md'. - -```bash -python generate_perf_report_jax.py --profile_path path/to/profile.xplane.pb --output_csvs_dir save/to/dir -``` - -Jax analysis, particularly reading the protobuf files, has been tested with tensorboard 2.19.0 and tensorboard-plugin-profile 2.19.0 and protobuf 5.29.2. -Other versions may not work - -Analyze Jax computations including GEMM analysis -Run this with the xplane.pb or json.gz and jit_train_step.gfx942_gpu_after_optimizations.txt -``` -from TraceLens.TraceLens import JaxAnalyses -import sys -import pandas as pd -filename_path = sys.argv[1] -averages, categorized, additional_events = JaxAnalyses.summarize_gpu_events(filename_path) -pd.set_option('display.max_rows', None) -print("Average utilization by type of kernel") -print(averages) -print("XLA computations (% for all GPUs)") -print(categorized) -print("Uncategorized XLA computations (% for all GPUs)") -print(additional_events) -if len(sys.argv)>2: - print("GEMMs") - print(JaxAnalyses.summarize_gpu_gemm_events(sys.argv[2])) -``` - -Standalone Jax GEMM analysis from protobuf (from profiler) or xla dump (jit_train_step.gfx942_gpu_after_optimizations.txt): -``` -from TraceLens.TraceLens import JaxAnalyses -import sys -import pandas as pd -pd.set_option('display.max_rows', None) -print("GEMMs") -filename = sys.argv[1] -if filename.endswith("pb"): - gemms = JaxAnalyses.summarize_gpu_gemm_events_from_pb(filename) -else: - gemms = JaxAnalyses.summarize_gpu_gemm_events_from_xla(filename) -print(gemms) -``` - -Detailed GEMM performance metrics -``` -from TraceLens import JaxAnalyses -import sys -import pandas as pd -pd.set_option('display.max_rows', None) -pd.set_option('display.max_columns', None) -pd.set_option('display.width', None) -filename = sys.argv[1] -# num_cus: MI300X - 304; MI210: 104 -gemms = JaxAnalyses.gemm_performance_from_pb(filename, arch = {"num_cus": 104}) -print(gemms) -``` - -Anylyze Jax communications -Run this with the xplane.pb or json.gz and jit_train_step.gfx942_gpu_after_optimizations-buffer-assignment.txt -``` -from TraceLens.TraceLens import JaxAnalyses -import sys -import pandas as pd -profile_path = sys.argv[1] -xla_path = sys.argv[2] -summarized_events = JaxAnalyses.summarize_gpu_communication_events(profile_path, xla_path) -for (df, bw_data, count_data, time_by_size, range_data) in filter(lambda x: len(x[0]) > 0, summarized_events.values()): - print(f"Stats for {df['base_collective'][0]}") - print("Bandwidth") - print(bw_data) - print("counts") - print(count_data) - print("buffer sizes") - print(time_by_size) - print("time_in_ranges") - print(range_data) -``` - -Trace to tree for Jax traces, based on the "Framework Name Scope" thread in the trace -``` -import TraceLens -import sys -data=TraceLens.util.DataLoader.load_data(sys.argv[1]) -events=data['traceEvents'] -metadata = TraceLens.util.TraceEventUtils.get_metadata(events) -categorizer = TraceLens.util.TraceEventUtils.prepare_event_categorizer(events) -real_events = TraceLens.util.TraceEventUtils.non_metadata_events(events) -tree = TraceLens.TraceToTree(real_events, linking_key='correlation', event_to_category=categorizer) -tree.build_tree(True) -tree.traverse_subtree_and_print(tree.get_UID2event(tree.cpu_root_nodes[1]), False) -``` \ No newline at end of file diff --git a/docs_original/locate_node.PNG b/docs_original/locate_node.PNG deleted file mode 100755 index 4f37a1a6b..000000000 Binary files a/docs_original/locate_node.PNG and /dev/null differ diff --git a/docs_original/origami_with_tracelens.md b/docs_original/origami_with_tracelens.md deleted file mode 100644 index 2e8132be1..000000000 --- a/docs_original/origami_with_tracelens.md +++ /dev/null @@ -1,142 +0,0 @@ - - -# Origami with TraceLens - -TraceLens can estimate **simulated GEMM and SDPA kernel times** using **Origami** (ROCm’s performance modeling library) when you provide a GPU architecture description and explicitly opt in. This is **optional**: `pip install TraceLens` does not install Origami. - ---- - -## What Origami does in TraceLens - -- **GEMM** and **SDPA** perf models call Origami’s Python bindings to predict a duration in microseconds for forward (and SDPA backward where applicable). -- Results show up in perf reports under columns such as **Origami Time (µs)**, **Origami TFLOPS/s**, **Origami TB/s**, and **Pct Origami** (relative to measured kernel busy time), when simulation uses Origami (see [Performance Report Column Definitions](perf_report_columns.md)). -- **Roofline** metrics from `--gpu_arch_json_path` are separate; they do not require Origami. Origami adds *simulated* timing on top when enabled. - ---- - -## Installation - -### Python package - -Install the published package (PyPI name **`rocm-origami`**): - -```bash -pip install rocm-origami -``` - -Note that installing Origami currently (as of April 2026) requries that ROCm is installed on the system because Origami can use the HIP library to detect the GPU currently installed in the system and use it as an architectural model. TraceLens does not currently use this functionality but it currently cannot be removed from Origami. - -### System environment - -Origami’s wheels/bindings expect a **ROCm** runtime on the machine (GPU not required for pure Python simulation in many cases, but library loading may depend on your setup). The project’s CI uses an AMD ROCm container and installs `rocm-origami` alongside TraceLens (see `.github/workflows/regression-tests.yml`). - -If `import origami` fails after `pip install`, check: - -- Python version and wheel compatibility for `rocm-origami`. -- `LD_LIBRARY_PATH` / ROCm install paths required by the Origami wheel you use. - ---- - -## When TraceLens uses Origami vs other backends - -Simulation is implemented in `TraceLens/PerfModel/perf_model.py` (`GEMM.get_simulation_time_func`): - -1. **Otherwise**, if **`enable_origami` is true** and the perf model has the needed architecture and parameters, TraceLens uses **Origami**. - -2. If **`enable_origami` is false**, TraceLens does **not** call Origami; simulated times are omitted for that path. - -So you need **both**: - -- A valid **GPU architecture** (see below), and -- **`enable_origami=True`** (CLI flag or Python API). - ---- - -## GPU architecture JSON - -Pass the same JSON file you use for roofline analysis with **`--gpu_arch_json_path`**. It must include fields Origami expects (for example GPU name, frequency, memory bandwidth, CU count), as consumed by `OrigamiHelper.get_hardware` in `TraceLens/PerfModel/origami_helper.py`. - -See [GPU Architecture Example](../examples/gpu_arch_example.md) and [Generate Performance Report](generate_perf_report.md) for the general format. - ---- - -## Command-line usage - -### PyTorch perf report - -```bash -TraceLens_generate_perf_report_pytorch \ - --profile_json_path path/to/profile.json.gz \ - --gpu_arch_json_path path/to/gpu_arch.json \ - --enable-origami \ - --output_csvs_dir ./out_csvs -``` - -Or: - -```bash -python -m TraceLens.Reporting.generate_perf_report_pytorch \ - --profile_json_path path/to/profile.json.gz \ - --gpu_arch_json_path path/to/gpu_arch.json \ - --enable-origami \ - --output_csvs_dir ./out_csvs -``` - -### vLLM-oriented PyTorch report - -Same pattern; the entry point mirrors the PyTorch script (`--enable-origami`). - -### JAX perf report - -```bash -TraceLens_generate_perf_report_jax \ - --profile_path path/to/trace.xplane.pb \ - --gpu_arch_json_path path/to/gpu_arch.json \ - --enable-origami \ - --output_csvs_dir ./out_csvs -``` - -### Standalone GEMM/SDPA simulator helper - -`TraceLens/PerfModel/run_perf_model.py`: - -```bash -python -m TraceLens.PerfModel.run_perf_model --op gemm ... --enable_origami -``` - ---- - -## Python API - -When building a **`TreePerfAnalyzer`** (or **`JaxTreePerfAnalyzer`**) in code, pass: - -```python -TreePerfAnalyzer.from_file( - profile_filepath="profile.json.gz", - arch=gpu_arch_dict, # or load JSON - enable_origami=True, -) -``` - -Reporting helpers such as **`generate_perf_report_pytorch`** accept **`enable_origami=True`** and forward it to the analyzer. - ---- - -## Troubleshooting - -| Symptom | What to check | -|--------|----------------| -| No Origami columns in CSVs | Confirm **`--enable-origami`** (or API **`enable_origami=True`**) **and** **`--gpu_arch_json_path`**. | -| Message on stderr about `origami` import | Install **`rocm-origami`** and fix ROCm/library paths. | -| Unsupported dtype warning | Origami path supports a fixed set of dtypes (e.g. fp16, bf16, fp32, fp64, fp8); others skip simulation. | - ---- - -## Related documentation - -- [Generate Performance Report](generate_perf_report.md) — CLI overview and roofline (`--gpu_arch_json_path`). -- [Performance Report Column Definitions](perf_report_columns.md) — sheet and column meanings. diff --git a/docs_original/perf_report_columns.md b/docs_original/perf_report_columns.md deleted file mode 100644 index 34e2b5efa..000000000 --- a/docs_original/perf_report_columns.md +++ /dev/null @@ -1,982 +0,0 @@ - - -# Performance Report Column Definitions - -This document provides detailed explanations of the columns in each sheet of the TraceLens performance report. - -## Table of Contents - -1. [gpu_timeline](#1-gpu_timeline) - High-level GPU time breakdown -2. [Operation Analysis](#2-operation-analysis) - - [2.1 ops (Base Data)](#21-ops-base-data) - - [2.2 ops_summary_by_category](#22-ops_summary_by_category-most-aggregated) - - [2.3 ops_summary](#23-ops_summary-by-operation-name) - - [2.4 ops_unique_args](#24-ops_unique_args-most-detailed) -3. [Performance Metrics Sheets (Roofline Analysis)](#3-performance-metrics-sheets-roofline-analysis) - - [Understanding the Metrics Pipeline](#understanding-the-metrics-pipeline) - - [Operation Parameters Reference](#operation-parameters-reference) - - [Why This Matters: Roofline Analysis](#why-this-matters-roofline-analysis) -4. [Collective Communication Analysis](#4-collective-communication-analysis) -5. [Additional Sheets](#5-additional-sheets) -6. [Common Analysis Workflows](#6-common-analysis-workflows) -7. [Related Documentation](#7-related-documentation) - ---- - -## Overview - -The performance report Excel file contains multiple sheets analyzing different aspects of GPU performance. The core sheets are: - -1. **gpu_timeline** - High-level GPU time breakdown -2. **ops** - Detailed per-operation data (base data) -3. **ops_summary_by_category** - Operations summarized by category -4. **ops_summary** - Operations summarized by name -5. **ops_unique_args** - Operations summarized by unique argument combinations - -Additional sheets may include op-specific analysis (GEMM, SDPA_fwd, CONV_fwd, etc.), kernel summary, short kernels, and collective analysis. - -**Unit Conventions**: -- **Time**: All times from the trace are in **microseconds (µs)** unless explicitly stated otherwise (e.g., `time ms` in `gpu_timeline` is in milliseconds) -- **Compute**: GFLOPS (billions of FLOPs), TFLOPS/s (trillions of FLOPs per second) -- **Memory**: MB (megabytes), GB/s (gigabytes per second), TB/s (terabytes per second) - ---- - -## 1. gpu_timeline - -**Purpose**: Provides a high-level breakdown of GPU time into computation, communication, memory copy, and idle time, accounting for overlaps between different types of operations. - -### Example Output - -Here's a typical `gpu_timeline` sheet from a distributed training workload: - -| type | time ms | percent | -| --------------------- | --------- | --------- | -| computation_time | 56305.19 | 99.30 | -| exposed_comm_time | 240.88 | 0.42 | -| exposed_memcpy_time | 14.44 | 0.03 | -| busy_time | 56560.52 | 99.75 | -| idle_time | 143.16 | 0.25 | -| total_time | 56703.68 | 100.00 | -| total_comm_time | 17203.43 | 30.34 | -| total_memcpy_time | 14.47 | 0.03 | - -**Analysis**: -- **99.30%** computation time → GPU is very efficiently utilized -- **0.42%** exposed communication → Excellent! Most communication (30.34% total - 0.42% exposed = 29.92%) overlaps with computation -- **0.25%** idle time → Minimal gaps, good kernel launch efficiency -- This workload demonstrates effective computation/communication overlap - -### Event Classification - -GPU events are classified into three categories based on their `cat` field and kernel name: - -``` -GPU Event -│ -├─ cat = "gpu_memcpy"? -│ └─ YES → Memory Copy (H2D, D2H, D2D) -│ -├─ cat = "kernel"? -│ └─ YES → Does kernel name contain "nccl"? -│ ├─ YES → Communication (AllReduce, AllGather, ReduceScatter, etc.) -│ └─ NO → Computation (GEMM, Conv, Elementwise, etc.) -│ -└─ cat = "gpu_memset"? - └─ YES → Computation (grouped here for simplicity; typically very short) -``` - -**Notes**: -1. "Computation" includes all kernels doing work (GEMM, convolution, elementwise, etc.), not just compute-bound operations. -2. Communication kernels (NCCL) include both synchronization delay and actual data transfer time. In single-rank traces, we cannot separate these components, so `total_comm_time` represents the total duration. See the NCCL Analyzer documentation for multi-rank analysis that can break down communication into sync and transfer phases. - -### How Time is Calculated - -**Step 1: Merge events by category** - -Events of the same category are merged across all GPU streams into non-overlapping intervals. - -Example - merging computation events: -``` -Stream 0: ──[Kernel A]────────[Kernel B]── -Stream 1: ─────[Kernel C]──[Kernel D]───── - ↓ -Merged: ──[──────────]──[──────────]── - (overlapping kernels merged into single intervals) -``` - -**Step 2: Create interval sets** - -After merging, we have four sets of non-overlapping intervals: -- `COMP` = merged computation intervals -- `COMM` = merged communication intervals -- `MEMCPY` = merged memcpy intervals -- `ALL_GPU` = merged intervals of ALL GPU events (computation + communication + memcpy) - -**Step 3: Apply set arithmetic** - -Exposed metrics are calculated by subtracting overlaps: - -``` -exposed_comm intervals = COMM - COMP -exposed_memcpy intervals = MEMCPY - COMP - COMM -``` - -In simple terms: -- **Exposed communication** = communication time that doesn't overlap with computation -- **Exposed memcpy** = memcpy time that doesn't overlap with computation or communication - -**Step 4: Sum interval durations** - -Time for each metric = sum of durations of all intervals in that set - -``` -computation_time = sum of durations in COMP -exposed_comm_time = sum of durations in (COMM - COMP) -exposed_memcpy_time = sum of durations in (MEMCPY - COMP - COMM) -busy_time = sum of durations in ALL_GPU -idle_time = total_time - busy_time -total_comm_time = sum of durations in COMM -total_memcpy_time = sum of durations in MEMCPY -total_time = end of last GPU event - start of first GPU event -``` - -**The equation**: -``` -computation_time + exposed_comm_time + exposed_memcpy_time + idle_time = total_time -``` - -### Columns - -| Column | Description | Values | -|--------|-------------|--------| -| `type` | Category of GPU time | `computation_time`, `exposed_comm_time`, `exposed_memcpy_time`, `busy_time`, `idle_time` (or `micro_idle_time` and `macro_idle_time`), `total_comm_time`, `total_memcpy_time`, `total_time` | -| `time ms` | Time in milliseconds | Actual duration for each category | -| `percent` | Percentage of total time | `(time ms / total_time) * 100` | - -### Interpreting Results - -- **High `computation_time`** (>80%): Workload is efficiently using GPU -- **High `exposed_comm_time`**: Communication not overlapped with computation → optimize with computation/communication overlap -- **High `exposed_memcpy_time`**: Memory transfers blocking work → optimize data movement -- **High `idle_time`**: GPU sitting idle → check for CPU bottlenecks, kernel launch overhead, or synchronization - -**Code Reference**: Generated by `TreePerfAnalyzer.get_df_gpu_timeline()` using `GPUEventAnalyser.compute_metrics()` - ---- - -## 2. Operation Analysis - -The following sections analyze CPU operations that launch GPU kernels. Each operation is analyzed for its direct GPU time impact. - -### 2.1. ops (Base Data) - -**Purpose**: Provides detailed per-operation data showing each CPU operation that launches GPU kernels. This is the base unsummarized data; the following sheets provide different ways of summarizing this information. - -**What you'll see**: Each row represents one CPU operation instance from your workload that launched GPU kernels, showing exactly which kernels it launched and how long they took. - -**Code Reference**: Generated by `TreePerfAnalyzer.get_df_kernel_launchers()` - ---- - -#### Understanding the Analysis Approach - -**Why analyze operations instead of just kernels?** - -When you look at a PyTorch trace, you might be tempted to analyze GPU kernels directly. However, this has significant limitations: - -- **Kernel names are cryptic and ambiguous**: A kernel named `Cijk_Ailk_Bljk_BBS_BH_Bias_HAS_SAV...` tells you almost nothing about what computation it's doing. The same kernel name can also map to different computations depending on input shape, dtype, and memory layout. - -- **Kernel names vary across platforms**: The same matrix multiply appears as `nvjet_*` or `cutlass_*` on NVIDIA GPUs, but as `Cijk_*` on AMD GPUs. This makes cross-platform comparison nearly impossible. - -- **Operations are stable and meaningful**: Names like `aten::addmm` (matrix multiply with add) or `aten::conv2d` are platform-independent and immediately tell you what computation is happening. Combined with argument information (shapes, dtypes, strides), they fully define the computation in a reproducible way. - -**Result**: By analyzing at the operation level, you get insights that are portable across platforms, interpretable without deep kernel knowledge, and reproducible. (See [Trace2Tree Motivation](./conceptual/trace2tree_motivation.md) for more details.) - ---- - -#### Which Operations Are Analyzed? - -We analyze **leaf operations** - the lowest-level CPU operations in the call stack that directly launch GPU kernels. This gives you the most granular view without double-counting. - -**Example call stack** (showing all layers from trace): - -``` -└── (python_function) nn.Module: Linear_75 - └── (python_function) forward @ linear.py:124 - └── (cpu_op) aten::linear - ├── (cpu_op) aten::to - │ └── (cpu_op) aten::_to_copy - │ └── (cpu_op) aten::copy_ ← Leaf op - │ └── (cuda_runtime) hipLaunchKernel ← Runtime layer - │ └── (kernel) elementwise_kernel... ← GPU kernel - ├── (cpu_op) aten::to - │ └── (cpu_op) aten::_to_copy - │ └── (cpu_op) aten::copy_ ← Leaf op - │ └── (cuda_runtime) hipLaunchKernel ← Runtime layer - │ └── (kernel) elementwise_kernel... ← GPU kernel - └── (cpu_op) aten::linear - └── (cpu_op) aten::addmm ← Leaf op - ├── (cuda_runtime) hipLaunchKernel ← Runtime layer - │ └── (kernel) elementwise_kernel... ← GPU kernel - └── (cuda_runtime) hipExtModuleLaunchKernel ← Runtime layer - └── (kernel) Cijk_Alik_Bljk_BBS_BH... ← GPU kernel -``` - -**Note**: The `cuda_runtime` label is used by PyTorch profiler even on ROCm platforms - it's just a naming convention. - -In this example: -- **Leaf operations analyzed**: `aten::copy_` (2 instances) and `aten::addmm` -- **Not included**: `aten::linear`, `aten::to`, `aten::_to_copy` (these are higher-level and don't directly launch kernels) -- **Why this matters**: If we included higher-level operations, we'd count the same GPU time multiple times. By focusing on leaf operations, each kernel's time is attributed to exactly one operation. - ---- - -#### What Are CPU Operations (cpu_op)? - -**Understanding PyTorch's execution flow**: - -When you write PyTorch code in Python (like `output = model(input)`), it goes through several layers before reaching the GPU. The trace captures all these layers: - -1. **Python frontend**: Your Python code calls `nn.Module` methods (labeled `python_function` in trace) -2. **Torch dispatcher**: Operations are routed through PyTorch's dispatcher, which selects the appropriate implementation based on device, dtype, and other factors (labeled `cpu_op` in trace) -3. **Runtime layer**: CUDA/HIP API calls to launch kernels (labeled `cuda_runtime` or `cuda_driver` in trace) -4. **GPU kernels**: The actual computation on the GPU (labeled `kernel`, `gpu_memcpy`, or `gpu_memset` in trace) - -**What "cpu_op" means**: Operations marked as "cpu_op" in the trace are **operations registered in the PyTorch dispatcher**. These can be registered in C++ or Python, and represent various levels of abstraction in the execution hierarchy. The key distinction is that we analyze the ones that directly launch kernels (the leaf operations in the call stack). - -**Why this matters**: The dispatcher layer is where PyTorch attaches rich metadata about the computation. - -**Note on the runtime layer**: In the trace hierarchy, there's also a runtime layer (labeled `cuda_runtime` or `cuda_driver` in the trace) between CPU operations and GPU kernels. This represents the CUDA/HIP API calls that actually launch kernels. Interestingly, PyTorch profiler uses `cuda_runtime`/`cuda_driver` naming convention even on ROCm/HIP platforms - it's just a naming convention inherited from CUDA. - ---- - -#### What Arguments Do Operations Contain? - -Each operation in the trace contains valuable argument information that fully characterizes the computation. **These arguments come directly from the JSON event's `args` field** that PyTorch profiler captures: - -**Example JSON event from trace**: -```json -{ - "name": "aten::addmm", - "cat": "cpu_op", - "ts": 1234567890, - "dur": 150, - "args": { - "Input Dims": [[1024, 512], [512, 256], [256]], - "Input type": ["c10::BFloat16", "c10::BFloat16", "c10::BFloat16"], - "Input Strides": [[512, 1], [256, 1], [1]], - "Concrete Inputs": ["", "", "1.0"] - }, - ... -} -``` - -These arguments are extracted and presented in the ops sheet: - -| Argument Type | Description | Example (from above aten::addmm) | Why It Matters | -|---------------|-------------|---------|----------------| -| **Input Dims** | Shape of input tensors | `((1024, 512), (512, 256), (256,))` | Different shapes → different performance (tiling, memory access patterns) | -| **Input type** | Data types | `('c10::BFloat16', 'c10::BFloat16', 'c10::BFloat16')` | FP32 vs BF16 vs FP16 affects speed and memory | -| **Input Strides** | Memory layout (row-major, column-major, etc.) | `((512, 1), (256, 1), (1,))` | Strided vs contiguous affects memory bandwidth | -| **Concrete Inputs** | Scalar/string arguments | `('', '', '1.0')` | Additional parameters that may affect algorithm selection | - -**Example - Why shapes matter**: -``` -aten::addmm with shape (1024, 512) × (512, 256): Fast (optimized tiling) -aten::addmm with shape (1023, 511) × (511, 255): Slower (odd dimensions, less efficient) -``` - -**Example - Why strides matter**: -``` -Same operation, same shape, different strides can have different performance: - stride=(512, 1) vs stride=(1, 512) -``` -Strides affect memory access patterns. Some strides work better for certain operations depending on how kernels access data. TraceLens captures strides so you can identify stride-dependent performance variations. - -**The key insight**: By capturing these arguments, you can group operations by their **unique argument combinations**. This enables targeted analysis: -- **Identify slow cases**: "Which specific input shapes are causing slowdowns?" -- **Reproduce issues**: "I can reproduce the slow case with these exact inputs" -- **Compare across platforms**: "Does this shape perform better on NVIDIA or AMD?" - -This combination of operation name + arguments forms a unique signature that can be aggregated in the summary sheets (ops_summary_by_category, ops_summary, ops_unique_args), enabling you to ask questions at different levels of granularity. - ---- - -#### How GPU Time Is Calculated - -For each operation, we calculate `total_direct_kernel_time_us` using **GPU Event Analyzer** to account for kernel overlaps. - -**Key insight**: When an operation launches multiple kernels, some may run concurrently on different GPU streams. We use GPU Event Analyzer to compute the actual "busy time" - the wall-clock time during which at least one of the operation's kernels is executing. - -**Example**: -``` -Operation: aten::addmm - ├── Kernel A: 100 µs (stream 0, executes 0-100 µs) - └── Kernel B: 80 µs (stream 1, executes 20-100 µs, overlaps with Kernel A) - -Naive sum: 100 + 80 = 180 µs ✗ (incorrect) -Actual busy time: 100 µs ✓ (kernels overlapped from 20-100 µs) -``` - -This gives you the actual wall-clock time each operation's kernels occupied the GPU. - ---- - -#### Understanding kernel_details - -Captures detailed information about each GPU kernel launched by an operation: kernel name, duration (in microseconds), and stream ID. - -**Example from aten::addmm**: -```python -[ - {'name': 'elementwise_kernel', 'dur': 2.3, 'stream': 7}, # dur in µs - {'name': 'Cijk_Alik_Bljk_BBS_BH_Bias_HAS_SAV...', 'dur': 45.2, 'stream': 7} # dur in µs -] - -# Analysis: addmm = elementwise (2.3 µs) + GEMM (45.2 µs) -# → GEMM is bottleneck (95% of time), both sequential (same stream) -# → Cijk_* = Tensile GEMM (AMD), cutlass_* would be NVIDIA -``` - -**Key insights from kernel_details**: -- **Backend implementation**: Kernel names reveal which library is used - - `Cijk_*`: Tensile GEMM kernels (AMD ROCm) - - `cutlass_*`: CUTLASS kernels (NVIDIA) - - `ck_tile::kentry<...FmhaFwd...>`: Composable Kernel (CK) Flash Attention (AMD) - - `void at::native::*`: PyTorch native kernels (element-wise, etc.) -- **Execution breakdown**: Shows performance contribution of each kernel in multi-kernel operations -- **Concurrency**: Same stream = sequential, different streams = potentially concurrent - ---- - -#### Columns - -| Column | Description | Details | -|--------|-------------|---------| -| `name` | Operation name | PyTorch operation name (e.g., `aten::addmm`, `aten::index_select`) | -| `op category` | Operation category | Categorized type (e.g., GEMM, CONV_fwd, SDPA_fwd, other) | -| `UID` | Unique event identifier | Unique ID for this specific operation instance in the trace | -| `total_direct_kernel_time` | GPU busy time in microseconds | Wall-clock time the GPU was busy executing this operation's kernels (accounts for overlaps using GPU Event Analyzer) | -| `direct_kernel_count` | Number of kernels launched | How many GPU kernels this operation launched | -| `Input Dims` | Input tensor shapes | Tuple of shapes for each input tensor, e.g., `((30522, 512), (), (141,))` | -| `Input type` | Input data types | Tuple of data types for each input, e.g., `('c10::BFloat16', 'Scalar', 'long int')` | -| `Input Strides` | Input tensor strides | Tuple of stride tuples for each input, e.g., `((512, 1), (), (1,))` | -| `Concrete Inputs` | Scalar/string arguments | Additional arguments passed to the operation, e.g., `('', '0', '')` | -| `kernel_details` | Kernel execution details | List of dicts with kernel name, duration (µs), and stream for each kernel launched, e.g., `[{'name': '...', 'dur': 3.136, 'stream': 7}]` | - ---- - -#### What This Data Represents - -**Flat tabular view of hierarchical trace**: The PyTorch trace has a hierarchical structure (Python → Operations → Runtime → Kernels). This sheet provides a **MECE (Mutually Exclusive, Collectively Exhaustive)** flat representation focusing on the leaf operations that launch GPU work. - -- **Granularity**: Each row is a single CPU operation instance from the trace -- **Completeness**: Only operations that launch GPU kernels are included (pure CPU operations are not shown) -- **Raw data**: This is unsummarized - you see every individual operation occurrence. The following sheets (ops_summary_by_category, ops_summary, ops_unique_args) provide different ways to aggregate this data. - ---- - -The following three sheets provide different levels of aggregation of the base `ops` data, following a **progressive disclosure of complexity**: - -### 2.2. ops_summary_by_category (Most Aggregated) - -**Purpose**: Highest-level summary - groups operations into broad computational categories (GEMM, Convolution, Attention, etc.). Use this to quickly identify which types of operations dominate your workload. - -### Columns - -| Column | Description | Calculation | -|--------|-------------|-------------| -| `op category` | Operation category | Possible values: GEMM, CONV_fwd, CONV_bwd, SDPA_fwd, SDPA_bwd, BN_fwd, BN_bwd, triton, elementwise, reduce, multi_tensor_apply, other | -| `Count` | Number of operations in this category | Count of unique CPU operations | -| `total_direct_kernel_time_ms` | Total GPU time in milliseconds | Sum of all GPU kernel durations launched by operations in this category | -| `Percentage (%)` | Percentage of total GPU time | `(total_direct_kernel_time_ms / sum of all kernel time) * 100` | -| `Cumulative Percentage (%)` | Running cumulative percentage | Sum of percentages from top to current row | - -**How operations are categorized**: - -Operations are automatically categorized based on name patterns and kernel characteristics. Some categories are detected from the operation name, while others are detected by inspecting the launched kernel names. - -| Category | Detection Method | Example Operations | -|----------|-----------------|-------------------| -| GEMM | Operation name | `aten::addmm`, `aten::mm`, `aten::bmm`, `aten::baddbmm` | -| CONV_fwd | Operation name | `aten::convolution`, `aten::miopen_convolution`, `aten::cudnn_convolution` | -| CONV_bwd | Operation name | `aten::convolution_backward` | -| SDPA_fwd | Operation name | `aten::_scaled_dot_product_flash_attention`, `aten::_flash_attention_forward` | -| SDPA_bwd | Operation name | `aten::_scaled_dot_product_flash_attention_backward` | -| BN_fwd | Operation name | `aten::batch_norm`, `aten::native_batch_norm`, `aten::cudnn_batch_norm` | -| BN_bwd | Operation name | `aten::native_batch_norm_backward`, `aten::cudnn_batch_norm_backward` | -| triton | Operation name | Operations starting with `triton` | -| elementwise | Kernel name inspection | Operations launching `at::native` elementwise kernels (e.g., `aten::relu`, `aten::add`, `aten::mul`) | -| reduce | Kernel name inspection | Operations launching `at::native` reduce kernels | -| multi_tensor_apply | Kernel name inspection | Operations launching `multi_tensor_apply` kernels | -| other | Default | Operations not matching any above patterns | - -**Use this when**: You want a high-level answer to "What types of operations are taking the most time?" - -**Code Reference**: Generated by `TreePerfAnalyzer.get_df_kernel_launchers_summary_by_category()` - ---- - -### 2.3. ops_summary (By Operation Name) - -**Purpose**: Mid-level summary - groups operations by their name (e.g., all `aten::addmm` calls together, all `aten::conv2d` calls together). Use this when you know which category is expensive and want to see which specific operations within that category are the culprits. - -### Columns - -| Column | Description | Calculation | -|--------|-------------|-------------| -| `name` | Operation name | PyTorch operation name (e.g., `aten::addmm`, `aten::_flash_attention_forward`) | -| `total_direct_kernel_time_sum` | Total GPU time in microseconds | Sum of all GPU kernel durations launched by this operation type | -| `Count` | Number of operations | Count of instances of this operation | -| `total_direct_kernel_time_ms` | Total GPU time in milliseconds | `total_direct_kernel_time_sum / 1000` | -| `Percentage (%)` | Percentage of total GPU time | `(total_direct_kernel_time_ms / sum of all kernel time) * 100` | -| `Cumulative Percentage (%)` | Running cumulative percentage | Sum of percentages from top to current row | -| `is_recompute` | *(Optional)* Whether this operation is part of activation recomputation | `True` or `False`. Only present when `--detect_recompute` is enabled. Same op name may appear in two rows (one for recompute, one for non-recompute). | - -**What this shows**: -- Each unique operation name gets one row (e.g., one row for all `aten::addmm` calls) -- All instances of the same operation are aggregated together, regardless of their input shapes, dtypes, or other arguments -- More granular than category view, but still hides variation due to different input arguments -- When `--detect_recompute` is enabled, rows are further split by recompute status - -**Use this when**: "I see GEMM is expensive - which specific GEMM operation is the problem: `aten::addmm`, `aten::mm`, or `aten::bmm`?" - -**Code Reference**: Generated by `TreePerfAnalyzer.get_df_kernel_launchers_summary()` - ---- - -### 2.4. ops_unique_args (Most Detailed) - -**Purpose**: Most detailed summary - groups operations by unique combinations of operation name AND input arguments (shapes, dtypes, strides, concrete inputs). Use this to identify which specific input patterns are causing performance issues. - -### Columns - -| Column | Description | Details | -|--------|-------------|---------| -| `name` | Operation name | PyTorch operation name | -| `Input Dims` | Input tensor dimensions | Shape of input tensors (e.g., `[[1, 512, 768], [768, 768]]`) | -| `Input type` | Input tensor data types | Data types of inputs (e.g., `['c10::BFloat16', 'c10::BFloat16']`) | -| `Input Strides` | Input tensor strides | Memory layout strides for each input tensor | -| `Concrete Inputs` | Scalar input values | Non-tensor inputs like kernel size, stride, padding | -| `operation_count` | Number of occurrences | How many times this exact operation+args combination appeared | -| `total_direct_kernel_time_sum` | Total GPU time in microseconds | Sum of GPU kernel time for all occurrences | -| `total_direct_kernel_time_mean` | Mean GPU time in microseconds | Average GPU time per occurrence | -| `total_direct_kernel_time_median` | Median GPU time in microseconds | Median GPU time across occurrences | -| `total_direct_kernel_time_std` | Standard deviation | Standard deviation of GPU time across occurrences (microseconds) | -| `total_direct_kernel_time_min` | Minimum GPU time | Fastest occurrence (microseconds) | -| `total_direct_kernel_time_max` | Maximum GPU time | Slowest occurrence (microseconds) | -| `ex_UID` | Example event UID | UID of one example event (for further analysis) | -| `kernel_details_summary` | Kernel execution details | Summary of which GPU kernels were launched and their statistics | -| `trunc_kernel_details` | Truncated kernel details | Shortened version of kernel_details_summary for readability | -| `Percentage (%)` | Percentage of total GPU time | `(total_direct_kernel_time_sum / sum of all kernel time) * 100` | -| `Cumulative Percentage (%)` | Running cumulative percentage | Sum of percentages from top to current row | -| `is_recompute` | *(Optional)* Whether this operation is part of activation recomputation | `True` or `False`. Only present when `--detect_recompute` is enabled. | - -**What makes arguments "unique"**: -- Same operation with different input shapes → different rows (e.g., `aten::addmm` with (1024, 512) vs (2048, 1024)) -- Same operation with different data types → different rows (e.g., BF16 vs FP32) -- Same operation with different strides → different rows (contiguous vs transposed) -- Each unique combination of (name + Input Dims + Input type + Input Strides + Concrete Inputs) gets one row - -**What this shows**: -- Statistics for each unique operation+args combination (count, mean, median, std, min, max) -- Which specific input patterns are slow vs fast -- Performance variation across different calls to the same operation - -**Use this when**: "I see `aten::addmm` is expensive - is it slow for all input shapes, or just specific ones? Are there outliers?" - ---- - -#### Understanding kernel_details_summary - -Provides aggregated kernel statistics across all occurrences of an operation+args combination. Shows which kernels are consistently launched and their performance distribution. All durations are in microseconds (µs). - -**Example for aten::addmm** (aggregated across 150 occurrences): -```python -[ - {'kernel_name': 'elementwise_kernel', 'count': 150, 'mean_duration_us': 2.3, 'std_dev_duration_us': 0.1}, - {'kernel_name': 'Cijk_Alik_Bljk_BBS_BH...', 'count': 150, 'mean_duration_us': 45.2, 'std_dev_duration_us': 3.5} -] - -# Analysis: addmm = elementwise (2.3 µs) + GEMM (45.2 µs) across 150 calls -# → GEMM is bottleneck, with moderate variance (std_dev=3.5) → check for outliers with ex_UID -# → elementwise is consistent (std_dev=0.1) -``` - -**Use this to identify**: -- **Bottlenecks**: Which kernel in a multi-kernel operation dominates time - - Example: In the GPT-3 XL analysis, `aten::addmm` consistently showed a Tensile GEMM kernel taking 97 µs (mean) while setup kernels took <5 µs each -- **Consistency**: Low std_dev = stable, high std_dev = investigate outliers with `ex_UID` - - Example: If `std_dev_duration_us` is 15.3 when `mean_duration_us` is 120.5, that's 12.7% variance - worth investigating -- **Backend recipe**: Which kernels are always launched together (e.g., elementwise + GEMM for addmm) -- **Call count validation**: Verify all occurrences launch the same kernels (same `count`) - -**Note**: `trunc_kernel_details` provides a shortened version for spreadsheet readability. - ---- - -**Deep-dive with `ex_UID`**: - -The `ex_UID` column provides a UID of one example event with this operation+arguments combination. You can use this to access the actual event object for deeper analysis: - -```python -# Get the event object -event = tree.get_UID2event(ex_UID) - -# Analyze call stack and context: -parent = tree.get_parent_event(event) # Get parent operation -children = tree.get_children_events(event) # Get child operations -gpu_events = tree.get_gpu_events(event) # Get launched GPU kernels -tree.traverse_parents_and_print(event) # See full call stack above this -tree.traverse_subtree_and_print(event) # See full call stack below this - -# Replay the operation for benchmarking: -from TraceLens import EventReplayer -replayer = EventReplayer(event, device='cuda') -replayer.replay() # Replays the exact operation with same inputs/args -``` - -This is useful when you want to investigate a specific slow case in detail or benchmark it in isolation. See [Event Replay Documentation](./EventReplay.md) for more details on replaying operations. - -**Replaying from Perf Report (No Trace Required)**: - -You can also replay operations directly from the Excel report without needing the original trace file: - -```python -import pandas as pd -import ast -from TraceLens import EventReplayer - -# Read ops from perf report -df = pd.read_excel('perf_report.xlsx', sheet_name='ops_unique_args') - -# Convert row to event format -def row_to_event(row): - return { - 'name': row['name'], - 'args': { - 'Input Dims': ast.literal_eval(row['Input Dims']), - 'Input Strides': ast.literal_eval(row['Input Strides']), - 'Input type': ast.literal_eval(row['Input type']), - 'Concrete Inputs': ast.literal_eval(row['Concrete Inputs']), - } - } - -# Replay an operation -row = df.iloc[0] # or filter by name/args -event = row_to_event(row) -replayer = EventReplayer(event, device='cuda') -replayer.replay() - -# Get standalone repro artifacts -repro_info = replayer.get_repro_info() # Returns serializable JSON -``` - -This is particularly useful for creating standalone reproducers or benchmarking specific operations without the full model or trace. See `examples/event_replayer_example.ipynb` for complete examples including batched replay. - -**Code Reference**: Generated by `TreePerfAnalyzer.get_df_kernel_launchers_unique_args()` - ---- - -## 3. Performance Metrics Sheets (Roofline Analysis) - -For certain operation categories (GEMM, CONV, SDPA, UnaryElementwise, BinaryElementwise), TraceLens generates additional sheets with **roofline model metrics**. These sheets help you understand how efficiently operations are using the GPU's computational and memory bandwidth capabilities. - -**Important Context**: While hardware counter profilers like `rocprof compute` and `nsight compute` reveal what the GPU actually executed—including effects of padding, redundant memory movement, and cache behavior—TraceLens focuses on the useful work dictated by operator semantics. Used together, these two perspectives provide a richer picture: hardware counters expose low-level execution characteristics, while TraceLens reveals the efficiency of the computation in context. - -### Understanding the Metrics Pipeline - -The performance metrics are calculated through a 4-step pipeline: - -``` -Trace Event → Parameter Extraction → Static Metrics + Runtime → Performance Metrics - ↑ ↑ - (Compute Model) (GPU Time) -``` - -#### Step 1: Event from Trace - -Starting with an operation event from the trace JSON: - -```json -{ - 'name': 'aten::addmm', - 'args': { - 'Input Dims': [[6144], [40960, 1536], [1536, 6144], [], []], - 'Input type': ['c10::BFloat16', 'c10::BFloat16', 'c10::BFloat16', 'Scalar', 'Scalar'], - 'Input Strides': [...] - } -} -``` - -#### Step 2: Parameter Extraction - -TraceLens uses operation-specific **performance models** to extract relevant parameters from the event's `args` field: - -```python -# For aten::addmm: A @ B + C where A is (M, K), B is (K, N), C is (M, N) -params = { - 'M': 40960, # From Input Dims[1][0] - 'N': 6144, # From Input Dims[2][1] - 'K': 1536, # From Input Dims[1][1] - 'bias': True, # C tensor present - 'dtype': 'c10::BFloat16' # From Input type -} -``` - -**Note**: Each operation type (GEMM, Conv, SDPA, etc.) has its own `get_param_details()` method that knows how to extract the relevant parameters from that operation's argument structure. See [Operation Parameters Reference](#operation-parameters-reference) below for details on what gets extracted for each operation type. - -#### Step 3: Parallel Calculation - -Once parameters are extracted, two pieces of information are computed in parallel: - -**A) Static Operation Metrics** (via Compute Model) - -The performance model calculates the **theoretical work** based on the operation type: - -```python -# GEMM compute model -GFLOPS = (2 * M * N * K) / 1e9 # 773.35 GFLOPS -Bytes_moved = (M*K + K*N + M*N) * bytes_per_elem # 618.01 MB (for BF16: 2 bytes/elem) -FLOPS_per_Byte = GFLOPS * 1e9 / Bytes_moved # 1193.38 (arithmetic intensity) -``` - -These metrics are **static** - they depend only on the operation parameters, not on actual execution. - -**B) Kernel Time** (from Trace2Tree) - -The actual GPU execution time is extracted by: -1. Finding all GPU kernels launched by this operation (using Trace2Tree's hierarchical analysis) -2. Using `GPUEventAnalyser` to compute total busy time, accounting for kernel overlaps - -```python -T = 1884 µs # Measured kernel execution time -``` - -#### Step 4: Runtime Performance Metrics - -Finally, static work metrics are combined with measured time to produce **runtime performance**: - -```python -TFLOPS/s = GFLOPS / (T / 1e6) = 773.35 / (1884 / 1e6) = 410.48 TFLOPS/s -TB/s = (Bytes_moved / 1e12) / (T / 1e6) = 0.34 TB/s -``` - -These metrics tell you: -- **TFLOPS/s**: How many trillion floating-point operations per second were achieved -- **TB/s**: How many terabytes per second of memory bandwidth were utilized - -### Why This Matters: Roofline Analysis - -The combination of **TFLOPS/s**, **TB/s**, and **FLOPS/Byte** (arithmetic intensity) allows you to perform **roofline analysis**: - -- **High FLOPS/Byte** (compute-intensive): Performance is limited by compute throughput (TFLOPS/s) - - Example: Large GEMM operations (e.g., 6144×2048 × 2048×8192) - - Goal: Maximize TFLOPS/s (approach GPU's peak compute) - -- **Low FLOPS/Byte** (memory-intensive): Performance is limited by memory bandwidth (TB/s) - - Example: Element-wise operations, small GEMMs (e.g., 2048×2048) - - Goal: Maximize TB/s (approach GPU's peak bandwidth) - -**The Roofline "Knee Point"**: The boundary between memory-bound and compute-bound is determined by the GPU's hardware characteristics: - -``` -Arithmetic Intensity Threshold = Peak FLOPS / Peak Bandwidth -``` - -For example: -- **MI325X**: 1300 TFLOPS / 6000 GB/s = **~217 FLOPs/Byte** -- **H100**: 1000 TFLOPS / 3350 GB/s = **~298 FLOPs/Byte** -- **MI300X**: 1300 TFLOPS / 5300 GB/s = **~245 FLOPs/Byte** - -Operations with **FLOPs/Byte below this threshold** are memory-bound; operations **above** are compute-bound. This is why small GEMMs (FLOPs/Byte ~8-50) are memory-bound, while large GEMMs (FLOPs/Byte ~100-300) are compute-bound on these GPUs. - -By comparing your achieved TFLOPS/s and TB/s against the GPU's theoretical peaks, you can identify optimization opportunities. - -**Interpreting Performance Numbers**: - -Understanding what "good" performance looks like requires comparing against theoretical peaks and max-achievable performance: - -| GPU | Peak Compute (BF16) | Peak Memory BW | Example Utilization | -|-----|---------------------|----------------|---------------------| -| MI325X | ~1.3 PFLOPS | ~6 TB/s | 500-800 TFLOPS/s = 38-62% of peak (typical for medium GEMMs) | -| H100 | ~1.0 PFLOPS | ~3.35 TB/s | 400-700 TFLOPS/s = 40-70% of peak | -| MI300X | ~1.3 PFLOPS | ~5.3 TB/s | Similar to MI325X | - -**Understanding Theoretical vs. Real-World Performance**: - -TraceLens uses idealized assumptions that represent upper bounds on performance. The actual roofline model has two key differences from the theoretical one: - -1. **Peak FLOPS**: The theoretical peak represents hardware limits, but real-world applications typically achieve lower performance due to realistic constraints. AMD's Max-Achievable FLOPS (MAF) methodology provides more realistic performance targets. For details, see: - - [Understanding Peak and Max-Achievable FLOPS](https://rocm.blogs.amd.com/software-tools-optimization/Understanding_Peak_and_Max-Achievable_FLOPS/README.html) - - [Measuring Max-Achievable FLOPS (Part 2)](https://rocm.blogs.amd.com/software-tools-optimization/measuring-max-achievable-flops-part2/README.html) - -2. **Arithmetic Intensity**: TraceLens assumes **100% cache hit rate** — each memory location accessed once from global memory, then perfectly cached. For example, in a GEMM, matrices A and B are counted only once even if elements are reused. In reality: - - Cache misses, evictions, and redundant loads cause **actual memory movement to be higher** - - This means **actual arithmetic intensity is lower** than TraceLens calculates - - **Impact**: An operation that TraceLens shows as compute-bound (high FLOPs/Byte) might actually be memory-bound in practice - - This is a **common practice** in performance modeling (see [NVIDIA's approach](https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#math-mem-bounds)) - - For typical **compute-bound operations** (large GEMMs, SDPA, Convolutions), this limitation has minimal practical impact since they remain compute-bound even with lower arithmetic intensity - -Hardware profilers like `rocprof compute` or `nsight compute` measure actual memory transactions and are needed to determine true arithmetic intensity and memory-boundedness. - -**Real Examples** (from GPT-3 XL on MI325X, BF16): -``` -Operation TFLOPS/s TB/s FLOPs/Byte Interpretation -────────────────────────────────────────────────────────────────────────────── -mm(256×256 × 256×256) 45 2.1 ~8 Memory-bound, 35% BW utilization -mm(2048×6144 × 6144×2048) 531 0.61 ~138 Compute-bound, 40% efficiency -mm(6144×2048 × 2048×2048) 624 0.71 ~117 Compute-bound, 48% efficiency -addmm(6144×2048 × 2048×8192) 762 0.59 ~203 Compute-bound, 58% efficiency -``` - -**Understanding compute-bound vs. memory-bound**: -- **Memory-bound** (FLOPs/Byte < ~50): Small GEMM with arithmetic intensity of ~8. Performance is limited by memory bandwidth (~2.1 TB/s), not compute. To optimize, focus on improving memory access patterns. -- **Compute-bound** (FLOPs/Byte > ~100): Large GEMMs with high arithmetic intensity. Performance is limited by compute throughput (531-762 TFLOPS/s). The 40% vs 58% efficiency difference reflects kernel optimization quality (tile sizes, wave occupancy), not the compute vs. memory boundedness. - -**Important**: Arithmetic intensity (FLOPs/Byte) determines whether an operation is compute-bound or memory-bound. The percentage of peak achieved indicates optimization quality within that constraint. See [NVIDIA's GEMM Performance Guide](https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html) for more details on this distinction. - -### What These Sheets Contain - -Performance metrics sheets are generated for operations with available performance models: -- **GEMM**: Matrix multiply operations (addmm, mm, bmm, baddbmm, etc.). See [GEMMs in AI Workloads](./conceptual/aimodels_gemms.md) for how model dimensions (batch size, sequence length, hidden dimension) map to GEMM shapes. -- **CONV_fwd / CONV_bwd**: Convolution operations -- **SDPA_fwd / SDPA_bwd**: Scaled dot-product attention -- **UnaryElementwise / BinaryElementwise**: Element-wise operations - -Each sheet contains: -- All columns from `ops_unique_args` (operation name, arguments, occurrences, etc.) -- **Static metrics**: GFLOPS, Data Moved (MB), FLOPS/Byte (calculated once from parameters) -- **Compute Spec**: Combined compute type and precision (e.g., `matrix_bf16`, `vector_fp32`). This indicates: - - `matrix_*`: Operations using matrix compute units (GEMM, CONV, SDPA) - - `vector_*`: Operations using vector compute units (elementwise) - - Precision: `fp8`, `fp16`, `bf16`, `fp32`, `fp64` -- **Roofline metrics** (when `--gpu_arch_json_path` provided): Roofline Time (µs), Pct Roofline - compares achieved time to theoretical roofline bound -- **Runtime metrics**: Kernel Time (µs), TFLOPS/s, TB/s (statistics across all occurrences) - - `mean`, `median`: Central tendency - use median if variance is high - - `std_dev`: Variability - high std_dev (>10% of mean) suggests inconsistent performance - - `min`, `max`: Range - large spread may indicate outliers worth investigating - -**Note**: These sheets contain all the columns from `ops_unique_args`, so you can replay operations from these sheets using the same approach described in [Replaying from Perf Report](#replaying-from-perf-report-no-trace-required). Simply read the desired performance metrics sheet (e.g., `GEMM`, `SDPA_fwd`) instead of `ops_unique_args`. - -**Code Reference**: Generated by `TreePerfAnalyzer.build_df_perf_metrics()` and `TreePerfAnalyzer.summarize_df_perf_metrics()` - ---- - -### Operation Parameters Reference - -TraceLens extracts operation-specific parameters from trace events to calculate theoretical FLOPs and memory traffic. Each operation type has its own parameter extraction logic based on what information is needed for its performance model. - -**Common sources**: -- `Input Dims`: Tensor shapes -- `Input type`: Data types -- `Concrete Inputs`: Scalar arguments (stride, padding, etc.) - -**Data type sizes**: `fp64` = 8 bytes, `fp32` = 4 bytes, `bf16/fp16` = 2 bytes, `fp8` = 1 byte - ---- - -#### GEMM (Matrix Multiply) - -| Parameter | Meaning | Used in | -|-----------|---------|---------| -| **M** | Number of rows in first matrix | All GEMM ops | -| **N** | Number of columns in second matrix | All GEMM ops | -| **K** | Inner dimension (A cols = B rows) | All GEMM ops | -| **B** | Batch size | bmm, baddbmm | -| **bias** | Whether bias is added | addmm, baddbmm | -| **dtype** | Data type (BFloat16, Float32, etc.) | All GEMM ops | - ---- - -#### Convolution - -| Parameter | Meaning | -|-----------|---------| -| **input_shape** | Input tensor shape (N, C_in, H, W, ...) | -| **filter_shape** | Filter/weight shape (C_out, C_in/groups, kH, kW, ...) | -| **stride** | Stride for convolution (e.g., (1, 1)) | -| **padding** | Padding applied (e.g., (0, 0)) | -| **dilation** | Dilation factor | -| **groups** | Number of groups for grouped convolution | -| **bias** | Whether bias is present | -| **dtype** | Data type | - ---- - -#### SDPA (Scaled Dot-Product Attention) - -| Parameter | Meaning | -|-----------|---------| -| **B** | Batch size | -| **N_Q** | Query sequence length | -| **N_KV** | Key/Value sequence length | -| **H_Q** | Number of query attention heads | -| **H_KV** | Number of key/value heads (for GQA/MQA) | -| **d_h_qk** | Head dimension for Q and K | -| **d_h_v** | Head dimension for V | -| **causal** | Whether causal masking is applied | -| **dropout** | Dropout probability | -| **dtype** | Data type | - -**Note**: Different attention implementations use different tensor layouts (BHND vs BNHD) - ---- - -#### Element-wise Operations - -| Parameter | Meaning | -|-----------|---------| -| **shape** | Tensor shape | -| **dtype** | Data type | - -For binary ops (add, mul, etc.), two shapes are extracted and broadcasting is handled automatically. - ---- - -**FLOPs and Bytes Calculation**: For the specific formulas used to calculate FLOPs and memory traffic from these parameters, see the performance model implementations in `TraceLens/PerfModel/perf_model.py`. - -**Extending with Custom Operations**: Performance metrics sheets are only generated for operations that have a registered performance model. If your workload includes custom operations (e.g., from Megatron, vLLM, or other libraries), you can extend TraceLens by: -1. Creating a performance model class for your operation (inherit from `GEMM`, `CONV`, `SDPA`, etc.) -2. Implementing `get_param_details()`, `flops()`, and `bytes()` methods -3. Passing an extension file via `--extension_file` when generating the report - -See `examples/megatron_extension.py` for a complete example of extending TraceLens with custom Megatron operations. - -**Note**: If your custom operation is frequently used across multiple projects, we can work to add it as a native operation in TraceLens. Please open an issue or reach out to discuss integration. - ---- - -## 4. Collective Communication Analysis - -**Purpose**: Analyzes NCCL collective communication operations from a single rank's perspective. Shows which collectives are taking the most time on this rank. - -**Sheet**: `coll_analysis` - -**Note**: This is single-rank analysis. For multi-rank synchronization analysis (skew, stragglers), see [NCCL Analyzer documentation](./NcclAnalyser.md) and [Multi-Rank Collective Report](./generate_multi_rank_collective_report_pytorch.md). - -### What This Sheet Contains - -This sheet provides aggregated statistics for collective operations seen by this rank, grouped by collective type, process group, data type, and message size. - -### Key Columns - -**Collective Identification**: -| Column | Description | -|--------|-------------| -| `rank` | Rank ID (single rank from this trace) | -| `Process Group Name` | Process group identifier | -| `Process Group Ranks` | List of ranks in this process group (e.g., `[0, 1, 2, 3, 4, 5, 6, 7]`) | -| `Collective name` | Type of collective operation (e.g., `allreduce`, `allgather`, `reduce_scatter`) | -| `Group size` | Number of ranks in the process group | -| `dtype` | Data type (e.g., `Float`, `BFloat16`) | -| `In msg nelems` | Number of elements in input message | -| `Out msg nelems` | Number of elements in output message | -| `In split size` | Input split configuration (for split collectives) | -| `Out split size` | Output split configuration (for split collectives) | -| `stream` | GPU stream ID where collective executes | - -**Message Sizes**: -| Column | Description | -|--------|-------------| -| `In msg size (MB)_first` | Input message size in megabytes | -| `Out msg size (MB)_first` | Output message size in megabytes | - -**Duration Statistics** (all times in microseconds): -| Column | Description | -|--------|-------------| -| `dur_sum` | Total duration across all occurrences | -| `dur_mean` | Mean duration per occurrence | -| `dur_std` | Standard deviation of duration | -| `dur_min` | Minimum duration | -| `dur_max` | Maximum duration | -| `operation_count` | Number of times this collective appeared | - -### Understanding the Metrics - -**Duration**: Time this rank spends in the collective operation (in microseconds). This includes: -- Synchronization delay waiting for other ranks -- Actual data transfer time - -**Important**: From a single rank's perspective, we cannot determine if time is spent waiting for other ranks (sync) or in actual communication. For multi-rank analysis with skew detection and straggler identification, use NCCL Analyzer with traces from all ranks. See [NCCL Analyzer documentation](./NcclAnalyser.md). - -### Interpreting Results - -**High total duration**: Sort by `dur_sum` to find the most time-consuming collective operations on this rank. - -**High variance** (large `dur_std` or difference between `dur_min` and `dur_max`): -- Indicates inconsistent performance -- May suggest network contention or varying degrees of synchronization delay -- Consider investigating with multi-rank traces to identify stragglers - -**Frequent operations** (high `operation_count`): -- Even if individual operations are fast, high frequency can add up (`dur_sum` accounts for this) -- Check if collectives can be batched or overlapped with computation - -**Example Analysis**: -``` -If you see: -- allreduce with dur_mean=3596 µs (3.6 ms), operation_count=2 -- But dur_std=35.5 µs (low variance) -→ Collective is consistent and moderately fast - -If you see: -- allreduce with dur_mean=192.6 µs, dur_std=313.7 µs, dur_max=1019.4 µs -→ High variance! One occurrence took 1ms while another took 60µs -→ Investigate with multi-rank trace to find if this rank or another is the straggler -``` - -**Code Reference**: Generated by `NcclAnalyser.build_df_summary_long()`. Enable with `--disable_coll_analysis` flag (enabled by default). - ---- - -## 5. Additional Sheets - -Depending on the command-line arguments used when generating the report, additional analysis sheets may be present: - -- **kernel_summary**: Per-kernel statistics aggregated by kernel name (enabled with `--kernel_summary`) - -- **short_kernel_histogram**, **short_kernels_summary**: Analysis of very short kernels (enabled with `--short_kernel_study`) - -- **unified_perf_summary**: Unified perf metrics for all ops with perf models or leaf ops that launch GPU kernels. Includes GFLOPS, TFLOPS/s, Data Moved, FLOPS/Byte, TB/s metrics aggregated by unique args. When `--detect_recompute` is enabled, an `is_recompute` column is added to split rows by recompute status. - ---- - -## 6. Common Analysis Workflows - -**Typical top-down analysis flow**: - -1. **Start with `gpu_timeline`**: Get high-level time breakdown (computation, communication, idle) - - High idle time → CPU bottleneck or kernel launch issues - - High exposed communication → Poor overlap with computation - -2. **Identify bottleneck categories** (`ops_summary_by_category`): Which operation types dominate? - - GEMM, Convolution, Attention, Elementwise, etc. - -3. **Find expensive operations** (`ops_summary`): Which specific operations take the most time? - -4. **Analyze variants** (`ops_unique_args`): Which input shapes/arguments are slow? - - Compare performance across different dimensions - - Check for stride or alignment issues - -5. **Deep dive with roofline** (GEMM, SDPA_fwd, etc.): Assess efficiency vs. hardware peaks - -6. **Check collectives** (`coll_analysis`): Analyze communication costs in distributed training - -7. **Replay operations**: Reproduce and benchmark specific cases in isolation - ---- - -## 7. Related Documentation - -- [Performance Report Generation Guide](./generate_perf_report.md) - How to generate reports -- [Trace2Tree Documentation](./Trace2Tree.md) - Understanding the call stack analysis -- [GEMMs in AI Workloads](./conceptual/aimodels_gemms.md) - Understanding how model dimensions map to GEMM shapes -- [TreePerf Examples](../examples/tree_perf_example.ipynb) - Example analysis workflows -- [Call Stack Analysis Example](../examples/call_stack_analysis.ipynb) - Real workflow for debugging performance -