TurboQuant 구현하기 (5편)

Google Research의 TurboQuant (arXiv:2504.19874) 논문을 PyTorch로 구현한 프로젝트입니다.

LLM 추론 시 가장 큰 메모리 병목인 KV Cache를 양자화 오버헤드 없이 압축합니다.

TurboQuant = PolarQuant (b-1 bit) + QJL (1 bit)
             주신호 압축              잔차 보정 (unbiased)

항목	기존 방법 (KIVI 등)	TurboQuant
양자화 상수 오버헤드	+1 ~ +2 bit	0 bit
3-bit 압축 시 실제 비트	4 ~ 5 bit	3 bit
내적 추정 편향	있음	없음
KV Cache 압축률	~4x	~6x

요구 사항

하드웨어

항목	최소	권장
GPU	CUDA 지원 NVIDIA GPU	RTX 3090 24GB
RAM	16GB	32GB
디스크	20GB (모델 캐시)	50GB

소프트웨어

항목	버전
OS	Ubuntu 20.04+ (WSL2 포함)
Python	3.10 ~ 3.12
NVIDIA 드라이버	595+
CUDA	13.0+ (드라이버 13.2 하위 호환)
PyTorch	2.10+ (cu130)
Transformers	4.40+
vLLM	0.8+ (serve_vllm.py 전용, 선택)
gcc	필수 (Unsloth/Triton 빌드용)

Windows에서는 WSL2 Ubuntu를 통해 실행합니다. PyTorch는 CUDA 런타임을 자체 포함하므로 별도 CUDA Toolkit 설치는 불필요합니다. WSL 필수: Unsloth 사용 시 C 컴파일러가 필요합니다: sudo apt update && sudo apt install -y gcc

DeepSeek 모델 가이드 (RTX 3090 24GB)

모델	파라미터	VRAM	24GB 가능
`DeepSeek-R1-Distill-Qwen-1.5B`	1.5B	~3GB	O
`DeepSeek-R1-Distill-Qwen-7B`	7B	~14GB	O (기본값)
`DeepSeek-R1-Distill-Qwen-14B`	14B	~9GB (4-bit)	O (`--load-in-4bit`)
`DeepSeek-R1-Distill-Llama-70B`	70B	~35GB+	X

프로젝트 구조

[Github] https://github.com/javafa/turboquant

turboquant/
├── requirements.txt             # Python 의존성
│
├── turboquant.py                # 핵심 라이브러리 (PolarQuant + QJL + TurboQuant)
├── compare.py                   # 원본 vs TurboQuant 비교 (품질 + 시간 + 메모리)
├── benchmark.py                 # 종합 벤치마크 (6가지 조합 x 메모리/속도 비교)
├── demo.py                      # HuggingFace 모델 KV cache 추출 + 압축 데모
├── serve_vllm.py                # vLLM OpenAI-compatible 서버 + 분석
│
└── scripts/
    ├── setup_wsl.sh             # WSL Ubuntu 환경 셋업
    └── setup_conda.bat          # Windows Conda 환경 셋업

파일 설명

파일	역할
`turboquant.py`	`PolarQuant`, `QJL`, `TurboQuant`, `TurboQuantKVCache` 클래스 구현. 배치 처리 지원.
`compare.py`	메인 스크립트. 모델 추론 후 원본 KV cache와 TurboQuant 압축 결과를 비교. Attention score 품질, 처리 시간, 메모리 절감을 한 번에 출력.
`benchmark.py`	종합 벤치마크. Baseline/BnB 4-bit/TurboQuant/Unsloth 등 6가지 조합의 VRAM, KV cache, 속도를 짧은/긴 프롬프트로 비교.
`demo.py`	HuggingFace 모델에서 KV cache를 추출하고 TurboQuant 압축률/정확도를 측정.
`serve_vllm.py`	vLLM 기반 OpenAI-compatible API 서버 + TurboQuant KV cache 분석.

2. 종합 벤치마크 (6가지 조합 비교)

비교 조합:

#	구성	가중치	KV Cache
1	Baseline (FP16)	FP16	FP16
2	BnB 4-bit	4-bit	FP16
3	TurboQuant 3-bit	FP16	3-bit
4	BnB 4-bit + TurboQuant	4-bit	3-bit
5	Unsloth 4-bit	4-bit	FP16
6	Unsloth 4-bit + TurboQuant	4-bit	3-bit

벤치마크 결과 (DeepSeek-R1-Distill-Qwen-1.5B, RTX 3090 24GB):

Short prompt (4 tokens 입력 → 64 tokens 생성):

Config                       | Model VRAM |  Peak VRAM |   KV Cache |      Speed |     Time
-----------------------------|------------|------------|------------|------------|----------
Baseline (FP16)              |    3022MB  |    3396MB  |    1.5 MB  |   39.4t/s  |    1.6s
BnB 4-bit                    |    1078MB  |    1578MB  |    1.5 MB  |   28.1t/s  |    2.3s
TurboQuant 3-bit             |    3022MB  |    3396MB  |  0.9MB(40%↓) |   48.5t/s  |    1.3s
BnB 4-bit + TurboQuant       |    1078MB  |    1578MB  |  0.9MB(40%↓) |   27.8t/s  |    2.3s
Unsloth 4-bit                |    1754MB  |    1974MB  |    1.5 MB  |   23.2t/s  |    2.8s
Unsloth 4-bit + TurboQuant   |    1754MB  |    1974MB  |  0.9MB(40%↓) |   37.2t/s  |    1.7s

