Python & Other Languages

Supports shared memory (cuda.shared.array), atomics, streams, and device functions (the CUDA target has moved to the numba-cuda package).

36.1 CuPy — NumPy on the GPU

import cupy as cp
x = cp.random.rand(10_000_000, dtype=cp.float32)
y = cp.sqrt(x) * 2 + 1               # runs on the GPU
print(float(y.sum()))
a = cp.asnumpy(y)                    # copy to host

# Custom element-wise and raw kernels
square = cp.ElementwiseKernel("float32 x", "float32 y", "y = x * x", "square")
add_kernel = cp.RawKernel(r'''
extern "C" __global__ void add(const float* a, const float* b, float* c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) c[i] = a[i] + b[i];
}''', "add")
add_kernel(((n + 255) // 256,), (256,), (a_gpu, b_gpu, c_gpu, cp.int32(n)))

36.2 Numba — write kernels in Python

Supports shared memory (cuda.shared.array), atomics, streams, and device functions (the CUDA target has moved to the numba-cuda package).

from numba import cuda
import numpy as np

@cuda.jit
def vec_add(a, b, c):
    i = cuda.grid(1)
    if i < c.size:
        c[i] = a[i] + b[i]

n = 1_000_000
a = cuda.to_device(np.ones(n, np.float32)); b = cuda.to_device(np.ones(n, np.float32))
c = cuda.device_array(n, np.float32)
vec_add[(n + 255) // 256, 256](a, b, c)
print(c.copy_to_host()[:5])

36.3 NVIDIA CUDA Python, cuda.core & Warp

cuda-python (cuda.bindings): low-level Pythonic access to the driver/runtime/NVRTC APIs; cuda.core provides a higher-level Pythonic API; cuda.cooperative / cuda.parallel expose CCCL algorithms in Python.

NVIDIA Warp: Python framework for simulation & graphics kernels (differentiable).

PyCUDA: older but popular wrapper with SourceModule.

36.4 PyTorch custom CUDA extensions

Or JIT: torch.utils.cpp_extension.load_inline(...). Use at::cuda::getCurrentCUDAStream() and tensor.data_ptr<float>() inside kernels.

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(name="my_ops",
      ext_modules=[CUDAExtension("my_ops", ["my_ops.cpp", "my_kernels.cu"])],
      cmdclass={"build_ext": BuildExtension})

36.5 OpenAI Triton — block-level GPU programming in Python

Triton handles coalescing, shared memory, and Tensor Core usage automatically — widely used for custom deep-learning kernels (e.g., FlashAttention variants, torch.compile backend).

import triton, triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

# add_kernel[(triton.cdiv(n, 1024),)](x, y, out, n, BLOCK=1024)

36.6 Other languages

CUDA Fortran (nvfortran), Julia (CUDA.jl), Rust (cudarc, rust-cuda), Java (JCuda, TornadoVM), C# (ILGPU, ManagedCuda), Go/Haskell bindings, MATLAB (gpuArray, CUDA MEX), R (via Rcpp + CUDA).