bj.choi - notes

[LLM 추론] Attention 비용과 KV Cache·GQA·MLA 최적화

Updated:

Prefill과 Decode의 병목을 구분하고, KV Cache·GQA·MLA가 추론 비용을 줄이는 원리를 서빙 관점에서 정리합니다.

Attention 비용과 KV Cache·GQA·MLA 최적화

LLM 서빙에서는 긴 프롬프트로 TTFT가 늘어나거나 동시 요청이 늘면 GPU 메모리가 먼저 부족해지기도 합니다. 둘 다 Attention과 관련 있지만 병목 구간은 다릅니다.

Prefill·Decode 병목 구분

Decoder-only LLM은 입력과 출력을 하나의 토큰열로 처리하지만 서빙 엔진은 실행 특성에 따라 두 구간으로 나눕니다.

  1. Prefill

    • 프롬프트 전체 처리
    • 모든 레이어의 KV Cache 생성
    • 첫 출력 토큰 예측
    • 핵심 연산: n×n, 즉 O(n²) Dense Attention score 계산
    • 주요 지표: TTFT

  2. Decode

    • 직전 출력 토큰 처리
    • 해당 토큰의 K·V를 Cache에 추가
    • 다음 출력 토큰 예측
    • 작은 배치의 주요 병목: 모델 가중치와 누적 KV Cache 읽기
    • 주요 지표: ITL·TPS

컨텍스트를 두 배로 늘리면 Dense Attention 계산량은 약 네 배로 커집니다. 전체 Prefill 시간은 MLP·선형 계층과 커널 효율에 따라서도 달라집니다. Decode는 한 번에 토큰 하나만 처리하지만 매 단계에서 HBM에 있는 가중치와 과거 K·V를 읽어야 합니다.

두 지표는 서로 다른 병목이므로 분리해서 확인해야 합니다.

KV Cache의 연산·메모리 교환

첫 출력 토큰은 확률분포에서 선택된 결과일 뿐, 아직 해당 토큰의 Q·K·V가 계산된 상태는 아닙니다. 다음 토큰을 예측하려면 선택된 토큰을 다시 모델에 넣고 모든 레이어를 통과시켜야 합니다.

  1. Prefill에서 프롬프트 전체의 K·V 계산
  2. 각 토큰·레이어의 K·V를 HBM에 저장
  3. Decode에서 새 토큰의 Q·K·V만 계산
  4. 새 토큰의 K·V를 기존 Cache에 추가
  5. 현재 Q와 프롬프트·과거 출력·현재 토큰의 K로 Attention score 계산

과거 K·V의 재계산은 줄지만 출력 토큰이 늘어날수록 저장 용량과 HBM 읽기량이 함께 증가합니다. Q와 Attention score는 지속적인 KV Cache로 보관하지 않고 해당 연산 중에만 임시 값으로 사용합니다.

KV Cache 크기는 다음 식으로 계산합니다.

2(K,V) × layers × KV heads × head dimension × tokens × batch × bytes

32레이어, KV head 8개, head dimension 128, BF16인 GQA 모델 기준

모델 가중치와 임시 메모리를 제외한 논리적인 KV 텐서 총량입니다. GPU별 실제 점유량은 병렬화 방식, Cache 자료형, 블록 할당, prefix 공유 여부에 따라 달라집니다.

GQA의 KV head 공유

Query head가 128개이고 KV head가 8개라면 Query head 16개가 K·V 한 벌을 공유합니다. Query는 head마다 다르므로 같은 K·V를 사용해도 Attention 결과까지 같아지지는 않습니다.

GQA는 KV head 수에 비례해 Cache를 줄이지만 K·V 관점을 공유하기 때문에 MHA보다 표현에 제약이 생길 수 있습니다.

MLA의 latent 캐싱

MLA는 KV head 수가 아니라 캐싱 대상을 줄입니다.

  1. K·V 생성 정보를 공통 저차원 latent c_KV로 압축
  2. 전체 K·V 대신 c_KV와 작은 RoPE용 K 저장
  3. K 복원 대신 현재 Q를 latent 공간으로 변환
  4. V 전체 복원 대신 latent를 가중합한 결과만 선형 변환

이 계산 순서 변경을 Matrix absorption이라고 합니다. Decode에서는 K의 up-projection 행렬을 Query projection에, V의 up-projection 행렬을 output projection에 흡수합니다. 전체 K·V 텐서를 반복해서 만들지 않는 방식입니다.

RoPE는 토큰 위치마다 다른 회전 변환을 적용하므로 고정 행렬에 흡수하기 어렵습니다. 위치 정보가 필요한 작은 K를 별도 경로로 분리해 Cache에 함께 저장하는 이유입니다.

동일한 레이어 수·Query head 수·head dimension·dtype·토큰 수를 가정한 개념 비교입니다. MLA의 절감 폭은 latent와 RoPE 차원에 따라 달라집니다.

구조 Cache 구성 절감 방식 주의점
MHA head별 전체 K·V 기준선 KV Cache가 가장 큼
GQA 그룹별 K·V KV head 수 감소 K·V 관점 공유
MLA 공통 latent + RoPE용 K 캐싱 대상을 latent로 변경 TP에서 latent 중복 가능

모델 구조와 병렬화 병행 확인

DPA(Data Parallel Attention)는 Attention 부분에 적용하는 병렬화 방식입니다. 각 DP-attention rank가 서로 다른 요청과 해당 KV Cache를 담당해 MLA Cache의 TP-rank 중복을 줄입니다. FFN·MoE 등 나머지 계층은 TP·EP로 분할될 수 있습니다.

운영 지표 기반 병목 진단

관측 신호 우선 확인할 구간 검토할 방향
긴 프롬프트에서 TTFT 급증 Prefill Attention I/O 효율이 높은 Attention 커널, 프롬프트 길이, Prefill 정책
작은 배치·짧은 컨텍스트에서 ITL 저하 모델 가중치 읽기 Weight 양자화, Continuous batching, TP·MoE
배치·컨텍스트 증가에 따라 ITL 저하 KV Cache 읽기 GQA·MLA 모델, KV 양자화, 배치 상한
긴 컨텍스트·높은 동시성에서 OOM KV Cache 용량·할당 효율 최대 길이·동시성, Paged KV 단편화 관리, 반복 prefix 공유

Continuous batching은 scheduling iteration마다 처리할 시퀀스와 토큰을 다시 구성합니다. Decode만 수행하는 iteration에서는 참여 중인 활성 시퀀스 수가 보통 batch size가 됩니다. Chunked Prefill을 지원하는 엔진은 Prefill과 Decode 토큰을 같은 iteration에 넣기도 합니다.

배치 상한은 전체 TPS와 요청별 ITL을 함께 측정해 결정해야 합니다.

병목별 최적화 선택

  1. 긴 프롬프트의 TTFT 문제는 Prefill Attention 계산 점검
  2. 작은 배치·짧은 컨텍스트의 ITL 문제는 모델 가중치 읽기 병목 진단
  3. 큰 배치·긴 컨텍스트의 ITL 문제는 KV Cache 읽기 병목 진단
  4. OOM 문제는 KV Cache 용량과 동시성 상한 점검
  5. 모델 구조에 맞는 GQA·MLA와 TP·DPA 구성 선택

KV Cache는 과거 K·V 재계산을 메모리와 맞바꾸고 GQA는 KV head 수를 줄이며 MLA는 저장 대상을 latent로 바꿉니다. 기술 이름보다 현재 병목을 먼저 측정해야 합니다.

Reference