[GH-PAGES] Updated website

2021-03-06 22:06:32 -05:00
parent 6f789b29ab
commit 32aaf8b469
17 changed files with 400 additions and 275 deletions
--- a/_downloads/034d953b6214fedce6ea03803c712b89/02-fused-softmax.ipynb
+++ b/_downloads/034d953b6214fedce6ea03803c712b89/02-fused-softmax.ipynb
@@ -15,7 +15,7 @@
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "\n# Fused Softmax\nIn this tutorial, you will write a fused softmax layer that outperform's PyTorch implementation and learn about:\n\n- The benefits of kernel fusion for bandwidth-bound operations.\n- The syntax and usage of reduction operators in Triton.\n- The automatic vectorization capabilities of the Triton compiler.\n"
+        "\n# Fused Softmax\nIn this tutorial, you will write a fused softmax operation (that outperforms PyTorch) and learn about:\n\n- The benefits of kernel fusion for bandwidth-bound operations.\n- The syntax and usage of reduction operators in Triton.\n- The automatic vectorization capabilities of the Triton compiler.\n"
      ]
    },
    {
@@ -40,21 +40,21 @@
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "When implemented naively in pytorch, computing :code:`y = naive_softmax(x)` for $x \\in R^{M \\times N}$ requires reading $7MN$ elements from DRAM and writing back $3MN + 2M$ elements.\nInstead, we want to write a custom \"fused\" pytorch operators that only reads X once and does all the necessary computations on-chip.\nThis would require reading and writing back only $MN$ bytes, so we could expect a theoretical speed-up of 5x.\nIn practice, though, we expect less because our kernel will spend some time computing exponentials and moving data around in shared memory.\n\n"
+        "When implemented naively in pytorch, computing :code:`y = naive_softmax(x)` for $x \\in R^{M \\times N}$ requires reading $7MN$ 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 X once and does all the necessary computations on-chip.\nIn this case, we would be reading and writing back only $MN$ bytes, so we could expect a theoretical speed-up of ~5x (i.e., $(10MN + 2M) / 2MN$).\nIn practice, though, we would be getting a bit less as our kernel computes exponentials and internally moves data around in shared memory.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "## Compute Kernel\nOur softmax kernel works as follows: each program loads a row of X and writes back a normalized row of Y. Note that one important limitation of Triton is that each block must have a power-of-two number of elements, which means that we need to guard the memory operations properly if we want to handle any possible input shapes:\n\n .. code-block:: C\n\n   __global__ void softmax(float* Y, float* X, int stride_xm, int stride_ym, int M, int N){\n     // row index\n     int    m             = get_program_id(0);\n     // column indices\n     int    n    [BLOCK] = 0 ... BLOCK;\n     // the memory address of all the elements\n     // that we want to load can be computed as follows\n     float* px   [BLOCK] = X + m*stride_xm + n;\n     // because BLOCK has to be a power of two\n     // (per Triton-C specs), it is important\n     // to guard each memory operation with predicates\n     // or we will read out of bounds\n     bool   check[BLOCK] = n < N;\n     float  x    [BLOCK] = check ? *px : -F32_INFINITY;\n     // syntax for reduction in Triton is:\n     // x[..., OPERATOR, ...]\n     //            ^\n     //           index\n     // The operators currently supported are {min, max, +}\n     float  z    [BLOCK] = x - x[max];\n     // The exponential in Triton is fast but approximate\n     // (i.e., like __expf in CUDA)\n     float  num  [BLOCK] = exp(z);\n     float  denom         = num[+];\n     // The result of the reduction is now stored in y\n     float  y    [BLOCK] = num / denom;\n     // We write it back\n     float* py   [BLOCK] = Y + m*stride_ym + n;\n     *?(check)py = y;\n   }\n\n"
+        "## Compute Kernel\nOur softmax kernel works as follows: each program loads a row of the input X, normalizes it and writes back the result to the output Y.\nNote that one important limitation of Triton is that each block must have a power-of-two number of elements,\nso we need to internally \"pad\" tiles and guard the memory operations properly if we want to handle any possible input shapes:\n\n .. code-block:: C\n\n   __global__ void softmax(float* Y, float* X, int stride_xm, int stride_ym, int M, int N){\n     // row index\n     int    m             = get_program_id(0);\n     // column indices\n     int    n    [BLOCK] = 0 ... BLOCK;\n     // the memory address of all the elements\n     // that we want to load can be computed as follows\n     float* px   [BLOCK] = X + m*stride_xm + n;\n     // because BLOCK has to be a power of two\n     // (per Triton-C specs), it is important\n     // to guard each memory operation with predicates\n     // or we will read out of bounds\n     bool   check[BLOCK] = n < N;\n     float  x    [BLOCK] = check ? *px : -F32_INFINITY;\n     // syntax for reduction in Triton is:\n     // x[:, :, OPERATOR, :, :]\n     //            ^\n     //           index\n     // where operator is in {min, max, +}\n     // for 1D vectors, this is just x[OPERATOR].\n     float  z    [BLOCK] = x - x[max];\n     // Note that exponentials in Triton are fast\n     // but approximate (i.e., think __expf in CUDA)\n     float  num  [BLOCK] = exp(z);\n     float  denom         = num[+];\n     // The result of the reduction is now stored in y\n     float  y    [BLOCK] = num / denom;\n     // We write it back\n     float* py   [BLOCK] = Y + m*stride_ym + n;\n     *?(check)py = y;\n   }\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "## Torch Bindings\nWe need to make sure that BLOCK is the smallest power of two\ngreater than the number of rows N of the input matrix.\nDifferent values of BLOCK will result in different kernels\n\n"
+        "## Torch Bindings\nHere our torch bindings is quite similar to that of the vector addition mentioned in the previous tutorial.\nWe just need to make sure that BLOCK is the smallest power of two greater than the number of columns N of the input matrix.\nThis means that different values of BLOCK will result in different kernels\n\n"
      ]
    },
    {
@@ -65,7 +65,14 @@
      },
      "outputs": [],
      "source": [
-        "import torch\nimport triton\n\n# Source code for the Triton kernel\n_src = \"\"\"\n__global__ void softmax(float* Y, float* X, int stride_ym, int stride_xm, int M, int N){\n    int    m             = get_program_id(0);\n    int    n    [BLOCK] = 0 ... BLOCK;\n    float* px   [BLOCK] = X + m*stride_xm + n;\n    bool   check[BLOCK] = n < N;\n    float  x    [BLOCK] = check ? *px : -F32_INFINITY;\n    float  z    [BLOCK] = x - x[max];\n    float  num  [BLOCK] = exp(z);\n    float  denom        = num[+];\n    float  y    [BLOCK] = num / denom;\n    float* py   [BLOCK] = Y + m*stride_ym + n;\n    *?(check)py = y; \n}\n\"\"\"\n\n\ndef next_power_of_2(n):\n    n -= 1\n    n |= n >> 1\n    n |= n >> 2\n    n |= n >> 4\n    n |= n >> 8\n    n |= n >> 16\n    n += 1\n    return n\n\n\n_kernels = dict()\n\n\ndef make_kernel(N, device):\n    BLOCK = next_power_of_2(N)\n    key = (BLOCK, device)\n    if key not in _kernels:\n        defines = {'BLOCK': BLOCK}\n        _kernels[key] = triton.kernel(_src, device=device, defines=defines)\n    return _kernels[key]\n\n\nclass _softmax(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x):\n        # constraints of the op\n        assert x.dtype == torch.float32\n        y = torch.empty_like(x)\n        # *create launch grid*:\n        # here we just launch a grid of M programs\n        M, N = y.shape\n        grid = lambda opt: (M, )\n        # *launch kernel*:\n        kernel = make_kernel(N, y.device)\n        kernel(y.data_ptr(), x.data_ptr(), y.stride(0), x.stride(0), M, N, grid=grid)\n        return y\n\n\nsoftmax = _softmax.apply"
+        "import torch\nimport triton\n\n# Source code for the Triton kernel\n_src = \"\"\"\n__global__ void softmax(float* Y, float* X, int stride_ym, int stride_xm, int M, int N){\n    int    m             = get_program_id(0);\n    int    n    [BLOCK] = 0 ... BLOCK;\n    float* px   [BLOCK] = X + m*stride_xm + n;\n    bool   check[BLOCK] = n < N;\n    float  x    [BLOCK] = check ? *px : -F32_INFINITY;\n    float  z    [BLOCK] = x - x[max];\n    float  num  [BLOCK] = exp(z);\n    float  denom        = num[+];\n    float  y    [BLOCK] = num / denom;\n    float* py   [BLOCK] = Y + m*stride_ym + n;\n    *?(check)py = y; \n}\n\"\"\"\n\n\n# helper function to get the smaller power-of-two larger than a given number\ndef next_power_of_2(n):\n    n -= 1\n    n |= n >> 1\n    n |= n >> 2\n    n |= n >> 4\n    n |= n >> 8\n    n |= n >> 16\n    n += 1\n    return n\n\n\n# kernel caching mechanism\ndef make_kernel(N, device):\n    cache = make_kernel.cache\n    # Now are kernels are indexed not only by the provided device but also\n    # by the rounded number of columns in the input matrix\n    BLOCK = next_power_of_2(N)\n    key = (BLOCK, device)\n    if key not in cache:\n        defines = {'BLOCK': BLOCK}\n        cache[key] = triton.kernel(_src, device=device, defines=defines)\n    return cache[key]\n\n\nmake_kernel.cache = dict()\n\n\nclass _softmax(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x):\n        # constraints of the op\n        assert x.dtype == torch.float32\n        y = torch.empty_like(x)\n        # The launch grid is simple: we have one kernel instance per row of the input matrix\n        M, N = y.shape\n        grid = lambda opt: (M, )\n        # Launch kernel\n        kernel = make_kernel(N, y.device)\n        kernel(y.data_ptr(), x.data_ptr(), y.stride(0), x.stride(0), M, N, grid=grid)\n        return y\n\n\nsoftmax = _softmax.apply"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We can use the above softmax function to compute the row-wise softmax of a given matrix.\n\n"
      ]
    },
    {
@@ -75,29 +82,11 @@
        "## Unit Test\n\n"
      ]
    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "collapsed": false
