Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

이 논문은 동적인 LLM 서빙 워크로드를 효율적으로 처리하기 위해 3D-DRAM 기반의 하이브리드 본딩 아키텍처인 'Helios'를 제안하며, 공간 인식 KV 캐시 할당 및 분산 타일 어텐션 실행을 통해 기존 설계 대비 성능과 에너지 효율성을 획기적으로 개선함을 보여줍니다.

핵심

1. 문제: 왜 지금의 AI는 지치나요? (기존 기술의 한계)

지금까지 AI(대형 언어 모델) 를 서비스할 때는 마치 거대한 도서관을 운영한다고 생각해보세요.

• 기존 방식 (GPU): 도서관에 책 (데이터) 이 쌓이면, 사서 (프로세서) 가 책장을 모두 뒤져서 필요한 책을 찾아옵니다. 책이 너무 많으면 사서가 책을 찾는 데만 시간이 걸려서 손님이 기다리는 시간이 길어집니다.
• 기존의 개선 시도 (NMP): 사서에게 책장 바로 옆에 작은 책상 (처리 엔진) 을 만들어주었습니다. 이렇게 하면 책 찾는 속도는 빨라졌지만, 책장 관리 방식이 너무 구식이었습니다.
  • 문제점 1 (고정된 책장): "한 손님이 한 번 들어오면, 그 손님이 읽을 모든 책 (키-밸류 캐시) 을 미리 한 책장에 꽉 채워놔야 한다"는 규칙이 있었습니다. 손님이 책을 10 권만 읽을지 1,000 권을 읽을지 모르는데, 미리 1,000 권 분량의 공간을 다 차지해버리는 것입니다.
  • 문제점 2 (비효율적인 이동): 손님이 다른 책장으로 이동할 때, 사서들이 서로 책장을 오가며 책을 옮기는 데 너무 많은 시간이 걸렸습니다.

결과적으로, 손님이 갑자기 몰리거나 (요청 폭주), 손님이 아주 긴 이야기를 할 때 (긴 문맥), 도서관은 혼란에 빠지고 서비스 속도가 느려집니다.

2. 해결책: 헬리오스 (Helios) 의 등장

헬리오스는 이 문제를 해결하기 위해 두 가지 혁신을 가져왔습니다.

A. "레고 블록" 같은 책장 관리 (동적 KV 캐시 할당)

헬리오스는 손님이 들어오자마자 미리 책을 다 채워두지 않습니다. 대신, 손님이 실제로 읽는 만큼만 '작은 책 (블록)'을 빌려줍니다.

• 비유: 기존 방식이 "한 손님을 위해 큰 방 전체를 예약해두는 것"이라면, 헬리오스는 **"손님이 앉는 의자 하나하나만 빌려주는 것"**과 같습니다.
• 효과:
  • 손님이 짧게 대화하면 작은 공간만 쓰고, 길게 대화하면 더 많은 공간을 채웁니다.
  • 도서관의 공간 낭비가 사라져서, 같은 크기의 도서관에 훨씬 더 많은 손님을 동시에 모실 수 있습니다.

B. "초고속 연결"과 "유연한 사서 팀" (하이브리드 본딩 & 분산 처리)

헬리오스는 책장 (메모리) 과 사서 (처리 장치) 를 3D 레고처럼 밀착시켜 붙였습니다 (하이브리드 본딩).

• 비유: 기존에는 사서와 책장이 떨어져 있어서 책이랑 사서 사이를 오가는 데 시간이 걸렸다면, 헬리오스는 사서가 책장 바로 위에 앉아있는 형태입니다.
• 유연한 팀워크: 사서들이 고정된 책장만 담당하는 게 아니라, 모든 사서가 서로 협력합니다.
  • 어떤 손님이 긴 이야기를 하면, 사서들이 모여서 그 이야기를 나누어 처리합니다.
  • 어떤 사서가 바쁘면, 옆 사서가 도와줍니다.
  • 이렇게 하면 어떤 손님이 오더라도 가장 효율적으로 일할 수 있는 사서 팀이 바로 구성됩니다.

3. 결과: 얼마나 빨라졌나요?

이 새로운 방식 (헬리오스) 을 적용한 결과, 기존 방식에 비해 놀라운 성과가 나왔습니다.

• 속도: 기존 칩보다 약 3.25 배 더 빠릅니다. (손님이 답변을 기다리는 시간이 획기적으로 줄어듦)
• 전력 효율: 같은 일을 하는 데 드는 전기는 약 3.36 배 더 절약됩니다. (전기세와 발열 문제 해결)
• 대기 시간: 특히 긴 대화를 할 때나, 손님이 몰릴 때 (P99 지수 기준) 대기 시간이 최대 76% 까지 단축되었습니다.