Long prompt (191 tokens 입력 → 256 tokens 생성):

Config                       | Model VRAM |  Peak VRAM |   KV Cache |      Speed |     Time
-----------------------------|------------|------------|------------|------------|----------
Baseline (FP16)              |    3391MB  |    3418MB  |   12.2 MB  |   27.0t/s  |    9.5s
BnB 4-bit                    |    1555MB  |    1595MB  |   12.2 MB  |   17.0t/s  |   15.1s
TurboQuant 3-bit             |    3399MB  |    3418MB  |  7.3MB(40%↓) |   28.5t/s  |    9.0s
BnB 4-bit + TurboQuant       |    1555MB  |    1595MB  |  7.3MB(40%↓) |   16.6t/s  |   15.4s
Unsloth 4-bit                |    1754MB  |    1979MB  |   12.2 MB  |   22.3t/s  |   11.5s
Unsloth 4-bit + TurboQuant   |    1754MB  |    1979MB  |  7.3MB(40%↓) |   23.6t/s  |   10.8s

Baseline(FP16) 대비 비교:

구성	VRAM	KV Cache	속도 (short)	속도 (long)
BnB 4-bit	-53.5%	동일	-28.6%	-37.3%
TurboQuant 3-bit	동일	-39.9%	+23.1%	+5.3%
BnB 4-bit + TurboQuant	-53.5%	-39.9%	-29.4%	-38.5%
Unsloth 4-bit	-42%	동일	-41.1%	-17.5%
Unsloth 4-bit + TurboQuant	-42%	-39.9%	-5.5%	-12.7%

TurboQuant는 유일하게 baseline보다 빠른 구성 (short +23%, long +5%)

메모리 최적화: BnB 4-bit + TurboQuant (VRAM -53%, KV -40%)

균형: Unsloth + TurboQuant (VRAM -42%, KV -40%, 속도 소폭 감소)

BnB 4-bit는 긴 시퀀스에서 dequantize 비용 누적으로 속도 저하 심화

Unsloth short 결과에는 Triton 커널 초기 컴파일 오버헤드가 포함됨

출력 항목:

섹션	내용
1. Loading Model	모델 로딩 시간, VRAM 사용량
2. Running Inference	원본 생성 결과, 속도 (tok/s)
3. Attention Score Comparison	코사인 유사도, Top-1/5 매칭, KL divergence (2/3/4-bit)
4. Timing Benchmark	원본 vs TurboQuant 쿼리당 처리 시간 (ms)
5. Generation-Level	Attention 패턴 Top-1 일치율
6. Memory Savings	FP16 vs TurboQuant 메모리 비교
7. Scaling Projection	긴 컨텍스트(1k~128k) 메모리 예측

CLI Options

benchmark.py

--models             all | 1.5b | 7b (기본: all)
--skip-baseline      FP16 baseline 건너뛰기
--short-tokens       짧은 프롬프트 생성 토큰 수 (기본: 64)
--long-tokens        긴 프롬프트 생성 토큰 수 (기본: 256)
--turboquant-bits    TurboQuant 비트 수 (기본: 3)

compare.py

--model              모델 ID (기본: deepseek-ai/DeepSeek-R1-Distill-Qwen-7B)
--bits               TurboQuant 비트 수, 2-8 (기본: 3)
--max-new-tokens     생성 토큰 수 (기본: 64)
--prompt             입력 프롬프트
--dtype              float16 | bfloat16 (기본: float16)
--load-in-4bit       4-bit 양자화 로딩 (14B 모델용)
--load-in-8bit       8-bit 양자화 로딩
--layers-to-compare  비교할 레이어 수 (기본: 4)

demo.py

--model              모델 ID
--bits               TurboQuant 비트 수 (기본: 3)
--max-new-tokens     생성 토큰 수 (기본: 64)
--prompt             입력 프롬프트
--dtype              float16 | bfloat16
--load-in-4bit       4-bit 양자화 로딩
--load-in-8bit       8-bit 양자화 로딩

serve_vllm.py

--model              모델 ID
--host               서버 호스트 (기본: 0.0.0.0)
--port               서버 포트 (기본: 8000)
--max-model-len      최대 시퀀스 길이 (기본: 4096)
--gpu-memory-utilization  GPU 메모리 활용률 (기본: 0.90)
--quantization       vLLM 양자화: awq, gptq 등
--analyze-only       서버 시작 없이 분석만 실행
--turboquant-bits    TurboQuant 비트 수 (기본: 3)

Algorithm Overview

[Compress]
  k (원본 키 벡터)
    ├── PolarQuant (b-1 bits): 랜덤 회전 → 극좌표 → 균일 양자화
    │     저장: (r, codes)         오버헤드: ~0.25 bit (r 하나)
    │
    ├── 잔차: residual = k - decompress(r, codes)
    │
    └── QJL (1 bit): sign(S @ residual)
          저장: sign_bits            오버헤드: 0 bit

[Estimate <q, k>]
  <q, k> ≈ <q, k_hat> + sqrt(pi/2) * (S@q)^T * sign_bits
            PolarQuant    QJL 잔차 보정 (unbiased)

References

TurboQuant: arXiv:2504.19874 — ICLR 2026
QJL: arXiv:2406.03482 — Quantized Johnson-Lindenstrauss
PolarQuant: arXiv:2502.02617 — AISTATS 2026
DeepSeek-R1: deepseek-ai/DeepSeek-R1