MLX 프로파일링 — 느린 연산 찾기 · MLX Quest

느린 연산을 찾아 고치고 다시 재기

MLX 추론이 기대보다 느릴 때 원인은 거의 ‘프레임워크 전체가 느리다’가 아니야. 보통은 특정 작업의 특정 연산이 병목인데, 측정하기 전까지 어느 연산인지 몰랐던 거지. 이 레슨에서는 그 병목을 찾는 실용적인 도구를 다뤄.

두 단계 접근

1단계 — MLX 수준 프로파일링. mlx-lm과 mlx는 MLX 안에서 시간이 어디에 쓰이는지 알 수 있는 메타데이터를 충분히 제공해. generation_tps, peak_memory, 평가별 실행 시간을 보고 느린 호출을 찾아.

2단계 — Metal 수준 프로파일링. MLX 수준 자료로 부족할 때, 예를 들어 행렬 곱셈 하나까지 좁혔는데 어느 Metal 커널이 느린지 알고 싶다면 Xcode의 Metal Debugger로 GPU 추적을 봐. 준비가 더 무겁지만 병목이 MLX API 아래에 있을 때는 유일한 길이야.

1단계 — mlx-lm 내장 메타데이터

lm.lesson2에서 다룬 stream_generate의 GenerationResponse 객체에는 generation_tps와 peak_memory가 들어 있어. 운영에서 기록하면 별도 관측 체계를 세우지 않아도 처리량과 메모리 흐름을 얻을 수 있어.

LLM이 아닌 작업에서는 의심되는 MLX 호출을 time.perf_counter()로 감싸고 경계에서 mx.eval()을 호출해. 뜨거운 반복문 안의 print()로 시간을 재지 마. core.lesson7에서 본 것처럼 터미널 입출력이 측정을 오염시켜.

2단계 — Xcode Metal Debugger

MLX 작업을 실행하는 동안 Xcode에서 Debug → Capture GPU Frame을 열어. 캡처에는 각 Metal 커널 호출과 실행 시간이 나타나고, 해당 커널을 만든 MLX 연산 이름도 볼 수 있어. 정보가 빽빽하지만 가장 오래 실행된 커널을 찾으면 대개 진짜 병목을 가리켜.

특정 MLX 연산이 커널 수준에서 왜 느린지 알아야 할 때 쓰는 도구야. 대부분의 사용자에게는 지나치게 깊고, 1단계만 제대로 익혀도 충분해.

고치고 다시 재는 반복

행동으로 이어지지 않는 프로파일링은 디버깅 구경에 불과해. 느린 연산을 찾았다면 보통 다음 가운데 하나로 고쳐.

자주 실행되는 함수에 mx.compile을 써(core.lesson7). 비용이 가장 적으면서 2~5배 빨라지는 경우가 많아.
아직 하지 않았다면 양자화해. 모델 크기는 메모리 대역폭에 영향을 주며, 이 대역폭이 Apple Silicon의 실제 병목일 때가 많아.
계산 모양을 MLX에 맞춰 바꿔. 작은 연산 여러 개보다 큰 연산 몇 개가 더 잘 합쳐져.
CPU로 돌아가는 연산을 찾으면 바꿔. MLX에서는 드물지만 PyTorch MPS에서는 흔해.
MLX를 올려. 프레임워크가 빠르게 변하고 커널 개선도 몇 주마다 배포돼.

고친 뒤에는 반드시 다시 재서 효과를 확인해. 무엇이 더 빠를지에 대한 직감은 믿을 수 없으니 재측정 단계를 건너뛰지 마.

Code

1단계 — 운영에서 generation_tps와 peak_memory 기록하기·python

from mlx_lm import load, stream_generate
import time

model, tok = load("mlx-community/Llama-3.2-1B-Instruct-4bit")

last = None
t0 = time.perf_counter()
for chunk in stream_generate(model, tok, prompt="Tell me about MLX:", max_tokens=100):
    last = chunk

elapsed = time.perf_counter() - t0
print(f"Generated {last.generation_tokens} tokens in {elapsed:.2f}s")
print(f"  generation_tps : {last.generation_tps:.1f}")
print(f"  peak_memory MB : {last.peak_memory / 1024 / 1024:.1f}")
print(f"  finish_reason  : {last.finish_reason}")

# Sample on M3 Ultra Studio with Llama-3.2-1B Q4 (verified 2026-05-03):
#   generation_tps : ~600-800
#   peak_memory MB : ~685

1단계 — MLX 호출 시간 정확히 재기·python

import mlx.core as mx
import time

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

def benchmark(fn, *args, n=10):
    # Warm up
    mx.eval(fn(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        r = fn(*args)
        mx.eval(r)
    return (time.perf_counter() - t0) / n * 1000   # ms per call

t = benchmark(lambda x: x @ x.T, x)
print(f"matmul 4096x4096 : {t:.2f} ms / call")

# Don't put print() in the timed loop — terminal I/O contaminates timing.
# Always wrap with mx.eval() to materialize before reading the clock.

Exercise

실제로 실행하는 MLX 작업에 1단계 프로파일링을 더해. 모든 생성에서 generation_tps와 peak_memory를 기록하고 현실적인 부하로 실행해 수치를 모아. 병목이 처리량인지, 메모리인지, MLX 밖의 전처리 지연 같은 다른 문제인지 찾아. 수치를 근거로 무엇부터 최적화할지 두 문장으로 적어.

MLX 프로파일링 — 느린 연산 찾기