{"slug": "performance-profiling-on-amd-gpus-part-7-ai-pytorch-r-edition", "title": "Performance Profiling on AMD GPUs - Part 7: AI/PyTorch® Edition", "summary": "AMD published Part 7 of its ROCm profiling guide series, an AI/PyTorch edition that applies the same profiling loop used on HIP and Fortran applications to a PyTorch training workload on AMD GPUs. The guide instructs readers to separate one-time warm-up costs — HIP context setup, rocBLAS/MIOpen initialization and Triton just-in-time compilation — from steady-state steps, then annotate timelines with ROCTx markers via torch.cuda.nvtx, trim them with rocprof-sys --selected-regions, summarize kernels with rocprofv3 and collect hardware counters with rocprof-compute using iteration multiplexing. It evaluates three optimizations — PyTorch-level fusion, torch.compile and custom Triton kernels — and teaches readers to distinguish a real win (fusion) from a change that only pays off under the compiler (Triton) by comparing kernel counts, timelines and throughput.", "body_md": "# Performance Profiling on AMD GPUs - Part 7: AI/PyTorch® Edition[#](#performance-profiling-on-amd-gpus-part-7-ai-pytorch-edition)\n\nThis blog, like the previous articles in the profiling guide series\n([Part 1](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/intro/README.html),\n[Part 2](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/novice/README.html),\n[Part 3](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/advanced/README.html),\n[Part 4](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/fortran_openmp/README.html),\n[Part 5](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/ai-assist-optimization/README.html),\nand [Part 6](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/att/README.html)),\napplies a systematic profiling workflow to\na GPU-accelerated application. Those posts establish the core loop on\nHIP and Fortran applications: baseline timing, kernel hotspots with `rocprofv3`, rooflines and\nhardware counters with `rocprof-compute`, and whole-system timelines with `rocprof-sys`. Here we apply\nthat loop to a PyTorch training workload, where one logical operator can lower to many\nGPU kernels and optimization often means changing how the framework dispatches work rather\nthan rewriting a kernel.\n\n**AI vs HPC note:** The profiling loop is unchanged, but one distinction from HPC drives everything\nthat follows. An HPC application is compiled ahead of execution, a PyTorch workload compiles during\nit. With `torch.compile`, kernels are traced and generated on the first iteration, and Triton kernels\nare just-in-time (JIT) compiled and auto-tuned the first time each tensor shape is seen; even in eager\nmode the first step pays one-time warm-up costs (HIP context, rocBLAS/MIOpen init, Triton JIT).\nSeparate warm-up from measurement and read timing and counters from steady-state steps only;\notherwise the profile reflects the compiler, not the model.\n\n1. **What you already know** :\n  - You run GPU-accelerated PyTorch code (training or stepping a model forward) and can match CLI knobs such as batch size, sequence length, and step count across runs.\n  - You understand that one logical operator (for example, attention or a layer norm) can map to many GPU kernels, and that fusion or custom kernels change kernel counts and kernel names.\n  - You recognize that a training step has both a *forward* and a*backward* pass, and that the\nbackward pass is often the larger of the two.\n  - Optional: you have compared runtime across hardware vendors and want AMD-specific traces and counters to explain the gap.\n2. **What this guide will teach you** :\n  - **Annotate from Python** with ROCTx markers via`torch.cuda.nvtx` (the namespace is CUDA-named,\nbut`torch.cuda` is PyTorch’s device API on a ROCm™ build as well, and ROCTx mirrors the NVTX\nrange calls, so the same annotations serve either runtime), then trim the timeline with`--selected-regions` so it maps onto the model and stays readable.\n  - **Analyze hardware counters** with`rocprof-compute` , using iteration multiplexing to collect the\nfull counter set in a single pass.\n  - **Evaluate three optimizations** (PyTorch-level fusion,`torch.compile` , and custom Triton\nkernels) and learn to tell a real win (fusion) from a change that only pays off under the\ncompiler (Triton) by comparing kernel counts, timelines, and throughput.\n\n## Profiling Basics[#](#profiling-basics)\n\nThe workflow set out below is the one we apply to optimize ROCm-based PyTorch workloads. We\nestablish a baseline, capture a whole-run system trace with `rocprof-sys-run`, annotate it with\nROCTx and narrow it with `rocprof-sys --selected-regions`, summarize kernels with `rocprofv3`, and\nanalyze bottlenecks with `rocprof-compute` using its iteration-multiplexing option. Each\noptimization is then measured against that baseline using the same set of tools. The rest of the\npost walks through this loop on the TinyOpenFold example.\n\n### The TinyOpenFold Example[#](#the-tinyopenfold-example)\n\n[OpenFold](https://github.com/aqlaboratory/openfold) is a trainable reproduction of AlphaFold 2. Its\nperformance-critical core is the **Evoformer**, a stack of blocks that each refine two tensors: a multiple\nsequence alignment (MSA) representation (B × N_seq × N_res × d_msa) and a pair representation\n(B × N_res × N_res × d_pair). Each block mixes the two, then refines the pair tensor with the\n*triangle* operations that give the pattern its name, a triangular multiplicative update and a\ntriangular self-attention, both of which update every residue pair (i, j) by routing information\nthrough a third residue k. **TinyOpenFold** shrinks this stack (fewer blocks, smaller channel\nwidths, synthetic batches) while keeping that pair/triangle pattern intact, so training runs are\nshort. The defaults are d_msa = 64, d_pair = 128, and n_blocks = 4, which together give\napproximately 2.64M parameters. For the full block-by-block breakdown, see the\nupstream [ARCHITECTURE.md](https://github.com/amd/HPCTrainingExamples/blob/main/MLExamples/TinyOpenFold/ARCHITECTURE.md).\n\nFor profiling, two properties of the pair tensor matter most: it is quadratic in sequence length, so it drives O(N_res²) activation memory, and the triangle stages that update it (multiplication and self-attention) drive O(N_res³) compute, N_res² pairs each reduced over all N_res choices of k. Those triangle stages are the dominant compute at moderate N_res and the first place to look in any trace. In addition to that intrinsic cost, the readable baseline pays for many small kernel launches and tail elementwise kernels, the launch-bound fragmentation that fusion (version 2) removes. The profiler separates the math-bound triangles from that launch-bound overhead.\n\nTinyOpenFold ships three implementations of the same model. When comparing performance between implementations, use the same training knob values.\n\n| Stage | Directory | Entry point | Role | \n|---|---|---|---|\n| **Version 1** | `version1_pytorch_baseline/` | `tiny_openfold_v1.py` | **Baseline** : clean, readable PyTorch. Many kernels per step, the reference for hotspots and operator-to-kernel mapping. | \n| **Version 2** | `version2_pytorch_fused/` | `tiny_openfold_v2.py` | **PyTorch fusion** : query, key, and value (QKV) and triangle projections fused; fewer dispatches, lower peak memory. | \n| **Version 3** | `version3_triton/` | `tiny_openfold_v3.py` | **Triton kernels** : custom fused LayerNorm plus fused scaled dot-product attention (SDPA) and QKV fusion, wrapped in autograd functions. | \n\n### Getting the Code[#](#getting-the-code)\n\nClone the canonical repository and enter the example directory:\n\n```\ngit clone https://github.com/amd/HPCTrainingExamples.git\ncd HPCTrainingExamples/MLExamples/TinyOpenFold\n```\n\nThe [project README](https://github.com/amd/HPCTrainingExamples/blob/main/MLExamples/TinyOpenFold/README.md)\ndocuments the quick start and scaling helpers, and\n[ARCHITECTURE.md](https://github.com/amd/HPCTrainingExamples/blob/main/MLExamples/TinyOpenFold/ARCHITECTURE.md)\nwalks through the parameter math block by block.\n\n### Environment Setup (ROCm + PyTorch)[#](#environment-setup-rocm-pytorch)\n\nWe validated this post on ROCm 7.13 with PyTorch 2.11 and Triton 3.6 on an AMD Instinct™ MI300X. There are two ways to obtain this stack, depending on whether the ROCm runtime comes from the Python wheels or from a module the cluster already provides. Either path supports every command in the rest of the post.\n\n- **Pip channel (used in this post).** Install PyTorch and Triton pre-packaged for ROCm 7.13 from\nAMD’s[TheRock](https://github.com/ROCm/TheRock) multi-arch pip channel. The wheels bundle a\nself-contained ROCm runtime as a dependency of`torch` , so no`/opt/rocm` and no ROCm module are\nneeded for training.\n- **System module.** Load a`rocm/7.13` module, if available, from the cluster and install the matching PyTorch,\nTriton, and other Python dependencies into a virtual environment against it.\n\nWe use the pip channel throughout for reproducibility on any gfx942 (MI300X) / gfx950 (MI355X) node:\n\n```\n# We used Python 3.12 for the blog. The `rocpd` trace tools (rocpd, rocpd2pftrace,\n# rocpd2otf2) support Python 3.10-3.13.\nmodule load python/3.12\n\ncd HPCTrainingExamples/MLExamples/TinyOpenFold\npython3 -m venv venv713\nsource venv713/bin/activate\npip3 install --upgrade pip\n\nINDEX_URL=https://rocm.nightlies.amd.com/whl-multi-arch/\n\n# PyTorch 2.11 + Triton 3.6 for ROCm 7.13 (torch pulls the ROCm libraries in automatically).\npip3 install --index-url $INDEX_URL torch torchvision torchaudio triton\n\n# ROCm SDK including the hardware profilers (rocprofv3, rocprof-compute, rocprof-sys).\n# The `profiler` extra adds them; pick the device arch(s) you run on (gfx942 = MI300X, gfx950 = MI355X).\npip3 install --index-url $INDEX_URL \"rocm[libraries,devel,profiler,device-gfx942]==7.13.*\"\n\n# otf2: Python dep for the rocpd OTF2 exporters (rocpd, rocpd2otf2); Requires Python <=3.13.\npip3 install otf2\n\n# Put the pip-installed ROCm binaries on PATH for this shell.\nexport ROCM_PATH=$(rocm-sdk path --root)\nexport PATH=$ROCM_PATH/bin:$PATH\nexport LD_LIBRARY_PATH=$ROCM_PATH/lib:$LD_LIBRARY_PATH\n```\n\nConfirm PyTorch sees the GPU:\n\n``` python\npython3 -c \"import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))\"\n```\n\nThen validate the example itself from the baseline directory:\n\n```\ncd version1_pytorch_baseline\npython3 tiny_openfold_v1.py --validate-setup\n```\n\nBecause this post focuses exclusively on single-GPU profiling, we pin a device with\n`ROCR_VISIBLE_DEVICES=0` (or `--device 0`) so that timing and counter collection remain\ndeterministic.\n\n## Step 1 — Establish a Baseline (Version 1)[#](#step-1-establish-a-baseline-version-1)\n\nAll of the commands in this section are run from inside the `version1_pytorch_baseline/` directory.\nThroughout the post we use\nthe **small** problem (64 residues, 16 MSA sequences, batch size 4) because it turns the loop over in a\nfew seconds. A **large** problem (128 residues, 32 MSA sequences, batch size 4) is used later for\ncomparison.\n\n### Application Timing[#](#application-timing)\n\nWe begin by measuring throughput with no profiler attached, which establishes the number against which every later change is compared.\n\n```\npython3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30 --device 0\n```\n\nFor the small problem, the baseline reports roughly 99 samples/sec, or about 40 ms per\nbatch[\\[1\\]](#id3). Record the corresponding number for the target hardware and compare against\nruns at the same configuration.\n\n### An Orthogonal Axis: `torch.compile`[#](#an-orthogonal-axis-torch-compile)\n\nThe cheapest optimization available, before restructuring the model, is to ask PyTorch to compile it.\nAll three versions accept `--enable-torch-compile`, which wraps the model in\n[`torch.compile`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html); its inductor\nbackend traces the model and generates fused Triton kernels automatically. Note that\n`torch.compile` is not one of the three numbered steps in this post. Because it requires no code\nchange and applies to every version, it constitutes an orthogonal axis that stacks on the baseline\n(Step 1), the manual fusion (Step 2), and the custom kernels (Step 3) alike.\nWe introduce it here and report its effect alongside every version in the final comparison.\n\n```\npython3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30 --device 0 \\\n  --enable-torch-compile\n```\n\nThis is the point at which the AI-vs-HPC difference described above becomes concrete. The first step is far slower than the rest, because that is when compilation and autotuning happen; the per-step time then drops and holds at a new steady state. In the printed per-step timing, step 0 pays the compile cost, and the later steps settle to roughly 174 samples/sec (~23 ms per batch), about a 1.75× speedup over eager for a one-line change. This behavior is the reason warm-up must be separated from measurement, since averaging the compile step into the timing would understate the gain.\n\nA more aggressive setting, `torch.compile(mode=\"max-autotune\")`, benchmarks several kernel\nimplementations and can enable CUDA graphs to cut launch overhead further; we leave that as an exercise\nfor the reader (see the\n[ROCm torch.compile blog](https://rocm.blogs.amd.com/artificial-intelligence/torch_compile/README.html#how-torch-compile-works)).\n\nBecause `torch.compile` does not reveal where time is spent, the remaining analysis runs in eager\nmode with the profilers.\n\n### Capture a Baseline System Trace (`rocprof-sys`)[#](#capture-a-baseline-system-trace-rocprof-sys)\n\nThe first look at a new workload should be a whole-run system trace which is a single timeline that shows\nPython and host threads, the HIP runtime, and GPU activity together. Capture it with `rocprof-sys-run`\nusing the `workload-trace` preset, the AI/ML preset for training and long-running GPU workloads. It\nenables the ROCm domains (HIP, kernel, marker, and memory), sizes the Perfetto buffer for a long run,\nand routes through rocPD to capture child-thread activity. It emits a Perfetto `.proto` timeline\nalongside a rocPD `.db`:\n\n```\nROCR_VISIBLE_DEVICES=0 rocprof-sys-run --preset=workload-trace -o rsys_v1_workload -- \\\n  python3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10 --device 0\n```\n\nThe resulting `rsys_v1_workload/<timestamp>/perfetto-trace-*.proto` can then be opened at\n[https://ui.perfetto.dev/](https://ui.perfetto.dev/). (ROCm\n[Optiq](https://rocm.docs.amd.com/projects/roc-optiq/en/latest/what-is-optiq.html) can now\nvisualize the accompanying rocPD `.db` directly.) The trace shows separate tracks for HIP APIs,\nkernel, and memory, with per-launch kernel detail. Figure 1a below gives the whole-run view, and\nFigure 1b zooms into a single training step. Together they show why a full-run trace of a training\nloop is dense and hard to navigate, with dozens of steps, each with many short kernels separated by\nlaunch gaps, and no labels tying a region of the timeline back to the model. The next two steps fix\nthat, first by annotating phases, then by trimming the capture.\n\nFigure 1a: Whole-run version 1 system trace using `rocprof-sys-run`.\n\nFigure 1b: Zoomed single step of the same trace shows many short kernels with launch gaps between them without phase labels.\n\n### Annotate Workload Phases (ROCTx)[#](#annotate-workload-phases-roctx)\n\nTo make the trace readable, mark the phases of a training step with ROCTx ranges. From Python,\nROCm records `torch.cuda.nvtx` ranges as ROCTx regions. Wrap the forward and backward passes\n(and any stage of interest, such as the attention or triangle blocks) with start and stop markers:\n\n``` python\nimport torch\n\nfor step in range(num_steps):\n    fwd = torch.cuda.nvtx.range_start(\"forward_pass\")\n    loss = model(batch)\n    torch.cuda.nvtx.range_end(fwd)\n\n    bwd = torch.cuda.nvtx.range_start(\"backward_pass\")\n    loss.backward()\n    torch.cuda.nvtx.range_end(bwd)\n    optimizer.step(); optimizer.zero_grad()\n```\n\nRe-run the same `rocprof-sys-run` capture on the annotated `tiny_openfold_v1_roctx.py`:\n\n```\nrocprof-sys-run --rocm=kernel,marker -T -o rsys_v1_annot -- \\\n  python3 tiny_openfold_v1_roctx.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10 --device 0\n```\n\nIn the resulting Perfetto trace (`rsys_v1_annot/<timestamp>/perfetto-trace-*.proto`), the named\nranges appear in the timeline with the GPU kernels of each phase underneath, which turns the dense\nbaseline trace into a per-phase breakdown.\n\nThe key takeaway from the annotated ranges, visible as the two labeled bands in Figure 2 below, is the split between the forward and backward passes. For TinyOpenFold the backward pass is roughly 55–60% of the step, so any optimization that touches both passes (like the fusion in version 2) compounds. The trace also localizes where that cost is incurred, since the triangle and transition general matrix multiplications (GEMMs) dominate, which bounds what later optimizations can achieve (see version 3).\n\nFigure 2: ROCTx-annotated version 1 trace (`rocprof-sys-run`) with `forward_pass`\nand `backward_pass` bands.\n\n### Narrow the Capture (`--selected-regions`)[#](#narrow-the-capture-selected-regions)\n\nWith phases labeled, the trace can be restricted to the region of interest instead of recording every\nstep. `rocprof-sys-run` keeps only the ROCTx ranges named through the `--selected-regions` flag,\nturning the dense whole-run timeline into a focused view of, say, the forward and backward passes:\n\n```\nrocprof-sys-run --rocm=kernel,marker --selected-regions=forward_pass,backward_pass -T -o rsys_v1_sel -- \\\n  python3 tiny_openfold_v1_roctx.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10 --device 0\n```\n\nThe resulting trace contains only the `forward_pass` and `backward_pass` regions, and the region names must\nmatch the start and stop ROCTx markers added above. The same regions can instead be set through the\n`ROCPROFSYS_SELECTED_REGIONS=forward_pass,backward_pass` environment variable. Figure 3 below shows\nthe result, the same capture as Figure 2 with the data-loading and optimizer steps dropped, so the\ntwo phases of interest fill the timeline.\n\nFigure 3: The same trace narrowed to `forward_pass` and `backward_pass` via\n`--selected-regions`.\n\n### Device Summary (`rocprofv3`)[#](#device-summary-rocprofv3)\n\nThe same region narrowing is available in `rocprofv3` for the kernel summary, so the hotspot\npercentages reflect the phases of interest rather than the whole run. We start with the full-run\nsummary and then narrow it.\n\n```\nROCR_VISIBLE_DEVICES=0 rocprofv3 --kernel-trace --stats -S -T -d out_v1 -o tinyfold -- \\\n  python3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10\n```\n\nOn ROCm 7.13 the `--stats` data lands in the `rocpd` database\n(`out_v1/tinyfold_results.db`) rather than a standalone CSV. Export the aggregated table from that\ndatabase with `rocpd2summary`:\n\n```\nrocpd2summary -i out_v1/tinyfold_results.db -f csv -d out_v1/csv -o tinyfold_summary\n```\n\nThis writes `out_v1/csv/tinyfold_summary_kernels_summary.csv` with one row per kernel\n(name, calls, total/avg duration, and percent of GPU time), the same numbers `-S` prints. Sorting by\nGPU time gives the baseline hot spots for the small problem:\n\n| Kernel | Calls | Total (ms) | % of GPU time | \n|---|---|---|---|\n| `vectorized_elementwise_kernel` | 19,468 | 84.4 | 26.5% | \n| `elementwise_kernel_manual_unroll` | 4,575 | 53.8 | 16.9% | \n| rocBLAS GEMM (`Cijk_…` MT32x16x128) | 1,800 | 31.1 | 9.7% | \n| rocBLAS GEMM (`Cijk_…` MT128x64x32) | 1,275 | 17.3 | 5.4% | \n| `layer_norm_grad_input_kernel` | 780 | 14.6 | 4.6% | \n\nThe summary exhibits a characteristic shape. The top slots are occupied by the elementwise and tail\nkernels, which represent the launch-bound fragmentation that results from nothing being fused (note\nthe 19,468 and 4,575 call counts over the 15 executed steps, 10 measured plus 5 warmup). Beneath\nthem sit the rocBLAS `Cijk_…` GEMMs, the compute-bound triangle and transition matmuls. It is\nprecisely that fragmentation which version 2 targets.\n\n`rocprofv3 --selected-regions` narrows the summary to the forward and backward passes, but it keys on\n`roctxProfilerResume()` and `roctxProfilerPause()` calls rather than region names. Collection\ntherefore stays off until the first resume call is reached. The annotated variant\n`tiny_openfold_v1_roctx_selected.py` wraps the same two\nspans with those calls (the `roctx` module ships with the SDK that should be prepended to `PYTHONPATH`):\n\n``` python\nimport roctx\ntid = roctx.getThreadId()\n\nfor step in range(num_steps):\n    fwd = torch.cuda.nvtx.range_start(\"forward_pass\")\n    roctx.profilerResume(tid)      # collection ON\n    loss = model(batch)\n    roctx.profilerPause(tid)       # collection OFF\n    torch.cuda.nvtx.range_end(fwd)\n\n    bwd = torch.cuda.nvtx.range_start(\"backward_pass\")\n    roctx.profilerResume(tid)\n    loss.backward()\n    roctx.profilerPause(tid)\n    torch.cuda.nvtx.range_end(bwd)\n    optimizer.step(); optimizer.zero_grad()\n```\n\nBecause only those two phases are collected, a few steps suffice:\n\n```\nrocprofv3 --selected-regions --kernel-trace --stats -S -T -d out_v1_sel -o selreg -- \\\n  python3 tiny_openfold_v1_roctx_selected.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 3 --device 0\n```\n\nNarrowing drops the data-loading, optimizer, and warm-up kernels (about 48% of dispatches), so the\n`--stats` percentages now describe the **forward/backward mix alone**, not the whole GPU run.\nHowever, the whole-run table above remains the reference for the overall hotspot distribution.\nSorting the narrowed summary by GPU time within those phases (`Cijk_*` GEMM tiles grouped into one\nbucket) reorders the picture:\n\n| Kernel | Calls | Total (µs) | % of GPU time (fwd+bwd) | \n|---|---|---|---|\n| rocBLAS GEMM (`Cijk_*` , all tiles) | 1,683 | 24,275 | 40.1% | \n| `vectorized_elementwise_kernel` | 2,895 | 15,262 | 25.2% | \n| `elementwise_kernel_manual_unroll` | 915 | 10,876 | 18.0% | \n| `layer_norm_grad_input_kernel` | 156 | 2,961 | 4.9% | \n| `vectorized_layer_norm_kernel` | 156 | 2,440 | 4.0% | \n\nWith the loader and optimizer kernels excluded, the triangle/transition GEMMs rise to the top at 40.1% of forward and backward time instead of the elementwise tail. This confirms that the GEMMs are the dominant compute inside the phases that matter, while the elementwise kernels remain the launch-bound overhead that fusion targets.\n\n### Identify the Bottleneck[#](#identify-the-bottleneck)\n\nTwo distinct signals in the profile point toward the next optimization to apply. The gaps between the many short kernels expose host-side launch overhead, and the HBM traffic from hardware counters quantifies memory movement, both of which fusion (version 2) reduces. We therefore begin by examining the hardware counters for the hottest kernel.\n\n## Analyze Hardware Usage (`rocprof-compute`)[#](#analyze-hardware-usage-rocprof-compute)\n\nSorting kernels by time shows where the time goes, whereas the hardware counters show how a kernel\nuses the GPU. The next step is to collect the full counter set for a top kernel identified in the\n`rocprofv3` summary. When done naively, this re-runs the application once per counter group\n(dozens of passes); `--iteration-multiplexing` avoids that by cycling the counters over\nthe kernel’s many calls within a single run. We collect the full counters for the top elementwise\nkernel and then analyze specific metric blocks.\n\n```\nrocprof-compute profile -n tinyfold_v1_vecelem_mux --device 0 \\\n  --iteration-multiplexing kernel_launch_params -k vectorized_elementwise_kernel -- \\\n  python3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 3 --device 0\nrocprof-compute analyze -p workloads/tinyfold_v1_vecelem_mux/MI300X_A1 -b 2.1.14 2.1.16 7.1 11.2 17.3 --view table\n```\n\nThe `-n` argument names the workload and `-k` restricts collection to the matching kernel; results land\nunder `workloads/tinyfold_v1_vecelem_mux/<device-arch>/` (for example `MI300X_A1`). The `-b`\nblocks select the metrics of interest: wavefront occupancy (2.1.16) and VALU active threads (2.1.14);\nwavefront launch stats such as total wavefronts, VGPR, SGPR, and LDS per workgroup (7.1);\ncompute-throughput utilization including VALU utilization and IPC (11.2); and L2 cache bandwidth and\nhit rate (17.3). We revisit these numbers for version 2 to see what fusion\ndoes to this kernel. The complete set of block IDs can be printed with\n`rocprof-compute profile --list-available-metrics`, as described in the\n[block filtering documentation](https://rocm.docs.amd.com/projects/rocprofiler-compute/en/latest/how-to/profile/mode.html#analysis-report-block-filtering).\n\n## Step 2 — Optimize with PyTorch Kernel Fusion (Version 2)[#](#step-2-optimize-with-pytorch-kernel-fusion-version-2)\n\nThe baseline established two facts, namely that many small kernels fragment the step and that the\nbackward pass dominates. Version 2 addresses the fragmentation without leaving PyTorch. All of the\ncommands in this section are run from inside the `version2_pytorch_fused/` directory.\n\n### What Version 2 Changed and Why[#](#what-version-2-changed-and-why)\n\nVersion 2 keeps the Evoformer math identical but restructures three hot spots so that the framework dispatches fewer, larger kernels:\n\n- **MSA QKV fusion** — the separate query, key, and value projections become a single combined\nprojection (3 kernels → 1). One larger GEMM does the work of three small ones, cutting launch\noverhead and giving the GPU a more efficient matrix shape.\n- **Flash-Attention-style attention** — this path avoids materializing the full attention matrix in\nHBM, which removes a round trip through memory for the largest intermediate tensors.\n- **Triangle gate/projection fusion** — this merge combines the gate and projection of each triangle\nstage (4 kernels → 2), again replacing several small launches with fewer, larger ones.\n\nThis restructuring also lowers the activation footprint of the step. Because the fused attention never materializes the full N×N attention matrix in HBM, version 2’s peak memory drops about 14% below version 1, on top of the throughput gain shown later in this section.\n\n### Verify with `rocprofv3` and `rocprof-sys`[#](#verify-with-rocprofv3-and-rocprof-sys)\n\nWe first run the same baseline timing as before.\n\n```\nROCR_VISIBLE_DEVICES=0 python3 tiny_openfold_v2.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30\n```\n\nFor the small problem, version 2 reports about 229 samples/sec (17.4 ms per batch), a\n+131% improvement over version 1 (roughly 2.3×), and at lower peak memory. This is the\npeak of these versions, and it comes entirely from dispatching fewer, larger kernels. Because\nversion 2 is already hand-fused, adding `--enable-torch-compile` moves it only to about 239\nsamples/sec and the compiler has little launch overhead left to remove:\n\n| Small problem | eager | + `torch.compile` | \n|---|---|---|\n| Version 1 baseline | 99 | 174 | \n| **Version 2 fused** | **229** (+131%) | **239** | \n\nNow re-run `rocprofv3` with the same discipline as version 1 and compare:\n\n```\nROCR_VISIBLE_DEVICES=0 rocprofv3 --kernel-trace --stats -S -T -d out_v2 -o tinyfold -- \\\n  python3 tiny_openfold_v2.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10\n```\n\nTo see the QKV fusion rather than infer it from counts, capture a `rocprof-sys-run` trace of each\nversion with the `msa_qkv` ROCTx region labeled, then zoom to that band in Perfetto. Both regions are\nalready annotated in the `*_roctx.py` variants:\n\n```\nROCR_VISIBLE_DEVICES=0 rocprof-sys-run --rocm=kernel,marker -T -o rsys_v1_qkv -- \\\n  python3 ../version1_pytorch_baseline/tiny_openfold_v1_roctx.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 5 --device 0\n\nROCR_VISIBLE_DEVICES=0 rocprof-sys-run --rocm=kernel,marker -T -o rsys_v2_qkv -- \\\n  python3 tiny_openfold_v2_roctx.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 5 --enable-all-fusion\n```\n\nOpen the resulting `rsys_v{1,2}_qkv/<timestamp>/perfetto-trace-*.proto` files and navigate to the\n`msa_qkv` region in each. The contrast is the fusion itself, as Figure 4a below shows the three\nseparate query, key, and value GEMM launches in version 1, and Figure 4b shows the same region in\nversion 2 after they collapse into a single wider GEMM.\n\nFigure 4a: Version 1 `msa_qkv` region—three separate Q/K/V GEMM launches.\n\nFigure 4b: Version 2 `msa_qkv` region—the three projections fused into one GEMM.\n\nThe trace confirms the kernel count dropped and time moved into a smaller set of larger kernels. On MI300X the total launches per step fall from roughly 3,700 to about 1,700, a ~54% reduction at matched accuracy (see the table below). Most of the reduction comes from the fused attention and triangle projections, which collapse many small launches into a few larger ones; the untouched GEMMs and elementwise tail account for the rest.\n\nThe per-step totals below are counted from the `rocpd` database (`rocpd summary` over each version’s\n`_results.db`, divided by the step count). Version 2’s fusion nearly halves the launch count,\ncorroborating its throughput win:\n\n| Version | Kernel launches / step | vs V1 | \n|---|---|---|\n| Version 1 | 3,702 | — | \n| Version 2 | 1,697 | −54% | \n\nSorting the version 2 summary by GPU time shows the shift directly:\n\n| Kernel | Calls | Total (ms) | % of GPU time | \n|---|---|---|---|\n| `vectorized_elementwise_kernel` | 5,820 | 12.9 | 7.1% | \n| `bwd_kernel_fuse` (fused Flash-Attention backward) | 240 | 12.8 | 7.0% | \n| rocBLAS GEMM (`Cijk_…` MT32x16x128) | 600 | 10.8 | 5.9% | \n| rocBLAS GEMM (`Cijk_…` MT128x64x32) | 675 | 9.2 | 5.1% | \n| rocBLAS GEMM (`Cijk_…` MT64x32x64) | 240 | 8.9 | 4.9% | \n\nCompared with version 1, the elementwise tail has collapsed. The `vectorized_elementwise_kernel`\nentry drops from 19,468 calls/84.4 ms (26.5%) to 5,820 calls/12.9 ms (7.1%), so time is now spread\nmore evenly and the profile is GEMM-dominated (three `Cijk_…` matmuls in the top five). The\n`bwd_kernel_fuse` entry represents the fused Flash-Attention backward kernel, which has relatively\nfew launches (240) but a high computational cost per call.\n\nWhat fusion does not touch are the triangle and transition GEMMs, which are already efficient rocBLAS calls. After version 2 those GEMMs are what remain, and that sets up the more nuanced lesson of version 3.\n\n### Confirm Kernel Fusion Saves Memory[#](#confirm-kernel-fusion-saves-memory)\n\nKernel fusion also lowers the activation footprint of each step. The metric that shows it is peak device\nallocation (`torch.cuda.max_memory_allocated`), which captures the intra-step activation high-water\nmark that fusion changes. The `--enable-memory-profiling` flag reports this value in the training\nsummary. We run version 1 and version 2 at matched configuration.\n\n```\nROCR_VISIBLE_DEVICES=0 python3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30 --device 0 --enable-memory-profiling\nROCR_VISIBLE_DEVICES=0 python3 tiny_openfold_v2.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30 --enable-all-fusion --enable-memory-profiling\n```\n\n| Run | Speed (samples/sec) | Peak memory (`max_memory_allocated` ) | vs V1 | \n|---|---|---|---|\n| Version 1 baseline | 97 | 2,266.9 MB | — | \n| Version 2 (all fusion) | 236 | 1,941.2 MB | −14.4% | \n| Version 2 (`--disable-all-fusion` , control) | 102 | 2,266.9 MB | 0.0% | \n\nThese speeds are a separate run from the timing numbers above.\n\nKernel fusion cuts peak memory 14.4% (−326 MB), because the Flash-Attention-style attention never\nmaterializes the full N×N attention matrix and so lowers the activation high-water mark. The\n`--disable-all-fusion` control lands byte-for-byte on version 1’s peak, which proves the saving comes\nfrom the fusion itself and not from version 2 being a different script.\n\n### Confirm the Effect at the Counter Level[#](#confirm-the-effect-at-the-counter-level)\n\nThe kernel-count table showed fusion cut launches; the hardware counters show what it did to the one elementwise kernel that survives in both versions. Re-run the same full-counter capture from Step 1 on version 2 and compare the two:\n\n```\nrocprof-compute profile -n tinyfold_v2_vecelem_mux --device 0 \\\n  --iteration-multiplexing kernel_launch_params -k vectorized_elementwise_kernel -- \\\n  python3 tiny_openfold_v2.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 3 --enable-all-fusion\nrocprof-compute analyze -p workloads/tinyfold_v2_vecelem_mux/MI300X_A1 -b 2.1.14 2.1.16 7.1 11.2 17.3 --view table\n```\n\n| Metric (`vectorized_elementwise_kernel` ) | Version 1 | Version 2 | Change | \n|---|---|---|---|\n| Total wavefronts (max) | 16,384 | 8,192 | ½ the waves | \n| L2 cache bandwidth achieved (avg, Gb/s) | 1,496 | 3,152 | ~2.1× higher | \n| L2 cache hit rate (%) | 44.4 | 51.5 | +7 points | \n| VALU utilization (%) | 0.43 | 0.18 | lower per call | \n| Wavefront occupancy (avg) | 203.6 | 108.9 | lower | \n| IPC issued | 1.02 | 1.00 | ~same | \n\nKernel fusion halves the wavefronts launched for the elementwise tail and drives about 2.1× higher L2 bandwidth utilization with a +7-point cache hit rate on the kernel that remains, the same work done in fewer, better-fed dispatches that hit cache more often. Taken together, these counters constitute the evidence behind version 2’s throughput gain. (Iteration-multiplexing metrics come from partially available counter data and are best read as relative comparisons, not absolute values.)\n\nIt should be noted that this is not a strict one-to-one kernel comparison. Kernel fusion absorbs the FP32 elementwise work (bias-add, gating, activations) into the fused attention/GEMM kernels, so version 2’s standalone elementwise tail is a different, smaller set of kernels than version 1’s, and the counters compare the tail as a whole, not an identical kernel.\n\n## Step 3 — Optimize with Triton Kernels (Version 3)[#](#step-3-optimize-with-triton-kernels-version-3)\n\nVersion 2 removed the launch overhead and reached the throughput peak. Version 3 adds custom Triton\nkernels, and in doing so demonstrates a central principle, namely that a hand-written kernel helps\nonly when it targets where the time is actually spent. All of the commands in this section are run\nfrom inside the `version3_triton/` directory.\n\n### What Version 3 Changed and Why[#](#what-version-3-changed-and-why)\n\n- **Fused LayerNorm** — a single Triton kernel computes mean, variance, normalization, and the\naffine scale in one pass over the data, instead of several PyTorch ops each reading and writing HBM.\n- **Fused attention (SDPA) + QKV fusion** — attention uses PyTorch’s fused scaled-dot-product\nattention (a real flash-attention kernel) with the query/key/value projections combined, each\nwrapped in a`torch.autograd.Function` so gradients flow correctly.\n\nBoth of these are memory-bound operations, and that fact is the crux of the matter. The model spends most of its time in the triangle and transition GEMMs, identical rocBLAS calls in all three versions that version 3 does not touch. Customizing the memory-bound LayerNorm and attention optimizes a small slice of the total runtime.\n\n### Warm Up Before Measuring[#](#warm-up-before-measuring)\n\nTriton compiles and autotunes its kernels on first use. A cold first step is far slower than steady\nstate and would unfairly penalize version 3. We therefore run the project’s correctness check\n(`python3 test_correctness.py`) and then take enough warmup steps to populate the Triton cache\nbefore recording any timing or counters.\n\n```\nROCR_VISIBLE_DEVICES=0 python3 tiny_openfold_v3.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30\n```\n\nThe measured result constitutes the central lesson of this section. For the small problem, eager\nversion 3 reports about 91 samples/sec (44.1 ms per batch), roughly 0.9× version 1, that is,\nslightly slower than the readable baseline, and well below version 2’s 229 samples/sec. Two effects\noffset one another here. The QKV fusion and fused SDPA reduce launch overhead, but wrapping each custom Triton kernel\nin a `torch.autograd.Function` adds approximately 0.1 ms per call, and across the ~36 LayerNorm\ncalls per step this Python overhead roughly cancels the savings. Optimizing a memory-bound slice of a\nGEMM-bound model cannot, on its own, move end-to-end time appreciably.\n\nThe custom kernels do pay off once the compiler removes that per-call overhead.\nWhen `--enable-torch-compile` is added, inductor fuses the graph around the custom kernels, though\nit still calls the Triton LayerNorm rather than replacing it.\n\n```\nROCR_VISIBLE_DEVICES=0 python3 tiny_openfold_v3.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 30 \\\n  --enable-torch-compile\n```\n\nCompiled version 3 reaches about 196 samples/sec (~2.2× over eager V3), which finally edges past compiled version 1 (~174 samples/sec) at the small problem, but still sits below version 2, and the edge does not hold at the large problem (see the final comparison table). The lesson is to always measure against the compiled baseline, not just eager, and not to assume a hand-written kernel wins until measurement says so.\n\n### Verify with `rocprofv3` and `rocprof-compute`[#](#verify-with-rocprofv3-and-rocprof-compute)\n\n```\nROCR_VISIBLE_DEVICES=0 rocprofv3 --kernel-trace --stats -S -T -d out_v3 -o tinyfold -- \\\n  python3 tiny_openfold_v3.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10\n```\n\nExpect new kernel-name patterns, since the Triton and SDPA kernels will not match the `aten::`-derived\nnames from versions 1 and 2. The summary can be filtered with `-T` and substring matching. Sorting\nthe version 3 summary by GPU time gives:\n\n| Kernel | Calls | Total (ms) | % of GPU time | \n|---|---|---|---|\n| `vectorized_elementwise_kernel` | 14,885 | 72.5 | 24.2% | \n| `elementwise_kernel_manual_unroll` | 6,435 | 67.6 | 22.5% | \n| rocBLAS GEMM (`Cijk_…` MT32x16x128) | 1,080 | 19.8 | 6.6% | \n| `reduce_kernel` | 1,095 | 16.1 | 5.4% | \n| rocBLAS GEMM (`Cijk_…` MT128x64x32) | 915 | 12.4 | 4.1% | \n| … |  |  |  | \n| `layernorm_dx_kernel` (Triton bwd, rank 10) | 780 | 9.5 | 3.2% | \n| `layernorm_kernel` (Triton fwd, rank 11) | 780 | 9.0 | 3.0% | \n\nThe two custom Triton LayerNorm kernels are version-3-specific contributions that rank below the\ntop five kernels. Together, they account for only approximately 6% of GPU time. This is the point\nat which the profiler explains the throughput\nresult. Unlike version 2, whose fusion collapsed the elementwise tail\n(`vectorized_elementwise_kernel` fell from 19,468 calls/84.4 ms to 5,820 calls/12.9 ms), eager\nversion 3 still looks like version 1, in that the same elementwise kernels top the list, because\nversion 3 only touches the memory-bound LayerNorm and attention, not the elementwise and GEMM bulk.\nThe custom Triton kernels appear *lower* at roughly 9 ms each, and the triangle and transition\nGEMMs still dominate, unchanged from versions 1 and 2. That is why\ncustomizing those kernels barely moves the total, and it shows in the launch count. Eager version 3\nleaves the elementwise and GEMM bulk untouched, so its per-step total stays near version 1:\n\n| Version | Kernel launches / step | vs V1 | \n|---|---|---|\n| Version 1 | 3,702 | — | \n| Version 2 | 1,697 | −54% | \n| Version 3 | 3,285 | −11% | \n\nTo judge the custom kernel on its own terms, compare it head-to-head with the native kernel it\nreplaced using the same full-counter workflow. The LayerNorm input-gradient exists in both versions\n(the Triton `layernorm_dx_kernel` in version 3 and the native `layer_norm_grad_input_kernel` in\nversion 1), so profiling each with `rocprof-compute` at identical grid and wave count gives a fair\nper-dispatch comparison:\n\n```\n# Version 3: the Triton input-gradient kernel\nrocprof-compute profile -n tinyfold_v3_roof_dx --device 0 --roof-only -k layernorm_dx_kernel -- \\\n  python3 tiny_openfold_v3.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10\n\n# Version 1: the native input-gradient kernel (run from version1_pytorch_baseline/)\nrocprof-compute profile -n tinyfold_v1_lngrad --device 0 --roof-only -k layer_norm_grad_input_kernel -- \\\n  python3 tiny_openfold_v1.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10 --device 0\n```\n\n| Metric (mean / dispatch) | V3 Triton `layernorm_dx_kernel` | V1 native `layer_norm_grad_input_kernel` | ratio (V3/V1) | \n|---|---|---|---|\n| Vector general-purpose registers (VGPR) / wave | 12 | 24 | 0.50× | \n| Local data share (LDS) / workgroup (bytes) | 8 | 16 | 0.50× | \n| Active GPU cycles (GRBM) | 264,138 | 404,504 | **0.65×** | \n| VALU F32 instructions | 1,855,173 | 1,354,675 | 1.37× | \n| LDS-active cycles | 725,937 | 3,498,614 | **0.21×** | \n| Total cache accesses | 700,731 | 2,681,938 | **0.26×** | \n| HBM read + write requests | 108,647 / 95,783 | 109,082 / 95,783 | ~1.00× | \n\nThe Triton kernel is genuinely more efficient per dispatch, with about 35% fewer active GPU cycles and roughly half the VGPR and LDS per wave. Its standout win is ~4–5× fewer LDS-active cycles and cache accesses (0.21×/0.26×), because the single-pass Triton LayerNorm keeps intermediates in registers instead of round-tripping through LDS and cache like the native multi-op path. It executes ~1.37× more VALU F32 instructions (the fused recompute), but that is cheap next to the cache and LDS savings. HBM traffic is essentially identical (~1.00×), as expected for a memory-bound op that must move the same bytes either way.\n\n[!Note] Two caveats on this table. First, these are two different kernels performing the same operation, not one kernel measured twice, since version 1 has no Triton kernel and version 3 replaced the native one, so a cross-implementation comparison is the only option. Second, it covers only the input-gradient (∂/∂x); PyTorch splits the full LayerNorm backward across additional kernels (the weight and bias gradients), which are not included here.\n\nThe custom kernel wins every counter that matters for a memory-bound op, yet it is only ~3% of GPU time, so those wins cannot move end-to-end throughput. The hardware counters explain why the big lever was fusion (version 2), not custom kernels.\n\n## Iterate and Compare Versions Fairly[#](#iterate-and-compare-versions-fairly)\n\nA profiling comparison is only meaningful if the two runs are otherwise identical. Before comparing any two versions:\n\n- Fix batch size, sequence length, MSA depth ( `--num-seqs` ), block count, automatic mixed precision\n(AMP) mode, and step count across all runs.\n- Pin a single device ( `ROCR_VISIBLE_DEVICES=0` or`--device 0` ) so timing is deterministic.\n- For version 3, separate warmup from measurement, and never compare a Triton cold start against a steady version 1 step.\n- Give each `rocprof-compute` run a distinct workload name (`tinyfold_v1_mux` ,`tinyfold_v2_mux` ,`tinyfold_v3_mux` ) so compare mode does not mix stale directories.\n\nWith matched configurations, the story is consistent across problem sizes. Version 2 fusion is the\npeak, eager version 3 sits near or just below version 1, and `torch.compile` gives roughly\n1.4–2.2× on version 1 and version 3 (adding little to the already-fused version 2). All numbers are\nsamples/sec on MI300X / ROCm 7.13[\\[1\\]](#id3):\n\n| Problem |  | Version 1 | Version 2 | Version 3 | \n|---|---|---|---|---|\n| **Small** (64 res, 16 MSA, batch 4) | eager | 99 | **229** (+131%) | 91 (−8%) | \n|  | + `torch.compile` | 174 | 239 | 196 | \n| **Large** (128 res, 32 MSA, batch 4) | eager | 63 | **106** (+68%) | 64 (+2%) | \n|  | + `torch.compile` | 95 | 106 | 91 | \n\nTwo profiling lessons follow directly from this table. First, the largest lever was framework-level fusion rather than the hand-written kernels, because fusion reduced launch overhead across the entire step, whereas the custom kernels addressed only a memory-bound slice of a GEMM-bound model. Second, one should always compare against the compiled baseline. Eager version 3 appears to regress, yet compiled version 3 edges past compiled version 1 at the small problem. The operative regime (launch-bound or compute-bound) determines how much any optimization can yield, and only measurement can establish which regime applies.\n\nTo produce a side-by-side `rocprof-compute` report between two versions, collect a full counter set for\neach (as shown above) and pass both workloads to `analyze`, as in the\n[advanced article](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/advanced/README.html).\nIf a hot kernel was renamed by fusion or Triton, match by dispatch order or a broader `-k` substring\nrather than an exact name.\n\n## Compare the Runtime Traces Across Versions[#](#compare-the-runtime-traces-across-versions)\n\nThe baseline system trace showed a fragmented step; the payoff is seeing that fragmentation shrink\nafter fusion. We re-run the same `rocprof-sys-run` capture with the `--selected-regions` narrowing\nfrom the baseline section on version 2 (and version 3), keeping the configuration identical, and\nreuse the annotated `*_roctx.py` variants so that the `forward_pass` and `backward_pass` regions\nare directly comparable.\n\n```\nROCR_VISIBLE_DEVICES=0 rocprof-sys-run --rocm=kernel,marker --selected-regions=forward_pass,backward_pass -T -o rsys_v2_sel -- \\\n  python3 tiny_openfold_v2_roctx.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10\n\nROCR_VISIBLE_DEVICES=0 rocprof-sys-run --rocm=kernel,marker --selected-regions=forward_pass,backward_pass -T -o rsys_v3_sel -- \\\n  python3 tiny_openfold_v3_roctx.py --batch-size 4 --seq-len 64 --num-seqs 16 --num-blocks 4 --num-steps 10\n```\n\nEach resulting Perfetto trace (`rsys_v{2,3}_sel/<timestamp>/perfetto-trace-*.proto`) can be opened at\n[https://ui.perfetto.dev/](https://ui.perfetto.dev/) alongside the version 1 capture from the baseline\nsection. Figures 5a, 5b, and 5c below show the same `forward_pass` and `backward_pass` regions\nacross all three versions at matched configuration, making the structural difference apparent from\nthe timeline shape. Figure 5a is the fragmented version 1 baseline, Figure 5b\nis the fused version 2, and Figure 5c is eager version 3.\n\nFigure 5a: Version 1 forward/backward passes—many short kernels with launch gaps.\n\nFigure 5b: Version 2—fewer, longer kernels and a visibly shorter region.\n\nFigure 5c: Eager version 3—fragmentation returns, resembling version 1 more than version 2.\n\nThe comparison renders the article’s central result in visual form. Version 2 fusion is the only change that shrinks both the kernel count and the wall-clock width of the region. Version 1 and eager version 3 look alike, dense with short launches and gaps, while version 2 stands out with fewer, longer kernels and a narrower span, matching the throughput numbers and the kernel-count deltas measured earlier.\n\nBecause it is a system trace, it also indicates when the GPU is the wrong place to look. A timeline\nthat shows the GPU idle between bursts of kernels points to a host bottleneck (Python overhead, the\ndata pipeline, or synchronization), and correlating the Python/process threads, HIP runtime, and GPU\nrows in one view separates a host-bound stall from a device-bound stretch. The same trace is the\nright starting point for extending TinyOpenFold to data-parallel multi-GPU runs, where collectives\nmust be distinguished from compute; the\n[advanced article](https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/advanced/README.html)\nwalks through merging per-rank traces.\n\n## Summary[#](#summary)\n\nThis post applied the ROCm profiling toolchain to a PyTorch training workload.\nROCm Systems Profiler (`rocprof-sys`) captured a whole-run system trace with its `workload-trace` preset, and\n`rocprofv3` produced a kernel summary, together exposing the bottleneck, a step broken down into\nthousands of small kernel launches on top of GEMM-bound triangle compute. Both tools support `--selected-regions` to narrow collection to\nthe phases of interest (forward and backward); restricting the summary this way removed the loader and\noptimizer noise and showed the triangle and transition GEMMs as the true hotspot inside those phases.\nROCm Compute Profiler (`rocprof-compute`) then quantified each optimization using hardware counters,\nits iteration-multiplexing option collecting the full counter set in a single pass and substantially\nreducing both the profiling time and the on-disk footprint.\n\nThe measurements obtained from these tools guided the work directly. PyTorch fusion (version 2) roughly halved the launches per\nstep for a ~2.3× gain while also cutting peak memory ~14%, confirmed by fewer wavefronts and higher cache bandwidth on\nthe surviving kernels, and proved to be the throughput peak. Custom Triton kernels (version 3) targeted\nmemory-bound operations that were only a small share of a GEMM-dominated step, so eager version 3\nmatched the baseline and overtook it only under `torch.compile` at the small problem, an orthogonal\noptimization that added a ~1.4–2.2× gain to version 1 and version 3 (and little to the already-fused\nversion 2).\n\nThe lesson is to profile before optimizing, and measure against the compiled baseline. Optimization effort is worth spending only on the hotspot regions the profiler timeline exposes.\n\n## Additional Resources[#](#additional-resources)\n\n- Performance Profiling on AMD GPUs Blog Posts:\n- TinyOpenFold example and tutorial:\n- OpenFold / AlphaFold 2:\n- Perfetto UI:\n- ROCm Optiq: \n  - [What is ROCm Optiq?](https://rocm.docs.amd.com/projects/roc-optiq/en/latest/what-is-optiq.html)\n  - Open source at [Optiq github repo](https://github.com/ROCm/roc-optiq)\n- `rocprofiler-sdk` :\n- `rocprof-compute` :\n- `rocprof-sys` :\n- PyTorch for ROCm:\n\nIf you have any questions or comments, please reach out to us on GitHub\n[Discussions](https://github.com/ROCm/rocm-blogs/discussions).\n\n## Disclaimers[#](#disclaimers)\n\nThe information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information. However, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes. THIS INFORMATION IS PROVIDED “AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.\n\nThird-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF SUCH THIRD-PARTY CONTENT.\n\nAMD, the AMD Arrow logo, AMD Instinct, ROCm, and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies.\n\nPyTorch, the PyTorch logo and any related marks are trademarks of The Linux Foundation.\n\nPython is a trademark of the Python Software Foundation.\n\n© 2026 Advanced Micro Devices, Inc. All rights reserved.", "url": "https://wpnews.pro/news/performance-profiling-on-amd-gpus-part-7-ai-pytorch-r-edition", "canonical_source": "https://rocm.blogs.amd.com/software-tools-optimization/profiling-guide/ai_tinyopenfold/README.html", "published_at": "2026-11-04 00:00:00+00:00", "updated_at": "2026-11-04 17:52:00.468338+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-tools", "mlops", "ai-infrastructure"], "entities": ["AMD", "PyTorch", "ROCm", "Triton", "rocprofv3", "rocprof-compute", "rocprof-sys", "torch.compile"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/performance-profiling-on-amd-gpus-part-7-ai-pytorch-r-edition", "markdown": "https://wpnews.pro/news/performance-profiling-on-amd-gpus-part-7-ai-pytorch-r-edition.md", "text": "https://wpnews.pro/news/performance-profiling-on-amd-gpus-part-7-ai-pytorch-r-edition.txt", "jsonld": "https://wpnews.pro/news/performance-profiling-on-amd-gpus-part-7-ai-pytorch-r-edition.jsonld"}}