{
  "video_id": "4c8qFBbsDb0",
  "channel_slug": "nvidiadeveloper",
  "channel_handle": "nvidiadeveloper",
  "title": "CuTe DSL for JAX Developers: Writing Custom GPU Kernels in Python",
  "duration_seconds": 816.0,
  "url": "https://www.youtube.com/watch?v=4c8qFBbsDb0",
  "upload_date": "",
  "transcript": "JAX on NVIDIA GPUs is already fast\nbut sooner or later you'll hit a wall\nMaybe you need a fused op XLA won’t generate\nMaybe you need a custom memory layout\nOr maybe you're building something totally non standard\nand you just need your own kernel\nIn this video\nI'll show you a practical workflow for writing\nhigh performance GPU kernels in Python\nusing the CUTLASS CuTe DSL\nand then calling those kernels from JAX\nlike they're native JAX ops\nWe'll go from: 'Here as a kernel' to\n'It runs inside @jax.jit'\nand then we'll go one step further:\nwe'll introduce ahead of time compilation\nusing JAX export so you can serialize an artifact\nand reuse it without retracing\nin the same Python process\nIf you've ever asked\nhow do I extend JAX with my own GPU kernels\nwithout leaving the JAX ecosystem, this is for you\nSo here's the plan:\nfirst, I'll give you the mental model\nwhat CuTe is doing and how it maps to threads,\nblocks and layouts\nThen we'll walk through a series of kernels\nstarting simple and getting more interesting:\nVector add: end-to-end, from kernel to JAX call\nSAXPY to show how to pass scalar parameters cleanly\nReLU\nand then a fused bias + ReLU to show why fusion matters\nA tiled GEMM example to show launch configuration\nand the limits of a simple implementation\nMulti-GPU\nhow CUTLASS call still composes with JAX sharding\nAnd then finally ahead of time compilation\nwith jax.export: serialize,\ndeserialize and run the same compiled computation later\nplus symbolic shapes\nI'll share the code and I'll call out the gotchas\ndevelopers run into\nwhen they try this code for the first time\nCUTLASS is an NVIDIA's open source\nCUDA template library\nfor building high-performance GPU kernels\nespecially things like GEMM and tensor operations\nIt gives you optimized building blocks\ninstead of writing everything from scratch\nInside CUTLASS, CuTe is the tensor and layout system\nIt lets you describe a tensor as data plus layout\nwhich makes tiling, indexing,\nand memory mapping much easier to reason about\nThe CuTe DSL brings those abstractions into Python\nSo, instead of writing complex C++ templates\nyou can write GPU kernels in Python\nwhile still targeting the same high performance\nCUTLASS infrastructure underneath\nMy colleague Michael\ngave a talk where he discussed the motivation\nbehind all these tools. I'm linking it here\nMake sure to check it out\nSo, in this video we are combining:\nCUTLASS performance primitives,\nCuTe tensor abstractions and JAX composability\nCuTe is best thought of as an index and layout DSL\nOn GPUs\nperformance often comes down to how threads map to data,\nhow data is laid out in memory\nand how you move data through registers\nand shared memory\nCuTe gives you a vocabulary for tensor equals\npointer plus layout\nInstead of manually computing offsets\nyou express indexing in terms of tensor coordinates\nand the layout handles the address math\nAnd the big practical takeaway is\nCuTe kernels are still CUDA kernels:\nThreads and blocks are real\nYou'll read the thread index,\nblock index and control launch shapes explicitly\nCuTe DSL kernels come in two layers\nThe layer 1: @cute.kernel defines the per-thread program\nthe instructions each thread executes\n@cute.jit defines the launcher: grid, block,\nand the CUDA stream to run on\nThat second point is critical for JAX integration\nThe launcher takes a CUDA stream as its first argument\nand that stream is managed by XLA\nso your kernel runs on the same execution timeline\nas the rest of your JAX program\nLet's start with the smallest complete example\nVector add: C equals A plus B\nIn the kernel, each thread loads one element of A\nand one element of B into registers,\nadds them and stores the result\nHere are a few things to notice:\nwe use thread index and a block\nindex to identify the thread and the block\nWe use register fragments to keep the math in registers\nand we use CuTe\ncopy helpers to move data between global\nmemory and registers efficiently\nThen in the launcher we peak:\na number of threads per block\nand a number of blocks in the grid\nso the mapping between threads\nand data is clean and predictable\nNow the key part: calling this from JAX\nThe trick is: we wrap the launcher with CUTLASS\nJAX cutlass_call that gives us a callable\nthat can execute as a JAX custom call\ninside @jax.jit\nTwo important details:\nWe must tell JAX the output shape and dtype up front\nso XLA can compile the graph\nAnd for this kernels indexing style\nwe reshaped the 1D vector to a 3D view\nelements_per_thread, threads_per_block, and num_blocks\nThat reshape is a layout reinterpretation, no data copy\nWe pad to a multiple of block size,\nreshape, call the kernel,\nand then reshape back and slice off the padding\nI want to pause for a second and name the core pattern\nyou will reuse\ncutlass_call bridges worlds\nOn one side\na CuTe launcher that wants a CUDA stream plus tensors\nOn the other side\na JAX function that wants to compose inside XLA\nso cutlass_call\nturns the JAX launcher into a JAX primitive\nor custom call\nthat means you can put it inside a larger\n@jax.jit function; you can combine it with other JAX ops,\nand XLA will schedule it as a part of the compiled program\nThe two recurring responsibilities you take on are\nFirst: make sure your kernel's expected tensor layout\nmatches how you present data from JAX\nAnd the second:\nprovide correct output shape and dtype\nso compilation is well-defined\nOnce you get those right,\neverything else becomes repeatable\nNext SAXPY\nout equals alpha times x plus y\nThis is the same shape story as vector add\npad, reshape to the 3D view,\ncall the kernel and reshape back\nWhat's new is the scalar alpha. In the CuTe kernel,\nalpha is just a parameter\nbut in JAX compilation\nthe question is:\nIs alpha dynamic or is it static at compile time?\nFor this demo we treat alpha as a static argument\nThat means it's known at trace-time\nand can be baked into the compiled code pass\nPractically\nwe do that by marking alpha as static in jax.jit\nand then passing it into cutlass_call\nas a keyword argument\nThis is a simple pattern but it matters a lot\nIt lets you write kernels that take configuration\nparameters or constants\nwithout paying a dynamic dispatch cost\nNow let's switch to a pattern you'll use constantly\nin machine learning: element-wise activations\nfor ReLU, instead of the 3D 'elements per thread' layout\nwe use a flat 1D indexing approach\nCompute a global\nlinear index from block and stride indexes,\nbounds check, load,\napply max from 0 and x and then store\nThis looks a lot like standard CUDA kernel style\nand the JAX wrapper becomes even simpler:\nflatten the input, call the kernel with N,\nreshape back\nThis gives you a good feel for how CuTe supports\nboth styles:\nmore structured layout-centric\nkernels and straightforward\nCUDA-like indexing kernels\nHere's where things get interesting\nfused bias plus ReLU\nIn deep learning, X + bias and then ReLU are common\nbut if you do them separately\nyou typically pay for two kernel launches\nand an intermediate write to global memory\nthen a read back\nFusion avoids that: in one kernel\nyou load x, load bias for the column; add,\napply ReLU and store once\nSo the message here isn't\n'CuTe is a replacement for every JAX op'\nThe message is: 'When you have a hotspot\nthat’s memory-traffic dominated\nor launch-overhead dominated,\nfusion is a huge lever\nIn the wrapper, notice we also pass width as a static parameter\nbecause width affects indexing\nand it's often helpful to treat it like a compile-time\nconstant\nLet's do an advanced example:\na tiled GEMM\nThe point of this section isn't to beat cuBLAS\ncalled in JAX under the hood in jnp.matmul\nIn fact, this simple educational GEMM won't\ncuBLAS is extremely optimized\nThe point here is to show the mechanics:\nhow you decide the sizes, how blocks map to output tiles,\nand how you pass problem sizes like\nM, N and K into the launcher from JAX\nIn the wrapper, we flatten a and b\nprovide M, N, K\ncall the kernel and reshape back to M and N\nOne of the biggest questions developers have is\n'If I add custom kernels, do I lose JAX's composability?'\nespecially multi-GPU\nIn this demo, the answer is 'No'\nA cutlass_call kernel can still participate in JAX sharing\nThe mental model is as follows:\nJAX shards your arrays across devices\nEach device\nruns the same custom call on its local shard\nand JAX handles coordination\nSo we create a device mesh,\ndefine sharding specs,\nand use shard map to run vector\nadd across multiple GPUs\nwithout changing the kernel code itself\nThat's a key selling point of doing this\nthrough JAX integration\nrather than launching CUDA kernels out-of-band\nNow let's talk AOT: ahead of time compilation\nNormally with JAX you write a jax.jit\nfunction, and the compilation happens\nin your current Python process\nThat's great for research and iteration,\nbut in production you often want a different workflow:\nCompile once, serialize an artifact,\nship it or cache it and load it later without retracing\nThat's where jax.export comes in\nIn this demo\nwe build a function that combines a CUTLASS custom call,\nan element-wise add kernel and a normal JAX\nop like sigma;\nthen we export it, serialize it to bytes,\nand deserialize it back into a callable artifact\nThere are two key details when exporting CUTLASS kernel\nFirst one:\nthe exported computation includes custom calls\nWe explicitly tell JAX\nthat these custom calls are safe to\ninclude in the exported output\nSecond: shapes matter\nIf we export with concrete shapes,\nthe artifact is locked to those dimensions\nIf we export with symbolic shapes,\nwe can reuse the same artifact across\nmultiple input sizes\nwithin the constraints of the kernel's\nlaunch assumptions\nSo the headline: you can build a JAX program\nthat contains a CuTe kernel,\nexport it as a stand-alone serialized artifact,\nreload it and run it later\nwithout needing the original Python function definition\nin the memory\nLet me close the technical portion\nwith a few practical tips\nthat will save you time\nFirst:\nalways be explicit about shape and dtype for outputs\nwhen building custom calls\nSecond:\ndecide early whether a parameter should be dynamic\nor static\nscalars like alpha or configuration values like width\nare often best treated as static\nThird: pay attention to padding and launch constraints\nIf your kernel assumes full blocks,\npad your inputs and slice the output\nand, finally, use this approach where it shines:\nfusion, special layouts,\ncustom data movement\nand kernels that standard libraries won't give you\nif you're trying to beat cuBLAS at GEMM\nyou're probably solving the wrong problem\nIf you want to go deeper on the conceptual side\nI linked the DevLab talk I've referenced earlier\nand I also linked the demo\nnotebook and the kernel source\nso you can run this end-to-end\nIf this is useful, like the video,\nsubscribe to NVIDIA Developer\nand drop a comment with the kernel\nyou wish you could write in JAX\nI'd love to see what people are trying to build",
  "transcript_chars": 10883,
  "ingested_at": "2026-05-15T10:49:41.071329+00:00",
  "source": "channel",
  "yt_meta": {
    "view_count": 1614,
    "like_count": 74,
    "channel_id": "UCBHcMCGaiJhv-ESTcWGJPcw",
    "categories": [
      "Science & Technology"
    ],
    "tags": [
      "CuTe DSL for JAX",
      "JAX custom GPU kernels",
      "NVIDIA GPU kernels in Python",
      "write GPU kernels in Python",
      "custom JAX operations",
      "JAX custom call",
      "JAX NVIDIA GPUs",
      "CUDA kernels for JAX",
      "Python GPU programming",
      "JAX kernel fusion",
      "SAXPY JAX kernel",
      "tiled GEMM CuTe DSL",
      "JAX multi GPU sharding",
      "JAX Ahead of Time compilation",
      "AOT compilation JAX",
      "custom CUDA kernels in JAX",
      "XLA custom call",
      "high performance GPU kernels",
      "JAX performance optimization",
      "GPU kernel fusion"
    ]
  }
}