GPU는 어떻게 머신러닝 연산을 수행할까?

도입하며

딥러닝 프레임워크에서 모델을 학습할 때, 보통 model.fit(x, y)이나 jax.jit(f)(x) 같은 고수준 API를 호출하여 실행한다. 하지만 이러한 호출이 실제 하드웨어 수준에서 어떤 절차를 거쳐 실행되는지 설명하기는 쉽지 않다. 예를 들어 JAX의 jax.numpy.dot 연산이 실행될 때, 하드웨어 수준에서는 어떤 일이 발생하는지 살펴보자.

jax_dot_example.py
import jax.numpy as jnp

c = jnp.dot(a, b)

겉보기에는 단순한 행렬 곱이지만, 실제로는 다음과 같은 단계가 순차적으로 수행된다. Python 인터프리터가 연산을 JAX로 전달하고, JAX는 이를 XLA 컴파일러에 넘긴다. 이후 XLA는 연산 그래프를 최적화하고 GPU용 커널을 생성하며, 드라이버는 생성된 커널을 GPU 명령 큐에 적재한다. 마지막으로 GPU 스케줄러는 커널을 수천 개의 연산 유닛에 분배하여 실행한다.

딥러닝 프레임워크와 GPGPU^{General-purpose computing on GPU} 기술은 위와 같은 복잡한 과정 대부분을 자동화하고 추상화하므로, 사용자 입장에서는 이를 의식하지 않아도 된다. 그러나 이러한 동작 원리를 전혀 이해하지 못하면 성능 최적화 혹은 문제 원인 분석이 어려워질 수 있다. 예를 들어 특정 배치 사이즈에서 OOM^{Out of Memory}이 발생하는 이유, 메모리가 충분하게 보임에도 메모리 단편화^{Fragmentation}로 인한 충돌이 발생하는 이유, 혹은 연산량이 작음에도 커널 실행 오버헤드로 인해 속도가 낮아지는 현상 등은 하드웨어 관점의 이해 없이는 분석하기 어렵다.

이 글에서는 Python 코드 실행 단계에서부터 하드웨어 자원까지의 처리 과정을 순차적으로 설명할 예정이다. 이 글을 모두 읽으면 아래와 같은 내용들을 얻어갈 수 있다.

• Python 코드에서 부터 하드웨어까지 전체 실행 과정
• GPU 내부 관점에서 명령 처리 방식
• 하드웨어의 특성을 최대한 활용한 성능 최적화

머신러닝과 딥러닝의 원리

머신러닝과 딥러닝 모델의 학습^Training과 추론^Inference은 일반적으로 텐서^Tensor [1] 연산의 연속으로 볼 수 있다. 백만~수십억 개의 파라미터를 갖는 모델이 입력을 받아 출력을 계산하는 과정은 다음과 같이 나타난다.

$$Y = f(W, X)$$

이는 행렬 곱 및 요소별 연산의 조합으로 구성된다.

import jax.numpy as jnp

# (2×3) 입력 텐서
x = jnp.array([[1.0, 2.0, 3.0],
               [4.0, 5.0, 6.0]])

# (3×2) 파라미터 텐서
w = jnp.array([[0.1, 0.2],
               [0.3, 0.4],
               [0.5, 0.6]])

# 행렬곱 → (2×2)
y = jnp.dot(x, w)
print(y)

병렬 연산의 중요성

대부분의 딥러닝 모델의 구조는 GEMM^{General Matrix Multiply} 행렬 곱 연산을 수행한다. 이는 BLAS의 Level 3 연산에 해당하며 다음과 같이 표현된다. [BLAS]

$$C = \alpha AB + \beta C$$

행렬 $A$가 $(M, K)$, $B$가 $(K, N)$일 경우 결과 $C$는 $(M, N)$이 된다.

여기서 BLAS는 선형 대수 연산을 Level 1~3으로 정의했다. 매트릭스간 연산인 GEMM은 BLAS상에서는 Level 3 연산에 해당된다. 다음에서 BLAS의 각 Level의 의미를 살펴보자.

• Level 1: 요소별 연산 $O(N)$
• Level 2: 벡터·행렬 연산 $O(MN + MK + KN)$
• Level 3: 행렬·행렬 연산 $O(MNK)$

Level 3 연산은 데이터 이동량 $O(MN + MK + KN)$ 대비 연산량이 $O(MNK)$으로 크다. 즉, 연산 밀도가 높기에 Compute Bound 특성을 갖는다. 이는 GPU의 연산 성능^FLOPS을 최대한 활용하기 유리하다. 아래 의사 코드는 데이터 이동량과 연산량이 각각 위 수식으로 도출되는 이유를 보여준다.

# A: (M, K)
# B: (K, N)
# C: (M, N)

for i in range(M):
    for j in range(N):
        acc = 0.0
        for k in range(K):
            acc += A[i][k] * B[k][j]  # O(MNK)
        C[i][j] = acc  # O(MN)

데이터 접근은 A, B, C 각각에 대해 대략 $M \times K$, $K \times N$, $M \times N$ 수준이 필요하므로 전체 이동량은 $O(MN + MK + KN)$이 된다. 연산량은 $M \times N \times K$이므로 $O(MNK)$이다.

연산 밀도가 높은 작업은 GPU의 FLOPS를 최대한 활용할 수 있으나, 반대로 연산량이 적고 데이터 이동량이 큰 연산은 메모리 대역폭에 의해 병목이 발생할 수 있다. 이에 따라 최신 딥러닝 컴파일러는 가능한 많은 연산을 GEMM으로 변환하거나, 요소별 연산을 융합^Fusion하여 메모리 접근을 최소화하는 방향으로 최적화를 수행한다.

역전파

