Triton kernel performance model

Triton kernel performance model#

2026-07-22

11 min read time

This topic explains how TraceLens computes GFLOPS, TB/s, and other performance metrics for torch.compile-generated Triton kernels.

The performance model is implemented in TraceLens/PerfModel/triton_compiled_perf_model.py.

How Triton kernel naming works#

When you run torch.compile(model), TorchDynamo captures the model’s operations as an FX graph, and the Inductor backend fuses groups of ATen ops into Triton kernels. Each kernel’s name encodes what it does:

triton_poi_fused__unsafe_view_mul_silu_2
│       │         │                    │
│       │         │                    └── index (2nd kernel in the graph)
│       │         └── fused ATen ops (alphabetical)
│       └── kernel type: poi=pointwise, red=reduction, per=persistent-reduction
└── generated by Triton/Inductor

Pointwise kernels (triton_poi_*) process each element independently. One loop dimension: xnumel (total elements).

Reduction kernels (triton_red_*, triton_per_*) reduce along one axis. Two loop dimensions: xnumel (outer) and rnumel (reduction axis).

Requirements and PyTorch version compatibility#

Trace capture requirements#

TraceLens uses a two-tier approach (V2 primary, V1 fallback) to extract perf metrics. For V2 to work, the trace must be captured with record_shapes=True:

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,   # required for V2 metadata
) as prof:
    model(inputs)

Note

Without record_shapes=True, the trace won’t contain Concrete Inputs, Input Dims, or Input type — and V2 falls back to V1.

PyTorch version compatibility#

V2 requires specific trace event fields. The following table shows verified availability (confirmed from actual traces captured on each version):

Trace Field

PyTorch 2.4

PyTorch 2.11

Used By

Concrete Inputs

Yes

Yes

V2: element counts (xnumel, rnumel)

Input Dims

Yes

Yes

V2: per-tensor byte calculation

Input type

Yes

Yes

V2: bytes per element

kernel_file

No

Yes

Not used (informational)

kernel_kwargs

No

Yes

Not used (informational)

num_warps

No

Yes

Not used (informational)

num_stages

No

Yes

Not used (informational)

V2 core fields are available since PyTorch 2.4. The extra fields (kernel_file, num_warps, and so on) were added later but aren’t required by TraceLens.

TORCHINDUCTOR_UNIQUE_KERNEL_NAMES#

Inductor controls kernel naming using TORCHINDUCTOR_UNIQUE_KERNEL_NAMES. When enabled, kernels get descriptive names that encode the kernel category and fused ops (for example, triton_poi_fused_add_mul_silu_0). When disabled, kernels are named generically (triton_0, triton_1).

Note

This is a compile-time setting — it must be set when the trace is captured, not when TraceLens analyzes it.

Default by PyTorch version:

PyTorch Version

Default

Notes

2.4 – 2.5

Disabled

Must opt in: TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1

2.6+

Enabled

Default changed to "1"

Source: torch/_inductor/config.py — verified at tags v2.4.0 and v2.6.0.

Impact on TraceLens:

Metric

Unique names enabled

Unique names disabled

V2 element counts (xnumel, rnumel)

Works

Works

V2 bytes (Data Moved)

Works

Works

FLOPs (_parse_kernel_name)

Works — parses ops from name

Returns [] — GFLOPS = 0

V1 cache fallback

Works

Works

For users on PyTorch 2.4–2.5, set the environment variable before trace capture to get FLOPs:

TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 python your_training_script.py

How it works#

Data flow#

Chrome Trace (.json.gz)
│
├── event["name"]
│   │
│   └── resolve_perf_model_class()
│       └── _match_triton_compiled() → TritonCompiledPerfModel
│
├── event["args"]                          GPU events
│   ├── Concrete Inputs                         │
│   ├── Input Dims                              │
│   └── Input type                              │
│         │                                     │
│    ┌────┴─────┐                               │
│    │ V2 path  │ ──── _meta_from_trace_args()  │
│    └────┬─────┘      ├── scalars → xnumel, rnumel
│         │            ├── dims + dtypes → total_bytes
│     success?         └── _parse_kernel_name() → aten_ops
│     ├── yes ── _meta ─────────┐               │
│     └── no                    │               │
│          │                    │               │
│     ┌────┴─────┐              │               │
│     │ V1 path  │              │               │
│     └────┬─────┘              │               │
│          │                    │               │
│   Inductor Cache (.py)        │               │
│   ├── Original ATen           │        GPUEventAnalyser
│   ├── size_hints              │          busy_time
│   └── signature               │               │
│          │                    │               │
│          └── _lookup() ── _meta               │
│                               │               │
│                          ┌────┴────┐    ┌─────┴─────┐
│                          │ flops() │    │ kernel    │
│                          │ bytes() │    │ time (µs) │
│                          └────┬────┘    └─────┬─────┘
│                               │               │
│                               └───────┬───────┘
│                                       │
│                                  TB/s, TFLOPS/s
│                          (tree_perf.py:compute_perf_metrics)

