PIM-SHERPA: Software Method for On-device LLM Inference by Resolving PIM Memory Attribute and Layout Inconsistencies

이 논문은 온디바이스 LLM 추론 시 발생하는 PIM 메모리 속성 및 레이아웃 불일치 문제를 해결하여 캐시 가능 영역과 비캐시 가능 영역 간의 모순을 완화하고, DRAM 더블 버퍼링 및 온라인 가중치 재배열 기법을 통해 메모리 용량을 약 48% 절감하면서도 이론적 최대 성능을 유지하는 소프트웨어 전용 방법론인 'PIM-SHERPA'를 제안합니다.

핵심

🚀 PIM-SHERPA: 스마트폰에서 거대 AI 를 더 똑똑하고 가볍게 만드는 '마법 사서'

이 논문은 스마트폰이나 엣지 기기에서 거대한 인공지능 (LLM, 예: Llama 3.2) 을 돌릴 때 발생하는 두 가지 치명적인 문제를 해결하는 새로운 소프트웨어 방법론을 소개합니다. 이름은 PIM-SHERPA입니다.

이걸 이해하기 위해 먼저 배경 지식을 쉽게 풀어서 설명해 드릴게요.

---

🧩 배경: AI 가 스마트폰에서 일하는 두 가지 방식

AI 가 질문에 답할 때는 크게 두 단계로 나뉩니다.

1. 준비 단계 (Prefill): 사용자가 "안녕"이라고 입력하면, AI 는 그 문장을 전체적으로 분석합니다. 이때는 계산이 매우 많고 (GPU 가 좋아함), 같은 단어들을 여러 번 반복해서 봅니다.
2. 답변 단계 (Decode): AI 가 한 글자씩 답변을 만들어냅니다. 이때는 계산은 적지만 메모리에서 데이터를 가져오는 속도가 매우 중요합니다.

여기서 **PIM (메모리 내 연산)**이라는 기술이 등장합니다.

• PIM 이란? 메모리 칩 안에 작은 계산기가 있어서, 데이터를 CPU 로 옮기지 않고 메모리 안에서 바로 계산하는 기술입니다. 데이터 이동이 줄어들어 속도가 8 배까지 빨라집니다.

---

⚠️ 문제: AI 가 "양면성"을 가진 두 가지 요구사항

문제는 이 두 단계가 서로 정반대의 조건을 원한다는 것입니다.

1. 준비 단계 (Prefill) 의 요구: "나는 계산할 때 같은 데이터를 자주 쓰니까, **캐시 (빠른 임시 창고)**에 넣어줘야 해!"
   • 왜? 캐시에 있으면 CPU 가 데이터를 빠르게 재사용할 수 있으니까요.
2. 답변 단계 (Decode/PIM) 의 요구: "나는 PIM 기술을 써서 메모리 안에서 계산하고 싶으니까, **캐시 밖 (비캐시 영역)**에 두어야 해!"
   • 왜? 만약 데이터가 캐시에 있으면, CPU 가 메모리 컨트롤러를 거치지 않고 캐시에서 바로 가져가버립니다. 그러면 PIM 칩이 "일할 기회"를 잃어버리게 되죠.

👉 결론: 같은 AI 모델의 '가중치 (지식)'가 한 번은 캐시 안에 있어야 하고, 한 번은 캐시 밖으로 나가야 합니다. 이걸 해결하기 위해 기존에는 모델을 두 개 복사해서 하나는 캐시용, 하나는 PIM 용으로 따로 저장했습니다. 하지만 스마트폰 메모리는 비싸고 좁기 때문에, 모델을 두 배로 늘리면 스마트폰이 터져버립니다 (메모리 부족).

---

🦸‍♂️ 해결책: PIM-SHERPA (소프트웨어 사서)

이 논문은 하드웨어를 바꾸지 않고, 소프트웨어만으로 이 문제를 해결합니다. 마치 지혜로운 사서 (Sherpa) 가 등산객 (AI) 을 위해 짐을 효율적으로 나르는 것처럼요.

저자들은 두 가지 방법을 제안합니다.

1. 방법 A: "이중 창고" (DRAM Double Buffering, DDB)

