- C++와 CUDA를 사용하여 라이브러리 없이 LLM 추론 엔진을 구축하는 방법
- 이를 통해 LLM 추론의 전체 스택을 이해하고, 다양한 최적화가 추론 속도에 미치는 영향을 실감할 수 있음
- 목표 : 단일 CPU + GPU 서버에서 단일 배치로 빠르게 추론할 수 있도록 모델을 구현하고 llama.cpp보다 빠른 토큰 처리 속도 달성
1. LLM 아키텍처 및 추론 개요
- 대부분의 주요 LLM은 연속적인 트랜스포머 블록을 사용하는 동일한 아키텍처를 따름.
- 모델 로딩은 커스터마이즈 가능한 트랜스포머 블록 클래스를 정의하고 이를 시퀀스로 구성하여 safetensors 가중치로 초기화하는 것임.
- 추론은 주로 단일 배치로 이루어지며, "디코드 단계"가 실행의 대부분을 차지함.
1.1 추론 개요
- 추론은 주어진 프롬프트 토큰을 모델에 전달하고, 이를 통해 KV 캐시를 채우는 프리필 단계와 반복적으로 모델을 전달하여 토큰을 생성하는 디코드 단계로 나뉨
-
Prefill 단계: 프롬프트 토큰을 처리하며 KV 캐시 초기화
-
Decode 단계: 한 번에 한 토큰씩 생성
-
KV 캐시: 이전 키/값 쌍을 저장하여 과거 컨텍스트와의 주의(attention)를 빠르게 계산함
- 모델의 포워드 패스는 임베딩 테이블을 사용하여 토큰 ID를 임베딩 벡터로 매핑하고, 트랜스포머 블록 시퀀스를 통해 상태를 변형시킴
1.2 병목과 벤치마크
-
병목: 현대 하드웨어에서는 메모리 대역폭이 제한 요소임
- 모델 추론 시 각 토큰을 생성하려면 전체 모델을 읽어야 하며, 연산보다 메모리 대역폭이 더 큰 제약을 가짐
- 모델 양자화는 추론 속도를 개선하는 데 효과적임
- 이론적인 최대 토큰 처리량은 하드웨어에 따라 다르며, 실제 성능은 여러 인퍼런스 엔진을 통해 확인할 수 있음
-
이론적 속도 한계:
- AMD EPYC 7702P: 최대 13.6 tok/s (FP16 기준)
- RTX 4090: 최대 67.1 tok/s (FP16 기준)
-
벤치마크:
- llama.cpp: CPU 8.7 tok/s, GPU 61 tok/s
- calm: GPU 66 tok/s
2. CPU 기반 추론
- CPU에서의 초기 구현은 단일 스레드로 이루어지며, FP32 가중치만 지원함
- 멀티스레딩을 통해 코드 병렬화를 시작하고, SIMD를 사용하여 성능을 향상시킬 수 있음
2.1 멀티스레딩
- OpenMP를 사용하여 행렬-벡터 곱셈(matmul)과 다중 헤드 어텐션을 병렬화하여 성능을 개선
- 최적화 결과: 속도 0.6 tok/s → 4.4 tok/s 개선
2.2 가중치 양자화 및 SIMD 최적화
-
양자화: FP32 가중치를 FP16으로 양자화하여 메모리 사용량 절반 감소하고 성능을 개선
-
SIMD: AVX2를 사용해 8개 FP32 값을 동시에 처리하도록 최적화
- 결과: 8.4 tok/s 달성
3. GPU 기반 추론
- 모델을 FP16으로 양자화하여 RTX 4090에 로드하고 GPU 추론 구현을 시작할 수 있음
- CUDA를 사용하여 C++ 함수(커널)를 GPU에서 병렬로 실행할 수 있음
3.1 CUDA로 단순 포팅
- CPU 연산을 1-1로 CUDA 커널로 변환하여 GPU 백엔드를 구현할 수 있음
- CUDA 커널은 비동기적으로 실행되지만, 동일한 스트림에서는 순차적으로 실행됨
-
문제점: 쓰레드 비효율로 인해 GPU 자원을 충분히 활용하지 못함 → 2.9 tok/s로 느림
3.2 더 나은 행렬 곱셈(matmul)
- 행렬 곱셈은 CPU에서 큰 런타임을 차지하며, OpenMP를 통해 최적화할 수 있음
- GPU에서는 블록당 1개의 행을 처리하도록 하여 스레드 활용도를 높일 수 있음
-
최적화 방법:
- 한 블록이 한 행을 처리하고, 블록 내 쓰레드가 협업해 계산
-
워프 단위 합산(warp reduction) 적용
- 결과: 51.7 tok/s로 속도 개선
3.3 커널 융합과 추가 최적화
- 커널을 융합하여 성능을 향상시킬 수 있음
-
커널 융합: 연속되는 연산을 하나의 커널로 합쳐서 메모리 접근과 연산 시간 최소화
- 메모리 접근 패턴 최적화 및 공간 재활용을 통해 56.1 tok/s 달성
3.4 Attention 최적화와 긴 컨텍스트 처리
-
문제점: 긴 컨텍스트에서는 어텐션 커널이 성능 병목 발생
-
해결책:
-
메모리 접근 최적화: 연속된 메모리 블록을 읽어오도록 재설계
-
atomicAdd 대신 공유 메모리 사용으로 누락된 소수점 값 문제 해결
- 최적화 결과:
-
짧은 컨텍스트: 63.8 tok/s (llama.cpp의 61.0 tok/s보다 빠름)
-
긴 컨텍스트: 58.8 tok/s 달성
3.5 KV 캐시 양자화와 컴파일러 최적화 문제
- KV 캐시를 FP16으로 양자화하면 성능 저하 발생 (컴파일러 최적화 부족)
-
해결책: 수동으로 루프를 언롤링하고 메모리 프리페칭을 적용
- 결과: FP32 대비 약 2배 속도 향상 및 긴 컨텍스트 성능 58.8 tok/s 유지
4. 향후 개선 방향
-
프롬프트 prefill 최적화: 여러 토큰을 동시에 처리하여 첫 토큰 생성 시간 단축
-
Attention 커널 융합: FlashAttention과 같은 최적화 기법 적용
-
더 높은 양자화: FP8, INT8, INT4 적용 및 활성화/캐시 양자화
-
커널 최적화: 메모리 대역폭과 연산 효율을 극대화하는 고급 기법 도입
-
라이브러리 사용: cuDNN, cuBLAS와 같은 라이브러리를 활용하여 최적화 시간을 단축
결과 요약:
- CPU와 GPU에서 다양한 최적화를 통해 63.8 tok/s 속도 달성
- llama.cpp와 calm에 근접하거나 더 나은 성능을 기록
- 라이브러리 없이 C++와 CUDA만으로 고성능 LLM 추론 엔진을 구현