Where numpy / PyTorch / MLX Hide These Primitives

The high-level libraries you use daily are this kernel underneath

Every line of vector arithmetic in numpy, PyTorch, and MLX eventually dispatches to a CUDA or Metal kernel that looks essentially like the one you just wrote. Understanding that — what gets generated, when, where — is the bridge from kernel-level coding to framework-level coding.

Layer	What dispatches	To which backend
NumPy `a + b` (CPU)	OpenBLAS / MKL kernel (SIMD-vectorized)	x86 AVX2/AVX512 or ARM NEON
PyTorch `a + b` on CUDA tensor	cuBLAS or hand-rolled CUDA kernel from ATen	NVCC-compiled CUDA
PyTorch `a + b` on MPS tensor (Mac)	MPS (Metal Performance Shaders) kernel	Metal compute
MLX `a + b` (Apple)	MLX-native Metal kernel	Metal compute
`torch.compile(f)` graph	Triton-generated kernel (custom code per shape)	NVCC / Triton

The thing each library hides is the 9-step workflow. a + b in PyTorch is internally:

Check if A and B share device and dtype.
Allocate output tensor on the same device.
Pick a kernel — element-wise add, broadcast add, fused with prior op via torch.compile.
Compute launch geometry (threads × blocks) from the tensor shape.
Launch the kernel.
Return the output tensor handle (no host copy yet — that happens on .cpu() or .numpy()).

That's why a snippet like x.cuda() + y.cuda() is fast even though it 'looks like' a Python expression — the heavy lifting is exactly the kernel you wrote in lessons 3–4.

Code

Three frameworks, one logical operation·python

import numpy as np                    # CPU baseline
import torch                          # CUDA on NVIDIA, MPS on Apple
import mlx.core as mx                 # Apple-native

N = 1 << 24

# numpy — runs through OpenBLAS/MKL on CPU
a_np = np.random.randn(N).astype(np.float32)
b_np = np.random.randn(N).astype(np.float32)
c_np = a_np + b_np

# PyTorch on GPU — your vec_add kernel under the hood
device = 'cuda' if torch.cuda.is_available() else 'mps'
a_t = torch.from_numpy(a_np).to(device)
b_t = torch.from_numpy(b_np).to(device)
c_t = a_t + b_t                       # dispatches to a CUDA / MPS kernel

# MLX on Apple Silicon — even more direct, lazy by default
a_m = mx.array(a_np)                  # zero-copy view onto same RAM
b_m = mx.array(b_np)
c_m = a_m + b_m
mx.eval(c_m)                          # MLX is lazy — eval forces it

Inspect what PyTorch dispatched (peek under the hood)·python

import torch
import torch._dynamo

torch._dynamo.config.verbose = True

@torch.compile
def add(a, b):
    return a + b

a = torch.randn(1024, device='cuda')
b = torch.randn(1024, device='cuda')
c = add(a, b)

# In verbose output you'll see Triton generate a kernel something like:
#   triton_kernel: tmp0 = tl.load(A); tmp1 = tl.load(B);
#                  tmp2 = tmp0 + tmp1; tl.store(C, tmp2)
# That's literally your vec_add, generated on the fly.

Exercise

Pick the framework you use most (numpy, PyTorch, or MLX) and find the source file for its element-wise add kernel. (Hints in the external links above.) Read just enough to confirm it's structurally the kernel you wrote — bounds check, global thread ID, one FMA per element. The 'wrapping' framework code is large, the kernel itself is 5–10 lines. Make a note of the file path; you'll come back when debugging perf.

Where numpy / PyTorch / MLX Hide These Primitives

The high-level libraries you use daily are this kernel underneath

Code

External links

Exercise

Progress

Comments 0