KV Cache 최적화 심층 분석: GQA·MLA·MHA 어텐션 메커니즘과 메모리 효율화 전략

import torch
import torch.nn as nn
import math

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, C = x.shape
        # Q, K, V 프로젝션
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # KV Cache 활용
        if kv_cache is not None:
            k_cache, v_cache = kv_cache
            k = torch.cat([k_cache, k], dim=2)
            v = torch.cat([v_cache, v], dim=2)

        # Attention 연산
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)

        output = output.transpose(1, 2).contiguous().view(B, T, C)
        return self.W_o(output), (k, v)

class MultiHeadAttention(nn.Module):
    """표준 Multi-Head Attention (MHA)
    각 헤드가 독립적인 KV 쌍을 유지"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)  # H * d_k 차원
        self.W_v = nn.Linear(d_model, d_model)  # H * d_k 차원
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # KV Cache 크기: [B, n_heads, S, d_k] x 2
        # 메모리: 2 * B * n_heads * S * d_k * sizeof(dtype)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), (k, v)

class MultiQueryAttention(nn.Module):
    """Multi-Query Attention (MQA)
    모든 Query 헤드가 단일 KV 헤드를 공유"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)       # H * d_k
        self.W_k = nn.Linear(d_model, self.d_k)       # 1 * d_k (단일 헤드)
        self.W_v = nn.Linear(d_model, self.d_k)       # 1 * d_k (단일 헤드)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, 1, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, 1, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # K, V를 모든 헤드에 브로드캐스트
        k = k.expand(-1, self.n_heads, -1, -1)
        v = v.expand(-1, self.n_heads, -1, -1)

        # KV Cache 크기: [B, 1, S, d_k] x 2
        # MHA 대비 1/H로 축소
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), (k[:, :1], v[:, :1])

class GroupedQueryAttention(nn.Module):
    """Grouped Query Attention (GQA)
    Query 헤드를 G개 그룹으로 나누어 각 그룹이 KV를 공유
    Llama 2/3에서 채택"""
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        self.n_heads = n_heads        # Query 헤드 수
        self.n_kv_heads = n_kv_heads  # KV 헤드 수 (그룹 수)
        self.d_k = d_model // n_heads
        self.n_rep = n_heads // n_kv_heads  # 그룹당 Query 헤드 수

        self.W_q = nn.Linear(d_model, n_heads * self.d_k)
        self.W_k = nn.Linear(d_model, n_kv_heads * self.d_k)
        self.W_v = nn.Linear(d_model, n_kv_heads * self.d_k)
        self.W_o = nn.Linear(d_model, d_model)

    def repeat_kv(self, x):
        """KV 헤드를 Query 헤드 수에 맞게 반복"""
        B, H_kv, S, D = x.shape
        if self.n_rep == 1:
            return x
        return (
            x[:, :, None, :, :]
            .expand(B, H_kv, self.n_rep, S, D)
            .reshape(B, self.n_heads, S, D)
        )

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_kv_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_kv_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # KV Cache 크기: [B, n_kv_heads, S, d_k] x 2
        new_cache = (k, v)

        # 브로드캐스트를 위해 KV 헤드를 반복
        k = self.repeat_kv(k)
        v = self.repeat_kv(v)

        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), new_cache

class MultiHeadLatentAttention(nn.Module):
    """Multi-Head Latent Attention (MLA)
    DeepSeek-V2/V3에서 사용
    KV를 저차원 잠재 공간으로 압축하여 캐시"""
    def __init__(self, d_model, n_heads, d_latent, rope_dim=64):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.d_latent = d_latent  # 잠재 차원 (d_model보다 훨씬 작음)
        self.rope_dim = rope_dim

        # Query 프로젝션
        self.W_q = nn.Linear(d_model, n_heads * self.d_k)

        # KV 압축 (Down-projection)
        self.W_dkv = nn.Linear(d_model, d_latent)

        # KV 복원 (Up-projection)
        self.W_uk = nn.Linear(d_latent, n_heads * self.d_k)
        self.W_uv = nn.Linear(d_latent, n_heads * self.d_k)

        # RoPE를 위한 별도 프로젝션
        self.W_kr = nn.Linear(d_model, rope_dim)

        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        # Query
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # KV 압축: d_model -> d_latent
        c_kv = self.W_dkv(x)  # [B, T, d_latent]

        # RoPE 키
        k_rope = self.W_kr(x)  # [B, T, rope_dim]

        if kv_cache is not None:
            c_kv_cached, k_rope_cached = kv_cache
            c_kv = torch.cat([c_kv_cached, c_kv], dim=1)
            k_rope = torch.cat([k_rope_cached, k_rope], dim=1)

        # KV Cache: 저차원 c_kv와 k_rope만 저장
        # 메모리: B * S * (d_latent + rope_dim) * sizeof(dtype)
        # MHA 대비: d_latent / (2 * n_heads * d_k) 비율로 축소
        new_cache = (c_kv, k_rope)

        # KV 복원: d_latent -> n_heads * d_k
        S = c_kv.shape[1]
        k = self.W_uk(c_kv).view(B, S, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_uv(c_kv).view(B, S, self.n_heads, self.d_k).transpose(1, 2)

        # Attention 연산
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), new_cache

