import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class CausalSelfAttentionWithKVCache(nn.Module):
    """KVキャッシュ付きCausal Self-Attention"""

    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Q, K, V の射影
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, kv_cache=None, use_cache=False):
        """
        Args:
            x: (batch_size, seq_len, d_model) - 新しいトークンの埋め込み
            kv_cache: タプル (cached_k, cached_v) または None
            use_cache: キャッシュを返すかどうか

        Returns:
            output: (batch_size, seq_len, d_model)
            new_kv_cache: (new_k, new_v) if use_cache else None
        """
        batch_size, seq_len, _ = x.size()

        # Q, K, V を計算
        Q = self.W_q(x)  # (B, seq_len, d_model)
        K = self.W_k(x)
        V = self.W_v(x)

        # ヘッドに分割: (B, seq_len, d_model) -> (B, n_heads, seq_len, d_k)
        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # KVキャッシュを結合
        if kv_cache is not None:
            cached_k, cached_v = kv_cache
            # キャッシュされたK, Vと新しいK, Vを結合
            K = torch.cat([cached_k, K], dim=2)
            V = torch.cat([cached_v, V], dim=2)

        # Attention計算
        # Q: (B, n_heads, seq_len, d_k)
        # K: (B, n_heads, total_len, d_k)
        total_len = K.size(2)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        # scores: (B, n_heads, seq_len, total_len)

        # 因果マスク（新しいトークンが過去のトークンのみを参照）
        if seq_len > 1:
            # 学習時や複数トークン処理時
            mask = torch.triu(
                torch.ones(seq_len, total_len, device=x.device),
                diagonal=total_len - seq_len + 1
            ).bool()
            scores = scores.masked_fill(mask, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Valueの加重和
        output = torch.matmul(attn_weights, V)  # (B, n_heads, seq_len, d_k)

        # ヘッドを結合
        output = output.transpose(1, 2).contiguous().view(
            batch_size, seq_len, self.d_model
        )
        output = self.W_o(output)

        if use_cache:
            return output, (K, V)
        return output, None

# パラメータ設定
d_model = 256
n_heads = 8
vocab_size = 1000

# モデルの初期化
attn = CausalSelfAttentionWithKVCache(d_model, n_heads)
embedding = nn.Embedding(vocab_size, d_model)

# 初期トークン（プロンプト）
torch.manual_seed(42)
prompt_tokens = torch.randint(0, vocab_size, (1, 5))  # (batch=1, seq=5)
print(f"プロンプト: {prompt_tokens.shape}")

# プロンプトの処理（キャッシュを構築）
x = embedding(prompt_tokens)
output, kv_cache = attn(x, kv_cache=None, use_cache=True)
print(f"初期KVキャッシュのKの形状: {kv_cache[0].shape}")
# (1, 8, 5, 32) = (batch, heads, seq_len, d_k)

# 自己回帰生成のシミュレーション
for step in range(10):
    # 新しいトークンを生成（ここではランダム）
    new_token = torch.randint(0, vocab_size, (1, 1))

    # 新しいトークンのみを処理
    new_x = embedding(new_token)
    output, kv_cache = attn(new_x, kv_cache=kv_cache, use_cache=True)

    print(f"ステップ {step+1}: KVキャッシュの長さ = {kv_cache[0].shape[2]}")

# 最終的なKVキャッシュの長さ
print(f"最終KVキャッシュのKの形状: {kv_cache[0].shape}")
# (1, 8, 15, 32) = 初期5 + 生成10 = 15トークン

プロンプト: torch.Size([1, 5])
初期KVキャッシュのKの形状: torch.Size([1, 8, 5, 32])
ステップ 1: KVキャッシュの長さ = 6
ステップ 2: KVキャッシュの長さ = 7
...
ステップ 10: KVキャッシュの長さ = 15
最終KVキャッシュのKの形状: torch.Size([1, 8, 15, 32])

import matplotlib.pyplot as plt
import numpy as np

# KVキャッシュのメモリ使用量をシミュレート
def calculate_kv_cache_memory(
    seq_len, d_model, n_layers, n_kv_heads, batch_size=1, dtype_bytes=2
):
    """KVキャッシュのメモリ使用量（MB）を計算"""
    d_k = d_model // n_kv_heads if n_kv_heads else d_model // 8
    memory = 2 * batch_size * seq_len * n_kv_heads * d_k * n_layers * dtype_bytes
    return memory / (1024 ** 2)  # MB

# パラメータ（Llama 2 7B相当）
d_model = 4096
n_layers = 32
n_kv_heads = 8  # GQA

seq_lengths = np.arange(0, 8192, 100)

# 異なるバッチサイズでのメモリ使用量
batch_sizes = [1, 4, 16, 64]

plt.figure(figsize=(10, 6))
for batch_size in batch_sizes:
    memory = [
        calculate_kv_cache_memory(seq, d_model, n_layers, n_kv_heads, batch_size)
        for seq in seq_lengths
    ]
    plt.plot(seq_lengths, memory, label=f'Batch size = {batch_size}', linewidth=2)

plt.xlabel('Sequence Length', fontsize=12)
plt.ylabel('KV Cache Memory (MB)', fontsize=12)
plt.title('KV Cache Memory Usage vs Sequence Length', fontsize=14)
plt.legend(fontsize=10)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()