Python에서 BLAS 호출

같은 C-level 호출, 세 Python 맛

대부분 시간엔 CUDA / Metal 안 쓰고 — 어딘가 밑에서 CUDA / Metal 호출 trigger하는 Python 써. 이 레슨이 Python 표현이 어느 BLAS routine으로 가는지 매핑.

Python 표현	BLAS 호출	어디
`np.dot(x, y)`	`sdot` / `ddot` (BLAS-1)	CPU의 OpenBLAS 또는 MKL
`np.dot(A, x)`	`sgemv` / `dgemv` (BLAS-2)	CPU의 OpenBLAS 또는 MKL
`np.dot(A, B)` 또는 `A @ B`	`sgemm` / `dgemm` (BLAS-3)	CPU의 OpenBLAS 또는 MKL
CUDA tensor의 `torch.matmul(A, B)`	`cublasSgemm` / `cublasGemmEx`	cuBLAS
MPS tensor의 `torch.matmul(A, B)`	`MPSMatrixMultiplication` encode/dispatch	MPS
`mx.matmul(A, B); mx.eval(C)`	MLX runtime 통한 `MPSMatrixMultiplication`	MPS

Python이 왜 안 느리게 만드나

Python 인터프리터가 BLAS 호출당 정확히 한 번 관여: 표현 파싱, C 확장에 dispatch. 실제 GEMM은 전부 C/CUDA/Metal 영역에서 돔. GPU 시간 1.5 ms 걸리는 4096³ GEMM에 Python 오버헤드는 microsecond — ~0.1%. 그래서 PyTorch의 A @ B가 손으로 쓴 C++ wrapper랑 같은 throughput.

뭐가 느리게 만드냐: 작은 호출 많이. 각 A @ B가 따로 dispatch. 큰 거 1번 대신 작은 matmul 1000번 하면 Python 오버헤드 millisecond로 누적. torch.compile, jax.jit, MLX의 lazy evaluation 다 작은 호출을 큰 거로 fuse하려고 존재.

Code

특정 BLAS path 강제 — 벤치마킹용으로 유용·python

# numpy: 어느 BLAS로 link됐는지 확인
import numpy as np
np.show_config()             # openblas_info / mkl_info / blis_info 출력

# 재현 가능한 timing 위해 single-threaded OpenBLAS 강제
import os
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS']     = '1'

# torch: device + 정밀도 명시적 픽
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()          # FP16 + Tensor Core로 cublasGemmEx
C_bf = A.bfloat16() @ B.bfloat16() # BF16 path (Hopper / Ada / Apple Silicon 다 지원)

직접 scipy.linalg.blas access — Python에서 C-level 통제·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  (기존 C 없으니 β 없음)
C = sgemm(alpha=1.0, a=A, b=B)
# A @ B 같은 결과인데 alpha + (선택적) beta + transpose를 copy 없이 통제.
# Fused linear-algebra pipeline에 유용.

MLX 직접 API — Apple 전용 작업에 깔끔·python

import mlx.core as mx

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

# Lazy — graph node, compute 아직 없음
C = mx.matmul(A, B)

# eval이 MPSMatrixMultiplication에 dispatch 강제
mx.eval(C)

# 가능할 때 op 여러 개 fuse:
D = mx.matmul(A, B) + mx.matmul(B, A)
mx.eval(D)   # MLX가 Metal 커널 하나로 fuse 가능

Exercise

셋 골라: NumPy, PyTorch CUDA, PyTorch MPS, MLX. 각각 4096³ FP32 GEMM 돌리고 캡처: (a) 쓴 Python 표현, (b) dispatch한 BLAS routine (꼼꼼히 하려면 strace/dtruss, 또는 그냥 이 레슨 표 신뢰), (c) 달성 TFLOP/s. 머신용 표 만들기. 표현에서 커널로 정확한 매핑이 Python-land랑 kernel-land 다리.

같은 C-level 호출, 세 Python 맛

Python이 왜 안 느리게 만드나

Code

External links

Exercise

Progress

댓글 0