-      },
-      "outputs": [],
-      "source": [
-        "x = torch.randn(1823, 781, device='cuda')\ny_tri = softmax(x)\ny_ref = torch.softmax(x, axis=1)\nprint(y_tri)\nprint(y_ref)\nprint(torch.allclose(y_tri, y_ref))"
-      ]
-    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "Seems to work!\n\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Benchmarking\n\n"
+        "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"
      ]
    },
    {
@@ -108,7 +97,39 @@
      },
      "outputs": [],
      "source": [
-        "import matplotlib.pyplot as plt\n\nM = 4096\nNs = [128 * i for i in range(2, 50)]\ntri_ms = []\nref_ms = []\ndef_ms = []\nfor N in Ns:\n    x = torch.randn(M, N, device='cuda', dtype=torch.float32)\n    gbps = lambda ms: x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)\n    tri_ms += [gbps(triton.testing.do_bench(lambda: softmax(x)))]\n    ref_ms += [gbps(triton.testing.do_bench(lambda: torch.softmax(x, axis=1)))]\n    def_ms += [gbps(triton.testing.do_bench(lambda: naive_softmax(x)))]\nplt.xlabel('N')\nplt.ylabel('Bandwidth (GB/s)')\nplt.plot(Ns, tri_ms, label='Triton')\nplt.plot(Ns, ref_ms, label='Torch')\nplt.plot(Ns, def_ms, label='Naive')\nplt.legend()\nplt.show()"
+        "torch.manual_seed(0)\nx = torch.randn(1823, 781, device='cuda')\ny_tri = softmax(x)\ny_ref = torch.softmax(x, axis=1)\nprint(torch.allclose(y_tri, y_ref))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "As expected, the results are identical.\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Benchmarking\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": [
+        "import matplotlib.pyplot as plt\n\nM = 4096\nNs = [256 * i for i in range(2, 50)]\ntri_bw = []\nref_bw = []\ndef_bw = []\nfor N in Ns:\n    x = torch.randn(M, N, device='cuda', dtype=torch.float32)\n    gbps = lambda ms: x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)\n    do_bench = lambda fn: gbps(triton.testing.do_bench(fn, warmup=10, rep=100, clear_l2=True))\n    tri_bw += [do_bench(lambda: softmax(x))]\n    ref_bw += [do_bench(lambda: torch.softmax(x, axis=1))]\n    def_bw += [do_bench(lambda: naive_softmax(x))]\nplt.xlabel('N')\nplt.ylabel('Bandwidth (GB/s)')\nplt.plot(Ns, tri_bw, label='Triton')\nplt.plot(Ns, ref_bw, label='Torch')\nplt.plot(Ns, def_bw, label='Naive')\nplt.legend()\nplt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "In the above plot, we can see that:\n\n - Triton is 4-5x faster than the naive implementation, which is consistent with our theoretical predictions.\n - Triton is significantly faster than :code:`torch.softmax` for very large input matrices. 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.\n   This means that -- when temporary data is too large to fit entirely in the GPU's cache -- it transfers almost twice the amount of data necessary.\n   Note that our Triton kernel is not only faster than PyTorch's CUDA kernel, it is also **easier to read, understand and maintain**.\n"
      ]
    }
  ],
--- a/_downloads/62d97d49a32414049819dd8bb8378080/01-vector-add.py
+++ b/_downloads/62d97d49a32414049819dd8bb8378080/01-vector-add.py
@@ -1,7 +1,7 @@
 """
 Vector Addition
 =================
-In this tutorial, you will write a simple, high-performance vector addition using Triton and learn about:
+In this tutorial, you will write a simple vector addition using Triton and learn about:

 - The basic syntax of the Triton programming language
 - The best practices for creating PyTorch custom operators using the :code:`triton.kernel` Python API
@@ -122,9 +122,15 @@ class _add(torch.autograd.Function):
 # Just like we standard PyTorch ops We use the :code:`.apply` method to create a callable object for our function
 add = _add.apply

+# %%
+# We can now use the above function to compute the sum of two `torch.tensor` objects:
+
 # %%
 # Unit Test
 # --------------------------
+#
+# Of course, the first thing that we should check is that whether kernel is correct. This is pretty easy to test, as shown below:
+
 torch.manual_seed(0)
 x = torch.rand(98432, device='cuda')
 y = torch.rand(98432, device='cuda')
@@ -134,17 +140,40 @@ print(za)
 print(zb)
 print(f'The maximum difference between torch and triton is ' f'{torch.max(torch.abs(za - zb))}')

+# %%
+# Seems like we're good to go!
+
 # %%
 # Benchmarking
 # --------------------------
-# We can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does
+# We can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does relative to PyTorch.

-warmup = 10
-rep = 200
-for N in [2**i for i in range(17, 26, 1)]:
-    x = torch.rand(N, device='cuda')
-    y = torch.rand(N, device='cuda')
-    triton_ms = triton.testing.do_bench(lambda: add(x, y), warmup=warmup, rep=rep)
-    torch_ms = triton.testing.do_bench(lambda: x + y, warmup=warmup, rep=rep)
-    # print the performance of triton and torch as well as the achieved bandwidth
-    print(f'{N} {triton_ms:.3f} {torch_ms:.3f}')
+import matplotlib.pyplot as plt
+
+# There are three tensors of 4N bytes each. So the bandwidth of a given kernel
+# is 12N / time_ms * 1e-6 GB/s
+gbps = lambda N, ms: 12 * N / ms * 1e-6
+# We want to benchmark small and large vector alike
+sizes = [2**i for i in range(12, 25, 1)]
+triton_bw = []
+torch_bw = []
+for N in sizes:
+    x = torch.rand(N, device='cuda', dtype=torch.float32)
+    y = torch.rand(N, device='cuda', dtype=torch.float32)
+    # Triton provide a do_bench utility function that can be used to benchmark
+    # arbitrary workloads. It supports a `warmup` parameter that is used to stabilize
+    # GPU clock speeds as well as a `rep` parameter that controls the number of times
+    # the benchmark is repeated. Importantly, we set `clear_l2 = True` to make sure
+    # that the L2 cache does not contain any element of x before each kernel call when
+    # N is small.
+    do_bench = lambda fn: gbps(N, triton.testing.do_bench(fn, warmup=10, rep=100, clear_l2=True))
+    triton_bw += [do_bench(lambda: add(x, y))]
+    torch_bw += [do_bench(lambda: x + y)]
+# We plot the results as a semi-log
+plt.semilogx(sizes, triton_bw, label='Triton')
+plt.semilogx(sizes, torch_bw, label='Torch')
+plt.legend()
+plt.show()
+
+# %%
+# Seems like our simple element-wise operation operates at peak bandwidth. While this is a fairly low bar for a custom GPU programming language, this is a good start before we move to more advanced operations.
--- a/_downloads/662999063954282841dc90b8945f85ce/tutorials_jupyter.zip
+++ b/_downloads/662999063954282841dc90b8945f85ce/tutorials_jupyter.zip
--- a/_downloads/763344228ae6bc253ed1a6cf586aa30d/tutorials_python.zip
+++ b/_downloads/763344228ae6bc253ed1a6cf586aa30d/tutorials_python.zip
--- a/_downloads/d91442ac2982c4e0cc3ab0f43534afbc/02-fused-softmax.py
+++ b/_downloads/d91442ac2982c4e0cc3ab0f43534afbc/02-fused-softmax.py
@@ -1,7 +1,7 @@
 """
 Fused Softmax
 =================
-In this tutorial, you will write a fused softmax layer that outperform's PyTorch implementation and learn about:
+In this tutorial, you will write a fused softmax operation (that outperforms PyTorch) and learn about:

 - The benefits of kernel fusion for bandwidth-bound operations.
 - The syntax and usage of reduction operators in Triton.
@@ -35,14 +35,16 @@ def naive_softmax(x):

 # %%
 # 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.
-# Instead, we want to write a custom "fused" pytorch operators that only reads X once and does all the necessary computations on-chip.
-# This would require reading and writing back only :math:`MN` bytes, so we could expect a theoretical speed-up of 5x.
-# In practice, though, we expect less because our kernel will spend some time computing exponentials and moving data around in shared memory.
+# 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.
+# In this case, we would be reading and writing back only :math:`MN` bytes, so we could expect a theoretical speed-up of ~5x (i.e., :math:`(10MN + 2M) / 2MN`).
+# In practice, though, we would be getting a bit less as our kernel computes exponentials and internally moves data around in shared memory.

 # %%
 # Compute Kernel
-# ----------------------------
-# Our softmax kernel works as follows: each program loads a row of X and writes back a normalized row of Y. Note that one important limitation of Triton is that each block must have a power-of-two number of elements, which means that we need to guard the memory operations properly if we want to handle any possible input shapes:
+# ----------------
+# Our softmax kernel works as follows: each program loads a row of the input 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" tiles and guard the memory operations properly if we want to handle any possible input shapes:
 #
 #  .. code-block:: C
 #
