Gflop/s에서 Tflop/s까지: Swift에서 행렬 곱셈 최적화
대규모 언어 모델(LLM) 훈련은 본질적으로 행렬 곱셈을 대규모로 수행하는 작업이다. 핵심은 z += x * y 연산을 수조 번 반복하는 루프이다. Apple Silicon에서 개발하는 사람들에게는 Swift의 높은 수준 안전성과 이러한 워크로드에 필요한 원시 성능 사이의 균형을 찾는 것이 도전 과제다.
최근 탐구에서 개발자 zdw는 Andrej Karpathy의 llm.c(GPT‑2 호환 모델의 순수 C 구현)를 Swift로 다시 작성하기로 했다. 목표는 단순히 C와 동등한 성능을 맞추는 것이 아니라, CPU와 SIMD 명령어부터 "비밀" AMX 코프로세서, Metal을 통한 GPU까지 M‑시리즈 칩에서 사용할 수 있는 모든 도구를 활용해 Swift를 한계까지 끌어올리는 것이었다.
시작점: 성능 격차
핵심 matmul_forward 함수를 C에서 기본 Swift로 옮겼을 때 초기 결과는 충격적이었다. Release 구성으로 런타임 assert를 제거하고 실행했음에도 불구하고, 기본 Swift 구현은 순수 C 버전보다 15~20배 느렸다.
| 모델 | Tokens/s | Training iterations/s | Training vs llm.c |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| 기본 Swift | 0.054 | 0.014 | 7.3% |
이는 대략 2.8 Gflop/s의 성능에 해당한다—1999년이라면 인상적이었겠지만 현대 LLM 워크로드에는 받아들일 수 없는 수준이다. 주요 원인으로 _ArrayBuffer.beginCOWMutation()이 지목되었다. Swift의 Copy‑on‑Write(COW) 고유성 검사 때문에 배열이 고유하더라도 막대한 오버헤드가 발생했다.
격차 해소: Swift 수준 최적화
COW 병목을 극복하기 위한 첫 단계는 Swift 6.2에 도입된 MutableSpan을 채택하는 것이었다. 이는 거의 제로 오버헤드로 메모리에 접근할 수 있는 신뢰할 수 있는 방법을 제공한다. 이로써 훈련 속도가 개선되었지만, Swift에는 C의 -ffast-math 플래그에 해당하는 직접적인 옵션이 없어 전방 패스가 여전히 느렸다. -ffast-math는 Fused Multiply‑Add(FMA) 명령을 활성화한다.
완화된 수학과 SIMD 활용
Swift-Numerics 라이브러리와 그 안의 Relaxed.multiplyAdd 함수를 사용함으로써 구현은 마침내 fmla(SIMD 벡터화 FMA) 명령을 활용할 수 있게 되었다. 이 변화만으로도 초당 토큰 수가 거의 10배 빨라졌다.
루프 언롤링과 인라인 배열
루프를 언롤링하도록 유도하는 최적화된 C 구현에 맞추기 위해 저자는 Swift 6.2의 InlineArray를 사용했다. 이를 통해 스택에 할당된 버퍼를 사용해 루프 내에서 힙 할당 배열의 높은 비용을 피할 수 있었다. 이 단계에서 "Fast Swift"는 C와 동등한 성능을 달성했으며, 실제로는 훈련 반복 횟수에서 C보다 약간 앞섰다(106.6% of llm.c).
확장: 멀티스레딩과 AMX
단일 스레드 성능은 해결됐지만, 다음 도약은 모든 CPU 코어를 활용하는 것이었다. DispatchQueue.concurrentPerform를 사용해 작업을 M3 Max의 16코어에 분산시켰다. 그러나 이로 인해 코드에 상당한 "시각적 혼란"이 발생했으며, Swift의 동시성 안전성 검사를 우회하기 위해 withUnsafeMutableBufferPointer와 @unchecked Sendable 래퍼를 사용해야 했다.
"비밀" 무기: AMX
표준 SIMD 외에도 Apple Silicon에는 AMX(Apple Matrix Coprocessor)가 내장되어 있다. Apple은 공식적으로 Accelerate 프레임워크를 통해만 이를 노출하지만, AMX_MATFP와 같은 역공학된 명령을 사용하면 16×16 타일을 직접 조작할 수 있다.
경고: AMX 명령을 직접 사용하는 것은 문서화되지 않았고 바이너리 호환성 파손 위험이 있기 때문에 프로덕션에서는 권장되지 않는다. Accelerate 프레임워크를 사용하는 것이 여전히 권장되는 경로이다.
AMX 명령을 구현함으로써 훈련 성능은 원래 llm.c 구현의 958.8%까지 끌어올렸다.
최종 단계: Metal과 GPU
Tflop/s 영역에 도달하기 위해 작업을 GPU로 옮겼으며, Metal을 사용해 컴퓨트 커널을 Metal/C++로 작성하고 Swift에서 호출 레이어를 구현했다.
- 기본 Metal: 단순 커널은 AMX보다 약간의 향상만 제공했다.
- 스레드형 Metal:
threadsPerThreadgroup을 최적화함으로써 큰 도약을 이뤘다(llm.c대비 2204.6%). - 타일링 Metal: 메모리 지역성을 개선하기 위해 타일링 커널을 구현(긴 행을 순회할 필요 감소)하면서 최종적으로 1 Tflop/s 장벽을 돌파했다.
최종 성능 비교
| 모델 | Tokens/s | Training iterations/s | Training vs llm.c |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| 멀티스레드 Swift | 4.356 | 1.014 | 558.5% |
| AMX | 5.884 | 1.678 | 958.8% |
| 타일링 Metal | 11.123 | 5.351 | 3057.7% |
기술적 통찰과 반론
FMA 논쟁
커뮤니티 토론에서 제기된 핵심 논점은 -ffast-math 사용이다. 저자는 FMA를 활성화하기 위해 이를 사용했지만, 일부 전문가들은 -ffast-math가 너무 광범위해 원치 않는 수치적 부정확성을 초래할 수 있다고 주장한다. 완전한 fast‑math 없이 FMA만을 원한다면 권장되는 대안은 -ffp-contract=fast이다.
GPU 소프트웨어 해자
"기본 Metal"에서 "타일링 Metal"로 전환하는 어려움은 NVIDIA의 CUDA와 같은 소프트웨어 생태계가 왜 그렇게 지배적인지를 보여준다. 최고 GPU 성능은 하드웨어만이 아니라 특정 데이터 형태에 맞춘 고도로 튜닝된 커널 라이브러리의 존재에 달려 있다.
결론
2.8 Gflop/s의 순진한 구현에서 시작해 저자는 1.1 Tflop/s(382배 향상)를 달성했다. 이 여정은 Swift가 C와 동등하거나 그보다 빠른 속도를 낼 수 있지만, 그 대가로 언어 고유의 우아함을 잃고 unsafe 포인터와 수동 메모리 관리로 코드를 복잡하게 만든다는 점을 강조한다. 프로덕션 애플리케이션에서는 이미 수년간 커널을 최적화해 온 기존 프레임워크(Accelerate, CoreML, MPSGraph)를 사용하는 것이 명확한 교훈이다.