딥러닝 모델의 학습 과정인 역전파^{Backpropagation}는 연쇄 법칙^{Chain Rule}[2]을 통해 손실 함수 $L$에 대한 파라미터 $W$의 기울기^Gradient를 구하는 과정이다. 예를 들어 $Y = W \cdot X$와 같은 선형 변환의 경우 기울기는 다음과 같이 주어진다.

$$\frac{\partial L}{\partial W} = \frac{\partial L}{\partial Y} \cdot X^T$$

역전파 과정에서는 순전파와 마찬가지로 행렬 곱 연산이 발생한다. 따라서 딥러닝 학습에서는 순전파^{Forward Pass}뿐만 아니라 역전파^{Backward Pass}에 포함된 GEMM 연산도 효율적으로 수행할 필요가 있다.

또한 역전파 계산에 필요한 입력 $X$와 출력 기울기 $dL/dY$를 보존해야 하므로, 활성화^Activation를 저장하는 과정에서 메모리 사용량이 증가한다. 이로 인해 GPU 메모리 용량이 주요 제약으로 작용하며, 활성화 체크포인트^{Activation Checkpointing} 기술이 활용된다. [3]

import jax
import jax.numpy as jnp

# y = Wx
def forward(w, x):
    return w @ x

# L2 loss
def loss(w, x):
    y = forward(w, x)
    return jnp.sum(y**2)

# ∂L/∂W = (∂L/∂Y) · X^T
grad_fn = jax.grad(loss)

in_dim, out_dim = 3, 2
w = jnp.ones((out_dim, in_dim))
x = jnp.array([1.0, 2.0, 3.0])

dw = grad_fn(w, x)
print(dw)

위 코드를 보면 역전파 또한 내부적으로 또 다른 행렬 곱을 포함한다는 점을 유추할 수 있다.

CPU의 한계와 GPU, NPU

딥러닝이나 그래픽 처리와 같은 특수 목적이 아닌 경우 일반적으로 연산은 CPU에서 수행된다. 그렇다면 딥러닝에서는 왜 GPU를 사용할까? 이를 이해하기 위해서는 각 아키텍처의 설계 목적과 철학을 살펴볼 필요가 있다. 다음 영상은 CPU와 GPU의 특징을 잘 보여준다.

CPU의 특징

CPU는 지연 시간^Latency을 최소화하고 복잡한 제어 흐름을 빠르게 처리하도록 설계된 범용 프로세서이다.

• 복잡한 제어 처리: CPU는 if~else 같은 조건 분기, 함수 호출, 다양한 명령 흐름을 빠르게 처리하기 위해 내부적으로 정교한 제어 장치를 갖추고 있다. 그래서 웹 브라우징, 운영체제, 게임 로직처럼 분기와 제어가 많은 작업에 강하다.
• 캐시 활용: CPU에는 L1, L2, L3 같은 여러 단계의 캐시가 있어서 자주 쓰는 데이터를 빠르게 꺼내 쓸 수 있다. 일반적인 프로그램은 같은 데이터를 반복해서 쓰는 경우가 많기에 자주 쓰는 데이터를 캐시에 저장하여 사용하면 유리하다.
• 벡터 연산이 가능하나 규모가 작음: CPU도 AVX 같은 벡터 연산 명령을 통해 한번에 여러 데이터를 계산할 수 있다. 하지만 한 번에 다룰 수 있는 데이터의 양(레지스터 크기)도 작고 코어 수도 적기 때문에, 딥러닝처럼 수천~수만 개 연산을 동시에 처리해야 하는 작업에는 한계가 있다. [4]

CPU의 SIMD^{Single Instruction Multiple Data} 기술은 하나의 명령어로 여러 데이터를 한꺼번에 처리하는 방식이다. 예를 들어 AVX-512 명령을 사용하면 512비트 레지스터 하나로 32비트 실수 16개를 동시에 계산할 수 있다. 다만 이런 방식은 여전히 CPU 내부의 소수 코어가 벡터 연산만 잠깐 가속하는 수준이기 때문에, 수천 개의 연산 코어가 동시에 일을 분담하는 GPU의 병렬 처리 능력과 비교하면 규모가 상대적으로 작다. 또한 CPU는 분기 처리나 복잡한 로직을 담당하기 위한 회로 비중이 커서, 순수 계산만을 위해 많은 연산 유닛을 넣기에는 구조적인 제한이 있다.

GPU

GPU는 단순한 명령을 수천 개 단순 연산을 수행하는 코어로 병렬 실행하여 처리량^Throughput을 극대화하도록 설계되었다.

• 수천 개의 연산 코어: CPU 코어처럼 복잡한 일을 잘하진 않지만, 훨씬 많은 코어를 한꺼번에 돌릴 수 있어서 대량의 연산을 동시에 처리할 수 있다.
• 고대역폭 메모리 HBM 활용: GPU는 연산만 빠른 것이 아니라 데이터를 읽고 쓰는 속도도 매우 빠르다. HBM 같은 고대역폭 메모리를 사용해서 CPU 메모리보다 수배 이상 빠르게 데이터를 공급할 수 있다. 딥러닝처럼 데이터를 계속 흘려보내야 하는 작업에 유리하다.
• 하나의 명령을 여러 스레드로 처리: GPU는 비슷한 연산을 한꺼번에 반복하는 구조에 최적화되어 있다. 예를 들어, 행렬의 모든 원소에 똑같이 곱하는 명령을 수십, 수백 개 스레드에 동시에 적용할 수 있다. 이런 방식이 GPU의 병렬성을 극대화한다.

