TPU, GPU를 넘어 AI 인프라의 새로운 지평을 열다

TPU는 딥러닝(Deep Learning)의 핵심 연산인 텐서(Tensor) 간의 행렬 곱 연산을 효율적으로 처리하기 위해 설계되었다. CPU는 범용적인 연산에 강점을 가지지만, 대규모 행렬 연산을 반복하기에는 비효율적이다. GPU는 병렬 연산에 특화되었지만, 그래픽 처리라는 범용 목적을 함께 고려해야 한다. TPU는 이러한 문제를 해결하기 위해, 시스톨릭 어레이(Systolic Array)와 온칩 메모리(On-chip Memory)를 활용하여 전력 효율을 극대화했다.

시스톨릭 어레이(Systolic Array)는 TPU의 핵심 구조로, 수백에서 수천 개의 MAC 유닛(Multiply-Accumulate Unit)을 바둑판처럼 배열하여 행렬 곱 연산을 수행한다. A의 값은 왼쪽에서 오른쪽으로, B의 값은 위에서 아래로 흘러가며, 각 MAC 유닛은 곱셈과 덧셈을 수행한다. 이 구조는 중간 계산 결과를 메모리에 저장하지 않고, 칩 내부에서 연속적으로 처리하므로 메모리 접근(Memory Access) 병목 현상을 줄이고 전력 소모를 감소시킨다.

CPU는 운영체제(Operating System)와 애플리케이션 로직(Application Logic)과 같은 복잡한 제어 흐름을 처리하는 데 최적화되어 있다. GPU는 대규모 행렬 연산에 강점을 가지며, 딥러닝 학습과 추론의 표준 가속기로 자리 잡았다. TPU는 딥러닝 연산에 특화된 전용 가속기로, 시스톨릭 어레이(Systolic Array)를 통해 대규모 행렬 곱을 효율적으로 처리한다. 결과적으로 CPU는 시스템 조율, GPU는 범용 AI 연산 가속, TPU는 대규모 AI 학습(Large-scale AI Training)에 특화된 역할을 수행한다.

JAX는 TPU의 하드웨어 특성을 소프트웨어(Software) 차원에서 효과적으로 활용하기 위해 설계된 프레임워크이다. JAX는 사용자가 작성한 수치 연산 코드를 계산 그래프(Computation Graph)로 변환하고, XLA(Accelerated Linear Algebra) 컴파일러를 통해 TPU에 최적화된 저수준 연산 형태로 컴파일한다. XLA는 연산 순서를 재배치하고, 불필요한 메모리 접근을 줄여 TPU가 시스톨릭 어레이(Systolic Array)를 최대한 효율적으로 활용하도록 돕는다. 즉, JAX는 TPU의 성능 극대화(Performance Maximization)를 위한 핵심 요소이다.

AI 모델의 대형화와 적용 범위 확장에 따라, 단일 가속기로 모든 요구를 만족시키기 어려워졌다. NPU(Neural Processing Unit)는 전력 소모와 지연 시간이 민감한 환경에서, 상대적으로 작은 모델을 빠르게 실행하는 데 초점을 맞춘다. DPU(Data Processing Unit)는 AI 학습과 추론 과정에서 발생하는 데이터 처리, I/O, 보안 작업을 전담하여 CPU의 부담을 줄이고, GPU와 TPU가 연산에 집중할 수 있도록 돕는다. 이러한 역할 분담(Role Division)을 통해 AI 인프라의 효율성과 확장성을 높이고 있다.