import torch
import torch.nn as nn
import math
# 基础自注意力模块 + 原生KV Cache
class BasicAttentionWithKVCache(nn.Module):
    def __init__(self, embed_dim=128, num_heads=4):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        # Q/K/V 投影层
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        
        # 初始化KV Cache (batch_size, num_heads, seq_len, head_dim)
        # 就像一张白纸，等着记录内容
        self.k_cache = None
        self.v_cache = None
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        # 步骤1：计算Q/K/V向量
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # 步骤2：KV Cache核心逻辑 - 追加缓存
        # 
        # 【关键逻辑】
        # 如果缓存里有内容（不是第一次），就把新的K/V拼到后面
        # 就像写日记：今天的内容接在昨天后面，不用重写昨天的
        #
        if self.k_cache is not None:
            k = torch.cat([self.k_cache, k], dim=2)  # 维度2是序列长度
            v = torch.cat([self.v_cache, v], dim=2)
        # 更新缓存 - 把拼接好的存起来，下次继续用
        self.k_cache = k
        self.v_cache = v
        
        # 步骤3：计算注意力分数
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn_probs = torch.softmax(attn_scores, dim=-1)
        
        # 步骤4：注意力输出
        attn_output = torch.matmul(attn_probs, v)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)
        return self.out_proj(attn_output)
# 测试代码
if __name__ == "__main__":
    print("=" * 60)
    print("        KV Cache 效果演示")
    print("=" * 60)
    
    # 初始化模型
    attention = BasicAttentionWithKVCache(embed_dim=128, num_heads=4)
    
    # ========================================
    # 场景1：输入Prompt（首次推理）
    # ========================================
    print("\n【场景1】你输入了一个问题：")
    print("    '今天天气怎么样，适合去哪里玩？'")
    print("    假设这句话被分成了5个token")
    
    # 模拟输入：batch_size=1, seq_len=5, embed_dim=128
    prompt_input = torch.randn(1, 5, 128)
    
    # 第一次推理（写入缓存）
    output1 = attention(prompt_input)
    
    print(f"\n    → 初始缓存形状 - K: {attention.k_cache.shape}")
    print("      含义: [1个批次, 4个注意力头, 5个token, 每头32维]")
    print("      通俗: 记住了你输入的5个字的内容")
    
    # ========================================
    # 场景2：生成第一个词（使用缓存）
    # ========================================
    print("\n【场景2】模型开始回答，生成第1个字：")
    print("    '今' → 模型想回答'今天天气不错...'")
    print("    这时候只需要处理这1个新字！")
    
    # 生成新词：seq_len=1（逐词生成）
    new_token = torch.randn(1, 1, 128)
    output2 = attention(new_token)
    
    print(f"\n    → 更新后缓存形状 - K: {attention.k_cache.shape}")
    print("      变化: 5 → 6 (增加了1个token)")
    print("      通俗: 新写的字自动接在后面，之前5个字不用重算")
    
    # ========================================
    # 场景3：继续生成（缓存持续增长）
    # ========================================
    print("\n【场景3】继续生成回答，第2、3、4个字...")
    
    for i in range(3):
        new_token = torch.randn(1, 1, 128)
        output = attention(new_token)
        current_len = attention.k_cache.shape[2]
        print(f"    生成第{i+2}个字 → 缓存长度: {current_len}")
    
    print(f"\n    → 最终缓存形状 - K: {attention.k_cache.shape}")
    print("      总共记住了9个token的内容")

import torch
# KV Cache 对称量化函数
def kv_cache_quantize(kv_cache, num_bits=8):
    # 计算INT8的表示范围: -128 ~ 127
    q_max = 2 ** (num_bits - 1) - 1   # 127
    q_min = - (2 ** (num_bits - 1))   # -128
    
    # 步骤1: 计算缩放因子
    # 找到绝对值最大值，确定映射比例
    scale = torch.max(torch.abs(kv_cache)) / q_max
    
    # 步骤2: 量化（浮点 → 整数）
    # 除以scale压缩到目标范围，四舍五入，截断越界值
    kv_quantized = torch.round(kv_cache / scale).clamp(q_min, q_max).to(torch.int8)
    
    return kv_quantized, scale
# 反量化函数（推理时使用）
def kv_cache_dequantize(kv_quantized, scale):
    return kv_quantized.to(torch.float16) * scale
# 测试量化效果
if __name__ == "__main__":
    print("=" * 60)
    print("       KV Cache 量化效果演示")
    print("=" * 60)
    
    # 模拟真实场景：batch=1, heads=32, seq=4096, dim=128
    shape = (1, 32, 4096, 128)
    print(f"\n【测试配置】")
    print(f"   KV Cache形状: {shape}")
    print(f"   总元素数: {torch.prod(torch.tensor(shape)).item():,}")
    
    # 模拟原生KV Cache: FP16
    kv_cache = torch.randn(*shape, dtype=torch.float16)
    fp16_size = kv_cache.element_size() * kv_cache.nelement() / 1024**3
    print(f"\n【原生FP16】")
    print(f"   显存占用: {fp16_size:.2f} GB")
    print(f"   数据类型: {kv_cache.dtype}")
    print(f"   每元素: {kv_cache.element_size()} bytes")
    
    # 量化为INT8
    print(f"\n【量化过程】")
    kv_quant, scale = kv_cache_quantize(kv_cache, num_bits=8)
    int8_size = kv_quant.element_size() * kv_quant.nelement() / 1024**3
    print(f"   量化位数: INT8")
    print(f"   缩放因子: {scale:.6f}")
    print(f"   量化后显存: {int8_size:.2f} GB")
    print(f"   压缩率: {fp16_size/int8_size:.1f}x")
    
    # 反量化验证精度
    print(f"\n【精度验证】")
    kv_dequant = kv_cache_dequantize(kv_quant, scale)
    mse = torch.mean((kv_cache.float() - kv_dequant.float())**2)
    max_error = torch.max(torch.abs(kv_cache.float() - kv_dequant.float()))
    print(f"   均方误差(MSE): {mse:.8f}")
    print(f"   最大误差: {max_error:.6f}")
    print(f"   相对误差: {(max_error/torch.max(torch.abs(kv_cache))).item()*100:.3f}%")
    
    # 展示具体数值示例
    print(f"\n【数值示例】前5个值对比:")
    sample_idx = (0, 0, 0, slice(0, 5))
    print(f"   原始值:    {kv_cache[sample_idx].tolist()}")
    print(f"   反量化后:  {kv_dequant[sample_idx].tolist()}")
    print(f"   差值:      {(kv_cache[sample_idx] - kv_dequant[sample_idx]).tolist()}")

import torch
# 动态淘汰函数：基于注意力权重剔除无效缓存
def dynamic_kv_evict(k_cache, v_cache, attn_probs, keep_ratio=0.7):
    """
    动态淘汰低注意力权重的KV缓存
    :param k_cache: 键缓存
    :param v_cache: 值缓存
    :param attn_probs: 注意力概率分布
    :param keep_ratio: 保留比例
    :return: 压缩后的KV缓存
    """
    # 计算每个Token的平均注意力权重
    token_attn_score = torch.mean(attn_probs, dim=(1, 2))  # [batch_size, seq_len]
    # 排序：按权重从高到低
    sorted_indices = torch.argsort(token_attn_score, descending=True)
    # 保留高权重Token
    keep_num = int(token_attn_score.shape[-1] * keep_ratio)
    keep_indices = sorted_indices[:, :keep_num]
    
    # 压缩缓存
    k_cache_evicted = k_cache[:, :, keep_indices[0], :]
    v_cache_evicted = v_cache[:, :, keep_indices[0], :]
    return k_cache_evicted, v_cache_evicted
# 测试动态淘汰
if __name__ == "__main__":
    # 模拟缓存：batch=1, heads=128, seq_len=1024, head_dim=128
    k_cache = torch.randn(1, 128, 1024, 128)
    v_cache = torch.randn(1, 128, 1024, 128)
    attn_probs = torch.randn(1, 128, 1, 1024)  # 注意力概率
    
    print("淘汰前缓存长度:", k_cache.shape[2])
    k_evict, v_evict = dynamic_kv_evict(k_cache, v_cache, attn_probs, keep_ratio=0.7)
    print("淘汰后缓存长度:", k_evict.shape[2])
    print("显存降低比例:", 1 - k_evict.nelement() / k_cache.nelement())