위 그림은 CPU와 GPU가 다른 목적을 염두에 두고 설계되었음을 보여 준다. CPU는 수가 적지만 복잡한 연산을 수행할 수 있는 코어와 큰 캐시, 정교한 제어 장치를 갖추어 개별 작업의 지연 시간을 줄이는 방향으로 구성된다. 반면 GPU는 비교적 단순한 코어를 대량으로 배치하고 제어 구조를 공유해 동시에 처리할 수 있는 작업량을 늘리는 데 중점을 둔다. 이러한 차이 때문에 GPU는 큰 규모의 데이터를 병렬로 처리하는 작업에서 높은 처리 효율을 보일 수 있다.

GPU는 SIMT^{Single Instruction Multiple Threads} 구조를 사용해 딥러닝 연산을 빠르게 처리한다. 앞서 살펴본 CPU의 SIMD와 비슷하게 하나의 명령을 여러 데이터에 적용하지만, GPU는 여러 스레드가 같은 명령을 공유한 상태로 동시에 실행되는 방식을 사용한다는 점이 다르다. 또 GPU에는 CPU보다 현저히 많은 연산 코어가 들어 있다. 이런 특징들로 인해 GPU는 병렬 연산에 많은 이점을 가진다. 데스크탑 용도의 CPU가 통상적으로 약 4~16 개의 연산 코어를 갖는 반면, GPU는 수천 개 이상의 연산 코어를 탑재하고 있어 대량의 연산을 동시에 처리할 수 있다.

NPU^{Neural Processing Unit}

NPU는 여기서 한 걸음 더 나아가 연산의 범용성을 포기하고 딥러닝 연산에만 특화된 프로세서다. CPU나 GPU는 일반 연산(스케줄링, 그래픽 연산 처리, 코드 로직 처리 등) 제어 로직을 최소화하고 칩의 대부분을 행렬 연산 유닛^{Matrix Multiply Unit}으로 채워 딥러닝 학습과 추론을 위한 전력 효율과 연산 밀도를 극대화한다. Google TPU, AWS Inferentia 등이 여기에 해당한다.

특징	CPU	GPU	NPU
목표	복잡한 로직의 빠른 순차 실행	대규모 병렬 데이터 처리	행렬 연산 전용 하드웨어 가속
코어 구조	적은 수의 복잡한 코어	수천 개의 단순 코어	수천 개의 단순 코어
메모리	DDR	HBM	HBM
제어 유닛	가장 큼	비교적 적음	가장 적음
유연성	매우 높음	높음	낮음
상호 연결	PCIe, QPI/UPI	PCIe, NVLink	독자 규격 (예: ICI) [5] [6]

GPU와 CPU의 관계

GPU는 단독으로 프로그램을 실행하는 장치가 아니라, CPU가 내리는 명령을 수행하는 보조 프로세서다. 프로그램의 흐름 제어, 데이터 준비, 어떤 커널을 언제 실행할지 등 전체적인 조율은 CPU가 담당한다.

데이터 전송

GPU가 계산하기 위해서는 먼저 데이터를 GPU 메모리^VRAM로 보내야 한다. 이때 시스템 메모리^RAM에서 GPU로 데이터가 이동하며, 보통 PCIe 버스를 통해 전송된다. [7]

이 과정은 직접 메모리 접근^{DMA: Direct Memory Access} 기술 덕분에 CPU가 직접 데이터를 옮기지 않아도 된다. 다만 PCIe 대역폭은 GPU가 내부에서 처리하는 속도(HBM 메모리)보다 훨씬 느리기 때문에, 전송량을 최소화하는 것이 성능 최적화에 중요하다. [8]

DMA^{Direct Memory Access}는 CPU가 데이터 전송 과정의 모든 바이트 이동을 직접 관장하지 않고, 전송 시작과 끝만 확인하면 되도록 하는 기술이다. CPU가 "이 데이터를 저쪽으로 옮겨"라고 DMA 컨트롤러에 명령하면, 실제 데이터 이동은 DMA 엔진이 수행하고 CPU는 다른 작업을 처리할 수 있다. 딥러닝 학습 시 수 GB에 달하는 가중치와 데이터를 매 에포크마다 GPU로 전송해야 하므로, 효율적인 DMA 활용은 전체 학습 속도에 중요한 영향을 미친다.

단일 GPU의 성능 한계를 극복하기 위해 멀티 GPU 환경을 구축할 때는, 느린 PCIe 버스 대신 GPU끼리 직접 통신할 수 있는 고속 인터커넥트(예: NVLink)가 필수적이다. 위 다이어그램은 GPU 코어들이 칩 내부에서 메시 혹은 링 버스 형태로 연결되거나, 외부적으로 고대역폭 링크를 통해 연결되는 구조를 나타낸다. 이러한 인터커넥트는 GPU 간의 메모리 복사 시간을 획기적으로 줄여주어 거대 언어 모델^LLM 학습과 같은 분산 처리 환경에서 병목 현상을 해결하는 핵심 기술이다. [NVLink] [InfiniBand]

명령 큐를 통한 CPU와 GPU 통신

CPU는 GPU에게 연산 처리나 데이터 복사 같은 명령을 순서대로 전달한다. 이때 GPU는 명령을 받자마자 바로 실행하는 게 아니라, 먼저 명령 큐^{Command Queue}에 적재된 요청을 하나씩 꺼내 비동기적으로 처리한다.

command_queue.cpp
// 명령 큐 동작에 대한 의사 코드

// CPU: 명령을 큐에 넣는 쪽
ring_buffer.push(command);

// GPU: 큐에서 명령을 꺼내 처리
while (ring_buffer.not_empty()) {
    execute(ring_buffer.pop());
}

