LP-GEMM: Integrating Layout Propagation into GEMM Operations

1. 문제 제기 (Problem)

현대 과학 컴퓨팅 및 머신러닝 (ML) 워크로드 (예: MLP, Transformer) 는 종종 서로 의존적인 일반 행렬 곱셈 (GEMM) 의 연속으로 구성됩니다.

• 기존 방식의 한계: 최신 BLAS 라이브러리 (OpenBLAS, MKL 등) 는 개별 GEMM 호출을 극도로 최적화하지만, BLAS API 의 제약으로 인해 각 호출이 독립적으로 작동하도록 설계되어 있습니다.
• 불필요한 오버헤드: 연속된 GEMM 연산에서, 이전 연산의 출력은 다음 연산의 입력이 됩니다. 그러나 기존 라이브러리는 각 GEMM 호출 시마다 입력 데이터를 패킹 (Packing) 하고, 연산 후 결과를 다시 언패킹 (Unpacking) 하여 표준 메모리 레이아웃으로 복원합니다.
• 결과: 이는 계산 자원을 낭비하는 중복적인 데이터 이동 (Redundant Packing/Unpacking) 을 초래하며, 특히 의존적인 GEMM 시퀀스 전체에서 성능 병목 현상을 유발합니다.

2. 방법론 (Methodology)

이 논문은 LP-GEMM (Layout Propagation GEMM) 을 제안하여 연속된 GEMM 연산 간의 데이터 레이아웃 전파 (Propagation) 를 가능하게 합니다.

• 핵심 아이디어: GEMM 커널을 세 가지 특수한 커널로 분해하여, 중간 연산 간의 언패킹/재패킹 단계를 제거합니다.
  1. Initial GEMM (ini-GEMM): 입력 데이터를 패킹하고, LP-GEMM 의 내부 레이아웃으로 결과를 저장합니다. (전파 시작)
  2. Intermediate GEMM (mid-GEMM): 이전 단계에서 패킹된 데이터를 입력으로 받아, 패킹 단계를 생략하고 동일한 레이아웃으로 결과를 유지합니다. (전파 지속)
  3. Ending GEMM (end-GEMM): 최종 결과를 표준 BLAS 레이아웃으로 언패킹하여 반환합니다. (전파 종료)
• 기술적 구현:
  • OpenBLAS 의 GotoBLAS 기반 최적화 (타일링, 패킹, 루프 언롤링) 를 기반으로 하되, 마이크로 커널 (µKernel) 의 계산 순서와 저장 순서를 일치시킵니다.
  • 기존 BLAS 는 계산 순서와 저장 순서가 달라 재배열이 필요했으나, LP-GEMM 은 이를 통일하여 메모리 국소성 (Spatial Locality) 을 극대화합니다.
  • Attention 레이어 적용: Transformer 모델의 Attention 레이어 (Q, K, V 계산 및 Softmax 등) 에서 레이아웃 전파를 유지하기 위해 RoPE(Rotary Position Embedding) 및 Softmax 연산을 레이아웃에 맞게 최적화했습니다.

3. 주요 기여 (Key Contributions)

1. LP-GEMM 커널 분해: OpenBLAS GEMM 을 ini-GEMM, mid-GEMM, end-GEMM 으로 분해하여 데이터 레이아웃 전파를 실현했습니다.
2. LLaMA 3.2 추론 파이프라인 구현: BLAS 레벨의 GEMM 호출만을 사용하여 LLaMA 3.2 의 추론 경로를 C++ 로 독립적으로 구현하여 실용성을 입증했습니다.
3. ML 워크로드 연산 최적화: AVX-512 아키텍처에서 Softmax, RoPE, RMSNorm 등 일반적인 ML 연산을 LP-GEMM 의 전파된 레이아웃에 맞게 최적화하여 구현했습니다.

4. 실험 결과 (Results)

저자는 Intel x86 (AVX-512) 과 RISC-V (RVV 1.0) 플랫폼에서 실험을 수행했습니다.

• 단일 GEMM 실행:
  • ini-GEMM 은 OpenBLAS 와 유사한 성능을 보였습니다.
  • mid-GEMM 과 end-GEMM 은 패킹 오버헤드를 제거하여 OpenBLAS 대비 중앙값 약 1.5 배, 최대 2 배 이상의 속도 향상을 보였습니다.
• Attention 레이어 (LLaMA 3.2):
  • x86: 작은 입력 크기의 Multi-Head Attention 에서 최대 2 배의 속도 향상을 달성했습니다.
  • RISC-V: OpenBLAS 의 비효율적인 언패킹 (비순차적 메모리 접근) 으로 인해 큰 행렬일수록 LP-GEMM 의 이점이 커져, 최대 5 배의 속도 향상을 기록했습니다.
• 연속 GEMM 시퀀스:
  • FlashGEMM 및 OpenBLAS 와 비교하여, 3 단계 연속 GEMM 벤치마크에서 대부분의 경우 LP-GEMM 이 더 높은 성능을 보여주었습니다.
  • 이는 레이아웃 전파가 개별 커널 최적화 이상의 효율성을 제공함을 의미합니다.

5. 의의 및 결론 (Significance & Conclusion)

• 효율성 증대: LP-GEMM 은 저수준의 마이크로아키텍처 튜닝이 아닌, 불필요한 데이터 이동 (패킹/언패킹) 제거를 통해 성능을 획기적으로 개선했습니다.
• 범용성: 이 접근법은 특정 하드웨어에 종속되지 않으며, OpenBLAS 기반의 오픈소스 구현체로 다양한 아키텍처 (x86, RISC-V) 에 적용 가능합니다.
• 미래 전망: LP-GEMM 은 "단일 커널 최적화"에서 "전체 워크로드 최적화 (End-to-End Optimization)"로의 패러다임 전환을 시사합니다. 향후 컨볼루션 등 다른 연산으로 레이아웃 전파를 확장하거나, 컴파일러 최적화 프레임워크와 통합하는 것이 중요한 연구 방향입니다.

요약하자면, LP-GEMM 은 BLAS API 의 제약으로 인해 발생하는 반복적인 데이터 재배열을 제거함으로써, 의존적인 행렬 곱셈 시퀀스 (특히 AI 모델) 의 실행 시간을 단축시키는 효과적이고 실용적인 솔루션입니다.