@@ -61,13 +63,14 @@ def naive_softmax(x):
 #      bool   check[BLOCK] = n < N;
 #      float  x    [BLOCK] = check ? *px : -F32_INFINITY;
 #      // syntax for reduction in Triton is:
-#      // x[..., OPERATOR, ...]
+#      // x[:, :, OPERATOR, :, :]
 #      //            ^
 #      //           index
-#      // The operators currently supported are {min, max, +}
+#      // where operator is in {min, max, +}
+#      // for 1D vectors, this is just x[OPERATOR].
 #      float  z    [BLOCK] = x - x[max];
-#      // The exponential in Triton is fast but approximate
-#      // (i.e., like __expf in CUDA)
+#      // Note that exponentials in Triton are fast
+#      // but approximate (i.e., think __expf in CUDA)
 #      float  num  [BLOCK] = exp(z);
 #      float  denom         = num[+];
 #      // The result of the reduction is now stored in y
@@ -79,10 +82,10 @@ def naive_softmax(x):

 # %%
 # Torch Bindings
-# ----------------------------
-# We need to make sure that BLOCK is the smallest power of two
-# greater than the number of rows N of the input matrix.
-# Different values of BLOCK will result in different kernels
+# ---------------
+# Here our torch bindings is quite similar to that of the vector addition mentioned in the previous tutorial.
+# We just need to make sure that BLOCK is the smallest power of two greater than the number of columns N of the input matrix.
+# This means that different values of BLOCK will result in different kernels

 import torch
 import triton