GPU 연산은 CPU 명령이 들어오기 전까지는 시작되지 않는다. GPU가 작업을 완료하기 전에 CPU가 결과를 사용하려 할 경우 아직 연산이 완료되지 않았을 가능성이 있다. 따라서 cudaDeviceSynchronize() 등의 API를 이용해 결과를 동기화하거나, CUDA Stream으로 실행 순서를 제어해야 한다.

CPU가 커널 실행 요청을 큐에 넣고 GPU가 이를 꺼내 실행하기까지 약 3~5 μs의 준비 시간이 요구된다. 만약 커널 자체가 1 μs 정도의 소규모 연산이라면 준비 시간이 실행 시간을 초과하여 성능이 오히려 저하될 수 있다. 이러한 이유로 실무에서는 CUDA Graph 및 커널 융합^Fusion과 같은 기법을 사용하여 커널 실행 횟수를 줄이는 전략이 중요하다.

딥러닝 프레임워크와 하드웨어 커널

PyTorch나 JAX 같은 딥러닝 프레임워크는 사용자가 Python으로 쓴 코드를 실제 하드웨어에서 실행될 수 있는 연산으로 바꿔주는 역할을 한다. 우리가 x + y 같은 간단한 코드를 쓰더라도, 프레임워크는 내부적으로 이를 연산 그래프로 만들고, 어떤 순서로 계산해야 할지, 무엇을 동시에 실행할 수 있을지 등을 판단한다.

연산자는 어떻게 GPU 커널에 연결될까?

딥러닝 프레임워크 내부에는 Dispatcher가 있어서 Python에서 호출된 함수가 어떤 하드웨어 커널을 써야 하는지 자동으로 매칭해 준다. 다음 예시를 보자.

torch.add(x, y)

위와 같은 코드를 만나면 프레임워크에서는 (1) 텐서 타입 [9] (2) 디바이스 [10] (3) 텐서 모양 [11] 등을 보고 알맞은 CUDA 커널을 선택한다. torch.matmul 같은 행렬 곱은 NVIDIA에서 미리 최적화한 cuBLAS 라이브러리를 자동으로 사용한다.

커널 융합과 JIT의 필요성

Python은 인터프리터 언어이므로, 작은 연산 단위가 반복될 경우 Python → CUDA 호출이 매번 발생하여 GPU의 실행 단위를 충분히 활용하지 못하게 된다. 이러한 오버헤드를 줄이기 위해 최신 딥러닝 프레임워크는 (1) 여러 연산을 하나의 커널로 묶는 커널 융합^{Kernel Fusion}과 (2) 자주 수행되는 계산을 컴파일하여 재사용하는 JIT^Just-In-Time 컴파일을 도입하고 있다. 이를 통해 Python 호출 비용을 줄이고 GPU 연산을 고속으로 수행할 수 있다.

JIT 컴파일러는 Python 연산을 런타임에 분석하여 중간 표현^{IR: Intermediate Representation}으로 변환한 뒤, 예컨대 Add → Mul → ReLU 순서의 연산을 하나의 커널로 합치도록 최적화한다. 이 과정은 메모리 접근 횟수를 줄이고 Python 디스패처 호출을 제거하여 전반적인 실행 효율을 높인다. 아래 예제를 보자.

out = x * y
out = out + z
out = jnp.maximum(out, 0.0)  # ReLU

• Python 수준에서 각각 GPU 커널 호출 발생
• 각 호출 사이에 전역 메모리 접근 필요
• 오버헤드가 커져 연산 자체보다는 호출과 메모리 접근이 지연을 유발

이러한 연산을 하나의 커널로 융합하면 입력을 한 번만 읽어 레지스터에서 연산을 수행하고, 최종 결과만 한 번 기록하면 된다. 이는 Memory Bound 상황을 완화하는 대표적인 소프트웨어 기법이다. 다음은 JAX에서의 예시이다.

@jax.jit
def fused(x, y, z):
    return jnp.maximum(x * y + z, 0.0)

JIT 컴파일은 다음 과정을 거친다.

1. Python 연산을 XLA의 HLO IR로 변환한다.
2. HLO 그래프 최적화 과정에서 mul + add + relu를 하나의 커널로 융합한다.
3. 전역 메모리 접근을 한 번으로 통합한다.
4. GPU 레지스터 내에서 연산을 수행한 후 결과만 기록한다.

최종적으로 컴파일된 커널은 CUDA → PTX → SASS 형태로 변환되어 GPU에서 실행되며, 내부적으로는 다음과 유사한 절차를 따른다. [PTX] [SASS]

1. load x, y, z
2. 레지스터에서 mul → add → relu 처리
3. store result

이와 같이 JIT과 커널 융합은 Python 호출 및 메모리 접근을 줄여 GPU의 연산 자원을 효과적으로 활용할 수 있도록 한다.

커널 관점에서 GPU의 처리 과정

커널이 실행되면 GPU 내부에서는 아주 많은 스레드가 동시에 동작한다. 이 스레드들은 Grid → Block → Warp라는 구조로 묶여서 실행된다.

GPU 스레드 구조

• Grid: 커널을 한 번 실행할 때 만들어지는 전체 스레드 집합이다. 즉, 이번 커널 실행에 필요한 모든 스레드를 의미한다. Grid는 GPU 전체에 분산된다.
• Block: Grid를 여러 묶음으로 나눈 것. Block 단위로 GPU의 연산 장치인 SM^{Streaming Multiprocessor}에 배치되며, Block 내부 스레드끼리는 빠르게 데이터를 공유할 수 있다.
• Warp: 실제 하드웨어 실행 단위. NVIDIA GPU에서는 32개 스레드가 하나의 Warp로 묶여 같은 명령을 동시에 실행한다.

