Tag: triton

Kapil Sharma from the PyTorch team has a great series of posts diving into the Triton compiler process: 1, 2, 3. As covered there, Triton lowers to a series of intermediate representations, and each level has a set of transformational passes that implement optimizations. TTIR is the generic Triton IR and leverages a number of standard MLIR passes like common subexpression elimination, as well as some Triton specific passes like managing broadcast ops. That’s then lowered to TTGIR, a GPU-specific IR¹

triton/third_party/nvidia/backend/compiler.py at rc/3.3.x · triton-lang/triton

The different backends configure the passes appropriate for their targets, so the Nvidia TTGIR configuration above details the passes for Nvidia hardware. Some are gated on the specific backend targeted, like warp specialization:

passes.ttgpuir.add_ws_task_partition(pm, opt.num_consumer_groups)
passes.ttgpuir.add_taskid_propagate(pm, opt.num_consumer_groups)
passes.ttgpuir.add_ws_data_partition(pm, opt.num_consumer_groups)
passes.ttgpuir.add_ws_code_partition(pm, opt.num_buffers_warp_spec, opt.num_consumer_groups,
opt.reg_dec_producer, opt.reg_inc_consumer)
passes.ttgpuir.add_pipeline(pm, opt.num_stages, dump_enabled)
passes.ttgpuir.add_ping_pong_sync(pm, opt.num_consumer_groups)
passes.ttgpuir.add_ws_lowering(pm, opt.num_consumer_groups)

Another example is using the Tensor Memory Accelerator for async loading on Hopper+

if capability // 10 >= 9:
   nvidia.passes.ttnvgpuir.add_tma_lowering(pm)
   nvidia.passes.ttnvgpuir.add_fence_insertion(pm)

For a quick recap of the why and how of Warp Specialization, check Colfax’s guide to optimizing a GEMM:

The most basic method by which GPU programmers can create overlapping is via excess warps (warps are groupings of 32 contiguous threads). Nvidia GPUs allow a large number of warps per SM (streaming multiprocessors), and can switch between them with minimal overhead. In particular, the warp schedulers can simply switch to another warp if one warp encounters a slow memory fetch. In order to give the warp schedulers more opportunity to hide latency, a technique called warp-specialization was introduced circa 2011 [1, 2]. With warp-specialization, some warps are dedicated to memory fetches (producers), while others are dedicated to compute (consumers), and named barriers are used for synchronization between them. The idea is that the warp schedulers can then more easily hide the latency of copy operations within compute (and vice-versa).

Even more generally than overlapping memory loads you can overlap other kinds of work. SMs have 1 Tensor Core to 32 ALUs (Cuda Core) (x4 on recent hardware). This means that you can overlap other kinds of work, stuff that isn’t dot products. It’s really common to want to load memory, do a matmul then apply a pointwise function like a relu or other activation function. You aim to keep the Tensor Core as busy as possible with a series of matmuls, and warp specialization lets you do that.

The transforms to implement this are implemented in triton/lib/Dialect/TritonGPU/Transforms at rc/3.3.x · triton-lang/triton

The first task partitions ops in the kernel. This task looks for load ops and dot product ops, and partitions them into producer (loads) and consumer (process) groups.

// Step 1. Select loads into the first task, which is the producer task by
// default. Place dots into the second task, which is the consumer.
// Only consider loads that are connected to a dot op in a loop.

The next task is a bookkeeping one to propagates the task IDs, so if there are unlabeled ops they are attached one of the partitions.

The WSDataPartition transform partitions the dot operations into consumer groups. It splits the dot products inside loops along M or N dimensions to be processable within a warp group, ensuring all dot operations are sliced and the slices labelled with task_ids.

Just to look at some of the numbers: A warp consists of 32 threads, and a sync op (for loading) uses a whole warp. A warp group is a set of 4 warps: this is pertinent because the TensorCore MMA prefers 4 warps working on 64x(64/128/256)xK tiles. Triton already has a WarpGroupDotOp that tries to set this up, and that’s one of the operations targeted in this pass. The pass splits a Triton CTA tile, which may be 128×256, so that each consumer warp group has a 64 row (or 256 column) chunk.

 if (sliceSizeM >= 64) {
      LLVM_DEBUG({ LDBG("partition along M\n"); });
      partitionDim = 0;
      partitionSize = sliceSizeM;
      partitionOperand = opndA;
    } else if (sliceSizeN >= 256) {
      LLVM_DEBUG({ LDBG("partition along N\n"); });
      partitionDim = 1;
      partitionSize = sliceSizeN;
      partitionOperand = opndB;
    } else {
      LDBG("partition not possible: " << sliceSizeM << " " << sliceSizeN);
      return false;
    }