class QuantizedKVCache:
    """KV Cache 양자화
    FP16 -> INT8로 변환하여 메모리 50% 절감"""
    def __init__(self, n_layers, n_heads, max_seq_len, d_k, dtype=torch.int8):
        self.n_layers = n_layers
        self.scales = {}  # 양자화 스케일 팩터
        self.zero_points = {}
        self.cache_k = {}
        self.cache_v = {}

    def quantize(self, tensor):
        """Per-channel INT8 양자화"""
        min_val = tensor.min(dim=-1, keepdim=True)[0]
        max_val = tensor.max(dim=-1, keepdim=True)[0]
        scale = (max_val - min_val) / 255.0
        zero_point = (-min_val / scale).round().clamp(0, 255)
        quantized = ((tensor / scale) + zero_point).round().clamp(0, 255).to(torch.uint8)
        return quantized, scale, zero_point

    def dequantize(self, quantized, scale, zero_point):
        """INT8 -> FP16 역양자화"""
        return (quantized.float() - zero_point) * scale

    def update(self, layer_idx, k, v):
        q_k, s_k, z_k = self.quantize(k)
        q_v, s_v, z_v = self.quantize(v)
        if layer_idx in self.cache_k:
            self.cache_k[layer_idx] = torch.cat([self.cache_k[layer_idx], q_k], dim=2)
            self.cache_v[layer_idx] = torch.cat([self.cache_v[layer_idx], q_v], dim=2)
        else:
            self.cache_k[layer_idx] = q_k
            self.cache_v[layer_idx] = q_v
        self.scales[layer_idx] = (s_k, s_v)
        self.zero_points[layer_idx] = (z_k, z_v)

    def get(self, layer_idx):
        k = self.dequantize(
            self.cache_k[layer_idx],
            self.scales[layer_idx][0],
            self.zero_points[layer_idx][0]
        )
        v = self.dequantize(
            self.cache_v[layer_idx],
            self.scales[layer_idx][1],
            self.zero_points[layer_idx][1]
        )
        return k, v

class H2OKVCache:
    """Heavy-Hitter Oracle (H2O) KV Cache 퇴거
    어텐션 점수가 높은 토큰만 유지"""
    def __init__(self, max_cache_size, n_heads, d_k):
        self.max_cache_size = max_cache_size
        self.attention_scores = None

    def update(self, k, v, attn_weights):
        """어텐션 가중치 기반으로 중요도를 누적하고 퇴거 수행"""
        B, H, S_new, _ = k.shape

        if self.attention_scores is None:
            self.attention_scores = attn_weights.sum(dim=2)  # [B, H, S]
        else:
            # 새 토큰에 대한 어텐션 점수 누적
            new_scores = attn_weights.sum(dim=2)
            self.attention_scores = torch.cat(
                [self.attention_scores, new_scores[:, :, -S_new:]], dim=2
            )

        # 캐시 크기 초과 시 퇴거
        current_size = self.attention_scores.shape[2]
        if current_size > self.max_cache_size:
            # 중요도가 낮은 토큰 식별 (첫 번째 토큰은 항상 유지)
            scores = self.attention_scores[:, :, 1:]  # 첫 토큰 제외
            _, indices = scores.topk(
                self.max_cache_size - 1, dim=2, sorted=False
            )
            indices = indices + 1  # 오프셋 보정
            # 첫 번째 토큰 인덱스 추가
            first_token = torch.zeros(B, H, 1, dtype=torch.long, device=k.device)
            keep_indices = torch.cat([first_token, indices], dim=2)

            # 선택된 토큰만 유지
            k = k.gather(2, keep_indices.unsqueeze(-1).expand(-1, -1, -1, k.shape[-1]))
            v = v.gather(2, keep_indices.unsqueeze(-1).expand(-1, -1, -1, v.shape[-1]))
            self.attention_scores = self.attention_scores.gather(2, keep_indices)

        return k, v

class SlidingWindowAttention(nn.Module):
    """Sliding Window Attention
    최근 W개 토큰에만 어텐션을 적용
    Mistral, Gemma 등에서 사용"""
    def __init__(self, d_model, n_heads, window_size=4096):
        super().__init__()
        self.window_size = window_size
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # 윈도우 크기로 KV Cache 제한
        if k.shape[2] > self.window_size:
            k = k[:, :, -self.window_size:]
            v = v[:, :, -self.window_size:]

        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), (k, v)