@@ -105,6 +108,7 @@ __global__ void softmax(float* Y, float* X, int stride_ym, int stride_xm, int M,
 """


+# helper function to get the smaller power-of-two larger than a given number
 def next_power_of_2(n):
    n -= 1
    n |= n >> 1
@@ -116,16 +120,20 @@ def next_power_of_2(n):
    return n


-_kernels = dict()
-
-
+# kernel caching mechanism
 def make_kernel(N, device):
+    cache = make_kernel.cache
+    # Now are kernels are indexed not only by the provided device but also
+    # by the rounded number of columns in the input matrix
    BLOCK = next_power_of_2(N)
    key = (BLOCK, device)
-    if key not in _kernels:
+    if key not in cache:
        defines = {'BLOCK': BLOCK}
-        _kernels[key] = triton.kernel(_src, device=device, defines=defines)
-    return _kernels[key]
+        cache[key] = triton.kernel(_src, device=device, defines=defines)
+    return cache[key]
+
+
+make_kernel.cache = dict()


 class _softmax(torch.autograd.Function):
@@ -134,11 +142,10 @@ class _softmax(torch.autograd.Function):
        # constraints of the op
        assert x.dtype == torch.float32
        y = torch.empty_like(x)
-        # *create launch grid*:
-        # here we just launch a grid of M programs
+        # The launch grid is simple: we have one kernel instance per row of the input matrix
        M, N = y.shape
        grid = lambda opt: (M, )
-        # *launch kernel*:
+        # Launch kernel
        kernel = make_kernel(N, y.device)
        kernel(y.data_ptr(), x.data_ptr(), y.stride(0), x.stride(0), M, N, grid=grid)
        return y
@@ -146,41 +153,58 @@ class _softmax(torch.autograd.Function):

 softmax = _softmax.apply

+# %%
+# We can use the above softmax function to compute the row-wise softmax of a given matrix.
+
 # %%
 # 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='cuda')
 y_tri = softmax(x)
 y_ref = torch.softmax(x, axis=1)
-print(y_tri)
-print(y_ref)
 print(torch.allclose(y_tri, y_ref))

-# %%
-# Seems to work!
+#%%
+# As expected, the results are identical.

 # %%
 # Benchmarking
-# ----------
+# -------------
+# 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.

 import matplotlib.pyplot as plt

 M = 4096
-Ns = [128 * i for i in range(2, 50)]
-tri_ms = []
-ref_ms = []
-def_ms = []
+Ns = [256 * i for i in range(2, 50)]
+tri_bw = []
+ref_bw = []
+def_bw = []
 for N in Ns:
    x = torch.randn(M, N, device='cuda', dtype=torch.float32)
    gbps = lambda ms: x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)
-    tri_ms += [gbps(triton.testing.do_bench(lambda: softmax(x)))]
-    ref_ms += [gbps(triton.testing.do_bench(lambda: torch.softmax(x, axis=1)))]
-    def_ms += [gbps(triton.testing.do_bench(lambda: naive_softmax(x)))]
+    do_bench = lambda fn: gbps(triton.testing.do_bench(fn, warmup=10, rep=100, clear_l2=True))
+    tri_bw += [do_bench(lambda: softmax(x))]
+    ref_bw += [do_bench(lambda: torch.softmax(x, axis=1))]
+    def_bw += [do_bench(lambda: naive_softmax(x))]
 plt.xlabel('N')
 plt.ylabel('Bandwidth (GB/s)')
-plt.plot(Ns, tri_ms, label='Triton')
-plt.plot(Ns, ref_ms, label='Torch')
-plt.plot(Ns, def_ms, label='Naive')
+plt.plot(Ns, tri_bw, label='Triton')
+plt.plot(Ns, ref_bw, label='Torch')
+plt.plot(Ns, def_bw, label='Naive')
 plt.legend()
-plt.show()
+plt.show()
+
+# %%
+# In the above plot, we can see that:
+#
+#  - Triton is 4-5x faster than the naive implementation, which is consistent with our theoretical predictions.
+#  - Triton is significantly faster than :code:`torch.softmax` for very large input matrices. 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 data necessary.
+#    Note that our Triton kernel is not only faster than PyTorch's CUDA kernel, it is also **easier to read, understand and maintain**.
--- a/_downloads/f191ee1e78dc52eb5f7cba88f71cef2f/01-vector-add.ipynb
+++ b/_downloads/f191ee1e78dc52eb5f7cba88f71cef2f/01-vector-add.ipynb
@@ -15,7 +15,7 @@
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "\n# Vector Addition\nIn this tutorial, you will write a simple, high-performance vector addition using Triton and learn about:\n\n- The basic syntax of the Triton programming language\n- The best practices for creating PyTorch custom operators using the :code:`triton.kernel` Python API\n- The best practices for validating and benchmarking custom ops against native reference implementations\n"
+        "\n# Vector Addition\nIn this tutorial, you will write a simple vector addition using Triton and learn about:\n\n- The basic syntax of the Triton programming language\n- The best practices for creating PyTorch custom operators using the :code:`triton.kernel` Python API\n- The best practices for validating and benchmarking custom ops against native reference implementations\n"
      ]
    },
    {
@@ -47,7 +47,14 @@
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "## Unit Test\n\n"
+        "We can now use the above function to compute the sum of two `torch.tensor` objects:\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Unit Test\n\nOf course, the first thing that we should check is that whether kernel is correct. This is pretty easy to test, as shown below:\n\n"
      ]
    },
    {
@@ -65,7 +72,14 @@
      "cell_type": "markdown",
      "metadata": {},
      "source": [
-        "## Benchmarking\nWe can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does\n\n"
+        "Seems like we're good to go!\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Benchmarking\nWe can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does relative to PyTorch.\n\n"
      ]
    },
    {
@@ -76,7 +90,14 @@
      },
      "outputs": [],
      "source": [
-        "warmup = 10\nrep = 200\nfor N in [2**i for i in range(17, 26, 1)]:\n    x = torch.rand(N, device='cuda')\n    y = torch.rand(N, device='cuda')\n    triton_ms = triton.testing.do_bench(lambda: add(x, y), warmup=warmup, rep=rep)\n    torch_ms = triton.testing.do_bench(lambda: x + y, warmup=warmup, rep=rep)\n    # print the performance of triton and torch as well as the achieved bandwidth\n    print(f'{N} {triton_ms:.3f} {torch_ms:.3f}')"
+        "import matplotlib.pyplot as plt\n\n# There are three tensors of 4N bytes each. So the bandwidth of a given kernel\n# is 12N / time_ms * 1e-6 GB/s\ngbps = lambda N, ms: 12 * N / ms * 1e-6\n# We want to benchmark small and large vector alike\nsizes = [2**i for i in range(12, 25, 1)]\ntriton_bw = []\ntorch_bw = []\nfor N in sizes:\n    x = torch.rand(N, device='cuda', dtype=torch.float32)\n    y = torch.rand(N, device='cuda', dtype=torch.float32)\n    # Triton provide a do_bench utility function that can be used to benchmark\n    # arbitrary workloads. It supports a `warmup` parameter that is used to stabilize\n    # GPU clock speeds as well as a `rep` parameter that controls the number of times\n    # the benchmark is repeated. Importantly, we set `clear_l2 = True` to make sure\n    # that the L2 cache does not contain any element of x before each kernel call when\n    # N is small.\n    do_bench = lambda fn: gbps(N, triton.testing.do_bench(fn, warmup=10, rep=100, clear_l2=True))\n    triton_bw += [do_bench(lambda: add(x, y))]\n    torch_bw += [do_bench(lambda: x + y)]\n# We plot the results as a semi-log\nplt.semilogx(sizes, triton_bw, label='Triton')\nplt.semilogx(sizes, torch_bw, label='Torch')\nplt.legend()\nplt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Seems like our simple element-wise operation operates at peak bandwidth. While this is a fairly low bar for a custom GPU programming language, this is a good start before we move to more advanced operations.\n"
      ]
    }
  ],
--- a/_images/sphx_glr_01-vector-add_001.png
+++ b/_images/sphx_glr_01-vector-add_001.png
--- a/_images/sphx_glr_01-vector-add_thumb.png
+++ b/_images/sphx_glr_01-vector-add_thumb.png
--- a/_images/sphx_glr_02-fused-softmax_001.png
+++ b/_images/sphx_glr_02-fused-softmax_001.png
--- a/_images/sphx_glr_02-fused-softmax_thumb.png
+++ b/_images/sphx_glr_02-fused-softmax_thumb.png
--- a/_sources/getting-started/tutorials/01-vector-add.rst.txt
+++ b/_sources/getting-started/tutorials/01-vector-add.rst.txt
@@ -20,7 +20,7 @@

 Vector Addition
 =================
-In this tutorial, you will write a simple, high-performance vector addition using Triton and learn about:
+In this tutorial, you will write a simple vector addition using Triton and learn about:

 - The basic syntax of the Triton programming language
 - The best practices for creating PyTorch custom operators using the :code:`triton.kernel` Python API
@@ -154,15 +154,22 @@ The only thing that matters when it comes to Triton and Torch is the :code:`trit



-.. GENERATED FROM PYTHON SOURCE LINES 126-128
+.. GENERATED FROM PYTHON SOURCE LINES 126-127
+
+We can now use the above function to compute the sum of two `torch.tensor` objects:
+
+.. GENERATED FROM PYTHON SOURCE LINES 129-133

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

-.. GENERATED FROM PYTHON SOURCE LINES 128-137
+Of course, the first thing that we should check is that whether kernel is correct. This is pretty easy to test, as shown below:
+
+.. GENERATED FROM PYTHON SOURCE LINES 133-143

 .. code-block:: default

+
    torch.manual_seed(0)
    x = torch.rand(98432, device='cuda')
    y = torch.rand(98432, device='cuda')
@@ -189,52 +196,67 @@ Unit Test



-.. GENERATED FROM PYTHON SOURCE LINES 138-141
+.. GENERATED FROM PYTHON SOURCE LINES 144-145
+
+Seems like we're good to go!
+
+.. GENERATED FROM PYTHON SOURCE LINES 147-150

 Benchmarking
 --------------------------
-We can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does
+We can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does relative to PyTorch.

-.. GENERATED FROM PYTHON SOURCE LINES 141-150
+.. GENERATED FROM PYTHON SOURCE LINES 150-178

 .. code-block:: default


-    warmup = 10
-    rep = 200
-    for N in [2**i for i in range(17, 26, 1)]:
-        x = torch.rand(N, device='cuda')
-        y = torch.rand(N, device='cuda')
-        triton_ms = triton.testing.do_bench(lambda: add(x, y), warmup=warmup, rep=rep)
-        torch_ms = triton.testing.do_bench(lambda: x + y, warmup=warmup, rep=rep)
-        # print the performance of triton and torch as well as the achieved bandwidth
-        print(f'{N} {triton_ms:.3f} {torch_ms:.3f}')
+    import matplotlib.pyplot as plt
+
+    # There are three tensors of 4N bytes each. So the bandwidth of a given kernel
+    # is 12N / time_ms * 1e-6 GB/s
+    gbps = lambda N, ms: 12 * N / ms * 1e-6
+    # We want to benchmark small and large vector alike
+    sizes = [2**i for i in range(12, 25, 1)]
+    triton_bw = []
+    torch_bw = []
+    for N in sizes:
+        x = torch.rand(N, device='cuda', dtype=torch.float32)
+        y = torch.rand(N, device='cuda', dtype=torch.float32)
+        # Triton provide a do_bench utility function that can be used to benchmark
+        # arbitrary workloads. It supports a `warmup` parameter that is used to stabilize
+        # GPU clock speeds as well as a `rep` parameter that controls the number of times
+        # the benchmark is repeated. Importantly, we set `clear_l2 = True` to make sure
+        # that the L2 cache does not contain any element of x before each kernel call when
+        # N is small.
+        do_bench = lambda fn: gbps(N, triton.testing.do_bench(fn, warmup=10, rep=100, clear_l2=True))
+        triton_bw += [do_bench(lambda: add(x, y))]
+        torch_bw += [do_bench(lambda: x + y)]
+    # We plot the results as a semi-log
+    plt.semilogx(sizes, triton_bw, label='Triton')
+    plt.semilogx(sizes, torch_bw, label='Torch')
+    plt.legend()
+    plt.show()



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

- Out:
-
- .. code-block:: none
-
-    131072 0.022 0.006
-    262144 0.021 0.005
-    524288 0.022 0.017
-    1048576 0.037 0.037
-    2097152 0.074 0.073
-    4194304 0.144 0.143
-    8388608 0.289 0.285
-    16777216 0.566 0.562
-    33554432 1.131 1.121
+.. image:: /getting-started/tutorials/images/sphx_glr_01-vector-add_001.png
+    :alt: 01 vector add
+    :class: sphx-glr-single-img





+.. GENERATED FROM PYTHON SOURCE LINES 179-179
+
+Seems like our simple element-wise operation operates at peak bandwidth. While this is a fairly low bar for a custom GPU programming language, this is a good start before we move to more advanced operations.
+
+
 .. rst-class:: sphx-glr-timing

-   **Total running time of the script:** ( 0 minutes  3.225 seconds)
+   **Total running time of the script:** ( 0 minutes  4.784 seconds)


 .. _sphx_glr_download_getting-started_tutorials_01-vector-add.py:
--- a/_sources/getting-started/tutorials/02-fused-softmax.rst.txt
+++ b/_sources/getting-started/tutorials/02-fused-softmax.rst.txt
@@ -20,7 +20,7 @@

 Fused Softmax
 =================
-In this tutorial, you will write a fused softmax layer that outperform's PyTorch implementation and learn about:
+In this tutorial, you will write a fused softmax operation (that outperforms PyTorch) and learn about:

 - The benefits of kernel fusion for bandwidth-bound operations.
 - The syntax and usage of reduction operators in Triton.
@@ -67,15 +67,17 @@ Let us consider instead the case of a simple (numerically stabilized) softmax op
 .. GENERATED FROM PYTHON SOURCE LINES 37-41

 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.
-Instead, we want to write a custom "fused" pytorch operators that only reads X once and does all the necessary computations on-chip.
-This would require reading and writing back only :math:`MN` bytes, so we could expect a theoretical speed-up of 5x.
-In practice, though, we expect less because our kernel will spend some time computing exponentials and moving data around in shared memory.
+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.
+In this case, we would be reading and writing back only :math:`MN` bytes, so we could expect a theoretical speed-up of ~5x (i.e., :math:`(10MN + 2M) / 2MN`).
+In practice, though, we would be getting a bit less as our kernel computes exponentials and internally moves data around in shared memory.

-.. GENERATED FROM PYTHON SOURCE LINES 43-79
+.. GENERATED FROM PYTHON SOURCE LINES 43-82

 Compute Kernel
----------------------------
-Our softmax kernel works as follows: each program loads a row of X and writes back a normalized row of Y. Note that one important limitation of Triton is that each block must have a power-of-two number of elements, which means that we need to guard the memory operations properly if we want to handle any possible input shapes:
+----------------
+Our softmax kernel works as follows: each program loads a row of the input 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" tiles and guard the memory operations properly if we want to handle any possible input shapes:

 .. code-block:: C

@@ -94,13 +96,14 @@ Our softmax kernel works as follows: each program loads a row of X and writes ba
     bool   check[BLOCK] = n < N;
     float  x    [BLOCK] = check ? *px : -F32_INFINITY;
     // syntax for reduction in Triton is:
-     // x[..., OPERATOR, ...]
+     // x[:, :, OPERATOR, :, :]
     //            ^
     //           index
-     // The operators currently supported are {min, max, +}
+     // where operator is in {min, max, +}
+     // for 1D vectors, this is just x[OPERATOR].
     float  z    [BLOCK] = x - x[max];
-     // The exponential in Triton is fast but approximate
-     // (i.e., like __expf in CUDA)
+     // Note that exponentials in Triton are fast
+     // but approximate (i.e., think __expf in CUDA)
     float  num  [BLOCK] = exp(z);
     float  denom         = num[+];
     // The result of the reduction is now stored in y
@@ -110,15 +113,15 @@ Our softmax kernel works as follows: each program loads a row of X and writes ba
     *?(check)py = y;
   }

-.. GENERATED FROM PYTHON SOURCE LINES 81-86
+.. GENERATED FROM PYTHON SOURCE LINES 84-89

 Torch Bindings
----------------------------
-We need to make sure that BLOCK is the smallest power of two
-greater than the number of rows N of the input matrix.
-Different values of BLOCK will result in different kernels
+---------------
+Here our torch bindings is quite similar to that of the vector addition mentioned in the previous tutorial.
+We just need to make sure that BLOCK is the smallest power of two greater than the number of columns N of the input matrix.
+This means that different values of BLOCK will result in different kernels

-.. GENERATED FROM PYTHON SOURCE LINES 86-149
+.. GENERATED FROM PYTHON SOURCE LINES 89-156

 .. code-block:: default

@@ -144,6 +147,7 @@ Different values of BLOCK will result in different kernels
    """


+    # helper function to get the smaller power-of-two larger than a given number
    def next_power_of_2(n):
        n -= 1
        n |= n >> 1
@@ -155,16 +159,20 @@ Different values of BLOCK will result in different kernels
        return n


-    _kernels = dict()
-
-
+    # kernel caching mechanism
    def make_kernel(N, device):
+        cache = make_kernel.cache
+        # Now are kernels are indexed not only by the provided device but also
+        # by the rounded number of columns in the input matrix
        BLOCK = next_power_of_2(N)
        key = (BLOCK, device)
-        if key not in _kernels:
+        if key not in cache:
            defines = {'BLOCK': BLOCK}
-            _kernels[key] = triton.kernel(_src, device=device, defines=defines)
-        return _kernels[key]
+            cache[key] = triton.kernel(_src, device=device, defines=defines)
+        return cache[key]
+
+
+    make_kernel.cache = dict()


    class _softmax(torch.autograd.Function):
@@ -173,11 +181,10 @@ Different values of BLOCK will result in different kernels
            # constraints of the op
            assert x.dtype == torch.float32
            y = torch.empty_like(x)
-            # *create launch grid*:
-            # here we just launch a grid of M programs
+            # The launch grid is simple: we have one kernel instance per row of the input matrix
            M, N = y.shape
            grid = lambda opt: (M, )
-            # *launch kernel*:
+            # Launch kernel
            kernel = make_kernel(N, y.device)
            kernel(y.data_ptr(), x.data_ptr(), y.stride(0), x.stride(0), M, N, grid=grid)
            return y
@@ -192,21 +199,29 @@ Different values of BLOCK will result in different kernels



-.. GENERATED FROM PYTHON SOURCE LINES 150-152
+.. GENERATED FROM PYTHON SOURCE LINES 157-158
+
+We can use the above softmax function to compute the row-wise softmax of a given matrix.
+
+.. GENERATED FROM PYTHON SOURCE LINES 160-162

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

-.. GENERATED FROM PYTHON SOURCE LINES 152-160
+.. GENERATED FROM PYTHON SOURCE LINES 164-166
+
+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 166-173

 .. code-block:: default


+    torch.manual_seed(0)
    x = torch.randn(1823, 781, device='cuda')
    y_tri = softmax(x)
    y_ref = torch.softmax(x, axis=1)
-    print(y_tri)
-    print(y_ref)
    print(torch.allclose(y_tri, y_ref))


@@ -219,47 +234,23 @@ Unit Test

 .. code-block:: none

-    tensor([[2.0935e-03, 6.4551e-04, 9.8605e-05,  ..., 3.3981e-04, 2.7386e-03,
-             9.1986e-05],
-            [7.0923e-04, 6.7521e-04, 5.1366e-04,  ..., 9.8392e-04, 2.6547e-04,
-             6.9062e-04],
-            [1.4032e-04, 5.8826e-04, 1.1694e-03,  ..., 6.6423e-04, 1.8178e-04,
-             6.7049e-04],
-            ...,
-            [1.1767e-03, 4.2703e-03, 6.0596e-04,  ..., 9.5274e-04, 1.1681e-03,
-             6.4924e-04],
-            [1.0772e-04, 7.4854e-04, 3.1912e-03,  ..., 2.4980e-04, 1.9012e-03,
-             5.2567e-04],
-            [2.8518e-03, 8.1899e-04, 7.7046e-04,  ..., 1.3403e-03, 5.3167e-04,
-             4.3268e-04]], device='cuda:0')
-    tensor([[2.0935e-03, 6.4551e-04, 9.8605e-05,  ..., 3.3981e-04, 2.7386e-03,
-             9.1986e-05],
-            [7.0923e-04, 6.7521e-04, 5.1366e-04,  ..., 9.8392e-04, 2.6547e-04,
-             6.9062e-04],
-            [1.4032e-04, 5.8826e-04, 1.1694e-03,  ..., 6.6423e-04, 1.8178e-04,
-             6.7049e-04],
-            ...,
-            [1.1767e-03, 4.2703e-03, 6.0596e-04,  ..., 9.5274e-04, 1.1681e-03,
-             6.4924e-04],
-            [1.0772e-04, 7.4854e-04, 3.1912e-03,  ..., 2.4980e-04, 1.9012e-03,
-             5.2567e-04],
-            [2.8518e-03, 8.1899e-04, 7.7046e-04,  ..., 1.3403e-03, 5.3167e-04,
-             4.3268e-04]], device='cuda:0')
    True




-.. GENERATED FROM PYTHON SOURCE LINES 161-162
+.. GENERATED FROM PYTHON SOURCE LINES 174-175

-Seems to work!
+As expected, the results are identical.

-.. GENERATED FROM PYTHON SOURCE LINES 164-166
+.. GENERATED FROM PYTHON SOURCE LINES 177-181

 Benchmarking
----------
+-------------
+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 166-186
+.. GENERATED FROM PYTHON SOURCE LINES 181-204

 .. code-block:: default

@@ -267,25 +258,28 @@ Benchmarking
    import matplotlib.pyplot as plt

    M = 4096
-    Ns = [128 * i for i in range(2, 50)]
-    tri_ms = []
-    ref_ms = []
-    def_ms = []
+    Ns = [256 * i for i in range(2, 50)]
+    tri_bw = []
+    ref_bw = []
+    def_bw = []
    for N in Ns:
        x = torch.randn(M, N, device='cuda', dtype=torch.float32)
        gbps = lambda ms: x.nelement() * x.element_size() * 1e-9 / (ms * 1e-3)
-        tri_ms += [gbps(triton.testing.do_bench(lambda: softmax(x)))]
-        ref_ms += [gbps(triton.testing.do_bench(lambda: torch.softmax(x, axis=1)))]
-        def_ms += [gbps(triton.testing.do_bench(lambda: naive_softmax(x)))]
+        do_bench = lambda fn: gbps(triton.testing.do_bench(fn, warmup=10, rep=100, clear_l2=True))
+        tri_bw += [do_bench(lambda: softmax(x))]
+        ref_bw += [do_bench(lambda: torch.softmax(x, axis=1))]
+        def_bw += [do_bench(lambda: naive_softmax(x))]
    plt.xlabel('N')
    plt.ylabel('Bandwidth (GB/s)')
-    plt.plot(Ns, tri_ms, label='Triton')
-    plt.plot(Ns, ref_ms, label='Torch')
-    plt.plot(Ns, def_ms, label='Naive')
+    plt.plot(Ns, tri_bw, label='Triton')
+    plt.plot(Ns, ref_bw, label='Torch')
+    plt.plot(Ns, def_bw, label='Naive')
    plt.legend()
    plt.show()


+
+
 .. image:: /getting-started/tutorials/images/sphx_glr_02-fused-softmax_001.png
    :alt: 02 fused softmax
    :class: sphx-glr-single-img
@@ -294,10 +288,19 @@ Benchmarking



+.. GENERATED FROM PYTHON SOURCE LINES 205-210
+
+In the above plot, we can see that:
+
+ - Triton is 4-5x faster than the naive implementation, which is consistent with our theoretical predictions.
+ - Triton is significantly faster than :code:`torch.softmax` for very large input matrices. 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 data 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:** ( 0 minutes  5.758 seconds)
+   **Total running time of the script:** ( 0 minutes  33.773 seconds)


 .. _sphx_glr_download_getting-started_tutorials_02-fused-softmax.py:
--- a/_sources/getting-started/tutorials/sg_execution_times.rst.txt
+++ b/_sources/getting-started/tutorials/sg_execution_times.rst.txt
@@ -5,10 +5,10 @@

 Computation times
 =================
-**00:08.983** total execution time for **getting-started_tutorials** files:
+**00:33.773** total execution time for **getting-started_tutorials** files:

 +-----------------------------------------------------------------------------------------+-----------+--------+
-| :ref:`sphx_glr_getting-started_tutorials_02-fused-softmax.py` (``02-fused-softmax.py``) | 00:05.758 | 0.0 MB |
+| :ref:`sphx_glr_getting-started_tutorials_02-fused-softmax.py` (``02-fused-softmax.py``) | 00:33.773 | 0.0 MB |
 +-----------------------------------------------------------------------------------------+-----------+--------+
-| :ref:`sphx_glr_getting-started_tutorials_01-vector-add.py` (``01-vector-add.py``)       | 00:03.225 | 0.0 MB |
+| :ref:`sphx_glr_getting-started_tutorials_01-vector-add.py` (``01-vector-add.py``)       | 00:00.000 | 0.0 MB |
 +-----------------------------------------------------------------------------------------+-----------+--------+
--- a/getting-started/tutorials/01-vector-add.html
+++ b/getting-started/tutorials/01-vector-add.html
@@ -179,7 +179,7 @@ to download the full example code</p>
 </div>
 <div class="sphx-glr-example-title section" id="vector-addition">
 <span id="sphx-glr-getting-started-tutorials-01-vector-add-py"></span><h1>Vector Addition<a class="headerlink" href="#vector-addition" title="Permalink to this headline">¶</a></h1>
-<p>In this tutorial, you will write a simple, high-performance vector addition using Triton and learn about:</p>
+<p>In this tutorial, you will write a simple vector addition using Triton and learn about:</p>
 <ul class="simple">
 <li><p>The basic syntax of the Triton programming language</p></li>
 <li><p>The best practices for creating PyTorch custom operators using the <code class="code docutils literal notranslate"><span class="pre">triton.kernel</span></code> Python API</p></li>
@@ -297,9 +297,11 @@ programming model for more details).</p>
 <span class="n">add</span> <span class="o">=</span> <span class="n">_add</span><span class="o">.</span><span class="n">apply</span>
 </pre></div>
 </div>