4. 요약: 왜 이것이 중요한가요?

지금까지 AI 서비스는 "손님이 많으면 느려지고, 긴 대화를 하면 비싸다"는 한계가 있었습니다.

헬리오스는 "책장 (메모리) 과 사서 (계산기) 를 밀착시키고, 손님의 필요에 따라 공간을 유연하게 나누어 쓰는" 방식을 도입했습니다. 이는 마치 스마트한 도서관이 되어, 손님이 아무리 많아도, 아무리 긴 이야기를 해도 빠르고 저렴하게 서비스를 제공할 수 있게 해줍니다.

이 기술이 상용화되면, 우리가 사용하는 AI 챗봇은 더 똑똑해지고, 더 많은 사람이 동시에 사용해도 멈추지 않게 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 의 온라인 서비스 배포는 매우 동적인 워크로드 (변동하는 요청 도착률, 다양한 시퀀스 길이) 를 처리해야 합니다. 기존 LLM 서빙 시스템은 계산 집약적 (Prefill) 과 메모리 집약적 (Decoding) 연산이 공존하는 환경에서 효율적인 처리가 어렵습니다.

기존의 Near-Memory Processing (NMP) 기반 가속기 설계는 다음과 같은 한계로 인해 동적인 LLM 서빙 워크로드를 효율적으로 처리하지 못합니다:

• 거친 KV 캐시 관리 (Coarse-grained KV Cache Management): 기존 NMP 는 각 어텐션 헤드의 전체 KV 캐시를 고정된 NMP 모듈 (예: HBM 큐브) 에 정적으로 할당합니다. 이로 인해 요청 길이가 가변적일 때 메모리 단편화 (Fragmentation) 가 발생하고, 컴퓨팅 부하 불균형이 초래되어 가속기의 유효 서빙 용량이 감소합니다.
• 유연하지 않은 어텐션 실행 흐름 (Inflexible Attention Execution Flow): 기존 설계는 전체 시퀀스에 대한 KV 캐시가 동일한 NMP 모듈 내에 있어야 하며, 어텐션 연산을 독립적인 GEMM(General Matrix Multiply) 단계로만 실행합니다. 이는 세분화된 KV 캐시 관리와 동적인 요청 처리를 방해합니다.
• 구조적 불일치: NMP 아키텍처의 비균일 메모리 추상화 (각 PE 가 로컬 메모리만 직접 접근 가능) 와 GPU 기반의 동적 KV 캐시 관리 알고리즘 간의 근본적인 격차가 존재합니다.

2. 제안된 방법론 (Methodology: Helios)

저자들은 이러한 문제를 해결하기 위해 Hybrid Bonding (HB) 기술을 기반으로 한 Helios라는 LLM 서빙 가속기를 제안합니다. Helios 는 하드웨어 아키텍처와 소프트웨어 시스템 설계를 통합 (Co-design) 하여 다음과 같은 혁신을 이루었습니다.

A. 아키텍처 설계 (Architecture)

• 하이브리드 본딩 (Hybrid Bonding) 활용: Cu-Cu 직접 결합을 통해 DRAM 다이와 로직 다이 (Processing Engines, PEs) 를 3D 적층합니다. 이는 기존 HBM 대비 높은 대역폭 (110,000/mm² I/O 밀도) 과 낮은 전력 소모를 제공합니다.
• 분산 타일링 어텐션 (Distributed Tiled Attention):
  • KV 캐시를 세분화된 블록 (예: 64 토큰 단위) 으로 나누고, 이를 모든 PE 에 분산 저장합니다.
  • 각 PE 는 로컬 KV 블록에 대해 Iterative Tiled Attention을 수행하며, 온라인 Softmax 와 부분 합 (Partial Sum) 을 계산합니다.
  • PE 간 통신을 통해 최종 어텐션 결과를 집계합니다.
• PE 아키텍처: 각 PE 는 행렬 연산 유닛, 온라인 Softmax 유닛, 벡터 유닛, 축소 (Reduction) 유닛을 포함하며, Mesh 토폴로지를 통해 PE 간 통신을 수행합니다.

B. 시스템 및 스케줄링 설계 (System & Scheduling)

