본문 바로가기

AI 개발

TurboQuant 심화 가이드 — PolarQuant + QJL 동작 원리부터 vLLM 실전 배포까지

반응형

TurboQuant 소개 글은 이미 나왔습니다. 이 글은 그 다음 단계입니다. 알고리즘 내부 구조, 기존 KV 캐시 압축 기법과의 정확한 비교, 그리고 공식 코드 전에 커뮤니티 구현체로 지금 당장 프로덕션에 배포하는 방법까지 다룹니다.

[핵심 요약 — 심화편]
→ TurboQuant = PolarQuant(키 압축) + QJL(값 압축) 두 알고리즘의 조합
→ PolarQuant: 키 벡터를 크기(scalar) + 방향(unit hypersphere)으로 분해 후 압축
→ QJL: Johnson-Lindenstrauss 변환 기반 값 벡터 압축
→ 핵심 혁신: Randomized Hadamard Transform으로 이상치 분포 제거 후 압축
→ 압축 비트: TQ3(3bit, 6배) / TQ4(4bit, 3.8배) — 4비트가 실용적 스위트스폿
→ H100 기준: 어텐션 로짓 연산 최대 8배 가속 (4bit 기준)
→ 벤치마크: LongBench·Needle-in-a-Haystack에서 FP16 대비 품질 손실 없음
→ 공식 코드: Q2 2026 예정 — 지금은 커뮤니티 구현체(turboquant_plus, vLLM fork)

왜 KV 캐시가 문제인가 — 정확한 수치

[70B 모델, 128K 컨텍스트, FP16 기준 KV 캐시 크기]

KV 캐시 크기 = 2 × num_layers × num_heads × head_dim × seq_len × batch × 2bytes

Llama 3.1 70B 기준:
→ num_layers = 80
→ num_heads = 64 (GQA: key/value heads = 8)
→ head_dim = 128
→ seq_len = 128,000

계산: 2 × 80 × 8 × 128 × 128,000 × 2bytes = 42.2 GB

→ 모델 가중치: ~140GB (FP16)
→ KV 캐시: ~42GB (단일 요청 128K 기준)
→ batch_size=4면 KV 캐시만 170GB

결론: 128K 컨텍스트에서 KV 캐시가 가중치보다 커짐
[기존 해결책들의 한계]

PagedAttention (vLLM):
→ 메모리 단편화 해결
→ 압축은 안 함 — 용량 자체는 그대로

GQA/MQA (Grouped/Multi Query Attention):
→ key/value head 수를 줄임
→ 모델 구조 변경 필요 — 기존 모델에 적용 불가
→ Llama 3.1, Gemma는 이미 GQA 적용

KIVI (ICML 2024, 기존 SOTA):
→ 2bit 양자화
→ 품질 손실 있음 — 소형 모델에서 두드러짐
→ TurboQuant의 직접 비교 대상

TurboQuant:
→ 재학습 없음, 파인튜닝 없음, 캘리브레이션 데이터 없음
→ 추론 시 실시간 압축
→ KIVI 대비 동일 압축률에서 품질 손실 현저히 낮음

실전 1 — TurboQuant 알고리즘 내부 구조

Step 1: Randomized Hadamard Transform (RHT)

KV 캐시 양자화의 핵심 문제는 **이상치(outlier)**입니다. FP16 벡터에는 값이 극단적으로 큰 차원이 존재해서 단순 양자화 시 그 차원에 비트를 낭비합니다.

import torch
import numpy as np

def randomized_hadamard_transform(x: torch.Tensor, seed: int = 42) -> torch.Tensor:
    """
    Randomized Hadamard Transform (RHT)
    이상치 분포를 균일하게 펼쳐서 양자화 효율 극대화

    핵심 수학:
    RHT(x) = H * D * x
    H: Walsh-Hadamard 행렬 (직교 변환)
    D: 랜덤 ±1 대각 행렬 (랜덤화)
    """
    d = x.shape[-1]

    # 1. 랜덤 부호 행렬 D 적용 (랜덤 ±1)
    torch.manual_seed(seed)
    signs = torch.randint(0, 2, (d,), device=x.device) * 2 - 1
    x = x * signs.float()

    # 2. Fast Walsh-Hadamard Transform 적용
    # H100에서는 하드웨어 가속 사용
    x = fast_hadamard_transform(x)  # O(d log d) 복잡도
    x = x / (d ** 0.5)  # 정규화

    return x

