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
- 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.
- 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.
- 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 |
- 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 |
- 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
- perf_model.py
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.