numpy / PyTorch / MLX는 이 primitive를 어디 숨기나

매일 쓰는 high-level 라이브러리가 결국 이 커널

numpy, PyTorch, MLX의 모든 vector arithmetic 줄이 결국 방금 쓴 거랑 본질적으로 같은 CUDA / Metal 커널에 dispatch돼. 뭐가 생성되고 언제 어디서 — 그걸 이해하는 게 커널 레벨 코딩에서 framework 레벨 코딩으로 가는 다리야.

Layer	뭐가 dispatch	어느 backend로
NumPy `a + b` (CPU)	OpenBLAS / MKL 커널 (SIMD-vectorized)	x86 AVX2/AVX512 또는 ARM NEON
PyTorch `a + b` on CUDA tensor	cuBLAS 또는 ATen의 hand-rolled CUDA 커널	NVCC-compiled CUDA
PyTorch `a + b` on MPS tensor (Mac)	MPS (Metal Performance Shaders) 커널	Metal compute
MLX `a + b` (Apple)	MLX-native Metal 커널	Metal compute
`torch.compile(f)` graph	Triton-generated 커널 (shape별 custom code)	NVCC / Triton

각 라이브러리가 숨기는 게 9단계 워크플로우. PyTorch에서 a + b는 내부적으로:

A랑 B가 device + dtype 공유하는지 체크.
같은 device에 출력 tensor 할당.
커널 픽 — element-wise add, broadcast add, torch.compile로 이전 op이랑 fuse 등.
tensor shape으로부터 launch geometry (thread × block) 계산.
커널 launch.
출력 tensor handle return (host copy 아직 없음 — .cpu()나 .numpy()에서 일어남).

그래서 x.cuda() + y.cuda() 같은 snippet이 'Python expression처럼 보여도' 빠른 거 — 무거운 일은 정확히 lesson 3–4에서 쓴 커널.

Code

세 framework, 한 논리적 op·python

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

N = 1 << 24

# numpy — CPU에서 OpenBLAS/MKL 통과
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 — 보닛 안은 vec_add 커널
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                       # CUDA / MPS 커널에 dispatch

# MLX on Apple Silicon — 더 직접적, default lazy
a_m = mx.array(a_np)                  # 같은 RAM에 zero-copy view
b_m = mx.array(b_np)
c_m = a_m + b_m
mx.eval(c_m)                          # MLX는 lazy — eval이 강제

PyTorch가 뭘 dispatch했는지 보기 (보닛 들춤)·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)

# verbose 출력에서 Triton이 대략 이런 커널 만드는 거 보여:
#   triton_kernel: tmp0 = tl.load(A); tmp1 = tl.load(B);
#                  tmp2 = tmp0 + tmp1; tl.store(C, tmp2)
# 말 그대로 vec_add, 즉석 생성.

Exercise

가장 많이 쓰는 framework (numpy, PyTorch, MLX) 골라서 element-wise add 커널의 소스 파일 찾아. (위 external link에 힌트.) 본질적으로 네가 쓴 커널이라는 거 확인할 정도만 읽어 — bounds check, global thread ID, 원소당 FMA 한 번. 'wrapping' framework 코드는 크고 커널 자체는 5–10 줄. 파일 path 적어둬; perf 디버깅할 때 다시 와.

numpy / PyTorch / MLX는 이 primitive를 어디 숨기나

매일 쓰는 high-level 라이브러리가 결국 이 커널

Code

External links

Exercise

Progress

댓글 0