[{"data":1,"prerenderedAt":30},["ShallowReactive",2],{"chapter:teenygrad\u002Fkernels-and-backends\u002Fwriting-a-kernel.json":3},{"project":4,"route":5,"title":6,"titleHtml":6,"navTitle":6,"part":7,"sourcePath":8,"editUrl":9,"html":10,"toc":11,"hasMermaid":22,"prev":23,"next":26},"teenygrad","\u002Fteenygrad\u002Fkernels-and-backends\u002Fwriting-a-kernel","Writing a Triton Kernel","Kernels & Backends","kernels-and-backends\u002Fwriting-a-kernel.md","https:\u002F\u002Fgithub.com\u002Fteenygrad\u002Fteenygrad\u002Fedit\u002Fmain\u002Fbooks\u002Fteenygrad\u002Fsrc\u002Fkernels-and-backends\u002Fwriting-a-kernel.md","\u003Cp>Kernel authoring has its own book: \u003Cstrong>\u003Ca href=\"\u002Fkernels\">Writing GPU Kernels in\nRust\u003C\u002Fa>\u003C\u002Fstrong>.\u003C\u002Fp>\n\u003Cp>It starts at what a GPU kernel is, assumes no CUDA experience, and goes as far\nas attaching a custom kernel to a model’s graph with a backward pass. This page\nis the two-minute version and a map into it.\u003C\u002Fp>\n\u003Ch2 id=\"the-shape-of-a-kernel\">The shape of a kernel\u003C\u002Fh2>\n\u003Cp>A kernel is a function generic over the\n\u003Ca href=\"\u002Fapi\u002Fteenygrad\u002Fteenygrad\u002Fteeny_triton\u002F\">\u003Ccode>teeny-triton\u003C\u002Fcode>\u003C\u002Fa> DSL’s\n\u003Ccode>Triton\u003C\u002Fcode> trait plus a dtype, marked with \u003Ca href=\"\u002Fteenygrad\u002Fkernels-and-backends\u002Fkernel-macro\">\u003Ccode>#[kernel]\u003C\u002Fcode>\u003C\u002Fa>:\u003C\u002Fp>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">#[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">kernel\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">]\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">pub\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> fn\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\"> vector_add\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Triton\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Num\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> const\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_SIZE\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>(\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    a_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    b_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    out_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    n_elements\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> where\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">I32Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> types\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> 1\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">I32Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Comparison\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> BoolTensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> =\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">BoolTensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> AddOffsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> 1\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">I32Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Output\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> =\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>>>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">{\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> pid \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">program_id\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Axis\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">X\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> offsets \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">arange\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\">0\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_SIZE\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> +\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> pid \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">*\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_SIZE\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> in_bounds \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">.\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">lt\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">n_elements\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#7F877D;font-style:italic\">    \u002F\u002F ... load, compute, store\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">}\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>You write what \u003Cstrong>one program\u003C\u002Fstrong> does with \u003Cstrong>one block\u003C\u002Fstrong> of the data — not what\none thread does with one element. \u003Ccode>program_id\u003C\u002Fcode> is how a program finds its slice.\u003C\u002Fp>\n\u003Ch2 id=\"how-it-compiles\">How it compiles\u003C\u002Fh2>\n\u003Cp>At \u003Ccode>teeny-triton\u003C\u002Fcode>’s own build time, \u003Ccode>build.rs\u003C\u002Fcode> reads the DSL source under\n\u003Ccode>src\u002Ftriton\u002F\u003C\u002Fcode> (plus \u003Ccode>teeny-core\u003C\u002Fcode>’s dtype definitions) and embeds it as a string\nconstant (\u003Ccode>teeny_triton::triton_lang::TRITON\u003C\u002Fcode>) — pure text processing, no\ncompiler invocation.\u003C\u002Fp>\n\u003Cp>At \u003Cem>your\u003C\u002Fem> kernel’s compile time (via \u003Ccode>teeny-compiler\u003C\u002Fcode>’s LLVM\u002FMLIR backend), that\nDSL text plus your kernel’s source is written out and compiled by the custom\n\u003Ccode>teenyc\u003C\u002Fcode> compiler — see\n\u003Ca href=\"\u002Fteenygrad\u002Fcompiler-internals\u002Fllvm-backend\">The LLVM\u002FMLIR Backend\u003C\u002Fa>.\u003C\u002Fp>\n\u003Cp>The consequence worth knowing up front: \u003Cstrong>your kernel function is never called\nby your program.\u003C\u002Fstrong> Its source text is the artefact.\u003C\u002Fp>\n\u003Ch2 id=\"where-to-go\">Where to go\u003C\u002Fh2>\n\u003Ctable>\n\u003Cthead>\n\u003Ctr>\n\u003Cth>You want\u003C\u002Fth>\n\u003Cth>Read\u003C\u002Fth>\n\u003C\u002Ftr>\n\u003C\u002Fthead>\n\u003Ctbody>\n\u003Ctr>\n\u003Ctd>To run a kernel today\u003C\u002Ftd>\n\u003Ctd>\u003Ca href=\"\u002Fkernels\u002Ffirst-kernel\u002Fvector-add\">Vector Add, End to End\u003C\u002Fa>\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>The programming model\u003C\u002Ftd>\n\u003Ctd>\u003Ca href=\"\u002Fkernels\u002Forientation\u002Fblock-not-thread\">You Program a Block, Not a Thread\u003C\u002Fa>\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>Why the source is captured as text\u003C\u002Ftd>\n\u003Ctd>\u003Ca href=\"\u002Fkernels\u002Forientation\u002Frust-to-ptx\">From Rust to PTX\u003C\u002Fa>\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>Loads, stores and masking\u003C\u002Ftd>\n\u003Ctd>\u003Ca href=\"\u002Fkernels\u002Ffirst-kernel\u002Floads-stores-masks\">Loads, Stores, and Masks\u003C\u002Fa>\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>Every DSL operation\u003C\u002Ftd>\n\u003Ctd>\u003Ca href=\"\u002Fkernels\u002Freference\u002Ftranslation-table\">Python Triton to Rust\u003C\u002Fa>\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>A kernel inside a model\u003C\u002Ftd>\n\u003Ctd>\u003Ca href=\"\u002Fkernels\u002Fin-a-model\u002Fgraph-op\">Your Kernel as a Graph Op\u003C\u002Fa>\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003C\u002Ftbody>\n\u003C\u002Ftable>\n\u003Cp>For a large, well-documented in-tree example, see\n\u003Ccode>kernels\u002Fteeny-kernels\u002Fsrc\u002Fnn\u002Fattention\u002Fflash_attn2.rs\u003C\u002Fcode> (Flash Attention 2,\nforward and backward).\u003C\u002Fp>\n",[12,16,19],{"id":13,"text":14,"level":15},"the-shape-of-a-kernel","The shape of a kernel",2,{"id":17,"text":18,"level":15},"how-it-compiles","How it compiles",{"id":20,"text":21,"level":15},"where-to-go","Where to go",false,{"title":24,"titleHtml":24,"route":25},"CPU, CUDA, and Vulkan Backends","\u002Fteenygrad\u002Fkernels-and-backends\u002Fbackends",{"title":27,"titleHtml":28,"route":29},"The #[kernel] Macro","The \u003Ccode>#[kernel]\u003C\u002Fcode> Macro","\u002Fteenygrad\u002Fkernels-and-backends\u002Fkernel-macro",1786271829003]