triton/master/_downloads/034d953b6214fedce6ea03803c712b89/02-fused-softmax.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Fused Softmax\nIn this tutorial, you will write a fused softmax operation that is significantly faster\nthan PyTorch's native op for a particular class of matrices: those whose rows can fit in\nthe GPU's SRAM.\nYou will learn about:\n\n- The benefits of kernel fusion for bandwidth-bound operations.\n- Reduction operators in Triton.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Motivations\nCustom GPU kernels for elementwise additions are educationally valuable but won't get you very far in practice.\nLet us consider instead the case of a simple (numerically stabilized) softmax operation:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\n\nimport triton\nimport triton.language as tl\n\n\n@torch.jit.script\ndef naive_softmax(x):\n    \"\"\"Compute row-wise softmax of X using native pytorch\n\n    We subtract the maximum element in order to avoid overflows. Softmax is invariant to\n    this shift.\n    \"\"\"\n    # read  MN elements ; write M  elements\n    x_max = x.max(dim=1)[0]\n    # read MN + M elements ; write MN elements\n    z = x - x_max[:, None]\n    # read  MN elements ; write MN elements\n    numerator = torch.exp(z)\n    # read  MN elements ; write M  elements\n    denominator = numerator.sum(dim=1)\n    # read MN + M elements ; write MN elements\n    ret = numerator / denominator[:, None]\n    # in total: read 5MN + 2M elements ; wrote 3MN + 2M elements\n    return ret"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "When implemented naively in PyTorch, computing :code:`y = naive_softmax(x)` for $x \\in R^{M \\times N}$\nrequires reading $5MN + 2M$ elements from DRAM and writing back $3MN + 2M$ elements.\nThis is obviously wasteful; we'd prefer to have a custom \"fused\" kernel that only reads\nX once and does all the necessary computations on-chip.\nDoing so would require reading and writing back only $MN$ bytes, so we could\nexpect a theoretical speed-up of ~4x (i.e., $(8MN + 4M) / 2MN$).\nThe `torch.jit.script` flags aims to perform this kind of \"kernel fusion\" automatically\nbut, as we will see later, it is still far from ideal.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute Kernel\nOur softmax kernel works as follows: each program loads a row of the input matrix X,\nnormalizes it and writes back the result to the output Y.\nNote that one important limitation of Triton is that each block must have a\npower-of-two number of elements, so we need to internally \"pad\" each row and guard the\nmemory operations properly if we want to handle any possible input shapes:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "@triton.jit\ndef softmax_kernel(\n    output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,\n    BLOCK_SIZE: tl.constexpr\n):\n    # The rows of the softmax are independent, so we parallelize across those\n    row_idx = tl.program_id(0)\n    # The stride represents how much we need to increase the pointer to advance 1 row\n    row_start_ptr = input_ptr + row_idx * input_row_stride\n    # The block size is the next power of two greater than n_cols, so we can fit each\n    # row in a single block\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    input_ptrs = row_start_ptr + col_offsets\n    # Load the row into SRAM, using a mask since BLOCK_SIZE may be > than n_cols\n    row = tl.load(input_ptrs, mask=col_offsets < n_cols, other=-float('inf'))\n    # Substract maximum for numerical stability\n    row_minus_max = row - tl.max(row, axis=0)\n    # Note that exponentials in Triton are fast but approximate (i.e., think __expf in CUDA)\n    numerator = tl.exp(row_minus_max)\n    denominator = tl.sum(numerator, axis=0)\n    softmax_output = numerator / denominator\n    # Write back output to DRAM\n    output_row_start_ptr = output_ptr + row_idx * output_row_stride\n    output_ptrs = output_row_start_ptr + col_offsets\n    tl.store(output_ptrs, softmax_output, mask=col_offsets < n_cols)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can create a helper function that enqueues the kernel and its (meta-)arguments for any given input tensor.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def softmax(x):\n    n_rows, n_cols = x.shape\n    # The block size is the smallest power of two greater than the number of columns in `x`\n    BLOCK_SIZE = triton.next_power_of_2(n_cols)\n    # Another trick we can use is to ask the compiler to use more threads per row by\n    # increasing the number of warps (`num_warps`) over which each row is distributed.\n    # You will see in the next tutorial how to auto-tune this value in a more natural\n    # way so you don't have to come up with manual heuristics yourself.\n    num_warps = 4\n    if BLOCK_SIZE >= 2048:\n        num_warps = 8\n    if BLOCK_SIZE >= 4096:\n        num_warps = 16\n    # Allocate output\n    y = torch.empty_like(x)\n    # Enqueue kernel. The 1D launch grid is simple: we have one kernel instance per row o\n    # f the input matrix\n    softmax_kernel[(n_rows,)](\n        y,\n        x,\n        x.stride(0),\n        y.stride(0),\n        n_cols,\n        num_warps=num_warps,\n        BLOCK_SIZE=BLOCK_SIZE,\n    )\n    return y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Unit Test\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We make sure that we test our kernel on a matrix with an irregular number of rows and columns.\nThis will allow us to verify that our padding mechanism works.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch.manual_seed(0)\nx = torch.randn(1823, 781, device='cuda')\ny_triton = softmax(x)\ny_torch = torch.softmax(x, axis=1)\nassert torch.allclose(y_triton, y_torch), (y_triton, y_torch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As expected, the results are identical.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Benchmark\nHere we will benchmark our operation as a function of the number of columns in the input matrix -- assuming 4096 rows.\nWe will then compare its performance against (1) :code:`torch.softmax` and (2) the :code:`naive_softmax` defined above.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "@triton.testing.perf_report(\n    triton.testing.Benchmark(\n        x_names=['N'],  # argument names to use as an x-axis for the plot\n        x_vals=[\n            128 * i for i in range(2, 100)\n        ],  # different possible values for `x_name`\n        line_arg='provider',  # argument name whose value corresponds to a different line in the plot\n        line_vals=[\n            'triton',\n            'torch-native',\n            'torch-jit',\n        ],  # possible values for `line_arg``\n        line_names=[\n            \"Triton\",\n            \"Torch (native)\",\n            \"Torch (jit)\",\n        ],  # label name for the lines\n        styles=[('blue', '-'), ('green', '-'), ('green', '--')],  # line styles\n        ylabel=\"GB/s\",  # label name for the y-axis\n        plot_name=\"softmax-performance\",  # name for the plot. Used also as a file name for saving the plot.\n        args={'M': 4096},  # values for function arguments not in `x_names` and `y_name`\n    )\n)\ndef benchmark(M, N, provider):\n    x = torch.randn(M, N, device='cuda', dtype=torch.float32)\n    if provider == 'torch-native':\n        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1))\n    if provider == 'triton':\n        ms, min_ms, max_ms = triton.testing.do_bench(lambda: softmax(x))\n    if provider == 'torch-jit':\n        ms, min_ms, max_ms = triton.testing.do_bench(lambda: naive_softmax(x))\n    gbps = lambda ms: 2 * x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)\n    return gbps(ms), gbps(max_ms), gbps(min_ms)\n\n\nbenchmark.run(show_plots=True, print_data=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In the above plot, we can see that:\n\n - Triton is 4x faster than the Torch JIT. This confirms our suspicions that the Torch JIT does not do any fusion here.\n - Triton is noticeably faster than :code:`torch.softmax` -- in addition to being **easier to read, understand and maintain**.\n   Note however that the PyTorch `softmax` operation is more general and will works on tensors of any shape.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}