The next pass is WSCodePartition. This is a big transform. It takes the task-sliced IR from the DataPartition and sets up the producer warp group to copy from global GPU mem to SMEM (or on blackwell TMEM). It also drops in barriers using product.acquire/cpmmit and consumer.wait/release to ensure proper ordering between the groups. The transform identifies “channels”, places where data is loaded (using load or descriptor_load for TMA on Hopper+) and associates the producer task_id with all the consumer task IDs that need to process that data.

Conceptually, the transform is turning the loops in the original kernel into something like this:

for k: # K-dimension tiles)
    ##### PRODUCER #####
    producer.acquire(token, idx, phase) # reserve smem[idx]
    async_copy_global_to_local(smem[idx]) # GMEM → SMEM[idx]
    producer.commit(token, idx) # make slot visible

    #####  CONSUMERS (run in parallel warps) #####
    ## Consumer 0 ##
    consumer.wait(token, idx, phase) # sleep until slot ready
    mma_sync(accum0, smem[idx], …) # read-only, do matmul
    consumer.release(token, idx) 

    ## Consumer 1 ##
    consumer.wait(token, idx, phase)
    mma_sync(accum1, smem[idx], …)
    consumer.release(token, idx)

    # Repeat for extra consumers. 

    # increment circular buffer
    idx   = (idx + 1) % numBuffers 
    # Toggle each time we hit producer, indicate old vs new data.
    phase = phase ^ (idx == 0)
}

The next pass is a generic pipeline pass. Each op is assigned a stage based on latency: e.g. slower ops like loads go into stage 0. This is then transformed with modulo scheduling. Any sync ops are converted to async variants (in lowerLoops), before writing out (in expandLoops) prologue, kernel and epilogue loops that contain all the relevant ops.

Finally the WSLowering pass takes the various operators we have (like producer.acquire) and replaces them with the hardware specific variants (e.g. wait_barrier). It also handles the bookkeeping like generating the warp group and task_ids from the warp ID:

OpBuilder builder(op);
Value _4 = builder.create<arith::ConstantIntOp>(loc, WARPS_PER_TASK, 32);
Value warpId = builder.create<ttng::GetCanonicalWarpIdOp>(loc);
Value asyncTaskId = builder.create<arith::DivUIOp>(loc, warpId, _4);
op.getResult().replaceAllUsesWith(asyncTaskId);

This is a wordy IR way of saying

warpId = gpu.thread_id().x / 32 
task_id = warpId / 4

Now the code is ready for lowering to regular PTX, and all of the warp-specific stuff is captured explicitly!

1. Note: Because of some project weirdness, warp specialization is quite different in the release branches from main, so I’ll refer to 3.3 from here on. It’s in very active development (by teams at Meta, OpenAI and Nvidia!) so the specifics are quite likely to change over coming releases! ↩︎

torch.compile offers some knobs for controlling the trade-off of execution performance with longer compile times. This is particularly useful for inference, where the same model will be running for a long time.

model_autotune = torch.compile(model, mode="max-autotune")

Passing the max-autotune option to instructs the compiler to test more options for the operations. The compiler has the option to use pre-built aten kernels, leverage kernels from libraries like CuDNN or Cutlass, or use templated Triton kernels. When autotuning, specific variants are tested on device with the shape information identified during tracing, and the fastest options are selected. Thanks to Triton templates, it can also use options like fusions where pointwise ops can be fused into a single kernel via a Triton template, saving kernel launch overhead.

The downside of this is that testing the options takes more time, so using max-autotune can lead to some very extended compile times. You also need a hefty enough GPU to get the benefit: is_big_gpu gates it on the number of SMs, so it works best on a 3090, V100 or above.

You can see a lot of the autotuning options in _inductor/config.py. Backends that are considered are set separately for GEMMs and convolution ops:

max_autotune_gemm_backends = os.environ.get(
    "TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS", "ATEN,TRITON,CPP"
).upper()

Each kernel has implementations using for the different backends which are added to possible choices. e.g. in _inductor/kernels/mm.py you can see calls to use_[backend]_template that verify whether the backend in question is a choice:

if is_nonzero and use_cutlass_template(layout, m, n, k):
        CUTLASS3xGemmTemplate.add_cutlass_gemm_choices(choices, layout, [mat1, mat2])

_inductor/select_algorithm.py does the actual benchmarking through the choices.

If you run autotuning, you’ll get some log output, and caches will be written to /tmp/torchinductor_yourusername.

We can try this out on a simple MLP:

import torch, time