그림은 GPU에서 스레드가 Grid → Block → Warp → Thread 순서로 계층적으로 구성되는 모습을 보여준다. 왼쪽의 Grid는 전체 작업 영역을 의미하고, 그 안의 작업은 여러 개의 Block으로 나뉜다. 오른쪽은 그 Block 하나를 확대한 모습으로, Block 내부의 스레드들이 여러 Warp로 다시 묶여 있는 구조를 나타낸다. 각 Warp 안의 스레드들은 아래에 표시된 Instructions를 함께 공유하며, 항상 같은 명령을 동시에 실행한다. 실제 하드웨어에서는 보통 한 Warp가 32개의 스레드로 구성되고, 이 Warp 단위로 스케줄링되어 SM^{Streaming Multiprocessor} 안의 코어들에서 병렬로 처리된다. 이런 계층적 구조 덕분에 GPU는 수십만 개 이상의 스레드를 효율적으로 조직하고 관리하면서, 동일한 코드를 거대한 데이터 집합에 한꺼번에 적용할 수 있다.

SM의 구조

앞서 살펴 보았듯이 각 Block은 GPU 안의 SM이라는 연산 장치에서 실행된다. SM 내부에는 아래와 같은 자원이 있다.

• 레지스터 파일^{Register File}: 각 SM에는 레지스터 파일이라 부르는 공간이 존재한다. (H100 기준으로 256KB) 이는 CPU의 레지스터 파일보다 훨씬 크다. 스레드마다 전용 레지스터가 할당되며, 컨텍스트 스위칭 시 레지스터를 저장, 복구할 필요가 없어 오버헤드가 매우 작다.
• 공유 메모리^{Shared Memory}: SM 내부에 위치하며 사용자 제어가 가능한 고속 메모리다. 주로 블록 내 스레드 간 통신에 사용된다.
• L1 캐시^Cache: 반복되는 데이터 접근을 재사용하기 위해 사용된다.

GPU가 빠른 이유 중 하나가 레지스터와 공유메모리가 상대적으로 크고, 대량 병렬 처리에 맞게 설계되어 있다는 점이다.

Warp 스케줄링

GPU는 메모리에서 데이터를 가져오는 동안 멈추지 않는다. 예를 들어 Warp A가 메모리에 접근하고자 하여 대기가 발생하면 Warp 스케줄러에 의해 Warp B 실행으로 전환된다. 즉, 한 Warp가 쉬면 다른 Warp가 바로 들어와서 계산하기 때문에 효율적이다. CPU처럼 문맥 전환^{Context Switching} 과정 없이 내부에서 바로 스케줄링이 일어난다는 점에서 GPU는 CPU보다 동시성 처리에 효과적이다.

이런 구조는 GPU가 데이터를 병렬로 수행하는데 장점도 있지만 단점 또한 존재한다. 만약 Warp 내의 스레드들이 if-else 문을 만나 서로 다른 경로로 분기한다면, GPU는 두 경로를 모두 순차적으로 실행해야 한다.

1. 먼저 if 조건이 참인 스레드들만 활성화^Active되고, 거짓인 스레드들은 비활성화^Inactive된 상태로 명령어를 실행한다.
2. 그다음 else 경로를 타는 스레드들만 활성화되어 명령어를 실행한다.

결과적으로 실행 시간이 약 두 배로 늘어나며, 하드웨어 활용률은 절반으로 떨어진다. 따라서 GPU 커널 코드에서는 분기문을 최소화하거나, Warp 내의 모든 스레드가 동일한 경로를 타도록 유도해야 한다.

TPU는 뭐가 다를까?

TPU^{Tensor Processing Unit}는 구글이 설계한 딥러닝 특화 칩으로 NPU에 해당한다. GPU가 그래픽이나 딥러닝에서 범용적으로 사용되는 가속기라면, TPU는 처음부터 딥러닝 계산만 빠르게 하도록 설계된 전문 칩이라고 보면 된다.

시스톨릭 배열

TPU는 시스톨릭 배열^{Systolic Array} 아키텍처를 사용한다. GPU는 계산할 때마다 레지스터나 메모리에서 데이터를 계속 꺼내 쓰지만, TPU는 데이터가 연산 유닛을 흐르면서 자연스럽게 계산이 이어지도록 설계되어 있다. 쉽게 설명하자면,

• 왼쪽에서 입력이 들어오고
• 아래로 partial sum(부분 결과)이 내려오고
• 연쇄적으로 계산이 누적

예시에서 보듯이 데이터가 심장 박동^Systolic처럼 흐르면서 연산이 이어진다고 해서 시스톨릭 배열^{Systolic Array}이라고 부른다. 이러한 흐름은 레지스터 접근 횟수를 획기적으로 줄여 전력 효율을 높여준다.

가중치를 미리 고정해놓는 방식

TPU는 가중치^Weight를 연산 유닛 안에 미리 고정해두고, 입력^Activation만 계속 흘려보내는 구조다. 이런 방식을 가중치 고정^{Weight Stationary}이라고 부른다. 이 방식은 같은 가중치를 계속 재사용할 수 있어서 메모리 대역폭을 크게 아껴준다. 즉, 딥러닝 계산에서 자주 비효율을 유발하는 I/O를 줄이는 구조다.

실행 방식

GPU는 Warp 스케줄링, 분기 처리, 동적 병렬화 등이 복잡하게 들어가 있지만, TPU는 그런 과정이 필요하지 않다. 대신 컴파일러^XLA가 정해준 순서를 그대로 실행한다. 그래서 수백~수천 개 TPU를 구동해야 하는 대규모 TPU Pod 환경에서도 오버헤드가 낮다. 칩간의 통신 방법도 다르다 TPU끼리는 ICI^{Inter-Chip Interconnect} 라는 전용 네트워크로 직접 연결돼 있고, 3D 토러스 구조로 묶을 수 있다. GPU의 PCIe/NVLink와는 완전히 다른 설계 철학을 가지고 있고 이런 방식이 대규모 모델을 학습해야 하는 딥러닝 학습 클러스터를 만들 때 특히 유리하다.

