Training an LLM in Swift, Part 1: Taking matrix mult from Gflop/s to Tflop/s

TL;DR Highlight

Apple Silicon에서 Swift로 직접 행렬 곱셈 커널을 구현하며 CPU, SIMD, AMX, GPU(Metal)를 단계별로 최적화해 Gflop/s에서 Tflop/s 수준까지 성능을 높이는 과정을 상세히 설명한 글이다. 프레임워크 없이 LLM 학습의 핵심 연산을 밑바닥부터 구현하고 싶은 개발자에게 Apple Silicon의 성능 한계를 체감할 수 있는 드문 자료다.

Who Should Read

Swift 또는 C로 저수준 성능 최적화에 관심 있는 개발자, 특히 Apple Silicon의 CPU/GPU 성능을 최대한 활용하고 싶거나 ML 연산의 내부 동작 원리를 직접 구현하며 이해하고 싶은 iOS/macOS 개발자.

Core Mechanics

이 글은 프레임워크나 라이브러리 없이 순수 Swift로 행렬 곱셈 커널을 10가지 방식으로 구현하며 최적화하는 실험 기록이다. 기준 구현체는 Andrej Karpathy의 llm.c(GPT2 호환 모델을 약 1000줄 C로 구현한 것)이며, 'Swift를 C보다 빠르게 만들기'가 목표다.
ML 연산의 본질은 결국 z += x * y를 엄청나게 많이 반복하는 것이다. 이를 '선형 텐서 프로젝션', '행렬 곱셈', '벡터 내적' 등 다양한 단위로 패키징할 수 있지만 핵심 루프는 동일하다.
Apple Silicon에는 성능을 낼 수 있는 계산 유닛이 여러 층으로 존재한다: 일반 CPU 코어 → SIMD(여러 데이터를 한 번에 처리하는 벡터 연산 유닛) → AMX(Apple 내부 행렬 곱셈 가속기, 비공개 API) → GPU(Metal). 각 단계마다 성능이 크게 달라진다.
AMX(Apple Matrix coprocessor)는 Apple Silicon 내부에 탑재된 행렬 연산 전용 가속기인데, 공식 문서가 없는 비공개 API다. 댓글에서는 M4 이상에서 ARM SME(Scalable Matrix Extension)를 통해 유사 기능에 접근할 수 있을 가능성이 언급됐다.
Metal(Apple의 GPU 프로그래밍 API)로 구현했을 때 약 1.1 Tflop/s를 달성했다. M3 Max GPU의 이론적 최대치가 약 15 Tflop/s인데, 이 종류의 행렬 곱셈 태스크에서 실질적 상한선은 3~5 Tflop/s 정도라고 저자가 설명한다.
C 버전에는 `-ffast-math` 컴파일러 플래그를 사용해 FMA(Fused Multiply-Add, 곱셈과 덧셈을 한 번의 연산으로 처리하는 기법)를 활성화했다. FMA는 정확도를 오히려 높이는 경우가 많지만, 기존 레거시 시스템과 결과가 달라질 수 있어 기본값이 아닌 상태로 유지되어 왔다.
이 글은 시리즈의 첫 번째 파트다. 향후 파트에서는 Apple이 제공하는 ML 프레임워크들(Core ML, Metal Performance Shaders, Accelerate 등)을 다룰 예정이며, 실제 프로덕션에서는 이 프레임워크들을 쓰는 게 맞다고 저자 스스로 언급한다.
샘플 앱은 행렬 곱셈 커널을 독립적으로 벤치마크하는 것이 아니라, 전체 LLM의 forward pass + backward pass 학습 반복 전체를 대상으로 측정한다. 그래서 실제 워크로드에 더 가까운 수치를 볼 수 있다.

Evidence

