Apple M4 Pro의 잠재력을 100% 끌어내기: ARM SME와 MpGEMM 최적화 분석

솔직히 말해서, 새로운 하드웨어가 나올 때마다 우리는 그 성능의 절반도 제대로 못 쓰고 있는 경우가 많습니다. 특히 Apple Silicon이 M 시리즈로 넘어오면서 하드웨어 스펙은 미친 듯이 올라갔지만, 정작 그 밑바닥에 있는 AMX 나 SME(Scalable Matrix Extension) 같은 특수 명령어 세트를 제대로 활용하는 소프트웨어는 드뭅니다.

대부분의 엔지니어들은 그냥 PyTorch나 TensorFlow가 알아서 해주겠거니 생각하거나, Apple의 Accelerate 프레임워크 가 신의 영역에 있다고 믿고 넘어갑니다. 그런데 최근 arXiv에 올라온 논문 하나가 제 눈길을 끌었습니다. 바로 MpGEMM 이라는 라이브러리를 통해, 벤더가 최적화해 둔 Apple Accelerate보다 GEMM(General Matrix Multiplication) 성능을 23%나 더 끌어올렸다는 내용입니다.

오늘은 이 논문(Demystifying ARM SME to Optimize General Matrix Multiplications)을 바탕으로, 도대체 어떻게 벤더 라이브러리를 이겼는지, 그리고 M4 Pro의 SME가 가지는 의미가 무엇인지 기술적으로 파고들어 보겠습니다.

ARM SME: 하드웨어는 준비되었으나 소프트웨어는?

ARM의 SME는 기존의 SVE(Scalable Vector Extension)에서 한 단계 더 나아가, 아예 행렬 연산을 위한 2D 타일 레지스터를 도입한 아키텍처입니다. 쉽게 말해 CPU 안에 작은 TPU 를 심어놓은 것과 비슷합니다.

문제는 이 SME를 다루는 복잡도입니다. 단순히 루프를 돌리는 게 아니라, 타일 레지스터에 데이터를 스트리밍 모드로 로드하고, 연산하고, 다시 저장하는 과정이 기존 SIMD와는 차원이 다릅니다. 논문의 저자들은 기존 선형 대수 라이브러리들이 이 SME의 잠재력을, 특히 Large Matrix 에서 제대로 활용하지 못하고 있다고 지적합니다.

MpGEMM이 Apple Accelerate를 이긴 비결

벤더가 제공하는 라이브러리(여기선 Apple Accelerate)는 보통 수십 명의 엔지니어가 어셈블리 레벨에서 깎고 깎아서 만든 결과물입니다. 이걸 오픈소스 리서치 코드가 1.23배나 앞섰다는 건, 아키텍처적인 접근 방식에서 뭔가 중요한 포인트를 짚었다는 뜻입니다.

논문에서 밝힌 핵심 최적화 기법은 크게 세 가지입니다.

1. On-the-fly Transposition (즉시 전치)

GEMM 연산에서 가장 골치 아픈 게 메모리 레이아웃입니다. Row-major냐 Column-major냐에 따라 캐시 히트율이 널뛰기를 하죠. 보통은 데이터를 패킹(Packing)하는 과정에서 메모리 복사가 일어나면서 병목이 생깁니다.

  • 기존 방식: 메모리에서 데이터를 읽어와서, 포맷을 맞추기 위해 버퍼에 다시 쓰고, 그걸 다시 레지스터로 로드.
  • MpGEMM: 데이터를 패킹하는 과정에서 SME의 특수 명령어 를 활용해 전치(Transposition)를 수행해버립니다. 즉, 메모리 대역폭 낭비를 최소화하면서 레지스터가 좋아하는 형태로 데이터를 먹여주는 겁니다.

2. Cache-aware Partitioning

M4 Pro 같은 최신 칩셋은 L2 캐시가 꽤 넉넉하지만, 행렬이 커지면 결국 캐시 미스가 발생합니다. 이 논문은 SME의 타일 크기에만 집착하지 않고, L2 캐시 사이즈에 딱 맞춰서 행렬을 쪼개는(Partitioning) 전략을 매우 공격적으로 가져갑니다. “데이터가 캐시에 없으면 연산 속도는 의미 없다”는 기본 원칙을 SME 문맥에 맞게 재해석한 겁니다.

3. Micro-kernels & Register Saturation

이 부분이 꽤 인상적인데, 사용 가능한 모든 타일 레지스터를 100% 가동하도록 마이크로 커널을 짰습니다. 단순히 연산 유닛만 돌리는 게 아니라, Multi-vector Load 명령어를 통해 데이터 공급이 연산 속도를 따라가도록 파이프라인을 꽉 채웠습니다. Apple M4 Pro에서 테스트했을 때, 이 전략이 DeepSeek나 LLaMA 같은 실제 LLM 워크로드에서 빛을 발했습니다.

Principal Engineer’s View: 그래서 실무에 쓸 수 있나?

제가 항상 팀원들에게 하는 말이 있습니다. “벤치마크 숫자만 보지 말고, 그 숫자가 나오는 맥락(Context)을 봐라.”

이 논문의 결과는 매우 고무적입니다. 특히 DeepSeek 나 LLaMA 같은 실제 모델 추론 워크로드에서 성능 향상을 입증했다는 점이 마음에 듭니다. 보통 이런 논문들은 정사각형 행렬(Square Matrix) 몇 개 돌려보고 “우리 빠르죠?” 하는 경우가 많은데, 실제 LLM에서 쓰이는 비정형적인 행렬 사이즈에서도 성능이 나온다는 건 꽤 의미가 큽니다.

하지만, 몇 가지 주의할 점은 있습니다:

  • 유지보수성: Apple Accelerate는 OS 업데이트와 함께 계속 최적화됩니다. 반면, MpGEMM 같은 리서치 코드는 특정 하드웨어 세대(M4)에 오버피팅(Overfitting) 되었을 가능성이 있습니다.
  • 안정성: 프로덕션 환경에서는 20%의 속도 향상보다 0.01%의 크래시 확률이 더 무섭습니다. 아직은 실험적인 코드라고 봐야 합니다.

결론

MpGEMM은 Apple Silicon의 SME 가 단순히 “있는 기능”이 아니라, 제대로만 건드리면 GPU 부하를 줄이면서도 엄청난 연산량을 CPU 단에서 처리할 수 있다는 것을 증명했습니다. M4 Pro를 가진 엔지니어라면, 그리고 온디바이스 AI 최적화에 목숨을 걸고 있다면, 이 논문의 구현 방식을 뜯어볼 가치는 충분합니다.

하드웨어는 이미 와 있습니다. 이제 우리가 코드를 그 수준으로 끌어올릴 차례입니다.

References