Two-tier metadata extraction#

TraceLens needs three pieces of information to compute metrics:

  1. Element counts — xnumel and rnumel

  2. Data types and shapes — to calculate bytes moved

  3. Fused ATen ops — to calculate FLOPs

These are obtained through a two-tier approach. V2 is tried first; V1 is the fallback:

# TritonCompiledPerfModel.__init__
self._meta = _meta_from_trace_args(event)   # V2: trace event args
if self._meta is None:
    self._meta = _lookup(self.name, kwargs.get("inductor_cache_dir"))  # V1: cache fallback

V2: trace-intrinsic (primary)#

Extracts metadata directly from the Chrome trace event’s args dict. Available in PyTorch 2.4+ traces (see PyTorch version compatibility). No disk I/O needed — the trace is self-contained.

{
  "name": "triton_poi_fused__unsafe_view_mul_silu_2",
  "args": {
    "Concrete Inputs": ["", "", "268435456"],
    "Input Dims":      [[8, 4096, 8192], [32768, 8192], []],
    "Input type":      ["c10::BFloat16", "c10::BFloat16", "Scalar"]
  }
}

Field

What it provides

Concrete Inputs

Scalar values including xnumel and rnumel (exact, not rounded). Non-integer values (floats, booleans) are skipped; only integer scalars are used for element counts.

Input Dims

Per-tensor shapes for exact byte calculation

Input type

C10 dtype strings (c10::BFloat16 = 2 bytes, c10::Float = 4 bytes)

V1: cache-based (fallback)#

Parses the Inductor wrapper .py files from the torch.compile cache directory. Used when trace args are missing (older PyTorch traces) or when the trace was collected without record_shapes=True.

Cache dirs searched (first match wins):

  1. --inductor_cache_dir CLI argument (explicit)

  2. $TORCHINDUCTOR_CACHE_DIR environment variable

  3. ~/.cache/torchinductor

  4. /tmp/torchinductor_<user>

Three fields are extracted using regex:

# 1. ATen ops (from comment above kernel)
# Original ATen: [aten.mul, aten.silu]

# 2. Element counts (from decorator)
size_hints={'x': 268435456}           # PyTorch 2.7+ dict format
size_hints=[268435456]                # older list format

# 3. Pointer dtypes (from triton_meta)
'signature': {'in_out_ptr0': '*bf16', 'in_ptr0': '*bf16', 'xnumel': 'i32'}

V1 compared with V2#

V2 (trace-intrinsic)

V1 (cache-based)

Data source

Chrome trace event["args"]

Inductor cache .py files on disk

PyTorch version

2.4+

Any

ATen ops source

Parsed from kernel name using DP segmentation against _known_aten_ops.py

Original ATen: [...] comment in cache file

Element counts

Exact (from Concrete Inputs)

Rounded to power-of-2 (size_hints)

Bytes calculation

Per-tensor: prod(shape) * bytes_per_elem

Per-pointer: ptr_bytes * xnumel

External dependency

None — trace is self-contained

Requires cache directory on disk

Works on shared traces

Yes

Only if cache is accessible

Inductor Triton kernel types#

Inductor generates Triton kernels in several categories, each with a distinct 3-character prefix derived from the category name (category[:3]). These 6 categories have been stable from PyTorch 2.3 through 2.11+ (defined in torch/_inductor/wrapper_benchmark.py):

Prefix

Category

Description

triton_poi_

pointwise

Element-wise ops (add, mul, relu, silu, and so on). One loop dimension: xnumel. Fused ops stay in registers — no intermediate global memory writes.

triton_red_

reduction

Reduction ops (sum, mean, batch norm, and so on) using a looped strategy. Two dimensions: xnumel (outer), rnumel (reduction axis). Used when the reduction axis is too large for persistent strategy.

triton_per_

persistent_reduction

Same reduction ops, but keeps the entire reduction axis in registers/SRAM. Used when rnumel is small enough (~1024 elements). Inductor’s should_use_persistent_reduction() heuristic decides.

