필수 LLM 최적화 기법(2)

LLM을 서빙할 때 트랜스포머 블록(Transformer block)의 주요 구성 요소는 어텐션(attention) 계층과 피드포워드 계층(feedforward layer, FFN)이다. FFN은 다층 퍼셉트론(multilayer perceptron, MLP) 계층으로 구성된다. 이번 글에서는 LLM 서빙의 어텐션을 확장하는 방법을 다룬다. 먼저 KV 캐시(KV cache)의 크기를 줄이는 방법을 살펴보고, 사용자 정의 GPU 커널(custom GPU kernel)이 어텐션 연산과 메모리 접근을 개선하는 방법을 설명한다.

1. 확장 가능한 어텐션과 커널 최적화

어텐션 메커니즘은 LLM 발전을 이끈 핵심 기술이다. 그러나 GPT-3에 쓰인 원래의 다중 헤드 어텐션(multi-head attention, MHA)만이 쿼리(query), 키(key), 값(value)을 대규모로 활용하는 유일한 방법은 아니며 가장 효율적인 방법도 아니다.

운영 워크로드와 LLM의 높은 연산 비용 때문에 모델 성능을 더 끌어내야 했다. 그 결과 멀티 쿼리 어텐션(multi-query attention, MQA), 그룹 쿼리 어텐션(grouped-query attention, GQA), 최근의 멀티 헤드 잠재 어텐션(multi-head latent attention, MLA)이 등장했다. 이 기법들은 모델 품질을 유지하면서 KV 캐시 크기를 크게 줄인다.

어텐션 구조를 살펴본 뒤에는 어텐션을 실행하는 GPU 커널을 다룬다. 커널은 범용 커널 융합(kernel fusion)에서 출발해 어텐션 전용, 하드웨어 최적화, 캐시 인식(cache-aware) 커널로 발전했다. 커널 융합의 원리를 먼저 이해한 뒤 FlashAttention 같은 구체적인 커널을 살펴본다. 마지막으로 또 하나의 핵심 혁신인 페이지드어텐션(PagedAttention)을 다룬다. PagedAttention은 GPU 메모리에 KV 캐시를 저장하고 관리하는 새로운 방식을 도입해 메모리 활용률을 높인다.

1-1. 확장 가능한 어텐션 메커니즘

KV 캐시 크기를 줄이는 것은 LLM 서빙 성능을 높이는 핵심 방법 중 하나다. 디코드 단계에서는 반복마다 KV 캐시를 고대역폭 메모리(HBM)에서 칩 내부 레지스터와 공유 메모리로 계속 옮긴다. KV 캐시가 작을수록 GPU 메모리 대역폭의 부담이 줄어든다. GPU 메모리에서 차지하는 공간도 작아진다. 따라서 더 큰 배치로 더 많은 요청을 병렬 실행해 처리량을 높일 수 있고, 제한된 GPU 메모리로 이전보다 긴 컨텍스트도 서빙할 수 있다.

그림 6-10의 맨 왼쪽에 있는 MHA(Multi-Head Attention)는 초기 모델 다수의 기반이 된 원래 방식이다. MHA에서는 각 쿼리의 어텐션을 계산할 때 서로 다른 키 헤드와 값 헤드가 하나씩 필요하다. 따라서 여기서 비교하는 네 가지 구성 중 KV 캐시가 가장 크고 효율도 가장 낮다.

그림 6-10. 다중 헤드 어텐션, 다중 쿼리 어텐션, 그룹 쿼리 어텐션, 다중 헤드 잠재 어텐션 비교.

성능을 높이기 위해 그림의 세 번째 구성인 MQA(multi-query attention)가 등장했다. MQA에서는 모든 쿼리가 하나의 키 헤드와 값 헤드를 공유한다. 7B 모델과 70B 모델은 일반적으로 각각 32개와 64개의 어텐션 헤드를 사용한다. MHA에는 각각 32개와 64개의 KV가 필요하지만 MQA에는 하나만 필요하다. 따라서 MQA는 KV 캐시 크기를 각각 32분의 1과 64분의 1로 줄일 수 있어 매우 큰 이점이다. 그러나 너무 공격적으로 공유하기 때문에 모델 정확도가 크게 떨어질 수 있다.

GQA는 MQA의 정확도 문제를 완화하고 성능과 정확도의 균형을 맞추기 위해 도입됐다. 쿼리 헤드를 여러 그룹으로 나누고 각 그룹이 하나의 키와 값을 공유한다. 연산 효율이 낮은 MHA와 정확도 손실이 큰 MQA 사이에서 좋은 절충안으로 입증됐으며, 현재 여러 모델 아키텍처에서 사용한다.

모델이 MHA, MQA, GQA 중 어느 방식을 사용하는지는 구성 파일에서 확인할 수 있다. Llama 2의 config.json에서는 KV 헤드 수와 어텐션 헤드 수가 같으므로 기본 MHA를 사용한다는 사실을 알 수 있다.

"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 32

반면 Llama 3 구성 파일에서는 num_key_value_heads가 8로 줄어 있다. KV 헤드 하나를 32 / 8 = 4개의 어텐션 헤드가 공유하는 GQA를 사용한다는 뜻이다.

