선택적 정밀도로 선형 어텐션 상태 압축하는 STEPQuant
원제 STEPQuant: When and Where Errors Matter in Delta-Rule Recurrent State Quantization
추천 106댓글 2
Key Point
선형 어텐션 모델의 메모리 병목인 재현 상태 양자화에서 시간과 공간 차원의 오류 영향을 고려한 첫 접근으로, 실제 배포 환경에서 메모리 사용량을 획기적으로 줄일 수 있다.
핵심 요약
- Linear attention은 KV 캐시 대신 고정 크기의 재현 상태를 사용하지만, 이 상태들이 메모리 병목이 되고 있다.
- 재현 상태를 저정밀도로 양자화하면 양자화 오류가 상태 업데이트 과정에서 누적되어 정확도가 크게 떨어진다.
- 연구진은 오류 영향이 두 가지 차원에 따라 달라짐을 발견했다: 시간적으로 오래된 메모리의 오류는 여러 디코딩 스텝에 걸쳐 지속되고, 공간적으로 서로 다른 키 행의 오류는 모델 출력에 미치는 영향이 다르며 상태 크기도 행과 열에 따라 변한다.
- STEPQuant는 시간·공간 정보를 바탕으로 정밀도를 할당하는 post-training 양자화 프레임워크로, 오류 크기와 메모리 수명에 따라 비트를 배분하고 키 행의 출력 영향도를 고려해 스케일을 결정한다.
- Qwen 3.8-27B와 Kimi-Linear-48B-A3B-Instruct 모델에서 6비트 STEPQuant는 FP32 정확도에 가깝고, 4비트에서 균등 INT8보다 우수하다.
- SGLang에 GPU 커널 최적화를 적용한 6비트 STEPQuant는 재현 상태를 5배 이상 압축하고 전체 서빙 메모리를 최대 68.7% 감소시킨다.
- 코드는 GitHub에 공개되었다.