Calling BLAS from Python

The same C-level call, three Python flavors

Most of the time you don't write CUDA / Metal — you write Python that triggers a CUDA / Metal call somewhere underneath. This lesson maps the Python expression to the BLAS routine that runs.

Python expression	BLAS call	Where
`np.dot(x, y)`	`sdot` / `ddot` (BLAS-1)	OpenBLAS or MKL on CPU
`np.dot(A, x)`	`sgemv` / `dgemv` (BLAS-2)	OpenBLAS or MKL on CPU
`np.dot(A, B)` or `A @ B`	`sgemm` / `dgemm` (BLAS-3)	OpenBLAS or MKL on CPU
`torch.matmul(A, B)` on CUDA tensors	`cublasSgemm` / `cublasGemmEx`	cuBLAS
`torch.matmul(A, B)` on MPS tensors	`MPSMatrixMultiplication` encode/dispatch	MPS
`mx.matmul(A, B); mx.eval(C)`	`MPSMatrixMultiplication` through MLX runtime	MPS

Why Python doesn't slow it down

The Python interpreter is involved exactly once per BLAS call: parsing your expression, dispatching to the C extension. The actual GEMM runs entirely in C/CUDA/Metal land. For a 4096³ GEMM that takes 1.5 ms of GPU time, Python overhead is microseconds — ~0.1%. That's why A @ B in PyTorch gives you the same throughput as a hand-written C++ wrapper.

What does slow you down: many small calls. Each A @ B dispatches separately. If you do 1000 small matmuls instead of 1 big one, Python overhead adds up to milliseconds. torch.compile, jax.jit, and MLX's lazy evaluation all exist to fuse small calls into bigger ones.

Code

Force a specific BLAS path — useful for benchmarking·python

# numpy: see what BLAS it's linked to
import numpy as np
np.show_config()             # prints openblas_info / mkl_info / blis_info

# Force single-threaded OpenBLAS for reproducible timing
import os
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS']     = '1'

# torch: explicitly pick the device + precision
import torch
A = torch.randn(4096, 4096, device='cuda', dtype=torch.float32)
B = torch.randn(4096, 4096, device='cuda', dtype=torch.float32)
C = A @ B                          # cublasSgemm

C16 = A.half() @ B.half()          # cublasGemmEx with FP16 + Tensor Core
C_bf = A.bfloat16() @ B.bfloat16() # BF16 path (Hopper / Ada / Apple Silicon all support)

Direct scipy.linalg.blas access — when you want C-level control from Python·python

from scipy.linalg.blas import sgemm
import numpy as np

A = np.random.randn(1024, 512).astype(np.float32)
B = np.random.randn(512, 768).astype(np.float32)

# C = α * A @ B  (no β here because no existing C)
C = sgemm(alpha=1.0, a=A, b=B)
# Same result as A @ B but you control alpha and (optionally) beta + transposes
# without copying. Useful for fused linear-algebra pipelines.

MLX direct API — clean for Apple-only work·python

import mlx.core as mx

A = mx.random.normal((4096, 4096))
B = mx.random.normal((4096, 4096))

# Lazy — graph node, no compute yet
C = mx.matmul(A, B)

# eval forces dispatch to MPSMatrixMultiplication
mx.eval(C)

# Multiple ops fuse when possible:
D = mx.matmul(A, B) + mx.matmul(B, A)
mx.eval(D)   # MLX may fuse into one Metal kernel

Exercise

Pick three of: NumPy, PyTorch CUDA, PyTorch MPS, MLX. For each, run a 4096³ FP32 GEMM and capture: (a) the Python expression you wrote, (b) the BLAS routine it dispatched to (use strace/dtruss if you want to be thorough, or just trust this lesson's table), (c) the achieved TFLOP/s. Build the table for your machine. The exact mapping from expression to kernel is the bridge between Python-land and kernel-land.

The same C-level call, three Python flavors

Why Python doesn't slow it down

Code

External links

Exercise

Progress

Comments 0