triton_tem_

template

Complex ops like matrix multiplication (mm, addmm, _scaled_mm) using pre-defined Triton templates with optional fused epilogues.

triton_for_

foreach

Fused operations across lists of tensors (for example, optimizer step updates applied to all parameters at once).

triton_spl_

split_scan

Split scan operations (for example, cumulative sums with decoupled lookback).

TraceLens coverage:

Category

Modeled by TraceLens?

Notes

triton_poi_

Yes

TritonCompiledPerfModel computes FLOPs and bytes

triton_red_

Yes

Same perf model, two-dimensional (xnumel × rnumel)

triton_per_

Yes

Treated identically to triton_red_

triton_tem_

No (not needed)

General matrix multiplication (GEMM) ops already modeled by aten_mm, aten_scaled_mm, and so on

triton_for_

No

Typically optimizer-step ops, not performance-critical

triton_spl_

No

Scan operations, uncommon in typical workloads

Naming pattern:

triton_{category}_fused_{op1}_{op2}_{...}_{index}

The fused ops are listed alphabetically. The index is a sequential counter within the compiled graph. When unique_kernel_names is disabled, the entire name collapses to triton_{index}.

Metric formulas#

FLOPs#

flops = sum(flops_per_elem[op] for op in fused_aten_ops) * xnumel * rnumel

The flops_per_elem table is in _FLOPS_PER_ELEM. Examples: add=1, mul=1, silu=4, gelu=8, exp=4, pow=2, rsqrt=2. Memory-only ops (clone, _to_copy, _unsafe_view) contribute 0.

Bytes#

V2 (trace args):

bytes = sum(prod(shape_i) * bytes_per_elem_i  for each tensor input)
      + ptr_bytes[0] * xnumel                  # only for pointwise (output write)

V1 (cache files):

# Reduction (rnumel > 1):
bytes = sum(ptr_bytes[:-1]) * xnumel * rnumel + ptr_bytes[-1] * xnumel

# Pointwise (rnumel == 1):
bytes = (sum(ptr_bytes) + sum(in_out_extra_bytes)) * xnumel

The in_out_extra_bytes accounts for in_out_ptr parameters that are both read and written — they appear once in the signature but move data twice.

Throughput#

All throughput metrics use the busy kernel time from GPUEventAnalyser (aggregated GPU execution time, not the CPU-side event["dur"]):

TB/s     = (bytes_moved / 1e12) / (kernel_time_us / 1e6)
TFLOPS/s = (gflops / 1e3)      / (kernel_time_us / 1e6)
FLOPS/Byte = flops / bytes_moved
Data Moved (MB) = bytes_moved / (1024 * 1024)

Worked example: pointwise kernel (SwiGLU)#

Kernel: triton_poi_fused__unsafe_view_mul_silu_2 Operation: silu(gate_proj(x)) * up_proj(x) in the SwiGLU MLP

Trace event args (V2 input)#

Concrete Inputs: ['', '', '268435456']
Input Dims:      [[8, 4096, 8192], [32768, 8192], []]
Input type:      ['c10::BFloat16', 'c10::BFloat16', 'Scalar']

Element count extraction#

Integer scalars from Concrete Inputs (empty strings and non-integer values like floats or booleans are skipped): [268435456]

Pointwise kernel → last integer scalar is xnumel:

  • xnumel = 268,435,456 (= 8 x 4,096 x 8,192)

  • rnumel = 1

Bytes calculation#

Step 1: Calculate bytes for each tensor input:

Input

Dims

Type

Bytes per elem

Elements

Bytes

in_out_ptr0

[8, 4096, 8192]

bf16

2

268,435,456

536,870,912

in_ptr0

[32768, 8192]

bf16

2

268,435,456

536,870,912

xnumel

[]

Scalar

—

—

skipped

 
input_bytes = (268,435,456 * 2) + (268,435,456 * 2) = 1,073,741,824

Step 2: Add output write for pointwise kernels.

In pointwise kernels, in_out_ptr0 is both read and written — the output overwrites the input buffer. The input bytes above already count the read. V2 adds one extra write of ptr_bytes[0] * xnumel:

output_write = 2 * 268,435,456 = 536,870,912

Total:

bytes_moved = 1,073,741,824 + 536,870,912
            = 1,610,612,736 bytes
            = 1,536 MB

FLOPs calculation#

Fused ops parsed from the kernel name: aten._unsafe_view, aten.mul, aten.silu

Op

FLOPs per element

aten._unsafe_view