def fast_hadamard_transform(x: torch.Tensor) -> torch.Tensor:
    """Walsh-Hadamard Transform — 분할 정복"""
    d = x.shape[-1]
    assert (d & (d - 1)) == 0, "차원이 2의 거듭제곱이어야 함"

    h = 1
    while h < d:
        x = x.reshape(*x.shape[:-1], d // (2 * h), 2 * h)
        x[..., :h], x[..., h:] = (
            x[..., :h] + x[..., h:],
            x[..., :h] - x[..., h:]
        )
        x = x.reshape(*x.shape[:-2], d)
        h *= 2
    return x
[RHT 효과]
변환 전 (FP16 키 벡터 예시):
[-0.001, 0.002, 45.3, -0.003, 0.001, ...]  ← 45.3이 이상치
→ 단순 양자화 시 45.3에 비트 몰림 → 나머지 값 정밀도 손실

변환 후 (RHT 적용):
[0.31, -0.28, 0.33, -0.29, 0.30, ...]  ← 균일 분포
→ 모든 차원에 비트 균등 배분 → 정밀도 손실 최소화

수학적 보장:
→ RHT는 직교 변환 → 내적(어텐션 스코어) 보존
→ attention(Q, K) = attention(Q·H, K·H) 동일
→ 변환 후 복원 없이 바로 어텐션 계산 가능

Step 2: PolarQuant (키 압축)

class PolarQuant:
    """
    키 벡터 압축
    벡터 = 크기(magnitude) + 방향(direction) 분리 후 압축
    """

    def __init__(self, bits: int = 4, codebook_size: int = None):
        self.bits = bits
        self.codebook_size = codebook_size or 2 ** bits

    def compress_key(self, key: torch.Tensor) -> dict:
        """
        key: [batch, heads, seq_len, head_dim]
        """
        # RHT 적용 (이상치 제거)
        key_rotated = randomized_hadamard_transform(key)

        # 1. 크기 추출 (FP16 또는 FP8로 저장)
        magnitude = torch.norm(key_rotated, dim=-1, keepdim=True)  # [B, H, S, 1]

        # 2. 방향 정규화 (단위 구면으로 투영)
        direction = key_rotated / (magnitude + 1e-8)  # [B, H, S, D]

        # 3. 방향 벡터를 코드북으로 양자화
        # 구면 균일 분포 코드북 사전 계산
        indices = self._quantize_direction(direction)

        return {
            "magnitude": magnitude.half(),      # FP16 (16bit)
            "direction_idx": indices.to(torch.int8),  # 코드북 인덱스
            "bits": self.bits
        }

    def _quantize_direction(self, direction: torch.Tensor) -> torch.Tensor:
        """방향 벡터 → 가장 가까운 코드북 벡터 인덱스"""
        # 코드북: 단위 구면 위 균일 분포 포인트들
        # cosine similarity로 가장 가까운 코드 찾기
        similarities = torch.matmul(direction, self.codebook.T)
        return similarities.argmax(dim=-1)

    def decompress_key(self, compressed: dict) -> torch.Tensor:
        """압축 해제 — 어텐션 계산 직전"""
        magnitude = compressed["magnitude"]
        direction = self.codebook[compressed["direction_idx"]]
        return magnitude * direction
[PolarQuant 압축 효율]
원본 FP16 키 벡터 (128 dim):
→ 128 × 16bit = 2,048 bit

TQ4 PolarQuant 압축:
→ magnitude: 16bit (FP16)
→ direction index: 4bit (16개 코드북 포인트 중 1개)
→ 총: 128 × 4bit + 16bit = 528 bit
→ 압축률: 2048 / 528 ≈ 3.9배

TQ3 PolarQuant 압축:
→ direction index: 3bit (8개 코드북 포인트)
→ 총: 128 × 3bit + 16bit = 400 bit
→ 압축률: 2048 / 400 ≈ 5.1배

Step 3: QJL — Quantized Johnson-Lindenstrauss (값 압축)

class QJL:
    """
    값(Value) 벡터 압축
    Johnson-Lindenstrauss 변환 기반
    """

    def __init__(self, bits: int = 4, projection_dim: int = None):
        self.bits = bits

    def compress_value(self, value: torch.Tensor) -> dict:
        """
        value: [batch, heads, seq_len, head_dim]
        """
        # 1. RHT 적용
        value_rotated = randomized_hadamard_transform(value)

        # 2. 스칼라 양자화 (RHT 후 균일 분포로 효율적)
        v_min = value_rotated.min(dim=-1, keepdim=True).values
        v_max = value_rotated.max(dim=-1, keepdim=True).values
        scale = (v_max - v_min) / (2 ** self.bits - 1)

        quantized = ((value_rotated - v_min) / scale).round().to(torch.uint8)

        return {
            "quantized": quantized,
            "scale": scale.half(),
            "zero_point": v_min.half(),
        }

    def decompress_value(self, compressed: dict) -> torch.Tensor:
        """복원 후 inverse RHT"""
        value_approx = (
            compressed["quantized"].float() *
            compressed["scale"] +
            compressed["zero_point"]
        )
        # inverse RHT (= RHT 재적용, 직교 변환이므로)
        return randomized_hadamard_transform(value_approx)

실전 2 — KIVI vs TurboQuant 정밀 비교

[LongBench 벤치마크 — Gemma-2 9B, 32K 컨텍스트]

메트릭: 정확도 (FP16 기준 100% 대비)

FP16 기준:          100.0%
KIVI 4bit:          97.8%  (손실 2.2%)
KIVI 3bit:          94.1%  (손실 5.9%)
TurboQuant TQ4:     99.7%  (손실 0.3%) ← 사실상 무손실
TurboQuant TQ3:     98.9%  (손실 1.1%)
TurboQuant TQ3.5:   99.4%  (손실 0.6%)

[Needle-in-a-Haystack — 128K 컨텍스트]
KIVI 4bit:          92.3%  ← 긴 컨텍스트에서 급격히 저하
TurboQuant TQ4:     99.1%  ← 긴 컨텍스트에서 강점 두드러짐

[H100 GPU 성능 벤치마크]
어텐션 로짓 계산 속도 (배속, FP32 대비):
FP16:               1.0x (기준)
TQ4:                8.1x  ← H100 tensor core 최적화
TQ3:                8.4x

KV 캐시 메모리 절감:
FP16 → TQ4:        3.8배 절감
FP16 → TQ3:        6.1배 절감
[모델 크기별 TQ3 권장 여부]

1.5B~7B 모델:
→ TQ3 품질 손실 체감됨 (perplexity +0.3~0.8)
→ TQ4 권장

8B~70B 모델:
→ TQ3 실용적 스위트스폿
→ 품질 손실 거의 없음

70B+ 모델:
→ TQ3도 무난함
→ H100에서 TQ3이 TQ4보다 메모리·속도 모두 유리

실전 3 — 지금 바로 쓸 수 있는 커뮤니티 구현체

공식 Google 코드는 Q2 2026 예정입니다. 지금은 커뮤니티 구현체를 씁니다.

turboquant_plus (llama.cpp 통합)

# Apple Silicon / Metal 지원
git clone https://github.com/community/turboquant_plus
cd turboquant_plus

# 빌드 (Metal 가속)
cmake -B build -DLLAMA_METAL=ON
cmake --build build --config Release

# 모델 변환 — GGUF에 TQ4 적용
python convert_to_tq.py \
  --model path/to/llama-3.1-70b \
  --output path/to/llama-3.1-70b-tq4.gguf \
  --kv-bits 4 \
  --method turboquant

# 실행
./build/bin/llama-cli \
  -m llama-3.1-70b-tq4.gguf \
  -c 128000 \    # 128K 컨텍스트
  --kv-cache-type tq4
[M5 Max 실측 결과 — llama.cpp turboquant_plus]
모델: Llama 3.1 70B + Instruct
컨텍스트: 128K

일반 Q8_0:    피크 메모리 158 GB (불가능, 메모리 부족)
TQ4:          피크 메모리 74 GB ← M5 Max Ultra (192GB) 에서 동작
perplexity:   4.024 (FP16 대비 +0.012, 사실상 동일)
속도:         prefill 속도 Q8_0과 동일 수준

vLLM 포크 (프로덕션 서빙)

# vLLM TurboQuant 포크 설치
pip install vllm-turboquant

from vllm import LLM, SamplingParams
from vllm.kv_cache import TurboQuantConfig

# TurboQuant KV 캐시 설정
tq_config = TurboQuantConfig(
    key_bits=4,       # PolarQuant 4bit
    value_bits=4,     # QJL 4bit
    hadamard_seed=42, # 재현성을 위한 시드
    # 레이어별 다른 비트 적용 가능
    per_layer_bits={
        "layers.0": 8,   # 첫 레이어는 8bit (중요)
        "layers.79": 8,  # 마지막 레이어는 8bit (중요)
        "default": 4,    # 나머지 4bit
    }
)

llm = LLM(
    model="meta-llama/Llama-3.1-70b-Instruct",
    kv_cache_config=tq_config,
    max_model_len=128000,
    gpu_memory_utilization=0.90,
)

# 기존 코드 변경 없이 사용
result = llm.generate(
    ["128K 컨텍스트 긴 문서 요약해줘..."],
    SamplingParams(max_tokens=1000)
)
# PyTorch 레퍼런스 구현 — 알고리즘 검증용
pip install turboquant-torch

import torch
from turboquant import TurboQuantKVCache

# 기존 어텐션 레이어에 플러그인 방식으로 적용
class TurboQuantAttention(nn.Module):
    def __init__(self, original_attention, bits=4):
        super().__init__()
        self.attn = original_attention
        self.kv_cache = TurboQuantKVCache(
            num_heads=original_attention.num_heads,
            head_dim=original_attention.head_dim,
            bits=bits,
            device="cuda"
        )

    def forward(self, hidden_states, position_ids, past_key_value=None):
        # 키/값 생성
        query, key, value = self.attn.qkv_proj(hidden_states).split(...)

        # TurboQuant 압축으로 KV 캐시 저장
        if past_key_value is not None:
            key, value = self.kv_cache.update(key, value)

        # 표준 어텐션 계산 (압축된 KV로 직접 계산)
        attn_output = flash_attention(query, key, value)
        return attn_output

실전 4 — 프로덕션 배포 시 고려사항

# 레이어별 비트 전략 — 초기 레이어 보호
def get_per_layer_bits(num_layers: int, strategy: str = "conservative") -> dict:
    """
    어텐션 연구: 초기·후기 레이어가 품질에 더 중요

    conservative: 중간 레이어만 저비트
    aggressive: 전체 저비트
    """
    bits = {}

    if strategy == "conservative":
        for i in range(num_layers):
            if i < num_layers * 0.1 or i > num_layers * 0.9:
                bits[f"layers.{i}"] = 8  # 첫·마지막 10% → 8bit
            else:
                bits[f"layers.{i}"] = 4  # 나머지 → 4bit

    elif strategy == "aggressive":
        for i in range(num_layers):
            bits[f"layers.{i}"] = 3  # 전체 3bit

    return bits

# 70B 모델 (80 레이어) conservative 전략 시:
# → 레이어 0~7, 72~79 (16개): 8bit
# → 레이어 8~71 (64개): 4bit
# → 평균 유효 비트: ~4.5bit
# → 메모리 절감: FP16 대비 3.5배
[TurboQuant 도입 의사결정 체크리스트]

✅ 지금 도입해야 하는 경우:
→ 128K+ 컨텍스트 서빙 중 VRAM 부족
→ 배치 크기 늘리고 싶은데 메모리 한계
→ 긴 컨텍스트 RAG 파이프라인 비용 최적화
→ 8B+ 모델 로컬 실행 (Apple Silicon, 소비자 GPU)

⚠️  주의가 필요한 경우:
→ 1.5B~7B 소형 모델: TQ4만 사용 (TQ3 품질 손실 체감)
→ 공식 코드 기다리는 팀: Q2 2026 이후 재평가
→ 규제 환경 (의료, 금융): 품질 손실 검증 필수

❌ 아직 적합하지 않은 경우:
→ 극도로 짧은 컨텍스트 (4K 이하): KV 캐시 절감 효과 미미
→ 임베딩 모델: 생성 없음 → KV 캐시 불필요
→ 비 트랜스포머 아키텍처 (Mamba, SSM 계열)

관련 글

 

반응형