+<p>We can now use the above function to compute the sum of two <cite>torch.tensor</cite> objects:</p>
 </div>
 <div class="section" id="unit-test">
 <h2>Unit Test<a class="headerlink" href="#unit-test" title="Permalink to this headline">¶</a></h2>
+<p>Of course, the first thing that we should check is that whether kernel is correct. This is pretty easy to test, as shown below:</p>
 <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
 <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">98432</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">)</span>
 <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">98432</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">)</span>
@@ -316,34 +318,42 @@ tensor([1.3713, 1.3076, 0.4940,  ..., 0.6682, 1.1984, 1.2696], device=&#39;cuda:
 The maximum difference between torch and triton is 0.0
 </pre></div>
 </div>
+<p>Seems like we’re good to go!</p>
 </div>
 <div class="section" id="benchmarking">
 <h2>Benchmarking<a class="headerlink" href="#benchmarking" title="Permalink to this headline">¶</a></h2>
-<p>We can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does</p>
-<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">warmup</span> <span class="o">=</span> <span class="mi">10</span>
-<span class="n">rep</span> <span class="o">=</span> <span class="mi">200</span>
-<span class="k">for</span> <span class="n">N</span> <span class="ow">in</span> <span class="p">[</span><span class="mi">2</span><span class="o">**</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">17</span><span class="p">,</span> <span class="mi">26</span><span class="p">,</span> <span class="mi">1</span><span class="p">)]:</span>
-    <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">)</span>
-    <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">)</span>
-    <span class="n">triton_ms</span> <span class="o">=</span> <span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">add</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">),</span> <span class="n">warmup</span><span class="o">=</span><span class="n">warmup</span><span class="p">,</span> <span class="n">rep</span><span class="o">=</span><span class="n">rep</span><span class="p">)</span>
-    <span class="n">torch_ms</span> <span class="o">=</span> <span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">x</span> <span class="o">+</span> <span class="n">y</span><span class="p">,</span> <span class="n">warmup</span><span class="o">=</span><span class="n">warmup</span><span class="p">,</span> <span class="n">rep</span><span class="o">=</span><span class="n">rep</span><span class="p">)</span>
-    <span class="c1"># print the performance of triton and torch as well as the achieved bandwidth</span>
-    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">N</span><span class="si">}</span><span class="s1"> </span><span class="si">{</span><span class="n">triton_ms</span><span class="si">:</span><span class="s1">.3f</span><span class="si">}</span><span class="s1"> </span><span class="si">{</span><span class="n">torch_ms</span><span class="si">:</span><span class="s1">.3f</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
+<p>We can now benchmark our custom op for vectors of increasing sizes to get a sense of how it does relative to PyTorch.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>
+
+<span class="c1"># There are three tensors of 4N bytes each. So the bandwidth of a given kernel</span>
+<span class="c1"># is 12N / time_ms * 1e-6 GB/s</span>
+<span class="n">gbps</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">N</span><span class="p">,</span> <span class="n">ms</span><span class="p">:</span> <span class="mi">12</span> <span class="o">*</span> <span class="n">N</span> <span class="o">/</span> <span class="n">ms</span> <span class="o">*</span> <span class="mf">1e-6</span>
+<span class="c1"># We want to benchmark small and large vector alike</span>
+<span class="n">sizes</span> <span class="o">=</span> <span class="p">[</span><span class="mi">2</span><span class="o">**</span><span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">12</span><span class="p">,</span> <span class="mi">25</span><span class="p">,</span> <span class="mi">1</span><span class="p">)]</span>
+<span class="n">triton_bw</span> <span class="o">=</span> <span class="p">[]</span>
+<span class="n">torch_bw</span> <span class="o">=</span> <span class="p">[]</span>
+<span class="k">for</span> <span class="n">N</span> <span class="ow">in</span> <span class="n">sizes</span><span class="p">:</span>
+    <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
+    <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
+    <span class="c1"># Triton provide a do_bench utility function that can be used to benchmark</span>
+    <span class="c1"># arbitrary workloads. It supports a `warmup` parameter that is used to stabilize</span>
+    <span class="c1"># GPU clock speeds as well as a `rep` parameter that controls the number of times</span>
+    <span class="c1"># the benchmark is repeated. Importantly, we set `clear_l2 = True` to make sure</span>
+    <span class="c1"># that the L2 cache does not contain any element of x before each kernel call when</span>
+    <span class="c1"># N is small.</span>
+    <span class="n">do_bench</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">fn</span><span class="p">:</span> <span class="n">gbps</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">warmup</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">rep</span><span class="o">=</span><span class="mi">100</span><span class="p">,</span> <span class="n">clear_l2</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
+    <span class="n">triton_bw</span> <span class="o">+=</span> <span class="p">[</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">add</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">))]</span>
+    <span class="n">torch_bw</span> <span class="o">+=</span> <span class="p">[</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">x</span> <span class="o">+</span> <span class="n">y</span><span class="p">)]</span>
+<span class="c1"># We plot the results as a semi-log</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">semilogx</span><span class="p">(</span><span class="n">sizes</span><span class="p">,</span> <span class="n">triton_bw</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Triton&#39;</span><span class="p">)</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">semilogx</span><span class="p">(</span><span class="n">sizes</span><span class="p">,</span> <span class="n">torch_bw</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Torch&#39;</span><span class="p">)</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">legend</span><span class="p">()</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
 </pre></div>
 </div>
-<p class="sphx-glr-script-out">Out:</p>
-<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>131072 0.022 0.006
-262144 0.021 0.005
-524288 0.022 0.017
-1048576 0.037 0.037
-2097152 0.074 0.073
-4194304 0.144 0.143
-8388608 0.289 0.285
-16777216 0.566 0.562
-33554432 1.131 1.121
-</pre></div>
-</div>
-<p class="sphx-glr-timing"><strong>Total running time of the script:</strong> ( 0 minutes  3.225 seconds)</p>
+<img alt="01 vector add" class="sphx-glr-single-img" src="../../_images/sphx_glr_01-vector-add_001.png" />
+<p>Seems like our simple element-wise operation operates at peak bandwidth. While this is a fairly low bar for a custom GPU programming language, this is a good start before we move to more advanced operations.</p>
+<p class="sphx-glr-timing"><strong>Total running time of the script:</strong> ( 0 minutes  4.784 seconds)</p>
 <div class="sphx-glr-footer class sphx-glr-footer-example docutils container" id="sphx-glr-download-getting-started-tutorials-01-vector-add-py">
 <div class="sphx-glr-download sphx-glr-download-python docutils container">
 <p><a class="reference download internal" download="" href="../../_downloads/62d97d49a32414049819dd8bb8378080/01-vector-add.py"><code class="xref download docutils literal notranslate"><span class="pre">Download</span> <span class="pre">Python</span> <span class="pre">source</span> <span class="pre">code:</span> <span class="pre">01-vector-add.py</span></code></a></p>
--- a/getting-started/tutorials/02-fused-softmax.html
+++ b/getting-started/tutorials/02-fused-softmax.html
@@ -179,7 +179,7 @@ to download the full example code</p>
 </div>
 <div class="sphx-glr-example-title section" id="fused-softmax">
 <span id="sphx-glr-getting-started-tutorials-02-fused-softmax-py"></span><h1>Fused Softmax<a class="headerlink" href="#fused-softmax" title="Permalink to this headline">¶</a></h1>
-<p>In this tutorial, you will write a fused softmax layer that outperform’s PyTorch implementation and learn about:</p>
+<p>In this tutorial, you will write a fused softmax operation (that outperforms PyTorch) and learn about:</p>
 <ul class="simple">
 <li><p>The benefits of kernel fusion for bandwidth-bound operations.</p></li>
 <li><p>The syntax and usage of reduction operators in Triton.</p></li>
@@ -209,13 +209,15 @@ Let us consider instead the case of a simple (numerically stabilized) softmax op
 </pre></div>
 </div>
 <p>When implemented naively in pytorch, computing <code class="code docutils literal notranslate"><span class="pre">y</span> <span class="pre">=</span> <span class="pre">naive_softmax(x)</span></code> for <span class="math notranslate nohighlight">\(x \in R^{M \times N}\)</span> requires reading <span class="math notranslate nohighlight">\(7MN\)</span> elements from DRAM and writing back <span class="math notranslate nohighlight">\(3MN + 2M\)</span> elements.
-Instead, we want to write a custom “fused” pytorch operators that only reads X once and does all the necessary computations on-chip.
-This would require reading and writing back only <span class="math notranslate nohighlight">\(MN\)</span> bytes, so we could expect a theoretical speed-up of 5x.
-In practice, though, we expect less because our kernel will spend some time computing exponentials and moving data around in shared memory.</p>
+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.
+In this case, we would be reading and writing back only <span class="math notranslate nohighlight">\(MN\)</span> bytes, so we could expect a theoretical speed-up of ~5x (i.e., <span class="math notranslate nohighlight">\((10MN + 2M) / 2MN\)</span>).
+In practice, though, we would be getting a bit less as our kernel computes exponentials and internally moves data around in shared memory.</p>
 </div>
 <div class="section" id="compute-kernel">
 <h2>Compute Kernel<a class="headerlink" href="#compute-kernel" title="Permalink to this headline">¶</a></h2>
-<p>Our softmax kernel works as follows: each program loads a row of X and writes back a normalized row of Y. Note that one important limitation of Triton is that each block must have a power-of-two number of elements, which means that we need to guard the memory operations properly if we want to handle any possible input shapes:</p>
+<p>Our softmax kernel works as follows: each program loads a row of the input 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” tiles and guard the memory operations properly if we want to handle any possible input shapes:</p>
 <blockquote>
 <div><div class="highlight-C notranslate"><div class="highlight"><pre><span></span><span class="n">__global__</span> <span class="kt">void</span> <span class="n">softmax</span><span class="p">(</span><span class="kt">float</span><span class="o">*</span> <span class="n">Y</span><span class="p">,</span> <span class="kt">float</span><span class="o">*</span> <span class="n">X</span><span class="p">,</span> <span class="kt">int</span> <span class="n">stride_xm</span><span class="p">,</span> <span class="kt">int</span> <span class="n">stride_ym</span><span class="p">,</span> <span class="kt">int</span> <span class="n">M</span><span class="p">,</span> <span class="kt">int</span> <span class="n">N</span><span class="p">){</span>
  <span class="c1">// row index</span>
@@ -232,13 +234,14 @@ In practice, though, we expect less because our kernel will spend some time comp
  <span class="kt">bool</span>   <span class="n">check</span><span class="p">[</span><span class="n">BLOCK</span><span class="p">]</span> <span class="o">=</span> <span class="n">n</span> <span class="o">&lt;</span> <span class="n">N</span><span class="p">;</span>
  <span class="kt">float</span>  <span class="n">x</span>    <span class="p">[</span><span class="n">BLOCK</span><span class="p">]</span> <span class="o">=</span> <span class="n">check</span> <span class="o">?</span> <span class="o">*</span><span class="nl">px</span> <span class="p">:</span> <span class="o">-</span><span class="n">F32_INFINITY</span><span class="p">;</span>
  <span class="c1">// syntax for reduction in Triton is:</span>
-  <span class="c1">// x[..., OPERATOR, ...]</span>
+  <span class="c1">// x[:, :, OPERATOR, :, :]</span>
  <span class="c1">//            ^</span>
  <span class="c1">//           index</span>
-  <span class="c1">// The operators currently supported are {min, max, +}</span>
+  <span class="c1">// where operator is in {min, max, +}</span>
+  <span class="c1">// for 1D vectors, this is just x[OPERATOR].</span>
  <span class="kt">float</span>  <span class="n">z</span>    <span class="p">[</span><span class="n">BLOCK</span><span class="p">]</span> <span class="o">=</span> <span class="n">x</span> <span class="o">-</span> <span class="n">x</span><span class="p">[</span><span class="n">max</span><span class="p">];</span>
-  <span class="c1">// The exponential in Triton is fast but approximate</span>
-  <span class="c1">// (i.e., like __expf in CUDA)</span>
+  <span class="c1">// Note that exponentials in Triton are fast</span>
+  <span class="c1">// but approximate (i.e., think __expf in CUDA)</span>
  <span class="kt">float</span>  <span class="n">num</span>  <span class="p">[</span><span class="n">BLOCK</span><span class="p">]</span> <span class="o">=</span> <span class="n">exp</span><span class="p">(</span><span class="n">z</span><span class="p">);</span>
  <span class="kt">float</span>  <span class="n">denom</span>         <span class="o">=</span> <span class="n">num</span><span class="p">[</span><span class="o">+</span><span class="p">];</span>
  <span class="c1">// The result of the reduction is now stored in y</span>
@@ -253,9 +256,9 @@ In practice, though, we expect less because our kernel will spend some time comp
 </div>
 <div class="section" id="torch-bindings">
 <h2>Torch Bindings<a class="headerlink" href="#torch-bindings" title="Permalink to this headline">¶</a></h2>
-<p>We need to make sure that BLOCK is the smallest power of two
-greater than the number of rows N of the input matrix.
-Different values of BLOCK will result in different kernels</p>
+<p>Here our torch bindings is quite similar to that of the vector addition mentioned in the previous tutorial.
+We just need to make sure that BLOCK is the smallest power of two greater than the number of columns N of the input matrix.
+This means that different values of BLOCK will result in different kernels</p>
 <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">import</span> <span class="nn">triton</span>

@@ -277,6 +280,7 @@ Different values of BLOCK will result in different kernels</p>
 <span class="s2">&quot;&quot;&quot;</span>


+<span class="c1"># helper function to get the smaller power-of-two larger than a given number</span>
 <span class="k">def</span> <span class="nf">next_power_of_2</span><span class="p">(</span><span class="n">n</span><span class="p">):</span>
    <span class="n">n</span> <span class="o">-=</span> <span class="mi">1</span>
    <span class="n">n</span> <span class="o">|=</span> <span class="n">n</span> <span class="o">&gt;&gt;</span> <span class="mi">1</span>
@@ -288,16 +292,20 @@ Different values of BLOCK will result in different kernels</p>
    <span class="k">return</span> <span class="n">n</span>


-<span class="n">_kernels</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
-
-
+<span class="c1"># kernel caching mechanism</span>
 <span class="k">def</span> <span class="nf">make_kernel</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">device</span><span class="p">):</span>
+    <span class="n">cache</span> <span class="o">=</span> <span class="n">make_kernel</span><span class="o">.</span><span class="n">cache</span>
+    <span class="c1"># Now are kernels are indexed not only by the provided device but also</span>
+    <span class="c1"># by the rounded number of columns in the input matrix</span>
    <span class="n">BLOCK</span> <span class="o">=</span> <span class="n">next_power_of_2</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>
    <span class="n">key</span> <span class="o">=</span> <span class="p">(</span><span class="n">BLOCK</span><span class="p">,</span> <span class="n">device</span><span class="p">)</span>
-    <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_kernels</span><span class="p">:</span>
+    <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">cache</span><span class="p">:</span>
        <span class="n">defines</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;BLOCK&#39;</span><span class="p">:</span> <span class="n">BLOCK</span><span class="p">}</span>
-        <span class="n">_kernels</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">triton</span><span class="o">.</span><span class="n">kernel</span><span class="p">(</span><span class="n">_src</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">defines</span><span class="o">=</span><span class="n">defines</span><span class="p">)</span>
-    <span class="k">return</span> <span class="n">_kernels</span><span class="p">[</span><span class="n">key</span><span class="p">]</span>
+        <span class="n">cache</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">triton</span><span class="o">.</span><span class="n">kernel</span><span class="p">(</span><span class="n">_src</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">defines</span><span class="o">=</span><span class="n">defines</span><span class="p">)</span>
+    <span class="k">return</span> <span class="n">cache</span><span class="p">[</span><span class="n">key</span><span class="p">]</span>
+
+
+<span class="n">make_kernel</span><span class="o">.</span><span class="n">cache</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>


 <span class="k">class</span> <span class="nc">_softmax</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">autograd</span><span class="o">.</span><span class="n">Function</span><span class="p">):</span>
