Transformer inference · roofline · KV capacity

LLM 추론: 계산량보다 메모리 이동이 느릴 때

Prefill과 decode는 같은 모델을 실행해도 행렬 크기와 데이터 재사용이 다릅니다. 메모리 용량, 대역폭과 연산량을 나눠 보고, 실제 요청 조건으로 병목을 판단합니다.

입력 처리와 토큰 생성은 다릅니다

Prefill은 입력 토큰들을 처리해 KV 캐시를 만듭니다. Decode는 앞서 생성한 토큰과 저장된 KV를 이용해 다음 토큰을 차례로 생성합니다. 일반적인 autoregressive 생성에서는 한 요청의 미래 토큰을 모두 미리 계산할 수 없습니다.

작은 배치의 decode에서는 가중치를 읽는 비용이 두드러질 수 있습니다. 배치를 늘리면 가중치 재사용은 좋아지지만 KV 읽기, 긴 문맥, 지연 요구도 함께 바뀝니다. 모든 decode가 항상 memory-bound이고 모든 prefill이 항상 compute-bound라고 고정하지 않습니다.

두 단계에서 다시 사용하는 데이터
  1. Prefill: 입력 토큰 묶음 → 초기 KV

  2. Decode: 새 토큰 + 기존 KV → 다음 토큰

  3. 측정: 첫 토큰 지연·토큰 간 지연·처리량

Transformer 추론 분석

FLOP/byte와 장치의 균형점을 비교합니다

Arithmetic intensity는 이동한 바이트당 연산량입니다. BF16 행렬 [B,D]와 [D,F]를 곱하고 입력·가중치·출력을 한 번씩 옮긴다는 단순 모델에서는 다음 식을 얻습니다. 캐시 재사용, 추가 연산, 통신은 포함하지 않은 하한 모델입니다.

FLOPs = 2 * B * D * F
bytes = 2 * (B * D + D * F + B * F)
intensity = FLOPs / bytes
small B relative to D,F: intensity approximately B FLOP/byte

Roofline의 연산·통신 모델

KV 크기를 가정에서 직접 계산합니다

모든 층에서 같은 길이의 dense KV를 저장하는 단순 예를 생각합니다. K와 V 두 벌, 층 수 L, KV head 수 H, head 차원 d, 원소당 바이트 s, 문맥 길이 T, 배치 B를 곱합니다. GQA에서는 query head 수가 아니라 KV head 수를 사용합니다.

가상의 L=32, H=8, d=128, s=2, B=1이면 토큰당 131072바이트, 즉 128KiB입니다. 4096토큰은 512MiB, 8192토큰은 1GiB입니다. 이는 특정 모델의 실측치가 아니며 가중치·활성값·임시 버퍼·메타데이터·할당 여유분은 제외한 값입니다.

KV bytes = 2 * L * H * d * s * T * B
2 * 32 * 8 * 128 * 2 = 131072 bytes/token
131072 * 4096 = 536870912 bytes = 512 MiB
131072 * 8192 = 1073741824 bytes = 1 GiB

KV와 attention의 추론 비용

GPU·TPU의 데이터 경로도 구분합니다

CUDA 문서의 LDGSTS는 compute capability 8.0 이상에서 global-to-shared 비동기 복사를 지원합니다. TMA는 9.0 이상에서 큰 다차원 전송을 위한 별도 경로입니다. 따라서 H100에서 처음으로 모든 비동기 복사가 가능해졌다고 설명하면 세대와 기능을 혼동합니다.

TPU의 VMEM은 TensorCore 가까이에 있는 on-chip scratchpad입니다. HBM에서 VMEM을 거쳐 행렬 연산 장치인 MXU로 보내는 경로와 Vector Unit의 원소별 연산을 구분합니다. 원소별 곱셈 애니메이션을 그대로 행렬 곱셈의 명령·주기 증거로 취급하지 않습니다.

CUDA 비동기 데이터 이동 · TPU의 메모리와 연산 경로

용량과 지연을 따로 측정합니다

긴 문맥에서 메모리가 부족하면 KV 크기와 동시 요청 수부터 확인합니다. 토큰 생성이 느리면 실제 이동량·대역폭·커널 시간·스케줄링을 봅니다. 메모리가 충분하다는 사실은 대역폭도 충분하다는 뜻이 아니며, 처리량 개선이 단일 요청 지연 개선을 보장하지도 않습니다.

이 글의 수치는 명시한 가정에 따른 산술 계산입니다. GPU·TPU 실행, 원 영상의 재생 검증, 특정 장치의 지연·성능 벤치마크를 수행한 결과가 아닙니다.