class SimpleMLP(torch.nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.linear1 = torch.nn.Linear(in_features, hidden_features)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(hidden_features, out_features)
    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

# Set up device and model
device = 'cuda'
model = SimpleMLP(in_features=1024, hidden_features=1024, out_features=1024).to(device)
x = torch.randn(256, 1024, device=device)  # batch of 256, 1024 features each

# Compile the model in default mode and max-autotune mode
model_default = torch.compile(model, mode="default")

# Warm-up runs (to trigger compilation)
torch.compiler.reset()
with torch.no_grad():
    model_default(x)
torch.cuda.synchronize()  # ensure warm-up completes

# Measure performance of default compiled model
start = torch.cuda.Event(enable_timing=True); end = torch.cuda.Event(enable_timing=True)
with torch.no_grad():
    start.record()
    for _ in range(50):
        _ = model_default(x)
    end.record()
torch.cuda.synchronize()
time_default_ms = start.elapsed_time(end) / 50.0
torch.compiler.reset()

model_autotune = torch.compile(model, mode="max-autotune")

with torch.no_grad():
    model_autotune(x)
torch.cuda.synchronize()  # ensure warm-up completes

# Measure performance of max-autotune compiled model
start = torch.cuda.Event(enable_timing=True); end = torch.cuda.Event(enable_timing=True)
with torch.no_grad():
    start.record()
    for _ in range(50):
        _ = model_autotune(x)
    end.record()
torch.cuda.synchronize()
time_autotune_ms = start.elapsed_time(end) / 50.0

print(f"Average inference time - torch.compile default: {time_default_ms:.3f} ms")
print(f"Average inference time - torch.compile max-autotune: {time_autotune_ms:.3f} ms")

Disappointedly, this is the result:

Average inference time - torch.compile default: 0.113 ms
Average inference time - torch.compile max-autotune: 3.251 ms

We can turn on logging with the TORCH_LOG env variable: some useful options are inductor, autotuning, and perf_hints.

TORCH_LOGS="perf_hints" python tune.py

You can control many more autotune options via the options flags, though its incompatible with passing a mode value. We can recreate the max-autotune mode, and turn on some useful tracing options like this (note that the options version uses an underscore, the mode a hypen!)

model_autotune = torch.compile(
    model,
    options={
        "max_autotune": True,
        "triton.cudagraphs": True,
        "coordinate_descent_tuning": True,
        "trace.enabled": True,
        "trace.graph_diagram": True,
    },
)

Options "trace.enabled": True, "trace.graph_diagram": True generate trace outputs, and output a nice diagram of the captured graph. Cudagraphs turned out to be the culprit here, which is common enough there is a non-cudagraph mode available to stop you having to remember all the options:

model_autotune = torch.compile(model, mode="max-autotune-no-cudagraphs")

As you can see here in the graphs of with and without, the slower version actually has an extra fusion performed!

Triton Autotuning

Triton also conducts autotuning, but it’s a little more explicit. When authoring a Triton kernel you can specify configurations. At compile time each config variant will be tested, the most performant one picked and the choice stored for future calls. A key value can be provided to indicate when to re-autotune based on changing inputs:

import os
import torch
import triton
import triton.language as tl

# Just to save passing this on the command line
os.environ["TRITON_PRINT_AUTOTUNING"] = "1"  

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 128}, num_warps=4,  num_stages=2),
        triton.Config({'BLOCK_SIZE': 256}, num_warps=8,  num_stages=2),
    ],
    key=['N']            # re‑tune only if the length N changes
)
@triton.jit
def vecadd_kernel(x_ptr, y_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
    pid   = tl.program_id(0)
    offs  = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask  = offs < N
    x     = tl.load(x_ptr  + offs, mask=mask, other=0.0)
    y     = tl.load(y_ptr  + offs, mask=mask, other=0.0)
    tl.store(out_ptr + offs, x + y, mask=mask)

def vec_add(x: torch.Tensor, y: torch.Tensor):
    assert x.is_cuda and y.is_cuda
    N   = x.numel()
    out = torch.empty_like(x)
    grid = (triton.cdiv(N, 128),)         # 128 = smallest BLOCK_SIZE we declared
    vecadd_kernel[grid](x, y, out, N)    
    return out

x = torch.randn(1 << 20, device="cuda")   # 1 048 576 elements
y = torch.randn_like(x)

_ = vec_add(x, y)  # first call → autotuning prints to stdout
_ = vec_add(x, y)  # second call → no autotuning, uses the best config found

Setting the env variable TRITON_PRINT_AUTOTUNING documents the process as it goes:

Autotuning kernel vecadd_kernel with config BLOCK_SIZE: 128, num_warps: 4, num_ctas: 1, num_stages: 2, num_buffers_warp_spec: 0, num_consumer_groups: 0, reg_dec_producer: 0, reg_inc_consumer: 0, maxnreg: None
Autotuning kernel vecadd_kernel with config BLOCK_SIZE: 256, num_warps: 8, num_ctas: 1, num_stages: 2, num_buffers_warp_spec: 0, num_consumer_groups: 0, reg_dec_producer: 0, reg_inc_consumer: 0, maxnreg: None
Triton autotuning for function vecadd_kernel finished after 0.44s; best config selected: BLOCK_SIZE: 128, num_warps: 4, num_ctas: 1, num_stages: 2, num_buffers_warp_spec: 0, num_consumer_groups: 0, reg_dec_producer: 0, reg_inc_consumer: 0, maxnreg: None;

You can use the same do_bench tester that the autotuner does, and see how the performance varies yourself:

import torch, triton, triton.testing as tt
import triton.language as tl

@triton.jit
def vecadd_kernel(x_ptr, y_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
    offs  = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask  = offs < N
    tl.store(out_ptr + offs,
             tl.load(x_ptr + offs, mask=mask) +
             tl.load(y_ptr + offs, mask=mask),
             mask=mask)

# tensors
N   = 1 << 20
x   = torch.randn(N, device='cuda')
y   = torch.randn_like(x)
out = torch.empty_like(x)

def bench(block_size, num_warps):
    grid = (triton.cdiv(N, block_size),)
    # tt.do_bench returns [median, p20, p80] in micro‑seconds
    return tt.do_bench(
        lambda: vecadd_kernel[grid](x, y, out, N, BLOCK_SIZE=block_size, num_warps=num_warps),
        warmup=5, rep=16, return_mode="all", quantiles=(0.5, 0.2, 0.8)
    )
timings = {
    "128/4": bench(128, 4),
    "256/8": bench(256, 8),
}

print("timings:", timings)

Running that gives shows that both kernels are basically equivalent, but the first one is slightly faster over the 16 runs.

timings: {'128/4': [0.01945599913597107, 0.01945599913597107, 0.02028159946203232], '256/8': [0.01945599913597107, 0.01945599913597107, 0.020479999482631683]}

There are a couple of different options to profile a Triton kernel.

Proton

Proton is the profiler that ships with Triton (profiler for triton). You can enable it and (optionally) activate/deactive around specific regions you want to profile. You have the ability to annotate functions with specific metrics as well.

 session = proton.start()  # Start profiling session

bias = torch.rand((256,), device='cuda', dtype=torch.float16)  # Bias vector
flops = 2 * M * N * K
bytes_accessed = A.element_size() * M*K + B.element_size() * K*N + C.element_size() * M*N  # rough bytes
    with proton.scope(f"fused_gemm_bias_relu [M={M}, N={N}, K={K}]", {"flops": flops, "bytes": bytes_accessed}):
        fused_gemm_bias_relu[grid](  
            A, B, C, bias, 
            M, N, K, 
            A.stride(0), A.stride(1), B.stride(0), B.stride(1), C.stride(0), C.stride(1),
            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K
        )

 proton.finalize() 

The output can be visualized with the built-in viewer:

proton-viewer -m time/ms,tflop/s ./proton.hatchet

0.040 6.645 ROOT
├─ 0.004 nan _ZN2at6native55_GLOBAL__N__11f7a751_22_DistributionUniform_[...]_15PhiloxCudaStateESH_SI_
└─ 0.037 7.284 fused_gemm_bias_relu [M=1024, N=256, K=512]
   └─ 0.037 nan fused_gemm_bias_relu

In this case you can see both the (trimmed!) generated name for the bias tensor set up as well as the name of my custom kernel.

nsight-compute

Nvidia also have a good range of tools for looking at performance. Note will need to enable access to the counters on device for this:

NVIDIA Development Tools Solutions – ERR_NVGPUCTRPERM: Permission issue with Performance Counters | NVIDIA Developer

On the offchance you’re doing this on WSL, https://peterchng.com/blog/2024/03/02/profiling-cuda-programs-on-wsl-2/ walks through the set up!

Nvidia ships nsight system which tracks larger system wide metrics, and nsight compute which is more focused on profiling execution. You can run it against a script like so:

ncu -o profile_results python test.py

The tool comes with a nice GUI for inspecting the results. It can show you the PTX or SASS source for the kernels, offers metrics like actively used registers (good for checking on register spilling), and calls out warnings on poor utilization or memory clashes.

Upcoming intra-kernel profiler

[tool][proton] Intra kernel profiling support by fywkevin · Pull Request #4861 · triton-lang/triton

There is an extension coming for Proton that enables profiling within kernels. This reserves a pre-allocated buffer on device and logs metrics locally, for reading out at the end of the execution. It outputs as a chrome trace for use within a wide range of dev tools. While this isn’t merged into mainline yet, you can see an example of the usage in the dev repo.