• 공간 인식 KV 캐시 할당 (Spatially-Aware KV Cache Allocation):
  • 요청의 KV 블록을 할당할 때 단순히 메모리 여유 공간만 고려하는 것이 아니라, PE 의 계산 부하 균형과 **PE 간 통신 비용 (NoC 트래픽)**을 동시에 최적화합니다.
  • 특정 PE(중앙부) 로의 트래픽 집중을 방지하고, 통신 경로를 고려하여 블록을 배치합니다.
• 동적 KV 캐시 관리: 요청의 수명 주기에 따라 KV 블록을 동적으로 할당 및 해제하여 메모리 단편화를 제거하고, 가변적인 요청 길이를 효율적으로 처리합니다.

C. 연산자 실행 흐름 (Operator Execution)

• ** intra-PE 최적화:** GEMM 연산과 벡터 연산 (Softmax 등) 을 겹쳐서 (Overlap) 실행하여 지연 시간을 숨깁니다.
• Inter-PE 통신:
  • Query 벡터는 모든 PE 에 복제 (All-Gather) 됩니다.
  • Key/Value 벡터는 해당 요청의 마지막 블록을 가진 PE 로만 전송되도록 최적화된 통신 패턴 (Reduce-Scatter 등) 을 사용합니다.
  • MLA(Multi-Head Latent Attention) 와 같은 최신 아키텍처를 지원하기 위해 행렬 흡수 (Matrix Absorption) 기법을 적용하여 중간 KV 텐서의 재구성을 방지합니다.

3. 주요 기여 (Key Contributions)

1. NMP 기반 동적 KV 캐시 관리의 한계 분석: 기존 NMP 설계가 동적인 LLM 서빙 워크로드에서 겪는 메모리 단편화와 부하 불균형 문제를 심층 분석했습니다.
2. Helios 아키텍처 제안: 하이브리드 본딩 기술을 활용하여 NMP 네이티브의 분산 타일링 어텐션을 가능하게 하는 하드웨어 설계를 제시했습니다.
3. 공간 인식 할당 알고리즘: 계산 부하 균형과 통신 오버헤드 최소화를 동시에 달성하는 동적 KV 캐시 할당 전략을 개발했습니다.
4. 성능 검증: 기존 GPU 및 NMP 설계 대비 압도적인 성능 개선을 입증했습니다.

4. 실험 결과 (Results)

다양한 LLM 모델 (OPT, LLaMA3, Mixtral, Qwen, DeepSeek) 과 워크로드 조건에서 평가한 결과입니다.

• 속도 향상 (Speedup): 기존 GPU 및 NMP 설계 대비 기하평균 3.25 배의 속도 향상을 달성했습니다.
• 에너지 효율 (Energy Efficiency): 기하평균 3.36 배 더 높은 에너지 효율을 보였습니다.
• 서빙 성능 (Serving Performance):
  • P50/P99 Time-Between-Tokens (TBT): 기존 설계 대비 최대 72% / 76% 감소 (지연 시간 단축).
  • 처리량 (Throughput): 기존 NMP 설계 대비 최대 7.33 배 높은 처리량을 기록했습니다.
• 메모리 단편화: 기존 NMP 설계가 겪는 메모리 단편화 비율이 Helios 에 비해 최대 23 배까지 높았으며, Helios 는 이를 0.2~15.0% 수준으로 억제하여 높은 부하에서도 성능 저하가 없음을 확인했습니다.
• 부하 균형: PE 간 계산 부하 불균형 비율을 7.1~13.3% 수준으로 유지하여 자원 활용도를 극대화했습니다.

5. 의의 및 결론 (Significance)

이 논문은 **하드웨어 - 소프트웨어 공동 설계 (Co-design)**의 중요성을 다시 한번 입증했습니다. 단순히 메모리 대역폭을 늘리는 것을 넘어, LLM 서빙의 **동적인 특성 (가변적 요청 길이, 유동적 부하)**에 맞춰 메모리 관리 전략과 어텐션 실행 흐름을 재설계함으로써, 3D DRAM 기반 NMP 가속기의 잠재력을 완전히 끌어올렸습니다.

Helios 는 대규모 LLM 서빙 서비스에서 발생하는 메모리 병목 현상과 부하 불균형 문제를 해결하여, 더 많은 사용자에게 더 빠르고 에너지 효율적인 AI 서비스를 제공할 수 있는 새로운 아키텍처 패러다임을 제시합니다. 특히, 하이브리드 본딩 기술을 활용한 3D 적층 구조와 세분화된 KV 캐시 관리의 결합은 차세대 LLM 가속기 설계의 중요한 방향성을 제시합니다.