직접 해 보기: 양자화 오차와 KV 캐시 메모리 계산하기

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455import numpy as np

# ---------- 1부: 양자화 오차 실험 ----------
def quantize_absmax(w, bits, group=None):
    """대칭(absmax) 양자화 후 바로 되돌린 값을 돌려준다. group=None이면 행렬 전체에 스케일 1개."""
    qmax = 2 ** (bits - 1) - 1              # 8비트 127, 4비트 7
    flat = w.reshape(-1) if group is None else w.reshape(-1, group)
    amax = np.abs(flat).max(axis=-1, keepdims=True)
    scale = qmax / np.maximum(amax, 1e-12)
    q = np.clip(np.round(flat * scale), -qmax, qmax)
    return (q / scale).reshape(w.shape)

def rel_err(a, b):
    return np.linalg.norm(a - b) / np.linalg.norm(a)

rng = np.random.default_rng(0)
W = rng.normal(0, 0.02, size=(256, 256))     # 보통 가중치
X = rng.normal(0, 1.0, size=(8, 256))        # 입력 8개
W_out = W.copy()
W_out[rng.integers(0, 256, 4), rng.integers(0, 256, 4)] = 1.0   # 이상치 4개 심기

print("== 1부: 양자화 오차 (출력 상대 오차 = ||XW - XWq|| / ||XW||) ==")
for name, M in [("보통 가중치", W), ("이상치 4개", W_out)]:
    ref = X @ M
    for bits in (8, 4):
        for group in (None, 64):
            Mq = quantize_absmax(M, bits, group)
            tag = "전체 1개 스케일" if group is None else f"{group}개 묶음 스케일"
            print(f"{name:7s} | INT{bits} | {tag:12s} | 출력 오차 {rel_err(ref, X @ Mq):.4f}")

# ---------- 2부: 메모리 계산기 ----------
GB = 1e9
def weight_mem(n_params, bits):           return n_params * bits / 8 / GB
def train_state_mem(n_params, bytes_per=16):  return n_params * bytes_per / GB
def kv_cache_mem(n_layer, n_kv, d_head, seq, batch, bytes_per=2):
    return 2 * n_layer * n_kv * d_head * seq * batch * bytes_per / GB

print("\n== 2부: 메모리 계산기 ==")
for n in (7e9, 13e9, 70e9):
    print(f"{n/1e9:>4.0f}B | FP16 {weight_mem(n,16):6.1f} GB | INT8 {weight_mem(n,8):6.1f} GB"
          f" | INT4 {weight_mem(n,4):6.1f} GB | 학습 상태(16B/파라미터) {train_state_mem(n):7.1f} GB")

# LLaMA 65B 구성(층 80, 헤드 64, 헤드 차원 128)으로 KV 캐시
for n_kv, label in ((64, "MHA"), (8, "GQA 8그룹(가정)"), (1, "MQA")):
    print(f"KV 캐시 {label:14s} | 길이 2048 · 배치 1 : {kv_cache_mem(80, n_kv, 128, 2048, 1):6.2f} GB"
          f" | 배치 32 : {kv_cache_mem(80, n_kv, 128, 2048, 32):7.2f} GB")

# ---------- 3부: 디코드 지연 하한 · 추측 디코딩 ----------
print("\n== 3부: 디코드 지연 하한과 추측 디코딩 ==")
bw = 2.039e12   # A100 80GB SXM 메모리 대역폭(바이트/초)
for n, bits in ((7e9, 16), (7e9, 4), (70e9, 16)):
    print(f"{n/1e9:.0f}B · {bits}비트 가중치 1회 읽기 = 토큰당 최소 {n*bits/8/bw*1e3:6.2f} ms")
for a in (0.6, 0.8, 0.9):
    row = " ".join(f"γ={g}:{(1-a**(g+1))/(1-a):.2f}" for g in (2, 4, 8))
    print(f"수락률 α={a} → 기대 토큰 수 {row}")

12345678910111213141516171819202122232425== 1부: 양자화 오차 (출력 상대 오차 = ||XW - XWq|| / ||XW||) ==
보통 가중치  | INT8 | 전체 1개 스케일    | 출력 오차 0.0107
보통 가중치  | INT8 | 64개 묶음 스케일   | 출력 오차 0.0059
보통 가중치  | INT4 | 전체 1개 스케일    | 출력 오차 0.1940
보통 가중치  | INT4 | 64개 묶음 스케일   | 출력 오차 0.1022
이상치 4개  | INT8 | 전체 1개 스케일    | 출력 오차 0.1015
이상치 4개  | INT8 | 64개 묶음 스케일   | 출력 오차 0.0089
이상치 4개  | INT4 | 전체 1개 스케일    | 출력 오차 0.9162
이상치 4개  | INT4 | 64개 묶음 스케일   | 출력 오차 0.1091

== 2부: 메모리 계산기 ==
   7B | FP16   14.0 GB | INT8    7.0 GB | INT4    3.5 GB | 학습 상태(16B/파라미터)   112.0 GB
  13B | FP16   26.0 GB | INT8   13.0 GB | INT4    6.5 GB | 학습 상태(16B/파라미터)   208.0 GB
  70B | FP16  140.0 GB | INT8   70.0 GB | INT4   35.0 GB | 학습 상태(16B/파라미터)  1120.0 GB
KV 캐시 MHA            | 길이 2048 · 배치 1 :   5.37 GB | 배치 32 :  171.80 GB
KV 캐시 GQA 8그룹(가정)    | 길이 2048 · 배치 1 :   0.67 GB | 배치 32 :   21.47 GB
KV 캐시 MQA            | 길이 2048 · 배치 1 :   0.08 GB | 배치 32 :    2.68 GB

== 3부: 디코드 지연 하한과 추측 디코딩 ==
7B · 16비트 가중치 1회 읽기 = 토큰당 최소   6.87 ms
7B · 4비트 가중치 1회 읽기 = 토큰당 최소   1.72 ms
70B · 16비트 가중치 1회 읽기 = 토큰당 최소  68.66 ms
수락률 α=0.6 → 기대 토큰 수 γ=2:1.96 γ=4:2.31 γ=8:2.47
수락률 α=0.8 → 기대 토큰 수 γ=2:2.44 γ=4:3.36 γ=8:4.33
수락률 α=0.9 → 기대 토큰 수 γ=2:2.71 γ=4:4.10 γ=8:6.13

12345678flowchart LR
  A[가짜 가중치 W · 입력 X 만들기] --> B[이상치 심은 W_out]
  A --> C[absmax 양자화<br/>INT8 · INT4 × 전체 · 64묶음]
  B --> C
  C --> D[출력 오차 비교]
  E[모델 크기 · 비트 · 길이 입력] --> F[메모리 계산기<br/>가중치 · 학습 상태 · KV 캐시]
  F --> G[내 GPU에 들어가나 판단]
  H[대역폭 · 수락률] --> I[지연 하한 · 기대 토큰 수]