• 비유: 요리사가 재료를 준비할 때, **현재 요리에 쓰는 냄비 (캐시)**와 **다음 요리를 미리 준비해 둔 접시 (작은 캐시 버퍼)**를 동시에 사용합니다.
• 원리:
  • AI 가 현재 층 (Layer) 을 계산하는 동안, 소프트웨어는 다음 층의 데이터를 미리 캐시 영역으로 가져옵니다.
  • 계산이 끝나는 순간, 다음 층의 데이터는 이미 준비되어 있어서 기다릴 필요가 없습니다.
  • 장점: 데이터 이동 (재배열) 시간이 계산 시간 뒤에 숨겨져서, 사용자는 속도 저하를 느끼지 못합니다.

2. 방법 B: "즉석 변환" (Online Weight Rearrangement, OWR)

• 비유: 사용자가 주문을 받자마자, 요리사가 바로 재료를 다듬어서 그릇에 담는 방식입니다.
• 원리:
  • 입력된 문장이 매우 길 때 (예: 긴 대화나 문서 요약), 계산 시간이 매우 깁니다.
  • 이때는 계산이 길어지는 동안, 데이터를 캐시 영역으로 옮기는 작업을 계산과 동시에 (혹은 그 틈을 타서) 수행합니다.
  • 장점: 복잡한 동기화 없이도, 긴 입력일 때는 계산 시간이 데이터 이동 시간을 완전히 덮어버려서 효율적입니다.

---

🌟 핵심 기술: "스위즐 (Swizzled) 메모리 복사"

이 두 방법의 핵심은 데이터를 옮길 때 단순히 복사하는 게 아니라, 모양을 바꿔서 복사한다는 점입니다.

• 문제: CPU 가 좋아하는 모양 (Host-friendly) 과 PIM 이 좋아하는 모양 (PIM-aware) 이 다릅니다. 마치 **서양식 식탁 (CPU)**과 **일본식 다다미 (PIM)**에 음식을 배치하는 방식이 다르듯이요.
• 해결: 데이터를 메모리에서 꺼내서, **PIM 이 좋아하는 모양으로 바로 정리 (Swizzle)**한 뒤 캐시 영역에 넣습니다.
• 효과: 이렇게 하면 CPU 는 원래대로 계산할 수 있고, PIM 은 다음 단계에서 바로 사용할 수 있습니다.

---

📊 결과: 얼마나 좋을까요?

이 방법을 적용한 실험 결과 (Llama 3.2 모델 기준):

1. 메모리 절약: 기존에 모델을 두 배로 복사하던 방식보다 약 48% 의 메모리를 아꼈습니다. (예: 12GB 램이 필요한 스마트폰에서 6GB 정도만 더 쓰면 됨)
2. 속도: 메모리 부족으로 모델을 줄일 필요 없이, 이론상 가능한 최대 속도와 거의 비슷한 성능을 냅니다.
3. 호환성: 하드웨어를 수정할 필요 없이, 기존 스마트폰 (갤럭시 S24+ 등) 에서 소프트웨어 업데이트만으로 작동합니다.

💡 요약

PIM-SHERPA는 스마트폰 AI 가 "캐시 안에 있어야 할지, 밖으로 나가야 할지" 고민하는 양면성 문제를 해결한 혁신적인 소프트웨어입니다.

• 기존 방식: "두 개의 AI 모델을 만들어서 메모리를 2 배로 쓰자" (비효율적, 메모리 부족)
• PIM-SHERPA: "하나의 모델을 가지고, 계산할 때는 캐시에서, PIM 할 때는 캐시 밖으로 유연하게 이동시키자" (효율적, 메모리 절약)

이 기술 덕분에 앞으로 우리 스마트폰은 더 큰 AI 모델을 더 빠르고, 더 적은 메모리로 구동할 수 있게 될 것입니다. 마치 무거운 짐을 나르는 사서가 등산로를 미리 파악해서, 등산객이 숨 쉴 틈도 없이 짐을 나르는 것과 같습니다! 🏔️🎒

기술 요약

1. 문제 정의 (Problem Statement)

