본문으로 건너뛰기

[vllm] W4A8 Marlin 커널 - 4비트 가중치 + 8비트 활성화 양자화

PR 링크: vllm-project/vllm#24722 상태: Merged | 변경: +4,363/-2,232

들어가며

Marlin은 vLLM의 고성능 양자화 GEMM 커널이다. 이 PR은 기존 W4A16(4비트 가중치, 16비트 활성화)에 더해 W4A8(4비트 가중치, 8비트 활성화) 지원을 추가한다. 활성화도 양자화함으로써 메모리 사용량과 연산량을 동시에 줄인다.

핵심 코드 분석

W4A8의 핵심은 가중치를 4비트로, 활성화(입력)를 8비트(FP8 또는 INT8)로 양자화하여 mixed-precision GEMM을 수행하는 것이다.

이 변경의 규모(+4,363/-2,232)는 Marlin 커널의 CUDA 코드 대폭 수정을 반영한다:

  1. 커널 템플릿 확장: 기존 W4A16 전용 커널을 W4A8도 지원하도록 일반화
  2. 디퀀타이제이션 로직: 4비트 가중치와 8비트 활성화의 디퀀타이제이션 파이프라인
  3. 텐서 코어 활용: INT8/FP8 입력에 대한 텐서 코어 매핑
  4. 스케일 관리: per-channel, per-group 스케일 팩터 지원

왜 이게 좋은가

  1. 메모리 절약: 활성화 8비트 양자화로 KV cache 외 연산 메모리도 절감
  2. 연산 가속: INT8/FP8 텐서 코어는 FP16 대비 2배 이상의 TOPS 제공
  3. 모델 품질 유지: W4A8은 W4A16 대비 품질 손실이 최소이면서 성능은 크게 향상
  4. Marlin 통합: vLLM의 기존 Marlin 인프라를 재사용하여 안정적인 구현

정리

6,595줄의 대규모 변경이지만, LLM 추론에서 양자화의 중요성을 고려하면 핵심적인 최적화다. W4A8은 메모리와 연산 모두에서 이점을 제공하며, 특히 GPU 메모리가 제한된 환경에서 더 큰 모델을 서빙할 수 있게 한다.

참고 자료


이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.

댓글

관련 포스트

PR Analysis 의 다른글