0 (memory-only)

aten.mul

1

aten.silu

4

Total

5

 
flops = 5 * xnumel * rnumel
      = 5 * 268,435,456 * 1
      = 1,342,177,280
GFLOPS = 1,342,177,280 / 1e9 = 1.342

Throughput metrics#

GPU kernel time from the Chrome trace: 505.3 us

TB/s     = (1,610,612,736 / 1e12) / (505.3 / 1e6) = 3.19 TB/s
TFLOPS/s = (1.342 / 1e3) / (505.3 / 1e6)          = 2.66 TFLOPS/s

Summary#

Metric

Value

Data Moved

1,536 MB

GFLOPS

1.342

FLOPS/Byte

0.833

TB/s

3.19

TFLOPS/s

2.66

Worked example: reduction kernel (RMSNorm)#

Kernel: triton_red_fused_add_mean_mul_pow_rsqrt_0 Operation: RMSNorm — x * rsqrt(mean(x^2) + eps) * weight

Trace event args (V2 input)#

Concrete Inputs: ['', '', '', '32768', '2048']
Input Dims:      [[8, 4096, 2048], [2048], [8, 4096, 2048], [], []]
Input type:      ['c10::BFloat16', 'c10::BFloat16', 'c10::BFloat16', 'Scalar', 'Scalar']

Element count extraction#

Integer scalars from Concrete Inputs (non-integer values skipped): [32768, 2048]

Reduction kernel → last two integer scalars are xnumel and rnumel:

  • xnumel = 32,768 (= 8 x 4,096 = batch x seq_len)

  • rnumel = 2,048 (= hidden_dim, the reduction axis)

Bytes calculation (V2)#

Input

Dims

Type

Bytes per elem

Elements

Bytes

in_ptr0

[8, 4096, 2048]

bf16

2

67,108,864

134,217,728

in_ptr1

[2048]

bf16

2

2,048

4,096

out_ptr1

[8, 4096, 2048]

bf16

2

67,108,864

134,217,728

No extra output write for reduction kernels (the output pointer is already listed as a separate tensor).

bytes_moved = 134,217,728 + 4,096 + 134,217,728
            = 268,439,552 bytes
            ≈ 256 MB

Bytes calculation (V1 — for comparison)#

V1 uses xnumel, rnumel, and per-pointer byte widths from the signature instead of per-tensor shapes:

signature: {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'out_ptr1': '*bf16',
            'xnumel': 'i32', 'r0_numel': 'i32'}
ptr_bytes = [2, 2, 2]    # three *bf16 pointers

Reduction formula: all inputs read xnumel * rnumel elements, the output (last pointer) writes xnumel elements:

bytes = sum(ptr_bytes[:-1]) * xnumel * rnumel + ptr_bytes[-1] * xnumel
      = (2 + 2) * 32,768 * 2,048 + 2 * 32,768
      = 268,435,456 + 65,536
      = 268,500,992 bytes
      ≈ 256 MB

Note

V1 and V2 give slightly different values because V1 assumes all input pointers access the full xnumel * rnumel grid, while V2 uses exact tensor shapes (the weight tensor in_ptr1 is only [2048], not [32768, 2048]).

FLOPs calculation#

Fused ops: aten.add, aten.mean, aten.mul, aten.pow, aten.rsqrt

Op

FLOPs per element

aten.add

1

aten.mean

1

aten.mul

1

aten.pow

2

aten.rsqrt

2

Total

7

 
flops = 7 * 32,768 * 2,048 = 469,762,048
GFLOPS = 0.470

Throughput metrics#

GPU kernel time: 79.4 us

TB/s     = (268,439,552 / 1e12) / (79.4 / 1e6)  = 3.38 TB/s
TFLOPS/s = (0.470 / 1e3) / (79.4 / 1e6)         = 5.92 TFLOPS/s

Results for the example transformer block#

Model: TransformerBlock(dim=2048, n_heads=16, mlp_ratio=4), batch=8, seq_len=4096, dtype=bf16. Traced with PyTorch 2.11+rocm7.2 on AMD Instinct™ 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.38

5.92

triton_red_fused__unsafe_view_add_mean_mul_pow_rsqrt_1 (RMSNorm+residual)

0.470

384

1.17

4.32

5.04

triton_poi_fused__unsafe_view_mul_silu_2 (SwiGLU)

1.342

1,536

0.83

3.19

2.66

triton_poi_fused__unsafe_view_add_3 (residual add)

0.067

512

0.12

3.81

0.48

References#

Contents