TurboQuant 실전 — llama.cpp와 HuggingFace에서 KV Cache 압축하기

pip install transformers torch accelerate
pip install turboquant  # TurboQuant Python 패키지

# HuggingFace에서 GGUF 모델 다운로드 예시
huggingface-cli download TheBloke/Llama-3-70B-GGUF \
  llama-3-70b.Q4_K_M.gguf --local-dir ./models

git clone https://github.com/TheTom/llama-cpp-turboquant.git
cd llama-cpp-turboquant
git checkout feature/turboquant-kv-cache
cmake -B build -DGGML_METAL=ON  # macOS (Apple Silicon)
# 또는
# cmake -B build -DGGML_CUDA=ON  # NVIDIA GPU
cmake --build build --config Release -j

./build/bin/llama-server -m model.gguf \
  --cache-type-k turbo3 --cache-type-v turbo3 \
  -c 262144 -fa on --port 8080

curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "llama-3-70b",
    "messages": [{"role": "user", "content": "Hello!"}],
    "max_tokens": 256
  }'

./build/bin/llama-bench -m model.gguf \
  -fa 1 -ngl 99 \
  -ctk f16 -ctv f16 \
  -ctk q8_0 -ctv q8_0 \
  -ctk turbo3 -ctv turbo3

from transformers import AutoTokenizer, AutoModelForCausalLM
from turboquant import TurboQuantCache

# 모델 로드
model_name = "meta-llama/Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
)

# TurboQuant 캐시 생성 (4비트)
cache = TurboQuantCache(bits=4)

# 텍스트 생성
inputs = tokenizer("대한민국의 수도는", return_tensors="pt").to(model.device)
outputs = model.generate(
    **inputs,
    past_key_values=cache,
    use_cache=True,
    max_new_tokens=512,
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

import torch
from turboquant import TurboQuantMSE

# TurboQuant 양자화기 생성
tq = TurboQuantMSE(dim=128, bits=4, device='cuda')

# 랜덤 벡터 생성 (실제로는 Key/Value 텐서)
vectors = torch.randn(1024, 128, device='cuda')

# 양자화
indices, norms = tq.quantize(vectors)

# 역양자화
vectors_hat = tq.dequantize(indices, norms)

# 오차 확인
mse = ((vectors - vectors_hat) ** 2).mean()
print(f"MSE: {mse:.6f}")
# 출력 예시: MSE: 0.002341

for bits in [2, 3, 4]:
    tq = TurboQuantMSE(dim=128, bits=bits, device='cuda')
    indices, norms = tq.quantize(vectors)
    vectors_hat = tq.dequantize(indices, norms)
    mse = ((vectors - vectors_hat) ** 2).mean()
    print(f"{bits}-bit MSE: {mse:.6f}")

# 출력 예시:
# 2-bit MSE: 0.018742
# 3-bit MSE: 0.005123
# 4-bit MSE: 0.002341

KV Cache 크기 = 2 × num_layers × num_heads × head_dim × seq_len × bytes_per_element

모델 파라미터:
- num_layers: 80
- num_heads: 64
- head_dim: 128
- seq_len: 131,072 (128K)

def calc_kv_cache_size(num_layers, num_heads, head_dim, seq_len, bytes_per_elem):
    """KV Cache 메모리 사용량 계산 (GB 단위)"""
    size_bytes = 2 * num_layers * num_heads * head_dim * seq_len * bytes_per_elem
    return size_bytes / (1024 ** 3)

# Llama 3 70B
configs = {
    "FP16":  2.0,
    "Q8_0":  1.0,
    "TQ4":   0.55,
    "TQ3":   0.41,
}

print("Llama 3 70B @ 128K context:")
for name, bpe in configs.items():
    size = calc_kv_cache_size(80, 64, 128, 131072, bpe)
    print(f"  {name:6s}: {size:6.1f} GB")

# 올바른 사용
./build/bin/llama-server -m model.gguf --cache-type-k turbo3 -fa on

# 잘못된 사용 (Flash Attention 없이)
./build/bin/llama-server -m model.gguf --cache-type-k turbo3
# 경고: FA 없이도 동작하지만, 성능이 크게 저하됨

# Key는 3비트, Value는 4비트로 차등 적용
./build/bin/llama-server -m model.gguf \
  --cache-type-k turbo3 --cache-type-v turbo4 \
  -c 131072 -fa on

# HuggingFace에서 residual window 설정
cache = TurboQuantCache(
    bits=3,
    residual_length=256,  # 최근 256 토큰은 FP16 유지
)

# Metal GPU 사용 확인
# Activity Monitor > GPU History에서 GPU 사용률 확인
# 또는 로그에서 Metal 관련 경고 확인
./build/bin/llama-server -m model.gguf --cache-type-k turbo3 -fa on 2>&1 | grep -i metal

./build/bin/llama-perplexity -m model.gguf \
  -f wikitext-2-raw/wiki.test.raw \
  --cache-type-k turbo3 --cache-type-v turbo3 \
  -fa on -c 4096

모델 가중치: GGUF Q4_K_M (4비트 가중치 양자화)
KV Cache:   TurboQuant turbo3 (3비트 KV 캐시)
결과:       소비자 GPU에서 70B 모델 + 128K+ 컨텍스트 실행 가능

FP16 모델 + FP16 KV Cache:
  모델: ~140 GB + KV Cache: ~27.3 GB = ~167.3 GB  (불가능)

Q4_K_M 모델 + FP16 KV Cache:
  모델: ~40 GB + KV Cache: ~27.3 GB = ~67.3 GB  (A100 80GB 1장)

Q4_K_M 모델 + TQ3 KV Cache:
  모델: ~40 GB + KV Cache: ~5.6 GB = ~45.6 GB  (A6000 48GB 1장)

Q4_K_M 모델 + TQ3 KV Cache (32K context):
  모델: ~40 GB + KV Cache: ~1.4 GB = ~41.4 GB  (RTX 4090 가능)

# Q4_K_M 가중치 + TurboQuant KV Cache
./build/bin/llama-server \
  -m llama-3-70b.Q4_K_M.gguf \
  --cache-type-k turbo3 --cache-type-v turbo3 \
  -c 131072 -fa on -ngl 99 \
  --port 8080