"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 8

KV 캐시 효율을 높인 또 다른 중요한 기술은 DeepSeek가 도입한 MLA다. 그림 6-10의 맨 오른쪽 구성이다. MLA(multi-head latent attention)는 KV 수를 단순히 줄이지 않고 압축한다는 점에서 다른 방식과 다르다. DeepSeek의 원 논문은 KV 캐시 크기가 그룹을 2.25개만 사용하는 GQA와 같으면서도 성능은 MHA보다 높다고 설명한다.

다만 지금까지 설명한 발전은 모두 모델 아키텍처 수준에서 이뤄진다. 어떤 아키텍처나 어텐션 메커니즘을 선택할지는 특정 사용 사례에 가장 적합한 모델 계열과 개별 모델을 고르는 문제와 연결된다. 전체 시스템을 함께 고려해야 하는 결정이다.

이런 아키텍처 변화는 모델 개발의 초점이 품질 개선에만 있지 않고 실제 제품화로 옮겨가고 있음을 보여준다. LLM이 AI 시스템의 핵심 구성 요소가 될수록 강력한 성능뿐 아니라 실제 시스템에서 서빙할 수 있을 만큼 실용적이고 확장 가능하며 비용 효율적인 아키텍처가 중요하다.

1-2. 커널 융합과 사용자 정의 어텐션 커널

이제 특화된 GPU 커널로 어텐션 연산을 더 빠르게 만드는 방법을 살펴보자. 커널은 GPU에서 실행되는 작고 특화된 프로그램이다. 행렬 곱셈, 소프트맥스(softmax) 등 LLM과 여러 딥러닝 모델에 필요한 핵심 연산을 수행한다. 모델 아키텍처, 하드웨어, LLM 워크로드에 맞춰 적절히 최적화한 전용 GPU 커널을 사용하면 GPU 활용률, 추론 속도, 처리량을 크게 높일 수 있다.

1) 커널 융합

커널 최적화의 핵심 기법 중 하나가 커널 융합이다. 일반 머신러닝과 LLM 모두에서 널리 사용한다. 곱셈과 덧셈처럼 여러 개의 개별 연산을 하나로 합쳐 메모리와 연산 장치 사이의 데이터 이동 오버헤드를 최소화한다.

레지스터나 공유 메모리에 이미 있는 데이터를 재사용하므로 GPU 전역 메모리(global memory)에 결과를 다시 쓴 뒤 재로드하는 왕복 작업이 필요 없다. 그림 6-11은 중간 단계에서 메모리에 쓰고 다시 읽지 않고 도형 변환을 한 번에 수행하는 커널 융합의 동작을 보여준다.

그림 6-11. 커널 융합을 사용하지 않은 경우(왼쪽)와 사용한 경우(오른쪽)의 메모리 및 연산 상호작용 비교.

2) FlashAttention

이제 어텐션 커널 하나를 자세히 살펴보자. FlashAttention은 고성능 어텐션 연산을 크게 발전시킨 기술이다. 어텐션은 빈번한 메모리 읽기와 쓰기 때문에 GPU 메모리 대역폭 병목을 겪는다. FlashAttention의 핵심은 알고리즘이 하드웨어 또는 메모리 입출력(I/O)을 인식하도록 만들어 I/O를 줄이고 이 제약을 완화하는 것이다.

GPU 전역 HBM은 GPU 메모리 계층에서 용량이 가장 크지만 데이터 이동은 가장 느리다. FlashAttention은 큰 행렬을 느린 전역 메모리에 구체화(materialize)하지 않고 더 작은 행렬로 나눈다. 이를 타일링(tiling) 또는 블로킹(blocking)이라고 한다.

모든 연산을 SRAM이나 레지스터처럼 더 빠른 메모리에서 수행하므로 GPU 전역 메모리를 사용할 필요가 줄어든다. FlashAttention 원 논문의 그림 6-12는 어텐션의 QKV 행렬 곱셈과 변환을 타일 단위로 반복 수행하는 과정을 보여준다. 모든 연산은 GPU SRAM에서 일어나며 최종 출력만 HBM에 저장한다.

FlashAttention 커널은 온라인 소프트맥스(online softmax) 같은 추가 아이디어와 함께 모든 어텐션 연산도 융합한다. 현재 사용할 수 있는 가장 우수한 어텐션 커널 중 하나이며 여러 서빙 프레임워크가 사용한다.

그림 6-12. FlashAttention이 반복 루프로 SRAM에서 QKV 연산을 수행하는 방식(왼쪽)과 융합된 FlashAttention의 성능 향상(오른쪽).

FlashAttention 2와 3은 원래 방식의 핵심 아이디어를 유지하면서 기술을 더 개선했다. 예를 들어 일반 행렬 곱셈(General Matrix Multiply, GEMM)과 소프트맥스 계산을 겹쳐 실행해 GPU 활용률을 높인다. 특히 H100 같은 최신 GPU에서 효과가 크다. 커널 최적화는 GPU 아키텍처, CUDA, 성능 프로파일링(performance profiling), 컴파일러 등 여러 분야의 전문 지식이 필요한 크고 중요한 연구 영역이다. 실무 개발 관점에서 기억할 핵심은 LLM을 서빙할 때 효율적인 커널을 활용해야 한다는 점이다.

