mx.compile — 숫자로 보이는 커널 결합 · MLX Quest

`mx.compile`은 무엇을 하나

mx.compile은 함수를 감싸 새 함수를 돌려줘. 첫 호출에서 연산을 그래프로 추적하고 하나 또는 더 적은 수의 Metal 커널로 컴파일해. 같은 입력 모양과 자료형으로 다시 부르면 추적을 건너뛰고 컴파일된 커널을 곧바로 실행해. 자주 호출되는 함수라면 커널 실행 횟수가 줄고, 메모리에 쓸 필요 없는 중간값도 더 단단하게 합쳐져.

Metal을 알아야 쓸 수 있는 기능은 아니야. 대신 언제 써야 이득인지 알아야 해.

직접 잰 결과

아래 코드는 (512, 512) 배열에 tanh 연산 20개를 연달아 적용하는 함수를 컴파일했을 때와 하지 않았을 때 각각 재. 내 사무실 Mac인 M3 Ultra Studio, mlx 0.31.2, 2026-05-03 기준 결과는 이랬어.

일반 실행 — 호출당 약 0.97 ms
컴파일 실행 — 호출당 약 0.31 ms
속도 향상 — 약 3.18배

같은 Mac에서 mlx 0.32.0으로 2026-08-07에 다시 돌렸더니 세 번 연속 2.84배, 1.64배, 2.82배가 나왔어. 그러니 3.18배는 상수가 아니라 표본 하나로 봐. 기계와 GPU를 함께 쓰는 다른 작업, 함수의 구체적인 입력 모양, 심지어 같은 조건의 실행 사이에서도 숫자는 달라져. 기대해야 할 건 소수점 자리까지 같은 결과가 아니라 눈에 띄는 차이라는 점이야.

언제 쓰고 언제 피하나

써 — 입력 모양이 일정하고 자주 반복되는 함수에 써. 학습 단계, 추론의 토큰 생성 단계처럼 같은 모양으로 수백만 번 불리는 함수가 딱 맞아.

쓰지 마 — 몇 번만 부르는 함수에는 쓰지 마. 추적과 컴파일 비용은 첫 호출에 내고 이후 여러 호출에서 나눠 회수해. 한 번 계산하고 끝난다면 비용이 절약보다 커.

조심해 — 입력 모양이 자주 달라지는 함수는 새 모양과 자료형 조합마다 다시 컴파일해. 모양이 계속 바뀌면 실행보다 재컴파일에 더 오래 걸릴 수도 있어. 한 번 컴파일해 천 번 실행하는 건 좋지만 천 번 컴파일하는 건 정반대야.

결과까지 바꾸는 마법은 아니야

mx.compile은 함수의 답을 바꾸지 않아. 출력은 컴파일하지 않은 버전과 비트 단위로 같거나 부동소수점 오차 안에서 같아. 모든 함수를 마법처럼 세 배 빠르게 만들지도 않아. 메모리 대역폭에 막힌 연산이나 이미 커널 하나로 끝나는 연산은 차이가 작거나 없을 수 있어. 위에서 세 배가 나온 건 tanh 연산 20개를 커널 하나로 합쳤기 때문이야. 행렬 곱셈 하나뿐인 함수라면 거의 차이가 안 날 거야.

Code

일반 실행과 컴파일 실행 — 직접 측정·python

import mlx.core as mx
import time


def heavy(x):
    y = x
    for _ in range(20):
        y = mx.tanh(y * 1.001 + 0.001)
    return y.sum()


heavy_compiled = mx.compile(heavy)

x = mx.random.normal((512, 512))

# Warm up both versions to amortize first-call costs (compile traces here).
mx.eval(heavy(x))
mx.eval(heavy_compiled(x))


def bench(fn, n=20):
    mx.eval(fn(x))   # warm
    t0 = time.perf_counter()
    for _ in range(n):
        r = fn(x)
        mx.eval(r)
    return (time.perf_counter() - t0) / n * 1000   # ms


t_plain = bench(heavy)
t_comp  = bench(heavy_compiled)

print(f'plain    : {t_plain:.3f} ms / call')
print(f'compiled : {t_comp:.3f} ms / call')
print(f'speedup  : {t_plain/t_comp:.2f}x')

# Verified on M3 Ultra Studio, mlx 0.31.2 (2026-05-03):
#   plain    : 0.969 ms / call
#   compiled : 0.305 ms / call
#   speedup  : 3.18x
#
# Re-verified on the same machine at mlx 0.32.0 (2026-08-07), three runs:
#   plain    : 0.982 / 0.711 / 0.928 ms per call
#   compiled : 0.346 / 0.433 / 0.329 ms per call
#   speedup  : 2.84x / 1.64x / 2.82x
# Same hardware, same code, and the speedup moved. Run it more than once
# before you quote a number from it.

Exercise

현재 Mac에서 측정 코드를 실행해. 다음에는 변수를 하나씩만 바꿔 다시 재봐. (a) 배열 모양을 (512, 512)에서 (4096, 4096)으로 바꾸면 속도 차이가 유지되고, 커지고, 줄어드는지 봐. (b) heavy 안의 반복 횟수를 20에서 1로 줄여 연산 하나만 남기면 속도 차이가 어떻게 되는지 봐. 합칠 작은 연산이 많고 입력 모양이 안정적일 때 컴파일이 가장 유리하다는 감각을 가져가.

mx.compile — 숫자로 보이는 커널 결합

`mx.compile`은 무엇을 하나

직접 잰 결과

언제 쓰고 언제 피하나

결과까지 바꾸는 마법은 아니야

Code

External links

Exercise

Progress

댓글 0

mx.compile은 무엇을 하나

직접 잰 결과

언제 쓰고 언제 피하나

결과까지 바꾸는 마법은 아니야

Code

External links

Exercise

Progress

`mx.compile`은 무엇을 하나