이 글이 Swift 성능 최적화에 대한 글이 드문 상황에서 매우 귀중한 자료라는 평가가 많았다. Swift 최적화 관련 문서나 튜토리얼 자체가 부족하기 때문에, LLM에 관심 없는 사람도 읽을 가치가 있다는 의견이 있었다.
AMX 인스트럭션이 정말로 '비밀'인지에 대한 의문이 제기됐다. M4 이상에서는 ARM SME(Scalable Matrix Extension) 인트린직을 통해 Swift에서 접근할 수 있을 가능성이 있다는 추측이 나왔지만, 직접 시도해본 사람은 없다고 했다.
Xcode에 포함된 clang이 올바른 컴파일러 플래그를 주면 OpenMP 코드도 컴파일 가능하다는 정보가 공유됐다. libomp.dylib와 헤더 파일이 필요하며, R 프로젝트(mac.r-project.org/openmp/)에서 배포하는 버전을 활용할 수 있다는 실용적인 팁이었다.
`-ffast-math` 사용에 대한 중요한 지적이 있었다. FMA만 활성화하려면 `-ffast-math` 대신 `-ffp-contract=fast`를 쓰는 게 맞다는 것이다. `-ffast-math`는 수치 정확도에 영향을 주는 여러 변환을 동시에 활성화하기 때문에 ML/AI나 그래픽 외 분야에서는 전문가도 함부로 쓰면 안 된다는 경고였다.
GPU 이론 성능과 실제 달성 가능 성능의 괴리에 대한 공감이 많았다. M3 Max가 이론상 15 Tflop/s지만 실제 행렬 곱셈에서 3~5 Tflop/s가 현실적 상한이라는 부분에서, CUDA가 강한 이유가 바로 이 수많은 소형 최적화 커널 때문이며 NVIDIA의 소프트웨어 해자(moat)가 여기서 나온다는 의견이 나왔다.

How to Apply

Apple Silicon Mac에서 Swift로 수치 연산 코드를 작성 중이라면, 이 글의 최적화 단계(순수 Swift → SIMD 벡터화 → 멀티스레딩 → Metal GPU)를 로드맵으로 삼아 자신의 코드에 단계적으로 적용해볼 수 있다. 각 단계에서 실제 측정 수치가 제공되므로 어느 단계에서 투자 대비 효과가 가장 큰지 판단하기 쉽다.
C/Swift 코드에 FMA(Fused Multiply-Add)를 활성화하고 싶다면, `-ffast-math` 대신 `-ffp-contract=fast` 플래그만 사용하는 것이 안전하다. 이렇게 하면 수치 정확도를 해치는 다른 최적화는 끄고 FMA만 켤 수 있어 ML 연산 성능을 높이면서 부작용을 최소화할 수 있다.
프레임워크 없이 LLM의 forward/backward pass를 직접 구현해보고 싶다면, 이 글의 기준 구현체인 Andrej Karpathy의 llm.c(github.com/karpathy/llm.c)를 참고 삼아 Swift로 포팅하는 실습을 할 수 있다. GPT2 호환 모델이 약 1000줄 C로 구현되어 있어 구조 파악이 어렵지 않다.
Metal로 GPU 커널을 직접 작성해 성능을 끌어올리려 할 때, 이 글이 실제 LLM 학습 워크로드 기준으로 1.1 Tflop/s라는 구체적 수치를 제시하므로 자신의 Metal 구현체 성능을 비교하는 기준점으로 활용할 수 있다.

Terminology

SIMDSingle Instruction Multiple Data의 약자. CPU가 하나의 명령으로 여러 데이터를 동시에 처리하는 벡터 연산 방식. 예를 들어 4개의 float를 한 번에 곱할 수 있다.

AMXApple Matrix coprocessor. Apple Silicon 칩 내부에 탑재된 행렬 연산 전용 하드웨어 가속기. 공식 문서가 없는 비공개 API라 직접 사용하기 어렵다.

FMAFused Multiply-Add. a * b + c 연산을 두 번(곱셈, 덧셈)이 아닌 한 번의 연산으로 처리하는 기법. 속도가 빠르고 오히려 정확도가 높아지는 경우가 많다.

forward passLLM에서 입력 데이터에 모델 가중치를 적용해 출력(예측값)을 만들어내는 과정. 추론(inference)과 동일한 개념이다.

backward passforward pass 결과의 오차를 계산하고, 그 오차가 각 가중치에 얼마나 기여했는지 역방향으로 계산해 가중치를 업데이트하는 과정. 학습의 핵심 단계다.

MetalApple이 제공하는 저수준 GPU 프로그래밍 API. NVIDIA의 CUDA에 대응하는 Apple Silicon용 GPU 컴퓨팅 인터페이스다.