# KV Cache 的典型存储方式
class KVCache:
    def __init__(self, num_layers, max_seq_len, num_kv_heads, head_dim, dtype):
        self.k_cache = torch.zeros(
            num_layers, max_seq_len, num_kv_heads, head_dim,
            dtype=dtype, device='cuda'
        )
        self.v_cache = torch.zeros(
            num_layers, max_seq_len, num_kv_heads, head_dim,
            dtype=dtype, device='cuda'
        )
        self.seq_len = 0

    def update(self, layer_idx, new_k, new_v):
        """追加新的 KV"""
        batch_len = new_k.shape[0]
        self.k_cache[layer_idx, self.seq_len:self.seq_len + batch_len] = new_k
        self.v_cache[layer_idx, self.seq_len:self.seq_len + batch_len] = new_v
        if layer_idx == 0:  # 只在第一层更新 seq_len
            self.seq_len += batch_len

    def get(self, layer_idx):
        """获取当前层的 KV"""
        return (
            self.k_cache[layer_idx, :self.seq_len],
            self.v_cache[layer_idx, :self.seq_len]
        )

def calculate_kv_cache_memory(
    num_layers=32,
    seq_len=2048,
    num_kv_heads=8,  # GQA
    head_dim=128,
    dtype_bytes=2,  # FP16
    batch_size=1
):
    """计算 KV Cache 内存占用"""
    memory_bytes = (
        2  # K and V
        * num_layers
        * seq_len
        * num_kv_heads
        * head_dim
        * dtype_bytes
        * batch_size
    )
    return memory_bytes

# LLaMA-7B, 2048 tokens, batch_size=1
memory = calculate_kv_cache_memory()
print(f"KV Cache Memory: {memory / 1024**2:.2f} MB")
# 约 256 MB

# LLaMA-70B, 8192 tokens, batch_size=8
memory_70b = calculate_kv_cache_memory(
    num_layers=80,
    seq_len=8192,
    num_kv_heads=8,
    head_dim=128,
    batch_size=8
)
print(f"KV Cache Memory (70B): {memory_70b / 1024**3:.2f} GB")
# 约 10 GB

# MHA: 32 个 Q heads, 32 个 KV heads
# GQA: 32 个 Q heads, 8 个 KV heads (每 4 个 Q 共享 1 个 KV)

class GQA(nn.Module):
    def __init__(self, hidden_size, num_q_heads, num_kv_heads):
        self.num_q_heads = num_q_heads      # 32
        self.num_kv_heads = num_kv_heads    # 8
        self.num_groups = num_q_heads // num_kv_heads  # 4

        self.q_proj = nn.Linear(hidden_size, num_q_heads * head_dim)
        self.k_proj = nn.Linear(hidden_size, num_kv_heads * head_dim)  # 更小
        self.v_proj = nn.Linear(hidden_size, num_kv_heads * head_dim)  # 更小

    def forward(self, x, kv_cache=None):
        q = self.q_proj(x)  # [batch, seq, 32 * head_dim]
        k = self.k_proj(x)  # [batch, seq, 8 * head_dim]
        v = self.v_proj(x)  # [batch, seq, 8 * head_dim]

        # 扩展 K, V 以匹配 Q 的头数
        k = k.repeat_interleave(self.num_groups, dim=2)
        v = v.repeat_interleave(self.num_groups, dim=2)

        return attention(q, k, v)

# FP16 -> INT8 量化
def quantize_kv_cache(kv_cache, scale=None):
    """将 KV Cache 从 FP16 量化到 INT8"""
    if scale is None:
        scale = kv_cache.abs().max() / 127.0

    quantized = (kv_cache / scale).round().clamp(-128, 127).to(torch.int8)
    return quantized, scale

def dequantize_kv_cache(quantized, scale):
    """反量化"""
    return quantized.to(torch.float16) * scale

class H2OCache:
    """Heavy-Hitter Oracle: 保留高注意力权重的 KV"""

    def __init__(self, max_size, heavy_ratio=0.2):
        self.max_size = max_size
        self.heavy_ratio = heavy_ratio

    def update(self, k, v, attn_weights):
        if self.cache_len + 1 > self.max_size:
            # 根据累积注意力权重淘汰
            importance = attn_weights.sum(dim=0)  # 累积注意力
            keep_idx = importance.topk(int(self.max_size * self.heavy_ratio)).indices

            # 保留重要的 + 最近的
            self.k_cache = torch.cat([self.k_cache[keep_idx], k])
            self.v_cache = torch.cat([self.v_cache[keep_idx], v])