4.1 CUDA

Writing and selecting CUDA kernels for inference, and cutting memory accesses through kernel fusion.

CUDA is how you write code to run on NVIDIA GPUs.

More formally, CUDA is NVIDIA’s proprietary computing platform and programming model for executing parallel tasks on GPUs. Both “platform” and “programming model” are broad definitions; it’s easier to look at CUDA in its component parts:

  • CUDA kernel: A user-defined function that executes parallelized code on the GPU.
  • CUDA graph: A directed acyclic graph (DAG) of kernels and other GPU operations for optimizing repeated workflows.
  • CUDA driver: A low-level interface between the application and the GPU hardware to manage memory and execution.
  • CUDA runtime: A developer-facing API for launching kernels and managing memory.

CUDA – which stands for Compute Unified Device Architecture, though the acronym is rarely expanded today – is the foundation for the entire generative AI ecosystem on NVIDIA GPUs.

CUDA is not a programming language. Instead, CUDA programs are specified in a programming language, most often C++, then compiled into separate CPU and GPU code by a compiler like nvcc.

A kernel is just a function that does some parallel computation. Whenever you see the phrase “CUDA kernel” you can replace it with “a piece of code written for NVIDIA GPUs.”

For a “Hello, World!” kernel, consider these six lines of C++ code which take in an array of length n and double each element in the array.

Figure 4.1: Example CUDA kernel doubles each value in an array.
Figure 4.1: Example CUDA kernel doubles each value in an array.

Ordinarily, on a CPU, this function would run in linear time, with each element in the array being doubled sequentially. But on a GPU, thousands of elements can be processed simultaneously, making this function much more efficient.

Writing CUDA kernels shifts inference engineering from thinking about algorithms to thinking about implementations. For example, the traditional attention algorithm that is central to generative AI can be expressed in a few dozen lines of code. However, FlashAttention, which is the same mathematical operation, takes tens of thousands of lines of code to implement the algorithm in a more memory-efficient manner for a specific GPU.

4.1.1 CUDA Kernels for Inference

Writing CUDA kernels does not mean building from scratch.

The prior art to build upon predates CUDA by decades. BLAS (Basic Linear Algebra Subprograms), first implemented for Fortran in the 1970s, is a specification for common linear algebra operations from dot products to matrix multiplication.

cuBLAS, the CUDA implementation for BLAS, brings this specification to CUDA in the form of pre-built kernels for essential linear algebra operations. Similarly, cuDNN (CUDA Deep Neural Network) provides primitives for neural networks.

Within BLAS, the most frequently used operation in inference is GEMM (General Matrix-Matrix Multiplication). Every linear layer in a model uses matrix multiplication, and cuBLAS provides a strong starting point.

But you aren’t limited to the cuBLAS implementation. As operations like GEMM are so essential for inference, you may need more fine-grained control. You might write different GEMM kernels for matrices of different shapes or to better run on specific GPU architectures.

CUTLASS is a template library that provides building blocks for writing high-performance kernels. For example, FlashAttention 3 uses CUTLASS. CuTe, another template library, introduces abstractions for tiled tensor operations for recent architectures. With tools like CUTLASS and CuTe, kernels can be written at a higher level of abstraction while retaining strong performance.

Another source for kernels is FlashInfer, a library with high-performance implementations of kernels for LLM inference, including many optimized attention kernels and fused sampling functions.

4.1.2 CUDA Kernel Selection

Most inference engineers will never need to write their own kernels. However, kernel selection – choosing the best kernel from a range of options – is an important part of inference optimization.

Kernel implementations are highly specialized. CUDA exposes low-level APIs for memory management and parallel computing, and kernel engineers make implementation decisions based on the exact specifications of the hardware they’re writing for.

It’s important to understand how closely tied kernel implementations are to specific hardware details. Kernels often have hard-coded values based on the memory bandwidth or the number or layout of Tensor Cores on a given GPU.

A kernel written for an H100 will likely fail to take advantage of the architecture and extra memory of a B200, while a kernel written for that B200 could be backwards incompatible with the previous-generation Hopper architecture. With each generation of GPUs, porting handwritten kernels to run optimally on the new architecture takes substantial engineering work.

Most kernel selection is automatic. Deep learning frameworks and inference engines have pre-configured kernels for various architectures, while PyTorch and TensorRT-LLM include automatic kernel selection in their compilation steps.

However, you might manually choose a few kernels for essential algorithms to speed up inference.

For example, most production-ready GEMM kernels come from cuBLAS. But when the DeepSeek lab released their updated version of DeepSeek-V3, they also released DeepGEMM, which provides more efficient GEMM kernels for running matrix multiplication in FP8 on the Hopper GPU architecture.

Manual kernel selection lets you insert a DeepGEMM kernel as a plugin to speed up a specific step in inference, like multiplying two matrices of precise dimensions. Just be careful to ensure compatibility. For example, if you were to upgrade to a B200 GPU, you’d have to either swap the kernel back, wait for Blackwell support in DeepGEMM (now supported as of the time of publication), or port the kernel yourself.

4.1.3 Reducing Memory Accesses with Kernel Fusion

Running two different kernels back-to-back on the same data results in wasted reads from and writes to memory.

As a simple example, imagine two kernels: multiply_by_2 and multiply_by_3. If these kernels were run back-to-back, the sequence would be:

  1. Read the input vector [1, 2, 3] from memory.
  2. Run multiply_by_2 on the vector.
  3. Save the output vector [2, 4, 6] to memory.
  4. Read the new input vector [2, 4, 6] from memory.
  5. Run multiply_by_3 on the vector.
  6. Save the final output vector [6, 12, 18] to memory.

The inefficiency is clear: steps three and four comprise an unnecessary round trip to memory. During decode, the bandwidth-bound phase of LLM inference, an inference engine can’t afford unnecessary reads from or writes to memory.

Kernel fusion is the process of taking two or more kernels and re-implementing them into a single kernel that handles both operations. In this example, the fused kernel would be multiply_by_6 and the new order of operations would be:

  1. Read the input vector [1, 2, 3] from memory.
  2. Run multiply_by_6 on the vector.
  3. Save the final output vector [6, 12, 18] to memory.

In practice, kernel fusion is far more complicated – functions are more complex, and data overlap isn’t as clean. But there are common patterns in kernel fusion for inference, like combining matrix multiplication, bias adding, and activation.

Figure 4.2: Kernel fusion reduces reads and writes between memory and compute within a GPU.
Figure 4.2: Kernel fusion reduces reads and writes between memory and compute within a GPU.

Kernel fusion can be an automatic or a manual process. Compilers can identify straightforward fusion opportunities and automatically create fused kernels. But more sophisticated algorithms, like FlashAttention, require handwritten fused kernels, which are used via plugins during inference.