[WIP] Synthetic op perf model resolution via SGLang capture linking#714
Merged
rkarhila-amd merged 25 commits intoJun 8, 2026
Conversation
## Summary
- Preserve top-level trace JSON fields as trace-level metadata while
keeping `traceEvents` as the event payload.
- Expose metadata through `TraceToTree` and `TreePerfAnalyzer` for
report consumers.
- Keep metadata optional so direct synthetic-event construction remains
backward compatible.
## Example usage
```python
from TraceLens.TreePerf.tree_perf import TreePerfAnalyzer
analyzer = TreePerfAnalyzer.from_file("profile.trace.json", rebuild_tree=False)
metadata = analyzer.trace_metadata
gpu_names = [
gpu.get("name", "unknown GPU")
for gpu in metadata.get("deviceProperties", [])
]
print(f"GPUs: {', '.join(gpu_names) if gpu_names else 'unknown'}")
print(f"GPU count: {len(gpu_names)}")
rank_info = metadata.get("distributedInfo") or {}
if rank_info:
print(f"Rank: {rank_info.get('rank')} / {rank_info.get('world_size')}")
```
Verified on a small real PyTorch Chrome trace. The exact GPU model is
trace-dependent:
```text
GPUs: AMD Instinct MI300X
GPU count: 1
Distributed info: not present
```
## Test plan
- `python -m pytest tests/test_trace_metadata.py`
- `python -m pytest tests/test_trace_metadata.py
tests/test_kernel_launchers.py -q`
Replaces #643 after cleaning up the branch name and commit history.
## Summary
- Moves perf-modeled torch-op categories onto perf model classes via
`category`, optional `bwd_category`, and `sheet_category` when legacy
sheet membership differs from runtime categorization.
- Keeps category-only ops in `CATEGORY_ONLY_OP_MAPPING`, so unified
reports can categorize ops without perf models without adding them to
legacy per-category perf-model sheets.
- Consolidates torch op categorization helpers into
`TraceLens.PerfModel.torch_op_mapping` and removes the old compatibility
maps/module (`dict_cat2names`, `dict_base_class2category`, and
`op_categories.py`).
- Derives legacy report-sheet membership from
`op_to_perf_model_class_map` via `build_sheet_category_to_op_names`.
## Behavior Changes
- `GroupedGemm` stays separate from `GEMM` as `GroupedGEMM_fwd` /
`GroupedGEMM_bwd`.
- Explicit SDPA backward names and other category-only ops still
categorize in unified reports, but are not added to legacy perf-model
sheets unless they have a perf model.
- Linked-forward backward metrics use `bwd_category` when the forward
perf model intentionally supports backward metrics.
- Legacy elementwise sheets remain split through `sheet_category` while
runtime categorization still reports `elementwise`.
- RMSNorm keeps its existing legacy `RMSNorm_fwd` sheet through an
explicit `sheet_category`.
## Compatibility Notes
Extensions should declare categories on perf model classes provided
through `perf_model_extension`, or use `op_category_extension =
{op_name: final_category}` for category-only ops. Direct legacy sheet
mutation through `dict_cat2names_extension` is no longer part of the
categorization flow.
## Test Plan
Passed locally:
- `python -m black --check TraceLens tests examples`
- `python -m pytest tests/test_torch_op_categorization_registry.py
tests/test_pseudo_ops_extension.py
tests/test_primus_op_categorization.py
tests/test_primus_turbo_grouped_gemm.py -q`
- `python -m pytest tests/test_inference_perf_report.py -q`
Also ran part of `python -m pytest tests/test_perf_report_regression.py
-q`; it was manually backgrounded after several passing cases. CI covers
the full regression environment.
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
<!-- Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> # Background Issue #619 — Enable perf metrics (GFLOPS, TFLOPS/s, TB/s) for Triton kernels generated by `torch.compile`. When `torch.compile` fuses ATen ops into Triton kernels, TraceLens previously showed blank metrics for these kernels. This PR adds `TritonCompiledPerfModel` to compute GFLOPS, Data Moved, and throughput metrics. # Reproducer — `examples/repro_triton_trace.py` Builds a `TransformerBlock` (RMSNorm + Attention + SwiGLU MLP, `dim=2048`, bf16), compiles with `torch.compile`, and captures a Chrome trace with `torch.profiler`. The trace contains 4 Triton kernels used to validate the perf model. # Two-Tier Approach - **V2 (Trace-Intrinsic, primary)** — Extracts metadata directly from Chrome trace event args (`Concrete Inputs`, `Input Dims`, `Input type`). No disk I/O, no cache dependency. Self-contained and portable. Requires PyTorch 2.4+. - **V1 (Cache-Based, fallback)** — Parses torchinductor cache `.py` files on disk. Used only when V2 fails (older traces lacking required fields). Cache dir can be specified explicitly via `--inductor_cache_dir`. - **If both fail** — The kernel appears in the report with timing but no perf metrics. ## Approaches Comparison <img width="1459" height="774" alt="image" src="https://github.com/user-attachments/assets/a4fc0b81-9d1e-4987-8700-a6d9b577b96e" /> ## Future work: where `TORCH_COMPILE_DEBUG=1` adds value | Feature | File | |---|---| | Verify FLOPs formula with exact shapes | `fx_graph_readable.py` | | Verify byte counting against exact load/store ops | `ir_post_fusion.txt` | | Graph break detection | `torchdynamo/debug.log` | | Recompile tracking | `torchdynamo/debug.log` | | Reproducing kernel issues in isolation | `fx_graph_runnable.py` | # How to Run ```bash python -m TraceLens.Reporting.generate_perf_report_pytorch \ --profile_json_path torch_trace_output_claude/triton_repro/trace.json.gz \ --output_xlsx_path torch_trace_output_claude/triton_repro/perf_report.xlsx \ --output_csvs_dir torch_trace_output_claude/triton_repro/csvs_output \ --enable_pseudo_ops --group_by_num_kernels ``` With explicit inductor cache (for V1 fallback on a different machine): ```bash python -m TraceLens.Reporting.generate_perf_report_pytorch \ --profile_json_path trace.json.gz \ --output_xlsx_path perf_report.xlsx \ --inductor_cache_dir /path/to/torchinductor/cache ``` # Results 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. | Kernel | GFLOPS | Data Moved (MB) | FLOPS/Byte | TB/s | TFLOPS/s | |--------|--------|-----------------|------------|------|----------| | `triton_red_fused_add_mean_mul_pow_rsqrt_0` (RMSNorm) | 0.470 | 256 | 1.75 | 3.48 | 6.10 | | `triton_red_fused__unsafe_view_add_mean_mul_pow_rsqrt_1` (RMSNorm+residual) | 0.470 | 384 | 1.17 | 4.44 | 5.18 | | `triton_poi_fused__unsafe_view_mul_silu_2` (SwiGLU) | 1.342 | 1,536 | 0.83 | 3.05 | 2.55 | | `triton_poi_fused__unsafe_view_add_3` (residual add) | 0.067 | 512 | 0.12 | 3.81 | 0.48 | # Key Design Decisions (addressing review feedback) ### 1. Broadened mapping API (no Triton-specific logic in `tree_perf.py`) Added `resolve_perf_model_class()` and `register_perf_model_matcher()` to `torch_op_mapping.py`. `TritonCompiledPerfModel` registers its own match rule. `tree_perf.py` uses a single uniform dispatch path — no prefix checks, no bare `except Exception`. `_has_perf_model()` uses a generic `can_model(event)` protocol so kernels that match the pattern but lack trace-arg metadata are correctly reported as unmodeled. Reusable for any future perf model with generated names. ### 2. Explicit `inductor_cache_dir` kwarg Added `--inductor_cache_dir` CLI argument following the `arch` / `python_path` / `enable_origami` convention. Threaded through `TreePerfAnalyzer` → `_perf_model_init_kwargs` → `TritonCompiledPerfModel` → `_cache_dirs()`. When provided, takes priority over env var and auto-detected paths. Conditionally passed via signature introspection (same pattern as `enable_origami`). ### 3. No torch import at analysis time Vendored the op registry as `TraceLens/PerfModel/_known_aten_ops.py` — an auto-generated Python module with 960 op names from torch 2.11.0+rocm7.2. Generated by `scripts/generate_known_aten_ops.py`, refreshed per torch version bump. Uses `.py` (not `.json`) so it's always included in the installed package. `_parse_kernel_name` imports the set at module level — same coverage, deterministic across analysis hosts. ### 4. Regression tests - **`tests/test_triton_compiled_perf_model.py`** — 6 unit tests covering `_parse_kernel_name`, `_meta_from_trace_args`, `flops()`, and `bytes()` with a RMSNorm kernel. Clear expected values; any change that alters output breaks loudly. - **`tests/traces/torch_compile_triton/`** — real trace + golden reference CSVs, auto-discovered by `test_perf_report_regression.py`. Validates the full pipeline end-to-end. # Code Changes Summary ### `tree_perf.py` — uniform dispatch, no Triton-specific logic - `compute_perf_metrics()`: replaced `op_to_perf_model_class_map.get()` + Triton prefix-check fallback with single call to `resolve_perf_model_class()` - `_has_perf_model()`: uses `resolve_perf_model_class()` + generic `can_model(event)` check - `build_df_unified_perf_table()`: removed 20-line opportunistic Triton block + bare `except Exception`; restored simple `row["perf_params"] = None` - `_perf_model_init_kwargs()`: conditionally passes `inductor_cache_dir` (same signature introspection pattern as `enable_origami`) ### `torch_op_mapping.py` — pattern-based matcher API ```python _perf_model_matchers: list = [] def register_perf_model_matcher(matcher): ... def resolve_perf_model_class(name): ... # dict first, then matchers def _match_triton_compiled(name): # lazy import, no torch pulled if name.startswith(("triton_poi_", "triton_red_", "triton_per_")): from ...triton_compiled_perf_model import TritonCompiledPerfModel return TritonCompiledPerfModel return None register_perf_model_matcher(_match_triton_compiled) ``` ### `triton_compiled_perf_model.py` — perf model class - `_parse_kernel_name()`: DP segmentation against vendored op registry (`_known_aten_ops.py`). Fallback to `_FLOPS_PER_ELEM` keys when registry unavailable. - `_meta_from_trace_args()`: V2 extraction from `Concrete Inputs` / `Input Dims` / `Input type`. Skips non-integer scalars (floats, booleans). - `_cache_dirs()` / `_lookup()`: V1 cache parsing. Priority: `inductor_cache_dir` kwarg → env var → defaults. - `can_model()`: static method — returns `True` only if V2 trace args are available and parseable. - `flops()` / `bytes()`: raises `NotImplementedError` if `_meta` is None (gated by `can_model`). # Files Changed | File | Change | |------|--------| | `TraceLens/PerfModel/triton_compiled_perf_model.py` | New perf model class (V2 + V1), kernel name parser, `can_model()` | | `TraceLens/PerfModel/torch_op_mapping.py` | Pattern-based matcher API + Triton registration | | `TraceLens/PerfModel/_known_aten_ops.py` | Vendored op registry (auto-generated) | | `TraceLens/PerfModel/__init__.py` | Export `resolve_perf_model_class` | | `TraceLens/TreePerf/tree_perf.py` | Uniform dispatch via `resolve_perf_model_class`, `inductor_cache_dir` plumbing | | `TraceLens/Reporting/generate_perf_report_pytorch.py` | `--inductor_cache_dir` CLI arg | | `scripts/generate_known_aten_ops.py` | Maintenance script for op registry | | `tests/test_triton_compiled_perf_model.py` | Unit tests | | `tests/traces/torch_compile_triton/` | Trace + golden reference CSVs | | `docs/triton_perf_model_walkthrough.md` | Walkthrough documentation | | `examples/repro_triton_trace.py` | Reproducer script | | `examples/read_triton_trace_args.py` | Trace inspection utility | # Known Limitations / Future Improvements | Item | Description | Priority | |------|-------------|----------| | `can_model` only checks V2 | Kernels that could be modeled via V1 (inductor cache) but lack trace args will report `has_perf_model=False`. V1 check would require disk I/O in a simple availability check. | Low — V2 is primary for modern traces | | `getpass.getuser()` in `_cache_dirs` | Can raise `KeyError` in containers without `$USER`/`/etc/passwd`. Bypassed when `--inductor_cache_dir` is provided. | Low — pre-existing, not introduced by this PR | | Scalar position assumption | `scalars[-2], scalars[-1]` assumes xnumel/rnumel are always the last integer scalars in `Concrete Inputs`. Needs verification against a wider set of traces. | Medium — investigate in follow-up | | V1 wrapper regex fragility | `content.find("''', device_str=")` and `"async_compile.triton"` sentinel may break if PyTorch changes Inductor's generated code format. | Low — V1 is fallback only | | More test cases | Current tests cover V2 + RMSNorm. Follow-ups: V1 wrapper-parsing path, multi-word ops, no-registry fallback. | Medium | --------- Co-authored-by: Matai <janmatai@ctr2-alola-ctrl-01.amd.com>
Reverting PR #603 as it is causing problematic Trace2Tree behavior
Prevent users from calling tracediff_comparison_extension when passing graph capture folder as the extension uses TraceDiff which does not support graph capture.
`get_kernel_launchers` calls `is_communication_string` once per kernel:
`tree_perf.py`:
```
for kernel in kernel_events:
# Skip nccl if not included
is_nccl = TraceEventUtils.is_communication_string(kernel.get("name", ""))
if is_nccl and not include_nccl:
continue
```
`get_communication_regexes` which was called by
`is_communication_string` was compiling the same regex object every
invocation, inflating runtime.
By avoiding recompilation, runtime has improved. Across 11 traces, trace
processing runtime has improved by up to 1.15x.
Add date and commit hash to generated wheel name
…rocess_gemm_ops (#665) ### Problem End-to-end xlsx report generation with a mixtral 8x7B FP8 Jax profile (B300, JAX 26.04, FP8) crashed at `process_gemm_ops`. ### Summary FP8 GEMMs with `damax_output=true` return a tuple output rather than a single tensor: `(f8e5m2[M,N]{1,0}, f32[], s8[W]{0})`: the result tensor, an amax scalar, and a workspace buffer. The existing parser split on `}`, assuming all tuple elements have `{}` layout annotations, but scalars like `f32[]` have none. This caused the split to lump multiple elements together, producing more than one `dtype` match and raising `ValueError: Did not find wide output`. Fixed by replacing the split with a regex tokeniser that extracts all `dtype[shape]{layout}` and `dtype[shape]` tokens from the tuple, filters out scalars, and takes the first non-scalar tensor as the result. Also fixed a related `KeyError`: some operand references in these FP8 GEMMs carry a `/*index=N*/` tuple-element prefix (e.g. `/*index=5*/%gte.remat.13`). The prefix wasn't being stripped before the `hlo_ops` dict lookup. The fix strips it with a regex and skips the operand gracefully if the base name still isn't present. ### Files Changed `TraceLens/util.py`, `JaxProfileProcessor.process_gemm_ops()` specifically --- ### Fix 1: Tuple output parsing (`damax_output=true`) ``` ## Before output_list = outputs[1:-2].split("},") # this code assumes that the first output is the one we care about # we should be able to make this an RE sizes_string = [ [i, d] for i in output_list for d in dtypes if i.startswith(d) ] if len(sizes_string) != 1: raise ValueError("Did not find wide output ", op) sizes_string = sizes_string[0] sizes_string[0] = ( sizes_string[0] + "}" ) # restore the } that was removed ## After # Extract all tensor tokens from the tuple using regex so that # scalars (e.g. f32[]) and workspace buffers (e.g. s8[N]{0}) # don't corrupt the split. Handles damax_output=true tuples like # (f8e5m2[M,N]{1,0}, f32[], s8[W]{0}). inner = outputs[1:-1] tokens = re.findall( r"[a-z0-9]+\[[^\]]*\]\{[^}]*\}|[a-z0-9]+\[[^\]]*\]", inner ) tensor_tokens = [ t for t in tokens if any(t.startswith(d) for d in dtypes) and not t.endswith("[]") ] if len(tensor_tokens) == 0: raise ValueError("Did not find wide output ", op) # Take the first tensor output; for FP8 GEMMs this is the result. sizes_string = [ tensor_tokens[0], next(d for d in dtypes if tensor_tokens[0].startswith(d)), ] ``` --- ### Fix 2: `/*index=N*/` operand reference prefix ``` ### Before elif "[" not in opid: output = hlo_ops[opid]["output"] if any( output.startswith(d) for d in dtypes + ["f8"] ) and not output.endswith("[]"): operand_list.append(hlo_ops[opid]["output"]) ### After elif "[" not in opid: # Strip /*index=N*/ prefix that appears in some FP8 operand refs hlo_key = re.sub(r"^/\*index=\d+\*/", "", opid).strip() if hlo_key not in hlo_ops: continue output = hlo_ops[hlo_key]["output"] if any( output.startswith(d) for d in dtypes + ["f8"] ) and not output.endswith("[]"): operand_list.append(hlo_ops[hlo_key]["output"]) ``` ### Tests Verified end-to-end xlsx report generation with a mixtral 8x7B FP8 Jax profile (B300, JAX 26.04, FP8). Previously crashed at `process_gemm_ops`, now produces complete report. <!-- Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> # Pull Request Template > **Note to AMDers:** > This is a public repository. Please do **not** upload any confidential or customer data. Make sure all such data has been anonymized or removed before making this PR. If you need to attach any private files or links, please insert a Internal OneDrive Link or a Jira Ticket Link instead.
# fix: filter scalar operands in pb-format path of process_gemm_ops and
add missing FP8 and pred dtype sizes to get_df_xla_perf
A couple of issues concerning FP8 GEMMs in JAX traces are causing
analysis reports to crash.
FP8 GEMMs (`__cublas$lt$matmul$f8`) include two `f32[]` scalar operands
representing input and weight scales. The XLA text-dump code path
correctly ignores these via a not `output.endswith("[]")`
guard, but the pb-format inline metadata path had no equivalent filter,
causing scalars to be appended to `operand_list`. This produces a count
of 4 operands, incorrectly tripping the bias
epilogue guard and raising a `ValueError` that crashes the entire report
before any output is written.
`get_df_xla_perf` computes total input bytes per XLA kernel event by
looking up bytes-per-element in a `dtype_to_bytes` dict. This dict is
missing all FP8 dtype variants (`f8e4m3fn`, `f8e4m3fnuz`,
`f8e5m2`, `f8e5m2fnuz`) and the boolean pred type. Any trace containing
XLA kernels with these operand types raises a `KeyError`, crashing the
report. While separate from the `process_gemm_ops` scalar-operand issue
(above), it is triggered by the same class of FP8 traces and also
crashes reports.
### Files changed
`TraceLens/util.py`
```
# process_gemm_ops(), ~line 332
# Before:
for opid in op["operands"]:
if "[" in opid and "]" in opid:
# pb format, shapes in operand list
operand_list.append(opid)
else:
output = hlo_ops[opid]["output"]
# After:
for opid in op["operands"]:
if "[" in opid and "]" in opid and not opid.split("[")[1].startswith("]"):
# pb format, shapes in operand list; exclude scalars e.g. f32[]
operand_list.append(opid)
elif "[" not in opid:
output = hlo_ops[opid]["output"]
```
`TraceLens/TreePerf/tree_perf.py`
```
# get_df_xla_perf(), ~line 3218
# Before:
dtype_to_bytes = {
"f32": 4, "bf16": 2, "s32": 4, "fp16": 2,
"u32": 4, "f16": 2, "u64": 8,
}
# After:
dtype_to_bytes = {
"f32": 4, "bf16": 2, "s32": 4, "fp16": 2,
"u32": 4, "f16": 2, "u64": 8,
"f8e4m3fn": 1, "f8e4m3fnuz": 1, "f8e5m2": 1, "f8e5m2fnuz": 1,
"pred": 1,
}
```
### Testing
Run against any FP8 JAX trace (e.g. Mixtral 8x7B FP8 .xplane.pb). Before
this fix the command crashes immediately; after it completes and
produces output.
I tested the existing `test_jax_perf_report.py` regression suite.
<!--
Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
See LICENSE for license information.
-->
# Pull Request Template
> **Note to AMDers:**
> This is a public repository. Please do **not** upload any confidential
or customer data. Make sure all such data has been anonymized or removed
before making this PR. If you need to attach any private files or links,
please insert a Internal OneDrive Link or a Jira Ticket Link instead.
…#698) When creating the Agent analysis for TraceLens, we created a set of utilities to pull GPU architecture files from known locations rather than explicitly providing them. This PR wires these utilities up to the individual reports in TraceLens
This PR adds a script under TraceUtils that finds common execution blocks between two traces. <!-- Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> --------- Co-authored-by: Deval Shah <devashah@amd.com>
This pull request adds support for SGLang version 0.5.12 to the inference analysis workflow, including updates to the build script, Docker image selection, and patch infrastructure. It also introduces enhancements to profiling and kernel shape discovery, enabling more detailed profiling capabilities and improved trace annotation. The main changes are grouped below. **Support for SGLang 0.5.12:** * Updated `build_docker_sglang.sh` to recognize and handle SGLang version 0.5.12, including version normalization, error messages, and base image resolution for MI300 and MI350/MI355 GPUs. [[1]](diffhunk://#diff-f985f887c1e41898c314a7efd4b32b4a9a35ac13f5cf42603328ac1ac9cdad88R26) [[2]](diffhunk://#diff-f985f887c1e41898c314a7efd4b32b4a9a35ac13f5cf42603328ac1ac9cdad88R39-R40) [[3]](diffhunk://#diff-f985f887c1e41898c314a7efd4b32b4a9a35ac13f5cf42603328ac1ac9cdad88R120-R122) [[4]](diffhunk://#diff-f985f887c1e41898c314a7efd4b32b4a9a35ac13f5cf42603328ac1ac9cdad88L125-R131) [[5]](diffhunk://#diff-f985f887c1e41898c314a7efd4b32b4a9a35ac13f5cf42603328ac1ac9cdad88R151-R156) [[6]](diffhunk://#diff-539cae780615a1cb686478205a499d43d5d955a42f1493bbd52b3af97aec7471R117-R118) * Added new patch directories and base image references for SGLang 0.5.12, ensuring compatibility with the latest version. [[1]](diffhunk://#diff-00debabc177bb5166e4529f6bebe30813c8e6160035a079c958501241e78dd06R1-R115) [[2]](diffhunk://#diff-06ee3359cb725dd1eea938585390696e2d1e734f1c031283e2eefa2eede91d13R1-R17) [[3]](diffhunk://#diff-453f2c82b26175f7623f6ad38d08ced2b7f4249b3c1b7ab41a28db0782c40de9R1-R13) [[4]](diffhunk://#diff-972a18ff5bbdf68023dc051002df9e26f2fb5ff182a8e1d101f3b80cbc109fe8R1-R24) [[5]](diffhunk://#diff-82dcb07a74e1667569a76ad8085d653758ced7aea4bc9a375faa3b4e7bca210dR1-R51) [[6]](diffhunk://#diff-5aca4c27a725b55de022495f5ee61d3ce32045940a933e606ce4895f9a1c90a5R1-R25) **Patch updates for SGLang 0.5.12:** * Added or updated patches for CUDA graph runner, fused MoE Triton kernels, HTTP server, IO structures, profiling utilities, and scheduler to ensure compatibility and leverage new profiling features in SGLang 0.5.12. [[1]](diffhunk://#diff-00debabc177bb5166e4529f6bebe30813c8e6160035a079c958501241e78dd06R1-R115) [[2]](diffhunk://#diff-06ee3359cb725dd1eea938585390696e2d1e734f1c031283e2eefa2eede91d13R1-R17) [[3]](diffhunk://#diff-453f2c82b26175f7623f6ad38d08ced2b7f4249b3c1b7ab41a28db0782c40de9R1-R13) [[4]](diffhunk://#diff-972a18ff5bbdf68023dc051002df9e26f2fb5ff182a8e1d101f3b80cbc109fe8R1-R24) [[5]](diffhunk://#diff-82dcb07a74e1667569a76ad8085d653758ced7aea4bc9a375faa3b4e7bca210dR1-R51) [[6]](diffhunk://#diff-5aca4c27a725b55de022495f5ee61d3ce32045940a933e606ce4895f9a1c90a5R1-R25) <!-- Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> --------- Co-authored-by: Deval Shah <devashah@amd.com>
This pull request updates the way profiler trace directories are set up in several patch files for different `vllm` versions. Instead of using the `capture_torch_profiler_dir` directly, the code now constructs the trace directory by appending `/capture_traces` to the `torch_profiler_dir` path. This ensures consistency in where profiler traces are stored across all supported versions. **Profiler trace directory update:** * Changed the assignment of `trace_dir` in all `config_vllm_v0.14.0.patch`, `config_vllm_v0.15.0.patch`, `config_vllm_v0.16.0.patch`, `config_vllm_v0.17.0.patch`, `config_vllm_v0.18.0.patch`, `config_vllm_v0.19.0.patch`, `config_vllm_v0.20.0.patch`, and `config_vllm_v0.21.0.patch` to use `self.vllm_config.profiler_config.torch_profiler_dir + "/capture_traces"` instead of `capture_torch_profiler_dir`. [[1]](diffhunk://#diff-e52de35a2f127f650a0c366525b4b165b1536da40a0fa588affe361a325bdb25L68-R68) [[2]](diffhunk://#diff-9e247fc5fd178a9e0c2cdc8069a08d05bdf8d6fa3221b0586a744469d5ad7f3dL62-R62) [[3]](diffhunk://#diff-1aff913653e81e2925c0141dd2b8b287184d4cd0bef7e8d78e06eaa77e13d028L62-R62) [[4]](diffhunk://#diff-573509070a76968c5ec8c695a9ce568ec994032c2ef8979d8fd7441fcdae7df1L62-R62) [[5]](diffhunk://#diff-96c8bc8d1ef177a538f5707e2fa92da3ec16f330c6fb3c73000977a8cd07eac3L62-R62) [[6]](diffhunk://#diff-d0a52249b1cb29a6c189770e443f0998d087a9b9ada0a1c5fc8b6340b84e885fL62-R62) [[7]](diffhunk://#diff-bbb65cd1e9dc0fa6e7e351de23da5e4c7365fa4aadd36bcfec64acdbbf4f01fbL61-R61) [[8]](diffhunk://#diff-8a129650614cc4d2ab18e798ab98dcb1763f00c24240f3b238fcec5923637dfeL61-R61) <!-- Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> Co-authored-by: Deval Shah <devashah@amd.com>
ops_summary.csv had non-deterministic category ordering for the sets in the "Categories" column. This would cause issues in TraceLens Agent evals. Now the set is sorted.
…636) ### Fix mixed steady-state window selection and naming in `split_inference_trace_annotation.py` - **Mixed-window selection**: filter candidate windows to those containing at least one prefill-bearing step (`context_requests > 0`, i.e. pure prefill or truly-mixed). Previously the selection's `pd_count` definition (any context > 0) and the naming's `num_prefilldecode` definition (context AND generation > 0) were inconsistent, so on traces where prefill steps are sparse the `mixed_steady_state_*` window was often selected with zero prefill activity. Falls back to the full candidate pool only when no window in the largest steady-state region contains any prefill step. - **Steady-state filename format**: unify all `--find-steady-state` outputs (mixed / decode_only / prefilldecode) on `prefill_X_prefilldecode_Y_decode_Z_bsA_concB`, including 1-step max-PD windows that previously fell back to the literal step name. The new `prefill_X_` prefix surfaces pure-prefill counts that were previously hidden under `num_prefilldecode = 0` in the filename. - `--store-single-iteration` and `--dummy` outputs are unchanged. <!-- Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> Co-authored-by: Deval Shah <devashah@amd.com>
This pull request adds support for identifying, creating, and reporting a new pseudo operation, `pseudo_mla_prefill_fwd`, for MLA fp8 prefill attention in the TraceLens performance analysis tool. The changes include logic to detect relevant events in trace trees, inject the new pseudo op, extract its parameters, and ensure it is included in performance reporting. Key changes: **MLA Prefill Pseudo Op Detection and Injection** - Added a new extension module, `mla_prefill_pseudo_ops.py`, which detects Python function events corresponding to MLA fp8 prefill attention, finds their associated CPU ops, and injects a `pseudo_mla_prefill_fwd` pseudo op into the trace tree. - Updated the pseudo op extension utility to automatically detect and apply the new MLA prefill pseudo op extension when relevant events are present in the trace. **Performance Model and Reporting Integration** - Added the `pseudo_mla_prefill_fwd` operation to the pseudo op mappings so it is recognized and handled by the performance model. - Implemented a new class, `pseudo_mla_prefill_fwd`, in `attention_perf_model_extensions.py`, which extracts relevant parameters from the event, including the new `d_h_v` dimension. - Included `pseudo_mla_prefill_fwd` in the list of attention operations for PyTorch inference performance reporting. **Additional Pseudo Op Mappings** - Added mappings for several new sglang and batched GEMM kernel pseudo ops to the performance model, improving overall op coverage. <!-- Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. See LICENSE for license information. --> --------- Co-authored-by: Deval Shah <devashah@amd.com>
Previously, `default_categorizer` would return `event.get(TraceLens.util.TraceEventUtils.TraceKeys.Category)` requiring a long attrbute chain lookup. By replacing this with a simple dict lookup, `event["cat"]`, runtime has been improved. There is no more `.get()` fallback since `_preprocess_and_index_events` has been modified to populate the `cat` field with `none` if the event doesn't have a category field. | Trace | Upstream (s) | Optimized (s) | Δ (s) | Δ % | |---|---|---|---|---| | trace1 | 89.6 | 83.9 | −5.7 | −6% | | trace2 | 0.48 | 0.42 | −0.06 | −12% | | trace3 | 73.9 | 68.1 | −5.8 | −8% | | trace4 | 74.6 | 72.3 | −2.3 | −3% | | trace5 | 111.6 | 108.0 | −3.6 | −3% | | trace6 | 1295.6 | 1245.5 | −50.1 | −4% | | trace7 | 55.7 | 50.3 | −5.4 | −10% | | trace8 | 263.7 | 232.2 | −31.5 | −12% | | trace9 | 219.9 | 198.7 | −21.2 | −10% | --------- Co-authored-by: Claude Sonnet 4 <noreply@anthropic.com>
Extend kernel-based categorization with resolve_perf_model_class, table-driven mappings, and SGLang capture trace enrichment so graph-replay synthetic ops can get Input Dims and perf model estimates when a capture folder is provided. Co-authored-by: Cursor <cursoragent@cursor.com>
rkarhila-amd
requested review from
a team,
ajassani,
devalshahamd and
gabeweisz
as code owners
June 5, 2026 13:02
Collaborator
|
Looks good to me - thanks for the categorization change |
Contributor
|
Can you provide some examples where this enrich function improves the categorization? |
TraceLens Agent Update (6/5)
Updated comparative reference tar balls.
**NOTE: Reference CSVs have been regenerated.** Previously, if events overlap with their siblings, these events would be dropped or orphaned as part of their own subtree. This PR adds these events to the subtree if the overlap is likely due to noise. By allowing for a 1 microsecond overlap, cuda_runtime events and kernels that were previously dropped from the subtree are now added back. The events with more than 1 microsecond of overlap can be safely dropped without losing any kernels. Across 20 traces, these are the number of events that are added back to subtree with this change: | Bucket | Range | Count | % | Categories | Interpretation | |--------|-------|-------|---|------------|----------------| | Small | < 1 us | 1,597 | 80.7% | cuda_runtime, cpu_op, some python_function | Sub-microsecond. Likely profiling clock jitter / instrumentation overhead | | Moderate | 1 us – 1 ms | 325 | 16.4% | python_function only | Real overlap. Event starts inside stack[-1] | | Large | > 1 ms | 58 | 2.9% | python_function only | Massive overlap -> PyCapsule/built-in methods | Reference CSVs have been updated since the "UID_first" values have been shifted by a constant value due to there being more events in the trace tree. The raw UID values come from the trace event metadata itself. However in the case of graph capture reports, when `merge_capture_trace_into_graph` merges the capture events into the graph tree, `update_subtree_uids_and_timestamps` regenerates all the UIDs of the capture trace event based on the number of capture trace events. | Trace | main | fix/sub_microsecond_guard | Δ (s) | % | |-------|-----:|-----:|-----:|----:| | trace1 | 82.23 | 82.83 | +0.60 | +0.7% | | trace2 | 0.42 | 0.43 | +0.01 | +2.3% | | trace3 | 68.87 | 68.27 | -0.60 | -0.9% | | trace4 | 71.71 | 71.41 | -0.30 | -0.4% | | trace5 | 107.04 | 106.03 | -1.01 | -0.9% | | trace6 | 1246.46 | 1247.37 | +0.91 | +0.1% | | trace7 | 49.82 | 50.09 | +0.27 | +0.5% | | trace8 | 228.86 | 227.73 | -1.13 | -0.5% | | trace9 | 198.32 | 199.69 | +1.37 | +0.7% | --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Integrate main's registry-based categorization, triton perf models, inductor_cache_dir, and trace metadata with WIP synthetic-op kernel resolution and SGLang capture linking. Co-authored-by: Cursor <cursoragent@cursor.com>
Build kernel_details before categorization in unified perf tables, re-categorize graph-replay synthetics in orchestrator_prepare, and extend classify_kernels rules for SGLang flash-attn, KV cache, MoE topk, and triton kernels. Co-authored-by: Cursor <cursoragent@cursor.com>
Fold graph-replay and classify_kernel paths into _kernel_name_fallback with ordered rule tuples, and delegate categorize_torch_op to the registry. Co-authored-by: Cursor <cursoragent@cursor.com>
Author
|
This is my WIP branch to be merged to my own PR 677. Sorry I didn't intend to ask for review before merging to 677. @devalshahamd: PR677 describes the use case. |
Enrich capture-linked synthetics with real tensor dims, compute roofline metrics via gpu_events for off-tree ops, and use per-kernel duration instead of hipGraphLaunch envelope time.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
resolve_perf_model_classwith table-driven kernel-to-perf-model mappings for graph-replay synthetic ops (requires Input Dims when kernel-resolved).trace_sglang_capture_link.pyto enrich synthetic ops from SGLangbs_<N>_rank0.json.gzcapture traces (graph launch index alignment, decode batch-size detection).TreePerfAnalyzerandgenerate_perf_report_pytorch_inferencewhen execution_details merge is unavailable.Test plan
pytest tests/test_synthetic_op_kernel_categorization.pypytest tests/test_sglang_capture_link.pyMade with Cursor