TPU에서 효율적인 생성 모델 추론을 위한 메모리 내 컴퓨팅 활용

생성 모델의 급속한 발전으로 인해 이러한 모델을 특수 하드웨어에 효율적으로 배치하는 것이 중요해졌습니다. 텐서 처리 장치(TPU)는 AI 작업 부담을 가속화하기 위해 설계되었지만, 높은 전력 소비로 인해 효율성을 개선하기 위한 혁신이 필요합니다. 메모리 내 컴퓨팅(CIM)은 우수한 면적과 에너지 효율성을 가진 유망한 패러다임으로 떠올랐습니다. 본 연구에서는 행렬 곱셈 장치(MXU)에서 기존의 디지털 시스톨릭 어레이를 대체하기 위해 디지털 CIM을 통합한 TPU 아키텍처를 제시합니다. 먼저 다양한 생성 모델 추론을 위한 CIM의 이점을 평가하기 위해 CIM 기반 TPU 아키텍처 모델과 시뮬레이터를 구축합니다. 관찰된 설계 통찰을 바탕으로, CIM 기반 TPU 아키텍처 설계 선택을 추가로 탐구합니다. 대형 언어 모델과 확산 변환기 추론을 위해 최대 44.2% 및 33.8%의 성능 향상과 MXU의 에너지 소비를 27.3배 줄일 수 있는 다양한 설계 선택을 통해 baseline TPUv4i 아키텍처와 비교하여 달성할 수 있습니다.