시스톨릭 배열은 데이터가 칩 안을 일정한 방향으로 흘러가면서 연산이 이루어지는 구조다. 이름처럼 심장^Systolic이 몸안에 혈액을 계속 순환시키는 것과 비슷한데, 각 셀(위 그림에서는 PE^{Process Element}라고 표시)은 곱셈과 덧셈 같은 기본 연산을 수행한 뒤 결과를 옆이나 아래 셀로 바로 넘긴다. 이렇게 하면 중간 결과를 계속 메모리에 저장했다가 다시 가져올 필요가 줄어들고, 데이터가 흐르듯 전달되면서 여러 번 재사용된다. 그만큼 불필요한 메모리 접근이 줄어들어 에너지 효율이 높아지고, 대규모 모델을 학습할 때 전력 비용을 크게 줄일 수 있다.

대규모 확장에 유리한 상호 연결 (Inter Chip Interconnect)

TPU는 개별 칩의 성능뿐 아니라, 수천 개 칩을 하나의 시스템처럼 묶었을 때 어떻게 확장할 수 있는지를 가장 중요하게 고려한다. 위 그림의 3D Torus 구조처럼 인접한 칩이 서로 직접 연결되어 있으면, 별도의 네트워크 스위치를 거치지 않고도 데이터를 고속으로 주고받을 수 있다. 이런 방식은 거대한 모델을 학습할 때 칩 간 통신에서 생길 수 있는 병목을 줄여 보다 선형적인 확장 성능을 기대할 수 있게 한다. TPU 3d Torus 토폴로지 참고 논문

결과적으로 GPU는 그래픽 처리부터 일반 병렬 연산까지 다양한 용도에 쓰일 수 있도록 설계된 반면, TPU는 딥러닝 계산에 맞춰 구조를 단순화하고 연산 효율을 극대화한 특화형 하드웨어다. 따라서 딥러닝 같은 대규모 행렬 연산에서는 매우 높은 효율을 보이지만, 그 외 용도에서는 GPU만큼 범용적으로 사용하기 어렵다.

XLA HLO

JAX와 TensorFlow는 XLA 컴파일러를 이용하여 하드웨어 성능을 최적화한다. 이 과정에서 프로그램은 HLO라는 중간 표현으로 변환되며, 하드웨어와 무관하게 여러 최적화가 적용된다. XLA HLO 참고 문서

• Fusion: 여러 연산을 하나의 커널로 합쳐 실행 오버헤드와 메모리 접근을 줄인다.
• Buffer Assignment: 컴파일 단계에서 텐서의 메모리 위치와 수명을 미리 정해 두어 실행 중 메모리 할당 및 해제 비용을 줄이고 메모리 단편화를 방지한다.

최적화가 끝난 HLO는 다시 각 하드웨어에 맞는 명령으로 바뀐다. 예를 들어 GPU에서는 보통 LLVM을 거쳐 PTX 코드로 변환된 후 최종적으로 GPU에서 실행 가능한 기계어(SASS) 형태로 컴파일된다.

XLA의 최적화 파이프라인은 크게 파싱(Parsing), HLO 최적화, 백엔드 코드 생성 순으로 진행된다. 사용자 코드를 분석하여 하드웨어 독립적인 HLO 그래프를 생성한 뒤, 연산 융합(Fusion), 공통 하위 표현식 제거(CSE), 버퍼 재사용 등 다양한 최적화를 수행한다. 이렇게 정제된 그래프는 최종적으로 GPU(PTX/SASS)나 TPU 기계어로 번역된다. 덕분에 개발자는 하드웨어의 복잡한 특성을 몰라도 논리적인 연산만 작성하면 최고의 성능을 얻을 수 있다.

그래서 JAX 모델로 부터 GPU 연산 과정은?

다시 이 글의 질문으로 돌아와 c = jnp.dot(a, b)가 호출될 때 실제로 어떤 단계가 수행되는지 살펴보자.

1. Python 실행 사용자가 jax.jit된 함수를 실행하면 JAX는 우선 Python 코드를 분석하여 연산 그래프를 만든다. 이때 실제 값이 아니라 shape, dtype 같은 메타 정보만 보고 연산 흐름을 파악한다.

2. XLA로 넘겨 최적화 생성된 그래프는 XLA 컴파일러에게 전달된다. XLA는 먼저 연산 그래프를 HLO라는 중간 표현으로 바꾼다. 이 단계에서 dot 뒤에 이어지는 연산(+, relu 등)을 하나의 큰 연산으로 합치거나, 불필요한 텐서를 만들지 않도록 최적화한다. HLO는 하드웨어에 독립적인 중간 코드라고 볼 수 있다.

hlo_example.hlo
// f32[128,128] 크기의 lhs, rhs 두 행렬을 곱하는 dot 연산이다.
// lhs_contracting_dims={1}, rhs_contracting_dims={0} 는
// 행렬 곱 A(M,K) · B(K,N)에서 K 축에 대해 합을 취한다는 의미이다.
%dot = f32[128,128]{1,0} dot(
        f32[128,128]{1,0} %lhs,
        f32[128,128]{1,0} %rhs),
  lhs_contracting_dims={1}, rhs_contracting_dims={0}

이처럼 XLA 단계에서 하드웨어 독립적인 최적화가 끝나면, 이후에 GPU가 이해할 수 있는 PTX 코드가 생성된다.

3. GPU가 이해할 코드 생성 XLA는 최종적으로 GPU가 이해하는 PTX를 생성한다. 다음은 생성된 저수준 코드의 단순화된 예시이다.

