Zachary, Novita AI의 알고리즘 전문가
연구 초점: 추론 가속
소개
지난 1년 동안 H2O가 등장한 이후로 KV 희소성에 관한 연구 논문이 급증했습니다. 그러나 실제 애플리케이션에서의 중요한 과제 중 하나는 학술 연구와 실제 구현 사이의 상당한 격차입니다. 예를 들어, vLLM과 같은 프레임워크는 페이지 기반 메모리 기술인 PagedAttention을 채택하는데, 이는 대부분의 희소성 알고리즘과 호환되지 않거나 이러한 알고리즘이 PagedAttention보다 성능이 낮습니다. 이러한 문제로 인해 희소성 알고리즘이 프로덕션 환경에서 효과적으로 활용되지 못하고 있습니다.
지난 1년간 KV 희소성 분야의 최근 학술 논문을 검토한 결과, 우리는 KV 희소성을 기반으로 vLLM 프레임워크를 수정했습니다. 이러한 수정은 Continuous Batching, FlashAttention 및 PagedAttention과 같은 vLLM 프레임워크의 고유한 최적화 기능과 일치하도록 이루어졌습니다. 최종적으로, 수정된 프레임워크를 SOTA 추론 프레임워크와 비교하고, 프로덕션에서 가장 일반적으로 배포되는 모델, 매개변수 및 하드웨어를 사용하여 추론 성능에서 1.5배의 속도 향상을 달성했습니다.
KV 희소성에 대해 논의하기 전에 LLM에서의 대규모 활성화 특성을 다루는 것이 필수적입니다. LLM에서는 소수의 활성화 값이 다른 값보다 훨씬 더 활성화되어 있으며, 때로는 덜 활성화된 값보다 100,000배 이상 차이가 나기도 합니다. 즉, 소수의 토큰이 중요한 역할을 합니다. 이 관찰은 중요한 토큰만 유지하여 추론 성능을 향상시킬 수 있는 KV 희소성 기술을 적용하기 위한 기초를 형성합니다.
Llama2 및 Llama3의 구체적인 분석 데이터는 이 관찰을 더욱 뒷받침하며, 그림 1과 그림 2에 나와 있습니다.
그림 1. LLaMA2-7B의 활성화 크기(z축). x축과 y축은 각각 시퀀스 및 특징 차원입니다. 이 특정 모델의 경우, 두 개의 고정된 특징 차원(1415, 2533)과 두 가지 유형의 토큰(시작 토큰, 첫 번째 마침표(.) 또는 줄바꿈 토큰(\))에서 거대한 크기의 활성화가 나타납니다.
Llama2와 달리 Llama3는 모든 레이어에서 대규모 활성화 특성을 나타내지 않습니다. 하위 레이어에서는 활성화가 더 균일하게 분포되는 반면, 중간 레이어는 국소적인 주의 패턴을 보입니다. 상위 레이어에서만 대규모 활성화 현상이 두드러집니다. 이 관찰은 Llama3와 같은 모델이 특정 활성화 패턴에 따라 레이어에 맞게 희소성을 조정하는 계층적 희소성 기술을 채택하는 이유를 설명합니다.
그림 2. LlaMa(Touvron et al., 2023a;b)의 레이어 간 검색 증강 생성의 주의 패턴은 하위 레이어에서 모델이 광범위한 스펙트럼 모드의 주의를 보여 모든 콘텐츠에 주의 점수를 균일하게 분배함을 보여줍니다. 중간 레이어에서는 각 문서 내에서 주의가 더 국소화되어 정제된 정보 수집을 나타냅니다(레이어 6과 10의 점선 빨간색 삼각형 모양). 이는 상위 레이어에서 절정에 달하며, "대규모 주의"가 몇 가지 주요 토큰(레이어 18 이후 집중된 주의 막대)에 집중되어 답변에 필요한 정보를 효율적으로 추출합니다.
시스템 모델
LLM에서 대규모 활성화 특성을 소개한 후, 이제 KV 희소성의 원리에 대한 기본적인 이해를 갖추었습니다. 다음으로 KV 희소성이 어떻게 구현되는지 논의하겠습니다. 대규모 모델의 추론 과정에서는 일반적으로 GPU 메모리 용량, 계산 성능 및 I/O 대역폭이라는 세 가지 주요 병목 현상이 발생합니다. 중복 계산을 피하기 위해 추론 중에 KV가 캐시되는 경우가 많아 GPU 메모리를 상당히 소비합니다. 한편, 계산 부하와 I/O 처리량은 각각 Prefill 및 Decode 단계에서 중요한 병목 현상이 됩니다.
vLLM 프레임워크를 수정한 것은 LLM의 대규모 활성화 특성을 기반으로 하여 레이어별 희소성(즉, 레이어마다 다른 수준의 희소성)을 구현했습니다. 이 접근 방식은 KV 캐싱과 관련된 오버헤드를 크게 줄여 LLM 모델의 추론 프로세스를 가속화합니다.
그림 3에 표시된 것처럼 추론 프로세스 중에 모델의 서로 다른 레이어에 KV 희소성을 적용합니다. 가지치기 전략을 사용하여 점수가 낮은 KV 쌍을 제거하고 점수가 높고 더 가까운 KV 쌍을 유지합니다. 이 접근 방식은 메모리 사용량과 계산 및 I/O 오버헤드를 줄여 궁극적으로 추론을 가속화합니다.
그림 3. KV 희소성이 모델에 미치는 영향
추론 가속 결과
대부분의 연구자들이 가장 우려하는 것은 vLLM 프레임워크에서 KV 희소성이 실제로 미치는 영향입니다. 세부 사항을 살펴보기 전에 먼저 성능 평가 결과를 제시합니다.
1. 추론 성능
실제 LLM 애플리케이션에서 입력/출력 길이가 4000/500인 구성이 가장 일반적으로 관찰됩니다. 또한 RTX 4090 GPU는 업계에서 가장 널리 사용되는 GPU 중 하나입니다. 이러한 조건을 바탕으로 배치 성능 비교 테스트를 수행했습니다. 대조군은 vLLM 0.6.1.p2이고, 실험군은 Novita AI Sparse 0.5.1(vLLM 0.5.1의 KV 희소성 수정 버전)입니다. 두 버전 간에 여러 차례 성능 비교 테스트가 수행되었으며, 주요 지표는 TTFT(첫 번째 토큰까지의 시간, 사용자 경험에 중요)와 처리량(실제 추론 속도)입니다.
최종 테스트 결과는 표 1에 나와 있습니다. KV 희소성을 적용하여 vLLM의 처리량을 약 1.58배 증가시킬 수 있었으며, TTFT(P50이 1초 미만)는 사용 가능한 수준으로 유지되었습니다. 또한 더 큰 배치 크기 시나리오에서 vLLM 0.6.1.p2는 동시성 수준 10에서 성능 한계에 도달한 반면, Novita AI Sparse 0.5.1은 동시성 수준 20에서도 안정적인 TTFT 성능을 유지할 수 있었습니다. 이는 프로덕션 환경에서 성능 안정성을 보장하는 KV 희소성의 견고성을 보여줍니다.
표 1: vLLM 0.6.1.p2와 Novita AI Sparse 0.5.1의 성능 비교
2. 모델 정확도
KV 희소성은 손실 압축 알고리즘이므로 모델의 성능을 평가하는 것도 마찬가지로 중요합니다. Llama3–8B 를 기반으로 MMLU 및 기타 벤치마크를 포함한 성능 평가를 수행했습니다. 결과에 따르면 정확도 손실은 약 3% 이내로 유지됩니다.
표 2에 표시된 대로 MMLU 및 인문학 벤치마크를 사용하여 모델 성능을 테스트했습니다.
표 2: MMLU 및 인문학 벤치마크를 사용한 Llama3–8B의 성능 평가 결과
둘째, 입력 길이가 7k에서 30k 사이인 실제 애플리케이션에서 흔히 발생하는 긴 텍스트 시나리오의 QA 작업 평가를 수행했습니다. 압축률은 표 3과 같이 10배 이상 유지되었습니다. 결과는 전체 정확도 손실이 약 10%임을 나타냅니다.
표 3: 입력 길이가 7k~30k인 긴 텍스트 QA 작업의 평가 결과
주요 기술
1. 레이어별 희소성 및 텐서 병렬화
vLLM 프레임워크는 일반적으로 Continuous Batching이라고 하는 반복 수준 스케줄링 전략을 사용합니다. 이 스케줄링 접근 방식은 배치의 모든 시퀀스가 생성을 완료할 때까지 기다리지 않고 반복 수준의 스케줄링을 용이하게 하여 정적 배치에 비해 GPU 활용도가 높습니다. 이 글의 시작 부분에서 언급했듯이 Llama3–8B와 같은 모델은 본질적으로 계층적 희소성과 같은 기능을 가지고 있습니다. 결과적으로 Continuous Batching을 기반으로 한 수정 사항은 그림 4와 같은 일련의 문제에 직면했습니다.
그림 4. Continuous batching을 사용하여 7개의 시퀀스 완료. 왼쪽에는 단일 반복 후 배치가 표시되고 오른쪽에는 여러 반복 후 배치가 표시됩니다. 시퀀스가 종료 토큰을 내보내는 즉시 해당 위치에 새 시퀀스를 도입합니다(예: 시퀀스 S5, S6 및 S7). 이 방법은 GPU가 모든 시퀀스가 완료될 때까지 기다렸다가 새 시퀀스를 시작할 필요가 없으므로 GPU 활용도를 높입니다.
우리가 해결한 주요 과제는 다양한 레이어에 걸친 메모리 관리였습니다. vLLM이 통합 페이징 메모리 관리 모델을 사용하고 큐 전략이 Full Cache 또는 Sliding Window만 수용할 수 있다는 점을 고려하여, vLLM의 핵심 구조를 수정하여 두 옵션을 모두 지원할 수 있도록 했습니다.
이 조정을 통해 각 레이어는 Full Cache 또는 Sliding Window 중에서 유연하게 선택할 수 있으며, 아래 다이어그램과 같이 레이어 간에 서로 다른 수준의 희소성을 허용합니다.
해결해야 할 첫 번째 과제는 다양한 레이어 간의 메모리 관리 문제입니다. vLLM은 모든 레이어에 대해 통합 페이징 메모리 관리 모델을 사용하고 큐 전략은 Full Cache 또는 Sliding Window만 허용하므로 vLLM의 기본 구조를 조정해야 합니다. 구체적으로는 Full Cache와 Sliding Window를 모두 지원하도록 하여 서로 다른 레이어가 선호하는 메모리 관리 전략을 독립적으로 선택할 수 있도록 합니다. 이 조정은 궁극적으로 서로 다른 레이어에 걸쳐 다양한 수준의 희소성을 용이하게 합니다. 해당 구조 다이어그램은 그림 5에 나와 있습니다.
그림 5. vLLM의 레이어별 메모리 관리 구조
전반적으로 조정은 세 가지 별도 단계로 분류할 수 있습니다.
서비스 초기화 단계: 서비스 시작 프로세스 중에 세 레이어에 할당된 메모리는 주로 남은 공간 계산과 함께 Metadata 구조의 계층적 초기화를 통해 초기화됩니다. 반면, 원래 vLLM은 완전한 Full Cache로 작동하며 다른 레이어에 대한 메모리 할당 로직을 고려하지 않습니다.
단계 준비 단계: vLLM에서 단계의 초기화 프로세스 중, 특히 블록 관리자 및 모델 실행기 메서드와 관련된 단계에서 KV 업데이트 메커니즘을 제어하기 위해 다른 레이어에 대해 다른 블록 테이블 관리 전략이 사용됩니다. 앞의 그림에서 설명한 대로 처음 세 레이어에는 Full Cache 메커니즘을 사용하여 모든 KV를 저장하고, 세 번째 레이어 이후의 레이어에는 Sliding Window 전략을 사용하여 가장 낮은 점수의 KV를 자동으로 교체합니다. 이 접근 방식을 통해 각 레이어의 요구 사항에 맞게 메모리 공간을 유연하게 할당하여 공간을 절약하면서 희소성 알고리즘의 효과를 극대화할 수 있습니다.
희소화 계산 단계: 이 단계는 일반적으로 모델 실행 후에 발생하며, 계산된 희소성 점수를 기반으로 희소화 연산이 실행됩니다. 이 단계에서 중요한 고려 사항은 Tensor Parallelism, CUDA 그래프 및 GQA와 같은 기술과의 호환성입니다. 실제 엔지니어링 구현에서는 텐서 병렬화에 특히 주의해야 합니다. 아래 그림과 같이 단일 머신 다중 GPU 시나리오에서는 적절한 분할 메커니즘을 사용하여 KV 희소화 연산을 각 GPU 내로 제한하여 GPU 간 통신을 피하는 것이 필수적입니다. 다행히 Attention 유닛은 ColumnLinear 기반 분할 접근 방식을 사용하므로 구현이 비교적 간단합니다.
그림 6. vLLM의 메모리 관리 및 희소화를 위한 3단계 프로세스
위의 논의는 주로 vLLM에서 계층적 희소성을 지원하기 위해 프레임워크 수준에서 이루어진 수정 사항에 초점을 맞춥니다. 다음으로 CUDA 수준에서 구현된 최적화에 주목하겠습니다.
2. Attention 수정
CUDA 수준의 수정 사항은 주로 Attention 계산 유닛, 특히 FlashAttention 및 PagedAttention에 대한 개선 사항에 초점을 맞춥니다. PagedAttention은 비교적 간단하므로 FlashAttention과 관련된 수정 사항을 강조하겠습니다.
먼저 Attention의 계산 공식을 제시합니다.
공식에서 볼 수 있듯이 softmax 계산 프로세스는 QK의 곱을 계산하는 것을 포함하며 2-패스 알고리즘(두 번의 패스 필요)입니다. 그러나 최종 계산 대상 O는 FlashAttention 알고리즘을 사용하여 단일 패스(한 번의 패스 필요)로 달성할 수 있습니다. FlashAttention의 구현 논리는 FlashAttention2 논문의 그림 7을 통해 설명할 수 있습니다. 간단히 말해, KV에 대한 블록 계산을 수행하기 위해 Q를 청크 단위로 순회하면서 루프가 끝날 때까지 O, P 및 rowmax 데이터를 점진적으로 업데이트하는 것입니다. 마지막으로 O를 ℓ로 나누면 FlashAttention의 1-패스 계산이 가능합니다.
그림 7. FlashAttention2의 구현 논리
KV 희소성을 구현할 때는 FlashAttention 계산 프로세스 중에 softmax에 필요한 값, 즉 rowmax/ℓ(전역 최대값 및 누적 값)이 이미 계산된다는 점에 유의해야 합니다. 그러나 FlashAttention은 이 두 값의 결과를 반환하지 않습니다.
FlashAttention에 대한 기본적인 이해를 바탕으로 PYRAMIDKV와 같은 논문에 제시된 희소화 점수 공식을 논의할 수 있습니다.
여기서 [n − α, n]은 명령 토큰의 범위입니다. 각 레이어 l 및 각 헤드 h에서 가장 높은 점수를 가진 상위 k l 개의 토큰이 선택되고 해당 KV 캐시가 유지됩니다. 다른 모든 KV 캐시는 폐기되며 생성 프로세스 전반에 걸쳐 후속 계산에 사용되지 않습니다.
구체적으로, 각 attention 헤드 h에서 KV 캐시에 유지하기 위해 i번째 토큰을 선택하는 점수 S를 다음과 같이 계산합니다.
위 공식에서 A는 softmax 점수를 나타내고, S는 가장 가까운 a개의 softmax 값의 점수 합계를 나타냅니다. 표시된 대로 softmax 자체는 2-패스 알고리즘이므로 FlashAttention에 직접 통합할 수 없습니다. 이는 또한 희소화 점수를 1-패스로 구현할 수 없음을 의미합니다.
앞서 언급했듯이 FlashAttention은 softmax 점수를 출력하지 않지만 rowmax/ℓ와 같은 중간 값을 계산합니다. 따라서 Online Softmax: 2-패스 알고리즘을 참조하여 rowmax/ℓ의 출력을 사용하여 KV 희소성에 대한 특정 점수 S를 달성할 수 있습니다. 실제로는 이 요구 사항을 충족하기 위해 그림 7에 설명된 FlashAttention을 개선하기만 하면 됩니다.
그림 7. KV 희소성 점수 계산을 위한 FlashAttention 수정
- FlashAttention에서 rowmax/ℓ 출력: FlashAttention은 rowmax/ℓ의 출력을 통합합니다. 이러한 값은 최대값과 누계를 나타내므로 데이터 크기가 매우 작아 FlashAttention의 전체 성능에 미치는 영향이 미미합니다. rowmax/ℓ를 재사용하면 두 번째 1-패스의 필요성을 피하여 전체 계산 효율성을 높일 수 있습니다.
- Online Softmax 단계 추가: FlashAttention을 기반으로 공식 p=e((kq)−row_max)/ℓ(1-패스 논리)를 사용하여 softmax 점수를 계산하는 Online Softmax 단계를 도입합니다. 그런 다음 행 방향 합산을 수행하여 희소화 점수 S를 계산할 수 있습니다.
Online Softmax 구현에서는 두 가지 중요한 세부 사항에 특히 주의해야 합니다.
- **루프 순서 조정 **: S는 k 차원을 기반으로 한 점수의 합계이므로 attention 메커니즘과 비교하여 Online Softmax의 계산을 q*k에서 *k**q로 조정해야 합니다. 이 조정은 열 합산을 행 합산으로 변환하여 모든 계산을 레지스터 내에서 완료할 수 있도록 하여 서로 다른 블록 간의 불필요한 통신을 방지합니다.
- **q의 시작 및 끝 위치 **: 인과 인코딩과 가장 가까운 a에 대한 합산 논리로 인해 모든 q 값을 k로 순회할 필요는 없습니다. 그림 7에서 0/1로 표시된 마스크 표시기를 참조하여 q의 시작 및 끝 위치를 기반으로 해당 블록에 계산이 필요한지 여부를 결정할 수 있으며, 이는 계산 리소스를 크게 절약합니다. 예를 들어, 입력 프롬프트의 길이가 5000이고 가장 가까운 a의 길이가 256인 경우 기존 계산은 5000*5000 연산이 필요하지만 시작 및 끝 위치를 기록하면 *5000**256 연산만 필요하므로 계산 효율성이 향상됩니다.
PagedAttention은 최종 점수 SSS를 계산한다는 공통 목표를 가지고 FlashAttention과 유사한 작업을 수행합니다. 여기서는 PagedAttention에 대한 자세한 설명을 제공하지 않습니다.
결론
vLLM 0.5.1을 기반으로 수정된 버전인 Novita AI Sparse 0.5.1은 현재 Llama3–8B 및 Llama3–70B와 같은 모델을 주로 지원합니다. CUDA 그래프를 사용하는 추론 모드에서 작동하며 주로 RTX 4090과 같은 소비자용 GPU에 배포됩니다. 엔지니어링 및 알고리즘 측면에서 최적화할 여지가 여전히 많이 남아 있습니다.
참고: vLLM-0.6.2에서는 다음 기능이 아직 지원되지 않습니다.
- Chunked-prefill
- Prefix caching
- FlashInfer 및 기타 FlashAttention이 아닌 백엔드
Novita AI Sparse 0.5.1에 관심이 있으시면 Novita AI를 방문하여 다른 사람들보다 먼저 경험해 보세요.
참고 문헌
[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference
[3] SnapKV: LLM Knows What You are Looking for Before Generation
[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models
[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion
[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation
