338 lines
11 KiB
ReStructuredText
338 lines
11 KiB
ReStructuredText
|
|
.. DO NOT EDIT.
|
|
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
|
|
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
|
|
.. "getting-started/tutorials/02-fused-softmax.py"
|
|
.. LINE NUMBERS ARE GIVEN BELOW.
|
|
|
|
.. only:: html
|
|
|
|
.. note::
|
|
:class: sphx-glr-download-link-note
|
|
|
|
Click :ref:`here <sphx_glr_download_getting-started_tutorials_02-fused-softmax.py>`
|
|
to download the full example code
|
|
|
|
.. rst-class:: sphx-glr-example-title
|
|
|
|
.. _sphx_glr_getting-started_tutorials_02-fused-softmax.py:
|
|
|
|
|
|
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.
|
|
You will learn about:
|
|
|
|
- The benefits of kernel fusion for bandwidth-bound operations.
|
|
- Reduction operators in Triton.
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 14-18
|
|
|
|
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:
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 18-46
|
|
|
|
.. code-block:: default
|
|
|
|
|
|
import torch
|
|
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
|
|
@torch.jit.script
|
|
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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 47-55
|
|
|
|
When implemented naively in PyTorch, computing :code:`y = naive_softmax(x)` for :math:`x \in R^{M \times N}`
|
|
requires reading :math:`5MN + 2M` elements from DRAM and writing back :math:`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 :math:`MN` bytes, so we could
|
|
expect a theoretical speed-up of ~4x (i.e., :math:`(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.
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 57-64
|
|
|
|
Compute Kernel
|
|
----------------
|
|
Our softmax kernel works as follows: each program loads a row of the input matrix X,
|
|
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:
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 64-93
|
|
|
|
.. code-block:: default
|
|
|
|
|
|
|
|
@triton.jit
|
|
def softmax_kernel(
|
|
output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,
|
|
BLOCK_SIZE: tl.constexpr
|
|
):
|
|
# The rows of the softmax are independent, so we parallelize across those
|
|
row_idx = tl.program_id(0)
|
|
# 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
|
|
row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf'))
|
|
# Substract maximum for numerical stability
|
|
row_minus_max = row - tl.max(row, axis=0)
|
|
# Note that exponentials in Triton are 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=col_offsets < n_cols)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 94-95
|
|
|
|
We can create a helper function that enqueues the kernel and its (meta-)arguments for any given input tensor.
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 95-125
|
|
|
|
.. code-block:: default
|
|
|
|
|
|
def softmax(x):
|
|
n_rows, n_cols = x.shape
|
|
# The block size 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 = 4
|
|
if BLOCK_SIZE >= 2048:
|
|
num_warps = 8
|
|
if BLOCK_SIZE >= 4096:
|
|
num_warps = 16
|
|
# Allocate output
|
|
y = torch.empty_like(x)
|
|
# Enqueue kernel. The 1D launch grid is simple: we have one kernel instance per row o
|
|
# f the input matrix
|
|
softmax_kernel[(n_rows,)](
|
|
y,
|
|
x,
|
|
x.stride(0),
|
|
y.stride(0),
|
|
n_cols,
|
|
num_warps=num_warps,
|
|
BLOCK_SIZE=BLOCK_SIZE,
|
|
)
|
|
return y
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 126-128
|
|
|
|
Unit Test
|
|
----------
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 130-132
|
|
|
|
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.
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 132-139
|
|
|
|
.. code-block:: default
|
|
|
|
|
|
torch.manual_seed(0)
|
|
x = torch.randn(1823, 781, device='cuda')
|
|
y_triton = softmax(x)
|
|
y_torch = torch.softmax(x, axis=1)
|
|
assert torch.allclose(y_triton, y_torch), (y_triton, y_torch)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 140-141
|
|
|
|
As expected, the results are identical.
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 143-147
|
|
|
|
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) :code:`torch.softmax` and (2) the :code:`naive_softmax` defined above.
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 147-186
|
|
|
|
.. code-block:: default
|
|
|
|
|
|
|
|
@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-native',
|
|
'torch-jit',
|
|
], # possible values for `line_arg``
|
|
line_names=[
|
|
"Triton",
|
|
"Torch (native)",
|
|
"Torch (jit)",
|
|
], # label name for the lines
|
|
styles=[('blue', '-'), ('green', '-'), ('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='cuda', dtype=torch.float32)
|
|
if provider == 'torch-native':
|
|
ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1))
|
|
if provider == 'triton':
|
|
ms, min_ms, max_ms = triton.testing.do_bench(lambda: softmax(x))
|
|
if provider == 'torch-jit':
|
|
ms, min_ms, max_ms = triton.testing.do_bench(lambda: naive_softmax(x))
|
|
gbps = lambda ms: 2 * x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)
|
|
return gbps(ms), gbps(max_ms), gbps(min_ms)
|
|
|
|
|
|
benchmark.run(show_plots=True, print_data=True)
|
|
|
|
|
|
|
|
|
|
.. image:: /getting-started/tutorials/images/sphx_glr_02-fused-softmax_001.png
|
|
:alt: 02 fused softmax
|
|
:class: sphx-glr-single-img
|
|
|
|
|
|
.. rst-class:: sphx-glr-script-out
|
|
|
|
Out:
|
|
|
|
.. code-block:: none
|
|
|
|
softmax-performance:
|
|
N Triton Torch (native) Torch (jit)
|
|
0 256.0 546.133347 546.133347 188.321838
|
|
1 384.0 614.400016 585.142862 151.703707
|
|
2 512.0 655.360017 606.814814 154.566038
|
|
3 640.0 706.206879 640.000002 160.000000
|
|
4 768.0 722.823517 664.216187 163.839992
|
|
.. ... ... ... ...
|
|
93 12160.0 812.359066 405.333344 198.936606
|
|
94 12288.0 812.429770 415.222812 199.096718
|
|
95 12416.0 812.498981 411.722274 198.904612
|
|
96 12544.0 810.925276 412.971190 199.012395
|
|
97 12672.0 811.007961 412.097543 199.167004
|
|
|
|
[98 rows x 4 columns]
|
|
|
|
|
|
|
|
|
|
.. GENERATED FROM PYTHON SOURCE LINES 187-192
|
|
|
|
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 :code:`torch.softmax` -- in addition to being **easier to read, understand and maintain**.
|
|
Note however that the PyTorch `softmax` operation is more general and will works on tensors of any shape.
|
|
|
|
|
|
.. rst-class:: sphx-glr-timing
|
|
|
|
**Total running time of the script:** ( 3 minutes 32.289 seconds)
|
|
|
|
|
|
.. _sphx_glr_download_getting-started_tutorials_02-fused-softmax.py:
|
|
|
|
|
|
.. only :: html
|
|
|
|
.. container:: sphx-glr-footer
|
|
:class: sphx-glr-footer-example
|
|
|
|
|
|
|
|
.. container:: sphx-glr-download sphx-glr-download-python
|
|
|
|
:download:`Download Python source code: 02-fused-softmax.py <02-fused-softmax.py>`
|
|
|
|
|
|
|
|
.. container:: sphx-glr-download sphx-glr-download-jupyter
|
|
|
|
:download:`Download Jupyter notebook: 02-fused-softmax.ipynb <02-fused-softmax.ipynb>`
|
|
|
|
|
|
.. only:: html
|
|
|
|
.. rst-class:: sphx-glr-signature
|
|
|
|
`Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_
|