@@ -306,11 +314,10 @@ Different values of BLOCK will result in different kernels</p>
        <span class="c1"># constraints of the op</span>
        <span class="k">assert</span> <span class="n">x</span><span class="o">.</span><span class="n">dtype</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">float32</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">empty_like</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
-        <span class="c1"># *create launch grid*:</span>
-        <span class="c1"># here we just launch a grid of M programs</span>
+        <span class="c1"># The launch grid is simple: we have one kernel instance per row of the input matrix</span>
        <span class="n">M</span><span class="p">,</span> <span class="n">N</span> <span class="o">=</span> <span class="n">y</span><span class="o">.</span><span class="n">shape</span>
        <span class="n">grid</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">opt</span><span class="p">:</span> <span class="p">(</span><span class="n">M</span><span class="p">,</span> <span class="p">)</span>
-        <span class="c1"># *launch kernel*:</span>
+        <span class="c1"># Launch kernel</span>
        <span class="n">kernel</span> <span class="o">=</span> <span class="n">make_kernel</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">y</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
        <span class="n">kernel</span><span class="p">(</span><span class="n">y</span><span class="o">.</span><span class="n">data_ptr</span><span class="p">(),</span> <span class="n">x</span><span class="o">.</span><span class="n">data_ptr</span><span class="p">(),</span> <span class="n">y</span><span class="o">.</span><span class="n">stride</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="n">x</span><span class="o">.</span><span class="n">stride</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="n">M</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">grid</span><span class="o">=</span><span class="n">grid</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">y</span>
@@ -319,75 +326,63 @@ Different values of BLOCK will result in different kernels</p>
 <span class="n">softmax</span> <span class="o">=</span> <span class="n">_softmax</span><span class="o">.</span><span class="n">apply</span>
 </pre></div>
 </div>
+<p>We can use the above softmax function to compute the row-wise softmax of a given matrix.</p>
 </div>
 <div class="section" id="unit-test">
 <h2>Unit Test<a class="headerlink" href="#unit-test" title="Permalink to this headline">¶</a></h2>
-<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">1823</span><span class="p">,</span> <span class="mi">781</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">)</span>
+<p>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.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
+<span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">1823</span><span class="p">,</span> <span class="mi">781</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">)</span>
 <span class="n">y_tri</span> <span class="o">=</span> <span class="n">softmax</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
 <span class="n">y_ref</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