최적화된 다른 커널로는 FlashInfer, xFormers, Triton이 있다. 여기서 Triton은 완전히 다른 개념인 Triton Inference Server와 혼동하면 안 된다. 이들은 모델 학습과 서빙을 빠르게 하는 고성능 어텐션을 제공하며 GPU 커널 최적화의 효과를 보여준다. 예를 들어, 다음과 같은 몇 줄의 코드와 환경 변수로 사용자 정의 FlashInfer 커널과 의존성을 설치해 vLLM에서 실행할 수 있다.

pip install vllm==0.8.5.post1
pip install flashinfer-python==0.2.2

export VLLM_ATTENTION_BACKEND=FLASHINFER
export VLLM_USE_FLASHINFER_SAMPLER=1
export VLLM_FLASHINFER_FORCE_TENSOR_CORES=1

SGLang에서도 플래그 하나를 조정하는 방식으로 비슷하게 설정한다.

--attention-backend {flashinfer|fa3|triton|torch_native|FlashMLA}

커널, 하드웨어, LLM 입력과 출력의 특성이 복잡하고 다양하므로 어떤 커널을 사용해야 하는지 단정하기는 어렵다. 적합한 커널을 찾으려면 대개 실험이 필요하다. vLLM과 SGLang 같은 여러 서빙 백엔드는 기본 어텐션 커널을 선택하는 내장 로직을 제공한다. SGLang은 A100, A40 같은 비 Hopper 장비에서는 FlashInfer를, H100, H200, H20 같은 NVIDIA Hopper 아키텍처 GPU에서는 FlashAttention 3을 기본값으로 사용한다. 실무에서는 먼저 권장 기본값으로 시작해 다른 최적화 기회를 적용한 다음, 추가 성능이 필요할 때 여러 커널을 실험해도 된다.

3) PagedAttention

어텐션을 개선하는 마지막 핵심 방법은 KV 캐시 메모리를 효율적으로 관리하는 것이다. 모델 서빙 중에는 새 KV 캐시를 계속 만들고 저장하며 오래된 캐시를 제거(eviction)한다. 최신 LLM이 더 긴 컨텍스트를 지원하면서 KV 캐시 크기도 계속 커진다.

LLM 서빙은 요청마다 입력 길이가 다르고 출력 길이도 미리 알 수 없으므로 큰 KV 캐시를 스케줄링하기가 어렵다. 전통적인 방식은 메모리 공간을 미리 할당하지만 실제로는 이 공간을 다 쓰지 못하는 경우가 많다. 그 결과 GPU 메모리 단편화(memory fragmentation)가 커지고 실제 메모리 활용률이 낮아진다.

이 문제를 해결하기 위해 PagedAttention이 등장했고, 이를 기반으로 인기 있는 오픈 소스 서빙 프레임워크 vLLM이 개발됐다. PagedAttention은 운영체제의 페이징(paging) 기법에서 아이디어를 얻었다. 페이징은 메모리를 페이지(page)라는 고정 크기 블록으로 나눠 흩어진 빈 공간을 쉽게 할당하고, 연속된 메모리가 없어도 긴 시퀀스를 처리할 수 있게 한다.

PagedAttention도 비슷하게 동작한다. 먼저 KV 캐시를 고정 크기 블록으로 나눈다. 그런 다음 조회 테이블(lookup table)을 이용해 쿼리 키를 특정 블록에 매핑한다. KV 캐시를 연속된 메모리에 저장할 필요가 없으며 필요할 때 각 블록에 개별적으로 접근하면 된다.

그림 6-13을 보면 전체 프롬프트와 완성(completion)의 조합이 오른쪽 실제 물리 메모리의 연속 공간에 저장되지 않는다. 블록 7, 블록 1, 블록 3의 서로 다른 세 블록에 저장된다. 각 블록에는 최대 4개의 토큰 또는 단어가 들어간다. 마지막 블록은 아직 생성 중이므로 2개만 들어 있다. 블록 테이블(block table)은 요청한 데이터가 들어 있는 물리 블록을 찾는 조회 테이블이다.

그림 6-13. PagedAttention을 사용하는 요청 생성 과정의 한 단계.

PagedAttention 원 논문은 이를 사용하지 않으면 KV 캐시 메모리 중 실제 토큰 상태를 저장하는 데 쓰이는 비율이 20.4-38.2%에 불과하지만, PagedAttention을 사용하면 KV 캐시 메모리 낭비가 거의 0에 가까워진다고 설명한다.

메모리 단편화를 크게 줄이기 때문에 PagedAttention과 그 변형은 연속 배치처럼 거의 모든 사용 환경에서 활성화하는 LLM 서빙의 기본 기능이 됐다. 구체적인 구현 방식은 vLLM 문서에서 더 자세히 확인할 수 있다.

2. 더 읽어 볼 내용

• 필수 LLM 최적화 기법(1)
• 필수 LLM 최적화 기법(3)