온디바이스 (On-device) 환경에서 대규모 언어 모델 (LLM) 을 실행할 때, 메모리 인-프로세싱 (PIM, Processing-in-Memory) 기술을 활용하면 메모리 대역폭 병목 현상을 해결할 수 있습니다. 특히 LLM 추론의 디코딩 (Decode) 단계는 메모리 대역폭에 민감하므로 PIM 에 적합합니다. 그러나 PIM 을 실제 제품 (예: LPDDR-PIM) 에 적용할 때 두 가지 주요 불일치 (Inconsistency) 문제가 발생합니다.

1. 메모리 속성 불일치 (Memory Attribute Inconsistency):

   • 프리필 (Prefill) 단계: 많은 토큰을 처리하므로 연산 집약적 (GEMM) 입니다. 캐시 재사용을 극대화하기 위해 가중치 (Weights) 를 캐시 가능 (Cacheable) 영역에 배치해야 합니다.
   • 디코딩 (Decode) 단계: 한 번에 하나의 토큰을 처리하므로 메모리 대역폭 의존적 (GEMV) 입니다. PIM 을 활성화하려면 DRAM 요청이 메모리 컨트롤러에 도달해야 하므로, 가중치가 캐시 불가능 (Non-cacheable) 영역에 있어야 합니다. (캐시 히트가 발생하면 PIM 실행이 차단됨)
   • 문제: 두 단계가 동일한 모델에서 연속적으로 수행되므로, 가중치에 대해 서로 상반된 메모리 속성이 요구됩니다.

2. 가중치 레이아웃 불일치 (Weight Layout Inconsistency):

   • 호스트 (Host) 친화적 레이아웃: 채널 및 뱅크 인터리빙을 최적화하기 위해 연속적인 메모리 접근을 선호합니다.
   • PIM 인지 레이아웃: PIM 블록 내 SIMD 연산을 최대화하기 위해 각 DRAM 뱅크 내에서 행 (Row) 단위로 데이터가 연속적으로 배치되어야 합니다.
   • 문제: 두 레이아웃이 상충되어, 한 가지 레이아웃만으로는 두 단계 모두에서 최적의 성능을 내기 어렵습니다.

기존 연구들은 이러한 불일치를 해결하기 위해 가중치 복제 (Weight Duplication) 방식을 사용하거나 (캐시 가능/불가능 버전 각각 저장), 하드웨어 변경 (메모리 컨트롤러 수정) 을 요구하여, 모바일 환경의 제한된 DRAM 용량을 소모하거나 적용성을 떨어뜨리는 문제가 있었습니다.

2. 제안 방법론: PIM-SHERPA

저자들은 하드웨어 수정 없이 소프트웨어만으로 위 두 가지 불일치를 해결하는 PIM-SHERPA를 제안합니다. 핵심 아이디어는 작은 캐시 가능 버퍼를 할당하고, 런타임에 가중치를 동적으로 이동 및 변환하는 것입니다.

주요 기법

1. DRAM 더블 버퍼링 (DRAM Double Buffering, DDB):

   • 원리: 현재 레이어의 GEMM 연산이 수행되는 동안, 다음 레이어의 가중치를 캐시 가능 버퍼로 미리 가져옵니다.
   • 동작: 두 개의 작은 버퍼 (Buffer 0, Buffer 1) 를 번갈아 사용합니다. 한 버퍼에서 연산이 진행되는 동안 다른 버퍼로 다음 레이어의 가중치를 PIM 인지 레이아웃에서 호스트 친화적 레이아웃으로 변환하며 복사 (Swizzled Memory Copy) 합니다.
   • 효과: 데이터 이동 (복사) 오버헤드를 연산 시간과 겹쳐 (Overlap) 숨깁니다.

2. 온라인 가중치 재배열 및 스위즐드 메모리 복사 (Online Weight Rearrangement with Swizzled Memory Copy, OWR):

   • 원리: 입력 시퀀스 길이가 길어 GEMM 연산 시간이 지배적인 경우, 연산 직전에 필요한 가중치만 단일 버퍼로 복사하여 변환합니다.
   • 동작: 복잡한 동기화 없이 순차적으로 복사 후 연산을 수행합니다.
   • 효과: 구현 복잡도를 낮추고, 긴 시퀀스 환경에서 복사 오버헤드가 전체 지연 시간에서 차지하는 비중을 줄입니다.

