Building CUDA/C++ extensions for PyTorch for advanced users

Interface for writing a CUDA/C++ extension

Writing a PyTorch extension starts at C++ level with a CUDA kernel and a host launch function. The CUDA kernel takes in flat array pointers pointing to a base memory location in vRAM. The flat array pointers denote the base memory location of the tensors involved in the particular calculation of interest. Additional size variables can be provided to the CUDA kernel, provided that the size variables are derived within the scope of the host launch function. For example, the tensor dimension sizes can be retrieved from the Torch::tensor class with the size() method. Within the scope of the CUDA kernel, the tensor calculation output is written to the relevant pointer location which denotes a memory offset wrt. the base memory location:

#include <torch/extension.h>
#include <cuda_runtime.h>


#define TILE_SIZE 32


// Forward matrix (2d tensor) multiplication kernel
__global__ void matmul_kernel(const float* __restrict__ A,
                              const float* __restrict__ B,
                              float* __restrict__ C,
                              int M, int K, int N) {

...

}


// Host launch function
torch::Tensor matmul(torch::Tensor A, torch::Tensor B) {
  int M = A.size(0);
  int K = A.size(1);
  int N = B.size(1);

  auto C = torch::empty({M, N}, torch::dtype(torch::kFloat32).device(A.device()));

  dim3 threads(TILE_SIZE, TILE_SIZE);
  dim3 blocks((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);

  // Shared memory: 4 tiles for forward kernel.
  size_t shared_mem_size = 4 * TILE_SIZE * TILE_SIZE * sizeof(float);

  matmul_kernel<<<blocks, threads, shared_mem_size>>>(
      A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, K, N);

  return C;
}

When writing a PyTorch extension, it is important to remember that you need to write kernels and launch functions which together make up the entire automatic differentation procedure for the tensor operation. For a gemm, this means that CUDA kernels and host launch functions need to be provided for the partial gradients wrt. both A and B. These partial derivatives are respectively: grad_A = grad_output * B^T, grad_B = A^T * grad_output. For more information regarding known combinations of forward and partial derivative backward combinations, we refer to Bob Carpenter’s automatic differentation handbook.

Configuring (bind) settings

For configuring bind settings, we take the mixed-precision 4-bit floating-point example for a PyTorch torch.matmul function. For more information regarding mixed-precision training, we refer to Advanced -> Mixed-precision calculation. In addition to the forward gemm kernel and partial gradient kernels, mixed-precision operations also require quantization. Quantization of the considered datatype is used to keep layers and functions operable with other layers and functions. We assume that the gemm, backward gradients, and quantization kernels / launch functions have been saved in the following files: gemm.cu, grad.cu, quantize.cu. In addition to the kernels and launch functions, a C++ binding file (bindings.cpp) should be supplied as well. The binding file uses the Python package pybind11 to expose the kernels and launch functions to configuration in Python:

#include <torch/extension.h>

// Forward declarations for functions defined in the CUDA files.
torch::Tensor fp4_quantize(torch::Tensor input);
torch::Tensor fp4_dequantize(torch::Tensor input);
torch::Tensor fp4_matmul(torch::Tensor A, torch::Tensor B);
torch::Tensor fp4_matmul_backward_A(torch::Tensor grad_output, torch::Tensor B);
torch::Tensor fp4_matmul_backward_B(torch::Tensor A, torch::Tensor grad_output);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fp4_quantize", &fp4_quantize, "Quantize tensor to FP4 using __nv_fp4_e2m1");
  m.def("fp4_dequantize", &fp4_dequantize, "Dequantize tensor from FP4");
  m.def("fp4_matmul", &fp4_matmul, "Optimized FP4 matmul with asynchronous copies, double buffering, and loop unrolling");
  m.def("fp4_matmul_backward_A", &fp4_matmul_backward_A, "Backward kernel for FP4 matmul (grad_A)");
  m.def("fp4_matmul_backward_B", &fp4_matmul_backward_B, "Backward kernel for FP4 matmul (grad_B)");
}

The bindings are imported into PyTorch by means of the torch.utils.cpp_extension.load function. Subsequently, custom PyTorch functions with support for automatic differentiation is created:

from torch.utils.cpp_extension import load
import torch
import os

this_dir = os.path.dirname(os.path.abspath(__file__))

fp4_ext = load(
    name="fp4_ext",
    sources=[
        f"{this_dir}/gemm.cu",
        f"{this_dir}/grad.cu",
        f"{this_dir}/quantize.cu",
        f"{this_dir}/bindings.cpp"
    ],
    verbose=False
)


class FP4MatMulFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, B):
        ctx.save_for_backward(A, B)
        out_fp4 = fp4_ext.fp4_matmul(A, B)
        return out_fp4

    @staticmethod
    def backward(ctx, grad_output):
        A, B = ctx.saved_tensors
        # Compute gradient with respect to A and B using custom backward kernels.
        grad_A = fp4_ext.fp4_matmul_backward_A(grad_output, B)
        grad_B = fp4_ext.fp4_matmul_backward_B(A, grad_output)
        return grad_A, grad_B


class FP4QuantizeFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        quantized = fp4_ext.fp4_quantize(input)
        return quantized.view_as(input)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output


class FP4DequantizeFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        dequantized = fp4_ext.fp4_dequantize(input)
        return dequantized.view_as(input)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output