import torch
from collections import OrderedDict
# 全局共享KV缓存池
class GlobalKVCachePool:
    def __init__(self, max_pool_size=10000):
        self.pool = OrderedDict()  # 有序字典：key=Token哈希, value=(k_cache, v_cache)
        self.max_pool_size = max_pool_size
    # 写入共享缓存
    def put(self, token_hash, k_cache, v_cache):
        if len(self.pool) >= self.max_pool_size:
            self.pool.popitem(last=False)  # 先进先出淘汰
        self.pool[token_hash] = (k_cache, v_cache)
    # 读取共享缓存
    def get(self, token_hash):
        return self.pool.get(token_hash, (None, None))
    # 检查缓存是否存在
    def exists(self, token_hash):
        return token_hash in self.pool
# 测试全局共享缓存
if __name__ == "__main__":
    print("=" * 65)
    print("       全局共享KV缓存池 - 效果演示")
    print("=" * 65)
    
    # 初始化缓存池
    global_pool = GlobalKVCachePool(max_pool_size=100)
    print(f"\n【初始化】缓存池容量: {global_pool.max_pool_size} 项\n")
    
    # 模拟系统Prompt（所有请求共享）
    sys_prompt_hash = hash("你是一个智能助手")
    sys_k = torch.randn(1, 128, 8, 128)
    sys_v = torch.randn(1, 128, 8, 128)
    cache_size_mb = sys_k.numel() * sys_k.element_size() * 2 / 1024 / 1024
    
    # 场景1：第一个请求写入系统Prompt缓存
    print("-" * 65)
    print("【场景1】请求1 - 首次系统Prompt | 提示词: '你是一个智能助手' ")
    print("-" * 65)
    print(f"系统Prompt缓存大小: {cache_size_mb:.2f} MB")
    global_pool.put(sys_prompt_hash, sys_k, sys_v)
    print(f"缓存池状态: {len(global_pool.pool)} 项")
    
    # 场景2：多个请求复用同一缓存
    print("\n" + "-" * 65)
    print("【场景2】请求2-5 - 复用系统Prompt缓存")
    print("-" * 65)
    prompts = [
        "你是一个智能助手",  # 与请求1相同
        "你是一个智能助手",  # 与请求1相同
        "你是一个智能助手",  # 与请求1相同
        "你是一个智能助手",  # 与请求1相同
    ]
    for i, prompt in enumerate(prompts, start=2):
        prompt_hash = hash(prompt)
        k, v = global_pool.get(prompt_hash)
        status = "✓ 命中" if k is not None else "✗ 未命中"
        saved = cache_size_mb if k is not None else 0
        print(f"请求{i}: {status} | 提示词: '{prompt}' | 节省: {saved:.2f} MB")
    
    # 场景3：不同用户不同Prompt
    print("\n" + "-" * 65)
    print("【场景3】用户A/B/C - 不同系统Prompt")
    print("-" * 65)
    users = [
        ("用户A", "你是编程助手"),
        ("用户B", "你是翻译专家"),
        ("用户C", "你是编程助手")  # 与用户A相同
    ]
    for user, prompt in users:
        prompt_hash = hash(prompt)
        k, v = global_pool.get(prompt_hash)
        if k is None:
            new_k = torch.randn(1, 128, 8, 128)
            new_v = torch.randn(1, 128, 8, 128)
            global_pool.put(prompt_hash, new_k, new_v)
            print(f"{user}: 新缓存写入 | Prompt: '{prompt[:8]}...'")
        else:
            print(f"{user}: ✓ 复用缓存 | Prompt: '{prompt[:8]}...' (与用户A相同)")
    
    # 场景4：统计对比
    print("\n" + "=" * 65)
    print("【效果统计】")
    print("=" * 65)
    total_saved = cache_size_mb * 4  # 4次复用
    print(f"缓存池项数: {len(global_pool.pool)}")
    print(f"系统Prompt复用次数: 4 次")
    print(f"累计节省显存: {total_saved:.2f} MB")
    print(f"传统方式需显存: {cache_size_mb * 5:.2f} MB (5份独立)")
    print(f"共享方式需显存: {cache_size_mb:.2f} MB (1份共享)")
    print(f"节省比例: {(1 - 1/5) * 100:.0f}%")
    
    print("\n" + "=" * 65)
    print("核心价值: N个请求共享1份缓存，显存节省 (N-1)/N")
    print("=" * 65)