3. 스위즐드 메모리 복사 (Swizzled Memory Copy, SMC):

   • PIM 인지 레이아웃 (비캐시 영역) 에서 호스트 친화적 레이아웃 (캐시 영역) 으로 데이터를 이동시키면서, DRAM 뱅크/행/열 매핑을 변환하는 소프트웨어 기반 복사 기법입니다.

3. 주요 기여 (Key Contributions)

• 문제 식별: PIM 기반 LLM 추론에서 프리필과 디코딩 단계 간의 메모리 속성 불일치와 가중치 레이아웃 불일치를 최초로 명확히 규명했습니다.
• 소프트웨어 솔루션: 하드웨어 변경 없이 작은 캐시 버퍼와 SMC 를 활용하여 위 불일치를 해결하는 두 가지 방법 (DDB, OWR) 을 제안했습니다.
• 메모리 용량 최적화: 기존 가중치 복제 방식 대비 DRAM 용량 소모를 획기적으로 줄였습니다.
• 실제 제품 환경 검증: 삼성 갤럭시 S24+ (Exynos 2400) 와 LPDDR5X-PIM 에뮬레이션을 통해 실제 온디바이스 환경에서의 유효성을 입증했습니다.

4. 실험 결과 (Results)

Llama 3.2 (1B, 3B) 모델을 Samsung Galaxy S24+ 에서 평가한 결과는 다음과 같습니다.

• DRAM 용량 절감:
  • 기존 가중치 복제 (Weight Duplication) 방식 대비 **약 47.8% ~ 49.7%**의 DRAM 용량을 절감했습니다.
  • 예: Llama 3.2 3B 모델의 경우, 복제 방식은 약 12GB 이상의 DRAM 이 필요하여 모바일 기기에서 실행 불가능했으나, PIM-SHERPA 는 약 6GB 만으로 실행 가능하게 만들었습니다.
• 성능 (TTFT - Time to First Token):
  • DDB: 입력 시퀀스 길이 (SL) 가 128 이상일 때, 이론적 최대치에 근접하는 성능을 내며 FACIL (하드웨어 수정 방식) 과 유사한 TTFT 를 달성했습니다.
  • OWR: SL 이 증가함에 따라 성능이 DDB 및 FACIL-O 에 근접했습니다 (SL 192 기준 16.7 TPS).
  • 비교: PIM-SHERPA 는 하드웨어 수정이 필요한 FACIL 과 유사한 성능을 내면서도 소프트웨어만으로 구현 가능합니다.
• 제한된 DRAM 환경:
  • DRAM 용량이 부족해 가중치 한 복사본만 저장 가능한 환경에서도 PIM-SHERPA 는 PIM 을 활용한 속도 향상 (Speedup) 을 제공하며, 캐시 불가능 영역에서의 단순 GEMM (NC-GEMM) 보다 훨씬 효율적이었습니다.

5. 의의 및 결론 (Significance)

• 실용성: 하드웨어 수정 없이 소프트웨어만으로 PIM 의 한계를 극복하여, 모바일 및 엣지 디바이스에서 고품질 LLM 추론을 가능하게 합니다.
• 상용화 촉진: 메모리 용량 제약이 심각한 모바일 환경에서 PIM 기술의 상용화를 위한 핵심 과제를 해결했습니다.
• 확장성: 제안된 기법은 GPU 나 NPU 시스템의 메모리 계층 구조에도 적용 가능하여 범용적인 솔루션으로 평가됩니다.

요약하자면, PIM-SHERPA는 온디바이스 LLM 추론에서 발생하는 메모리 속성 및 레이아웃 충돌 문제를 소프트웨어 기법 (DDB, OWR) 으로 해결함으로써, 하드웨어 변경 없이도 PIM 의 이점을 극대화하고 메모리 효율성을 획기적으로 개선한 획기적인 연구입니다.