ptx_example.s
// 전역 메모리에서 a[i], b[i]를 읽어와 곱한 뒤 c[i]에 저장하는 동작.

.visible .entry vec_mul(
    .param .u64 a_ptr,   // float* a
    .param .u64 b_ptr,   // float* b
    .param .u64 c_ptr    // float* c
) {
    .reg .f32  %fa, %fb, %fc;
    .reg .u64  %ra, %rb, %rc;

    // 파라미터에서 포인터 로드
    ld.param.u64 %ra, [a_ptr];
    ld.param.u64 %rb, [b_ptr];
    ld.param.u64 %rc, [c_ptr];

    // 전역 메모리에서 값 읽기
    ld.global.f32 %fa, [%ra];
    ld.global.f32 %fb, [%rb];

    // 곱셈 연산
    mul.f32 %fc, %fa, %fb;

    // 결과 쓰기
    st.global.f32 [%rc], %fc;
    ret;
}

4. GPU 드라이버가 PTX를 기계어로 컴파일 GPU 드라이버는 PTX를 현재 실행 중인 GPU 아키텍처에 맞는 SASS로 컴파일한다. 딥러닝의 행렬 곱 연산은 조건이 맞을 경우 일반 FP 스칼라 FMA^FFMA 경로 대신 텐서 코어용 HMMA 명령으로 변환된다.

sass_example.s
// A와 B를 읽어 행렬 곱하여 C에 저장하는 동작.

// A 타일 로드 (global → register)
LDG.E.SYS R4, [R2];    // R2: A 타일의 주소, R4: A 값 로드

// B 타일 로드 (global → register)
LDG.E.SYS R8, [R3];    // R3: B 타일의 주소, R8: B 값 로드

// Tensor Core에서 행렬 곱 타일 연산
// R0 = R4 × R8 + R0 (행렬 타일 누적)
HMMA.16816.F32 R0, R4, R8, R0;

// 결과를 전역 메모리에 저장
STG.E.SYS [R6], R0;    // R6: C 타일이 저장될 주소, R0: 결과 값

5. GPU 메모리에 로딩 컴파일된 SASS 코드는 GPU 메모리에 로드되며, 커널 실행 준비가 끝난다.

6. 데이터를 GPU로 복사 입력 a, b는 CPU 메모리에서 GPU 메모리로 PCIe, DMA 등의 작업을 거쳐 옮겨진다. 그리고 GPU 내부에서 사용할 수 있도록 준비된다.

7. 커널 실행 CPU는 컴파일된 커널을 실행하라는 명령을 GPU 명령 큐에 넣는다. 이 명령에는 그리드 크기, 블록 크기, 공유 메모리 크기 등의 실행 파라미터가 포함된다. GPU는 이 커널을 수천 개의 코어에 나누어 실행한다. 이 과정에서 일전에 다루었던 SM 할당, Warp 스케줄링이 수행된다.

8. 텐서 코어 연산 수행 SASS의 HMMA 명령어가 실행되어 텐서 코어에서 행렬 곱셈이 수행된다. 텐서 코어는 한 사이클에 $4 \times 4 \times 4$ 행렬 곱셈을 수행할 수 있다. [12]

9. 결과 쓰기 연산 결과는 레지스터나 공유 메모리에 누적된다. 최종 결과를 전역 메모리에 쓰기 전에 가능한 한 칩 내부 메모리에서 연산을 완료한다. 모든 연산이 끝나면 최종 결과 c가 HBM에 기록된다. 대부분 바로 다음 연산의 입력으로 사용되므로 CPU로 다시 돌아가지 않는 경우가 많다.

10. 최종 결과 CPU 메모리로 복사 필요한 경우 최종 결과를 다시 PCIe, DMA 등의 방법으로 CPU 메모리로 복사한다. 하지만 딥러닝 학습 중에는 대부분의 데이터가 GPU에 머물며 다음 연산의 입력으로 사용된다.

이 다이어그램은 Python 코드가 실제 GPU 하드웨어에서 실행되기까지의 주요 단계를 정리한 것이다. 먼저 Python과 JAX 단계에서 작성된 연산은 XLA를 통해 중간 표현으로 변환되고, LLVM과 PTX 단계를 거치면서 GPU가 이해할 수 있는 형태의 명령으로 컴파일된다. 이후 컴파일된 코드와 필요한 데이터는 PCIe를 통해 GPU로 전송되며, GPU 내부에서는 워프 스케줄러와 실행 파이프라인을 거쳐 다수의 연산 유닛에서 병렬로 실행된다.

딥러닝 학습 과정에서는 중간 결과가 대부분 GPU 메모리에 계속 남아 다음 연산으로 전달되며, 필요한 경우에만 CPU로 되돌아간다. 이를 통해 CPU와 GPU가 각자의 역할을 분담하고, 연산은 GPU에서 지속적으로 수행될 수 있다.

끝으로

우리는 이 글을 통해 jnp.dot과 같은 단순한 함수 호출이 실제로는 GPU 상에서 수천 개의 연산 유닛을 활용하여 수행된다는 사실을 살펴보았다. Python 인터프리터 수준에서 작성된 코드는 JIT과 XLA를 거치며 하드웨어 친화적인 형태로 변환되고, 최종적으로 GPU에서 실행 가능한 기계어로 컴파일된다. 이러한 과정은 모두 자동화되어 있지만, 내부적으로는 복잡한 최적화와 하드웨어 제어가 포함되어 있다.

