2021-03-06 17:35:11 -05:00
<!DOCTYPE html>
< html class = "writer-html5" lang = "en" >
< head >
< meta charset = "utf-8" / >
< meta name = "viewport" content = "width=device-width, initial-scale=1.0" / >
< title > Fused Softmax — Triton documentation< / title >
< link rel = "stylesheet" href = "../../_static/css/theme.css" type = "text/css" / >
< link rel = "stylesheet" href = "../../_static/pygments.css" type = "text/css" / >
< link rel = "stylesheet" href = "../../_static/gallery.css" type = "text/css" / >
< link rel = "stylesheet" href = "../../_static/gallery-binder.css" type = "text/css" / >
< link rel = "stylesheet" href = "../../_static/gallery-dataframe.css" type = "text/css" / >
< link rel = "stylesheet" href = "../../_static/gallery-rendered-html.css" type = "text/css" / >
<!-- [if lt IE 9]>
< script src = "../../_static/js/html5shiv.min.js" > < / script >
<![endif]-->
< script type = "text/javascript" id = "documentation_options" data-url_root = "../../" src = "../../_static/documentation_options.js" > < / script >
< script src = "../../_static/jquery.js" > < / script >
< script src = "../../_static/underscore.js" > < / script >
< script src = "../../_static/doctools.js" > < / script >
< script async = "async" src = "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS-MML_HTMLorMML" > < / script >
< script type = "text/javascript" src = "../../_static/js/theme.js" > < / script >
< link rel = "index" title = "Index" href = "../../genindex.html" / >
< link rel = "search" title = "Search" href = "../../search.html" / >
< link rel = "prev" title = "Vector Addition" href = "01-vector-add.html" / >
< / head >
< body class = "wy-body-for-nav" >
< div class = "wy-grid-for-nav" >
< nav data-toggle = "wy-nav-shift" class = "wy-nav-side" >
< div class = "wy-side-scroll" >
< div class = "wy-side-nav-search" >
< a href = "../../index.html" class = "icon icon-home" > Triton
< / a >
< div role = "search" >
< form id = "rtd-search-form" class = "wy-form" action = "../../search.html" method = "get" >
< input type = "text" name = "q" placeholder = "Search docs" / >
< input type = "hidden" name = "check_keywords" value = "yes" / >
< input type = "hidden" name = "area" value = "default" / >
< / form >
< / div >
< / div >
< div class = "wy-menu wy-menu-vertical" data-spy = "affix" role = "navigation" aria-label = "main navigation" >
< p class = "caption" > < span class = "caption-text" > Getting Started< / span > < / p >
< ul class = "current" >
< li class = "toctree-l1" > < a class = "reference internal" href = "../installation.html" > Installation< / a > < / li >
< li class = "toctree-l1 current" > < a class = "reference internal" href = "index.html" > Tutorials< / a > < ul class = "current" >
< li class = "toctree-l2" > < a class = "reference internal" href = "01-vector-add.html" > Vector Addition< / a > < / li >
< li class = "toctree-l2 current" > < a class = "current reference internal" href = "#" > Fused Softmax< / a > < ul >
< li class = "toctree-l3" > < a class = "reference internal" href = "#motivations" > Motivations< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "#compute-kernel" > Compute Kernel< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "#torch-bindings" > Torch Bindings< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "#unit-test" > Unit Test< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "#benchmarking" > Benchmarking< / a > < / li >
< / ul >
< / li >
< / ul >
< / li >
< / ul >
< / div >
< / div >
< / nav >
< section data-toggle = "wy-nav-shift" class = "wy-nav-content-wrap" >
< nav class = "wy-nav-top" aria-label = "top navigation" >
< i data-toggle = "wy-nav-top" class = "fa fa-bars" > < / i >
< a href = "../../index.html" > Triton< / a >
< / nav >
< div class = "wy-nav-content" >
< div class = "rst-content" >
< div role = "navigation" aria-label = "breadcrumbs navigation" >
< ul class = "wy-breadcrumbs" >
< li > < a href = "../../index.html" class = "icon icon-home" > < / a > » < / li >
< li > < a href = "index.html" > Tutorials< / a > » < / li >
< li > Fused Softmax< / li >
< li class = "wy-breadcrumbs-aside" >
< a href = "../../_sources/getting-started/tutorials/02-fused-softmax.rst.txt" rel = "nofollow" > View page source< / a >
< / li >
< / ul >
< hr / >
< / div >
< div role = "main" class = "document" itemscope = "itemscope" itemtype = "http://schema.org/Article" >
< div itemprop = "articleBody" >
< div class = "sphx-glr-download-link-note admonition note" >
< p class = "admonition-title" > Note< / p >
< p > Click < a class = "reference internal" href = "#sphx-glr-download-getting-started-tutorials-02-fused-softmax-py" > < span class = "std std-ref" > here< / span > < / a >
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 >
2021-03-06 22:06:32 -05:00
< p > In this tutorial, you will write a fused softmax operation (that outperforms PyTorch) and learn about:< / p >
2021-03-06 17:35:11 -05:00
< 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 >
< li > < p > The automatic vectorization capabilities of the Triton compiler.< / p > < / li >
< / ul >
< div class = "section" id = "motivations" >
< h2 > Motivations< a class = "headerlink" href = "#motivations" title = "Permalink to this headline" > ¶< / a > < / h2 >
< p > Custom GPU kernels for elementwise additions are educationally valuable but won’
Let us consider instead the case of a simple (numerically stabilized) softmax operation:< / p >
< div class = "highlight-default notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "kn" > import< / span > < span class = "nn" > torch< / span >
< span class = "c1" > # Compute the row-wise softmax of x< / span >
< span class = "k" > def< / span > < span class = "nf" > naive_softmax< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > ):< / span >
< span class = "c1" > # read MN elements ; write M elements< / span >
< span class = "n" > x_max< / span > < span class = "o" > =< / span > < span class = "n" > torch< / span > < span class = "o" > .< / span > < span class = "n" > max< / 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 = "mi" > 0< / span > < span class = "p" > ]< / span >
< span class = "c1" > # read 2MN elements ; write MN elements< / span >
< span class = "n" > z< / span > < span class = "o" > =< / span > < span class = "n" > x< / span > < span class = "o" > -< / span > < span class = "n" > x_max< / span > < span class = "p" > [:,< / span > < span class = "kc" > None< / span > < span class = "p" > ]< / span >
< span class = "c1" > # read MN elements ; write MN elements< / span >
< span class = "n" > numerator< / span > < span class = "o" > =< / span > < span class = "n" > torch< / span > < span class = "o" > .< / span > < span class = "n" > exp< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "c1" > # read MN elements ; write M elements< / span >
< span class = "n" > denominator< / span > < span class = "o" > =< / span > < span class = "n" > torch< / span > < span class = "o" > .< / span > < span class = "n" > sum< / span > < span class = "p" > (< / span > < span class = "n" > numerator< / 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 = "c1" > # read 2MN elements ; write MN elements< / span >
< span class = "n" > ret< / span > < span class = "o" > =< / span > < span class = "n" > numerator< / span > < span class = "o" > /< / span > < span class = "n" > denominator< / span > < span class = "p" > [:,< / span > < span class = "kc" > None< / span > < span class = "p" > ]< / span >
< span class = "c1" > # in total: read 7MN elements ; wrote 3MN + 2M elements< / span >
< span class = "k" > return< / span > < span class = "n" > ret< / span >
< / 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.
2021-03-06 22:06:32 -05:00
This is obviously wasteful; we’
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 >
2021-03-06 17:35:11 -05:00
< / div >
< div class = "section" id = "compute-kernel" >
< h2 > Compute Kernel< a class = "headerlink" href = "#compute-kernel" title = "Permalink to this headline" > ¶< / a > < / h2 >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< 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 >
< span class = "kt" > int< / span > < span class = "n" > m< / span > < span class = "o" > =< / span > < span class = "n" > get_program_id< / span > < span class = "p" > (< / span > < span class = "mi" > 0< / span > < span class = "p" > );< / span >
< span class = "c1" > // column indices< / span >
< span class = "kt" > int< / span > < span class = "n" > n< / span > < span class = "p" > [< / span > < span class = "n" > BLOCK< / span > < span class = "p" > ]< / span > < span class = "o" > =< / span > < span class = "mi" > 0< / span > < span class = "p" > ...< / span > < span class = "n" > BLOCK< / span > < span class = "p" > ;< / span >
< span class = "c1" > // the memory address of all the elements< / span >
< span class = "c1" > // that we want to load can be computed as follows< / span >
< span class = "kt" > float< / span > < span class = "o" > *< / span > < span class = "n" > px< / 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" > m< / span > < span class = "o" > *< / span > < span class = "n" > stride_xm< / span > < span class = "o" > +< / span > < span class = "n" > n< / span > < span class = "p" > ;< / span >
< span class = "c1" > // because BLOCK has to be a power of two< / span >
< span class = "c1" > // (per Triton-C specs), it is important< / span >
< span class = "c1" > // to guard each memory operation with predicates< / span >
< span class = "c1" > // or we will read out of bounds< / span >
< 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" > < < / 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 >
2021-03-06 22:06:32 -05:00
< span class = "c1" > // x[:, :, OPERATOR, :, :]< / span >
2021-03-06 17:35:11 -05:00
< span class = "c1" > // ^< / span >
< span class = "c1" > // index< / span >
2021-03-06 22:06:32 -05:00
< span class = "c1" > // where operator is in {min, max, +}< / span >
< span class = "c1" > // for 1D vectors, this is just x[OPERATOR].< / span >
2021-03-06 17:35:11 -05:00
< 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 >
2021-03-06 22:06:32 -05:00
< span class = "c1" > // Note that exponentials in Triton are fast< / span >
< span class = "c1" > // but approximate (i.e., think __expf in CUDA)< / span >
2021-03-06 17:35:11 -05:00
< 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 >
< span class = "kt" > float< / span > < span class = "n" > y< / span > < span class = "p" > [< / span > < span class = "n" > BLOCK< / span > < span class = "p" > ]< / span > < span class = "o" > =< / span > < span class = "n" > num< / span > < span class = "o" > /< / span > < span class = "n" > denom< / span > < span class = "p" > ;< / span >
< span class = "c1" > // We write it back< / span >
< span class = "kt" > float< / span > < span class = "o" > *< / span > < span class = "n" > py< / span > < span class = "p" > [< / span > < span class = "n" > BLOCK< / span > < span class = "p" > ]< / span > < span class = "o" > =< / span > < span class = "n" > Y< / span > < span class = "o" > +< / span > < span class = "n" > m< / span > < span class = "o" > *< / span > < span class = "n" > stride_ym< / span > < span class = "o" > +< / span > < span class = "n" > n< / span > < span class = "p" > ;< / span >
< span class = "o" > *?< / span > < span class = "p" > (< / span > < span class = "n" > check< / span > < span class = "p" > )< / span > < span class = "n" > py< / span > < span class = "o" > =< / span > < span class = "n" > y< / span > < span class = "p" > ;< / span >
< span class = "p" > }< / span >
< / pre > < / div >
< / div >
< / div > < / blockquote >
< / div >
< div class = "section" id = "torch-bindings" >
< h2 > Torch Bindings< a class = "headerlink" href = "#torch-bindings" title = "Permalink to this headline" > ¶< / a > < / h2 >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< 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 >
< span class = "c1" > # Source code for the Triton kernel< / span >
< span class = "n" > _src< / span > < span class = "o" > =< / span > < span class = "s2" > " " " < / span >
< span class = "s2" > __global__ void softmax(float* Y, float* X, int stride_ym, int stride_xm, int M, int N){< / span >
< span class = "s2" > int m = get_program_id(0);< / span >
< span class = "s2" > int n [BLOCK] = 0 ... BLOCK;< / span >
< span class = "s2" > float* px [BLOCK] = X + m*stride_xm + n;< / span >
< span class = "s2" > bool check[BLOCK] = n < N;< / span >
< span class = "s2" > float x [BLOCK] = check ? *px : -F32_INFINITY;< / span >
< span class = "s2" > float z [BLOCK] = x - x[max];< / span >
< span class = "s2" > float num [BLOCK] = exp(z);< / span >
< span class = "s2" > float denom = num[+];< / span >
< span class = "s2" > float y [BLOCK] = num / denom;< / span >
< span class = "s2" > float* py [BLOCK] = Y + m*stride_ym + n;< / span >
< span class = "s2" > *?(check)py = y;< / span >
< span class = "s2" > }< / span >
< span class = "s2" > " " " < / span >
2021-03-06 22:06:32 -05:00
< span class = "c1" > # helper function to get the smaller power-of-two larger than a given number< / span >
2021-03-06 17:35:11 -05:00
< 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" > > > < / span > < span class = "mi" > 1< / span >
< span class = "n" > n< / span > < span class = "o" > |=< / span > < span class = "n" > n< / span > < span class = "o" > > > < / span > < span class = "mi" > 2< / span >
< span class = "n" > n< / span > < span class = "o" > |=< / span > < span class = "n" > n< / span > < span class = "o" > > > < / span > < span class = "mi" > 4< / span >
< span class = "n" > n< / span > < span class = "o" > |=< / span > < span class = "n" > n< / span > < span class = "o" > > > < / span > < span class = "mi" > 8< / span >
< span class = "n" > n< / span > < span class = "o" > |=< / span > < span class = "n" > n< / span > < span class = "o" > > > < / span > < span class = "mi" > 16< / span >
< span class = "n" > n< / span > < span class = "o" > +=< / span > < span class = "mi" > 1< / span >
< span class = "k" > return< / span > < span class = "n" > n< / span >
2021-03-06 22:06:32 -05:00
< span class = "c1" > # kernel caching mechanism< / span >
2021-03-06 17:35:11 -05:00
< 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 >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< 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 >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< span class = "n" > defines< / span > < span class = "o" > =< / span > < span class = "p" > {< / span > < span class = "s1" > ' BLOCK' < / span > < span class = "p" > :< / span > < span class = "n" > BLOCK< / span > < span class = "p" > }< / span >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< 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 >
< span class = "nd" > @staticmethod< / span >
< span class = "k" > def< / span > < span class = "nf" > forward< / span > < span class = "p" > (< / span > < span class = "n" > ctx< / span > < span class = "p" > ,< / span > < span class = "n" > x< / span > < span class = "p" > ):< / span >
< 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 >
2021-03-06 22:06:32 -05:00
< span class = "c1" > # The launch grid is simple: we have one kernel instance per row of the input matrix< / span >
2021-03-06 17:35:11 -05:00
< 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 >
2021-03-06 22:06:32 -05:00
< span class = "c1" > # Launch kernel< / span >
2021-03-06 17:35:11 -05:00
< 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 >
< 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 >
2021-03-06 22:06:32 -05:00
< p > We can use the above softmax function to compute the row-wise softmax of a given matrix.< / p >
2021-03-06 17:35:11 -05:00
< / div >
< div class = "section" id = "unit-test" >
< h2 > Unit Test< a class = "headerlink" href = "#unit-test" title = "Permalink to this headline" > ¶< / a > < / h2 >
2021-03-06 22:06:32 -05:00
< 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" > ' cuda' < / span > < span class = "p" > )< / span >
2021-03-06 17:35:11 -05:00
< 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" > 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 >
2021-03-06 22:06:32 -05:00
< div class = "sphx-glr-script-out highlight-none notranslate" > < div class = "highlight" > < pre > < span > < / span > True
2021-03-06 17:35:11 -05:00
< / pre > < / div >
< / div >
2021-03-06 22:06:32 -05:00
< p > As expected, the results are identical.< / p >
2021-03-06 17:35:11 -05:00
< / div >
< div class = "section" id = "benchmarking" >
< h2 > Benchmarking< a class = "headerlink" href = "#benchmarking" title = "Permalink to this headline" > ¶< / a > < / h2 >
2021-03-06 22:06:32 -05:00
< p > Here we will benchmark our operation as a function of the number of columns in the input matrix –
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 >
2021-03-06 17:35:11 -05:00
< 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 >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< 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" > ' cuda' < / 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 >
2021-03-06 22:06:32 -05:00
< 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 >
2021-03-06 17:35:11 -05:00
< span class = "n" > plt< / span > < span class = "o" > .< / span > < span class = "n" > xlabel< / span > < span class = "p" > (< / span > < span class = "s1" > ' N' < / 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" > ' Bandwidth (GB/s)' < / span > < span class = "p" > )< / span >
2021-03-06 22:06:32 -05:00
< 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" > ' Triton' < / 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" > ' Torch' < / 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" > ' Naive' < / span > < span class = "p" > )< / span >
2021-03-06 17:35:11 -05:00
< 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" / >
2021-03-06 22:06:32 -05:00
< 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 – ’ –
Note that our Triton kernel is not only faster than PyTorch’ < 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 >
2021-03-06 17:35:11 -05:00
< 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 >
< / div >
< div class = "sphx-glr-download sphx-glr-download-jupyter docutils container" >
< p > < a class = "reference download internal" download = "" href = "../../_downloads/034d953b6214fedce6ea03803c712b89/02-fused-softmax.ipynb" > < code class = "xref download docutils literal notranslate" > < span class = "pre" > Download< / span > < span class = "pre" > Jupyter< / span > < span class = "pre" > notebook:< / span > < span class = "pre" > 02-fused-softmax.ipynb< / span > < / code > < / a > < / p >
< / div >
< / div >
< p class = "sphx-glr-signature" > < a class = "reference external" href = "https://sphinx-gallery.github.io" > Gallery generated by Sphinx-Gallery< / a > < / p >
< / div >
< / div >
< / div >
< / div >
< footer >
< div class = "rst-footer-buttons" role = "navigation" aria-label = "footer navigation" >
< a href = "01-vector-add.html" class = "btn btn-neutral float-left" title = "Vector Addition" accesskey = "p" rel = "prev" > < span class = "fa fa-arrow-circle-left" aria-hidden = "true" > < / span > Previous< / a >
< / div >
< hr / >
< div role = "contentinfo" >
< p >
© Copyright 2020, Philippe Tillet.
< / p >
< / div >
Built with < a href = "https://www.sphinx-doc.org/" > Sphinx< / a > using a
< a href = "https://github.com/readthedocs/sphinx_rtd_theme" > theme< / a >
provided by < a href = "https://readthedocs.org" > Read the Docs< / a > .
< / footer >
< / div >
< / div >
< / section >
< / div >
< script type = "text/javascript" >
jQuery(function () {
SphinxRtdTheme.Navigation.enable(true);
});
< / script >
< / body >
< / html >