class PagedKVCache:
    """vLLM PagedAttention 개념 구현
    KV Cache를 페이지 단위로 관리"""
    def __init__(self, block_size=16, n_blocks=1024, n_heads=32, d_k=128):
        self.block_size = block_size
        self.n_blocks = n_blocks
        self.n_heads = n_heads
        self.d_k = d_k

        # 물리 블록 풀 (사전 할당)
        self.k_blocks = torch.zeros(n_blocks, n_heads, block_size, d_k)
        self.v_blocks = torch.zeros(n_blocks, n_heads, block_size, d_k)

        # 빈 블록 관리
        self.free_blocks = list(range(n_blocks))

        # 시퀀스별 블록 테이블 (가상 -> 물리 매핑)
        self.block_tables = {}  # seq_id -> list of physical block indices
        self.seq_lengths = {}

    def allocate(self, seq_id):
        """새 시퀀스에 대해 첫 블록 할당"""
        if not self.free_blocks:
            raise RuntimeError("OOM: No free blocks available")
        block_idx = self.free_blocks.pop(0)
        self.block_tables[seq_id] = [block_idx]
        self.seq_lengths[seq_id] = 0

    def append_token(self, seq_id, k, v):
        """시퀀스에 새 토큰의 KV 추가"""
        seq_len = self.seq_lengths[seq_id]
        block_idx_in_seq = seq_len // self.block_size
        offset = seq_len % self.block_size

        # 현재 블록이 가득 차면 새 블록 할당
        if offset == 0 and block_idx_in_seq > 0:
            if not self.free_blocks:
                raise RuntimeError("OOM: No free blocks available")
            new_block = self.free_blocks.pop(0)
            self.block_tables[seq_id].append(new_block)

        # 물리 블록에 KV 저장
        physical_block = self.block_tables[seq_id][block_idx_in_seq]
        self.k_blocks[physical_block, :, offset] = k
        self.v_blocks[physical_block, :, offset] = v
        self.seq_lengths[seq_id] += 1

    def get_kv(self, seq_id):
        """시퀀스의 전체 KV Cache 조합"""
        blocks = self.block_tables[seq_id]
        seq_len = self.seq_lengths[seq_id]

        k_list, v_list = [], []
        for i, block_idx in enumerate(blocks):
            if i == len(blocks) - 1:
                # 마지막 블록은 실제 토큰 수만큼만
                remaining = seq_len % self.block_size or self.block_size
                k_list.append(self.k_blocks[block_idx, :, :remaining])
                v_list.append(self.v_blocks[block_idx, :, :remaining])
            else:
                k_list.append(self.k_blocks[block_idx])
                v_list.append(self.v_blocks[block_idx])

        return torch.cat(k_list, dim=1), torch.cat(v_list, dim=1)

    def free(self, seq_id):
        """시퀀스 완료 후 블록 반환"""
        for block_idx in self.block_tables[seq_id]:
            self.free_blocks.append(block_idx)
        del self.block_tables[seq_id]
        del self.seq_lengths[seq_id]

# 1. vLLM 서버 재시작 시 메모리 제한 설정
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --tensor-parallel-size 8 \
  --max-model-len 65536 \
  --gpu-memory-utilization 0.9 \
  --max-num-seqs 8 \
  --enforce-eager

# 2. KV Cache 양자화 활성화
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --tensor-parallel-size 8 \
  --kv-cache-dtype fp8_e5m2 \
  --max-model-len 128000

# 3. 동적 배치 크기 제한
# max-num-seqs를 줄여 동시 처리 요청 수 제한

# 1. INT4 -> INT8로 양자화 정밀도 상향
# vLLM에서는 kv-cache-dtype 변경
# --kv-cache-dtype fp8_e5m2  (FP8 사용)

# 2. 핵심 레이어만 고정밀도 유지 (혼합 양자화)
# 초기 레이어와 마지막 레이어는 FP16, 중간 레이어는 INT8
kv_cache_config = {
    "default_dtype": "int8",
    "high_precision_layers": [0, 1, 2, 77, 78, 79],  # 처음/마지막 3개 레이어
    "high_precision_dtype": "fp16"
}

# 3. 벤치마크 재실행하여 품질 확인
# MMLU, HumanEval, GSM8K 등 표준 벤치마크로 검증

# 1. 프리픽스 캐싱 일시 비활성화
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --enable-prefix-caching false

# 2. 프리픽스 해시 키에 세션 ID 포함하도록 수정
# 동일 프리픽스라도 다른 세션이면 별도 캐시 사용

# 3. 프리픽스 캐싱 재활성화 시 TTL 설정
# --prefix-cache-ttl 300  (5분 후 자동 만료)