지연 실행 — 생각을 한 번만 바꾸면 돼 · MLX Quest

바꿔야 할 생각은 하나뿐이야

NumPy와 PyTorch의 즉시 실행 모드에서는 배열 코드 한 줄이 그 자리에서 계산돼. y = x * 2 + 1은 결과를 만들어 저장하고, 나중에 y를 쓰지 않으면 그 계산은 낭비가 돼. MLX에서는 같은 줄이 계산 방법을 설명하는 작은 그래프를 만들어. 실제 계산은 명시적으로든 암묵적으로든 값을 요구할 때 일어나.

생각의 변화는 이게 전부야. MLX 지연 실행의 나머지 특징은 모두 이 결정에서 따라와.

실제 계산은 언제 시작되나

명시적으로: 하나 이상의 배열을 mx.eval(...)에 넘겨. 값을 실제로 만드는 표준 방법이야.

암묵적으로: 구체적인 값을 읽어야 하는 모든 동작이 계산을 시작해. print(array)는 그래프가 아니라 숫자를 보여줘야 하니 실행을 강제하고, 스칼라의 .item()도 같아. NumPy 배열이나 Python 목록으로 바꾸는 일도 값을 알아야 하므로 계산을 시작해. 반면 배열 색인은 새 지연 그래프 연산을 만들 뿐이고, 그 결과를 출력하거나 mx.eval·.item() 등으로 구체화할 때 실행돼.

REPL에서 작은 예제를 만질 때 MLX가 NumPy처럼 느껴지는 이유가 이거야. 출력할 때마다 조용히 계산이 끝나서 지연 실행이 눈에 안 보여. 여러 연산을 묶어 GPU 호출 수를 줄이고 하나로 합치려 할 때 비로소 차이가 중요해져.

굳이 늦춰서 얻는 것

커널 결합 — MLX가 연산 그래프를 보면 여러 연산을 하나의 Metal 커널로 합칠 수 있어. 레슨 7의 mx.compile이 이 능력에 기대고, 즉시 실행은 할 수 없는 일이야.
함수 변환 — mx.grad와 mx.vmap이 깔끔하게 동작해. 부작용이 있는 호출 순서가 아니라 함수 전체를 그래프로 받기 때문이야. JAX도 같은 선택을 했어.
메모리 절약 — 끝내 읽히지 않는 중간값은 아예 계산하지 않을 수 있어.

디버깅할 때 생기는 차이

출력문이 암묵적으로 계산을 끝내주니 REPL에서는 NumPy처럼 한 단계씩 값을 확인할 수 있어. 문제는 출력문이 없는 운영 코드가 예상보다 훨씬 긴 연산 사슬을 미룰 수 있다는 거야. 오류도 사슬 끝의 mx.eval에 도달해서야 나타날 수 있어. 조사할 때는 의심스러운 중간값에 mx.eval을 넣고, REPL에서는 같은 목적으로 print를 써.

Code

즉시 실행처럼 보이지만 지연 실행으로 동작해·python

import mlx.core as mx

x = mx.array([1.0, 2.0, 3.0])
y = x * 2 + 1                  # builds a graph; no kernel has run
# At this point, y is a graph node, not concrete numbers.

# Implicit eval — print materializes
print('y :', y)                 # → array([3, 5, 7], dtype=float32)

명시적 eval — 출력문 없는 긴 연산에 써·python

import mlx.core as mx

a = mx.random.normal((1024, 1024))
b = mx.random.normal((1024, 1024))

# All three lines below build a graph; nothing executes yet.
c = a @ b
d = c + a
e = mx.tanh(d)

# Force the whole graph to execute, in one go (better fusion opportunity).
mx.eval(e)

print('e mean:', float(e.mean()))   # implicit eval, but e is already materialized

한 mx.eval 호출에 여러 출력·python

import mlx.core as mx

x = mx.array([1.0, 2.0, 3.0])
sq = x ** 2
sm = sq.sum()
mn = sq.mean()

# Materialize all three at once — MLX can schedule them efficiently together.
mx.eval(sq, sm, mn)
print('sq:', sq, 'sum:', float(sm), 'mean:', float(mn))

Exercise

(512, 512) 난수 배열에 행렬 곱셈, 덧셈, tanh, sum, mean 같은 다섯 연산을 이어 붙여. 두 버전의 시간을 재봐. 하나는 끝에 mx.eval을 한 번만 쓰고, 다른 하나는 매 줄 뒤에 mx.eval을 써. time.perf_counter()로 각각 세 번 실행해 중앙값을 구해. 현재 Mac에서도 차이가 뚜렷해? 첫 번째만 연산 결합의 이득을 받을 수 있으니 그래야 해.

지연 실행 — 생각을 한 번만 바꾸면 돼