AI 연산을 위한 하드웨어는 점점 세분화되고 특정 용도에 맞춰 설계되는 흐름이 이어지고 있다. Google TPU v7은 여러 칩을 묶어 대규모 Pod로 확장하는 구성을 중심에 두고 개발되었고, NVIDIA의 차세대 아키텍처인 Blackwell 역시 새로운 정밀도(FP8)와 메모리 구조를 도입해 대형 모델 학습 환경을 겨냥하고 있다.

Blackwell은 Transformer Engine을 통해 FP8을 하드웨어에서 직접 처리하도록 설계되었으며, 이후이어지는 Rubin 플랫폼은 전용 CPU(Vera)와 HBM4 메모리를 통합해 데이터 이동에서 발생하는 병목을 줄이는 방향으로 발표되었다. 이러한 흐름은 이제 단일 칩 성능 향상만으로는 충분하지 않으며, 랙 단위 혹은 데이터센터 전체를 하나의 컴퓨팅 시스템으로 바라보는 방향으로 설계가 이동하고 있음을 보여 준다.

소프트웨어 역시 이러한 변화를 반영하고 있다. JAX와 같은 컴파일러 기반 프레임워크는 하드웨어 복잡성을 추상화하면서도, pallas, triton과 같은 도구를 통해 보다 저수준의 제어를 가능하게 한다. 산업 현장에서는 이러한 저수준 최적화를 제품 성능 향상에 적극적으로 활용하고 있다. DeepSeek PTX 최적화 사례

딥러닝 프레임워크는 높은 수준으로 추상화되어 있어 전체 실행 과정을 완전히 이해하기는 쉽지 않다. 그러나 극한의 성능을 요구하는 경우에는 결국 하드웨어의 동작 원리와 물리적 제약을 고려할 필요가 있다. 작성된 코드는 결국 하드웨어 상의 실제 동작으로 이어진다는 점을 염두에 두어야 한다. 전체 과정을 이해하고 적절히 제어할 수 있을 때 비로소 하드웨어의 잠재력을 충분히 활용할 수 있다.

결국 GPU는 어떻게 머신러닝 연산을 수행할까?에 대한 대답은 다음과 같을 것이다.

GPU는 어떻게 머신러닝 연산을 수행하는가라는 질문에 대해, 핵심적인 요점은 다음과 같이 정리할 수 있다.

• Python 코드가 SASS 기계어로 변환되는 전체 파이프라인: Python → JAX → XLA → HLO → PTX → SASS.
• GPU 내부에서 Warp 스케줄링을 통해 메모리 지연을 숨기고 연산 밀도를 높임.
• PCIe 버스를 통한 Host-to-Device 전송 및 DMA 기반 직접 메모리 접근.

---

각주

---

• 1: 텐서^Tensor: 다차원 배열로, 딥러닝 모델의 입력 데이터, 파라미터, 출력 데이터 등을 나타낸다. [↩︎]
• BLAS: 기본 선형 대수 연산 프로그램^{Basic Linear Algebra Subprograms} 매트릭스 혹은 벡터 간의 최적화된 연산을 수행할 수 있도록 도와준다. [↩︎]
• 2: 연쇄 법칙^{Chain Rule}: 함수의 합성에 대한 미분은 각 함수의 미분을 곱한 것과 같다. [↩︎]
• 3: 활성화 체크포인트^{Activation Checkpointing}: 메모리 제약으로 인한 활성화 텐서를 저장하지 않고, 역전파 과정에서 필요할 때만 계산하여 메모리 사용량을 줄이는 기술이다. [↩︎]
• 4: AVX^{Advanced Vector Extensions}: 인텔과 AMD CPU에서 지원하는 명령어 집합으로, 한 번에 여러 개의 데이터를 처리할 수 있는 벡터 연산 기능을 제공한다. [↩︎]
• 5: ICI^{Inter-Chip Interconnect}: 구글 TPU 간의 초고속 데이터 전송을 위해 설계된 전용 네트워크 인터페이스. 3D Torus 토폴로지 구성을 가능하게 한다. [↩︎]
• 6: NPU 인터커넥트: NPU는 벤더마다 고유의 인터커넥트 기술을 사용하는 경우가 많다. 예를 들어 AWS Trainium은 NeuronLink를 사용한다. [↩︎]
• 7: PCIe^{Peripheral Component Interconnect Express}: 컴퓨터 메인보드에 주변 장치(그래픽 카드, SSD 등)를 연결하기 위한 고속 직렬 인터페이스 표준이다. [↩︎]
• 8: DMA^{Direct Memory Access}: CPU가 직접 메모리에 접근하지 않고, 하드웨어가 직접 메모리에 접근하는 방식을 의미한다. [↩︎]
• NVLink: NVIDIA가 개발한 고대역폭 GPU 간 연결 기술로, PCIe보다 훨씬 빠른 속도로 GPU 간 데이터를 전송할 수 있다. [↩︎]
• InfiniBand: 데이터 센터 및 고성능 컴퓨팅^HPC 환경에서 주로 사용되는 고속 네트워크 통신 표준으로, 대역폭이 넓고 지연 시간^Latency이 매우 낮은 것이 특징이다. [↩︎]
• 9: float32, float16 등 [↩︎]
• 10: CPU인지 GPU인지 [↩︎]
• 11: Shape, Stride [↩︎]
• PTX: NVIDIA GPU가 실행할 수 있는 저수준 연산 단계라고 보면 된다. CUDA가 고수준 연산이라면 PTX는 저수준 연산 단계라고 볼 수 있고 GPU 드라이버에 의해 PTX는 최종적으로 SASS로 변환된다. [↩︎]
• SASS: GPU가 실행할 수 있는 기계어 코드를 의미한다. JIT 컴파일러에 의해 컴파일되어 실행한다. [↩︎]
• 12: Tensor Core: FP16/FP32 행렬 연산에 특화된 하드웨어 블록이다. [↩︎]