Note
Go to the end to download the full example code.
Fused Softmax¶
In this tutorial, you will write a fused softmax operation that is significantly faster than PyTorch’s native op for a particular class of matrices: those whose rows can fit in the GPU’s SRAM.
In doing so, you will learn about:
The benefits of kernel fusion for bandwidth-bound operations.
Reduction operators in Triton.
Motivations¶
Custom GPU kernels for elementwise additions are educationally valuable but won’t get you very far in practice. Let us consider instead the case of a simple (numerically stabilized) softmax operation:
import torch
import triton
import triton.language as tl
from triton.runtime import driver
DEVICE = triton.runtime.driver.active.get_active_torch_device()
def is_hip():
return triton.runtime.driver.active.get_current_target().backend == "hip"
def is_cdna():
return is_hip() and triton.runtime.driver.active.get_current_target().arch in ('gfx940', 'gfx941', 'gfx942',
'gfx90a', 'gfx908')
def naive_softmax(x):
"""Compute row-wise softmax of X using native pytorch
We subtract the maximum element in order to avoid overflows. Softmax is invariant to
this shift.
"""
# read MN elements ; write M elements
x_max = x.max(dim=1)[0]
# read MN + M elements ; write MN elements
z = x - x_max[:, None]
# read MN elements ; write MN elements
numerator = torch.exp(z)
# read MN elements ; write M elements
denominator = numerator.sum(dim=1)
# read MN + M elements ; write MN elements
ret = numerator / denominator[:, None]
# in total: read 5MN + 2M elements ; wrote 3MN + 2M elements
return ret
When implemented naively in PyTorch, computing y = naive_softmax(x) for \(x \in R^{M \times N}\)
requires reading \(5MN + 2M\) elements from DRAM and writing back \(3MN + 2M\) elements.
This is obviously wasteful; we’d prefer to have a custom “fused” kernel that only reads
X once and does all the necessary computations on-chip.
Doing so would require reading and writing back only \(MN\) bytes, so we could
expect a theoretical speed-up of ~4x (i.e., \((8MN + 4M) / 2MN\)).
The torch.jit.script flags aims to perform this kind of “kernel fusion” automatically
but, as we will see later, it is still far from ideal.
Compute Kernel¶
Our softmax kernel works as follows: each program loads a set of rows of the input matrix X strided by number of programs, normalizes it and writes back the result to the output Y.
Note that one important limitation of Triton is that each block must have a power-of-two number of elements, so we need to internally “pad” each row and guard the memory operations properly if we want to handle any possible input shapes:
@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_rows, n_cols, BLOCK_SIZE: tl.constexpr,
num_stages: tl.constexpr):
# starting row of the program
row_start = tl.program_id(0)
row_step = tl.num_programs(0)
for row_idx in tl.range(row_start, n_rows, row_step, num_stages=num_stages):
# The stride represents how much we need to increase the pointer to advance 1 row
row_start_ptr = input_ptr + row_idx * input_row_stride
# The block size is the next power of two greater than n_cols, so we can fit each
# row in a single block
col_offsets = tl.arange(0, BLOCK_SIZE)
input_ptrs = row_start_ptr + col_offsets
# Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols
mask = col_offsets < n_cols
row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
# Subtract maximum for numerical stability
row_minus_max = row - tl.max(row, axis=0)
# Note that exponentiation in Triton is fast but approximate (i.e., think __expf in CUDA)
numerator = tl.exp(row_minus_max)
denominator = tl.sum(numerator, axis=0)
softmax_output = numerator / denominator
# Write back output to DRAM
output_row_start_ptr = output_ptr + row_idx * output_row_stride
output_ptrs = output_row_start_ptr + col_offsets
tl.store(output_ptrs, softmax_output, mask=mask)
We can create a helper function that enqueues the kernel and its (meta-)arguments for any given input tensor.
properties = driver.active.utils.get_device_properties(DEVICE.index)
NUM_SM = properties["multiprocessor_count"]
NUM_REGS = properties["max_num_regs"]
SIZE_SMEM = properties["max_shared_mem"]
WARP_SIZE = properties["warpSize"]
target = triton.runtime.driver.active.get_current_target()
kernels = {}
def softmax(x):
n_rows, n_cols = x.shape
# The block size of each loop iteration is the smallest power of two greater than the number of columns in `x`
BLOCK_SIZE = triton.next_power_of_2(n_cols)
# Another trick we can use is to ask the compiler to use more threads per row by
# increasing the number of warps (`num_warps`) over which each row is distributed.
# You will see in the next tutorial how to auto-tune this value in a more natural
# way so you don't have to come up with manual heuristics yourself.
num_warps = 8
# Number of software pipelining stages.
num_stages = 4 if SIZE_SMEM > 200000 else 2
# Allocate output
y = torch.empty_like(x)
# pre-compile kernel to get register usage and compute thread occupancy.
kernel = softmax_kernel.warmup(y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE,
num_stages=num_stages, num_warps=num_warps, grid=(1, ))
kernel._init_handles()
n_regs = kernel.n_regs
size_smem = kernel.metadata.shared
if is_hip():
# NUM_REGS represents the number of regular purpose registers. On CDNA architectures this is half of all registers available.
# However, this is not always the case. In most cases all registers can be used as regular purpose registers.
# ISA SECTION (3.6.4 for CDNA3)
# VGPRs are allocated out of two pools: regular VGPRs and accumulation VGPRs. Accumulation VGPRs are used
# with matrix VALU instructions, and can also be loaded directly from memory. A wave may have up to 512 total
# VGPRs, 256 of each type. When a wave has fewer than 512 total VGPRs, the number of each type is flexible - it is
# not required to be equal numbers of both types.
if is_cdna():
NUM_GPRS = NUM_REGS * 2
# MAX_NUM_THREADS represents maximum number of resident threads per multi-processor.
# When we divide this number with WARP_SIZE we get maximum number of waves that can
# execute on a CU (multi-processor) in parallel.
MAX_NUM_THREADS = properties["max_threads_per_sm"]
max_num_waves = MAX_NUM_THREADS // WARP_SIZE
occupancy = min(NUM_GPRS // WARP_SIZE // n_regs, max_num_waves) // num_warps
else:
occupancy = NUM_REGS // (n_regs * WARP_SIZE * num_warps)
occupancy = min(occupancy, SIZE_SMEM // size_smem)
num_programs = NUM_SM * occupancy
num_programs = min(num_programs, n_rows)
# Create a number of persistent programs.
kernel[(num_programs, 1, 1)](
y,
x,
x.stride(0),
y.stride(0),
n_rows,
n_cols,
)
return y
Unit Test¶
We make sure that we test our kernel on a matrix with an irregular number of rows and columns. This will allow us to verify that our padding mechanism works.
torch.manual_seed(0)
x = torch.randn(1823, 781, device=DEVICE)
y_triton = softmax(x)
y_torch = torch.softmax(x, axis=1)
assert torch.allclose(y_triton, y_torch), (y_triton, y_torch)
As expected, the results are identical.
Benchmark¶
Here we will benchmark our operation as a function of the number of columns in the input matrix – assuming 4096 rows.
We will then compare its performance against (1) torch.softmax and (2) the naive_softmax defined above.
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=['N'], # argument names to use as an x-axis for the plot
x_vals=[128 * i for i in range(2, 100)], # different possible values for `x_name`
line_arg='provider', # argument name whose value corresponds to a different line in the plot
line_vals=['triton', 'torch'], # possible values for `line_arg``
line_names=[
"Triton",
"Torch",
], # label name for the lines
styles=[('blue', '-'), ('green', '-')], # line styles
ylabel="GB/s", # label name for the y-axis
plot_name="softmax-performance", # name for the plot. Used also as a file name for saving the plot.
args={'M': 4096}, # values for function arguments not in `x_names` and `y_name`
))
def benchmark(M, N, provider):
x = torch.randn(M, N, device=DEVICE, dtype=torch.float32)
stream = getattr(torch, DEVICE.type).Stream()
getattr(torch, DEVICE.type).set_stream(stream)
if provider == 'torch':
ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1))
if provider == 'triton':
ms = triton.testing.do_bench(lambda: softmax(x))
gbps = lambda ms: 2 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
return gbps(ms)
benchmark.run(show_plots=True, print_data=True)
softmax-performance:
N Triton Torch
0 256.0 471.424827 691.798453
1 384.0 615.079869 805.984266
2 512.0 750.274944 916.124374
3 640.0 818.378659 945.206405
4 768.0 892.654907 1015.145273
5 896.0 946.337105 1060.407404
6 1024.0 999.443894 1112.085580
7 1152.0 1095.252432 613.984071
8 1280.0 1138.299978 669.848616
9 1408.0 1164.017567 720.800222
10 1536.0 1192.525407 779.856515
11 1664.0 1217.177752 813.934385
12 1792.0 1239.332759 857.480468
13 1920.0 1252.474291 908.756434
14 2048.0 1280.069838 958.774182
15 2176.0 1239.433631 979.659355
16 2304.0 1247.761109 1012.206004
17 2432.0 1278.509228 1053.455128
18 2560.0 1289.007869 1088.444044
19 2688.0 1287.306153 1100.124616
20 2816.0 1306.450954 1132.555213
21 2944.0 1308.963851 1164.589075
22 3072.0 1328.726719 1185.969562
23 3200.0 1337.289290 1191.670691
24 3328.0 1345.335760 1227.729491
25 3456.0 1356.153227 1245.061721
26 3584.0 1351.351008 1258.456953
27 3712.0 1363.329242 1268.075776
28 3840.0 1371.651701 1302.478376
29 3968.0 1377.028724 1312.465033
30 4096.0 1383.239622 1324.697413
31 4224.0 1336.185069 1160.506709
32 4352.0 1337.186178 1171.761479
33 4480.0 1352.246591 1184.820061
34 4608.0 1365.647523 1196.269814
35 4736.0 1360.954474 1196.409698
36 4864.0 1373.752031 1217.874702
37 4992.0 1370.069986 1237.749977
38 5120.0 1377.004616 1247.818237
39 5248.0 1377.664869 1260.722562
40 5376.0 1375.240426 1281.777893
41 5504.0 1386.455844 1296.453905
42 5632.0 1384.552270 1312.449720
43 5760.0 1398.264105 1326.043202
44 5888.0 1387.661642 1343.727758
45 6016.0 1398.786927 1355.933924
46 6144.0 1412.156947 1373.348319
47 6272.0 1416.525785 1374.979250
48 6400.0 1415.883636 1387.088301
49 6528.0 1419.067162 1395.926371
50 6656.0 1421.104254 1401.688610
51 6784.0 1412.452383 1412.739161
52 6912.0 1430.227401 1421.593451
53 7040.0 1420.411159 1433.135892
54 7168.0 1427.999402 1436.441142
55 7296.0 1432.994948 1443.234449
56 7424.0 1430.016482 1447.591727
57 7552.0 1428.233139 1453.301807
58 7680.0 1432.569848 1458.304969
59 7808.0 1432.159579 1466.414786
60 7936.0 1433.671541 1469.480213
61 8064.0 1437.912610 1471.413973
62 8192.0 1440.151724 1485.740131
63 8320.0 1386.921525 1402.073965
64 8448.0 1381.158507 1404.303608
65 8576.0 1392.322937 1397.719044
66 8704.0 1390.809928 1398.734510
67 8832.0 1387.124772 1402.026963
68 8960.0 1399.912948 1414.461639
69 9088.0 1407.038780 1418.233156
70 9216.0 1405.001796 1425.122909
71 9344.0 1399.885687 1425.606734
72 9472.0 1400.700169 1431.404883
73 9600.0 1392.575611 1435.040357
74 9728.0 1400.912346 1445.077234
75 9856.0 1414.215554 1443.391821
76 9984.0 1398.355313 1449.212535
77 10112.0 1414.047128 1455.536752
78 10240.0 1419.752373 1465.529093
79 10368.0 1411.736714 1461.317807
80 10496.0 1415.264552 1468.267799
81 10624.0 1408.512639 1467.672714
82 10752.0 1403.474228 1472.718987
83 10880.0 1406.134012 1479.630736
84 11008.0 1420.682049 1476.522222
85 11136.0 1421.928237 1483.508182
86 11264.0 1428.475578 1488.376386
87 11392.0 1411.653781 1489.595182
88 11520.0 1422.504962 1492.618316
89 11648.0 1431.282172 1496.709178
90 11776.0 1431.660039 1501.693476
91 11904.0 1442.758616 1503.381644
92 12032.0 1425.280961 1506.980696
93 12160.0 1420.666943 1512.620639
94 12288.0 1434.262221 1392.210664
95 12416.0 1450.401569 1391.525977
96 12544.0 1442.717940 1394.315694
97 12672.0 1449.240213 1392.694059
- In the above plot, we can see that:
Triton is 4x faster than the Torch JIT. This confirms our suspicions that the Torch JIT does not do any fusion here.
Triton is noticeably faster than
torch.softmax– in addition to being easier to read, understand and maintain. Note however that the PyTorch softmax operation is more general and will work on tensors of any shape.
Total running time of the script: (0 minutes 23.191 seconds)