생성 시점의 활성화에 맞춘 LLM 가지치기로 디코딩 가속
원제 SparseDecoding: Decoding-Aware Pruning for Accurate and Efficient LLM Inference
추천 29댓글 2
Key Point
디코딩 단계의 실제 토큰 생성 환경과 가지치기 과정의 불일치를 해결하면서 SpMV 연산에 최적화된 커널까지 제공해, LLM 추론 속도 개선에 직접적으로 적용 가능한 실무적 해법을 제시한다.
핵심 요약
- LLM 디코딩 단계는 메모리 대역폭이 성능 병목이 되는데, Hessian 기반 가지치기로 읽는 파라미터 수를 줄여 지연 시간을 개선할 수 있다.
- 기존 방법들은 자연 텍스트에서 계산한 Hessian으로 모델을 가지치기하지만, 실제 디코딩 중에는 모델이 자체 생성 토큰을 처리하면서 분포 차이(distribution shift)가 발생한다.
- 이로 인해 생성 시 활성화 분포가 가지치기 시 사용된 분포에서 벗어나 모델 성능이 저하되는 문제가 있다.
- 기존 LLM 가지치기 방법들은 행렬 곱셈(SpMM)에 최적화되어 있으나, 디코딩을 지배하는 행렬-벡터 연산(SpMV)을 충분히 지원하지 않는다.
- SparseDecoding은 디코딩 인식형 가지치기 프레임워크로, 알고리즘 측면에서 프리필 단계를 제외한 밀집 모델의 자동회귀 생성 중 수집한 활성화로 보정 행렬을 구성해 디코딩 활성화와 가지치기 목표를 정렬한다.
- 시스템 측면에서는 비트마스크 인덱싱과 고정 스텝 순회를 사용한 최적화된 N:M 스파스 행렬-벡터 커널을 개발했다.
- Llama-3.1-8B, Llama-3.3-70B, Qwen3-14B/32B 같은 대형 모델들에서 고정 텍스트 보정보다 우수한 성능을 보이며 A100 GPU에서 최대 1.48배의 end-to-end 디코딩 속도 향상을 달성했다.