triton/_sources/getting-started/tutorials/02-fused-softmax.rst.txt

.. 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-43

.. code-block:: default


    import torch


    @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 2MN 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 2MN elements ; write MN elements
        ret = numerator / denominator[:, None]
        # in total: read 7MN elements ; wrote 3MN + 2M elements
        return ret









.. GENERATED FROM PYTHON SOURCE LINES 44-52

When implemented naively in PyTorch, computing :code:`y = naive_softmax(x)` for :math:`x \in R^{M \times N}`
requires reading :math:`7MN` 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 ~5x (i.e., :math:`(10MN + 2M) / 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 54-61

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 61-94

.. code-block:: default


    import triton
    import triton.language as tl


    @triton.jit
    def softmax_kernel(
        output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, **meta
    ):
        # The rows of the softmax are independent, so we parallelize across those
        row_idx = tl.program_id(0)
        BLOCK_SIZE = meta['BLOCK_SIZE']
        # 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 95-96

We can create a helper function that enqueues the kernel and its (meta-)arguments for any given input tensor.

.. GENERATED FROM PYTHON SOURCE LINES 96-139

.. code-block:: default



    def next_power_of_2(n):
        """Return the smallest power of 2 greater than or equal to n"""
        n -= 1
        n |= n >> 1
        n |= n >> 2
        n |= n >> 4
        n |= n >> 8
        n |= n >> 16
        n += 1
        return n


    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 = 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 140-142

Unit Test
----------

.. GENERATED FROM PYTHON SOURCE LINES 144-146

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 146-153

.. 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)
    print(torch.allclose(y_triton, y_torch))





.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none

    True




.. GENERATED FROM PYTHON SOURCE LINES 154-155

As expected, the results are identical.

.. GENERATED FROM PYTHON SOURCE LINES 157-161

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 161-200

.. 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  512.000001      546.133347   186.181817
    1     384.0  585.142862      585.142862   153.600004
    2     512.0  630.153853      606.814814   154.566038
    3     640.0  660.645170      640.000002   160.000000
    4     768.0  702.171410      664.216187   163.839992
    ..      ...         ...             ...          ...
    93  12160.0  812.359066      406.179533   199.038365
    94  12288.0  812.429770      415.222812   199.298541
    95  12416.0  810.840807      412.149375   198.854847
    96  12544.0  810.925276      412.971190   199.209928
    97  12672.0  809.389265      412.097543   199.167004

    [98 rows x 4 columns]




.. GENERATED FROM PYTHON SOURCE LINES 201-207

In the above plot, we can see that:

 - Triton is 2-3x faster than the Torch JIT.
 - Triton is even faster than :code:`torch.softmax`. My guess from looking at the source-code of the `PyTorch kernel <https://github.com/pytorch/pytorch/blob/9409a3a39b7149bb2d833a89e0c944109bef7c27/caffe2/operators/softmax_ops.cu#L240>`_ is that PyTorch only partially fuses the computation of the softmax.
   This means that -- when temporary data is too large to fit entirely in the GPU's cache -- it transfers almost twice the amount of memory necessary.
   Note that our Triton kernel is not only faster than PyTorch's CUDA kernel, it is also **easier to read, understand and maintain**.


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 1 minutes  13.186 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>`_