-<span class="nb">print</span><span class="p">(</span><span class="n">y_tri</span><span class="p">)</span>
-<span class="nb">print</span><span class="p">(</span><span class="n">y_ref</span><span class="p">)</span>
 <span class="nb">print</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">allclose</span><span class="p">(</span><span class="n">y_tri</span><span class="p">,</span> <span class="n">y_ref</span><span class="p">))</span>
 </pre></div>
 </div>
 <p class="sphx-glr-script-out">Out:</p>
-<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>tensor([[2.0935e-03, 6.4551e-04, 9.8605e-05,  ..., 3.3981e-04, 2.7386e-03,
-         9.1986e-05],
-        [7.0923e-04, 6.7521e-04, 5.1366e-04,  ..., 9.8392e-04, 2.6547e-04,
-         6.9062e-04],
-        [1.4032e-04, 5.8826e-04, 1.1694e-03,  ..., 6.6423e-04, 1.8178e-04,
-         6.7049e-04],
-        ...,
-        [1.1767e-03, 4.2703e-03, 6.0596e-04,  ..., 9.5274e-04, 1.1681e-03,
-         6.4924e-04],
-        [1.0772e-04, 7.4854e-04, 3.1912e-03,  ..., 2.4980e-04, 1.9012e-03,
-         5.2567e-04],
-        [2.8518e-03, 8.1899e-04, 7.7046e-04,  ..., 1.3403e-03, 5.3167e-04,
-         4.3268e-04]], device=&#39;cuda:0&#39;)
-tensor([[2.0935e-03, 6.4551e-04, 9.8605e-05,  ..., 3.3981e-04, 2.7386e-03,
-         9.1986e-05],
-        [7.0923e-04, 6.7521e-04, 5.1366e-04,  ..., 9.8392e-04, 2.6547e-04,
-         6.9062e-04],
-        [1.4032e-04, 5.8826e-04, 1.1694e-03,  ..., 6.6423e-04, 1.8178e-04,
-         6.7049e-04],
-        ...,
-        [1.1767e-03, 4.2703e-03, 6.0596e-04,  ..., 9.5274e-04, 1.1681e-03,
-         6.4924e-04],
-        [1.0772e-04, 7.4854e-04, 3.1912e-03,  ..., 2.4980e-04, 1.9012e-03,
-         5.2567e-04],
-        [2.8518e-03, 8.1899e-04, 7.7046e-04,  ..., 1.3403e-03, 5.3167e-04,
-         4.3268e-04]], device=&#39;cuda:0&#39;)
-True
+<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>True
 </pre></div>
 </div>
-<p>Seems to work!</p>
+<p>As expected, the results are identical.</p>
 </div>
 <div class="section" id="benchmarking">
 <h2>Benchmarking<a class="headerlink" href="#benchmarking" title="Permalink to this headline">¶</a></h2>
+<p>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 class="code docutils literal notranslate"><span class="pre">torch.softmax</span></code> and (2) the <code class="code docutils literal notranslate"><span class="pre">naive_softmax</span></code> defined above.</p>
 <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>

 <span class="n">M</span> <span class="o">=</span> <span class="mi">4096</span>
-<span class="n">Ns</span> <span class="o">=</span> <span class="p">[</span><span class="mi">128</span> <span class="o">*</span> <span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">50</span><span class="p">)]</span>
-<span class="n">tri_ms</span> <span class="o">=</span> <span class="p">[]</span>
-<span class="n">ref_ms</span> <span class="o">=</span> <span class="p">[]</span>
-<span class="n">def_ms</span> <span class="o">=</span> <span class="p">[]</span>
+<span class="n">Ns</span> <span class="o">=</span> <span class="p">[</span><span class="mi">256</span> <span class="o">*</span> <span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">50</span><span class="p">)]</span>
+<span class="n">tri_bw</span> <span class="o">=</span> <span class="p">[]</span>
+<span class="n">ref_bw</span> <span class="o">=</span> <span class="p">[]</span>
+<span class="n">def_bw</span> <span class="o">=</span> <span class="p">[]</span>
 <span class="k">for</span> <span class="n">N</span> <span class="ow">in</span> <span class="n">Ns</span><span class="p">:</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">M</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
    <span class="n">gbps</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">ms</span><span class="p">:</span> <span class="n">x</span><span class="o">.</span><span class="n">nelement</span><span class="p">()</span> <span class="o">*</span> <span class="n">x</span><span class="o">.</span><span class="n">element_size</span><span class="p">()</span> <span class="o">*</span> <span class="mf">1e-9</span> <span class="o">/</span> <span class="p">(</span><span class="n">ms</span> <span class="o">*</span> <span class="mf">1e-3</span><span class="p">)</span>
-    <span class="n">tri_ms</span> <span class="o">+=</span> <span class="p">[</span><span class="n">gbps</span><span class="p">(</span><span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">softmax</span><span class="p">(</span><span class="n">x</span><span class="p">)))]</span>
-    <span class="n">ref_ms</span> <span class="o">+=</span> <span class="p">[</span><span class="n">gbps</span><span class="p">(</span><span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)))]</span>
-    <span class="n">def_ms</span> <span class="o">+=</span> <span class="p">[</span><span class="n">gbps</span><span class="p">(</span><span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">naive_softmax</span><span class="p">(</span><span class="n">x</span><span class="p">)))]</span>
+    <span class="n">do_bench</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">fn</span><span class="p">:</span> <span class="n">gbps</span><span class="p">(</span><span class="n">triton</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">do_bench</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">warmup</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">rep</span><span class="o">=</span><span class="mi">100</span><span class="p">,</span> <span class="n">clear_l2</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
+    <span class="n">tri_bw</span> <span class="o">+=</span> <span class="p">[</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">softmax</span><span class="p">(</span><span class="n">x</span><span class="p">))]</span>
+    <span class="n">ref_bw</span> <span class="o">+=</span> <span class="p">[</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">))]</span>
+    <span class="n">def_bw</span> <span class="o">+=</span> <span class="p">[</span><span class="n">do_bench</span><span class="p">(</span><span class="k">lambda</span><span class="p">:</span> <span class="n">naive_softmax</span><span class="p">(</span><span class="n">x</span><span class="p">))]</span>
 <span class="n">plt</span><span class="o">.</span><span class="n">xlabel</span><span class="p">(</span><span class="s1">&#39;N&#39;</span><span class="p">)</span>
 <span class="n">plt</span><span class="o">.</span><span class="n">ylabel</span><span class="p">(</span><span class="s1">&#39;Bandwidth (GB/s)&#39;</span><span class="p">)</span>
-<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">Ns</span><span class="p">,</span> <span class="n">tri_ms</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Triton&#39;</span><span class="p">)</span>
-<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">Ns</span><span class="p">,</span> <span class="n">ref_ms</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Torch&#39;</span><span class="p">)</span>
-<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">Ns</span><span class="p">,</span> <span class="n">def_ms</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Naive&#39;</span><span class="p">)</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">Ns</span><span class="p">,</span> <span class="n">tri_bw</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Triton&#39;</span><span class="p">)</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">Ns</span><span class="p">,</span> <span class="n">ref_bw</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Torch&#39;</span><span class="p">)</span>
+<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">Ns</span><span class="p">,</span> <span class="n">def_bw</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s1">&#39;Naive&#39;</span><span class="p">)</span>
 <span class="n">plt</span><span class="o">.</span><span class="n">legend</span><span class="p">()</span>
 <span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
 </pre></div>
 </div>
 <img alt="02 fused softmax" class="sphx-glr-single-img" src="../../_images/sphx_glr_02-fused-softmax_001.png" />
-<p class="sphx-glr-timing"><strong>Total running time of the script:</strong> ( 0 minutes  5.758 seconds)</p>
+<p>In the above plot, we can see that:</p>
+<blockquote>
+<div><ul class="simple">
+<li><p>Triton is 4-5x faster than the naive implementation, which is consistent with our theoretical predictions.</p></li>
+<li><p>Triton is significantly faster than <code class="code docutils literal notranslate"><span class="pre">torch.softmax</span></code> for very large input matrices. My guess from looking at the source-code of the <a class="reference external" href="https://github.com/pytorch/pytorch/blob/9409a3a39b7149bb2d833a89e0c944109bef7c27/caffe2/operators/softmax_ops.cu#L240">PyTorch kernel</a> 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 data necessary.
+Note that our Triton kernel is not only faster than PyTorch’s CUDA kernel, it is also <strong>easier to read, understand and maintain</strong>.</p></li>
+</ul>
+</div></blockquote>
+<p class="sphx-glr-timing"><strong>Total running time of the script:</strong> ( 0 minutes  33.773 seconds)</p>
 <div class="sphx-glr-footer class sphx-glr-footer-example docutils container" id="sphx-glr-download-getting-started-tutorials-02-fused-softmax-py">
 <div class="sphx-glr-download sphx-glr-download-python docutils container">
 <p><a class="reference download internal" download="" href="../../_downloads/d91442ac2982c4e0cc3ab0f43534afbc/02-fused-softmax.py"><code class="xref download docutils literal notranslate"><span class="pre">Download</span> <span class="pre">Python</span> <span class="pre">source</span> <span class="pre">code:</span> <span class="pre">02-fused-softmax.py</span></code></a></p>
--- a/getting-started/tutorials/sg_execution_times.html
+++ b/getting-started/tutorials/sg_execution_times.html
@@ -160,7 +160,7 @@
            
  <div class="section" id="computation-times">
 <span id="sphx-glr-getting-started-tutorials-sg-execution-times"></span><h1>Computation times<a class="headerlink" href="#computation-times" title="Permalink to this headline">¶</a></h1>
-<p><strong>00:08.983</strong> total execution time for <strong>getting-started_tutorials</strong> files:</p>
+<p><strong>00:33.773</strong> total execution time for <strong>getting-started_tutorials</strong> files:</p>
 <table class="docutils align-default">
 <colgroup>
 <col style="width: 82%" />
@@ -169,11 +169,11 @@
 </colgroup>
 <tbody>
 <tr class="row-odd"><td><p><a class="reference internal" href="02-fused-softmax.html#sphx-glr-getting-started-tutorials-02-fused-softmax-py"><span class="std std-ref">Fused Softmax</span></a> (<code class="docutils literal notranslate"><span class="pre">02-fused-softmax.py</span></code>)</p></td>
-<td><p>00:05.758</p></td>
+<td><p>00:33.773</p></td>
 <td><p>0.0 MB</p></td>
 </tr>
 <tr class="row-even"><td><p><a class="reference internal" href="01-vector-add.html#sphx-glr-getting-started-tutorials-01-vector-add-py"><span class="std std-ref">Vector Addition</span></a> (<code class="docutils literal notranslate"><span class="pre">01-vector-add.py</span></code>)</p></td>
-<td><p>00:03.225</p></td>
+<td><p>00:00.000</p></td>
 <td><p>0.0 MB</p></td>
 </tr>
 </tbody>
--- a/searchindex.js
+++ b/searchindex.js