[{"data":1,"prerenderedAt":36},["ShallowReactive",2],{"chapter:vision-rs\u002Fkernels-and-performance\u002Fcustom-kernels.json":3},{"project":4,"route":5,"title":6,"titleHtml":6,"navTitle":6,"part":7,"sourcePath":8,"editUrl":9,"html":10,"toc":11,"hasMermaid":28,"prev":29,"next":32},"vision-rs","\u002Fvision-rs\u002Fkernels-and-performance\u002Fcustom-kernels","Custom Kernels","Kernels & Performance","kernels-and-performance\u002Fcustom-kernels.md","https:\u002F\u002Fgithub.com\u002Fteenygrad\u002Fvision-rs\u002Fedit\u002Fmain\u002Fbook\u002Fsrc\u002Fkernels-and-performance\u002Fcustom-kernels.md","\u003Cp>vision-rs writes several GPU kernels directly, on top of teenygrad’s\nTriton-style DSL (see\n\u003Ca href=\"\u002Fteenygrad\u002Fkernels-and-backends\u002Fwriting-a-kernel\">Writing a Triton Kernel\u003C\u002Fa>\nand \u003Ca href=\"\u002Fteenygrad\u002Fkernels-and-backends\u002Fkernel-macro\">The \u003Ccode>#[kernel]\u003C\u002Fcode> Macro\u003C\u002Fa>\nin the teenygrad book for the underlying mechanics), rather than composing\nthem purely from teenygrad’s built-in ops. All of them live under\n\u003Ccode>vision_rs::models::yolo::kernels\u003C\u002Fcode>.\u003C\u002Fp>\n\u003Ch2 id=\"flash-attention-2\">Flash Attention 2\u003C\u002Fh2>\n\u003Cp>\u003Ccode>kernels::attention::flash_attn2\u003C\u002Fcode> implements the standard Flash Attention 2\nalgorithm: online softmax over \u003Ccode>[BATCH*N_HEADS, N_CTX, HEAD_DIM]\u003C\u002Fcode>-layout\ntensors, one CTA per \u003Ccode>(batch_head, query_row)\u003C\u002Fcode> pair, so the full\n\u003Ccode>N_CTX_Q × N_CTX_K\u003C\u002Fcode> attention matrix is never materialised (memory is\n\u003Ccode>O(N_CTX × HEAD_DIM)\u003C\u002Fcode> per CTA rather than \u003Ccode>O(N_CTX²)\u003C\u002Fcode>). \u003Ccode>HEAD_DIM\u003C\u002Fcode> is a\ncompile-time const generic and must be a power of two. Forward, \u003Ccode>dQ\u003C\u002Fcode>\nbackward, and \u003Ccode>dK\u003C\u002Fcode>\u002F\u003Ccode>dV\u003C\u002Fcode> backward are separate kernels (\u003Ccode>#[kernel]\u003C\u002Fcode>-annotated\nfns), each with their own grid shape.\u003C\u002Fp>\n\u003Ch2 id=\"position-sensitive-attention-psa\">Position-Sensitive Attention (PSA)\u003C\u002Fh2>\n\u003Cp>\u003Ccode>kernels::attention::psa\u003C\u002Fcode> doesn’t reimplement attention — it wraps Flash\nAttention 2 with the data-rearrangement kernels the YOLO26 \u003Ccode>PSABlock\u003C\u002Fcode> needs\naround it:\u003C\u002Fp>\n\u003Cul>\n\u003Cli>\u003Cstrong>\u003Ccode>PsaPackQkv\u003C\u002Fcode>\u003C\u002Fstrong> — repacks a QKV conv’s NCHW output (\u003Ccode>[B, qkv_h, H, W]\u003C\u002Fcode>,\nchannels laid out per-head as \u003Ccode>[Q | K | V_lo | V_hi]\u003C\u002Fcode>, each \u003Ccode>KEY_DIM\u003C\u002Fcode>\nwide) into the \u003Ccode>[4, BH, N, KEY_DIM]\u003C\u002Fcode> layout Flash Attention 2 expects.\u003C\u002Fli>\n\u003Cli>\u003Cstrong>\u003Ccode>PsaExtractV\u003C\u002Fcode>\u003C\u002Fstrong> — pulls the V section back out in NCHW for the residual\npath.\u003C\u002Fli>\n\u003Cli>\u003Cstrong>\u003Ccode>PsaMergeAttn\u003C\u002Fcode>\u003C\u002Fstrong> — merges attention output back into NCHW \u003Ccode>[B, c, H, W]\u003C\u002Fcode>.\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Cp>The \u003Ccode>V_lo\u003C\u002Fcode>\u002F\u003Ccode>V_hi\u003C\u002Fcode> split exists because \u003Ccode>head_dim = 2 * key_dim\u003C\u002Fcode> in the\nultralytics PSABlock; rather than run Flash Attention 2 once with\n\u003Ccode>HEAD_DIM = head_dim\u003C\u002Fcode>, PSA runs it \u003Cem>twice\u003C\u002Fem> with \u003Ccode>HEAD_DIM = key_dim\u003C\u002Fcode> (once\nper half), avoiding a HEAD_DIM value that isn’t a clean power-of-two\nmultiple of the underlying key dimension in all configurations.\u003C\u002Fp>\n\u003Cp>Each of \u003Ccode>PsaPackQkvOp\u003C\u002Fcode>\u002F\u003Ccode>PsaExtractVOp\u003C\u002Fcode>\u002F\u003Ccode>PsaMergeAttnOp\u003C\u002Fcode>\u002F\u003Ccode>FlashAttn2PsaOp\u003C\u002Fcode> is\na thin \u003Ccode>CustomOp\u003C\u002Fcode> wrapper (see\n\u003Ca href=\"\u002Fteenygrad\u002Fnn-layers\u002Fbuilding-models\">Building Models\u003C\u002Fa>\nin the teenygrad book) that records a graph node directly — no separate\nlowering middleware needed, since \u003Ccode>lower()\u003C\u002Fcode> just hands the pre-built\n\u003Ccode>Arc<RuntimeOp>\u003C\u002Fcode> straight through.\u003C\u002Fp>\n\u003Ch2 id=\"detect-decode\">Detect-decode\u003C\u002Fh2>\n\u003Cp>\u003Ccode>kernels::detect_decode\u003C\u002Fcode> converts a model’s raw LTRB box predictions plus a\nprecomputed anchor grid into decoded \u003Ccode>[cx, cy, w, h]\u003C\u002Fcode> boxes in one fused\nkernel pass, rather than doing the anchor-grid arithmetic on the host.\n\u003Ccode>DetectDecodeOp\u003C\u002Fcode> carries the anchor grid (\u003Ccode>anchor_x\u003C\u002Fcode>, \u003Ccode>anchor_y\u003C\u002Fcode>,\n\u003Ccode>strides\u003C\u002Fcode>) as graph-node state; \u003Ccode>DetectDecodeRuntimeOp\u003C\u002Fcode> uploads that data to\ndevice parameter buffers at model-load time via\n\u003Ccode>RuntimeOp::param_init_data\u003C\u002Fcode>, so it’s a one-time setup cost, not a\nper-inference one.\u003C\u002Fp>\n\u003Ch2 id=\"loss-kernels\">Loss kernels\u003C\u002Fh2>\n\u003Cp>\u003Ccode>kernels::loss\u003C\u002Fcode> has the CUDA forward\u002Fbackward kernels\n\u003Ca href=\"\u002Fvision-rs\u002Fcore-concepts\u002Ftraining\">\u003Ccode>Yolo26Loss\u003C\u002Fcode>\u003C\u002Fa> dispatches:\u003C\u002Fp>\n\u003Cul>\n\u003Cli>\u003Cstrong>\u003Ccode>ciou\u003C\u002Fcode>\u003C\u002Fstrong> — fused CIoU loss: takes predicted\u002Ftarget \u003Ccode>[4, N]\u003C\u002Fcode> XYWH boxes,\nproduces per-anchor loss plus the intermediate \u003Ccode>iou\u003C\u002Fcode>\u002F\u003Ccode>v\u003C\u002Fcode>\u002F\u003Ccode>alpha\u003C\u002Fcode> values\nthe backward pass needs (saved-activation pattern, avoiding\nrecomputation). One CTA per \u003Ccode>BLOCK_N\u003C\u002Fcode>-wide anchor tile.\u003C\u002Fli>\n\u003Cli>\u003Cstrong>\u003Ccode>cls\u003C\u002Fcode>\u003C\u002Fstrong> — the classification loss (BCE-based) forward\u002Fbackward kernels.\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Ch2 id=\"compiling-and-inspecting-kernel-source\">Compiling and inspecting kernel source\u003C\u002Fh2>\n\u003Cp>Every \u003Ccode>#[kernel]\u003C\u002Fcode>-annotated function generates a struct (e.g.\n\u003Ccode>FlashAttention2Forward<D>\u003C\u002Fcode>) implementing teenygrad’s \u003Ccode>Kernel\u003C\u002Fcode> trait, with\n\u003Ccode>.source()\u003C\u002Fcode>\u002F\u003Ccode>.name\u003C\u002Fcode> giving you the generated Rust source and entry-point\nname — useful when debugging a kernel change, or when writing a snapshot\ntest against the generated MLIR\u002Fsource (see the \u003Ccode>test_*\u003C\u002Fcode> files under\n\u003Ccode>tests\u002F\u003C\u002Fcode> for examples of both).\u003C\u002Fp>\n",[12,16,19,22,25],{"id":13,"text":14,"level":15},"flash-attention-2","Flash Attention 2",2,{"id":17,"text":18,"level":15},"position-sensitive-attention-psa","Position-Sensitive Attention (PSA)",{"id":20,"text":21,"level":15},"detect-decode","Detect-decode",{"id":23,"text":24,"level":15},"loss-kernels","Loss kernels",{"id":26,"text":27,"level":15},"compiling-and-inspecting-kernel-source","Compiling and inspecting kernel source",false,{"title":30,"titleHtml":30,"route":31},"Training","\u002Fvision-rs\u002Fcore-concepts\u002Ftraining",{"title":33,"titleHtml":34,"route":35},"The teenyc Toolchain","The \u003Ccode>teenyc\u003C\u002Fcode> Toolchain","\u002Fvision-rs\u002Fkernels-and-performance\u002Fteenyc-toolchain",1786271830496]